mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-24 00:52:24 +00:00
fix(router): wrap every Responses and Messages fallback hop for mid-stream failover
The /v1/responses and /v1/messages streaming wrappers only ever wrapped the primary's stream, so a hop reached through the regular fallback chain had no mid-stream handler: its failure re-raised, or the outer wrapper retried the same entry with a fresh attempted set and never reached the rest of the list. Every attempt of the chain now runs through a per-endpoint attempt function that wraps its own stream, mirroring chat completions, and the per-request fallback and retry overrides ride a frozen carrier so each hop's re-entry still sees them after the retry layer pops them.
This commit is contained in:
mateo-berri
parent
db7d52eeda
commit
a2ae80ec9b
4 changed files with 328 additions and 104 deletions
|
|
@ -184,14 +184,17 @@ from litellm.router_utils.cooldown_handlers import (
|
|||
is_caller_timeout_408,
|
||||
)
|
||||
from litellm.router_utils.fallback_event_handlers import (
|
||||
MID_STREAM_FALLBACK_CONTROLS_KEY,
|
||||
AttemptedFallbackTargets,
|
||||
_check_non_standard_fallback_format,
|
||||
carry_over_pre_routing_selection,
|
||||
clear_pre_routing_selection,
|
||||
fallback_lookup_groups,
|
||||
fallbacks_disabled_for_request,
|
||||
get_fallback_model_group_for_lookup_groups,
|
||||
get_pre_routing_selection,
|
||||
has_unattempted_fallback_target,
|
||||
mid_stream_fallback_hop_kwargs,
|
||||
per_request_fallback_controls,
|
||||
record_disable_fallbacks,
|
||||
record_pre_routing_selection,
|
||||
run_async_fallback,
|
||||
|
|
@ -3305,12 +3308,7 @@ class Router:
|
|||
content_policy_fallbacks: Final[list | None] = initial_kwargs.get(
|
||||
"content_policy_fallbacks", self.content_policy_fallbacks
|
||||
)
|
||||
# Re-enter via the per-attempt helper so the fallback chain
|
||||
# picks deployments through
|
||||
# _ageneric_api_call_with_fallbacks_helper.
|
||||
# original_generic_function is preserved by the caller so
|
||||
# the helper knows what underlying API to invoke per attempt.
|
||||
initial_kwargs["original_function"] = self._ageneric_api_call_with_fallbacks_helper
|
||||
initial_kwargs["original_function"] = self._ageneric_api_call_with_fallbacks_responses_attempt
|
||||
if e.is_pre_first_chunk or not e.generated_content:
|
||||
# No content generated before the error — retry with the
|
||||
# original input. Adding a continuation prompt would
|
||||
|
|
@ -5141,22 +5139,28 @@ class Router:
|
|||
request_kwargs=None,
|
||||
)
|
||||
|
||||
async def _ageneric_api_call_with_fallbacks(self, model: str, original_function: Callable, **kwargs):
|
||||
async def _ageneric_api_call_with_fallbacks(
|
||||
self, model: str, original_function: Callable, attempt_function: Callable | None = None, **kwargs
|
||||
):
|
||||
"""
|
||||
Helper function to make a generic LLM API call through the router, this allows you to use retries/fallbacks with litellm router
|
||||
|
||||
attempt_function runs every attempt of the chain instead of the plain helper, so a streaming
|
||||
endpoint can wrap each attempt's stream with its own mid-stream fallback handling.
|
||||
"""
|
||||
try:
|
||||
kwargs["model"] = model
|
||||
kwargs["original_generic_function"] = original_function
|
||||
kwargs["original_function"] = self._ageneric_api_call_with_fallbacks_helper
|
||||
kwargs["original_function"] = attempt_function or self._ageneric_api_call_with_fallbacks_helper
|
||||
if attempt_function is not None:
|
||||
controls: Final = per_request_fallback_controls(kwargs)
|
||||
kwargs[MID_STREAM_FALLBACK_CONTROLS_KEY] = controls # rebind-ok: forwarded to every hop
|
||||
self._update_kwargs_before_fallbacks(model=model, kwargs=kwargs, metadata_variable_name="litellm_metadata")
|
||||
verbose_router_logger.debug(
|
||||
"Inside ageneric_api_call_with_fallbacks() - model: %s; kwargs: %s", model, kwargs
|
||||
)
|
||||
response: Final = await self.async_function_with_fallbacks(**kwargs)
|
||||
return response
|
||||
|
||||
return response
|
||||
except Exception as e:
|
||||
asyncio.create_task(
|
||||
send_llm_exception_alert(
|
||||
|
|
@ -5277,61 +5281,42 @@ class Router:
|
|||
self, original_function: Callable, **kwargs: Any
|
||||
) -> Union["ResponsesAPIResponse", "BaseResponsesAPIStreamingIterator"]:
|
||||
"""
|
||||
_ageneric_api_call_with_fallbacks for the Responses API, with the
|
||||
addition of mid-stream fallback handling.
|
||||
|
||||
When stream=True and the underlying call returns a
|
||||
BaseResponsesAPIStreamingIterator, wrap it with
|
||||
_aresponses_streaming_iterator so MidStreamFallbackError raised
|
||||
during iteration triggers the Router's cross-provider fallback chain.
|
||||
_ageneric_api_call_with_fallbacks for the Responses API, with every attempt's stream
|
||||
carrying its own mid-stream fallback handling
|
||||
(see _ageneric_api_call_with_fallbacks_responses_attempt).
|
||||
"""
|
||||
return await self._ageneric_api_call_with_fallbacks(
|
||||
original_function=original_function,
|
||||
attempt_function=self._ageneric_api_call_with_fallbacks_responses_attempt,
|
||||
**kwargs,
|
||||
)
|
||||
|
||||
async def _ageneric_api_call_with_fallbacks_responses_attempt(
|
||||
self,
|
||||
model: str,
|
||||
original_generic_function: Callable,
|
||||
**kwargs: object, # kwargs-ok: forwarded verbatim to the per-attempt helper, shape varies per call site
|
||||
) -> Union["ResponsesAPIResponse", "BaseResponsesAPIStreamingIterator"]:
|
||||
"""
|
||||
One attempt of the Responses API fallback chain. A streaming result is wrapped with
|
||||
_aresponses_streaming_iterator over this attempt's own kwargs, so a fallback hop that
|
||||
fails mid-stream resumes the original group's chain instead of re-raising; the name keeps
|
||||
_get_router_metadata_variable_name resolving to litellm_metadata for every hop.
|
||||
"""
|
||||
from litellm.litellm_core_utils.core_helpers import safe_deep_copy
|
||||
from litellm.responses.streaming_iterator import (
|
||||
BaseResponsesAPIStreamingIterator,
|
||||
)
|
||||
|
||||
# Snapshot the request kwargs before _ageneric_api_call_with_fallbacks
|
||||
# mutates them. A shallow copy alone is not enough: the primary
|
||||
# attempt mutates nested dicts in place — notably `litellm_metadata`,
|
||||
# which `_update_kwargs_with_deployment` populates with
|
||||
# deployment-specific fields (`deployment`, `model_info`, `api_base`,
|
||||
# tags, etc.). Without an explicit copy of that dict, the shallow
|
||||
# copy would still share its reference, leaking primary-deployment
|
||||
# metadata into the mid-stream fallback request.
|
||||
#
|
||||
# We avoid deep-copying the full kwargs because it can contain
|
||||
# non-deepcopyable objects (logging handles, async clients, etc.);
|
||||
# `safe_deep_copy` deep-copies the metadata dicts key-by-key with a
|
||||
# fallback to the original reference for any non-picklable value.
|
||||
# The original_generic_function is preserved so the per-attempt
|
||||
# helper knows which underlying API to call on fallback.
|
||||
# The pre-routing hook stamps its tier selection into this bucket during the primary
|
||||
# attempt; seeding it before the snapshot gives both the live kwargs and the copy a
|
||||
# bucket, so the post-call carry-over below always has somewhere to read and write.
|
||||
kwargs.setdefault("litellm_metadata", {}) # mutable-ok: shared bucket # rebind-ok: stamp must be readable here
|
||||
|
||||
fallback_kwargs: Final[dict[str, object]] = kwargs.copy()
|
||||
if isinstance(fallback_kwargs.get("litellm_metadata"), dict):
|
||||
fallback_kwargs["litellm_metadata"] = safe_deep_copy(fallback_kwargs["litellm_metadata"])
|
||||
if isinstance(fallback_kwargs.get("metadata"), dict):
|
||||
fallback_kwargs["metadata"] = safe_deep_copy(fallback_kwargs["metadata"])
|
||||
fallback_kwargs["original_generic_function"] = original_function
|
||||
|
||||
response: Final = await self._ageneric_api_call_with_fallbacks(original_function=original_function, **kwargs)
|
||||
|
||||
# The snapshot predates the pre-routing hook, so the tier it stamped into the live kwargs
|
||||
# is carried over write-or-clear: a stale or caller-supplied selection left in the copy
|
||||
# would key the mid-stream fallback lookup off a tier this attempt never routed to.
|
||||
clear_pre_routing_selection(fallback_kwargs)
|
||||
live_pre_routing_selection: Final = get_pre_routing_selection(kwargs)
|
||||
if live_pre_routing_selection is not None:
|
||||
record_pre_routing_selection(fallback_kwargs, live_pre_routing_selection)
|
||||
|
||||
controls: Final = kwargs.pop(MID_STREAM_FALLBACK_CONTROLS_KEY, None)
|
||||
hop_kwargs: Final = mid_stream_fallback_hop_kwargs(
|
||||
model=model, original_generic_function=original_generic_function, controls=controls, kwargs=kwargs
|
||||
)
|
||||
response: Final = await self._ageneric_api_call_with_fallbacks_helper(
|
||||
model=model, original_generic_function=original_generic_function, **kwargs
|
||||
)
|
||||
carry_over_pre_routing_selection(live_kwargs=kwargs, snapshot=hop_kwargs)
|
||||
if kwargs.get("stream") and isinstance(response, BaseResponsesAPIStreamingIterator):
|
||||
return await self._aresponses_streaming_iterator(
|
||||
response=response,
|
||||
initial_kwargs=fallback_kwargs,
|
||||
)
|
||||
return await self._aresponses_streaming_iterator(response=response, initial_kwargs=hop_kwargs)
|
||||
return response
|
||||
|
||||
async def _aanthropic_messages_streaming_iterator(
|
||||
|
|
@ -5560,7 +5545,7 @@ class Router:
|
|||
content_policy_fallbacks: Final[list | None] = initial_kwargs.get( # mutable-ok: matches the param below
|
||||
"content_policy_fallbacks", self.content_policy_fallbacks
|
||||
)
|
||||
initial_kwargs["original_function"] = self._ageneric_api_call_with_fallbacks_helper
|
||||
initial_kwargs["original_function"] = self._ageneric_api_call_with_fallbacks_anthropic_messages_attempt
|
||||
self._update_kwargs_before_fallbacks(
|
||||
model=model_group,
|
||||
kwargs=initial_kwargs,
|
||||
|
|
@ -5614,46 +5599,41 @@ class Router:
|
|||
**kwargs: object, # kwargs-ok: forwarded verbatim to original_function, shape varies per call site
|
||||
) -> Union["AnthropicMessagesResponse", AsyncIterator[bytes]]:
|
||||
"""
|
||||
_ageneric_api_call_with_fallbacks for anthropic_messages, with the
|
||||
addition of mid-stream fallback handling (see
|
||||
_aanthropic_messages_streaming_iterator). Parity with
|
||||
_ageneric_api_call_with_fallbacks for anthropic_messages, with every attempt's stream
|
||||
carrying its own mid-stream fallback handling
|
||||
(see _ageneric_api_call_with_fallbacks_anthropic_messages_attempt). Parity with
|
||||
_aresponses_with_streaming_fallbacks for the Responses API.
|
||||
"""
|
||||
from litellm.litellm_core_utils.core_helpers import safe_deep_copy
|
||||
|
||||
# Snapshot the request kwargs before the primary attempt mutates them
|
||||
# in place: _update_kwargs_with_deployment writes deployment-specific
|
||||
# fields (deployment, model_info, api_base, tags, ...) into the
|
||||
# SAME litellm_metadata/metadata dicts a shallow .copy() would still
|
||||
# share, leaking primary-deployment metadata into the mid-stream
|
||||
# fallback request. safe_deep_copy avoids deep-copying the full
|
||||
# kwargs (which can hold non-deepcopyable logging handles/clients).
|
||||
# The pre-routing hook stamps its tier selection into this bucket during the primary
|
||||
# attempt; seeding it before the snapshot gives both the live kwargs and the copy a
|
||||
# bucket, so the post-call carry-over below always has somewhere to read and write.
|
||||
kwargs.setdefault("litellm_metadata", {}) # mutable-ok: shared bucket # rebind-ok: stamp must be readable here
|
||||
|
||||
fallback_kwargs: Final[dict[str, object]] = kwargs.copy() # mutable-ok: mutated below before re-entry
|
||||
if isinstance(fallback_kwargs.get("litellm_metadata"), dict):
|
||||
fallback_kwargs["litellm_metadata"] = safe_deep_copy(fallback_kwargs["litellm_metadata"])
|
||||
if isinstance(fallback_kwargs.get("metadata"), dict):
|
||||
fallback_kwargs["metadata"] = safe_deep_copy(fallback_kwargs["metadata"])
|
||||
fallback_kwargs["original_generic_function"] = original_function
|
||||
|
||||
response: Final = await self._ageneric_api_call_with_fallbacks(original_function=original_function, **kwargs)
|
||||
|
||||
# The snapshot predates the pre-routing hook, so the tier it stamped into the live kwargs
|
||||
# is carried over write-or-clear: a stale or caller-supplied selection left in the copy
|
||||
# would key the mid-stream fallback lookup off a tier this attempt never routed to.
|
||||
clear_pre_routing_selection(fallback_kwargs)
|
||||
live_pre_routing_selection: Final = get_pre_routing_selection(kwargs)
|
||||
if live_pre_routing_selection is not None:
|
||||
record_pre_routing_selection(fallback_kwargs, live_pre_routing_selection)
|
||||
return await self._ageneric_api_call_with_fallbacks(
|
||||
original_function=original_function,
|
||||
attempt_function=self._ageneric_api_call_with_fallbacks_anthropic_messages_attempt,
|
||||
**kwargs,
|
||||
)
|
||||
|
||||
async def _ageneric_api_call_with_fallbacks_anthropic_messages_attempt(
|
||||
self,
|
||||
model: str,
|
||||
original_generic_function: Callable,
|
||||
**kwargs: object, # kwargs-ok: forwarded verbatim to the per-attempt helper, shape varies per call site
|
||||
) -> Union["AnthropicMessagesResponse", AsyncIterator[bytes]]:
|
||||
"""
|
||||
One attempt of the anthropic_messages fallback chain. A streaming result is wrapped with
|
||||
_aanthropic_messages_streaming_iterator over this attempt's own kwargs, so a fallback hop
|
||||
that fails mid-stream resumes the original group's chain instead of re-raising; the name
|
||||
keeps _get_router_metadata_variable_name resolving to litellm_metadata for every hop.
|
||||
"""
|
||||
controls: Final = kwargs.pop(MID_STREAM_FALLBACK_CONTROLS_KEY, None)
|
||||
hop_kwargs: Final = mid_stream_fallback_hop_kwargs(
|
||||
model=model, original_generic_function=original_generic_function, controls=controls, kwargs=kwargs
|
||||
)
|
||||
response: Final = await self._ageneric_api_call_with_fallbacks_helper(
|
||||
model=model, original_generic_function=original_generic_function, **kwargs
|
||||
)
|
||||
carry_over_pre_routing_selection(live_kwargs=kwargs, snapshot=hop_kwargs)
|
||||
if kwargs.get("stream") and hasattr(response, "__aiter__"):
|
||||
return await self._aanthropic_messages_streaming_iterator(
|
||||
response=cast("AsyncIterator[bytes]", response), # cast-ok: stream=True always returns a byte iterator
|
||||
initial_kwargs=fallback_kwargs,
|
||||
initial_kwargs=hop_kwargs,
|
||||
)
|
||||
return response
|
||||
|
||||
|
|
|
|||
|
|
@ -1,6 +1,6 @@
|
|||
import hashlib
|
||||
import json
|
||||
from collections.abc import Mapping, Sequence
|
||||
from collections.abc import Callable, Mapping, Sequence
|
||||
from dataclasses import dataclass
|
||||
from datetime import datetime
|
||||
from enum import Enum
|
||||
|
|
@ -10,7 +10,7 @@ from typing import TYPE_CHECKING, Any, Final
|
|||
import litellm
|
||||
from litellm._logging import verbose_router_logger
|
||||
from litellm.integrations.custom_logger import CustomLogger
|
||||
from litellm.litellm_core_utils.core_helpers import get_metadata_variable_name_from_kwargs
|
||||
from litellm.litellm_core_utils.core_helpers import get_metadata_variable_name_from_kwargs, safe_deep_copy
|
||||
from litellm.litellm_core_utils.sensitive_data_masker import mask_sensitive_structure
|
||||
from litellm.router_utils.add_retry_fallback_headers import (
|
||||
add_fallback_headers_to_response,
|
||||
|
|
@ -284,6 +284,76 @@ def get_pre_routing_selection(kwargs: Mapping[str, object]) -> str | None:
|
|||
return next((selected for selected in selections if isinstance(selected, str) and selected), None)
|
||||
|
||||
|
||||
def carry_over_pre_routing_selection(live_kwargs: Mapping[str, object], snapshot: Mapping[str, object]) -> None:
|
||||
"""
|
||||
Replace whatever selection the snapshot carries with the one the pre-routing hook stamped
|
||||
into the live kwargs while routing this attempt, so a mid-stream fallback keys its lookup
|
||||
off the tier this attempt actually routed to.
|
||||
"""
|
||||
clear_pre_routing_selection(snapshot)
|
||||
live_selection: Final = get_pre_routing_selection(live_kwargs)
|
||||
if live_selection is not None:
|
||||
record_pre_routing_selection(snapshot, live_selection)
|
||||
|
||||
|
||||
MID_STREAM_FALLBACK_CONTROLS_KEY: Final = "_mid_stream_fallback_controls"
|
||||
_PER_REQUEST_FALLBACK_CONTROL_KEYS: Final = (
|
||||
"fallbacks",
|
||||
"context_window_fallbacks",
|
||||
"content_policy_fallbacks",
|
||||
"num_retries",
|
||||
"model_group_retry_policy",
|
||||
)
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class MidStreamFallbackControls:
|
||||
"""
|
||||
The per-request fallback and retry overrides every streaming attempt must see again.
|
||||
|
||||
async_function_with_retries pops them before the attempt function runs, so without this
|
||||
carrier a fallback hop's own mid-stream re-entry would fall back to the router-level settings.
|
||||
"""
|
||||
|
||||
overrides: Mapping[str, object]
|
||||
|
||||
|
||||
_NO_FALLBACK_CONTROLS: Final = MidStreamFallbackControls(MappingProxyType({}))
|
||||
|
||||
|
||||
def per_request_fallback_controls(kwargs: Mapping[str, object]) -> MidStreamFallbackControls:
|
||||
return MidStreamFallbackControls(
|
||||
MappingProxyType({key: kwargs[key] for key in _PER_REQUEST_FALLBACK_CONTROL_KEYS if key in kwargs})
|
||||
)
|
||||
|
||||
|
||||
def mid_stream_fallback_hop_kwargs(
|
||||
model: str,
|
||||
original_generic_function: Callable[..., object],
|
||||
controls: object,
|
||||
kwargs: Mapping[str, object],
|
||||
) -> dict[str, object]: # mutable-ok: the streaming iterators rewrite it in place when they re-enter the chain
|
||||
"""
|
||||
The kwargs one streaming attempt re-enters the fallback chain with if its stream fails.
|
||||
|
||||
A shallow copy keeps ``attempted_targets`` shared with the outer chain, so entries this
|
||||
request already tried are never retried; the metadata buckets are copied key by key because
|
||||
the attempt writes deployment-specific fields into them in place.
|
||||
"""
|
||||
hop_controls: Final = controls if isinstance(controls, MidStreamFallbackControls) else _NO_FALLBACK_CONTROLS
|
||||
copied_buckets: Final = MappingProxyType(
|
||||
{name: safe_deep_copy(kwargs[name]) for name in _ROUTER_METADATA_BUCKETS if isinstance(kwargs.get(name), dict)}
|
||||
)
|
||||
return { # mutable-ok: handed to the streaming iterator as its initial_kwargs, which it rewrites on re-entry
|
||||
**kwargs,
|
||||
**copied_buckets,
|
||||
**hop_controls.overrides,
|
||||
MID_STREAM_FALLBACK_CONTROLS_KEY: hop_controls,
|
||||
"model": model,
|
||||
"original_generic_function": original_generic_function,
|
||||
}
|
||||
|
||||
|
||||
DISABLE_FALLBACKS_METADATA_KEY: Final = "_disable_fallbacks"
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -251,7 +251,7 @@ async def test_aresponses_with_streaming_fallbacks_non_streaming_passthrough():
|
|||
|
||||
with patch.object(
|
||||
router,
|
||||
"_ageneric_api_call_with_fallbacks",
|
||||
"_ageneric_api_call_with_fallbacks_helper",
|
||||
new=AsyncMock(return_value=plain_response),
|
||||
):
|
||||
out = await router._aresponses_with_streaming_fallbacks(
|
||||
|
|
@ -278,7 +278,7 @@ async def test_aresponses_with_streaming_fallbacks_wraps_streaming_iterator():
|
|||
|
||||
with patch.object(
|
||||
router,
|
||||
"_ageneric_api_call_with_fallbacks",
|
||||
"_ageneric_api_call_with_fallbacks_helper",
|
||||
new=AsyncMock(return_value=streaming_iter),
|
||||
), patch.object(
|
||||
router,
|
||||
|
|
@ -294,6 +294,135 @@ async def test_aresponses_with_streaming_fallbacks_wraps_streaming_iterator():
|
|||
mock_wrap.assert_awaited_once()
|
||||
|
||||
|
||||
# -------- every fallback entry stays reachable across hops --------
|
||||
|
||||
|
||||
def _make_three_tier_router(**router_kwargs) -> Router:
|
||||
return Router(
|
||||
model_list=[
|
||||
{"model_name": "primary", "litellm_params": {"model": "openai/primary-model", "api_key": "sk-test"}},
|
||||
{"model_name": "fb1", "litellm_params": {"model": "openai/fb1-model", "api_key": "sk-test"}},
|
||||
{"model_name": "fb2", "litellm_params": {"model": "openai/fb2-model", "api_key": "sk-test"}},
|
||||
],
|
||||
num_retries=0,
|
||||
**router_kwargs,
|
||||
)
|
||||
|
||||
|
||||
def _mid_stream_failure(model: str):
|
||||
import litellm
|
||||
from litellm.exceptions import MidStreamFallbackError
|
||||
|
||||
return MidStreamFallbackError(
|
||||
message="stream dropped",
|
||||
model=model,
|
||||
llm_provider="openai",
|
||||
original_exception=litellm.InternalServerError(message="stream dropped", llm_provider="openai", model=model),
|
||||
is_pre_first_chunk=True,
|
||||
)
|
||||
|
||||
|
||||
def _scripted_responses_stream(events: list, error: Exception | None = None):
|
||||
from litellm.responses.streaming_iterator import BaseResponsesAPIStreamingIterator
|
||||
|
||||
class _ScriptedStream(BaseResponsesAPIStreamingIterator):
|
||||
def __init__(self) -> None:
|
||||
self._events = list(events)
|
||||
self._hidden_params: dict = {}
|
||||
self.completed_response = None
|
||||
|
||||
def __aiter__(self):
|
||||
return self
|
||||
|
||||
async def __anext__(self):
|
||||
if self._events:
|
||||
return self._events.pop(0)
|
||||
if error is not None:
|
||||
raise error
|
||||
raise StopAsyncIteration
|
||||
|
||||
async def aclose(self) -> None:
|
||||
return None
|
||||
|
||||
return _ScriptedStream()
|
||||
|
||||
|
||||
def _three_tier_original(calls: list, primary_fails_pre_stream: bool):
|
||||
import litellm
|
||||
|
||||
completed_event = _make_completed_event(1, 1, 2)
|
||||
|
||||
async def fake_original(**kwargs):
|
||||
model = kwargs["model"]
|
||||
calls.append(model)
|
||||
if model == "openai/primary-model":
|
||||
if primary_fails_pre_stream:
|
||||
raise litellm.InternalServerError(message="primary down", llm_provider="openai", model=model)
|
||||
return _scripted_responses_stream([], _mid_stream_failure(model))
|
||||
if model == "openai/fb1-model":
|
||||
return _scripted_responses_stream([], _mid_stream_failure(model))
|
||||
return _scripted_responses_stream([completed_event])
|
||||
|
||||
return fake_original, completed_event
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_aresponses_pre_stream_primary_failure_then_hop_stream_failure_reaches_second_entry():
|
||||
"""Regression: fallbacks=[{"primary": ["fb1", "fb2"]}]. The primary fails before streaming,
|
||||
fb1 is reached through the regular fallback chain and then fails mid-stream. Only the
|
||||
primary's stream used to be wrapped, so fb1's mid-stream failure either re-raised or
|
||||
re-tried fb1 itself; fb2 was unreachable."""
|
||||
router = _make_three_tier_router(fallbacks=[{"primary": ["fb1", "fb2"]}])
|
||||
calls: list = []
|
||||
fake_original, completed_event = _three_tier_original(calls, primary_fails_pre_stream=True)
|
||||
|
||||
stream = await router._aresponses_with_streaming_fallbacks(
|
||||
original_function=fake_original, model="primary", stream=True, input="hi"
|
||||
)
|
||||
collected = [event async for event in stream]
|
||||
|
||||
assert calls == ["openai/primary-model", "openai/fb1-model", "openai/fb2-model"]
|
||||
assert collected == [completed_event]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_aresponses_two_consecutive_mid_stream_failures_reach_second_entry():
|
||||
"""Regression: the primary and fb1 both fail mid-stream; fb2 must still be tried."""
|
||||
router = _make_three_tier_router(fallbacks=[{"primary": ["fb1", "fb2"]}])
|
||||
calls: list = []
|
||||
fake_original, completed_event = _three_tier_original(calls, primary_fails_pre_stream=False)
|
||||
|
||||
stream = await router._aresponses_with_streaming_fallbacks(
|
||||
original_function=fake_original, model="primary", stream=True, input="hi"
|
||||
)
|
||||
collected = [event async for event in stream]
|
||||
|
||||
assert calls == ["openai/primary-model", "openai/fb1-model", "openai/fb2-model"]
|
||||
assert collected == [completed_event]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_aresponses_per_request_fallbacks_survive_into_hop_streams():
|
||||
"""Regression: a request-level fallbacks list (key or team router_settings) is popped
|
||||
before each attempt runs, so a hop's mid-stream re-entry used to see only the router's
|
||||
own (empty) list and gave up after fb1."""
|
||||
router = _make_three_tier_router()
|
||||
calls: list = []
|
||||
fake_original, completed_event = _three_tier_original(calls, primary_fails_pre_stream=False)
|
||||
|
||||
stream = await router._aresponses_with_streaming_fallbacks(
|
||||
original_function=fake_original,
|
||||
model="primary",
|
||||
stream=True,
|
||||
input="hi",
|
||||
fallbacks=[{"primary": ["fb1", "fb2"]}],
|
||||
)
|
||||
collected = [event async for event in stream]
|
||||
|
||||
assert calls == ["openai/primary-model", "openai/fb1-model", "openai/fb2-model"]
|
||||
assert collected == [completed_event]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_aresponses_fallback_on_in_stream_error_event():
|
||||
"""A retriable in-stream error event (429) must trigger the router's mid-stream
|
||||
|
|
|
|||
|
|
@ -4234,7 +4234,7 @@ async def test_aresponses_streaming_iterator_fallback():
|
|||
call_kwargs = mock_fallback_utils.call_args.kwargs
|
||||
fbk = call_kwargs["kwargs"]
|
||||
# Bound methods compare equal when they share the same instance + __func__.
|
||||
assert fbk["original_function"] == router._ageneric_api_call_with_fallbacks_helper
|
||||
assert fbk["original_function"] == router._ageneric_api_call_with_fallbacks_responses_attempt
|
||||
assert fbk["original_generic_function"] is litellm.aresponses
|
||||
assert call_kwargs["model_group"] == "anthropic/claude-sonnet-4-6"
|
||||
assert call_kwargs["disable_fallbacks"] is False
|
||||
|
|
@ -13819,7 +13819,7 @@ async def test_aanthropic_messages_with_streaming_fallbacks_non_streaming_passth
|
|||
|
||||
with patch.object(
|
||||
router,
|
||||
"_ageneric_api_call_with_fallbacks",
|
||||
"_ageneric_api_call_with_fallbacks_helper",
|
||||
new=AsyncMock(return_value=plain_response),
|
||||
):
|
||||
out = await router._aanthropic_messages_with_streaming_fallbacks(
|
||||
|
|
@ -13843,7 +13843,7 @@ async def test_aanthropic_messages_with_streaming_fallbacks_wraps_streaming_iter
|
|||
with (
|
||||
patch.object(
|
||||
router,
|
||||
"_ageneric_api_call_with_fallbacks",
|
||||
"_ageneric_api_call_with_fallbacks_helper",
|
||||
new=AsyncMock(return_value=streaming_iter),
|
||||
),
|
||||
patch.object(
|
||||
|
|
@ -14128,7 +14128,7 @@ async def test_aanthropic_messages_with_streaming_fallbacks_deep_copies_nested_m
|
|||
):
|
||||
with patch.object(
|
||||
router,
|
||||
"_ageneric_api_call_with_fallbacks",
|
||||
"_ageneric_api_call_with_fallbacks_helper",
|
||||
new=AsyncMock(side_effect=fake_original),
|
||||
):
|
||||
await router._aanthropic_messages_with_streaming_fallbacks(
|
||||
|
|
@ -14162,7 +14162,7 @@ async def test_aanthropic_messages_with_streaming_fallbacks_deep_copies_metadata
|
|||
):
|
||||
with patch.object(
|
||||
router,
|
||||
"_ageneric_api_call_with_fallbacks",
|
||||
"_ageneric_api_call_with_fallbacks_helper",
|
||||
new=AsyncMock(side_effect=fake_original),
|
||||
):
|
||||
await router._aanthropic_messages_with_streaming_fallbacks(
|
||||
|
|
@ -14177,6 +14177,51 @@ async def test_aanthropic_messages_with_streaming_fallbacks_deep_copies_metadata
|
|||
assert "deployment" not in fallback_kwargs["metadata"]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_anthropic_messages_hop_stream_failure_reaches_second_fallback_entry():
|
||||
"""Regression: fallbacks=[{"primary": ["fb1", "fb2"]}]. The primary fails before
|
||||
streaming, fb1 is reached through the regular fallback chain and then sends an
|
||||
error frame mid-stream. Only the primary's stream used to be wrapped, so the outer
|
||||
wrapper re-tried fb1 with a fresh attempted set and forwarded fb1's error frame to
|
||||
the client on an HTTP 200; fb2 was unreachable."""
|
||||
router = Router(
|
||||
model_list=[
|
||||
{"model_name": "primary", "litellm_params": {"model": "anthropic/primary-model", "api_key": "sk-test"}},
|
||||
{"model_name": "fb1", "litellm_params": {"model": "anthropic/fb1-model", "api_key": "sk-test"}},
|
||||
{"model_name": "fb2", "litellm_params": {"model": "anthropic/fb2-model", "api_key": "sk-test"}},
|
||||
],
|
||||
num_retries=0,
|
||||
fallbacks=[{"primary": ["fb1", "fb2"]}],
|
||||
)
|
||||
calls: list = []
|
||||
|
||||
async def fake_original(**kwargs):
|
||||
model = kwargs["model"]
|
||||
calls.append(model)
|
||||
if model == "anthropic/primary-model":
|
||||
raise litellm.InternalServerError(message="primary down", llm_provider="anthropic", model=model)
|
||||
if model == "anthropic/fb1-model":
|
||||
return _AnthropicMessagesFakeByteStream(
|
||||
[_anthropic_messages_message_start_chunk(), _anthropic_messages_overloaded_error_chunk()]
|
||||
)
|
||||
return _AnthropicMessagesFakeByteStream(
|
||||
[_anthropic_messages_message_start_chunk(), _anthropic_messages_content_chunk("from fb2")]
|
||||
)
|
||||
|
||||
stream = await router._aanthropic_messages_with_streaming_fallbacks(
|
||||
original_function=fake_original,
|
||||
model="primary",
|
||||
stream=True,
|
||||
messages=[{"role": "user", "content": "hi"}],
|
||||
max_tokens=10,
|
||||
)
|
||||
body = b"".join([chunk async for chunk in stream])
|
||||
|
||||
assert calls == ["anthropic/primary-model", "anthropic/fb1-model", "anthropic/fb2-model"]
|
||||
assert b"from fb2" in body
|
||||
assert b"overloaded_error" not in body
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_anthropic_messages_fallback_triggers_after_lifecycle_only_frame():
|
||||
"""Regression: Anthropic routinely sends a message_start lifecycle frame
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue