diff --git a/litellm/proxy/anthropic_endpoints/endpoints.py b/litellm/proxy/anthropic_endpoints/endpoints.py index c450b655a2..abea9e6fee 100644 --- a/litellm/proxy/anthropic_endpoints/endpoints.py +++ b/litellm/proxy/anthropic_endpoints/endpoints.py @@ -154,8 +154,13 @@ async def anthropic_response( # noqa: PLR0915 response = responses[1] + # Extract model_id from request metadata (set by router during routing) + litellm_metadata = data.get("litellm_metadata", {}) or {} + model_info = litellm_metadata.get("model_info", {}) or {} + model_id = model_info.get("id", "") or "" + + # Get other metadata from hidden_params hidden_params = getattr(response, "_hidden_params", {}) or {} - model_id = hidden_params.get("model_id", None) or "" cache_key = hidden_params.get("cache_key", None) or "" api_base = hidden_params.get("api_base", None) or "" response_cost = hidden_params.get("response_cost", None) or "" @@ -216,12 +221,32 @@ async def anthropic_response( # noqa: PLR0915 str(e) ) ) + + # Extract model_id from request metadata (same as success path) + litellm_metadata = data.get("litellm_metadata", {}) or {} + model_info = litellm_metadata.get("model_info", {}) or {} + model_id = model_info.get("id", "") or "" + + # Get headers + headers = ProxyBaseLLMRequestProcessing.get_custom_headers( + user_api_key_dict=user_api_key_dict, + call_id=data.get("litellm_call_id", ""), + model_id=model_id, + version=version, + response_cost=0, + model_region=getattr(user_api_key_dict, "allowed_model_region", ""), + request_data=data, + timeout=getattr(e, "timeout", None), + litellm_logging_obj=None, + ) + error_msg = f"{str(e)}" raise ProxyException( message=getattr(e, "message", error_msg), type=getattr(e, "type", "None"), param=getattr(e, "param", "None"), code=getattr(e, "status_code", 500), + headers=headers, ) diff --git a/litellm/proxy/common_request_processing.py b/litellm/proxy/common_request_processing.py index 1cdeb3b99e..0143a6e6ce 100644 --- a/litellm/proxy/common_request_processing.py +++ b/litellm/proxy/common_request_processing.py @@ -344,6 +344,7 @@ class ProxyBaseLLMRequestProcessing: user_max_tokens: Optional[int] = None, user_api_base: Optional[str] = None, model: Optional[str] = None, + llm_router: Optional[Router] = None, ) -> Tuple[dict, LiteLLMLoggingObj]: start_time = datetime.now() # start before calling guardrail hooks @@ -498,6 +499,7 @@ class ProxyBaseLLMRequestProcessing: user_api_base=user_api_base, model=model, route_type=route_type, + llm_router=llm_router, ) tasks = [] @@ -536,6 +538,13 @@ class ProxyBaseLLMRequestProcessing: hidden_params = getattr(response, "_hidden_params", {}) or {} model_id = hidden_params.get("model_id", None) or "" + + # Fallback: extract model_id from litellm_metadata if not in hidden_params + if not model_id: + litellm_metadata = self.data.get("litellm_metadata", {}) or {} + model_info = litellm_metadata.get("model_info", {}) or {} + model_id = model_info.get("id", "") or "" + cache_key = hidden_params.get("cache_key", None) or "" api_base = hidden_params.get("api_base", None) or "" response_cost = hidden_params.get("response_cost", None) or "" @@ -756,11 +765,19 @@ class ProxyBaseLLMRequestProcessing: _litellm_logging_obj: Optional[LiteLLMLoggingObj] = self.data.get( "litellm_logging_obj", None ) + + # Attempt to get model_id from logging object + # + # Note: We check the direct model_info path first (not nested in metadata) because that's where the router sets it. + # The nested metadata path is only a fallback for cases where model_info wasn't set at the top level. + model_id = self.maybe_get_model_id(_litellm_logging_obj) + custom_headers = ProxyBaseLLMRequestProcessing.get_custom_headers( user_api_key_dict=user_api_key_dict, call_id=( _litellm_logging_obj.litellm_call_id if _litellm_logging_obj else None ), + model_id=model_id, version=version, response_cost=0, model_region=getattr(user_api_key_dict, "allowed_model_region", ""), @@ -1073,3 +1090,50 @@ class ProxyBaseLLMRequestProcessing: obj.setdefault("usage", {})["cost"] = cost_val return obj return None + + def maybe_get_model_id(self, _logging_obj: Optional[LiteLLMLoggingObj]) -> Optional[str]: + """ + Get model_id from logging object or request metadata. + + The router sets model_info.id when selecting a deployment. This tries multiple locations + where the ID might be stored depending on the request lifecycle stage. + """ + model_id = None + if _logging_obj: + # 1. Try getting from litellm_params (updated during call) + if ( + hasattr(_logging_obj, "litellm_params") + and _logging_obj.litellm_params + ): + # First check direct model_info path (set by router.py with selected deployment) + model_info = _logging_obj.litellm_params.get("model_info") or {} + model_id = model_info.get("id", None) + + # Fallback to nested metadata path + if not model_id: + metadata = _logging_obj.litellm_params.get("metadata") or {} + model_info = metadata.get("model_info") or {} + model_id = model_info.get("id", None) + + # 2. Fallback to kwargs (initial) + if not model_id: + _kwargs = getattr(_logging_obj, "kwargs", None) + if _kwargs: + litellm_params = _kwargs.get("litellm_params", {}) + # First check direct model_info path + model_info = litellm_params.get("model_info") or {} + model_id = model_info.get("id", None) + + # Fallback to nested metadata path + if not model_id: + metadata = litellm_params.get("metadata") or {} + model_info = metadata.get("model_info") or {} + model_id = model_info.get("id", None) + + # 3. Final fallback to self.data["litellm_metadata"] (for routes like /v1/responses that populate data before error) + if not model_id: + litellm_metadata = self.data.get("litellm_metadata", {}) or {} + model_info = litellm_metadata.get("model_info", {}) or {} + model_id = model_info.get("id", None) + + return model_id diff --git a/litellm/responses/streaming_iterator.py b/litellm/responses/streaming_iterator.py index 8eecc3e821..0407776029 100644 --- a/litellm/responses/streaming_iterator.py +++ b/litellm/responses/streaming_iterator.py @@ -8,7 +8,9 @@ import httpx import litellm from litellm.constants import STREAM_SSE_DONE_STRING from litellm.litellm_core_utils.asyncify import run_async_function +from litellm.litellm_core_utils.core_helpers import process_response_headers from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj +from litellm.litellm_core_utils.llm_response_utils.get_api_base import get_api_base from litellm.litellm_core_utils.thread_pool_executor import executor from litellm.llms.base_llm.responses.transformation import BaseResponsesAPIConfig from litellm.responses.utils import ResponsesAPIRequestUtils @@ -51,6 +53,23 @@ class BaseResponsesAPIStreamingIterator: self.litellm_metadata = litellm_metadata self.custom_llm_provider = custom_llm_provider + # set hidden params for response headers (e.g., x-litellm-model-id) + # This matches ths stream wrapper in litellm/litellm_core_utils/streaming_handler.py + _api_base = get_api_base( + model=model or "", + optional_params=self.logging_obj.model_call_details.get( + "litellm_params", {} + ), + ) + _model_info: Dict = litellm_metadata.get("model_info", {}) if litellm_metadata else {} + self._hidden_params = { + "model_id": _model_info.get("id", None), + "api_base": _api_base, + } + self._hidden_params["additional_headers"] = process_response_headers( + self.response.headers or {} + ) # GUARANTEE OPENAI HEADERS IN RESPONSE + def _process_chunk(self, chunk) -> Optional[ResponsesAPIStreamingResponse]: """Process a single chunk of data from the stream""" if not chunk: diff --git a/tests/test_litellm/proxy/test_model_id_header_propagation.py b/tests/test_litellm/proxy/test_model_id_header_propagation.py new file mode 100644 index 0000000000..cc4e7c084d --- /dev/null +++ b/tests/test_litellm/proxy/test_model_id_header_propagation.py @@ -0,0 +1,250 @@ +""" +Test that x-litellm-model-id header is propagated correctly on error responses. + +This test suite verifies the `maybe_get_model_id` method +which is responsible for extracting model_id from different locations +depending on the request lifecycle stage. +""" + +import pytest +from unittest.mock import MagicMock + +from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessing +from litellm.proxy._types import UserAPIKeyAuth + + +def test_maybe_get_model_id_from_litellm_params(): + """ + Test extraction of model_id from logging_obj.litellm_params (used by /v1/chat/completions). + """ + # Create a ProxyBaseLLMRequestProcessing instance + processor = ProxyBaseLLMRequestProcessing(data={}) + + # Create a mock logging object with model_info in litellm_params + mock_logging_obj = MagicMock() + mock_logging_obj.litellm_params = { + "model_info": { + "id": "test-model-id-from-litellm-params" + } + } + + # Test extraction + model_id = processor.maybe_get_model_id(mock_logging_obj) + + assert model_id == "test-model-id-from-litellm-params" + + +def test_maybe_get_model_id_from_litellm_params_nested(): + """ + Test extraction of model_id from nested metadata in logging_obj.litellm_params. + """ + processor = ProxyBaseLLMRequestProcessing(data={}) + + # Create a mock logging object with model_info nested in metadata + mock_logging_obj = MagicMock() + mock_logging_obj.litellm_params = { + "metadata": { + "model_info": { + "id": "test-model-id-nested" + } + } + } + + # Test extraction + model_id = processor.maybe_get_model_id(mock_logging_obj) + + assert model_id == "test-model-id-nested" + + +def test_maybe_get_model_id_from_kwargs(): + """ + Test extraction of model_id from logging_obj.kwargs (fallback path). + """ + processor = ProxyBaseLLMRequestProcessing(data={}) + + # Create a mock logging object with model_info in kwargs + mock_logging_obj = MagicMock() + mock_logging_obj.litellm_params = None + mock_logging_obj.kwargs = { + "litellm_params": { + "model_info": { + "id": "test-model-id-from-kwargs" + } + } + } + + # Test extraction + model_id = processor.maybe_get_model_id(mock_logging_obj) + + assert model_id == "test-model-id-from-kwargs" + + +def test_maybe_get_model_id_from_data(): + """ + Test extraction of model_id from self.data (used by /v1/messages and /v1/responses). + """ + # Create a processor with model_info in data + processor = ProxyBaseLLMRequestProcessing(data={ + "litellm_metadata": { + "model_info": { + "id": "test-model-id-from-data" + } + } + }) + + # Create a mock logging object without model_info + mock_logging_obj = MagicMock() + mock_logging_obj.litellm_params = {} + mock_logging_obj.kwargs = {} + + # Test extraction - should fall back to self.data + model_id = processor.maybe_get_model_id(mock_logging_obj) + + assert model_id == "test-model-id-from-data" + + +def test_maybe_get_model_id_no_logging_obj(): + """ + Test extraction of model_id when logging_obj is None (should use self.data). + """ + # Create a processor with model_info in data + processor = ProxyBaseLLMRequestProcessing(data={ + "litellm_metadata": { + "model_info": { + "id": "test-model-id-no-logging-obj" + } + } + }) + + # Test extraction with None logging_obj + model_id = processor.maybe_get_model_id(None) + + assert model_id == "test-model-id-no-logging-obj" + + +def test_maybe_get_model_id_not_found(): + """ + Test extraction of model_id when it's not available anywhere (should return None). + """ + processor = ProxyBaseLLMRequestProcessing(data={}) + + # Create a mock logging object without model_info anywhere + mock_logging_obj = MagicMock() + mock_logging_obj.litellm_params = {} + mock_logging_obj.kwargs = {} + + # Test extraction - should return None + model_id = processor.maybe_get_model_id(mock_logging_obj) + + assert model_id is None + + +def test_maybe_get_model_id_priority_litellm_params_over_data(): + """ + Test that model_id from logging_obj.litellm_params takes priority over self.data. + """ + # Create a processor with model_info in both places + processor = ProxyBaseLLMRequestProcessing(data={ + "litellm_metadata": { + "model_info": { + "id": "model-id-from-data" + } + } + }) + + # Create a mock logging object with model_info + mock_logging_obj = MagicMock() + mock_logging_obj.litellm_params = { + "model_info": { + "id": "model-id-from-litellm-params" + } + } + + # Test extraction - should prefer litellm_params + model_id = processor.maybe_get_model_id(mock_logging_obj) + + assert model_id == "model-id-from-litellm-params" + + +def test_get_custom_headers_includes_model_id(): + """ + Test that get_custom_headers includes x-litellm-model-id when model_id is provided. + """ + # Create mock user_api_key_dict with all required attributes + mock_user_api_key_dict = MagicMock() + mock_user_api_key_dict.user_id = "test-user" + mock_user_api_key_dict.team_id = "test-team" + mock_user_api_key_dict.tpm_limit = 1000 + mock_user_api_key_dict.rpm_limit = 100 + + # Call get_custom_headers with a model_id + headers = ProxyBaseLLMRequestProcessing.get_custom_headers( + user_api_key_dict=mock_user_api_key_dict, + model_id="test-model-123", + cache_key="test-cache-key", + api_base="https://api.example.com", + version="1.0.0", + response_cost=0.001, + request_data={}, + hidden_params={} + ) + + # Verify model_id is in headers + assert "x-litellm-model-id" in headers + assert headers["x-litellm-model-id"] == "test-model-123" + + +def test_get_custom_headers_without_model_id(): + """ + Test that get_custom_headers works correctly when model_id is None or empty. + """ + # Create mock user_api_key_dict with all required attributes + mock_user_api_key_dict = MagicMock() + mock_user_api_key_dict.user_id = "test-user" + mock_user_api_key_dict.team_id = "test-team" + mock_user_api_key_dict.tpm_limit = 1000 + mock_user_api_key_dict.rpm_limit = 100 + + # Call get_custom_headers without a model_id + headers = ProxyBaseLLMRequestProcessing.get_custom_headers( + user_api_key_dict=mock_user_api_key_dict, + model_id=None, + cache_key="test-cache-key", + api_base="https://api.example.com", + version="1.0.0", + response_cost=0.001, + request_data={}, + hidden_params={} + ) + + # x-litellm-model-id should not be in headers (or should be empty/None) + if "x-litellm-model-id" in headers: + assert headers["x-litellm-model-id"] in [None, ""] + + +def test_get_custom_headers_with_empty_string_model_id(): + """ + Test that get_custom_headers handles empty string model_id correctly. + """ + # Create mock user_api_key_dict with all required attributes + mock_user_api_key_dict = MagicMock() + mock_user_api_key_dict.user_id = "test-user" + mock_user_api_key_dict.team_id = "test-team" + mock_user_api_key_dict.tpm_limit = 1000 + mock_user_api_key_dict.rpm_limit = 100 + + # Call get_custom_headers with empty string model_id + headers = ProxyBaseLLMRequestProcessing.get_custom_headers( + user_api_key_dict=mock_user_api_key_dict, + model_id="", + cache_key="test-cache-key", + api_base="https://api.example.com", + version="1.0.0", + response_cost=0.001, + request_data={}, + hidden_params={} + ) + + # x-litellm-model-id should not be in headers (or should be empty) + if "x-litellm-model-id" in headers: + assert headers["x-litellm-model-id"] == ""