Neo-ZQYY/apps/backend/app/trace/sse_wrapper.py

# -*- coding: utf-8 -*-
"""
SSE 流式响应追踪辅助函数

提供一组轻量函数，在 SSE event_generator 内部调用，
追踪流式响应全过程（SSE_START → AI_CALL → SSE_EVENT → SSE_END / AI_ERROR）。

所有函数在无活跃 TraceContext 时静默跳过，不影响业务逻辑。
SSE_EVENT span 每 10 个 token 记录一次，避免 span 爆炸。
"""

from __future__ import annotations

from datetime import datetime

from app.trace.context import (
    SpanType,
    TraceSpan,
    TraceType,
    get_current_trace,
)


def record_sse_start(
    endpoint: str,
    user_id: int | None = None,
    chat_id: str = "",
) -> None:
    """记录 SSE_START span：流开始。

    同时将 trace_type 切换为 "sse"。
    """
    ctx = get_current_trace()
    if ctx is None:
        return

    # 切换 trace 类型为 SSE
    ctx.trace_type = TraceType.SSE

    ctx.add_span(TraceSpan(
        span_type=SpanType.SSE_START,
        module="trace.sse_wrapper",
        function="record_sse_start",
        description_zh=f"SSE 流开始: endpoint={endpoint}, chat_id={chat_id}",
        description_en=f"SSE stream started: endpoint={endpoint}, chat_id={chat_id}",
        params={"endpoint": endpoint, "user_id": user_id, "chat_id": chat_id},
        result_summary="",
        duration_ms=0.0,
        timestamp=datetime.now().isoformat(),
    ))


def record_ai_call(
    app_id: str,
    prompt_length: int,
    session_id: str = "",
) -> None:
    """记录 AI_CALL span：DashScope API 调用。"""
    ctx = get_current_trace()
    if ctx is None:
        return

    ctx.add_span(TraceSpan(
        span_type=SpanType.AI_CALL,
        module="trace.sse_wrapper",
        function="record_ai_call",
        description_zh=f"AI 调用: app_id={app_id}, prompt 长度={prompt_length}",
        description_en=f"AI call: app_id={app_id}, prompt_length={prompt_length}",
        params={
            "app_id": app_id,
            "prompt_length": prompt_length,
            "session_id": session_id,
        },
        result_summary="",
        duration_ms=0.0,
        timestamp=datetime.now().isoformat(),
    ))


def record_sse_token(token_count: int, total_tokens: int) -> None:
    """记录 SSE_EVENT span：每 10 个 token 记录一次，避免 span 爆炸。

    仅当 total_tokens 是 10 的倍数时才记录 span。
    """
    ctx = get_current_trace()
    if ctx is None:
        return

    # 每 10 个 token 记录一次
    if total_tokens % 10 != 0:
        return

    ctx.add_span(TraceSpan(
        span_type=SpanType.SSE_EVENT,
        module="trace.sse_wrapper",
        function="record_sse_token",
        description_zh=f"SSE token 流: 本批 {token_count} token, 累计 {total_tokens}",
        description_en=f"SSE token stream: batch {token_count}, cumulative {total_tokens}",
        params={"token_count": token_count, "total_tokens": total_tokens},
        result_summary=f"cumulative={total_tokens}",
        duration_ms=0.0,
        timestamp=datetime.now().isoformat(),
    ))


def record_sse_end(
    total_tokens: int,
    total_duration_ms: float,
    completed: bool = True,
) -> None:
    """记录 SSE_END span：流结束。"""
    ctx = get_current_trace()
    if ctx is None:
        return

    ctx.add_span(TraceSpan(
        span_type=SpanType.SSE_END,
        module="trace.sse_wrapper",
        function="record_sse_end",
        description_zh=f"SSE 流结束: 总 token={total_tokens}, 耗时={total_duration_ms:.0f}ms, 完成={completed}",
        description_en=f"SSE stream ended: total_tokens={total_tokens}, duration={total_duration_ms:.0f}ms, completed={completed}",
        params={
            "total_tokens": total_tokens,
            "total_duration_ms": total_duration_ms,
            "completed": completed,
        },
        result_summary=f"tokens={total_tokens}, completed={completed}",
        duration_ms=total_duration_ms,
        timestamp=datetime.now().isoformat(),
    ))


def record_ai_error(error_type: str, message: str) -> None:
    """记录 AI_ERROR span：AI 调用失败。"""
    ctx = get_current_trace()
    if ctx is None:
        return

    ctx.add_span(TraceSpan(
        span_type=SpanType.AI_ERROR,
        module="trace.sse_wrapper",
        function="record_ai_error",
        description_zh=f"AI 调用失败: {error_type} - {message}",
        description_en=f"AI call failed: {error_type} - {message}",
        params={"error_type": error_type, "message": message},
        result_summary=f"{error_type}: {message}",
        duration_ms=0.0,
        timestamp=datetime.now().isoformat(),
    ))