# -*- coding: utf-8 -*- """ SSE 核心工具函数 —— 流式事件生成器、文本提取、响应清理等通用辅助 """ import json import re import time import traceback from typing import AsyncGenerator from tools.tool_names_cn import TOOL_NAME_CN as _TOOL_CN from db.chat_store import save_message, get_messages # ── 工具名 → 中文描述映射 ── def _cn_tool_desc(tool_name: str) -> str: """将工具函数名映射为中文描述短语（用于 SSE message 字段）。 映射表统一维护在 tools/tool_names_cn.py 中，新增工具只需改那一处。 """ return _TOOL_CN.get(tool_name, tool_name) # ── 通用 SSE 流式事件生成器 ── async def sse_event_generator( agent, user_message: str, thread_id: str, session_id: str, config: dict, agent_cn_name: str = "智能助手", save_to_db: bool = True, ) -> AsyncGenerator[str, None]: """ 通用 SSE 流式事件生成器。 使用 stream_mode=["updates", "messages"] 获取完整执行图景， 前端可根据 event type 区分：thinking / executing / tool_call / token / done / error。 Args: save_to_db: 是否将消息保存到数据库。点选解读等非对话场景应设为 False。 """ # 保存用户消息（仅对话场景写库） if save_to_db: save_message(session_id, "user", user_message) # 用于收集完整回复 full_response = "" start_time = time.time() try: # ── 发送 agent_start 事件，前端据此创建 Agent 卡片 ── yield f"data: {json.dumps({'type': 'agent_start', 'agent': agent_cn_name, 'cn_agent': agent_cn_name, 'message': f'「{agent_cn_name}」开始处理...'}, ensure_ascii=False)}\n\n" # 构建消息列表（包含历史消息 + 当前消息） if save_to_db: messages = [{"role": "user", "content": msg["content"]} for msg in get_messages(session_id, limit=50)] else: messages = [] messages.append({"role": "user", "content": user_message}) # 使用 agent.astream() 异步流式调用，多模式获取完整图景 async for chunk in agent.astream( {"messages": messages}, stream_mode=["updates", "messages"], config=config, version="v2", ): # 判断事件来源：主代理 vs 子代理 is_subagent = any(s.startswith("tools:") for s in chunk.get("ns", [])) source = "subagent" if is_subagent else "main" # ── updates 模式：步骤级事件（thinking / executing）── if chunk["type"] == "updates": for node_name in chunk["data"]: if node_name == "model_request": # 代理正在思考/推理 yield f"data: {json.dumps({'type': 'thinking', 'source': source, 'node': node_name, 'cn_agent': agent_cn_name, 'message': '正在分析您的问题...'}, ensure_ascii=False)}\n\n" elif node_name == "tools": # 代理正在执行工具调用，提取工具名称 tools_data = chunk["data"].get(node_name, {}) tool_names = [] for msg in tools_data.get("messages", []): if hasattr(msg, "name"): tool_names.append(msg.name) elif isinstance(msg, dict) and msg.get("name"): tool_names.append(msg["name"]) # 生成中文描述 cn_names = [_cn_tool_desc(t) for t in tool_names] if tool_names else ["工具调用"] msg_text = f"正在{'、'.join(cn_names)}..." yield f"data: {json.dumps({'type': 'executing', 'source': source, 'cn_agent': agent_cn_name, 'tools': tool_names, 'cn_tools': cn_names, 'message': msg_text}, ensure_ascii=False)}\n\n" # ── messages 模式：token 级事件（token / tool_call / tool_result / updated_todo_list）── elif chunk["type"] == "messages": token_data = chunk["data"] if isinstance(token_data, (list, tuple)) and len(token_data) >= 1: msg_obj, _metadata = token_data[0], token_data[1] if len(token_data) > 1 else {} else: msg_obj = token_data # 检测工具调用（tool_call_chunks 在流式传输中逐步到达） if hasattr(msg_obj, "tool_call_chunks") and msg_obj.tool_call_chunks: for tc in msg_obj.tool_call_chunks: if tc.get("name"): cn = _cn_tool_desc(tc["name"]) yield f"data: {json.dumps({'type': 'tool_call', 'source': source, 'cn_agent': agent_cn_name, 'tool': tc['name'], 'cn_tool': cn, 'message': f'调用工具：{cn}'}, ensure_ascii=False)}\n\n" # 检测工具结果 is_tool_msg = hasattr(msg_obj, "type") and msg_obj.type == "tool" if is_tool_msg: tool_name = getattr(msg_obj, "name", "unknown") if tool_name == "write_todos": # write_todos 特殊处理：提取 JSON 并发送 updated_todo_list 事件 tool_content = _extract_text_content(msg_obj) or "" todo_list = _parse_todo_list(tool_content) if todo_list: yield f"data: {json.dumps({'type': 'updated_todo_list', 'cn_agent': agent_cn_name, 'source': source, 'todos': todo_list, 'message': '任务进度已更新'}, ensure_ascii=False)}\n\n" else: # 解析失败时也至少发一个 tool_result 事件，并打印诊断日志 print(f"[write_todos] 解析失败，原始内容: {tool_content[:300]}") yield f"data: {json.dumps({'type': 'tool_result', 'source': source, 'cn_agent': agent_cn_name, 'tool': tool_name, 'cn_tool': _cn_tool_desc(tool_name), 'message': '任务进度已更新'}, ensure_ascii=False)}\n\n" else: cn = _cn_tool_desc(tool_name) yield f"data: {json.dumps({'type': 'tool_result', 'source': source, 'cn_agent': agent_cn_name, 'tool': tool_name, 'cn_tool': cn, 'message': f'{cn} 完成'}, ensure_ascii=False)}\n\n" # 提取并流式输出文本内容（工具结果不当作 token 输出，仅输出 AI 生成的文本） if not is_tool_msg: content = _extract_text_content(msg_obj) if content and isinstance(content, str): full_response += content yield f"data: {json.dumps({'type': 'token', 'source': source, 'content': content}, ensure_ascii=False)}\n\n" # 保存助手回复（仅对话场景写库） if save_to_db and full_response.strip(): cleaned = _clean_response(full_response) save_message(session_id, "assistant", cleaned) print(full_response) # 发送 agent_done + done 事件 duration_ms = int((time.time() - start_time) * 1000) yield f"data: {json.dumps({'type': 'agent_done', 'agent': agent_cn_name, 'cn_agent': agent_cn_name, 'duration_ms': duration_ms, 'message': f'「{agent_cn_name}」完成（{duration_ms}ms）'}, ensure_ascii=False)}\n\n" yield f"data: {json.dumps({'type': 'done', 'thread_id': thread_id, 'session_id': session_id, 'duration_ms': duration_ms, 'message': f'回答完成（{duration_ms}ms）'}, ensure_ascii=False)}\n\n" except Exception as e: traceback.print_exc() yield f"data: {json.dumps({'type': 'error', 'message': str(e)}, ensure_ascii=False)}\n\n" # ── 文本提取 / 解析 / 清理 ── def _extract_text_content(msg_obj) -> str | None: """从 LangChain 消息对象中提取文本内容（兼容 dict 和对象两种形式）。""" if isinstance(msg_obj, dict): return msg_obj.get("content") elif hasattr(msg_obj, "content"): raw = getattr(msg_obj, "content", None) if isinstance(raw, str): return raw elif isinstance(raw, list): # 多模态内容块：合并所有 text 类型的块 parts = [] for block in raw: if isinstance(block, dict) and block.get("type") == "text": parts.append(block.get("text", "")) elif hasattr(block, "type") and getattr(block, "type", "") == "text": parts.append(getattr(block, "text", "")) return "".join(parts) if parts else None return None def _parse_todo_list(text: str) -> list | None: """从 write_todos 输出中提取 todo 列表。 write_todos 输出格式: "Updated todo list to [{'content': '...', 'status': '...'}, ...]" 返回 JSON-serializable list of dicts，失败返回 None。 """ import ast # 贪婪匹配最外层 [...]（内容中的 [ ] 在字符串字面量内，ast.literal_eval 可正确处理） match = re.search(r"\[.*\]", text) if not match: return None try: python_list = ast.literal_eval(match.group()) if isinstance(python_list, list): return python_list except (ValueError, SyntaxError, TypeError): pass return None def _clean_response(text: str) -> str: """清理 Agent 回复中的格式噪音""" # 移除 ANSI 转义序列 text = re.sub(r'\x1b\[[0-9;]*m', '', text) text = re.sub(r'\x1b\[[0-9;]*[a-zA-Z]', '', text) text = re.sub(r'\x1b\][^\x07]*\x07', '', text) # 移除工具调用残留（大括号 JSON 块如果独立成行则移除） text = re.sub(r'^\s*\{[^}]*\}\s*$', '', text, flags=re.MULTILINE) return text.strip()