AIRegulation-DocAnalysis/backend/app/shared/model_usage_tracker.py

"""In-memory registry that tracks per-model call outcomes and token usage.

This module lives in `app/shared` — the same cross-cutting-support tier as
`bootstrap.py` — because it is not business logic: it exists purely so the
System Status page can show which AI models (main LLM, HyDE LLM, embedding,
reranker) are configured, whether their most recent call succeeded, and how
many tokens they have consumed since this process started. Tracking here
must never disrupt a real user-facing call: every public method swallows its
own exceptions and logs a warning instead of raising.
"""

from __future__ import annotations

import threading
from dataclasses import dataclass
from datetime import datetime, timezone
from functools import lru_cache

from loguru import logger


@dataclass
class ModelUsageEntry:
    """Represent accumulated usage/connection state for one provider+model pair."""

    provider: str
    model: str
    total_tokens: int = 0
    prompt_tokens: int = 0
    completion_tokens: int = 0
    call_count_ok: int = 0
    call_count_error: int = 0
    last_called_at: datetime | None = None
    last_latency_ms: int | None = None
    last_error: str | None = None

    @property
    def status(self) -> str:
        """Derive never_called/ok/error from call history.

        The "disabled" status (reranker only, when turned off in settings) is
        NOT decided here: this dataclass has no access to live settings. The
        API route layer (Task 6) applies that override on top of this value,
        so config always wins over stale historical data.
        """
        if self.last_called_at is None:
            return "never_called"
        return "error" if self.last_error else "ok"


class ModelUsageTracker:
    """Thread-safe in-memory registry of per-model call/usage stats.

    Keyed by "{provider}:{model}" rather than by business role (main LLM /
    HyDE / embedding / reranker) so that any future call site is captured
    automatically, even before anyone teaches this class about its role.
    """

    def __init__(self) -> None:
        """Initialize an empty registry guarded by a single lock."""
        self._entries: dict[str, ModelUsageEntry] = {}
        # One coarse lock is enough: record() runs at most a few times per
        # request, and snapshot() is only read by the low-traffic status page.
        self._lock = threading.Lock()

    def record(
        self,
        *,
        provider: str,
        model: str,
        success: bool,
        usage: dict | None = None,
        latency_ms: int | None = None,
        error: str | None = None,
    ) -> None:
        """Record the outcome of one call to provider/model.

        Never raises: any internal failure is logged and swallowed so a bug
        in observability code cannot break a real LLM/embedding/reranker call.
        """
        try:
            key = f"{provider}:{model}"
            usage = usage if isinstance(usage, dict) else {}
            with self._lock:
                entry = self._entries.setdefault(key, ModelUsageEntry(provider=provider, model=model))
                entry.total_tokens += int(usage.get("total_tokens", 0) or 0)
                entry.prompt_tokens += int(usage.get("prompt_tokens", 0) or 0)
                entry.completion_tokens += int(usage.get("completion_tokens", 0) or 0)
                if success:
                    entry.call_count_ok += 1
                    entry.last_error = None
                else:
                    entry.call_count_error += 1
                    entry.last_error = error or "unknown error"
                entry.last_called_at = datetime.now(timezone.utc)
                entry.last_latency_ms = latency_ms
        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:
            return dict(self._entries)

    def get(self, provider: str, model: str) -> ModelUsageEntry | None:
        """Return the entry for one provider/model pair, or None if never recorded."""
        return self.snapshot().get(f"{provider}:{model}")


@lru_cache
def get_model_usage_tracker() -> ModelUsageTracker:
    """Return the process-wide singleton tracker (mirrors get_settings()/get_llm_factory())."""
    return ModelUsageTracker()