mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-12 23:01:41 +00:00
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:
parent
97947c2542
commit
491f36be50
7 changed files with 200 additions and 6 deletions
|
|
@ -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.
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
|
|
|||
|
|
@ -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(
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
|
|
|||
|
|
@ -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:
|
||||
|
|
|
|||
|
|
@ -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 {}
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue