fix(proxy): price the streamed usage frame from the served model, as the spend row is

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

litellm/litellm_core_utils/streaming_handler.py

@@ -2500,7 +2500,9 @@ class CustomStreamWrapper:
         return any(get(field) for field in cls._CONTINUATION_DISQUALIFYING_DELTA_FIELDS)
 
     def _new_usage_chunk(self) -> ModelResponseStream:
-        usage_chunk: Final = self.model_response_creator()
+        usage_chunk: Final = self.model_response_creator(
+            hidden_params=_provider_hidden_params(None, self._provider_response_model)
+        )
         if self.caller_requested_stream_usage:
             usage_chunk.choices = []
         return usage_chunk

litellm/proxy/usage_cost_egress.py

@@ -10,6 +10,7 @@ from typing import Final, TypeAlias, TypeVar
 
 from pydantic import BaseModel
 
+from litellm.litellm_core_utils.hidden_params import get_hidden_params
 from litellm.types.llms.openai import ResponseAPIUsage, ResponsesAPIResponse
 from litellm.types.utils import ModelResponse, ModelResponseStream, Usage
 
@@ -47,7 +48,8 @@ def with_client_stream_usage_cost(chunk: BaseModel, price: Pricer) -> BaseModel:
         usage: Final = getattr(chunk, "usage", None)
         if not isinstance(usage, Usage):
             return chunk
-        priced: Final = ModelResponse(usage=usage, service_tier=getattr(chunk, "service_tier", None))
+        priced: Final = ModelResponse(model=chunk.model, usage=usage, service_tier=getattr(chunk, "service_tier", None))
+        priced.hidden_params = dict(get_hidden_params(chunk) or {})
         cost: Final = _client_usage_cost(getattr(usage, "cost", None), price(priced))
         return chunk if cost is None else _with_usage_cost(chunk, usage, cost)
     if getattr(chunk, "type", None) not in _RESPONSES_USAGE_EVENTS:

tests/unit/litellm_core_utils/test_streaming_handler.py

@@ -6118,6 +6118,28 @@ def test_chunk_creator_keeps_provider_model_private_across_stream():
     )
 
 
+def test_the_usage_chunk_carries_the_served_model_so_its_price_matches_the_assembled_response():
+    from litellm.router_utils.add_retry_fallback_headers import get_hidden_params_dict
+
+    wrapper = CustomStreamWrapper(
+        completion_stream=None,
+        model="requested-route",
+        logging_obj=MagicMock(),
+        custom_llm_provider="openai",
+    )
+    wrapper.chunk_creator(
+        chunk=ModelResponseStream(
+            id="chunk-1",
+            model="selected-model",
+            choices=[StreamingChoices(finish_reason=None, index=0, delta=Delta(content="hello"))],
+        )
+    )
+
+    usage_chunk = wrapper._new_usage_chunk()
+
+    assert get_hidden_params_dict(usage_chunk)["provider_response_model"] == "selected-model"
+
+
 def test_assembled_stream_uses_later_provider_model_for_cost(
     monkeypatch: pytest.MonkeyPatch,
 ):

tests/unit/proxy/test_usage_cost_egress.py

@@ -103,16 +103,26 @@ def test_a_content_chunk_is_not_priced() -> None:
     assert priced == []
 
 
-def test_the_stream_pricer_sees_the_upstream_cost_and_the_served_tier_so_the_bill_matches_the_spend_row() -> None:
-    seen: list[tuple[object, object]] = []
+def test_the_stream_pricer_sees_the_served_model_tier_and_upstream_cost_so_the_bill_matches_the_spend_row() -> None:
+    seen: list[tuple[object, object, object, object]] = []
 
     def price(reply: object) -> float | None:
-        seen.append((getattr(getattr(reply, "usage", None), "cost", None), getattr(reply, "service_tier", None)))
+        seen.append(
+            (
+                getattr(reply, "model", None),
+                getattr(reply, "hidden_params", {}).get("provider_response_model"),
+                getattr(reply, "service_tier", None),
+                getattr(getattr(reply, "usage", None), "cost", None),
+            )
+        )
         return OPERATOR_BILL
 
-    with_client_stream_usage_cost(_usage_chunk(UPSTREAM_COST, service_tier="default"), price)
+    chunk: Final = _usage_chunk(UPSTREAM_COST, service_tier="default")
+    chunk.hidden_params = {"provider_response_model": "gpt-4o-mini"}
 
-    assert seen == [(UPSTREAM_COST, "default")]
+    with_client_stream_usage_cost(chunk, price)
+
+    assert seen == [("omni-model", "gpt-4o-mini", "default", UPSTREAM_COST)]
 
 
 @pytest.mark.parametrize("upstream_cost", [UPSTREAM_COST, None])