ailine/backend/app/main_graph/main_graph_builder.py

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