Aether/_deprecated_py_src/services/request/candidate.py

"""
请求候选记录服务 - 管理候选队列
"""

import uuid
from datetime import datetime, timezone

from sqlalchemy.orm import Session

from src.core.batch_committer import get_batch_committer
from src.core.logger import logger
from src.models.database import RequestCandidate


class RequestCandidateService:
    """请求候选记录服务"""

    @staticmethod
    def _persist_candidate_update(db: Session, *, immediate: bool) -> None:
        if immediate:
            db.commit()
            return
        db.flush()
        get_batch_committer().mark_dirty(db)

    @staticmethod
    def create_candidate(
        db: Session,
        request_id: str,
        candidate_index: int,
        candidate_id: str | None = None,
        retry_index: int = 0,  # 新增：重试序号
        user_id: str | None = None,
        api_key_id: str | None = None,
        username: str | None = None,
        api_key_name: str | None = None,
        provider_id: str | None = None,
        endpoint_id: str | None = None,
        key_id: str | None = None,
        status: str = "available",
        skip_reason: str | None = None,
        is_cached: bool = False,
        extra_data: dict | None = None,
        required_capabilities: dict | None = None,
    ) -> RequestCandidate:
        """
        创建候选记录

        Args:
            db: 数据库会话
            request_id: 请求ID
            candidate_index: 候选序号
            retry_index: 重试序号（从0开始）
            user_id: 用户ID
            api_key_id: API Key ID
            username: 用户名快照
            api_key_name: API Key 名称快照
            provider_id: Provider ID
            endpoint_id: Endpoint ID
            key_id: API Key ID
            status: 候选状态 ('available', 'used', 'skipped', 'success', 'failed')
            skip_reason: 跳过原因
            is_cached: 是否为缓存亲和性候选
            extra_data: 额外数据
            required_capabilities: 请求需要的能力标签
        """
        candidate = RequestCandidate(
            id=str(candidate_id or uuid.uuid4()),
            request_id=request_id,
            candidate_index=candidate_index,
            retry_index=retry_index,  # 新增
            user_id=user_id,
            api_key_id=api_key_id,
            username=username,
            api_key_name=api_key_name,
            provider_id=provider_id,
            endpoint_id=endpoint_id,
            key_id=key_id,
            status=status,
            skip_reason=skip_reason,
            is_cached=is_cached,
            extra_data=extra_data or {},
            required_capabilities=required_capabilities,
            created_at=datetime.now(timezone.utc),
        )
        db.add(candidate)
        db.flush()  # 只flush，不立即 commit
        # 标记为批量提交（非关键数据，可延迟）
        get_batch_committer().mark_dirty(db)
        return candidate

    @staticmethod
    def mark_candidate_started(db: Session, candidate_id: str) -> None:
        """
        标记候选开始执行

        Args:
            db: 数据库会话
            candidate_id: 候选ID
        """
        candidate = db.query(RequestCandidate).filter(RequestCandidate.id == candidate_id).first()
        if candidate:
            candidate.status = "pending"
            candidate.started_at = datetime.now(timezone.utc)
            # 中间态改为 flush：最终 success/failed 仍会立即提交，
            # 但开始执行这一跳不再单独制造一次事务往返。
            RequestCandidateService._persist_candidate_update(db, immediate=False)

    @staticmethod
    def update_candidate_status(db: Session, candidate_id: str, status: str) -> None:
        """
        更新候选状态（通用方法）

        Args:
            db: 数据库会话
            candidate_id: 候选ID
            status: 新状态（pending, available, success, failed, skipped）
        """
        candidate = db.query(RequestCandidate).filter(RequestCandidate.id == candidate_id).first()
        if candidate:
            candidate.status = status
            # 如果状态变更为 pending，记录开始时间
            if status == "pending" and not candidate.started_at:
                candidate.started_at = datetime.now(timezone.utc)
            RequestCandidateService._persist_candidate_update(
                db, immediate=status not in {"pending", "streaming"}
            )

    @staticmethod
    def mark_candidate_streaming(
        db: Session,
        candidate_id: str,
        concurrent_requests: int | None = None,
    ) -> None:
        """
        标记候选为流式传输中

        用于流式请求：连接建立成功后，流开始传输时调用。
        此时请求尚未完成，需要等流传输完毕后再调用 mark_candidate_success。

        注意：streaming 阶段不设置 status_code，最终状态码由
        mark_candidate_success / mark_candidate_failed 在流结束时写入。

        Args:
            db: 数据库会话
            candidate_id: 候选ID
            concurrent_requests: 并发请求数
        """
        candidate = db.query(RequestCandidate).filter(RequestCandidate.id == candidate_id).first()
        if candidate:
            candidate.status = "streaming"
            candidate.concurrent_requests = concurrent_requests
            # streaming 状态不设置 finished_at 和 status_code，因为请求还在进行中
            RequestCandidateService._persist_candidate_update(db, immediate=False)

    @staticmethod
    def mark_candidate_success(
        db: Session,
        candidate_id: str,
        status_code: int,
        latency_ms: int,
        concurrent_requests: int | None = None,
        extra_data: dict | None = None,
    ) -> None:
        """
        标记候选执行成功

        Args:
            db: 数据库会话
            candidate_id: 候选ID
            status_code: HTTP 状态码
            latency_ms: 延迟（毫秒）
            concurrent_requests: 并发请求数
            extra_data: 额外数据
        """
        candidate = db.query(RequestCandidate).filter(RequestCandidate.id == candidate_id).first()
        if candidate:
            candidate.status = "success"
            candidate.status_code = status_code
            candidate.latency_ms = latency_ms
            candidate.concurrent_requests = concurrent_requests
            candidate.finished_at = datetime.now(timezone.utc)
            # 成功时清空错误字段（可能是整流重试后成功，之前记录过错误）
            candidate.error_type = None
            candidate.error_message = None
            if extra_data:
                candidate.extra_data = {**(candidate.extra_data or {}), **extra_data}
            # 关键状态更新：立即提交，不使用批量提交
            # 原因：前端需要实时看到请求成功/失败状态
            db.commit()

    @staticmethod
    def mark_candidate_failed(
        db: Session,
        candidate_id: str,
        error_type: str,
        error_message: str,
        status_code: int | None = None,
        latency_ms: int | None = None,
        concurrent_requests: int | None = None,
        extra_data: dict | None = None,
    ) -> None:
        """
        标记候选执行失败

        Args:
            db: 数据库会话
            candidate_id: 候选ID
            error_type: 错误类型
            error_message: 错误消息
            status_code: HTTP 状态码（如果有）
            latency_ms: 延迟（毫秒）
            concurrent_requests: 并发请求数
            extra_data: 额外数据
        """
        candidate = db.query(RequestCandidate).filter(RequestCandidate.id == candidate_id).first()
        if candidate:
            candidate.status = "failed"
            candidate.error_type = error_type
            candidate.error_message = error_message
            candidate.status_code = status_code
            candidate.latency_ms = latency_ms
            candidate.concurrent_requests = concurrent_requests
            candidate.finished_at = datetime.now(timezone.utc)
            if extra_data:
                candidate.extra_data = {**(candidate.extra_data or {}), **extra_data}
            # 关键状态更新：立即提交，不使用批量提交
            # 原因：前端需要实时看到请求成功/失败状态
            db.commit()

    @staticmethod
    def mark_candidate_cancelled(
        db: Session,
        candidate_id: str,
        status_code: int = 499,
        latency_ms: int | None = None,
        concurrent_requests: int | None = None,
        extra_data: dict | None = None,
    ) -> None:
        """
        标记候选被客户端取消

        客户端主动断开连接不算系统失败，使用 cancelled 状态。

        Args:
            db: 数据库会话
            candidate_id: 候选ID
            status_code: HTTP 状态码（通常是 499）
            latency_ms: 延迟（毫秒）
            concurrent_requests: 并发请求数
            extra_data: 额外数据
        """
        candidate = db.query(RequestCandidate).filter(RequestCandidate.id == candidate_id).first()
        if candidate:
            candidate.status = "cancelled"
            candidate.status_code = status_code
            candidate.latency_ms = latency_ms
            candidate.concurrent_requests = concurrent_requests
            candidate.finished_at = datetime.now(timezone.utc)
            if extra_data:
                candidate.extra_data = {**(candidate.extra_data or {}), **extra_data}
            db.commit()

    @staticmethod
    def mark_candidate_skipped(
        db: Session,
        candidate_id: str,
        skip_reason: str | None = None,
        *,
        status_code: int | None = None,
        concurrent_requests: int | None = None,
        extra_data: dict | None = None,
    ) -> None:
        """
        标记候选为已跳过

        Args:
            db: 数据库会话
            candidate_id: 候选ID
            skip_reason: 跳过原因
            status_code: HTTP 状态码（可选）
            concurrent_requests: 并发请求数（这里实际记录 RPM 计数）
            extra_data: 额外数据（合并写入）
        """
        candidate = db.query(RequestCandidate).filter(RequestCandidate.id == candidate_id).first()
        if candidate:
            candidate.status = "skipped"
            candidate.skip_reason = skip_reason
            candidate.finished_at = datetime.now(timezone.utc)

            if status_code is not None:
                candidate.status_code = int(status_code)
            if concurrent_requests is not None:
                candidate.concurrent_requests = int(concurrent_requests)

            if extra_data:
                base = candidate.extra_data if isinstance(candidate.extra_data, dict) else {}
                candidate.extra_data = {**base, **extra_data}

            db.flush()  # 只 flush，不立即 commit
            get_batch_committer().mark_dirty(db)

    @staticmethod
    def get_candidates_by_request_id(db: Session, request_id: str) -> list[RequestCandidate]:
        """
        获取请求的所有候选记录

        Args:
            db: 数据库会话
            request_id: 请求ID

        Returns:
            候选记录列表，按 candidate_index 排序
        """
        return (
            db.query(RequestCandidate)
            .filter(RequestCandidate.request_id == request_id)
            .order_by(RequestCandidate.candidate_index)
            .all()
        )

    @staticmethod
    def get_candidate_stats_by_provider(db: Session, provider_id: str, limit: int = 100) -> dict:
        """
        获取 Provider 的候选统计

        Args:
            db: 数据库会话
            provider_id: Provider ID
            limit: 最近记录数量限制

        Returns:
            统计信息字典
        """
        candidates = (
            db.query(RequestCandidate)
            .filter(RequestCandidate.provider_id == provider_id)
            .order_by(RequestCandidate.created_at.desc())
            .limit(limit)
            .all()
        )

        total_candidates = len(candidates)
        success_count = sum(1 for c in candidates if c.status == "success")
        failed_count = sum(1 for c in candidates if c.status == "failed")
        cancelled_count = sum(1 for c in candidates if c.status == "cancelled")
        skipped_count = sum(1 for c in candidates if c.status == "skipped")
        pending_count = sum(1 for c in candidates if c.status == "pending")
        available_count = sum(1 for c in candidates if c.status == "available")

        # 计算失败率（只统计已完成的候选，即成功或失败的，cancelled 不算失败）
        completed_count = success_count + failed_count
        failure_rate = (failed_count / completed_count * 100) if completed_count > 0 else 0

        return {
            "total_attempts": total_candidates,  # 前端使用 total_attempts 字段
            "success_count": success_count,
            "failed_count": failed_count,
            "cancelled_count": cancelled_count,  # 客户端取消数
            "skipped_count": skipped_count,
            "pending_count": pending_count,
            "available_count": available_count,  # 尚未被调度的候选数
            "failure_rate": round(failure_rate, 2),
        }

    @staticmethod
    def calculate_candidate_ttfb(
        db: Session,
        candidate_id: str,
        request_start_time: float,
        global_first_byte_time_ms: int,
    ) -> int:
        """
        计算候选自身的首字节时间 (TTFB)

        请求链路追踪中的 TTFB 应该是"该候选自身"的首字时间，
        而不是整个请求从开始到收到首字节的时间。

        Args:
            db: 数据库会话
            candidate_id: 候选 ID
            request_start_time: 请求开始时间（Unix timestamp，秒）
            global_first_byte_time_ms: 全局首字节时间（相对于 request_start_time 的毫秒数）

        Returns:
            候选自身的 TTFB（毫秒），如果计算失败则返回 global_first_byte_time_ms
        """
        try:
            candidate = (
                db.query(RequestCandidate).filter(RequestCandidate.id == candidate_id).first()
            )
            if candidate and candidate.started_at:
                started_at = candidate.started_at
                if started_at.tzinfo is None:
                    started_at = started_at.replace(tzinfo=timezone.utc)
                # 使用整数毫秒计算，避免浮点精度问题
                request_start_epoch_ms = round(request_start_time * 1000)
                started_at_epoch_ms = round(started_at.timestamp() * 1000)
                first_byte_epoch_ms = request_start_epoch_ms + global_first_byte_time_ms
                return max(0, int(first_byte_epoch_ms - started_at_epoch_ms))
        except Exception as e:
            logger.debug("计算候选 TTFB 失败: {}", e)
        return global_first_byte_time_ms