ailine/app/nodes/finalize.py

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

from typing import Any, Dict
from langgraph.config import get_stream_writer

# 本地模块
from app.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

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

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