vent_agent_2.0/db/billing_store.py

"""计费存取：单价表、用量流水、每日额度。

口径说明（与 1.5 的关键差异）：用量一律取模型返回的**真实 token 数**，
不在这一层做估算。stream 层拿不到真实 usage 时才回落估算，并在
``billing_usage`` 上留下 ``usage_source`` 痕迹（见 api/context_meter.py）。
"""

from __future__ import annotations

from typing import Any
import threading

_usage_lock = threading.RLock()

from db import db_conn
from obs import get_logger

logger = get_logger(__name__)


# ============================================================
# 单价表
# ============================================================

def set_pricing(
    model: str, input_price: float, output_price: float, currency: str = "CNY"
) -> None:
    """写入/覆盖某模型的单价（每 1K token）。"""
    existing = db_conn.query_one("SELECT model FROM billing_pricing WHERE model = ?", (model,))
    stamp = db_conn.fmt_dt()
    if existing:
        db_conn.execute(
            "UPDATE billing_pricing SET input_price = ?, output_price = ?, currency = ?, updated_at = ?"
            " WHERE model = ?",
            (input_price, output_price, currency, stamp, model),
        )
    else:
        db_conn.execute(
            "INSERT INTO billing_pricing (model, input_price, output_price, currency, updated_at)"
            " VALUES (?, ?, ?, ?, ?)",
            (model, input_price, output_price, currency, stamp),
        )


def get_pricing(model: str) -> dict[str, Any] | None:
    """取某模型单价；未定价返回 None。"""
    return db_conn.query_one("SELECT * FROM billing_pricing WHERE model = ?", (model,))


def list_pricing() -> list[dict[str, Any]]:
    """列出全部单价。"""
    return db_conn.query("SELECT * FROM billing_pricing ORDER BY model ASC")


# ============================================================
# 用量流水
# ============================================================

def record_usage(
    user_name: str,
    session_id: str,
    kind: str,
    model: str,
    input_tokens: int,
    output_tokens: int,
    cost: float,
    usage_day: str,
    trace_id: str = "",
    usage_id: str | None = None,
) -> str:
    """落一条用量流水，返回流水 ID。"""
    usage_id = usage_id or db_conn.new_id()
    total = int(input_tokens) + int(output_tokens)
    with _usage_lock, db_conn.transaction():
        if db_conn.query_one("SELECT id FROM billing_usage WHERE id = ?", (usage_id,)):
            return usage_id
        db_conn.execute(
            "INSERT INTO billing_usage (id, user_name, session_id, kind, model, input_tokens,"
            " output_tokens, total_tokens, cost, usage_day, trace_id, created_at)"
            " VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)",
            (
                usage_id,
                user_name,
                session_id,
                kind,
                model,
                int(input_tokens),
                int(output_tokens),
                total,
                float(cost),
                usage_day,
                trace_id,
                db_conn.fmt_dt(),
            ),
        )
        _accumulate_quota(user_name, usage_day, total)
    return usage_id


def usage_for_day(user_name: str, usage_day: str) -> dict[str, Any]:
    """汇总某用户某天的用量。"""
    row = db_conn.query_one(
        "SELECT COUNT(*) AS requests,"
        " COALESCE(SUM(input_tokens), 0) AS input_tokens,"
        " COALESCE(SUM(output_tokens), 0) AS output_tokens,"
        " COALESCE(SUM(total_tokens), 0) AS total_tokens,"
        " COALESCE(SUM(cost), 0) AS cost"
        " FROM billing_usage WHERE user_name = ? AND usage_day = ?",
        (user_name, usage_day),
    ) or {}
    return {
        "user_name": user_name,
        "usage_day": usage_day,
        "requests": int(row.get("requests", 0) or 0),
        "input_tokens": int(row.get("input_tokens", 0) or 0),
        "output_tokens": int(row.get("output_tokens", 0) or 0),
        "total_tokens": int(row.get("total_tokens", 0) or 0),
        "cost": float(row.get("cost", 0) or 0),
    }


def recent_usage(user_name: str, limit: int = 20) -> list[dict[str, Any]]:
    """最近若干条用量流水。"""
    return db_conn.query(
        "SELECT * FROM billing_usage WHERE user_name = ?"
        " ORDER BY created_at DESC, id DESC LIMIT ?",
        (user_name, limit),
    )


# ============================================================
# 每日额度
# ============================================================

def _ensure_quota_row(user_name: str, usage_day: str) -> dict[str, Any]:
    with _usage_lock, db_conn.transaction():
        return _ensure_quota_row_locked(user_name, usage_day)


def _ensure_quota_row_locked(user_name: str, usage_day: str) -> dict[str, Any]:
    """确保当日额度行存在并返回它。额度上限取自配置（0 = 不限）。"""
    import config

    row = db_conn.query_one(
        "SELECT * FROM daily_quota WHERE user_name = ? AND usage_day = ?",
        (user_name, usage_day),
    )
    if row:
        return row
    limit_tokens = int(config.DAILY_TOKEN_LIMIT)
    db_conn.execute(
        "INSERT INTO daily_quota (user_name, usage_day, limit_tokens, used_tokens, updated_at)"
        " VALUES (?, ?, ?, 0, ?)",
        (user_name, usage_day, limit_tokens, db_conn.fmt_dt()),
    )
    return {
        "user_name": user_name,
        "usage_day": usage_day,
        "limit_tokens": limit_tokens,
        "used_tokens": 0,
    }


def _accumulate_quota(user_name: str, usage_day: str, tokens: int) -> None:
    """把本轮 token 累加到当日额度行。"""
    _ensure_quota_row(user_name, usage_day)
    db_conn.execute(
        "UPDATE daily_quota SET used_tokens = used_tokens + ?, updated_at = ?"
        " WHERE user_name = ? AND usage_day = ?",
        (int(tokens), db_conn.fmt_dt(), user_name, usage_day),
    )


def get_quota(user_name: str, usage_day: str) -> dict[str, Any]:
    """读当日额度与已用。"""
    row = _ensure_quota_row(user_name, usage_day)
    limit_tokens = int(row.get("limit_tokens", 0) or 0)
    used = int(row.get("used_tokens", 0) or 0)
    return {
        "user_name": user_name,
        "usage_day": usage_day,
        "limit_tokens": limit_tokens,
        "used_tokens": used,
        "remaining_tokens": max(limit_tokens - used, 0) if limit_tokens > 0 else None,
    }


def set_quota_limit(user_name: str, usage_day: str, limit_tokens: int) -> dict[str, Any]:
    """调整某用户某天的额度上限（0 = 不限）。"""
    _ensure_quota_row(user_name, usage_day)
    db_conn.execute(
        "UPDATE daily_quota SET limit_tokens = ?, updated_at = ?"
        " WHERE user_name = ? AND usage_day = ?",
        (int(limit_tokens), db_conn.fmt_dt(), user_name, usage_day),
    )
    return get_quota(user_name, usage_day)