refactor: 改用 LangGraph 原生 create_react_agent + astream_events

backend/app/main_graph/main_graph_builder.py

@@ -1,31 +1,22 @@
 """
-极简 Agent 主图 - 回归 LangGraph 标准模式
-
-架构：
-START → [init_state] → [记忆] → [Agent] ⇄ [Tools] → [Finalize] → END
-                                        ↑________↓
+极简 Agent 主图 - 用 LangGraph 原生 create_react_agent + 记忆节点
 """
 
+from langgraph.prebuilt import create_react_agent
 from langgraph.graph import StateGraph, START, END
-from langgraph.prebuilt import ToolNode
-from langchain_core.runnables.config import RunnableConfig
-from typing import Dict, Any, Optional
-
-from .state import AgentState
-from .nodes.memory_trigger import memory_trigger_node, set_mem0_client
-from .nodes.summarize import create_summarize_node
-from .nodes.agent import create_agent_node
-from backend.app.tools import ALL_TOOLS
+from ..state import AgentState
+from ..nodes.memory_trigger import memory_trigger_node, set_mem0_client
 from backend.app.logger import info, warning
+from backend.app.tools import ALL_TOOLS
 
 
 def build_agent_graph(
     chat_services: dict,
     mem0_client=None,
     max_steps: int = 10
-) -> StateGraph:
+):
     """
-    构建极简 Agent 图
+    构建包含记忆节点的 react agent 图
 
     Args:
         chat_services: 模型服务字典
@@ -33,153 +24,56 @@ def build_agent_graph(
         max_steps: 最大步数限制
 
     Returns:
-        StateGraph: 构建好的图
+        编译好的 graph
     """
-
-    graph = StateGraph(AgentState)
+    # 获取主模型
+    primary_model = chat_services.get("primary", next(iter(chat_services.values())))
 
     # ========== 设置全局客户端 ==========
     if mem0_client:
         set_mem0_client(mem0_client)
 
-    # ========== 创建核心节点 ==========
-
-    # 1. Agent 节点（绑定工具的 LLM）
-    llm = chat_services.get("primary", list(chat_services.values())[0] if chat_services else None)
-    if llm is None:
-        raise ValueError("No LLM service provided")
-
-    llm_with_tools = llm.bind_tools(ALL_TOOLS)
-    agent_node = create_agent_node(llm_with_tools, llm)
-
-    # 2. Tool 节点（LangGraph 内置）
-    tool_node = ToolNode(ALL_TOOLS)
-
-    # 3. 记忆/总结节点（保留现有）
-    retrieve_memory_node = None
-    summarize_node = None
-    if mem0_client:
-        try:
-            from .nodes.retrieve_memory import create_retrieve_memory_node
-            retrieve_memory_node = create_retrieve_memory_node(mem0_client)
-            summarize_node = create_summarize_node(mem0_client)
-        except Exception as e:
-            info(f"[Graph Builder] 记忆节点初始化失败: {e}")
-
-    # ========== 添加节点 ==========
-
-    # 1. 初始化节点（重置步数）
-    async def init_state_node(state: AgentState) -> Dict[str, Any]:
-        """初始化状态：重置步数计数器"""
+    # ========== 1. 初始化节点：重置步数 ==========
+    async def init_state_node(state: AgentState):
         info("[Init State] 初始化状态，重置步数")
         return {
             "current_step": 0
         }
 
-    graph.add_node("init_state", init_state_node)
+    # ========== 2. 记忆节点（可选） ==========
+    retrieve_memory_node = None
+    if mem0_client:
+        try:
+            from ..nodes.retrieve_memory import create_retrieve_memory_node
+            retrieve_memory_node = create_retrieve_memory_node(mem0_client)
+        except Exception as e:
+            info(f"[Graph Builder] 记忆节点初始化失败: {e}")
 
