106 lines
3.8 KiB
Python
106 lines
3.8 KiB
Python
"""``cost_usd`` / ``model_usage`` are per-turn deltas of the SDK's running totals."""
|
|
|
|
import tempfile
|
|
from datetime import UTC, datetime, timedelta
|
|
from pathlib import Path
|
|
from typing import Any
|
|
|
|
from sqlmodel import col, select
|
|
|
|
from beaver_gateway.storage import Database, Usage, append_usage, backfill_usage_deltas
|
|
|
|
|
|
def _mu(cost: float, out: int, cache_read: int) -> dict[str, Any]:
|
|
return {
|
|
"claude-opus-5": {
|
|
"costUSD": cost,
|
|
"inputTokens": 1,
|
|
"outputTokens": out,
|
|
"cacheReadInputTokens": cache_read,
|
|
"cacheCreationInputTokens": 0,
|
|
"webSearchRequests": 0,
|
|
"contextWindow": 1_000_000,
|
|
"provider": "firstParty",
|
|
}
|
|
}
|
|
|
|
|
|
def _row(
|
|
ts: datetime, sid: str | None, cost: float, out: int, cache_read: int
|
|
) -> Usage:
|
|
return Usage(
|
|
ts=ts,
|
|
agent_name="d",
|
|
session_id=sid,
|
|
model="claude-opus-5",
|
|
output_tokens=out,
|
|
cache_read_tokens=cache_read,
|
|
cost_usd=cost,
|
|
model_usage=_mu(cost, out, cache_read),
|
|
)
|
|
|
|
|
|
async def _db() -> Database:
|
|
path = Path(tempfile.mkdtemp(prefix="beaver-usage-")) / "u.db"
|
|
db = Database(f"sqlite:///{path}")
|
|
await db.create_all()
|
|
return db
|
|
|
|
|
|
async def _all(db: Database) -> list[Usage]:
|
|
async with db.session() as session:
|
|
return list((await session.exec(select(Usage).order_by(col(Usage.id)))).all())
|
|
|
|
|
|
async def test_append_usage_stores_deltas_and_keeps_totals() -> None:
|
|
db = await _db()
|
|
t0 = datetime.now(UTC).replace(tzinfo=None)
|
|
async with db.session() as session:
|
|
await append_usage(session, _row(t0, "s1", 0.5, 100, 1000))
|
|
await append_usage(
|
|
session, _row(t0 + timedelta(minutes=1), "s1", 0.8, 160, 2500)
|
|
)
|
|
# the claude process was respawned: the counters start over
|
|
await append_usage(session, _row(t0 + timedelta(minutes=2), "s1", 0.2, 30, 400))
|
|
# another session is a separate counter
|
|
await append_usage(session, _row(t0 + timedelta(minutes=3), "s2", 0.7, 50, 900))
|
|
# no session id - nothing to diff against
|
|
await append_usage(session, _row(t0 + timedelta(minutes=4), None, 0.3, 10, 100))
|
|
rows = await _all(db)
|
|
assert [r.cost_usd for r in rows] == [0.5, 0.3, 0.2, 0.7, 0.3]
|
|
assert [r.cost_usd_total for r in rows] == [0.5, 0.8, 0.2, 0.7, 0.3]
|
|
second = rows[1].model_usage
|
|
assert second is not None
|
|
opus = second["claude-opus-5"]
|
|
assert (opus["costUSD"], opus["outputTokens"], opus["cacheReadInputTokens"]) == (
|
|
0.3,
|
|
60,
|
|
1500,
|
|
)
|
|
assert (opus["contextWindow"], opus["provider"]) == (1_000_000, "firstParty")
|
|
assert rows[1].model_usage_total == _mu(0.8, 160, 2500)
|
|
assert rows[2].model_usage is not None
|
|
assert rows[2].model_usage["claude-opus-5"]["outputTokens"] == 30
|
|
await db.dispose()
|
|
|
|
|
|
async def test_backfill_rewrites_legacy_cumulative_rows_once() -> None:
|
|
db = await _db()
|
|
t0 = datetime.now(UTC).replace(tzinfo=None)
|
|
async with db.session() as session:
|
|
# legacy rows: the raw running total sits in cost_usd / model_usage
|
|
for i, cost in enumerate((1.0, 1.5, 2.5, 0.4)):
|
|
session.add(_row(t0 + timedelta(minutes=i), "s1", cost, 10 * (i + 1), 100))
|
|
session.add(_row(t0 + timedelta(minutes=9), "s9", 0.9, 5, 50))
|
|
await session.commit()
|
|
async with db.session() as session:
|
|
assert await backfill_usage_deltas(session) == 5
|
|
rows = await _all(db)
|
|
assert [r.cost_usd for r in rows] == [1.0, 0.5, 1.0, 0.4, 0.9]
|
|
assert [r.cost_usd_total for r in rows] == [1.0, 1.5, 2.5, 0.4, 0.9]
|
|
assert rows[2].model_usage is not None
|
|
assert rows[2].model_usage["claude-opus-5"]["outputTokens"] == 10
|
|
async with db.session() as session:
|
|
assert await backfill_usage_deltas(session) == 0
|
|
await db.dispose()
|