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 json

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 = "",
        usage_id: str | None = None,
    ) -> 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 = usage_id or db_conn.new_id()

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

    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:
            payload = {"user_name": self.user_name, "session_id": self.session_id,
                       "kind": self.kind, "model": self.model, "usage_day": self.usage_day,
                       "trace_id": self.trace_id, "usage_id": self._usage_id,
                       "input_tokens": self.input_tokens, "output_tokens": self.output_tokens,
                       "source": self.usage_source}
            with billing_store._usage_lock, db_conn.transaction():
                if not db_conn.query_one("SELECT id FROM billing_pending WHERE id = ?", (self._usage_id,)):
                    db_conn.execute("INSERT INTO billing_pending (id, payload, created_at) VALUES (?, ?, ?)",
                                    (self._usage_id, json.dumps(payload), db_conn.fmt_dt()))
        except Exception:
            logger.exception("无法保存待结算记录 usage_id=%s", self._usage_id)
            return {"settled": False, "reason": "outbox_error"}

        try:
            cost, currency = pricing.compute_cost(
                self.model, self.input_tokens, self.output_tokens
            )
        except Exception as exc:
            logger.exception(
                "计价失败未入账 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,
            )
            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,
                usage_id=self._usage_id,
            )
        except Exception as exc:
            logger.exception(
                "入账失败 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,
            )
            return {"settled": False, "reason": "store_error", "error": str(exc)}

        self._settled = True
        self._usage_id = usage_id
        try:
            db_conn.execute("DELETE FROM billing_pending WHERE id = ?", (usage_id,))
        except Exception:
            logger.exception("清理待结算记录失败，下次重试将按流水 ID 去重")
        if pricing.enabled():
            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,
                extra={"kv": {"event": "billing.settled", "usage_id": usage_id,
                              "session_id": self.session_id, "cost": cost, "currency": currency}})
        else:
            logger.info('Token 用量已记录', extra={'kv': {'event': 'usage.recorded',
                'usage_id': usage_id, 'session_id': self.session_id, 'model': self.model,
                'input_tokens': self.input_tokens, 'output_tokens': self.output_tokens,
                'billing_enabled': False}})
        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,
            "billing_enabled": pricing.enabled(),
        }


def retry_pending() -> int:
    """重试持久化的结算；相同流水 ID 不会再次累计额度。"""
    settled = 0
    for row in db_conn.query("SELECT * FROM billing_pending ORDER BY created_at LIMIT 100"):
        try:
            payload = json.loads(row["payload"])
            inputs = payload.pop("input_tokens")
            outputs = payload.pop("output_tokens")
            source = payload.pop("source")
            meter = UsageMeter(**payload)
            meter.add(inputs, outputs, source)
            from obs import bind_trace_id, unbind_trace_id, log_context
            token = bind_trace_id(meter.trace_id)
            try:
                with log_context(run_id=meter._usage_id, session_id=meter.session_id, kind=meter.kind):
                    settled += bool(meter.settle().get("settled"))
            finally:
                unbind_trace_id(token)
        except Exception:
            logger.exception("待结算记录恢复失败 id=%s", row["id"])
    return settled