vent_deep_agent/api/sse_core.py

# -*- coding: utf-8 -*-
"""
SSE 核心工具函数 —— 流式事件生成器、文本提取、响应清理等通用辅助
"""

import json
import re
import time
import traceback
from typing import AsyncGenerator

from langgraph.types import Command  # 程序化 interrupt() 恢复时需启用

from tools.tool_names_cn import TOOL_NAME_CN as _TOOL_CN
from db.chat_store import save_message, get_messages
from tools.context_tracker import capture_context_usage


# ── 工具名 → 中文描述映射 ──


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,
    original_user_message: str | None = None,
    interrupt_before: list | None = None,
    interrupt_after: list | None = None,
) -> AsyncGenerator[str, None]:
    """
    通用 SSE 流式事件生成器。

    使用 stream_mode=["updates", "messages"] 获取完整执行图景，
    前端可根据 event type 区分：thinking / executing / tool_call / token / done / error。

    Args:
        save_to_db: 是否将消息保存到数据库。点选解读等非对话场景应设为 False。
        original_user_message: 用户的原始消息（不含系统注入的前缀）。
            若提供，则 DB 保存此原始消息；user_message 仍作为 Agent 的输入。
    """
    # 保存用户消息（仅对话场景写库），优先使用原始消息
    if save_to_db:
        save_message(session_id, "user", original_user_message or user_message)

    # 用于收集完整回复
    full_response = ""
    start_time = time.time()
    _token_usage = {"prompt": 0, "completion": 0}  # 从流中捕获的实际 token 用量

    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,
            interrupt_before=interrupt_before,
            interrupt_after=interrupt_after,
            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":
                        # 代理正在思考/推理
                        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:
                    # 先检查推理/思考内容（DeepSeek 等模型的 reasoning_content）
                    reasoning = _extract_reasoning_content(msg_obj)
                    if reasoning:
                        yield f"data: {json.dumps({'type': 'thinking_token', 'source': source, 'content': reasoning}, ensure_ascii=False)}\n\n"
                    # 再提取常规文本内容
                    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"

                # ── 提取 token 用量元数据（最后一个消息块通常携带 usage）──
                _capture_token_meta(msg_obj, _token_usage)

        # 保存助手回复（仅对话场景写库）
        if save_to_db and full_response.strip():
            cleaned = _clean_response(full_response)
            save_message(session_id, "assistant", cleaned)
            print(full_response)

        # ── 更新上下文用量 ──
        if save_to_db and _token_usage["prompt"] > 0:
            await capture_context_usage(
                session_id=session_id,
                prompt_tokens=_token_usage["prompt"],
                completion_tokens=_token_usage["completion"],
                current_message=user_message,
            )

        # 发送 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"

        # 检查 LangGraph 中断状态（Human-in-the-Loop）
        interrupt_info = _check_interrupt(agent, config)
        if interrupt_info:
            yield f"data: {json.dumps({'type': 'interrupt', **interrupt_info}, ensure_ascii=False)}\n\n"
            # 中断时不发送 done 事件，等待用户审批后恢复
            return

        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


# ── 调试开关：设为 True 时打印推理内容摘要（排查完毕后关闭）──
_DEBUG_REASONING = False


