Merge 5e8d58dff5 into e9b1c11098

litellm/proxy/common_request_processing.py

@@ -2899,25 +2899,36 @@ class ProxyBaseLLMRequestProcessing:
         version: str | None,
         callback_headers: Mapping[str, str],
     ) -> Mapping[str, str]:
-        """The streaming response headers describing `hidden_params`' deployment."""
+        """Build stream headers without publishing a cost before the stream finishes."""
+        raw_cache_key: Final[object] = hidden_params.get("cache_key")
+        raw_api_base: Final[object] = hidden_params.get("api_base")
+        raw_fastest_response_batch_completion: Final[object] = hidden_params.get("fastest_response_batch_completion")
+        raw_additional_headers: Final[object] = hidden_params.get("additional_headers")
+        headers: Final[Mapping[str, str]] = ProxyBaseLLMRequestProcessing.get_custom_headers(
+            user_api_key_dict=user_api_key_dict,
+            call_id=logging_obj.litellm_call_id,
+            model_id=self._get_model_id_from_response(hidden_params, self.data),
+            cache_key=raw_cache_key if isinstance(raw_cache_key, str) else "",
+            api_base=raw_api_base if isinstance(raw_api_base, str) else "",
+            version=version,
+            response_cost=None,
+            model_region=getattr(user_api_key_dict, "allowed_model_region", ""),
+            fastest_response_batch_completion=(
+                raw_fastest_response_batch_completion
+                if isinstance(raw_fastest_response_batch_completion, bool)
+                else None
+            ),
+            request_data=self.data,
+            hidden_params=hidden_params,
+            litellm_logging_obj=logging_obj,
+            **(raw_additional_headers if isinstance(raw_additional_headers, Mapping) else MappingProxyType({})),
+        )
+        merged_headers: Final[Mapping[str, str]] = {**headers, **callback_headers}
         return MappingProxyType(
             {
-                **ProxyBaseLLMRequestProcessing.get_custom_headers(
-                    user_api_key_dict=user_api_key_dict,
-                    call_id=logging_obj.litellm_call_id,
-                    model_id=self._get_model_id_from_response(hidden_params, self.data),
-                    cache_key=hidden_params.get("cache_key") or "",
-                    api_base=hidden_params.get("api_base") or "",
-                    version=version,
-                    response_cost=_response_cost_header_value(hidden_params),
-                    model_region=getattr(user_api_key_dict, "allowed_model_region", ""),
-                    fastest_response_batch_completion=hidden_params.get("fastest_response_batch_completion"),
-                    request_data=self.data,
-                    hidden_params=hidden_params,
-                    litellm_logging_obj=logging_obj,
-                    **(hidden_params.get("additional_headers") or MappingProxyType({})),
-                ),
-                **callback_headers,
+                key: value
+                for key, value in merged_headers.items()
+                if not key.lower().startswith("x-litellm-response-cost")
             }
         )
 

tests/unit/proxy/test_common_request_processing.py

@@ -1507,6 +1507,48 @@ class TestProxyBaseLLMRequestProcessing:
         assert "x-litellm-response-cost-discount-amount" in headers
         assert float(headers["x-litellm-response-cost-discount-amount"]) == 0.000005
 
+    def test_stream_response_headers_omit_incomplete_cost(self):
+        processing = ProxyBaseLLMRequestProcessing.__new__(ProxyBaseLLMRequestProcessing)
+        processing.data = {}
+        user = MagicMock(spec=UserAPIKeyAuth)
+        user.tpm_limit = None
+        user.rpm_limit = None
+        user.max_budget = None
+        user.spend = 0
+        logging_obj = LiteLLMLoggingObj(
+            model="gpt-4",
+            messages=[{"role": "user", "content": "test"}],
+            stream=True,
+            call_type="completion",
+            start_time=None,
+            litellm_call_id="test-call-id",
+            function_id="test-function-id",
+        )
+        logging_obj.set_cost_breakdown(
+            input_cost=0.0,
+            output_cost=0.0,
+            total_cost=0.001,
+            cost_for_built_in_tools_cost_usd_dollar=0.0,
+            original_cost=0.0,
+            margin_percent=0.1,
+            margin_total_amount=0.001,
+        )
+
+        headers = processing._stream_response_headers(
+            hidden_params={"response_cost": 0.001},
+            user_api_key_dict=user,
+            logging_obj=logging_obj,
+            version=None,
+            callback_headers={
+                "x-litellm-response-cost": "0.067",
+                "x-custom-header": "custom-value",
+            },
+        )
+
+        assert "x-litellm-call-id" in headers
+        assert headers["x-custom-header"] == "custom-value"
+        assert not any(key.startswith("x-litellm-response-cost") for key in headers)
+
     def test_get_custom_headers_without_discount_info(self):
         """
         Test that when no discount is applied, discount headers are not included.