scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-14 23:21:35 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 156

fix(streaming): include reasoning tokens in estimated usage

Browse source
This commit is contained in:
Calvin Div 2026-08-29 00:01:40 +08:00
parent ca0b951a43
commit 68ca902ba9
2 changed files with 40 additions and 1 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

4
litellm/litellm_core_utils/streaming_chunk_builder_utils.py
View file

@ -981,14 +981,16 @@ class ChunkProcessor:
except Exception: # don't allow this failing to block a complete streaming response from being returned
print_verbose("token_counter failed, assuming prompt tokens is 0")
returned_usage.prompt_tokens = 0
returned_usage.completion_tokens = (
resolved_completion_tokens: Final = (
completion_tokens
or token_counter(
model=model,
text=completion_output,
count_response_tokens=True, # count_response_tokens is a Flag to tell token counter this is a response, No need to add extra tokens we do for input messages
)
+ (reasoning_tokens or 0)
)
returned_usage.completion_tokens = resolved_completion_tokens
returned_usage.total_tokens = returned_usage.prompt_tokens + returned_usage.completion_tokens
if cache_creation_input_tokens is not None:

37
tests/test_litellm/litellm_core_utils/test_streaming_chunk_builder_utils.py
View file

@ -1391,6 +1391,43 @@ def test_count_reasoning_tokens_counts_visible_reasoning():
assert processor.count_reasoning_tokens(response) > 0
def test_stream_chunk_builder_includes_reasoning_in_estimated_completion_tokens():
reasoning_content = "let me count the primes under thirty"
visible_content = "10"
chunks = [
ModelResponseStream(
id="chatcmpl-estimated-reasoning",
model="claude-opus-4-8",
choices=[
StreamingChoices(
finish_reason=None,
index=0,
delta=Delta(role="assistant", reasoning_content=reasoning_content),
)
],
),
ModelResponseStream(
id="chatcmpl-estimated-reasoning",
model="claude-opus-4-8",
choices=[
StreamingChoices(
finish_reason="stop",
index=0,
delta=Delta(content=visible_content),
)
],
),
]
response = stream_chunk_builder(chunks=chunks, messages=[{"role": "user", "content": "count"}])
assert response is not None
reasoning_tokens = response.usage.completion_tokens_details.reasoning_tokens
assert reasoning_tokens is not None
assert response.usage.completion_tokens > reasoning_tokens
assert response.usage.total_tokens == response.usage.prompt_tokens + response.usage.completion_tokens
@pytest.mark.parametrize(
"estimated_reasoning_tokens, expected_reasoning_tokens, expected_text_tokens",
[(40, 40, 60), (250, 100, 0)],