ailine/backend/app/main_graph/main_graph_builder.py

"""
极简 Agent 主图 - 回归 LangGraph 标准模式

架构：
START → [init_state] → [记忆] → [Agent] ⇄ [Tools] → [Finalize] → END
                                        ↑________↓
"""

from langgraph.graph import StateGraph, START, END
from langgraph.prebuilt import ToolNode
from langchain_core.runnables.config import RunnableConfig
from typing import Dict, Any, Optional

from .state import AgentState
from .nodes.memory_trigger import memory_trigger_node, set_mem0_client
from .nodes.summarize import create_summarize_node
from .nodes.agent import create_agent_node
from backend.app.tools import ALL_TOOLS
from backend.app.logger import info, warning


def build_agent_graph(
    chat_services: dict,
    mem0_client=None,
    max_steps: int = 10
) -> StateGraph:
    """
    构建极简 Agent 图

    Args:
        chat_services: 模型服务字典
        mem0_client: 记忆客户端（可选）
        max_steps: 最大步数限制

    Returns:
        StateGraph: 构建好的图
    """

    graph = StateGraph(AgentState)

    # ========== 设置全局客户端 ==========
    if mem0_client:
        set_mem0_client(mem0_client)

    # ========== 创建核心节点 ==========

    # 1. Agent 节点（绑定工具的 LLM）
    llm = chat_services.get("primary", list(chat_services.values())[0] if chat_services else None)
    if llm is None:
        raise ValueError("No LLM service provided")

    llm_with_tools = llm.bind_tools(ALL_TOOLS)
    agent_node = create_agent_node(llm_with_tools, llm)

    # 2. Tool 节点（LangGraph 内置）
    tool_node = ToolNode(ALL_TOOLS)

    # 3. 记忆/总结节点（保留现有）
    retrieve_memory_node = None
    summarize_node = None
    if mem0_client:
        try:
            from .nodes.retrieve_memory import create_retrieve_memory_node
            retrieve_memory_node = create_retrieve_memory_node(mem0_client)
            summarize_node = create_summarize_node(mem0_client)
        except Exception as e:
            info(f"[Graph Builder] 记忆节点初始化失败: {e}")

    # ========== 添加节点 ==========

    # 1. 初始化节点（重置步数）
    async def init_state_node(state: AgentState) -> Dict[str, Any]:
        """初始化状态：重置步数计数器"""
        info("[Init State] 初始化状态，重置步数")
        return {
            "current_step": 0
        }

    graph.add_node("init_state", init_state_node)

    # 2. 记忆阶段
    if retrieve_memory_node:
        graph.add_node("retrieve_memory", retrieve_memory_node)
    graph.add_node("memory_trigger", memory_trigger_node)

    # 3. 核心 Agent 循环
    graph.add_node("agent", agent_node)
    graph.add_node("tools", tool_node)

    # 4. 完成阶段
    if summarize_node:
        graph.add_node("summarize", summarize_node)

    # 简单的完成节点
    async def finalize_node_simple(state: AgentState, config: Optional[RunnableConfig] = None) -> Dict[str, Any]:
        """简单的完成节点，只发送完成事件"""
        info("[Finalize] 进入完成节点")

        try:
            from backend.app.main_graph.config import get_stream_writer
            writer = get_stream_writer()

            # 提取最后的回复
            final_reply = ""
            if state.messages:
                last_msg = state.messages[-1]
                final_reply = last_msg.content if hasattr(last_msg, "content") else str(last_msg)

            if writer and hasattr(writer, "__call__"):
                try:
                    writer({
                        "type": "custom",
                        "data": {
                            "type": "done",
                            "token_usage": state.last_token_usage,
                            "elapsed_time": state.last_elapsed_time,
                            "final_result": final_reply
                        }
                    })
                    info("🏁 [完成事件] 已发送完成事件")
                except Exception as e:
                    warning(f"⚠️ [完成事件] 发送失败 (非致命): {e}")
        except Exception as e:
            warning(f"⚠️ [完成事件] 处理失败 (非致命): {e}")

        return {}

    graph.add_node("finalize", finalize_node_simple)

    # ========== 添加边 ==========

    # 1. 初始化
    graph.add_edge(START, "init_state")

    # 2. 记忆阶段
    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")

    # 3. 进入 Agent
    graph.add_edge("memory_trigger", "agent")

    # 4. 核心循环：Agent ⇄ Tools
    def should_continue(state: AgentState) -> str:
        """判断是继续调用工具还是结束"""
        messages = state.messages
        last_message = messages[-1] if messages else None

        # 检查是否有 tool_calls
        if last_message and hasattr(last_message, "tool_calls") and last_message.tool_calls:
            return "tools"

        # 否则结束
        return "finalize"

    graph.add_conditional_edges(
        "agent",
        should_continue,
        {
            "tools": "tools",
            "finalize": "finalize"
        }
    )

    # Tools 执行完回到 Agent
    graph.add_edge("tools", "agent")

    # 5. 完成阶段
    if summarize_node:
        def should_summarize(state: AgentState) -> str:
            if state.turns_since_last_summary >= 5:
                return "summarize"
            return "finalize"

        # 总结逻辑暂简化：先 finalize
        graph.add_edge("agent", "finalize")
    else:
        graph.add_edge("agent", "finalize")

    graph.add_edge("finalize", END)

    info("✅ [图构建] 极简 Agent 图构建完成")
    return graph