diff --git a/litellm/litellm_core_utils/streaming_chunk_builder_utils.py b/litellm/litellm_core_utils/streaming_chunk_builder_utils.py index 0fa9ef55aa5..10fc77f7b04 100644 --- a/litellm/litellm_core_utils/streaming_chunk_builder_utils.py +++ b/litellm/litellm_core_utils/streaming_chunk_builder_utils.py @@ -991,7 +991,7 @@ class ChunkProcessor: prompt_tokens_details, cache_creation_token_details ) - if usage_chunk.dropped_non_finite_cost(): + if Usage.dropped_non_finite_cost(usage_chunk): cost = None elif usage_chunk_dict["cost"] is not None: cost = usage_chunk_dict["cost"] diff --git a/litellm/types/utils.py b/litellm/types/utils.py index 2b960ff2f0e..21a43f9cad8 100644 --- a/litellm/types/utils.py +++ b/litellm/types/utils.py @@ -2076,11 +2076,11 @@ class Usage(SafeAttributeModel, CompletionUsage): else: # maintain openai compatibility in usage object if possible del self.server_tool_use + dropped_non_finite_cost: Final = cost is not None and _carries_non_finite_number(cost) if cost is None: del self.cost - elif _carries_non_finite_number(cost): + elif dropped_non_finite_cost: verbose_logger.warning("Dropping a non-finite provider usage cost") - self._dropped_non_finite_cost = True del self.cost else: self.cost = cost @@ -2098,9 +2098,11 @@ class Usage(SafeAttributeModel, CompletionUsage): for k, v in params.items(): setattr(self, k, v) + self._dropped_non_finite_cost = dropped_non_finite_cost - def dropped_non_finite_cost(self) -> bool: - return self._dropped_non_finite_cost + @staticmethod + def dropped_non_finite_cost(usage: "Usage") -> bool: + return usage._dropped_non_finite_cost def __contains__(self, key) -> bool: # Define custom behavior for the 'in' operator diff --git a/tests/unit/litellm_core_utils/test_streaming_chunk_builder_utils.py b/tests/unit/litellm_core_utils/test_streaming_chunk_builder_utils.py index 88e21bb262d..6125b30d984 100644 --- a/tests/unit/litellm_core_utils/test_streaming_chunk_builder_utils.py +++ b/tests/unit/litellm_core_utils/test_streaming_chunk_builder_utils.py @@ -1102,6 +1102,22 @@ def test_a_non_finite_final_stream_cost_does_not_fall_back_to_an_earlier_chunk_c assert getattr(response.usage, "cost", None) is None assert response.usage.total_tokens == 27 + +@pytest.mark.parametrize("provider_key", ["dropped_non_finite_cost", "_dropped_non_finite_cost"]) +def test_a_provider_usage_key_named_like_the_dropped_cost_flag_keeps_the_stream_cost(provider_key: str): + chunk: Final = ModelResponseStream( + id="chatcmpl-flag-key", + created=1745513206, + model="sonar", + choices=[StreamingChoices(finish_reason="stop", index=0, delta=Delta(content="ok"))], + usage={"prompt_tokens": 26, "completion_tokens": 1, "total_tokens": 27, "cost": 0.00371, provider_key: True}, + ) + + response: Final = stream_chunk_builder(chunks=[chunk]) + + assert response is not None + assert getattr(response.usage, "cost", None) == 0.00371 + def test_stream_chunk_builder_tolerates_trailing_chunk_without_choices(): """Regression for https://github.com/BerriAI/litellm/issues/32051