mirror of
https://github.com/BerriAI/litellm.git
synced 2026-11-10 08:28:25 +00:00
Merge branch 'fix-responses-sse-framing' of https://github.com/HootRock/litellm into HEAD
# Conflicts: # tests/unit/proxy/response_api_endpoints/test_endpoints.py # tests/unit/proxy/test_common_request_processing.py
This commit is contained in:
mateo-berri
commit
5d4db205ce
2 changed files with 96 additions and 96 deletions
|
|
@ -17,7 +17,6 @@ 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
|
||||
|
|
@ -25,6 +24,7 @@ from litellm.proxy.proxy_server import app
|
|||
from litellm.types.llms.openai import ResponseAPIUsage, ResponsesAPIResponse
|
||||
from litellm.types.router import RouterErrors
|
||||
from tests._master_key import MASTER_KEY
|
||||
from tests._support.clock_loops import AutojumpClockLoop
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
|
|
@ -178,77 +178,72 @@ async def test_streaming_upstream_errors_keep_the_client_protocol(
|
|||
assert "error" in events[-1]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("path", ["/v1/responses", "/v1/chat/completions"])
|
||||
async def test_a_failure_after_keepalive_pings_keeps_the_client_protocol(monkeypatch: pytest.MonkeyPatch, path: str):
|
||||
import litellm.proxy.proxy_server as ps
|
||||
from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
|
||||
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"
|
||||
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"
|
||||
|
||||
first_ping_sent: Final = asyncio.Event()
|
||||
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}})
|
||||
|
||||
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=[
|
||||
{
|
||||
"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=rate_limit_after_first_ping
|
||||
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,
|
||||
)
|
||||
async with httpx.AsyncClient(
|
||||
transport=httpx.ASGITransport(app_marking_first_ping), 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"}),
|
||||
},
|
||||
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
|
||||
)
|
||||
frames: Final = tuple(frame for frame in result.text.split("\n\n") if frame)
|
||||
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"]
|
||||
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 = {
|
||||
|
|
|
|||
|
|
@ -3,9 +3,9 @@ import contextvars
|
|||
import copy
|
||||
import datetime
|
||||
import json
|
||||
from collections.abc import AsyncIterable, Mapping
|
||||
from collections.abc import AsyncIterable, Awaitable, Mapping
|
||||
from types import MappingProxyType, SimpleNamespace
|
||||
from typing import AsyncGenerator, Awaitable, Callable, Final, Iterator, Literal, Optional, Sequence
|
||||
from typing import AsyncGenerator, Callable, Final, Iterator, Literal, Optional, Sequence
|
||||
from urllib.parse import unquote_plus
|
||||
from unittest.mock import AsyncMock, MagicMock, patch
|
||||
|
||||
|
|
@ -81,6 +81,7 @@ from litellm.router import Router
|
|||
from litellm.litellm_core_utils.streaming_handler import CustomStreamWrapper
|
||||
from litellm.utils import ModelResponseListIterator
|
||||
from litellm.router_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():
|
||||
|
|
@ -10404,14 +10405,16 @@ 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 _rate_limit_failure_once_released(released: asyncio.Event) -> JSONResponse:
|
||||
await released.wait()
|
||||
async def _slow_rate_limit_failure() -> Response:
|
||||
await asyncio.sleep(0.2)
|
||||
raise HTTPException(status_code=429, detail="rate limited")
|
||||
|
||||
|
||||
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"}})
|
||||
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(
|
||||
|
|
@ -10419,43 +10422,45 @@ _POLICY_BLOCK_DETAIL: Final = MappingProxyType(
|
|||
)
|
||||
|
||||
|
||||
async def _dict_detail_failure_once_released(released: asyncio.Event) -> JSONResponse:
|
||||
await released.wait()
|
||||
async def _slow_dict_detail_failure() -> Response:
|
||||
await asyncio.sleep(0.2)
|
||||
raise HTTPException(status_code=400, detail=dict(_POLICY_BLOCK_DETAIL))
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize(
|
||||
"produce,expected_error",
|
||||
[
|
||||
(_rate_limit_failure_once_released, {"message": "rate limited", "code": "429"}),
|
||||
(_guardrail_block_body_once_released, {"message": "blocked", "type": "guardrail_error", "code": "400"}),
|
||||
(_slow_rate_limit_failure, {"message": "rate limited", "code": "429"}),
|
||||
(_slow_guardrail_block_body, {"message": "blocked", "type": "guardrail_error", "code": "400"}),
|
||||
(
|
||||
_dict_detail_failure_once_released,
|
||||
_slow_dict_detail_failure,
|
||||
{"message": "blocked by policy", "code": "400", "provider_specific_fields": dict(_POLICY_BLOCK_DETAIL)},
|
||||
),
|
||||
],
|
||||
)
|
||||
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]
|
||||
def test_ttft_keepalive_ends_a_late_responses_failure_with_an_error_event_and_no_done(
|
||||
produce: Callable[[], Awaitable[Response]], expected_error: Mapping[str, object]
|
||||
) -> None:
|
||||
released: Final = asyncio.Event()
|
||||
response: Final = await open_sse_before_first_byte(
|
||||
produce(released), ping_interval_seconds=0.05, late_framing=RESPONSES_LATE_SSE_FRAMING
|
||||
)
|
||||
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]]
|
||||
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 {key: event["error"].get(key) for key in expected_error} == expected_error
|
||||
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 {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()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue