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:
wangwei
2026-07-23 13:42:33 +08:00
co-authored by Copilot
parent f2bd0deeb3
commit 7adc050968
2 changed files with 34 additions and 7 deletions
@@ -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