Aether/src/api/handlers/base/stream_telemetry.py

"""
流式遥测记录器 - 从 ChatHandlerBase 提取的统计记录逻辑

职责：
1. 记录流式请求的成功/失败统计
2. 更新 Usage 状态
3. 更新候选记录状态
"""

import asyncio
import time
from typing import Any, Dict, Optional

from sqlalchemy.orm import Session

from src.api.handlers.base.base_handler import MessageTelemetry
from src.api.handlers.base.stream_context import StreamContext
from src.config.settings import config
from src.core.logger import logger
from src.database import get_db
from src.models.database import ApiKey, User


class StreamTelemetryRecorder:
    """
    流式遥测记录器

    负责在流式请求完成后记录统计信息。
    从 ChatHandlerBase 中提取的 _record_stream_stats 逻辑。
    """

    def __init__(
        self,
        request_id: str,
        user_id: str,
        api_key_id: str,
        client_ip: str,
        format_id: str,
    ):
        """
        初始化遥测记录器

        Args:
            request_id: 请求 ID
            user_id: 用户 ID
            api_key_id: API Key ID
            client_ip: 客户端 IP
            format_id: API 格式标识
        """
        self.request_id = request_id
        self.user_id = user_id
        self.api_key_id = api_key_id
        self.client_ip = client_ip
        self.format_id = format_id

    async def record_stream_stats(
        self,
        ctx: StreamContext,
        original_headers: Dict[str, str],
        original_request_body: Dict[str, Any],
        start_time: float,
    ) -> None:
        """
        记录流式统计信息

        Args:
            ctx: 流式上下文
            original_headers: 原始请求头
            original_request_body: 原始请求体
            start_time: 请求开始时间 (time.time())
        """
        bg_db = None

        try:
            # 在流结束后计算响应时间，与首字时间使用相同的时间基准
            # 注意：不要把统计延迟（stream_stats_delay）算进响应时间里
            response_time_ms = int((time.time() - start_time) * 1000)

            await asyncio.sleep(config.stream_stats_delay)  # 等待流完全关闭

            if not ctx.provider_name:
                await self._update_usage_status_on_error(
                    response_time_ms=response_time_ms,
                    error_message="Provider name not available",
                )
                return

            db_gen = get_db()
            bg_db = next(db_gen)

            try:
                user = bg_db.query(User).filter(User.id == self.user_id).first()
                api_key_obj = bg_db.query(ApiKey).filter(ApiKey.id == self.api_key_id).first()

                if not user or not api_key_obj:
                    logger.warning(
                        f"[{self.request_id}] User or ApiKey not found, updating status directly"
                    )
                    await self._update_usage_status_directly(
                        bg_db,
                        status="completed" if ctx.is_success() else "failed",
                        response_time_ms=response_time_ms,
                        status_code=ctx.status_code,
                    )
                    return

                bg_telemetry = MessageTelemetry(
                    bg_db, user, api_key_obj, self.request_id, self.client_ip
                )

                actual_request_body = ctx.provider_request_body or original_request_body
                response_body = ctx.build_response_body(response_time_ms)

                if ctx.is_success():
                    await self._record_success(
                        bg_telemetry,
                        ctx,
                        original_headers,
                        actual_request_body,
                        response_body,
                        response_time_ms,
                    )
                else:
                    await self._record_failure(
                        bg_telemetry,
                        ctx,
                        original_headers,
                        actual_request_body,
                        response_body,
                        response_time_ms,
                    )

                # 更新候选记录状态
                await self._update_candidate_status(bg_db, ctx, response_time_ms)

            finally:
                if bg_db:
                    bg_db.close()

        except Exception as e:
            logger.exception("记录流式统计信息时出错")
            await self._update_usage_status_on_error(
                response_time_ms=response_time_ms,
                error_message=f"记录统计信息失败: {str(e)[:200]}",
            )

    async def _record_success(
        self,
        telemetry: MessageTelemetry,
        ctx: StreamContext,
        original_headers: Dict[str, str],
        actual_request_body: Dict[str, Any],
        response_body: Dict[str, Any],
        response_time_ms: int,
    ) -> None:
        """记录成功的请求"""
        await telemetry.record_success(
            provider=ctx.provider_name or "unknown",
            model=ctx.model,
            input_tokens=ctx.input_tokens,
            output_tokens=ctx.output_tokens,
            response_time_ms=response_time_ms,
            first_byte_time_ms=ctx.first_byte_time_ms,  # 传递首字时间
            status_code=ctx.status_code,
            request_headers=original_headers,
            request_body=actual_request_body,
            response_headers=ctx.response_headers,
            response_body=response_body,
            cache_creation_tokens=ctx.cache_creation_tokens,
            cache_read_tokens=ctx.cached_tokens,
            is_stream=True,
            provider_request_headers=ctx.provider_request_headers,
            api_format=ctx.api_format,
            provider_id=ctx.provider_id,
            provider_endpoint_id=ctx.endpoint_id,
            provider_api_key_id=ctx.key_id,
            target_model=ctx.mapped_model,
        )

        logger.debug(f"{self.format_id} 流式响应完成")
        logger.info(ctx.get_log_summary(self.request_id, response_time_ms))

    async def _record_failure(
        self,
        telemetry: MessageTelemetry,
        ctx: StreamContext,
        original_headers: Dict[str, str],
        actual_request_body: Dict[str, Any],
        response_body: Dict[str, Any],
        response_time_ms: int,
    ) -> None:
        """记录失败的请求"""
        await telemetry.record_failure(
            provider=ctx.provider_name or "unknown",
            model=ctx.model,
            response_time_ms=response_time_ms,
            status_code=ctx.status_code,
            error_message=ctx.error_message or f"HTTP {ctx.status_code}",
            request_headers=original_headers,
            request_body=actual_request_body,
            is_stream=True,
            api_format=ctx.api_format,
            provider_request_headers=ctx.provider_request_headers,
            input_tokens=ctx.input_tokens,
            output_tokens=ctx.output_tokens,
            cache_creation_tokens=ctx.cache_creation_tokens,
            cache_read_tokens=ctx.cached_tokens,
            response_body=response_body,
            target_model=ctx.mapped_model,
        )

        logger.debug(f"{self.format_id} 流式响应中断")
        log_summary = ctx.get_log_summary(self.request_id, response_time_ms)
        # 对于失败日志，添加缓存信息
        logger.info(f"{log_summary} cache:{ctx.cached_tokens}")

    async def _update_candidate_status(
        self,
        db: Session,
        ctx: StreamContext,
        response_time_ms: int,
    ) -> None:
        """更新候选记录状态"""
        if not ctx.attempt_id:
            return

        from src.services.request.candidate import RequestCandidateService

        if ctx.is_success():
            RequestCandidateService.mark_candidate_success(
                db=db,
                candidate_id=ctx.attempt_id,
                status_code=ctx.status_code,
                latency_ms=response_time_ms,
                extra_data={
                    "stream_completed": True,
                    "data_count": ctx.data_count,
                },
            )
        else:
            error_type = "client_disconnected" if ctx.status_code == 499 else "stream_error"
            RequestCandidateService.mark_candidate_failed(
                db=db,
                candidate_id=ctx.attempt_id,
                error_type=error_type,
                error_message=ctx.error_message or f"HTTP {ctx.status_code}",
                status_code=ctx.status_code,
                latency_ms=response_time_ms,
                extra_data={
                    "stream_completed": False,
                    "data_count": ctx.data_count,
                },
            )

    async def _update_usage_status_on_error(
        self,
        response_time_ms: int,
        error_message: str,
    ) -> None:
        """在记录失败时更新 Usage 状态"""
        try:
            db_gen = get_db()
            error_db = next(db_gen)
            try:
                await self._update_usage_status_directly(
                    error_db,
                    status="failed",
                    response_time_ms=response_time_ms,
                    status_code=500,
                    error_message=error_message,
                )
            finally:
                error_db.close()
        except Exception as inner_e:
            logger.error(f"[{self.request_id}] 更新 Usage 状态失败: {inner_e}")

    async def _update_usage_status_directly(
        self,
        db: Session,
        status: str,
        response_time_ms: int,
        status_code: int = 200,
        error_message: Optional[str] = None,
    ) -> None:
        """直接更新 Usage 表的状态字段"""
        try:
            from src.models.database import Usage

            usage = db.query(Usage).filter(Usage.request_id == self.request_id).first()
            if usage:
                setattr(usage, "status", status)
                setattr(usage, "status_code", status_code)
                setattr(usage, "response_time_ms", response_time_ms)
                if error_message:
                    setattr(usage, "error_message", error_message)
                db.commit()
                logger.debug(f"[{self.request_id}] Usage 状态已更新: {status}")
        except Exception as e:
            logger.error(f"[{self.request_id}] 直接更新 Usage 状态失败: {e}")