254 lines
8.1 KiB
Python
254 lines
8.1 KiB
Python
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)
|