彻底重构状态系统:整合所有旧状态到 MainGraphState,修复所有节点
Some checks failed
构建并部署 AI Agent 服务 / deploy (push) Failing after 6m35s

This commit is contained in:
2026-05-01 23:20:31 +08:00
parent 9a58eb8e6d
commit 9386b9fa7a
7 changed files with 155 additions and 193 deletions

View File

@@ -1,13 +1,12 @@
"""
整合后的完整主图构建器 - 结合旧图和新图的优点
Main Graph Builder - Integrated Full Version (Old + New)
整合后的完整主图构建器 - 所有节点都直接操作 MainGraphState
"""
from app.main_graph.graph import StateGraph, START, END
from typing import Dict, Any, Optional
from langchain_core.runnables.config import RunnableConfig
from app.main_graph.state import MainGraphState, CurrentAction, MessagesState
from app.main_graph.state import MainGraphState
from app.main_graph.nodes.react_nodes import (
init_state_node,
react_reason_node,
@@ -28,16 +27,21 @@ from app.memory.mem0_client import Mem0Client
from app.logger import info, debug
# ========== 全局变量(用于传递 mem0_client==========
# 这样就不用改旧节点的签名了
_global_mem0_client: Optional[Mem0Client] = None
def set_global_mem0_client(client: Mem0Client):
"""设置全局的 mem0_client"""
global _global_mem0_client
_global_mem0_client = client
set_mem0_client(client) # 同时设置给 memory_trigger_node
# ========== 检查是否需要总结 ==========
def should_summarize(state: MainGraphState) -> str:
"""
检查是否需要总结对话(对话足够长时)
Args:
state: 当前图状态
Returns:
"summarize""finalize"
"""
if state.turns_since_last_summary >= 5: # 每5轮对话总结一次
return "summarize"
else:
return "finalize"
# ========== 子图包装器(处理子图错误传递)==========
@@ -93,65 +97,18 @@ def wrap_subgraph_for_error_handling(subgraph, name: str):
return wrapped_node
# ========== 检查是否需要总结 ==========
def should_summarize(state: MainGraphState) -> str:
"""
检查是否需要总结对话(对话足够长时)
Args:
state: 当前图状态
Returns:
"summarize""finalize"
"""
messages = getattr(state, 'messages', [])
if len(messages) >= 4:
return "summarize"
else:
return "finalize"
# ========== 兼容层:让旧节点工作在新状态上 ==========
def adapt_old_node_for_new_state(old_node):
"""
适配旧节点(期望 MessagesState到新状态 MainGraphState
Args:
old_node: 旧节点函数
Returns: 适配后的节点函数
"""
async def adapted_node(state: MainGraphState, config: RunnableConfig) -> Dict[str, Any]:
# 把 MainGraphState 转换为 MessagesState旧节点期望的格式
old_state: MessagesState = {
"messages": state.messages,
"llm_calls": getattr(state, 'llm_calls', 0),
"memory_context": getattr(state, 'memory_context', ""),
"system_prompt": getattr(state, 'system_prompt', "")
}
# 调用旧节点
result = await old_node(old_state, config)
# 把结果更新回 MainGraphState
if "memory_context" in result:
state.memory_context = result["memory_context"]
if "llm_calls" in result:
state.llm_calls = result["llm_calls"]
return result
return adapted_node
# ========== 主图构建 ==========
def build_react_main_graph(llm=None, tools=None, mem0_client=None) -> StateGraph:
"""
构建整合后的完整主图(简化版:先让系统工作起来)
构建整合后的完整主图
完整流程:
START
retrieve_memory (从Mem0检索长期记忆)
memory_trigger (记忆触发器)
init_state (初始化)
react_reason (推理) ←───────────────────────┐
@@ -165,6 +122,10 @@ def build_react_main_graph(llm=None, tools=None, mem0_client=None) -> StateGraph
├─ handle_error → (重试或结束) ────────────┤
└─ llm_call (大模型调用) ←────────────────┘
检查:需要总结吗?
├─ 是 → summarize (提交给Mem0存储)
└─ 否 → (跳过)
finalize (发送完成事件)
END
@@ -172,7 +133,7 @@ def build_react_main_graph(llm=None, tools=None, mem0_client=None) -> StateGraph
# 创建图
graph = StateGraph(MainGraphState)
# 设置全局 mem0_client (暂时不用记忆功能)
# 设置全局 mem0_client
if mem0_client:
set_global_mem0_client(mem0_client)
@@ -181,8 +142,20 @@ def build_react_main_graph(llm=None, tools=None, mem0_client=None) -> StateGraph
if llm is not None:
llm_node = create_llm_call_node(llm, tools or [])
retrieve_memory_node = None
summarize_node = None
if mem0_client:
retrieve_memory_node = create_retrieve_memory_node(mem0_client)
summarize_node = create_summarize_node(mem0_client)
# ========== 添加节点 ==========
# 简化:先不用记忆检索相关节点
# 第一阶段:记忆检索
if retrieve_memory_node:
graph.add_node("retrieve_memory", retrieve_memory_node)
graph.add_node("memory_trigger", memory_trigger_node)
# 第二阶段React 循环推理
graph.add_node("init_state", init_state_node)
graph.add_node("react_reason", react_reason_node)
graph.add_node("rag_retrieve", rag_retrieve_node)
@@ -210,15 +183,25 @@ def build_react_main_graph(llm=None, tools=None, mem0_client=None) -> StateGraph
wrap_subgraph_for_error_handling(news_analysis_graph.compile(), "news_analysis")
)
# 完成节点
# 第三阶段:完成处理
if summarize_node:
graph.add_node("summarize", summarize_node)
graph.add_node("finalize", finalize_node)
# ========== 添加边 ==========
# 简化:直接从 START 到 init_state
graph.add_edge(START, "init_state")
# 第一阶段:记忆检索
if retrieve_memory_node:
graph.add_edge(START, "retrieve_memory")
graph.add_edge("retrieve_memory", "memory_trigger")
else:
graph.add_edge(START, "memory_trigger")
# 进入第二阶段
graph.add_edge("memory_trigger", "init_state")
graph.add_edge("init_state", "react_reason")
# 条件路由
# 第二阶段React 循环推理
graph.add_conditional_edges(
"react_reason",
route_by_reasoning,
@@ -241,14 +224,27 @@ def build_react_main_graph(llm=None, tools=None, mem0_client=None) -> StateGraph
graph.add_edge("news_analysis_subgraph", "react_reason")
graph.add_edge("handle_error", "react_reason")
# llm_call 之后直接到 finalize
# 第三阶段:llm_call 后进入完成处理
if llm_node is not None:
graph.add_edge("llm_call", "finalize")
if summarize_node:
# 检查是否需要总结
graph.add_conditional_edges(
"llm_call",
should_summarize,
{
"summarize": "summarize",
"finalize": "finalize"
}
)
graph.add_edge("summarize", "finalize")
else:
# 没有 summarize 节点,直接 finalize
graph.add_edge("llm_call", "finalize")
# 完成
graph.add_edge("finalize", END)
info("✅ [图构建] 整合后的完整主图构建完成(简化版)")
info("✅ [图构建] 整合后的完整主图构建完成")
return graph
@@ -265,6 +261,5 @@ def build_main_graph() -> StateGraph:
__all__ = [
"build_react_main_graph",
"build_main_graph",
"wrap_subgraph_for_error_handling",
"set_global_mem0_client"
]
"wrap_subgraph_for_error_handling"
]