From 30553f7cb6569c65dfd57cca72b8b712f4dc6d94 Mon Sep 17 00:00:00 2001 From: Dhruv Yadav Date: Sun, 2 Nov 2025 02:02:00 +0530 Subject: [PATCH 1/4] fix openrouter streaming usage --- .../streaming_chunk_builder_utils.py | 18 ++++- .../litellm_core_utils/streaming_handler.py | 39 +++++++--- .../streaming_chunk_builder_utils.py | 1 + litellm/types/utils.py | 10 ++- .../test_streaming_chunk_builder_utils.py | 43 +++++++++-- .../test_streaming_handler.py | 77 +++++++++++++++++++ 6 files changed, 169 insertions(+), 19 deletions(-) diff --git a/litellm/litellm_core_utils/streaming_chunk_builder_utils.py b/litellm/litellm_core_utils/streaming_chunk_builder_utils.py index 1935372e5df..bcd1286468b 100644 --- a/litellm/litellm_core_utils/streaming_chunk_builder_utils.py +++ b/litellm/litellm_core_utils/streaming_chunk_builder_utils.py @@ -468,6 +468,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 @@ -477,6 +478,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( @@ -503,6 +506,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]: @@ -540,9 +544,13 @@ class ChunkProcessor: web_search_requests: Optional[int] = None completion_tokens_details: Optional[CompletionTokensDetails] = None prompt_tokens_details: Optional[PromptTokensDetailsWrapper] = None + cost: Optional[float] = None + for chunk in chunks: usage_chunk: Optional[Usage] = None - if "usage" in chunk: + 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) @@ -601,6 +609,9 @@ class ChunkProcessor: prompt_tokens_details = usage_chunk_dict["prompt_tokens_details"] + if usage_chunk_dict["cost"] is not None: + cost = usage_chunk_dict["cost"] + return UsagePerChunk( prompt_tokens=prompt_tokens, completion_tokens=completion_tokens, @@ -610,6 +621,7 @@ class ChunkProcessor: web_search_requests=web_search_requests, completion_tokens_details=completion_tokens_details, prompt_tokens_details=prompt_tokens_details, + cost=cost, ) def calculate_usage( @@ -649,6 +661,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( @@ -717,6 +730,9 @@ class ChunkProcessor: 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()) diff --git a/litellm/litellm_core_utils/streaming_handler.py b/litellm/litellm_core_utils/streaming_handler.py index 1bb2b99c015..343ffb91e53 100644 --- a/litellm/litellm_core_utils/streaming_handler.py +++ b/litellm/litellm_core_utils/streaming_handler.py @@ -1024,10 +1024,12 @@ class CustomStreamWrapper: if self.custom_llm_provider == "bedrock" and "trace" in model_response: return model_response - # Default - return StopIteration - if hasattr(model_response, "usage"): + # 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: self.chunks.append(model_response) - raise StopIteration + return model_response + return # flush any remaining holding chunk if len(self.holding_chunk) > 0: if model_response.choices[0].delta.content is None: @@ -1140,7 +1142,11 @@ class CustomStreamWrapper: not isinstance(chunk, dict) or "provider_specific_fields" not in chunk ): - raise StopIteration + if self.custom_llm_provider != "openrouter": + raise StopIteration + else: + # OpenRouter: continue processing - usage will come in later chunks + pass anthropic_response_obj: GChunk = cast(GChunk, chunk) completion_obj["content"] = anthropic_response_obj["text"] if anthropic_response_obj["is_finished"]: @@ -1573,12 +1579,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 @@ -1878,6 +1888,10 @@ class CustomStreamWrapper: if hasattr( response, "usage" ): # remove usage from chunk, only send on final chunk + usage_to_preserve = response.usage + if usage_to_preserve: + response._hidden_params["usage"] = usage_to_preserve + # Convert the object to a dictionary obj_dict = response.model_dump() @@ -2355,12 +2369,19 @@ 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 + + if latest_usage_chunk and hasattr(latest_usage_chunk, "cost") and latest_usage_chunk.cost is not None: + return latest_usage_chunk returned_usage_chunk = Usage( prompt_tokens=prompt_tokens, diff --git a/litellm/types/litellm_core_utils/streaming_chunk_builder_utils.py b/litellm/types/litellm_core_utils/streaming_chunk_builder_utils.py index a1f89dac5cf..c9f9d4e6baa 100644 --- a/litellm/types/litellm_core_utils/streaming_chunk_builder_utils.py +++ b/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] diff --git a/litellm/types/utils.py b/litellm/types/utils.py index 82557513a8a..2ac3c1f1110 100644 --- a/litellm/types/utils.py +++ b/litellm/types/utils.py @@ -1826,16 +1826,19 @@ 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 @@ -1844,6 +1847,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) 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 c86e146b0ef..8ee4a656635 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 @@ -402,7 +402,7 @@ def test_stream_chunk_builder_litellm_usage_chunks(): def test_get_model_from_chunks_azure_model_router(): """ Test that _get_model_from_chunks finds the actual model from Azure Model Router chunks. - + Azure Model Router returns the request model (e.g., 'azure-model-router') in the first chunk, but subsequent chunks contain the actual model (e.g., 'gpt-4.1-nano-2025-04-14'). This is important for accurate cost calculation. @@ -413,24 +413,24 @@ def test_get_model_from_chunks_azure_model_router(): {"model": "gpt-4.1-nano-2025-04-14", "id": "chatcmpl-123", "choices": []}, {"model": "gpt-4.1-nano-2025-04-14", "id": "chatcmpl-123", "choices": []}, ] - + result = ChunkProcessor._get_model_from_chunks( chunks=chunks, first_chunk_model="azure-model-router" ) - + # Should return the actual model, not the request model assert result == "gpt-4.1-nano-2025-04-14" - + # Test when all chunks have the same model (non-router case) chunks_same_model = [ {"model": "gpt-4", "id": "chatcmpl-456", "choices": []}, {"model": "gpt-4", "id": "chatcmpl-456", "choices": []}, ] - + result_same = ChunkProcessor._get_model_from_chunks( chunks=chunks_same_model, first_chunk_model="gpt-4" ) - + # Should return the first chunk's model when all are the same assert result_same == "gpt-4" @@ -511,7 +511,7 @@ def test_stream_chunk_builder_anthropic_web_search(): assert usage.prompt_tokens == 50 assert usage.completion_tokens == 27 - assert usage.total_tokens == 77 + assert usage.total_tokens == 77 assert usage.server_tool_use['web_search_requests'] == 2 @@ -603,3 +603,32 @@ def test_stream_chunk_builder_dict_snapshot_preserves_hidden_provider_fields(): response = stream_chunk_builder(chunks=[chunk_dict]) assert response is not None assert response._hidden_params["provider_specific_fields"]["traffic_type"] == "default" + + +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) + setattr(chunk2_usage, "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 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 20e064ef8f4..f15034e7865 100644 --- a/tests/test_litellm/litellm_core_utils/test_streaming_handler.py +++ b/tests/test_litellm/litellm_core_utils/test_streaming_handler.py @@ -1161,6 +1161,83 @@ 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) + setattr(chunk2_usage, "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 + + +@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) + setattr(chunk3_usage, "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_handle_special_delta_attributes( initialized_custom_stream_wrapper: CustomStreamWrapper, ): From 387a94e4228dc4c71322d8bf7207edbf6818eac1 Mon Sep 17 00:00:00 2001 From: Dhruv Yadav Date: Tue, 31 Mar 2026 15:06:42 +0530 Subject: [PATCH 2/4] propagate streaming usage.cost to _hidden_params for cost tracking Without this, the provider-reported cost (e.g. from OpenRouter) was available on usage.cost but never reached litellm's cost calculator, which reads from _hidden_params["additional_headers"]. Also cleans up setattr usage in tests since Usage already has a cost field. --- .../litellm_core_utils/streaming_handler.py | 20 ++++++ .../test_streaming_chunk_builder_utils.py | 3 +- .../test_streaming_handler.py | 70 +++++++++++++++++-- 3 files changed, 87 insertions(+), 6 deletions(-) diff --git a/litellm/litellm_core_utils/streaming_handler.py b/litellm/litellm_core_utils/streaming_handler.py index 343ffb91e53..c6b93e91b15 100644 --- a/litellm/litellm_core_utils/streaming_handler.py +++ b/litellm/litellm_core_utils/streaming_handler.py @@ -1930,6 +1930,16 @@ class CustomStreamWrapper: response = self.model_response_creator() if complete_streaming_response is not None: + # Propagate provider-reported cost (e.g. OpenRouter) + # to _hidden_params so the cost calculator picks it up + _final_usage = getattr(complete_streaming_response, "usage", None) + if _final_usage is not None and hasattr(_final_usage, "cost") and _final_usage.cost is not None: + if "additional_headers" not in complete_streaming_response._hidden_params: + complete_streaming_response._hidden_params["additional_headers"] = {} + complete_streaming_response._hidden_params["additional_headers"][ + "llm_provider-x-litellm-response-cost" + ] = float(_final_usage.cost) + setattr( response, "usage", @@ -2156,6 +2166,16 @@ class CustomStreamWrapper: response = self.model_response_creator() if complete_streaming_response is not None: + # Propagate provider-reported cost (e.g. OpenRouter) + # to _hidden_params so the cost calculator picks it up + _final_usage = getattr(complete_streaming_response, "usage", None) + if _final_usage is not None and hasattr(_final_usage, "cost") and _final_usage.cost is not None: + if "additional_headers" not in complete_streaming_response._hidden_params: + complete_streaming_response._hidden_params["additional_headers"] = {} + complete_streaming_response._hidden_params["additional_headers"][ + "llm_provider-x-litellm-response-cost" + ] = float(_final_usage.cost) + setattr( response, "usage", 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 8ee4a656635..f27955691a0 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 @@ -615,8 +615,7 @@ def test_cost_field_in_usage_chunks(): usage=chunk1_usage, ) - chunk2_usage = Usage(completion_tokens=5, prompt_tokens=10, total_tokens=15) - setattr(chunk2_usage, "cost", 0.00025) + chunk2_usage = Usage(completion_tokens=5, prompt_tokens=10, total_tokens=15, cost=0.00025) chunk2 = ModelResponseStream( id="chatcmpl-1", created=1745513207, 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 f15034e7865..6b1410a1b82 100644 --- a/tests/test_litellm/litellm_core_utils/test_streaming_handler.py +++ b/tests/test_litellm/litellm_core_utils/test_streaming_handler.py @@ -1173,8 +1173,7 @@ def test_calculate_total_usage_with_cost(): usage=chunk1_usage, ) - chunk2_usage = Usage(completion_tokens=5, prompt_tokens=10, total_tokens=15) - setattr(chunk2_usage, "cost", 0.00025) + chunk2_usage = Usage(completion_tokens=5, prompt_tokens=10, total_tokens=15, cost=0.00025) chunk2 = ModelResponseStream( id="test-1", created=1745513207, @@ -1209,8 +1208,7 @@ async def test_openrouter_streaming_cost_after_finish_reason(logging_obj: Loggin choices=[StreamingChoices(finish_reason="stop", index=0, delta=Delta(content=""))], usage=None, ) - chunk3_usage = Usage(completion_tokens=5, prompt_tokens=10, total_tokens=15) - setattr(chunk3_usage, "cost", 0.00025) + chunk3_usage = Usage(completion_tokens=5, prompt_tokens=10, total_tokens=15, cost=0.00025) chunk3 = ModelResponseStream( id="chatcmpl-or", created=1742056049, @@ -1238,6 +1236,70 @@ async def test_openrouter_streaming_cost_after_finish_reason(logging_obj: Loggin 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 + + # Simulate the propagation logic from streaming_handler + _final_usage = getattr(complete_response, "usage", None) + if _final_usage is not None and hasattr(_final_usage, "cost") and _final_usage.cost is not None: + if "additional_headers" not in complete_response._hidden_params: + complete_response._hidden_params["additional_headers"] = {} + complete_response._hidden_params["additional_headers"][ + "llm_provider-x-litellm-response-cost" + ] = float(_final_usage.cost) + + 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, ): From df69af27d68e07ac2e11212591383a00a614f6ea Mon Sep 17 00:00:00 2001 From: Dhruv Yadav Date: Tue, 31 Mar 2026 15:11:03 +0530 Subject: [PATCH 3/4] apply black formatting --- .../streaming_chunk_builder_utils.py | 2 +- .../litellm_core_utils/streaming_handler.py | 57 ++++++-- .../test_streaming_chunk_builder_utils.py | 34 +++-- .../test_streaming_handler.py | 129 +++++++++++++----- 4 files changed, 163 insertions(+), 59 deletions(-) diff --git a/litellm/litellm_core_utils/streaming_chunk_builder_utils.py b/litellm/litellm_core_utils/streaming_chunk_builder_utils.py index bcd1286468b..19a729b6d02 100644 --- a/litellm/litellm_core_utils/streaming_chunk_builder_utils.py +++ b/litellm/litellm_core_utils/streaming_chunk_builder_utils.py @@ -545,7 +545,7 @@ class ChunkProcessor: completion_tokens_details: Optional[CompletionTokensDetails] = None prompt_tokens_details: Optional[PromptTokensDetailsWrapper] = None cost: Optional[float] = None - + for chunk in chunks: usage_chunk: Optional[Usage] = None if hasattr(chunk, "usage") and chunk.usage is not None: diff --git a/litellm/litellm_core_utils/streaming_handler.py b/litellm/litellm_core_utils/streaming_handler.py index c6b93e91b15..432f94540af 100644 --- a/litellm/litellm_core_utils/streaming_handler.py +++ b/litellm/litellm_core_utils/streaming_handler.py @@ -1026,7 +1026,10 @@ class CustomStreamWrapper: # 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: + if ( + hasattr(model_response, "usage") + and model_response.usage is not None + ): self.chunks.append(model_response) return model_response return @@ -1933,12 +1936,23 @@ class CustomStreamWrapper: # Propagate provider-reported cost (e.g. OpenRouter) # to _hidden_params so the cost calculator picks it up _final_usage = getattr(complete_streaming_response, "usage", None) - if _final_usage is not None and hasattr(_final_usage, "cost") and _final_usage.cost is not None: - if "additional_headers" not in complete_streaming_response._hidden_params: - complete_streaming_response._hidden_params["additional_headers"] = {} - complete_streaming_response._hidden_params["additional_headers"][ - "llm_provider-x-litellm-response-cost" - ] = float(_final_usage.cost) + if ( + _final_usage is not None + and hasattr(_final_usage, "cost") + and _final_usage.cost is not None + ): + if ( + "additional_headers" + not in complete_streaming_response._hidden_params + ): + complete_streaming_response._hidden_params[ + "additional_headers" + ] = {} + complete_streaming_response._hidden_params[ + "additional_headers" + ]["llm_provider-x-litellm-response-cost"] = float( + _final_usage.cost + ) setattr( response, @@ -2169,12 +2183,23 @@ class CustomStreamWrapper: # Propagate provider-reported cost (e.g. OpenRouter) # to _hidden_params so the cost calculator picks it up _final_usage = getattr(complete_streaming_response, "usage", None) - if _final_usage is not None and hasattr(_final_usage, "cost") and _final_usage.cost is not None: - if "additional_headers" not in complete_streaming_response._hidden_params: - complete_streaming_response._hidden_params["additional_headers"] = {} - complete_streaming_response._hidden_params["additional_headers"][ - "llm_provider-x-litellm-response-cost" - ] = float(_final_usage.cost) + if ( + _final_usage is not None + and hasattr(_final_usage, "cost") + and _final_usage.cost is not None + ): + if ( + "additional_headers" + not in complete_streaming_response._hidden_params + ): + complete_streaming_response._hidden_params[ + "additional_headers" + ] = {} + complete_streaming_response._hidden_params[ + "additional_headers" + ]["llm_provider-x-litellm-response-cost"] = float( + _final_usage.cost + ) setattr( response, @@ -2400,7 +2425,11 @@ def calculate_total_usage(chunks: List[ModelResponse]) -> Usage: if "completion_tokens" in usage: completion_tokens = usage.get("completion_tokens", 0) or 0 - if latest_usage_chunk and hasattr(latest_usage_chunk, "cost") and latest_usage_chunk.cost is not None: + if ( + latest_usage_chunk + and hasattr(latest_usage_chunk, "cost") + and latest_usage_chunk.cost is not None + ): return latest_usage_chunk returned_usage_chunk = Usage( 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 f27955691a0..1f8f63cf7ed 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 @@ -183,7 +183,11 @@ def test_get_combined_thinking_content_preserves_interleaved_blocks(): make_chunk(role="assistant", content=None), make_chunk( thinking_blocks=[ - {"type": "thinking", "thinking": "Step 1 analysis...", "signature": None} + { + "type": "thinking", + "thinking": "Step 1 analysis...", + "signature": None, + } ] ), make_chunk( @@ -201,7 +205,11 @@ def test_get_combined_thinking_content_preserves_interleaved_blocks(): ), make_chunk( thinking_blocks=[ - {"type": "thinking", "thinking": "Step 2 analysis...", "signature": None} + { + "type": "thinking", + "thinking": "Step 2 analysis...", + "signature": None, + } ] ), make_chunk( @@ -512,7 +520,7 @@ def test_stream_chunk_builder_anthropic_web_search(): assert usage.prompt_tokens == 50 assert usage.completion_tokens == 27 assert usage.total_tokens == 77 - assert usage.server_tool_use['web_search_requests'] == 2 + assert usage.server_tool_use["web_search_requests"] == 2 def test_sort_chunks_handles_dict_hidden_params_created_at(): @@ -602,7 +610,9 @@ def test_stream_chunk_builder_dict_snapshot_preserves_hidden_provider_fields(): response = stream_chunk_builder(chunks=[chunk_dict]) assert response is not None - assert response._hidden_params["provider_specific_fields"]["traffic_type"] == "default" + assert ( + response._hidden_params["provider_specific_fields"]["traffic_type"] == "default" + ) def test_cost_field_in_usage_chunks(): @@ -611,21 +621,29 @@ def test_cost_field_in_usage_chunks(): id="chatcmpl-1", created=1745513206, model="openrouter/claude", - choices=[StreamingChoices(finish_reason=None, index=0, delta=Delta(content="Hi"))], + 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_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=""))], + 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") + usage = processor.calculate_usage( + chunks=[chunk1, chunk2], model="openrouter/claude", completion_output="Hi" + ) assert hasattr(usage, "cost") assert usage.cost == 0.00025 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 6b1410a1b82..e5e5ae4f534 100644 --- a/tests/test_litellm/litellm_core_utils/test_streaming_handler.py +++ b/tests/test_litellm/litellm_core_utils/test_streaming_handler.py @@ -770,7 +770,9 @@ async def test_vertex_streaming_bad_request_not_midstream(logging_obj: Logging): from litellm.llms.vertex_ai.common_utils import VertexAIError async def _raise_bad_request(**kwargs): - raise VertexAIError(status_code=400, message="invalid maxOutputTokens", headers=None) + raise VertexAIError( + status_code=400, message="invalid maxOutputTokens", headers=None + ) response = CustomStreamWrapper( completion_stream=None, @@ -788,7 +790,9 @@ async def test_vertex_streaming_bad_request_not_midstream(logging_obj: Logging): @pytest.mark.asyncio -async def test_vertex_streaming_rate_limit_triggers_midstream_fallback(logging_obj: Logging): +async def test_vertex_streaming_rate_limit_triggers_midstream_fallback( + logging_obj: Logging, +): """Ensure Vertex 429 rate-limit errors raise MidStreamFallbackError, not RateLimitError. Regression test for https://github.com/BerriAI/litellm/issues/20870 @@ -797,7 +801,9 @@ async def test_vertex_streaming_rate_limit_triggers_midstream_fallback(logging_o from litellm.llms.vertex_ai.common_utils import VertexAIError async def _raise_rate_limit(**kwargs): - raise VertexAIError(status_code=429, message="Resource exhausted.", headers=None) + raise VertexAIError( + status_code=429, message="Resource exhausted.", headers=None + ) response = CustomStreamWrapper( completion_stream=None, @@ -825,7 +831,9 @@ def test_sync_streaming_rate_limit_triggers_midstream_fallback(logging_obj: Logg from litellm.llms.vertex_ai.common_utils import VertexAIError def _raise_rate_limit(**kwargs): - raise VertexAIError(status_code=429, message="Resource exhausted.", headers=None) + raise VertexAIError( + status_code=429, message="Resource exhausted.", headers=None + ) response = CustomStreamWrapper( completion_stream=None, @@ -850,7 +858,9 @@ def test_sync_streaming_bad_request_not_midstream(logging_obj: Logging): from litellm.llms.vertex_ai.common_utils import VertexAIError def _raise_bad_request(**kwargs): - raise VertexAIError(status_code=400, message="invalid maxOutputTokens", headers=None) + raise VertexAIError( + status_code=400, message="invalid maxOutputTokens", headers=None + ) response = CustomStreamWrapper( completion_stream=None, @@ -1169,16 +1179,22 @@ def test_calculate_total_usage_with_cost(): id="test-1", created=1745513206, model="openrouter/test", - choices=[StreamingChoices(finish_reason=None, index=0, delta=Delta(content="Hi"))], + 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_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=""))], + choices=[ + StreamingChoices(finish_reason="stop", index=0, delta=Delta(content="")) + ], usage=chunk2_usage, ) @@ -1198,26 +1214,38 @@ async def test_openrouter_streaming_cost_after_finish_reason(logging_obj: Loggin id="chatcmpl-or", created=1742056047, model="openrouter/claude", - choices=[StreamingChoices(finish_reason=None, index=0, delta=Delta(content="Hi", role="assistant"))], + 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=""))], + 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_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=""))], + choices=[ + StreamingChoices(finish_reason=None, index=0, delta=Delta(content="")) + ], usage=chunk3_usage, ) - completion_stream = ModelResponseListIterator(model_responses=[chunk1, chunk2, chunk3]) + completion_stream = ModelResponseListIterator( + model_responses=[chunk1, chunk2, chunk3] + ) response = CustomStreamWrapper( completion_stream=completion_stream, model="openrouter/claude", @@ -1248,22 +1276,32 @@ def test_openrouter_streaming_cost_propagates_to_hidden_params(): id="chatcmpl-or", created=1742056047, model="openrouter/claude", - choices=[StreamingChoices(finish_reason=None, index=0, delta=Delta(content="Hi", role="assistant"))], + 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=""))], + 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), + 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 @@ -1278,7 +1316,11 @@ def test_openrouter_streaming_cost_propagates_to_hidden_params(): # Simulate the propagation logic from streaming_handler _final_usage = getattr(complete_response, "usage", None) - if _final_usage is not None and hasattr(_final_usage, "cost") and _final_usage.cost is not None: + if ( + _final_usage is not None + and hasattr(_final_usage, "cost") + and _final_usage.cost is not None + ): if "additional_headers" not in complete_response._hidden_params: complete_response._hidden_params["additional_headers"] = {} complete_response._hidden_params["additional_headers"][ @@ -1296,7 +1338,9 @@ def test_openrouter_streaming_cost_propagates_to_hidden_params(): # 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) + provider_cost = get_response_cost_from_hidden_params( + complete_response._hidden_params + ) assert provider_cost == 0.00025 @@ -1502,6 +1546,7 @@ def _build_chunks(pattern: list[str], N: int) -> list[ModelResponseStream]: chunks.append(_make_chunk(p)) return chunks + _REPETITION_TEST_CASES = [ # Basic cases pytest.param( @@ -1558,7 +1603,14 @@ _REPETITION_TEST_CASES = [ id="last_chunk_different_no_raise", ), pytest.param( - ["same"] * (litellm.REPEATED_STREAMING_CHUNK_LIMIT // 2 + 1) + ["different_mid"] + ["same"] * (litellm.REPEATED_STREAMING_CHUNK_LIMIT - litellm.REPEATED_STREAMING_CHUNK_LIMIT // 2 + 1), + ["same"] * (litellm.REPEATED_STREAMING_CHUNK_LIMIT // 2 + 1) + + ["different_mid"] + + ["same"] + * ( + litellm.REPEATED_STREAMING_CHUNK_LIMIT + - litellm.REPEATED_STREAMING_CHUNK_LIMIT // 2 + + 1 + ), False, id="middle_chunk_different_no_raise", ), @@ -1568,7 +1620,9 @@ _REPETITION_TEST_CASES = [ id="last_two_different_no_raise", ), pytest.param( - ["diff"] * litellm.REPEATED_STREAMING_CHUNK_LIMIT + ["same"] * litellm.REPEATED_STREAMING_CHUNK_LIMIT + ["diff"], + ["diff"] * litellm.REPEATED_STREAMING_CHUNK_LIMIT + + ["same"] * litellm.REPEATED_STREAMING_CHUNK_LIMIT + + ["diff"], True, id="in_between_same_and_diff_raise", ), @@ -1594,6 +1648,8 @@ def test_raise_on_model_repetition( for chunk in chunks: wrapper.chunks.append(chunk) wrapper.raise_on_model_repetition() + + def test_usage_chunk_after_finish_reason_updates_hidden_params(logging_obj): """ Test that provider-reported usage from a post-finish_reason chunk @@ -1675,12 +1731,13 @@ def test_usage_chunk_after_finish_reason_updates_hidden_params(logging_obj): last_chunk = collected[-1] hidden_usage = last_chunk._hidden_params.get("usage") assert hidden_usage is not None, "Expected usage in _hidden_params" - assert hidden_usage.prompt_tokens == 20, ( - f"Expected prompt_tokens=20 from provider, got {hidden_usage.prompt_tokens}" - ) - assert hidden_usage.completion_tokens == 135, ( - f"Expected completion_tokens=135 from provider, got {hidden_usage.completion_tokens}" - ) + assert ( + hidden_usage.prompt_tokens == 20 + ), f"Expected prompt_tokens=20 from provider, got {hidden_usage.prompt_tokens}" + assert ( + hidden_usage.completion_tokens == 135 + ), f"Expected completion_tokens=135 from provider, got {hidden_usage.completion_tokens}" + @pytest.mark.asyncio async def test_custom_stream_wrapper_aclose(): @@ -1754,9 +1811,9 @@ def test_content_not_dropped_when_finish_reason_already_set( result = initialized_custom_stream_wrapper.chunk_creator(chunk=content_chunk) - assert result is not None, ( - "chunk_creator() returned None — content was dropped (issue #22098)" - ) + assert ( + result is not None + ), "chunk_creator() returned None — content was dropped (issue #22098)" assert result.choices[0].delta.content == "world!" @@ -1808,14 +1865,14 @@ def test_tool_use_not_dropped_when_finish_reason_already_set( result = initialized_custom_stream_wrapper.chunk_creator(chunk=tool_chunk) - assert result is not None, ( - "chunk_creator() returned None — tool_use data was dropped" - ) + assert ( + result is not None + ), "chunk_creator() returned None — tool_use data was dropped" tool_calls = result.choices[0].delta.tool_calls - assert tool_calls is not None and len(tool_calls) > 0, ( - "tool_calls should contain at least one tool call" - ) + assert ( + tool_calls is not None and len(tool_calls) > 0 + ), "tool_calls should contain at least one tool call" assert tool_calls[0].id == "call_1" assert tool_calls[0].function.name == "get_weather" From d1e9c521ebc4a64df654587c6a3b6de5626fb681 Mon Sep 17 00:00:00 2001 From: Dhruv Yadav Date: Tue, 31 Mar 2026 15:30:03 +0530 Subject: [PATCH 4/4] address review feedback: remove provider-specific guard, fix double-append, DRY cost propagation MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - Remove hardcoded openrouter check from core handler; the general return-instead-of-raise in return_processed_chunk_logic already handles post-finish-reason usage chunks for any provider. - Remove self.chunks.append in return_processed_chunk_logic to avoid double-appending (caller already appends). - Build proper Usage in calculate_total_usage instead of returning raw chunk early — preserves cost while ensuring token fields are always reconstructed. - Extract cost propagation into _propagate_usage_cost_to_hidden_params static method, used by both sync and async paths. - Test calls the real method instead of copy-pasting production logic. --- .../litellm_core_utils/streaming_handler.py | 84 +++++++------------ .../test_streaming_handler.py | 14 +--- 2 files changed, 33 insertions(+), 65 deletions(-) diff --git a/litellm/litellm_core_utils/streaming_handler.py b/litellm/litellm_core_utils/streaming_handler.py index 432f94540af..d2d66444189 100644 --- a/litellm/litellm_core_utils/streaming_handler.py +++ b/litellm/litellm_core_utils/streaming_handler.py @@ -1030,7 +1030,6 @@ class CustomStreamWrapper: hasattr(model_response, "usage") and model_response.usage is not None ): - self.chunks.append(model_response) return model_response return # flush any remaining holding chunk @@ -1145,11 +1144,7 @@ class CustomStreamWrapper: not isinstance(chunk, dict) or "provider_specific_fields" not in chunk ): - if self.custom_llm_provider != "openrouter": - raise StopIteration - else: - # OpenRouter: continue processing - usage will come in later chunks - pass + raise StopIteration anthropic_response_obj: GChunk = cast(GChunk, chunk) completion_obj["content"] = anthropic_response_obj["text"] if anthropic_response_obj["is_finished"]: @@ -1829,6 +1824,23 @@ 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": # noqa: PLR0915 cache_hit = False if ( @@ -1933,26 +1945,9 @@ class CustomStreamWrapper: response = self.model_response_creator() if complete_streaming_response is not None: - # Propagate provider-reported cost (e.g. OpenRouter) - # to _hidden_params so the cost calculator picks it up - _final_usage = getattr(complete_streaming_response, "usage", None) - if ( - _final_usage is not None - and hasattr(_final_usage, "cost") - and _final_usage.cost is not None - ): - if ( - "additional_headers" - not in complete_streaming_response._hidden_params - ): - complete_streaming_response._hidden_params[ - "additional_headers" - ] = {} - complete_streaming_response._hidden_params[ - "additional_headers" - ]["llm_provider-x-litellm-response-cost"] = float( - _final_usage.cost - ) + self._propagate_usage_cost_to_hidden_params( + complete_streaming_response + ) setattr( response, @@ -2180,26 +2175,9 @@ class CustomStreamWrapper: response = self.model_response_creator() if complete_streaming_response is not None: - # Propagate provider-reported cost (e.g. OpenRouter) - # to _hidden_params so the cost calculator picks it up - _final_usage = getattr(complete_streaming_response, "usage", None) - if ( - _final_usage is not None - and hasattr(_final_usage, "cost") - and _final_usage.cost is not None - ): - if ( - "additional_headers" - not in complete_streaming_response._hidden_params - ): - complete_streaming_response._hidden_params[ - "additional_headers" - ] = {} - complete_streaming_response._hidden_params[ - "additional_headers" - ]["llm_provider-x-litellm-response-cost"] = float( - _final_usage.cost - ) + self._propagate_usage_cost_to_hidden_params( + complete_streaming_response + ) setattr( response, @@ -2425,19 +2403,19 @@ def calculate_total_usage(chunks: List[ModelResponse]) -> Usage: if "completion_tokens" in usage: completion_tokens = usage.get("completion_tokens", 0) or 0 - if ( - latest_usage_chunk - and hasattr(latest_usage_chunk, "cost") - and latest_usage_chunk.cost is not None - ): - return latest_usage_chunk - returned_usage_chunk = Usage( prompt_tokens=prompt_tokens, completion_tokens=completion_tokens, total_tokens=prompt_tokens + completion_tokens, ) + if ( + latest_usage_chunk + and hasattr(latest_usage_chunk, "cost") + and latest_usage_chunk.cost is not None + ): + returned_usage_chunk.cost = latest_usage_chunk.cost + return returned_usage_chunk 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 e5e5ae4f534..94b0487ef1e 100644 --- a/tests/test_litellm/litellm_core_utils/test_streaming_handler.py +++ b/tests/test_litellm/litellm_core_utils/test_streaming_handler.py @@ -1314,18 +1314,8 @@ def test_openrouter_streaming_cost_propagates_to_hidden_params(): assert hasattr(complete_response.usage, "cost") assert complete_response.usage.cost == 0.00025 - # Simulate the propagation logic from streaming_handler - _final_usage = getattr(complete_response, "usage", None) - if ( - _final_usage is not None - and hasattr(_final_usage, "cost") - and _final_usage.cost is not None - ): - if "additional_headers" not in complete_response._hidden_params: - complete_response._hidden_params["additional_headers"] = {} - complete_response._hidden_params["additional_headers"][ - "llm_provider-x-litellm-response-cost" - ] = float(_final_usage.cost) + # 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 (