ailine/backend/app/main_graph/subgraph_wrapper.py

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

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

from .state import MainGraphState, ErrorRecord, ErrorSeverity
from ..logger import info


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

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

    Returns: 包装后的节点函数
    """
    async def wrapped_node(state: MainGraphState, config: Optional[Dict[str, Any]] = None) -> MainGraphState:
        # 发送子图开始事件
        if config:
            try:
                from langchain_core.callbacks.manager import adispatch_custom_event
                callbacks = config.get("callbacks")
                if callbacks:
                    await adispatch_custom_event(
                        "react_reasoning",
                        {
                            "step": state.reasoning_step,
                            "action": f"{name}_subgraph_start",
                            "confidence": 1.0,
                            "reasoning": f"开始执行 {name} 子图..."
                        },
                        callbacks=callbacks
                    )
            except Exception as e:
                info(f"[{name}_subgraph] 无法发送开始事件: {e}")

        try:
            # 调用子图
            result = subgraph.invoke(state)

            # 更新主图状态
            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()
            })

            # 发送子图完成事件
            if config:
                try:
                    from langchain_core.callbacks.manager import adispatch_custom_event
                    callbacks = config.get("callbacks")
                    if callbacks:
                        await adispatch_custom_event(
                            "react_reasoning",
                            {
                                "step": state.reasoning_step,
                                "action": f"{name}_subgraph_complete",
                                "confidence": 1.0,
                                "reasoning": f"{name} 子图执行完成"
                            },
                            callbacks=callbacks
                        )
                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

            # 发送子图错误事件
            if config:
                try:
                    from langchain_core.callbacks.manager import adispatch_custom_event
                    callbacks = config.get("callbacks")
                    if callbacks:
                        await adispatch_custom_event(
                            "react_reasoning",
                            {
                                "step": state.reasoning_step,
                                "action": f"{name}_subgraph_error",
                                "confidence": 1.0,
                                "reasoning": f"{name} 子图执行失败: {str(e)}"
                            },
                            callbacks=callbacks
                        )
                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"
        ),
    }