smqjh-agent/smqjh-agent-runtime/app/agent_graph.py

from __future__ import annotations

import json
import re
import time
from contextvars import ContextVar
from datetime import datetime, timezone
from typing import Any, Callable, Literal

from langgraph.graph import END, START, StateGraph

from .config import RuntimeConfig
from .mcp_client import McpClient
from .models import (
    AgentResultTable,
    AgentRunResponse,
    AgentState,
    AgentTraceEvent,
    ChatMessage,
    McpToolDescriptor,
    PlannedToolCall,
    ToolObservation,
)


Route = Literal["execute_tools", "summarize", "chat"]
TraceSink = Callable[[AgentTraceEvent], None]

_TRACE_SINK: ContextVar[TraceSink | None] = ContextVar("smqjh_agent_trace_sink", default=None)


class SmqjhAgentGraph:
    def __init__(self, config: RuntimeConfig) -> None:
        self.config = config
        self.mcp = McpClient(config.mcp.url, config.mcp.token)
        self.graph = self._build_graph()

    def health(self) -> Any:
        return self.mcp.status()

    def tools(self) -> list[McpToolDescriptor]:
        return self._load_runnable_tools()

    def run(
        self,
        message: str,
        history: list[ChatMessage] | None = None,
        emit: TraceSink | None = None,
    ) -> AgentRunResponse:
        token = _TRACE_SINK.set(emit)
        try:
            state: AgentState = {
                "message": message,
                "history": history or [],
                "visitedSignatures": [],
                "steps": 0,
                "trace": [_trace("system", "已收到请求", message)],
            }
            result = self.graph.invoke(state)
            final = result.get("final")
            if final:
                return final

            observations = result.get("observations", [])
            trace_events = result.get("trace", [])
            return {
                "content": "这次没有生成有效回答，请稍后重试。",
                "model": "none",
                "usedMcp": False,
                "steps": int(result.get("steps", 0)),
                "toolCalls": observations,
                "tables": _extract_tables(observations),
                "trace": trace_events,
            }
        finally:
            _TRACE_SINK.reset(token)

    def _build_graph(self):
        workflow = StateGraph(AgentState)
        workflow.add_node("load_tools", self._load_tools_node)
        workflow.add_node("plan", self._plan_node)
        workflow.add_node("execute_tools", self._execute_tools_node)
        workflow.add_node("summarize", self._summarize_node)
        workflow.add_node("chat", self._chat_node)
        workflow.add_edge(START, "load_tools")
        workflow.add_edge("load_tools", "plan")
        workflow.add_conditional_edges(
            "plan",
            self._route_after_plan,
            {
                "execute_tools": "execute_tools",
                "summarize": "summarize",
                "chat": "chat",
            },
        )
        workflow.add_edge("execute_tools", "plan")
        workflow.add_edge("summarize", END)
        workflow.add_edge("chat", END)
        return workflow.compile()

    def _load_tools_node(self, state: AgentState) -> AgentState:
        start = _trace("system", "正在连接 MCP", "加载可用业务工具和数据库只读查询能力")
        try:
            tools = self._load_runnable_tools()
            done = _trace("system", "MCP 工具加载完成", f"可用工具 {len(tools)} 个")
            return {
                "tools": tools,
                "trace": _append_trace(state, start, done),
            }
        except Exception as error:
            failed = _trace("error", "MCP 工具加载失败", _error_message(error))
            return {
                "tools": [],
                "trace": _append_trace(state, start, failed),
            }

    def _plan_node(self, state: AgentState) -> AgentState:
        if int(state.get("steps", 0)) >= self.config.max_steps:
            return {
                "plannedToolCalls": [],
                "trace": _append_trace(
                    state,
                    _trace("plan", "达到最大工具步数", f"maxSteps={self.config.max_steps}"),
                ),
            }

        start = _trace("plan", "正在分析需求", "DeepSeek 正在判断是否需要查库或调用业务工具")
        try:
            result = self.mcp.call_tool(
                "smqjh.ai.tool.plan",
                {
                    "request": self._assistant_request(state),
                    "tools": state.get("tools", []),
                    "observations": state.get("observations", []),
                },
            )
            record = _as_dict(result.get("structuredContent"))
            raw_calls = record.get("toolCalls")
            raw_call_items = raw_calls if isinstance(raw_calls, list) else []
            visited = set(state.get("visitedSignatures", []))
            planned = [
                call
                for call in (_normalize_tool_call(item, state.get("message", "")) for item in raw_call_items)
                if call and _tool_signature(call["name"], call["arguments"]) not in visited
            ][:3]
            done = _trace(
                "plan",
                "DeepSeek 已规划工具调用" if planned else "DeepSeek 判断无需继续调用工具",
                _reason_text(record),
                {"toolCalls": planned},
            )
            return {
                "plannedToolCalls": planned,
                "trace": _append_trace(state, start, done),
            }
        except Exception as error:
            failed = _trace("error", "DeepSeek 工具规划失败", _error_message(error))
            return {
                "plannedToolCalls": [],
                "trace": _append_trace(state, start, failed),
            }

    def _execute_tools_node(self, state: AgentState) -> AgentState:
        by_name = {tool.get("name", ""): tool for tool in state.get("tools", [])}
        observations: list[ToolObservation] = []
        visited = set(state.get("visitedSignatures", []))
        trace_events = list(state.get("trace", []))

        for call in state.get("plannedToolCalls", []):
            requested_name = call["name"]
            arguments = call["arguments"]
            resolved_name = _resolve_tool_name(requested_name, by_name)
            signature = _tool_signature(resolved_name or requested_name, arguments)
            visited.add(signature)

            if not resolved_name:
                observation: ToolObservation = {
                    "name": requested_name,
                    "source": "mcp",
                    "arguments": arguments,
                    "ok": False,
                    "error": "DeepSeek 选择了未注册 MCP 工具",
                }
                observations.append(observation)
                trace_events.append(_trace("error", f"工具未注册：{requested_name}", observation["error"]))
                continue

            trace_events.append(_trace("tool", f"正在执行：{resolved_name}", _tool_title(resolved_name)))
            started = time.perf_counter()
            try:
                result = self.mcp.call_tool(resolved_name, arguments)
                observation = {
                    "name": resolved_name,
                    "source": "mcp",
                    "arguments": arguments,
                    "ok": True,
                    "result": result.get("structuredContent", result),
                    "durationMs": int((time.perf_counter() - started) * 1000),
                }
                observations.append(observation)
                trace_events.append(
                    _trace(
                        "tool",
                        f"工具执行完成：{resolved_name}",
                        f"{observation.get('durationMs', 0)}ms",
                        {"arguments": arguments, "result": _summarize_result(observation.get("result"))},
                    )
                )
            except Exception as error:
                observation = {
                    "name": resolved_name,
                    "source": "mcp",
                    "arguments": arguments,
                    "ok": False,
                    "error": _error_message(error),
                    "durationMs": int((time.perf_counter() - started) * 1000),
                }
                observations.append(observation)
                trace_events.append(
                    _trace(
                        "error",
                        f"工具执行失败：{resolved_name}",
                        observation.get("error", ""),
                        {"arguments": arguments},
                    )
                )

        return {
            "observations": [*state.get("observations", []), *observations],
            "plannedToolCalls": [],
            "visitedSignatures": list(visited),
            "steps": int(state.get("steps", 0)) + 1,
            "trace": trace_events,
        }

    def _summarize_node(self, state: AgentState) -> AgentState:
        observations = state.get("observations", [])
        tables = _extract_tables(observations)
        start = _trace("summary", "正在整理结果", "将工具结果整理成结论、依据和表格")
        try:
            result = self.mcp.call_tool(
                "smqjh.ai.tool.summarize",
                {
                    "request": self._assistant_request(state),
                    "observations": observations,
                },
            )
            record = _as_dict(result.get("structuredContent"))
            content = str(record.get("content") or "").strip() or _fallback_summary(observations, tables)
            model = str(record.get("model") or "mcp-deepseek")
            done = _trace("summary", "结果总结完成", f"{len(tables)} 个表格")
            trace_events = _append_trace(state, start, done)
            return {
                "final": {
                    "content": content,
                    "model": model,
                    "usedMcp": True,
                    "steps": int(state.get("steps", 0)),
                    "toolCalls": observations,
                    "tables": tables,
                    "trace": trace_events,
                },
                "trace": trace_events,
            }
        except Exception as error:
            failed = _trace("error", "结果总结失败，已使用工具结果兜底", _error_message(error))
            trace_events = _append_trace(state, start, failed)
            return {
                "final": {
                    "content": _fallback_summary(observations, tables),
                    "model": "tool-only",
                    "usedMcp": True,
                    "steps": int(state.get("steps", 0)),
                    "toolCalls": observations,
                    "tables": tables,
                    "trace": trace_events,
                },
                "trace": trace_events,
            }

    def _chat_node(self, state: AgentState) -> AgentState:
        start = _trace("chat", "正在生成回复", "没有继续调用业务工具，进入普通对话回答")
        try:
            result = self.mcp.call_tool("smqjh.ai.chat", {"request": self._assistant_request(state)})
            record = _as_dict(result.get("structuredContent"))
            content = str(record.get("content") or "").strip() or "DeepSeek 没有返回有效内容。"
            model = str(record.get("model") or "mcp-deepseek")
            done = _trace("chat", "普通对话完成", model)
            trace_events = _append_trace(state, start, done)
            return {
                "final": {
                    "content": content,
                    "model": model,
                    "usedMcp": True,
                    "steps": int(state.get("steps", 0)),
                    "toolCalls": [],
                    "tables": [],
                    "trace": trace_events,
                },
                "trace": trace_events,
            }
        except Exception as error:
            failed = _trace("error", "普通对话失败", _error_message(error))
            trace_events = _append_trace(state, start, failed)
            return {
                "final": {
                    "content": f"暂时无法完成回答：{_error_message(error)}",
                    "model": "none",
                    "usedMcp": False,
                    "steps": int(state.get("steps", 0)),
                    "toolCalls": [],
                    "tables": [],
                    "trace": trace_events,
                },
                "trace": trace_events,
            }

    def _route_after_plan(self, state: AgentState) -> Route:
        if state.get("plannedToolCalls") and int(state.get("steps", 0)) < self.config.max_steps:
            return "execute_tools"
        if state.get("observations"):
            return "summarize"
        return "chat"

    def _assistant_request(self, state: AgentState) -> dict[str, Any]:
        return {
            "message": state.get("message", ""),
            "history": state.get("history", []),
            "environmentName": self.config.environment_name,
            "baseUrl": self.config.base_url,
            "authenticated": True,
            "username": "web-agent",
        }

    def _load_runnable_tools(self) -> list[McpToolDescriptor]:
        return [tool for tool in self.mcp.list_tools() if not str(tool.get("name", "")).startswith("smqjh.ai.")]


