From 8841ee8d1395d3a4f76ae25e46ce284b4e9277c3 Mon Sep 17 00:00:00 2001 From: Claude Date: Wed, 7 Oct 2026 01:42:58 +0000 Subject: [PATCH] feat(router): stream chat completions through the Rust backend The call driver reads an attempt's stream until its first content chunk and hands back a ReplayStream that replays what it buffered. A MidStreamFallbackError before content goes straight to the fallback chain without retrying the group, as Python's chat stream wrapper does; any other error before content is raised as is, and after content a mid-stream fallback error surfaces as its original exception. A native test runs the same failing-before-first-chunk fallback chain through both backends and gets the same content from the same sequence of model groups. Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01Aw1xktpUnix5tWhvSzey8t --- litellm-rust/crates/router/POC.md | 6 +- litellm-rust/crates/router/src/engine.rs | 13 ++- litellm-rust/crates/router/src/failure.rs | 13 +++ litellm-rust/crates/router/tests/engine.rs | 53 ++++++++++- litellm/router_backends/rust_call.py | 87 ++++++++++++++++++- litellm/router_backends/rust_router.py | 4 +- .../router/test_rust_router.py | 72 ++++++++++++++- 7 files changed, 238 insertions(+), 10 deletions(-) diff --git a/litellm-rust/crates/router/POC.md b/litellm-rust/crates/router/POC.md index 3a270278266..2e04bfec42f 100644 --- a/litellm-rust/crates/router/POC.md +++ b/litellm-rust/crates/router/POC.md @@ -107,6 +107,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 -Next: opt more suites into `router_backend` (`tests/unit/test_router_order_fallback.py` and the fallback/cooldown tests in `tests/unit/router_utils` that go through `Router(...)`), then streaming (the driver buffers to the first content chunk and reports `StreamFailedBeforeContent`; decide whether chat keeps Python's fallback-only rule), then the generic operations (`aresponses`, `aanthropic_messages`, `aimage_generation`), then `upsert_deployment`/`delete_deployment`/`update_settings` through `Registry::replace` +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 -Known gaps to close or decline: chat streaming pre-first-chunk failures (Python goes straight to fallbacks, no in-group retry), `include_fallback_errors`, request-level `model_group_retry_policy`/`specific_deployment`/`priority`, CustomLogger fallback-event hooks, `router_cooldown_event_callback`, the 200-key `InMemoryCache` eviction, fallback keys containing `/` (provider-prefixed matching needs the cost map) +Next: opt more suites into `router_backend` (`tests/unit/test_router_order_fallback.py` and the fallback/cooldown tests in `tests/unit/router_utils` that go through `Router(...)`), then sync streaming, then the generic operations (`aresponses`, `aanthropic_messages`, `aimage_generation`), then `upsert_deployment`/`delete_deployment`/`update_settings` through `Registry::replace` + +Known gaps to close or decline: `include_fallback_errors`, request-level `model_group_retry_policy`/`specific_deployment`/`priority`, CustomLogger fallback-event hooks, `router_cooldown_event_callback`, 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 2a96c68e219..bd91ebdbd4a 100644 --- a/litellm-rust/crates/router/src/engine.rs +++ b/litellm-rust/crates/router/src/engine.rs @@ -10,7 +10,7 @@ use std::{ use crate::{ cooldown::{self, status_is_retryable}, - failure::{Classified, ExceptionClass, Failure, Raised}, + failure::{Classified, ExceptionClass, Failure, Raised, StreamFailure}, fallback::{chain_for_groups, generic_targets}, host::{ Attempt, HopStamp, Invoked, MockFailure, Op, RetryStamp, RouterHost, Target, TypedFallback, @@ -308,7 +308,8 @@ impl<'a, H: RouterHost> Run<'a, H> { Ok(success) => return Ok(success), Err(stop) => self.failed(stop).map_err(Stop::Host)?, }; - if failure.classified.guardrail_intervention { + if failure.classified.guardrail_intervention || failure.classified.stream_failure.is_some() + { return Err(Stop::Failed(failure)); } if self.call.num_retries.is_none() @@ -350,6 +351,9 @@ impl<'a, H: RouterHost> Run<'a, H> { } Err(stop) => self.failed(stop).map_err(Stop::Host)?, }; + if failure.classified.stream_failure.is_some() { + return Err(Stop::Failed(failure)); + } self.log_retry(hop, &failure); let remaining = num_retries - attempt - 1; healthy = self.healthy_counts(&hop.group).await.0; @@ -530,7 +534,10 @@ impl<'a, H: RouterHost> Run<'a, H> { hop: &Hop, failure: Box>, ) -> Attempted { - if self.call.disable_fallbacks || failure.classified.guardrail_intervention { + if self.call.disable_fallbacks + || failure.classified.guardrail_intervention + || failure.classified.stream_failure == Some(StreamFailure::Terminal) + { return Err(Stop::Failed(failure)); } let top_level = hop.depth == 0; diff --git a/litellm-rust/crates/router/src/failure.rs b/litellm-rust/crates/router/src/failure.rs index 5dfedd71393..3c345b0f58c 100644 --- a/litellm-rust/crates/router/src/failure.rs +++ b/litellm-rust/crates/router/src/failure.rs @@ -19,6 +19,17 @@ pub enum ExceptionClass { ApiConnection, } +/// How a streamed attempt failed before its first content chunk reached the caller. +#[derive(Clone, Copy, Debug, Deserialize, PartialEq, Eq)] +#[serde(rename_all = "snake_case")] +pub enum StreamFailure { + /// A failure the stream can fall back from (`MidStreamFallbackError`). Chat goes straight + /// to the fallback chain, without retrying in the group, as Python's stream wrapper does. + BeforeContent, + /// Any other stream error, which reaches the caller as raised. + Terminal, +} + /// What the host learned about an exception an attempt raised. The host computes the facts /// that depend on Python objects (header parsing, logging state); routing reads only these. #[derive(Clone, Debug, Default, Deserialize, PartialEq)] @@ -45,6 +56,8 @@ pub struct Classified { pub exempt_from_cooldown: bool, /// `type(e) in litellm.LITELLM_EXCEPTION_TYPES`, which gates the exhausted-retry stamp. pub exact_litellm_type: bool, + #[serde(default)] + pub stream_failure: Option, } impl Classified { diff --git a/litellm-rust/crates/router/tests/engine.rs b/litellm-rust/crates/router/tests/engine.rs index 60897747341..c3d5c1c38c4 100644 --- a/litellm-rust/crates/router/tests/engine.rs +++ b/litellm-rust/crates/router/tests/engine.rs @@ -5,7 +5,7 @@ use std::{ use litellm_router::{ engine::{Engine, Override, RouteError, RouterCall}, - failure::{Classified, ExceptionClass, Raised}, + failure::{Classified, ExceptionClass, Raised, StreamFailure}, host::{Attempt, Invoked, Op, RouterHost, Target}, operation::Operation, pyrepr::PyNumber, @@ -20,6 +20,7 @@ use rstest::rstest; enum Scripted { Ok, Fail(i64, ExceptionClass), + Stream(StreamFailure), } /// Answers each deployment's attempts from its own script; an unscripted attempt succeeds. @@ -113,6 +114,13 @@ impl RouterHost for ScriptedHost { error: format!("{id}#{index}"), classified: classified(status, class), }, + Scripted::Stream(failure) => Invoked::Failure { + error: format!("{id}#{index}"), + classified: Classified { + stream_failure: Some(failure), + ..classified(503, ExceptionClass::ServiceUnavailable) + }, + }, }) } @@ -421,3 +429,46 @@ async fn exhausted_retries_stamp_the_latest_error() { .any(|op| op == format!("stamp({}, 2)", raised(&error))) ); } + +#[rstest] +#[tokio::test] +async fn a_stream_failing_before_content_falls_back_without_retrying_the_group() { + let mut settings = settings(); + settings.fallbacks = Some(vec![chain("g", &["h"])]); + let engine = engine( + vec![ + deployment("a", "g"), + deployment("b", "g"), + deployment("c", "h"), + ], + settings, + ); + let before_content = Scripted::Stream(StreamFailure::BeforeContent); + let host = ScriptedHost::default() + .script("a", &[before_content.clone()]) + .script("b", &[before_content]); + + let routed = engine.route(&host, call("g")).await.ok().unwrap(); + + assert_eq!(host.attempts().len(), 2); + assert_eq!(host.attempts()[1], "c"); + assert_eq!(routed.outcome.model_group, "h"); + assert!(host.sleeps.lock().unwrap().is_empty()); +} + +#[rstest] +#[tokio::test] +async fn a_terminal_stream_failure_is_raised_without_retries_or_fallbacks() { + let mut settings = settings(); + settings.fallbacks = Some(vec![chain("g", &["h"])]); + let engine = engine(vec![deployment("a", "g"), deployment("c", "h")], settings); + let host = ScriptedHost::default().script("a", &[Scripted::Stream(StreamFailure::Terminal)]); + + let Err(RouteError::Failed(failed)) = engine.route(&host, call("g")).await else { + panic!("a terminal stream failure is raised"); + }; + + assert_eq!(host.attempts(), ["a"]); + assert_eq!(raised(&failed.error), "a#1"); + assert!(failed.ops.is_empty()); +} diff --git a/litellm/router_backends/rust_call.py b/litellm/router_backends/rust_call.py index 5edc1910ff4..0c09de05c95 100644 --- a/litellm/router_backends/rust_call.py +++ b/litellm/router_backends/rust_call.py @@ -17,12 +17,18 @@ from typing import Final, Literal, Protocol, TypeAlias, cast from pydantic import BaseModel, ConfigDict, TypeAdapter import litellm +from litellm.exceptions import MidStreamFallbackError from litellm.integrations.custom_guardrail import is_guardrail_intervention from litellm.litellm_core_utils.exception_mapping_utils import ( _get_response_headers, # pyright: ignore[reportPrivateUsage] # the header reader the cooldown callback uses ) from litellm.litellm_core_utils.secret_redaction import redact_string from litellm.litellm_core_utils.sensitive_data_masker import mask_sensitive_structure +from litellm.litellm_core_utils.streaming_handler import CustomStreamWrapper +from litellm.router_backends.python_router import ( + FallbackAwareStreamWrapper, + _stream_chunks_have_generated_content, # pyright: ignore[reportPrivateUsage] # the content rule Python's stream wrapper applies +) from litellm.router_strategy.complexity_router.context_compaction import compact_to_fit from litellm.router_utils.add_retry_fallback_headers import ( add_fallback_headers_to_response, @@ -39,7 +45,7 @@ from litellm.router_utils.cooldown_handlers import ( is_caller_timeout_408, ) from litellm.types.router import RouterErrors, RouterRateLimitError -from litellm.types.utils import ModelResponse +from litellm.types.utils import ModelResponse, ModelResponseStream from litellm.utils import ( _get_retry_after_from_exception_header, # pyright: ignore[reportPrivateUsage] # the Retry-After parser retries and cooldowns use ) @@ -200,6 +206,73 @@ def classify(error: BaseException, kwargs: Mapping[str, object], callbacks_ran: } +class _StreamFailed(Exception): + """A stream that failed before its first content chunk.""" + + def __init__(self, error: BaseException, kind: Literal["before_content", "terminal"]) -> None: + super().__init__(str(error)) + self.error: Final = error + self.kind: Final = kind + + +class ReplayStream(FallbackAwareStreamWrapper): + """The attempt's stream, replaying the chunks read while waiting for its first content. + + After content has reached the caller the stream cannot fall back, so a mid-stream + fallback error surfaces as the error that caused it, as the Python router's wrapper does. + """ + + def __init__(self, inner: CustomStreamWrapper, buffered: Sequence[ModelResponseStream]) -> None: + super().__init__( # pyright: ignore[reportUnknownMemberType] # CustomStreamWrapper's untyped constructor + completion_stream=cast(object, inner.completion_stream), # cast-ok: untyped attribute + model=cast(object, inner.model), # cast-ok: untyped attribute + custom_llm_provider=inner.custom_llm_provider, + logging_obj=inner.logging_obj, + _response_headers=getattr(inner, "_response_headers", None), + ) + self._inner: Final = inner + self._replay: Final = iter(tuple(buffered)) + self.chunks = cast(list[ModelResponseStream], inner.chunks) # cast-ok: untyped attribute + self._hidden_params = dict(cast(Mapping[str, object], inner._hidden_params)) # pyright: ignore[reportPrivateUsage] # cast-ok: the stream's own untyped metadata + + def __aiter__(self) -> ReplayStream: + return self + + async def __anext__(self) -> ModelResponseStream: + buffered: Final = next(self._replay, None) + if buffered is not None: + return buffered + try: + return await self._inner.__anext__() + except MidStreamFallbackError as error: + if error.original_exception is not None: + raise error.original_exception from error + raise + + async def aclose(self) -> None: + await self._inner.aclose() + + +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 + + +async def _replay_from_first_content(stream: CustomStreamWrapper) -> ReplayStream: + """Reads `stream` until a chunk carries content, so a failure before then can still go to + another deployment; the chunks read so far replay ahead of the rest.""" + buffered: Final[list[ModelResponseStream]] = [] # mutable-ok: chunks read before the first content + try: + async 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.""" @@ -244,6 +317,12 @@ class RoutedCall: return self._mock(parsed, kwargs) try: response: Final = await self._acompletion(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: @@ -382,8 +461,14 @@ class RoutedCall: message="Response output was blocked.", model=attempt.model_group, llm_provider="" ) ) + if isinstance(response, CustomStreamWrapper) and _not_fetched(response): + await response.fetch_stream() normalizer.success_calls[model_name] += 1 + if isinstance(response, CustomStreamWrapper): + return await _replay_from_first_content(response) return response + except _StreamFailed: + raise except _AttemptFailed as failed: self._stamp(failed.error, deployment, kwargs, model_name) raise diff --git a/litellm/router_backends/rust_router.py b/litellm/router_backends/rust_router.py index f9abcb54ffb..86dd42e3d29 100644 --- a/litellm/router_backends/rust_router.py +++ b/litellm/router_backends/rust_router.py @@ -271,8 +271,8 @@ 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: - raise NotImplementedError("the Rust router backend does not support streaming 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 62adfe0d934..663f93d9b9f 100644 --- a/tests/test_litellm_rust/router/test_rust_router.py +++ b/tests/test_litellm_rust/router/test_rust_router.py @@ -6,6 +6,8 @@ import random from collections.abc import Callable, Mapping, Sequence from typing import Final, Protocol +from unittest.mock import MagicMock, patch + import pytest import litellm @@ -15,7 +17,9 @@ from litellm.router_backends.selection import select_backend from litellm.rust_bridge.catalog import Route, RouteRule from litellm.rust_bridge.configuration import Rollout from litellm.types.router import RouterRateLimitError -from litellm.types.utils import ModelResponse +from litellm.exceptions import MidStreamFallbackError +from litellm.litellm_core_utils.streaming_handler import CustomStreamWrapper +from litellm.types.utils import ModelResponse, ModelResponseStream pytestmark = pytest.mark.requires_rust_extension @@ -178,3 +182,69 @@ async def test_the_facade_serves_a_supported_config_from_the_rust_backend() -> N assert backend.get_model_names() == ["g"] with pytest.raises(NotImplementedError, match=r"Router\.upsert_deployment"): _ = backend.upsert_deployment + + +class _FailingStream(CustomStreamWrapper): + """A provider stream that dies before its first chunk, as a 500 mid-stream does.""" + + 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: + raise MidStreamFallbackError( + message=f"provider 500 from {self.model}", + model=str(self.model), + llm_provider="openai", + is_pre_first_chunk=True, + original_exception=litellm.InternalServerError( + message=f"provider 500 from {self.model}", model=str(self.model), llm_provider="openai" + ), + ) + + +class _OkStream(_FailingStream): + def __init__(self, model: str) -> None: + super().__init__(model) + self._replies: Final = iter( + ( + ModelResponseStream(choices=[{"index": 0, "delta": {"role": "assistant"}}]), + ModelResponseStream(choices=[{"index": 0, "delta": {"content": f"ok-from-{model}"}}]), + ) + ) + + async def __anext__(self) -> ModelResponseStream: + reply: Final = next(self._replies, None) + if reply is None: + 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) + + +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"}}, + ), + "fallbacks": [{"primary": ["fb1", "fb2"]}], + "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: + 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)) + + assert observed[0] == observed[1] + assert observed[1] == ("ok-from-openai/fb2-model", ("primary", "fb1", "fb2"))