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 @@ async def async_post_call_response_headers_hook( 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 @@ async def async_post_call_response_headers_hook( - 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 @@ async def base_process_llm_request( 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 @@ async def base_process_llm_request( 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 @@ async def _handle_llm_api_exception( 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 7fe0ce6d6f5..457c552c9e4 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -7367,6 +7367,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 e6da95bb78f..fff86757116 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 @@ def get_cache( ### 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. @@ -1975,6 +1989,9 @@ async def post_call_response_headers_hook( """ 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 +2002,22 @@ async def post_call_response_headers_hook( _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: @@ -1999,6 +2026,30 @@ async def post_call_response_headers_hook( ) 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 9c397aaaaeb..cfdad62692e 100644 --- a/litellm/responses/main.py +++ b/litellm/responses/main.py @@ -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( @@ -778,6 +781,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 @@ def __init__( 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