diff --git a/litellm/proxy/common_request_processing.py b/litellm/proxy/common_request_processing.py index abc99a3be23..6cdab2daa82 100644 --- a/litellm/proxy/common_request_processing.py +++ b/litellm/proxy/common_request_processing.py @@ -115,7 +115,7 @@ from litellm.proxy.common_utils.openai_error_payload import ( openai_error_param, openai_error_type, ) -from litellm.proxy.common_utils.responses_stream_errors import ResponsesStreamErrorState +from litellm.proxy.common_utils.responses_stream_errors import ResponsesStreamErrorState, sse_error_object from litellm.proxy.common_utils.sse_keepalive import ( SSE_COMMENT_PING_BYTES, SSE_STREAM_START_TAIL, @@ -912,7 +912,7 @@ async def _enforce_guardrail_added_tag_budgets( def _sse_data_payload(event_line: str | bytes) -> str | None: text: Final = event_line.decode("utf-8") if isinstance(event_line, bytes) else event_line return next( - (line.removeprefix("data: ").strip() for line in text.splitlines() if line.startswith("data: ")), + (line.removeprefix("data: ").strip() for line in text.split("\n") if line.startswith("data: ")), None, ) @@ -1196,13 +1196,14 @@ def _openai_late_body_frames(produced: Response) -> tuple[bytes, ...]: def _responses_late_failure_frames(exc: Exception) -> tuple[str, ...]: - frame: Final = ResponsesStreamErrorState().format_failure(exc) - return () if frame is None else (frame,) + _, error_obj = sse_error_payload(exc) + return (ResponsesStreamErrorState().format_error(error_obj),) def _responses_late_body_frames(produced: Response) -> tuple[bytes, ...]: - frame: Final = ResponsesStreamErrorState().format_error_body(bytes(produced.body), produced.status_code) - return () if frame is None else (frame.encode(),) + body: Final = bytes(produced.body) + error_obj: Final = sse_error_object(body or b"{}") or {"message": body.decode(), "code": str(produced.status_code)} + return (ResponsesStreamErrorState().format_error(error_obj).encode(),) OPENAI_LATE_SSE_FRAMING: Final = LateSseFraming(failure=_openai_late_failure_frames, body=_openai_late_body_frames) @@ -3423,8 +3424,8 @@ class ProxyBaseLLMRequestProcessing: Streaming ``/v1/responses`` short-circuits out of ``base_process_llm_request`` before the non-streaming ownership - tail runs. The OpenAI SDK closes the connection at ``data: [DONE]`` - and starlette cancels the body task on disconnect, so a write that + tail runs. The OpenAI SDK can close the connection right after the terminal + event and starlette cancels the body task on disconnect, so a write that waits for the generator to finish never lands. The iterator sets ``completed_response`` before it hands over its terminal chunk, so the ``LiteLLM_ManagedObjectTable`` row is written the moment it diff --git a/litellm/proxy/common_utils/responses_stream_errors.py b/litellm/proxy/common_utils/responses_stream_errors.py index 312bf472594..132155d9492 100644 --- a/litellm/proxy/common_utils/responses_stream_errors.py +++ b/litellm/proxy/common_utils/responses_stream_errors.py @@ -1,10 +1,11 @@ +import json import time from collections.abc import Mapping from http import HTTPStatus from types import MappingProxyType from typing import Final -from pydantic import BaseModel, ConfigDict, field_validator +from pydantic import BaseModel, ConfigDict, ValidationError, field_validator from litellm._logging import redact_internal_details_from_client_message from litellm._uuid import uuid @@ -57,10 +58,17 @@ class _FailureDetails(LiteLLMBaseModel): return cls.model_validate({field: getattr(exception, field, None) for field in cls.model_fields}) -class _ErrorBody(LiteLLMBaseModel): +class _ErrorPayload(LiteLLMBaseModel): model_config = ConfigDict(frozen=True) - error: _FailureDetails = _FailureDetails() + error: Mapping[str, object] + + +def sse_error_object(data: str | bytes) -> Mapping[str, object] | None: + try: + return _ErrorPayload.model_validate_json(data).error + except ValidationError: + return None def _original_failure(exception: Exception) -> Exception: @@ -142,17 +150,26 @@ class ResponsesStreamErrorState: 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: + payload: Final = json.dumps( + { + "type": ResponsesAPIStreamEvents.ERROR.value, + "sequence_number": self.sequence_number + 1, + "error": error, + }, + 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: - original: Final = _original_failure(exception) - return self._format_failed(_failure_details(original), str(original)) - - def format_error_body(self, body: bytes, status_code: int) -> str | None: - error: Final = _ErrorBody.model_validate_json(body or b"{}").error - return self._format_failed(error.model_copy(update={"status_code": status_code}), body.decode()) - - def _format_failed(self, details: _FailureDetails, fallback_message: str) -> str | None: if self.terminal_emitted: return None + original: Final = _original_failure(exception) + details: Final = _failure_details(original) response: Final = ResponsesAPIResponse.model_validate( MappingProxyType( { @@ -165,7 +182,7 @@ class ResponsesStreamErrorState: "error": MappingProxyType( { "code": _response_error_code(details), - "message": redact_internal_details_from_client_message(details.message or fallback_message), + "message": redact_internal_details_from_client_message(details.message or str(original)), } ), } diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 15832d81b7f..f6cfc88300c 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -10619,7 +10619,7 @@ async def async_data_generator( yield (raw_sse_buffer if raw_sse_buffer.endswith(_SSE_FRAME_DELIMITERS) else raw_sse_buffer + "\n\n") if error_message is not None: - yield error_message + 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 diff --git a/litellm/proxy/response_polling/background_streaming.py b/litellm/proxy/response_polling/background_streaming.py index 5f6a7432afb..975d700688e 100644 --- a/litellm/proxy/response_polling/background_streaming.py +++ b/litellm/proxy/response_polling/background_streaming.py @@ -19,7 +19,7 @@ from typing_extensions import ReadOnly, TypedDict from litellm._logging import verbose_proxy_logger 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_LATE_SSE_FRAMING, ProxyBaseLLMRequestProcessing from litellm.proxy.response_polling.polling_handler import ResponsePollingHandler from litellm.types.llms.openai import ResponsesAPIStatus @@ -68,6 +68,7 @@ class _StreamEvent(TypedDict, total=False): delta: ReadOnly[str] part: ReadOnly[object] response: ReadOnly[_TerminalResponse] + error: ReadOnly[_JsonDict] class _StreamEventParser: @@ -164,6 +165,7 @@ async def background_streaming_task( user_api_base=user_api_base, version=version, skip_pre_call_logic=True, + late_framing=RESPONSES_LATE_SSE_FRAMING, ) # Process streaming response following OpenAI events format @@ -325,6 +327,10 @@ async def background_streaming_task( status="in_progress", ) + elif event_type == "error": + terminal_status = "failed" + terminal_error = event.get("error") + elif event_type in ( "response.completed", "response.failed", diff --git a/tests/unit/proxy/proxy_server/test_streaming_helpers.py b/tests/unit/proxy/proxy_server/test_streaming_helpers.py index 835411608ed..b01487af5d1 100644 --- a/tests/unit/proxy/proxy_server/test_streaming_helpers.py +++ b/tests/unit/proxy/proxy_server/test_streaming_helpers.py @@ -146,6 +146,37 @@ async def test_responses_stream_dispatches_named_events_without_a_chat_done_sent assert payload["type"] == event_type assert json.loads(frames[1].split("data: ", 1)[1])["delta"] == "hello" +@pytest.mark.asyncio +async def test_a_guardrail_error_string_mid_responses_stream_ends_it_as_a_named_error_event() -> 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} + ) + guardrail_error: Final = {"message": "response withheld", "type": "guardrail_scan_error", "code": 500} + + 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", + "sequence_number": 1, + "error": guardrail_error, + } + + @pytest.mark.parametrize("binary", [False, True]) @pytest.mark.parametrize( ("event_type", "event_line"), diff --git a/tests/unit/proxy/response_api_endpoints/test_endpoints.py b/tests/unit/proxy/response_api_endpoints/test_endpoints.py index 5945352c61e..d31ddaf0adf 100644 --- a/tests/unit/proxy/response_api_endpoints/test_endpoints.py +++ b/tests/unit/proxy/response_api_endpoints/test_endpoints.py @@ -17,6 +17,7 @@ import respx from fastapi import HTTPException from fastapi.testclient import TestClient from httpx import Response +from starlette.types import Message, Receive, Scope, Send import litellm from litellm.proxy._types import ProxyException @@ -188,10 +189,20 @@ async def test_a_failure_after_keepalive_pings_keeps_the_client_protocol(monkeyp 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) + first_ping_sent: Final = asyncio.Event() + + async def rate_limit_after_first_ping(_: httpx.Request) -> httpx.Response: + await first_ping_sent.wait() return httpx.Response(429, json={"error": {"message": message, "type": "rate_limit_error", "code": None}}) + async def app_marking_first_ping(scope: Scope, receive: Receive, send: Send) -> None: + async def send_and_mark(asgi_message: Message) -> None: + await send(asgi_message) + if asgi_message["type"] == "http.response.body" and asgi_message.get("body", b"").startswith(b": ping"): + first_ping_sent.set() + + await app(scope, receive, send_and_mark) + router: Final = litellm.Router( model_list=[ { @@ -212,9 +223,11 @@ async def test_a_failure_after_keepalive_pings_keeps_the_client_protocol(monkeyp 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 + side_effect=rate_limit_after_first_ping ) - async with httpx.AsyncClient(transport=httpx.ASGITransport(app), base_url="http://testserver") as client: + async with httpx.AsyncClient( + transport=httpx.ASGITransport(app_marking_first_ping), base_url="http://testserver" + ) as client: result: Final = await client.post( path, json={ @@ -232,10 +245,10 @@ async def test_a_failure_after_keepalive_pings_keeps_the_client_protocol(monkeyp assert message in json.loads(frames[-2].removeprefix("data: "))["error"]["message"] else: assert "data: [DONE]" not in result.text - assert frames[-1].startswith("event: response.failed\ndata: "), result.text - failed_error: Final = json.loads(frames[-1].split("data: ", 1)[1])["response"]["error"] - assert failed_error["code"] == "rate_limit_exceeded" - assert message in failed_error["message"] + 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"] _COMPACTION_ITEM: Final = { diff --git a/tests/unit/proxy/response_polling/test_response_polling_handler.py b/tests/unit/proxy/response_polling/test_response_polling_handler.py index 63e38ce533a..8affc8a6e19 100644 --- a/tests/unit/proxy/response_polling/test_response_polling_handler.py +++ b/tests/unit/proxy/response_polling/test_response_polling_handler.py @@ -21,7 +21,7 @@ from typing import Any, Dict, 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 @@ -1522,6 +1522,38 @@ 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 = 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 = Mock() + mock_response.body_iterator = _body_iterator() + handler = AsyncMock(spec=ResponsePollingHandler) + kwargs = _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 = handler.update_state.call_args_list[-1] + assert final_call.kwargs["status"] == "failed" + assert final_call.kwargs["error"]["message"] == "rate limited" + assert final_call.kwargs["error"]["code"] == "429" + @pytest.mark.asyncio async def test_response_incomplete_sets_incomplete_status_and_details(self): """Test that a response.incomplete stream event results in incomplete status""" diff --git a/tests/unit/proxy/test_common_request_processing.py b/tests/unit/proxy/test_common_request_processing.py index 310b526365f..ecc3ff860d8 100644 --- a/tests/unit/proxy/test_common_request_processing.py +++ b/tests/unit/proxy/test_common_request_processing.py @@ -5,7 +5,7 @@ import datetime import json from collections.abc import AsyncIterable, Mapping from types import MappingProxyType, SimpleNamespace -from typing import AsyncGenerator, Callable, Final, Iterator, Literal, Optional, Sequence +from typing import AsyncGenerator, Awaitable, Callable, Final, Iterator, Literal, Optional, Sequence from urllib.parse import unquote_plus from unittest.mock import AsyncMock, MagicMock, patch @@ -2174,6 +2174,19 @@ class TestCommonRequestProcessingHelpers: 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", {}) + + assert isinstance(response, JSONResponse) + assert response.status_code == 400 + assert json.loads(response.body) == {"error": error} + async def test_create_streaming_response_first_chunk_not_error(self): async def mock_generator(): yield 'data: {"content": "first part"}\n\n' @@ -10391,41 +10404,58 @@ 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(): - await asyncio.sleep(0.2) +async def _rate_limit_failure_once_released(released: asyncio.Event) -> JSONResponse: + await released.wait() raise HTTPException(status_code=429, detail="rate limited") -async def _slow_guardrail_block_body(): - await asyncio.sleep(0.2) +async def _guardrail_block_body_once_released(released: asyncio.Event) -> JSONResponse: + await released.wait() 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 _dict_detail_failure_once_released(released: asyncio.Event) -> JSONResponse: + await released.wait() + raise HTTPException(status_code=400, detail=dict(_POLICY_BLOCK_DETAIL)) + + @pytest.mark.asyncio @pytest.mark.parametrize( - "produce,expected_code,expected_message", + "produce,expected_error", [ - (_slow_rate_limit_failure, "rate_limit_exceeded", "rate limited"), - (_slow_guardrail_block_body, "invalid_request_error", "blocked"), + (_rate_limit_failure_once_released, {"message": "rate limited", "code": "429"}), + (_guardrail_block_body_once_released, {"message": "blocked", "type": "guardrail_error", "code": "400"}), + ( + _dict_detail_failure_once_released, + {"message": "blocked by policy", "code": "400", "provider_specific_fields": dict(_POLICY_BLOCK_DETAIL)}, + ), ], ) -async def test_ttft_keepalive_ends_a_late_responses_failure_with_response_failed_and_no_done( - produce, expected_code, expected_message -): +async def test_ttft_keepalive_ends_a_late_responses_failure_with_an_error_event_and_no_done( + produce: Callable[[asyncio.Event], Awaitable[JSONResponse]], expected_error: Mapping[str, object] +) -> None: + released: Final = asyncio.Event() response: Final = await open_sse_before_first_byte( - produce(), ping_interval_seconds=0.05, late_framing=RESPONSES_LATE_SSE_FRAMING + produce(released), ping_interval_seconds=0.05, late_framing=RESPONSES_LATE_SSE_FRAMING ) - collected: Final = await _drain(response) + assert isinstance(response, StreamingResponse) + frames: Final = aiter(response.body_iterator) + first_frame: Final = await anext(frames) + released.set() + collected: Final = [first_frame, *[frame async for frame in frames]] assert collected[0] == TTFT_PING assert b"data: [DONE]\n\n" not in collected event_line, data_line = collected[-1].decode().strip().split("\n") - failed: Final = json.loads(data_line.removeprefix("data: ")) - assert event_line == "event: response.failed" - assert failed["type"] == "response.failed" - assert failed["response"]["status"] == "failed" - assert failed["response"]["error"]["code"] == expected_code - assert expected_message in failed["response"]["error"]["message"] + event: Final = json.loads(data_line.removeprefix("data: ")) + assert event_line == "event: error" + assert event["type"] == "error" + assert {key: event["error"].get(key) for key in expected_error} == expected_error @pytest.mark.asyncio