mirror of
https://github.com/BerriAI/litellm.git
synced 2026-08-28 05:25:59 +00:00
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:
parent
ebc3845dac
commit
70e47f4897
2 changed files with 164 additions and 0 deletions
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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(
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue