Merge pull request #32255 from BerriAI/litellm_fix_openrouter_streaming_usage_cost

fix(streaming): use provider-reported usage cost for OpenRouter streams

litellm/litellm_core_utils/streaming_chunk_builder_utils.py

@@ -467,6 +467,7 @@ class ChunkProcessor:
         cache_read_input_tokens: Optional[int] = None
         completion_tokens_details: Optional[CompletionTokensDetails] = None
         prompt_tokens_details: Optional[PromptTokensDetailsWrapper] = None
+        cost: Optional[float] = None
 
         if "prompt_tokens" in usage_chunk:
             prompt_tokens = usage_chunk.get("prompt_tokens", 0) or 0
@@ -476,6 +477,8 @@ class ChunkProcessor:
             cache_creation_input_tokens = usage_chunk.get("cache_creation_input_tokens")
         if "cache_read_input_tokens" in usage_chunk:
             cache_read_input_tokens = usage_chunk.get("cache_read_input_tokens")
+        if "cost" in usage_chunk:
+            cost = usage_chunk.get("cost")
         if hasattr(usage_chunk, "completion_tokens_details"):
             if isinstance(usage_chunk.completion_tokens_details, dict):
                 completion_tokens_details = CompletionTokensDetails(**usage_chunk.completion_tokens_details)
@@ -494,6 +497,7 @@ class ChunkProcessor:
             "cache_read_input_tokens": cache_read_input_tokens,
             "completion_tokens_details": completion_tokens_details,
             "prompt_tokens_details": prompt_tokens_details,
+            "cost": cost,
         }
 
     def count_reasoning_tokens(self, response: ModelResponse) -> Optional[int]:
@@ -512,6 +516,22 @@ class ChunkProcessor:
 
         return reasoning_tokens
 
+    @staticmethod
+    def _extract_usage_chunk(chunk: dict[str, Any] | ModelResponse | ModelResponseStream) -> Usage | None:
+        usage_chunk: Usage | dict[str, Any] | None = None
+        if hasattr(chunk, "usage") and chunk.usage is not None:
+            usage_chunk = chunk.usage
+        elif "usage" in chunk:
+            usage_chunk = chunk["usage"]
+        elif (isinstance(chunk, ModelResponse) or isinstance(chunk, ModelResponseStream)) and hasattr(
+            chunk, "_hidden_params"
+        ):
+            usage_chunk = chunk._hidden_params.get("usage", None)
+
+        if isinstance(usage_chunk, dict):
+            return Usage(**usage_chunk)
+        return usage_chunk
+
     def _calculate_usage_per_chunk(
         self,
         chunks: List[Union[Dict[str, Any], ModelResponse]],
@@ -548,18 +568,12 @@ class ChunkProcessor:
         # is last-wins, so without preserving this separately the 1h breakdown is
         # lost and 1h cache writes get billed at the 5m rate.
         cache_creation_token_details: Optional[CacheCreationTokenDetails] = None
+        cost: Optional[float] = None
+
         for chunk in chunks:
-            usage_chunk: Optional[Usage] = None
-            if "usage" in chunk:
-                usage_chunk = chunk["usage"]
-            elif (isinstance(chunk, ModelResponse) or isinstance(chunk, ModelResponseStream)) and hasattr(
-                chunk, "_hidden_params"
-            ):
-                usage_chunk = chunk._hidden_params.get("usage", None)
+            usage_chunk = self._extract_usage_chunk(chunk)
 
             if usage_chunk is not None:
-                if isinstance(usage_chunk, dict):
-                    usage_chunk = Usage(**usage_chunk)
                 usage_chunk_dict = self._usage_chunk_calculation_helper(usage_chunk)
                 if usage_chunk_dict["prompt_tokens"] is not None and usage_chunk_dict["prompt_tokens"] > 0:
                     prompt_tokens = usage_chunk_dict["prompt_tokens"]
@@ -610,6 +624,9 @@ class ChunkProcessor:
                     prompt_tokens_details, cache_creation_token_details
                 )
 
