Agent/agent/llm/pricing.py

"""
LLM 定价计算器

使用策略模式，支持：
1. YAML 配置文件定义模型价格
2. 不同 token 类型的差异化定价（input/output/reasoning/cache）
3. 自动匹配模型（支持通配符）
4. 费用计算

定价单位：美元 / 1M tokens
"""

import os
import re
from dataclasses import dataclass, field
from pathlib import Path
from typing import Dict, Any, Optional, List
import yaml

from .usage import TokenUsage


@dataclass
class ModelPricing:
    """
    单个模型的定价配置

    所有价格单位：美元 / 1M tokens
    """
    model: str                          # 模型名称（支持通配符 *）
    input_price: float = 0.0            # 输入 token 价格
    output_price: float = 0.0           # 输出 token 价格

    # 可选的差异化定价
    reasoning_price: Optional[float] = None      # 推理 token 价格（默认 = output_price）
    cache_creation_price: Optional[float] = None # 缓存创建价格（默认 = input_price * 1.25）
    cache_read_price: Optional[float] = None     # 缓存读取价格（默认 = input_price * 0.1）

    # 元数据
    provider: Optional[str] = None      # 提供商
    description: Optional[str] = None   # 描述

    def get_reasoning_price(self) -> float:
        """获取推理 token 价格"""
        return self.reasoning_price if self.reasoning_price is not None else self.output_price

    def get_cache_creation_price(self) -> float:
        """获取缓存创建价格"""
        return self.cache_creation_price if self.cache_creation_price is not None else self.input_price * 1.25

    def get_cache_read_price(self) -> float:
        """获取缓存读取价格"""
        return self.cache_read_price if self.cache_read_price is not None else self.input_price * 0.1

    def calculate_cost(self, usage: TokenUsage) -> float:
        """
        计算费用

        Args:
            usage: Token 使用量

        Returns:
            费用（美元）
        """
        cost = 0.0

        # 基础输入费用
        # 如果有缓存，需要分开计算
        if usage.cache_read_tokens or usage.cache_creation_tokens:
            # 普通输入 = 总输入 - 缓存读取（缓存读取部分单独计价）
            regular_input = usage.input_tokens - usage.cache_read_tokens
            cost += (regular_input / 1_000_000) * self.input_price
            cost += (usage.cache_read_tokens / 1_000_000) * self.get_cache_read_price()
            cost += (usage.cache_creation_tokens / 1_000_000) * self.get_cache_creation_price()
        else:
            cost += (usage.input_tokens / 1_000_000) * self.input_price

        # 输出费用
        # 如果有 reasoning tokens，需要分开计算
        if usage.reasoning_tokens:
            # 普通输出 = 总输出 - reasoning（reasoning 部分单独计价）
            regular_output = usage.output_tokens - usage.reasoning_tokens
            cost += (regular_output / 1_000_000) * self.output_price
            cost += (usage.reasoning_tokens / 1_000_000) * self.get_reasoning_price()
        else:
            cost += (usage.output_tokens / 1_000_000) * self.output_price

        return cost

    def matches(self, model_name: str) -> bool:
        """
        检查模型名称是否匹配

        支持通配符：
        - "gpt-4*" 匹配 "gpt-4", "gpt-4-turbo", "gpt-4o" 等
        - "claude-3-*" 匹配 "claude-3-opus", "claude-3-sonnet" 等
        """
        pattern = self.model.replace("*", ".*")
        return bool(re.match(f"^{pattern}$", model_name, re.IGNORECASE))

    @classmethod
    def from_dict(cls, data: Dict[str, Any]) -> "ModelPricing":
        """从字典创建"""
        return cls(
            model=data["model"],
            input_price=data.get("input_price", 0.0),
            output_price=data.get("output_price", 0.0),
            reasoning_price=data.get("reasoning_price"),
            cache_creation_price=data.get("cache_creation_price"),
            cache_read_price=data.get("cache_read_price"),
            provider=data.get("provider"),
            description=data.get("description"),
        )


