diff --git a/litellm/litellm_core_utils/llm_response_utils/response_metadata.py b/litellm/litellm_core_utils/llm_response_utils/response_metadata.py index 051a8fdf256..a93116aba98 100644 --- a/litellm/litellm_core_utils/llm_response_utils/response_metadata.py +++ b/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, diff --git a/litellm/litellm_core_utils/streaming_handler.py b/litellm/litellm_core_utils/streaming_handler.py index 484edd869cb..8921ec2689c 100644 --- a/litellm/litellm_core_utils/streaming_handler.py +++ b/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 diff --git a/tests/unit/test_cost_calculator.py b/tests/unit/test_cost_calculator.py index c536522e736..ad014ad7d8b 100644 --- a/tests/unit/test_cost_calculator.py +++ b/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"), [