80 lines
2.5 KiB
Python
80 lines
2.5 KiB
Python
"""
|
||
极简 Agent 主图 - 用 LangGraph 原生 create_react_agent + 记忆节点
|
||
"""
|
||
|
||
from langgraph.prebuilt import create_react_agent
|
||
from langgraph.graph import StateGraph, START, END
|
||
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
|
||
):
|
||
"""
|
||
构建包含记忆节点的 react agent 图
|
||
|
||
Args:
|
||
chat_services: 模型服务字典
|
||
mem0_client: 记忆客户端（可选）
|
||
max_steps: 最大步数限制
|
||
|
||
Returns:
|
||
编译好的 graph
|
||
"""
|
||
# 获取主模型
|
||
primary_model = chat_services.get("primary", next(iter(chat_services.values())))
|
||
|
||
# ========== 设置全局客户端 ==========
|
||
if mem0_client:
|
||
set_mem0_client(mem0_client)
|
||
|
||
# ========== 1. 初始化节点：重置步数 ==========
|
||
async def init_state_node(state: AgentState):
|
||
info("[Init State] 初始化状态，重置步数")
|
||
return {
|
||
"current_step": 0
|
||
}
|
||
|
||
# ========== 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}")
|
||
|
||
# ========== 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)
|
||
|
||
# 直接把 create_react_agent 的可运行对象作为节点
|
||
graph.add_node("agent", agent_runnable)
|
||
|
||
# ========== 边的连接 ==========
|
||
graph.add_edge(START, "init_state")
|
||
|
||
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")
|
||
|
||
graph.add_edge("memory_trigger", "agent")
|
||
graph.add_edge("agent", END)
|
||
|
||
info("✅ [Graph Builder] 极简 Agent 图构建完成（用 create_react_agent）")
|
||
return graph
|