diff --git a/litellm/litellm_core_utils/streaming_chunk_builder_utils.py b/litellm/litellm_core_utils/streaming_chunk_builder_utils.py index f909111a05c..387716ddad8 100644 --- a/litellm/litellm_core_utils/streaming_chunk_builder_utils.py +++ b/litellm/litellm_core_utils/streaming_chunk_builder_utils.py @@ -471,6 +471,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 @@ -480,6 +481,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( @@ -506,6 +509,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]: @@ -543,9 +547,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) @@ -604,6 +612,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, @@ -613,6 +624,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( @@ -652,6 +664,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( @@ -720,6 +733,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 dee3e2dfb4c..b20574551aa 100644 --- a/litellm/litellm_core_utils/streaming_handler.py +++ b/litellm/litellm_core_utils/streaming_handler.py @@ -1024,10 +1024,14 @@ 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: @@ -1577,12 +1581,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 @@ -1820,6 +1828,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 ( @@ -1882,6 +1907,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() @@ -1920,6 +1949,10 @@ 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", @@ -2146,6 +2179,10 @@ 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", @@ -2363,12 +2400,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, @@ -2376,6 +2417,13 @@ def calculate_total_usage(chunks: List[ModelResponse]) -> Usage: 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/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 d0bc9b78941..42b4580b7b5 100644 --- a/litellm/types/utils.py +++ b/litellm/types/utils.py @@ -1829,16 +1829,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 @@ -1847,6 +1850,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..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( @@ -402,7 +410,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 +421,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,8 +519,8 @@ 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.total_tokens == 77 + assert usage.server_tool_use["web_search_requests"] == 2 def test_sort_chunks_handles_dict_hidden_params_created_at(): @@ -602,4 +610,42 @@ 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(): + 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 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 47a77c110b0..ea52a1e8a72 100644 --- a/tests/test_litellm/litellm_core_utils/test_streaming_handler.py +++ b/tests/test_litellm/litellm_core_utils/test_streaming_handler.py @@ -1171,6 +1171,169 @@ 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 + + +@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, ):