vent_agent_2.0/billing/guard.py

"""计费预检。

策略是 **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),
    }