From 1679c9d1ba2b4066d347cc6ba4cc8ee458e152b6 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sat, 7 Nov 2026 01:02:40 -0800 Subject: [PATCH] fix(streaming): keep OpenAIError for aborts on OpenAI-compatible streams Streams parsed by the OpenAI streaming handler raised the base exception on an upstream abort, so failure callbacks lost the OpenAIError class they got before. The handler now names its abort error class, and the Bedrock OpenAI-compatible decoder passes it through --- litellm/llms/base_llm/base_model_iterator.py | 10 ++++++---- litellm/llms/bedrock/chat/invoke_handler.py | 4 +++- litellm/llms/openai/chat/gpt_transformation.py | 4 +++- .../llms/bedrock/chat/test_invoke_handler.py | 17 +++++++++++++++++ .../chat/test_openai_gpt_transformation.py | 11 +++++++++++ 5 files changed, 40 insertions(+), 6 deletions(-) diff --git a/litellm/llms/base_llm/base_model_iterator.py b/litellm/llms/base_llm/base_model_iterator.py index f6563d3e7a2..2f6700837d5 100644 --- a/litellm/llms/base_llm/base_model_iterator.py +++ b/litellm/llms/base_llm/base_model_iterator.py @@ -1,6 +1,6 @@ import json from abc import abstractmethod -from typing import TYPE_CHECKING, Final, cast +from typing import TYPE_CHECKING, ClassVar, Final, cast from pydantic import BaseModel, ConfigDict, ValidationError @@ -100,16 +100,18 @@ class _StreamedChunk(BaseModel): choices: tuple[_StreamedChoice, ...] | None = None -def raise_on_upstream_abort(chunk: object) -> None: +def raise_on_upstream_abort(chunk: object, error: type[BaseLLMException] = BaseLLMException) -> None: try: choices: Final = _StreamedChunk.model_validate(chunk).choices or () except ValidationError: return if any(choice.finish_reason == "abort" for choice in choices): - raise BaseLLMException(status_code=500, message=UPSTREAM_ABORT_MESSAGE) + raise error(status_code=500, message=UPSTREAM_ABORT_MESSAGE) class BaseModelResponseIterator: + upstream_abort_error: ClassVar[type[BaseLLMException]] = BaseLLMException + def __init__(self, streaming_response, sync_stream: bool, json_mode: bool | None = False): self.streaming_response = streaming_response self.response_iterator = self.streaming_response @@ -179,7 +181,7 @@ class BaseModelResponseIterator: tool_use=None, ) elif stripped_json_chunk: - raise_on_upstream_abort(stripped_json_chunk) + raise_on_upstream_abort(stripped_json_chunk, self.upstream_abort_error) return self.chunk_parser(chunk=stripped_json_chunk) else: return GenericStreamingChunk( diff --git a/litellm/llms/bedrock/chat/invoke_handler.py b/litellm/llms/bedrock/chat/invoke_handler.py index 70b2787db9e..dd56481af5a 100644 --- a/litellm/llms/bedrock/chat/invoke_handler.py +++ b/litellm/llms/bedrock/chat/invoke_handler.py @@ -1022,5 +1022,7 @@ class AmazonOpenAICompatibleStreamDecoder(AWSEventStreamDecoder): ) def _chunk_parser(self, chunk_data: dict[str, object]) -> ModelResponseStream: - base_model_iterator.raise_on_upstream_abort(chunk_data) + base_model_iterator.raise_on_upstream_abort( + chunk_data, self.openai_model_response_iterator.upstream_abort_error + ) return self.openai_model_response_iterator.chunk_parser(chunk=chunk_data) diff --git a/litellm/llms/openai/chat/gpt_transformation.py b/litellm/llms/openai/chat/gpt_transformation.py index 4aadf1dcbca..2257f105c60 100644 --- a/litellm/llms/openai/chat/gpt_transformation.py +++ b/litellm/llms/openai/chat/gpt_transformation.py @@ -6,7 +6,7 @@ import json import os from collections.abc import AsyncIterator, Coroutine, Iterator, Mapping, Sequence from types import MappingProxyType -from typing import TYPE_CHECKING, Any, Final, Literal, Optional, cast, overload +from typing import TYPE_CHECKING, Any, ClassVar, Final, Literal, Optional, cast, overload from urllib.parse import urlparse import httpx @@ -1019,6 +1019,8 @@ class OpenAIUnknownModelConfig(OpenAIGPTConfig): class OpenAIChatCompletionStreamingHandler(BaseModelResponseIterator): + upstream_abort_error: ClassVar[type[BaseLLMException]] = OpenAIError + def _map_reasoning_to_reasoning_content(self, choices: list) -> list: """ Map 'reasoning' field to 'reasoning_content' field in delta. diff --git a/tests/unit/llms/bedrock/chat/test_invoke_handler.py b/tests/unit/llms/bedrock/chat/test_invoke_handler.py index 926368d79b7..e55ce9988c8 100644 --- a/tests/unit/llms/bedrock/chat/test_invoke_handler.py +++ b/tests/unit/llms/bedrock/chat/test_invoke_handler.py @@ -29,6 +29,7 @@ from litellm.llms.bedrock.chat.invoke_handler import ( ) from litellm.llms.bedrock.common_utils import BedrockError, get_bedrock_stream_event_statuses from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler +from litellm.llms.openai.common_utils import OpenAIError from litellm.types.llms.bedrock import BedrockInvokeNovaRequest from litellm.types.utils import ModelResponseStream from tests.unit.llms.bedrock.slow_upstream import ( @@ -1054,6 +1055,22 @@ def test_iter_bytes_raises_when_the_stream_ends_mid_message() -> None: assert "first bytes=" not in exc_info.value.message +def test_iter_bytes_raises_the_openai_error_when_upstream_aborts() -> None: + decoder: Final = AmazonOpenAICompatibleStreamDecoder(model="moonshot.kimi-k2-thinking", sync_stream=True) + raw_stream: Final = b"".join( + _bedrock_event_stream_frame(chunk) + for chunk in ( + _openai_stream_chunk({"role": "assistant", "content": "Partial answer"}), + _openai_stream_chunk({}, "abort"), + ) + ) + + with pytest.raises(OpenAIError, match="finish_reason=abort") as exc_info: + _chunks: Final = list(decoder.iter_bytes(iter([raw_stream]))) + + assert exc_info.value.status_code == 500 + + def test_iter_bytes_yields_a_complete_stream_without_raising() -> None: decoder: Final = AmazonOpenAICompatibleStreamDecoder(model="moonshot.kimi-k2-thinking", sync_stream=True) diff --git a/tests/unit/llms/openai/chat/test_openai_gpt_transformation.py b/tests/unit/llms/openai/chat/test_openai_gpt_transformation.py index 263c2f7618d..12c92ec7456 100644 --- a/tests/unit/llms/openai/chat/test_openai_gpt_transformation.py +++ b/tests/unit/llms/openai/chat/test_openai_gpt_transformation.py @@ -348,6 +348,17 @@ class TestOpenAIChatCompletionStreamingHandler: assert excinfo.value.status_code == 400 assert "not multimodal" in excinfo.value.message + def test_stream_ending_in_abort_raises_an_openai_error(self): + from litellm.llms.openai.common_utils import OpenAIError + + line = "data: " + json.dumps({"choices": [{"index": 0, "delta": {}, "finish_reason": "abort"}]}) + handler = OpenAIChatCompletionStreamingHandler(streaming_response=iter((line,)), sync_stream=True) + + with pytest.raises(OpenAIError, match="finish_reason=abort") as excinfo: + list(handler) + + assert excinfo.value.status_code == 500 + def test_chunk_parser_error_payload_without_usable_code_maps_to_500(self): """OpenAI-style error payloads may carry a string code (e.g. "invalid_api_key") or none at all; those must map to 500, not crash."""