Kris
fb8bd90329
1. 调整多个服务文件的导入顺序以符合规范 2. 新增Postgres插件状态存储实现，支持CRUD和过期清理 3. 在维护任务中添加插件过期状态自动清理逻辑 4. 新增统计收集器，支持请求/错误计数、响应时间统计和指标上报
102 lines
3.2 KiB
Python
102 lines
3.2 KiB
Python
from __future__ import annotations
|
|
|
|
import asyncio
|
|
from collections.abc import Awaitable, Callable
|
|
from dataclasses import dataclass, field
|
|
from typing import TYPE_CHECKING, Any
|
|
|
|
if TYPE_CHECKING:
|
|
from yuxi.channels.models import ChannelAccountSnapshot
|
|
|
|
|
|
@dataclass
|
|
class ChatAbortEntry:
|
|
task: Any = None
|
|
event: asyncio.Event = field(default_factory=asyncio.Event)
|
|
|
|
def abort(self) -> None:
|
|
self.event.set()
|
|
if self.task is not None and not self.task.done():
|
|
self.task.cancel()
|
|
|
|
def is_aborted(self) -> bool:
|
|
return self.event.is_set()
|
|
|
|
|
|
@dataclass
|
|
class ChatRunBuffer:
|
|
run_id: str = ""
|
|
_chunks: list[str] = field(default_factory=list)
|
|
_finished: bool = False
|
|
|
|
def append_chunk(self, chunk: str) -> None:
|
|
self._chunks.append(chunk)
|
|
|
|
def get_full_text(self) -> str:
|
|
return "".join(self._chunks)
|
|
|
|
def get_chunks(self) -> list[str]:
|
|
return list(self._chunks)
|
|
|
|
def is_finished(self) -> bool:
|
|
return self._finished
|
|
|
|
def mark_finished(self) -> None:
|
|
self._finished = True
|
|
|
|
def chunk_count(self) -> int:
|
|
return len(self._chunks)
|
|
|
|
|
|
@dataclass
|
|
class GatewayRequestContext:
|
|
runtime_config: dict = field(default_factory=dict)
|
|
runtime_snapshot: ChannelAccountSnapshot | None = None
|
|
|
|
start_channel: Callable[[str, dict], Awaitable[None]] | None = None
|
|
stop_channel: Callable[[str], Awaitable[None]] | None = None
|
|
mark_channel_logged_out: Callable[[str, str], Awaitable[None]] | None = None
|
|
|
|
broadcast_fn: Callable[[str, Any], Awaitable[None]] | None = None
|
|
node_send_to_session_fn: Callable[[str, str, Any], Awaitable[None]] | None = None
|
|
|
|
get_snapshot: Callable[[str], Awaitable[ChannelAccountSnapshot]] | None = None
|
|
|
|
dedupe: Any = None
|
|
|
|
chat_abort_controllers: dict[str, ChatAbortEntry] = field(default_factory=dict)
|
|
chat_run_buffers: dict[str, ChatRunBuffer] = field(default_factory=dict)
|
|
|
|
def get_runtime_config(self, key: str, default=None):
|
|
return self.runtime_config.get(key, default)
|
|
|
|
def set_runtime_config(self, key: str, value) -> None:
|
|
self.runtime_config[key] = value
|
|
|
|
def add_chat_run(self, run_id: str, entry: ChatAbortEntry, buffer: ChatRunBuffer | None = None) -> None:
|
|
self.chat_abort_controllers[run_id] = entry
|
|
if buffer is not None:
|
|
self.chat_run_buffers[run_id] = buffer
|
|
|
|
def remove_chat_run(self, run_id: str) -> None:
|
|
self.chat_abort_controllers.pop(run_id, None)
|
|
self.chat_run_buffers.pop(run_id, None)
|
|
|
|
def abort_chat_run(self, run_id: str) -> bool:
|
|
entry = self.chat_abort_controllers.get(run_id)
|
|
if entry is None:
|
|
return False
|
|
entry.abort()
|
|
return True
|
|
|
|
def get_chat_buffer(self, run_id: str) -> ChatRunBuffer | None:
|
|
return self.chat_run_buffers.get(run_id)
|
|
|
|
async def broadcast(self, event: str, payload: Any = None) -> None:
|
|
if self.broadcast_fn is not None:
|
|
await self.broadcast_fn(event, payload)
|
|
|
|
async def node_send_to_session(self, session_id: str, event: str, payload: Any = None) -> None:
|
|
if self.node_send_to_session_fn is not None:
|
|
await self.node_send_to_session_fn(session_id, event, payload)
|