scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-13 23:11:40 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 100

fix(proxy): extend response headers hook to streaming, TTS, image gen, and pass-through

Browse source
This commit is contained in:
michelligabriele 2026-03-20 19:05:42 +01:00
parent d3891e6eae
commit 7badcbbca6
No known key found for this signature in database
10 changed files with 189 additions and 10 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

10
litellm/proxy/image_endpoints/endpoints.py
View file

@ -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(

48
litellm/proxy/pass_through_endpoints/pass_through_endpoints.py
View file

@ -901,6 +901,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,
@ -911,10 +925,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,
)
@ -939,6 +950,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,
@ -949,10 +974,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,
)
@ -1030,6 +1052,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,

10
litellm/proxy/proxy_server.py
View file

@ -8101,6 +8101,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", "")

3
litellm/proxy/utils.py
View file

@ -2172,7 +2172,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"),

108
tests/test_litellm/proxy/hooks/test_post_call_response_headers_hook.py
View file

@ -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"
)

5
tests/test_litellm/proxy/image_endpoints/test_endpoints.py
View file

@ -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"

6
tests/test_litellm/proxy/pass_through_endpoints/test_llm_pass_through_endpoints.py
View file

@ -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(

6
tests/test_litellm/proxy/pass_through_endpoints/test_pass_through_endpoints.py
View file

@ -896,6 +896,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()
@ -1551,6 +1554,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()

2
tests/test_litellm/proxy/pass_through_endpoints/test_passthrough_post_call_guardrails.py
View file

@ -129,6 +129,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(
@ -154,6 +155,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(

1
tests/test_litellm/proxy/test_audio_speech_prometheus_hooks.py
View file

@ -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"