scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-24 00:52:24 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 90

feat(router): expose x-litellm-fallback-model-used header in responses

Browse source 
When a fallback model is used instead of the primary model, stamp
x-litellm-fallback-model-used on the response's additional_headers so
callers can tell which model actually served the request.

- add `fallback_model` param to `add_fallback_headers_to_response()`
- capture effective fallback model name in `run_async_fallback()` before
  the recursive call and pass it through to the header helper
- add unit tests covering header presence/absence and edge cases

Fixes https://github.com/BerriAI/litellm/issues/25503

Co-Authored-By: Claude Sonnet 4.6 (1M context) <noreply@anthropic.com>
This commit is contained in:
Ankit Jhalaria 2026-04-14 13:47:54 -07:00
parent e9e86ed956
commit 29d8288c58
4 changed files with 215 additions and 1 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

7
litellm/router_utils/add_retry_fallback_headers.py
View file

@ -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)

6
litellm/router_utils/fallback_event_handlers.py
View file

@ -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(

0
tests/test_litellm/router/__init__.py Normal file
View file

203
tests/test_litellm/router/test_fallback_headers.py Normal file
View file

@ -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",
)