+                if usage_chunk_dict["cost"] is not None:
+                    cost = usage_chunk_dict["cost"]
+
         prompt_tokens_details = self._attach_cache_creation_token_details(
             prompt_tokens_details, cache_creation_token_details
         )
@@ -629,6 +646,7 @@ class ChunkProcessor:
             web_search_requests=web_search_requests,
             completion_tokens_details=completion_tokens_details,
             prompt_tokens_details=prompt_tokens_details,
+            cost=cost,
         )
 
     @staticmethod
@@ -727,6 +745,7 @@ class ChunkProcessor:
         prompt_tokens_details: Optional[PromptTokensDetailsWrapper] = calculated_usage_per_chunk[
             "prompt_tokens_details"
         ]
+        cost: Optional[float] = calculated_usage_per_chunk["cost"]
 
         try:
             returned_usage.prompt_tokens = prompt_tokens or token_counter(model=model, messages=messages)
@@ -784,6 +803,9 @@ class ChunkProcessor:
             else:
                 returned_usage.prompt_tokens_details.web_search_requests = web_search_requests
 
+        if cost is not None:
+            setattr(returned_usage, "cost", cost)
+
         # Return a new usage object with the new values
 
         returned_usage = Usage(**returned_usage.model_dump())

litellm/litellm_core_utils/streaming_handler.py

@@ -962,10 +962,11 @@ class CustomStreamWrapper:
                 if self.custom_llm_provider == "bedrock" and "trace" in model_response:
                     return model_response
 
-                # Default - return StopIteration
-                if hasattr(model_response, "usage"):
-                    self.chunks.append(model_response)
-                raise StopIteration
+                # Don't raise StopIteration here - some providers (like OpenRouter)
+                # send usage/cost data in chunks after the finish_reason chunk
+                if hasattr(model_response, "usage") and model_response.usage is not None:
+                    return model_response
+                return
             # flush any remaining holding chunk
             if len(self.holding_chunk) > 0:
                 if model_response.choices[0].delta.content is None:
@@ -1474,12 +1475,16 @@ class CustomStreamWrapper:
 
                 self.tool_call = True
 
+            if hasattr(chunk, "usage") and chunk.usage is not None:
+                model_response.usage = chunk.usage
+
             ## RETURN ARG