def _normalize_tool_call(value: Any, message: str) -> PlannedToolCall | None:
    record = _as_dict(value)
    name = str(record.get("name") or "").strip()
    if not name:
        return None
    return {
        "name": name,
        "arguments": _repair_tool_arguments(name, _as_dict(record.get("arguments")), message),
    }


def _repair_tool_arguments(name: str, arguments: dict[str, Any], message: str) -> dict[str, Any]:
    if name == "smqjh.database.smart.query" and not isinstance(arguments.get("question"), str):
        return {**arguments, "question": message}
    if name == "smqjh.product.lookup.summary":
        current = str(arguments.get("productKeyword") or "").strip()
        return {**arguments, "productKeyword": current or _extract_product_keyword(message)}
    if name in {
        "smqjh.settlement.monthly.plan",
        "smqjh.settlement.monthly.preview",
        "settlement.monthly.plan",
        "settlement.monthly.preview",
    }:
        enterprise = str(arguments.get("enterprise") or "").strip() or _extract_settlement_enterprise(message)
        month = str(arguments.get("month") or "").strip() or _extract_month(message)
        return {**arguments, "enterprise": enterprise, "month": month}
    return arguments


def _extract_product_keyword(message: str) -> str:
    text = message
    stop_words = [
        "帮我",
        "麻烦",
        "查询一下",
        "查询",
        "查看",
        "当前",
        "业务系统",
        "系统里面",
        "系统里",
        "后台",
        "我方",
        "我们的",
        "商品库",
        "商品表",
        "商品描述是什么",
        "商品描述",
        "描述是什么",
        "描述",
        "价格是多少",
        "价格",
        "定价",
        "是多少",
        "是什么",
        "呢",
    ]
    for token in stop_words:
        text = text.replace(token, " ")
    text = re.sub(r"[，。！？、；：,.!?;:]", " ", text)
    return " ".join(text.split())


