feat: record streaming token usage in TrackedLLMClient.stream_chat
Implement manual generator driving using next()/StopIteration to capture the return value (trailing usage dict) from inner stream_chat() implementations, enabling token tracking for streaming LLM calls. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
This commit is contained in:
@@ -56,18 +56,24 @@ class TrackedLLMClient:
|
||||
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.
|
||||
"""Delegate to the wrapped client's stream_chat(), recording call outcome and usage.
|
||||
|
||||
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.
|
||||
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:
|
||||
for chunk in self._inner.stream_chat(messages, *args, **kwargs):
|
||||
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)
|
||||
@@ -77,6 +83,7 @@ class TrackedLLMClient:
|
||||
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,
|
||||
)
|
||||
|
||||
Reference in New Issue
Block a user