mirror of
https://github.com/BerriAI/litellm.git
synced 2026-11-12 08:42:18 +00:00
fix(proxy)!: preserve Responses SSE event framing (#47980)
* fix(proxy): preserve Responses SSE event framing
* style(proxy): format Responses stream terminator condition
* fix(proxy): validate Responses guardrail event names
* fix(proxy): keep Responses first-event errors and keepalive failures in Responses framing
* test(proxy): use controlled clocks for Responses keepalive failures
* fix(proxy): end late Responses failures with an error event the SDK raises on
Errors that arrive after keepalive pings, and guardrail error strings yielded
mid-stream, now end /v1/responses with a named `error` event carrying the
same error object chat sends, so the OpenAI SDK still raises APIError with a
clean message and provider_specific_fields. Background polling passes the
same framing and records an `error` event as a failed response. The first
data line of a frame is found by splitting on newlines only, so a U+2028 in
an error message no longer hides a first-event error.
* fix(proxy): match OpenAI's Responses error event and stream late passthrough blocks as block responses
* fix(proxy): keep the reason when a background Responses request fails
A background poll of a request a post-call guardrail blocked saved a
null code and message, because the guardrail's error event nests them.
A request whose first chunk was an error polled as completed with no
error. Both now poll failed with the same code and message a late
failure gets.
* test(proxy): type the first-chunk poll test parameters
* fix(proxy): keep first-chunk error detection to lone error frames and handle string poll errors
* fix(proxy): keep the HTTP status on late Responses errors and end Responses streams on error events
A late error body without a code now takes the response's status, an upstream error event counts as the stream's terminal event, and a Responses stream that fails before its first chunk ends with an error event instead of data: [DONE].
* test(proxy): type the first-chunk failure test's parameters
* fix(proxy): keep cursor chat framing and JSON-encode Responses hook errors
* fix(proxy): encode non-dict Mapping Responses chunks as JSON objects
* fix(proxy): encode mappings and models nested in Responses stream chunks as JSON objects
* test(proxy): freeze the hook-output fixture
* fix(proxy): keep masked mappings and odd keys safe in Responses stream frames
Only read-only dict views and pydantic models nested in a hook chunk or error event expand into JSON objects. Other mappings, such as httpx.Headers, keep their own string form so masked values stay masked. Keys JSON cannot take become strings, and a chunk JSON still cannot encode goes out as the plain data frame it was before.
* fix(proxy): never fail a Responses stream frame on a value JSON cannot encode
Only dicts and read-only dict views expand into JSON, at the top level too, so a header mapping yielded as a whole chunk keeps its masked string form. Any failure while encoding a chunk falls back to the plain data frame the base sent, and an error event that cannot encode still goes out with its code and message, so a late block never drops the connection.
* fix(proxy): keep what a nested value's repr hides out of Responses stream frames
* fix(proxy): print values in Responses stream frames the way each base path did
* fix(proxy): send string and number subclasses whose repr shows their value as that value in Responses stream frames
* fix(proxy): print container subclasses and SDK models in Responses stream frames the way base did
* fix(proxy): read hook container storage the way repr does and send error-event models as str()
* fix(proxy): frame streamed Responses reads and send unreadable hook values as base did
Track pre-framed and raw SSE frames so a late failure never follows a terminal event, accept CRLF and CR line breaks in a first named error, frame GET /v1/responses/{id}?stream=true as Responses events, print a whole hook proxy chunk once and fall back to its text, ignore a non-integer status_code in a hook error frame, read a late-block dict subclass through items() when that view differs from its storage, and keep the order and inline form base's repr printed for nested mappings
* fix(proxy): match base on Responses stream edge cases
Read a decimal-string or integral-float status_code as its integer again, so
"429" and 429.0 still map to rate_limit_exceeded. Print each hook value's repr
once and reuse it, so a repr that changes between calls cannot reach the
client. Track only pre-framed sequence numbers from 0 to 2**63-1, so a huge
one cannot break the failure event. Build nested JSON with one frame per level
again, which restores the nesting depth sent as JSON. Count a pre-framed frame
as terminal from its data type, and a named error frame only when its data
carries an error, the way the OpenAI SDKs read them.
* fix(proxy): print hook proxies once and keep base's chunk tracking
Read-only mappings in hook chunks now emit the text from their first print, hook
chunks are read without their error field, only a non-empty error object makes a
framed error event terminal, and framed sequence numbers stay below 2**63-1 while
hook chunks keep base's unbounded tracking
* fix(proxy): print each read-only mapping item once and poll framed error events as failed
A MappingProxyType in a hook chunk or error detail now expands from the
mapping it wraps, so each item is printed once, the way base printed it.
A value whose repr changes between calls stays masked, a value whose repr
works once is kept, and deep nesting keeps one stack frame per level
The background reader takes each frame's type from the same terminal rule
the stream uses, so an event: error frame with an error object polls
failed even when its data names another type or none
* fix(proxy): print self-holding hook values once and poll how the stream ended
A container that holds itself now goes out as the marker its repr prints, like {...}, [...] or mappingproxy({...}), and the dict behind a read-only mapping joins the cycle check, so no item is printed a second time. A dict whose item repr adds a key keeps going after the items already shown, the way dict's repr does, instead of printing the whole dict again
A whole read-only mapping chunk that cannot go out as JSON is printed once by the formatter, so a print that fails fails the stream as it did on the base
The background poll reads the event type from the block it parses, as on the base, and keeps the first error frame the stream treats as terminal anywhere in a chunk. That frame fails the poll only when no response.completed, response.failed or response.incomplete follows it
* fix(proxy): name hook chunk events from the sent JSON and keep polls on frames they cannot read
A mapping chunk now takes its event: line from the JSON it sends, so a masked type value never reaches that line, and a chunk the event model cannot read goes out as a plain data frame instead of failing the stream
A dict whose item repr deletes or swaps keys keeps printing every key not yet shown, and an error code past the integer digit limit reads as no status
The background poll skips a frame whose shape it cannot read, as base skipped frames it could not parse, stops reading a chunk at [DONE], and lets only the last terminal frame in a chunk decide its status
A Responses stream's first chunk answers as a JSON error only for codes 400 and up
* fix(proxy): read every frame of a polled chunk and name model chunks from the sent JSON
A background poll parsed only the first data frame of each chunk, so a terminal event behind an error or delta frame in the same chunk was lost, and the poll stored a status its own stream never ended with. It now reads every frame up to [DONE]
A terminal event whose response is a string or null kept no status and polled completed. Its status now comes from the event type, and final output items that are not objects are skipped, so the rest of the final output is still stored
Model chunks now take their event line from the JSON they send, as mapping chunks already do, so a serializer that drops type no longer sends event: response.completed over {}
* fix(proxy): read polled chunks as SSE events and keep draining past a mid-chunk [DONE]
A background poll split a chunk into data lines, so one event with two data lines parsed as two payloads, and a frame with an integer past the digit limit raised a ValueError the parse guard did not catch, which failed the poll and dropped its spend row. The poll now joins each event's data lines the way an SSE reader does and skips any frame it cannot parse
A [DONE] behind another frame in one chunk ended the poll and left the stream unbilled. Only a chunk that opens with [DONE] ends the poll now, as on main, and the poll stops reading a chunk at its first terminal event, so two terminal events in one chunk poll the first one as main did
* fix(proxy): JSON-encode only plain hook chunks on Responses streams
A hook chunk whose whole tree is builtin JSON types streams as its named
event. Any other chunk keeps the base data frame, so nothing a custom
repr hides reaches the wire. Error event details encode the way the chat
error frame does.
* fix(proxy): skip containers already visited when checking hook chunks for JSON
* fix(responses): keep streaming past an event with no type
The Responses stream iterator read event.type directly when it noted a
yielded event, so an event whose JSON had no type, like a custom hook's
dict that the proxy now sends as JSON, raised AttributeError and ended
the stream. Read it with getattr, as the iterator's other checks do.
* fix(proxy): keep hook dicts with no type on the base Responses frame
A hook dict with no type is not a Responses event, so it gets no event
name either way. Sending it as JSON only made litellm SDKs that read
type on every event stop the stream. Keep such dicts on the frame the
base sends, so every SDK version reads them as before.
* fix(responses): keep a background poll's first terminal status and take a first-chunk error code from its HTTP status
---------
Co-authored-by: mateo-berri <277851410+mateo-berri@users.noreply.github.com>
This commit is contained in:
HootRock
GitHub
parent
b3c4a1c3e4
commit
8f19607259
22 changed files with 1793 additions and 102 deletions
|
|
@ -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
|
||||
)
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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:
|
||||
|
|
|
|||
|
|
@ -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):
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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),
|
||||
)
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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:
|
||||
|
|
|
|||
|
|
@ -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.
|
||||
|
|
|
|||
|
|
@ -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(
|
||||
|
|
|
|||
|
|
@ -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}
|
||||
|
||||
|
|
|
|||
|
|
@ -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(
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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"
|
||||
|
|
|
|||
|
|
@ -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"])
|
||||
|
||||
|
|
|
|||
|
|
@ -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()
|
||||
|
|
|
|||
|
|
@ -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
|
||||
# ---------------------------------------------------------------------------
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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"""
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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"
|
||||
|
|
|
|||
|
|
@ -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()
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue