ailine/backend/app/main_graph/main_graph_builder.py

"""
极简 Agent 主图 - 简化版本！
因为完整的 ReAct 循环已经在 agent.py 里了！
"""

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.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 图（ReAct 循环在 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. Agent 节点（包含完整 ReAct 循环） ==========
    llm_with_tools = primary_model.bind_tools(ALL_TOOLS)
    agent_node_fn = create_agent_node(llm_with_tools, primary_model)

    # ========== 4. 完成节点 ==========
    async def finalize_node_simple(state: AgentState):
        info("[Finalize] 进入完成节点")
        return {}

    # ========== 5. 构建图 ==========
    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("finalize", finalize_node_simple)

    # ========== 6. 边的连接 ==========
    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", "finalize")
    graph.add_edge("finalize", END)

    info("✅ [Graph Builder] 简化 Agent 图构建完成（ReAct 在节点内）")
    return graph