112 lines
3.1 KiB
Python
112 lines
3.1 KiB
Python
"""
|
||
Agent 主图 - 标准 LangGraph 结构
|
||
将 ReAct 循环分离为独立的 agent/tools 节点
|
||
"""
|
||
|
||
from typing import Literal
|
||
from langgraph.graph import StateGraph, START, END
|
||
from backend.app.main_graph.state import AgentState
|
||
from backend.app.main_graph.nodes.memory_trigger import memory_trigger_node, set_mem0_client
|
||
from backend.app.main_graph.nodes.agent import create_agent_node
|
||
from backend.app.main_graph.nodes.tools import tools_node
|
||
from backend.app.main_graph.nodes.finalize import finalize_node
|
||
from backend.app.logger import info
|
||
|
||
|
||
def build_agent_graph(
|
||
chat_services: dict,
|
||
mem0_client=None,
|
||
max_steps: int = 10
|
||
):
|
||
"""
|
||
构建标准 Agent 图
|
||
|
||
节点:
|
||
- init_state: 初始化状态
|
||
- memory_trigger: 记忆触发
|
||
- agent: 单步推理
|
||
- tools: 工具执行
|
||
- finalize: 后处理
|
||
|
||
边:
|
||
- START -> init_state -> memory_trigger -> agent
|
||
- agent -> (条件边) -> tools 或 finalize
|
||
- tools -> agent (循环)
|
||
- finalize -> END
|
||
|
||
Args:
|
||
chat_services: 模型服务字典
|
||
mem0_client: 记忆客户端（可选）
|
||
max_steps: 最大步数限制
|
||
|
||
Returns:
|
||
构建好的 StateGraph（未编译）
|
||
"""
|
||
# 设置全局客户端
|
||
if mem0_client:
|
||
set_mem0_client(mem0_client)
|
||
|
||
# ========== 1. init_state 节点 ==========
|
||
async def init_state_node(state: AgentState):
|
||
info("[Init State] 初始化状态")
|
||
return {
|
||
"current_step": 0,
|
||
"max_steps": max_steps,
|
||
"tool_call_history": [],
|
||
"tool_result_history": [],
|
||
"tools_used": [],
|
||
"stop": False,
|
||
"stop_reason": "",
|
||
}
|
||
|
||
# ========== 2. Agent 节点 ==========
|
||
agent_node_fn = create_agent_node(chat_services)
|
||
|
||
# ========== 3. 条件边函数 ==========
|
||
def should_continue(state: AgentState) -> Literal["tools", "finalize"]:
|
||
"""根据 agent 节点输出决定下一步"""
|
||
# 手动停止标志
|
||
if getattr(state, "stop", False):
|
||
return "finalize"
|
||
|
||
# 检查是否有工具调用
|
||
last_msg = state.messages[-1]
|
||
if hasattr(last_msg, 'tool_calls') and last_msg.tool_calls:
|
||
return "tools"
|
||
|
||
return "finalize"
|
||
|
||
# ========== 4. 构建图 ==========
|
||
graph = StateGraph(AgentState)
|
||
|
||
# 添加节点
|
||
graph.add_node("init_state", init_state_node)
|
||
graph.add_node("memory_trigger", memory_trigger_node)
|
||
graph.add_node("agent", agent_node_fn)
|
||
graph.add_node("tools", tools_node)
|
||
graph.add_node("finalize", finalize_node)
|
||
|
||
# 边的连接
|
||
graph.add_edge(START, "init_state")
|
||
graph.add_edge("init_state", "memory_trigger")
|
||
graph.add_edge("memory_trigger", "agent")
|
||
|
||
# 条件边: agent -> tools 或 finalize
|
||
graph.add_conditional_edges(
|
||
"agent",
|
||
should_continue,
|
||
{
|
||
"tools": "tools",
|
||
"finalize": "finalize"
|
||
}
|
||
)
|
||
|
||
# 循环回边: tools -> agent
|
||
graph.add_edge("tools", "agent")
|
||
|
||
# 结束
|
||
graph.add_edge("finalize", END)
|
||
|
||
info("✅ [Graph Builder] 标准 Agent 图构建完成")
|
||
return graph
|