from asyncio import Task from typing import Dict, List, Sequence, Tuple from rich.traceback import install from sqlmodel import select import asyncio from src.common.logger import get_logger from src.common.database.database import get_db_session from src.common.database.database_model import Messages from src.common.data_models.mai_message_data_model import MaiMessage, UserInfo from src.common.data_models.message_component_data_model import ( TextComponent, ImageComponent, EmojiComponent, AtComponent, ReplyComponent, VoiceComponent, ForwardNodeComponent, StandardMessageComponents, ) install(extra_lines=3) logger = get_logger("chat_message") class MsgIDMapping: """回复消息内容缓存。""" def __init__(self) -> None: """初始化消息 ID 到内容的映射缓存。""" self.mapping: Dict[str, Tuple[str | Task[str], UserInfo]] = {} class SessionMessage(MaiMessage): async def process( self, *, enable_heavy_media_analysis: bool = True, enable_voice_transcription: bool = True, ) -> None: """处理消息内容并转化为纯文本。 Args: enable_heavy_media_analysis: 是否同步执行图片与表情包描述生成。 enable_voice_transcription: 是否同步执行语音转写。 """ id_content_map = MsgIDMapping() tasks = [ self.process_single_component( component, id_content_map, enable_heavy_media_analysis=enable_heavy_media_analysis, enable_voice_transcription=enable_voice_transcription, ) for component in self.raw_message.components ] results = await asyncio.gather(*tasks, return_exceptions=True) processed_texts: List[str] = [] for result in results: if isinstance(result, BaseException): logger.error(f"处理消息组件时发生错误: {result}") else: processed_texts.append(result) self.processed_plain_text = " ".join(processed_texts) async def process_single_component( self, component: StandardMessageComponents, id_content_map: MsgIDMapping, recursion_depth: int = 0, *, enable_heavy_media_analysis: bool = True, enable_voice_transcription: bool = True, ) -> str: """按类型处理单个消息组件。 Args: component: 待处理的消息组件。 id_content_map: 回复消息解析缓存。 recursion_depth: 当前递归深度。 enable_heavy_media_analysis: 是否同步执行图片与表情包描述生成。 enable_voice_transcription: 是否同步执行语音转写。 Returns: str: 组件对应的文本表示。 """ if isinstance(component, TextComponent): return component.text elif isinstance(component, ImageComponent): return await self.process_image_component( component, enable_heavy_media_analysis=enable_heavy_media_analysis, ) elif isinstance(component, EmojiComponent): return await self.process_emoji_component( component, enable_heavy_media_analysis=enable_heavy_media_analysis, ) elif isinstance(component, AtComponent): return await self.process_at_component(component) elif isinstance(component, VoiceComponent): return await self.process_voice_component( component, enable_voice_transcription=enable_voice_transcription, ) elif isinstance(component, ReplyComponent): return await self.process_reply_component(component, id_content_map) elif isinstance(component, ForwardNodeComponent): return await self.process_forward_component( component, id_content_map, recursion_depth=recursion_depth + 1, enable_heavy_media_analysis=enable_heavy_media_analysis, enable_voice_transcription=enable_voice_transcription, ) else: raise NotImplementedError(f"暂时不支持的消息组件类型: {type(component)}") async def process_image_component( self, component: ImageComponent, *, enable_heavy_media_analysis: bool = True, ) -> str: """处理图片组件。 Args: component: 图片组件。 enable_heavy_media_analysis: 是否同步执行图片描述生成。 Returns: str: 图片组件对应的文本表示。 """ if component.content: # 先检查是否处理过 return component.content from src.chat.image_system.image_manager import image_manager # 获取描述 try: desc = await image_manager.get_image_description( image_bytes=component.binary_data, wait_for_build=enable_heavy_media_analysis, ) except Exception: desc = None # 失败置空 content = f"[图片：{desc}]" if desc else "[图片]" component.content = content component.binary_data = b"" # 处理完就丢掉二进制数据，节省内存 return content async def process_emoji_component( self, component: EmojiComponent, *, enable_heavy_media_analysis: bool = True, ) -> str: """处理表情包组件。 Args: component: 表情包组件。 enable_heavy_media_analysis: 是否同步执行表情包描述生成。 Returns: str: 表情包组件对应的文本表示。 """ if component.content: # 先检查是否处理过 return component.content from src.chat.emoji_system.emoji_manager import emoji_manager # 获取表情包描述 try: tuple_content = await emoji_manager.get_emoji_description( emoji_bytes=component.binary_data, wait_for_build=enable_heavy_media_analysis, ) except Exception: tuple_content = None # 失败置空 if tuple_content: desc, _ = tuple_content content = f"[表情包: {desc}]" else: content = "[表情包]" component.content = content component.binary_data = b"" # 处理完就丢掉二进制数据，节省内存 return content async def process_at_component(self, component: AtComponent) -> str: # 如果已经有昵称或备注了，直接使用 if component.target_user_cardname: return f"@{component.target_user_cardname}" elif component.target_user_nickname: return f"@{component.target_user_nickname}" from src.common.utils.utils_person import PersonUtils # 查询用户信息 if person_info := PersonUtils.get_person_info_by_user_id_and_platform(component.target_user_id, self.platform): component.target_user_nickname = component.target_user_nickname or person_info.user_nickname if self.message_info.group_info and person_info.group_cardname_list: for group_card in person_info.group_cardname_list: if group_card.group_id == self.message_info.group_info.group_id: component.target_user_cardname = group_card.group_cardname break if component.target_user_cardname: # 优先使用群备注 return f"@{component.target_user_cardname}" elif component.target_user_nickname: # 其次使用昵称 return f"@{component.target_user_nickname}" else: # 最后使用用户ID return f"@{component.target_user_id}" async def process_voice_component( self, component: VoiceComponent, *, enable_voice_transcription: bool = True, ) -> str: """处理语音组件。 Args: component: 语音组件。 enable_voice_transcription: 是否同步执行语音转写。 Returns: str: 语音组件对应的文本表示。 """ if component.content: # 先检查是否处理过 return component.content if not enable_voice_transcription: component.content = "[语音消息]" return component.content from src.common.utils.utils_voice import get_voice_text text = await get_voice_text(component.binary_data) content = "[语音消息，转录失败]" if text is None else f"[语音: {text}]" component.content = content return content async def process_reply_component( self, component: ReplyComponent, id_content_map: MsgIDMapping, ) -> str: if component.target_message_content: return component.target_message_content if result_item := id_content_map.mapping.get(component.target_message_id): # ID映射缓存优先 content, sender_info = result_item if isinstance(content, Task): # 如果是Task，说明是转发组件传入的占位结果，需要等待其完成 content = await content # 获取最终结果 id_content_map.mapping[component.target_message_id] = (content, sender_info) # 更新为实际内容 component.target_message_content = content tgt_msg_s_name = sender_info.user_cardname or sender_info.user_nickname or sender_info.user_id component.target_message_sender_cardname = sender_info.user_cardname component.target_message_sender_nickname = sender_info.user_nickname component.target_message_sender_id = sender_info.user_id return f"[回复了{tgt_msg_s_name}的消息: {content}]" else: # 尝试从数据库根据消息id查找消息内容 try: with get_db_session() as session: statement = select(Messages).filter_by(message_id=component.target_message_id).limit(1) if db_msg := session.exec(statement).first(): component.target_message_content = db_msg.processed_plain_text component.target_message_sender_cardname = db_msg.user_cardname component.target_message_sender_nickname = db_msg.user_nickname component.target_message_sender_id = db_msg.user_id tgt_msg_s_name = db_msg.user_cardname or db_msg.user_nickname or db_msg.user_id return f"[回复了{tgt_msg_s_name}的消息: {db_msg.processed_plain_text}]" except Exception as e: logger.error(f"查询回复消息时发生错误: {e}") return "[回复了一条消息，但原消息已无法访问]" async def process_forward_component( self, component: ForwardNodeComponent, id_content_map: MsgIDMapping, recursion_depth: int = 0, *, enable_heavy_media_analysis: bool = True, enable_voice_transcription: bool = True, ) -> str: """处理合并转发组件。 Args: component: 合并转发组件。 id_content_map: 回复消息解析缓存。 recursion_depth: 当前递归深度。 enable_heavy_media_analysis: 是否同步执行图片与表情包描述生成。 enable_voice_transcription: 是否同步执行语音转写。 Returns: str: 合并转发组件对应的文本表示。 """ task_list: List[Task] = [] node_user_info_list: List[UserInfo] = [] for node in component.forward_components: task = asyncio.create_task( self._process_multiple_components( node.content, id_content_map, recursion_depth + 1, enable_heavy_media_analysis=enable_heavy_media_analysis, enable_voice_transcription=enable_voice_transcription, ) ) node_user_info = UserInfo(node.user_id or "未知用户", node.user_nickname, node.user_cardname) # 传入ID缓存映射，方便Reply组件获取并等待处理结果 id_content_map.mapping[node.message_id] = (task, node_user_info) task_list.append(task) node_user_info_list.append(node_user_info) results = await asyncio.gather(*task_list, return_exceptions=True) # 并行处理节点内容 forward_texts = [] for idx, result in enumerate(results): if isinstance(result, BaseException): logger.error(f"处理转发消息组件时发生错误: {result}") else: usr_info = node_user_info_list[idx] msg_sender_name = usr_info.user_cardname or usr_info.user_nickname or usr_info.user_id or "未知用户" forward_texts.append(f"{'-' * recursion_depth * 2} 【{msg_sender_name}】: {result}") return "【合并转发消息: \n" + "\n".join(forward_texts) + "\n】" async def _process_multiple_components( self, components: Sequence[StandardMessageComponents], id_content_map: MsgIDMapping, recursion_depth: int = 0, *, enable_heavy_media_analysis: bool = True, enable_voice_transcription: bool = True, ) -> str: """并行处理多个消息组件。 Args: components: 待处理的组件序列。 id_content_map: 回复消息解析缓存。 recursion_depth: 当前递归深度。 enable_heavy_media_analysis: 是否同步执行图片与表情包描述生成。 enable_voice_transcription: 是否同步执行语音转写。 Returns: str: 多个组件拼接后的文本表示。 """ tasks = [ self.process_single_component( component, id_content_map, recursion_depth, enable_heavy_media_analysis=enable_heavy_media_analysis, enable_voice_transcription=enable_voice_transcription, ) for component in components ] results = await asyncio.gather(*tasks, return_exceptions=True) # 并行处理多个组件 processed_texts: List[str] = [] for result in results: if isinstance(result, BaseException): logger.error(f"处理消息组件时发生错误: {result}") else: processed_texts.append(result) return " ".join(processed_texts)