From 5d132981ad7e09918f92b8f56a0771b5663f12ce Mon Sep 17 00:00:00 2001 From: wangwei Date: Thu, 23 Jul 2026 14:18:22 +0800 Subject: [PATCH] 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> --- .../storage/postgres_model_usage_store.py | 142 ++++++++++++++++++ backend/app/shared/model_usage_tracker.py | 10 ++ .../test_model_usage_persistence.py | 108 +++++++++++++ .../observability/test_model_usage_tracker.py | 30 ++++ 4 files changed, 290 insertions(+) create mode 100644 backend/app/infrastructure/storage/postgres_model_usage_store.py create mode 100644 backend/tests/observability/test_model_usage_persistence.py diff --git a/backend/app/infrastructure/storage/postgres_model_usage_store.py b/backend/app/infrastructure/storage/postgres_model_usage_store.py new file mode 100644 index 0000000..5561dfd --- /dev/null +++ b/backend/app/infrastructure/storage/postgres_model_usage_store.py @@ -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() diff --git a/backend/app/shared/model_usage_tracker.py b/backend/app/shared/model_usage_tracker.py index 9b4ce9c..dddb928 100644 --- a/backend/app/shared/model_usage_tracker.py +++ b/backend/app/shared/model_usage_tracker.py @@ -97,6 +97,16 @@ class ModelUsageTracker: except Exception as exc: # noqa: BLE001 - tracking must never break a real call 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]: """Return a shallow copy of all tracked entries, safe to mutate by the caller.""" with self._lock: diff --git a/backend/tests/observability/test_model_usage_persistence.py b/backend/tests/observability/test_model_usage_persistence.py new file mode 100644 index 0000000..cf7f393 --- /dev/null +++ b/backend/tests/observability/test_model_usage_persistence.py @@ -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() diff --git a/backend/tests/observability/test_model_usage_tracker.py b/backend/tests/observability/test_model_usage_tracker.py index d09db11..0c82e08 100644 --- a/backend/tests/observability/test_model_usage_tracker.py +++ b/backend/tests/observability/test_model_usage_tracker.py @@ -77,3 +77,33 @@ def test_snapshot_returns_independent_copy(): def test_get_model_usage_tracker_returns_singleton(): """get_model_usage_tracker() always returns the same process-wide instance.""" 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