diff --git a/tests/local_testing/test_stream_chunk_builder.py b/tests/local_testing/test_stream_chunk_builder.py index 63907eb7d5e..8224773aa4c 100644 --- a/tests/local_testing/test_stream_chunk_builder.py +++ b/tests/local_testing/test_stream_chunk_builder.py @@ -168,56 +168,6 @@ def test_stream_chunk_builder_litellm_tool_call_regular_message(): # test_stream_chunk_builder_litellm_tool_call_regular_message() -def test_stream_chunk_builder_litellm_usage_chunks(): - """ - Checks if stream_chunk_builder is able to correctly rebuild with given metadata from streaming chunks - """ - from litellm.types.utils import Usage - - messages = [ - {"role": "user", "content": "Tell me the funniest joke you know."}, - { - "role": "assistant", - "content": "Why did the chicken cross the road?\nYou will not guess this one I bet\n", - }, - {"role": "user", "content": "I do not know, why?"}, - {"role": "assistant", "content": "uhhhh\n\n\nhmmmm.....\nthinking....\n"}, - {"role": "user", "content": "\nI am waiting...\n\n...\n"}, - ] - - usage: litellm.Usage = Usage( - completion_tokens=27, - prompt_tokens=50, - total_tokens=82, - completion_tokens_details=None, - prompt_tokens_details=None, - ) - - gemini_pt = usage.prompt_tokens - - # make a streaming gemini call - try: - response = completion( - model="gemini/gemini-1.5-flash", - messages=messages, - stream=True, - complete_response=True, - stream_options={"include_usage": True}, - ) - except litellm.InternalServerError as e: - pytest.skip(f"Skipping test due to internal server error - {str(e)}") - - usage: litellm.Usage = response.usage - - stream_rebuilt_pt = usage.prompt_tokens - - # assert prompt tokens are the same - - assert ( - gemini_pt == stream_rebuilt_pt - ), f"Stream builder is not able to rebuild usage correctly. Got={stream_rebuilt_pt}, expected={gemini_pt}" - - def test_stream_chunk_builder_litellm_mixed_calls(): response = stream_chunk_builder(stream_chunk_testdata.chunks) assert ( diff --git a/tests/test_litellm/litellm_core_utils/test_streaming_chunk_builder_utils.py b/tests/test_litellm/litellm_core_utils/test_streaming_chunk_builder_utils.py index facb5cef3cd..f6636874336 100644 --- a/tests/test_litellm/litellm_core_utils/test_streaming_chunk_builder_utils.py +++ b/tests/test_litellm/litellm_core_utils/test_streaming_chunk_builder_utils.py @@ -243,3 +243,85 @@ def test_cache_read_input_tokens_retained(): assert usage.cache_creation_input_tokens == 4 assert usage.cache_read_input_tokens == 11775 assert usage.prompt_tokens_details.cached_tokens == 11775 + + +def test_stream_chunk_builder_litellm_usage_chunks(): + """ + Validate ChunkProcessor.calculate_usage uses provided usage fields from streaming chunks + and reconstructs prompt and completion tokens without making any upstream API calls. + """ + # Prepare two mocked streaming chunks with usage split across them + chunk1 = ModelResponseStream( + id="chatcmpl-mocked-usage-1", + created=1745513206, + model="gemini/gemini-2.5-flash-lite", + object="chat.completion.chunk", + system_fingerprint=None, + choices=[ + StreamingChoices( + finish_reason=None, + index=0, + delta=Delta( + provider_specific_fields=None, + content="", + role=None, + function_call=None, + tool_calls=None, + audio=None, + ), + logprobs=None, + ) + ], + provider_specific_fields=None, + stream_options={"include_usage": True}, + usage=Usage( + completion_tokens=0, + prompt_tokens=50, + total_tokens=50, + completion_tokens_details=None, + prompt_tokens_details=None, + ), + ) + + chunk2 = ModelResponseStream( + id="chatcmpl-mocked-usage-1", + created=1745513207, + model="gemini/gemini-2.5-flash-lite", + object="chat.completion.chunk", + system_fingerprint=None, + choices=[ + StreamingChoices( + finish_reason="stop", + index=0, + delta=Delta( + provider_specific_fields=None, + content=None, + role=None, + function_call=None, + tool_calls=None, + audio=None, + ), + logprobs=None, + ) + ], + provider_specific_fields=None, + stream_options={"include_usage": True}, + usage=Usage( + completion_tokens=27, + prompt_tokens=0, + total_tokens=27, + completion_tokens_details=None, + prompt_tokens_details=None, + ), + ) + + chunks = [chunk1, chunk2] + processor = ChunkProcessor(chunks=chunks) + + usage = processor.calculate_usage( + chunks=chunks, model="gemini/gemini-2.5-flash-lite", completion_output="" + ) + + assert usage.prompt_tokens == 50 + assert usage.completion_tokens == 27 + assert usage.total_tokens == 77