314 lines
11 KiB
Python
314 lines
11 KiB
Python
"""
|
||
整合后的完整主图构建器 - 结合旧图和新图的优点
|
||
Main Graph Builder - Integrated Full Version (Old + New)
|
||
"""
|
||
|
||
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.nodes.react_nodes import (
|
||
init_state_node,
|
||
react_reason_node,
|
||
web_search_node,
|
||
error_handling_node,
|
||
route_by_reasoning
|
||
)
|
||
from app.main_graph.nodes.llm_call import create_llm_call_node
|
||
from app.main_graph.nodes.rag_nodes import rag_retrieve_node
|
||
from app.main_graph.nodes.retrieve_memory import create_retrieve_memory_node
|
||
from app.main_graph.nodes.memory_trigger import memory_trigger_node, set_mem0_client
|
||
from app.main_graph.nodes.summarize import create_summarize_node
|
||
from app.main_graph.nodes.finalize import finalize_node
|
||
from app.subgraphs.contact import build_contact_subgraph
|
||
from app.subgraphs.dictionary import build_dictionary_subgraph
|
||
from app.subgraphs.news_analysis import build_news_analysis_subgraph
|
||
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 wrap_subgraph_for_error_handling(subgraph, name: str):
|
||
"""
|
||
包装子图，使其错误能传递给主图
|
||
|
||
Args:
|
||
subgraph: 编译好的子图
|
||
name: 子图名称（用于错误标识）
|
||
|
||
Returns: 包装后的节点函数
|
||
"""
|
||
def wrapped_node(state: MainGraphState) -> MainGraphState:
|
||
try:
|
||
# 调用子图
|
||
result = subgraph.invoke(state)
|
||
|
||
# 更新主图状态
|
||
if name == "contact":
|
||
state.contact_result = result
|
||
elif name == "dictionary":
|
||
state.dictionary_result = result
|
||
elif name == "news_analysis":
|
||
state.news_result = result
|
||
|
||
# 标记成功
|
||
state.success = True
|
||
return state
|
||
|
||
except Exception as e:
|
||
# 捕获子图错误，传递给主图
|
||
from app.main_graph.state import ErrorRecord, ErrorSeverity
|
||
from datetime import datetime
|
||
|
||
error_record = ErrorRecord(
|
||
error_type=f"{name}SubgraphError",
|
||
error_message=str(e),
|
||
severity=ErrorSeverity.WARNING,
|
||
source=f"{name}_subgraph",
|
||
timestamp=datetime.now().isoformat(),
|
||
retry_count=0,
|
||
max_retries=1,
|
||
context={"user_query": state.user_query}
|
||
)
|
||
state.errors.append(error_record)
|
||
state.current_error = error_record
|
||
state.current_phase = "error_handling"
|
||
state.success = False
|
||
|
||
return state
|
||
|
||
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 (推理) ←──────────────────────┐
|
||
↓ │
|
||
条件路由 │
|
||
├─ rag_retrieve →─────────────────────────┤
|
||
├─ contact_subgraph →─────────────────────┤
|
||
├─ dictionary_subgraph →──────────────────┤
|
||
├─ news_analysis_subgraph →───────────────┤
|
||
├─ web_search →───────────────────────────┤
|
||
├─ handle_error → (重试或结束) ───────────┤
|
||
└─ llm_call (大模型调用) ←────────────────┘
|
||
↓
|
||
检查：需要总结？
|
||
├─ 是 → summarize (提交给Mem0存储) ← 来自旧图
|
||
└─ 否 → (跳过)
|
||
↓
|
||
finalize (发送完成事件) ← 来自旧图
|
||
↓
|
||
END
|
||
"""
|
||
# 创建图
|
||
graph = StateGraph(MainGraphState)
|
||
|
||
# 设置全局 mem0_client
|
||
if mem0_client:
|
||
set_global_mem0_client(mem0_client)
|
||
|
||
# 创建节点
|
||
llm_node = None
|
||
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", adapt_old_node_for_new_state(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)
|
||
graph.add_node("web_search", web_search_node)
|
||
graph.add_node("handle_error", error_handling_node)
|
||
|
||
if llm_node is not None:
|
||
graph.add_node("llm_call", llm_node)
|
||
|
||
# 子图节点
|
||
contact_graph = build_contact_subgraph()
|
||
dictionary_graph = build_dictionary_subgraph()
|
||
news_analysis_graph = build_news_analysis_subgraph()
|
||
|
||
graph.add_node(
|
||
"contact_subgraph",
|
||
wrap_subgraph_for_error_handling(contact_graph.compile(), "contact")
|
||
)
|
||
graph.add_node(
|
||
"dictionary_subgraph",
|
||
wrap_subgraph_for_error_handling(dictionary_graph.compile(), "dictionary")
|
||
)
|
||
graph.add_node(
|
||
"news_analysis_subgraph",
|
||
wrap_subgraph_for_error_handling(news_analysis_graph.compile(), "news_analysis")
|
||
)
|
||
|
||
# 第三阶段：完成处理（来自旧图）
|
||
if summarize_node:
|
||
graph.add_node("summarize", adapt_old_node_for_new_state(summarize_node))
|
||
graph.add_node("finalize", finalize_node)
|
||
|
||
# ========== 添加边 ==========
|
||
|
||
# 第一阶段：记忆检索
|
||
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,
|
||
{
|
||
"rag_retrieve": "rag_retrieve",
|
||
"web_search": "web_search",
|
||
"contact_subgraph": "contact_subgraph",
|
||
"dictionary_subgraph": "dictionary_subgraph",
|
||
"news_analysis_subgraph": "news_analysis_subgraph",
|
||
"handle_error": "handle_error",
|
||
"llm_call": "llm_call"
|
||
}
|
||
)
|
||
|
||
# 循环边：检索/搜索/子图/错误处理后 → 回到推理
|
||
graph.add_edge("rag_retrieve", "react_reason")
|
||
graph.add_edge("web_search", "react_reason")
|
||
graph.add_edge("contact_subgraph", "react_reason")
|
||
graph.add_edge("dictionary_subgraph", "react_reason")
|
||
graph.add_edge("news_analysis_subgraph", "react_reason")
|
||
graph.add_edge("handle_error", "react_reason")
|
||
|
||
# 第三阶段：llm_call 后进入完成处理
|
||
if llm_node is not None:
|
||
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("✅ [图构建] 整合后的完整主图构建完成")
|
||
|
||
return graph
|
||
|
||
|
||
# ========== 兼容性：保留旧的函数名 ==========
|
||
def build_main_graph() -> StateGraph:
|
||
"""
|
||
兼容性函数：旧代码调用 build_main_graph() 时返回 React 版本
|
||
"""
|
||
return build_react_main_graph()
|
||
|
||
|
||
# ========== 导出 ==========
|
||
__all__ = [
|
||
"build_react_main_graph",
|
||
"build_main_graph",
|
||
"wrap_subgraph_for_error_handling",
|
||
"set_global_mem0_client"
|
||
]
|