feat: add PostgresModelUsageStore and ModelUsageTracker.seed()
- Add seed() method to ModelUsageTracker for bulk-loading persisted entries at startup - Create PostgresModelUsageStore for persistence of model usage counters to Postgres - Store only current cumulative snapshots (no historical time-series) - Use standard CREATE TABLE IF NOT EXISTS idiom matching other Postgres stores - Add comprehensive mocked unit tests for both components Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
This commit is contained in:
wangwei
@@ -0,0 +1,142 @@
|
||||
"""Postgres-backed persistence for cumulative AI model usage counters.
|
||||
|
||||
Keeps ModelUsageTracker (an in-memory, process-lifetime-only registry defined
|
||||
in app/shared/model_usage_tracker.py) from losing its counters on every
|
||||
backend restart. This store only ever persists the *current cumulative
|
||||
snapshot* per provider+model — not a historical time-series log — matching
|
||||
the "durable counters" scope decided in
|
||||
docs/superpowers/specs/2026-07-23-status-model-usage-hardening-design.md.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from contextlib import contextmanager
|
||||
|
||||
import psycopg2
|
||||
import psycopg2.extras
|
||||
from psycopg2.pool import ThreadedConnectionPool
|
||||
|
||||
from app.config.settings import settings
|
||||
from app.shared.model_usage_tracker import ModelUsageEntry
|
||||
|
||||
# Table creation follows the same CREATE TABLE IF NOT EXISTS idiom used by
|
||||
# every other Postgres store in this codebase — no migration framework.
|
||||
_CREATE_TABLE = """
|
||||
CREATE TABLE IF NOT EXISTS model_usage_stats (
|
||||
provider VARCHAR(64) NOT NULL,
|
||||
model VARCHAR(128) NOT NULL,
|
||||
total_tokens BIGINT NOT NULL DEFAULT 0,
|
||||
prompt_tokens BIGINT NOT NULL DEFAULT 0,
|
||||
completion_tokens BIGINT NOT NULL DEFAULT 0,
|
||||
call_count_ok BIGINT NOT NULL DEFAULT 0,
|
||||
call_count_error BIGINT NOT NULL DEFAULT 0,
|
||||
last_called_at TIMESTAMPTZ,
|
||||
last_latency_ms INTEGER,
|
||||
last_error TEXT,
|
||||
updated_at TIMESTAMPTZ NOT NULL DEFAULT NOW(),
|
||||
PRIMARY KEY (provider, model)
|
||||
);
|
||||
"""
|
||||
|
||||
_UPSERT = """
|
||||
INSERT INTO model_usage_stats
|
||||
(provider, model, total_tokens, prompt_tokens, completion_tokens,
|
||||
call_count_ok, call_count_error, last_called_at, last_latency_ms, last_error, updated_at)
|
||||
VALUES
|
||||
(%(provider)s, %(model)s, %(total_tokens)s, %(prompt_tokens)s, %(completion_tokens)s,
|
||||
%(call_count_ok)s, %(call_count_error)s, %(last_called_at)s, %(last_latency_ms)s, %(last_error)s, NOW())
|
||||
ON CONFLICT (provider, model) DO UPDATE SET
|
||||
total_tokens = EXCLUDED.total_tokens,
|
||||
prompt_tokens = EXCLUDED.prompt_tokens,
|
||||
completion_tokens = EXCLUDED.completion_tokens,
|
||||
call_count_ok = EXCLUDED.call_count_ok,
|
||||
call_count_error = EXCLUDED.call_count_error,
|
||||
last_called_at = EXCLUDED.last_called_at,
|
||||
last_latency_ms = EXCLUDED.last_latency_ms,
|
||||
last_error = EXCLUDED.last_error,
|
||||
updated_at = NOW();
|
||||
"""
|
||||
|
||||
|
||||
class PostgresModelUsageStore:
|
||||
"""Load and flush ModelUsageTracker snapshots to/from a Postgres table."""
|
||||
|
||||
def __init__(self) -> None:
|
||||
"""Open a small connection pool and ensure the table exists."""
|
||||
self._pool = ThreadedConnectionPool(
|
||||
minconn=1,
|
||||
maxconn=3,
|
||||
host=settings.postgres_host,
|
||||
port=settings.postgres_port,
|
||||
user=settings.postgres_user,
|
||||
password=settings.postgres_password,
|
||||
dbname=settings.postgres_db,
|
||||
)
|
||||
self._ensure_schema()
|
||||
|
||||
def _ensure_schema(self) -> None:
|
||||
"""Create the model_usage_stats table if it does not already exist."""
|
||||
with self._conn() as conn:
|
||||
with conn.cursor() as cur:
|
||||
cur.execute(_CREATE_TABLE)
|
||||
conn.commit()
|
||||
|
||||
@contextmanager
|
||||
def _conn(self):
|
||||
"""Borrow a pooled connection and always return it, even on error."""
|
||||
conn = self._pool.getconn()
|
||||
try:
|
||||
yield conn
|
||||
finally:
|
||||
self._pool.putconn(conn)
|
||||
|
||||
def load_all(self) -> dict[str, ModelUsageEntry]:
|
||||
"""Return every persisted row as {"provider:model": ModelUsageEntry}."""
|
||||
with self._conn() as conn:
|
||||
with conn.cursor(cursor_factory=psycopg2.extras.RealDictCursor) as cur:
|
||||
cur.execute("SELECT * FROM model_usage_stats")
|
||||
rows = cur.fetchall()
|
||||
entries: dict[str, ModelUsageEntry] = {}
|
||||
for row in rows:
|
||||
entry = ModelUsageEntry(
|
||||
provider=row["provider"],
|
||||
model=row["model"],
|
||||
total_tokens=row["total_tokens"],
|
||||
prompt_tokens=row["prompt_tokens"],
|
||||
completion_tokens=row["completion_tokens"],
|
||||
call_count_ok=row["call_count_ok"],
|
||||
call_count_error=row["call_count_error"],
|
||||
last_called_at=row["last_called_at"],
|
||||
last_latency_ms=row["last_latency_ms"],
|
||||
last_error=row["last_error"],
|
||||
)
|
||||
entries[f"{entry.provider}:{entry.model}"] = entry
|
||||
return entries
|
||||
|
||||
def flush(self, entries: dict[str, ModelUsageEntry]) -> None:
|
||||
"""Upsert the current cumulative snapshot of every tracked entry.
|
||||
|
||||
A no-op for an empty snapshot — avoids opening a connection for nothing
|
||||
(e.g. before any LLM/embedding/reranker call has happened yet).
|
||||
"""
|
||||
if not entries:
|
||||
return
|
||||
with self._conn() as conn:
|
||||
with conn.cursor() as cur:
|
||||
for entry in entries.values():
|
||||
cur.execute(
|
||||
_UPSERT,
|
||||
{
|
||||
"provider": entry.provider,
|
||||
"model": entry.model,
|
||||
"total_tokens": entry.total_tokens,
|
||||
"prompt_tokens": entry.prompt_tokens,
|
||||
"completion_tokens": entry.completion_tokens,
|
||||
"call_count_ok": entry.call_count_ok,
|
||||
"call_count_error": entry.call_count_error,
|
||||
"last_called_at": entry.last_called_at,
|
||||
"last_latency_ms": entry.last_latency_ms,
|
||||
"last_error": entry.last_error,
|
||||
},
|
||||
)
|
||||
conn.commit()
|
||||
Reference in New Issue
Block a user