ailine/backend/app/main_graph/utils/subgraph_builder.py

"""
React 模式主图构建器 - 完整循环推理版本
Main Graph Builder - Full React Mode with Loop Reasoning
"""

from app.main_graph.graph import StateGraph, START, END
from typing import Dict, Any

from app.main_graph.state import MainGraphState, CurrentAction
from .react_nodes import (
    init_state_node,
    react_reason_node,
    error_handling_node,
    final_response_node,
    route_by_reasoning
)
from .rag_nodes import rag_retrieve_node
from app.subgraphs.contact import build_contact_subgraph
from app.subgraphs.dictionary import build_dictionary_subgraph
from app.subgraphs.news_analysis import build_news_analysis_subgraph


# ========== 子图包装器（处理子图错误传递） ==========
def wrap_subgraph_for_error_handling(subgraph, name: str):
    """
    包装子图，使其错误能传递给主图
    
    Args:
        subgraph: 编译好的子图
        name: 子图名称（用于错误标识）
    
    Returns: 包装后的节点函数
    """
    def wrapped_node(state: MainGraphState) -> MainGraphState:
        try:
            # 调用子图
            result = subgraph.invoke(state)
            
            # 更新主图状态
            if name == "contact":
                state.contact_result = result
            elif name == "dictionary":
                state.dictionary_result = result
            elif name == "news_analysis":
                state.news_result = result
            
            # 标记成功
            state.success = True
            return state
            
        except Exception as e:
            # 捕获子图错误，传递给主图
            from app.main_graph.state import ErrorRecord, ErrorSeverity
            from datetime import datetime
            
            error_record = ErrorRecord(
                error_type=f"{name}SubgraphError",
                error_message=str(e),
                severity=ErrorSeverity.WARNING,
                source=f"{name}_subgraph",
                timestamp=datetime.now().isoformat(),
                retry_count=0,
                max_retries=1,
                context={"user_query": state.user_query}
            )
            state.errors.append(error_record)
            state.current_error = error_record
            state.current_phase = "error_handling"
            state.success = False
            
            return state
    
    return wrapped_node


# ========== 主图构建 ==========
def build_react_main_graph() -> StateGraph:
    """
    构建完整的 React 模式主图
    
    流程：
    START
      ↓
    init_state (初始化)
      ↓
    react_reason (推理) ←──────────────┐
      ↓                                │
    条件路由                          │
      ├─→ rag_retrieve →───────────────┤
      ├─→ contact_subgraph →───────────┤
      ├─→ dictionary_subgraph →────────┤
      ├─→ news_analysis_subgraph →─────┤
      ├─→ handle_error → (重试或结束) ──┤
      └─→ final_response
      ↓
    END
    """
    # 创建图
    graph = StateGraph(MainGraphState)
    
    # ========== 添加节点 ==========
    
    # 1. 初始化节点
    graph.add_node("init_state", init_state_node)
    
    # 2. React 推理节点
    graph.add_node("react_reason", react_reason_node)
    
    # 3. RAG 检索节点
    graph.add_node("rag_retrieve", rag_retrieve_node)
    
    # 4. 错误处理节点
    graph.add_node("handle_error", error_handling_node)
    
    # 5. 最终回答节点
    graph.add_node("final_response", final_response_node)
    
    # ========== 添加子图节点 ==========
    
    # 构建并包装子图（带错误处理）
    contact_graph = build_contact_subgraph()
    dictionary_graph = build_dictionary_subgraph()
    news_analysis_graph = build_news_analysis_subgraph()
    
    graph.add_node(
        "contact_subgraph", 
        wrap_subgraph_for_error_handling(contact_graph.compile(), "contact")
    )
    graph.add_node(
        "dictionary_subgraph", 
        wrap_subgraph_for_error_handling(dictionary_graph.compile(), "dictionary")
    )
    graph.add_node(
        "news_analysis_subgraph", 
        wrap_subgraph_for_error_handling(news_analysis_graph.compile(), "news_analysis")
    )
    
    # ========== 添加边 ==========
    
    # 1. START → init_state
    graph.add_edge(START, "init_state")
    
    # 2. init_state → react_reason
    graph.add_edge("init_state", "react_reason")
    
    # 3. 条件路由：react_reason → 各分支
    graph.add_conditional_edges(
        "react_reason",
        route_by_reasoning,
        {
            # 检索分支 → 检索后回到推理
            "rag_retrieve": "rag_retrieve",
            
            # 子图分支 → 子图后回到推理
            "contact_subgraph": "contact_subgraph",
            "dictionary_subgraph": "dictionary_subgraph",
            "news_analysis_subgraph": "news_analysis_subgraph",
            
            # 错误处理分支
            "handle_error": "handle_error",
            
            # 最终回答分支
            "final_response": "final_response",
        }
    )
    
    # 4. 循环边：检索/子图/错误处理 后 → 回到推理
    graph.add_edge("rag_retrieve", "react_reason")
    graph.add_edge("contact_subgraph", "react_reason")
    graph.add_edge("dictionary_subgraph", "react_reason")
    graph.add_edge("news_analysis_subgraph", "react_reason")
    graph.add_edge("handle_error", "react_reason")  # 错误处理后可能重试
    
    # 5. 最终边：final_response → END
    graph.add_edge("final_response", END)
    
    return graph


# ========== 兼容性：保留旧的函数名 ==========
def build_main_graph() -> StateGraph:
    """
    兼容性函数：旧代码调用 build_main_graph() 时返回 React 版本
    """
    return build_react_main_graph()


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