vent_agent_2.0/billing/pricing.py

"""模型单价与费用计算。

单价含义：每 1K token 的价格（与 config 表 billing_pricing 的列注释一致）。
骨架自带一组种子单价，启动时写入 billing_pricing（不覆盖已存在的人工调价）。
"""

from __future__ import annotations

import config
from db import billing_store, db_conn
from obs import get_logger

logger = get_logger(__name__)


class UnpricedModelError(RuntimeError):
    """模型在 billing_pricing 里没有单价，且不允许按 0 计费。"""


# 种子单价：仅为了让骨架开箱可算费，**不是**真实报价，上线前必须按实际
# 供应商账单校准。单位：元 / 1K token。
SEED_PRICING: dict[str, tuple[float, float]] = {
    "deepseek-flash": (0.0005, 0.0015),
    "deepseek-v4-flash": (0.0005, 0.0015),
    "deepseek-v4-pro": (0.002, 0.008),
    "deepseek-chat": (0.001, 0.002),
    "deepseek-reasoner": (0.001, 0.002),
    "gpt-4o-mini": (0.001, 0.004),
}


def seed_default_pricing() -> int:
    """把种子单价写入单价表；已存在的模型不动。返回新写入条数。"""
    created = 0
    for model, (input_price, output_price) in SEED_PRICING.items():
        if billing_store.get_pricing(model) is None:
            billing_store.set_pricing(model, input_price, output_price)
            created += 1
    # 当前配置的主模型必须在表内，否则首轮就会 fail-closed
    for model in {config.MODEL}:
        if model and billing_store.get_pricing(model) is None:
            billing_store.set_pricing(model, 0.0, 0.0)
            created += 1
            logger.warning(
                "模型 %s 未在种子单价表中，已按 0 价登记——上线前请补真实单价", model
            )
    if created:
        logger.info("单价表初始化：新写入 %d 条", created)
    return created


def compute_cost(model: str, input_tokens: int, output_tokens: int) -> tuple[float, str]:
    """按单价表算本轮费用。

    返回 ``(cost, currency)``。模型未定价时：
    - 默认 fail-closed —— 抛 :class:`UnpricedModelError`，避免「静默免费」
    - ``BILLING_ALLOW_UNPRICED=1`` 时降级为 0 价并告警（便于本地开发）
    """
    pricing = billing_store.get_pricing(model)
    if pricing is None:
        if config.get_bool("BILLING_ALLOW_UNPRICED", False):
            logger.warning("模型 %s 未定价，按 0 计费（BILLING_ALLOW_UNPRICED=1）", model)
            return 0.0, "CNY"
        raise UnpricedModelError(
            f"模型 {model!r} 未在 billing_pricing 表中定价；"
            f"请补单价，或设 BILLING_ALLOW_UNPRICED=1 临时按 0 计费"
        )
    input_price = float(pricing.get("input_price") or 0)
    output_price = float(pricing.get("output_price") or 0)
    cost = (int(input_tokens) / 1000.0) * input_price + (
        int(output_tokens) / 1000.0
    ) * output_price
    return round(cost, 8), str(pricing.get("currency") or "CNY")


def pricing_table() -> list[dict]:
    """当前单价表（供 /api/usage 展示）。"""
    rows = billing_store.list_pricing()
    return [
        {
            "model": row.get("model"),
            "input_price": float(row.get("input_price") or 0),
            "output_price": float(row.get("output_price") or 0),
            "currency": row.get("currency"),
            "updated_at": row.get("updated_at"),
        }
        for row in rows
    ]


__all__ = [
    "SEED_PRICING",
    "UnpricedModelError",
    "compute_cost",
    "pricing_table",
    "seed_default_pricing",
    "db_conn",
]