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:
@@ -83,3 +83,23 @@ def test_stream_chat_records_call_without_token_usage():
|
||||
entry = tracker.get("deepseek", "deepseek-v4-flash")
|
||||
assert entry.call_count_ok == 1
|
||||
assert entry.total_tokens == 0
|
||||
|
||||
|
||||
def test_stream_chat_records_usage_from_generator_return_value():
|
||||
"""stream_chat() must forward the inner generator's returned usage dict to record()."""
|
||||
inner = _make_inner()
|
||||
|
||||
def fake_stream(*args, **kwargs):
|
||||
yield "chunk-1"
|
||||
yield "chunk-2"
|
||||
return {"prompt_tokens": 6, "completion_tokens": 2, "total_tokens": 8}
|
||||
|
||||
inner.stream_chat.side_effect = fake_stream
|
||||
tracker = ModelUsageTracker()
|
||||
|
||||
chunks = list(TrackedLLMClient(inner, tracker).stream_chat([{"role": "user", "content": "hi"}]))
|
||||
|
||||
assert chunks == ["chunk-1", "chunk-2"]
|
||||
entry = tracker.get("deepseek", "deepseek-v4-flash")
|
||||
assert entry.total_tokens == 8
|
||||
assert entry.call_count_ok == 1
|
||||
|
||||
Reference in New Issue
Block a user