初始化提交

backend/app/ai_provider/core/billing.py

@@ -0,0 +1,253 @@
+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)