def _extract_reasoning_content(msg_obj) -> str | None:
    """从 LangChain 消息对象中提取推理/思考内容（DeepSeek 等模型的 reasoning_content）。

    支持三种来源（按优先级）：
    1. msg_obj.additional_kwargs["reasoning_content"]（OpenAI 兼容流式 delta）
    2. msg_obj.reasoning_content（LangChain 直接属性）
    3. content 中的 thinking 类型块

    Returns:
        推理文本字符串，无推理内容时返回 None
    """
    # 方式 1：additional_kwargs 中的 reasoning_content（最常见）
    if hasattr(msg_obj, "additional_kwargs") and isinstance(msg_obj.additional_kwargs, dict):
        reasoning = msg_obj.additional_kwargs.get("reasoning_content", "")
        if reasoning:
            return reasoning

    # 方式 2：直接属性 reasoning_content
    reasoning = getattr(msg_obj, "reasoning_content", None)
    if reasoning:
        return reasoning

    # 方式 3：content 为 list 时，提取 thinking 类型的块
    if hasattr(msg_obj, "content"):
        raw = getattr(msg_obj, "content", None)
        if isinstance(raw, list):
            parts = []
            for block in raw:
                if isinstance(block, dict) and block.get("type") == "thinking":
                    parts.append(block.get("thinking", ""))
                elif hasattr(block, "type") and getattr(block, "type", "") == "thinking":
                    parts.append(getattr(block, "thinking", ""))
            if parts:
                return "".join(parts)

    return None


def _capture_token_meta(msg_obj, usage_ref: dict):
    """从 LangChain 消息对象中提取 LLM token 用量元数据。

    优先读取 usage_metadata（langchain ≥0.3），
    其次读取 response_metadata.usage（OpenAI 兼容格式）。
    结果写入 usage_ref dict（原地修改）。
    """
    # 方式 1：usage_metadata（langchain 标准字段）
    um = getattr(msg_obj, "usage_metadata", None)
    if um and isinstance(um, dict):
        inp = um.get("input_tokens", 0)
        out = um.get("output_tokens", 0)
        if inp or out:
            usage_ref["prompt"] = inp
            usage_ref["completion"] = out
            return

    # 方式 2：response_metadata.usage（OpenAI / DeepSeek 兼容）
    rm = getattr(msg_obj, "response_metadata", None)
    if rm and isinstance(rm, dict):
        usage = rm.get("usage", {}) or rm.get("token_usage", {})
        if isinstance(usage, dict):
            inp = usage.get("prompt_tokens", 0)
            out = usage.get("completion_tokens", 0)
            if inp or out:
                usage_ref["prompt"] = inp
                usage_ref["completion"] = out
                return


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


# ── LangGraph Human-in-the-Loop 中断处理 ──


def _check_interrupt(agent, config: dict) -> dict | None:
    """
    检查 LangGraph 状态是否被中断。

    支持两种中断检测：
    1. 程序化中断（节点内调用 interrupt()）→ state.interrupts 非空
    2. interrupt_before / interrupt_after 暂停 → state.next 非空

    当中断发生时返回中断信息 dict，否则返回 None。

    返回格式:
        {"node": "tools", "message": "智能体准备执行工具，等待审批...",
         "interrupts": [...], "plan": "计划文本（程序化中断时）"}
    """
    try:
        state = agent.get_state(config)
    except Exception:
        return None

    if state is None:
        return None

    # 1. 检测程序化中断（节点内调用 interrupt()）
    interrupts = getattr(state, "interrupts", None)
    if not interrupts:
        values = getattr(state, "values", {}) or {}
        interrupts = values.get("__interrupt__", [])

    # 2. 检测 interrupt_before / interrupt_after 暂停
    #    图暂停时 state.next 非空（有待执行节点）；完成时为空
    next_nodes = getattr(state, "next", None)
    is_paused_by_interrupt_config = bool(next_nodes) if next_nodes is not None else False

    if not interrupts and not is_paused_by_interrupt_config:
        return None

    # 提取中断节点名
    if interrupts:
        interrupt_data = interrupts[0] if interrupts else {}
        # 处理 interrupt() 返回值可能是 Interrupt 对象或 dict
        if hasattr(interrupt_data, "value"):
            interrupt_data = interrupt_data.value
        if isinstance(interrupt_data, dict):
            node = interrupt_data.get("type", "tools")
            plan = interrupt_data.get("plan", "")
            message = interrupt_data.get("message", "智能体暂停执行，等待您的审批...")
        else:
            node = str(next_nodes or "tools")
            plan = ""
            message = "智能体暂停执行，等待您的审批..."
    else:
        node = next_nodes
        plan = ""
        message = "智能体暂停执行，等待您的审批..."
    if isinstance(node, (list, tuple)):
        node = node[0] if node else "unknown"

    result = {
        "node": str(node),
        "message": message,
        "interrupts": [str(i) for i in interrupts] if interrupts else [],
    }
    if plan:
        result["plan"] = plan
    return result


