"""计费预检。 策略是 **fail-closed**(沿用 1.5 的取舍):预检自身出问题时不放行, 因为放行的代价是「无法计费的请求被白跑」,而拦截的代价只是用户重试一次。 两类失败要区分开,因为它们对应不同的 HTTP 语义: - :class:`QuotaExceededError` → 429(用户问题,重试无用,等额度重置) - :class:`BillingConfigError` → 500(服务端配置问题,运维要处理) """ from __future__ import annotations from typing import Any import config from db import billing_store, db_conn from obs import get_logger logger = get_logger(__name__) class QuotaExceededError(RuntimeError): """当日 token 额度已用尽。""" def __init__(self, user_name: str, usage_day: str, quota: dict[str, Any]) -> None: self.user_name = user_name self.usage_day = usage_day self.quota = quota super().__init__( f"用户 {user_name} 在 {usage_day} 的 token 额度已用尽" f"({quota.get('used_tokens', 0)}/{quota.get('limit_tokens', 0)})" ) class BillingConfigError(RuntimeError): """计费依赖不可用(数据库异常、单价缺失等),需要运维介入。""" def precheck(user_name: str, usage_day: str | None = None) -> dict[str, Any]: """对话开始前检查额度。通过则返回当前额度快照。""" day = usage_day or db_conn.day_key() try: quota = billing_store.get_quota(user_name, day) from billing import pricing pricing.compute_cost(config.MODEL, 0, 0) except Exception as exc: raise BillingConfigError(f"额度预检失败(计费库不可用):{type(exc).__name__}: {exc}") from exc limit_tokens = int(quota.get("limit_tokens") or 0) used_tokens = int(quota.get("used_tokens") or 0) # limit_tokens == 0 表示不限量(默认),此时只做读取不拦截 if limit_tokens > 0 and used_tokens >= limit_tokens: logger.warning( "额度超限 user=%s day=%s used=%d limit=%d", user_name, day, used_tokens, limit_tokens ) raise QuotaExceededError(user_name, day, quota) return quota def snapshot(user_name: str, usage_day: str | None = None) -> dict[str, Any]: """只读额度快照,不拦截(供 /api/usage/today 使用)。异常转为不可用标记。""" day = usage_day or db_conn.day_key() try: quota = billing_store.get_quota(user_name, day) usage = billing_store.usage_for_day(user_name, day) except Exception as exc: logger.warning("读取用量失败:%s: %s", type(exc).__name__, exc) return { "usage_day": day, "available": False, "error": f"{type(exc).__name__}: {exc}", } return { "usage_day": day, "available": True, "limit_tokens": quota.get("limit_tokens", 0), "used_tokens": quota.get("used_tokens", 0), "remaining_tokens": quota.get("remaining_tokens"), # total_tokens 与 used_tokens 同值,保留两个键是因为消费方用法不一: # 前端看「额度」用 used/limit,统计看「用量」用 total_tokens——与 # billing_store.usage_for_day 的字段名对齐,避免同一个概念两个名字。 "total_tokens": quota.get("used_tokens", 0), "requests": usage.get("requests", 0), "input_tokens": usage.get("input_tokens", 0), "output_tokens": usage.get("output_tokens", 0), "cost": usage.get("cost", 0.0), "default_limit": int(config.DAILY_TOKEN_LIMIT), }