fix(proxy): make async_post_call_response_headers_hook consistent across all endpoints

The response headers hook had 5 gaps that prevented callbacks from
reliably extracting routing metadata across endpoint types:

1. Hook never fired for /audio/transcriptions (endpoint bypasses
   base_process_llm_request)
2. custom_llm_provider not accessible in hook data for any endpoint
3. custom_llm_provider not stamped in ResponsesAPIResponse._hidden_params
   (unlike chat completions)
4. model_info under inconsistent keys (metadata vs litellm_metadata)
5. request_headers always None at all call sites

This adds a litellm_call_info parameter to the hook that normalizes
routing metadata (custom_llm_provider, model_info, api_base, model_id)
regardless of endpoint type. Also stamps custom_llm_provider on
Responses API responses, adds the hook call to the transcription
handler, and passes request_headers at all call sites.

Supersedes PR #21385.
This commit is contained in:
michelligabriele 2026-03-06 18:14:28 +01:00 committed by shivam
parent 97947c2542
commit 491f36be50
No known key found for this signature in database
7 changed files with 200 additions and 6 deletions

View file

@ -377,6 +377,7 @@ class CustomLogger: # https://docs.litellm.ai/docs/observability/custom_callbac
user_api_key_dict: UserAPIKeyAuth,
response: Any,
request_headers: Optional[Dict[str, str]] = None,
litellm_call_info: Optional[Dict[str, Any]] = None,
) -> Optional[Dict[str, str]]:
"""
Called after an LLM API call (success or failure) to allow injecting custom HTTP response headers.
@ -386,6 +387,11 @@ class CustomLogger: # https://docs.litellm.ai/docs/observability/custom_callbac
- user_api_key_dict: UserAPIKeyAuth - The user API key dictionary.
- response: Any - The response object (None for failure cases).
- request_headers: Optional[Dict[str, str]] - The original request headers.
- litellm_call_info: Optional[Dict[str, Any]] - Normalized routing metadata:
- custom_llm_provider: str - The LLM provider (e.g. "openai", "azure")
- model_info: dict - The model_info from router config
- api_base: str - The API base URL used
- model_id: str - The deployment model ID
Returns:
- Optional[Dict[str, str]]: A dictionary of headers to inject into the HTTP response.

View file

@ -934,6 +934,7 @@ class ProxyBaseLLMRequestProcessing:
data=self.data,
user_api_key_dict=user_api_key_dict,
response=response,
request_headers=dict(request.headers),
)
if callback_headers:
custom_headers.update(callback_headers)
@ -1042,6 +1043,7 @@ class ProxyBaseLLMRequestProcessing:
data=self.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)
@ -1210,6 +1212,7 @@ class ProxyBaseLLMRequestProcessing:
data=self.data,
user_api_key_dict=user_api_key_dict,
response=None,
request_headers=self.data.get("proxy_server_request", {}).get("headers", {}),
)
if callback_headers:
headers.update(callback_headers)

View file

@ -7298,6 +7298,16 @@ async def audio_transcriptions(
)
)
# 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(

View file

@ -1975,6 +1975,9 @@ class ProxyLogging:
"""
merged_headers: Dict[str, str] = {}
try:
# Build litellm_call_info — normalized routing metadata for callbacks
litellm_call_info = self._build_litellm_call_info(data=data, response=response)
for callback in litellm.callbacks:
_callback: Optional[CustomLogger] = None
if isinstance(callback, str):
@ -1985,12 +1988,22 @@ class ProxyLogging:
_callback = callback # type: ignore
if _callback is not None and isinstance(_callback, CustomLogger):
result = await _callback.async_post_call_response_headers_hook(
data=data,
user_api_key_dict=user_api_key_dict,
response=response,
request_headers=request_headers,
)
try:
result = await _callback.async_post_call_response_headers_hook(
data=data,
user_api_key_dict=user_api_key_dict,
response=response,
request_headers=request_headers,
litellm_call_info=litellm_call_info,
)
except TypeError:
# Backwards compat: callback doesn't accept litellm_call_info
result = await _callback.async_post_call_response_headers_hook(
data=data,
user_api_key_dict=user_api_key_dict,
response=response,
request_headers=request_headers,
)
if result is not None:
merged_headers.update(result)
except Exception as e:
@ -1999,6 +2012,30 @@ class ProxyLogging:
)
return merged_headers
@staticmethod
def _build_litellm_call_info(
data: dict, response: Any
) -> Dict[str, Any]:
"""
Build a normalized dict of routing metadata from response._hidden_params
and data, abstracting away the metadata vs litellm_metadata split.
"""
hidden_params = getattr(response, "_hidden_params", {}) or {}
# model_info: check both metadata keys (chat uses "metadata", responses uses "litellm_metadata")
model_info = (
data.get("metadata", {}).get("model_info")
or data.get("litellm_metadata", {}).get("model_info")
or {}
)
return {
"custom_llm_provider": hidden_params.get("custom_llm_provider"),
"model_info": model_info,
"api_base": hidden_params.get("api_base"),
"model_id": hidden_params.get("model_id"),
}
def is_a2a_streaming_response(self, response: dict) -> bool:
expected_keys = ["jsonrpc", "id", "result"]
return all(key in response for key in expected_keys)

