test_stream_chunk_builder_litellm_usage_chunks

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 (

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