From 5c464c45e99ac710b1ef72264fa7f2a171dc2875 Mon Sep 17 00:00:00 2001 From: Claude Date: Wed, 7 Oct 2026 03:23:33 +0000 Subject: [PATCH] feat(router): stream sync chat completions through the Rust backend A stream failing before content now settles the call at the hop it failed in, as Python's stream wrappers do: the outer retry and fallback layers no longer retry it, try their next target or rewrite its outcome, and the caller gets the provider error rather than the MidStreamFallbackError. The sync driver buffers and replays streams too and follows the async rule, since Python's sync wrapper reruns the same group and can recurse until RecursionError. The sync path no longer updates the call counters, matching Python Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01Aw1xktpUnix5tWhvSzey8t --- .../crates/router/POC-behavior-spec.md | 3 + litellm-rust/crates/router/POC.md | 4 +- litellm-rust/crates/router/src/engine.rs | 90 +++++++---- litellm-rust/crates/router/tests/engine.rs | 54 +++++++ litellm/router_backends/rust_call.py | 64 ++++++-- litellm/router_backends/rust_router.py | 2 - .../router/test_rust_router.py | 143 +++++++++++++++--- 7 files changed, 296 insertions(+), 64 deletions(-) diff --git a/litellm-rust/crates/router/POC-behavior-spec.md b/litellm-rust/crates/router/POC-behavior-spec.md index c39fbe5f868..083508f5826 100644 --- a/litellm-rust/crates/router/POC-behavior-spec.md +++ b/litellm-rust/crates/router/POC-behavior-spec.md @@ -532,3 +532,6 @@ Response (`set_response_headers`, PR:11756-11784, HDR): 23. `_update_usage` writes a per-id counter that nothing on the hot path reads (PR:8746-8763, 15102-15113) 24. `_get_router_metadata_variable_name` substring-matches `"file"`, so any handler name containing it switches buckets (BU:161-171) 25. Exhausted-retry stamping uses `type(e) in LITELLM_EXCEPTION_TYPES`, so subclasses (e.g. `MidStreamFallbackError`) never get `max_retries`/`num_retries` (PR:8173) +26. Sync chat streams: a stream failing before content reruns `function_with_fallbacks` on the same group instead of walking the fallback chain, so fallbacks are only reached when the rerun's call itself raises, and a group that keeps failing recurses until `RecursionError` (PR:3624-3735). NOT replicated: the Rust backend applies the async rule (section 9) to sync streams too +27. `_combine_fallback_usage` only adds the failed stream's partial usage when `hasattr(usage, "usage")`, which a `Usage` never has, so partial usage is never added to the fallback's (PR:2942-2947). Matching it means the Rust backend adds nothing +28. Python's sync `_completion` updates none of `total_calls`/`success_calls`/`fail_calls`, while `_acompletion` updates all three (PR:2652-2760 vs 3780-3940) diff --git a/litellm-rust/crates/router/POC.md b/litellm-rust/crates/router/POC.md index a960b6cd3b2..ca9fabe11a2 100644 --- a/litellm-rust/crates/router/POC.md +++ b/litellm-rust/crates/router/POC.md @@ -175,8 +175,8 @@ Design notes made while implementing: - Step 5 (in progress): `litellm/router_backends/selection.py::pinned_backend` fixes the backend decision; `tests/unit/router_backends/backends.py::router_backend` runs an opted-in module on both backends (Rust pinned to `RUST_REQUIRED`), with known Rust-side failures as strict xfails in `RUST_GAPS`. Opted in so far: `tests/unit/test_router_per_deployment_num_retries.py`, `tests/unit/test_router_exception_redaction.py`. `tests/test_litellm_rust/router/test_redis_parity.py` runs both backends against a live `redis-server` and compares every key, raw value (timestamp masked) and TTL: cooldowns, `allowed_fails` counters and failure RPM keys match byte for byte, and a Rust router honors cooldowns a Python router wrote -Streaming (async chat): `RoutedCall` reads the stream until the first chunk with content (Python's `_stream_chunks_have_generated_content`) and returns a `ReplayStream` (a `FallbackAwareStreamWrapper`) that replays the buffered chunks. A `MidStreamFallbackError` before content is classified `stream_failure: before_content` and goes straight to the fallback chain without in-group retries, matching Python's chat wrapper (spec section 9) rather than retrying in the group; any other pre-content stream error is `terminal` and raised as is; after content a mid-stream fallback error surfaces as its original exception. Not done: the sync stream path, and combining the failed attempt's partial usage into the fallback's +Streaming (chat, both drivers): `RoutedCall` reads the stream until the first chunk with content (Python's `_stream_chunks_have_generated_content`) and returns a `ReplayStream` (a `FallbackAwareStreamWrapper`, sync and async) that replays the buffered chunks. A `MidStreamFallbackError` before content is classified `stream_failure: before_content` and goes to that hop's own fallback chain without in-group retries; any other pre-content stream error is `terminal`. Either way the result settles the call: Python's outer retry and fallback layers had already returned the stream, so the engine's `Stop::Settled` and `Success::settled` keep them from retrying, trying the next target or rewriting the outcome. A stream error reaching the caller is unwrapped to its `original_exception`, as Python's wrappers do. The sync driver uses the async rule on purpose, because Python's sync wrapper reruns the same group and can recurse until `RecursionError` (spec section 11 item 26). Partial usage is not added to the fallback's, matching Python, whose check never passes (item 27). Like Python, the sync path no longer touches `total_calls`/`success_calls`/`fail_calls` (item 28) -Next, following the revised POC scope: finish increment 1 (sync streaming, partial usage on fallback, `aresponses`, `aanthropic_messages`, and more suites opted into `router_backend`, starting with `tests/unit/test_router_order_fallback.py` and the fallback/cooldown tests in `tests/unit/router_utils` that go through `Router(...)`), then increments 2, 3 and 4. Then sketch the options for increment 3 +Next, following the revised POC scope: finish increment 1 (`aresponses`, `aanthropic_messages`, and more suites opted into `router_backend`, starting with `tests/unit/test_router_order_fallback.py` and the fallback/cooldown tests in `tests/unit/router_utils` that go through `Router(...)`), then increments 2, 3 and 4. Then sketch the options for increment 3 Known gaps to close or decline: CustomLogger fallback-event hooks and `router_cooldown_event_callback` (increment 3), the 200-key `InMemoryCache` eviction, fallback keys containing `/` (provider-prefixed matching needs the cost map) diff --git a/litellm-rust/crates/router/src/engine.rs b/litellm-rust/crates/router/src/engine.rs index bd91ebdbd4a..90e0ac4a5fd 100644 --- a/litellm-rust/crates/router/src/engine.rs +++ b/litellm-rust/crates/router/src/engine.rs @@ -157,10 +157,12 @@ impl Engine { outcome: success.outcome, ops, }), - Err(Stop::Failed(failure)) => Err(RouteError::Failed(Box::new(Failed { - error: failure.raised, - ops, - }))), + Err(Stop::Failed(failure) | Stop::Settled(failure)) => { + Err(RouteError::Failed(Box::new(Failed { + error: failure.raised, + ops, + }))) + } Err(Stop::Host(fault)) => Err(RouteError::Host(fault)), } } @@ -169,13 +171,38 @@ impl Engine { struct Success { response: R, outcome: Outcome, + /// Served by the fallback a failed stream started. Python's outer layers had already + /// returned that stream, so they leave this outcome as the fallback reported it. + settled: bool, } enum Stop { Failed(Box>), + /// The failure of a stream's own fallback. Python's outer layers had already returned the + /// stream, so they neither retry nor fall back from it. + Settled(Box>), Host(F), } +/// The failure the retry layer handles; anything else passes through it. +fn failed(stop: Stop) -> Result>, Stop> { + match stop { + Stop::Failed(failure) => Ok(failure), + other => Err(other), + } +} + +fn settle(result: Attempted) -> Attempted { + match result { + Ok(success) => Ok(Success { + settled: true, + ..success + }), + Err(Stop::Failed(failure)) => Err(Stop::Settled(failure)), + Err(stop) => Err(stop), + } +} + type Attempted = Result, Stop>; type Step<'a, R, E, F> = Pin, Stop>> + Send + 'a>>; @@ -242,9 +269,8 @@ impl<'a, H: RouterHost> Run<'a, H> { None => self.retries(&hop).await, }; match first { - Ok(success) => Ok(success), - Err(Stop::Host(fault)) => Err(Stop::Host(fault)), Err(Stop::Failed(failure)) => self.fallback(&hop, failure).await, + other => other, } }) } @@ -283,6 +309,7 @@ impl<'a, H: RouterHost> Run<'a, H> { max_retries: None, attempted_fallbacks: 0, }, + settled: false, }), } } @@ -306,10 +333,12 @@ impl<'a, H: RouterHost> Run<'a, H> { }; let failure = match first { Ok(success) => return Ok(success), - Err(stop) => self.failed(stop).map_err(Stop::Host)?, + Err(stop) => failed(stop)?, }; - if failure.classified.guardrail_intervention || failure.classified.stream_failure.is_some() - { + if let Some(kind) = failure.classified.stream_failure { + return self.stream_failed(hop, failure, kind).await; + } + if failure.classified.guardrail_intervention { return Err(Stop::Failed(failure)); } if self.call.num_retries.is_none() @@ -345,14 +374,16 @@ impl<'a, H: RouterHost> Run<'a, H> { stamp.max_retries = num_retries; let failure = match self.attempt(hop, stamp, &skipped).await { Ok(mut success) => { - success.outcome.attempted_retries = attempt + 1; - success.outcome.max_retries = Some(num_retries); + if !success.settled { + success.outcome.attempted_retries = attempt + 1; + success.outcome.max_retries = Some(num_retries); + } return Ok(success); } - Err(stop) => self.failed(stop).map_err(Stop::Host)?, + Err(stop) => failed(stop)?, }; - if failure.classified.stream_failure.is_some() { - return Err(Stop::Failed(failure)); + if let Some(kind) = failure.classified.stream_failure { + return self.stream_failed(hop, failure, kind).await; } self.log_retry(hop, &failure); let remaining = num_retries - attempt - 1; @@ -378,10 +409,18 @@ impl<'a, H: RouterHost> Run<'a, H> { Err(Stop::Failed(latest)) } - fn failed(&self, stop: Stop) -> Result>, H::Fault> { - match stop { - Stop::Failed(failure) => Ok(failure), - Stop::Host(fault) => Err(fault), + /// `_acompletion_streaming_iterator`: the attempt returned its stream, so the retry and + /// fallback layers above it are done. A failure before content goes to this hop's own + /// fallback chain without retrying in the group; any other stream error is raised as is. + async fn stream_failed( + &mut self, + hop: &Hop, + failure: Box>, + kind: StreamFailure, + ) -> Attempted { + match kind { + StreamFailure::Terminal => Err(Stop::Settled(failure)), + StreamFailure::BeforeContent => settle(self.fallback(hop, failure).await), } } @@ -500,6 +539,7 @@ impl<'a, H: RouterHost> Run<'a, H> { max_retries: None, attempted_fallbacks: 0, }, + settled: false, }) } Invoked::Failure { error, classified } => { @@ -534,10 +574,7 @@ impl<'a, H: RouterHost> Run<'a, H> { hop: &Hop, failure: Box>, ) -> Attempted { - if self.call.disable_fallbacks - || failure.classified.guardrail_intervention - || failure.classified.stream_failure == Some(StreamFailure::Terminal) - { + if self.call.disable_fallbacks || failure.classified.guardrail_intervention { return Err(Stop::Failed(failure)); } let top_level = hop.depth == 0; @@ -610,12 +647,11 @@ impl<'a, H: RouterHost> Run<'a, H> { return Err(Stop::Failed(failure)); } Ok(Some(targets)) => match self.run_fallbacks(hop, &targets, &failure).await { - Ok(success) => return Ok(success), - Err(Stop::Host(fault)) => return Err(Stop::Host(fault)), Err(Stop::Failed(last)) => FallbackResult::Failed { attempted: Some(targets), last: last.raised, }, + other => return other, }, }; if top_level { @@ -678,16 +714,18 @@ impl<'a, H: RouterHost> Run<'a, H> { .await { Ok(mut success) => { - success.outcome.attempted_fallbacks = depth; + if !success.settled { + success.outcome.attempted_fallbacks = depth; + } return Ok(success); } - Err(Stop::Host(fault)) => return Err(Stop::Host(fault)), Err(Stop::Failed(failure)) => { if self.call.shared_logging && self.callbacks_logged { self.cooldown_failed_hop(&failure).await; } last = *failure; } + Err(stop) => return Err(stop), } } Err(Stop::Failed(Box::new(last))) diff --git a/litellm-rust/crates/router/tests/engine.rs b/litellm-rust/crates/router/tests/engine.rs index b3376aced5d..44a62f200cf 100644 --- a/litellm-rust/crates/router/tests/engine.rs +++ b/litellm-rust/crates/router/tests/engine.rs @@ -472,3 +472,57 @@ async fn a_terminal_stream_failure_is_raised_without_retries_or_fallbacks() { assert_eq!(raised(&failed.error), "a#1"); assert!(failed.ops.is_empty()); } + +#[rstest] +#[tokio::test] +async fn a_terminal_stream_failure_in_a_fallback_hop_ends_the_call() { + let mut settings = settings(); + settings.fallbacks = Some(vec![chain("g", &["h", "i"])]); + let engine = engine( + vec![ + deployment("a", "g"), + deployment("c", "h"), + deployment("d", "i"), + ], + settings, + ); + let host = ScriptedHost::default() + .script("a", &[BAD_REQUEST]) + .script("c", &[Scripted::Stream(StreamFailure::Terminal)]); + + let Err(RouteError::Failed(failed)) = engine.route(&host, call("g")).await else { + panic!("the fallback's stream error is raised"); + }; + + assert_eq!(host.attempts(), ["a", "c"]); + assert_eq!(raised(&failed.error), "c#2"); +} + +#[rstest] +#[tokio::test] +async fn a_fallback_hop_stream_falls_back_on_its_own_chain_and_reports_its_depth() { + let mut settings = settings(); + settings.fallbacks = Some(vec![chain("g", &["h"]), chain("h", &["i"])]); + let engine = engine( + vec![ + deployment("a", "g"), + deployment("c", "h"), + deployment("d", "i"), + ], + settings, + ); + let host = ScriptedHost::default() + .script("a", &[BAD_REQUEST]) + .script("c", &[Scripted::Stream(StreamFailure::BeforeContent)]); + + let routed = engine.route(&host, call("g")).await.ok().unwrap(); + + assert_eq!(host.attempts(), ["a", "c", "d"]); + assert_eq!( + ( + routed.outcome.model_group.as_str(), + routed.outcome.attempted_fallbacks + ), + ("i", 2) + ); +} diff --git a/litellm/router_backends/rust_call.py b/litellm/router_backends/rust_call.py index 0c09de05c95..c34b3cc2490 100644 --- a/litellm/router_backends/rust_call.py +++ b/litellm/router_backends/rust_call.py @@ -11,7 +11,7 @@ from __future__ import annotations import asyncio import time -from collections.abc import Callable, Mapping, MutableMapping, Sequence +from collections.abc import Callable, Iterator, Mapping, MutableMapping, Sequence from typing import Final, Literal, Protocol, TypeAlias, cast from pydantic import BaseModel, ConfigDict, TypeAdapter @@ -245,16 +245,35 @@ class ReplayStream(FallbackAwareStreamWrapper): try: return await self._inner.__anext__() except MidStreamFallbackError as error: - if error.original_exception is not None: - raise error.original_exception from error - raise + raise _unwrapped(error) from error + + def __iter__(self) -> Iterator[ModelResponseStream]: + return self + + def __next__(self) -> ModelResponseStream: + buffered: Final = next(self._replay, None) + if buffered is not None: + return buffered + try: + return self._inner.__next__() + except MidStreamFallbackError as error: + raise _unwrapped(error) from error async def aclose(self) -> None: await self._inner.aclose() +def _unwrapped(error: BaseException) -> BaseException: + """The provider error behind a mid-stream fallback error, which Python's stream wrappers raise in its place.""" + if isinstance(error, MidStreamFallbackError) and error.original_exception is not None: + return error.original_exception + return error + + def _not_fetched(stream: CustomStreamWrapper) -> bool: - return cast(object, stream.completion_stream) is None and stream.make_call is not None # cast-ok: untyped attribute + return ( + cast(object, stream.completion_stream) is None and cast(object, stream.make_call) is not None + ) # cast-ok: untyped attributes async def _replay_from_first_content(stream: CustomStreamWrapper) -> ReplayStream: @@ -273,6 +292,20 @@ async def _replay_from_first_content(stream: CustomStreamWrapper) -> ReplayStrea return ReplayStream(stream, buffered) +def _replay_from_first_content_sync(stream: CustomStreamWrapper) -> ReplayStream: + buffered: Final[list[ModelResponseStream]] = [] # mutable-ok: chunks read before the first content + try: + for chunk in stream: + buffered.append(chunk) + if _stream_chunks_have_generated_content([chunk]): + break + except MidStreamFallbackError as error: + raise _StreamFailed(error, "before_content") from error + except Exception as error: + raise _StreamFailed(error, "terminal") from error + return ReplayStream(stream, buffered) + + class _AttemptFailed(Exception): """An attempt's error that litellm did not raise, so its failure callbacks never ran for it.""" @@ -336,6 +369,12 @@ class RoutedCall: return self._mock(parsed, kwargs) try: response: Final = self._completion(parsed, kwargs) + except _StreamFailed as failed: + return ( + "error", + failed.error, + {**classify(failed.error, kwargs, callbacks_ran=True), "stream_failure": failed.kind}, + ) except _AttemptFailed as failed: return ("error", failed.error, classify(failed.error, kwargs, callbacks_ran=False)) except Exception as error: @@ -358,7 +397,7 @@ class RoutedCall: def failure(self, error: object, ops: Sequence[Mapping[str, object]]) -> BaseException: self._apply(ops) - return self._exception(error) + return _unwrapped(self._exception(error)) def _prepare(self, attempt: Attempt) -> Kwargs: self._apply(attempt.ops) @@ -486,7 +525,6 @@ class RoutedCall: normalizer: Final = self._normalizer deployment: Final = self._deployment(attempt) messages: Final = self._kwargs.get("messages") - model_name: str | None = None # rebind-ok: set once the deployment's params are read try: normalizer._drop_unsupported_classifier_reasoning_effort( # pyright: ignore[reportPrivateUsage] # the attempt body PythonRouter runs deployment=deployment, # pyright: ignore[reportArgumentType] # router deployment dict @@ -498,7 +536,6 @@ class RoutedCall: ) kwargs.setdefault("messages", messages) normalizer._update_kwargs_with_deployment(deployment=deployment, kwargs=kwargs) # pyright: ignore[reportPrivateUsage] # as above - model_name = str(litellm_params["model"]) # rebind-ok: read by the failure counter below potential_client: Final = normalizer._get_client(deployment=deployment, kwargs=kwargs) # pyright: ignore[reportPrivateUsage] # as above dynamic_api_key: Final = kwargs.get("api_key") client: Final = ( @@ -526,13 +563,18 @@ class RoutedCall: message="Response output was blocked.", model=attempt.model_group, llm_provider="" ) ) - normalizer.success_calls[model_name] += 1 + if isinstance(response, CustomStreamWrapper) and _not_fetched(response): + response.fetch_sync_stream() + if isinstance(response, CustomStreamWrapper): + return _replay_from_first_content_sync(response) return response + except _StreamFailed: + raise except _AttemptFailed as failed: - self._stamp(failed.error, deployment, kwargs, model_name) + self._stamp(failed.error, deployment, kwargs, None) raise except Exception as error: - self._stamp(error, deployment, kwargs, model_name) + self._stamp(error, deployment, kwargs, None) raise def _stamp( diff --git a/litellm/router_backends/rust_router.py b/litellm/router_backends/rust_router.py index 86dd42e3d29..5ac98c24e67 100644 --- a/litellm/router_backends/rust_router.py +++ b/litellm/router_backends/rust_router.py @@ -271,8 +271,6 @@ class RustRouter: unsupported: Final = [name for name in _UNSUPPORTED_REQUEST_KWARGS if name in kwargs] if unsupported: raise NotImplementedError(f"the Rust router backend does not support request option {unsupported[0]!r} yet") - if kwargs.get("stream") is True and operation == "completion": - raise NotImplementedError("the Rust router backend does not support sync streaming yet") clientside: Final = _IS_CLIENTSIDE_CREDENTIAL(kwargs) if clientside: raise NotImplementedError("the Rust router backend does not support client-side credentials yet") diff --git a/tests/test_litellm_rust/router/test_rust_router.py b/tests/test_litellm_rust/router/test_rust_router.py index 663f93d9b9f..71931d16a1e 100644 --- a/tests/test_litellm_rust/router/test_rust_router.py +++ b/tests/test_litellm_rust/router/test_rust_router.py @@ -172,6 +172,22 @@ def test_sync_completion_retries_and_answers(build: Callable[[Mapping[str, objec assert _content(response) == "from a" +def test_sync_calls_leave_the_call_counters_as_python_does() -> None: + arguments: Final = { + "model_list": (_deployment("g", "a"), _deployment("h", "b", Exception("broke"))), + "num_retries": 0, + } + counters: Final[list[tuple[object, ...]]] = [] # mutable-ok: one reading per backend + for build in BACKENDS: + backend = build(arguments, 1) + backend.completion("g", MESSAGES) + with pytest.raises(litellm.InternalServerError): + backend.completion("h", MESSAGES) + counters.append(tuple(dict(getattr(backend, name)) for name in ("total_calls", "success_calls", "fail_calls"))) + + assert counters[1] == counters[0] + + async def test_the_facade_serves_a_supported_config_from_the_rust_backend() -> None: backend: Final = select_backend( (), {"model_list": [_deployment("g", "a")]}, rules=(RouteRule(Route.ROUTER, Rollout.RUST_REQUIRED),) @@ -190,10 +206,7 @@ class _FailingStream(CustomStreamWrapper): def __init__(self, model: str) -> None: super().__init__(completion_stream=object(), model=model, custom_llm_provider="openai", logging_obj=MagicMock()) - def __aiter__(self) -> _FailingStream: - return self - - async def __anext__(self) -> ModelResponseStream: + def _fail(self) -> ModelResponseStream: raise MidStreamFallbackError( message=f"provider 500 from {self.model}", model=str(self.model), @@ -204,6 +217,18 @@ class _FailingStream(CustomStreamWrapper): ), ) + def __aiter__(self) -> _FailingStream: + return self + + async def __anext__(self) -> ModelResponseStream: + return self._fail() + + def __iter__(self) -> _FailingStream: + return self + + def __next__(self) -> ModelResponseStream: + return self._fail() + class _OkStream(_FailingStream): def __init__(self, model: str) -> None: @@ -221,30 +246,102 @@ class _OkStream(_FailingStream): raise StopAsyncIteration return reply - -async def _fake_stream(**kwargs: object) -> CustomStreamWrapper: # kwargs-ok: litellm.acompletion's surface - model: Final = str(kwargs["model"]) - return _OkStream(model) if "fb2" in model else _FailingStream(model) + def __next__(self) -> ModelResponseStream: + reply: Final = next(self._replies, None) + if reply is None: + raise StopIteration + return reply -async def test_a_stream_failing_before_content_walks_the_fallback_chain_on_both_backends() -> None: - arguments: Final = { - "model_list": ( - {"model_name": "primary", "litellm_params": {"model": "openai/primary-model", "api_key": "k"}}, - {"model_name": "fb1", "litellm_params": {"model": "openai/fb1-model", "api_key": "k"}}, - {"model_name": "fb2", "litellm_params": {"model": "openai/fb2-model", "api_key": "k"}}, +def _streams(healthy: str) -> Callable[..., CustomStreamWrapper]: + def stream(**kwargs: object) -> CustomStreamWrapper: # kwargs-ok: litellm.completion's surface + model: Final = str(kwargs["model"]) + return _OkStream(model) if healthy in model else _FailingStream(model) + + return stream + + +def _async_streams(healthy: str) -> Callable[..., object]: + sync: Final = _streams(healthy) + + async def stream(**kwargs: object) -> CustomStreamWrapper: # kwargs-ok: litellm.acompletion's surface + return sync(**kwargs) + + return stream + + +def _stream_groups( + fallbacks: list[dict[str, list[str]]], +) -> Mapping[str, object]: # mutable-ok: Router's fallbacks shape + return { + "model_list": tuple( + {"model_name": group, "litellm_params": {"model": f"openai/{group}-model", "api_key": "k"}} + for group in ("primary", "fb1", "fb2") ), - "fallbacks": [{"primary": ["fb1", "fb2"]}], + "fallbacks": fallbacks, "num_retries": 2, } - observed: Final[list[tuple[str, tuple[str, ...]]]] = [] # mutable-ok: one observation per backend - for build in BACKENDS: - backend = build(arguments, 1) - with patch("litellm.acompletion", side_effect=_fake_stream) as acompletion: + + +StreamOutcome = tuple[str, tuple[str, ...], Mapping[str, object]] + + +async def _astream(backend: Backend, healthy: str) -> StreamOutcome: + with patch("litellm.acompletion", side_effect=_async_streams(healthy)) as acompletion: + try: response = await backend.acompletion("primary", MESSAGES, stream=True) content = "".join([chunk.choices[0].delta.content or "" async for chunk in response]) # pyright: ignore[reportAttributeAccessIssue,reportUnknownMemberType,reportUnknownVariableType] # stream chunks - groups = tuple(str(call.kwargs["metadata"]["model_group"]) for call in acompletion.call_args_list) - observed.append((content, groups)) + except litellm.InternalServerError as error: + content = f"raised {type(error).__name__}: {error}" + response = None + groups: Final = tuple(str(call.kwargs["metadata"]["model_group"]) for call in acompletion.call_args_list) + return content, groups, _headers(response) - assert observed[0] == observed[1] - assert observed[1] == ("ok-from-openai/fb2-model", ("primary", "fb1", "fb2")) + +@pytest.mark.parametrize( + ("fallbacks", "healthy", "expected"), + ( + pytest.param( + [{"primary": ["fb1", "fb2"]}], + "fb2", + ("ok-from-openai/fb2-model", ("primary", "fb1", "fb2")), + id="chain-of-one-group", + ), + pytest.param( + [{"primary": ["fb1"]}, {"fb1": ["fb2"]}], + "fb2", + ("ok-from-openai/fb2-model", ("primary", "fb1", "fb2")), + id="fallback-hop-uses-its-own-chain", + ), + pytest.param( + [{"primary": ["fb1", "fb2"]}], + "nowhere", + ( + "raised InternalServerError: litellm.InternalServerError: provider 500 from openai/fb2-model", + ("primary", "fb1", "fb2"), + ), + id="every-group-fails", + ), + ), +) +async def test_a_stream_failing_before_content_falls_back_the_same_way_on_both_backends( + fallbacks: list[dict[str, list[str]]], # mutable-ok: Router's fallbacks shape + healthy: str, + expected: tuple[str, tuple[str, ...]], +) -> None: + python: Final = await _astream(_python(_stream_groups(fallbacks), 1), healthy) + rust: Final = await _astream(_rust(_stream_groups(fallbacks), 1), healthy) + + assert rust == python + assert rust[:2] == expected + + +def test_a_sync_stream_failing_before_content_falls_back_like_the_async_path() -> None: + """Python's sync wrapper reruns the same group instead of falling back, and recurses until + `RecursionError` while the group keeps failing; the Rust backend applies the async rule.""" + with patch("litellm.completion", side_effect=_streams("fb2")) as completion: + response = _rust(_stream_groups([{"primary": ["fb1", "fb2"]}]), 1).completion("primary", MESSAGES, stream=True) + content = "".join(chunk.choices[0].delta.content or "" for chunk in response) # pyright: ignore[reportAttributeAccessIssue,reportUnknownMemberType,reportUnknownVariableType,reportGeneralTypeIssues] # stream chunks + + assert content == "ok-from-openai/fb2-model" + assert [call.kwargs["metadata"]["model_group"] for call in completion.call_args_list] == ["primary", "fb1", "fb2"]