diff --git a/litellm/litellm_core_utils/streaming_handler.py b/litellm/litellm_core_utils/streaming_handler.py index 99b1c1a2ab7..07f979a063f 100644 --- a/litellm/litellm_core_utils/streaming_handler.py +++ b/litellm/litellm_core_utils/streaming_handler.py @@ -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( diff --git a/tests/test_litellm/litellm_core_utils/test_streaming_handler.py b/tests/test_litellm/litellm_core_utils/test_streaming_handler.py index 101935cac0a..b7091bec3c8 100644 --- a/tests/test_litellm/litellm_core_utils/test_streaming_handler.py +++ b/tests/test_litellm/litellm_core_utils/test_streaming_handler.py @@ -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, ):