diff --git a/litellm/litellm_core_utils/streaming_handler.py b/litellm/litellm_core_utils/streaming_handler.py index 956e202a65a..079b29d5509 100644 --- a/litellm/litellm_core_utils/streaming_handler.py +++ b/litellm/litellm_core_utils/streaming_handler.py @@ -46,6 +46,7 @@ from litellm.litellm_core_utils.model_response_utils import ( ) from litellm.litellm_core_utils.redact_messages import LiteLLMLoggingObject, drop_token_ids from litellm.litellm_core_utils.thread_pool_executor import executor +from litellm.llms.base_llm.chat.transformation import BaseLLMException from litellm.llms.ollama.common_utils import OllamaResponseIdNormalizer from litellm.types.llms.base import LiteLLMBaseModel from litellm.types.router import GenericLiteLLMParams @@ -2534,7 +2535,10 @@ class CustomStreamWrapper: def fetch_sync_stream(self): if self.completion_stream is None and self.make_call is not None: # Call make_call to get the completion stream - self.completion_stream = self.make_call(client=litellm.module_level_client) + try: + self.completion_stream = self.make_call(client=litellm.module_level_client) + except (BaseLLMException, httpx.HTTPError) as e: + self._raise_open_error(e) self._stream_iter = self.completion_stream.__iter__() return self.completion_stream @@ -2542,7 +2546,10 @@ class CustomStreamWrapper: async def fetch_stream(self): if self.completion_stream is None and self.make_call is not None: # Call make_call to get the completion stream - self.completion_stream = await self.make_call(client=litellm.module_level_aclient) + try: + self.completion_stream = await self.make_call(client=litellm.module_level_aclient) + except (BaseLLMException, httpx.HTTPError) as e: + self._raise_open_error(e) self._stream_iter = self.completion_stream.__aiter__() return self.completion_stream @@ -2864,6 +2871,12 @@ class CustomStreamWrapper: if not self._emitted_disqualifying_content and self._delta_disqualifies_continuation(delta): self._emitted_disqualifying_content = True + def _raise_open_error(self, error: BaseLLMException | httpx.HTTPError) -> NoReturn: + try: + raise self._map_stream_exception(error) + except Exception as mapped: + raise mapped from error + def _map_stream_exception(self, original_exception: Exception) -> Exception: return exception_type( model=self.model, diff --git a/tests/unit/litellm_core_utils/test_streaming_handler.py b/tests/unit/litellm_core_utils/test_streaming_handler.py index cb0ef9a9a33..beb1b58ec27 100644 --- a/tests/unit/litellm_core_utils/test_streaming_handler.py +++ b/tests/unit/litellm_core_utils/test_streaming_handler.py @@ -38,9 +38,10 @@ from litellm.litellm_core_utils.streaming_handler import ( _ProviderChunkParsed, ) from litellm.llms.anthropic.chat.handler import ModelResponseIterator as AnthropicModelResponseIterator -from litellm.llms.custom_httpx.http_handler import HTTPHandler +from litellm.llms.custom_httpx.http_handler import HTTPHandler, MaskedHTTPStatusError from litellm.llms.hosted_vllm.chat.transformation import HostedVLLMChatConfig from litellm.llms.openai.chat.gpt_transformation import OpenAIChatCompletionStreamingHandler +from litellm.llms.vertex_ai.common_utils import VertexAIError from litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import ( ModelResponseIterator as GeminiModelResponseIterator, ) @@ -1433,6 +1434,102 @@ def test_sync_streaming_bad_request_not_midstream(logging_obj: Logging): assert "invalid maxOutputTokens" in str(excinfo.value) +def _gemini_stream_error(status_code: int, sync_mode: bool) -> Exception: + if not sync_mode: + return VertexAIError(status_code=status_code, message="upstream said no", headers=None) + request: Final = httpx.Request("POST", "https://generativelanguage.googleapis.com/v1beta/models/gemini-2.5-flash") + body: Final = b'{"error": {"message": "upstream said no"}}' + response: Final = httpx.Response(status_code, request=request, content=body) + return MaskedHTTPStatusError( + httpx.HTTPStatusError("upstream said no", request=request, response=response), + message=body, # pyright: ignore[reportArgumentType] # the sync stream handler passes the raw body bytes + text=body, # pyright: ignore[reportArgumentType] # the sync stream handler passes the raw body bytes + ) + + +def _gemini_stream_opened_by(error: BaseException, sync_mode: bool, logging_obj: Logging) -> CustomStreamWrapper: + def make_sync_call(client: object) -> NoReturn: + raise error + + async def make_call(client: object) -> NoReturn: + raise error + + return CustomStreamWrapper( + completion_stream=None, + model="gemini-2.5-flash", + logging_obj=logging_obj, + custom_llm_provider="gemini", + make_call=make_sync_call if sync_mode else make_call, + ) + + +async def _fetch(stream: CustomStreamWrapper, sync_mode: bool) -> object: + return stream.fetch_sync_stream() if sync_mode else await stream.fetch_stream() + + +@pytest.mark.parametrize("sync_mode", [True, False]) +@pytest.mark.parametrize( + ("status_code", "expected_error"), + [ + (400, litellm.BadRequestError), + (429, litellm.RateLimitError), + (503, litellm.ServiceUnavailableError), + ], +) +async def test_a_gemini_stream_that_fails_to_open_raises_the_mapped_error_from_fetch( + logging_obj: Logging, + sync_mode: bool, + status_code: int, + expected_error: type[litellm.BadRequestError | litellm.RateLimitError | litellm.ServiceUnavailableError], +) -> None: + provider_error: Final = _gemini_stream_error(status_code, sync_mode) + stream: Final = _gemini_stream_opened_by(provider_error, sync_mode, logging_obj) + with pytest.raises(expected_error, match="upstream said no") as excinfo: + await _fetch(stream, sync_mode) + assert (excinfo.value.status_code, excinfo.value.llm_provider) == (status_code, "gemini") + assert excinfo.value.__cause__ is provider_error + + +@pytest.mark.parametrize("sync_mode", [True, False]) +async def test_a_stream_open_error_that_is_already_mapped_is_raised_as_is( + logging_obj: Logging, sync_mode: bool +) -> None: + already_mapped: Final = litellm.BadRequestError( + message="already mapped", model="gemini-2.5-flash", llm_provider="gemini" + ) + stream: Final = _gemini_stream_opened_by(already_mapped, sync_mode, logging_obj) + with pytest.raises(litellm.BadRequestError) as excinfo: + await _fetch(stream, sync_mode) + assert excinfo.value is already_mapped + assert excinfo.value.__cause__ is None + + +async def test_a_cancelled_stream_open_stays_cancelled(logging_obj: Logging) -> None: + stream: Final = _gemini_stream_opened_by(asyncio.CancelledError(), False, logging_obj) + with pytest.raises(asyncio.CancelledError): + await stream.fetch_stream() + + +@pytest.mark.parametrize("sync_mode", [True, False]) +async def test_a_stream_that_opens_is_opened_once(logging_obj: Logging, sync_mode: bool) -> None: + def make_sync_call(client: object) -> ModelResponseListIterator: + return ModelResponseListIterator(model_responses=[]) + + async def make_call(client: object) -> ModelResponseListIterator: + return ModelResponseListIterator(model_responses=[]) + + stream: Final = CustomStreamWrapper( + completion_stream=None, + model="gemini-2.5-flash", + logging_obj=logging_obj, + custom_llm_provider="gemini", + make_call=make_sync_call if sync_mode else make_call, + ) + opened: Final = await _fetch(stream, sync_mode) + assert isinstance(opened, ModelResponseListIterator) + assert await _fetch(stream, sync_mode) is opened + + def _bedrock_error_event(exception_type: str): """A mocked botocore event-stream error event: status_code is botocore's hard-coded 400, with the real type in the :exception-type header.""" @@ -2113,37 +2210,32 @@ async def test_in_body_stream_error_500_wraps_for_midstream_fallback( assert "internal engine crash" in str(excinfo.value) -@pytest.mark.asyncio -async def test_async_streaming_read_timeout_triggers_midstream_fallback( - logging_obj: Logging, -): - """A mid-stream httpx.ReadTimeout must wrap into MidStreamFallbackError so - the Router's FallbackStreamWrapper can switch to a fallback model. +@pytest.mark.parametrize("sync_mode", [True, False]) +async def test_a_read_timeout_while_opening_the_stream_triggers_midstream_fallback( + logging_obj: Logging, sync_mode: bool +) -> None: + read_timeout: Final = httpx.ReadTimeout("Timeout on reading data from socket") - Previously __anext__ caught httpx.TimeoutException and re-raised it raw, - which bypassed _handle_stream_fallback_error and prevented stream_timeout - from triggering fallbacks the way connection-phase timeout does. - """ - import httpx + def make_sync_call(client: object) -> NoReturn: + raise read_timeout - from litellm.exceptions import MidStreamFallbackError + async def make_call(client: object) -> NoReturn: + raise read_timeout - async def _raise_read_timeout(**kwargs): - raise httpx.ReadTimeout("Timeout on reading data from socket") - - response = CustomStreamWrapper( + response: Final = CustomStreamWrapper( completion_stream=None, model="gpt-4", logging_obj=logging_obj, custom_llm_provider="openai", - make_call=_raise_read_timeout, + make_call=make_sync_call if sync_mode else make_call, ) with pytest.raises(MidStreamFallbackError) as excinfo: - await response.__anext__() + next(response) if sync_mode else await response.__anext__() assert excinfo.value.is_pre_first_chunk is True assert isinstance(excinfo.value.original_exception, litellm.Timeout) + assert excinfo.value.original_exception.__cause__ is read_timeout def _raw_read_timeout() -> NoReturn: diff --git a/tests/unit/test_router/core/test_streaming_characterization_chat.py b/tests/unit/test_router/core/test_streaming_characterization_chat.py index 590fafede8e..8f727c5f0ee 100644 --- a/tests/unit/test_router/core/test_streaming_characterization_chat.py +++ b/tests/unit/test_router/core/test_streaming_characterization_chat.py @@ -10,14 +10,17 @@ while one that raises later fails while the caller iterates. import asyncio from collections.abc import Callable, Mapping -from typing import Final +from typing import Final, NoReturn +import httpx import pytest import litellm from litellm.exceptions import MidStreamFallbackError from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLogging from litellm.litellm_core_utils.streaming_handler import CustomStreamWrapper +from litellm.llms.custom_httpx.http_handler import MaskedHTTPStatusError +from litellm.llms.vertex_ai.common_utils import VertexAIError from litellm.router_core import RouterCore from litellm.types.utils import Usage from tests.unit.test_router.core.scripted_host import ( @@ -35,6 +38,7 @@ from tests.unit.test_router.core.streaming_helpers import ( collect_texts, collect_texts_sync, drain_into, + request_logging, ) MakeCore = Callable[..., RouterCore] @@ -107,6 +111,51 @@ async def test_a_first_chunk_bad_request_is_raised_from_the_call_without_retry( assert host.deployments_called() == ["a"] +def _gemini_stream_that_fails_to_open(sync: bool) -> Callable[..., CustomStreamWrapper]: + request: Final = httpx.Request("POST", "https://generativelanguage.googleapis.com/v1beta/models/gemini-2.5-flash") + body: Final = b'{"error": {"code": 400, "message": "gemini said no", "status": "INVALID_ARGUMENT"}}' + sync_error: Final = MaskedHTTPStatusError( + httpx.HTTPStatusError("gemini said no", request=request, response=httpx.Response(400, request=request)), + message=body, # pyright: ignore[reportArgumentType] # the sync stream handler passes the raw body bytes + text=body, # pyright: ignore[reportArgumentType] # the sync stream handler passes the raw body bytes + ) + + def make_sync_call(client: object) -> NoReturn: + raise sync_error + + async def make_call(client: object) -> NoReturn: + raise VertexAIError(status_code=400, message=body.decode(), headers=None) + + def answer(**_request: object) -> CustomStreamWrapper: # kwargs-ok: the host passes the request kwargs + return CustomStreamWrapper( + completion_stream=None, + model="gemini-2.5-flash", + custom_llm_provider="gemini", + logging_obj=request_logging(), + make_call=make_sync_call if sync else make_call, + ) + + return answer + + +async def test_a_gemini_stream_that_fails_to_open_raises_the_mapped_bad_request( + make_core: MakeCore, host: ScriptedHost +) -> None: + host.answer("openai/a", _gemini_stream_that_fails_to_open(sync=False)) + core: Final = make_core(model_list=[deployment("a", "a")], num_retries=0) + with pytest.raises(litellm.BadRequestError, match="gemini said no"): + await open_stream(core) + + +def test_a_sync_gemini_stream_that_fails_to_open_raises_the_mapped_bad_request( + make_core: MakeCore, host: ScriptedHost +) -> None: + host.answer("openai/a", _gemini_stream_that_fails_to_open(sync=True)) + core: Final = make_core(model_list=[deployment("a", "a")], num_retries=0) + with pytest.raises(litellm.BadRequestError, match="gemini said no"): + core.completion(model="a", messages=MESSAGES, stream=True) # pyright: ignore[reportAttributeAccessIssue] # endpoint method + + async def test_the_first_chunk_already_read_is_still_delivered_to_the_caller( make_core: MakeCore, host: ScriptedHost ) -> None: