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:
@@ -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()
|
||||||
@@ -97,6 +97,16 @@ class ModelUsageTracker:
|
|||||||
except Exception as exc: # noqa: BLE001 - tracking must never break a real call
|
except Exception as exc: # noqa: BLE001 - tracking must never break a real call
|
||||||
logger.warning("ModelUsageTracker.record failed for {}:{} - {}", provider, model, exc)
|
logger.warning("ModelUsageTracker.record failed for {}:{} - {}", provider, model, exc)
|
||||||
|
|
||||||
|
def seed(self, entries: dict[str, ModelUsageEntry]) -> None:
|
||||||
|
"""Bulk-load persisted entries (called once at startup, before any traffic).
|
||||||
|
|
||||||
|
Unlike record(), this replaces entries wholesale rather than
|
||||||
|
accumulating deltas — it exists to restore counters saved by a
|
||||||
|
previous process run, not to record a new call.
|
||||||
|
"""
|
||||||
|
with self._lock:
|
||||||
|
self._entries.update(entries)
|
||||||
|
|
||||||
def snapshot(self) -> dict[str, ModelUsageEntry]:
|
def snapshot(self) -> dict[str, ModelUsageEntry]:
|
||||||
"""Return a shallow copy of all tracked entries, safe to mutate by the caller."""
|
"""Return a shallow copy of all tracked entries, safe to mutate by the caller."""
|
||||||
with self._lock:
|
with self._lock:
|
||||||
|
|||||||
@@ -0,0 +1,108 @@
|
|||||||
|
"""Unit tests for PostgresModelUsageStore, using a mocked psycopg2 pool.
|
||||||
|
|
||||||
|
Mirrors the mocking pattern in backend/tests/perception/test_postgres_event_store.py
|
||||||
|
— no real database is needed.
|
||||||
|
"""
|
||||||
|
|
||||||
|
from __future__ import annotations
|
||||||
|
|
||||||
|
import sys
|
||||||
|
from datetime import datetime, timezone
|
||||||
|
from unittest.mock import MagicMock, patch
|
||||||
|
|
||||||
|
# Patch psycopg2 before importing the module under test.
|
||||||
|
mock_psycopg2 = MagicMock()
|
||||||
|
mock_psycopg2.extras = MagicMock()
|
||||||
|
sys.modules.setdefault("psycopg2", mock_psycopg2)
|
||||||
|
sys.modules.setdefault("psycopg2.extras", mock_psycopg2.extras)
|
||||||
|
sys.modules.setdefault("psycopg2.pool", MagicMock())
|
||||||
|
|
||||||
|
from app.shared.model_usage_tracker import ModelUsageEntry
|
||||||
|
|
||||||
|
|
||||||
|
def _cursor_returning(rows):
|
||||||
|
"""Build a MagicMock standing in for a psycopg2 cursor context manager."""
|
||||||
|
cursor = MagicMock()
|
||||||
|
cursor.__enter__ = lambda s: s
|
||||||
|
cursor.__exit__ = MagicMock(return_value=False)
|
||||||
|
cursor.fetchall.return_value = rows
|
||||||
|
return cursor
|
||||||
|
|
||||||
|
|
||||||
|
@patch("app.infrastructure.storage.postgres_model_usage_store.PostgresModelUsageStore._ensure_schema")
|
||||||
|
@patch("app.infrastructure.storage.postgres_model_usage_store.ThreadedConnectionPool")
|
||||||
|
def test_load_all_returns_entries_keyed_by_provider_model(mock_pool_class, mock_ensure):
|
||||||
|
"""load_all() must turn each row into a ModelUsageEntry keyed by 'provider:model'."""
|
||||||
|
row = {
|
||||||
|
"provider": "deepseek",
|
||||||
|
"model": "deepseek-v4-flash",
|
||||||
|
"total_tokens": 100,
|
||||||
|
"prompt_tokens": 60,
|
||||||
|
"completion_tokens": 40,
|
||||||
|
"call_count_ok": 5,
|
||||||
|
"call_count_error": 1,
|
||||||
|
"last_called_at": datetime(2026, 7, 23, tzinfo=timezone.utc),
|
||||||
|
"last_latency_ms": 250,
|
||||||
|
"last_error": None,
|
||||||
|
}
|
||||||
|
mock_pool = MagicMock()
|
||||||
|
mock_pool_class.return_value = mock_pool
|
||||||
|
conn = MagicMock()
|
||||||
|
conn.__enter__ = lambda s: s
|
||||||
|
conn.__exit__ = MagicMock(return_value=False)
|
||||||
|
conn.cursor.return_value = _cursor_returning([row])
|
||||||
|
mock_pool.getconn.return_value = conn
|
||||||
|
|
||||||
|
from app.infrastructure.storage.postgres_model_usage_store import PostgresModelUsageStore
|
||||||
|
store = PostgresModelUsageStore()
|
||||||
|
entries = store.load_all()
|
||||||
|
|
||||||
|
assert "deepseek:deepseek-v4-flash" in entries
|
||||||
|
entry = entries["deepseek:deepseek-v4-flash"]
|
||||||
|
assert isinstance(entry, ModelUsageEntry)
|
||||||
|
assert entry.total_tokens == 100
|
||||||
|
assert entry.call_count_error == 1
|
||||||
|
|
||||||
|
|
||||||
|
@patch("app.infrastructure.storage.postgres_model_usage_store.PostgresModelUsageStore._ensure_schema")
|
||||||
|
@patch("app.infrastructure.storage.postgres_model_usage_store.ThreadedConnectionPool")
|
||||||
|
def test_flush_upserts_every_entry(mock_pool_class, mock_ensure):
|
||||||
|
"""flush() must execute one UPSERT per tracked entry and commit once."""
|
||||||
|
mock_pool = MagicMock()
|
||||||
|
mock_pool_class.return_value = mock_pool
|
||||||
|
conn = MagicMock()
|
||||||
|
conn.__enter__ = lambda s: s
|
||||||
|
conn.__exit__ = MagicMock(return_value=False)
|
||||||
|
cursor = MagicMock()
|
||||||
|
cursor.__enter__ = lambda s: s
|
||||||
|
cursor.__exit__ = MagicMock(return_value=False)
|
||||||
|
conn.cursor.return_value = cursor
|
||||||
|
mock_pool.getconn.return_value = conn
|
||||||
|
|
||||||
|
from app.infrastructure.storage.postgres_model_usage_store import PostgresModelUsageStore
|
||||||
|
store = PostgresModelUsageStore()
|
||||||
|
entries = {
|
||||||
|
"deepseek:deepseek-v4-flash": ModelUsageEntry(
|
||||||
|
provider="deepseek", model="deepseek-v4-flash", total_tokens=100, call_count_ok=5,
|
||||||
|
),
|
||||||
|
}
|
||||||
|
|
||||||
|
store.flush(entries)
|
||||||
|
|
||||||
|
assert cursor.execute.call_count == 1
|
||||||
|
conn.commit.assert_called_once()
|
||||||
|
|
||||||
|
|
||||||
|
@patch("app.infrastructure.storage.postgres_model_usage_store.PostgresModelUsageStore._ensure_schema")
|
||||||
|
@patch("app.infrastructure.storage.postgres_model_usage_store.ThreadedConnectionPool")
|
||||||
|
def test_flush_with_no_entries_does_not_touch_the_database(mock_pool_class, mock_ensure):
|
||||||
|
"""flush({}) must be a no-op — no point opening a connection for nothing."""
|
||||||
|
mock_pool = MagicMock()
|
||||||
|
mock_pool_class.return_value = mock_pool
|
||||||
|
|
||||||
|
from app.infrastructure.storage.postgres_model_usage_store import PostgresModelUsageStore
|
||||||
|
store = PostgresModelUsageStore()
|
||||||
|
|
||||||
|
store.flush({})
|
||||||
|
|
||||||
|
mock_pool.getconn.assert_not_called()
|
||||||
@@ -77,3 +77,33 @@ def test_snapshot_returns_independent_copy():
|
|||||||
def test_get_model_usage_tracker_returns_singleton():
|
def test_get_model_usage_tracker_returns_singleton():
|
||||||
"""get_model_usage_tracker() always returns the same process-wide instance."""
|
"""get_model_usage_tracker() always returns the same process-wide instance."""
|
||||||
assert get_model_usage_tracker() is get_model_usage_tracker()
|
assert get_model_usage_tracker() is get_model_usage_tracker()
|
||||||
|
|
||||||
|
|
||||||
|
def test_seed_populates_registry_from_persisted_entries():
|
||||||
|
"""seed() must bulk-load entries (e.g. from Postgres at startup) into the registry."""
|
||||||
|
tracker = ModelUsageTracker()
|
||||||
|
persisted = {
|
||||||
|
"deepseek:deepseek-v4-flash": ModelUsageEntry(
|
||||||
|
provider="deepseek", model="deepseek-v4-flash", total_tokens=500, call_count_ok=20,
|
||||||
|
),
|
||||||
|
}
|
||||||
|
|
||||||
|
tracker.seed(persisted)
|
||||||
|
|
||||||
|
entry = tracker.get("deepseek", "deepseek-v4-flash")
|
||||||
|
assert entry.total_tokens == 500
|
||||||
|
assert entry.call_count_ok == 20
|
||||||
|
|
||||||
|
|
||||||
|
def test_seed_then_record_accumulates_on_top_of_seeded_value():
|
||||||
|
"""A call recorded after seeding must add to the seeded total, not replace it."""
|
||||||
|
tracker = ModelUsageTracker()
|
||||||
|
tracker.seed({
|
||||||
|
"deepseek:deepseek-v4-flash": ModelUsageEntry(
|
||||||
|
provider="deepseek", model="deepseek-v4-flash", total_tokens=500,
|
||||||
|
),
|
||||||
|
})
|
||||||
|
|
||||||
|
tracker.record(provider="deepseek", model="deepseek-v4-flash", success=True, usage={"total_tokens": 10})
|
||||||
|
|
||||||
|
assert tracker.get("deepseek", "deepseek-v4-flash").total_tokens == 510
|
||||||
|
|||||||
Reference in New Issue
Block a user