from __future__ import annotations from concurrent.futures import ThreadPoolExecutor from dataclasses import dataclass from datetime import datetime, timezone from typing import Literal from langgraph.graph import END, START, StateGraph from app.config import Settings from app.decision_engine import DecisionEngine from app.schemas import RouteName, ToolResult from app.state import AgentState from app.tools import SQLQueryTool, VectorSearchTool, WebSearchTool def trace_event(state: AgentState, node: str, detail: dict) -> list[dict]: return [ *state.get("trace", []), { "node": node, "at": datetime.now(timezone.utc).isoformat(), "detail": detail, }, ] @dataclass class GraphDependencies: settings: Settings engine: DecisionEngine vector_tool: VectorSearchTool sql_tool: SQLQueryTool web_tool: WebSearchTool def build_graph(deps: GraphDependencies): def normalize_query(state: AgentState) -> dict: normalized = " ".join(state["original_query"].strip().split()) return { "current_query": normalized, "retrieval_round": state.get("retrieval_round", 0), "trace": trace_event(state, "normalize_query", {"query": normalized}), } def route_query(state: AgentState) -> dict: decision = deps.engine.route( state["current_query"], deps.settings.max_retrieval_rounds ) return { "route_decision": decision, "trace": trace_event( state, "route_query", { "intent": decision.intent, "routes": [route.value for route in decision.routes], "reason_code": decision.reason_code, }, ), } def after_route( state: AgentState, ) -> Literal["clarify", "refuse", "generate", "plan"]: routes = set(state["route_decision"].routes) # 退出类路由优先级高于检索，防止复合输出意外触发工具。 if RouteName.CLARIFY in routes: return "clarify" if RouteName.REFUSE in routes: return "refuse" if RouteName.DIRECT_ANSWER in routes: return "generate" return "plan" def plan_query(state: AgentState) -> dict: plan = deps.engine.plan(state["current_query"], state["route_decision"]) return { "retrieval_plan": plan, "trace": trace_event( state, "plan_query", {"steps": [step.model_dump(mode="json") for step in plan.steps]}, ), } def execute_step(step) -> ToolResult: # Executor 只识别注册过的枚举工具；模型不能构造任意函数名。 if step.tool == RouteName.MILVUS_SEARCH: return deps.vector_tool.invoke( query=step.query, policy_type=str(step.arguments.get("policy_type", "")), top_k=int(step.arguments.get("top_k", 5)), ) if step.tool == RouteName.SQL_QUERY: return deps.sql_tool.invoke( user_id=str(step.arguments.get("user_id", "")), days=int(step.arguments.get("days", 30)), product_keyword=str(step.arguments.get("product_keyword", "")), ) if step.tool == RouteName.WEB_SEARCH: requested = int(step.arguments.get("max_results", 5)) # 即使模型给出更大的 max_results，也不能突破程序配置上限。 return deps.web_tool.invoke( query=step.query, max_results=min(max(requested, 1), deps.settings.web_search_max_results), ) return ToolResult( status="error", tool=step.tool.value, error_code="UNSUPPORTED_TOOL", error_message=f"未注册工具：{step.tool.value}", ) def execute_plan(state: AgentState) -> dict: steps = state["retrieval_plan"].steps if not steps: results: list[ToolResult] = [] elif len(steps) == 1: results = [execute_step(steps[0])] else: # 本项目的计划均为无依赖读任务；依赖型 DAG 需按拓扑批次调度。 with ThreadPoolExecutor(max_workers=min(len(steps), 4)) as executor: results = list(executor.map(execute_step, steps)) # ToolResult 用于控制流，Evidence 用于答案生成，二者不能混为一体。 evidence = [item for result in results for item in result.evidence] errors = [ { "tool": result.tool, "code": result.error_code, "message": result.error_message, } for result in results if result.status == "error" ] # 将规划 Query、受控参数和实际结果并列记录，便于观察真实调用。 executed_queries = [ { "step_id": step.id, "tool": result.tool, "query": step.query, "arguments": step.arguments, "status": result.status, "latency_ms": result.latency_ms, } for step, result in zip(steps, results, strict=True) ] return { "tool_results": results, "executed_queries": executed_queries, "evidence": evidence, "errors": [*state.get("errors", []), *errors], "trace": trace_event( state, "execute_plan", { "tools": [result.tool for result in results], "statuses": [result.status for result in results], "evidence_count": len(evidence), "executed_queries": executed_queries, }, ), } def grade_evidence(state: AgentState) -> dict: current_round = state.get("retrieval_round", 0) max_rounds = state["route_decision"].max_rounds grade = deps.engine.grade( query=state["current_query"], decision=state["route_decision"], evidence=state.get("evidence", []), current_round=current_round, max_rounds=max_rounds, min_score=deps.settings.min_evidence_score, ) # 模型只能建议 Rewrite，是否还有预算必须由程序层决定。 # 达到上限后强制 Stop，避免模型持续 Rewrite 导致图无限循环。 if ( grade.recommended_action == "rewrite_query" and current_round >= max_rounds ): grade = grade.model_copy( update={ "sufficient": False, "recommended_action": "stop", "reason": ( f"{grade.reason}；已达到最大检索轮数 {max_rounds}，" "程序层强制停止。" ), } ) return { "quality_grade": grade, "trace": trace_event( state, "grade_evidence", grade.model_dump(mode="json"), ), } def after_grade(state: AgentState) -> Literal["generate", "rewrite", "generate_partial"]: # 条件边只消费结构化动作，不解析 grader 的自然语言 reason。 action = state["quality_grade"].recommended_action if action == "accept": return "generate" if action == "rewrite_query": return "rewrite" return "generate_partial" def rewrite_query(state: AgentState) -> dict: rewritten = deps.engine.rewrite( state["current_query"], state["quality_grade"] ) next_round = state.get("retrieval_round", 0) + 1 return { "current_query": rewritten, "retrieval_round": next_round, "trace": trace_event( state, "rewrite_query", {"round": next_round, "rewritten_query": rewritten}, ), } def generate_answer(state: AgentState) -> dict: answer = deps.engine.answer( state["original_query"], state.get("evidence", []), partial=False ) return { "final_answer": answer, "termination_reason": "evidence_accepted" if state.get("evidence") else "direct_answer", "trace": trace_event(state, "generate_answer", {"partial": False}), } def generate_partial_answer(state: AgentState) -> dict: answer = deps.engine.answer( state["original_query"], state.get("evidence", []), partial=True ) return { "final_answer": answer, "termination_reason": "retrieval_budget_exhausted", "trace": trace_event(state, "generate_partial_answer", {"partial": True}), } def clarify(state: AgentState) -> dict: return { "final_answer": "请补充用户编号、商品名称或订单范围后再查询。", "termination_reason": "clarification_required", "trace": trace_event(state, "clarify", {}), } def refuse(state: AgentState) -> dict: return { "final_answer": "当前请求涉及受限数据，系统拒绝执行。", "termination_reason": "security_policy", "trace": trace_event(state, "refuse", {}), } builder = StateGraph(AgentState) builder.add_node("normalize_query", normalize_query) builder.add_node("route_query", route_query) builder.add_node("plan_query", plan_query) builder.add_node("execute_plan", execute_plan) builder.add_node("grade_evidence", grade_evidence) builder.add_node("rewrite_query", rewrite_query) builder.add_node("generate_answer", generate_answer) builder.add_node("generate_partial_answer", generate_partial_answer) builder.add_node("clarify", clarify) builder.add_node("refuse", refuse) builder.add_edge(START, "normalize_query") builder.add_edge("normalize_query", "route_query") builder.add_conditional_edges( "route_query", after_route, { "clarify": "clarify", "refuse": "refuse", "generate": "generate_answer", "plan": "plan_query", }, ) builder.add_edge("plan_query", "execute_plan") builder.add_edge("execute_plan", "grade_evidence") builder.add_conditional_edges( "grade_evidence", after_grade, { "generate": "generate_answer", "rewrite": "rewrite_query", "generate_partial": "generate_partial_answer", }, ) builder.add_edge("rewrite_query", "plan_query") builder.add_edge("generate_answer", END) builder.add_edge("generate_partial_answer", END) builder.add_edge("clarify", END) builder.add_edge("refuse", END) return builder.compile()