ailine/backend/app/main_graph/subgraph_wrapper.py

"""
子图包装器 - 为子图添加错误处理和事件追踪
"""

from typing import Dict, Any, Optional
from datetime import datetime

from langchain_core.runnables.config import RunnableConfig

from .state import MainGraphState, ErrorRecord, ErrorSeverity
from backend.app.logger import info
from .nodes._utils import dispatch_custom_event, make_react_event


def wrap_subgraph_for_error_handling(subgraph, name: str):
    """
    包装子图，使其错误能传递给主图

    Args:
        subgraph: 编译好的子图
        name: 子图名称（用于错误标识）

    Returns: 包装后的节点函数
    """
    async def wrapped_node(state: MainGraphState, config: Optional[RunnableConfig] = None) -> MainGraphState:
        # 发送子图开始事件
        try:
            await dispatch_custom_event(
                "react_reasoning",
                make_react_event(
                    state.reasoning_step,
                    f"{name}_subgraph_start",
                    1.0,
                    f"开始执行 {name} 子图..."
                ),
                config
            )
        except Exception as e:
            info(f"[{name}_subgraph] 无法发送开始事件: {e}")

        try:
            # 调用子图（异步，传 config）
            result = await subgraph.ainvoke(state, config=config)

            # 更新主图状态
            subgraph_result = None
            if name == "contact":
                state.contact_result = result
                subgraph_result = result.get("final_result", "")
            elif name == "dictionary":
                state.dictionary_result = result
                subgraph_result = result.get("final_result", "")
            elif name == "news_analysis":
                state.news_result = result
                subgraph_result = result.get("final_result", "")

            # 设置最终结果
            if subgraph_result:
                state.final_result = subgraph_result
            else:
                state.final_result = "子图执行完成"

            # 标记成功
            state.success = True
            state.current_phase = "done"
            state.reasoning_history.append({
                "step": state.reasoning_step,
                "action": "subgraph_completed",
                "confidence": 1.0,
                "reasoning": f"{name}子图执行完成",
                "timestamp": datetime.now().isoformat()
            })

            # 发送子图完成事件
            try:
                await dispatch_custom_event(
                    "react_reasoning",
                    make_react_event(
                        state.reasoning_step,
                        f"{name}_subgraph_complete",
                        1.0,
                        f"{name} 子图执行完成"
                    ),
                    config
                )
            except Exception as e:
                info(f"[{name}_subgraph] 无法发送完成事件: {e}")

            return state

        except Exception as e:
            # 捕获子图错误，传递给主图
            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

            # 发送子图错误事件
            try:
                await dispatch_custom_event(
                    "react_reasoning",
                    make_react_event(
                        state.reasoning_step,
                        f"{name}_subgraph_error",
                        1.0,
                        f"{name} 子图执行失败: {str(e)}"
                    ),
                    config
                )
            except Exception as e:
                info(f"[{name}_subgraph] 无法发送错误事件: {e}")

            return state

    return wrapped_node


def create_subgraph_nodes(contact_graph, dictionary_graph, news_analysis_graph) -> Dict[str, Any]:
    """
    创建所有子图节点的字典

    Args:
        contact_graph: 联系人子图
        dictionary_graph: 词典子图
        news_analysis_graph: 新闻分析子图

    Returns:
        子图节点字典 {name: wrapped_node}
    """
    return {
        "contact_subgraph": wrap_subgraph_for_error_handling(
            contact_graph.compile(), "contact"
        ),
        "dictionary_subgraph": wrap_subgraph_for_error_handling(
            dictionary_graph.compile(), "dictionary"
        ),
        "news_analysis_subgraph": wrap_subgraph_for_error_handling(
            news_analysis_graph.compile(), "news_analysis"
        ),
    }