111 lines
3.4 KiB
Python
111 lines
3.4 KiB
Python
"""
|
||
极简 Agent 主图 - 自己的节点结构，更好控制流式
|
||
"""
|
||
|
||
from langgraph.graph import StateGraph, START, END
|
||
from langgraph.prebuilt import ToolNode
|
||
from ..state import AgentState
|
||
from ..nodes.memory_trigger import memory_trigger_node, set_mem0_client
|
||
from ..nodes.agent import create_agent_node
|
||
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
|
||
):
|
||
"""
|
||
构建包含记忆节点的 Agent 图
|
||
|
||
Args:
|
||
chat_services: 模型服务字典
|
||
mem0_client: 记忆客户端（可选）
|
||
max_steps: 最大步数限制
|
||
|
||
Returns:
|
||
构建好的 StateGraph（未编译）
|
||
"""
|
||
# 获取主模型
|
||
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,
|
||
"max_steps": max_steps
|
||
}
|
||
|
||
# ========== 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. 核心节点 ==========
|
||
llm_with_tools = primary_model.bind_tools(ALL_TOOLS)
|
||
agent_node_fn = create_agent_node(llm_with_tools, primary_model)
|
||
tool_node_fn = ToolNode(ALL_TOOLS)
|
||
|
||
# ========== 4. 条件边判断函数 ==========
|
||
def should_continue(state: AgentState):
|
||
"""判断是继续调用工具还是结束"""
|
||
messages = state.messages
|
||
last_message = messages[-1] if messages else None
|
||
|
||
if last_message and hasattr(last_message, 'tool_calls') and last_message.tool_calls:
|
||
return "tools"
|
||
|
||
return "finalize"
|
||
|
||
# ========== 5. 完成节点 ==========
|
||
async def finalize_node_simple(state: AgentState):
|
||
info("[Finalize] 进入完成节点")
|
||
return {}
|
||
|
||
# ========== 6. 构建图 ==========
|
||
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)
|
||
graph.add_node("agent", agent_node_fn)
|
||
graph.add_node("tools", tool_node_fn)
|
||
graph.add_node("finalize", finalize_node_simple)
|
||
|
||
# ========== 7. 边的连接 ==========
|
||
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_conditional_edges(
|
||
"agent",
|
||
should_continue,
|
||
{
|
||
"tools": "tools",
|
||
"finalize": "finalize"
|
||
}
|
||
)
|
||
|
||
graph.add_edge("tools", "agent")
|
||
graph.add_edge("finalize", END)
|
||
|
||
info("✅ [Graph Builder] 极简 Agent 图构建完成")
|
||
return graph
|