mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-14 04:09:09 +00:00
test(integration): audit cells for prompt_cache_breakpoint through the chat-to-Responses bridge
Adds the /audit cells for the bridge's prompt_cache_breakpoint carry: valid and malformed markers on every carrying block kind and on tool and assistant content, the drop_params on and off contract, the hook-injected system marker, the Anthropic SDK path through the Responses adapter, upstream errors, idempotent spend rows, and the chaos cells (mixed burst, upstream outage, slow streams, worker SIGKILL and a proxy restart mid burst), all against a scripted Responses endpoint with the request it received read back by response id
This commit is contained in:
mateo-berri
parent
eca0a96a29
commit
96f5d1754e
3 changed files with 1109 additions and 0 deletions
202
tests/integration/_support/prompt_cache_breakpoint.py
Normal file
202
tests/integration/_support/prompt_cache_breakpoint.py
Normal file
|
|
@ -0,0 +1,202 @@
|
|||
from __future__ import annotations
|
||||
|
||||
import os
|
||||
import re
|
||||
import uuid
|
||||
from collections.abc import Iterator, Mapping, Sequence
|
||||
from contextlib import contextmanager
|
||||
from dataclasses import dataclass
|
||||
from pathlib import Path
|
||||
from typing import Final, Literal, TypeAlias, assert_never
|
||||
from urllib.parse import urlsplit
|
||||
|
||||
import psutil
|
||||
import psycopg
|
||||
from integration._support import responses_vendor as rv
|
||||
from integration._support.client import eventually, object_value, string_value
|
||||
from integration._support.database import ROWS
|
||||
from integration._support.openai_wire import answering_model_discovery, responses_reply
|
||||
from integration._support.wire import Reply, Request, Wire
|
||||
from psycopg.rows import DictRow, dict_row
|
||||
from pydantic import JsonValue
|
||||
|
||||
MODEL: Final = "openai/responses/gpt-6.1-sol"
|
||||
EXPLICIT: Final[Mapping[str, JsonValue]] = {"mode": "explicit"}
|
||||
EXPLICIT_30M: Final[Mapping[str, JsonValue]] = {"mode": "explicit", "ttl": "30m"}
|
||||
NO_CACHE: Final[Mapping[str, JsonValue]] = {"cache": {"no-cache": True}}
|
||||
INJECTION: Final[Mapping[str, JsonValue]] = {
|
||||
"cache_control_injection_points": [{"location": "message", "role": "system"}],
|
||||
"prompt_cache_options": {"mode": "explicit"},
|
||||
}
|
||||
_SCRIPTED_FAILURE: Final = re.compile(r"fail-(\d{3})")
|
||||
_STARTED_WORKER: Final = re.compile(r"Started server process \[(\d+)\]")
|
||||
|
||||
Kind: TypeAlias = Literal["text", "image_url", "file", "input_audio"]
|
||||
KINDS: Final[tuple[Kind, ...]] = ("text", "image_url", "file", "input_audio")
|
||||
WIRE_TYPE: Final[Mapping[Kind, str]] = {
|
||||
"text": "input_text",
|
||||
"image_url": "input_image",
|
||||
"file": "input_file",
|
||||
"input_audio": "input_text",
|
||||
}
|
||||
|
||||
|
||||
def _scripted(request: Request) -> Reply:
|
||||
body: Final = rv.JSON_OBJECT.validate_json(request.body)
|
||||
text: Final = request.body.decode()
|
||||
marker: Final = rv.newest_marker(text)
|
||||
failure: Final = _SCRIPTED_FAILURE.search(text)
|
||||
if failure is not None:
|
||||
return rv.error(int(failure.group(1)), f"scripted {failure.group(1)} marker-{marker}", "scripted_failure")
|
||||
return responses_reply(
|
||||
f"resp_{marker or uuid.uuid4().hex}",
|
||||
string_value(body["model"]),
|
||||
rv.answer(marker),
|
||||
stream=body.get("stream") is True,
|
||||
)
|
||||
|
||||
|
||||
respond: Final = answering_model_discovery(_scripted)
|
||||
|
||||
|
||||
def response_id(marker: str) -> str:
|
||||
return f"resp_{marker}"
|
||||
|
||||
|
||||
def prompt(marker: str) -> str:
|
||||
return f"Say marker-{marker}"
|
||||
|
||||
|
||||
def text(value: str) -> dict[str, JsonValue]:
|
||||
return {"type": "text", "text": value}
|
||||
|
||||
|
||||
def marked(block: Mapping[str, JsonValue], marker: JsonValue) -> dict[str, JsonValue]:
|
||||
return {**block, "prompt_cache_breakpoint": marker}
|
||||
|
||||
|
||||
def block(kind: Kind, value: str) -> dict[str, JsonValue]:
|
||||
match kind:
|
||||
case "text":
|
||||
return text(value)
|
||||
case "image_url":
|
||||
return {"type": "image_url", "image_url": {"url": "https://example.com/breakpoint.png"}}
|
||||
case "file":
|
||||
return {"type": "file", "file": {"file_id": "file-breakpoint"}}
|
||||
case "input_audio":
|
||||
return {"type": "input_audio", "input_audio": {"data": "Zm9v", "format": "wav"}}
|
||||
case _:
|
||||
assert_never(kind)
|
||||
|
||||
|
||||
def drained_posts(wire: Wire) -> tuple[Request, ...]:
|
||||
return tuple(request for request in wire.drain() if request.method == "POST")
|
||||
|
||||
|
||||
def with_marker(posts: Sequence[Request], marker: str) -> tuple[Request, ...]:
|
||||
return tuple(request for request in posts if f"marker-{marker}" in request.body.decode())
|
||||
|
||||
|
||||
def posted(wire: Wire, marker: str) -> Request:
|
||||
matching: Final = with_marker(drained_posts(wire), marker)
|
||||
assert len(matching) == 1, [request.body for request in matching]
|
||||
(request,) = matching
|
||||
assert request.target == "/v1/responses", request.target
|
||||
return request
|
||||
|
||||
|
||||
def body_of(request: Request) -> dict[str, JsonValue]:
|
||||
return rv.JSON_OBJECT.validate_json(request.body)
|
||||
|
||||
|
||||
def input_items(request: Request) -> list[dict[str, JsonValue]]:
|
||||
return rv.ITEMS.validate_python(body_of(request)["input"])
|
||||
|
||||
|
||||
def content_of(items: Sequence[Mapping[str, JsonValue]], role: str) -> list[dict[str, JsonValue]]:
|
||||
messages: Final = tuple(item for item in items if item.get("type") == "message" and item.get("role") == role)
|
||||
assert len(messages) == 1, items
|
||||
return rv.ITEMS.validate_python(messages[0]["content"])
|
||||
|
||||
|
||||
def single_block(items: Sequence[Mapping[str, JsonValue]], role: str) -> dict[str, JsonValue]:
|
||||
blocks: Final = content_of(items, role)
|
||||
assert len(blocks) == 1, blocks
|
||||
return blocks[0]
|
||||
|
||||
|
||||
def instruction_block(items: Sequence[Mapping[str, JsonValue]]) -> dict[str, JsonValue]:
|
||||
messages: Final = tuple(
|
||||
item for item in items if item.get("type") == "message" and item.get("role") in ("system", "developer")
|
||||
)
|
||||
assert len(messages) == 1, items
|
||||
blocks: Final = rv.ITEMS.validate_python(messages[0]["content"])
|
||||
assert len(blocks) == 1, blocks
|
||||
return blocks[0]
|
||||
|
||||
|
||||
def function_output(items: Sequence[Mapping[str, JsonValue]], call_id: str) -> list[dict[str, JsonValue]]:
|
||||
outputs: Final = tuple(
|
||||
item for item in items if item.get("type") == "function_call_output" and item.get("call_id") == call_id
|
||||
)
|
||||
assert len(outputs) == 1, items
|
||||
return rv.ITEMS.validate_python(outputs[0]["output"])
|
||||
|
||||
|
||||
def assert_marker(block_on_wire: Mapping[str, JsonValue], expected: JsonValue) -> None:
|
||||
if expected is None:
|
||||
assert "prompt_cache_breakpoint" not in block_on_wire, block_on_wire
|
||||
return
|
||||
assert block_on_wire.get("prompt_cache_breakpoint") == expected, block_on_wire
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class SpendLogs:
|
||||
connection: psycopg.Connection[DictRow]
|
||||
|
||||
def rows_for(self, model: str) -> list[dict[str, JsonValue]]:
|
||||
cursor: Final = self.connection.execute(
|
||||
'SELECT litellm_call_id, request_id, status FROM "LiteLLM_SpendLogs" WHERE model_group = %s', (model,)
|
||||
)
|
||||
return ROWS.validate_python(cursor.fetchall())
|
||||
|
||||
def landed(
|
||||
self, model: str, call_id: str, identity: str | None, *, status: str = "success", seconds: float = 70
|
||||
) -> dict[str, JsonValue]:
|
||||
rows: Final = eventually(
|
||||
lambda: self.rows_for(model),
|
||||
lambda found: any(row["litellm_call_id"] == call_id for row in found),
|
||||
seconds=seconds,
|
||||
)
|
||||
matching: Final = tuple(row for row in rows if row["litellm_call_id"] == call_id)
|
||||
assert len(matching) == 1, rows
|
||||
(row,) = matching
|
||||
assert row["status"] == status, row
|
||||
assert identity is None or rv.same_response(string_value(row["request_id"]), identity), (row, identity)
|
||||
return row
|
||||
|
||||
|
||||
@contextmanager
|
||||
def spend_logs() -> Iterator[SpendLogs]:
|
||||
with psycopg.connect(os.environ["DATABASE_URL"], row_factory=dict_row, autocommit=True) as connection:
|
||||
connection.execute("SET default_transaction_read_only = on")
|
||||
yield SpendLogs(connection)
|
||||
|
||||
|
||||
def model_id(entries: Sequence[JsonValue], model: str) -> str:
|
||||
matching: Final = tuple(entry for entry in entries if object_value(entry).get("model_name") == model)
|
||||
assert len(matching) == 1, entries
|
||||
return string_value(object_value(object_value(matching[0])["model_info"])["id"])
|
||||
|
||||
|
||||
def started_worker_pids(log: Path) -> tuple[int, ...]:
|
||||
return tuple(int(found.group(1)) for found in _STARTED_WORKER.finditer(log.read_text()))
|
||||
|
||||
|
||||
def open_upstream_connections(pid: int, upstream: str) -> int:
|
||||
port: Final = urlsplit(upstream).port
|
||||
return sum(
|
||||
1
|
||||
for connection in psutil.Process(pid).net_connections(kind="tcp")
|
||||
if connection.status == psutil.CONN_ESTABLISHED and connection.raddr and connection.raddr.port == port
|
||||
)
|
||||
|
|
@ -0,0 +1,397 @@
|
|||
import asyncio
|
||||
import dataclasses
|
||||
import signal
|
||||
import threading
|
||||
import uuid
|
||||
from collections import Counter
|
||||
from collections.abc import Callable, Iterator, Mapping, Sequence
|
||||
from dataclasses import dataclass
|
||||
from pathlib import Path
|
||||
from queue import SimpleQueue
|
||||
from types import MappingProxyType
|
||||
from typing import Final, Literal, TypeAlias
|
||||
from urllib.parse import urlsplit
|
||||
|
||||
import httpx
|
||||
import psutil
|
||||
import pytest
|
||||
import yaml
|
||||
from integration._support import prompt_cache_breakpoint as pcb
|
||||
from integration._support import responses_vendor as rv
|
||||
from integration._support.client import Gateway, eventually, gateway_from_environment, string_value
|
||||
from integration._support.process import OwnedProxy, graceful_stop_seconds, owned_proxy_process
|
||||
from integration._support.wire import Reply, Request, Wire, wire_server
|
||||
from pydantic import JsonValue
|
||||
|
||||
pytestmark: Final = pytest.mark.timeout(2 * graceful_stop_seconds() + 120)
|
||||
|
||||
_GLOBAL_UNSET: Final = "bridge-breakpoint-global-unset"
|
||||
_GLOBAL_FALSE: Final = "bridge-breakpoint-global-false"
|
||||
_ENDPOINTS: Final = ("chat", "messages", "responses")
|
||||
|
||||
Endpoint: TypeAlias = Literal["chat", "messages", "responses"]
|
||||
_RecordProperty: TypeAlias = Callable[[str, object], None]
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class _Call:
|
||||
endpoint: Endpoint
|
||||
stream: bool
|
||||
marker: str
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class _Served:
|
||||
call: _Call
|
||||
status: int
|
||||
text: str
|
||||
call_id: str
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class _GlobalRig:
|
||||
wire: Wire
|
||||
proxy: OwnedProxy
|
||||
|
||||
@property
|
||||
def gateway(self) -> Gateway:
|
||||
return self.proxy.gateway
|
||||
|
||||
|
||||
def _global_config(directory: Path, api_base: str) -> Path:
|
||||
stock: Final = rv.JSON_OBJECT.validate_python(
|
||||
yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text())
|
||||
)
|
||||
deployment: Final[Mapping[str, JsonValue]] = {
|
||||
"model": pcb.MODEL,
|
||||
"api_base": api_base,
|
||||
"api_key": "integration-provider-key",
|
||||
}
|
||||
config: Final[Mapping[str, JsonValue]] = {
|
||||
**stock,
|
||||
"model_list": [
|
||||
{"model_name": _GLOBAL_UNSET, "litellm_params": dict(deployment)},
|
||||
{"model_name": _GLOBAL_FALSE, "litellm_params": {**deployment, "drop_params": False}},
|
||||
],
|
||||
"litellm_settings": {**rv.JSON_OBJECT.validate_python(stock["litellm_settings"]), "drop_params": True},
|
||||
"router_settings": {**rv.JSON_OBJECT.validate_python(stock.get("router_settings") or {}), "num_retries": 0},
|
||||
}
|
||||
path: Final = directory / "bridge-breakpoint-global.yaml"
|
||||
path.write_text(yaml.safe_dump(config))
|
||||
return path
|
||||
|
||||
|
||||
@pytest.fixture(scope="module")
|
||||
def global_rig(tmp_path_factory: pytest.TempPathFactory) -> Iterator[_GlobalRig]:
|
||||
directory: Final = tmp_path_factory.mktemp("bridge-breakpoint-global")
|
||||
with wire_server(pcb.respond) as wire, gateway_from_environment() as gateway:
|
||||
config: Final = _global_config(directory, f"{wire.url}/v1")
|
||||
with owned_proxy_process(gateway, directory, {}, config=config, workers=2) as owned:
|
||||
yield _GlobalRig(wire, owned)
|
||||
|
||||
|
||||
@pytest.fixture(scope="module")
|
||||
def spend() -> Iterator[pcb.SpendLogs]:
|
||||
with pcb.spend_logs() as logs:
|
||||
yield logs
|
||||
|
||||
|
||||
def _path(endpoint: Endpoint) -> str:
|
||||
match endpoint:
|
||||
case "chat":
|
||||
return "/v1/chat/completions"
|
||||
case "messages":
|
||||
return "/v1/messages"
|
||||
case "responses":
|
||||
return "/v1/responses"
|
||||
|
||||
|
||||
def _body(model: str, call: _Call, breakpoint: JsonValue) -> Mapping[str, JsonValue]:
|
||||
common: Final[Mapping[str, JsonValue]] = {"model": model, "stream": call.stream, **pcb.NO_CACHE}
|
||||
text: Final = pcb.marked(pcb.text(pcb.prompt(call.marker)), breakpoint)
|
||||
match call.endpoint:
|
||||
case "chat":
|
||||
return {**common, "messages": [{"role": "user", "content": [text]}]}
|
||||
case "messages":
|
||||
return {**common, "max_tokens": 64, "messages": [{"role": "user", "content": [text]}]}
|
||||
case "responses":
|
||||
return {
|
||||
**common,
|
||||
"input": [{"type": "message", "role": "user", "content": [{**text, "type": "input_text"}]}],
|
||||
}
|
||||
|
||||
|
||||
def _calls(count: int, endpoints: tuple[Endpoint, ...], *, stream: bool | None = None) -> tuple[_Call, ...]:
|
||||
return tuple(
|
||||
_Call(endpoints[index % len(endpoints)], index % 2 == 1 if stream is None else stream, uuid.uuid4().hex)
|
||||
for index in range(count)
|
||||
)
|
||||
|
||||
|
||||
async def _send(client: httpx.AsyncClient, key: str, model: str, call: _Call, breakpoint: JsonValue) -> _Served:
|
||||
async with client.stream(
|
||||
"POST",
|
||||
_path(call.endpoint),
|
||||
json=_body(model, call, breakpoint),
|
||||
headers={"Authorization": f"Bearer {key}", "anthropic-version": "2023-06-01"},
|
||||
) as response:
|
||||
raw: Final = await response.aread()
|
||||
return _Served(call, response.status_code, raw.decode(), response.headers["x-litellm-call-id"])
|
||||
|
||||
|
||||
async def _burst(
|
||||
gateway: Gateway,
|
||||
model: str,
|
||||
calls: tuple[_Call, ...],
|
||||
*,
|
||||
breakpoint: JsonValue = pcb.EXPLICIT,
|
||||
tolerate_transport_errors: bool = False,
|
||||
) -> tuple[_Served, ...]:
|
||||
async with httpx.AsyncClient(base_url=str(gateway.client.base_url), timeout=60, trust_env=False) as client:
|
||||
results: Final = await asyncio.gather(
|
||||
*(_send(client, gateway.key, model, call, breakpoint) for call in calls),
|
||||
return_exceptions=tolerate_transport_errors,
|
||||
)
|
||||
for result in results:
|
||||
assert not isinstance(result, BaseException) or isinstance(result, httpx.TransportError), repr(result)
|
||||
return tuple(result for result in results if isinstance(result, _Served))
|
||||
|
||||
|
||||
def _frames(text: str) -> tuple[Mapping[str, JsonValue], ...]:
|
||||
return tuple(rv.JSON_OBJECT.validate_json(line[6:]) for line in text.splitlines() if line.startswith("data: {"))
|
||||
|
||||
|
||||
def _upstream_id_shown_to_caller(served: _Served) -> str | None:
|
||||
if served.call.endpoint == "messages":
|
||||
return None
|
||||
if not served.call.stream:
|
||||
return string_value(rv.JSON_OBJECT.validate_json(served.text)["id"])
|
||||
frames: Final = _frames(served.text)
|
||||
if served.call.endpoint == "responses":
|
||||
(completed,) = [frame for frame in frames if frame.get("type") == "response.completed"]
|
||||
return string_value(rv.JSON_OBJECT.validate_python(completed["response"])["id"])
|
||||
return string_value(frames[0]["id"])
|
||||
|
||||
|
||||
def _assert_answered_in_its_own_shape(served: _Served) -> None:
|
||||
assert served.status == 200, served.text
|
||||
assert set(rv.MARKER.findall(served.text)) == {served.call.marker}, served.text
|
||||
assert served.text.startswith(("event:", "data:")) == served.call.stream, served.text
|
||||
assert served.text.startswith("{") != served.call.stream, served.text
|
||||
assert ("response.completed" in served.text) == (served.call.stream and served.call.endpoint == "responses")
|
||||
shown: Final = _upstream_id_shown_to_caller(served)
|
||||
assert shown is None or rv.same_response(shown, pcb.response_id(served.call.marker)), served.text
|
||||
|
||||
|
||||
def _marked_once(posts: Sequence[Request], calls: Sequence[_Call], expected: JsonValue) -> None:
|
||||
by_marker: Final = {marker: request for request in posts if (marker := rv.newest_marker(request.body.decode()))}
|
||||
assert len(by_marker) == len(posts), [request.body for request in posts]
|
||||
assert set(by_marker) == {call.marker for call in calls}, sorted(by_marker)
|
||||
for call in calls:
|
||||
block: Final = pcb.single_block(pcb.input_items(by_marker[call.marker]), "user")
|
||||
assert block["type"] == "input_text" and block["text"] == pcb.prompt(call.marker), block
|
||||
pcb.assert_marker(block, expected)
|
||||
|
||||
|
||||
def _assert_each_lands_once(
|
||||
spend: pcb.SpendLogs, model: str, failed: Sequence[_Served], served: Sequence[_Served]
|
||||
) -> None:
|
||||
expected: Final = len(failed) + len(served)
|
||||
rows: Final = eventually(lambda: spend.rows_for(model), lambda found: len(found) >= expected, seconds=70)
|
||||
by_call: Final = {string_value(row["litellm_call_id"]): row for row in rows}
|
||||
assert len(by_call) == len(rows) == expected, rows
|
||||
for item in failed:
|
||||
assert by_call[item.call_id]["status"] == "failure", (item.call_id, rows)
|
||||
for item in served:
|
||||
row: Final = by_call[item.call_id]
|
||||
assert row["status"] == "success", (item.call_id, row)
|
||||
shown: Final = _upstream_id_shown_to_caller(item)
|
||||
assert shown is None or rv.same_response(string_value(row["request_id"]), shown), (row, shown)
|
||||
|
||||
|
||||
def _health(gateway: Gateway, model: str) -> Mapping[str, JsonValue]:
|
||||
response: Final = gateway.request("GET", f"/health?model={model}", None)
|
||||
assert response.status_code in (200, 503), response.text
|
||||
return rv.JSON_OBJECT.validate_json(response.text)
|
||||
|
||||
|
||||
def _free_port() -> int:
|
||||
with wire_server(pcb.respond) as probe:
|
||||
port: Final = urlsplit(probe.url).port
|
||||
assert port is not None, probe.url
|
||||
return port
|
||||
|
||||
|
||||
def _chat(gateway: Gateway, model: str, marker: str, breakpoint: JsonValue) -> httpx.Response:
|
||||
return gateway.request("POST", "/v1/chat/completions", dict(_body(model, _Call("chat", False, marker), breakpoint)))
|
||||
|
||||
|
||||
def _completion(response: httpx.Response, marker: str) -> str:
|
||||
assert response.status_code == 200, response.text
|
||||
body: Final = rv.JSON_OBJECT.validate_json(response.text)
|
||||
assert rv.same_response(string_value(body["id"]), pcb.response_id(marker)), body
|
||||
assert rv.answer(marker) in response.text, response.text
|
||||
return response.headers["x-litellm-call-id"]
|
||||
|
||||
|
||||
def _user_block_on_wire(wire: Wire, marker: str) -> dict[str, JsonValue]:
|
||||
block: Final = pcb.single_block(pcb.input_items(pcb.posted(wire, marker)), "user")
|
||||
assert block["type"] == "input_text" and block["text"] == pcb.prompt(marker), block
|
||||
return block
|
||||
|
||||
|
||||
@pytest.mark.parametrize("model", (_GLOBAL_UNSET, _GLOBAL_FALSE), ids=("deployment-unset", "deployment-false"))
|
||||
def test_global_drop_params_drops_a_malformed_marker(global_rig: _GlobalRig, model: str, spend: pcb.SpendLogs) -> None:
|
||||
marker: Final = uuid.uuid4().hex
|
||||
call_id: Final = _completion(_chat(global_rig.gateway, model, marker, "yes"), marker)
|
||||
pcb.assert_marker(_user_block_on_wire(global_rig.wire, marker), None)
|
||||
spend.landed(model, call_id, pcb.response_id(marker))
|
||||
control: Final = uuid.uuid4().hex
|
||||
control_id: Final = _completion(_chat(global_rig.gateway, model, control, pcb.EXPLICIT), control)
|
||||
pcb.assert_marker(_user_block_on_wire(global_rig.wire, control), pcb.EXPLICIT)
|
||||
spend.landed(model, control_id, pcb.response_id(control))
|
||||
|
||||
|
||||
async def test_mixed_burst_carries_every_marker_once(gateway: Gateway, spend: pcb.SpendLogs) -> None:
|
||||
calls: Final = _calls(24, _ENDPOINTS)
|
||||
with wire_server(pcb.respond) as wire, gateway.scenario() as scenario:
|
||||
model: Final = scenario.model(model=pcb.MODEL, api_base=f"{wire.url}/v1", drop_params=True)
|
||||
served: Final = await _burst(gateway, model, calls)
|
||||
assert len(served) == 24
|
||||
for item in served:
|
||||
_assert_answered_in_its_own_shape(item)
|
||||
_marked_once(pcb.drained_posts(wire), calls, pcb.EXPLICIT)
|
||||
_assert_each_lands_once(spend, model, (), served)
|
||||
|
||||
|
||||
async def test_upstream_outage_fails_cleanly_and_the_restarted_upstream_serves_marked_calls(
|
||||
gateway: Gateway, spend: pcb.SpendLogs
|
||||
) -> None:
|
||||
port: Final = _free_port()
|
||||
while_down: Final = _calls(12, _ENDPOINTS)
|
||||
after: Final = _calls(12, _ENDPOINTS)
|
||||
with gateway.scenario() as scenario:
|
||||
model: Final = scenario.model(model=pcb.MODEL, api_base=f"http://127.0.0.1:{port}/v1", drop_params=True)
|
||||
failed: Final = await _burst(gateway, model, while_down)
|
||||
assert len(failed) == 12
|
||||
for item in failed:
|
||||
assert item.status >= 500, (item.status, item.text)
|
||||
assert "answer marker" not in item.text and "event:" not in item.text, item.text
|
||||
down: Final = _health(gateway, model)
|
||||
assert (down["healthy_count"], down["unhealthy_count"]) == (0, 1), down
|
||||
with wire_server(pcb.respond, port=port) as wire:
|
||||
_health(gateway, model)
|
||||
probes: Final = pcb.drained_posts(wire)
|
||||
assert [rv.newest_marker(request.body.decode()) for request in probes] == [None], probes
|
||||
served: Final = await _burst(gateway, model, after)
|
||||
assert len(served) == 12
|
||||
for item in served:
|
||||
_assert_answered_in_its_own_shape(item)
|
||||
_marked_once(pcb.drained_posts(wire), after, pcb.EXPLICIT)
|
||||
_assert_each_lands_once(spend, model, failed, served)
|
||||
|
||||
|
||||
def _slow(request: Request) -> Reply:
|
||||
reply: Final = pcb.respond(request)
|
||||
return dataclasses.replace(reply, pause_between_chunks=0.4) if reply.chunks else reply
|
||||
|
||||
|
||||
async def test_concurrent_slow_streams_each_complete_with_one_upstream_call(
|
||||
gateway: Gateway, spend: pcb.SpendLogs
|
||||
) -> None:
|
||||
calls: Final = _calls(6, ("chat",), stream=True)
|
||||
with wire_server(_slow) as wire, gateway.scenario() as scenario:
|
||||
model: Final = scenario.model(model=pcb.MODEL, api_base=f"{wire.url}/v1", drop_params=True)
|
||||
served: Final = await _burst(gateway, model, calls)
|
||||
assert len(served) == 6
|
||||
for item in served:
|
||||
_assert_answered_in_its_own_shape(item)
|
||||
_marked_once(pcb.drained_posts(wire), calls, pcb.EXPLICIT)
|
||||
_assert_each_lands_once(spend, model, (), served)
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class _Held:
|
||||
release: threading.Event
|
||||
markers: SimpleQueue[str]
|
||||
|
||||
def respond(self, request: Request) -> Reply:
|
||||
marker: Final = rv.newest_marker(request.body.decode()) if request.method == "POST" else None
|
||||
if marker is None:
|
||||
return pcb.respond(request)
|
||||
self.markers.put(marker)
|
||||
if not self.release.wait(timeout=60):
|
||||
return rv.error(504, "the burst was never released", "held")
|
||||
return pcb.respond(request)
|
||||
|
||||
|
||||
def _worker_pids(owned: OwnedProxy) -> tuple[int, ...]:
|
||||
return eventually(lambda: pcb.started_worker_pids(owned.log), lambda pids: len(pids) == 2, seconds=30)
|
||||
|
||||
|
||||
async def _hold_burst(
|
||||
held: _Held, candidate: Gateway, model: str, calls: tuple[_Call, ...]
|
||||
) -> asyncio.Task[tuple[_Served, ...]]:
|
||||
burst: Final = asyncio.create_task(_burst(candidate, model, calls, tolerate_transport_errors=True))
|
||||
await asyncio.to_thread(eventually, held.markers.qsize, lambda size: size == len(calls), 60)
|
||||
return burst
|
||||
|
||||
|
||||
async def test_worker_sigkill_mid_burst_leaves_the_sibling_answering(
|
||||
gateway: Gateway, tmp_path: Path, spend: pcb.SpendLogs
|
||||
) -> None:
|
||||
calls: Final = _calls(20, ("chat",), stream=False)
|
||||
held: Final = _Held(threading.Event(), SimpleQueue())
|
||||
with wire_server(held.respond) as wire:
|
||||
config: Final = _global_config(tmp_path, f"{wire.url}/v1")
|
||||
with owned_proxy_process(gateway, tmp_path, {}, config=config, workers=2) as owned:
|
||||
candidate: Final = owned.gateway
|
||||
workers: Final = _worker_pids(owned)
|
||||
burst: Final = await _hold_burst(held, candidate, _GLOBAL_UNSET, calls)
|
||||
held_by: Final = MappingProxyType({pid: pcb.open_upstream_connections(pid, wire.url) for pid in workers})
|
||||
assert sum(held_by.values()) == 20, held_by
|
||||
victim_pid, survivor_pid = sorted(workers, key=held_by.__getitem__)
|
||||
victim: Final = psutil.Process(victim_pid)
|
||||
victim.suspend()
|
||||
victim.send_signal(signal.SIGKILL)
|
||||
held.release.set()
|
||||
served: Final = await burst
|
||||
assert held_by[survivor_pid] >= 10, held_by
|
||||
assert len(served) == held_by[survivor_pid], (held_by, len(served))
|
||||
for item in served:
|
||||
_assert_answered_in_its_own_shape(item)
|
||||
_marked_once(pcb.drained_posts(wire), calls, pcb.EXPLICIT)
|
||||
follow_up: Final = uuid.uuid4().hex
|
||||
call_id: Final = _completion(_chat(candidate, _GLOBAL_UNSET, follow_up, "yes"), follow_up)
|
||||
pcb.assert_marker(_user_block_on_wire(wire, follow_up), None)
|
||||
spend.landed(_GLOBAL_UNSET, call_id, pcb.response_id(follow_up))
|
||||
|
||||
|
||||
async def test_proxy_restart_mid_burst_never_lands_a_served_call_twice(
|
||||
gateway: Gateway, tmp_path: Path, record_property: _RecordProperty, spend: pcb.SpendLogs
|
||||
) -> None:
|
||||
calls: Final = _calls(20, ("chat",), stream=False)
|
||||
held: Final = _Held(threading.Event(), SimpleQueue())
|
||||
with wire_server(held.respond) as wire:
|
||||
config: Final = _global_config(tmp_path, f"{wire.url}/v1")
|
||||
with owned_proxy_process(gateway, tmp_path, {}, config=config, workers=2) as first:
|
||||
_worker_pids(first)
|
||||
burst: Final = await _hold_burst(held, first.gateway, _GLOBAL_UNSET, calls)
|
||||
first.process.terminate()
|
||||
held.release.set()
|
||||
served: Final = await burst
|
||||
for item in served:
|
||||
_assert_answered_in_its_own_shape(item)
|
||||
second_directory: Final = tmp_path / "second"
|
||||
second_directory.mkdir()
|
||||
with owned_proxy_process(gateway, second_directory, {}, config=config, workers=2) as second:
|
||||
follow_up: Final = uuid.uuid4().hex
|
||||
call_id: Final = _completion(_chat(second.gateway, _GLOBAL_UNSET, follow_up, pcb.EXPLICIT), follow_up)
|
||||
pcb.assert_marker(_user_block_on_wire(wire, follow_up), pcb.EXPLICIT)
|
||||
spend.landed(_GLOBAL_UNSET, call_id, pcb.response_id(follow_up))
|
||||
counts: Final = Counter(string_value(row["litellm_call_id"]) for row in spend.rows_for(_GLOBAL_UNSET))
|
||||
assert all(count == 1 for count in counts.values()), counts
|
||||
landed: Final = sum(1 for item in served if item.call_id in counts)
|
||||
record_property("served", len(served))
|
||||
record_property("landed", landed)
|
||||
record_property("lost_responses", len(calls) - len(served))
|
||||
|
|
@ -0,0 +1,510 @@
|
|||
import asyncio
|
||||
import uuid
|
||||
from collections.abc import Iterator, Mapping, Sequence
|
||||
from dataclasses import dataclass
|
||||
from typing import Final, Literal, TypeAlias
|
||||
|
||||
import anthropic
|
||||
import httpx
|
||||
import openai
|
||||
import pytest
|
||||
from integration._support import prompt_cache_breakpoint as pcb
|
||||
from integration._support import responses_vendor as rv
|
||||
from integration._support.client import Gateway, eventually, gateway_from_environment, string_value
|
||||
from integration._support.wire import Request, Wire, wire_server
|
||||
from pydantic import JsonValue
|
||||
|
||||
pytestmark: Final = pytest.mark.timeout(120)
|
||||
|
||||
Mode: TypeAlias = Literal["on", "off"]
|
||||
|
||||
_UNKNOWN_KEY: Final[Mapping[str, JsonValue]] = {"mode": "explicit", "note": "kept"}
|
||||
_MALFORMED: Final[tuple[tuple[str, JsonValue], ...]] = (
|
||||
("string", "yes"),
|
||||
("int", 1),
|
||||
("list", ["explicit"]),
|
||||
("empty-string", ""),
|
||||
("5kb-string", "x" * 5000),
|
||||
("empty-object", {}),
|
||||
("bogus-mode", {"mode": "bogus"}),
|
||||
("bad-ttl", {"mode": "explicit", "ttl": "1h"}),
|
||||
)
|
||||
_MALFORMED_IDS: Final = tuple(name for name, _ in _MALFORMED)
|
||||
_MALFORMED_VALUES: Final = tuple(value for _, value in _MALFORMED)
|
||||
_CASES: Final[tuple[tuple[str, Mode, JsonValue, JsonValue], ...]] = (
|
||||
("valid-on", "on", pcb.EXPLICIT, pcb.EXPLICIT),
|
||||
("valid-off", "off", pcb.EXPLICIT, pcb.EXPLICIT),
|
||||
("malformed-on", "on", "yes", None),
|
||||
("malformed-off", "off", "yes", "yes"),
|
||||
)
|
||||
_CASE_IDS: Final = tuple(case[0] for case in _CASES)
|
||||
_CASE_VALUES: Final = tuple(case[1:] for case in _CASES)
|
||||
_ADAPTER_CASES: Final[tuple[tuple[str, Mode, JsonValue, JsonValue], ...]] = (
|
||||
("valid-on", "on", pcb.EXPLICIT, pcb.EXPLICIT),
|
||||
("valid-off", "off", pcb.EXPLICIT, pcb.EXPLICIT),
|
||||
("malformed-on", "on", "yes", "yes"),
|
||||
("malformed-off", "off", "yes", "yes"),
|
||||
)
|
||||
_ADAPTER_CASE_IDS: Final = tuple(case[0] for case in _ADAPTER_CASES)
|
||||
_ADAPTER_CASE_VALUES: Final = tuple(case[1:] for case in _ADAPTER_CASES)
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class _Bridge:
|
||||
gateway: Gateway
|
||||
wire: Wire
|
||||
on: str
|
||||
off: str
|
||||
injecting_on: str
|
||||
injecting_off: str
|
||||
spend: pcb.SpendLogs
|
||||
|
||||
def model(self, mode: Mode) -> str:
|
||||
return self.on if mode == "on" else self.off
|
||||
|
||||
def injecting(self, mode: Mode) -> str:
|
||||
return self.injecting_on if mode == "on" else self.injecting_off
|
||||
|
||||
@property
|
||||
def api_base(self) -> str:
|
||||
return f"{self.wire.url}/v1"
|
||||
|
||||
|
||||
@pytest.fixture(scope="module")
|
||||
def bridge() -> Iterator[_Bridge]:
|
||||
with (
|
||||
wire_server(pcb.respond) as wire,
|
||||
gateway_from_environment() as gateway,
|
||||
gateway.scenario() as scenario,
|
||||
pcb.spend_logs() as spend,
|
||||
):
|
||||
api_base: Final = f"{wire.url}/v1"
|
||||
yield _Bridge(
|
||||
gateway,
|
||||
wire,
|
||||
scenario.model(model=pcb.MODEL, api_base=api_base, drop_params=True),
|
||||
scenario.model(model=pcb.MODEL, api_base=api_base),
|
||||
scenario.model(model=pcb.MODEL, api_base=api_base, drop_params=True, **pcb.INJECTION),
|
||||
scenario.model(model=pcb.MODEL, api_base=api_base, **pcb.INJECTION),
|
||||
spend,
|
||||
)
|
||||
|
||||
|
||||
def _v1(gateway: Gateway) -> str:
|
||||
return str(gateway.client.base_url).rstrip("/") + "/v1"
|
||||
|
||||
|
||||
def _user(marker: str, breakpoint: JsonValue) -> dict[str, JsonValue]:
|
||||
return {"role": "user", "content": [pcb.marked(pcb.text(pcb.prompt(marker)), breakpoint)]}
|
||||
|
||||
|
||||
def _chat(
|
||||
bridge: _Bridge, model: str, messages: Sequence[JsonValue], *, stream: bool = False, key: str | None = None
|
||||
) -> httpx.Response:
|
||||
return bridge.gateway.request(
|
||||
"POST",
|
||||
"/v1/chat/completions",
|
||||
{"model": model, "messages": list(messages), "stream": stream, **pcb.NO_CACHE},
|
||||
key=key,
|
||||
)
|
||||
|
||||
|
||||
def _completion(response: httpx.Response, marker: str) -> str:
|
||||
assert response.status_code == 200, response.text
|
||||
body: Final = rv.JSON_OBJECT.validate_json(response.text)
|
||||
assert rv.same_response(string_value(body["id"]), pcb.response_id(marker)), body
|
||||
(choice,) = rv.ITEMS.validate_python(body["choices"])
|
||||
assert rv.JSON_OBJECT.validate_python(choice["message"])["content"] == rv.answer(marker), body
|
||||
return response.headers["x-litellm-call-id"]
|
||||
|
||||
|
||||
def _wire_body(request: Request, *, stream: bool = False) -> dict[str, JsonValue]:
|
||||
body: Final = pcb.body_of(request)
|
||||
assert body["model"] == "gpt-6.1-sol", body
|
||||
assert (body.get("stream") is True) is stream, body
|
||||
return body
|
||||
|
||||
|
||||
def _user_block_on_wire(bridge: _Bridge, marker: str, *, stream: bool = False) -> dict[str, JsonValue]:
|
||||
request: Final = pcb.posted(bridge.wire, marker)
|
||||
block: Final = pcb.single_block(pcb.input_items(request), "user")
|
||||
_wire_body(request, stream=stream)
|
||||
assert block["type"] == "input_text" and block["text"] == pcb.prompt(marker), block
|
||||
return block
|
||||
|
||||
|
||||
def test_openai_sdk_sends_a_valid_marker_through_the_bridge(bridge: _Bridge) -> None:
|
||||
marker: Final = uuid.uuid4().hex
|
||||
with openai.OpenAI(api_key=bridge.gateway.key, base_url=_v1(bridge.gateway), max_retries=0) as client:
|
||||
raw: Final = client.chat.completions.with_raw_response.create(
|
||||
model=bridge.on, messages=[_user(marker, pcb.EXPLICIT)], extra_body=dict(pcb.NO_CACHE)
|
||||
)
|
||||
completion: Final = raw.parse()
|
||||
assert rv.same_response(completion.id, pcb.response_id(marker)), completion
|
||||
assert completion.choices[0].message.content == rv.answer(marker), completion
|
||||
pcb.assert_marker(_user_block_on_wire(bridge, marker), pcb.EXPLICIT)
|
||||
bridge.spend.landed(bridge.on, raw.headers["x-litellm-call-id"], pcb.response_id(marker))
|
||||
|
||||
|
||||
def test_openai_sdk_stream_carries_the_system_list_marker(bridge: _Bridge) -> None:
|
||||
marker: Final = uuid.uuid4().hex
|
||||
system: Final[dict[str, JsonValue]] = {"role": "system", "content": [pcb.marked(pcb.text("sys"), pcb.EXPLICIT)]}
|
||||
with openai.OpenAI(api_key=bridge.gateway.key, base_url=_v1(bridge.gateway), max_retries=0) as client:
|
||||
raw: Final = client.chat.completions.with_raw_response.create(
|
||||
model=bridge.on,
|
||||
messages=[system, {"role": "user", "content": pcb.prompt(marker)}],
|
||||
stream=True,
|
||||
extra_body=dict(pcb.NO_CACHE),
|
||||
)
|
||||
chunks: Final = tuple(raw.parse())
|
||||
assert chunks and rv.same_response(chunks[0].id, pcb.response_id(marker)), chunks
|
||||
streamed: Final = "".join(chunk.choices[0].delta.content or "" for chunk in chunks if chunk.choices)
|
||||
assert streamed == rv.answer(marker), chunks
|
||||
request: Final = pcb.posted(bridge.wire, marker)
|
||||
_wire_body(request, stream=True)
|
||||
(system_block,) = pcb.content_of(pcb.input_items(request), "system")
|
||||
assert system_block["type"] == "input_text" and system_block["text"] == "sys", system_block
|
||||
pcb.assert_marker(system_block, pcb.EXPLICIT)
|
||||
bridge.spend.landed(bridge.on, raw.headers["x-litellm-call-id"], pcb.response_id(marker))
|
||||
|
||||
|
||||
async def test_async_openai_sdk_keeps_the_ttl_without_drop_params(bridge: _Bridge) -> None:
|
||||
marker: Final = uuid.uuid4().hex
|
||||
async with openai.AsyncOpenAI(api_key=bridge.gateway.key, base_url=_v1(bridge.gateway), max_retries=0) as client:
|
||||
raw: Final = await client.chat.completions.with_raw_response.create(
|
||||
model=bridge.off, messages=[_user(marker, pcb.EXPLICIT_30M)], extra_body=dict(pcb.NO_CACHE)
|
||||
)
|
||||
completion: Final = raw.parse()
|
||||
assert rv.same_response(completion.id, pcb.response_id(marker)), completion
|
||||
assert completion.choices[0].message.content == rv.answer(marker), completion
|
||||
pcb.assert_marker(_user_block_on_wire(bridge, marker), pcb.EXPLICIT_30M)
|
||||
bridge.spend.landed(bridge.off, raw.headers["x-litellm-call-id"], pcb.response_id(marker))
|
||||
|
||||
|
||||
@pytest.mark.parametrize("mode", ("on", "off"))
|
||||
@pytest.mark.parametrize("breakpoint", (pcb.EXPLICIT, pcb.EXPLICIT_30M), ids=("explicit", "ttl"))
|
||||
def test_valid_marker_shapes_reach_the_wire_unchanged(bridge: _Bridge, mode: Mode, breakpoint: JsonValue) -> None:
|
||||
marker: Final = uuid.uuid4().hex
|
||||
call_id: Final = _completion(_chat(bridge, bridge.model(mode), [_user(marker, breakpoint)]), marker)
|
||||
pcb.assert_marker(_user_block_on_wire(bridge, marker), breakpoint)
|
||||
bridge.spend.landed(bridge.model(mode), call_id, pcb.response_id(marker))
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("mode", "expected"), (("on", pcb.EXPLICIT), ("off", _UNKNOWN_KEY)), ids=("normalized-on", "verbatim-off")
|
||||
)
|
||||
def test_marker_with_an_unknown_key(bridge: _Bridge, mode: Mode, expected: JsonValue) -> None:
|
||||
marker: Final = uuid.uuid4().hex
|
||||
call_id: Final = _completion(_chat(bridge, bridge.model(mode), [_user(marker, _UNKNOWN_KEY)]), marker)
|
||||
pcb.assert_marker(_user_block_on_wire(bridge, marker), expected)
|
||||
bridge.spend.landed(bridge.model(mode), call_id, pcb.response_id(marker))
|
||||
|
||||
|
||||
def _second_block_on_wire(bridge: _Bridge, marker: str) -> dict[str, JsonValue]:
|
||||
request: Final = pcb.posted(bridge.wire, marker)
|
||||
_wire_body(request)
|
||||
first, second = pcb.content_of(pcb.input_items(request), "user")
|
||||
assert first == {"type": "input_text", "text": pcb.prompt(marker)}, first
|
||||
return second
|
||||
|
||||
|
||||
@pytest.mark.parametrize("mode", ("on", "off"))
|
||||
@pytest.mark.parametrize("kind", pcb.KINDS)
|
||||
def test_valid_marker_is_carried_on_every_block_kind(bridge: _Bridge, kind: pcb.Kind, mode: Mode) -> None:
|
||||
marker: Final = uuid.uuid4().hex
|
||||
content: Final[list[JsonValue]] = [
|
||||
pcb.text(pcb.prompt(marker)),
|
||||
pcb.marked(pcb.block(kind, "second"), pcb.EXPLICIT),
|
||||
]
|
||||
call_id: Final = _completion(_chat(bridge, bridge.model(mode), [{"role": "user", "content": content}]), marker)
|
||||
second: Final = _second_block_on_wire(bridge, marker)
|
||||
assert second["type"] == pcb.WIRE_TYPE[kind], second
|
||||
pcb.assert_marker(second, pcb.EXPLICIT)
|
||||
bridge.spend.landed(bridge.model(mode), call_id, pcb.response_id(marker))
|
||||
|
||||
|
||||
@pytest.mark.parametrize("kind", pcb.KINDS)
|
||||
def test_malformed_marker_is_dropped_on_every_block_kind(bridge: _Bridge, kind: pcb.Kind) -> None:
|
||||
marker: Final = uuid.uuid4().hex
|
||||
content: Final[list[JsonValue]] = [pcb.text(pcb.prompt(marker)), pcb.marked(pcb.block(kind, "second"), "yes")]
|
||||
call_id: Final = _completion(_chat(bridge, bridge.on, [{"role": "user", "content": content}]), marker)
|
||||
second: Final = _second_block_on_wire(bridge, marker)
|
||||
assert second["type"] == pcb.WIRE_TYPE[kind], second
|
||||
pcb.assert_marker(second, None)
|
||||
bridge.spend.landed(bridge.on, call_id, pcb.response_id(marker))
|
||||
|
||||
|
||||
@pytest.mark.parametrize(("mode", "breakpoint", "expected"), _CASE_VALUES, ids=_CASE_IDS)
|
||||
def test_tool_output_marker(bridge: _Bridge, mode: Mode, breakpoint: JsonValue, expected: JsonValue) -> None:
|
||||
marker: Final = uuid.uuid4().hex
|
||||
messages: Final[list[JsonValue]] = [
|
||||
{"role": "user", "content": pcb.prompt(marker)},
|
||||
{
|
||||
"role": "assistant",
|
||||
"content": None,
|
||||
"tool_calls": [{"id": "call_1", "type": "function", "function": {"name": "lookup", "arguments": "{}"}}],
|
||||
},
|
||||
{"role": "tool", "tool_call_id": "call_1", "content": [pcb.marked(pcb.text("found it"), breakpoint)]},
|
||||
]
|
||||
call_id: Final = _completion(_chat(bridge, bridge.model(mode), messages), marker)
|
||||
request: Final = pcb.posted(bridge.wire, marker)
|
||||
_wire_body(request)
|
||||
(output,) = pcb.function_output(pcb.input_items(request), "call_1")
|
||||
assert output["type"] == "input_text" and output["text"] == "found it", output
|
||||
pcb.assert_marker(output, expected)
|
||||
bridge.spend.landed(bridge.model(mode), call_id, pcb.response_id(marker))
|
||||
|
||||
|
||||
@pytest.mark.parametrize(("mode", "breakpoint", "expected"), _CASE_VALUES, ids=_CASE_IDS)
|
||||
def test_assistant_list_marker(bridge: _Bridge, mode: Mode, breakpoint: JsonValue, expected: JsonValue) -> None:
|
||||
marker: Final = uuid.uuid4().hex
|
||||
messages: Final[list[JsonValue]] = [
|
||||
{"role": "user", "content": pcb.prompt(marker)},
|
||||
{"role": "assistant", "content": [pcb.marked(pcb.text("earlier answer"), breakpoint)]},
|
||||
{"role": "user", "content": "and again"},
|
||||
]
|
||||
call_id: Final = _completion(_chat(bridge, bridge.model(mode), messages), marker)
|
||||
request: Final = pcb.posted(bridge.wire, marker)
|
||||
_wire_body(request)
|
||||
(earlier,) = pcb.content_of(pcb.input_items(request), "assistant")
|
||||
assert earlier["type"] == "output_text" and earlier["text"] == "earlier answer", earlier
|
||||
pcb.assert_marker(earlier, expected)
|
||||
bridge.spend.landed(bridge.model(mode), call_id, pcb.response_id(marker))
|
||||
|
||||
|
||||
def test_injected_system_marker_survives_a_trailing_audio_block(bridge: _Bridge) -> None:
|
||||
marker: Final = uuid.uuid4().hex
|
||||
system: Final[dict[str, JsonValue]] = {"role": "system", "content": [pcb.text("sys"), pcb.block("input_audio", "")]}
|
||||
messages: Final[list[JsonValue]] = [system, {"role": "user", "content": pcb.prompt(marker)}]
|
||||
call_id: Final = _completion(_chat(bridge, bridge.injecting_on, messages), marker)
|
||||
request: Final = pcb.posted(bridge.wire, marker)
|
||||
body: Final = _wire_body(request)
|
||||
assert body["prompt_cache_options"] == {"mode": "explicit"}, body
|
||||
first, audio = pcb.content_of(pcb.input_items(request), "system")
|
||||
assert first == {"type": "input_text", "text": "sys"}, first
|
||||
assert audio["type"] == "input_text", audio
|
||||
assert string_value(audio["text"]).startswith("{'type': 'input_audio'"), audio
|
||||
pcb.assert_marker(audio, pcb.EXPLICIT)
|
||||
bridge.spend.landed(bridge.injecting_on, call_id, pcb.response_id(marker))
|
||||
|
||||
|
||||
def test_injected_marker_on_a_string_system_message(bridge: _Bridge) -> None:
|
||||
marker: Final = uuid.uuid4().hex
|
||||
messages: Final[list[JsonValue]] = [
|
||||
{"role": "system", "content": "Answer briefly"},
|
||||
{"role": "user", "content": pcb.prompt(marker)},
|
||||
]
|
||||
call_id: Final = _completion(_chat(bridge, bridge.injecting_off, messages), marker)
|
||||
request: Final = pcb.posted(bridge.wire, marker)
|
||||
body: Final = _wire_body(request)
|
||||
assert body["prompt_cache_options"] == {"mode": "explicit"}, body
|
||||
system_block: Final = pcb.single_block(pcb.input_items(request), "system")
|
||||
assert system_block == {"type": "input_text", "text": "Answer briefly", "prompt_cache_breakpoint": pcb.EXPLICIT}
|
||||
bridge.spend.landed(bridge.injecting_off, call_id, pcb.response_id(marker))
|
||||
|
||||
|
||||
def _anthropic(bridge: _Bridge) -> anthropic.Anthropic:
|
||||
return anthropic.Anthropic(base_url=str(bridge.gateway.client.base_url), api_key=bridge.gateway.key, max_retries=0)
|
||||
|
||||
|
||||
def _anthropic_text(message: anthropic.types.Message, marker: str) -> None:
|
||||
(content,) = message.content
|
||||
assert content.type == "text" and content.text == rv.answer(marker), message
|
||||
|
||||
|
||||
@pytest.mark.parametrize(("mode", "breakpoint", "expected"), _ADAPTER_CASE_VALUES, ids=_ADAPTER_CASE_IDS)
|
||||
def test_anthropic_sdk_marker_on_user_text(
|
||||
bridge: _Bridge, mode: Mode, breakpoint: JsonValue, expected: JsonValue
|
||||
) -> None:
|
||||
marker: Final = uuid.uuid4().hex
|
||||
with _anthropic(bridge) as client:
|
||||
raw: Final = client.messages.with_raw_response.create(
|
||||
model=bridge.model(mode),
|
||||
max_tokens=64,
|
||||
messages=[{"role": "user", "content": [pcb.marked(pcb.text(pcb.prompt(marker)), breakpoint)]}],
|
||||
extra_body=dict(pcb.NO_CACHE),
|
||||
)
|
||||
_anthropic_text(raw.parse(), marker)
|
||||
pcb.assert_marker(_user_block_on_wire(bridge, marker), expected)
|
||||
bridge.spend.landed(bridge.model(mode), raw.headers["x-litellm-call-id"], None)
|
||||
|
||||
|
||||
def test_anthropic_sdk_system_string_gets_the_injected_marker(bridge: _Bridge) -> None:
|
||||
marker: Final = uuid.uuid4().hex
|
||||
with _anthropic(bridge) as client:
|
||||
raw: Final = client.messages.with_raw_response.create(
|
||||
model=bridge.injecting_off,
|
||||
max_tokens=64,
|
||||
system="Answer briefly",
|
||||
messages=[{"role": "user", "content": pcb.prompt(marker)}],
|
||||
extra_body=dict(pcb.NO_CACHE),
|
||||
)
|
||||
_anthropic_text(raw.parse(), marker)
|
||||
request: Final = pcb.posted(bridge.wire, marker)
|
||||
body: Final = _wire_body(request)
|
||||
assert body["prompt_cache_options"] == {"mode": "explicit"}, body
|
||||
instruction: Final = pcb.instruction_block(pcb.input_items(request))
|
||||
assert instruction == {"type": "input_text", "text": "Answer briefly", "prompt_cache_breakpoint": pcb.EXPLICIT}
|
||||
bridge.spend.landed(bridge.injecting_off, raw.headers["x-litellm-call-id"], None)
|
||||
|
||||
|
||||
def test_native_responses_request_never_enters_the_bridge(bridge: _Bridge) -> None:
|
||||
marker: Final = uuid.uuid4().hex
|
||||
block: Final[dict[str, JsonValue]] = {"type": "input_text", "text": pcb.prompt(marker)}
|
||||
response: Final = bridge.gateway.request(
|
||||
"POST",
|
||||
"/v1/responses",
|
||||
{
|
||||
"model": bridge.on,
|
||||
"input": [{"type": "message", "role": "user", "content": [pcb.marked(block, _UNKNOWN_KEY)]}],
|
||||
**pcb.NO_CACHE,
|
||||
},
|
||||
)
|
||||
assert response.status_code == 200, response.text
|
||||
body: Final = rv.JSON_OBJECT.validate_json(response.text)
|
||||
assert rv.same_response(string_value(body["id"]), pcb.response_id(marker)), body
|
||||
assert rv.answer(marker) in response.text, response.text
|
||||
request: Final = pcb.posted(bridge.wire, marker)
|
||||
_wire_body(request)
|
||||
on_wire: Final = pcb.single_block(pcb.input_items(request), "user")
|
||||
assert on_wire == pcb.marked(block, _UNKNOWN_KEY), on_wire
|
||||
bridge.spend.landed(bridge.on, response.headers["x-litellm-call-id"], pcb.response_id(marker))
|
||||
|
||||
|
||||
@pytest.mark.parametrize("breakpoint", _MALFORMED_VALUES, ids=_MALFORMED_IDS)
|
||||
def test_malformed_marker_is_dropped_under_drop_params(bridge: _Bridge, breakpoint: JsonValue) -> None:
|
||||
marker: Final = uuid.uuid4().hex
|
||||
call_id: Final = _completion(_chat(bridge, bridge.on, [_user(marker, breakpoint)]), marker)
|
||||
pcb.assert_marker(_user_block_on_wire(bridge, marker), None)
|
||||
bridge.spend.landed(bridge.on, call_id, pcb.response_id(marker))
|
||||
|
||||
|
||||
@pytest.mark.parametrize("breakpoint", _MALFORMED_VALUES, ids=_MALFORMED_IDS)
|
||||
def test_malformed_marker_passes_verbatim_without_drop_params(bridge: _Bridge, breakpoint: JsonValue) -> None:
|
||||
marker: Final = uuid.uuid4().hex
|
||||
call_id: Final = _completion(_chat(bridge, bridge.off, [_user(marker, breakpoint)]), marker)
|
||||
pcb.assert_marker(_user_block_on_wire(bridge, marker), breakpoint)
|
||||
bridge.spend.landed(bridge.off, call_id, pcb.response_id(marker))
|
||||
|
||||
|
||||
def test_two_marked_blocks_are_both_carried(bridge: _Bridge) -> None:
|
||||
marker: Final = uuid.uuid4().hex
|
||||
content: Final[list[JsonValue]] = [
|
||||
pcb.marked(pcb.text(pcb.prompt(marker)), pcb.EXPLICIT),
|
||||
pcb.marked(pcb.text("and more"), pcb.EXPLICIT_30M),
|
||||
]
|
||||
call_id: Final = _completion(_chat(bridge, bridge.on, [{"role": "user", "content": content}]), marker)
|
||||
request: Final = pcb.posted(bridge.wire, marker)
|
||||
_wire_body(request)
|
||||
first, second = pcb.content_of(pcb.input_items(request), "user")
|
||||
assert first == {"type": "input_text", "text": pcb.prompt(marker), "prompt_cache_breakpoint": pcb.EXPLICIT}
|
||||
assert second == {"type": "input_text", "text": "and more", "prompt_cache_breakpoint": pcb.EXPLICIT_30M}
|
||||
bridge.spend.landed(bridge.on, call_id, pcb.response_id(marker))
|
||||
|
||||
|
||||
def test_wrong_key_is_refused_before_the_wire(bridge: _Bridge) -> None:
|
||||
marker: Final = uuid.uuid4().hex
|
||||
response: Final = _chat(bridge, bridge.on, [_user(marker, pcb.EXPLICIT)], key="sk-wrong")
|
||||
assert response.status_code == 401, response.text
|
||||
assert pcb.with_marker(pcb.drained_posts(bridge.wire), marker) == (), marker
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("mode", "status"), (("on", 400), ("off", 400), ("on", 401)), ids=("400-on", "400-off", "401-on")
|
||||
)
|
||||
def test_upstream_error_reaches_the_caller_once(bridge: _Bridge, mode: Mode, status: int) -> None:
|
||||
marker: Final = uuid.uuid4().hex
|
||||
failing: Final[dict[str, JsonValue]] = {
|
||||
"role": "user",
|
||||
"content": [pcb.marked(pcb.text(f"{pcb.prompt(marker)} fail-{status}"), pcb.EXPLICIT)],
|
||||
}
|
||||
response: Final = _chat(bridge, bridge.model(mode), [failing])
|
||||
assert response.status_code == status, response.text
|
||||
assert f"scripted {status} marker-{marker}" in response.text, response.text
|
||||
(request,) = pcb.with_marker(pcb.drained_posts(bridge.wire), marker)
|
||||
pcb.assert_marker(pcb.single_block(pcb.input_items(request), "user"), pcb.EXPLICIT)
|
||||
bridge.spend.landed(bridge.model(mode), response.headers["x-litellm-call-id"], None, status="failure")
|
||||
follow_up: Final = uuid.uuid4().hex
|
||||
call_id: Final = _completion(_chat(bridge, bridge.model(mode), [_user(follow_up, pcb.EXPLICIT)]), follow_up)
|
||||
pcb.assert_marker(_user_block_on_wire(bridge, follow_up), pcb.EXPLICIT)
|
||||
bridge.spend.landed(bridge.model(mode), call_id, pcb.response_id(follow_up))
|
||||
|
||||
|
||||
def test_null_drop_params_on_the_deployment_means_off(bridge: _Bridge) -> None:
|
||||
marker: Final = uuid.uuid4().hex
|
||||
with bridge.gateway.scenario() as scenario:
|
||||
model: Final = scenario.model(model=pcb.MODEL, api_base=bridge.api_base, drop_params=None)
|
||||
call_id: Final = _completion(_chat(bridge, model, [_user(marker, "yes")]), marker)
|
||||
pcb.assert_marker(_user_block_on_wire(bridge, marker), "yes")
|
||||
bridge.spend.landed(model, call_id, pcb.response_id(marker))
|
||||
|
||||
|
||||
@pytest.mark.parametrize("mode", ("on", "off"))
|
||||
@pytest.mark.parametrize("shape", ("null", "missing"))
|
||||
def test_null_or_missing_marker_sends_a_plain_block(bridge: _Bridge, shape: str, mode: Mode) -> None:
|
||||
marker: Final = uuid.uuid4().hex
|
||||
block: Final = pcb.marked(pcb.text(pcb.prompt(marker)), None) if shape == "null" else pcb.text(pcb.prompt(marker))
|
||||
call_id: Final = _completion(_chat(bridge, bridge.model(mode), [{"role": "user", "content": [block]}]), marker)
|
||||
assert _user_block_on_wire(bridge, marker) == {"type": "input_text", "text": pcb.prompt(marker)}
|
||||
bridge.spend.landed(bridge.model(mode), call_id, pcb.response_id(marker))
|
||||
|
||||
|
||||
async def _send_marked(client: httpx.AsyncClient, key: str, model: str, marker: str) -> httpx.Response:
|
||||
return await client.post(
|
||||
"/v1/chat/completions",
|
||||
json={"model": model, "messages": [_user(marker, "yes")], **pcb.NO_CACHE},
|
||||
headers={"Authorization": f"Bearer {key}"},
|
||||
)
|
||||
|
||||
|
||||
def _probe_marker(bridge: _Bridge, model: str) -> JsonValue:
|
||||
marker: Final = uuid.uuid4().hex
|
||||
_completion(_chat(bridge, model, [_user(marker, "yes")]), marker)
|
||||
return _user_block_on_wire(bridge, marker).get("prompt_cache_breakpoint")
|
||||
|
||||
|
||||
@pytest.mark.timeout(180)
|
||||
async def test_flipping_drop_params_mid_burst_keeps_every_marked_request_answered(bridge: _Bridge) -> None:
|
||||
gateway: Final = bridge.gateway
|
||||
with gateway.scenario() as scenario:
|
||||
model: Final = scenario.model(model=pcb.MODEL, api_base=bridge.api_base, drop_params=True)
|
||||
identity: Final = pcb.model_id(gateway.get("/model/info")["data"], model)
|
||||
markers: Final = tuple(uuid.uuid4().hex for _ in range(20))
|
||||
async with httpx.AsyncClient(base_url=str(gateway.client.base_url), timeout=60, trust_env=False) as client:
|
||||
burst: Final = asyncio.gather(*(_send_marked(client, gateway.key, model, marker) for marker in markers))
|
||||
updated: Final = await asyncio.to_thread(
|
||||
gateway.request,
|
||||
"POST",
|
||||
"/model/update",
|
||||
{
|
||||
"model_name": model,
|
||||
"litellm_params": {"model": pcb.MODEL, "drop_params": False},
|
||||
"model_info": {"id": identity},
|
||||
},
|
||||
)
|
||||
responses: Final = await burst
|
||||
assert updated.status_code == 200, updated.text
|
||||
for marker, response in zip(markers, responses, strict=True):
|
||||
_completion(response, marker)
|
||||
posts: Final = pcb.drained_posts(bridge.wire)
|
||||
for marker in markers:
|
||||
(request,) = pcb.with_marker(posts, marker)
|
||||
seen: Final = pcb.single_block(pcb.input_items(request), "user").get("prompt_cache_breakpoint")
|
||||
assert seen in (None, "yes"), request.body
|
||||
flipped: Final = eventually(lambda: _probe_marker(bridge, model), lambda seen: seen == "yes", seconds=70)
|
||||
assert flipped == "yes"
|
||||
for marker, response in zip(markers, responses, strict=True):
|
||||
bridge.spend.landed(model, response.headers["x-litellm-call-id"], pcb.response_id(marker))
|
||||
|
||||
|
||||
def test_three_identical_marked_requests_are_each_sent_and_logged(bridge: _Bridge) -> None:
|
||||
marker: Final = uuid.uuid4().hex
|
||||
responses: Final = tuple(_chat(bridge, bridge.on, [_user(marker, pcb.EXPLICIT)]) for _ in range(3))
|
||||
call_ids: Final = tuple(_completion(response, marker) for response in responses)
|
||||
assert len(set(call_ids)) == 3, call_ids
|
||||
posts: Final = pcb.with_marker(pcb.drained_posts(bridge.wire), marker)
|
||||
assert len(posts) == 3, [request.body for request in posts]
|
||||
for request in posts:
|
||||
pcb.assert_marker(pcb.single_block(pcb.input_items(request), "user"), pcb.EXPLICIT)
|
||||
for call_id in call_ids:
|
||||
bridge.spend.landed(bridge.on, call_id, pcb.response_id(marker))
|
||||
Loading…
Add table
Reference in a new issue