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."""