diff --git a/litellm/router.py b/litellm/router.py index 276a949d944..0fc9ecc9bb5 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -1579,6 +1579,9 @@ class Router: logging_obj=model_response.logging_obj, ) self._async_generator = async_generator + # Preserve hidden params (including litellm_overhead_time_ms) from original response + if hasattr(model_response, "_hidden_params"): + self._hidden_params = model_response._hidden_params.copy() def __aiter__(self): return self @@ -7002,9 +7005,9 @@ class Router: raise ValueError("Deployment not found") ## GET BASE MODEL - base_model = deployment.get("model_info", {}).get("base_model", None) + base_model = (deployment.get("model_info") or {}).get("base_model", None) if base_model is None: - base_model = deployment.get("litellm_params", {}).get("base_model", None) + base_model = (deployment.get("litellm_params") or {}).get("base_model", None) model = base_model @@ -7019,7 +7022,7 @@ class Router: raise ValueError( f"Deployment missing valid litellm_params. " f"Got: {type(litellm_params_data).__name__}, " - f"deployment_id: {deployment.get('model_info', {}).get('id', 'unknown')}" + f"deployment_id: {(deployment.get('model_info') or {}).get('id', 'unknown')}" ) _model, custom_llm_provider, _, _ = litellm.get_llm_provider( model=litellm_params.model, @@ -7039,10 +7042,10 @@ class Router: if potential_models is not None: for potential_model in potential_models: try: - if potential_model.get("model_info", {}).get( + if (potential_model.get("model_info") or {}).get( "id" - ) == deployment.get("model_info", {}).get("id"): - model = potential_model.get("litellm_params", {}).get( + ) == (deployment.get("model_info") or {}).get("id"): + model = (potential_model.get("litellm_params") or {}).get( "model" ) break @@ -7063,9 +7066,10 @@ class Router: model_info = litellm.get_model_info(model=model_info_name) ## CHECK USER SET MODEL INFO - user_model_info = deployment.get("model_info", {}) + user_model_info = deployment.get("model_info") or {} - model_info.update(user_model_info) + if model_info is not None: + model_info.update(user_model_info) return model_info diff --git a/tests/test_litellm/proxy/test_common_request_processing.py b/tests/test_litellm/proxy/test_common_request_processing.py index 977304f732b..bf794478f10 100644 --- a/tests/test_litellm/proxy/test_common_request_processing.py +++ b/tests/test_litellm/proxy/test_common_request_processing.py @@ -1,4 +1,6 @@ import copy +import datetime +from typing import AsyncGenerator from unittest.mock import AsyncMock, MagicMock import pytest @@ -1348,3 +1350,202 @@ class TestOverrideOpenAIResponseModel: # Verify the model was not changed assert response_obj.model == fallback_model + + +class TestStreamingOverheadHeader: + """ + Tests that x-litellm-overhead-duration-ms is emitted in streaming responses. + + Regression tests for: streaming requests not including overhead header. + """ + + def test_get_custom_headers_includes_overhead_when_set(self): + """ + get_custom_headers() returns x-litellm-overhead-duration-ms + when litellm_overhead_time_ms is in hidden_params. + """ + mock_user_api_key_dict = MagicMock(spec=UserAPIKeyAuth) + mock_user_api_key_dict.tpm_limit = None + mock_user_api_key_dict.rpm_limit = None + mock_user_api_key_dict.max_budget = None + mock_user_api_key_dict.spend = 0.0 + mock_user_api_key_dict.allowed_model_region = None + + hidden_params = { + "litellm_overhead_time_ms": 42.5, + "_response_ms": 500.0, + "model_id": "test-model-id", + "api_base": "https://api.openai.com", + } + + headers = ProxyBaseLLMRequestProcessing.get_custom_headers( + user_api_key_dict=mock_user_api_key_dict, + call_id="test-call-id", + model_id="test-model-id", + cache_key="", + api_base="https://api.openai.com", + version="1.0.0", + response_cost=0.001, + model_region="", + hidden_params=hidden_params, + ) + + assert "x-litellm-overhead-duration-ms" in headers + assert headers["x-litellm-overhead-duration-ms"] == "42.5" + + def test_get_custom_headers_omits_overhead_when_none(self): + """ + get_custom_headers() omits x-litellm-overhead-duration-ms + when litellm_overhead_time_ms is not in hidden_params. + """ + mock_user_api_key_dict = MagicMock(spec=UserAPIKeyAuth) + mock_user_api_key_dict.tpm_limit = None + mock_user_api_key_dict.rpm_limit = None + mock_user_api_key_dict.max_budget = None + mock_user_api_key_dict.spend = 0.0 + mock_user_api_key_dict.allowed_model_region = None + + hidden_params = { + "_response_ms": 500.0, + "model_id": "test-model-id", + } + + headers = ProxyBaseLLMRequestProcessing.get_custom_headers( + user_api_key_dict=mock_user_api_key_dict, + call_id="test-call-id", + model_id="test-model-id", + cache_key="", + api_base="https://api.openai.com", + version="1.0.0", + response_cost=0.001, + model_region="", + hidden_params=hidden_params, + ) + + # Should be absent (None gets filtered by exclude_values) + assert "x-litellm-overhead-duration-ms" not in headers + + def test_update_response_metadata_sets_overhead_on_stream_wrapper(self): + """ + update_response_metadata() sets litellm_overhead_time_ms on + a streaming response's _hidden_params when llm_api_duration_ms is available. + """ + from litellm.litellm_core_utils.llm_response_utils.response_metadata import ( + update_response_metadata, + ) + + # Mock the logging object with llm_api_duration_ms set + mock_logging_obj = MagicMock() + mock_logging_obj.model_call_details = { + "llm_api_duration_ms": 200.0, + "litellm_params": {}, + } + mock_logging_obj.caching_details = None + mock_logging_obj.callback_duration_ms = None + mock_logging_obj.litellm_call_id = "test-call-id" + mock_logging_obj._response_cost_calculator = MagicMock(return_value=0.001) + + # Simulate a streaming result object with _hidden_params (like CustomStreamWrapper) + stream_result = MagicMock() + stream_result._hidden_params = { + "model_id": "test-model-id", + "api_base": "https://api.openai.com", + "additional_headers": {}, + } + + start_time = datetime.datetime.now() - datetime.timedelta(milliseconds=300) + end_time = datetime.datetime.now() + + update_response_metadata( + result=stream_result, + logging_obj=mock_logging_obj, + model="gpt-4o", + kwargs={}, + start_time=start_time, + end_time=end_time, + ) + + assert "litellm_overhead_time_ms" in stream_result._hidden_params + overhead = stream_result._hidden_params["litellm_overhead_time_ms"] + assert overhead is not None + assert isinstance(overhead, float) + # overhead = total_response_ms (~300ms) - llm_api_duration_ms (200ms) = ~100ms + assert overhead > 0 + + @pytest.mark.asyncio + async def test_streaming_response_includes_overhead_header(self): + """ + StreamingResponse returned by create_response() includes + x-litellm-overhead-duration-ms in its headers. + """ + + async def mock_generator() -> AsyncGenerator[str, None]: + yield 'data: {"id":"chatcmpl-test","choices":[{"delta":{"content":"hi"}}]}\n\n' + yield "data: [DONE]\n\n" + + headers = { + "x-litellm-overhead-duration-ms": "42.5", + "x-litellm-call-id": "test-call-id", + "x-litellm-model-id": "test-model-id", + } + + response = await create_response( + generator=mock_generator(), + media_type="text/event-stream", + headers=headers, + ) + + assert isinstance(response, StreamingResponse) + assert response.headers.get("x-litellm-overhead-duration-ms") == "42.5" + + def test_streaming_overhead_header_in_custom_headers_from_stream_hidden_params( + self, + ): + """ + Verifies that when get_custom_headers() is called with a streaming + response's hidden_params (containing litellm_overhead_time_ms), + the x-litellm-overhead-duration-ms header is correctly populated. + + This tests the critical path: update_response_metadata sets the value + → get_custom_headers reads it → StreamingResponse header is set. + """ + mock_user_api_key_dict = MagicMock(spec=UserAPIKeyAuth) + mock_user_api_key_dict.tpm_limit = None + mock_user_api_key_dict.rpm_limit = None + mock_user_api_key_dict.max_budget = None + mock_user_api_key_dict.spend = 0.0 + mock_user_api_key_dict.allowed_model_region = None + + # This is what CustomStreamWrapper._hidden_params looks like after + # update_response_metadata() has been called on it + hidden_params = { + "model_id": "openai-gpt4o-deployment", + "api_base": "https://api.openai.com", + "additional_headers": {}, + "litellm_overhead_time_ms": 55.3, # set by update_response_metadata + "_response_ms": 280.0, + "litellm_call_id": "test-call-id", + "response_cost": 0.002, + "cache_key": None, + "fastest_response_batch_completion": None, + "callback_duration_ms": None, + } + + custom_headers = ProxyBaseLLMRequestProcessing.get_custom_headers( + user_api_key_dict=mock_user_api_key_dict, + call_id="test-call-id", + model_id=hidden_params.get("model_id"), + cache_key=hidden_params.get("cache_key") or "", + api_base=hidden_params.get("api_base") or "", + version="1.0.0", + response_cost=hidden_params.get("response_cost"), + model_region="", + hidden_params=hidden_params, + ) + + # The overhead header must be present and correct + assert "x-litellm-overhead-duration-ms" in custom_headers, ( + "x-litellm-overhead-duration-ms header must be emitted during streaming. " + "It was missing — this is the streaming overhead header regression." + ) + assert custom_headers["x-litellm-overhead-duration-ms"] == "55.3" diff --git a/tests/test_litellm/test_router.py b/tests/test_litellm/test_router.py index 5542cdf8be4..5732deda6fb 100644 --- a/tests/test_litellm/test_router.py +++ b/tests/test_litellm/test_router.py @@ -1297,6 +1297,61 @@ async def test_acompletion_streaming_iterator_edge_cases(): print("✓ Edge case tests passed!") +@pytest.mark.asyncio +async def test_acompletion_streaming_iterator_preserves_hidden_params(): + """ + Regression test: FallbackStreamWrapper must copy _hidden_params from the + original CustomStreamWrapper so that x-litellm-overhead-duration-ms (and + other hidden params) are present in the proxy response headers for streaming. + """ + from unittest.mock import MagicMock + + router = litellm.Router( + model_list=[ + { + "model_name": "gpt-4", + "litellm_params": {"model": "gpt-4", "api_key": "fake-key"}, + } + ], + ) + + # Simulate a CustomStreamWrapper that already has timing metadata set by + # update_response_metadata (litellm_overhead_time_ms, _response_ms, etc.) + mock_response = MagicMock() + mock_response.model = "gpt-4" + mock_response.custom_llm_provider = "openai" + mock_response.logging_obj = MagicMock() + mock_response._hidden_params = { + "litellm_overhead_time_ms": 12.34, + "_response_ms": 500.0, + "litellm_call_id": "test-call-id", + "api_base": "https://api.openai.com", + "additional_headers": {}, + } + + # Make the mock iterable (yields nothing — we only care about hidden_params) + async def _empty(): + return + yield # make it an async generator + + mock_response.__aiter__ = lambda self: _empty().__aiter__() + + result = await router._acompletion_streaming_iterator( + model_response=mock_response, + messages=[{"role": "user", "content": "hi"}], + initial_kwargs={"model": "gpt-4", "stream": True}, + ) + + # The returned FallbackStreamWrapper must carry the original _hidden_params + assert hasattr(result, "_hidden_params"), "result must have _hidden_params" + assert result._hidden_params.get("litellm_overhead_time_ms") == 12.34, ( + "litellm_overhead_time_ms must be preserved — " + "this is what drives x-litellm-overhead-duration-ms in streaming responses" + ) + assert result._hidden_params.get("litellm_call_id") == "test-call-id" + assert result._hidden_params.get("_response_ms") == 500.0 + + @pytest.mark.asyncio async def test_async_function_with_fallbacks_common_utils(): """Test the async_function_with_fallbacks_common_utils method""" @@ -1858,7 +1913,7 @@ def test_get_deployment_credentials_with_provider_resolves_credential_name(): litellm_credential_name to actual credential values (for UI-created models). """ from litellm.types.utils import CredentialItem - + # Setup credential list with a test credential litellm.credential_list = [ CredentialItem(