def _extract_settlement_enterprise(message: str) -> str:
    for enterprise in ("铜仁移动", "招商银行贵阳分行", "中数未来"):
        if enterprise in message:
            return enterprise
    return ""


def _extract_month(message: str) -> str:
    explicit = re.search(r"(20\d{2})[-/年](0?[1-9]|1[0-2])", message)
    if explicit:
        return f"{explicit.group(1)}-{int(explicit.group(2)):02d}"
    month_only = re.search(r"(?<!\d)(0?[1-9]|1[0-2])\s*月份?", message)
    if month_only:
        return f"{datetime.now().year}-{int(month_only.group(1)):02d}"
    return ""


def _resolve_tool_name(name: str, by_name: dict[str, McpToolDescriptor]) -> str | None:
    candidates = [
        name,
        name.removeprefix("smqjh.") if name.startswith("smqjh.") else f"smqjh.{name}",
        "smqjh.product.lookup.summary" if name == "product.lookup.summary" else "",
        "smqjh.order.count.query" if name == "order.count.query" else "",
        "smqjh.database.smart.query" if name == "database.smart.query" else "",
        "smqjh.database.readonly.query" if name == "database.readonly.query" else "",
        "smqjh.settlement.monthly.preview" if name == "settlement.monthly.preview" else "",
        "smqjh.settlement.monthly.plan" if name == "settlement.monthly.plan" else "",
        "smqjh.cloud.health" if name == "cloud.health" else "",
    ]
    return next((candidate for candidate in candidates if candidate and candidate in by_name), None)