async def resume_stream(
    agent,
    session_id: str,
    thread_id: str,
    config: dict,
    agent_cn_name: str = "智能助手",
    action: str = "approve",
    interrupt_before: list | None = None,
    interrupt_after: list | None = None,
):
    """
    恢复被中断的 LangGraph 流式执行。

    参数:
        action: "approve" 或 "reject"

    Yields:
        SSE 事件字符串
    """
    if action == "reject":
        yield f"data: {json.dumps({'type': 'error', 'message': '用户拒绝了工具执行'}, ensure_ascii=False)}\n\n"
        yield f"data: {json.dumps({'type': 'done', 'thread_id': thread_id, 'session_id': session_id, 'message': '已取消执行'}, ensure_ascii=False)}\n\n"
        return

    full_response = ""
    start_time = time.time()

    try:
        # 检测中断类型：程序化中断需用 Command(resume=...)，配置式中断传 None
        state = agent.get_state(config)
        has_programmatic = bool(getattr(state, "interrupts", None)) if state else False
        stream_input = Command(resume={"action": action}) if has_programmatic else None

        # 恢复执行
        async for chunk in agent.astream(
            stream_input,
            stream_mode=["updates", "messages"],
            config=config,
            interrupt_before=interrupt_before,
            interrupt_after=interrupt_after,
            version="v2",
        ):
            is_subagent = any(s.startswith("tools:") for s in chunk.get("ns", []))

            if chunk["type"] == "updates":
                for node_name in chunk["data"]:
                    if node_name == "model":
                        yield f"data: {json.dumps({'type': 'thinking', 'source': 'main', '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 ["工具调用"]
                        yield f"data: {json.dumps({'type': 'executing', 'source': 'main', 'cn_agent': agent_cn_name, 'tools': tool_names, 'cn_tools': cn_names, 'message': f"正在{'、'.join(cn_names)}..."}, ensure_ascii=False)}\n\n"

            elif chunk["type"] == "messages":
                token_data = chunk["data"]
                if isinstance(token_data, (list, tuple)) and len(token_data) >= 1:
                    msg_obj = token_data[0]
                else:
                    msg_obj = token_data

                # 工具结果
                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":
                        yield f"data: {json.dumps({'type': 'tool_result', 'source': 'main', 'cn_agent': agent_cn_name, 'tool': tool_name, 'cn_tool': _cn_tool_desc(tool_name), 'message': f'{_cn_tool_desc(tool_name)} 完成'}, ensure_ascii=False)}\n\n"

                # AI 文本
                if not is_tool_msg:
                    # 先检查推理/思考内容（DeepSeek 等模型的 reasoning_content）
                    reasoning = _extract_reasoning_content(msg_obj)
                    if reasoning:
                        yield f"data: {json.dumps({'type': 'thinking_token', 'source': 'main', 'content': reasoning}, ensure_ascii=False)}\n\n"
                    # 再提取常规文本内容
                    content = _extract_text_content(msg_obj)
                    if content and isinstance(content, str):
                        full_response += content
                        yield f"data: {json.dumps({'type': 'token', 'source': 'main', 'content': content}, ensure_ascii=False)}\n\n"

        # 保存助手回复
        if full_response.strip():
            cleaned = _clean_response(full_response)
            save_message(session_id, "assistant", cleaned)

        # 再次检查中断（可能有多次中断）
        interrupt_info = _check_interrupt(agent, config)
        if interrupt_info:
            yield f"data: {json.dumps({'type': 'interrupt', **interrupt_info}, ensure_ascii=False)}\n\n"
            return

        duration_ms = int((time.time() - start_time) * 1000)
        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"