from __future__ import annotations from dataclasses import dataclass, field from datetime import date, datetime from pydantic import BaseModel from app.ai_provider.types.chat import TokenUsage @dataclass class CostBreakdown: input_cost: float = 0.0 output_cost: float = 0.0 total: float = 0.0 currency: str = "CNY" class CostCalculator: """根据价格信息和 token 数量计算成本。 保持厂商原始单位 —— 不做跨币种换算。 - 普通 input token 按 ``input_per_1m`` 计费。 - 缓存命中的 prompt token 优先按 ``input_cache_read_per_1m`` 计费 （若有配置），剩余按普通 input 费率。 - 若未配置 cache 费率，所有 prompt token 统一按 input_per_1m 计费。 """ @staticmethod def calculate( pricing: PricingInfo | None, prompt_tokens: int, completion_tokens: int, prompt_tokens_cached: int = 0, ) -> CostBreakdown: if pricing is None: return CostBreakdown(total=0.0) # 缓存 token：优先走缓存费率 cached = min(prompt_tokens_cached, prompt_tokens) uncached = prompt_tokens - cached input_cost = 0.0 if uncached > 0 and pricing.input_per_1m is not None: input_cost += (uncached / 1_000_000) * pricing.input_per_1m if cached > 0 and pricing.input_cache_read_per_1m is not None: input_cost += (cached / 1_000_000) * pricing.input_cache_read_per_1m elif cached > 0 and pricing.input_per_1m is not None: # 无缓存费率时统一按普通 input 计 input_cost += (cached / 1_000_000) * pricing.input_per_1m output_cost = 0.0 if pricing.output_per_1m is not None and completion_tokens > 0: output_cost = (completion_tokens / 1_000_000) * pricing.output_per_1m return CostBreakdown( input_cost=round(input_cost, 8), output_cost=round(output_cost, 8), total=round(input_cost + output_cost, 8), currency=pricing.currency, ) @dataclass class UsageRecord: model_name: str provider: str prompt_tokens: int = 0 completion_tokens: int = 0 total_tokens: int = 0 currency: str = "CNY" cost: float = 0.0 timestamp: datetime = field(default_factory=datetime.now) endpoint_name: str | None = None status: str = "success" error_message: str | None = None class BillingTracker: """记录用量并按模型计算成本。""" def __init__(self) -> None: self._records: list[UsageRecord] = [] def record( self, model_name: str, provider: str, usage: TokenUsage | None, pricing: PricingInfo | None = None, endpoint_name: str | None = None, status: str = "success", error_message: str | None = None, ) -> UsageRecord: prompt_tokens = usage.prompt_tokens if usage else 0 completion_tokens = usage.completion_tokens if usage else 0 total = usage.total_tokens if usage else 0 cached = usage.prompt_tokens_cached if usage else 0 cost_info = CostCalculator.calculate(pricing, prompt_tokens, completion_tokens, cached) record = UsageRecord( model_name=model_name, provider=provider, prompt_tokens=prompt_tokens, completion_tokens=completion_tokens, total_tokens=total, currency=cost_info.currency, cost=cost_info.total, endpoint_name=endpoint_name, status=status, error_message=error_message, ) self._records.append(record) return record def get_records( self, since: datetime | None = None, model_name: str | None = None, ) -> list[UsageRecord]: records = self._records if since: records = [r for r in records if r.timestamp >= since] if model_name: records = [r for r in records if r.model_name == model_name] return records def total_cost(self, currency: str | None = None) -> float: if currency: return sum(r.cost for r in self._records if r.currency == currency) return sum(r.cost for r in self._records) class BudgetController: """强制执行每日和每月预算限制。""" def __init__( self, daily_budget_cny: float = 0.0, monthly_budget_cny: float = 0.0, ) -> None: self.daily_budget_cny = daily_budget_cny self.monthly_budget_cny = monthly_budget_cny self._daily_cost_cny: float = 0.0 self._monthly_cost_cny: float = 0.0 self._current_date: date = date.today() def add_cost_cny(self, cost: float) -> None: self._daily_cost_cny += cost self._monthly_cost_cny += cost def allow_request(self) -> bool: return not ( (self.daily_budget_cny > 0 and self._daily_cost_cny >= self.daily_budget_cny) or (self.monthly_budget_cny > 0 and self._monthly_cost_cny >= self.monthly_budget_cny) ) def reset_daily(self) -> None: self._daily_cost_cny = 0.0 self._current_date = date.today() def reset_monthly(self) -> None: self._monthly_cost_cny = 0.0 class BillingFactory: """计费组件工厂。""" @staticmethod def create_tracker() -> BillingTracker: return BillingTracker() @staticmethod def create_budget_controller( daily_budget_cny: float = 0.0, monthly_budget_cny: float = 0.0, ) -> BudgetController: return BudgetController( daily_budget_cny=daily_budget_cny, monthly_budget_cny=monthly_budget_cny, ) # ── 定价信息 ────────────────────────────────────────────────────────── class PricingInfo(BaseModel): """模型定价信息。保持厂商原始单位，不转换不混合。 各厂商按各自币种/单位记录，展示时原生显示。 """ currency: str = "CNY" input_per_1m: float | None = None output_per_1m: float | None = None input_cache_read_per_1m: float | None = None class PricingRegistry: """模型定价注册表。以模型逻辑名为键查找定价。 与 ModelRegistry 解耦 —— 改价不影响模型注册，无定价的模型 也可以正常使用（不计费）。 """ def __init__(self, pricing_map: dict[str, PricingInfo] | None = None) -> None: self._pricing: dict[str, PricingInfo] = dict(pricing_map or {}) def get(self, model_name: str) -> PricingInfo | None: """获取指定模型的定价信息。无定价时返回 None。""" return self._pricing.get(model_name) def list_all(self) -> dict[str, PricingInfo]: """返回全部定价数据的副本。""" return dict(self._pricing) class PricingFactory: """定价注册表工厂。""" @staticmethod def from_dict(data: dict[str, dict]) -> PricingRegistry: """从 YAML 解析后的 dict 构建 PricingRegistry。 YAML 按厂商嵌套组织（自文档化），工厂将其摊平为按模型逻辑名索引： ```yaml pricing: deepseek: deepseek-v4-flash: currency: CNY input_per_1m: 1.0 anthropic: claude-sonnet: currency: USD input_per_1m: 3.0 ``` """ pricing_map: dict[str, PricingInfo] = {} for _provider_name, models_data in data.items(): if not isinstance(models_data, dict): continue for model_name, raw in models_data.items(): if raw is None: continue # 跳过所有字段均为 null 的条目（如 mimo-v2.5-asr） values = [v for v in raw.values() if not isinstance(v, bool)] if all(v is None for v in values): continue pricing_map[model_name] = PricingInfo(**raw) return PricingRegistry(pricing_map)