scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-11-14 09:03:16 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 214

fix(streaming): surface upstream abort on azure_ai, databricks, Vertex Model Garden and text completion streams

Browse source
This commit is contained in:
mateo-berri 2026-10-29 09:05:42 -07:00
parent 4212f1af27
commit fe92f14602
8 changed files with 201 additions and 53 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

20
litellm/litellm_core_utils/streaming_handler.py
View file

@ -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,

28
litellm/llms/base_llm/base_model_iterator.py
View file

@ -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(

1
litellm/llms/bedrock/chat/invoke_handler.py
View file

@ -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)

3
litellm/llms/databricks/streaming_utils.py
View file

@ -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(

9
litellm/llms/openai/chat/gpt_transformation.py
View file

@ -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]] = {

157
tests/unit/litellm_core_utils/test_streaming_handler.py
View file

@ -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

11
tests/unit/llms/base_llm/test_base_model_iterator.py
View file

@ -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}}'

25
tests/unit/llms/bedrock/chat/test_invoke_handler.py
View file

@ -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]