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

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

from typing import Any, Dict
from app.main_graph.config import get_stream_writer

# 本地模块
from app.main_graph.state import MessagesState
from app.utils.logging import log_state_change
from app.logger import info, error

from langchain_core.runnables.config import RunnableConfig


def _get_attr(state, attr_name, default=None):
    """通用方法：兼容 dict 和 dataclass 两种状态格式"""
    if isinstance(state, dict):
        return state.get(attr_name, default)
    else:
        return getattr(state, attr_name, default)


async def finalize_node(state, config: RunnableConfig) -> Dict[str, Any]:
    """
    完成事件节点 - 发送完成事件，包含token使用情况和耗时信息
    
    Args:
        state: 当前对话状态（兼容 dict 和 dataclass）
        config: 运行时配置
    
    Returns:
        空字典（完成节点，无状态更新）
    """
    log_state_change("finalize", state, "进入")

    try:
        # 获取流式写入器并发送完成事件
        writer = get_stream_writer()
        writer({
            "type": "custom", 
            "data": {
                "type": "done",
                "token_usage": _get_attr(state, "last_token_usage", {}),
                "elapsed_time": _get_attr(state, "last_elapsed_time", 0.0)
            }
        })
        info("🏁 [完成事件] 已发送完成事件，包含token使用情况和耗时信息")
    except Exception as e:
        error(f"❌ [完成事件] 发送完成事件时发生异常: {e}")
    
    log_state_change("finalize", state, "离开")
    return {}