fix(cost): update stream wrapper response_cost upon exhaustion (#49369)

litellm/litellm_core_utils/llm_response_utils/response_metadata.py

@@ -168,11 +168,16 @@ class ResponseMetadata:
         ## ADD OTHER HIDDEN PARAMS
         model_info: Final = kwargs.get("model_info", {}) or {}
         model_id: Final = model_info.get("id", None)
+        is_unexhausted_stream: Final = hasattr(self.result, "completion_stream") and not getattr(
+            self.result, "sent_last_chunk", False
+        )
         new_params: Final = {
             "litellm_call_id": getattr(logging_obj, "litellm_call_id", None),
             "api_base": get_api_base(model=model or "", optional_params=kwargs),
             "model_id": model_id,
-            "response_cost": logging_obj.response_cost_calculator(
+            "response_cost": None
+            if is_unexhausted_stream
+            else logging_obj.response_cost_calculator(
                 result=self.result,
                 litellm_model_name=model,
                 router_model_id=model_id,

litellm/litellm_core_utils/streaming_handler.py

@@ -2094,6 +2094,34 @@ class CustomStreamWrapper:
             dict[str, object], getattr(chunk, HIDDEN_PARAMS_ATTR)
         )
 
+    def _update_final_response_cost(
+        self,
+        complete_streaming_response: "_CompleteStreamingResponse | None",
+        usage_chunk: "ModelResponseStream | None" = None,
+    ) -> None:
+        if complete_streaming_response is not None:
+            hidden = getattr(complete_streaming_response, "_hidden_params", None)
+            final_cost = hidden.get("response_cost") if isinstance(hidden, dict) else None
+            if final_cost is None and self.logging_obj is not None:
+                try:
+                    final_cost = self.logging_obj.response_cost_calculator(
+                        result=complete_streaming_response,
+                        litellm_model_name=self.model,
+                    )
+                except Exception:
+                    final_cost = None
+            self._hidden_params["response_cost"] = final_cost
+            if self._last_returned_hidden_params is not None:
+                self._last_returned_hidden_params["response_cost"] = final_cost
+            if (
+                usage_chunk is not None
+                and hasattr(usage_chunk, "_hidden_params")
+                and isinstance(usage_chunk._hidden_params, dict)
+            ):
+                usage_chunk._hidden_params["response_cost"] = final_cost
+        else:
+            self._hidden_params["response_cost"] = None
+
     def __next__(self) -> "ModelResponseStream":
         if self._closed:
             raise StopIteration
@@ -2237,6 +2265,8 @@ class CustomStreamWrapper:
                     if final_usage is not None:
                         self._last_returned_hidden_params["usage"] = final_usage
 
+                self._update_final_response_cost(complete_streaming_response, response)
+
                 if self.sent_stream_usage is False and self.send_stream_usage is True:
                     self.sent_stream_usage = True
                     return response
@@ -2438,6 +2468,8 @@ class CustomStreamWrapper:
                 if final_usage is not None:
                     self._last_returned_hidden_params["usage"] = final_usage
 
+            self._update_final_response_cost(complete_streaming_response, response)
+
             if self.sent_stream_usage is False and self.send_stream_usage is True:
                 self.sent_stream_usage = True
                 return response

tests/unit/test_cost_calculator.py

@@ -9528,6 +9528,56 @@ def test_streamed_custom_pricing_cost_matches_non_streamed_cost() -> None:
     assert streamed_cost == pytest.approx(non_streamed._hidden_params["response_cost"])
 
 
+@pytest.mark.usefixtures("isolated_pricing_registry")
+@pytest.mark.parametrize("async_mode", [False, True])
+@pytest.mark.asyncio
+@respx.mock
+async def test_stream_wrapper_response_cost_updated_after_exhaustion(async_mode: bool) -> None:
+    respx.post("https://api.openai.com/v1/chat/completions").mock(side_effect=_custom_pricing_upstream)
+    request: Final = {
+        "model": f"openai/{_CUSTOM_PRICING_MODEL}",
+        "messages": [{"role": "user", "content": "Hi"}],
+        "api_key": "sk-test",
+        "stream": True,
+        "stream_options": {"include_usage": True},
+        **_CUSTOM_PRICING_RATES,
+    }
+
+    stream = await litellm.acompletion(**request) if async_mode else litellm.completion(**request)
+    assert stream._hidden_params.get("response_cost") is None
+
+    if async_mode:
+        chunks = [chunk async for chunk in stream]
+    else:
+        chunks = list(stream)
+
+    assert stream._hidden_params["response_cost"] == pytest.approx(_CUSTOM_PRICING_EXPECTED_COST)
+    assert chunks[-1]._hidden_params["response_cost"] == pytest.approx(_CUSTOM_PRICING_EXPECTED_COST)
+
+
+@pytest.mark.usefixtures("isolated_pricing_registry")
+@pytest.mark.asyncio
+@respx.mock
+async def test_stream_wrapper_unpriced_model_leaves_response_cost_none() -> None:
+    respx.post("https://api.openai.com/v1/chat/completions").mock(side_effect=_custom_pricing_upstream)
+    request: Final = {
+        "model": "openai/totally-unpriced-model-xyz",
+        "messages": [{"role": "user", "content": "Hi"}],
+        "api_key": "sk-test",
+        "stream": True,
+        "stream_options": {"include_usage": True},
+    }
+
+    stream = await litellm.acompletion(**request)
+    assert stream._hidden_params.get("response_cost") is None
+
+    async for _ in stream:
+        pass
+
+    assert stream._hidden_params.get("response_cost") is None
+
+
+
 @pytest.mark.parametrize(
     ("model", "provider"),
     [