Kris
1f78c44b03
这是一个批量整理提交，包含以下主要改动： 1. 删除多处冗余的空行和未使用的导入 2. 修复文件末尾缺少换行符的问题 3. 调整部分模块的导入顺序与代码排版 4. 修复部分配置默认值与策略逻辑 5. 新增多个功能模块与辅助工具 6. 完善异常处理与日志记录 7. 修复速率限制、消息缓存、权限校验等逻辑bug 8. 废弃部分旧有API与配置项并添加警告提示
294 lines
12 KiB
Python
294 lines
12 KiB
Python
from __future__ import annotations
|
|
|
|
import asyncio
|
|
|
|
from yuxi.channels.services.context import ChatAbortEntry, ChatRunBuffer
|
|
from yuxi.channels.models import ChannelMessage, ChannelResponse
|
|
from yuxi.channels.policy.context_policy import ContextCommand, ContextPolicy
|
|
from yuxi.channels.policy.dedup_policy import DedupPolicy
|
|
from yuxi.channels.policy.group_chat_policy import GroupChatPolicy
|
|
from yuxi.channels.policy.schedule_policy import SchedulePolicy
|
|
from yuxi.channels.policy.welcome_policy import WelcomePolicy
|
|
from yuxi.channels.policy.media_policy import MediaPolicy
|
|
from yuxi.channels.policy.voice_policy import VoicePolicy
|
|
from yuxi.channels.session_mapper import SessionMapper, VIRTUAL_DEPARTMENT_ID
|
|
from yuxi.utils.logging_config import logger
|
|
|
|
|
|
class _ChannelUser:
|
|
__slots__ = ("id", "department_id", "username", "user_id")
|
|
|
|
def __init__(self, uid: str, dept_id: int):
|
|
self.id = uid
|
|
self.department_id = dept_id
|
|
self.username = f"channel_user_{uid}"
|
|
self.user_id = uid
|
|
|
|
|
|
class MessageRouter:
|
|
def __init__(
|
|
self,
|
|
channel_manager=None,
|
|
dedup_policy: DedupPolicy | None = None,
|
|
context_policy: ContextPolicy | None = None,
|
|
group_chat_policy: GroupChatPolicy | None = None,
|
|
welcome_policy: WelcomePolicy | None = None,
|
|
schedule_policy: SchedulePolicy | None = None,
|
|
media_policy: MediaPolicy | None = None,
|
|
voice_policy: VoicePolicy | None = None,
|
|
):
|
|
self._channel_manager = channel_manager
|
|
self.dedup_policy = dedup_policy or DedupPolicy()
|
|
self.context_policy = context_policy or ContextPolicy()
|
|
self.group_chat_policy = group_chat_policy or GroupChatPolicy()
|
|
self.welcome_policy = welcome_policy or WelcomePolicy()
|
|
self.schedule_policy = schedule_policy or SchedulePolicy()
|
|
self.media_policy = media_policy or MediaPolicy()
|
|
self.voice_policy = voice_policy or VoicePolicy()
|
|
self.chat_abort_controllers: dict[str, ChatAbortEntry] = {}
|
|
self.chat_run_buffers: dict[str, ChatRunBuffer] = {}
|
|
|
|
async def route_inbound(self, message: ChannelMessage) -> None:
|
|
identity = message.identity
|
|
|
|
if await self.dedup_policy.check_and_remember(message):
|
|
logger.debug(f"Dropping duplicate message from {identity.channel_id}")
|
|
return
|
|
|
|
context_result = self.context_policy.parse(message)
|
|
if context_result.handled:
|
|
await self._handle_context_command(message, context_result.command, context_result.args)
|
|
return
|
|
|
|
if not self.schedule_policy.is_working_hours():
|
|
reply = self.schedule_policy.get_off_hours_reply()
|
|
if reply:
|
|
response = ChannelResponse(identity=identity, content=reply)
|
|
await self._send_response(identity.channel_id, response)
|
|
return
|
|
|
|
is_at_bot = bool(message.mentions and message.mentions.is_bot_mentioned)
|
|
if not self.group_chat_policy.should_respond(message, is_at_bot):
|
|
return
|
|
|
|
from yuxi.storage.postgres.manager import pg_manager
|
|
from yuxi.repositories.channel_message_record_repository import ChannelMessageRecordRepository
|
|
|
|
async with pg_manager.get_async_session_context() as db:
|
|
session_mapper = SessionMapper(db)
|
|
internal_user_id = await session_mapper.resolve_user(message)
|
|
thread_id = await session_mapper.resolve_thread(message, internal_user_id)
|
|
|
|
msg_record_repo = ChannelMessageRecordRepository(db)
|
|
agent_config_id = self._resolve_agent_config_id(message)
|
|
record = await msg_record_repo.create_record(message, agent_config_id=agent_config_id)
|
|
|
|
if self.welcome_policy.mark_welcomed(internal_user_id):
|
|
welcome_response = ChannelResponse(
|
|
identity=identity,
|
|
content=self.welcome_policy.get_welcome_message(),
|
|
)
|
|
await self._send_response(identity.channel_id, welcome_response)
|
|
|
|
run_id = f"{message.identity.channel_id}:{message.identity.channel_message_id}"
|
|
|
|
try:
|
|
task = asyncio.ensure_future(
|
|
self._invoke_agent(
|
|
db=db,
|
|
query=message.content,
|
|
thread_id=thread_id,
|
|
internal_user_id=internal_user_id,
|
|
agent_config_id=agent_config_id,
|
|
message=message,
|
|
)
|
|
)
|
|
self.chat_abort_controllers[run_id] = ChatAbortEntry(task=task)
|
|
|
|
response_content = await task
|
|
|
|
self.chat_abort_controllers.pop(run_id, None)
|
|
|
|
response = ChannelResponse(identity=identity, content=response_content)
|
|
await self._send_response(identity.channel_id, response)
|
|
await msg_record_repo.mark_success(record.id, response)
|
|
|
|
except asyncio.CancelledError:
|
|
logger.info(f"Chat aborted for run {run_id}")
|
|
self.chat_abort_controllers.pop(run_id, None)
|
|
error_response = ChannelResponse(
|
|
identity=identity,
|
|
content="对话已被中断。",
|
|
)
|
|
await self._send_response(identity.channel_id, error_response)
|
|
|
|
except Exception as e:
|
|
logger.error(f"Agent invocation failed: {e}")
|
|
error_response = ChannelResponse(
|
|
identity=identity,
|
|
content="\u62b1\u6b49\uff0c\u5904\u7406\u4f60\u7684\u6d88\u606f\u65f6\u51fa\u9519\u4e86\uff0c\u8bf7\u7a0d\u540e\u518d\u8bd5\u3002",
|
|
)
|
|
await self._send_response(identity.channel_id, error_response)
|
|
await msg_record_repo.mark_error(record.id, str(e))
|
|
|
|
async def route_outbound(self, agent_result, channel_id: str, identity) -> None:
|
|
response = ChannelResponse(
|
|
identity=identity,
|
|
content=agent_result.response_text,
|
|
attachments=getattr(agent_result, "attachments", []),
|
|
)
|
|
await self._send_response(channel_id, response)
|
|
|
|
def abort_chat(self, run_id: str) -> bool:
|
|
entry = self.chat_abort_controllers.get(run_id)
|
|
if entry is None:
|
|
return False
|
|
entry.abort()
|
|
return True
|
|
|
|
async def _handle_context_command(self, message: ChannelMessage, command: ContextCommand, args: str) -> None:
|
|
identity = message.identity
|
|
|
|
if command == ContextCommand.RESET:
|
|
from yuxi.storage.postgres.manager import pg_manager
|
|
|
|
async with pg_manager.get_async_session_context() as db:
|
|
session_mapper = SessionMapper(db)
|
|
internal_user_id = await session_mapper.resolve_user(message)
|
|
await session_mapper.reset_thread(message, internal_user_id)
|
|
|
|
response = ChannelResponse(
|
|
identity=identity,
|
|
content="\u5bf9\u8bdd\u4e0a\u4e0b\u6587\u5df2\u91cd\u7f6e\uff0c\u65b0\u7684\u4f1a\u8bdd\u5df2\u521b\u5efa",
|
|
)
|
|
await self._send_response(identity.channel_id, response)
|
|
|
|
elif command == ContextCommand.HISTORY:
|
|
response = ChannelResponse(
|
|
identity=identity, content="\u5386\u53f2\u8bb0\u5f55\u529f\u80fd\u6682\u672a\u5b9e\u73b0"
|
|
)
|
|
await self._send_response(identity.channel_id, response)
|
|
|
|
elif command == ContextCommand.CONTEXT:
|
|
response = ChannelResponse(
|
|
identity=identity, content="\u4e0a\u4e0b\u6587\u4fe1\u606f\u529f\u80fd\u6682\u672a\u5b9e\u73b0"
|
|
)
|
|
await self._send_response(identity.channel_id, response)
|
|
|
|
elif command == ContextCommand.SUMMARY:
|
|
response = ChannelResponse(
|
|
identity=identity, content="\u5bf9\u8bdd\u6458\u8981\u529f\u80fd\u6682\u672a\u5b9e\u73b0"
|
|
)
|
|
await self._send_response(identity.channel_id, response)
|
|
|
|
async def _invoke_agent(
|
|
self,
|
|
db,
|
|
query: str,
|
|
thread_id: str,
|
|
internal_user_id: str,
|
|
agent_config_id: int,
|
|
message: ChannelMessage,
|
|
) -> str:
|
|
from yuxi.services.chat_service import stream_agent_chat
|
|
|
|
channel_config = self._get_channel_config(message.identity.channel_id)
|
|
department_id = channel_config.get("department_id", VIRTUAL_DEPARTMENT_ID)
|
|
|
|
channel_user = _ChannelUser(internal_user_id, department_id)
|
|
|
|
run_id = f"{message.identity.channel_id}:{message.identity.channel_message_id}"
|
|
buffer = ChatRunBuffer(run_id=run_id)
|
|
self.chat_run_buffers[run_id] = buffer
|
|
|
|
adapter = None
|
|
if self._channel_manager:
|
|
adapter = self._channel_manager._adapters.get(message.identity.channel_id)
|
|
supports_streaming = getattr(adapter, "supports_streaming", False)
|
|
|
|
try:
|
|
async for chunk in stream_agent_chat(
|
|
query=query,
|
|
agent_config_id=agent_config_id,
|
|
thread_id=thread_id,
|
|
meta={
|
|
"source": "channel",
|
|
"channel_id": message.identity.channel_id,
|
|
"channel_type": message.identity.channel_type.value,
|
|
},
|
|
image_content=None,
|
|
current_user=channel_user,
|
|
db=db,
|
|
):
|
|
import json
|
|
|
|
try:
|
|
data = json.loads(chunk.decode("utf-8").strip())
|
|
if data.get("status") == "loading" and data.get("response"):
|
|
buffer.append_chunk(data["response"])
|
|
if supports_streaming and adapter:
|
|
await adapter.send_stream_chunk(
|
|
chat_id=message.identity.channel_chat_id or "",
|
|
message_id=message.identity.channel_message_id or "",
|
|
chunk_text=data["response"],
|
|
finished=False,
|
|
)
|
|
except (json.JSONDecodeError, UnicodeDecodeError):
|
|
logger.debug(f"Non-JSON stream chunk: {chunk[:100]!r}")
|
|
except Exception:
|
|
logger.debug(f"Unexpected stream parse error for chunk: {chunk[:100]!r}")
|
|
|
|
buffer.mark_finished()
|
|
|
|
if supports_streaming and adapter:
|
|
await adapter.send_stream_chunk(
|
|
chat_id=message.identity.channel_chat_id or "",
|
|
message_id=message.identity.channel_message_id or "",
|
|
chunk_text="",
|
|
finished=True,
|
|
)
|
|
|
|
return buffer.get_full_text()
|
|
finally:
|
|
self.chat_run_buffers.pop(run_id, None)
|
|
|
|
def _resolve_agent_config_id(self, message: ChannelMessage) -> int:
|
|
channel_id = message.identity.channel_id
|
|
content = message.content.strip()
|
|
|
|
agent_config_id = message.metadata.get("agent_config_id")
|
|
if agent_config_id is not None:
|
|
return int(agent_config_id)
|
|
|
|
channel_config = self._get_channel_config(channel_id)
|
|
|
|
if content.startswith("/"):
|
|
cmd = content.split()[0].lower()
|
|
cmd_routing = channel_config.get("command_routing", {})
|
|
if cmd in cmd_routing:
|
|
return int(cmd_routing[cmd])
|
|
|
|
channel_default = channel_config.get("agent_config_id")
|
|
if channel_default is not None:
|
|
return int(channel_default)
|
|
|
|
global_default = self._get_global_default_agent_id()
|
|
if global_default is not None:
|
|
return global_default
|
|
|
|
return 1
|
|
|
|
def _get_channel_config(self, channel_id: str) -> dict:
|
|
if self._channel_manager and hasattr(self._channel_manager, "_channels_config"):
|
|
return self._channel_manager._channels_config.get(channel_id, {})
|
|
return {}
|
|
|
|
def _get_global_default_agent_id(self) -> int | None:
|
|
from yuxi import config as conf
|
|
|
|
return getattr(conf, "default_agent_id", None)
|
|
|
|
async def _send_response(self, channel_id: str, response: ChannelResponse) -> None:
|
|
if self._channel_manager and hasattr(self._channel_manager, "send_outbound"):
|
|
await self._channel_manager.send_outbound(channel_id, response)
|