class PricingCalculator:
    """
    定价计算器

    从 YAML 配置加载定价表，计算 LLM 调用费用
    """

    def __init__(self, config_path: Optional[str] = None):
        """
        初始化定价计算器

        Args:
            config_path: 定价配置文件路径，默认查找：
                1. 环境变量 AGENT_PRICING_CONFIG
                2. ./pricing.yaml
                3. ./config/pricing.yaml
                4. 使用内置默认配置
        """
        self._pricing_map: Dict[str, ModelPricing] = {}
        self._patterns: List[ModelPricing] = []  # 带通配符的定价

        # 加载配置
        config_path = self._resolve_config_path(config_path)
        if config_path and Path(config_path).exists():
            self._load_from_file(config_path)
        else:
            self._load_defaults()

    def _resolve_config_path(self, config_path: Optional[str]) -> Optional[str]:
        """解析配置文件路径"""
        if config_path:
            return config_path

        # 检查环境变量
        if env_path := os.getenv("AGENT_PRICING_CONFIG"):
            return env_path

        # 获取 agent 包的根目录（agent/llm/pricing.py -> agent/）
        agent_dir = Path(__file__).parent.parent
        project_root = agent_dir.parent  # 项目根目录

        # 检查默认位置（按优先级）
        search_paths = [
            # 1. 当前工作目录
            Path("pricing.yaml"),
            Path("config/pricing.yaml"),
            # 2. 项目根目录
            project_root / "pricing.yaml",
            project_root / "config" / "pricing.yaml",
            # 3. agent 包目录
            agent_dir / "pricing.yaml",
            agent_dir / "config" / "pricing.yaml",
        ]

        for path in search_paths:
            if path.exists():
                print(f"[Pricing] Loaded config from: {path}")
                return str(path)

        return None

    def _load_from_file(self, config_path: str) -> None:
        """从 YAML 文件加载配置"""
        with open(config_path, "r", encoding="utf-8") as f:
            config = yaml.safe_load(f)

        for item in config.get("models", []):
            pricing = ModelPricing.from_dict(item)
            if "*" in pricing.model:
                self._patterns.append(pricing)
            else:
                self._pricing_map[pricing.model.lower()] = pricing

    def _load_defaults(self) -> None:
        """加载内置默认定价"""
        defaults = self._get_default_pricing()
        for item in defaults:
            pricing = ModelPricing.from_dict(item)
            if "*" in pricing.model:
                self._patterns.append(pricing)
            else:
                self._pricing_map[pricing.model.lower()] = pricing

    def _get_default_pricing(self) -> List[Dict[str, Any]]:
        """
        内置默认定价表

        价格来源：各提供商官网（2024-12 更新）
        单位：美元 / 1M tokens
        """
        return [
            # ===== OpenAI =====
            {"model": "gpt-4o", "input_price": 2.50, "output_price": 10.00, "provider": "openai"},
            {"model": "gpt-4o-mini", "input_price": 0.15, "output_price": 0.60, "provider": "openai"},
            {"model": "gpt-4-turbo", "input_price": 10.00, "output_price": 30.00, "provider": "openai"},
            {"model": "gpt-4", "input_price": 30.00, "output_price": 60.00, "provider": "openai"},
            {"model": "gpt-3.5-turbo", "input_price": 0.50, "output_price": 1.50, "provider": "openai"},
            # o1/o3 系列（reasoning tokens 单独计价）
            {"model": "o1", "input_price": 15.00, "output_price": 60.00, "reasoning_price": 60.00, "provider": "openai"},
            {"model": "o1-mini", "input_price": 3.00, "output_price": 12.00, "reasoning_price": 12.00, "provider": "openai"},
            {"model": "o1-preview", "input_price": 15.00, "output_price": 60.00, "reasoning_price": 60.00, "provider": "openai"},
            {"model": "o3-mini", "input_price": 1.10, "output_price": 4.40, "reasoning_price": 4.40, "provider": "openai"},

            # ===== Anthropic Claude =====
            {"model": "claude-3-5-sonnet-20241022", "input_price": 3.00, "output_price": 15.00, "provider": "anthropic"},
            {"model": "claude-3-5-haiku-20241022", "input_price": 0.80, "output_price": 4.00, "provider": "anthropic"},
            {"model": "claude-3-opus-20240229", "input_price": 15.00, "output_price": 75.00, "provider": "anthropic"},
            {"model": "claude-3-sonnet-20240229", "input_price": 3.00, "output_price": 15.00, "provider": "anthropic"},
            {"model": "claude-3-haiku-20240307", "input_price": 0.25, "output_price": 1.25, "provider": "anthropic"},
            # Claude 通配符
            {"model": "claude-3-5-sonnet*", "input_price": 3.00, "output_price": 15.00, "provider": "anthropic"},
            {"model": "claude-3-opus*", "input_price": 15.00, "output_price": 75.00, "provider": "anthropic"},
            {"model": "claude-sonnet-4*", "input_price": 3.00, "output_price": 15.00, "provider": "anthropic"},
            {"model": "claude-opus-4*", "input_price": 15.00, "output_price": 75.00, "provider": "anthropic"},

            # ===== Google Gemini =====
            {"model": "gemini-2.0-flash", "input_price": 0.10, "output_price": 0.40, "provider": "google"},
            {"model": "gemini-2.0-flash-thinking", "input_price": 0.10, "output_price": 0.40, "reasoning_price": 0.40, "provider": "google"},
            {"model": "gemini-1.5-pro", "input_price": 1.25, "output_price": 5.00, "provider": "google"},
            {"model": "gemini-1.5-flash", "input_price": 0.075, "output_price": 0.30, "provider": "google"},
            {"model": "gemini-2.5-pro", "input_price": 1.25, "output_price": 10.00, "reasoning_price": 10.00, "provider": "google"},
            # Gemini 通配符
            {"model": "gemini-2.0*", "input_price": 0.10, "output_price": 0.40, "provider": "google"},
            {"model": "gemini-1.5*", "input_price": 1.25, "output_price": 5.00, "provider": "google"},
            {"model": "gemini-2.5*", "input_price": 1.25, "output_price": 10.00, "provider": "google"},

            # ===== DeepSeek =====
            {"model": "deepseek-chat", "input_price": 0.14, "output_price": 0.28, "provider": "deepseek"},
            {"model": "deepseek-reasoner", "input_price": 0.55, "output_price": 2.19, "reasoning_price": 2.19, "provider": "deepseek"},
            {"model": "deepseek-r1*", "input_price": 0.55, "output_price": 2.19, "reasoning_price": 2.19, "provider": "deepseek"},

            # ===== OpenRouter 转发（使用原模型价格）=====
            {"model": "anthropic/claude-3-5-sonnet", "input_price": 3.00, "output_price": 15.00, "provider": "openrouter"},
            {"model": "anthropic/claude-3-opus", "input_price": 15.00, "output_price": 75.00, "provider": "openrouter"},
            {"model": "anthropic/claude-sonnet-4*", "input_price": 3.00, "output_price": 15.00, "provider": "openrouter"},
            {"model": "anthropic/claude-opus-4*", "input_price": 15.00, "output_price": 75.00, "provider": "openrouter"},
            {"model": "openai/gpt-4o", "input_price": 2.50, "output_price": 10.00, "provider": "openrouter"},
            {"model": "openai/o1*", "input_price": 15.00, "output_price": 60.00, "reasoning_price": 60.00, "provider": "openrouter"},
            {"model": "google/gemini*", "input_price": 1.25, "output_price": 5.00, "provider": "openrouter"},
            {"model": "deepseek/deepseek-r1*", "input_price": 0.55, "output_price": 2.19, "reasoning_price": 2.19, "provider": "openrouter"},

            # ===== Yescode 代理 =====
            {"model": "claude-sonnet-4.5", "input_price": 3.00, "output_price": 15.00, "cache_creation_price": 3.75, "cache_read_price": 0.30, "provider": "yescode"},
        ]

    def get_pricing(self, model: str) -> Optional[ModelPricing]:
        """
        获取模型定价

        Args:
            model: 模型名称

        Returns:
            ModelPricing 或 None（未找到）
        """
        model_lower = model.lower()

        # 精确匹配
        if model_lower in self._pricing_map:
            return self._pricing_map[model_lower]

        # 通配符匹配
        for pattern in self._patterns:
            if pattern.matches(model):
                return pattern

        return None

    def calculate_cost(
        self,
        model: str,
        usage: TokenUsage,
        fallback_input_price: float = 1.0,
        fallback_output_price: float = 2.0
    ) -> float:
        """
        计算费用

        Args:
            model: 模型名称
            usage: Token 使用量
            fallback_input_price: 未找到定价时的默认输入价格
            fallback_output_price: 未找到定价时的默认输出价格

        Returns:
            费用（美元）
        """
        pricing = self.get_pricing(model)

        if pricing:
            return pricing.calculate_cost(usage)

        # 使用 fallback 价格
        fallback = ModelPricing(
            model=model,
            input_price=fallback_input_price,
            output_price=fallback_output_price
        )
        return fallback.calculate_cost(usage)

    def add_pricing(self, pricing: ModelPricing) -> None:
        """动态添加定价"""
        if "*" in pricing.model:
            self._patterns.append(pricing)
        else:
            self._pricing_map[pricing.model.lower()] = pricing

    def list_models(self) -> List[str]:
        """列出所有已配置的模型"""
        models = list(self._pricing_map.keys())
        models.extend(p.model for p in self._patterns)
        return sorted(models)


# 全局单例
_calculator: Optional[PricingCalculator] = None


def get_pricing_calculator() -> PricingCalculator:
    """获取全局定价计算器"""
    global _calculator
    if _calculator is None:
        _calculator = PricingCalculator()
    return _calculator


def calculate_cost(model: str, usage: TokenUsage) -> float:
    """
    便捷函数：计算费用

    Args:
        model: 模型名称
        usage: Token 使用量

    Returns:
        费用（美元）
    """
    return get_pricing_calculator().calculate_cost(model, usage)