diff --git a/litellm/llms/openai/responses/guardrail_translation/handler.py b/litellm/llms/openai/responses/guardrail_translation/handler.py index 20b7506ae0a..f81957401ff 100644 --- a/litellm/llms/openai/responses/guardrail_translation/handler.py +++ b/litellm/llms/openai/responses/guardrail_translation/handler.py @@ -276,6 +276,7 @@ _ITEM_ID_KEYED_TOOL_CALL_TYPES: Final = frozenset({_TOOL_SEARCH_CALL_TYPE, _WEB_ _OUTPUT_TOOL_CALL_ITEM_TYPES: Final = _TOOL_CALL_ITEM_TYPES | {_TOOL_SEARCH_CALL_TYPE} _SCANNED_TOOL_SEARCH_NAME: Final = "tool_search" _TOOL_SEARCH_ARGUMENTS_ADAPTER: Final = TypeAdapter(dict[str, object]) +_STREAM_EVENT_TYPE_ADAPTER: Final = TypeAdapter(str) _TOOL_CALL_PAYLOAD_FIELDS: Final[Mapping[str, str]] = MappingProxyType( {"function_call": "arguments", "custom_tool_call": "input"} ) @@ -1991,7 +1992,6 @@ class OpenAIResponsesHandler(BaseTranslation): new output item under the same response id, and close with a ``response.completed`` carrying only the replacement item. - The proxy's data generator appends ``data: [DONE]`` itself. """ events: Final = ( self._block_continuation_events(exc, responses_so_far or ()) @@ -1999,6 +1999,7 @@ class OpenAIResponsesHandler(BaseTranslation): else self._standalone_block_events(exc) ) return tuple( + f"event: {_STREAM_EVENT_TYPE_ADAPTER.validate_python(stream_item_field(event, 'type'))}\n" f"data: {event.model_dump_json(exclude_none=True, exclude_unset=True, serialize_as_any=True)}\n\n".encode() for event in events ) diff --git a/litellm/proxy/common_request_processing.py b/litellm/proxy/common_request_processing.py index 4d7c67972f2..55f2ee94d0e 100644 --- a/litellm/proxy/common_request_processing.py +++ b/litellm/proxy/common_request_processing.py @@ -7,6 +7,7 @@ from collections.abc import AsyncGenerator, AsyncIterator, Awaitable, Callable, from dataclasses import dataclass from datetime import datetime, timezone from functools import lru_cache +from http import HTTPStatus from itertools import chain from types import MappingProxyType from typing import ( @@ -52,7 +53,7 @@ from litellm.constants import ( STREAM_SSE_KEEPALIVE_PING_BYTES, UNSAFE_PROXY_RESPONSE_HEADERS, ) -from litellm.integrations.custom_guardrail import CustomGuardrail +from litellm.integrations.custom_guardrail import CustomGuardrail, ModifyResponseException from litellm.integrations.otel.runtime import phase_event from litellm.litellm_core_utils.add_retry_fallback_headers import get_hidden_params_dict from litellm.litellm_core_utils.agentic_followup_failure import is_agentic_followup_failure @@ -97,6 +98,7 @@ from litellm.litellm_core_utils.streaming_handler import ( backfill_missing_cache_usage_fields, ) from litellm.litellm_core_utils.tool_fields import DropPolicy, MalformedField, drop_policy, resolve_tool_fields +from litellm.llms.openai.responses.guardrail_translation.handler import OpenAIResponsesHandler from litellm.proxy._types import LiteLLMRoutes, ProxyErrorTypes, ProxyException, UserAPIKeyAuth from litellm.proxy.anthropic_endpoints.answered_status import is_provider_overload, with_openai_overload_status from litellm.proxy.auth.auth_checks import ( @@ -126,6 +128,11 @@ from litellm.proxy.common_utils.openai_error_payload import ( openai_error_type, relayed_provider_specific_fields, ) +from litellm.proxy.common_utils.responses_stream_errors import ( + SSE_LINE_BREAK, + ResponsesStreamErrorState, + sse_error_object, +) from litellm.proxy.common_utils.sse_keepalive import ( SSE_COMMENT_PING_BYTES, SSE_STREAM_START_TAIL, @@ -1024,11 +1031,32 @@ async def _enforce_guardrail_added_tag_budgets( ) from e -async def _parse_event_data_for_error(event_line: str | bytes) -> int | None: +SseDataPayload: TypeAlias = Callable[[str | bytes], str | None] + + +def _sse_text(event_line: str | bytes) -> str: + return event_line.decode("utf-8") if isinstance(event_line, bytes) else event_line + + +def _sse_data_payload(event_line: str | bytes) -> str | None: + text: Final = _sse_text(event_line) + return text.removeprefix("data: ").strip() if text.startswith("data: ") else None + + +def _responses_event_data_payload(event_line: str | bytes) -> str | None: + match SSE_LINE_BREAK.split(_sse_text(event_line).strip()): + case ["event: error", data_line] if data_line.startswith("data: "): + return data_line.removeprefix("data: ").strip() + case _: + return _sse_data_payload(event_line) + + +async def _parse_event_data_for_error( + event_line: str | bytes, data_payload: SseDataPayload = _sse_data_payload +) -> int | None: """Parses an event line and returns an error code if present, else None.""" - event_line = event_line.decode("utf-8") if isinstance(event_line, bytes) else event_line - if event_line.startswith("data: "): - json_str: Final = event_line[len("data: ") :].strip() + json_str: Final = data_payload(event_line) + if json_str is not None: if not json_str or json_str == "[DONE]": # handle empty data or [DONE] message return None try: @@ -1059,18 +1087,17 @@ async def _parse_event_data_for_error(event_line: str | bytes) -> int | None: return None -def _extract_error_from_sse_chunk(event_line: str | bytes) -> dict: +def _extract_error_from_sse_chunk(event_line: str | bytes, data_payload: SseDataPayload = _sse_data_payload) -> dict: """ Extract error dictionary from SSE format chunk. Args: event_line: SSE format event line, e.g. "data: {"error": {...}}\n\n" + data_payload: Reads the JSON payload out of the event line Returns: Error dictionary in OpenAI API format """ - event_line = event_line.decode("utf-8") if isinstance(event_line, bytes) else event_line - # Default error format default_error: Final = { "message": "Unknown error", @@ -1079,8 +1106,8 @@ def _extract_error_from_sse_chunk(event_line: str | bytes) -> dict: "code": "500", } - if event_line.startswith("data: "): - json_str: Final = event_line[len("data: ") :].strip() + json_str: Final = data_payload(event_line) + if json_str is not None: if not json_str or json_str == "[DONE]": return default_error @@ -1288,11 +1315,27 @@ def _sse_error_frames(error_obj: Mapping[str, object]) -> tuple[str, str]: return f"data: {json.dumps({'error': error_obj}, default=str)}\n\n", "data: [DONE]\n\n" +def _responses_error_frames(error_obj: Mapping[str, object]) -> tuple[str, ...]: + return (ResponsesStreamErrorState().format_error(error_obj),) + + +@dataclass(frozen=True, slots=True) +class LateSseFraming: + """How a response the keepalive already opened ends once its upstream fails late.""" + + failure: Callable[[Exception], tuple[str, ...]] + body: Callable[[Response], tuple[bytes, ...]] + + def _openai_late_failure_frames(exc: Exception) -> tuple[str, ...]: _, error_obj = sse_error_payload(exc) return _sse_error_frames(error_obj) +def _openai_late_body_frames(produced: Response) -> tuple[bytes, ...]: + return b"data: " + (bytes(produced.body) or b"{}") + b"\n\n", b"data: [DONE]\n\n" + + def _anthropic_late_failure_frames(exc: Exception) -> tuple[str, ...]: answered: Final = proxy_exception_from_route_error(with_answered_status(exc)) return ( @@ -1303,6 +1346,57 @@ def _anthropic_late_failure_frames(exc: Exception) -> tuple[str, ...]: ) +def _responses_late_failure_frames(exc: Exception) -> tuple[str, ...]: + if isinstance(exc, ModifyResponseException): + return tuple(chunk.decode() for chunk in OpenAIResponsesHandler().build_block_sse_chunks(exc)) + return _responses_error_frames(sse_error_payload(exc)[1]) + + +def _responses_late_body_frames(produced: Response) -> tuple[bytes, ...]: + body: Final = bytes(produced.body) + error_obj: Final = sse_error_object(body or b"{}") or {"message": body.decode()} + with_status: Final = {**error_obj, "code": error_obj.get("code") or str(produced.status_code)} + return tuple(frame.encode() for frame in _responses_error_frames(with_status)) + + +OPENAI_LATE_SSE_FRAMING: Final = LateSseFraming(failure=_openai_late_failure_frames, body=_openai_late_body_frames) +ANTHROPIC_LATE_SSE_FRAMING: Final = LateSseFraming( + failure=_anthropic_late_failure_frames, body=_openai_late_body_frames +) +RESPONSES_LATE_SSE_FRAMING: Final = LateSseFraming( + failure=_responses_late_failure_frames, body=_responses_late_body_frames +) + + +@dataclass(frozen=True, slots=True) +class FirstChunkSseFraming: + """How a stream's first chunk is read for an error, and how the stream ends when reading it fails.""" + + data_payload: SseDataPayload + error_frames: Callable[[Mapping[str, object]], tuple[str, ...]] + lowest_error_status: int + + +OPENAI_FIRST_CHUNK_FRAMING: Final = FirstChunkSseFraming( + data_payload=_sse_data_payload, error_frames=_sse_error_frames, lowest_error_status=HTTPStatus.CONTINUE +) +RESPONSES_FIRST_CHUNK_FRAMING: Final = FirstChunkSseFraming( + data_payload=_responses_event_data_payload, + error_frames=_responses_error_frames, + lowest_error_status=HTTPStatus.BAD_REQUEST, +) + + +@dataclass(frozen=True, slots=True) +class SseFraming: + first_chunk: FirstChunkSseFraming + late: LateSseFraming + + +OPENAI_SSE_FRAMING: Final = SseFraming(first_chunk=OPENAI_FIRST_CHUNK_FRAMING, late=OPENAI_LATE_SSE_FRAMING) +RESPONSES_SSE_FRAMING: Final = SseFraming(first_chunk=RESPONSES_FIRST_CHUNK_FRAMING, late=RESPONSES_LATE_SSE_FRAMING) + + def _sse_stream_headers(headers: Mapping[str, str]) -> Mapping[str, str]: """`headers` plus the two that stop reverse proxies from buffering SSE (issue #28384).""" return MappingProxyType({**headers, **_TTFT_KEEPALIVE_HEADERS}) @@ -1331,6 +1425,7 @@ async def create_response( request: Request | None = None, refresh_headers: Callable[[], Awaitable[Mapping[str, str]]] | None = None, general_settings: Mapping[str, object] = _NO_GENERAL_SETTINGS, + framing: FirstChunkSseFraming = OPENAI_FIRST_CHUNK_FRAMING, ) -> StreamingResponse | JSONResponse: """ Create streaming response, checking if the first chunk is an error. @@ -1366,8 +1461,10 @@ async def create_response( if first_chunk_value is not None: try: - error_code_from_chunk: Final = await _parse_event_data_for_error(first_chunk_value) - if error_code_from_chunk is not None: + error_code_from_chunk: Final = await _parse_event_data_for_error( + first_chunk_value, framing.data_payload + ) + if error_code_from_chunk is not None and error_code_from_chunk >= framing.lowest_error_status: # First chunk is an error, stream hasn't really started yet # Should return standard JSON error response instead of SSE format final_status_code = error_code_from_chunk @@ -1378,7 +1475,9 @@ async def create_response( # Parse error content error_dict: Final = with_call_id( - JSON_OBJECT.validate_python(_extract_error_from_sse_chunk(first_chunk_value)), + JSON_OBJECT.validate_python( + _extract_error_from_sse_chunk(first_chunk_value, framing.data_payload) + ), error_body_call_id(general_settings, resolved_headers.get(LITELLM_CALL_ID_HEADER)), ) @@ -1420,7 +1519,7 @@ async def create_response( error_status, error_obj = sse_error_payload(e) async def error_gen_message() -> AsyncGenerator[str, None]: - for frame in _sse_error_frames(error_obj): + for frame in framing.error_frames(error_obj): yield frame return StreamingResponse( @@ -1503,13 +1602,13 @@ async def _aclose_late_response(produced: Response) -> None: verbose_proxy_logger.debug("error closing relayed streaming generator: %s", exc) -async def _relay_late_response(produced: Response) -> AsyncGenerator[bytes, None]: +async def _relay_late_response(produced: Response, framing: LateSseFraming) -> AsyncGenerator[bytes, None]: """Replay a Response that was built after a keepalive had already opened the wire.""" if not isinstance(produced, StreamingResponse): # The status line is already on the wire, so a non-streaming body, an error # body included, can only reach the client as an SSE frame. - yield b"data: " + (bytes(produced.body) or b"{}") + b"\n\n" - yield b"data: [DONE]\n\n" + for frame in framing.body(produced): + yield frame return try: @@ -1551,7 +1650,7 @@ async def open_sse_before_first_byte( ping_interval_seconds: float | str | None, media_type: str = "text/event-stream", on_late_failure: Callable[[Exception], Awaitable[HTTPException | None]] | None = None, - late_failure_frames: Callable[[Exception], tuple[str, ...]] = _openai_late_failure_frames, + late_framing: LateSseFraming = OPENAI_LATE_SSE_FRAMING, ) -> _LateResponseT | StreamingResponse: """Write SSE keepalive comments while the upstream LLM call is still in flight. @@ -1597,10 +1696,10 @@ async def open_sse_before_first_byte( # would never fire and the failure would go unaudited. The hook # also gets to sanitize what reaches the client, by returning or # raising a replacement, so its answer decides the frame. - for frame in late_failure_frames(await _sanitized_late_failure(exc, on_late_failure)): + for frame in late_framing.failure(await _sanitized_late_failure(exc, on_late_failure)): yield frame.encode() return - async for chunk in _relay_late_response(produced): + async for chunk in _relay_late_response(produced, late_framing): yield chunk finally: if not produce_task.done(): @@ -3005,6 +3104,7 @@ class ProxyBaseLLMRequestProcessing: is_streaming_request: bool | None = False, contents: list[object] | None = None, skip_pre_call_logic: bool = False, + sse_framing: SseFraming = OPENAI_SSE_FRAMING, ) -> object: """Run the request, sending SSE keepalives while the upstream is still silent. @@ -3046,10 +3146,11 @@ class ProxyBaseLLMRequestProcessing: is_streaming_request=is_streaming_request, contents=contents, skip_pre_call_logic=skip_pre_call_logic, + first_chunk_framing=sse_framing.first_chunk, ), ping_interval_seconds=ttft_keepalive_interval(self.data, llm_router), on_late_failure=_audit_late_failure, - late_failure_frames=_anthropic_late_failure_frames if answers_as_anthropic else _openai_late_failure_frames, + late_framing=ANTHROPIC_LATE_SSE_FRAMING if answers_as_anthropic else sse_framing.late, ) async def _process_llm_request( @@ -3073,6 +3174,7 @@ class ProxyBaseLLMRequestProcessing: is_streaming_request: bool | None = False, contents: list[object] | None = None, # Add contents parameter skip_pre_call_logic: bool = False, + first_chunk_framing: FirstChunkSseFraming = OPENAI_FIRST_CHUNK_FRAMING, ) -> Any: """ Common request processing logic for both chat completions and responses API endpoints @@ -3356,6 +3458,7 @@ class ProxyBaseLLMRequestProcessing: request=request, refresh_headers=refresh_stream_headers, general_settings=general_settings, + framing=first_chunk_framing, ) ### CALL HOOKS ### - modify outgoing data @@ -3595,8 +3698,8 @@ class ProxyBaseLLMRequestProcessing: Streaming ``/v1/responses`` short-circuits out of ``base_process_llm_request`` before the non-streaming ownership - tail runs. The OpenAI SDK closes the connection at ``data: [DONE]`` - and starlette cancels the body task on disconnect, so a write that + tail runs. The OpenAI SDK can close the connection right after the terminal + event and starlette cancels the body task on disconnect, so a write that waits for the generator to finish never lands. The iterator sets ``completed_response`` before it hands over its terminal chunk, so the ``LiteLLM_ManagedObjectTable`` row is written the moment it diff --git a/litellm/proxy/common_utils/responses_stream_errors.py b/litellm/proxy/common_utils/responses_stream_errors.py index 3409b4b6657..5582e985d4b 100644 --- a/litellm/proxy/common_utils/responses_stream_errors.py +++ b/litellm/proxy/common_utils/responses_stream_errors.py @@ -1,10 +1,17 @@ +import json +import re import time -from collections.abc import Mapping +from collections.abc import Iterator, Mapping, Sequence +from contextlib import suppress +from dataclasses import dataclass +from enum import Enum from http import HTTPStatus +from itertools import accumulate, chain, groupby, islice, repeat, takewhile from types import MappingProxyType from typing import Final -from pydantic import BaseModel, ConfigDict, field_validator +from pydantic import BaseModel, ConfigDict, TypeAdapter, ValidationError, field_validator +from typing_extensions import TypeIs # noqa: TID251 # the guards narrow hook values to object items, isinstance cannot from litellm._logging import redact_internal_details_from_client_message from litellm._uuid import uuid @@ -29,6 +36,13 @@ class _StreamEvent(LiteLLMBaseModel): response: _ResponseIdentity | None = None +class _FramedEvent(_StreamEvent): + error: object = None + + +_LAX_INT: Final = TypeAdapter(int) + + class _FailureDetails(LiteLLMBaseModel): model_config = ConfigDict(frozen=True) @@ -37,6 +51,16 @@ class _FailureDetails(LiteLLMBaseModel): type: str | None = None status_code: int | None = None + @field_validator("status_code", mode="before") + @classmethod + def normalize_status_code(cls, value: object) -> int | None: + if isinstance(value, bool): + return None + try: + return _LAX_INT.validate_python(value) + except ValidationError: + return None + @field_validator("message", mode="before") @classmethod def normalize_message(cls, value: object) -> str | None: @@ -57,6 +81,91 @@ class _FailureDetails(LiteLLMBaseModel): return cls.model_validate({field: getattr(exception, field, None) for field in cls.model_fields}) +class _ErrorPayload(LiteLLMBaseModel): + model_config = ConfigDict(frozen=True) + + error: Mapping[str, object] | str + + def error_object(self) -> Mapping[str, object]: + return {"message": self.error} if isinstance(self.error, str) else self.error + + +_JSON_LEAF_TYPES: Final = frozenset((str, int, float, bool, type(None))) +_MAX_PLAIN_JSON_DEPTH: Final = 100 + + +def _is_dict(value: object) -> TypeIs[Mapping[object, object]]: # guard-ok: dict items as object + return isinstance(value, dict) + + +def _is_sequence(value: object) -> TypeIs[Sequence[object]]: # guard-ok: items as object + return isinstance(value, (list, tuple)) + + +def _is_plain_node(value: object) -> bool: + kind: Final = type(value) + if kind in _JSON_LEAF_TYPES or kind in (list, tuple): + return True + if isinstance(value, Enum): + return isinstance(value, (str, int)) and kind.__repr__ is Enum.__repr__ + return kind is dict and _is_dict(value) and all(type(key) is str for key in value) + + +def _children(value: object) -> tuple[object, ...]: + if _is_dict(value): + return tuple(value.values()) + return tuple(value) if _is_sequence(value) else () + + +@dataclass(frozen=True, slots=True) +class _PlainJsonLevel: + values: tuple[object, ...] + seen: frozenset[int] + + +def _next_level(level: _PlainJsonLevel, _: None) -> _PlainJsonLevel: + unseen: Final = {id(value): value for value in level.values if id(value) not in level.seen} + return _PlainJsonLevel( + values=tuple(chain.from_iterable(map(_children, unseen.values()))), + seen=level.seen.union(unseen), + ) + + +def _settles(level: _PlainJsonLevel) -> bool: + return not level.values or not all(map(_is_plain_node, level.values)) + + +def _is_plain_json(chunk: object) -> bool: + first: Final = _PlainJsonLevel(values=(chunk,), seen=frozenset()) + levels: Final = islice(accumulate(repeat(None), _next_level, initial=first), _MAX_PLAIN_JSON_DEPTH) + return next((not level.values for level in levels if _settles(level)), False) + + +def _carries_a_type(chunk: object) -> bool: + return _is_dict(chunk) and type(chunk) is dict and "type" in chunk + + +def responses_chunk_data(chunk: object) -> str | None: + try: + return json.dumps(chunk, allow_nan=False) if _carries_a_type(chunk) and _is_plain_json(chunk) else None + except (ValueError, RecursionError): + return None + + +def sse_error_object(data: str | bytes) -> Mapping[str, object] | None: + try: + return _ErrorPayload.model_validate_json(data).error_object() + except ValidationError: + return None + + +def event_error_object(event: Mapping[str, object]) -> Mapping[str, object] | None: + try: + return _ErrorPayload.model_validate(event).error_object() or None + except ValidationError: + return None + + def _original_failure(exception: Exception) -> Exception: current = exception # rebind-ok: the recursion gate requires iterative wrapper traversal while isinstance(current, MidStreamFallbackError) and current.original_exception is not None: @@ -95,6 +204,16 @@ def _status_error_code(status_code: int | None) -> str: return _CLIENT_ERROR_CODES.get(status_code, "invalid_request_error") +def _status_of(code: str | int | None) -> int | None: + if isinstance(code, int): + return code + if code is None or not code.isdecimal(): + return None + with suppress(ValueError): + return int(code) + return None + + def _response_error_code(details: _FailureDetails) -> str: for value in (details.code, details.type): if value == "insufficient_quota": @@ -106,6 +225,120 @@ def _response_error_code(details: _FailureDetails) -> str: return _status_error_code(details.status_code) +_STATUS_PHRASES: Final = MappingProxyType({int(status): status.phrase for status in HTTPStatus}) + + +def _status_phrase(status_code: int | None) -> str: + return _STATUS_PHRASES.get(status_code, "Unknown error") if status_code is not None else "Unknown error" + + +def responses_error_fields(error: Mapping[str, object]) -> Mapping[str, object]: + details: Final = _FailureDetails.model_validate(error) + status_code: Final = _status_of(details.code) + return { + "code": _response_error_code(details.model_copy(update={"status_code": status_code})), + "message": details.message or _status_phrase(status_code), + } + + +def _error_event( + fields: Mapping[str, object], error: Mapping[str, object], sequence_number: int +) -> Mapping[str, object]: + return { + "type": ResponsesAPIStreamEvents.ERROR.value, + **fields, + "param": error.get("param"), + "sequence_number": sequence_number, + "error": error, + } + + +SSE_LINE_BREAK: Final = re.compile(r"\r\n|\r|\n") +_TERMINAL_EVENT_TYPES: Final = frozenset(("response.completed", "response.failed", "response.incomplete", "error")) +_MAX_SEQUENCE_NUMBER: Final = 2**63 - 1 + + +def _field_values(lines: Sequence[str], name: str) -> tuple[str, ...]: + prefix: Final = f"{name}:" + return tuple(line.removeprefix(prefix).removeprefix(" ") for line in lines if line.startswith(prefix)) + + +def _parsed_event(data: str) -> _FramedEvent: + try: + return _FramedEvent.model_validate_json(data) + except ValidationError: + return _FramedEvent() + + +def _in_range_sequence(event: _FramedEvent) -> _FramedEvent: + sequence_number: Final = event.sequence_number + if sequence_number is None or 0 <= sequence_number < _MAX_SEQUENCE_NUMBER: + return event + return event.model_copy(update={"sequence_number": None}) + + +def _names_an_error_object(event: _FramedEvent, lines: Sequence[str]) -> bool: + return _is_dict(event.error) and bool(event.error) and _field_values(lines, "event")[-1:] == ("error",) + + +def _block_data(lines: Sequence[str]) -> str: + return "\n".join(_field_values(lines, "data")) + + +def _block_event(lines: Sequence[str]) -> _StreamEvent: + parsed: Final = _in_range_sequence(_parsed_event(_block_data(lines))) + if parsed.type in _TERMINAL_EVENT_TYPES or not _names_an_error_object(parsed, lines): + return parsed + return parsed.model_copy(update={"type": "error"}) + + +def _blocks(frames: str) -> Iterator[tuple[str, ...]]: + return (tuple(lines) for filled, lines in groupby(SSE_LINE_BREAK.split(frames), key=bool) if filled) + + +def _frame_events(frames: str) -> tuple[_StreamEvent, ...]: + return tuple(_block_event(lines) for lines in _blocks(frames)) + + +def _before_done(lines: Sequence[str]) -> bool: + return _block_data(lines).strip() != "[DONE]" + + +def _terminal_blocks(frame: str) -> tuple[tuple[tuple[str, ...], _StreamEvent], ...]: + read: Final = ((lines, _block_event(lines)) for lines in takewhile(_before_done, _blocks(frame))) + 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: + return None + lines, event = terminal[-1] + return _block_data(lines) if event.type == "error" else None + + +def _observed_event(chunk: object) -> _StreamEvent | None: + if not isinstance(chunk, (BaseModel, Mapping)): + return None + try: + return _StreamEvent.model_validate(chunk) + except ValidationError: + return None + + +def _wire_event(frame: str | bytes) -> _StreamEvent | None: + text: Final = frame.decode("utf-8", errors="replace") if isinstance(frame, bytes) else frame + try: + return _StreamEvent.model_validate_json(text.removeprefix("data: ")) + except ValidationError: + return None + + class ResponsesStreamErrorState: def __init__(self) -> None: self.response_id: str | None = None @@ -116,12 +349,9 @@ class ResponsesStreamErrorState: self._pending_event: _StreamEvent | None = None def observe_chunk(self, chunk: object) -> None: - self._pending_event = _StreamEvent.model_validate(chunk) if isinstance(chunk, (BaseModel, Mapping)) else None + self._pending_event = _observed_event(chunk) - def mark_emitted(self, frame: str | bytes) -> str | bytes: - event: Final = self._pending_event - if event is None: - return frame + def _record(self, event: _StreamEvent) -> None: if event.sequence_number is not None: self.sequence_number = max(self.sequence_number, event.sequence_number) if event.response is not None: @@ -129,9 +359,34 @@ class ResponsesStreamErrorState: self.model = event.response.model or self.model if event.response.created_at is not None: self.created_at = event.response.created_at - if event.type in ("response.completed", "response.failed", "response.incomplete"): + if event.type in _TERMINAL_EVENT_TYPES: self.terminal_emitted = True - return frame + + def observe_frames(self, frames: str) -> str: + for event in _frame_events(frames): + self._record(event) + return frames + + def mark_emitted(self, frame: str | bytes) -> str | bytes: + observed: Final = self._pending_event + if observed is None: + return frame + event: Final = _wire_event(frame) + self._record(observed if event is None else event) + if event is None or not event.type or "\n" in event.type or "\r" in event.type: + return frame + prefix: Final = f"event: {event.type}\n" + return prefix.encode("utf-8") + frame if isinstance(frame, bytes) else prefix + frame + + def format_error(self, error: Mapping[str, object]) -> str: + fields: Final = responses_error_fields(error) + sequence_number: Final = self.sequence_number + 1 + payload: Final = json.dumps(_error_event(fields, error, sequence_number), default=str) + return f"event: error\ndata: {payload}\n\n" + + def format_error_message(self, frame: str) -> str: + error: Final = sse_error_object(frame.removeprefix("data: ")) + return frame if error is None else self.format_error(error) def format_failure(self, exception: Exception) -> str | None: if self.terminal_emitted: diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 848bc246404..fe633238e20 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -601,7 +601,7 @@ from litellm.proxy.common_utils.periodic_reload_schedule import ( from litellm.proxy.common_utils.proxy_state import ProxyState from litellm.proxy.common_utils.registry_read_through import resync_model_deployments_by_ids from litellm.proxy.common_utils.reset_budget_job import ResetBudgetJob -from litellm.proxy.common_utils.responses_stream_errors import ResponsesStreamErrorState +from litellm.proxy.common_utils.responses_stream_errors import ResponsesStreamErrorState, responses_chunk_data from litellm.proxy.common_utils.scheduled_job_stagger import ( apply_scheduled_job_stagger, attach_job_timing_logger, @@ -10550,6 +10550,15 @@ def _format_streaming_sse_chunk(chunk: object) -> str | bytes: return f"data: {chunk}\n\n" +def _format_responses_sse_chunk(chunk: object) -> str | bytes: + data: Final = responses_chunk_data(chunk) + return _format_streaming_sse_chunk(chunk=chunk) if data is None else f"data: {data}\n\n" + + +def _observed_frames(frames: str, error_state: ResponsesStreamErrorState | None) -> str: + return frames if error_state is None else error_state.observe_frames(frames) + + _SSE_FRAME_DELIMITERS: Final = ("\r\n\r\n", "\n\n", "\r\r") _OPENAI_STREAM_DONE_FRAME: Final = "data: [DONE]\n\n" _MAX_RAW_SSE_BUFFER_CHARS: Final = 8 * 1024 * 1024 @@ -10809,7 +10818,9 @@ async def async_data_generator( stream_completed = False client_disconnected = False error_state: Final = ResponsesStreamErrorState() if responses_stream_errors else None - ends_with_openai_done: Final = not request_data.get("_litellm_skip_openai_stream_done") + ends_with_openai_done: Final = not responses_stream_errors and not request_data.get( + "_litellm_skip_openai_stream_done" + ) needs_iterator_wrap: Final = proxy_logging_obj.needs_iterator_wrap() stream_iterator: Final = _object_stream( proxy_logging_obj.async_post_call_streaming_iterator_hook( @@ -10944,12 +10955,14 @@ async def async_data_generator( frame, raw_sse_buffer = _pop_complete_sse_frame(raw_sse_buffer) if frame is None: break - yield frame + yield _observed_frames(frame, error_state) if len(raw_sse_buffer) > _MAX_RAW_SSE_BUFFER_CHARS: raise ValueError("Raw SSE stream exceeded maximum buffered size without a frame delimiter") raw_passthrough = True elif chunk.startswith(("data:", "event:", ":")): - yield (chunk if chunk.endswith(_SSE_FRAME_DELIMITERS) else chunk + "\n\n") + yield _observed_frames( + chunk if chunk.endswith(_SSE_FRAME_DELIMITERS) else chunk + "\n\n", error_state + ) raw_passthrough = True elif isinstance(chunk, str) and is_raw_sse_stream: raw_sse_buffer += chunk @@ -10957,7 +10970,7 @@ async def async_data_generator( frame, raw_sse_buffer = _pop_complete_sse_frame(raw_sse_buffer) if frame is None: break - yield frame + yield _observed_frames(frame, error_state) if len(raw_sse_buffer) > _MAX_RAW_SSE_BUFFER_CHARS: raise ValueError("Raw SSE stream exceeded maximum buffered size without a frame delimiter") raw_passthrough = True @@ -10968,7 +10981,7 @@ async def async_data_generator( if not raw_passthrough: try: if error_state is not None: - yield error_state.mark_emitted(_format_streaming_sse_chunk(chunk=chunk)) + yield error_state.mark_emitted(_format_responses_sse_chunk(chunk=chunk)) else: yield _format_streaming_sse_chunk(chunk=chunk) except Exception as e: @@ -10991,11 +11004,14 @@ async def async_data_generator( ProxyLogging.fire_deferred_stream_logging(request_data) if raw_sse_buffer: - yield (raw_sse_buffer if raw_sse_buffer.endswith(_SSE_FRAME_DELIMITERS) else raw_sse_buffer + "\n\n") + yield _observed_frames( + raw_sse_buffer if raw_sse_buffer.endswith(_SSE_FRAME_DELIMITERS) else raw_sse_buffer + "\n\n", + error_state, + ) if error_message is not None: - yield error_message - # OpenAI-compatible streams terminate with data: [DONE]; Google GenAI (?alt=sse) does not. + yield error_message if error_state is None else error_state.format_error_message(error_message) + # Chat Completions streams terminate with data: [DONE]; Responses and Google GenAI (?alt=sse) streams do not. if ends_with_openai_done: yield _OPENAI_STREAM_DONE_FRAME except (asyncio.CancelledError, GeneratorExit): diff --git a/litellm/proxy/response_api_endpoints/endpoints.py b/litellm/proxy/response_api_endpoints/endpoints.py index c4cd7c2d411..50d1c7c5b59 100644 --- a/litellm/proxy/response_api_endpoints/endpoints.py +++ b/litellm/proxy/response_api_endpoints/endpoints.py @@ -61,6 +61,7 @@ from litellm.proxy.auth.user_api_key_auth import ( user_api_key_auth_websocket, ) from litellm.proxy.common_request_processing import ( + RESPONSES_SSE_FRAMING, ProxyBaseLLMRequestProcessing, cost_headers, create_response, @@ -483,6 +484,7 @@ async def responses_api( if requests_background(_PROXY_SERVER_REQUEST_ADAPTER.validate_python(data)) else native_data_generator ), + sse_framing=RESPONSES_SSE_FRAMING, model=None, user_model=user_model, user_temperature=user_temperature, @@ -541,7 +543,6 @@ async def responses_api( async def _blocked_stream() -> AsyncGenerator[str, None]: for chunk in block_chunks: yield chunk.decode() - yield "data: [DONE]\n\n" return await create_response(generator=_blocked_stream(), media_type="text/event-stream", headers={}) return build_blocked_response(e) @@ -908,6 +909,7 @@ async def get_response( response_id, read_settles_billing, ), + sse_framing=RESPONSES_SSE_FRAMING, model=None, user_model=user_model, user_temperature=user_temperature, diff --git a/litellm/proxy/response_polling/background_streaming.py b/litellm/proxy/response_polling/background_streaming.py index 574152048d4..76f713ae963 100644 --- a/litellm/proxy/response_polling/background_streaming.py +++ b/litellm/proxy/response_polling/background_streaming.py @@ -11,6 +11,7 @@ https://platform.openai.com/docs/api-reference/responses-streaming import asyncio import json from collections.abc import AsyncIterable, Callable, Mapping, Sequence +from http import HTTPStatus from time import perf_counter from typing import TYPE_CHECKING, Final, Protocol, TypeAlias, runtime_checkable @@ -19,9 +20,15 @@ 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 ProxyBaseLLMRequestProcessing +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, +) from litellm.proxy.response_polling.polling_handler import ResponsePollingHandler from litellm.types.llms.openai import ResponsesAPIStatus @@ -70,6 +77,7 @@ class _StreamEvent(TypedDict, total=False): delta: ReadOnly[str] part: ReadOnly[object] response: ReadOnly[_TerminalResponse] + error: ReadOnly[_JsonDict] class _StreamEventParser: @@ -82,8 +90,29 @@ class _StreamsBody(Protocol): def body_iterator(self) -> AsyncIterable[str | bytes | memoryview]: ... -def _sse_frame_data(frame: str) -> str | None: - return next((line[6:].strip() for line in split_sse_lines(frame) if line.startswith("data: ")), None) +def _terminal_response(event: _StreamEvent) -> _TerminalResponse: + match event.get("response"): + case Mapping() as response: + return response + case _: + return {} + + +def _event_error(event: Mapping[str, object]) -> Mapping[str, object]: + match event.get("error"): + case str() as message if message: + return {**event, "message": message} + case _: + return event_error_object(event) or event + + +def _first_chunk_error(response: object) -> Mapping[str, object] | None: + if isinstance(response, _StreamsBody) or not isinstance(response, Response): + return None + if response.status_code < HTTPStatus.BAD_REQUEST: + return None + error_obj: Final = sse_error_object(bytes(response.body)) or {} + return responses_error_fields({**error_obj, "code": error_obj.get("code") or str(response.status_code)}) async def _never_receive() -> Message: @@ -176,6 +205,7 @@ async def background_streaming_task( user_api_base=user_api_base, version=version, skip_pre_call_logic=True, + sse_framing=RESPONSES_SSE_FRAMING, ) # Process streaming response following OpenAI events format @@ -210,6 +240,7 @@ async def background_streaming_task( None # Will be set by response.completed/failed/incomplete/cancelled ) terminal_error = None + frame_failure: Mapping[str, object] | None = None # rebind-ok: set once by the first error frame _event_to_status: Final[Mapping[str, ResponsesAPIStatus]] = { "response.completed": "completed", "response.failed": "failed", @@ -246,10 +277,15 @@ async def background_streaming_task( if isinstance(chunk, bytes): chunk = chunk.decode("utf-8") - if isinstance(chunk, str) and (chunk_data := _sse_frame_data(chunk)) is not None: - if chunk_data == "[DONE]": - break + if ( + isinstance(chunk, str) + and frame_failure is None + and (error_data := frame_error_data(chunk)) is not None + ): + frame_failure = responses_error_fields(_event_error(_StreamEventParser.parse(error_data))) + payloads, stream_done = frame_payloads(chunk) if isinstance(chunk, str) else ((), False) + for chunk_data in payloads: try: event: _StreamEvent = _StreamEventParser.parse(chunk_data) event_type = event.get("type", "") @@ -271,7 +307,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 @@ -338,7 +374,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 @@ -352,16 +388,11 @@ async def background_streaming_task( status="in_progress", ) - elif event_type in ( - "response.completed", - "response.failed", - "response.incomplete", - "response.cancelled", - ): + elif terminal_status is None and event_type in _event_to_status: # Terminal event - extract all ResponsesAPIResponse fields # https://platform.openai.com/docs/api-reference/responses-streaming finished_at = perf_counter() - response_data = event.get("response", {}) + response_data = _terminal_response(event) terminal_status = response_data.get( "status", _event_to_status.get(event_type, "completed"), @@ -395,22 +426,30 @@ async def background_streaming_task( if "output" in response_data: final_output = response_data.get("output", []) for item in final_output: - item_id = item.get("id") - if item_id: - output_items[item_id] = item + match item: + case {"id": item_id} if item_id: + output_items[item_id] = item + case _: + continue state_dirty = True - # Flush state to Redis if interval elapsed - await flush_state_if_needed() - - except json.JSONDecodeError 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 # Final flush to ensure all accumulated state is saved await flush_state_if_needed(force=True) - # Use the terminal status from the stream, default to "completed" - final_status: Final = terminal_status or "completed" + early_error: Final = _first_chunk_error(response) + stream_failure: Final = frame_failure if terminal_status is None else None + failure: Final = stream_failure if early_error is None else early_error + final_error: Final = terminal_error if failure is None else failure + final_status: Final = (terminal_status or "completed") if failure is None else "failed" if finished_at is None: finished_at = perf_counter() # rebind-ok: close the interval when the stream has no terminal event @@ -419,7 +458,7 @@ async def background_streaming_task( status=final_status, timing=timing_snapshot(), usage=usage_data, - error=terminal_error, + error=final_error, reasoning=reasoning_data, tool_choice=tool_choice_data, tools=tools_data, @@ -441,7 +480,7 @@ async def background_streaming_task( "Finished background streaming for %s, status=%s, error=%s, incomplete_details=%s, output_items=%s", polling_id, final_status, - terminal_error, + final_error, incomplete_details_data, len(output_items), ) diff --git a/litellm/proxy/response_polling/polling_handler.py b/litellm/proxy/response_polling/polling_handler.py index d66773cd048..0a13e3a9e99 100644 --- a/litellm/proxy/response_polling/polling_handler.py +++ b/litellm/proxy/response_polling/polling_handler.py @@ -101,7 +101,7 @@ class ResponsePollingHandler: polling_id: str, status: ResponsesAPIStatus | None = None, usage: dict | None = None, - error: dict | None = None, + error: Mapping[str, object] | None = None, incomplete_details: dict | None = None, reasoning: dict | None = None, tool_choice: object | None = None, diff --git a/litellm/responses/streaming_iterator.py b/litellm/responses/streaming_iterator.py index 0c002a770e8..4273647567e 100644 --- a/litellm/responses/streaming_iterator.py +++ b/litellm/responses/streaming_iterator.py @@ -1177,7 +1177,8 @@ class BaseResponsesAPIStreamingIterator: def _note_yielded_event(self, event: ResponsesAPIStreamingResponse) -> None: self._yielded_first_chunk = True - if event.type not in PRE_OUTPUT_LIFECYCLE_EVENT_TYPES: + event_type: Final[object] = getattr(event, "type", None) + if event_type not in PRE_OUTPUT_LIFECYCLE_EVENT_TYPES: self._output_started = True def _fallback_error(self, original: Exception) -> MidStreamFallbackError: diff --git a/tests/e2e/llm_translation/test_containers_e2e.py b/tests/e2e/llm_translation/test_containers_e2e.py index 0a494d11108..95dc68341e8 100644 --- a/tests/e2e/llm_translation/test_containers_e2e.py +++ b/tests/e2e/llm_translation/test_containers_e2e.py @@ -32,7 +32,7 @@ second regression, since the global-credential fallback then reaches the container anyway. The streaming cell repeats the flow with ``stream=True`` and uploads right -after the last event. The OpenAI SDK closes the connection at ``[DONE]``, so an +after the last event. The OpenAI SDK closes the connection after the last frame, so an ownership row written after the stream is cancelled with the body task and every follow-up container call 403s (LIT-8612); the row has to land before the ``response.completed`` frame goes out. diff --git a/tests/integration/observability/test_guardrail_effects.py b/tests/integration/observability/test_guardrail_effects.py index 1220a65528c..5d3ab4474ad 100644 --- a/tests/integration/observability/test_guardrail_effects.py +++ b/tests/integration/observability/test_guardrail_effects.py @@ -1200,9 +1200,7 @@ def test_responses_pre_call_denial_streams_sse_with_typed_message_item(gateway: response.headers["content-type"], response.text, ) - lines: Final = tuple(line for line in response.text.split("\n") if line.startswith("data: ")) - assert lines[-1] == "data: [DONE]", response.text - events: Final = tuple(json.loads(line.removeprefix("data: ")) for line in lines[:-1]) + events: Final = _blocked_stream_events(response.text) kinds: Final = tuple(event["type"] for event in events) assert tuple(kind for kind in kinds if kind != "response.output_text.delta") == ( "response.created", @@ -1289,9 +1287,13 @@ def _responses_output_denial_config(tmp_path: Path, identity: str, model: str) - def _blocked_stream_events(text: str) -> tuple[dict[str, object], ...]: - lines: Final = tuple(line for line in text.split("\n") if line.startswith("data: ")) - assert lines[-1] == "data: [DONE]", text - return tuple(json.loads(line.removeprefix("data: ")) for line in lines[:-1]) + assert "data: [DONE]" not in text, text + lines: Final = tuple(line for line in text.split("\n") if line.startswith(("event: ", "data: "))) + events: Final[tuple[dict[str, object], ...]] = tuple( + json.loads(line.removeprefix("data: ")) for line in lines[1::2] + ) + assert lines[::2] == tuple(f"event: {event['type']}" for event in events), text + return events def _dead_api_base() -> str: @@ -2798,7 +2800,7 @@ def test_client_disconnect_after_a_finished_responses_tool_call_scans_the_text_a with _disconnect_rig(gateway, tmp_path, shape="text_then_tool_call") as rig: try: received: Final = _stream_and_close( - rig, "/v1/responses", _tool_call_request(rig, "/v1/responses"), "response.output_item.done" + rig, "/v1/responses", _tool_call_request(rig, "/v1/responses"), '"type":"response.output_item.done"' ) assert rig.token() in received, received scans: Final = eventually( diff --git a/tests/integration/providers/test_bedrock_mantle_claude_chat_chaos.py b/tests/integration/providers/test_bedrock_mantle_claude_chat_chaos.py index bcbf6d0f6f3..2325b705790 100644 --- a/tests/integration/providers/test_bedrock_mantle_claude_chat_chaos.py +++ b/tests/integration/providers/test_bedrock_mantle_claude_chat_chaos.py @@ -261,7 +261,7 @@ async def test_slow_native_streams_reach_every_caller_whole_and_are_logged_once( assert len(served) == 10 for item in served: _assert_answered_with_its_own_marker(item) - assert item.text.rstrip().endswith("data: [DONE]"), item.text + assert item.text.rstrip().endswith("data: [DONE]") == (item.call.endpoint == "chat"), item.text _assert_native_requests_for(wire.drain(), calls) assert _spend_status_by_call(served) == {item.call_id: "success" for item in served} diff --git a/tests/integration/providers/test_bedrock_mantle_claude_chat_wire.py b/tests/integration/providers/test_bedrock_mantle_claude_chat_wire.py index 2eed92fb873..7c2f60e58a5 100644 --- a/tests/integration/providers/test_bedrock_mantle_claude_chat_wire.py +++ b/tests/integration/providers/test_bedrock_mantle_claude_chat_wire.py @@ -854,7 +854,12 @@ def test_responses_api_stream_on_a_mantle_claude_id_bridges_to_the_native_stream "POST", "/v1/responses", {"model": model, "input": _prompt(marker), "stream": True} ) assert response.status_code == 200, response.text - events: Final = _sse_objects(response.text) + assert "data: [DONE]" not in response.text, response.text + events: Final = tuple( + _JSON_OBJECT.validate_json(line.removeprefix("data: ")) + for line in response.text.splitlines() + if line.startswith("data: ") + ) kinds: Final = [event["type"] for event in events] assert (kinds[0], kinds[-1]) == ("response.created", "response.completed"), kinds assert "".join( diff --git a/tests/integration/streaming/test_chat_stream_guardrail_error_events.py b/tests/integration/streaming/test_chat_stream_guardrail_error_events.py index a8ebe510b8e..b2390de7737 100644 --- a/tests/integration/streaming/test_chat_stream_guardrail_error_events.py +++ b/tests/integration/streaming/test_chat_stream_guardrail_error_events.py @@ -482,8 +482,7 @@ def _responses_body(marker: str, *directives: str) -> dict[str, JsonValue]: def _failed_response_error(response: httpx.Response) -> dict[str, JsonValue]: assert response.status_code == 200, response.text - lines: Final = _data_lines(response.text) - assert lines[-1] == "[DONE]" and lines.count("[DONE]") == 1, response.text + assert "[DONE]" not in _data_lines(response.text), response.text events: Final = _payloads(response.text) kinds: Final = tuple(string_value(event["type"]) for event in events) assert kinds[0] == "response.created" and kinds[-1] == "response.failed", response.text diff --git a/tests/integration/streaming/test_responses_bridge_stream_errors.py b/tests/integration/streaming/test_responses_bridge_stream_errors.py index 8a5655be0f8..6b0e35f9c56 100644 --- a/tests/integration/streaming/test_responses_bridge_stream_errors.py +++ b/tests/integration/streaming/test_responses_bridge_stream_errors.py @@ -474,9 +474,10 @@ def test_native_responses_stream_forwards_the_failed_response_on_both_legs(gatew assert response.status_code == 200, response.text frames: Final = data_frames(response.text) assert [frame["type"] for frame in frames] == ["response.created", "response.failed"], response.text + assert [name for name, _ in sse_events(response.text)] == ["response.created", "response.failed"], response.text failed: Final = object_value(frames[-1]["response"]) assert failed["status"] == "failed", failed assert object_value(failed["error"])["code"] == "rate_limit_exceeded", failed - assert response.text.rstrip().endswith("data: [DONE]"), response.text + assert "data: [DONE]" not in response.text, response.text assert len(wire.drain()) == 1 assert spend_row(response.headers["x-litellm-call-id"])["status"] == "failure" diff --git a/tests/integration/streaming/test_responses_top_level_error_events.py b/tests/integration/streaming/test_responses_top_level_error_events.py index 2cbc4570dec..c50f16129cb 100644 --- a/tests/integration/streaming/test_responses_top_level_error_events.py +++ b/tests/integration/streaming/test_responses_top_level_error_events.py @@ -258,8 +258,10 @@ def _responses_frames(gateway: Gateway, stream: tuple[bytes, ...]) -> tuple[http ) _assert_one_attempt(served.wire, marker) assert served.response.status_code == 200, served.response.text - assert served.response.text.rstrip().endswith("data: [DONE]"), served.response.text + assert "data: [DONE]" not in served.response.text, served.response.text frames: Final = data_frames(served.response.text) + names: Final = tuple(name for name, _ in sse_events(served.response.text)) + assert names == tuple(item["type"] for item in frames), served.response.text assert frames[-1]["type"] == "response.failed", served.response.text return served.response, object_value(object_value(frames[-1]["response"])["error"]) diff --git a/tests/unit/llms/openai/responses/test_openai_responses_guardrail_handler.py b/tests/unit/llms/openai/responses/test_openai_responses_guardrail_handler.py index a5bb49f0793..7daf773d55b 100644 --- a/tests/unit/llms/openai/responses/test_openai_responses_guardrail_handler.py +++ b/tests/unit/llms/openai/responses/test_openai_responses_guardrail_handler.py @@ -4620,7 +4620,14 @@ class TestBuildBlockSseChunks: def _payloads(self, chunks): import json - return [json.loads(chunk.decode().removeprefix("data: ").strip()) for chunk in chunks] + payloads: Final = [ + json.loads(next(line[6:] for line in chunk.decode().splitlines() if line.startswith("data: "))) + for chunk in chunks + ] + assert [chunk.decode().splitlines()[0] for chunk in chunks] == [ + f"event: {payload['type']}" for payload in payloads + ] + return payloads def test_standalone_block_emits_complete_synthetic_stream(self): handler = OpenAIResponsesHandler() diff --git a/tests/unit/proxy/proxy_server/test_streaming_helpers.py b/tests/unit/proxy/proxy_server/test_streaming_helpers.py index 80bed119e79..fdf58ec2fd7 100644 --- a/tests/unit/proxy/proxy_server/test_streaming_helpers.py +++ b/tests/unit/proxy/proxy_server/test_streaming_helpers.py @@ -18,7 +18,9 @@ import asyncio import hashlib import json import logging -from collections.abc import AsyncIterator, Callable, Mapping +from collections.abc import AsyncIterator, Callable, Iterable, Iterator, Mapping +from datetime import datetime, timezone +from enum import Enum from types import MappingProxyType from typing import Final, Literal from unittest.mock import AsyncMock, MagicMock @@ -28,7 +30,7 @@ import pytest from fastapi import HTTPException, Response from fastapi.responses import JSONResponse, StreamingResponse from openai import APIError as OpenAIAPIError -from pydantic import BaseModel +from pydantic import BaseModel, model_serializer import litellm import litellm.proxy.proxy_server as ps @@ -36,9 +38,11 @@ from litellm.constants import RETURN_RAW_MODEL_NAME_METADATA_KEY from litellm.integrations.custom_logger import CustomLogger from litellm.proxy._types import ProxyException, UserAPIKeyAuth from litellm.proxy.common_request_processing import create_response +from litellm.proxy.common_utils.responses_stream_errors import ResponsesStreamErrorState from litellm.proxy.proxy_server import ( _apply_streaming_chunk_hooks, _format_fallback_metadata_sse_event, + _format_responses_sse_chunk, _format_streaming_sse_chunk, _get_client_requested_model_for_streaming, _get_streaming_fallback_metadata, @@ -51,11 +55,16 @@ from litellm.proxy.proxy_server import ( select_data_generator, ) from litellm.types.llms.openai import ( + ErrorEvent, + OutputTextDeltaEvent, ResponseCompletedEvent, ResponseCreatedEvent, ResponseFailedEvent, + ResponseIncompleteEvent, ResponsesAPIResponse, + ResponsesAPIStreamEvents, ) +from litellm.types.proxy.litellm_pre_call_utils import RedactedDict from litellm.types.utils import ( ChatCompletionDeltaToolCall, Delta, @@ -106,6 +115,353 @@ async def _async_iter_raises(exc: Exception): raise exc +@pytest.mark.asyncio +@pytest.mark.parametrize("terminal_type", ["response.completed", "response.failed", "response.incomplete", "error"]) +async def test_responses_stream_dispatches_named_events_without_a_chat_done_sentinel(terminal_type: str) -> None: + response: Final = ResponsesAPIResponse(id="resp_framing", created_at=1, model="local-model", output=[]) + created: Final = ResponseCreatedEvent.model_validate( + {"type": "response.created", "sequence_number": 0, "response": response} + ) + delta: Final = OutputTextDeltaEvent.model_validate( + { + "type": "response.output_text.delta", + "sequence_number": 1, + "item_id": "msg_framing", + "output_index": 0, + "content_index": 0, + "delta": "hello", + } + ) + terminal_models: Final = { + "response.completed": ResponseCompletedEvent, + "response.failed": ResponseFailedEvent, + "response.incomplete": ResponseIncompleteEvent, + "error": ErrorEvent, + } + terminal: Final = terminal_models[terminal_type].model_validate( + {"type": terminal_type, "sequence_number": 2, "response": response} + ) + + async def upstream() -> AsyncIterator[BaseModel]: + for event in (created, delta, terminal): + yield event + + frames: Final = tuple( + [ + frame.decode() if isinstance(frame, bytes) else frame + async for frame in select_data_generator( + response=upstream(), user_api_key_dict=_user_auth(), request_data={}, responses_stream_errors=True + ) + ] + ) + assert len(frames) == 3 + for frame, event_type in zip( + frames, ("response.created", "response.output_text.delta", terminal_type), strict=True + ): + assert frame.startswith(f"event: {event_type}\n") + payload: Final = json.loads( + next(line.removeprefix("data: ") for line in frame.splitlines() if line.startswith("data: ")) + ) + assert payload["type"] == event_type + assert json.loads(frames[1].split("data: ", 1)[1])["delta"] == "hello" + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "as_chunk", + [lambda event: event, lambda event: event.model_dump()], + ids=["model", "dict"], +) +async def test_responses_stream_sends_no_failure_event_after_an_upstream_error_event( + as_chunk: Callable[[ErrorEvent], BaseModel | Mapping[str, object]], +) -> None: + response: Final = ResponsesAPIResponse(id="resp_error_then_raise", created_at=1, model="local-model", output=[]) + created: Final = ResponseCreatedEvent.model_validate( + {"type": "response.created", "sequence_number": 0, "response": response} + ) + error_event: Final = ErrorEvent.model_validate( + {"type": "error", "sequence_number": 1, "code": "server_error", "message": "upstream failed", "param": None} + ) + + async def upstream() -> AsyncIterator[BaseModel | Mapping[str, object]]: + yield created + yield as_chunk(error_event) + raise litellm.APIError(status_code=500, message="Stream close failed", llm_provider="openai", model="local-model") + + frames: Final = tuple( + [ + frame.decode() if isinstance(frame, bytes) else frame + async for frame in select_data_generator( + response=upstream(), user_api_key_dict=_user_auth(), request_data={}, responses_stream_errors=True + ) + ] + ) + + assert tuple(frame.split("\n", 1)[0] for frame in frames) == ("event: response.created", "event: error") + assert json.loads(frames[-1].split("data: ", 1)[1])["message"] == "upstream failed" + + +async def _responses_frames(chunks: Iterable[object]) -> tuple[str, ...]: + async def upstream() -> AsyncIterator[object]: + for chunk in chunks: + yield chunk + + return tuple( + [ + frame.decode() if isinstance(frame, bytes) else frame + async for frame in select_data_generator( + response=upstream(), user_api_key_dict=_user_auth(), request_data={}, responses_stream_errors=True + ) + ] + ) + + +_DUMPED_DELTA: Final = OutputTextDeltaEvent( + type=ResponsesAPIStreamEvents.OUTPUT_TEXT_DELTA, + item_id="msg_1", + output_index=0, + content_index=0, + delta="hi", + sequence_number=3, +) + + +@pytest.mark.asyncio +async def test_responses_stream_sends_a_hook_chunk_dumped_from_an_event_as_that_named_event() -> None: + frames: Final = await _responses_frames((_DUMPED_DELTA.model_dump(),)) + + assert len(frames) == 1 + event_line, data_line = frames[0].strip().split("\n") + assert event_line == "event: response.output_text.delta" + assert json.loads(data_line.removeprefix("data: ")) == json.loads(_DUMPED_DELTA.model_dump_json()) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "chunk", + [{"type": 5}, {"type": "response.completed", "response": "hook-text"}], + ids=["non_string_type", "string_response"], +) +async def test_responses_stream_sends_a_hook_chunk_its_event_model_cannot_read_as_a_data_frame( + chunk: Mapping[str, object], +) -> None: + frames: Final = await _responses_frames((chunk,)) + + assert frames == (f"data: {json.dumps(chunk)}\n\n",) + + +class _TypeDroppedOnTheWire(BaseModel, frozen=True): + type: str = "response.completed" + sequence_number: int = 3 + + @model_serializer + def _without_type(self) -> Mapping[str, int]: + return {"sequence_number": self.sequence_number} + + +@pytest.mark.asyncio +async def test_responses_stream_names_a_model_chunk_from_the_json_it_sends() -> None: + frames: Final = await _responses_frames((_TypeDroppedOnTheWire(),)) + + assert len(frames) == 1 + assert frames[0].startswith("data: ") + assert json.loads(frames[0].removeprefix("data: ")) == {"sequence_number": 3} + + +class _MaskedReprString(str): + def __repr__(self) -> str: + return "'***'" + + +class _MaskedReprModel(BaseModel, frozen=True): + token: str + + def __repr__(self) -> str: + return "_MaskedReprModel(token=[secure])" + + +class _MaskedCount(int): + def __repr__(self) -> str: + return "_MaskedCount([secure])" + + +class _RigColor(Enum): + RED = "red" + + +def _self_referencing_chunk() -> Mapping[str, object]: + chunk: Final[dict[str, object]] = {"type": "response.output_text.delta"} + chunk["self"] = chunk # mutable-ok: a self-reference needs the dict to exist before it can hold itself + return chunk + + +def _chunk_holding_itself_twice() -> Mapping[str, object]: + chunk: Final[dict[str, object]] = {"type": "response.output_text.delta"} + chunk.update(left=chunk, right=chunk) # mutable-ok: a self-reference needs the dict to exist before it can hold itself + return chunk + + +_REDACTED_HEADERS: Final = RedactedDict({"authorization": "Bearer hook-canary"}) + + +@pytest.mark.parametrize( + "chunk", + [ + {"type": _MaskedReprString("hook-canary"), "delta": "hi"}, + {"type": "response.output_text.delta", "headers": _REDACTED_HEADERS}, + RedactedDict({"type": "response.output_text.delta", "authorization": "Bearer hook-canary"}), + MappingProxyType({"type": "response.output_text.delta"}), + {"type": "response.output_text.delta", "item": _MaskedReprModel(token="hook-canary")}, + {"type": "response.output_text.delta", "count": _MaskedCount(7)}, + {"type": "response.output_text.delta", "color": _RigColor.RED}, + {"type": "response.output_text.delta", ("a", "b"): 1}, + {"type": "response.output_text.delta", "score": float("nan")}, + _self_referencing_chunk(), + _chunk_holding_itself_twice(), + ], + ids=[ + "masked_string", + "redacted_dict_value", + "redacted_dict_chunk", + "mapping_proxy_chunk", + "model_with_own_repr", + "int_subclass_with_own_repr", + "enum_without_a_json_value", + "tuple_key", + "nan", + "holds_itself", + "holds_itself_twice", + ], +) +def test_responses_sse_chunk_that_json_cannot_show_as_printed_keeps_the_base_frame(chunk: Mapping[object, object]) -> None: + assert _format_responses_sse_chunk(chunk=chunk) == _format_streaming_sse_chunk(chunk=chunk) + + +def test_responses_sse_chunk_without_a_type_keeps_the_base_frame() -> None: + chunk: Final = {"note": "hook output without a type"} + + assert _format_responses_sse_chunk(chunk=chunk) == _format_streaming_sse_chunk(chunk=chunk) + + +def test_responses_error_event_encodes_its_detail_the_way_the_chat_error_frame_does() -> None: + error: Final = { + "message": "blocked", + "provider_specific_fields": {"guardrail": "rig-guard", "at": datetime(2026, 10, 30, tzinfo=timezone.utc)}, + } + + frame: Final = ResponsesStreamErrorState().format_error(error) + + assert json.loads(frame.split("data: ", 1)[1])["error"] == json.loads(json.dumps(error, default=str)) + + +@pytest.mark.parametrize("status_code", ["abc", 1.5, True, None], ids=["string", "float", "bool", "null"]) +def test_responses_hook_error_frame_keeps_its_message_whatever_status_code_it_carries(status_code: object) -> None: + hook_frame: Final = json.dumps( + {"error": {"message": "rig status", "type": "guardrail_scan_error", "status_code": status_code}} + ) + + frame: Final = ResponsesStreamErrorState().format_error_message(f"data: {hook_frame}\n\n") + + assert frame.startswith("event: error\n") + assert json.loads(frame.split("data: ", 1)[1])["message"] == "rig status" + + +class _StatusCarryingError(Exception): + def __init__(self, status_code: object) -> None: + super().__init__("rig status") + self.status_code: Final = status_code + + +@pytest.mark.parametrize( + ("status_code", "code"), + [ + (429, "rate_limit_exceeded"), + ("429", "rate_limit_exceeded"), + (429.0, "rate_limit_exceeded"), + ("abc", "server_error"), + (429.5, "server_error"), + ], + ids=["int", "decimal_string", "integral_float", "word", "fractional_float"], +) +def test_responses_failure_event_reads_a_numeric_status_code_the_way_base_did(status_code: object, code: str) -> None: + frame: Final = ResponsesStreamErrorState().format_failure(_StatusCarryingError(status_code)) + + assert frame is not None + failure: Final = ResponseFailedEvent.model_validate_json(frame.split("data: ", 1)[1]) + assert failure.response.error is not None + assert failure.response.error["code"] == code + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("guardrail_error", "error_object"), + [ + ( + {"message": "response withheld", "type": "guardrail_scan_error", "code": 500}, + {"message": "response withheld", "type": "guardrail_scan_error", "code": 500}, + ), + ("response withheld", {"message": "response withheld"}), + ], +) +async def test_a_guardrail_error_string_mid_responses_stream_ends_it_as_a_named_error_event( + guardrail_error: Mapping[str, object] | str, error_object: Mapping[str, object] +) -> None: + response: Final = ResponsesAPIResponse(id="resp_hook_error", created_at=1, model="local-model", output=[]) + created: Final = ResponseCreatedEvent.model_validate( + {"type": "response.created", "sequence_number": 0, "response": response} + ) + + async def upstream() -> AsyncIterator[object]: + yield created + yield f"data: {json.dumps({'error': guardrail_error})}\n\n" + + frames: Final = tuple( + [ + frame.decode() if isinstance(frame, bytes) else frame + async for frame in select_data_generator( + response=upstream(), user_api_key_dict=_user_auth(), request_data={}, responses_stream_errors=True + ) + ] + ) + event_line, data_line = frames[-1].strip().split("\n") + + assert "data: [DONE]" not in "".join(frames) + assert event_line == "event: error" + assert json.loads(data_line.removeprefix("data: ")) == { + "type": "error", + "code": "server_error", + "message": "response withheld", + "param": None, + "sequence_number": 1, + "error": error_object, + } + + +@pytest.mark.parametrize("binary", [False, True]) +@pytest.mark.parametrize( + ("event_type", "event_line"), + [ + (None, ""), + ("", ""), + ("response.delta\ninjected", ""), + ("response.delta\rinjected", ""), + ("response.delta\r\ninjected", ""), + ("response.vendor_delta", "event: response.vendor_delta\n"), + ], +) +def test_responses_event_names_preserve_payload_without_injecting_sse_lines( + event_type: str | None, event_line: str, binary: bool +) -> None: + payload: Final = {"type": event_type, "sequence_number": 0} + data_line: Final = f"data: {json.dumps(payload)}\n\n" + frame: Final = data_line.encode("utf-8") if binary else data_line + expected: Final = (event_line + data_line).encode("utf-8") if binary else event_line + data_line + state: Final = ResponsesStreamErrorState() + state.observe_chunk(payload) + + assert state.mark_emitted(frame) == expected + + class _FakeStream: def __init__(self, chunks, hidden_params=None): self._chunks = chunks @@ -1430,14 +1786,15 @@ async def test_responses_stream_failure_sends_the_failure_hook_replacement(monke ] ) failure: Final = ResponseFailedEvent.model_validate_json( - next(line.removeprefix("data: ") for line in frames[-2].splitlines() if line.startswith("data: ")) + next(line.removeprefix("data: ") for line in frames[-1].splitlines() if line.startswith("data: ")) ) assert _BLOCK_MESSAGE not in "".join(frames) assert failure.response.error is not None assert failure.response.error["code"] == "permission_error" assert "Request blocked by policy" in failure.response.error["message"] - assert frames[-1] == "data: [DONE]\n\n" + assert frames[-1].startswith("event: response.failed\n") + assert "data: [DONE]\n\n" not in frames @pytest.mark.asyncio @@ -1633,8 +1990,8 @@ async def test_responses_stream_keeps_tool_deltas_and_only_emits_a_valid_termina for frame in event_frames ) - assert decoded[-1] == "data: [DONE]\n\n" - assert len(decoded) == len(event_frames) + 1 + assert len(decoded) == len(event_frames) + assert "data: [DONE]\n\n" not in decoded assert payloads[0]["response"]["id"] == "resp_visible" assert payloads[1] == tool_delta.model_dump() assert len(payloads) == 3 @@ -1660,6 +2017,127 @@ async def test_responses_stream_keeps_tool_deltas_and_only_emits_a_valid_termina assert "error" not in payloads[-1] +@pytest.mark.asyncio +@pytest.mark.parametrize("raw_sse_stream", [False, True]) +@pytest.mark.parametrize("newline", ["\n", "\r\n"]) +@pytest.mark.parametrize( + "event_name,data,terminal", + [ + ( + "response.output_text.delta", + {"type": "response.output_text.delta", "sequence_number": 4, "delta": "hi"}, + False, + ), + ( + "response.completed", + {"type": "response.completed", "sequence_number": 4, "response": {"id": "resp_raw", "model": "gpt-6"}}, + True, + ), + ("error", {"error": {"message": "denied", "code": "529"}}, True), + ("response.completed", {"sequence_number": 4, "note": "typeless"}, False), + ("error", {"sequence_number": 4, "note": "no error object"}, False), + ("error", {"sequence_number": 4, "error": "denied"}, False), + ("error", {"sequence_number": 4, "error": ["denied"]}, False), + ], +) +async def test_a_responses_stream_failing_after_a_pre_framed_event_adds_a_failure_only_without_a_terminal( + raw_sse_stream: bool, newline: str, event_name: str, data: Mapping[str, object], terminal: bool +) -> None: + pre_framed: Final = f"event: {event_name}{newline}data: {json.dumps(data)}{newline}{newline}" + + async def upstream() -> AsyncIterator[bytes]: + yield pre_framed.encode() + raise litellm.APIError(status_code=500, message="Stream closed", llm_provider="openai", model="gpt-6") + + frames: Final = [ + frame.decode() if isinstance(frame, bytes) else frame + async for frame in select_data_generator( + response=upstream(), + user_api_key_dict=_user_auth(), + request_data={"_litellm_raw_sse_stream": True} if raw_sse_stream else {}, + responses_stream_errors=True, + ) + ] + + assert frames[0] == pre_framed + assert len(frames) == (1 if terminal else 2) + if not terminal: + failure: Final = ResponseFailedEvent.model_validate_json(frames[1].split("data: ", 1)[1]) + assert frames[1].startswith("event: response.failed\n") + assert failure.sequence_number == 5 + + +@pytest.mark.parametrize( + "sequence_number", + ["9" * 4300, "-7", str(2**63), str(2**63 - 1)], + ids=["4300_digits", "negative", "past_int64", "int64_max"], +) +def test_responses_error_event_after_an_out_of_range_pre_framed_sequence_number_stays_parseable( + sequence_number: str, +) -> None: + state: Final = ResponsesStreamErrorState() + state.observe_frames( + 'data: {"type": "response.output_text.delta", "sequence_number": 4}\n\n' + f'data: {{"type": "response.output_text.delta", "sequence_number": {sequence_number}}}\n\n' + ) + + error_frame: Final = state.format_error_message('data: {"error": {"message": "scan failed"}}\n\n') + failure_frame: Final = state.format_failure(RuntimeError("stream closed")) + + error_event: Final = json.loads(error_frame.split("data: ", 1)[1]) + assert failure_frame is not None + assert error_event["message"] == "scan failed" + assert error_event["sequence_number"] == 5 + assert json.loads(failure_frame.split("data: ", 1)[1])["sequence_number"] == 5 + + +class _UnreadableErrorModel(BaseModel, frozen=True): + type: str = "response.output_text.delta" + sequence_number: int = 4 + + @property + def error(self) -> object: + raise RuntimeError("error is not readable") + + +class _UnreadableErrorMapping(Mapping[str, object]): + def __init__(self, data: Mapping[str, object]) -> None: + self.data: Final = data + + def __getitem__(self, key: str) -> object: + if key == "error": + raise RuntimeError("error is not readable") + return self.data[key] + + def __iter__(self) -> Iterator[str]: + return iter(self.data) + + def __len__(self) -> int: + return len(self.data) + + +@pytest.mark.parametrize( + "chunk,next_sequence_number", + [ + (_UnreadableErrorModel(), 5), + (_UnreadableErrorMapping({"type": "response.output_text.delta", "sequence_number": 4}), 5), + ({"type": "response.output_text.delta", "sequence_number": 2**63}, 2**63 + 1), + ], + ids=["unreadable_error_property", "unreadable_error_item", "past_int64"], +) +def test_responses_failure_after_a_hook_chunk_continues_its_sequence_number_like_base( + chunk: object, next_sequence_number: int +) -> None: + state: Final = ResponsesStreamErrorState() + state.observe_chunk(chunk) + state.mark_emitted(_format_responses_sse_chunk(chunk=chunk)) + + failure_frame: Final = state.format_failure(RuntimeError("stream closed")) + + assert failure_frame is not None + assert json.loads(failure_frame.split("data: ", 1)[1])["sequence_number"] == next_sequence_number + + # --------------------------------------------------------------------------- # select_data_generator # --------------------------------------------------------------------------- diff --git a/tests/unit/proxy/response_api_endpoints/test_endpoints.py b/tests/unit/proxy/response_api_endpoints/test_endpoints.py index b7d0d9e62cd..24ad068e6df 100644 --- a/tests/unit/proxy/response_api_endpoints/test_endpoints.py +++ b/tests/unit/proxy/response_api_endpoints/test_endpoints.py @@ -9,7 +9,7 @@ import datetime import json import re import unittest -from collections.abc import Callable, Mapping +from collections.abc import AsyncIterator, Callable, Mapping from typing import Any, Final, Literal, cast from unittest.mock import AsyncMock, MagicMock, call, patch @@ -26,6 +26,7 @@ import litellm from litellm.caching.caching import DualCache from litellm.constants import EMPTY_MAPPING from litellm.integrations.custom_guardrail import CustomGuardrail +from litellm.integrations.custom_logger import CustomLogger from litellm.proxy._types import ProxyException, UserAPIKeyAuth from litellm.proxy.common_utils.proxy_rate_limit_error import ProxyRateLimitError from litellm.proxy.proxy_server import app @@ -36,6 +37,7 @@ from litellm.types.llms.openai import ResponseAPIUsage, ResponsesAPIResponse from litellm.types.router import RouterErrors from litellm.types.utils import CallTypesLiteral from tests._master_key import MASTER_KEY +from tests._support.clock_loops import AutojumpClockLoop @pytest.mark.asyncio @@ -154,8 +156,9 @@ async def test_streaming_upstream_errors_keep_the_client_protocol( assert result.status_code == 200, result.text assert message in result.text if path == "/v1/responses": - assert frames[-1] == "data: [DONE]", result.text - assert frames[-2].startswith("event: response.failed\n"), result.text + assert all("data: [DONE]" not in frame for frame in frames), result.text + assert frames[-1].startswith("event: response.failed\n"), result.text + assert [frame.splitlines()[0] for frame in frames] == [f"event: {event['type']}" for event in events] if partial: assert [event["type"] for event in events] == [ "response.created", @@ -188,6 +191,173 @@ async def test_streaming_upstream_errors_keep_the_client_protocol( assert "error" in events[-1] +class _FirstFrameHook(CustomLogger): + def __init__(self, first: str | bytes | Exception) -> None: + super().__init__() + self.first: Final = first + + async def async_post_call_streaming_iterator_hook( + self, user_api_key_dict: object, response: object, request_data: object + ) -> AsyncIterator[str | bytes]: + if isinstance(self.first, Exception): + raise self.first + yield self.first + + +_DATA_ERROR_FRAME: Final = 'data: {"error": {"message": "denied", "type": "guardrail_scan_error", "code": 403}}\n\n' +_NAMED_ERROR_FRAME: Final = ( + b'event: error\ndata: {"type": "error", "error": {"type": "overloaded_error", "message": "denied", "code": "529"}}\n\n' +) + + +async def _post_through_first_frame_hook( + monkeypatch: pytest.MonkeyPatch, path: str, first: str | bytes | Exception +) -> httpx.Response: + import litellm.proxy.proxy_server as ps + from litellm.proxy.auth.user_api_key_auth import user_api_key_auth + + model: Final = "gpt-6-astra" + upstream_url: Final = "https://streaming.example/v1" + is_chat: Final = path == "/v1/chat/completions" + router: Final = litellm.Router( + model_list=[ + { + "model_name": model, + "litellm_params": {"model": "openai/" + model, "api_base": upstream_url, "api_key": "fixture-key"}, + } + ], + num_retries=0, + ) + monkeypatch.setattr(ps, "llm_router", router) + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + monkeypatch.setattr(litellm, "callbacks", [_FirstFrameHook(first)]) + monkeypatch.setitem(app.dependency_overrides, user_api_key_auth, _auth_override) + with respx.mock as transport: + transport.post(upstream_url + ("/chat/completions" if is_chat else "/responses")).respond( + 200, content="data: [DONE]\n\n", headers={"Content-Type": "text/event-stream"} + ) + async with httpx.AsyncClient(transport=httpx.ASGITransport(app), base_url="http://testserver") as client: + return await client.post( + path, + json={ + "model": model, + "stream": True, + **({"messages": [{"role": "user", "content": "hello"}]} if is_chat else {"input": "hello"}), + }, + ) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "path,first_frame,expected_status", + [ + ("/v1/responses", _DATA_ERROR_FRAME, 403), + ("/v1/chat/completions", _DATA_ERROR_FRAME, 403), + ("/v1/chat/completions", _NAMED_ERROR_FRAME, 200), + ("/cursor/chat/completions", _NAMED_ERROR_FRAME, 200), + ("/v1/responses", _NAMED_ERROR_FRAME, 529), + ("/v1/responses", _NAMED_ERROR_FRAME.replace(b"event: error", b"event: response.created"), 200), + ("/v1/chat/completions", _DATA_ERROR_FRAME + "data: [DONE]\n\n", 200), + ("/v1/responses", _DATA_ERROR_FRAME + "data: [DONE]\n\n", 200), + ], +) +async def test_only_a_lone_first_error_frame_answers_with_its_status( + monkeypatch: pytest.MonkeyPatch, path: str, first_frame: str | bytes, expected_status: int +) -> None: + result: Final = await _post_through_first_frame_hook(monkeypatch, path, first_frame) + + assert result.status_code == expected_status, result.text + if expected_status == 200: + assert result.headers["content-type"].startswith("text/event-stream"), result.text + else: + assert result.json()["error"]["message"] == "denied" + + +@pytest.mark.asyncio +@pytest.mark.parametrize("path", ["/cursor/chat/completions", "/v1/chat/completions"]) +async def test_a_hook_failure_before_the_first_chunk_answers_in_chat_completions_shape( + monkeypatch: pytest.MonkeyPatch, path: str +) -> None: + message: Final = "rate limited before the first chunk" + result: Final = await _post_through_first_frame_hook( + monkeypatch, path, HTTPException(status_code=429, detail=message) + ) + frames: Final = tuple(frame for frame in result.text.split("\n\n") if frame) + + assert result.status_code == 429, result.text + assert frames[-1] == "data: [DONE]", result.text + assert message in json.loads(frames[0].removeprefix("data: "))["error"]["message"] + assert len(frames) == 2, result.text + + +@pytest.mark.parametrize("path", ["/v1/responses", "/v1/chat/completions"]) +def test_a_failure_after_keepalive_pings_keeps_the_client_protocol(monkeypatch: pytest.MonkeyPatch, path: str) -> None: + async def request_and_check() -> None: + import litellm.proxy.proxy_server as ps + from litellm.proxy.auth.user_api_key_auth import user_api_key_auth + + model: Final = "gpt-6-astra" + message: Final = "Upstream is rate limiting this key" + upstream_url: Final = "https://streaming.example/v1" + is_chat: Final = path == "/v1/chat/completions" + + async def slow_rate_limit(_: httpx.Request) -> httpx.Response: + await asyncio.sleep(0.3) + return httpx.Response(429, json={"error": {"message": message, "type": "rate_limit_error", "code": None}}) + + router: Final = litellm.Router( + model_list=[ + { + "model_name": model, + "litellm_params": { + "model": "openai/" + model, + "api_base": upstream_url, + "api_key": "fixture-key", + "max_retries": 0, + }, + } + ], + num_retries=0, + ) + monkeypatch.setattr(ps, "llm_router", router) + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + monkeypatch.setattr(litellm, "sse_keepalive_ping_interval_seconds", 0.05) + monkeypatch.setitem(app.dependency_overrides, user_api_key_auth, _auth_override) + with respx.mock as transport: + transport.post(upstream_url + ("/chat/completions" if is_chat else "/responses")).mock( + side_effect=slow_rate_limit + ) + async with httpx.AsyncClient(transport=httpx.ASGITransport(app), base_url="http://testserver") as client: + result: Final = await client.post( + path, + json={ + "model": model, + "stream": True, + **({"messages": [{"role": "user", "content": "hello"}]} if is_chat else {"input": "hello"}), + }, + ) + frames: Final = tuple(frame for frame in result.text.split("\n\n") if frame) + + assert result.status_code == 200, result.text + assert frames[0] == ": ping", result.text + if is_chat: + assert frames[-1] == "data: [DONE]", result.text + assert message in json.loads(frames[-2].removeprefix("data: "))["error"]["message"] + else: + assert "data: [DONE]" not in result.text + assert frames[-1].startswith("event: error\ndata: "), result.text + error: Final = json.loads(frames[-1].split("data: ", 1)[1])["error"] + assert error["code"] == "429" + assert message in error["message"] + + loop: Final = AutojumpClockLoop() + try: + loop.run_until_complete(request_and_check()) + finally: + loop.run_until_complete(loop.shutdown_asyncgens()) + loop.close() + + _COMPACTION_ITEM: Final = { "type": "compaction", "encrypted_content": json.dumps({"type": "compaction", "content": "the cat is named Biscuit"}), @@ -1478,6 +1648,25 @@ async def test_a_streamed_read_cut_before_its_terminal_event_ends_in_a_responses assert "before any terminal event" in events[-1]["response"]["error"]["message"] +def _resumed_wire() -> str: + return "".join(f"event: {event['type']}\ndata: {json.dumps(event)}\n\n" for event in _resumed_events("resp_read")) + + +@pytest.mark.asyncio +async def test_a_streamed_read_whose_first_event_is_a_named_error_answers_with_its_status( + monkeypatch: pytest.MonkeyPatch, +): + monkeypatch.setattr(litellm, "callbacks", [_FirstFrameHook(_NAMED_ERROR_FRAME)]) + result, _ = await _read_through_proxy( + monkeypatch, + {"stream": "true"}, + httpx.Response(200, content=_resumed_wire(), headers={"Content-Type": "text/event-stream"}), + ) + + assert result.status_code == 529, result.text + assert result.json()["error"]["message"] == "denied" + + @pytest.mark.asyncio async def test_a_plain_read_through_the_proxy_still_answers_json(monkeypatch: pytest.MonkeyPatch): result, upstream_requests = await _read_through_proxy( @@ -5885,9 +6074,11 @@ class TestGuardrailBlockedResponsesShape: assert response.status_code == 200, response.text assert response.headers["content-type"].startswith("text/event-stream") frames = _sse_data_frames(response.text) - assert frames[-1] == "[DONE]" - events = [json.loads(frame) for frame in frames[:-1]] + assert "[DONE]" not in frames + events = [json.loads(frame) for frame in frames] types = [event["type"] for event in events] + assert types[-1] == "response.completed" + assert [line[7:] for line in response.text.splitlines() if line.startswith("event: ")] == types assert "response.created" in types completed = [event for event in events if event["type"] == "response.completed"] assert len(completed) == 1 @@ -5914,7 +6105,8 @@ class TestGuardrailBlockedResponsesShape: assert response.status_code == 200, response.text frames = _sse_data_frames(response.text) - completed = [json.loads(frame) for frame in frames[:-1] if json.loads(frame)["type"] == "response.completed"] + assert "[DONE]" not in frames + completed = [json.loads(frame) for frame in frames if json.loads(frame)["type"] == "response.completed"] usage = completed[0]["response"]["usage"] assert usage["input_tokens"] == 14 assert usage["output_tokens"] == 20 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 d6df7a6c69a..01fb8eac842 100644 --- a/tests/unit/proxy/response_polling/test_response_polling_handler.py +++ b/tests/unit/proxy/response_polling/test_response_polling_handler.py @@ -16,12 +16,13 @@ following the OpenAI Response API format. import asyncio import json +from collections.abc import Mapping from datetime import datetime, timezone from typing import Any, Dict, Final, Optional from unittest.mock import AsyncMock, Mock, patch import pytest -from fastapi import Request +from fastapi import HTTPException, Request from pydantic import ValidationError from litellm.proxy.response_polling.polling_handler import ResponsePollingHandler @@ -1524,6 +1525,395 @@ class TestBackgroundStreamingTerminalEvents: assert final_call.kwargs["status"] == "failed" assert final_call.kwargs["error"] == error_payload + @pytest.mark.asyncio + async def test_a_late_failure_after_keepalive_pings_sets_failed_status_and_error(self): + from litellm.proxy.common_request_processing import RESPONSES_LATE_SSE_FRAMING + from litellm.proxy.response_polling.background_streaming import ( + background_streaming_task, + ) + + late_frames: Final = RESPONSES_LATE_SSE_FRAMING.failure(HTTPException(status_code=429, detail="rate limited")) + + async def _body_iterator(): + yield b": ping\n\n" + for frame in late_frames: + yield frame.encode() + + mock_response: Final = Mock() + mock_response.body_iterator = _body_iterator() + handler: Final = AsyncMock(spec=ResponsePollingHandler) + kwargs: Final = _make_background_streaming_kwargs("poll_late_failure", handler) + + with patch( # test-quality-ok: the processor is built inside the task, same idiom as the sibling tests + "litellm.proxy.response_polling.background_streaming.ProxyBaseLLMRequestProcessing" + ) as MockProcessor: + MockProcessor.return_value.base_process_llm_request = AsyncMock( + return_value=mock_response + ) + await background_streaming_task(**kwargs) + + final_call: Final = handler.update_state.call_args_list[-1] + assert final_call.kwargs["status"] == "failed" + assert final_call.kwargs["error"] == {"code": "rate_limit_exceeded", "message": "rate limited"} + + @pytest.mark.asyncio + @pytest.mark.parametrize( + ("data_type", "nested_error", "expected_error"), + [ + ( + {"type": "error"}, + {"type": "guardrail_error", "code": "400", "message": "output blocked by policy"}, + {"code": "invalid_request_error", "message": "output blocked by policy"}, + ), + ({"type": "error"}, "output blocked by policy", {"code": "server_error", "message": "output blocked by policy"}), + ( + {}, + {"code": "400", "message": "output blocked by policy"}, + {"code": "invalid_request_error", "message": "output blocked by policy"}, + ), + ( + {"type": "response.output_text.delta"}, + {"code": "400", "message": "output blocked by policy"}, + {"code": "invalid_request_error", "message": "output blocked by policy"}, + ), + ], + ids=["typed_object", "typed_string", "untyped_object", "delta_typed_object"], + ) + async def test_a_nested_only_error_event_keeps_its_reason( + self, + data_type: Mapping[str, str], + nested_error: Mapping[str, str] | str, + expected_error: Mapping[str, str], + ): + from litellm.proxy.response_polling.background_streaming import ( + background_streaming_task, + ) + + error_event: Final = {**data_type, "sequence_number": 9, "error": nested_error} + + async def _body_iterator(): + yield b'data: {"type": "response.in_progress"}\n\n' + yield f"event: error\ndata: {json.dumps(error_event)}\n\n".encode() + + mock_response: Final = Mock() + mock_response.body_iterator = _body_iterator() + handler: Final = AsyncMock(spec=ResponsePollingHandler) + kwargs: Final = _make_background_streaming_kwargs("poll_nested_error", handler) + + with patch( # test-quality-ok: the processor is built inside the task, same idiom as the sibling tests + "litellm.proxy.response_polling.background_streaming.ProxyBaseLLMRequestProcessing" + ) as MockProcessor: + MockProcessor.return_value.base_process_llm_request = AsyncMock( + return_value=mock_response + ) + await background_streaming_task(**kwargs) + + final_call: Final = handler.update_state.call_args_list[-1] + assert final_call.kwargs["status"] == "failed" + assert final_call.kwargs["error"] == expected_error + + @pytest.mark.asyncio + @pytest.mark.parametrize( + ("chunks", "expected_status", "expected_error", "expected_usage"), + [ + ( + ( + 'data: {"type": "response.output_text.delta", "delta": "hi"}\n\n' + 'event: error\ndata: {"error": {"code": "400", "message": "output blocked by policy"}}\n\n', + ), + "failed", + {"code": "invalid_request_error", "message": "output blocked by policy"}, + None, + ), + ( + ( + 'event: error\ndata: {"error": {"code": "400", "message": "output blocked by policy"}}\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\n' + 'data: {"type": "response.completed", "response": {"status": "completed", "usage": {"input_tokens": 3}}}\n\n', + ), + "completed", + None, + {"input_tokens": 3}, + ), + ( + ( + 'event: error\ndata: {"error": {"code": "' + "9" * 4400 + '", "message": "huge code"}}\n\n', + 'data: {"type": "response.completed", "response": {"status": "completed", "usage": {"input_tokens": 3}}}\n\n', + ), + "completed", + None, + {"input_tokens": 3}, + ), + ( + ( + 'data: {"type": "response.output_item.done", "output_index": 0, "item": "hook-text"}\n\n', + 'data: {"type": "response.completed", "response": {"status": "completed", "usage": {"input_tokens": 3}}}\n\n', + ), + "completed", + None, + {"input_tokens": 3}, + ), + ( + ( + 'event: error\ndata: {"error": {"code": "400", "message": "output blocked by policy"}}\n\n' + 'event: response.completed\ndata: {"type": "response.completed", "response": {"status": "completed"}}\n\n', + ), + "completed", + None, + None, + ), + ( + ( + 'data: {"type": "response.output_text.delta", "delta": "hi"}\n\n', + 'data: [DONE]\n\nevent: error\ndata: {"error": {"code": "400", "message": "output blocked by policy"}}\n\n', + ), + "completed", + None, + None, + ), + (('data: {"type": "response.failed", "response": "hook-text"}\n\n',), "failed", None, None), + (('data: {"type": "response.failed", "response": null}\n\n',), "failed", None, None), + ( + ( + 'event: error\ndata: {"error": {"code": "400", "message": "output blocked by policy"}}\n\n' + 'event: response.incomplete\ndata: {"type": "response.incomplete", ' + '"response": {"status": "incomplete", "usage": {"input_tokens": 3}}}\n\n', + ), + "incomplete", + None, + {"input_tokens": 3}, + ), + ( + ( + 'event: error\ndata: {"error": {"code": "400", "message": "output blocked by policy"}}\n\n', + 'data: {"type": "response.output_text.delta", "delta": "hi"}\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\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, + ), + ( + ( + '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: {"type": "response.failed", "response": {"status": "failed", ' + '"error": {"code": "server_error", "message": "first"}, "usage": {"input_tokens": 3}}}\n\n', + 'data: {"type": "response.completed", "response": {"status": "completed", "usage": {"input_tokens": 9}}}\n\n', + ), + "failed", + {"code": "server_error", "message": "first"}, + {"input_tokens": 3}, + ), + ], + ids=[ + "error_frame_after_a_delta", + "error_frame_then_completed", + "unspaced_delta_then_completed", + "error_code_past_the_int_digit_limit_then_completed", + "output_item_done_with_a_string_item_then_completed", + "error_frame_then_completed_in_one_chunk", + "error_frame_after_done_in_one_chunk", + "failed_event_whose_response_is_a_string", + "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", + "completed_then_failed_in_the_next_chunk", + "failed_then_completed_in_the_next_chunk", + ], + ) + async def test_a_poll_stores_what_its_stream_ended_with( + self, + chunks: tuple[str, ...], + expected_status: str, + expected_error: Mapping[str, str] | None, + expected_usage: Mapping[str, int] | None, + ): + from litellm.proxy.response_polling.background_streaming import ( + background_streaming_task, + ) + + async def _body_iterator(): + yield b'data: {"type": "response.in_progress"}\n\n' + for chunk in chunks: + yield chunk.encode() + + mock_response: Final = Mock() + mock_response.body_iterator = _body_iterator() + handler: Final = AsyncMock(spec=ResponsePollingHandler) + kwargs: Final = _make_background_streaming_kwargs("poll_stream_end", handler) + + with patch( # test-quality-ok: the processor is built inside the task, same idiom as the sibling tests + "litellm.proxy.response_polling.background_streaming.ProxyBaseLLMRequestProcessing" + ) as MockProcessor: + MockProcessor.return_value.base_process_llm_request = AsyncMock( + return_value=mock_response + ) + await background_streaming_task(**kwargs) + + final_call: Final = handler.update_state.call_args_list[-1] + assert final_call.kwargs["status"] == expected_status + assert final_call.kwargs["error"] == expected_error + assert final_call.kwargs["usage"] == expected_usage + + @pytest.mark.asyncio + async def test_a_poll_stores_the_final_output_items_beside_one_that_is_not_an_object(self): + from litellm.proxy.response_polling.background_streaming import ( + background_streaming_task, + ) + + final_item: Final = {"id": "msg_1", "content": [{"type": "output_text", "text": "final"}]} + draft_item: Final = {"id": "msg_1", "content": [{"type": "output_text", "text": ""}]} + completed: Final = {"type": "response.completed", "response": {"status": "completed", "output": ["hook-text", final_item]}} + + async def _body_iterator(): + yield f"data: {json.dumps({'type': 'response.output_item.added', 'item': draft_item})}\n\n".encode() + yield b'data: {"type": "response.output_text.delta", "item_id": "msg_1", "delta": "draft"}\n\n' + yield f"data: {json.dumps(completed)}\n\n".encode() + + mock_response: Final = Mock() + mock_response.body_iterator = _body_iterator() + handler: Final = AsyncMock(spec=ResponsePollingHandler) + kwargs: Final = _make_background_streaming_kwargs("poll_final_output", handler) + + with patch( # test-quality-ok: the processor is built inside the task, same idiom as the sibling tests + "litellm.proxy.response_polling.background_streaming.ProxyBaseLLMRequestProcessing" + ) as MockProcessor: + MockProcessor.return_value.base_process_llm_request = AsyncMock( + return_value=mock_response + ) + await background_streaming_task(**kwargs) + + stored_outputs: Final = [call.kwargs["output"] for call in handler.update_state.call_args_list if "output" in call.kwargs] + assert stored_outputs[-1] == [final_item] + + @pytest.mark.asyncio + @pytest.mark.parametrize( + ("status_code", "body", "expected_error"), + [ + ( + 404, + json.dumps( + {"error": {"message": "model not found", "type": "invalid_request_error", "param": None, "code": "404"}} + ), + {"code": "not_found_error", "message": "model not found"}, + ), + (429, "rate limited", {"code": "rate_limit_exceeded", "message": "Too Many Requests"}), + ( + 403, + json.dumps({"error": {"type": "guardrail_scan_error", "code": 403}}), + {"code": "permission_error", "message": "Forbidden"}, + ), + ( + 404, + json.dumps({"error": {"message": "model not found", "type": "invalid_request_error", "param": None}}), + {"code": "not_found_error", "message": "model not found"}, + ), + ], + ) + async def test_a_first_chunk_error_polls_failed_like_a_late_one( + self, status_code: int, body: str, expected_error: Mapping[str, str] + ): + from fastapi import Response + + from litellm.proxy.response_polling.background_streaming import ( + background_streaming_task, + ) + + first_chunk_error: Final = Response(status_code=status_code, content=body) + handler: Final = AsyncMock(spec=ResponsePollingHandler) + kwargs: Final = _make_background_streaming_kwargs("poll_first_chunk_error", handler) + + with patch( # test-quality-ok: the processor is built inside the task, same idiom as the sibling tests + "litellm.proxy.response_polling.background_streaming.ProxyBaseLLMRequestProcessing" + ) as MockProcessor: + MockProcessor.return_value.base_process_llm_request = AsyncMock( + return_value=first_chunk_error + ) + await background_streaming_task(**kwargs) + + final_call: Final = handler.update_state.call_args_list[-1] + assert final_call.kwargs["status"] == "failed" + assert final_call.kwargs["error"] == expected_error + + @pytest.mark.asyncio + async def test_a_non_streamed_success_response_still_polls_completed(self): + from fastapi.responses import JSONResponse + + from litellm.proxy.response_polling.background_streaming import ( + background_streaming_task, + ) + + handler: Final = AsyncMock(spec=ResponsePollingHandler) + kwargs: Final = _make_background_streaming_kwargs("poll_non_streamed_success", handler) + + with patch( # test-quality-ok: the processor is built inside the task, same idiom as the sibling tests + "litellm.proxy.response_polling.background_streaming.ProxyBaseLLMRequestProcessing" + ) as MockProcessor: + MockProcessor.return_value.base_process_llm_request = AsyncMock( + return_value=JSONResponse(status_code=200, content={"id": "resp_123", "status": "completed"}) + ) + await background_streaming_task(**kwargs) + + final_call: Final = handler.update_state.call_args_list[-1] + assert final_call.kwargs["status"] == "completed" + assert final_call.kwargs["error"] is None + @pytest.mark.asyncio async def test_response_incomplete_sets_incomplete_status_and_details(self): """Test that a response.incomplete stream event results in incomplete status""" diff --git a/tests/unit/proxy/test_common_request_processing.py b/tests/unit/proxy/test_common_request_processing.py index 5c27ad8ce77..c59cde6d630 100644 --- a/tests/unit/proxy/test_common_request_processing.py +++ b/tests/unit/proxy/test_common_request_processing.py @@ -4,7 +4,7 @@ import copy import datetime import gzip import json -from collections.abc import AsyncIterable, Mapping +from collections.abc import AsyncIterable, Awaitable, Mapping from types import MappingProxyType, SimpleNamespace from typing import AsyncGenerator, Callable, Final, Iterator, Literal, Optional, Sequence from urllib.parse import unquote_plus @@ -41,6 +41,7 @@ from litellm.constants import ( STREAM_SSE_KEEPALIVE_PING_BYTES, ) from litellm.integrations.custom_logger import CustomLogger +from litellm.integrations.custom_guardrail import ModifyResponseException from litellm.integrations.opentelemetry import UserAPIKeyAuth from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache from litellm.proxy.common_request_processing import ( @@ -53,11 +54,15 @@ from litellm.proxy.common_request_processing import ( _ClientDisconnectedBeforeFirstChunk, attach_guardrail_information, _extract_error_from_sse_chunk, + FirstChunkSseFraming, _get_cost_breakdown_from_logging_obj, CostBreakdownHeaderValues, _has_attribute_error_in_chain, + OPENAI_FIRST_CHUNK_FRAMING, + RESPONSES_FIRST_CHUNK_FRAMING, include_guardrail_response_requested, is_azure_model_router_request, + RESPONSES_LATE_SSE_FRAMING, open_sse_before_first_byte, resolve_litellm_call_id, ttft_keepalive_interval, @@ -90,6 +95,7 @@ from litellm.router import Router from litellm.litellm_core_utils.streaming_handler import CustomStreamWrapper from litellm.utils import ModelResponseListIterator from litellm.litellm_core_utils.add_retry_fallback_headers import prepare_response_for_header_attachment +from tests._support.clock_loops import AutojumpClockLoop def test_attach_guardrail_information_copies_recorded_entries_onto_model_response(): @@ -2258,6 +2264,14 @@ class TestProxyBaseLLMRequestProcessing: ) +def _fail_stream_setup() -> str: + raise RuntimeError("stream setup failed") + + +async def _stream_failing_before_its_first_chunk() -> AsyncGenerator[str, None]: + yield _fail_stream_setup() + + @pytest.mark.asyncio class TestCommonRequestProcessingHelpers: async def consume_stream(self, streaming_response: StreamingResponse) -> list: @@ -2336,6 +2350,74 @@ class TestCommonRequestProcessingHelpers: assert body["error"]["code"] == 403 assert body["error"]["message"] == "forbidden" + @pytest.mark.parametrize("newline", ["\n", "\r\n", "\r"]) + async def test_a_responses_stream_whose_first_event_is_a_named_error_answers_as_json(self, newline: str): + error: Final = {"type": "guardrail_error", "code": "400", "message": "blocked", "param": None} + event: Final = json.dumps({"type": "error", "sequence_number": 0, "error": error}) + + async def responses_stream(): + yield f"event: error{newline}data: {event}{newline}{newline}" + + response: Final = await create_response( + responses_stream(), "text/event-stream", {}, framing=RESPONSES_FIRST_CHUNK_FRAMING + ) + + assert isinstance(response, JSONResponse) + assert response.status_code == 400 + assert json.loads(response.body) == {"error": error} + + async def test_a_first_event_error_with_a_unicode_line_separator_in_its_message_answers_as_json(self): + error: Final = {"type": "guardrail_error", "code": "400", "message": "blocked\u2028by policy", "param": None} + + async def responses_stream(): + event: Final = json.dumps({"type": "error", "sequence_number": 0, "error": error}, ensure_ascii=False) + yield f"event: error\ndata: {event}\n\n" + + response: Final = await create_response( + responses_stream(), "text/event-stream", {}, framing=RESPONSES_FIRST_CHUNK_FRAMING + ) + + assert isinstance(response, JSONResponse) + assert response.status_code == 400 + assert json.loads(response.body) == {"error": error} + + @pytest.mark.parametrize( + "framing,frame_starts", + [ + (OPENAI_FIRST_CHUNK_FRAMING, ('data: {"error": ', "data: [DONE]")), + (RESPONSES_FIRST_CHUNK_FRAMING, ('event: error\ndata: {"type": "error", "code": "server_error"',)), + ], + ) + async def test_a_stream_that_fails_before_its_first_chunk_ends_in_its_own_api_shape( + self, framing: FirstChunkSseFraming, frame_starts: tuple[str, ...] + ) -> None: + response: Final = await create_response( + _stream_failing_before_its_first_chunk(), "text/event-stream", {}, framing=framing + ) + frames: Final = await self.consume_stream(response) + + assert response.status_code == 500 + assert len(frames) == len(frame_starts) + assert all(frame.startswith(start) for frame, start in zip(frames, frame_starts, strict=True)) + + async def test_a_responses_stream_whose_first_frame_names_a_success_code_streams_on(self): + frames: Final = ( + 'data: {"error": {"message": "note", "type": "guardrail_scan_error", "code": 200}}\n\n', + 'event: response.completed\ndata: {"type": "response.completed", "sequence_number": 1}\n\n', + ) + + async def responses_stream(): + for frame in frames: + yield frame + + response: Final = await create_response( + responses_stream(), "text/event-stream", {}, framing=RESPONSES_FIRST_CHUNK_FRAMING + ) + + assert isinstance(response, StreamingResponse) + assert response.status_code == status.HTTP_200_OK + assert tuple(await self.consume_stream(response)) == frames + async def test_create_streaming_response_first_chunk_not_error(self): async def mock_generator(): yield 'data: {"content": "first part"}\n\n' @@ -10905,6 +10987,104 @@ async def test_ttft_keepalive_relays_a_late_non_streaming_body_as_an_sse_frame() assert collected[-1] == b"data: [DONE]\n\n" +async def _slow_rate_limit_failure() -> Response: + await asyncio.sleep(0.2) + raise HTTPException(status_code=429, detail="rate limited") + + +async def _slow_guardrail_block_body() -> Response: + await asyncio.sleep(0.2) + return JSONResponse( + status_code=400, content={"error": {"message": "blocked", "type": "guardrail_error", "code": "400"}} + ) + + +_POLICY_BLOCK_DETAIL: Final = MappingProxyType( + {"error": {"message": "blocked by policy", "type": "guardrail_error"}, "guardrail": "content-filter"} +) + + +async def _slow_codeless_rate_limit_body() -> Response: + await asyncio.sleep(0.2) + return JSONResponse(status_code=429, content={"error": {"message": "slow down", "code": None}}) + + +async def _slow_dict_detail_failure() -> Response: + await asyncio.sleep(0.2) + raise HTTPException(status_code=400, detail=dict(_POLICY_BLOCK_DETAIL)) + + +@pytest.mark.parametrize( + "produce,expected_code,expected_error", + [ + (_slow_rate_limit_failure, "rate_limit_exceeded", {"message": "rate limited", "code": "429"}), + (_slow_codeless_rate_limit_body, "rate_limit_exceeded", {"message": "slow down", "code": "429"}), + ( + _slow_guardrail_block_body, + "invalid_request_error", + {"message": "blocked", "type": "guardrail_error", "code": "400"}, + ), + ( + _slow_dict_detail_failure, + "invalid_request_error", + {"message": "blocked by policy", "code": "400", "provider_specific_fields": dict(_POLICY_BLOCK_DETAIL)}, + ), + ], +) +def test_ttft_keepalive_ends_a_late_responses_failure_with_an_error_event_and_no_done( + produce: Callable[[], Awaitable[Response]], expected_code: str, expected_error: Mapping[str, object] +) -> None: + async def collect_and_check() -> None: + response: Final = await open_sse_before_first_byte( + produce(), ping_interval_seconds=0.05, late_framing=RESPONSES_LATE_SSE_FRAMING + ) + collected: Final = await _drain(response) + + assert collected[0] == TTFT_PING + assert b"data: [DONE]\n\n" not in collected + event_line, data_line = collected[-1].decode().strip().split("\n") + event: Final = json.loads(data_line.removeprefix("data: ")) + assert event_line == "event: error" + assert event["type"] == "error" + assert (event["code"], event["message"]) == (expected_code, expected_error["message"]) + assert {key: event["error"].get(key) for key in expected_error} == expected_error + + loop: Final = AutojumpClockLoop() + try: + loop.run_until_complete(collect_and_check()) + finally: + loop.run_until_complete(loop.shutdown_asyncgens()) + loop.close() + + +async def _slow_passthrough_block() -> Response: + await asyncio.sleep(0.2) + raise ModifyResponseException(message="blocked by policy", model="gpt-4o", request_data={}) + + +def test_ttft_keepalive_ends_a_late_responses_passthrough_block_with_the_block_response() -> None: + async def collect_and_check() -> None: + response: Final = await open_sse_before_first_byte( + _slow_passthrough_block(), ping_interval_seconds=0.05, late_framing=RESPONSES_LATE_SSE_FRAMING + ) + collected: Final = await _drain(response) + frames: Final = tuple(chunk.decode() for chunk in collected if chunk != TTFT_PING) + names: Final = tuple(frame.split("\n", 1)[0] for frame in frames) + completed: Final = json.loads(frames[-1].split("\n", 1)[1].removeprefix("data: ")) + + assert collected[0] == TTFT_PING + assert (names[0], names[-1]) == ("event: response.created", "event: response.completed") + assert "event: error" not in names + assert completed["response"]["output"][0]["content"][0]["text"] == "blocked by policy" + + loop: Final = AutojumpClockLoop() + try: + loop.run_until_complete(collect_and_check()) + finally: + loop.run_until_complete(loop.shutdown_asyncgens()) + loop.close() + + @pytest.mark.asyncio async def test_ttft_keepalive_closes_the_upstream_stream_it_relayed(): """Starlette never calls the produced response, so its own cleanup never runs diff --git a/tests/unit/proxy/test_response_model_sanitization.py b/tests/unit/proxy/test_response_model_sanitization.py index b3ad53221eb..2f0bcd56e79 100644 --- a/tests/unit/proxy/test_response_model_sanitization.py +++ b/tests/unit/proxy/test_response_model_sanitization.py @@ -734,7 +734,10 @@ def test_restamp_streaming_chunk_rewrites_nested_model_in_guardrail_block_frames )[0] for frame in frames ] - events = [json.loads(frame.decode().removeprefix("data: ")) for frame in client_frames] + events = [ + json.loads(next(line[6:] for line in frame.decode().splitlines() if line.startswith("data: "))) + for frame in client_frames + ] assert all(frame.endswith(b"\n\n") for frame in client_frames) assert events[-1]["type"] == "response.completed" diff --git a/tests/unit/responses/test_streaming_iterator.py b/tests/unit/responses/test_streaming_iterator.py index 5072f3d3108..b1859ddf33c 100644 --- a/tests/unit/responses/test_streaming_iterator.py +++ b/tests/unit/responses/test_streaming_iterator.py @@ -25,6 +25,7 @@ from litellm.integrations.custom_logger import CustomLogger from litellm.litellm_core_utils.core_helpers import UPSTREAM_USAGE_COST, billed_usage_cost from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.llms.base_llm.responses.transformation import BaseResponsesAPIConfig +from litellm.llms.openai.responses.transformation import OpenAIResponsesAPIConfig from litellm.responses.streaming_iterator import ( CachedResponsesAPIStreamingIterator, MockResponsesAPIStreamingIterator, @@ -415,6 +416,20 @@ async def test_complete_stream_still_ends_normally(trailer): assert logging_obj.async_failure_handler.await_count == 0 +@pytest.mark.asyncio +async def test_event_without_a_type_reaches_the_caller_instead_of_ending_the_stream(): + iterator: Final = _make_iterator( + sse_events=[_CREATED_EVENT, _sse_event({"note": "hook output without a type"})], + logging_obj=_failure_tracking_logging_obj(), + config=OpenAIResponsesAPIConfig(), + ) + + await iterator.__anext__() + untyped: Final = await iterator.__anext__() + + assert untyped.model_extra == {"note": "hook output without a type"} + + @pytest.mark.parametrize("trailing_error", _TRAILING_ERRORS, ids=type) def test_sync_transport_error_before_any_output_raises_fallback_error(trailing_error): logging_obj: Final = _failure_tracking_logging_obj()