scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-14 23:21:35 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 102

fix(router): send fallback metadata when streaming (#30914)

Browse source 
When a streaming request triggers a fallback, there was previously no way to
know it happened. This commit addresses this in a few ways:

1. The response now correctly populates the fallback headers
    (`x-litellm-attempted-fallbacks`) so callers know a fallback happened.
2. The correct model ID is passed in the streaming chunks.
3. A streaming chunk with the fallback error can be optionally sent back
    to the client (opt-in) by passing `include_fallback_errors: true` in
    the request.

The format of the fallback errors while streaming is intentionally OpenAI
compatible to not break existing libraries that parse these events. It was
tested with Vercel's AI SDK (ai-sdk.dev). It is also opt-in, so it is not
delieved unexpectedly to callers by default.
This commit is contained in:
Tal Marian 2026-06-22 15:30:45 +03:00 • committed by GitHub
parent e5ce7abefb
commit 7fa04dbfd9
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
8 changed files with 1048 additions and 58 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

144
litellm/proxy/proxy_server.py
View file

@ -106,6 +106,10 @@ from litellm.proxy.common_utils.callback_utils import (
process_callback,
)
from litellm.proxy.common_utils.realtime_utils import _realtime_request_body
from litellm.router_utils.add_retry_fallback_headers import (
get_fallback_errors_from_headers,
get_hidden_params_dict,
)
from litellm.types.utils import (
ModelResponse,
ModelResponseStream,
@ -7085,57 +7089,120 @@ def _get_client_requested_model_for_streaming(request_data: dict) -> str:
return requested_model if isinstance(requested_model, str) else ""
def _is_positive_int_like(value: Any) -> bool:
try:
return int(value) > 0
except (TypeError, ValueError):
return False
def _should_include_fallback_errors(request_data: dict[str, object]) -> bool:
return request_data.get("include_fallback_errors") is True
def _get_streaming_fallback_metadata(
response_obj: object,
) -> tuple[bool, str | None, list[dict[str, object]]]:
additional_headers = get_hidden_params_dict(response_obj).get("additional_headers")
if not isinstance(additional_headers, dict):
return False, None, []
if not _is_positive_int_like(
additional_headers.get("x-litellm-attempted-fallbacks")
):
return False, None, []
fallback_model = additional_headers.get("x-litellm-model-group")
fallback_errors = get_fallback_errors_from_headers(additional_headers)
if isinstance(fallback_model, str) and fallback_model:
return True, fallback_model, fallback_errors
return True, None, fallback_errors
def _format_fallback_metadata_sse_event(
*,
fallback_model: str | None,
fallback_errors: list[dict[str, object]],
) -> str:
import time
payload = {
"id": "litellm-fallback-metadata",
"object": "chat.completion.chunk",
"created": int(time.time()),
"model": fallback_model or "",
"choices": [],
"litellm_fallback": {
"fallback_model": fallback_model,
"errors": fallback_errors,
},
}
return f"data: {json.dumps(payload)}\n\n"
def _restamp_streaming_chunk_model(
*,
chunk: Any,
requested_model_from_client: str,
request_data: dict,
model_mismatch_logged: bool,
) -> Tuple[Any, bool]:
fallback_was_attempted: bool = False,
fallback_model_from_metadata: str | None = None,
) -> tuple[Any, bool]:
target_model = (
fallback_model_from_metadata
if fallback_was_attempted
else requested_model_from_client
)
# Always return the client-requested model name (not provider-prefixed internal identifiers)
# on streaming chunks.
# On fallback, use the public OpenAI-compatible model name. This keeps
# provider-prefixed internal identifiers from leaking into the public API.
#
# Note: This warning is intentionally verbose. A mismatch is a useful signal that an
# internal provider/deployment identifier is leaking into the public API, and helps
# maintainers/operators catch regressions while preserving OpenAI-compatible output.
if not requested_model_from_client or not isinstance(chunk, (BaseModel, dict)):
if not target_model or not isinstance(chunk, (BaseModel, dict)):
return chunk, model_mismatch_logged
# For Azure Model Router, preserve the actual model used in each chunk
if _is_azure_model_router_request(requested_model_from_client):
if not fallback_was_attempted and _is_azure_model_router_request(
requested_model_from_client
):
return chunk, model_mismatch_logged
# For fastest_response batch completions, preserve the winning model's name
# instead of stamping the comma-separated list the client sent.
if request_data.get("fastest_response", False):
if not fallback_was_attempted and request_data.get("fastest_response", False):
return chunk, model_mismatch_logged
downstream_model = (
chunk.get("model") if isinstance(chunk, dict) else getattr(chunk, "model", None)
)
if downstream_model == requested_model_from_client:
if downstream_model == target_model:
return chunk, model_mismatch_logged
if not model_mismatch_logged and downstream_model != requested_model_from_client:
if not model_mismatch_logged and downstream_model != target_model:
verbose_proxy_logger.debug(
"litellm_call_id=%s: streaming chunk model mismatch - requested=%r downstream=%r. Overriding model to requested.",
"litellm_call_id=%s: streaming chunk model mismatch - target=%r downstream=%r fallback_was_attempted=%s. Overriding chunk model to target.",
request_data.get("litellm_call_id"),
requested_model_from_client,
target_model,
downstream_model,
fallback_was_attempted,
)
model_mismatch_logged = True
if isinstance(chunk, dict):
chunk["model"] = requested_model_from_client
chunk["model"] = target_model
return chunk, model_mismatch_logged
try:
setattr(chunk, "model", requested_model_from_client)
chunk.model = target_model
except Exception as e:
verbose_proxy_logger.error(
"litellm_call_id=%s: failed to override chunk.model=%r on chunk_type=%s. error=%s",
request_data.get("litellm_call_id"),
requested_model_from_client,
target_model,
type(chunk),
str(e),
exc_info=True,
@ -7294,7 +7361,14 @@ async def async_data_generator(
requested_model_from_client = _get_client_requested_model_for_streaming(
request_data=request_data
)
(
fallback_was_attempted,
fallback_model_from_metadata,
fallback_errors,
) = _get_streaming_fallback_metadata(response)
model_mismatch_logged = False
fallback_metadata_event_sent = False
include_fallback_errors = _should_include_fallback_errors(request_data)
# Use a running string instead of list + join to avoid O(n^2) overhead.
# Previously "".join(str_so_far_parts) was called every chunk, re-joining
# the entire accumulated response. String += is O(n) amortized total.
@ -7332,13 +7406,37 @@ async def async_data_generator(
str_so_far=_str_so_far,
)
# Mid-stream fallbacks surface metadata on individual chunks rather than
# the response wrapper. Keep scanning chunks until a fallback model is
# resolved, then latch it for the rest of the stream.
if fallback_model_from_metadata is None:
(
chunk_fallback_was_attempted,
chunk_fallback_model,
chunk_fallback_errors,
) = _get_streaming_fallback_metadata(chunk)
if chunk_fallback_was_attempted:
fallback_was_attempted = True
fallback_model_from_metadata = chunk_fallback_model
fallback_errors = fallback_errors or chunk_fallback_errors
pending_fallback_event = (
include_fallback_errors
and fallback_was_attempted
and fallback_errors
and not fallback_metadata_event_sent
)
chunk, model_mismatch_logged = _restamp_streaming_chunk_model(
chunk=chunk,
requested_model_from_client=requested_model_from_client,
request_data=request_data,
model_mismatch_logged=model_mismatch_logged,
fallback_was_attempted=fallback_was_attempted,
fallback_model_from_metadata=fallback_model_from_metadata,
)
raw_passthrough = False
if isinstance(chunk, BaseModel):
chunk = _serialize_streaming_chunk(chunk)
elif isinstance(chunk, bytes):
@ -7354,14 +7452,14 @@ async def async_data_generator(
raise ValueError(
"Raw SSE stream exceeded maximum buffered size without a frame delimiter"
)
continue
if chunk.startswith(("data:", "event:", ":")):
raw_passthrough = True
elif chunk.startswith(("data:", "event:", ":")):
yield (
chunk
if chunk.endswith(_SSE_FRAME_DELIMITERS)
else chunk + "\n\n"
)
continue
raw_passthrough = True
elif isinstance(chunk, str) and is_raw_sse_stream:
raw_sse_buffer += chunk
while True:
@ -7373,15 +7471,23 @@ async def async_data_generator(
raise ValueError(
"Raw SSE stream exceeded maximum buffered size without a frame delimiter"
)
continue
raw_passthrough = True
elif isinstance(chunk, str) and chunk.startswith("data: "):
error_message = chunk
break
try:
yield _format_streaming_sse_chunk(chunk=chunk)
except Exception as e:
yield f"data: {str(e)}\n\n"
if not raw_passthrough:
try:
yield _format_streaming_sse_chunk(chunk=chunk)
except Exception as e:
yield f"data: {str(e)}\n\n"
if pending_fallback_event:
yield _format_fallback_metadata_sse_event(
fallback_model=fallback_model_from_metadata,
fallback_errors=fallback_errors,
)
fallback_metadata_event_sent = True
stream_completed = True
if not needs_iterator_wrap:

87
litellm/router.py
View file

@ -40,7 +40,6 @@ import anyio
import httpx
import openai
from openai import AsyncOpenAI
from pydantic import BaseModel
from typing_extensions import overload
import litellm
@ -81,8 +80,10 @@ from litellm.router_strategy.lowest_tpm_rpm_v2 import LowestTPMLoggingHandler_v2
from litellm.router_strategy.simple_shuffle import simple_shuffle
from litellm.router_strategy.tag_based_routing import get_deployments_for_tag
from litellm.router_utils.add_retry_fallback_headers import (
_HiddenParamsHost,
add_fallback_headers_to_response,
add_retry_headers_to_response,
get_hidden_params_dict,
)
from litellm.router_utils.batch_utils import (
_get_router_metadata_variable_name,
@ -2165,6 +2166,36 @@ class Router:
)
setattr(fallback_item, "usage", combined_usage)
@staticmethod
def _prepare_fallback_hidden_params(
fallback_response: object,
) -> tuple[dict[str, object], dict[str, object]]:
fallback_hidden_params = get_hidden_params_dict(fallback_response)
fallback_headers = fallback_hidden_params.get("additional_headers")
if not isinstance(fallback_headers, dict):
return fallback_hidden_params, {}
return fallback_hidden_params, cast("dict[str, object]", fallback_headers)
@staticmethod
def _apply_fallback_hidden_params_to_item(
fallback_item: object,
prepared_fallback_hidden_params: tuple[dict[str, object], dict[str, object]],
) -> None:
if fallback_item is None or not hasattr(fallback_item, "_hidden_params"):
return
fallback_hidden_params, fallback_headers = prepared_fallback_hidden_params
item_hidden_params = get_hidden_params_dict(fallback_item)
item_headers = item_hidden_params.get("additional_headers")
if not isinstance(item_headers, dict):
item_headers = {}
cast(_HiddenParamsHost, fallback_item)._hidden_params = {
**item_hidden_params,
**fallback_hidden_params,
"additional_headers": {**item_headers, **fallback_headers},
}
async def _acompletion_streaming_iterator(
self,
model_response: CustomStreamWrapper,
@ -2257,12 +2288,22 @@ class Router:
model_group=model_group,
args=(),
kwargs=initial_kwargs,
include_fallback_errors=initial_kwargs.get(
"include_fallback_errors", False
)
is True,
)
)
# If fallback returns a streaming response, iterate over it
if hasattr(fallback_response, "__aiter__"):
prepared_fallback_hidden_params = (
Router._prepare_fallback_hidden_params(fallback_response)
)
async for fallback_item in fallback_response: # type: ignore
Router._apply_fallback_hidden_params_to_item(
fallback_item, prepared_fallback_hidden_params
)
if (
fallback_item
and isinstance(fallback_item, ModelResponseStream)
@ -2686,11 +2727,21 @@ class Router:
model_group=model_group,
args=(),
kwargs=initial_kwargs,
include_fallback_errors=initial_kwargs.get(
"include_fallback_errors", False
)
is True,
)
)
if hasattr(fallback_response, "__aiter__"):
prepared_fallback_hidden_params = (
Router._prepare_fallback_hidden_params(fallback_response)
)
async for fallback_item in fallback_response: # type: ignore
Router._apply_fallback_hidden_params_to_item(
fallback_item, prepared_fallback_hidden_params
)
if partial_usage is not None:
Router._combine_responses_fallback_usage(
fallback_item, partial_usage
@ -2815,7 +2866,13 @@ class Router:
)
if hasattr(fallback_response, "__iter__"):
prepared_fallback_hidden_params = (
Router._prepare_fallback_hidden_params(fallback_response)
)
for fallback_item in fallback_response:
Router._apply_fallback_hidden_params_to_item(
fallback_item, prepared_fallback_hidden_params
)
if (
fallback_item
and isinstance(fallback_item, ModelResponseStream)
@ -2972,6 +3029,7 @@ class Router:
**kwargs,
}
input_kwargs.pop("silent_model", None)
input_kwargs.pop("include_fallback_errors", None)
_response = litellm.acompletion(**input_kwargs)
@ -6478,6 +6536,7 @@ class Router:
model_group: Optional[str],
args: tuple,
kwargs: dict,
include_fallback_errors: bool = False,
):
"""
Common utilities for async_function_with_fallbacks
@ -6501,6 +6560,8 @@ class Router:
input_kwargs["max_fallbacks"] = self.max_fallbacks
if "fallback_depth" not in input_kwargs:
input_kwargs["fallback_depth"] = 0
if include_fallback_errors:
input_kwargs["include_fallback_errors"] = True
# ORDER-BASED FALLBACKS: prepend higher order levels to the fallback list
# Skip for error types that have their own dedicated fallback handlers
@ -6759,6 +6820,7 @@ class Router:
If it fails after num_retries, fall back to another model group
"""
model_group: Optional[str] = kwargs.get("model")
include_fallback_errors = kwargs.get("include_fallback_errors", False) is True
disable_fallbacks: Optional[bool] = kwargs.pop("disable_fallbacks", False)
fallbacks: Optional[List] = kwargs.get("fallbacks", self.fallbacks)
context_window_fallbacks: Optional[List] = kwargs.get(
@ -6802,6 +6864,7 @@ class Router:
model_group,
args,
kwargs,
include_fallback_errors=include_fallback_errors,
)
def _handle_mock_testing_fallbacks(
@ -9725,17 +9788,19 @@ class Router:
# - if healthy_deployments > 1, return model group rate limit headers
# - else return the model's rate limit headers
"""
if (
isinstance(response, BaseModel)
and hasattr(response, "_hidden_params")
and isinstance(response._hidden_params, dict) # type: ignore
):
response._hidden_params.setdefault("additional_headers", {}) # type: ignore
response._hidden_params["additional_headers"][ # type: ignore
"x-litellm-model-group"
] = model_group
if response is not None and hasattr(response, "_hidden_params"):
hidden_params = getattr(response, "_hidden_params", {}) or {}
if hasattr(hidden_params, "model_dump"):
hidden_params = hidden_params.model_dump()
if not isinstance(hidden_params, dict):
return response
response._hidden_params = hidden_params
additional_headers = response._hidden_params["additional_headers"] # type: ignore
additional_headers = hidden_params.get("additional_headers")
if not isinstance(additional_headers, dict):
additional_headers = {}
hidden_params["additional_headers"] = additional_headers
additional_headers["x-litellm-model-group"] = model_group
# Lift QualityRouter routing decision into response headers for
# transparency. The decision is stashed in request_kwargs.metadata

118
litellm/router_utils/add_retry_fallback_headers.py
View file

@ -1,44 +1,99 @@
from typing import Any, Optional, Union
import json
from typing import Protocol, TypedDict, cast
from pydantic import BaseModel
from litellm.types.utils import HiddenParams
class FallbackErrorInfo(TypedDict):
message: str
type: str
param: str | None
code: str | None
def _add_headers_to_response(response: Any, headers: dict) -> Any:
class _HiddenParamsHost(Protocol):
_hidden_params: dict[str, object]
def get_hidden_params_dict(response: object) -> dict[str, object]:
hidden_params: object = cast(object, getattr(response, "_hidden_params", None))
if isinstance(hidden_params, BaseModel):
return cast("dict[str, object]", hidden_params.model_dump())
if isinstance(hidden_params, dict):
return cast("dict[str, object]", hidden_params)
return {}
def _ensure_additional_headers_dict(
hidden_params: dict[str, object],
) -> dict[str, object]:
additional_headers = hidden_params.get("additional_headers")
if isinstance(additional_headers, dict):
return cast("dict[str, object]", additional_headers)
return {}
def get_fallback_error_info(error: Exception) -> FallbackErrorInfo:
message = cast(object, getattr(error, "message", str(error)))
error_type = cast(object, getattr(error, "type", error.__class__.__name__))
param = cast(object, getattr(error, "param", None))
code = cast(object, getattr(error, "status_code", getattr(error, "code", None)))
return FallbackErrorInfo(
message=str(message),
type=str(error_type),
param=str(param) if param is not None else None,
code=str(code) if code is not None else None,
)
def _coerce_error_dicts(items: list[object]) -> list[dict[str, object]]:
return [cast("dict[str, object]", item) for item in items if isinstance(item, dict)]
def get_fallback_errors_from_headers(
additional_headers: dict[str, object],
) -> list[dict[str, object]]:
existing_errors = additional_headers.get("x-litellm-fallback-errors")
if isinstance(existing_errors, list):
return _coerce_error_dicts(cast("list[object]", existing_errors))
if isinstance(existing_errors, str):
try:
parsed_errors: object = cast(object, json.loads(existing_errors))
except json.JSONDecodeError:
return []
if isinstance(parsed_errors, list):
return _coerce_error_dicts(cast("list[object]", parsed_errors))
return []
def _add_headers_to_response(response: object, headers: dict[str, object]) -> object:
"""
Helper function to add headers to a response's hidden params
"""
if response is None or not isinstance(response, BaseModel):
if response is None:
return response
hidden_params: Optional[Union[dict, HiddenParams]] = getattr(
response, "_hidden_params", {}
)
if not isinstance(response, BaseModel) and not hasattr(response, "_hidden_params"):
return response
if hidden_params is None:
hidden_params_dict = {}
elif isinstance(hidden_params, HiddenParams):
hidden_params_dict = hidden_params.model_dump()
else:
hidden_params_dict = hidden_params
hidden_params = get_hidden_params_dict(response)
additional_headers = _ensure_additional_headers_dict(hidden_params)
additional_headers.update(headers)
hidden_params["additional_headers"] = additional_headers
hidden_params_dict.setdefault("additional_headers", {})
hidden_params_dict["additional_headers"].update(headers)
setattr(response, "_hidden_params", hidden_params_dict)
cast(_HiddenParamsHost, response)._hidden_params = hidden_params
return response
def add_retry_headers_to_response(
response: Any,
response: object,
attempted_retries: int,
max_retries: Optional[int] = None,
) -> Any:
max_retries: int | None = None,
) -> object:
"""
Add retry headers to the request
"""
retry_headers = {
retry_headers: dict[str, object] = {
"x-litellm-attempted-retries": attempted_retries,
}
if max_retries is not None:
@ -48,9 +103,10 @@ def add_retry_headers_to_response(
def add_fallback_headers_to_response(
response: Any,
response: object,
attempted_fallbacks: int,
) -> Any:
fallback_errors: list[FallbackErrorInfo] | None = None,
) -> object:
"""
Add fallback headers to the response
@ -64,7 +120,19 @@ 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[str, object] = {
"x-litellm-attempted-fallbacks": attempted_fallbacks,
}
return _add_headers_to_response(response, fallback_headers)
response = _add_headers_to_response(response, fallback_headers)
if fallback_errors is None or response is None:
return response
hidden_params = get_hidden_params_dict(response)
additional_headers = _ensure_additional_headers_dict(hidden_params)
merged_errors = get_fallback_errors_from_headers(additional_headers) + [
cast("dict[str, object]", error) for error in fallback_errors
]
additional_headers["x-litellm-fallback-errors"] = json.dumps(merged_errors)
hidden_params["additional_headers"] = additional_headers
cast(_HiddenParamsHost, response)._hidden_params = hidden_params
return response

7
litellm/router_utils/fallback_event_handlers.py
View file

@ -6,6 +6,7 @@ from litellm._logging import verbose_router_logger
from litellm.integrations.custom_logger import CustomLogger
from litellm.router_utils.add_retry_fallback_headers import (
add_fallback_headers_to_response,
get_fallback_error_info,
)
from litellm.types.router import LiteLLMParamsTypedDict
@ -90,6 +91,7 @@ async def run_async_fallback(
original_exception: Exception,
max_fallbacks: int,
fallback_depth: int,
include_fallback_errors: bool = False,
**kwargs,
) -> Any:
"""
@ -118,6 +120,7 @@ async def run_async_fallback(
raise original_exception
error_from_fallbacks = original_exception
fallback_errors = (get_fallback_error_info(original_exception),)
for mg in fallback_model_group:
if mg == original_model_group:
@ -143,6 +146,9 @@ async def run_async_fallback(
response = add_fallback_headers_to_response(
response=response,
attempted_fallbacks=fallback_depth,
fallback_errors=(
list(fallback_errors) if include_fallback_errors else None
),
)
# callback for successfull_fallback_event():
await log_success_fallback_event(
@ -153,6 +159,7 @@ async def run_async_fallback(
return response
except Exception as e:
error_from_fallbacks = e
fallback_errors = fallback_errors + (get_fallback_error_info(e),)
await log_failure_fallback_event(
original_model_group=original_model_group,
kwargs=kwargs,

350
tests/test_litellm/proxy/proxy_server/test_streaming_helpers.py
View file

@ -16,8 +16,6 @@ Pins covered:
from __future__ import annotations
import json
from typing import Any, AsyncIterator
from unittest.mock import AsyncMock, MagicMock
import pytest
@ -26,8 +24,11 @@ from litellm.proxy._types import UserAPIKeyAuth
from litellm.proxy.proxy_server import (
_apply_streaming_chunk_hooks,
_fast_serialize_simple_model_response_stream,
_format_fallback_metadata_sse_event,
_format_streaming_sse_chunk,
_get_client_requested_model_for_streaming,
_get_streaming_fallback_metadata,
_is_positive_int_like,
_restamp_streaming_chunk_model,
_serialize_streaming_chunk,
async_assistants_data_generator,
@ -71,6 +72,15 @@ async def _async_iter_raises(exc: Exception):
raise exc
class _FakeStream:
def __init__(self, chunks, hidden_params=None):
self._chunks = chunks
self._hidden_params = hidden_params or {}
def __aiter__(self):
return _async_iter(self._chunks)
# ---------------------------------------------------------------------------
# data_generator
# ---------------------------------------------------------------------------
@ -274,6 +284,34 @@ def test_restamp_streaming_chunk_model_overrides_model_on_dict():
assert logged is True
def test_restamp_streaming_chunk_model_uses_fallback_model_from_metadata():
chunk = _simple_chunk(model="openai/internal-fallback")
new_chunk, logged = _restamp_streaming_chunk_model(
chunk=chunk,
requested_model_from_client="primary-model",
request_data={"litellm_call_id": "id-1"},
model_mismatch_logged=False,
fallback_was_attempted=True,
fallback_model_from_metadata="fallback-model",
)
assert new_chunk.model == "fallback-model"
assert logged is True
def test_restamp_streaming_chunk_model_preserves_fallback_model_without_group():
chunk = _simple_chunk(model="openai/internal-fallback")
new_chunk, logged = _restamp_streaming_chunk_model(
chunk=chunk,
requested_model_from_client="primary-model",
request_data={},
model_mismatch_logged=False,
fallback_was_attempted=True,
fallback_model_from_metadata=None,
)
assert new_chunk.model == "openai/internal-fallback"
assert logged is False
def test_restamp_streaming_chunk_model_invalid_chunk_type_unchanged():
"""For a non-BaseModel, non-dict chunk the helper returns it as-is
along with the original ``model_mismatch_logged`` flag."""
@ -288,6 +326,147 @@ def test_restamp_streaming_chunk_model_invalid_chunk_type_unchanged():
assert logged is False
def test_is_positive_int_like_invalid_and_edge_values():
assert _is_positive_int_like(None) is False
assert _is_positive_int_like("not-a-number") is False
assert _is_positive_int_like(0) is False
assert _is_positive_int_like(-1) is False
assert _is_positive_int_like("1") is True
assert _is_positive_int_like(2) is True
def test_get_streaming_fallback_metadata_reads_headers():
fallback_errors = [
{
"message": "litellm.RateLimitError: upstream limited request",
"type": "RateLimitError",
"param": None,
"code": "429",
}
]
stream = _FakeStream(
[],
hidden_params={
"additional_headers": {
"x-litellm-attempted-fallbacks": "1",
"x-litellm-model-group": "fallback-model",
"x-litellm-fallback-errors": json.dumps(fallback_errors),
}
},
)
assert _get_streaming_fallback_metadata(stream) == (
True,
"fallback-model",
fallback_errors,
)
def test_get_streaming_fallback_metadata_no_additional_headers():
stream = _FakeStream([], hidden_params={})
assert _get_streaming_fallback_metadata(stream) == (False, None, [])
def test_get_streaming_fallback_metadata_zero_fallback_count():
stream = _FakeStream(
[],
hidden_params={
"additional_headers": {"x-litellm-attempted-fallbacks": 0}
},
)
assert _get_streaming_fallback_metadata(stream) == (False, None, [])
def test_get_streaming_fallback_metadata_no_model_group_returns_none_model():
stream = _FakeStream(
[],
hidden_params={
"additional_headers": {
"x-litellm-attempted-fallbacks": 1,
}
},
)
was_attempted, fallback_model, errors = _get_streaming_fallback_metadata(stream)
assert was_attempted is True
assert fallback_model is None
assert errors == []
def test_restamp_streaming_chunk_model_azure_router_preserves_model():
chunk = _simple_chunk(model="azure_ai/internal-deployment")
new_chunk, logged = _restamp_streaming_chunk_model(
chunk=chunk,
requested_model_from_client="azure_ai/model-router",
request_data={},
model_mismatch_logged=False,
)
assert new_chunk.model == "azure_ai/internal-deployment"
assert logged is False
def test_restamp_streaming_chunk_model_fastest_response_preserves_model():
chunk = _simple_chunk(model="winning-model")
new_chunk, logged = _restamp_streaming_chunk_model(
chunk=chunk,
requested_model_from_client="gpt-4,claude-3",
request_data={"fastest_response": True},
model_mismatch_logged=False,
)
assert new_chunk.model == "winning-model"
assert logged is False
def test_restamp_streaming_chunk_model_setattr_exception_logs_and_returns():
from pydantic import ConfigDict
class FrozenChunk(_simple_chunk().__class__):
model_config = ConfigDict(frozen=True)
chunk = FrozenChunk(
id="chatcmpl-test",
choices=[],
created=0,
model="openai/internal-x",
object="chat.completion.chunk",
)
new_chunk, logged = _restamp_streaming_chunk_model(
chunk=chunk,
requested_model_from_client="gpt-4",
request_data={"litellm_call_id": "test-id"},
model_mismatch_logged=False,
)
assert new_chunk.model == "openai/internal-x"
assert logged is True
def test_format_fallback_metadata_sse_event():
fallback_errors = [
{
"message": "litellm.RateLimitError: upstream limited request",
"type": "RateLimitError",
"param": None,
"code": "429",
}
]
event = _format_fallback_metadata_sse_event(
fallback_model="fallback-model",
fallback_errors=fallback_errors,
)
assert isinstance(event, str)
assert event.startswith("data: ")
payload = json.loads(event.removeprefix("data: ").removesuffix("\n\n"))
assert payload["choices"] == []
assert payload["litellm_fallback"] == {
"fallback_model": "fallback-model",
"errors": fallback_errors,
}
assert payload["id"] == "litellm-fallback-metadata"
assert payload["object"] == "chat.completion.chunk"
assert payload["model"] == "fallback-model"
assert isinstance(payload["created"], int)
# ---------------------------------------------------------------------------
# _fast_serialize_simple_model_response_stream
# ---------------------------------------------------------------------------
@ -473,7 +652,7 @@ async def test_async_data_generator_yields_sse_chunks_and_done(monkeypatch):
# First chunk is bytes (fast path) wrapped via _format_streaming_sse_chunk.
first = out[0]
assert isinstance(first, bytes)
payload = json.loads(first.removeprefix(b"data: ").rstrip(b"\n\n"))
payload = json.loads(first.removeprefix(b"data: ").removesuffix(b"\n\n"))
assert normalize(payload) == {
"id": "<VOLATILE>",
"object": "chat.completion.chunk",
@ -488,6 +667,171 @@ async def test_async_data_generator_yields_sse_chunks_and_done(monkeypatch):
}
@pytest.mark.asyncio
async def test_async_data_generator_uses_response_fallback_metadata(monkeypatch):
_patch_logging_flags(monkeypatch)
response = _FakeStream(
[_simple_chunk(model="openai/internal-fallback", content="hello")],
hidden_params={
"additional_headers": {
"x-litellm-attempted-fallbacks": 1,
"x-litellm-model-group": "fallback-model",
}
},
)
out = []
async for line in async_data_generator(
response=response,
user_api_key_dict=_user_auth(),
request_data={"model": "primary-model", "include_fallback_errors": True},
):
out.append(line)
first = out[0]
assert isinstance(first, bytes)
payload = json.loads(first.removeprefix(b"data: ").removesuffix(b"\n\n"))
assert payload["model"] == "fallback-model"
@pytest.mark.asyncio
async def test_async_data_generator_uses_chunk_fallback_metadata(monkeypatch):
_patch_logging_flags(monkeypatch)
chunk = _simple_chunk(model="openai/internal-fallback", content="hello")
chunk._hidden_params = {
"additional_headers": {
"x-litellm-attempted-fallbacks": 1,
"x-litellm-model-group": "fallback-model",
}
}
out = []
async for line in async_data_generator(
response=_async_iter([chunk]),
user_api_key_dict=_user_auth(),
request_data={"model": "primary-model"},
):
out.append(line)
first = out[0]
assert isinstance(first, bytes)
payload = json.loads(first.removeprefix(b"data: ").removesuffix(b"\n\n"))
assert payload["model"] == "fallback-model"
@pytest.mark.asyncio
async def test_async_data_generator_switches_model_mid_stream_on_fallback(monkeypatch):
"""Pre-fallback chunks keep the client-requested model; once a chunk carries
fallback metadata the model latches to the fallback group for the rest of the
stream. This pins the client-visible mid-stream model change."""
_patch_logging_flags(monkeypatch)
primary_chunk = _simple_chunk(model="openai/internal-primary", content="hi")
fallback_chunk = _simple_chunk(model="openai/internal-fallback", content="there")
fallback_chunk._hidden_params = {
"additional_headers": {
"x-litellm-attempted-fallbacks": 1,
"x-litellm-model-group": "fallback-model",
}
}
out = []
async for line in async_data_generator(
response=_async_iter([primary_chunk, fallback_chunk]),
user_api_key_dict=_user_auth(),
request_data={"model": "primary-model"},
):
out.append(line)
first_payload = json.loads(out[0].removeprefix(b"data: ").removesuffix(b"\n\n"))
second_payload = json.loads(out[1].removeprefix(b"data: ").removesuffix(b"\n\n"))
assert first_payload["model"] == "primary-model"
assert second_payload["model"] == "fallback-model"
@pytest.mark.asyncio
async def test_async_data_generator_emits_fallback_error_metadata_event(monkeypatch):
_patch_logging_flags(monkeypatch)
fallback_errors = [
{
"message": "litellm.RateLimitError: upstream limited request",
"type": "RateLimitError",
"param": None,
"code": "429",
}
]
response = _FakeStream(
[_simple_chunk(model="openai/internal-fallback", content="hello")],
hidden_params={
"additional_headers": {
"x-litellm-attempted-fallbacks": 1,
"x-litellm-model-group": "fallback-model",
"x-litellm-fallback-errors": json.dumps(fallback_errors),
}
},
)
out = []
async for line in async_data_generator(
response=response,
user_api_key_dict=_user_auth(),
request_data={"model": "primary-model", "include_fallback_errors": True},
):
out.append(line)
assert isinstance(out[0], bytes)
chunk_payload = json.loads(out[0].removeprefix(b"data: ").removesuffix(b"\n\n"))
assert chunk_payload["model"] == "fallback-model"
assert isinstance(out[1], str)
assert out[1].startswith("data: ")
metadata_payload = json.loads(out[1].removeprefix("data: ").removesuffix("\n\n"))
assert metadata_payload["choices"] == []
assert metadata_payload["litellm_fallback"] == {
"fallback_model": "fallback-model",
"errors": fallback_errors,
}
assert metadata_payload["id"] == "litellm-fallback-metadata"
assert metadata_payload["object"] == "chat.completion.chunk"
assert metadata_payload["model"] == "fallback-model"
assert isinstance(metadata_payload["created"], int)
@pytest.mark.asyncio
async def test_async_data_generator_skips_fallback_error_event_without_opt_in(
monkeypatch,
):
_patch_logging_flags(monkeypatch)
fallback_errors = [
{
"message": "litellm.RateLimitError: upstream limited request",
"type": "RateLimitError",
"param": None,
"code": "429",
}
]
response = _FakeStream(
[_simple_chunk(model="openai/internal-fallback", content="hello")],
hidden_params={
"additional_headers": {
"x-litellm-attempted-fallbacks": 1,
"x-litellm-model-group": "fallback-model",
"x-litellm-fallback-errors": json.dumps(fallback_errors),
}
},
)
out = []
async for line in async_data_generator(
response=response,
user_api_key_dict=_user_auth(),
request_data={"model": "primary-model"},
):
out.append(line)
assert isinstance(out[0], bytes)
payload = json.loads(out[0].removeprefix(b"data: ").removesuffix(b"\n\n"))
assert payload["model"] == "fallback-model"
@pytest.mark.asyncio
async def test_async_data_generator_mid_stream_exception_yields_error_payload(
monkeypatch,

142
tests/test_litellm/router_utils/test_add_retry_fallback_headers.py Normal file
View file

@ -0,0 +1,142 @@
import json
from pydantic import BaseModel
from litellm.router_utils.add_retry_fallback_headers import (
add_fallback_headers_to_response,
add_retry_headers_to_response,
get_fallback_errors_from_headers,
get_hidden_params_dict,
)
class StreamingWrapper:
def __init__(self):
self._hidden_params = {"additional_headers": {"x-existing": "keep"}}
def test_add_fallback_headers_to_streaming_wrapper():
response = StreamingWrapper()
result = add_fallback_headers_to_response(
response=response,
attempted_fallbacks=1,
)
assert result is response
assert response._hidden_params["additional_headers"] == {
"x-existing": "keep",
"x-litellm-attempted-fallbacks": 1,
}
def test_add_fallback_headers_serializes_fallback_errors():
response = StreamingWrapper()
fallback_errors = [
{
"message": "litellm.RateLimitError: upstream limited request",
"type": "RateLimitError",
"param": None,
"code": "429",
}
]
result = add_fallback_headers_to_response(
response=response,
attempted_fallbacks=1,
fallback_errors=fallback_errors,
)
assert result is response
assert response._hidden_params["additional_headers"][
"x-litellm-attempted-fallbacks"
] == 1
assert (
json.loads(
response._hidden_params["additional_headers"]["x-litellm-fallback-errors"]
)
== fallback_errors
)
def test_add_retry_headers_to_streaming_wrapper():
response = StreamingWrapper()
result = add_retry_headers_to_response(
response=response,
attempted_retries=2,
max_retries=3,
)
assert result is response
assert response._hidden_params["additional_headers"] == {
"x-existing": "keep",
"x-litellm-attempted-retries": 2,
"x-litellm-max-retries": 3,
}
def test_get_hidden_params_dict_with_pydantic_model_hidden_params():
class InnerHiddenParams(BaseModel):
additional_headers: dict = {}
class Response:
def __init__(self):
self._hidden_params = InnerHiddenParams(
additional_headers={"x-custom": "value"}
)
result = get_hidden_params_dict(Response())
assert result == {"additional_headers": {"x-custom": "value"}}
def test_get_hidden_params_dict_with_no_hidden_params():
class PlainResponse:
pass
assert get_hidden_params_dict(PlainResponse()) == {}
def test_add_fallback_headers_when_no_existing_additional_headers():
class NoHeadersWrapper:
def __init__(self):
self._hidden_params = {}
response = NoHeadersWrapper()
result = add_fallback_headers_to_response(response=response, attempted_fallbacks=2)
assert result is response
assert response._hidden_params["additional_headers"]["x-litellm-attempted-fallbacks"] == 2
def test_add_fallback_headers_returns_none_when_response_is_none():
result = add_fallback_headers_to_response(response=None, attempted_fallbacks=1)
assert result is None
def test_add_fallback_headers_returns_unchanged_when_response_has_no_hidden_params():
class PlainObject:
pass
obj = PlainObject()
result = add_fallback_headers_to_response(response=obj, attempted_fallbacks=1)
assert result is obj
assert not hasattr(obj, "_hidden_params")
def test_get_fallback_errors_from_headers_existing_list_passthrough():
errors = [{"message": "err", "type": "T", "param": None, "code": "400"}]
result = get_fallback_errors_from_headers({"x-litellm-fallback-errors": errors})
assert result == errors
def test_get_fallback_errors_from_headers_invalid_json_returns_empty():
result = get_fallback_errors_from_headers(
{"x-litellm-fallback-errors": "not-valid-json-{"}
)
assert result == []
def test_get_fallback_errors_from_headers_missing_key_returns_empty():
result = get_fallback_errors_from_headers({})
assert result == []

94
tests/test_litellm/router_utils/test_fallback_event_handlers.py Normal file
View file

@ -0,0 +1,94 @@
import json
import pytest
from litellm.router_utils.fallback_event_handlers import run_async_fallback
class StreamingWrapper:
def __init__(self):
self._hidden_params = {"additional_headers": {}}
class FakeRouter:
def log_retry(self, kwargs, e):
return kwargs
async def async_function_with_fallbacks(self, *args, **kwargs):
return StreamingWrapper()
class AlwaysFailRouter:
def log_retry(self, kwargs, e):
return kwargs
async def async_function_with_fallbacks(self, *args, **kwargs):
raise RuntimeError("fallback model also failed")
@pytest.mark.asyncio
async def test_run_async_fallback_adds_errors_when_opted_in():
response = await run_async_fallback(
litellm_router=FakeRouter(),
fallback_model_group=["fallback-model"],
original_model_group="primary-model",
original_exception=RuntimeError("upstream limited request"),
max_fallbacks=3,
fallback_depth=0,
include_fallback_errors=True,
)
additional_headers = response._hidden_params["additional_headers"]
assert additional_headers["x-litellm-attempted-fallbacks"] == 1
assert json.loads(additional_headers["x-litellm-fallback-errors"]) == [
{
"message": "upstream limited request",
"type": "RuntimeError",
"param": None,
"code": None,
}
]
@pytest.mark.asyncio
async def test_run_async_fallback_omits_errors_without_opt_in():
response = await run_async_fallback(
litellm_router=FakeRouter(),
fallback_model_group=["fallback-model"],
original_model_group="primary-model",
original_exception=RuntimeError("upstream limited request"),
max_fallbacks=3,
fallback_depth=0,
)
additional_headers = response._hidden_params["additional_headers"]
assert additional_headers["x-litellm-attempted-fallbacks"] == 1
assert "x-litellm-fallback-errors" not in additional_headers
@pytest.mark.asyncio
async def test_run_async_fallback_raises_when_all_fallbacks_fail():
with pytest.raises(RuntimeError, match="fallback model also failed"):
await run_async_fallback(
litellm_router=AlwaysFailRouter(),
fallback_model_group=["fallback-model"],
original_model_group="primary-model",
original_exception=RuntimeError("original request failed"),
max_fallbacks=3,
fallback_depth=0,
include_fallback_errors=True,
)
@pytest.mark.asyncio
async def test_run_async_fallback_skips_original_model_group():
response = await run_async_fallback(
litellm_router=FakeRouter(),
fallback_model_group=["primary-model", "fallback-model"],
original_model_group="primary-model",
original_exception=RuntimeError("original failed"),
max_fallbacks=3,
fallback_depth=0,
)
assert response._hidden_params["additional_headers"]["x-litellm-attempted-fallbacks"] == 1

164
tests/test_litellm/test_router_streaming_fallback_metadata.py Normal file
View file

@ -0,0 +1,164 @@
import json
from unittest.mock import MagicMock
import pytest
import litellm
from litellm.router import Router
from litellm.router_utils.add_retry_fallback_headers import get_hidden_params_dict
def test_apply_fallback_hidden_params_copies_from_fallback_response():
fallback_errors = [
{
"message": "litellm.RateLimitError: upstream limited request",
"type": "RateLimitError",
"param": None,
"code": "429",
}
]
chunk = litellm.ModelResponseStream(
id="test",
model="openai/internal-fallback",
choices=[],
)
chunk._hidden_params = {
"additional_headers": {"x-existing-chunk-header": "keep"},
"model_id": "chunk-model-id",
}
fallback_response = MagicMock()
fallback_response._hidden_params = {
"additional_headers": {
"x-litellm-attempted-fallbacks": 1,
"x-litellm-model-group": "fallback-model",
"x-litellm-fallback-errors": json.dumps(fallback_errors),
},
"api_base": "https://fallback.example",
}
Router._apply_fallback_hidden_params_to_item(
fallback_item=chunk,
prepared_fallback_hidden_params=Router._prepare_fallback_hidden_params(
fallback_response
),
)
assert chunk._hidden_params["api_base"] == "https://fallback.example"
assert chunk._hidden_params["model_id"] == "chunk-model-id"
assert chunk._hidden_params["additional_headers"] == {
"x-existing-chunk-header": "keep",
"x-litellm-attempted-fallbacks": 1,
"x-litellm-model-group": "fallback-model",
"x-litellm-fallback-errors": json.dumps(fallback_errors),
}
def _two_group_fallback_router() -> Router:
return litellm.Router(
model_list=[
{
"model_name": "primary-model",
"litellm_params": {"model": "openai/gpt-fake", "api_key": "sk-fake"},
},
{
"model_name": "fallback-model",
"litellm_params": {"model": "openai/gpt-fake-2", "api_key": "sk-fake"},
},
],
fallbacks=[{"primary-model": ["fallback-model"]}],
)
def _additional_headers(response: object) -> dict:
return get_hidden_params_dict(response).get("additional_headers", {})
@pytest.mark.asyncio
async def test_include_fallback_errors_propagates_through_router():
router = _two_group_fallback_router()
response = await router.acompletion(
model="primary-model",
messages=[{"role": "user", "content": "Hello"}],
mock_testing_fallbacks=True,
mock_response="fallback success",
include_fallback_errors=True,
)
headers = _additional_headers(response)
assert headers["x-litellm-attempted-fallbacks"] == 1
errors = json.loads(headers["x-litellm-fallback-errors"])
assert isinstance(errors, list) and len(errors) >= 1
assert set(errors[0].keys()) == {"message", "type", "param", "code"}
@pytest.mark.asyncio
async def test_router_omits_fallback_errors_without_opt_in():
router = _two_group_fallback_router()
response = await router.acompletion(
model="primary-model",
messages=[{"role": "user", "content": "Hello"}],
mock_testing_fallbacks=True,
mock_response="fallback success",
)
headers = _additional_headers(response)
assert headers["x-litellm-attempted-fallbacks"] == 1
assert "x-litellm-fallback-errors" not in headers
def test_prepare_fallback_hidden_params_no_additional_headers():
class FakeResponse:
_hidden_params = {"api_base": "http://example.com"}
hidden_params, headers = Router._prepare_fallback_hidden_params(FakeResponse())
assert hidden_params == {"api_base": "http://example.com"}
assert headers == {}
def test_apply_fallback_hidden_params_to_item_none_item():
Router._apply_fallback_hidden_params_to_item(
None, ({"api_base": "http://fallback.example"}, {"x-custom": "value"})
)
def test_apply_fallback_hidden_params_to_item_no_existing_additional_headers():
class FakeChunk:
_hidden_params = {"model_id": "test-id"}
chunk = FakeChunk()
Router._apply_fallback_hidden_params_to_item(
chunk,
(
{"api_base": "http://fallback.example"},
{"x-litellm-attempted-fallbacks": 1},
),
)
assert chunk._hidden_params["api_base"] == "http://fallback.example"
assert chunk._hidden_params["model_id"] == "test-id"
assert chunk._hidden_params["additional_headers"] == {
"x-litellm-attempted-fallbacks": 1
}
@pytest.mark.asyncio
async def test_set_response_headers_adds_model_group_to_streaming_wrapper():
class StreamingWrapper:
def __init__(self):
self._hidden_params = {"additional_headers": {"x-existing": "keep"}}
router = litellm.Router(model_list=[])
response = StreamingWrapper()
result = await router.set_response_headers(
response=response,
model_group="fallback-model",
)
assert result is response
assert response._hidden_params["additional_headers"] == {
"x-existing": "keep",
"x-litellm-model-group": "fallback-model",
}