feat: seed and periodically persist model usage stats to Postgres

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
This commit is contained in:
wangwei
2026-07-23 14:32:21 +08:00
co-authored by Copilot
parent 5d132981ad
commit 29f79d7434
2 changed files with 183 additions and 0 deletions
+72
View File
@@ -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)