diff --git a/litellm/proxy/common_utils/responses_stream_errors.py b/litellm/proxy/common_utils/responses_stream_errors.py index ca948f65bd2..8dccd938a90 100644 --- a/litellm/proxy/common_utils/responses_stream_errors.py +++ b/litellm/proxy/common_utils/responses_stream_errors.py @@ -484,6 +484,11 @@ def _terminal_blocks(frame: str) -> tuple[tuple[tuple[str, ...], _StreamEvent], return tuple((lines, event) for lines, event in read if event.type in _TERMINAL_EVENT_TYPES) +def frame_payloads(frame: str) -> tuple[tuple[str, ...], bool]: + payloads: Final = tuple(_block_data(lines).strip() for lines in _blocks(frame) if _field_values(lines, "data")) + return tuple(takewhile(lambda payload: payload != "[DONE]", payloads)), payloads[:1] == ("[DONE]",) + + def frame_error_data(frame: str) -> str | None: terminal: Final = _terminal_blocks(frame) if not terminal: diff --git a/litellm/proxy/response_polling/background_streaming.py b/litellm/proxy/response_polling/background_streaming.py index 9d211fe4e71..73190ad8df8 100644 --- a/litellm/proxy/response_polling/background_streaming.py +++ b/litellm/proxy/response_polling/background_streaming.py @@ -12,7 +12,6 @@ import asyncio import json from collections.abc import AsyncIterable, Callable, Mapping, Sequence from http import HTTPStatus -from itertools import takewhile from time import perf_counter from typing import TYPE_CHECKING, Final, Protocol, TypeAlias, runtime_checkable @@ -21,12 +20,12 @@ from starlette.types import Message from typing_extensions import ReadOnly, TypedDict from litellm._logging import verbose_proxy_logger -from litellm.llms.custom_httpx.sse_lines import split_sse_lines from litellm.proxy.auth.user_api_key_auth import UserAPIKeyAuth from litellm.proxy.common_request_processing import RESPONSES_SSE_FRAMING, ProxyBaseLLMRequestProcessing from litellm.proxy.common_utils.responses_stream_errors import ( event_error_object, frame_error_data, + frame_payloads, responses_error_fields, sse_error_object, ) @@ -91,11 +90,6 @@ class _StreamsBody(Protocol): def body_iterator(self) -> AsyncIterable[str | bytes | memoryview]: ... -def _payloads_before_done(frame: str) -> tuple[tuple[str, ...], bool]: - payloads: Final = tuple(line[6:].strip() for line in split_sse_lines(frame) if line.startswith("data: ")) - return tuple(takewhile(lambda payload: payload != "[DONE]", payloads)), "[DONE]" in payloads - - def _terminal_response(event: _StreamEvent) -> _TerminalResponse: match event.get("response"): case Mapping() as response: @@ -289,7 +283,7 @@ async def background_streaming_task( ): frame_failure = responses_error_fields(_event_error(_StreamEventParser.parse(error_data))) - payloads, stream_done = _payloads_before_done(chunk) if isinstance(chunk, str) else ((), False) + payloads, stream_done = frame_payloads(chunk) if isinstance(chunk, str) else ((), False) for chunk_data in payloads: try: event: _StreamEvent = _StreamEventParser.parse(chunk_data) @@ -312,7 +306,7 @@ async def background_streaming_task( # Process different event types based on OpenAI streaming spec if event_type == "response.output_item.added": # New output item added - item = event.get("item", {}) + item = event.get("item", _OutputItem()) item_id = item.get("id") if item_id: output_items[item_id] = item @@ -379,7 +373,7 @@ async def background_streaming_task( elif event_type == "response.output_item.done": # Output item completed - use final item data - item = event.get("item", {}) + item = event.get("item", _OutputItem()) item_id = item.get("id") if item_id: output_items[item_id] = item @@ -443,10 +437,12 @@ async def background_streaming_task( continue state_dirty = True - except (json.JSONDecodeError, AttributeError, TypeError) as e: + except (ValueError, AttributeError, TypeError) as e: verbose_proxy_logger.warning("Failed to parse streaming chunk: %s", e) else: await flush_state_if_needed() + if event_type in _event_to_status: + break if stream_done: break diff --git a/tests/unit/proxy/response_polling/test_response_polling_handler.py b/tests/unit/proxy/response_polling/test_response_polling_handler.py index 43184f95700..752dfc5f124 100644 --- a/tests/unit/proxy/response_polling/test_response_polling_handler.py +++ b/tests/unit/proxy/response_polling/test_response_polling_handler.py @@ -1701,6 +1701,41 @@ class TestBackgroundStreamingTerminalEvents: None, {"input_tokens": 3}, ), + ( + ( + 'data: {"type": "response.output_text.delta", "delta": "hi"}\n\n' + 'data: {"type": "hook.huge", "n": ' + "9" * 4400 + "}\n\n", + 'data: {"type": "response.completed", "response": {"status": "completed", "usage": {"input_tokens": 3}}}\n\n', + ), + "completed", + None, + {"input_tokens": 3}, + ), + ( + ( + 'data: {"type": "response.output_text.delta", "delta": "hi"}\n\ndata: [DONE]\n\n', + 'data: {"type": "response.completed", "response": {"status": "completed", "usage": {"input_tokens": 3}}}\n\n', + ), + "completed", + None, + {"input_tokens": 3}, + ), + ( + ( + 'data: {"type": "response.completed", "response": {"status": "completed", "usage": {"input_tokens": 3}}}\n\n' + 'data: {"type": "response.failed", "response": {"status": "failed", ' + '"error": {"code": "server_error", "message": "late"}}}\n\n', + ), + "completed", + None, + {"input_tokens": 3}, + ), + ( + ('data: event: response.incomplete\ndata: {"type": "response.incomplete", "response": {"status": "incomplete"}}\n\n',), + "completed", + None, + None, + ), ], ids=[ "error_frame_after_a_delta", @@ -1714,6 +1749,10 @@ class TestBackgroundStreamingTerminalEvents: "failed_event_whose_response_is_null", "error_frame_and_incomplete_in_one_chunk", "error_frame_then_completed_behind_a_delta_in_one_chunk", + "frame_past_the_int_digit_limit_behind_a_delta_then_completed", + "done_behind_a_delta_in_one_chunk_then_completed", + "completed_then_failed_in_one_chunk", + "event_line_sent_as_a_data_line_of_one_event", ], ) async def test_a_poll_stores_what_its_stream_ended_with(