Aether/_deprecated_py_src/services/usage/_billing_integration.py

from __future__ import annotations

from typing import Any

from src.core.api_format.signature import normalize_signature_key
from src.services.billing.token_normalization import normalize_input_tokens_for_billing
from src.services.usage._recording_helpers import (
    build_usage_params,
    deserialize_body_if_json,
    sanitize_request_metadata,
)
from src.services.usage._types import UsageCostInfo, UsageRecordParams


class UsageBillingIntegrationMixin:
    """计费集成方法 -- 准备用量记录的共享逻辑"""

    @classmethod
    async def _prepare_usage_record(
        cls,
        params: UsageRecordParams,
    ) -> tuple[dict[str, Any], float]:
        """准备用量记录的共享逻辑

        此方法提取了 record_usage 和 record_usage_async 的公共处理逻辑：
        - 获取费率倍数
        - 计算成本
        - 构建 Usage 参数

        Args:
            params: 用量记录参数数据类

        Returns:
            (usage_params 字典, total_cost 总成本)
        """
        # 计费口径以 Provider 为准（优先 endpoint_api_format）
        billing_api_format: str | None = None
        if params.endpoint_api_format:
            try:
                billing_api_format = normalize_signature_key(str(params.endpoint_api_format))
            except Exception:
                billing_api_format = None
        if billing_api_format is None and params.api_format:
            try:
                billing_api_format = normalize_signature_key(str(params.api_format))
            except Exception:
                billing_api_format = None

        input_tokens_for_billing = normalize_input_tokens_for_billing(
            billing_api_format,
            params.input_tokens,
            params.cache_read_input_tokens,
        )

        # 获取费率倍数和是否免费套餐（传递 api_format 支持按格式配置的倍率）
        actual_rate_multiplier, is_free_tier = await cls._get_rate_multiplier_and_free_tier(
            params.db, params.provider_api_key_id, params.provider_id, billing_api_format
        )

        metadata = dict(params.metadata or {})
        is_failed_request = params.status_code >= 400 or params.error_message is not None

        # Helper: compute billing task_type (billing domain)
        billing_task_type = (params.request_type or "").lower()
        if billing_task_type not in {"chat", "cli", "video", "image", "audio"}:
            billing_task_type = "chat"

        # 使用新计费系统计算费用
        from src.services.billing.service import BillingService

        request_count = 0 if is_failed_request else 1
        has_cache_tokens = bool(
            params.cache_creation_input_tokens > 0 or params.cache_read_input_tokens > 0
        )
        effective_cache_ttl_minutes = params.cache_ttl_minutes

        # 主链路很多场景不会显式传 cache_ttl_minutes，这里补全以确保 1h/5m TTL 差异化计价生效。
        if effective_cache_ttl_minutes is None and has_cache_tokens and params.provider_api_key_id:
            try:
                from src.models.database import ProviderAPIKey

                key_ttl = (
                    params.db.query(ProviderAPIKey.cache_ttl_minutes)
                    .filter(ProviderAPIKey.id == params.provider_api_key_id)
                    .scalar()
                )
                if key_ttl is not None:
                    key_ttl_int = int(key_ttl)
                    if key_ttl_int >= 0:
                        effective_cache_ttl_minutes = key_ttl_int
            except Exception:
                # Best-effort fallback below.
                pass

        # 无法从 key 获取时，尽量从 5m/1h 细分回推（主要覆盖 Claude cache_creation）。
        if effective_cache_ttl_minutes is None and has_cache_tokens:
            t5m = int(params.cache_creation_input_tokens_5m or 0)
            t1h = int(params.cache_creation_input_tokens_1h or 0)
            if t1h > 0 and t5m == 0:
                effective_cache_ttl_minutes = 60
            elif t5m > 0 and t1h == 0:
                effective_cache_ttl_minutes = 5
            elif t1h > 0:
                # 混合场景优先按长 TTL 计，避免 1h 缓存被按 5m 误计。
                effective_cache_ttl_minutes = 60

        dims: dict[str, Any] = {
            "input_tokens": input_tokens_for_billing,
            "output_tokens": params.output_tokens,
            "cache_creation_input_tokens": params.cache_creation_input_tokens,
            "cache_read_input_tokens": params.cache_read_input_tokens,
            "request_count": request_count,
        }
        if effective_cache_ttl_minutes is not None:
            dims["cache_ttl_minutes"] = effective_cache_ttl_minutes
        # If tiered pricing is disabled, force first tier by using tier-key=0.
        if not params.use_tiered_pricing:
            dims["total_input_context"] = 0

        billing = BillingService(params.db)
        result = billing.calculate(
            task_type=billing_task_type,
            model=params.model,
            provider_id=params.provider_id or "",
            dimensions=dims,
            strict_mode=None,
        )
        snap = result.snapshot

        breakdown = snap.cost_breakdown or {}
        input_cost = float(breakdown.get("input_cost", 0.0))
        output_cost = float(breakdown.get("output_cost", 0.0))
        cache_creation_cost = float(breakdown.get("cache_creation_cost", 0.0))
        cache_read_cost = float(breakdown.get("cache_read_cost", 0.0))
        request_cost = float(breakdown.get("request_cost", 0.0))
        cache_cost = cache_creation_cost + cache_read_cost
        total_cost = float(snap.total_cost or 0.0)

        rv = snap.resolved_variables or {}

        def _as_float(v: Any, d: float | None) -> float | None:
            try:
                if v is None:
                    return d
                return float(v)
            except Exception:
                return d

        input_price = _as_float(rv.get("input_price_per_1m"), 0.0) or 0.0
        output_price = _as_float(rv.get("output_price_per_1m"), 0.0) or 0.0
        cache_creation_price = _as_float(rv.get("cache_creation_price_per_1m"), None)
        cache_read_price = _as_float(rv.get("cache_read_price_per_1m"), None)
        request_price = _as_float(rv.get("price_per_request"), None)

        # Audit snapshot (pruned later by sanitize_request_metadata)
        metadata["billing_snapshot"] = snap.to_dict()

        # Best-effort prune metadata to reduce DB/memory pressure.
        metadata = sanitize_request_metadata(metadata)

        # 构建 Usage 参数
        request_body = deserialize_body_if_json(params.request_body)
        provider_request_body = deserialize_body_if_json(params.provider_request_body)
        response_body = deserialize_body_if_json(params.response_body)
        client_response_body = deserialize_body_if_json(params.client_response_body)
        usage_params = build_usage_params(
            db=params.db,
            user=params.user,
            api_key=params.api_key,
            provider=params.provider,
            model=params.model,
            input_tokens=input_tokens_for_billing,
            output_tokens=params.output_tokens,
            cache_creation_input_tokens=params.cache_creation_input_tokens,
            cache_read_input_tokens=params.cache_read_input_tokens,
            cache_creation_input_tokens_5m=params.cache_creation_input_tokens_5m,
            cache_creation_input_tokens_1h=params.cache_creation_input_tokens_1h,
            request_type=params.request_type,
            api_format=params.api_format,
            api_family=params.api_family,
            endpoint_kind=params.endpoint_kind,
            endpoint_api_format=params.endpoint_api_format,
            has_format_conversion=params.has_format_conversion,
            is_stream=params.is_stream,
            response_time_ms=params.response_time_ms,
            first_byte_time_ms=params.first_byte_time_ms,
            status_code=params.status_code,
            error_message=params.error_message,
            metadata=metadata,
            request_headers=params.request_headers,
            request_body=request_body,
            provider_request_headers=params.provider_request_headers,
            provider_request_body=provider_request_body,
            response_headers=params.response_headers,
            client_response_headers=params.client_response_headers,
            response_body=response_body,
            client_response_body=client_response_body,
            request_id=params.request_id,
            provider_id=params.provider_id,
            provider_endpoint_id=params.provider_endpoint_id,
            provider_api_key_id=params.provider_api_key_id,
            status=params.status,
            target_model=params.target_model,
            cost=UsageCostInfo(
                input_cost=input_cost,
                output_cost=output_cost,
                cache_creation_cost=cache_creation_cost,
                cache_read_cost=cache_read_cost,
                cache_cost=cache_cost,
                request_cost=request_cost,
                total_cost=total_cost,
                input_price=input_price,
                output_price=output_price,
                cache_creation_price=cache_creation_price,
                cache_read_price=cache_read_price,
                request_price=request_price,
                actual_rate_multiplier=actual_rate_multiplier,
                is_free_tier=is_free_tier,
            ),
        )

        return usage_params, total_cost

    @classmethod
    async def _prepare_usage_records_batch(
        cls,
        params_list: list[UsageRecordParams],
    ) -> list[tuple[dict[str, Any], float, Exception | None]]:
        """批量并行准备用量记录（性能优化）

        并行调用 _prepare_usage_record，提高批量处理效率。

        Args:
            params_list: 用量记录参数列表

        Returns:
            列表，每项为 (usage_params, total_cost, exception)
            如果处理成功，exception 为 None
        """
        import asyncio

        async def prepare_single(
            params: UsageRecordParams,
        ) -> tuple[dict[str, Any], float, Exception | None]:
            try:
                usage_params, total_cost = await cls._prepare_usage_record(params)
                return (usage_params, total_cost, None)
            except Exception as e:
                return ({}, 0.0, e)

        if not params_list:
            return []

        # 避免一次性创建过多 task（并且 _prepare_usage_record 内部也可能包含并行调用）
        # 这里采用分批 gather 来限制并发量。
        chunk_size = 50
        results: list[tuple[dict[str, Any], float, Exception | None]] = []
        for i in range(0, len(params_list), chunk_size):
            chunk = params_list[i : i + chunk_size]
            chunk_results = await asyncio.gather(*(prepare_single(p) for p in chunk))
            results.extend(chunk_results)
        return results