"""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 only. Token usage is NOT recorded here: none of the current provider stream_chat() implementations parse a trailing usage chunk from the gateway (see the design doc's Known Limitations), so accumulating a token count here would silently be wrong. Only call success/failure and latency are tracked for streaming calls. """ start = time.time() error: Optional[str] = None try: for chunk in self._inner.stream_chat(messages, *args, **kwargs): 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, 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)