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)