fix(streaming): keep a dropped non-finite final usage cost from falling back to an earlier chunk's cost

litellm/litellm_core_utils/streaming_chunk_builder_utils.py

@@ -984,7 +984,9 @@ class ChunkProcessor:
                     prompt_tokens_details, cache_creation_token_details
                 )
 
-                if usage_chunk_dict["cost"] is not None:
+                if usage_chunk.dropped_non_finite_cost():
+                    cost = None
+                elif usage_chunk_dict["cost"] is not None:
                     cost = usage_chunk_dict["cost"]
 
         prompt_tokens_details = attach_cache_creation_token_details(prompt_tokens_details, cache_creation_token_details)

litellm/types/utils.py

@@ -1949,7 +1949,7 @@ class ServerToolUse(LiteLLMBaseModel):
         return getattr(self, key)
 
 
-_COST_BREAKDOWN_ADAPTER: Final = TypeAdapter(Mapping[str, object])
+_COST_BREAKDOWN_ADAPTER: Final = TypeAdapter(Mapping[object, object])
 
 
 def _is_non_finite_float(value: object) -> bool:
@@ -1969,6 +1969,7 @@ class Usage(SafeAttributeModel, CompletionUsage):
     _cache_read_input_tokens: int = PrivateAttr(
         0
     )  # hidden param for prompt caching. Might change, once openai introduces their equivalent.
+    _dropped_non_finite_cost: bool = PrivateAttr(False)
 
     server_tool_use: ServerToolUse | None = None
     cost: float | None = None
@@ -2079,6 +2080,7 @@ class Usage(SafeAttributeModel, CompletionUsage):
             del self.cost
         elif _carries_non_finite_number(cost):
             verbose_logger.warning("Dropping a non-finite provider usage cost")
+            self._dropped_non_finite_cost = True
             del self.cost
         else:
             self.cost = cost
@@ -2097,6 +2099,9 @@ class Usage(SafeAttributeModel, CompletionUsage):
         for k, v in params.items():
             setattr(self, k, v)
 
+    def dropped_non_finite_cost(self) -> bool:
+        return self._dropped_non_finite_cost
+
     def __contains__(self, key) -> bool:
         # Define custom behavior for the 'in' operator
         return hasattr(self, key)

tests/unit/litellm_core_utils/llm_response_utils/test_convert_dict_to_response.py

@@ -2997,3 +2997,19 @@ def test_a_non_finite_provider_usage_cost_is_dropped_so_the_response_is_strict_j
     assert "cost" not in usage
     assert usage["total_tokens"] == 4
 
+
+def test_a_finite_provider_cost_breakdown_with_non_string_keys_is_kept():
+    response: Final = convert_to_model_response_object(
+        response_object={
+            "id": "chatcmpl-int-key-cost",
+            "object": "chat.completion",
+            "created": 1,
+            "model": "gpt-4o",
+            "choices": [{"index": 0, "message": {"role": "assistant", "content": "ok"}, "finish_reason": "stop"}],
+            "usage": {"prompt_tokens": 3, "completion_tokens": 1, "total_tokens": 4, "cost": {1: 0.5}},
+        },
+        model_response_object=ModelResponse(),
+    )
+
+    assert response.usage.cost == {1: 0.5}
+

tests/unit/litellm_core_utils/test_streaming_chunk_builder_utils.py

@@ -1077,6 +1077,31 @@ def test_cost_field_in_usage_chunks():
     assert usage.completion_tokens == 5
 
 
+
+def _perplexity_usage_chunk(cost: Mapping[str, float], finish_reason: str | None) -> ModelResponseStream:
+    return ModelResponseStream(
+        id="chatcmpl-pplx",
+        created=1745513206,
+        model="sonar",
+        choices=[StreamingChoices(finish_reason=finish_reason, index=0, delta=Delta(content="ok"))],
+        usage={"prompt_tokens": 26, "completion_tokens": 1, "total_tokens": 27, "cost": dict(cost)},
+    )
+
+
+@pytest.mark.parametrize("poisoned_field", ["request_cost", "total_cost"])
+def test_a_non_finite_final_stream_cost_does_not_fall_back_to_an_earlier_chunk_cost(poisoned_field: str):
+    final_cost: Final = {"input_tokens_cost": 0.0012, "output_tokens_cost": 1e-05, "request_cost": 0.0025}
+    chunks: Final = [
+        _perplexity_usage_chunk({"input_tokens_cost": 0, "output_tokens_cost": 0, "total_cost": 0}, None),
+        _perplexity_usage_chunk({**final_cost, "total_cost": 0.00371, poisoned_field: float("nan")}, "stop"),
+    ]
+
+    response: Final = stream_chunk_builder(chunks=chunks)
+
+    assert response is not None
+    assert getattr(response.usage, "cost", None) is None
+    assert response.usage.total_tokens == 27
+
 def test_stream_chunk_builder_tolerates_trailing_chunk_without_choices():
     """Regression for https://github.com/BerriAI/litellm/issues/32051