View file

@ -507,6 +507,9 @@ async def aresponses(
litellm_metadata=kwargs.get("litellm_metadata", {}),
custom_llm_provider=custom_llm_provider,
)
# Stamp custom_llm_provider so callbacks can identify the provider
# (mirrors litellm/main.py:1371 for chat completions)
response._hidden_params["custom_llm_provider"] = custom_llm_provider
if response is None:
raise ValueError(
@ -773,6 +776,9 @@ def responses(
litellm_metadata=kwargs.get("litellm_metadata", {}),
custom_llm_provider=custom_llm_provider,
)
# Stamp custom_llm_provider so callbacks can identify the provider
# (mirrors litellm/main.py:1371 for chat completions)
response._hidden_params["custom_llm_provider"] = custom_llm_provider
return response
except Exception as e:

View file

@ -79,6 +79,7 @@ class BaseResponsesAPIStreamingIterator:
self._hidden_params = {
"model_id": _model_info.get("id", None),
"api_base": _api_base,
"custom_llm_provider": custom_llm_provider,
}
self._hidden_params["additional_headers"] = process_response_headers(
self.response.headers or {}

View file

@ -195,3 +195,134 @@ async def test_default_hook_returns_none():
response=None,
)
assert result is None
# --- Tests for litellm_call_info parameter ---
class CallInfoInspectorLogger(CustomLogger):
"""Logger that captures litellm_call_info for inspection."""
def __init__(self):
self.called = False
self.received_call_info = None
async def async_post_call_response_headers_hook(
self,
data: dict,
user_api_key_dict: UserAPIKeyAuth,
response: Any,
request_headers: Optional[Dict[str, str]] = None,
litellm_call_info: Optional[Dict[str, Any]] = None,
) -> Optional[Dict[str, str]]:
self.called = True
self.received_call_info = litellm_call_info
return None
@pytest.mark.asyncio
async def test_litellm_call_info_from_hidden_params():
"""Test that litellm_call_info is built from response._hidden_params."""
inspector = CallInfoInspectorLogger()
class MockResponse:
_hidden_params = {
"custom_llm_provider": "openai",
"api_base": "https://api.openai.com",
"model_id": "model-abc",
}
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": "gpt-4", "metadata": {"model_info": {"id": "model-abc", "provider": "HubSpot"}}},
user_api_key_dict=UserAPIKeyAuth(api_key="test-key"),
response=MockResponse(),
)
assert inspector.called is True
assert inspector.received_call_info is not None
assert inspector.received_call_info["custom_llm_provider"] == "openai"
assert inspector.received_call_info["api_base"] == "https://api.openai.com"
assert inspector.received_call_info["model_id"] == "model-abc"
assert inspector.received_call_info["model_info"]["provider"] == "HubSpot"
@pytest.mark.asyncio
async def test_litellm_call_info_from_litellm_metadata():
"""Test that litellm_call_info finds model_info under litellm_metadata (responses API path)."""
inspector = CallInfoInspectorLogger()
class MockResponse:
_hidden_params = {
"custom_llm_provider": "azure",
"api_base": "https://east.openai.azure.com",
"model_id": "deploy-xyz",
}
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": "gpt-4", "litellm_metadata": {"model_info": {"id": "deploy-xyz"}}},
user_api_key_dict=UserAPIKeyAuth(api_key="test-key"),
response=MockResponse(),
)
assert inspector.received_call_info["model_info"]["id"] == "deploy-xyz"
assert inspector.received_call_info["custom_llm_provider"] == "azure"
@pytest.mark.asyncio
async def test_litellm_call_info_with_none_response():
"""Test that litellm_call_info handles None response (failure path)."""
inspector = CallInfoInspectorLogger()
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": "gpt-4", "metadata": {}},
user_api_key_dict=UserAPIKeyAuth(api_key="test-key"),
response=None,
)
assert inspector.called is True
assert inspector.received_call_info is not None
assert inspector.received_call_info["custom_llm_provider"] is None
assert inspector.received_call_info["model_info"] == {}
@pytest.mark.asyncio
async def test_litellm_call_info_backwards_compatible():
"""Test that existing callbacks without litellm_call_info parameter still work."""
# HeaderInjectorLogger doesn't accept litellm_call_info — must not crash
injector = HeaderInjectorLogger(headers={"x-test": "1"})
class MockResponse:
_hidden_params = {"custom_llm_provider": "openai", "api_base": "https://api.openai.com", "model_id": "m1"}
with patch("litellm.callbacks", [injector]):
from litellm.proxy.utils import ProxyLogging
from litellm.caching.caching import DualCache
proxy_logging = ProxyLogging(user_api_key_cache=DualCache())
result = await proxy_logging.post_call_response_headers_hook(
data={"model": "gpt-4", "metadata": {}},
user_api_key_dict=UserAPIKeyAuth(api_key="test-key"),
response=MockResponse(),
)
assert result == {"x-test": "1"}
assert injector.called is True