-    # 2. 记忆阶段
+    # ========== 3. 创建 react agent 子图 ==========
+    agent_runnable = create_react_agent(primary_model, ALL_TOOLS)
+
+    # ========== 4. 构建主图 ==========
+    graph = StateGraph(AgentState)
+
+    graph.add_node("init_state", init_state_node)
     if retrieve_memory_node:
         graph.add_node("retrieve_memory", retrieve_memory_node)
     graph.add_node("memory_trigger", memory_trigger_node)
 
-    # 3. 核心 Agent 循环
-    graph.add_node("agent", agent_node)
-    graph.add_node("tools", tool_node)
+    # 直接把 create_react_agent 的可运行对象作为节点
+    graph.add_node("agent", agent_runnable)
 
-    # 4. 完成阶段
-    if summarize_node:
-        graph.add_node("summarize", summarize_node)
-
-    # 简单的完成节点
-    async def finalize_node_simple(state: AgentState, config: Optional[RunnableConfig] = None) -> Dict[str, Any]:
-        """简单的完成节点，只发送完成事件"""
-        info("[Finalize] 进入完成节点")
-
-        try:
-            from backend.app.main_graph.config import get_stream_writer
-            writer = get_stream_writer()
-
-            # 提取最后的回复
-            final_reply = ""
-            if state.messages:
-                last_msg = state.messages[-1]
-                final_reply = last_msg.content if hasattr(last_msg, "content") else str(last_msg)
-
-            if writer and hasattr(writer, "__call__"):
-                try:
-                    writer({
-                        "type": "custom",
-                        "data": {
-                            "type": "done",
-                            "token_usage": state.last_token_usage,
-                            "elapsed_time": state.last_elapsed_time,
-                            "final_result": final_reply
-                        }
-                    })
-                    info("🏁 [完成事件] 已发送完成事件")
-                except Exception as e:
-                    warning(f"⚠️ [完成事件] 发送失败 (非致命): {e}")
-        except Exception as e:
-            warning(f"⚠️ [完成事件] 处理失败 (非致命): {e}")
-
-        return {}
-
-    graph.add_node("finalize", finalize_node_simple)
-
-    # ========== 添加边 ==========
-
-    # 1. 初始化
+    # ========== 边的连接 ==========
     graph.add_edge(START, "init_state")
 
-    # 2. 记忆阶段
     if retrieve_memory_node:
         graph.add_edge("init_state", "retrieve_memory")
         graph.add_edge("retrieve_memory", "memory_trigger")
     else:
         graph.add_edge("init_state", "memory_trigger")
 
-    # 3. 进入 Agent
     graph.add_edge("memory_trigger", "agent")
+    graph.add_edge("agent", END)
 
-    # 4. 核心循环：Agent ⇄ Tools
-    def should_continue(state: AgentState) -> str:
-        """判断是继续调用工具还是结束"""
-        messages = state.messages
-        last_message = messages[-1] if messages else None
-
-        # 检查是否有 tool_calls
-        if last_message and hasattr(last_message, "tool_calls") and last_message.tool_calls:
-            return "tools"
-
-        # 否则结束
-        return "finalize"
-
-    graph.add_conditional_edges(
-        "agent",
-        should_continue,
-        {
-            "tools": "tools",
-            "finalize": "finalize"
-        }
-    )
-
-    # Tools 执行完回到 Agent
-    graph.add_edge("tools", "agent")
-
-    # 5. 完成阶段
-    if summarize_node:
-        def should_summarize(state: AgentState) -> str:
-            if state.turns_since_last_summary >= 5:
-                return "summarize"
-            return "finalize"
-
-        # 总结逻辑暂简化：先 finalize
-        graph.add_edge("agent", "finalize")
-    else:
-        graph.add_edge("agent", "finalize")
-
-    graph.add_edge("finalize", END)
-
-    info("✅ [图构建] 极简 Agent 图构建完成")
+    info("✅ [Graph Builder] 极简 Agent 图构建完成（用 create_react_agent）")
     return graph