mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-04 02:31:27 +00:00
53 lines
2.5 KiB
Python
53 lines
2.5 KiB
Python
import json
|
|
import time
|
|
import uuid
|
|
from collections.abc import Callable
|
|
from typing import Final
|
|
|
|
import pytest
|
|
from integration._support.client import Gateway
|
|
from integration._support.wire import Reply, Request, wire_server
|
|
from integration.streaming.test_stream_contracts import text_stream
|
|
|
|
UPSTREAM_SILENCE_SECONDS: Final = 3.0
|
|
KEEPALIVE_SECONDS: Final = 1
|
|
|
|
|
|
def _reply_after_silence(identity: str) -> Callable[[Request], Reply]:
|
|
def respond(_request: Request) -> Reply:
|
|
time.sleep(UPSTREAM_SILENCE_SECONDS)
|
|
return Reply(content_type="text/event-stream", chunks=text_stream(identity))
|
|
|
|
return respond
|
|
|
|
|
|
@pytest.mark.covers("streaming.keepalive.sse_pings_fill_silent_time_to_first_token")
|
|
def test_stream_emits_sse_ping_comments_before_the_first_data_frame_while_upstream_is_silent(
|
|
gateway: Gateway,
|
|
) -> None:
|
|
identity: Final = "stream-ttft-keepalive-" + uuid.uuid4().hex
|
|
with gateway.scenario() as scenario:
|
|
with wire_server(_reply_after_silence(identity)) as wire:
|
|
model: Final = scenario.model(api_base=wire.url + "/v1", keepalive_seconds=KEEPALIVE_SECONDS)
|
|
with gateway.client.stream(
|
|
"POST",
|
|
"/v1/chat/completions",
|
|
json={"model": model, "messages": [{"role": "user", "content": identity}], "stream": True},
|
|
headers={"Authorization": f"Bearer {gateway.key}"},
|
|
) as response:
|
|
assert response.status_code == 200, response.read().decode()
|
|
frames: Final = tuple(line for line in response.iter_lines() if line)
|
|
first_data: Final = next(index for index, line in enumerate(frames) if line.startswith("data:"))
|
|
assert first_data >= 1, f"No keepalive reached the client before the first data frame: {frames}"
|
|
assert frames[:first_data] == (": ping",) * first_data, frames
|
|
assert frames[-1] == "data: [DONE]", frames
|
|
deltas: Final = tuple(json.loads(line.removeprefix("data: ")) for line in frames[first_data:-1])
|
|
assert (
|
|
"".join(choice["delta"].get("content", "") for chunk in deltas for choice in chunk["choices"])
|
|
== "Hello 雪 café"
|
|
), frames
|
|
requests: Final = wire.drain()
|
|
assert len(requests) == 1
|
|
outbound: Final = json.loads(requests[0].body)
|
|
assert outbound["model"] == "gpt-4o-mini" and outbound["stream"] is True, outbound
|
|
assert "keepalive_seconds" not in outbound, outbound
|