def _extract_tables(observations: list[ToolObservation]) -> list[AgentResultTable]:
    tables: list[AgentResultTable] = []
    for observation in observations:
        if not observation.get("ok"):
            continue
        tables.extend(_extract_tables_from_value(_tool_title(observation.get("name", "")), observation.get("result")))
    return tables[:6]


def _extract_tables_from_value(title: str, value: Any) -> list[AgentResultTable]:
    tables: list[AgentResultTable] = []
    seen: set[int] = set()

    def visit(node: Any, current_title: str, depth: int) -> None:
        if depth > 5 or not isinstance(node, dict):
            return
        node_id = id(node)
        if node_id in seen:
            return
        seen.add(node_id)
        add_rows_table(current_title, node, "rows")
        add_rows_table(current_title, node, "comparisonRows")
        for key in ("data", "structuredContent", "result", "record"):
            if key in node:
                visit(node[key], current_title, depth + 1)

    def add_rows_table(current_title: str, record: dict[str, Any], rows_key: str) -> None:
        rows = record.get(rows_key)
        if not isinstance(rows, list) or not all(isinstance(row, dict) for row in rows):
            return
        row_records = [_normalize_row(row) for row in rows[:200]]
        column_key = "comparisonColumns" if rows_key == "comparisonRows" else "columns"
        provided_columns = record.get(column_key)
        if isinstance(provided_columns, list) and provided_columns:
            columns = [str(column) for column in provided_columns]
        else:
            columns = []
            for row in row_records:
                for key in row:
                    if key not in columns:
                        columns.append(key)
        if not columns:
            return
        tables.append(
            {
                "title": _build_table_title(current_title, record, rows_key),
                "columns": columns,
                "rows": row_records,
            }
        )

    visit(value, title, 0)
    return tables


