scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-10 22:41:41 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 81

fix(streaming): preserve usage details in usage-only chunks

Browse source
This commit is contained in:
Mohammad Javad Naderi 2026-08-16 14:18:54 +03:30
parent 973329e986
commit 901b08f524
2 changed files with 44 additions and 5 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

6
litellm/litellm_core_utils/streaming_handler.py
View file

@ -1451,11 +1451,7 @@ class CustomStreamWrapper:
setattr(
model_response,
"usage",
litellm.Usage(
prompt_tokens=response_obj["usage"].get("prompt_tokens", None) or None,
completion_tokens=response_obj["usage"].get("completion_tokens", None) or None,
total_tokens=response_obj["usage"].get("total_tokens", None) or None,
),
litellm.Usage(**response_obj["usage"]),
)
elif isinstance(response_obj["usage"], Usage):
setattr(

43
tests/test_litellm/litellm_core_utils/test_streaming_handler.py
View file

@ -2,6 +2,7 @@ import json
import os
import sys
import time
from types import SimpleNamespace
from unittest.mock import AsyncMock, MagicMock, Mock, patch
import pytest
@ -618,6 +619,48 @@ def test_streaming_handler_with_stop_chunk(
assert returned_chunk is None
def test_usage_only_openai_chunk_preserves_usage_details(
initialized_custom_stream_wrapper: CustomStreamWrapper,
):
initialized_custom_stream_wrapper.custom_llm_provider = "openai"
initialized_custom_stream_wrapper.stream_options = {"include_usage": True}
chunk = SimpleNamespace(
id="chatcmpl-test",
model="glm-4.7",
choices=[],
usage={
"completion_tokens": 28,
"prompt_tokens": 3721,
"total_tokens": 3749,
"completion_tokens_details": {
"accepted_prediction_tokens": None,
"audio_tokens": 0,
"reasoning_tokens": 0,
"rejected_prediction_tokens": None,
"image_tokens": 0,
},
"prompt_tokens_details": {
"audio_tokens": 0,
"cache_write_tokens": 0,
"cached_tokens": 3237,
"video_tokens": 0,
},
"cost": 0.0003777566,
"is_byok": False,
"cost_details": {
"upstream_inference_cost": 0.0003777566,
"upstream_inference_prompt_cost": 0.0003494206,
"upstream_inference_completions_cost": 0.000028336,
},
},
)
response = initialized_custom_stream_wrapper.chunk_creator(chunk)
assert response.usage.prompt_tokens_details.cached_tokens == 3237
assert response.usage.completion_tokens_details.reasoning_tokens == 0
def test_finish_reason_chunk_preserves_non_openai_attributes(
initialized_custom_stream_wrapper: CustomStreamWrapper,
):