87 lines
3.5 KiB
Python
87 lines
3.5 KiB
Python
|
"""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)
|