From 7c5e2e83898105d565ef2f2032fb389b3523d77a Mon Sep 17 00:00:00 2001 From: michelligabriele Date: Thu, 12 Mar 2026 16:51:00 +0100 Subject: [PATCH] fix(proxy): make async_post_call_response_headers_hook consistent across all endpoints (#22985) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * 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. * fix(proxy): address review feedback — safer backwards compat and None guards - Replace try/except TypeError with inspect.signature() check for litellm_call_info backwards compatibility. This avoids masking real TypeErrors inside callback implementations and prevents double invocation with inconsistent parameters. - Use (data.get("key") or {}) instead of data.get("key", {}) to guard against keys that exist with an explicit None value, which would cause AttributeError on the subsequent .get() call. * fix(proxy): cache inspect.signature result for callback compat check Move the inspect.signature() call into a module-level helper with a dict cache keyed by callback identity. Avoids repeated introspection per request per callback in the hot path. * fix(proxy): use class identity for signature cache key Key the _CALLBACK_ACCEPTS_CALL_INFO cache by id(type(cb)) instead of id(cb) to avoid stale entries from Python address reuse after GC. All instances of the same callback class share the same method signature, so class identity is both safer and more cache-efficient. --- litellm/integrations/custom_logger.py | 6 + litellm/proxy/common_request_processing.py | 3 + litellm/proxy/proxy_server.py | 10 ++ litellm/proxy/utils.py | 63 ++++++++- litellm/responses/main.py | 6 + litellm/responses/streaming_iterator.py | 1 + .../test_post_call_response_headers_hook.py | 131 ++++++++++++++++++ 7 files changed, 214 insertions(+), 6 deletions(-) diff --git a/litellm/integrations/custom_logger.py b/litellm/integrations/custom_logger.py index c244363e389..b00584edbc3 100644 --- a/litellm/integrations/custom_logger.py +++ b/litellm/integrations/custom_logger.py @@ -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. diff --git a/litellm/proxy/common_request_processing.py b/litellm/proxy/common_request_processing.py index ce39ecf52dc..951339ccd2b 100644 --- a/litellm/proxy/common_request_processing.py +++ b/litellm/proxy/common_request_processing.py @@ -920,6 +920,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) @@ -1028,6 +1029,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) @@ -1196,6 +1198,7 @@ class ProxyBaseLLMRequestProcessing: data=self.data, user_api_key_dict=user_api_key_dict, response=None, + request_headers=(self.data.get("proxy_server_request") or {}).get("headers", {}), ) if callback_headers: headers.update(callback_headers) diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index e6bb3ee412e..4627e1b8a0a 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -7517,6 +7517,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( diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index 2f9d27568e3..ba6c6fa16e0 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -1,6 +1,7 @@ import asyncio import copy import hashlib +import inspect import json import os import smtplib @@ -285,6 +286,19 @@ class InternalUsageCache: ### LOGGING ### + +# Cache for inspect.signature checks — avoids repeated introspection per request +_CALLBACK_ACCEPTS_CALL_INFO: Dict[int, bool] = {} + + +def _accepts_litellm_call_info(cb: CustomLogger) -> bool: + key = id(type(cb)) + if key not in _CALLBACK_ACCEPTS_CALL_INFO: + sig = inspect.signature(cb.async_post_call_response_headers_hook) + _CALLBACK_ACCEPTS_CALL_INFO[key] = "litellm_call_info" in sig.parameters + return _CALLBACK_ACCEPTS_CALL_INFO[key] + + class ProxyLogging: """ Logging/Custom Handlers for proxy. @@ -1978,6 +1992,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): @@ -1988,12 +2005,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, - ) + if _accepts_litellm_call_info(_callback): + 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, + ) + else: + # 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: @@ -2002,6 +2029,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") or {}).get("model_info") + or (data.get("litellm_metadata") or {}).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) diff --git a/litellm/responses/main.py b/litellm/responses/main.py index 3f2065fe346..e6160d95d49 100644 --- a/litellm/responses/main.py +++ b/litellm/responses/main.py @@ -508,6 +508,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( @@ -782,6 +785,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: diff --git a/litellm/responses/streaming_iterator.py b/litellm/responses/streaming_iterator.py index 705756cadd3..5c5a955ae00 100644 --- a/litellm/responses/streaming_iterator.py +++ b/litellm/responses/streaming_iterator.py @@ -83,6 +83,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 {} 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 6a12366fdd3..3399a34e075 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 @@ -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