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:
@@ -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
|
||||
|
||||
Reference in New Issue
Block a user