diff --git a/litellm/router_utils/add_retry_fallback_headers.py b/litellm/router_utils/add_retry_fallback_headers.py index 6b921a0db8a..8305acce725 100644 --- a/litellm/router_utils/add_retry_fallback_headers.py +++ b/litellm/router_utils/add_retry_fallback_headers.py @@ -50,6 +50,7 @@ def add_retry_headers_to_response( def add_fallback_headers_to_response( response: Any, attempted_fallbacks: int, + fallback_model: Optional[str] = None, ) -> Any: """ Add fallback headers to the response @@ -57,6 +58,8 @@ def add_fallback_headers_to_response( Args: response: The response to add the headers to attempted_fallbacks: The number of fallbacks attempted + fallback_model: The model that was actually used for the successful fallback. + Set to None when the primary model succeeded (no fallback occurred). Returns: The response with the headers added @@ -64,7 +67,9 @@ def add_fallback_headers_to_response( Note: It's intentional that we don't add max_fallbacks in response headers Want to avoid bloat in the response headers for performance. """ - fallback_headers = { + fallback_headers: dict = { "x-litellm-attempted-fallbacks": attempted_fallbacks, } + if fallback_model is not None: + fallback_headers["x-litellm-fallback-model-used"] = fallback_model return _add_headers_to_response(response, fallback_headers) diff --git a/litellm/router_utils/fallback_event_handlers.py b/litellm/router_utils/fallback_event_handlers.py index 62e706a0cf5..247d1e971d6 100644 --- a/litellm/router_utils/fallback_event_handlers.py +++ b/litellm/router_utils/fallback_event_handlers.py @@ -130,6 +130,11 @@ async def run_async_fallback( kwargs["model"] = mg elif isinstance(mg, dict): kwargs.update(mg) + # Capture the effective fallback model name before the recursive call + # so we can stamp it on the response header regardless of further fallbacks. + effective_fallback_model: Optional[str] = ( + mg if isinstance(mg, str) else kwargs.get("model") + ) kwargs.setdefault("metadata", {}).update( {"model_group": kwargs.get("model", None)} ) # update model_group used, if fallbacks are done @@ -143,6 +148,7 @@ async def run_async_fallback( response = add_fallback_headers_to_response( response=response, attempted_fallbacks=fallback_depth, + fallback_model=effective_fallback_model, ) # callback for successfull_fallback_event(): await log_success_fallback_event( diff --git a/tests/test_litellm/router/__init__.py b/tests/test_litellm/router/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/test_litellm/router/test_fallback_headers.py b/tests/test_litellm/router/test_fallback_headers.py new file mode 100644 index 00000000000..b49a63e9f09 --- /dev/null +++ b/tests/test_litellm/router/test_fallback_headers.py @@ -0,0 +1,203 @@ +""" +Tests for x-litellm-fallback-model-used header exposure. + +Covers issue: https://github.com/BerriAI/litellm/issues/25503 + +When a fallback model is used instead of the primary model, the response should +include an x-litellm-fallback-model-used header so callers can tell which model +actually served the request. +""" + +from typing import Any +from unittest.mock import AsyncMock, MagicMock, patch + +import pytest + +from litellm.router_utils.add_retry_fallback_headers import ( + add_fallback_headers_to_response, +) +from litellm.router_utils.fallback_event_handlers import run_async_fallback + + +# --------------------------------------------------------------------------- +# Unit tests for add_fallback_headers_to_response +# --------------------------------------------------------------------------- + + +class TestAddFallbackHeadersToResponse: + """Unit tests for add_fallback_headers_to_response.""" + + def _make_response(self) -> MagicMock: + """Return a minimal pydantic-like response mock.""" + from pydantic import BaseModel + + class _FakeResponse(BaseModel): + model: str = "gpt-4" + _hidden_params: dict = {} + + resp = _FakeResponse() + return resp + + def test_fallback_model_header_set_when_fallback_occurred(self): + """x-litellm-fallback-model-used is present when a fallback model is provided.""" + resp = self._make_response() + result = add_fallback_headers_to_response( + response=resp, + attempted_fallbacks=1, + fallback_model="claude-3-haiku", + ) + headers = result._hidden_params.get("additional_headers", {}) + assert headers.get("x-litellm-fallback-model-used") == "claude-3-haiku" + + def test_fallback_model_header_absent_when_no_fallback(self): + """x-litellm-fallback-model-used is NOT set when primary model succeeded.""" + resp = self._make_response() + result = add_fallback_headers_to_response( + response=resp, + attempted_fallbacks=0, + fallback_model=None, + ) + headers = result._hidden_params.get("additional_headers", {}) + assert "x-litellm-fallback-model-used" not in headers + + def test_attempted_fallbacks_header_always_set(self): + """x-litellm-attempted-fallbacks is always present regardless of fallback_model.""" + resp = self._make_response() + result = add_fallback_headers_to_response( + response=resp, + attempted_fallbacks=2, + ) + headers = result._hidden_params.get("additional_headers", {}) + assert headers.get("x-litellm-attempted-fallbacks") == 2 + + def test_fallback_model_default_is_none(self): + """Calling add_fallback_headers_to_response without fallback_model does not error.""" + resp = self._make_response() + # Should not raise + result = add_fallback_headers_to_response( + response=resp, + attempted_fallbacks=0, + ) + headers = result._hidden_params.get("additional_headers", {}) + assert "x-litellm-fallback-model-used" not in headers + + def test_returns_none_unchanged(self): + """If response is None, it is returned unchanged without error.""" + result = add_fallback_headers_to_response( + response=None, + attempted_fallbacks=1, + fallback_model="gpt-3.5-turbo", + ) + assert result is None + + def test_fallback_model_header_with_multiple_fallback_depths(self): + """Header captures the model used even when multiple fallback depths occurred.""" + resp = self._make_response() + result = add_fallback_headers_to_response( + response=resp, + attempted_fallbacks=3, + fallback_model="gpt-3.5-turbo", + ) + headers = result._hidden_params.get("additional_headers", {}) + assert headers.get("x-litellm-fallback-model-used") == "gpt-3.5-turbo" + assert headers.get("x-litellm-attempted-fallbacks") == 3 + + +# --------------------------------------------------------------------------- +# Unit tests for run_async_fallback (verifies fallback_model propagation) +# --------------------------------------------------------------------------- + + +class TestRunAsyncFallbackHeaderPropagation: + """Tests that run_async_fallback stamps x-litellm-fallback-model-used on success.""" + + @pytest.mark.asyncio + async def test_fallback_model_header_stamped_on_successful_string_fallback(self): + """ + When a string fallback model succeeds, x-litellm-fallback-model-used + should be set to that model's name. + """ + from pydantic import BaseModel + + class _FakeResponse(BaseModel): + model: str = "claude-3-haiku" + _hidden_params: dict = {} + + fake_response = _FakeResponse() + + mock_router = MagicMock() + mock_router.log_retry = MagicMock(side_effect=lambda kwargs, e: kwargs) + mock_router.async_function_with_fallbacks = AsyncMock( + return_value=fake_response + ) + + result = await run_async_fallback( + litellm_router=mock_router, + fallback_model_group=["claude-3-haiku"], + original_model_group="gpt-4", + original_exception=Exception("primary failed"), + max_fallbacks=3, + fallback_depth=0, + model="gpt-4", + ) + + headers = result._hidden_params.get("additional_headers", {}) + assert headers.get("x-litellm-fallback-model-used") == "claude-3-haiku" + + @pytest.mark.asyncio + async def test_fallback_model_header_not_present_without_fallback(self): + """ + When the primary model succeeds (no fallback), x-litellm-fallback-model-used + should NOT appear in the response headers. + """ + from pydantic import BaseModel + + class _FakeResponse(BaseModel): + model: str = "gpt-4" + _hidden_params: dict = {} + + fake_response = _FakeResponse() + result = add_fallback_headers_to_response( + response=fake_response, + attempted_fallbacks=0, + fallback_model=None, + ) + headers = result._hidden_params.get("additional_headers", {}) + assert "x-litellm-fallback-model-used" not in headers + + @pytest.mark.asyncio + async def test_all_fallbacks_fail_raises_exception(self): + """When all fallback models fail, the last exception is re-raised.""" + mock_router = MagicMock() + mock_router.log_retry = MagicMock(side_effect=lambda kwargs, e: kwargs) + mock_router.async_function_with_fallbacks = AsyncMock( + side_effect=Exception("fallback also failed") + ) + + with pytest.raises(Exception, match="fallback also failed"): + await run_async_fallback( + litellm_router=mock_router, + fallback_model_group=["claude-3-haiku"], + original_model_group="gpt-4", + original_exception=Exception("primary failed"), + max_fallbacks=3, + fallback_depth=0, + model="gpt-4", + ) + + @pytest.mark.asyncio + async def test_max_fallback_depth_raises_original_exception(self): + """When max_fallbacks is reached, the original exception is re-raised.""" + original_exc = Exception("original failure") + mock_router = MagicMock() + + with pytest.raises(Exception, match="original failure"): + await run_async_fallback( + litellm_router=mock_router, + fallback_model_group=["claude-3-haiku"], + original_model_group="gpt-4", + original_exception=original_exc, + max_fallbacks=3, + fallback_depth=3, # already at max + model="gpt-4", + )