diff --git a/litellm/llms/openai/responses/guardrail_translation/handler.py b/litellm/llms/openai/responses/guardrail_translation/handler.py index 559daf767c2..c8c1b318fe0 100644 --- a/litellm/llms/openai/responses/guardrail_translation/handler.py +++ b/litellm/llms/openai/responses/guardrail_translation/handler.py @@ -263,6 +263,7 @@ _TOOL_SEARCH_CALL_TYPE: Final = "tool_search_call" _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"} ) @@ -1936,7 +1937,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 ()) @@ -1944,6 +1944,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 c9b6fff3a00..15c7325eaf8 100644 --- a/litellm/proxy/common_request_processing.py +++ b/litellm/proxy/common_request_processing.py @@ -51,7 +51,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.bug_report import ( @@ -92,6 +92,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.auth.auth_checks import ( can_key_call_resolved_model, @@ -119,6 +120,7 @@ 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 ResponsesStreamErrorState, sse_error_object from litellm.proxy.common_utils.sse_keepalive import ( SSE_COMMENT_PING_BYTES, SSE_STREAM_START_TAIL, @@ -1004,11 +1006,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_text(event_line).strip().split("\n"): + 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: @@ -1039,18 +1062,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", @@ -1059,8 +1081,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 @@ -1268,6 +1290,70 @@ 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 _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) +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, ...]] + + +OPENAI_FIRST_CHUNK_FRAMING: Final = FirstChunkSseFraming(data_payload=_sse_data_payload, error_frames=_sse_error_frames) +RESPONSES_FIRST_CHUNK_FRAMING: Final = FirstChunkSseFraming( + data_payload=_responses_event_data_payload, error_frames=_responses_error_frames +) + + +@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}) @@ -1296,6 +1382,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. @@ -1331,7 +1418,9 @@ 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) + 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: # First chunk is an error, stream hasn't really started yet # Should return standard JSON error response instead of SSE format @@ -1343,7 +1432,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)), ) @@ -1385,7 +1476,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( @@ -1468,13 +1559,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: @@ -1516,6 +1607,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_framing: LateSseFraming = OPENAI_LATE_SSE_FRAMING, ) -> _LateResponseT | StreamingResponse: """Write SSE keepalive comments while the upstream LLM call is still in flight. @@ -1561,11 +1653,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. - _, error_obj = sse_error_payload(await _sanitized_late_failure(exc, on_late_failure)) - for frame in _sse_error_frames(error_obj): + 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(): @@ -2963,6 +3054,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. @@ -3003,9 +3095,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_framing=sse_framing.late, ) async def _process_llm_request( @@ -3029,6 +3123,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 @@ -3312,6 +3407,7 @@ class ProxyBaseLLMRequestProcessing: request=request, refresh_headers=refresh_stream_headers, general_settings=general_settings, + framing=first_chunk_framing, ) ### CALL HOOKS ### - modify outgoing data @@ -3551,8 +3647,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..951650543a7 100644 --- a/litellm/proxy/common_utils/responses_stream_errors.py +++ b/litellm/proxy/common_utils/responses_stream_errors.py @@ -1,10 +1,14 @@ +import json import time -from collections.abc import Mapping +from collections import Counter, OrderedDict, defaultdict +from collections.abc import Iterable, Iterator, Mapping, Sequence +from enum import Enum from http import HTTPStatus from types import MappingProxyType from typing import Final -from pydantic import BaseModel, ConfigDict, field_validator +from pydantic import BaseModel, ConfigDict, 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 @@ -57,6 +61,167 @@ 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 + + +_DROPPED: Final = object() + + +class _BaseForm(Enum): + PRINTED = "printed" + DUMPED = "dumped" + + +def _prints_every_item(kind: type[object], form: _BaseForm) -> bool: + repr_owner: Final = next((owner for owner in kind.__mro__ if "__repr__" in owner.__dict__), object) + if form is _BaseForm.DUMPED: + return kind.__str__ is object.__str__ and repr_owner in (dict, list, tuple, OrderedDict, defaultdict, Counter) + return ( + kind.__str__ is object.__str__ + and kind.__format__ is object.__format__ + and (kind in (OrderedDict, Counter) or repr_owner in (dict, list, tuple, defaultdict)) + ) + + +def _is_shown_dict(value: object, form: _BaseForm) -> TypeIs[Mapping[object, object]]: # guard-ok: dict items as object + kind: Final = type(value) + return isinstance(value, dict) and _prints_every_item(kind, form) + + +def _is_shown_sequence(value: object, form: _BaseForm) -> TypeIs[Sequence[object]]: # guard-ok: items as object + kind: Final = type(value) + return isinstance(value, (list, tuple)) and _prints_every_item(kind, form) + + +def _is_proxy(value: object) -> TypeIs[Mapping[object, object]]: # guard-ok: mappingproxy cannot be subclassed + return isinstance(value, MappingProxyType) + + +def _printed_proxy_items(proxy: Mapping[object, object], printed: str) -> Iterable[tuple[object, object]] | None: + snapshot: Final = dict(proxy) + shown: Final = repr(snapshot) + return snapshot.items() if printed in (shown, f"mappingproxy({shown})") else None + + +def _repr_prints_its_fields(model: BaseModel) -> bool: + kind: Final = type(model) + return kind.__repr__ is BaseModel.__repr__ and kind.__repr_str__ is BaseModel.__repr_str__ + + +def _plain_scalar(value: str | float) -> str | float: + if isinstance(value, str): + return str.__str__(value) + if isinstance(value, int): + return int.__int__(value) + return float.__float__(value) + + +def _repr_shows_its_value(value: str | float) -> bool: + plain: Final = _plain_scalar(value) + if isinstance(value, Enum): + return repr(value) == f"<{type(value).__name__}.{value.name}: {plain!r}>" + return repr(value) == repr(plain) + + +def _text(value: object, form: _BaseForm) -> object: + try: + if isinstance(value, (str, int, float)) and (form is _BaseForm.DUMPED or _repr_shows_its_value(value)): + return _plain_scalar(value) + return repr(value) if form is _BaseForm.PRINTED else str(value) + except Exception: # noqa: BLE001 # hook values run their own __repr__ and __str__, so a failure drops the value + return _DROPPED + + +def _json_key(key: object, form: _BaseForm) -> object: + return key if key is None or type(key) in (str, int, float, bool) else _text(key, form) + + +def _shown_sequence(value: object, form: _BaseForm) -> Iterator[object] | None: + if not _is_shown_sequence(value, form): + return None + if form is _BaseForm.PRINTED and isinstance(value, list): + return list[object].__iter__(value) + if form is _BaseForm.PRINTED and isinstance(value, tuple): + return tuple[object, ...].__iter__(value) + return iter(value) + + +def _shown_items(value: object, form: _BaseForm) -> Iterable[tuple[object, object]] | None: + if not _is_shown_dict(value, form): + return None + if form is _BaseForm.PRINTED and isinstance(value, dict): + return dict[object, object].items(value) + return value.items() + + +def _expansion(value: object, form: _BaseForm) -> tuple[Iterable[tuple[object, object]], _BaseForm] | None: + shown_items: Final = _shown_items(value, form) + if shown_items is not None: + return shown_items, form + if _is_proxy(value): + proxy_items: Final = _printed_proxy_items(value, repr(value) if form is _BaseForm.PRINTED else str(value)) + return None if proxy_items is None else (proxy_items, _BaseForm.PRINTED) + if form is _BaseForm.PRINTED and isinstance(value, BaseModel) and _repr_prints_its_fields(value): + return value.__repr_args__(), _BaseForm.PRINTED + return None + + +def _json_safe(value: object, form: _BaseForm, path: frozenset[int]) -> object: + if value is None or type(value) in (str, int, float, bool): + return value + if id(value) in path: + return _text(value, form) + inner_path: Final = path | frozenset((id(value),)) + try: + sequence: Final = _shown_sequence(value, form) + if sequence is not None: + shown_items: Final = [_json_safe(item, form, inner_path) for item in sequence] + return [item for item in shown_items if item is not _DROPPED] + expansion: Final = _expansion(value, form) + if expansion is None: + return _text(value, form) + items, inner = expansion + pairs: Final = [(_json_key(key, inner), _json_safe(item, inner, inner_path)) for key, item in items] + return {key: item for key, item in pairs if key is not _DROPPED and item is not _DROPPED} + except Exception: # noqa: BLE001 # hook values run their own printers and can nest past the recursion limit + return _text(value, form) + + +def responses_chunk_json(chunk: object) -> str | None: + if _is_proxy(chunk) and _printed_proxy_items(chunk, format(chunk, "")) is None: + return None + if not _is_shown_dict(chunk, _BaseForm.PRINTED) and not _is_proxy(chunk): + return None + shown: Final = _json_safe(chunk, _BaseForm.PRINTED, frozenset()) + return None if isinstance(shown, str) or shown is _DROPPED else json.dumps(shown) + + +def _dumped_detail(error: Mapping[str, object]) -> Mapping[str, object]: + path: Final = frozenset((id(error),)) + shown: Final = [(key, _json_safe(value, _BaseForm.DUMPED, path)) for key, value in error.items()] + return {key: value for key, value in shown if value is not _DROPPED} + + +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 +260,12 @@ 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 + return int(code) if code is not None and code.isdecimal() else None + + def _response_error_code(details: _FailureDetails) -> str: for value in (details.code, details.type): if value == "insufficient_quota": @@ -106,6 +277,34 @@ 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, + } + + class ResponsesStreamErrorState: def __init__(self) -> None: self.response_id: str | None = None @@ -129,9 +328,22 @@ 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 ("response.completed", "response.failed", "response.incomplete", "error"): self.terminal_emitted = True - return frame + if 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, _dumped_detail(error), sequence_number)) + 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 b9b15d332e9..e50bb9d2891 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -554,7 +554,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_json from litellm.proxy.common_utils.scheduled_job_stagger import ( apply_scheduled_job_stagger, attach_job_timing_logger, @@ -10266,6 +10266,11 @@ 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_json(chunk) + return _format_streaming_sse_chunk(chunk=chunk) if data is None else f"data: {data}\n\n" + + _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 @@ -10505,7 +10510,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( @@ -10664,7 +10671,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: @@ -10690,8 +10697,8 @@ async def async_data_generator( yield (raw_sse_buffer if raw_sse_buffer.endswith(_SSE_FRAME_DELIMITERS) else raw_sse_buffer + "\n\n") 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 089cb75a795..c2e84f585ca 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, @@ -482,6 +483,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, @@ -540,7 +542,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) diff --git a/litellm/proxy/response_polling/background_streaming.py b/litellm/proxy/response_polling/background_streaming.py index 1eb76f8f558..99c988317bc 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 typing import TYPE_CHECKING, Final, Protocol, TypeAlias, runtime_checkable from fastapi import Request, Response @@ -20,7 +21,12 @@ 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, + responses_error_fields, + sse_error_object, +) from litellm.proxy.response_polling.polling_handler import ResponsePollingHandler from litellm.types.llms.openai import ResponsesAPIStatus @@ -69,6 +75,7 @@ class _StreamEvent(TypedDict, total=False): delta: ReadOnly[str] part: ReadOnly[object] response: ReadOnly[_TerminalResponse] + error: ReadOnly[_JsonDict] class _StreamEventParser: @@ -85,6 +92,22 @@ 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 _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 + return responses_error_fields(sse_error_object(bytes(response.body)) or {"code": response.status_code}) + + async def _never_receive() -> Message: await asyncio.Event().wait() raise AssertionError("unreachable") @@ -165,6 +188,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 @@ -326,6 +350,10 @@ async def background_streaming_task( status="in_progress", ) + elif event_type == "error": + terminal_status = "failed" + terminal_error = responses_error_fields(_event_error(event)) + elif event_type in ( "response.completed", "response.failed", @@ -382,14 +410,15 @@ async def background_streaming_task( # 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) + final_error: Final = terminal_error if early_error is None else early_error + final_status: Final = "failed" if early_error is not None else terminal_status or "completed" await polling_handler.update_state( polling_id=polling_id, status=final_status, usage=usage_data, - error=terminal_error, + error=final_error, reasoning=reasoning_data, tool_choice=tool_choice_data, tools=tools_data, @@ -411,7 +440,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 fe224a52cfa..f46ff249688 100644 --- a/litellm/proxy/response_polling/polling_handler.py +++ b/litellm/proxy/response_polling/polling_handler.py @@ -98,7 +98,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/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 d4f9533fcdc..9c73fcc3b6e 100644 --- a/tests/integration/providers/test_bedrock_mantle_claude_chat_wire.py +++ b/tests/integration/providers/test_bedrock_mantle_claude_chat_wire.py @@ -846,7 +846,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 63bdadceebc..6d361492dd6 100644 --- a/tests/integration/streaming/test_responses_top_level_error_events.py +++ b/tests/integration/streaming/test_responses_top_level_error_events.py @@ -256,8 +256,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 d732d922375..2662e736911 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 @@ -4429,7 +4429,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 c37475bd539..6deaac8983a 100644 --- a/tests/unit/proxy/proxy_server/test_streaming_helpers.py +++ b/tests/unit/proxy/proxy_server/test_streaming_helpers.py @@ -16,8 +16,14 @@ Pins covered: from __future__ import annotations import asyncio +import functools import json -from collections.abc import AsyncIterator, Callable +from collections import Counter, OrderedDict +from collections.abc import AsyncIterator, Callable, ItemsView, Iterable, Iterator, Mapping +from datetime import datetime, timezone +from enum import Enum +from http import HTTPStatus +from types import MappingProxyType from typing import Final, Literal from unittest.mock import AsyncMock, MagicMock @@ -26,7 +32,8 @@ import pytest from fastapi import HTTPException, Response from fastapi.responses import JSONResponse, StreamingResponse from openai import APIError as OpenAIAPIError -from pydantic import BaseModel +from openai.types.responses import ResponseOutputMessage, ResponseOutputText +from pydantic import BaseModel, Field, PrivateAttr import litellm import litellm.proxy.proxy_server as ps @@ -34,10 +41,12 @@ 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, _fast_serialize_simple_model_response_stream, _format_fallback_metadata_sse_event, + _format_responses_sse_chunk, _format_streaming_sse_chunk, _get_client_requested_model_for_streaming, _get_streaming_fallback_metadata, @@ -50,11 +59,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 Delta, ModelResponseStream, StreamingChoices, Usage from tests._support.clock_loops import AutojumpClockLoop @@ -92,6 +106,571 @@ 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(), lambda event: MappingProxyType(event.model_dump())], + ids=["model", "dict", "mapping_proxy"], +) +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" + + +class _HookOutputItem(BaseModel, frozen=True): + id: str + type: str + + +@pytest.mark.parametrize( + "item", + [ + MappingProxyType({"id": "msg_nested", "type": "message"}), + _HookOutputItem(id="msg_nested", type="message"), + ], + ids=["mapping_proxy", "model"], +) +def test_responses_sse_chunk_encodes_values_nested_in_a_hook_chunk_as_json_objects( + item: Mapping[str, object] | BaseModel, +) -> None: + frame: Final = _format_responses_sse_chunk( + chunk=MappingProxyType({"type": "response.output_item.added", "output_index": 0, "item": item}) + ) + + assert isinstance(frame, str) + assert json.loads(frame.removeprefix("data: "))["item"] == {"id": "msg_nested", "type": "message"} + + +def test_responses_error_event_encodes_nested_error_fields_as_json_objects() -> None: + frame: Final = ResponsesStreamErrorState().format_error( + { + "message": "overloaded", + "provider_specific_fields": MappingProxyType({"upstream": MappingProxyType({"type": "overloaded_error"})}), + } + ) + + payload: Final = json.loads(frame.split("data: ", 1)[1]) + assert payload["error"]["provider_specific_fields"] == {"upstream": {"type": "overloaded_error"}} + + +_UPSTREAM_HEADERS: Final = httpx.Headers({"authorization": "Bearer hook-canary", "x-request-id": "req_1"}) + + +def _responses_frame_payload(nested: object) -> Mapping[str, object]: + chunk_frame: Final = _format_responses_sse_chunk( + chunk=MappingProxyType({"type": "response.output_text.delta", "nested": nested}) + ) + error_frame: Final = ResponsesStreamErrorState().format_error({"message": "blocked", "nested": nested}) + assert isinstance(chunk_frame, str) + chunk_payload: Final = json.loads(chunk_frame.removeprefix("data: ")) + error_payload: Final = json.loads(error_frame.split("data: ", 1)[1]) + assert chunk_payload["nested"] == error_payload["error"]["nested"] + return chunk_payload + + +class _HiddenFieldModel(BaseModel, frozen=True): + user: str + password: str = Field(repr=False) + + +class _MaskedReprModel(BaseModel, frozen=True): + token: str + + def __repr__(self) -> str: + return "_MaskedReprModel(token=[secure])" + + +class _MaskedReprStrModel(BaseModel, frozen=True): + token: str + + def __repr_str__(self, join_str: str) -> str: + return "token=[secure]" + + +class _MaskedReprArgsModel(BaseModel, frozen=True): + token: str + + def __repr_args__(self) -> Iterable[tuple[str | None, object]]: + return (("token", "[secure]"),) + + +class _MaskedStrModel(BaseModel, frozen=True): + token: str + + def __str__(self) -> str: + return "_MaskedStrModel(token=[secure])" + + +class _MaskedReprString(str): + def __repr__(self) -> str: + return "_MaskedReprString([secure])" + + +_REDACTED_HEADERS: Final = RedactedDict({"authorization": "Bearer hook-canary"}) + +_MASKED_IN_EVERY_FRAME: Final = ( + _UPSTREAM_HEADERS, + _REDACTED_HEADERS, + MappingProxyType({"headers": _REDACTED_HEADERS}), + MappingProxyType(_REDACTED_HEADERS), + _HiddenFieldModel(user="rig", password="hook-canary"), + _MaskedReprStrModel(token="hook-canary"), + _MaskedReprArgsModel(token="hook-canary"), +) +_MASKED_IN_EVERY_FRAME_IDS: Final = ( + "header_mapping", + "redacted_dict", + "proxy_over_dict_holding_redacted_dict", + "proxy_wrapping_redacted_dict", + "repr_false_field", + "masked_repr_str_model", + "masked_repr_args_model", +) + + +class _UnmaskingDict(dict[str, str]): + def items(self) -> ItemsView[str, str]: # pyright: ignore[reportIncompatibleMethodOverride] # rig hides storage + return {key: "hook-canary" for key in dict.keys(self)}.items() + + +class _UnmaskingList(list[str]): + def __iter__(self) -> Iterator[str]: + return iter(("hook-canary",)) + + +class _MostCommonMaskedCounter(Counter[str]): + def most_common(self, n: int | None = None) -> list[tuple[str, int]]: + return [(key, count) for key, count in super().most_common(n) if key != "hook-canary"] + + +@pytest.mark.parametrize( + "nested", + [ + *_MASKED_IN_EVERY_FRAME, + _MaskedReprModel(token="hook-canary"), + _MaskedReprString("hook-canary"), + _UnmaskingDict(token="enc:MASKED"), + _UnmaskingList(("enc:MASKED",)), + _MostCommonMaskedCounter(("shown", "hook-canary")), + ], + ids=[ + *_MASKED_IN_EVERY_FRAME_IDS, + "masked_repr_model", + "masked_repr_string", + "dict_items_differ_from_storage", + "list_iter_differs_from_storage", + "counter_most_common_masks", + ], +) +def test_responses_sse_chunk_never_shows_what_a_nested_value_hides_from_its_repr(nested: object) -> None: + frame: Final = _format_responses_sse_chunk( + chunk=MappingProxyType({"type": "response.output_text.delta", "nested": nested}) + ) + + assert isinstance(frame, str) + assert json.loads(frame.removeprefix("data: "))["type"] == "response.output_text.delta" + assert "hook-canary" not in frame + + +@pytest.mark.parametrize( + "nested", + [*_MASKED_IN_EVERY_FRAME, _MaskedStrModel(token="hook-canary")], + ids=[*_MASKED_IN_EVERY_FRAME_IDS, "masked_str_model"], +) +def test_responses_error_event_never_shows_what_a_nested_value_hides_from_its_string_form(nested: object) -> None: + frame: Final = ResponsesStreamErrorState().format_error( + {"message": "blocked", "provider_specific_fields": {"guardrail": "rig-guard", "detail": nested}} + ) + + payload: Final = json.loads(frame.split("data: ", 1)[1]) + assert payload["error"]["provider_specific_fields"]["guardrail"] == "rig-guard" + assert "hook-canary" not in frame + + +class _OwnFormatDict(dict[str, str]): + def __format__(self, format_spec: str) -> str: + return json.dumps({"type": "response.output_text.delta", "token": "MASKED"}) + + +class _MaskedStrDict(dict[str, str]): + def __str__(self) -> str: + return "{'type': 'response.output_text.delta', 'token': 'MASKED'}" + + +@pytest.mark.parametrize( + "chunk", + [ + _OwnFormatDict(type="response.output_text.delta", token="hook-canary"), + MappingProxyType(_MaskedStrDict(type="response.output_text.delta", token="hook-canary")), + ], + ids=["own_format", "proxy_over_masked_str"], +) +def test_responses_sse_chunk_sends_a_whole_chunk_as_the_text_its_printer_gave(chunk: object) -> None: + frame: Final = _format_responses_sse_chunk(chunk=chunk) + + assert isinstance(frame, str) + assert "hook-canary" not in frame + + +class _RigTuple(tuple[object, ...]): + pass + + +class _RevealingReprModel(BaseModel, frozen=True): + a: int + _note: str = PrivateAttr(default="hook-canary") + + def __repr__(self) -> str: + return f"_RevealingReprModel(note={self._note})" + + +@pytest.mark.parametrize( + "detail", + [ + OrderedDict(member=_MaskedStrModel(token="hook-canary")), + _RigTuple((_MaskedStrModel(token="hook-canary"),)), + _HookOutputItem(id="msg_nested", type="message"), + _RevealingReprModel(a=1), + ], + ids=["dict_subclass", "tuple_subclass", "model", "model_with_own_repr"], +) +def test_responses_error_event_sends_a_detail_the_way_base_json_dumps_did(detail: object) -> None: + frame: Final = ResponsesStreamErrorState().format_error( + {"message": "blocked", "provider_specific_fields": {"detail": detail}} + ) + + payload: Final = json.loads(frame.split("data: ", 1)[1]) + assert payload["error"]["provider_specific_fields"]["detail"] == json.loads(json.dumps(detail, default=str)) + + +class _RigColor(Enum): + RED = 1 + + +class _Score(float): + pass + + +class _MaskedCount(int): + def __repr__(self) -> str: + return "_MaskedCount([secure])" + + +_WHEN: Final = datetime(2026, 10, 30, 12, 0, tzinfo=timezone.utc) + + +@pytest.mark.parametrize( + ("nested", "printed", "dumped"), + [ + (_WHEN, repr(_WHEN), str(_WHEN)), + (_RigColor.RED, repr(_RigColor.RED), str(_RigColor.RED)), + (_MaskedCount(7), repr(_MaskedCount(7)), 7), + ], + ids=["datetime", "enum", "number_subclass_with_its_own_repr"], +) +def test_responses_stream_frames_print_values_json_cannot_take_as_base_printed_them( + nested: object, printed: str, dumped: object +) -> None: + chunk_frame: Final = _format_responses_sse_chunk( + chunk=MappingProxyType({"type": "response.output_text.delta", "nested": nested}) + ) + error_frame: Final = ResponsesStreamErrorState().format_error({"message": "blocked", "nested": nested}) + + assert isinstance(chunk_frame, str) + assert json.loads(chunk_frame.removeprefix("data: "))["nested"] == printed + assert json.loads(error_frame.split("data: ", 1)[1])["error"]["nested"] == dumped + + +class _RigCode(str, Enum): + BLOCKED = "blocked_by_rig" + + +@pytest.mark.parametrize( + ("nested", "plain"), + [(_RigCode.BLOCKED, "blocked_by_rig"), (HTTPStatus.BAD_REQUEST, 400), (_Score(0.5), 0.5)], + ids=["str_enum", "int_enum", "float_subclass"], +) +def test_responses_stream_frames_send_string_and_number_subclasses_as_their_plain_value( + nested: object, plain: object +) -> None: + assert _responses_frame_payload(nested)["nested"] == plain + + +_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.parametrize("mapping", [dict, MappingProxyType], ids=["dict", "mapping_proxy"]) +def test_responses_sse_chunk_a_hook_dumped_from_an_event_streams_as_that_event( + mapping: Callable[[Mapping[str, object]], Mapping[str, object]], +) -> None: + frame: Final = _format_responses_sse_chunk(chunk=mapping(_DUMPED_DELTA.model_dump())) + + assert isinstance(frame, str) + assert json.loads(frame.removeprefix("data: ")) == json.loads(_DUMPED_DELTA.model_dump_json()) + + +_COMPLETED_WITH_SDK_OUTPUT: Final = ResponseCompletedEvent( + type=ResponsesAPIStreamEvents.RESPONSE_COMPLETED, + sequence_number=4, + response=ResponsesAPIResponse( + id="resp_1", + object="response", + created_at=0, + model="gpt-x", + status="completed", + output=[ + ResponseOutputMessage( + id="msg_1", + type="message", + role="assistant", + status="completed", + content=[ResponseOutputText(type="output_text", text="hi", annotations=[])], + ) + ], + ), +) + + +def test_responses_sse_chunk_a_hook_copied_from_an_event_streams_as_that_event() -> None: + frame: Final = _format_responses_sse_chunk(chunk=dict(_COMPLETED_WITH_SDK_OUTPUT)) + + assert isinstance(frame, str) + assert json.loads(frame.removeprefix("data: ")) == json.loads(_COMPLETED_WITH_SDK_OUTPUT.model_dump_json()) + + +@pytest.mark.parametrize("mapping", [dict, MappingProxyType], ids=["dict", "mapping_proxy"]) +def test_responses_stream_frames_turn_mapping_keys_json_cannot_take_into_strings( + mapping: Callable[[Mapping[object, object]], Mapping[object, object]], +) -> None: + assert _responses_frame_payload(mapping({("a", "b"): 1}))["nested"] == {repr(("a", "b")): 1} + + +def test_responses_sse_chunk_turns_top_level_keys_json_cannot_take_into_strings() -> None: + frame: Final = _format_responses_sse_chunk( + chunk=MappingProxyType({"type": "response.output_text.delta", ("a", "b"): "x"}) + ) + + assert isinstance(frame, str) + assert json.loads(frame.removeprefix("data: ")) == {"type": "response.output_text.delta", repr(("a", "b")): "x"} + + +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 test_responses_stream_frames_print_a_value_that_holds_itself_as_its_text() -> None: + nested: Final = _self_referencing_chunk() + + assert _responses_frame_payload(nested)["nested"] == {"type": "response.output_text.delta", "self": repr(nested)} + + +class _UnprintableValue: + def __repr__(self) -> str: + raise RuntimeError("no string form") + + +def test_responses_stream_frames_leave_out_only_the_nested_value_with_no_string_form() -> None: + chunk_frame: Final = _format_responses_sse_chunk( + chunk={"type": "response.output_text.delta", "meta": {"guardrail": "rig-guard", "bad": _UnprintableValue()}} + ) + error_frame: Final = ResponsesStreamErrorState().format_error( + { + "message": "blocked", + "code": "400", + "type": "guardrail_error", + "provider_specific_fields": {"guardrail": "rig-guard", "bad": _UnprintableValue()}, + } + ) + + assert isinstance(chunk_frame, str) + assert json.loads(chunk_frame.removeprefix("data: "))["meta"] == {"guardrail": "rig-guard"} + assert json.loads(error_frame.split("data: ", 1)[1])["error"] == { + "message": "blocked", + "code": "400", + "type": "guardrail_error", + "provider_specific_fields": {"guardrail": "rig-guard"}, + } + + +def test_responses_error_event_keeps_a_detail_nested_past_the_recursion_limit() -> None: + nested: Final = functools.reduce(lambda inner, _: {"x": inner}, range(5000), {"leaf": 1}) + frame: Final = ResponsesStreamErrorState().format_error( + {"message": "blocked", "provider_specific_fields": {"guardrail": "rig-guard", "detail": nested}} + ) + + payload: Final = json.loads(frame.split("data: ", 1)[1]) + assert payload["error"]["provider_specific_fields"]["guardrail"] == "rig-guard" + assert "'leaf': 1" in frame + + +@pytest.mark.parametrize( + "chunk", + [ + _UPSTREAM_HEADERS, + RedactedDict({"type": "response.output_text.delta", "authorization": "Bearer hook-canary"}), + MappingProxyType(RedactedDict({"type": "response.output_text.delta", "authorization": "Bearer hook-canary"})), + ], + ids=["header_mapping_chunk", "redacted_dict_chunk", "proxy_wrapping_redacted_dict_chunk"], +) +def test_responses_sse_chunk_that_is_not_a_plain_mapping_streams_as_the_plain_data_frame( + chunk: Mapping[str, object], +) -> None: + assert _format_responses_sse_chunk(chunk=chunk) == _format_streaming_sse_chunk(chunk=chunk) + +@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 @@ -1165,14 +1744,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 @@ -1368,8 +1948,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 diff --git a/tests/unit/proxy/response_api_endpoints/test_endpoints.py b/tests/unit/proxy/response_api_endpoints/test_endpoints.py index 07f85cb9628..e32ab1840ad 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.proxy_server import app from litellm.proxy.utils import ProxyLogging @@ -35,6 +36,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 @@ -153,8 +155,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", @@ -187,6 +190,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"}), @@ -5756,9 +5926,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 @@ -5785,7 +5957,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 087b273548a..1aa998db19f 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,145 @@ 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( + ("nested_error", "expected_error"), + [ + ( + {"type": "guardrail_error", "code": "400", "message": "output blocked by policy"}, + {"code": "invalid_request_error", "message": "output blocked by policy"}, + ), + ("output blocked by policy", {"code": "server_error", "message": "output blocked by policy"}), + ], + ) + async def test_a_nested_only_error_event_keeps_its_reason( + self, 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 = {"type": "error", "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( + ("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"}, + ), + ], + ) + 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 fe6b894b368..ccffb82b1f6 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 @@ -40,6 +40,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 ( @@ -52,11 +53,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, @@ -88,6 +93,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(): @@ -2256,6 +2262,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: @@ -2334,6 +2348,54 @@ class TestCommonRequestProcessingHelpers: assert body["error"]["code"] == 403 assert body["error"]["message"] == "forbidden" + async def test_a_responses_stream_whose_first_event_is_a_named_error_answers_as_json(self): + error: Final = {"type": "guardrail_error", "code": "400", "message": "blocked", "param": None} + + async def responses_stream(): + yield f"event: error\ndata: {json.dumps({'type': 'error', 'sequence_number': 0, 'error': error})}\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} + + 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_create_streaming_response_first_chunk_not_error(self): async def mock_generator(): yield 'data: {"content": "first part"}\n\n' @@ -10610,6 +10672,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 f52736453c1..6f3c5d8430b 100644 --- a/tests/unit/proxy/test_response_model_sanitization.py +++ b/tests/unit/proxy/test_response_model_sanitization.py @@ -763,7 +763,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"