feat: add TrackedLLMClient decorator for transparent usage recording
Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
This commit is contained in:
@@ -0,0 +1,86 @@
|
||||
"""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)
|
||||
Reference in New Issue
Block a user