def _fallback_summary(observations: list[ToolObservation], tables: list[AgentResultTable]) -> str:
    ok_count = len([item for item in observations if item.get("ok")])
    failed = [item for item in observations if not item.get("ok")]
    lines = [f"已执行 {len(observations)} 个 MCP 工具，成功 {ok_count} 个。"]
    if tables:
        lines.append("结果已整理为下方表格。")
    if failed:
        lines.append("失败工具：" + "；".join(f"{item.get('name')}：{item.get('error') or '调用失败'}" for item in failed))
    evidence = _find_evidence(observations)
    if evidence:
        lines.append(f"依据：{evidence}")
    return "\n\n".join(lines)


def _find_evidence(observations: list[ToolObservation]) -> str:
    for observation in observations:
        sql = _find_key_value(observation.get("result"), "executedSql", 0)
        if isinstance(sql, str) and sql.strip():
            return sql.strip()
        evidence = _find_key_value(observation.get("result"), "evidence", 0)
        if isinstance(evidence, str) and evidence.strip():
            return evidence.strip()
    return ""


def _find_key_value(value: Any, key: str, depth: int) -> Any:
    if depth > 4 or not isinstance(value, dict):
        return None
    if key in value:
        return value[key]
    for child in value.values():
        result = _find_key_value(child, key, depth + 1)
        if result is not None:
            return result
    return None


def _trace(phase: str, title: str, detail: str = "", data: Any = None) -> AgentTraceEvent:
    event: AgentTraceEvent = {
        "at": datetime.now(timezone.utc).isoformat(),
        "phase": phase,  # type: ignore[typeddict-item]
        "title": title,
    }
    if detail:
        event["detail"] = detail
    if data is not None:
        event["data"] = data
    sink = _TRACE_SINK.get()
    if sink:
        try:
            sink(event)
        except Exception:
            pass
    return event


def _append_trace(state: AgentState, *events: AgentTraceEvent) -> list[AgentTraceEvent]:
    return [*state.get("trace", []), *events]


def _reason_text(record: dict[str, Any]) -> str:
    reason = record.get("reason")
    return reason.strip() if isinstance(reason, str) else ""


def _summarize_result(value: Any) -> Any:
    record = _as_dict(value)
    if isinstance(record.get("rowCount"), int):
        return {"rowCount": record.get("rowCount"), "title": record.get("title"), "summary": record.get("summary")}
    if isinstance(record.get("rows"), list):
        return {"rows": len(record["rows"])}
    return value


def _build_table_title(base_title: str, record: dict[str, Any], rows_key: str) -> str:
    title = str(record.get("title") or base_title).strip()
    count = f"（{record['rowCount']} 行）" if "rowCount" in record else ""
    if rows_key == "comparisonRows":
        suffix = str(record.get("comparisonTitle") or "明细表").strip()
        return f"{title}：{suffix}"
    return f"{title}{count}"


def _tool_title(name: str) -> str:
    titles = {
        "smqjh.config.get": "运行配置",
        "smqjh.cloud.health": "网关连通检查",
        "smqjh.schema.search": "业务表搜索结果",
        "smqjh.schema.getTable": "业务表说明",
        "smqjh.schema.businessRules": "业务规则",
        "smqjh.database.readonly.query": "数据库只读查询结果",
        "smqjh.database.smart.query": "智能数据库查询结果",
        "smqjh.order.count.query": "订单统计结果",
        "smqjh.product.lookup.summary": "商品资料查询结果",
        "smqjh.settlement.enterprise.list": "月结企业清单",
        "smqjh.settlement.monthly.plan": "企业月结计划",
        "smqjh.settlement.monthly.preview": "企业月结预览",
    }
    return titles.get(name, name)


def _tool_signature(name: str, arguments: dict[str, Any]) -> str:
    return f"{name}:{json.dumps(arguments, ensure_ascii=False, sort_keys=True, default=str)}"


def _as_dict(value: Any) -> dict[str, Any]:
    return value if isinstance(value, dict) else {}


def _normalize_row(row: dict[str, Any]) -> dict[str, str]:
    return {str(key): "" if value is None else str(value) for key, value in row.items()}


def _error_message(error: Exception) -> str:
    return str(error) or error.__class__.__name__