scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-11-05 07:37:50 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 233

fix(proxy): keep first-chunk error detection to lone error frames and handle string poll errors

Browse source
This commit is contained in:
mateo-berri 2026-10-28 18:33:43 -07:00
parent e6290016c1
commit 4bfca9cfd4
6 changed files with 166 additions and 23 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

43
litellm/proxy/common_request_processing.py
View file

@ -910,17 +910,31 @@ async def _enforce_guardrail_added_tag_budgets(
) from e
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 = event_line.decode("utf-8") if isinstance(event_line, bytes) else event_line
return next(
(line.removeprefix("data: ").strip() for line in text.split("\n") if line.startswith("data: ")),
None,
)
text: Final = _sse_text(event_line)
return text.removeprefix("data: ").strip() if text.startswith("data: ") else None
async def _parse_event_data_for_error(event_line: str | bytes) -> int | None:
def _responses_event_data_payload(event_line: str | bytes) -> str | None:
match _sse_text(event_line).strip().split("\n"):
case ["event: error", data_line] if data_line.startswith("data: "):
return data_line.removeprefix("data: ").strip()
case _:
return _sse_data_payload(event_line)
async def _parse_event_data_for_error(
event_line: str | bytes, data_payload: SseDataPayload = _sse_data_payload
) -> int | None:
"""Parses an event line and returns an error code if present, else None."""
json_str: Final = _sse_data_payload(event_line)
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
@ -952,12 +966,13 @@ 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
@ -970,7 +985,7 @@ def _extract_error_from_sse_chunk(event_line: str | bytes) -> dict:
"code": "500",
}
json_str: Final = _sse_data_payload(event_line)
json_str: Final = data_payload(event_line)
if json_str is not None:
if not json_str or json_str == "[DONE]":
return default_error
@ -1243,6 +1258,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,
first_chunk_payload: SseDataPayload = _sse_data_payload,
) -> StreamingResponse | JSONResponse:
"""
Create streaming response, checking if the first chunk is an error.
@ -1278,7 +1294,7 @@ async def create_response(
if first_chunk_value is not None:
try:
error_code_from_chunk: Final = await _parse_event_data_for_error(first_chunk_value)
error_code_from_chunk: Final = await _parse_event_data_for_error(first_chunk_value, first_chunk_payload)
if error_code_from_chunk is not None:
# First chunk is an error, stream hasn't really started yet
# Should return standard JSON error response instead of SSE format
@ -1290,7 +1306,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, first_chunk_payload)
),
error_body_call_id(general_settings, resolved_headers.get(LITELLM_CALL_ID_HEADER)),
)
@ -3217,6 +3235,9 @@ class ProxyBaseLLMRequestProcessing:
request=request,
refresh_headers=refresh_stream_headers,
general_settings=general_settings,
first_chunk_payload=(
_responses_event_data_payload if route_type == "aresponses" else _sse_data_payload
),
)
### CALL HOOKS ### - modify outgoing data

19
litellm/proxy/common_utils/responses_stream_errors.py
View file

@ -71,6 +71,13 @@ def sse_error_object(data: str | bytes) -> Mapping[str, object] | None:
return None
def event_error_object(event: Mapping[str, object]) -> Mapping[str, object] | None:
try:
return _ErrorPayload.model_validate(event).error 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:
@ -126,11 +133,19 @@ 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_of(details.code)})),
"message": details.message,
"code": _response_error_code(details.model_copy(update={"status_code": status_code})),
"message": details.message or _status_phrase(status_code),
}

16
litellm/proxy/response_polling/background_streaming.py
View file

@ -21,7 +21,11 @@ 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 RESPONSES_LATE_SSE_FRAMING, ProxyBaseLLMRequestProcessing
from litellm.proxy.common_utils.responses_stream_errors import responses_error_fields, sse_error_object
from litellm.proxy.common_utils.responses_stream_errors import (
event_error_object,
responses_error_fields,
sse_error_object,
)
from litellm.proxy.response_polling.polling_handler import ResponsePollingHandler
from litellm.types.llms.openai import ResponsesAPIStatus
@ -87,6 +91,14 @@ def _sse_frame_data(frame: str) -> str | None:
return next((line[6:].strip() for line in frame.splitlines() if line.startswith("data: ")), None)
def _event_error(event: Mapping[str, object]) -> Mapping[str, object]:
match event.get("error"):
case str() as message if message:
return {**event, "message": message}
case _:
return event_error_object(event) or event
def _first_chunk_error(response: object) -> Mapping[str, object] | None:
if isinstance(response, _StreamsBody) or not isinstance(response, Response):
return None
@ -339,7 +351,7 @@ async def background_streaming_task(
elif event_type == "error":
terminal_status = "failed"
terminal_error = responses_error_fields(event.get("error") or event)
terminal_error = responses_error_fields(_event_error(event))
elif event_type in (
"response.completed",

76
tests/unit/proxy/response_api_endpoints/test_endpoints.py
View file

@ -7,7 +7,7 @@ import contextlib
import dataclasses
import json
import unittest
from collections.abc import Mapping
from collections.abc import AsyncIterator, Mapping
from typing import Any, Final, Literal, cast
from unittest.mock import AsyncMock, MagicMock, call, patch
@ -19,6 +19,7 @@ from fastapi.testclient import TestClient
from httpx import Response
import litellm
from litellm.integrations.custom_logger import CustomLogger
from litellm.proxy._types import ProxyException
from litellm.proxy.proxy_server import app
from litellm.types.llms.openai import ResponseAPIUsage, ResponsesAPIResponse
@ -178,6 +179,79 @@ async def test_streaming_upstream_errors_keep_the_client_protocol(
assert "error" in events[-1]
class _FirstFrameHook(CustomLogger):
def __init__(self, frame: str | bytes) -> None:
super().__init__()
self.frame: Final = frame
async def async_post_call_streaming_iterator_hook(
self, user_api_key_dict: object, response: object, request_data: object
) -> AsyncIterator[str | bytes]:
yield self.frame
_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'
)
@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),
("/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:
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_frame)])
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:
result: Final = await client.post(
path,
json={
"model": model,
"stream": True,
**({"messages": [{"role": "user", "content": "hello"}]} if is_chat else {"input": "hello"}),
},
)
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.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:

26
tests/unit/proxy/response_polling/test_response_polling_handler.py
View file

@ -1555,12 +1555,23 @@ class TestBackgroundStreamingTerminalEvents:
assert final_call.kwargs["error"] == {"code": "rate_limit_exceeded", "message": "rate limited"}
@pytest.mark.asyncio
async def test_a_nested_only_error_event_keeps_its_reason(self):
@pytest.mark.parametrize(
("nested_error", "expected_error"),
[
(
{"type": "guardrail_error", "code": "400", "message": "output blocked by policy"},
{"code": "invalid_request_error", "message": "output blocked by policy"},
),
("output blocked by policy", {"code": "server_error", "message": "output blocked by policy"}),
],
)
async def test_a_nested_only_error_event_keeps_its_reason(
self, nested_error: Mapping[str, str] | str, expected_error: Mapping[str, str]
):
from litellm.proxy.response_polling.background_streaming import (
background_streaming_task,
)
nested_error: Final = {"type": "guardrail_error", "code": "400", "message": "output blocked by policy"}
error_event: Final = {"type": "error", "sequence_number": 9, "error": nested_error}
async def _body_iterator():
@ -1582,7 +1593,7 @@ class TestBackgroundStreamingTerminalEvents:
final_call: Final = handler.update_state.call_args_list[-1]
assert final_call.kwargs["status"] == "failed"
assert final_call.kwargs["error"] == {"code": "invalid_request_error", "message": "output blocked by policy"}
assert final_call.kwargs["error"] == expected_error
@pytest.mark.asyncio
@pytest.mark.parametrize(
@ -1595,11 +1606,16 @@ class TestBackgroundStreamingTerminalEvents:
),
{"code": "not_found_error", "message": "model not found"},
),
(429, "Too Many Requests", {"code": "rate_limit_exceeded", "message": None}),
(429, "rate limited", {"code": "rate_limit_exceeded", "message": "Too Many Requests"}),
(
403,
json.dumps({"error": {"type": "guardrail_scan_error", "code": 403}}),
{"code": "permission_error", "message": "Forbidden"},
),
],
)
async def test_a_first_chunk_error_polls_failed_like_a_late_one(
self, status_code: int, body: str, expected_error: Mapping[str, str | None]
self, status_code: int, body: str, expected_error: Mapping[str, str]
):
from fastapi import Response

9
tests/unit/proxy/test_common_request_processing.py
View file

@ -49,6 +49,7 @@ from litellm.proxy.common_request_processing import (
_get_cost_breakdown_from_logging_obj,
CostBreakdownHeaderValues,
_has_attribute_error_in_chain,
_responses_event_data_payload,
include_guardrail_response_requested,
is_azure_model_router_request,
RESPONSES_LATE_SSE_FRAMING,
@ -2170,7 +2171,9 @@ class TestCommonRequestProcessingHelpers:
async def responses_stream():
yield f"event: error\ndata: {json.dumps({'type': 'error', 'sequence_number': 0, 'error': error})}\n\n"
response: Final = await create_response(responses_stream(), "text/event-stream", {})
response: Final = await create_response(
responses_stream(), "text/event-stream", {}, first_chunk_payload=_responses_event_data_payload
)
assert isinstance(response, JSONResponse)
assert response.status_code == 400
@ -2183,7 +2186,9 @@ class TestCommonRequestProcessingHelpers:
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", {})
response: Final = await create_response(
responses_stream(), "text/event-stream", {}, first_chunk_payload=_responses_event_data_payload
)
assert isinstance(response, JSONResponse)
assert response.status_code == 400