ailine/backend/app/main_graph/main_graph_builder.py

"""
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