ailine/backend/app/main_graph/main_graph_builder.py

"""
主图构建器 - 构建整合后的完整主图
"""

from langgraph.graph import StateGraph, START, END
from typing import Dict, Any

from .state import MainGraphState
from .nodes.reasoning import react_reason_node
from .nodes.web_search import web_search_node
from .nodes.error_handling import error_handling_node
from .nodes.routing import init_state_node, route_by_reasoning, should_summarize
from .nodes.hybrid_router import (
    hybrid_router_node,
    route_from_hybrid_decision,
    check_fast_path_success,
)
from .nodes.fast_paths import (
    fast_chitchat_node,
    fast_rag_node,
    fast_tool_node,
)
from .nodes.llm_call import create_dynamic_llm_call_node
from .nodes.rag_nodes import rag_retrieve_node
from .nodes.retrieve_memory import create_retrieve_memory_node
from .nodes.memory_trigger import memory_trigger_node, set_mem0_client
from .nodes.summarize import create_summarize_node
from .nodes.finalize import finalize_node
from ..subgraphs.contact import build_contact_subgraph
from ..subgraphs.dictionary import build_dictionary_subgraph
from ..subgraphs.news_analysis import build_news_analysis_subgraph
from ..logger import info

from .subgraph_wrapper import create_subgraph_nodes


# ========== 主图构建 ==========

def build_react_main_graph(
    chat_services: dict,
    tools=None,
    mem0_client=None,
    use_hybrid_router: bool = True
) -> StateGraph:
    """
    构建整合后的完整主图（支持混合路由 + 动态模型选择）

    Args:
        chat_services: 模型名称 -> ChatModel 实例 的字典
        tools: 工具列表
        mem0_client: Mem0 客户端实例
        use_hybrid_router: 是否使用混合路由（快速路径 + React 循环）

    Returns:
        StateGraph: 构建好的图
    """
    # 创建图
    graph = StateGraph(MainGraphState)

    # 设置全局 mem0_client
    if mem0_client:
        set_mem0_client(mem0_client)

    # ========== 创建节点 ==========

    # LLM 调用节点
    llm_node = create_dynamic_llm_call_node(chat_services, tools or [])

    # 记忆节点
    retrieve_memory_node = None
    summarize_node = None
    if mem0_client:
        retrieve_memory_node = create_retrieve_memory_node(mem0_client)
        summarize_node = create_summarize_node(mem0_client)

    # 子图节点
    contact_graph = build_contact_subgraph()
    dictionary_graph = build_dictionary_subgraph()
    news_analysis_graph = build_news_analysis_subgraph()
    subgraph_nodes = create_subgraph_nodes(
        contact_graph, dictionary_graph, news_analysis_graph
    )

    # ========== 添加节点到图 ==========

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

    # 阶段 2: 初始化
    graph.add_node("init_state", init_state_node)

    # 阶段 3: 混合路由（可选）
    if use_hybrid_router:
        graph.add_node("hybrid_router", hybrid_router_node)
        graph.add_node("fast_chitchat", fast_chitchat_node)
        graph.add_node("fast_rag", fast_rag_node)
        graph.add_node("fast_tool", fast_tool_node)

    # 阶段 4: React 循环推理（始终保留）
    graph.add_node("react_reason", react_reason_node)
    graph.add_node("rag_retrieve", rag_retrieve_node)
    graph.add_node("web_search", web_search_node)
    graph.add_node("handle_error", error_handling_node)

    if llm_node is not None:
        graph.add_node("llm_call", llm_node)

    # 子图节点
    for node_name, node_func in subgraph_nodes.items():
        graph.add_node(node_name, node_func)

    # 阶段 5: 完成处理
    if summarize_node:
        graph.add_node("summarize", summarize_node)
    graph.add_node("finalize", finalize_node)

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

    # 阶段 1: 记忆检索
    _add_memory_edges(graph, retrieve_memory_node)

    # 阶段 2: 初始化
    graph.add_edge("memory_trigger", "init_state")

    # 阶段 3: 路由分支
    _add_routing_edges(graph, use_hybrid_router, llm_node)

    # 阶段 4: React 循环边
    _add_react_loop_edges(graph, subgraph_nodes)

    # 阶段 5: 完成阶段
    _add_finalize_edges(graph, llm_node, summarize_node)

    info(f"✅ [图构建] 整合后的完整主图构建完成（混合路由: {use_hybrid_router}）")

    return graph


def _add_memory_edges(graph: StateGraph, retrieve_memory_node) -> None:
    """添加记忆检索阶段的边"""
    if retrieve_memory_node:
        graph.add_edge(START, "retrieve_memory")
        graph.add_edge("retrieve_memory", "memory_trigger")
    else:
        graph.add_edge(START, "memory_trigger")


def _add_routing_edges(graph: StateGraph, use_hybrid_router: bool, llm_node) -> None:
    """添加路由阶段的边"""
    if use_hybrid_router:
        graph.add_edge("init_state", "hybrid_router")

        # 混合路由条件分支
        graph.add_conditional_edges(
            "hybrid_router",
            route_from_hybrid_decision,
            {
                "fast_chitchat": "fast_chitchat",
                "fast_rag": "fast_rag",
                "fast_tool": "fast_tool",
                "react_loop": "react_reason"
            }
        )

        # 快速路径的完成检查
        for fast_node in ["fast_chitchat", "fast_rag", "fast_tool"]:
            graph.add_conditional_edges(
                fast_node,
                check_fast_path_success,
                {
                    "llm_call": "llm_call",
                    "escalate": "react_reason"
                }
            )

        info(f"✅ [图构建] 混合路由模式已启用")
    else:
        graph.add_edge("init_state", "react_reason")
        info(f"✅ [图构建] 纯 React 模式")


def _add_react_loop_edges(graph: StateGraph, subgraph_nodes: Dict[str, Any]) -> None:
    """添加 React 循环阶段的边"""
    subgraph_names = list(subgraph_nodes.keys())

    # React 推理的条件分支
    graph.add_conditional_edges(
        "react_reason",
        route_by_reasoning,
        {
            "rag_retrieve": "rag_retrieve",
            "web_search": "web_search",
            **{name: name for name in subgraph_names},
            "handle_error": "handle_error",
            "llm_call": "llm_call"
        }
    )

    # 循环边（回到 react_reason）
    loop_back_nodes = ["rag_retrieve", "web_search", "handle_error"] + subgraph_names
    for node_name in loop_back_nodes:
        graph.add_edge(node_name, "react_reason")


def _add_finalize_edges(graph: StateGraph, llm_node, summarize_node) -> None:
    """添加完成阶段的边"""
    if llm_node is not None:
        if summarize_node:
            graph.add_conditional_edges(
                "llm_call",
                should_summarize,
                {
                    "summarize": "summarize",
                    "finalize": "finalize"
                }
            )
            graph.add_edge("summarize", "finalize")
        else:
            graph.add_edge("llm_call", "finalize")

    graph.add_edge("finalize", END)


# ========== 导出 ==========
__all__ = [
    "build_react_main_graph",
]