Aether/_deprecated_py_src/services/billing/usage_mapper.py

"""
Usage 字段映射器

将不同 API 格式的原始 usage 数据映射为标准化格式。

支持的格式：
- openai:*: OpenAI compatible (Chat/CLI)
- claude:*: Anthropic Messages (Chat/CLI)
- gemini:*: Google Gemini (Chat/CLI)
"""

from typing import Any

from src.services.billing.models import StandardizedUsage


class UsageMapper:
    """
    Usage 字段映射器

    将不同 API 格式的 usage 统一映射为 StandardizedUsage。

    示例:
        # OpenAI 格式
        raw_usage = {
            "prompt_tokens": 100,
            "completion_tokens": 50,
            "prompt_tokens_details": {"cached_tokens": 20},
            "completion_tokens_details": {"reasoning_tokens": 10}
        }
        usage = UsageMapper.map(raw_usage, "OPENAI")

        # Claude 格式
        raw_usage = {
            "input_tokens": 100,
            "output_tokens": 50,
            "cache_creation_input_tokens": 30,
            "cache_read_input_tokens": 20
        }
        usage = UsageMapper.map(raw_usage, "CLAUDE")
    """

    # =========================================================================
    # 字段映射配置
    # 格式: "source_path" -> "target_field"
    # source_path 支持点号分隔的嵌套路径
    # =========================================================================

    # OpenAI 格式字段映射
    OPENAI_MAPPING: dict[str, str] = {
        "prompt_tokens": "input_tokens",
        "completion_tokens": "output_tokens",
        "prompt_tokens_details.cached_tokens": "cache_read_tokens",
        "completion_tokens_details.reasoning_tokens": "reasoning_tokens",
    }

    # Claude 格式字段映射
    CLAUDE_MAPPING: dict[str, str] = {
        "input_tokens": "input_tokens",
        "output_tokens": "output_tokens",
        "cache_creation_input_tokens": "cache_creation_tokens",
        "cache_read_input_tokens": "cache_read_tokens",
    }

    # Gemini 格式字段映射
    GEMINI_MAPPING: dict[str, str] = {
        "promptTokenCount": "input_tokens",
        "candidatesTokenCount": "output_tokens",
        "cachedContentTokenCount": "cache_read_tokens",
        # Gemini 的 usageMetadata 格式
        "usageMetadata.promptTokenCount": "input_tokens",
        "usageMetadata.candidatesTokenCount": "output_tokens",
        "usageMetadata.cachedContentTokenCount": "cache_read_tokens",
    }

    @classmethod
    def map(
        cls,
        raw_usage: dict[str, Any],
        api_format: str,
        extra_mapping: dict[str, str] | None = None,
    ) -> StandardizedUsage:
        """
        将原始 usage 映射为标准化格式

        Args:
            raw_usage: 原始 usage 字典
            api_format: API 格式 ("OPENAI", "CLAUDE", "GEMINI" 等)
            extra_mapping: 额外的字段映射（用于自定义扩展）

        Returns:
            标准化的 usage 对象
        """
        if not raw_usage:
            return StandardizedUsage()

        # 获取对应格式的字段映射
        mapping = cls._get_mapping(api_format)

        # 合并额外映射
        if extra_mapping:
            mapping = {**mapping, **extra_mapping}

        result = StandardizedUsage()

        # 执行映射
        for source_path, target_field in mapping.items():
            value = cls._get_nested_value(raw_usage, source_path)
            if value is not None:
                result.set(target_field, value)

        return result

    @classmethod
    def map_from_response(
        cls,
        response: dict[str, Any],
        api_format: str,
    ) -> StandardizedUsage:
        """
        从完整响应中提取并映射 usage

        不同 API 格式的 usage 位置可能不同：
        - OpenAI: response["usage"]
        - Claude: response["usage"] 或 message_delta 中
        - Gemini: response["usageMetadata"]

        Args:
            response: 完整的 API 响应
            api_format: API 格式

        Returns:
            标准化的 usage 对象
        """
        format_norm = (api_format or "").strip().lower()
        api_family = format_norm.split(":", 1)[0] if ":" in format_norm else format_norm

        # 提取 usage 部分
        usage_data: dict[str, Any] = {}

        if api_family == "gemini":
            # Gemini: usageMetadata
            usage_data = response.get("usageMetadata", {})
            if not usage_data:
                # 尝试从 candidates 中获取
                candidates = response.get("candidates", [])
                if candidates:
                    usage_data = candidates[0].get("usageMetadata", {})
        else:
            # OpenAI/Claude: usage
            usage_data = response.get("usage", {})

        return cls.map(usage_data, api_format)

    @classmethod
    def _get_mapping(cls, api_format: str) -> dict[str, str]:
        """获取对应格式的字段映射"""
        format_norm = (api_format or "").strip().lower()
        api_family = format_norm.split(":", 1)[0] if ":" in format_norm else format_norm

        if api_family == "openai":
            return cls.OPENAI_MAPPING
        if api_family == "gemini":
            return cls.GEMINI_MAPPING
        # 默认 Claude（也覆盖未知/空值）
        return cls.CLAUDE_MAPPING

    @classmethod
    def _get_nested_value(cls, data: dict[str, Any], path: str) -> Any:
        """
        获取嵌套字段值

        支持点号分隔的路径，如 "prompt_tokens_details.cached_tokens"

        Args:
            data: 数据字典
            path: 字段路径

        Returns:
            字段值，不存在则返回 None
        """
        if not data or not path:
            return None

        keys = path.split(".")
        value: Any = data

        for key in keys:
            if isinstance(value, dict):
                value = value.get(key)
                if value is None:
                    return None
            else:
                return None

        return value

    @classmethod
    def register_format(cls, format_name: str, mapping: dict[str, str]) -> None:
        """
        注册新的格式映射

        Args:
            format_name: 格式名称（会自动转为大写）
            mapping: 字段映射
        """
        cls.FORMAT_MAPPINGS[format_name.upper()] = mapping

    @classmethod
    def get_supported_formats(cls) -> list:
        """获取所有支持的格式"""
        return list(cls.FORMAT_MAPPINGS.keys())


# =========================================================================
# 便捷函数
# =========================================================================


def map_usage(
    raw_usage: dict[str, Any],
    api_format: str,
) -> StandardizedUsage:
    """
    便捷函数：将原始 usage 映射为标准化格式

    Args:
        raw_usage: 原始 usage 字典
        api_format: API 格式

    Returns:
        StandardizedUsage 对象
    """
    return UsageMapper.map(raw_usage, api_format)


def map_usage_from_response(
    response: dict[str, Any],
    api_format: str,
) -> StandardizedUsage:
    """
    便捷函数：从响应中提取并映射 usage

    Args:
        response: API 响应
        api_format: API 格式

    Returns:
        StandardizedUsage 对象
    """
    return UsageMapper.map_from_response(response, api_format)