diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 3f8b01cc865..47976f11698 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -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: diff --git a/litellm/router.py b/litellm/router.py index e54eadfb872..acce1c58a8b 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -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 diff --git a/litellm/router_utils/add_retry_fallback_headers.py b/litellm/router_utils/add_retry_fallback_headers.py index 6b921a0db8a..0b927714ca9 100644 --- a/litellm/router_utils/add_retry_fallback_headers.py +++ b/litellm/router_utils/add_retry_fallback_headers.py @@ -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 diff --git a/litellm/router_utils/fallback_event_handlers.py b/litellm/router_utils/fallback_event_handlers.py index eb756e3cf8b..3635a403105 100644 --- a/litellm/router_utils/fallback_event_handlers.py +++ b/litellm/router_utils/fallback_event_handlers.py @@ -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, diff --git a/tests/test_litellm/proxy/proxy_server/test_streaming_helpers.py b/tests/test_litellm/proxy/proxy_server/test_streaming_helpers.py index 33de1ede917..865d733b954 100644 --- a/tests/test_litellm/proxy/proxy_server/test_streaming_helpers.py +++ b/tests/test_litellm/proxy/proxy_server/test_streaming_helpers.py @@ -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": "", "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, diff --git a/tests/test_litellm/router_utils/test_add_retry_fallback_headers.py b/tests/test_litellm/router_utils/test_add_retry_fallback_headers.py new file mode 100644 index 00000000000..2aee0f0a4ef --- /dev/null +++ b/tests/test_litellm/router_utils/test_add_retry_fallback_headers.py @@ -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 == [] diff --git a/tests/test_litellm/router_utils/test_fallback_event_handlers.py b/tests/test_litellm/router_utils/test_fallback_event_handlers.py new file mode 100644 index 00000000000..1eeeba7b7b0 --- /dev/null +++ b/tests/test_litellm/router_utils/test_fallback_event_handlers.py @@ -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 diff --git a/tests/test_litellm/test_router_streaming_fallback_metadata.py b/tests/test_litellm/test_router_streaming_fallback_metadata.py new file mode 100644 index 00000000000..d2d742a7401 --- /dev/null +++ b/tests/test_litellm/test_router_streaming_fallback_metadata.py @@ -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", + }