AIRegulation-DocAnalysis/backend/app/services/llm/tracked_client.py

"""Transparent decorator around BaseLLMClient implementations.

Records per-call token usage, latency, and success/failure into a
ModelUsageTracker without changing any caller-visible behavior.
"""

from __future__ import annotations

import time
from typing import Any, Dict, List, Optional

from app.shared.model_usage_tracker import ModelUsageTracker

from .base_client import BaseLLMClient, LLMResponse
from .tool_types import Tool


class TrackedLLMClient:
    """Wrap any BaseLLMClient and record its usage into a ModelUsageTracker.

    Deliberately does NOT subclass BaseLLMClient: that ABC declares abstract
    methods (_init_client, get_available_models) with no meaningful override
    here, and subclassing would make Python refuse to instantiate this class
    ("Can't instantiate abstract class") before __getattr__ ever got a chance
    to forward the call. Plain composition + __getattr__ delegation works
    because every caller in this codebase only ever uses duck-typed access:
    .chat(), .stream_chat(), .get_available_models(), .close(), .config.
    """

    def __init__(self, inner: BaseLLMClient, tracker: ModelUsageTracker) -> None:
        """Store the wrapped client and the tracker to report into."""
        self._inner = inner
        self._tracker = tracker

    def chat(
        self,
        messages: List[Dict[str, str]],
        max_tokens: Optional[int] = None,
        temperature: Optional[float] = None,
        tools: Optional[List[Tool]] = None,
        **kwargs: Any,
    ) -> LLMResponse:
        """Delegate to the wrapped client's chat(), then record the outcome."""
        start = time.time()
        response = self._inner.chat(messages, max_tokens, temperature, tools, **kwargs)
        # Key by the *configured* model, not response.model, so lookups driven
        # by settings (llm_model / hyde_llm_model) always match what we recorded.
        self._tracker.record(
            provider=self._inner.config.provider.value,
            model=self._inner.config.model,
            success=response.is_success,
            usage=response.usage,
            latency_ms=int((time.time() - start) * 1000),
            error=response.error,
        )
        return response

    def stream_chat(self, messages: List[Dict[str, str]], *args: Any, **kwargs: Any):
        """Delegate to the wrapped client's stream_chat(), recording call outcome and usage.

        Drives the inner generator manually (instead of a plain `for` loop) so
        it can capture the generator's return value via StopIteration.value —
        the trailing token-usage dict the inner client captures from a
        stream_options.include_usage chunk, if the gateway sent one.
        """
        start = time.time()
        error: Optional[str] = None
        usage: Optional[Dict[str, int]] = None
        gen = self._inner.stream_chat(messages, *args, **kwargs)
        try:
            while True:
                try:
                    chunk = next(gen)
                except StopIteration as stop:
                    usage = stop.value
                    break
                yield chunk
        except Exception as exc:  # noqa: BLE001 - report, then re-raise unchanged
            error = str(exc)
            raise
        finally:
            self._tracker.record(
                provider=self._inner.config.provider.value,
                model=self._inner.config.model,
                success=error is None,
                usage=usage,
                latency_ms=int((time.time() - start) * 1000),
                error=error,
            )

    def __getattr__(self, name: str) -> Any:
        """Forward any other attribute/method access to the wrapped client."""
        return getattr(self._inner, name)