vent_deep_agent/api/sse_core.py

# -*- 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()