ailine/backend/app/main_graph/nodes/finalize.py

"""
完成事件节点模块
负责发送完成事件，包含token使用情况和耗时信息
"""

from typing import Any, Dict

# 本地模块
from ...main_graph.state import MainGraphState
from ...utils.logging import log_state_change
from ...logger import info, warning

from langchain_core.runnables.config import RunnableConfig


async def finalize_node(state: MainGraphState, config: RunnableConfig) -> Dict[str, Any]:
    """
    完成事件节点 - 发送完成事件，包含token使用情况和耗时信息

    Args:
        state: 当前对话状态
        config: 运行时配置

    Returns:
        更新后的状态（包含 final_result）
    """
    log_state_change("finalize", state, "进入")

    # 确保 final_result 被传递出去
    result = {
        "final_result": state.final_result,
        "success": state.success,
        "current_phase": "done"
    }

    try:
        # 获取流式写入器并发送完成事件
        from backend.app.main_graph.config import get_stream_writer
        writer = get_stream_writer()

        # 只在 writer 存在且不是 noop 时才发送
        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": state.final_result
                    }
                })
                info("🏁 [完成事件] 已发送完成事件")
            except Exception as e:
                warning(f"⚠️ [完成事件] 发送完成事件失败 (非致命): {e}")
    except Exception as e:
        warning(f"⚠️ [完成事件] 处理失败 (非致命): {e}")

    log_state_change("finalize", state, "离开")
    return result