vent_agent_2.0/billing/meter.py

"""用量入账。

:class:`UsageMeter` 是**请求作用域**的：一次 /api/chat 请求建一个实例，
在 SSE 生成器的 finally 里调 ``settle()``。两条纪律：

1. **幂等**——``settle()`` 只生效一次。SSE 生成器有正常结束、异常、客户端
   断连三条路径，全都汇到同一个 finally，不幂等就会重复扣费。
2. **归属时刻固定**——用量归属日取请求开始时算出的 ``usage_day``，不取
   ``settle()`` 执行时刻。否则跨零点的长请求会被算到第二天。
"""

from __future__ import annotations

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

logger = get_logger(__name__)


class UsageMeter:
    """一次请求的用量累加与入账。"""

    def __init__(
        self,
        user_name: str,
        session_id: str,
        kind: str,
        model: str,
        usage_day: str | None = None,
        trace_id: str = "",
    ) -> None:
        self.user_name = user_name
        self.session_id = session_id
        self.kind = kind
        self.model = model or config.MODEL
        # 归属日由请求开始时刻决定，故此处在构造时定死
        self.usage_day = usage_day or db_conn.day_key()
        self.trace_id = trace_id
        self.input_tokens = 0
        self.output_tokens = 0
        self.usage_source = "provider"
        self._settled = False
        self._usage_id: str | None = None

    # -------- 累加 --------

    def add(self, input_tokens: int, output_tokens: int, source: str = "provider") -> None:
        """累加一轮的 token。``source`` 为 provider | estimate。"""
        self.input_tokens += max(int(input_tokens), 0)
        self.output_tokens += max(int(output_tokens), 0)
        # 只要出现过估算，整条流水的口径就降级标注为 estimate，避免口径混杂
        if source == "estimate" and self.usage_source == "provider":
            self.usage_source = "estimate"

    # -------- 入账 --------

    @property
    def settled(self) -> bool:
        """是否已入账。"""
        return self._settled

    @property
    def total_tokens(self) -> int:
        return self.input_tokens + self.output_tokens

    def settle(self) -> dict:
        """把累计用量写入流水。重复调用直接返回上次结果。

        计费失败不向上抛——请求已经跑完，计费问题不该让用户看到报错。
        失败时记 ERROR 日志并把结果标为 ``settled=False``，留待对账。
        """
        if self._settled:
            return {"settled": True, "usage_id": self._usage_id, "duplicated": True}
        if self.total_tokens == 0:
            self._settled = True
            return {"settled": False, "reason": "no_usage"}

        try:
            cost, currency = pricing.compute_cost(
                self.model, self.input_tokens, self.output_tokens
            )
        except Exception as exc:
            logger.error(
                "计价失败未入账 user=%s model=%s tokens=%d/%d err=%s: %s",
                self.user_name, self.model, self.input_tokens, self.output_tokens,
                type(exc).__name__, exc,
            )
            self._settled = True
            return {"settled": False, "reason": "pricing_error", "error": str(exc)}

        try:
            usage_id = billing_store.record_usage(
                user_name=self.user_name,
                session_id=self.session_id,
                kind=self.kind,
                model=self.model,
                input_tokens=self.input_tokens,
                output_tokens=self.output_tokens,
                cost=cost,
                usage_day=self.usage_day,
                trace_id=self.trace_id,
            )
        except Exception as exc:
            logger.error(
                "入账失败 user=%s session=%s tokens=%d/%d err=%s: %s",
                self.user_name, self.session_id, self.input_tokens, self.output_tokens,
                type(exc).__name__, exc,
            )
            self._settled = True
            return {"settled": False, "reason": "store_error", "error": str(exc)}

        self._settled = True
        self._usage_id = usage_id
        logger.info(
            "入账完成 user=%s kind=%s model=%s tokens=%d/%d cost=%.6f%s source=%s",
            self.user_name, self.kind, self.model,
            self.input_tokens, self.output_tokens, cost, currency, self.usage_source,
        )
        return {
            "settled": True,
            "usage_id": usage_id,
            "input_tokens": self.input_tokens,
            "output_tokens": self.output_tokens,
            "cost": cost,
            "currency": currency,
            "usage_source": self.usage_source,
        }