mirror of
https://github.com/BerriAI/litellm.git
synced 2026-11-10 08:28:25 +00:00
fix(proxy): keep cursor chat framing and JSON-encode Responses hook errors
This commit is contained in:
parent
5003ef1f31
commit
ca873a1327
7 changed files with 99 additions and 41 deletions
|
|
@ -1250,6 +1250,16 @@ RESPONSES_FIRST_CHUNK_FRAMING: Final = FirstChunkSseFraming(
|
|||
)
|
||||
|
||||
|
||||
@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})
|
||||
|
|
@ -2945,7 +2955,7 @@ class ProxyBaseLLMRequestProcessing:
|
|||
is_streaming_request: bool | None = False,
|
||||
contents: list[object] | None = None,
|
||||
skip_pre_call_logic: bool = False,
|
||||
late_framing: LateSseFraming = OPENAI_LATE_SSE_FRAMING,
|
||||
sse_framing: SseFraming = OPENAI_SSE_FRAMING,
|
||||
) -> object:
|
||||
"""Run the request, sending SSE keepalives while the upstream is still silent.
|
||||
|
||||
|
|
@ -2986,10 +2996,11 @@ class ProxyBaseLLMRequestProcessing:
|
|||
is_streaming_request=is_streaming_request,
|
||||
contents=contents,
|
||||
skip_pre_call_logic=skip_pre_call_logic,
|
||||
first_chunk_framing=sse_framing.first_chunk,
|
||||
),
|
||||
ping_interval_seconds=ttft_keepalive_interval(self.data, llm_router),
|
||||
on_late_failure=_audit_late_failure,
|
||||
late_framing=late_framing,
|
||||
late_framing=sse_framing.late,
|
||||
)
|
||||
|
||||
async def _process_llm_request(
|
||||
|
|
@ -3013,6 +3024,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
|
||||
|
|
@ -3306,9 +3318,7 @@ class ProxyBaseLLMRequestProcessing:
|
|||
request=request,
|
||||
refresh_headers=refresh_stream_headers,
|
||||
general_settings=general_settings,
|
||||
framing=(
|
||||
RESPONSES_FIRST_CHUNK_FRAMING if route_type == "aresponses" else OPENAI_FIRST_CHUNK_FRAMING
|
||||
),
|
||||
framing=first_chunk_framing,
|
||||
)
|
||||
|
||||
### CALL HOOKS ### - modify outgoing data
|
||||
|
|
|
|||
|
|
@ -61,19 +61,22 @@ class _FailureDetails(LiteLLMBaseModel):
|
|||
class _ErrorPayload(LiteLLMBaseModel):
|
||||
model_config = ConfigDict(frozen=True)
|
||||
|
||||
error: Mapping[str, object]
|
||||
error: Mapping[str, object] | str
|
||||
|
||||
def error_object(self) -> Mapping[str, object]:
|
||||
return {"message": self.error} if isinstance(self.error, str) else self.error
|
||||
|
||||
|
||||
def sse_error_object(data: str | bytes) -> Mapping[str, object] | None:
|
||||
try:
|
||||
return _ErrorPayload.model_validate_json(data).error
|
||||
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 or None
|
||||
return _ErrorPayload.model_validate(event).error_object() or None
|
||||
except ValidationError:
|
||||
return None
|
||||
|
||||
|
|
|
|||
|
|
@ -10202,6 +10202,12 @@ def _format_streaming_sse_chunk(chunk: object) -> str | bytes:
|
|||
return f"data: {chunk}\n\n"
|
||||
|
||||
|
||||
def _format_responses_sse_chunk(chunk: object) -> str | bytes:
|
||||
if isinstance(chunk, Mapping):
|
||||
return f"data: {json.dumps(chunk, default=str)}\n\n"
|
||||
return _format_streaming_sse_chunk(chunk=chunk)
|
||||
|
||||
|
||||
_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
|
||||
|
|
@ -10602,7 +10608,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:
|
||||
|
|
|
|||
|
|
@ -52,7 +52,7 @@ from litellm.proxy.auth.user_api_key_auth import (
|
|||
user_api_key_auth_websocket,
|
||||
)
|
||||
from litellm.proxy.common_request_processing import (
|
||||
RESPONSES_LATE_SSE_FRAMING,
|
||||
RESPONSES_SSE_FRAMING,
|
||||
ProxyBaseLLMRequestProcessing,
|
||||
cost_headers,
|
||||
create_response,
|
||||
|
|
@ -438,7 +438,7 @@ async def responses_api(
|
|||
if requests_background(_PROXY_SERVER_REQUEST_ADAPTER.validate_python(data))
|
||||
else native_data_generator
|
||||
),
|
||||
late_framing=RESPONSES_LATE_SSE_FRAMING,
|
||||
sse_framing=RESPONSES_SSE_FRAMING,
|
||||
model=None,
|
||||
user_model=user_model,
|
||||
user_temperature=user_temperature,
|
||||
|
|
|
|||
|
|
@ -21,7 +21,7 @@ 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 RESPONSES_LATE_SSE_FRAMING, ProxyBaseLLMRequestProcessing
|
||||
from litellm.proxy.common_request_processing import RESPONSES_SSE_FRAMING, ProxyBaseLLMRequestProcessing
|
||||
from litellm.proxy.common_utils.responses_stream_errors import (
|
||||
event_error_object,
|
||||
responses_error_fields,
|
||||
|
|
@ -188,7 +188,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,
|
||||
sse_framing=RESPONSES_SSE_FRAMING,
|
||||
)
|
||||
|
||||
# Process streaming response following OpenAI events format
|
||||
|
|
|
|||
|
|
@ -17,7 +17,7 @@ from __future__ import annotations
|
|||
|
||||
import asyncio
|
||||
import json
|
||||
from collections.abc import AsyncIterator, Callable
|
||||
from collections.abc import AsyncIterator, Callable, Mapping
|
||||
from typing import Final, Literal
|
||||
from unittest.mock import AsyncMock, MagicMock
|
||||
|
||||
|
|
@ -148,7 +148,8 @@ async def test_responses_stream_dispatches_named_events_without_a_chat_done_sent
|
|||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_responses_stream_sends_no_failure_event_after_an_upstream_error_event() -> None:
|
||||
@pytest.mark.parametrize("as_mapping", [False, True])
|
||||
async def test_responses_stream_sends_no_failure_event_after_an_upstream_error_event(as_mapping: bool) -> 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}
|
||||
|
|
@ -157,9 +158,9 @@ async def test_responses_stream_sends_no_failure_event_after_an_upstream_error_e
|
|||
{"type": "error", "sequence_number": 1, "code": "server_error", "message": "upstream failed", "param": None}
|
||||
)
|
||||
|
||||
async def upstream() -> AsyncIterator[BaseModel]:
|
||||
async def upstream() -> AsyncIterator[BaseModel | Mapping[str, object]]:
|
||||
yield created
|
||||
yield error_event
|
||||
yield error_event.model_dump() if as_mapping else error_event
|
||||
raise litellm.APIError(status_code=500, message="Stream close failed", llm_provider="openai", model="local-model")
|
||||
|
||||
frames: Final = tuple(
|
||||
|
|
@ -172,14 +173,26 @@ async def test_responses_stream_sends_no_failure_event_after_an_upstream_error_e
|
|||
)
|
||||
|
||||
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"
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_a_guardrail_error_string_mid_responses_stream_ends_it_as_a_named_error_event() -> None:
|
||||
@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}
|
||||
)
|
||||
guardrail_error: Final = {"message": "response withheld", "type": "guardrail_scan_error", "code": 500}
|
||||
|
||||
async def upstream() -> AsyncIterator[object]:
|
||||
yield created
|
||||
|
|
@ -203,7 +216,7 @@ async def test_a_guardrail_error_string_mid_responses_stream_ends_it_as_a_named_
|
|||
"message": "response withheld",
|
||||
"param": None,
|
||||
"sequence_number": 1,
|
||||
"error": guardrail_error,
|
||||
"error": error_object,
|
||||
}
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -181,14 +181,16 @@ async def test_streaming_upstream_errors_keep_the_client_protocol(
|
|||
|
||||
|
||||
class _FirstFrameHook(CustomLogger):
|
||||
def __init__(self, frame: str | bytes) -> None:
|
||||
def __init__(self, first: str | bytes | Exception) -> None:
|
||||
super().__init__()
|
||||
self.frame: Final = frame
|
||||
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]:
|
||||
yield self.frame
|
||||
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'
|
||||
|
|
@ -197,22 +199,9 @@ _NAMED_ERROR_FRAME: Final = (
|
|||
)
|
||||
|
||||
|
||||
@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:
|
||||
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
|
||||
|
||||
|
|
@ -230,14 +219,14 @@ async def test_only_a_lone_first_error_frame_answers_with_its_status(
|
|||
)
|
||||
monkeypatch.setattr(ps, "llm_router", router)
|
||||
monkeypatch.setattr(litellm, "disable_aiohttp_transport", True)
|
||||
monkeypatch.setattr(litellm, "callbacks", [_FirstFrameHook(first_frame)])
|
||||
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:
|
||||
result: Final = await client.post(
|
||||
return await client.post(
|
||||
path,
|
||||
json={
|
||||
"model": model,
|
||||
|
|
@ -246,6 +235,26 @@ async def test_only_a_lone_first_error_frame_answers_with_its_status(
|
|||
},
|
||||
)
|
||||
|
||||
|
||||
@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
|
||||
|
|
@ -253,6 +262,23 @@ async def test_only_a_lone_first_error_frame_answers_with_its_status(
|
|||
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:
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue