fix(router): surface original provider exception, not the internal wrapper, when mid-stream fallback gives up

When content has already streamed and MidStreamFallbackError carries
original_exception (e.g. RateLimitError), both the async and sync
streaming iterators bare-re-raised the wrapper itself, so the client
lost the specific error type/code/provider_specific_fields instead of
seeing the real provider error. The fallback-failure path a few lines
below already unwraps to original_exception for the same reason; apply
the same pattern here.

Also extend _stream_chunks_have_generated_content to recognize audio,
images, and annotations deltas as generated content, matching
is_chunk_non_empty's existing annotations check and Delta's treatment
of audio/images as first-class content fields — a stream carrying only
one of these before failing was not recognized as already-streamed,
so the router could still restart it via fallback after the client had
received real content.
This commit is contained in:
Deepanshu 2026-08-04 14:33:12 -04:00
parent ebc3845dac
commit 70e47f4897
2 changed files with 164 additions and 0 deletions

View file

@ -320,6 +320,9 @@ def _stream_chunks_have_generated_content(chunks: Sequence[ModelResponseStream])
or delta.get("reasoning_content")
or delta.get("thinking_blocks")
or delta.get("reasoning_items")
or delta.get("audio")
or delta.get("images")
or delta.get("annotations")
):
return True
return False
@ -2115,6 +2118,8 @@ class Router:
if not e.is_pre_first_chunk and (
e.generated_content or _stream_chunks_have_generated_content(model_response.chunks)
):
if e.original_exception is not None:
raise e.original_exception from e
raise
from litellm.main import stream_chunk_builder
@ -2658,6 +2663,8 @@ class Router:
if not e.is_pre_first_chunk and (
e.generated_content or _stream_chunks_have_generated_content(model_response.chunks)
):
if e.original_exception is not None:
raise e.original_exception from e
raise
from litellm.main import stream_chunk_builder

View file

@ -1838,6 +1838,85 @@ async def test_acompletion_streaming_iterator():
print("\n=== All tests passed! ===")
@pytest.mark.asyncio
async def test_acompletion_streaming_iterator_reraises_original_exception_when_available():
"""Async: when the mid-stream MidStreamFallbackError wraps a real provider
exception (original_exception), the router must re-raise that original
exception instead of the internal wrapper, so the client sees the
specific error type/code (e.g. RateLimitError) rather than a generic
MidStreamFallbackError."""
from unittest.mock import MagicMock
from litellm.exceptions import MidStreamFallbackError, RateLimitError
router = litellm.Router(
model_list=[
{
"model_name": "gpt-4",
"litellm_params": {"model": "gpt-4", "api_key": "fake-key"},
}
],
set_verbose=True,
)
messages = [{"role": "user", "content": "Test"}]
initial_kwargs = {"model": "gpt-4", "stream": True}
original_exception = RateLimitError(
message="rate limited",
llm_provider="vertex_ai",
model="gpt-4",
)
error = MidStreamFallbackError(
message="rate limited",
model="gpt-4",
llm_provider="openai",
original_exception=original_exception,
generated_content="Hello",
)
mock_chunks = [
MagicMock(choices=[MagicMock(delta=MagicMock(content="Hello"))]),
MagicMock(choices=[MagicMock(delta=MagicMock(content=" there"))]),
]
class AsyncIteratorWithError:
def __init__(self, items, error_after_index):
self.items = items
self.index = 0
self.error_after_index = error_after_index
def __aiter__(self):
return self
async def __anext__(self):
if self.index >= len(self.items):
raise StopAsyncIteration
if self.index == self.error_after_index:
raise error
item = self.items[self.index]
self.index += 1
return item
mock_error_response = AsyncIteratorWithError(mock_chunks, 1)
setattr(mock_error_response, "model", "gpt-4")
setattr(mock_error_response, "custom_llm_provider", "openai")
setattr(mock_error_response, "logging_obj", MagicMock())
result = await router._acompletion_streaming_iterator(
model_response=mock_error_response,
messages=messages,
initial_kwargs=initial_kwargs,
)
with pytest.raises(RateLimitError) as exc_info:
async for _ in result:
pass
assert exc_info.value is original_exception
assert exc_info.value.type == "throttling_error"
assert exc_info.value.code == "429"
@pytest.mark.asyncio
async def test_acompletion_streaming_iterator_edge_cases():
"""Test edge cases for _acompletion_streaming_iterator."""
@ -2132,6 +2211,70 @@ def test_completion_streaming_iterator_reraises_mid_chunk_error():
list(result)
def test_completion_streaming_iterator_reraises_original_exception_when_available():
"""Sync: when the mid-chunk MidStreamFallbackError wraps a real provider
exception (original_exception), the router must re-raise that original
exception instead of the internal wrapper, so the client sees the
specific error type/code (e.g. RateLimitError) rather than a generic
MidStreamFallbackError."""
from unittest.mock import MagicMock
from litellm.exceptions import MidStreamFallbackError, RateLimitError
router = litellm.Router(
model_list=[
{
"model_name": "gpt-4",
"litellm_params": {"model": "gpt-4", "api_key": "fake-key"},
}
],
)
messages = [{"role": "user", "content": "Test"}]
initial_kwargs = {"model": "gpt-4", "stream": True}
original_exception = RateLimitError(
message="rate limited",
llm_provider="vertex_ai",
model="gpt-4",
)
mid_chunk_error = MidStreamFallbackError(
message="rate limited",
model="gpt-4",
llm_provider="openai",
original_exception=original_exception,
generated_content="Hello, I am",
is_pre_first_chunk=False,
)
class SyncIteratorMidChunkError:
def __init__(self):
self.model = "gpt-4"
self.custom_llm_provider = "openai"
self.logging_obj = MagicMock()
self.chunks = []
def __iter__(self):
return self
def __next__(self):
raise mid_chunk_error
mock_response = SyncIteratorMidChunkError()
result = router._completion_streaming_iterator(
model_response=mock_response,
messages=messages,
initial_kwargs=initial_kwargs,
)
with pytest.raises(RateLimitError) as exc_info:
list(result)
assert exc_info.value is original_exception
assert exc_info.value.type == "throttling_error"
assert exc_info.value.code == "429"
def test_completion_streaming_iterator_reraises_mid_chunk_error_with_no_text_content():
"""Sync: a reasoning-only chunk sets is_pre_first_chunk=False without populating
generated_content (which only tracks text deltas). The re-raise guard must still
@ -6354,6 +6497,20 @@ def test_stream_chunks_have_generated_content_detects_text_and_non_text():
reasoning_items_chunk = _chunk(reasoning_items_delta)
assert _stream_chunks_have_generated_content([reasoning_items_chunk]) is True
audio_delta = Delta(audio={"data": "abc123", "expires_at": 1234567890, "transcript": "hello"})
audio_chunk = _chunk(audio_delta)
assert _stream_chunks_have_generated_content([audio_chunk]) is True
images_delta = Delta(images=[{"image_url": {"url": "https://example.com/img.png"}, "index": 0, "type": "image_url"}])
images_chunk = _chunk(images_delta)
assert _stream_chunks_have_generated_content([images_chunk]) is True
annotations_delta = Delta(
annotations=[{"type": "url_citation", "url_citation": {"url": "https://example.com"}}]
)
annotations_chunk = _chunk(annotations_delta)
assert _stream_chunks_have_generated_content([annotations_chunk]) is True
def test_get_configured_token_limits_reads_deployment_model_info():
router = litellm.Router(