From fe60f9d0f144bd27ecddf5cb7f9b88187b461d99 Mon Sep 17 00:00:00 2001 From: michelligabriele Date: Tue, 9 Jun 2026 22:10:23 +0200 Subject: [PATCH] fix(proxy): extend response headers hook to streaming, TTS, image gen, and pass-through (#24232) * fix(proxy): extend response headers hook to streaming, TTS, image gen, and pass-through * test: mock post_call_response_headers_hook in audio speech route tests --- litellm/proxy/image_endpoints/endpoints.py | 10 ++ .../pass_through_endpoints.py | 48 ++++++-- litellm/proxy/proxy_server.py | 10 ++ litellm/proxy/utils.py | 3 +- .../test_post_call_response_headers_hook.py | 108 +++++++++++++++++- .../proxy/image_endpoints/test_endpoints.py | 5 + .../test_llm_pass_through_endpoints.py | 6 + .../test_pass_through_endpoints.py | 20 ++++ .../test_passthrough_post_call_guardrails.py | 2 + .../proxy/proxy_server/test_routes_audio.py | 2 + .../test_audio_speech_prometheus_hooks.py | 1 + 11 files changed, 205 insertions(+), 10 deletions(-) diff --git a/litellm/proxy/image_endpoints/endpoints.py b/litellm/proxy/image_endpoints/endpoints.py index fe8b7c6fdc..c217116e45 100644 --- a/litellm/proxy/image_endpoints/endpoints.py +++ b/litellm/proxy/image_endpoints/endpoints.py @@ -173,6 +173,16 @@ async def image_generation( ) ) + # Call response headers hook (matches base_process_llm_request behavior) + callback_headers = await proxy_logging_obj.post_call_response_headers_hook( + data=data, + user_api_key_dict=user_api_key_dict, + response=response, + request_headers=dict(request.headers), + ) + if callback_headers: + fastapi_response.headers.update(callback_headers) + return response except Exception as e: await proxy_logging_obj.post_call_failure_hook( diff --git a/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py b/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py index 45e264b1cd..d2b848e3c3 100644 --- a/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py +++ b/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py @@ -1099,6 +1099,20 @@ async def pass_through_request( # noqa: PLR0915 status_code=e.response.status_code, detail=await e.response.aread() ) + # Call response headers hook for streaming pass-through + _response_headers = HttpPassThroughEndpointHelpers.get_response_headers( + headers=response.headers, + litellm_call_id=litellm_call_id, + ) + callback_headers = await proxy_logging_obj.post_call_response_headers_hook( + data=_parsed_body or {}, + user_api_key_dict=user_api_key_dict, + response=response, + request_headers=dict(request.headers), + ) + if callback_headers: + _response_headers.update(callback_headers) + return StreamingResponse( PassThroughStreamingHandler.chunk_processor( response=response, @@ -1109,10 +1123,7 @@ async def pass_through_request( # noqa: PLR0915 passthrough_success_handler_obj=pass_through_endpoint_logging, url_route=str(url), ), - headers=HttpPassThroughEndpointHelpers.get_response_headers( - headers=response.headers, - litellm_call_id=litellm_call_id, - ), + headers=_response_headers, status_code=response.status_code, ) @@ -1151,6 +1162,20 @@ async def pass_through_request( # noqa: PLR0915 status_code=e.response.status_code, detail=await e.response.aread() ) + # Call response headers hook for detected streaming pass-through + _response_headers = HttpPassThroughEndpointHelpers.get_response_headers( + headers=response.headers, + litellm_call_id=litellm_call_id, + ) + callback_headers = await proxy_logging_obj.post_call_response_headers_hook( + data=_parsed_body or {}, + user_api_key_dict=user_api_key_dict, + response=response, + request_headers=dict(request.headers), + ) + if callback_headers: + _response_headers.update(callback_headers) + return StreamingResponse( PassThroughStreamingHandler.chunk_processor( response=response, @@ -1161,10 +1186,7 @@ async def pass_through_request( # noqa: PLR0915 passthrough_success_handler_obj=pass_through_endpoint_logging, url_route=str(url), ), - headers=HttpPassThroughEndpointHelpers.get_response_headers( - headers=response.headers, - litellm_call_id=litellm_call_id, - ), + headers=_response_headers, status_code=response.status_code, ) @@ -1303,6 +1325,16 @@ async def pass_through_request( # noqa: PLR0915 api_base=str(url._uri_reference), ) + # Call response headers hook + callback_headers = await proxy_logging_obj.post_call_response_headers_hook( + data=_parsed_body or {}, + user_api_key_dict=user_api_key_dict, + response=response, + request_headers=dict(request.headers), + ) + if callback_headers: + custom_headers.update(callback_headers) + response_headers = HttpPassThroughEndpointHelpers.get_response_headers( headers=response.headers, custom_headers=custom_headers, diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 213f682b8f..e50e58838e 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -9191,6 +9191,16 @@ async def audio_speech( hidden_params=hidden_params, ) + # Call response headers hook (matches audio_transcription behavior) + callback_headers = await proxy_logging_obj.post_call_response_headers_hook( + data=data, + user_api_key_dict=user_api_key_dict, + response=response, + request_headers=dict(request.headers), + ) + if callback_headers: + custom_headers.update(callback_headers) + # Determine media type based on model type media_type = "audio/mpeg" # Default for OpenAI TTS request_model = data.get("model", "") diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index 5ad42b5e1b..2bba8bbd60 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -2496,7 +2496,8 @@ class ProxyLogging: ) return { - "custom_llm_provider": hidden_params.get("custom_llm_provider"), + "custom_llm_provider": hidden_params.get("custom_llm_provider") + or getattr(response, "custom_llm_provider", None), "model_info": model_info, "api_base": hidden_params.get("api_base"), "model_id": hidden_params.get("model_id"), diff --git a/tests/test_litellm/proxy/hooks/test_post_call_response_headers_hook.py b/tests/test_litellm/proxy/hooks/test_post_call_response_headers_hook.py index 8d8dd2d428..660b0b0162 100644 --- a/tests/test_litellm/proxy/hooks/test_post_call_response_headers_hook.py +++ b/tests/test_litellm/proxy/hooks/test_post_call_response_headers_hook.py @@ -13,7 +13,6 @@ from unittest.mock import patch sys.path.insert(0, os.path.abspath("../../../..")) -import litellm from litellm.integrations.custom_logger import CustomLogger from litellm.proxy._types import UserAPIKeyAuth @@ -336,3 +335,110 @@ async def test_litellm_call_info_backwards_compatible(): assert result == {"x-test": "1"} assert injector.called is True + + +# --- Tests for custom_llm_provider fallback (streaming response types) --- + + +@pytest.mark.asyncio +async def test_litellm_call_info_fallback_to_response_attribute(): + """Test that _build_litellm_call_info falls back to response.custom_llm_provider + when _hidden_params doesn't contain it (streaming response types).""" + inspector = CallInfoInspectorLogger() + + class MockStreamResponse: + """Mimics CustomStreamWrapper: custom_llm_provider as attribute, + _hidden_params without it.""" + + custom_llm_provider = "bedrock" + _hidden_params = { + "model_id": "model-xyz", + "api_base": "https://bedrock.us-east-1.amazonaws.com", + } + + with patch("litellm.callbacks", [inspector]): + from litellm.proxy.utils import ProxyLogging + from litellm.caching.caching import DualCache + + proxy_logging = ProxyLogging(user_api_key_cache=DualCache()) + + await proxy_logging.post_call_response_headers_hook( + data={ + "model": "claude-3", + "metadata": {"model_info": {"id": "model-xyz"}}, + }, + user_api_key_dict=UserAPIKeyAuth(api_key="test-key"), + response=MockStreamResponse(), + ) + + assert inspector.called is True + assert inspector.received_call_info is not None + assert inspector.received_call_info["custom_llm_provider"] == "bedrock" + assert ( + inspector.received_call_info["api_base"] + == "https://bedrock.us-east-1.amazonaws.com" + ) + assert inspector.received_call_info["model_id"] == "model-xyz" + + +@pytest.mark.asyncio +async def test_litellm_call_info_fallback_no_hidden_params(): + """Test that _build_litellm_call_info works when response has no _hidden_params + at all (LiteLLMCompletionStreamingIterator case).""" + inspector = CallInfoInspectorLogger() + + class MockIteratorResponse: + """Mimics LiteLLMCompletionStreamingIterator: custom_llm_provider as attribute, + no _hidden_params attribute at all.""" + + custom_llm_provider = "vertex_ai" + + with patch("litellm.callbacks", [inspector]): + from litellm.proxy.utils import ProxyLogging + from litellm.caching.caching import DualCache + + proxy_logging = ProxyLogging(user_api_key_cache=DualCache()) + + await proxy_logging.post_call_response_headers_hook( + data={"model": "gemini-pro", "metadata": {}}, + user_api_key_dict=UserAPIKeyAuth(api_key="test-key"), + response=MockIteratorResponse(), + ) + + assert inspector.called is True + assert inspector.received_call_info is not None + assert inspector.received_call_info["custom_llm_provider"] == "vertex_ai" + assert inspector.received_call_info["api_base"] is None + assert inspector.received_call_info["model_id"] is None + + +@pytest.mark.asyncio +async def test_litellm_call_info_hidden_params_takes_priority(): + """Test that _hidden_params.custom_llm_provider takes priority over + the response attribute when both are present.""" + inspector = CallInfoInspectorLogger() + + class MockResponse: + custom_llm_provider = "attribute_value" + _hidden_params = { + "custom_llm_provider": "hidden_params_value", + "api_base": "https://example.com", + "model_id": "m1", + } + + with patch("litellm.callbacks", [inspector]): + from litellm.proxy.utils import ProxyLogging + from litellm.caching.caching import DualCache + + proxy_logging = ProxyLogging(user_api_key_cache=DualCache()) + + await proxy_logging.post_call_response_headers_hook( + data={"model": "test", "metadata": {}}, + user_api_key_dict=UserAPIKeyAuth(api_key="test-key"), + response=MockResponse(), + ) + + assert ( + inspector.received_call_info["custom_llm_provider"] + == "hidden_params_value" + ) diff --git a/tests/test_litellm/proxy/image_endpoints/test_endpoints.py b/tests/test_litellm/proxy/image_endpoints/test_endpoints.py index 8fec05abe9..91a011a823 100644 --- a/tests/test_litellm/proxy/image_endpoints/test_endpoints.py +++ b/tests/test_litellm/proxy/image_endpoints/test_endpoints.py @@ -43,11 +43,15 @@ async def test_image_generation_prompt_rerouting(monkeypatch): async def fake_post_call_success_hook(*, data, user_api_key_dict, response): return response + async def fake_post_call_response_headers_hook(**kwargs): + return {"x-callback-test": "value"} + fake_proxy_logger = SimpleNamespace( pre_call_hook=fake_pre_call_hook, update_request_status=fake_update_request_status, post_call_failure_hook=fake_post_call_failure_hook, post_call_success_hook=fake_post_call_success_hook, + post_call_response_headers_hook=fake_post_call_response_headers_hook, ) captured_route_request_data: Dict[str, Any] = {} @@ -110,3 +114,4 @@ async def test_image_generation_prompt_rerouting(monkeypatch): assert pre_call_input["messages"][0]["content"] == "original prompt" assert captured_route_request_data["prompt"] == "sanitized prompt" assert "messages" not in captured_route_request_data + assert response.headers.get("x-callback-test") == "value" diff --git a/tests/test_litellm/proxy/pass_through_endpoints/test_llm_pass_through_endpoints.py b/tests/test_litellm/proxy/pass_through_endpoints/test_llm_pass_through_endpoints.py index c9be00afed..b75fc27e21 100644 --- a/tests/test_litellm/proxy/pass_through_endpoints/test_llm_pass_through_endpoints.py +++ b/tests/test_litellm/proxy/pass_through_endpoints/test_llm_pass_through_endpoints.py @@ -1803,6 +1803,9 @@ class TestForwardHeaders: mock_logging_obj.pre_call_hook = AsyncMock(return_value=mock_request_body) mock_logging_obj.post_call_success_hook = AsyncMock() mock_logging_obj.post_call_failure_hook = AsyncMock() + mock_logging_obj.post_call_response_headers_hook = AsyncMock( + return_value={} + ) # Call pass_through_request with forward_headers=True result = await pass_through_request( @@ -1901,6 +1904,9 @@ class TestForwardHeaders: mock_logging_obj.pre_call_hook = AsyncMock(return_value=mock_request_body) mock_logging_obj.post_call_success_hook = AsyncMock() mock_logging_obj.post_call_failure_hook = AsyncMock() + mock_logging_obj.post_call_response_headers_hook = AsyncMock( + return_value={} + ) # Call pass_through_request with forward_headers=False (default) result = await pass_through_request( diff --git a/tests/test_litellm/proxy/pass_through_endpoints/test_pass_through_endpoints.py b/tests/test_litellm/proxy/pass_through_endpoints/test_pass_through_endpoints.py index 9b9d5e22a4..31f1c1c57d 100644 --- a/tests/test_litellm/proxy/pass_through_endpoints/test_pass_through_endpoints.py +++ b/tests/test_litellm/proxy/pass_through_endpoints/test_pass_through_endpoints.py @@ -957,6 +957,9 @@ async def test_pass_through_request_contains_proxy_server_request_in_kwargs(): return_value={"test": "data"} ) mock_proxy_logging.post_call_failure_hook = AsyncMock() + mock_proxy_logging.post_call_response_headers_hook = AsyncMock( + return_value={"x-callback-test": "value"} + ) # Setup mock for http response mock_response = MagicMock() @@ -1069,6 +1072,9 @@ async def test_pass_through_request_streaming_marks_logging_obj_as_stream(): return_value={"model": "claude-3", "stream": True} ) mock_proxy_logging.post_call_failure_hook = AsyncMock() + mock_proxy_logging.post_call_response_headers_hook = AsyncMock( + return_value={"x-callback-test": "value"} + ) upstream_response = MagicMock() upstream_response.status_code = 200 @@ -1134,6 +1140,9 @@ async def test_pass_through_request_sse_response_marks_logging_obj_as_stream(): return_value={"model": "claude-3"} ) mock_proxy_logging.post_call_failure_hook = AsyncMock() + mock_proxy_logging.post_call_response_headers_hook = AsyncMock( + return_value={"x-callback-test": "value"} + ) upstream_response = MagicMock() upstream_response.status_code = 200 @@ -1975,6 +1984,9 @@ async def test_pass_through_request_query_params_forwarding(): mock_proxy_logging.pre_call_hook = AsyncMock( return_value=test_body ) + mock_proxy_logging.post_call_response_headers_hook = AsyncMock( + return_value={"x-callback-test": "value"} + ) # Setup mock for http response mock_response = MagicMock() @@ -2703,6 +2715,10 @@ async def test_pass_through_request_non_streaming_uses_content_for_state_raw_bod "litellm.proxy.proxy_server.proxy_logging_obj.pre_call_hook", new=AsyncMock(side_effect=_hook_mutates_body), ), + patch( + "litellm.proxy.proxy_server.proxy_logging_obj.post_call_response_headers_hook", + new=AsyncMock(return_value={}), + ), patch( "litellm.proxy.pass_through_endpoints.pass_through_endpoints.pass_through_endpoint_logging.pass_through_async_success_handler", new=AsyncMock(), @@ -2763,6 +2779,10 @@ async def test_pass_through_request_streaming_uses_content_for_state_raw_body(): "litellm.proxy.proxy_server.proxy_logging_obj.pre_call_hook", new=AsyncMock(side_effect=lambda **kw: kw["data"]), ), + patch( + "litellm.proxy.proxy_server.proxy_logging_obj.post_call_response_headers_hook", + new=AsyncMock(return_value={}), + ), patch( "litellm.proxy.pass_through_endpoints.pass_through_endpoints.pass_through_endpoint_logging.pass_through_async_success_handler", new=AsyncMock(), diff --git a/tests/test_litellm/proxy/pass_through_endpoints/test_passthrough_post_call_guardrails.py b/tests/test_litellm/proxy/pass_through_endpoints/test_passthrough_post_call_guardrails.py index eafe71e106..a2f7476abd 100644 --- a/tests/test_litellm/proxy/pass_through_endpoints/test_passthrough_post_call_guardrails.py +++ b/tests/test_litellm/proxy/pass_through_endpoints/test_passthrough_post_call_guardrails.py @@ -130,6 +130,7 @@ class TestPassthroughPostCallGuardrails: mock_proxy_logging.post_call_success_hook = AsyncMock( return_value=_GEMINI_RESPONSE ) + mock_proxy_logging.post_call_response_headers_hook = AsyncMock(return_value={}) with _common_patches(mock_proxy_logging, mock_response): await pass_through_request( @@ -155,6 +156,7 @@ class TestPassthroughPostCallGuardrails: mock_proxy_logging = MagicMock() mock_proxy_logging.pre_call_hook = AsyncMock(return_value={}) mock_proxy_logging.post_call_success_hook = AsyncMock() + mock_proxy_logging.post_call_response_headers_hook = AsyncMock(return_value={}) with _common_patches(mock_proxy_logging, mock_response): result = await pass_through_request( diff --git a/tests/test_litellm/proxy/proxy_server/test_routes_audio.py b/tests/test_litellm/proxy/proxy_server/test_routes_audio.py index d88bcf136e..74542a3eaf 100644 --- a/tests/test_litellm/proxy/proxy_server/test_routes_audio.py +++ b/tests/test_litellm/proxy/proxy_server/test_routes_audio.py @@ -26,6 +26,7 @@ def patched_speech(monkeypatch): MagicMock( pre_call_hook=AsyncMock(side_effect=lambda **kw: kw["data"]), post_call_failure_hook=AsyncMock(), + post_call_response_headers_hook=AsyncMock(return_value={}), update_request_status=AsyncMock(), ), ) @@ -61,6 +62,7 @@ def patched_speech_error(monkeypatch): MagicMock( pre_call_hook=AsyncMock(side_effect=lambda **kw: kw["data"]), post_call_failure_hook=AsyncMock(), + post_call_response_headers_hook=AsyncMock(return_value={}), update_request_status=AsyncMock(), ), ) diff --git a/tests/test_litellm/proxy/test_audio_speech_prometheus_hooks.py b/tests/test_litellm/proxy/test_audio_speech_prometheus_hooks.py index 5a13a4dc53..99f6f3a9b7 100644 --- a/tests/test_litellm/proxy/test_audio_speech_prometheus_hooks.py +++ b/tests/test_litellm/proxy/test_audio_speech_prometheus_hooks.py @@ -68,6 +68,7 @@ async def test_audio_speech_success_does_not_call_post_call_success_hook( mock_logging.post_call_failure_hook = mock_failure_hook mock_logging.pre_call_hook = mock_pre_call mock_logging.update_request_status = mock_update_status + mock_logging.post_call_response_headers_hook = AsyncMock(return_value={}) async def _mock_route_request(*, data, route_type, llm_router, user_model): assert route_type == "aspeech"