fix(storage,api): usage cost as per-turn delta of the sdk running total, backfill
This commit is contained in:
@@ -6,9 +6,10 @@ write goes through ``core/conversations``; the frontend only shapes JSON.
|
||||
bus as SSE with the same keepalive the markdown frontend uses, so a proxy
|
||||
never sees a silent socket.
|
||||
|
||||
Usage figures come from the ``usage`` table (one row per turn, API-price
|
||||
``cost_usd`` and per-model ``model_usage`` from the SDK's
|
||||
``ResultMessage``); subscription quotas come from ``rate_limits``
|
||||
Usage figures come from the ``usage`` table (one row per turn; API-price
|
||||
``cost_usd`` and per-model ``model_usage`` are per-turn deltas of the SDK's
|
||||
cumulative ``ResultMessage`` counters, see ``storage.append_usage``);
|
||||
subscription quotas come from ``rate_limits``
|
||||
(``RateLimitEvent``). The quota covers the whole subscription, so
|
||||
``/api/limits`` puts the gateway's own spend for the window next to it
|
||||
for calibration by eye.
|
||||
|
||||
@@ -9,12 +9,14 @@ from beaver_gateway.storage.db import (
|
||||
Database,
|
||||
append_audit,
|
||||
append_usage,
|
||||
backfill_usage_deltas,
|
||||
create_token,
|
||||
list_active_tokens,
|
||||
list_audit_records,
|
||||
list_tokens,
|
||||
revoke_token,
|
||||
touch_token,
|
||||
usage_deltas,
|
||||
)
|
||||
from beaver_gateway.storage.models import (
|
||||
AuditLog,
|
||||
@@ -45,10 +47,12 @@ __all__ = [
|
||||
"Usage",
|
||||
"append_audit",
|
||||
"append_usage",
|
||||
"backfill_usage_deltas",
|
||||
"create_token",
|
||||
"list_active_tokens",
|
||||
"list_audit_records",
|
||||
"list_tokens",
|
||||
"revoke_token",
|
||||
"touch_token",
|
||||
"usage_deltas",
|
||||
]
|
||||
|
||||
@@ -0,0 +1,37 @@
|
||||
"""Rewrite ``usage.cost_usd`` / ``model_usage`` as per-turn deltas.
|
||||
|
||||
Until 2026-09-01 the gateway stored ``ResultMessage.total_cost_usd`` verbatim,
|
||||
which is cumulative for the claude process - summing the column overstated a
|
||||
day by an order of magnitude. Run once after deploying the delta-aware
|
||||
:func:`beaver_gateway.storage.append_usage`::
|
||||
|
||||
python -m beaver_gateway.storage.backfill_usage
|
||||
|
||||
Reads ``DATABASE_URL`` like the gateway (``.env`` included); safe to rerun.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
|
||||
from dotenv import load_dotenv
|
||||
|
||||
from beaver_gateway.settings import Settings
|
||||
from beaver_gateway.storage.db import Database, backfill_usage_deltas
|
||||
|
||||
|
||||
async def _main() -> None:
|
||||
load_dotenv(override=False)
|
||||
settings = Settings() # ty: ignore[missing-argument]
|
||||
db = Database(settings.database_url)
|
||||
try:
|
||||
await db.create_all()
|
||||
async with db.session() as session:
|
||||
changed = await backfill_usage_deltas(session)
|
||||
finally:
|
||||
await db.dispose()
|
||||
print(f"usage rows rewritten: {changed}")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
asyncio.run(_main())
|
||||
@@ -20,7 +20,7 @@ from typing import TYPE_CHECKING, Any
|
||||
|
||||
from sqlalchemy import inspect, text
|
||||
from sqlalchemy.ext.asyncio import create_async_engine
|
||||
from sqlmodel import SQLModel, select
|
||||
from sqlmodel import SQLModel, col, select
|
||||
from sqlmodel.ext.asyncio.session import AsyncSession
|
||||
|
||||
from beaver_gateway.storage.models import AuditLog, Token, Usage
|
||||
@@ -227,18 +227,152 @@ async def list_audit_records(
|
||||
|
||||
|
||||
async def append_usage(session: AsyncSession, row: Usage) -> None:
|
||||
"""Persist a turn's usage with ``cost_usd`` / ``model_usage`` as per-turn deltas.
|
||||
|
||||
``row`` arrives with the SDK's cumulative ``total_cost_usd`` and
|
||||
``model_usage`` in ``cost_usd`` / ``model_usage``; they move to the
|
||||
``*_total`` columns and the stored values become the difference against
|
||||
the previous row of the same session (see :func:`usage_deltas`).
|
||||
"""
|
||||
prev = await _last_usage(session, row.session_id)
|
||||
usage_deltas(row, prev)
|
||||
session.add(row)
|
||||
await session.commit()
|
||||
|
||||
|
||||
async def _last_usage(session: AsyncSession, session_id: str | None) -> Usage | None:
|
||||
if not session_id:
|
||||
return None
|
||||
stmt = (
|
||||
select(Usage)
|
||||
.where(col(Usage.session_id) == session_id)
|
||||
.order_by(col(Usage.ts).desc(), col(Usage.id).desc())
|
||||
.limit(1)
|
||||
)
|
||||
return (await session.exec(stmt)).first()
|
||||
|
||||
|
||||
# Counters inside ``ResultMessage.model_usage[model]``; everything else there
|
||||
# (provider, costBasis, contextWindow, maxOutputTokens, ...) is a constant.
|
||||
_MODEL_USAGE_COUNTERS = (
|
||||
"inputTokens",
|
||||
"outputTokens",
|
||||
"cacheReadInputTokens",
|
||||
"cacheCreationInputTokens",
|
||||
"webSearchRequests",
|
||||
"costUSD",
|
||||
)
|
||||
|
||||
|
||||
def usage_deltas(row: Usage, prev: Usage | None) -> None:
|
||||
"""Turn the SDK's cumulative cost counters on ``row`` into per-turn deltas.
|
||||
|
||||
``total_cost_usd`` and ``model_usage`` in a ``ResultMessage`` are running
|
||||
totals of the claude process: they grow with every turn of the session and
|
||||
start over when the process is respawned (rotation, resume after a
|
||||
restart). A row fresh from the SDK (``cost_usd_total is None``) first moves
|
||||
the raw values into ``cost_usd_total`` / ``model_usage_total``; then
|
||||
``cost_usd`` / ``model_usage`` become ``current - previous`` against
|
||||
``prev`` (the last row of the same session), or the raw value when there
|
||||
is no previous row or the counter went backwards - a fresh process.
|
||||
Idempotent on rows that already carry their totals, which is what the
|
||||
backfill relies on.
|
||||
"""
|
||||
if row.cost_usd_total is None and row.model_usage_total is None:
|
||||
row.cost_usd_total = row.cost_usd
|
||||
row.model_usage_total = row.model_usage
|
||||
prev_cost = prev.cost_usd_total if prev is not None else None
|
||||
prev_models = prev.model_usage_total if prev is not None else None
|
||||
row.cost_usd = _counter_delta(row.cost_usd_total, prev_cost)
|
||||
row.model_usage = _model_usage_delta(row.model_usage_total, prev_models)
|
||||
|
||||
|
||||
def _counter_delta(current: float | None, previous: float | None) -> float | None:
|
||||
if current is None:
|
||||
return None
|
||||
if previous is None or current < previous:
|
||||
return current
|
||||
return round(current - previous, 6)
|
||||
|
||||
|
||||
def _model_usage_delta(
|
||||
current: dict[str, Any] | None, previous: dict[str, Any] | None
|
||||
) -> dict[str, Any] | None:
|
||||
if current is None:
|
||||
return None
|
||||
out: dict[str, Any] = {}
|
||||
for model, fields in current.items():
|
||||
if not isinstance(fields, dict):
|
||||
out[model] = fields
|
||||
continue
|
||||
before = (previous or {}).get(model)
|
||||
before = before if isinstance(before, dict) else {}
|
||||
delta = dict(fields)
|
||||
for key in _MODEL_USAGE_COUNTERS:
|
||||
if key not in fields:
|
||||
continue
|
||||
cur, prev = fields.get(key), before.get(key)
|
||||
if not isinstance(cur, int | float) or isinstance(cur, bool):
|
||||
continue
|
||||
prev_num = (
|
||||
prev
|
||||
if isinstance(prev, int | float) and not isinstance(prev, bool)
|
||||
else None
|
||||
)
|
||||
value = _counter_delta(cur, prev_num)
|
||||
delta[key] = (
|
||||
round(value or 0, 6) if isinstance(cur, float) else int(value or 0)
|
||||
)
|
||||
out[model] = delta
|
||||
return out
|
||||
|
||||
|
||||
async def backfill_usage_deltas(session: AsyncSession) -> int:
|
||||
"""Rewrite ``cost_usd`` / ``model_usage`` of every row as per-turn deltas.
|
||||
|
||||
Rows written before the ``*_total`` columns existed hold the SDK's
|
||||
cumulative values in ``cost_usd`` / ``model_usage``; they are moved to the
|
||||
totals first. Rows are replayed per session in time order. Idempotent:
|
||||
a second run rewrites nothing. Returns the number of rows changed.
|
||||
"""
|
||||
stmt = select(Usage).order_by(col(Usage.session_id), col(Usage.ts), col(Usage.id))
|
||||
rows = (await session.exec(stmt)).all()
|
||||
last_by_session: dict[str, Usage] = {}
|
||||
changed = 0
|
||||
for row in rows:
|
||||
before = (
|
||||
row.cost_usd,
|
||||
row.model_usage,
|
||||
row.cost_usd_total,
|
||||
row.model_usage_total,
|
||||
)
|
||||
prev = last_by_session.get(row.session_id) if row.session_id else None
|
||||
usage_deltas(row, prev)
|
||||
after = (
|
||||
row.cost_usd,
|
||||
row.model_usage,
|
||||
row.cost_usd_total,
|
||||
row.model_usage_total,
|
||||
)
|
||||
if after != before:
|
||||
session.add(row)
|
||||
changed += 1
|
||||
if row.session_id:
|
||||
last_by_session[row.session_id] = row
|
||||
await session.commit()
|
||||
return changed
|
||||
|
||||
|
||||
__all__ = [
|
||||
"Database",
|
||||
"append_audit",
|
||||
"append_usage",
|
||||
"backfill_usage_deltas",
|
||||
"create_token",
|
||||
"list_active_tokens",
|
||||
"list_audit_records",
|
||||
"list_tokens",
|
||||
"revoke_token",
|
||||
"touch_token",
|
||||
"usage_deltas",
|
||||
]
|
||||
|
||||
@@ -287,7 +287,17 @@ class TranscriptEntry(SQLModel, table=True):
|
||||
|
||||
|
||||
class Usage(SQLModel, table=True):
|
||||
"""Per-turn token accounting from ``ResultMessage.usage``."""
|
||||
"""Per-turn token accounting from ``ResultMessage``.
|
||||
|
||||
Token columns come from ``ResultMessage.usage`` and are sums over the
|
||||
turn's API calls. ``cost_usd`` and ``model_usage`` are **per-turn deltas**:
|
||||
the SDK reports ``total_cost_usd`` and ``model_usage`` cumulatively for
|
||||
the claude process (growing across the turns of a session, reset when
|
||||
the process is respawned), so :func:`append_usage` stores the raw values
|
||||
in ``cost_usd_total`` / ``model_usage_total`` and the difference against
|
||||
the previous row of the same session here. Summing ``cost_usd`` over rows
|
||||
is therefore meaningful; summing ``cost_usd_total`` is not.
|
||||
"""
|
||||
|
||||
__tablename__ = "usage"
|
||||
|
||||
@@ -311,6 +321,14 @@ class Usage(SQLModel, table=True):
|
||||
default=None,
|
||||
sa_column=Column(JSON().with_variant(JSONB(), "postgresql"), nullable=True),
|
||||
)
|
||||
"""Per-model tokens, cost and web searches - per-turn delta (see class doc)."""
|
||||
cost_usd_total: float | None = Field(default=None)
|
||||
"""``ResultMessage.total_cost_usd`` verbatim: cumulative for the claude process."""
|
||||
model_usage_total: dict[str, Any] | None = Field(
|
||||
default=None,
|
||||
sa_column=Column(JSON().with_variant(JSONB(), "postgresql"), nullable=True),
|
||||
)
|
||||
"""``ResultMessage.model_usage`` verbatim: cumulative for the claude process."""
|
||||
|
||||
|
||||
class RateLimit(SQLModel, table=True):
|
||||
|
||||
Reference in New Issue
Block a user