diff --git a/backend/app/shared/bootstrap.py b/backend/app/shared/bootstrap.py index 69fa1ba..36c07c4 100644 --- a/backend/app/shared/bootstrap.py +++ b/backend/app/shared/bootstrap.py @@ -2,9 +2,12 @@ from __future__ import annotations +import asyncio from functools import lru_cache from typing import Callable +from loguru import logger + from app.application.agent import AgentConversationService, AgentSessionService from app.application.agent.agentic_service import AgenticConversationService from app.application.documents import DocumentCommandService, DocumentQueryService @@ -36,6 +39,7 @@ from app.infrastructure.storage.minio_binary_store import MinioDocumentBinarySto from app.infrastructure.storage.postgres_document_processing_store import PostgresDocumentProcessingStore from app.infrastructure.storage.postgres_document_repository import PostgresDocumentRepository from app.infrastructure.storage.postgres_parse_artifact_store import PostgresParseArtifactStore +from app.infrastructure.storage.postgres_model_usage_store import PostgresModelUsageStore from app.infrastructure.vectorstore.bm25_retriever import BM25Retriever from app.infrastructure.vectorstore.cross_encoder_reranker import OpenAICompatibleReranker from app.infrastructure.vectorstore.dense_retriever import DenseRetriever @@ -43,6 +47,7 @@ from app.infrastructure.vectorstore.milvus_vector_index import MilvusVectorIndex from app.services.llm.llm_factory import LLMFactory from app.domain.compliance.ports import ComplianceRepository from app.infrastructure.compliance.repository import PostgresComplianceRepository +from app.shared.model_usage_tracker import get_model_usage_tracker # Keep shared wiring centralized so dependency construction remains consistent. @@ -162,6 +167,14 @@ def get_parse_artifact_store(): return None +@lru_cache +def get_model_usage_store(): + """Return the Postgres model-usage store, or None when postgres backend is not enabled.""" + if settings.document_repository_backend == "postgres": + return PostgresModelUsageStore() + return None + + @lru_cache def get_document_processing_store(): """Return document processing store for the active repository backend.""" @@ -412,8 +425,67 @@ def get_user_store(): def preload_runtime_dependencies() -> None: """Warm dependencies that are safe and useful to preload during startup.""" LLMFactory.preload_clients(["qwen", "deepseek"]) + _start_model_usage_persistence() def cleanup_runtime_dependencies() -> None: """Release runtime dependencies that expose explicit cleanup hooks.""" LLMFactory.cleanup() + _stop_model_usage_persistence() + + +_model_usage_flush_task: "asyncio.Task | None" = None + + +def _start_model_usage_persistence() -> None: + """Seed ModelUsageTracker from Postgres and start its periodic flush loop. + + No-op when document_repository_backend != "postgres" — ModelUsageTracker + then keeps behaving exactly as it always has: purely in-memory, reset on + every restart. Never raises: persistence must not block app startup. + """ + global _model_usage_flush_task + try: + store = get_model_usage_store() + except Exception as exc: # noqa: BLE001 - persistence must never block startup + logger.warning("Failed to initialize model usage persistence: {}", exc) + return + if store is None: + return + + tracker = get_model_usage_tracker() + try: + tracker.seed(store.load_all()) + except Exception as exc: # noqa: BLE001 - a bad load must not block startup + logger.warning("Failed to load persisted model usage stats: {}", exc) + + async def _flush_loop() -> None: + """Snapshot the tracker into Postgres every 60 seconds until cancelled.""" + while True: + await asyncio.sleep(60) + try: + store.flush(tracker.snapshot()) + except Exception as exc: # noqa: BLE001 - one bad cycle must not kill the loop + logger.warning("Failed to flush model usage stats: {}", exc) + + _model_usage_flush_task = asyncio.create_task(_flush_loop()) + + +def _stop_model_usage_persistence() -> None: + """Cancel the periodic flush task and perform one best-effort final flush.""" + global _model_usage_flush_task + if _model_usage_flush_task is not None: + _model_usage_flush_task.cancel() + _model_usage_flush_task = None + + try: + store = get_model_usage_store() + except Exception as exc: # noqa: BLE001 - shutdown must not crash on this + logger.warning("Failed to access model usage store during shutdown: {}", exc) + return + if store is None: + return + try: + store.flush(get_model_usage_tracker().snapshot()) + except Exception as exc: # noqa: BLE001 - shutdown must not crash on a flush failure + logger.warning("Failed final model usage flush: {}", exc) diff --git a/backend/tests/observability/test_model_usage_bootstrap.py b/backend/tests/observability/test_model_usage_bootstrap.py new file mode 100644 index 0000000..67163aa --- /dev/null +++ b/backend/tests/observability/test_model_usage_bootstrap.py @@ -0,0 +1,111 @@ +"""Unit tests for the model-usage persistence wiring in app.shared.bootstrap. + +get_model_usage_store()'s settings-gating is tested the same way +tests/test_reranker_bootstrap.py tests get_reranker() — by patching +"app.shared.bootstrap.settings" wholesale, matching this codebase's +established convention for testing @lru_cache settings-gated factories. +The remaining tests isolate _start_model_usage_persistence() / +_stop_model_usage_persistence() from get_model_usage_store() entirely (via +monkeypatch on the module-level function), so no real database or event loop +is needed anywhere in this file — asyncio.create_task itself is also mocked. +""" + +from __future__ import annotations + +import sys +from unittest.mock import MagicMock, patch + +# Patch psycopg2 before importing anything that transitively imports it, in +# case this file is collected before test_model_usage_persistence.py. +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 import bootstrap +from app.shared.model_usage_tracker import ModelUsageEntry, ModelUsageTracker + + +def test_get_model_usage_store_returns_none_when_not_postgres_backend(): + """get_model_usage_store() must be None unless document_repository_backend == 'postgres'.""" + bootstrap.get_model_usage_store.cache_clear() + + with patch("app.shared.bootstrap.settings") as mock_settings: + mock_settings.document_repository_backend = "json" + result = bootstrap.get_model_usage_store() + + bootstrap.get_model_usage_store.cache_clear() + assert result is None + + +def test_get_model_usage_store_returns_instance_when_postgres_backend(): + """get_model_usage_store() must return a PostgresModelUsageStore when enabled. + + ThreadedConnectionPool is mocked so no real connection is attempted; the + postgres_host/port/user/password/db values PostgresModelUsageStore reads + come from app.config.settings.settings directly (not from the + app.shared.bootstrap.settings reference mocked below), so they don't need + to be set here — only document_repository_backend gates this factory. + """ + bootstrap.get_model_usage_store.cache_clear() + + with patch("psycopg2.pool.ThreadedConnectionPool"), \ + patch( + "app.infrastructure.storage.postgres_model_usage_store.PostgresModelUsageStore._ensure_schema" + ), \ + patch("app.shared.bootstrap.settings") as mock_settings: + mock_settings.document_repository_backend = "postgres" + result = bootstrap.get_model_usage_store() + + bootstrap.get_model_usage_store.cache_clear() + from app.infrastructure.storage.postgres_model_usage_store import PostgresModelUsageStore + assert isinstance(result, PostgresModelUsageStore) + + +def test_start_model_usage_persistence_seeds_tracker_and_starts_flush_loop(monkeypatch): + """When a store is available, startup must seed the tracker and schedule the flush task.""" + fake_store = MagicMock() + fake_store.load_all.return_value = { + "deepseek:deepseek-v4-flash": ModelUsageEntry( + provider="deepseek", model="deepseek-v4-flash", total_tokens=99, + ), + } + tracker = ModelUsageTracker() + monkeypatch.setattr(bootstrap, "get_model_usage_store", lambda: fake_store) + monkeypatch.setattr(bootstrap, "get_model_usage_tracker", lambda: tracker) + + with patch("asyncio.create_task") as mock_create_task: + bootstrap._start_model_usage_persistence() + # Close the coroutine object passed to the mock so pytest doesn't warn + # about "coroutine was never awaited" — it was never meant to run here. + mock_create_task.call_args[0][0].close() + + assert tracker.get("deepseek", "deepseek-v4-flash").total_tokens == 99 + mock_create_task.assert_called_once() + + bootstrap._stop_model_usage_persistence() # reset the module-level task handle + + +def test_start_model_usage_persistence_is_a_no_op_without_a_store(monkeypatch): + """No store configured (json backend) — startup must not touch asyncio or the tracker.""" + monkeypatch.setattr(bootstrap, "get_model_usage_store", lambda: None) + + with patch("asyncio.create_task") as mock_create_task: + bootstrap._start_model_usage_persistence() + + mock_create_task.assert_not_called() + + +def test_stop_model_usage_persistence_cancels_task_and_flushes(monkeypatch): + """Shutdown must cancel the running flush task and perform one final flush.""" + fake_store = MagicMock() + monkeypatch.setattr(bootstrap, "get_model_usage_store", lambda: fake_store) + fake_task = MagicMock() + bootstrap._model_usage_flush_task = fake_task + + bootstrap._stop_model_usage_persistence() + + fake_task.cancel.assert_called_once() + fake_store.flush.assert_called_once() + assert bootstrap._model_usage_flush_task is None