-            return self.return_processed_chunk_logic(
+            result = self.return_processed_chunk_logic(
                 completion_obj=completion_obj,
                 model_response=model_response,  # type: ignore
                 response_obj=response_obj,
             )
+            return result
 
         except StopIteration:
             raise StopIteration
@@ -1686,6 +1691,21 @@ class CustomStreamWrapper:
             model_response.choices[0].finish_reason = "tool_calls"
         return model_response
 
+    @staticmethod
+    def _propagate_usage_cost_to_hidden_params(
+        response: "ModelResponse",
+    ) -> None:
+        """
+        If the assembled response carries a provider-reported cost on
+        usage.cost, copy it into _hidden_params so litellm's cost
+        calculator uses it instead of a token-based estimate.
+        """
+        _usage = getattr(response, "usage", None)
+        if _usage is not None and hasattr(_usage, "cost") and _usage.cost is not None:
+            if "additional_headers" not in response._hidden_params:
+                response._hidden_params["additional_headers"] = {}
+            response._hidden_params["additional_headers"]["llm_provider-x-litellm-response-cost"] = float(_usage.cost)
+
     def __next__(self) -> "ModelResponseStream":
         cache_hit = False
         if self.custom_llm_provider is not None and self.custom_llm_provider == "cached_response":
@@ -1741,6 +1761,10 @@ class CustomStreamWrapper:
                     # hasattr(response, "usage") is always True — must check
                     # `is not None` to avoid running this path on every chunk.
                     if getattr(response, "usage", None) is not None:
+                        usage_to_preserve = response.usage
+                        if usage_to_preserve:
+                            response._hidden_params["usage"] = usage_to_preserve
+
                         obj_dict = response.model_dump()
 
                         if "usage" in obj_dict:
@@ -1789,6 +1813,8 @@ class CustomStreamWrapper:
 
                 response = self.model_response_creator()
                 if complete_streaming_response is not None:
+                    self._propagate_usage_cost_to_hidden_params(complete_streaming_response)
+
                     setattr(
                         response,
                         "usage",
@@ -1999,6 +2025,8 @@ class CustomStreamWrapper:
 
                 response = self.model_response_creator()
                 if complete_streaming_response is not None:
+                    self._propagate_usage_cost_to_hidden_params(complete_streaming_response)
+
                     setattr(
                         response,
                         "usage",
@@ -2228,12 +2256,16 @@ def calculate_total_usage(chunks: List[ModelResponse]) -> Usage:
     """Assume most recent usage chunk has total usage uptil then."""
     prompt_tokens: int = 0
     completion_tokens: int = 0
+    latest_usage_chunk = None
+
     for chunk in chunks:
         if "usage" in chunk and chunk["usage"] is not None:
-            if "prompt_tokens" in chunk["usage"]:
-                prompt_tokens = chunk["usage"].get("prompt_tokens", 0) or 0
-            if "completion_tokens" in chunk["usage"]:
-                completion_tokens = chunk["usage"].get("completion_tokens", 0) or 0
+            usage = chunk["usage"]
+            latest_usage_chunk = usage
+            if "prompt_tokens" in usage:
+                prompt_tokens = usage.get("prompt_tokens", 0) or 0
+            if "completion_tokens" in usage:
+                completion_tokens = usage.get("completion_tokens", 0) or 0
 
     returned_usage_chunk = Usage(
         prompt_tokens=prompt_tokens,
@@ -2241,6 +2273,15 @@ def calculate_total_usage(chunks: List[ModelResponse]) -> Usage:
         total_tokens=prompt_tokens + completion_tokens,
     )
 
+    if latest_usage_chunk is not None:
+        latest_cost = (
+            latest_usage_chunk.get("cost")
+            if isinstance(latest_usage_chunk, dict)
+            else getattr(latest_usage_chunk, "cost", None)
+        )
+        if latest_cost is not None:
+            returned_usage_chunk.cost = latest_cost
+
     return returned_usage_chunk
 
 

litellm/types/litellm_core_utils/streaming_chunk_builder_utils.py

@@ -14,3 +14,4 @@ class UsagePerChunk(TypedDict):
     web_search_requests: Optional[int]
     completion_tokens_details: Optional[CompletionTokensDetails]
     prompt_tokens_details: Optional[PromptTokensDetailsWrapper]
+    cost: Optional[float]

litellm/types/utils.py

@@ -1798,14 +1798,17 @@ class ModelResponseStream(ModelResponseBase):
         else:
             created = created
 
+        usage_to_set = None
         if "usage" in kwargs and kwargs["usage"] is not None:
             if isinstance(kwargs["usage"], dict):
-                kwargs["usage"] = Usage(**kwargs["usage"])
+                usage_to_set = Usage(**kwargs["usage"])
+                kwargs["usage"] = usage_to_set
             elif isinstance(kwargs["usage"], BaseModel):
                 dump = (
                     kwargs["usage"].model_dump() if hasattr(kwargs["usage"], "model_dump") else kwargs["usage"].dict()
                 )
-                kwargs["usage"] = Usage(**dump)
+                usage_to_set = Usage(**dump)
+                kwargs["usage"] = usage_to_set
 
         kwargs["id"] = id
         kwargs["created"] = created
@@ -1814,6 +1817,9 @@ class ModelResponseStream(ModelResponseBase):
 
         super().__init__(**kwargs)
 
+        if usage_to_set is not None:
+            self.usage = usage_to_set
+
     def __contains__(self, key):
         # Define custom behavior for the 'in' operator
         return hasattr(self, key)

tests/test_litellm/litellm_core_utils/test_streaming_chunk_builder_utils.py

@@ -956,3 +956,39 @@ def test_stream_chunk_builder_propagates_vertex_ai_metadata_from_dict_chunks():
     assert response.model_dump()["vertex_ai_grounding_metadata"] == [
         {"webSearchQueries": ["test query"]}
     ]
+
+
+def test_cost_field_in_usage_chunks():
+    chunk1_usage = Usage(completion_tokens=1, prompt_tokens=10, total_tokens=11)
+    chunk1 = ModelResponseStream(
+        id="chatcmpl-1",
+        created=1745513206,
+        model="openrouter/claude",
+        choices=[
+            StreamingChoices(finish_reason=None, index=0, delta=Delta(content="Hi"))
+        ],
+        usage=chunk1_usage,
+    )
+
+    chunk2_usage = Usage(
+        completion_tokens=5, prompt_tokens=10, total_tokens=15, cost=0.00025
+    )
+    chunk2 = ModelResponseStream(
+        id="chatcmpl-1",
+        created=1745513207,
+        model="openrouter/claude",
+        choices=[
+            StreamingChoices(finish_reason="stop", index=0, delta=Delta(content=""))
+        ],
+        usage=chunk2_usage,
+    )
+
+    processor = ChunkProcessor(chunks=[chunk1, chunk2])
+    usage = processor.calculate_usage(
+        chunks=[chunk1, chunk2], model="openrouter/claude", completion_output="Hi"
+    )
+
+    assert hasattr(usage, "cost")
+    assert usage.cost == 0.00025
+    assert usage.prompt_tokens == 10
+    assert usage.completion_tokens == 5

tests/test_litellm/litellm_core_utils/test_streaming_handler.py

@@ -1393,6 +1393,190 @@ def test_has_any_special_delta_attributes(
     assert result is False
 
 
+def test_calculate_total_usage_with_cost():
+    from litellm.litellm_core_utils.streaming_handler import calculate_total_usage
+
+    chunk1_usage = Usage(completion_tokens=1, prompt_tokens=10, total_tokens=11)
+    chunk1 = ModelResponseStream(
+        id="test-1",
+        created=1745513206,
+        model="openrouter/test",
+        choices=[
+            StreamingChoices(finish_reason=None, index=0, delta=Delta(content="Hi"))
+        ],
+        usage=chunk1_usage,
+    )
+
+    chunk2_usage = Usage(
+        completion_tokens=5, prompt_tokens=10, total_tokens=15, cost=0.00025
+    )
+    chunk2 = ModelResponseStream(
+        id="test-1",
+        created=1745513207,
+        model="openrouter/test",
+        choices=[
+            StreamingChoices(finish_reason="stop", index=0, delta=Delta(content=""))
+        ],
+        usage=chunk2_usage,
+    )
+
+    usage = calculate_total_usage([chunk1, chunk2])
+
+    assert hasattr(usage, "cost")
+    assert usage.cost == 0.00025
+    assert usage.prompt_tokens == 10
+    assert usage.completion_tokens == 5
+
+
+def test_calculate_total_usage_with_dict_usage_cost():
+    """Regression: dict-shaped `usage` with a `cost` key must still surface
+    provider cost even though `hasattr` on a dict does not consult its keys."""
+    from litellm.litellm_core_utils.streaming_handler import calculate_total_usage
+
+    chunk = {
+        "usage": {
+            "prompt_tokens": 10,
+            "completion_tokens": 5,
+            "total_tokens": 15,
+            "cost": 0.00025,
+        }
+    }
+
+    usage = calculate_total_usage([chunk])
+
+    assert usage.prompt_tokens == 10
+    assert usage.completion_tokens == 5
+    assert getattr(usage, "cost", None) == 0.00025
+
+
+@pytest.mark.asyncio
+async def test_openrouter_streaming_cost_after_finish_reason(logging_obj: Logging):
+    from litellm.utils import ModelResponseListIterator
+
+    chunk1 = ModelResponseStream(
+        id="chatcmpl-or",
+        created=1742056047,
+        model="openrouter/claude",
+        choices=[
+            StreamingChoices(
+                finish_reason=None, index=0, delta=Delta(content="Hi", role="assistant")
+            )
+        ],
+        usage=None,
+    )
+    chunk2 = ModelResponseStream(
+        id="chatcmpl-or",
+        created=1742056048,
+        model="openrouter/claude",
+        choices=[
+            StreamingChoices(finish_reason="stop", index=0, delta=Delta(content=""))
+        ],
+        usage=None,
+    )
+    chunk3_usage = Usage(
+        completion_tokens=5, prompt_tokens=10, total_tokens=15, cost=0.00025
+    )
+    chunk3 = ModelResponseStream(
+        id="chatcmpl-or",
+        created=1742056049,
+        model="openrouter/claude",
+        choices=[
+            StreamingChoices(finish_reason=None, index=0, delta=Delta(content=""))
+        ],
+        usage=chunk3_usage,
+    )
+
+    completion_stream = ModelResponseListIterator(
+        model_responses=[chunk1, chunk2, chunk3]
+    )
+    response = CustomStreamWrapper(
+        completion_stream=completion_stream,
+        model="openrouter/claude",
+        custom_llm_provider="openrouter",
+        logging_obj=logging_obj,
+        stream_options={"include_usage": True},
+    )
+
+    collected_chunks = []
+    async for chunk in response:
+        collected_chunks.append(chunk)
+
+    usage_chunks = [c for c in collected_chunks if hasattr(c, "usage") and c.usage]
+    assert len(usage_chunks) > 0
+    assert hasattr(usage_chunks[-1].usage, "cost")
+    assert usage_chunks[-1].usage.cost == 0.00025
+
+
+def test_openrouter_streaming_cost_propagates_to_hidden_params():
+    """
+    Verify that provider-reported cost from usage.cost flows into
+    _hidden_params["additional_headers"]["llm_provider-x-litellm-response-cost"]
+    on the complete streaming response, so litellm's cost calculator uses it.
+    """
+    import litellm
+
+    chunk1 = ModelResponseStream(
+        id="chatcmpl-or",
+        created=1742056047,
+        model="openrouter/claude",
+        choices=[
+            StreamingChoices(
+                finish_reason=None, index=0, delta=Delta(content="Hi", role="assistant")
+            )
+        ],
+        usage=None,
+    )
+    chunk2 = ModelResponseStream(
+        id="chatcmpl-or",
+        created=1742056048,
+        model="openrouter/claude",
+        choices=[
+            StreamingChoices(finish_reason="stop", index=0, delta=Delta(content=""))
+        ],
+        usage=None,
+    )
+    chunk3 = ModelResponseStream(
+        id="chatcmpl-or",
+        created=1742056049,
+        model="openrouter/claude",
+        choices=[
+            StreamingChoices(finish_reason=None, index=0, delta=Delta(content=""))
+        ],
+        usage=Usage(
+            completion_tokens=5, prompt_tokens=10, total_tokens=15, cost=0.00025
+        ),
+    )
+
+    # Build the complete response as stream_chunk_builder does
+    complete_response = litellm.stream_chunk_builder(
+        chunks=[chunk1, chunk2, chunk3],
+        messages=[{"role": "user", "content": "test"}],
+    )
+
+    assert complete_response is not None
+    assert hasattr(complete_response.usage, "cost")
+    assert complete_response.usage.cost == 0.00025
+
+    # Use the real propagation method from CustomStreamWrapper
+    CustomStreamWrapper._propagate_usage_cost_to_hidden_params(complete_response)
+
+    assert "additional_headers" in complete_response._hidden_params
+    assert (
+        complete_response._hidden_params["additional_headers"][
+            "llm_provider-x-litellm-response-cost"
+        ]
+        == 0.00025
+    )
+
+    # Verify the cost calculator would pick this up
+    from litellm.cost_calculator import get_response_cost_from_hidden_params
+
+    provider_cost = get_response_cost_from_hidden_params(
+        complete_response._hidden_params
+    )
+    assert provider_cost == 0.00025
+
+
 def test_handle_special_delta_attributes(
     initialized_custom_stream_wrapper: CustomStreamWrapper,
 ):