From fe92f14602ef26fa1a160f80c1ef3a74f263d393 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Thu, 29 Oct 2026 09:05:42 -0700 Subject: [PATCH] fix(streaming): surface upstream abort on azure_ai, databricks, Vertex Model Garden and text completion streams --- .../litellm_core_utils/streaming_handler.py | 20 ++- litellm/llms/base_llm/base_model_iterator.py | 28 +++- litellm/llms/bedrock/chat/invoke_handler.py | 1 + litellm/llms/databricks/streaming_utils.py | 3 + .../llms/openai/chat/gpt_transformation.py | 9 +- .../test_streaming_handler.py | 157 +++++++++++++----- .../llms/base_llm/test_base_model_iterator.py | 11 ++ .../llms/bedrock/chat/test_invoke_handler.py | 25 +++ 8 files changed, 201 insertions(+), 53 deletions(-) diff --git a/litellm/litellm_core_utils/streaming_handler.py b/litellm/litellm_core_utils/streaming_handler.py index 5542249d2c7..d06432e8cd6 100644 --- a/litellm/litellm_core_utils/streaming_handler.py +++ b/litellm/litellm_core_utils/streaming_handler.py @@ -23,6 +23,7 @@ from typing_extensions import NotRequired, TypedDict import litellm from litellm import verbose_logger from litellm._uuid import uuid +from litellm.exceptions import InternalServerError from litellm.litellm_core_utils.asyncify import asyncify from litellm.litellm_core_utils.core_helpers import chunk_usage, provider_reported_usage_cost from litellm.litellm_core_utils.hidden_params import ( @@ -36,6 +37,7 @@ from litellm.litellm_core_utils.model_response_utils import ( ) from litellm.litellm_core_utils.redact_messages import LiteLLMLoggingObject from litellm.litellm_core_utils.thread_pool_executor import executor +from litellm.llms.base_llm.base_model_iterator import UPSTREAM_ABORT_MESSAGE from litellm.llms.ollama.common_utils import OllamaResponseIdNormalizer from litellm.types.llms.base import LiteLLMBaseModel from litellm.types.router import GenericLiteLLMParams @@ -817,13 +819,23 @@ class CustomStreamWrapper: except Exception as e: raise e - def handle_openai_text_completion_chunk(self, chunk): + def _upstream_abort_error(self, model: str) -> InternalServerError: + return InternalServerError( + message=UPSTREAM_ABORT_MESSAGE, + model=model, + llm_provider=self.custom_llm_provider, + ) + + def handle_openai_text_completion_chunk(self, chunk: object): try: text = "" is_finished = False finish_reason = None usage = None choices: Final[Sequence[_TextCompletionChoiceLike]] = getattr(chunk, "choices", []) + if any(choice.finish_reason == "abort" for choice in choices): + chunk_model: Final[str] = getattr(chunk, "model", "") + raise self._upstream_abort_error(chunk_model) if len(choices) > 0: text = choices[0].text if choices[0].finish_reason is not None: @@ -1575,11 +1587,7 @@ class CustomStreamWrapper: if isinstance(chunk, ChatCompletionChunk): self._openai_handler_receives_chat_chunks = True if any(str(choice.finish_reason) == "abort" for choice in chunk.choices or ()): - raise litellm.InternalServerError( - message="Upstream generation aborted (finish_reason=abort)", - model=chunk.model, - llm_provider=self.custom_llm_provider, - ) + raise self._upstream_abort_error(chunk.model) if self.custom_llm_provider in [ LlmProviders.AZURE.value, LlmProviders.AZURE_AI.value, diff --git a/litellm/llms/base_llm/base_model_iterator.py b/litellm/llms/base_llm/base_model_iterator.py index 3afc685b3dd..f6563d3e7a2 100644 --- a/litellm/llms/base_llm/base_model_iterator.py +++ b/litellm/llms/base_llm/base_model_iterator.py @@ -2,8 +2,11 @@ import json from abc import abstractmethod from typing import TYPE_CHECKING, Final, cast +from pydantic import BaseModel, ConfigDict, ValidationError + import litellm from litellm.constants import STREAM_SSE_DONE_STRING +from litellm.llms.base_llm.chat.transformation import BaseLLMException if TYPE_CHECKING: import httpx @@ -84,6 +87,28 @@ def convert_model_response_to_streaming( raise ValueError(f"Failed to convert ModelResponse to ModelResponseStream: {model_response}. Error: {e}") +UPSTREAM_ABORT_MESSAGE: Final = "Upstream generation aborted (finish_reason=abort)" + + +class _StreamedChoice(BaseModel): + model_config = ConfigDict(frozen=True) + finish_reason: object = None + + +class _StreamedChunk(BaseModel): + model_config = ConfigDict(frozen=True) + choices: tuple[_StreamedChoice, ...] | None = None + + +def raise_on_upstream_abort(chunk: object) -> 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) + + class BaseModelResponseIterator: def __init__(self, streaming_response, sync_stream: bool, json_mode: bool | None = False): self.streaming_response = streaming_response @@ -142,7 +167,7 @@ class BaseModelResponseIterator: def _handle_string_chunk(self, str_line: str) -> GenericStreamingChunk | ModelResponseStream: # chunk is a str at this point - stripped_json_chunk: Final = BaseModelResponseIterator._string_to_dict_parser(str_line=str_line) + stripped_json_chunk: Final = BaseModelResponseIterator.string_to_dict_parser(str_line=str_line) stripped_sse_chunk: Final = litellm.CustomStreamWrapper.strip_sse_data_from_chunk(str_line) if stripped_sse_chunk is not None and stripped_sse_chunk.strip() == STREAM_SSE_DONE_STRING: return GenericStreamingChunk( @@ -154,6 +179,7 @@ class BaseModelResponseIterator: tool_use=None, ) elif stripped_json_chunk: + raise_on_upstream_abort(stripped_json_chunk) 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 a6b34010297..c06f8458f17 100644 --- a/litellm/llms/bedrock/chat/invoke_handler.py +++ b/litellm/llms/bedrock/chat/invoke_handler.py @@ -1017,4 +1017,5 @@ class AmazonOpenAICompatibleStreamDecoder(AWSEventStreamDecoder): ) def _chunk_parser(self, chunk_data: dict[str, object]) -> ModelResponseStream: + base_model_iterator.raise_on_upstream_abort(chunk_data) return self.openai_model_response_iterator.chunk_parser(chunk=chunk_data) diff --git a/litellm/llms/databricks/streaming_utils.py b/litellm/llms/databricks/streaming_utils.py index cc85012b803..e3600fd61ac 100644 --- a/litellm/llms/databricks/streaming_utils.py +++ b/litellm/llms/databricks/streaming_utils.py @@ -5,6 +5,7 @@ from pydantic import ConfigDict, TypeAdapter import litellm from litellm import verbose_logger +from litellm.llms.base_llm.base_model_iterator import raise_on_upstream_abort from litellm.types.llms.openai import ( ChatCompletionToolCallChunk, ChatCompletionToolCallFunctionChunk, @@ -105,6 +106,7 @@ class ModelResponseIterator: chunk: Final = self._next_sse_data() try: json_chunk: Final = _JSON_CHUNK.validate_python(json.loads(chunk)) + raise_on_upstream_abort(json_chunk) return self.chunk_parser(chunk=json_chunk) except ValueError as e: verbose_logger.debug( @@ -132,6 +134,7 @@ class ModelResponseIterator: raise StopAsyncIteration try: json_chunk: Final = _JSON_CHUNK.validate_python(json.loads(chunk)) + raise_on_upstream_abort(json_chunk) return self.chunk_parser(chunk=json_chunk) except ValueError as e: verbose_logger.debug( diff --git a/litellm/llms/openai/chat/gpt_transformation.py b/litellm/llms/openai/chat/gpt_transformation.py index 22af51aaabe..434586257ed 100644 --- a/litellm/llms/openai/chat/gpt_transformation.py +++ b/litellm/llms/openai/chat/gpt_transformation.py @@ -119,7 +119,6 @@ def translate_developer_role_for_openai_compatible_endpoint( _OBJECT_MAPPING: Final[TypeAdapter[Mapping[str, object]]] = TypeAdapter(Mapping[str, object]) -_CHOICE_MAPPINGS: Final = TypeAdapter(tuple[Mapping[str, object], ...]) _DECODED_JSON: Final = TypeAdapter(object) @@ -1012,13 +1011,7 @@ class OpenAIChatCompletionStreamingHandler(BaseModelResponseIterator): status_code=error_status_code, message=error_message, ) - raw_choices: Final = chunk.get("choices") or [] - if any(choice.get("finish_reason") == "abort" for choice in _CHOICE_MAPPINGS.validate_python(raw_choices)): - raise OpenAIError( - status_code=500, - message="Upstream generation aborted (finish_reason=abort)", - ) - choices: Final = self._map_reasoning_to_reasoning_content(raw_choices) + choices: Final = self._map_reasoning_to_reasoning_content(chunk.get("choices") or []) fingerprint: Final = _OBJECT_MAPPING.validate_python(chunk).get("system_fingerprint") kwargs: Final[dict[str, Any]] = { diff --git a/tests/unit/litellm_core_utils/test_streaming_handler.py b/tests/unit/litellm_core_utils/test_streaming_handler.py index f4502937fce..2439a491795 100644 --- a/tests/unit/litellm_core_utils/test_streaming_handler.py +++ b/tests/unit/litellm_core_utils/test_streaming_handler.py @@ -15,6 +15,7 @@ from unittest.mock import AsyncMock, MagicMock, Mock, patch import httpx import pytest from openai import AsyncOpenAI, AsyncStream, OpenAI, Stream +from openai.types import Completion from openai.types.chat import ChatCompletionChunk from pydantic import BaseModel, TypeAdapter import litellm @@ -30,6 +31,9 @@ from litellm.litellm_core_utils.streaming_handler import ( _ProviderChunkParsed, ) from litellm.llms.custom_httpx.http_handler import HTTPHandler +from litellm.llms.azure_ai.chat.transformation import AzureAIChatCompletionResponseIterator +from litellm.llms.databricks.chat.transformation import DatabricksChatResponseIterator +from litellm.llms.databricks.streaming_utils import ModelResponseIterator as OpenAILikeModelResponseIterator from litellm.llms.openai.chat.gpt_transformation import OpenAIChatCompletionStreamingHandler from litellm.responses.mcp.chat_completions_handler import _add_mcp_metadata_to_response from litellm.types.utils import ( @@ -1482,8 +1486,45 @@ async def test_bedrock_validation_error_raises_directly(logging_obj: Logging): assert getattr(excinfo.value, "status_code", None) == 400 +_RAW_SSE_ITERATORS: Final = MappingProxyType( + { + "http_handler": ("hosted_vllm", OpenAIChatCompletionStreamingHandler), + "azure_ai": ("azure_ai", AzureAIChatCompletionResponseIterator), + "databricks": ("databricks", DatabricksChatResponseIterator), + "openai_like": ("vertex_ai", OpenAILikeModelResponseIterator), + } +) + + +def _chat_stream_wrapper( + transport: str, raw_response: httpx.Response, client: OpenAI | AsyncOpenAI, logging_obj: Logging +) -> CustomStreamWrapper: + if transport == "sdk": + return CustomStreamWrapper( + completion_stream=( + AsyncStream(cast_to=ChatCompletionChunk, response=raw_response, client=client) + if isinstance(client, AsyncOpenAI) + else Stream(cast_to=ChatCompletionChunk, response=raw_response, client=client) + ), + model="test-model", + logging_obj=logging_obj, + custom_llm_provider="openai", + ) + provider, iterator = _RAW_SSE_ITERATORS[transport] + async_stream: Final = isinstance(client, AsyncOpenAI) + return CustomStreamWrapper( + completion_stream=iterator( + streaming_response=raw_response.aiter_lines() if async_stream else raw_response.iter_lines(), + sync_stream=not async_stream, + ), + model="test-model", + logging_obj=logging_obj, + custom_llm_provider=provider, + ) + + @pytest.mark.parametrize("async_stream", [False, True]) -@pytest.mark.parametrize("transport", ["sdk", "http_handler"]) +@pytest.mark.parametrize("transport", ["sdk", *_RAW_SSE_ITERATORS]) @pytest.mark.parametrize("prefix", ["", "Partial answer"]) @pytest.mark.parametrize("abort_index", [0, 1]) @pytest.mark.parametrize("abort_content", ["", "unfinished"]) @@ -1525,24 +1566,7 @@ async def test_aborted_stream_raises_instead_of_finishing_successfully( request=httpx.Request("POST", "https://provider.example/v1/chat/completions"), ) client: Final = AsyncOpenAI(api_key="unused") if async_stream else OpenAI(api_key="unused") - stream: Final = ( - ( - AsyncStream(cast_to=ChatCompletionChunk, response=raw_response, client=client) - if isinstance(client, AsyncOpenAI) - else Stream(cast_to=ChatCompletionChunk, response=raw_response, client=client) - ) - if transport == "sdk" - else OpenAIChatCompletionStreamingHandler( - streaming_response=raw_response.aiter_lines() if async_stream else raw_response.iter_lines(), - sync_stream=not async_stream, - ) - ) - response: Final = CustomStreamWrapper( - completion_stream=stream, - model="test-model", - logging_obj=logging_obj, - custom_llm_provider="openai" if transport == "sdk" else "hosted_vllm", - ) + response: Final = _chat_stream_wrapper(transport, raw_response, client, logging_obj) try: with pytest.raises(MidStreamFallbackError, match="abort") as excinfo: @@ -1559,7 +1583,7 @@ async def test_aborted_stream_raises_instead_of_finishing_successfully( @pytest.mark.parametrize("async_stream", [False, True]) -@pytest.mark.parametrize("transport", ["sdk", "http_handler"]) +@pytest.mark.parametrize("transport", ["sdk", "http_handler", "azure_ai", "openai_like"]) async def test_stream_with_null_choices_chunk_finishes_normally( async_stream: bool, transport: str, logging_obj: Logging ) -> None: @@ -1592,24 +1616,7 @@ async def test_stream_with_null_choices_chunk_finishes_normally( request=httpx.Request("POST", "https://provider.example/v1/chat/completions"), ) client: Final = AsyncOpenAI(api_key="unused") if async_stream else OpenAI(api_key="unused") - stream: Final = ( - ( - AsyncStream(cast_to=ChatCompletionChunk, response=raw_response, client=client) - if isinstance(client, AsyncOpenAI) - else Stream(cast_to=ChatCompletionChunk, response=raw_response, client=client) - ) - if transport == "sdk" - else OpenAIChatCompletionStreamingHandler( - streaming_response=raw_response.aiter_lines() if async_stream else raw_response.iter_lines(), - sync_stream=not async_stream, - ) - ) - response: Final = CustomStreamWrapper( - completion_stream=stream, - model="test-model", - logging_obj=logging_obj, - custom_llm_provider="openai" if transport == "sdk" else "hosted_vllm", - ) + response: Final = _chat_stream_wrapper(transport, raw_response, client, logging_obj) try: chunks: Final = [chunk async for chunk in response] if async_stream else list(response) @@ -1624,6 +1631,80 @@ async def test_stream_with_null_choices_chunk_finishes_normally( assert [choice.finish_reason for choice in streamed_choices if choice.finish_reason] == ["stop"] +def _text_completion_stream_wrapper( + final_finish_reason: str, client: OpenAI | AsyncOpenAI, logging_obj: Logging +) -> CustomStreamWrapper: + body: Final = ( + "".join( + "data: " + + json.dumps( + { + "id": "cmpl-test", + "object": "text_completion", + "created": 1, + "model": "test-model", + "choices": [{"index": 0, "text": text, "logprobs": None, "finish_reason": finish_reason}], + } + ) + + "\n\n" + for text, finish_reason in (("Partial answer", None), ("", final_finish_reason)) + ) + + "data: [DONE]\n\n" + ) + raw_response: Final = httpx.Response( + 200, + content=body, + headers={"content-type": "text/event-stream"}, + request=httpx.Request("POST", "https://provider.example/v1/completions"), + ) + return CustomStreamWrapper( + completion_stream=( + AsyncStream(cast_to=Completion, response=raw_response, client=client) + if isinstance(client, AsyncOpenAI) + else Stream(cast_to=Completion, response=raw_response, client=client) + ), + model="test-model", + logging_obj=logging_obj, + custom_llm_provider="text-completion-openai", + ) + + +@pytest.mark.parametrize("async_stream", [False, True]) +async def test_aborted_text_completion_stream_raises(async_stream: bool, logging_obj: Logging) -> None: + client: Final = AsyncOpenAI(api_key="unused") if async_stream else OpenAI(api_key="unused") + response: Final = _text_completion_stream_wrapper("abort", client, logging_obj) + + try: + with pytest.raises(MidStreamFallbackError, match="abort") as excinfo: + _chunks: Final = [chunk async for chunk in response] if async_stream else list(response) + + assert excinfo.value.generated_content == "Partial answer" + assert getattr(excinfo.value.original_exception, "status_code", None) == 500 + finally: + if isinstance(client, AsyncOpenAI): + await client.close() + else: + client.close() + + +@pytest.mark.parametrize("async_stream", [False, True]) +async def test_stopped_text_completion_stream_finishes_normally(async_stream: bool, logging_obj: Logging) -> None: + client: Final = AsyncOpenAI(api_key="unused") if async_stream else OpenAI(api_key="unused") + response: Final = _text_completion_stream_wrapper("stop", client, logging_obj) + + try: + chunks: Final = [chunk async for chunk in response] if async_stream else list(response) + finally: + if isinstance(client, AsyncOpenAI): + await client.close() + else: + client.close() + + streamed_choices: Final = tuple(itertools.chain.from_iterable(chunk.choices for chunk in chunks)) + assert "".join(choice.delta.content or "" for choice in streamed_choices) == "Partial answer" + assert [choice.finish_reason for choice in streamed_choices if choice.finish_reason] == ["stop"] + + def _hosted_vllm_stream_wrapper(logging_obj: Logging, error_payload: dict) -> CustomStreamWrapper: """A CustomStreamWrapper over the real OpenAI-compatible line iterator, fed an HTTP 200 SSE body that carries an in-body error payload the way diff --git a/tests/unit/llms/base_llm/test_base_model_iterator.py b/tests/unit/llms/base_llm/test_base_model_iterator.py index 2971405885d..d7a6653e09b 100644 --- a/tests/unit/llms/base_llm/test_base_model_iterator.py +++ b/tests/unit/llms/base_llm/test_base_model_iterator.py @@ -428,6 +428,17 @@ def test_done_marker_ends_the_stream(line: str) -> None: ] +@pytest.mark.parametrize("choices", ["abort", 1, [None], [{"finish_reason": "stop"}, "abort"]]) +def test_chunk_with_unreadable_choices_reaches_the_provider_parser(choices: object) -> None: + line: Final = "data: " + json.dumps({"choices": choices}) + + chunks: Final = list(BaseModelResponseIterator(streaming_response=iter((line,)), sync_stream=True)) + + assert chunks == [ + GenericStreamingChunk(text="", is_finished=False, finish_reason="", usage=None, index=0, tool_use=None) + ] + + def test_error_frame_mentioning_done_raises_instead_of_ending_the_stream() -> None: line: Final = 'data: {"error": {"message": "prompt too long [DONE]", "code": 400}}' diff --git a/tests/unit/llms/bedrock/chat/test_invoke_handler.py b/tests/unit/llms/bedrock/chat/test_invoke_handler.py index 996cfec6e22..ab58c62ca76 100644 --- a/tests/unit/llms/bedrock/chat/test_invoke_handler.py +++ b/tests/unit/llms/bedrock/chat/test_invoke_handler.py @@ -966,6 +966,31 @@ async def test_moonshot_invoke_async_stream_yields_openai_shaped_chunks(_aws_tes _assert_moonshot_stream_content([chunk async for chunk in stream]) +def test_moonshot_invoke_stream_raises_when_upstream_aborts(_aws_test_credentials: None) -> None: + 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"), + ) + ) + client: Final = HTTPHandler( + client=httpx.Client(transport=httpx.MockTransport(lambda request: httpx.Response(200, content=raw_stream))) + ) + + stream: Final = litellm.completion( + model="bedrock/invoke/moonshot.kimi-k2-thinking", + messages=[{"role": "user", "content": "weather as json"}], + stream=True, + client=client, + ) + + assert isinstance(stream, CustomStreamWrapper) + with pytest.raises(MidStreamFallbackError, match="abort") as excinfo: + _chunks: Final = list(stream) + assert excinfo.value.generated_content == "Partial answer" + + def _truncated_frame() -> bytes: return _bedrock_event_stream_frame(_openai_stream_chunk({"role": "assistant"}))[:-8]