Merge pull request #41075 from BerriAI/litellm_integration_providers

test: provider wire contracts, streaming and recovery
This commit is contained in:
yuneng-jiang 2026-09-16 17:20:43 -07:00 committed by GitHub
commit 545df49374
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
14 changed files with 967 additions and 4 deletions

View file

@ -21,3 +21,7 @@ Fixtures must contain synthetic data only. Keep private incident records and sou
Database cases own their temporary schemas, roles, constraints and proxy processes. They prove reader-versus-writer execution with PostgreSQL lock observations, exercise real transaction wait limits and verify rollback after a reached database failure
Accounting cases compare persisted input and output cost components against literal rates, including zero and default prices. Cache state models assert actual upstream calls, response identity and every persisted charge. Generated accounting tests have a 180-second test limit to accommodate the asynchronous spend writer; CircleCI keeps the whole shard capped at 11 minutes
Provider contracts exercise actual TCP requests with synthetic credentials and local protocol peers. The S3 verifier uses independently implemented equations, a published known-answer vector, a fixed signing clock and deliberately invalid signed requests. Bedrock cases clear ambient AWS credential sources and check the literal model path, loaded role references, STS requests and bearer-only behavior
Streaming checks send real HTTP transfer chunks, including one-byte partitions, fragmented tools, incomplete transfers and a cancellation barrier. They assert meaningful text, tool arguments, final usage and persisted cost. The Redis recovery case owns a separate database and Redis process, uses the supported one-second circuit-breaker recovery setting, waits for the real subscriber and verifies response data in Redis after restart. CircleCI reuses its existing Redis image for that extra process; it never pulls an image during tests

View file

@ -46,13 +46,13 @@ def stop_root_process(process: subprocess.Popen[bytes]) -> bool:
@contextmanager
def owned_proxy(gateway: Gateway, directory: Path, overrides: Mapping[str, str]) -> Iterator[Gateway]:
def owned_proxy(gateway: Gateway, directory: Path, overrides: Mapping[str, str], *, config: Path | None = None, remove_environment: tuple[str, ...] = ()) -> Iterator[Gateway]:
with socket.socket() as reserve:
reserve.bind(("127.0.0.1", 0))
port: Final = reserve.getsockname()[1]
root: Final = Path(__file__).resolve().parents[3]
environment: Final = {
**os.environ,
**{name: value for name, value in os.environ.items() if name not in remove_environment},
"LITELLM_MASTER_KEY": gateway.key,
"LITELLM_SALT_KEY": os.environ.get("LITELLM_SALT_KEY", "sk-integration-salt"),
"STORE_MODEL_IN_DB": "True",
@ -67,7 +67,7 @@ def owned_proxy(gateway: Gateway, directory: Path, overrides: Mapping[str, str])
"-m",
"integration._support.proxy",
"--config",
"tests/integration/proxy_config.yaml",
str(config or "tests/integration/proxy_config.yaml"),
"--host",
"127.0.0.1",
"--port",

View file

@ -0,0 +1,116 @@
import os
import shutil
import signal
import socket
import subprocess
import time
import uuid
from collections.abc import Iterator
from contextlib import contextmanager
from dataclasses import dataclass
from pathlib import Path
from typing import Final, TextIO
from redis import Redis
from redis.exceptions import ConnectionError as RedisConnectionError
@dataclass
class OwnedRedis:
host: str
port: int
command: tuple[str, ...]
log: TextIO
pid_file: str
process: subprocess.Popen | None = None
server_pid: int | None = None
def start(self) -> None:
assert self.process is None
self.process = subprocess.Popen(self.command, stdout=self.log, stderr=subprocess.STDOUT, start_new_session=True)
deadline: Final = time.monotonic() + 8
with Redis(host=self.host, port=self.port, socket_connect_timeout=0.2, socket_timeout=0.2) as client:
while True:
assert self.process.poll() is None, "Owned Redis exited before readiness"
try:
if client.ping():
actual: Final = int(client.info("server")["process_id"])
expected: Final = self.process.pid if self.command[0] != "docker" else int(subprocess.check_output(["docker", "exec", "redis-cache", "cat", self.pid_file], timeout=2))
assert actual == expected, "Redis readiness reached a different process"
self.server_pid = actual
return
except RedisConnectionError:
pass
assert time.monotonic() < deadline, "Owned Redis readiness deadline exceeded"
time.sleep(0.05)
def stop(self) -> None:
assert self.process is not None
failure = None
forced = False
try:
if self.process.poll() is None:
with Redis(host=self.host, port=self.port, socket_connect_timeout=1, socket_timeout=1) as client:
assert int(client.info("server")["process_id"]) == self.server_pid, "Redis ownership changed before shutdown"
client.shutdown(nosave=True)
except Exception as error:
failure = error
finally:
try:
self.process.wait(timeout=5)
except subprocess.TimeoutExpired:
forced = True
self.signal(signal.SIGTERM)
try:
self.process.wait(timeout=5)
except subprocess.TimeoutExpired:
self.signal(signal.SIGKILL)
self.process.wait(timeout=3)
self.process = None
self.server_pid = None
with Redis(host=self.host, port=self.port, socket_connect_timeout=0.2, socket_timeout=0.2) as client:
try:
client.ping()
except RedisConnectionError:
stopped = True
else:
stopped = False
assert stopped, "Owned Redis still serves after shutdown"
assert failure is None and not forced, f"Owned Redis required shutdown recovery: {failure!r}"
def signal(self, action: signal.Signals) -> None:
assert self.process is not None
if self.command[0] != "docker":
self.process.send_signal(action)
return
pid: Final = int(subprocess.check_output(["docker", "exec", "redis-cache", "cat", self.pid_file], timeout=2))
command: Final = subprocess.check_output(["docker", "exec", "redis-cache", "cat", f"/proc/{pid}/cmdline"], timeout=2)
assert self.pid_file.encode() in command, "Redis process ownership changed"
subprocess.run(["docker", "exec", "redis-cache", "kill", f"-{int(action)}", str(pid)], check=True, timeout=2)
@contextmanager
def owned_redis(directory: Path) -> Iterator[OwnedRedis]:
binary: Final = shutil.which("redis-server")
if binary:
with socket.socket() as reservation:
reservation.bind(("127.0.0.1", 0))
port = reservation.getsockname()[1]
host = "127.0.0.1"
prefix = (binary,)
else:
host = subprocess.check_output(["docker", "inspect", "--format", "{{range .NetworkSettings.Networks}}{{.IPAddress}}{{end}}", "redis-cache"], text=True).strip()
assert host, "CircleCI owned Redis container has no address"
port = 16379
prefix = ("docker", "exec", "redis-cache", "redis-server")
output: Final = Path(os.environ.get("INTEGRATION_RESULTS_DIR", str(directory)))
output.mkdir(parents=True, exist_ok=True)
with (output / "owned-redis-recovery.log").open("w") as log:
pid_file: Final = str(directory / "owned-redis.pid") if binary else f"/tmp/integration-redis-{uuid.uuid4().hex}.pid"
server: Final = OwnedRedis(host, port, (*prefix, "--port", str(port), "--set-proc-title", "no", "--pidfile", pid_file, "--bind", "0.0.0.0" if not binary else "127.0.0.1", "--protected-mode", "no", "--save", "", "--appendonly", "no"), log, pid_file)
try:
server.start()
yield server
finally:
if server.process is not None:
server.stop()

View file

@ -0,0 +1,25 @@
import hashlib
import hmac
from collections.abc import Mapping
from typing import Final
def encoded_path(value: str) -> str:
safe: Final = b"ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789-_.~/"
return "".join(chr(byte) if byte in safe else f"%{byte:02X}" for byte in value.encode("utf-8"))
def signature(
method: str, path: str, headers: Mapping[str, str], signed: str, body: bytes, secret: str, scope: str,
) -> tuple[str, str]:
"""AWS SigV4 equations, independent of botocore and LiteLLM's signer."""
canonical_headers: Final = "".join(name + ":" + " ".join(headers[name].split()) + "\n" for name in signed.split(";"))
canonical: Final = "\n".join((method, path, "", canonical_headers, signed, hashlib.sha256(body).hexdigest()))
canonical_hash: Final = hashlib.sha256(canonical.encode()).hexdigest()
date, region, service, terminator = scope.split("/")
assert terminator == "aws4_request"
key = ("AWS4" + secret).encode()
for part in (date, region, service, terminator):
key = hmac.new(key, part.encode(), hashlib.sha256).digest()
to_sign: Final = "\n".join(("AWS4-HMAC-SHA256", headers["x-amz-date"], scope, canonical_hash))
return canonical_hash, hmac.new(key, to_sign.encode(), hashlib.sha256).hexdigest()

View file

@ -31,6 +31,12 @@ INTERNAL_FIELDS: Final = frozenset(
)
def error_type(status: int) -> str:
if status == 429:
return "rate_limit_error"
return "invalid_request_error" if status < 500 else "server_error"
@dataclass(frozen=True, slots=True)
class Observation:
path: str
@ -66,7 +72,7 @@ class Provider:
status: Final = script.popleft()
if status != 200:
return JSONResponse(
{"error": {"message": "Controlled provider failure", "type": "api_error", "code": str(status)}},
{"error": {"message": "Controlled provider failure", "type": error_type(status), "code": str(status)}},
status_code=status,
)
return await chat_completions(request)

View file

@ -0,0 +1,113 @@
from __future__ import annotations
import threading
from collections.abc import Callable, Iterator, Mapping
from contextlib import contextmanager
from dataclasses import dataclass
from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
from queue import SimpleQueue
from typing import Final
@dataclass(frozen=True, slots=True)
class Request:
method: str
target: str
headers: Mapping[str, str]
body: bytes
@dataclass(frozen=True, slots=True)
class Reply:
status: int = 200
body: bytes = b"{}"
content_type: str = "application/json"
chunks: tuple[bytes, ...] | None = None
abort_after: int | None = None
gate_after_first: threading.Event | None = None
@dataclass(frozen=True, slots=True)
class Wire:
url: str
received: SimpleQueue[Request]
disconnected: SimpleQueue[str]
def drain(self) -> tuple[Request, ...]:
return tuple(self.received.get_nowait() for _ in range(self.received.qsize()))
@contextmanager
def wire_server(respond: Callable[[Request], Reply]) -> Iterator[Wire]:
"""Owned TCP peer; requests traverse the real HTTP client and serialization."""
received: Final[SimpleQueue[Request]] = SimpleQueue()
errors: Final[SimpleQueue[Exception]] = SimpleQueue()
disconnected: Final[SimpleQueue[str]] = SimpleQueue()
class Handler(BaseHTTPRequestHandler):
protocol_version = "HTTP/1.1"
timeout = 5
def respond(self) -> None:
request: Final = Request(
self.command, self.path,
{name.lower(): value for name, value in self.headers.items()},
self.rfile.read(int(self.headers.get("content-length", "0"))),
)
received.put(request)
try:
reply = respond(request)
except Exception as error:
errors.put(error)
reply = Reply(status=500)
self.send_response(reply.status)
self.send_header("content-type", reply.content_type)
if reply.chunks is None:
self.send_header("content-length", str(len(reply.body)))
else:
self.send_header("transfer-encoding", "chunked")
self.send_header("connection", "close")
self.end_headers()
try:
if reply.chunks is None:
self.wfile.write(reply.body)
else:
for index, chunk in enumerate(reply.chunks):
if reply.abort_after == index:
break
self.wfile.write(b"%x\r\n%s\r\n" % (len(chunk), chunk))
self.wfile.flush()
if index == 0 and reply.gate_after_first is not None:
assert reply.gate_after_first.wait(timeout=5), "Stream barrier was never released"
else:
self.wfile.write(b"0\r\n\r\n")
self.wfile.flush()
except (BrokenPipeError, ConnectionResetError):
disconnected.put(request.target)
except Exception as error:
errors.put(error)
self.close_connection = True
do_POST = respond
do_PUT = respond
do_GET = respond
do_DELETE = respond
def log_message(self, format: str, *args: object) -> None:
pass
class OwnedHTTPServer(ThreadingHTTPServer):
daemon_threads = False
with OwnedHTTPServer(("127.0.0.1", 0), Handler) as server:
thread: Final = threading.Thread(target=server.serve_forever, kwargs={"poll_interval": 0.05})
thread.start()
try:
yield Wire(f"http://127.0.0.1:{server.server_port}", received, disconnected)
finally:
server.shutdown()
thread.join(timeout=6)
assert not thread.is_alive(), "Owned HTTP server survived cleanup"
server.server_close()
failure: Final = None if errors.empty() else errors.get_nowait()
assert failure is None, f"Owned HTTP peer failed: {failure!r}"

View file

@ -103,6 +103,53 @@
],
"tests/integration/spend/test_cache_and_quota.py::test_repeated_hits_keep_response_identity_and_create_distinct_zero_cost_rows": [
"quota_management.response_cache.repeated_hits_preserve_identity_and_single_charge"
],
"tests/integration/providers/test_s3_wire.py::test_sigv4_verifier_matches_published_put_and_rejects_corruption": [
"other.provider_wire.s3.verifier_known_answer_and_negative_controls"
],
"tests/integration/providers/test_s3_wire.py::test_s3_sync_and_async_uploads_pass_independent_wire_verification": [
"other.provider_wire.s3.sync_async_reserved_keys_are_signed_and_accepted"
],
"tests/integration/providers/test_bedrock_auth_wire.py::test_bearer_only_sdk_sync_async_requests_do_not_require_aws_credentials": [
"other.provider_wire.bedrock.bearer_sdk_skips_credential_chain"
],
"tests/integration/providers/test_bedrock_auth_wire.py::test_bearer_environment_reference_loads_from_db_and_yaml_and_survives_reload": [
"other.provider_wire.bedrock.bearer_db_yaml_survives_reload"
],
"tests/integration/streaming/test_stream_contracts.py::test_generated_tcp_partitions_preserve_unicode_text_identity_and_final_usage": [
"other.streaming.byte_partitions.preserve_text_identity_and_usage"
],
"tests/integration/streaming/test_stream_contracts.py::test_fragmented_tool_names_and_arguments_keep_each_call_identity": [
"other.streaming.tools.fragmented_calls_keep_independent_arguments"
],
"tests/integration/streaming/test_stream_contracts.py::test_proxy_stream_usage_visibility_keeps_exact_persisted_charge": [
"other.streaming.usage.client_visibility_preserves_persisted_accounting"
],
"tests/integration/streaming/test_stream_contracts.py::test_truncated_http_stream_is_an_error_and_next_stream_succeeds": [
"other.streaming.failure.truncated_transport_raises_and_control_recovers"
],
"tests/integration/streaming/test_stream_contracts.py::test_client_cancellation_releases_the_actual_provider_connection": [
"other.streaming.cancellation.closes_actual_provider_connection"
],
"tests/integration/routing/test_observed_routing.py::test_retry_counts_and_public_errors_match_actual_provider_attempts": [
"other.routing.retries.several_attempts_reach_success_without_hidden_retries",
"other.routing.errors.nonretryable_and_exhausted_failures_remain_errors"
],
"tests/integration/routing/test_observed_routing.py::test_loaded_fallback_selects_expected_deployment_and_keeps_response_identity": [
"other.routing.fallback.loaded_configuration_selects_only_permitted_target"
],
"tests/integration/routing/test_observed_routing.py::test_saved_deployment_target_update_changes_wire_and_preserves_control": [
"other.routing.alias_update.persisted_target_changes_only_selected_route"
],
"tests/integration/providers/test_bedrock_role_configuration.py::test_role_reference_from_db_and_yaml_reaches_real_sts_http_and_bedrock": [
"other.provider_wire.bedrock.db_yaml_role_reference_reaches_sts_and_signed_request"
],
"tests/integration/routing/test_redis_recovery.py::test_owned_redis_outage_recovers_requests_and_real_response_cache": [
"other.routing.redis.owned_outage_recovers_serving_and_response_cache"
],
"tests/integration/providers/test_anthropic_wire.py::test_anthropic_tool_history_and_cache_tokens_keep_wire_and_accounting_contracts": [
"other.provider_wire.anthropic.tool_history_system_cache_and_internal_fields",
"quota_management.spend_tracking.cache_tokens.disjoint_classes_use_explicit_rates"
]
}
}

View file

@ -0,0 +1,61 @@
import json
import uuid
from typing import Final
import pytest
from integration._support.client import Gateway, eventually, object_value
from integration._support.database import read_rows
from integration._support.wire import Reply, Request, wire_server
@pytest.mark.covers("other.provider_wire.anthropic.tool_history_system_cache_and_internal_fields", "quota_management.spend_tracking.cache_tokens.disjoint_classes_use_explicit_rates")
def test_anthropic_tool_history_and_cache_tokens_keep_wire_and_accounting_contracts(gateway: Gateway) -> None:
identity: Final = "anthropic-wire-" + uuid.uuid4().hex
tool_schema: Final = {"type": "object", "properties": {"x": {"type": "integer"}, "y": {"type": "integer"}}, "required": ["x", "y"]}
def respond(request: Request) -> Reply:
assert request.method == "POST" and request.target == "/v1/messages"
assert request.headers["x-api-key"] == "synthetic-anthropic-key"
body: Final = json.loads(request.body)
assert body["model"] == "claude-sonnet-4-5-20250929"
assert body["system"] == [{"type": "text", "text": "synthetic policy", "cache_control": {"type": "ephemeral"}}]
assert body["tools"][0]["name"] == "add" and body["tools"][0]["input_schema"] == tool_schema
assert body["max_tokens"] == 16
assert not {"timeout", "stream_chunk_size", "litellm_params", "litellm_metadata", "rpm", "tpm"}.intersection(body)
messages: Final = body["messages"]
assert [message["role"] for message in messages] == ["user", "assistant", "user"]
assert messages[0]["content"] == [{"type": "text", "text": "first"}]
assert messages[1]["content"] == [{"type": "tool_use", "id": "history-call", "name": "add", "input": {"x": 1, "y": 2}}]
assert messages[2]["content"] == [{"type": "tool_result", "tool_use_id": "history-call", "content": "3"}, {"type": "text", "text": "next"}]
return Reply(body=json.dumps({"id": identity, "type": "message", "role": "assistant", "model": "claude-sonnet-4-5-20250929", "content": [{"type": "tool_use", "id": "next-call", "name": "add", "input": {"x": 3, "y": 4}}], "stop_reason": "tool_use", "stop_sequence": None, "usage": {"input_tokens": 10, "output_tokens": 4, "cache_read_input_tokens": 5, "cache_creation_input_tokens": 7}}).encode())
with wire_server(respond) as wire, gateway.scenario() as scenario:
model: Final = scenario.model(model="anthropic/claude-sonnet-4-5-20250929", api_base=wire.url, api_key="synthetic-anthropic-key", input_cost_per_token=0.001, output_cost_per_token=0.002, cache_read_input_token_cost=0.0001, cache_creation_input_token_cost=0.002)
response: Final = gateway.request("POST", "/v1/chat/completions", {
"model": model, "max_tokens": 16, "timeout": 5,
"messages": [
{"role": "system", "content": [{"type": "text", "text": "synthetic policy", "cache_control": {"type": "ephemeral"}}]},
{"role": "user", "content": "first"},
{"role": "assistant", "tool_calls": [{"id": "history-call", "type": "function", "function": {"name": "add", "arguments": '{"x":1,"y":2}'}}]},
{"role": "tool", "tool_call_id": "history-call", "content": "3"},
{"role": "user", "content": "next"},
],
"tools": [{"type": "function", "function": {"name": "add", "parameters": tool_schema}}],
})
assert response.status_code == 200, response.text
body: Final = response.json()
assert body["id"].startswith("chatcmpl-")
assert body["choices"][0]["finish_reason"] == "tool_calls"
tool: Final = body["choices"][0]["message"]["tool_calls"][0]
assert tool["id"] == "next-call" and tool["function"]["name"] == "add"
assert json.loads(tool["function"]["arguments"]) == {"x": 3, "y": 4}
assert body["usage"]["prompt_tokens"] == 22 and body["usage"]["completion_tokens"] == 4
assert len(wire.drain()) == 1
rows: Final = eventually(lambda: read_rows('SELECT spend, prompt_tokens, completion_tokens, metadata FROM "LiteLLM_SpendLogs" WHERE request_id=%s', (body["id"],)), lambda values: len(values) == 1, seconds=70)
assert float(rows[0]["spend"]) == pytest.approx(10 * 0.001 + 5 * 0.0001 + 7 * 0.002 + 4 * 0.002)
assert rows[0]["prompt_tokens"] == 22 and rows[0]["completion_tokens"] == 4
metadata: Final = rows[0]["metadata"]
parsed: Final = json.loads(metadata) if isinstance(metadata, str) else object_value(metadata)
assert parsed["cost_breakdown"]["input_cost"] == pytest.approx(0.0245)
assert parsed["cost_breakdown"]["output_cost"] == pytest.approx(0.008)

View file

@ -0,0 +1,99 @@
import asyncio
import json
import os
import uuid
from pathlib import Path
from typing import Final
import pytest
import yaml
from integration._support.client import Gateway
from integration._support.process import owned_proxy
from integration._support.wire import Reply, Request, wire_server
MODEL: Final = "bedrock/converse/anthropic.claude-3-haiku-20240307-v1:0"
TOKEN: Final = "synthetic-bedrock-bearer"
RESPONSE: Final = json.dumps({
"output": {"message": {"role": "assistant", "content": [{"text": "bedrock wire control"}]}},
"stopReason": "end_turn", "usage": {"inputTokens": 11, "outputTokens": 4, "totalTokens": 15},
"metrics": {"latencyMs": 1},
}).encode()
def bearer_peer(request: Request) -> Reply:
assert request.method == "POST" and request.target == "/model/anthropic.claude-3-haiku-20240307-v1%3A0/converse"
assert request.headers["authorization"] == f"Bearer {TOKEN}"
assert "x-amz-security-token" not in request.headers
body: Final = json.loads(request.body)
assert body["messages"] == [{"role": "user", "content": [{"text": "synthetic bearer request"}]}]
assert body["system"] == [{"text": "synthetic system"}]
assert body["inferenceConfig"]["maxTokens"] == 16
assert not {"timeout", "stream_chunk_size", "litellm_params", "litellm_metadata", "api_key"}.intersection(body)
return Reply(body=RESPONSE)
@pytest.mark.covers("other.provider_wire.bedrock.bearer_sdk_skips_credential_chain")
async def test_bearer_only_sdk_sync_async_requests_do_not_require_aws_credentials(monkeypatch: pytest.MonkeyPatch, tmp_path: Path) -> None:
import litellm
empty: Final = tmp_path / "empty-aws-config"
empty.write_text("")
for name in tuple(name for name in os.environ if name.startswith("AWS_")):
monkeypatch.delenv(name, raising=False)
for name, value in {"AWS_CONFIG_FILE": str(empty), "AWS_SHARED_CREDENTIALS_FILE": str(empty), "AWS_EC2_METADATA_DISABLED": "true", "LITELLM_RUST": "false"}.items():
monkeypatch.setenv(name, value)
with wire_server(bearer_peer) as wire:
with pytest.raises(litellm.APIConnectionError, match=r"config profile .* could not be found"):
await asyncio.to_thread(litellm.completion, model=MODEL, aws_profile_name="integration-profile-must-not-be-read", aws_region_name="us-east-1", aws_bedrock_runtime_endpoint=wire.url, messages=[{"role": "user", "content": "synthetic credential control"}], timeout=5, num_retries=0)
assert wire.drain() == ()
for source in ("argument", "environment"):
if source == "environment":
monkeypatch.setenv("AWS_BEARER_TOKEN_BEDROCK", TOKEN)
parameters: Final = {
"model": MODEL, "api_key": TOKEN if source == "argument" else None,
"aws_region_name": "us-east-1", "aws_profile_name": "integration-profile-must-not-be-read",
"aws_bedrock_runtime_endpoint": wire.url, "timeout": 5, "num_retries": 0,
"messages": [{"role": "system", "content": "synthetic system"}, {"role": "user", "content": "synthetic bearer request"}],
"max_tokens": 16,
}
for asynchronous in (False, True):
result: Final = await litellm.acompletion(**parameters) if asynchronous else await asyncio.to_thread(litellm.completion, **parameters)
assert result.choices[0].message.content == "bedrock wire control"
assert result.choices[0].finish_reason == "stop"
assert result.usage.prompt_tokens == 11 and result.usage.completion_tokens == 4
assert len(wire.drain()) == 1
@pytest.mark.covers("other.provider_wire.bedrock.bearer_db_yaml_survives_reload")
def test_bearer_environment_reference_loads_from_db_and_yaml_and_survives_reload(gateway: Gateway, tmp_path: Path) -> None:
empty: Final = tmp_path / "empty-aws-config"
empty.write_text("")
with wire_server(bearer_peer) as wire:
parameters: Final = {
"model": MODEL, "api_key": "os.environ/INTEGRATION_BEARER_TOKEN", "aws_region_name": "us-east-1",
"aws_profile_name": "integration-profile-must-not-be-read", "aws_bedrock_runtime_endpoint": wire.url,
}
alias: Final = f"integration-yaml-{uuid.uuid4().hex}"
configuration: Final = yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text())
configuration["model_list"] = [{"model_name": alias, "litellm_params": parameters, "model_info": {"id": alias}}]
path: Final = tmp_path / "bedrock.yaml"
path.write_text(yaml.safe_dump(configuration))
overrides: Final = {"INTEGRATION_BEARER_TOKEN": TOKEN, "AWS_CONFIG_FILE": str(empty), "AWS_SHARED_CREDENTIALS_FILE": str(empty), "AWS_EC2_METADATA_DISABLED": "true", "LITELLM_RUST": "false"}
with owned_proxy(gateway, tmp_path, overrides, config=path, remove_environment=tuple(name for name in os.environ if name.startswith("AWS_"))) as candidate, candidate.scenario() as scenario:
database_model: Final = scenario.model(**parameters)
for generation in range(2):
for model in (alias, database_model):
response: Final = candidate.request("POST", "/v1/chat/completions", {
"model": model, "messages": [{"role": "system", "content": "synthetic system"}, {"role": "user", "content": "synthetic bearer request"}],
"max_tokens": 16, "cache": {"no-cache": True},
})
assert response.status_code == 200, response.text
assert response.json()["choices"][0]["message"]["content"] == "bedrock wire control"
assert response.json()["usage"]["total_tokens"] == 15
assert len(wire.drain()) == 1, f"Expected actual provider call after reload {generation}"
if generation == 0:
entries: Final = candidate.get("/model/info")["data"]
target: Final = next(entry for entry in entries if entry["model_name"] == database_model)
response: Final = candidate.request("PATCH", f"/model/{target['model_info']['id']}/update", {"model_info": {"description": "bearer reload"}})
assert response.status_code == 200, response.text

View file

@ -0,0 +1,75 @@
import json
import os
import uuid
from pathlib import Path
from typing import Final
from urllib.parse import parse_qs
import pytest
import yaml
from integration._support.client import Gateway
from integration._support.process import owned_proxy
from integration._support.wire import Reply, Request, wire_server
from integration.providers.test_bedrock_auth_wire import MODEL, RESPONSE
@pytest.mark.covers("other.provider_wire.bedrock.db_yaml_role_reference_reaches_sts_and_signed_request")
def test_role_reference_from_db_and_yaml_reaches_real_sts_http_and_bedrock(gateway: Gateway, tmp_path: Path) -> None:
role: Final = "arn:aws:iam::123456789012:role/integration-" + uuid.uuid4().hex
assumed_key: Final = "ASIAINTEGRATION000001"
assumed_token: Final = "synthetic-assumed-session-token"
def sts(request: Request) -> Reply:
parameters: Final = parse_qs(request.body.decode())
action: Final = parameters["Action"][0]
assert request.method == "POST" and action in {"GetCallerIdentity", "AssumeRole"}
if action == "GetCallerIdentity":
result = "<GetCallerIdentityResult><Arn>arn:aws:iam::123456789012:user/integration-source</Arn><UserId>integration-source</UserId><Account>123456789012</Account></GetCallerIdentityResult>"
else:
assert parameters["RoleArn"] == [role]
assert parameters["RoleSessionName"][0] in {"integration-yaml-session", "integration-db-session"}
result = f"<AssumeRoleResult><Credentials><AccessKeyId>{assumed_key}</AccessKeyId><SecretAccessKey>synthetic-assumed-secret-key-for-testing</SecretAccessKey><SessionToken>{assumed_token}</SessionToken><Expiration>2035-01-01T00:00:00Z</Expiration></Credentials><AssumedRoleUser><Arn>arn:aws:sts::123456789012:assumed-role/integration/session</Arn><AssumedRoleId>integration:session</AssumedRoleId></AssumedRoleUser><PackedPolicySize>0</PackedPolicySize></AssumeRoleResult>"
return Reply(content_type="text/xml", body=f'<{action}Response xmlns="https://sts.amazonaws.com/doc/2011-06-15/">{result}<ResponseMetadata><RequestId>synthetic-sts-request</RequestId></ResponseMetadata></{action}Response>'.encode())
def bedrock(request: Request) -> Reply:
assert request.method == "POST" and request.target == "/model/anthropic.claude-3-haiku-20240307-v1%3A0/converse"
assert f"Credential={assumed_key}/" in request.headers["authorization"]
assert request.headers["x-amz-security-token"] == assumed_token
assert json.loads(request.body)["messages"][0]["content"][0]["text"] == "synthetic role request"
return Reply(body=RESPONSE)
with wire_server(sts) as authority, wire_server(bedrock) as provider:
parameters: Final = {
"model": MODEL, "aws_region_name": "us-east-1", "aws_role_name": "os.environ/INTEGRATION_ROLE_ARN",
"aws_session_name": "integration-yaml-session", "aws_bedrock_runtime_endpoint": provider.url,
"aws_sts_endpoint": authority.url,
}
alias: Final = "integration-role-yaml-" + uuid.uuid4().hex
configuration: Final = yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text())
configuration["model_list"] = [{"model_name": alias, "litellm_params": parameters, "model_info": {"id": alias}}]
path: Final = tmp_path / "roles.yaml"
path.write_text(yaml.safe_dump(configuration))
empty: Final = tmp_path / "empty-aws-config"
empty.write_text("")
overrides: Final = {
"INTEGRATION_ROLE_ARN": role, "AWS_ACCESS_KEY_ID": "AKIAINTEGRATION000001", "AWS_SECRET_ACCESS_KEY": "synthetic-source-secret-key-for-testing",
"AWS_CONFIG_FILE": str(empty), "AWS_SHARED_CREDENTIALS_FILE": str(empty), "AWS_EC2_METADATA_DISABLED": "true",
"AWS_ENDPOINT_URL_STS": authority.url, "AWS_DEFAULT_REGION": "us-east-1", "LITELLM_RUST": "false",
}
with owned_proxy(gateway, tmp_path, overrides, config=path, remove_environment=tuple(name for name in os.environ if name.startswith("AWS_"))) as candidate, candidate.scenario() as scenario:
database_model: Final = scenario.model(**{**parameters, "api_key": None, "aws_session_name": "integration-db-session"})
for generation in range(2):
for model in (alias, database_model):
response: Final = candidate.request("POST", "/v1/chat/completions", {"model": model, "messages": [{"role": "user", "content": "synthetic role request"}], "cache": {"no-cache": True}})
assert response.status_code == 200, response.text
assert response.json()["choices"][0]["message"]["content"] == "bedrock wire control"
assert response.json()["usage"]["total_tokens"] == 15
assert len(provider.drain()) == 1
if generation == 0:
target: Final = next(entry for entry in candidate.get("/model/info")["data"] if entry["model_name"] == database_model)
response: Final = candidate.request("PATCH", f"/model/{target['model_info']['id']}/update", {"model_info": {"description": "role reload"}})
assert response.status_code == 200, response.text
assumed: Final = tuple(parse_qs(request.body.decode()) for request in authority.drain() if parse_qs(request.body.decode())["Action"] == ["AssumeRole"])
assert {entry["RoleSessionName"][0] for entry in assumed} == {"integration-yaml-session", "integration-db-session"}
assert all(entry["RoleArn"] == [role] for entry in assumed)

View file

@ -0,0 +1,111 @@
import asyncio
import base64
import hashlib
import hmac
import json
from datetime import datetime
from typing import Final
import httpx
import pytest
from integration._support.sigv4 import encoded_path, signature
from integration._support.wire import Reply, Request, wire_server
ACCESS: Final = "AKIAIOSFODNN7EXAMPLE"
SECRET: Final = "wJalrXUtnFEMI/K7MDENG/bPxRfiCYEXAMPLEKEY"
@pytest.mark.covers("other.provider_wire.s3.verifier_known_answer_and_negative_controls")
def test_sigv4_verifier_matches_published_put_and_rejects_corruption() -> None:
# Public AWS example credentials and PUT vector, not an active account:
# https://docs.aws.amazon.com/AmazonS3/latest/developerguide/sig-v4-header-based-auth.html
headers: Final = {
"date": "Fri, 24 May 2013 00:00:00 GMT", "host": "examplebucket.s3.amazonaws.com",
"x-amz-content-sha256": "44ce7dd67c959e0d3524ffac1771dfbba87d2b6b4b4e99e42034a8b803f8b072",
"x-amz-date": "20130524T000000Z", "x-amz-storage-class": "REDUCED_REDUNDANCY",
}
signed: Final = "date;host;x-amz-content-sha256;x-amz-date;x-amz-storage-class"
expected: Final = (
"9e0e90d9c76de8fa5b200d8c849cd5b8dc7a3be3951ddb7f6a76b4158342019d",
"98ad721746da40c64f1a55b78f14c238d841ea1380cd77a1b5971af0ece108bd",
)
actual: Final = signature("PUT", "/test%24file.text", headers, signed, b"Welcome to Amazon S3.", SECRET, "20130524/us-east-1/s3/aws4_request")
assert actual == expected
assert signature("PUT", "/test$file.text", headers, signed, b"Welcome to Amazon S3.", SECRET, "20130524/us-east-1/s3/aws4_request") != expected
assert encoded_path("/bucket/a=b+c/d e/雪.json") == "/bucket/a%3Db%2Bc/d%20e/%E9%9B%AA.json"
@pytest.mark.covers("other.provider_wire.s3.sync_async_reserved_keys_are_signed_and_accepted")
async def test_s3_sync_and_async_uploads_pass_independent_wire_verification(monkeypatch: pytest.MonkeyPatch) -> None:
from litellm.integrations.s3_v2 import S3Logger
from litellm.types.integrations.s3_v2 import s3BatchLoggingElement
monkeypatch.setattr("botocore.auth.get_current_datetime", lambda: datetime(2026, 9, 14))
payload: Final = {"id": "synthetic-event", "content": "synthetic snow 雪"}
expected_path = ""
def verify(request: Request) -> Reply:
if request.method != "PUT" or request.target != expected_path:
return Reply(status=403)
try:
authorization: Final = request.headers.get("authorization", "")
assert authorization.startswith("AWS4-HMAC-SHA256 ")
fields: Final = dict(part.split("=", 1) for part in authorization.removeprefix("AWS4-HMAC-SHA256 ").split(", "))
access, scope = fields["Credential"].split("/", 1)
assert access == ACCESS and scope == "20260914/us-east-1/s3/aws4_request"
assert request.headers["x-amz-date"] == "20260914T000000Z"
signed: Final = fields["SignedHeaders"].split(";")
assert signed == sorted(set(signed))
assert {"host", "content-md5", "x-amz-date"}.issubset(signed)
assert {name for name in request.headers if name.startswith("x-amz-") and name != "x-amz-content-sha256"}.issubset(signed)
assert request.headers["content-md5"] == base64.b64encode(hashlib.md5(request.body, usedforsecurity=False).digest()).decode()
assert request.headers["x-amz-content-sha256"] == hashlib.sha256(request.body).hexdigest()
expected: Final = signature("PUT", request.target, request.headers, fields["SignedHeaders"], request.body, SECRET, scope)[1]
return Reply(status=200 if hmac.compare_digest(expected, fields["Signature"]) else 403)
except (AssertionError, KeyError, ValueError):
return Reply(status=403)
with wire_server(verify) as wire:
prior: Final = asyncio.all_tasks()
logger: Final = S3Logger(s3_bucket_name="integration-bucket", s3_region_name="us-east-1", s3_endpoint_url=wire.url,
s3_aws_access_key_id=ACCESS, s3_aws_secret_access_key=SECRET, s3_callback_params_override={})
owned: Final = asyncio.all_tasks() - prior
assert len(owned) == 1
try:
for mode in ("sync", "async"):
for key in ("plain.json", "a=b+c/d e/雪.json", "percent%2Fplus+.json"):
expected_path = encoded_path(f"/integration-bucket/{key}")
element: Final = s3BatchLoggingElement(payload=payload, s3_object_key=key, s3_object_download_filename="event.json")
if mode == "sync":
await asyncio.to_thread(logger.upload_data_to_s3, element)
else:
await logger.async_upload_data_to_s3(element)
requests: Final = wire.drain()
assert len(requests) == 1, "Upload must be accepted on its first actual PUT"
request: Final = requests[0]
assert request.target == expected_path
assert json.loads(request.body) == payload
assert verify(request).status == 200
with httpx.Client(timeout=5, trust_env=False) as client:
corrupt: Final = {**request.headers, "authorization": request.headers["authorization"][:-1] + ("0" if request.headers["authorization"][-1] != "0" else "1")}
assert client.put(wire.url + expected_path, content=request.body, headers=corrupt).status_code == 403
assert client.put(wire.url + expected_path + "-wrong", content=request.body, headers=request.headers).status_code == 403
assert client.put(wire.url + expected_path, content=request.body + b" ", headers={name: value for name, value in request.headers.items() if name != "content-length"}).status_code == 403
fields: Final = dict(part.split("=", 1) for part in request.headers["authorization"].removeprefix("AWS4-HMAC-SHA256 ").split(", "))
for signed, scope, md5 in (
(fields["SignedHeaders"].replace("host;", ""), "20260914/us-east-1/s3/aws4_request", request.headers["content-md5"]),
(fields["SignedHeaders"], "20260914/us-west-2/s3/aws4_request", request.headers["content-md5"]),
(fields["SignedHeaders"], "20260914/us-east-1/s3/aws4_request", "AAAAAAAAAAAAAAAAAAAAAA=="),
):
candidate_headers: Final = {**request.headers, "content-md5": md5}
digest: Final = signature("PUT", request.target, candidate_headers, signed, request.body, SECRET, scope)[1]
candidate_headers["authorization"] = f"AWS4-HMAC-SHA256 Credential={ACCESS}/{scope}, SignedHeaders={signed}, Signature={digest}"
assert client.put(wire.url + expected_path, content=request.body, headers=candidate_headers).status_code == 403
assert len(wire.drain()) == 6
finally:
for task in owned:
task.cancel()
await asyncio.gather(*owned, return_exceptions=True)
assert all(task.done() for task in owned)

View file

@ -0,0 +1,98 @@
import json
import uuid
from pathlib import Path
from typing import Final
import httpx
import pytest
import yaml
from integration._support.client import Gateway, object_value
from integration._support.wire import Reply, Request, wire_server
@pytest.mark.covers("other.routing.retries.several_attempts_reach_success_without_hidden_retries", "other.routing.errors.nonretryable_and_exhausted_failures_remain_errors")
def test_retry_counts_and_public_errors_match_actual_provider_attempts(gateway: Gateway) -> None:
with httpx.Client(base_url=gateway.upstream_url, timeout=5, trust_env=False) as upstream, gateway.scenario() as scenario:
original: Final = object_value(gateway.get("/router/settings")["current_values"])["num_retries"]
provider_model: Final = "errors-" + uuid.uuid4().hex
model: Final = scenario.model(model=f"openai/{provider_model}", input_cost_per_token=0, output_cost_per_token=0)
def remove() -> None:
response: Final = upstream.delete(f"/__scripts/{provider_model}")
assert response.status_code in (200, 404)
assert upstream.get(f"/__scripts/{provider_model}").status_code == 404
scenario.cleanups.callback(remove)
try:
for index, (retries, statuses, status, attempts) in enumerate(((2, [500, 500, 200], 200, 3), (2, [400, 200], 400, 1), (1, [429, 429, 200], 429, 2), (1, [500, 500, 200], 500, 2))):
gateway.post("/config/update", {"router_settings": {"num_retries": retries}})
assert object_value(gateway.get("/router/settings")["current_values"])["num_retries"] == retries
upstream.post(f"/__scripts/{provider_model}", json={"statuses": statuses}).raise_for_status()
upstream.get("/__observations").raise_for_status()
response: Final = gateway.request("POST", "/v1/chat/completions", {"model": model, "messages": [{"role": "user", "content": f"{provider_model} {index}"}]})
assert response.status_code == status, response.text
requests: Final = upstream.get("/__observations").json()["requests"]
assert len(requests) == attempts
assert all(request["body"]["model"] == provider_model for request in requests)
assert upstream.get(f"/__scripts/{provider_model}").json()["remaining"] == statuses[attempts:]
if status == 200:
assert response.json()["usage"]["total_tokens"] == 40
else:
error: Final = response.json()["error"]
assert isinstance(error["message"], str) and "Controlled provider failure" in error["message"]
assert str(error["code"]) == str(status)
assert error["type"] == {400: "invalid_request_error", 429: "throttling_error", 500: "internal_server_error"}[status]
assert error["param"] is None
assert "Traceback" not in response.text and "File \"" not in response.text
finally:
gateway.post("/config/update", {"router_settings": {"num_retries": original}})
assert object_value(gateway.get("/router/settings")["current_values"])["num_retries"] == original
@pytest.mark.covers("other.routing.fallback.loaded_configuration_selects_only_permitted_target")
def test_loaded_fallback_selects_expected_deployment_and_keeps_response_identity(tmp_path: Path) -> None:
from litellm import Router
def respond(request: Request) -> Reply:
model: Final = json.loads(request.body)["model"]
assert model in {"primary-wire", "fallback-wire", "unrelated-wire"}
if model == "primary-wire":
return Reply(status=500, body=b'{"error":{"message":"synthetic primary unavailable","type":"api_error","code":"500"}}')
return Reply(body=json.dumps({"id": "response-" + model, "object": "chat.completion", "created": 1, "model": model, "choices": [{"index": 0, "message": {"role": "assistant", "content": "served " + model}, "finish_reason": "stop"}], "usage": {"prompt_tokens": 11, "completion_tokens": 4, "total_tokens": 15}}).encode())
with wire_server(respond) as wire:
path: Final = tmp_path / "fallback.yaml"
path.write_text(yaml.safe_dump({"model_list": [{"model_name": alias, "litellm_params": {"model": "openai/" + upstream, "api_key": "synthetic-routing-key", "api_base": wire.url + "/v1"}} for alias, upstream in (("primary", "primary-wire"), ("fallback", "fallback-wire"), ("unrelated", "unrelated-wire"))], "router_settings": {"num_retries": 0, "disable_cooldowns": True, "fallbacks": [{"primary": ["fallback"]}]}}))
loaded: Final = yaml.safe_load(path.read_text())
router: Final = Router(model_list=loaded["model_list"], **loaded["router_settings"])
try:
result: Final = router.completion(model="primary", messages=[{"role": "user", "content": "fallback control"}])
assert result.id == "response-fallback-wire"
assert result.choices[0].message.content == "served fallback-wire"
assert result.choices[0].finish_reason == "stop"
assert result.usage.prompt_tokens == 11 and result.usage.completion_tokens == 4
assert tuple(json.loads(request.body)["model"] for request in wire.drain()) == ("primary-wire", "fallback-wire")
control: Final = router.completion(model="unrelated", messages=[{"role": "user", "content": "independent route"}])
assert control.id == "response-unrelated-wire"
assert tuple(json.loads(request.body)["model"] for request in wire.drain()) == ("unrelated-wire",)
finally:
router.reset()
@pytest.mark.covers("other.routing.alias_update.persisted_target_changes_only_selected_route")
def test_saved_deployment_target_update_changes_wire_and_preserves_control(gateway: Gateway) -> None:
with httpx.Client(base_url=gateway.upstream_url, timeout=5, trust_env=False) as upstream, gateway.scenario() as scenario:
prefix: Final = "target-" + uuid.uuid4().hex
model: Final = scenario.model(model="openai/" + prefix + "-first", input_cost_per_token=0, output_cost_per_token=0)
other: Final = scenario.model(model="openai/" + prefix + "-control", input_cost_per_token=0, output_cost_per_token=0)
target: Final = next(entry for entry in gateway.get("/model/info")["data"] if entry["model_name"] == model)
for generation, suffix in enumerate(("first", "second")):
if generation:
response: Final = gateway.request("PATCH", f"/model/{target['model_info']['id']}/update", {"litellm_params": {"model": "openai/" + prefix + "-second"}})
assert response.status_code == 200, response.text
upstream.get("/__observations").raise_for_status()
for alias in (model, other):
assert gateway.chat(alias, text=f"{prefix} generation {generation}")["usage"]["total_tokens"] == 40
requests: Final = upstream.get("/__observations").json()["requests"]
assert [request["body"]["model"] for request in requests] == [prefix + "-" + suffix, prefix + "-control"]

View file

@ -0,0 +1,59 @@
import os
import uuid
from pathlib import Path
from typing import Final
from urllib.parse import urlsplit, urlunsplit
import httpx
import psycopg
import pytest
from psycopg import sql
from redis import Redis
from integration._support.client import Gateway, eventually
from integration._support.process import owned_proxy
from integration._support.redis_process import owned_redis
@pytest.mark.covers("other.routing.redis.owned_outage_recovers_serving_and_response_cache")
def test_owned_redis_outage_recovers_requests_and_real_response_cache(gateway: Gateway, tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None:
original: Final = os.environ["DATABASE_URL"]
identity: Final = "integration_recovery_" + uuid.uuid4().hex
parsed: Final = urlsplit(original)
database_url: Final = urlunsplit((parsed.scheme, parsed.netloc, "/" + identity, "", ""))
with psycopg.connect(original, autocommit=True) as admin:
admin.execute(sql.SQL("CREATE DATABASE {}").format(sql.Identifier(identity)))
try:
with owned_redis(tmp_path) as cache, monkeypatch.context() as environment:
environment.setenv("DATABASE_URL", database_url)
with owned_proxy(gateway, tmp_path, {"DATABASE_URL": database_url, "REDIS_HOST": cache.host, "REDIS_PORT": str(cache.port), "REDIS_CIRCUIT_BREAKER_RECOVERY_TIMEOUT": "1"}) as candidate, candidate.scenario() as scenario, httpx.Client(base_url=gateway.upstream_url, timeout=5, trust_env=False) as upstream:
model: Final = scenario.model()
key: Final = scenario.key(models=[model])
for generation in ("before", "after"):
with Redis(host=cache.host, port=cache.port, socket_timeout=1) as client:
eventually(client.ping, bool)
eventually(lambda: client.pubsub_numsub("litellm_proxy.auth_cache_invalidation")[0][1], lambda count: count >= 1, seconds=8)
upstream.get("/__observations").raise_for_status()
first: Final = candidate.chat(model, key=key, text=identity + generation)
second: Final = candidate.chat(model, key=key, text=identity + generation)
assert first["id"] == second["id"]
assert first["choices"] == second["choices"] and first["usage"]["total_tokens"] == 40
assert len(upstream.get("/__observations").json()["requests"]) == 1
with Redis(host=cache.host, port=cache.port, socket_timeout=1) as client:
eventually(
lambda first=first: tuple(client.get(name) for name in client.scan_iter() if client.type(name) == b"string"),
lambda values, first=first: any(str(first["id"]).encode() in value for value in values if value is not None),
seconds=10,
)
if generation == "before":
cache.stop()
upstream.get("/__observations").raise_for_status()
during: Final = candidate.chat(model, key=key, text=identity + "during")
assert during["usage"]["total_tokens"] == 40
assert len(upstream.get("/__observations").json()["requests"]) == 1
cache.start()
with psycopg.connect(database_url) as fresh:
assert fresh.execute('SELECT count(*) FROM "LiteLLM_VerificationToken"').fetchone()[0] >= 1
finally:
admin.execute(sql.SQL("DROP DATABASE {}").format(sql.Identifier(identity)))
assert admin.execute("SELECT datname FROM pg_database WHERE datname=%s", (identity,)).fetchall() == []

View file

@ -0,0 +1,149 @@
import asyncio
import json
import threading
import uuid
from typing import Final
import pytest
from hypothesis import Phase, example, given, settings, strategies as st
from openai import OpenAI
from integration._support.client import Gateway, eventually
from integration._support.database import read_rows
from integration._support.wire import Reply, wire_server
def frame(identity: str, delta: dict, *, finish: str | None = None) -> bytes:
value: Final = {"id": identity, "object": "chat.completion.chunk", "created": 1, "model": "gpt-4o-mini", "choices": [{"index": 0, "delta": delta, "finish_reason": finish}]}
return b"data: " + json.dumps(value, ensure_ascii=False).encode() + b"\n\n"
def text_stream(identity: str) -> tuple[bytes, ...]:
usage: Final = {"id": identity, "object": "chat.completion.chunk", "created": 1, "model": "gpt-4o-mini", "choices": [], "usage": {"prompt_tokens": 11, "completion_tokens": 4, "total_tokens": 15}}
return (frame(identity, {"role": "assistant", "content": "Hello "}), frame(identity, {"content": "雪 café"}), frame(identity, {}, finish="stop"), b"data: " + json.dumps(usage).encode() + b"\n\n", b"data: [DONE]\n\n")
@pytest.mark.covers("other.streaming.byte_partitions.preserve_text_identity_and_usage")
def test_generated_tcp_partitions_preserve_unicode_text_identity_and_final_usage() -> None:
import litellm
body: Final = b"".join(text_stream("stream-partition-control"))
@settings(max_examples=20, deadline=None, database=None, phases=(Phase.explicit, Phase.generate, Phase.shrink))
@example(cuts=tuple(range(1, len(body))))
@example(cuts=())
@given(cuts=st.lists(st.integers(min_value=1, max_value=len(body) - 1), max_size=35, unique=True).map(tuple))
def check(cuts: tuple[int, ...]) -> None:
boundaries: Final = (0, *sorted(cuts), len(body))
pieces: Final = tuple(body[left:right] for left, right in zip(boundaries, boundaries[1:]))
with wire_server(lambda request: Reply(content_type="text/event-stream", chunks=pieces)) as wire:
stream: Final = litellm.completion(model="openai/gpt-4o-mini", api_base=wire.url + "/v1", api_key="synthetic-stream-key", messages=[{"role": "user", "content": "partition control"}], stream=True, stream_options={"include_usage": True}, timeout=5, num_retries=0)
try:
chunks: Final = tuple(stream)
finally:
asyncio.run(stream.aclose())
assert "".join(choice.delta.content or "" for chunk in chunks for choice in chunk.choices) == "Hello 雪 café"
assert {chunk.id for chunk in chunks} == {"stream-partition-control"}
assert [choice.finish_reason for chunk in chunks for choice in chunk.choices if choice.finish_reason] == ["stop"]
usages: Final = tuple(chunk.usage for chunk in chunks if getattr(chunk, "usage", None) is not None)
assert len(usages) == 1
assert usages[0].prompt_tokens == 11 and usages[0].completion_tokens == 4
assert len(wire.drain()) == 1
check()
@pytest.mark.covers("other.streaming.tools.fragmented_calls_keep_independent_arguments")
def test_fragmented_tool_names_and_arguments_keep_each_call_identity() -> None:
import litellm
identity: Final = "stream-tools-control"
deltas: Final = (
{"role": "assistant", "tool_calls": [{"index": 0, "id": "call-add", "type": "function", "function": {"name": "ad", "arguments": ""}}, {"index": 1, "id": "call-multiply", "type": "function", "function": {"name": "multi", "arguments": ""}}]},
{"tool_calls": [{"index": 1, "function": {"name": "ply", "arguments": '{"x":3,'}}, {"index": 0, "function": {"arguments": '{"x":1,'}}]},
{"tool_calls": [{"index": 0, "function": {"name": "d", "arguments": '"y":2}'}}, {"index": 1, "function": {"arguments": '"y":4}'}}]},
)
frames: Final = (*tuple(frame(identity, delta) for delta in deltas), frame(identity, {}, finish="tool_calls"), b"data: [DONE]\n\n")
with wire_server(lambda request: Reply(content_type="text/event-stream", chunks=frames)) as wire:
stream: Final = litellm.completion(model="openai/gpt-4o-mini", api_base=wire.url + "/v1", api_key="synthetic-stream-key", messages=[{"role": "user", "content": "tool control"}], stream=True, timeout=5, num_retries=0)
try:
chunks: Final = tuple(stream)
finally:
asyncio.run(stream.aclose())
events: Final = tuple((choice.index, tool) for chunk in chunks for choice in chunk.choices for tool in (choice.delta.tool_calls or ()))
for index, name, call_id, arguments in ((0, "add", "call-add", {"x": 1, "y": 2}), (1, "multiply", "call-multiply", {"x": 3, "y": 4})):
selected: Final = tuple(tool for choice, tool in events if (choice, tool.index) == (0, index))
assert "".join(tool.id or "" for tool in selected) == call_id
assert "".join(tool.function.name or "" for tool in selected) == name
assert json.loads("".join(tool.function.arguments or "" for tool in selected)) == arguments
assert {tool.index for _, tool in events} == {0, 1}
assert [choice.finish_reason for chunk in chunks for choice in chunk.choices if choice.finish_reason] == ["tool_calls"]
assert len(wire.drain()) == 1
@pytest.mark.covers("other.streaming.usage.client_visibility_preserves_persisted_accounting")
def test_proxy_stream_usage_visibility_keeps_exact_persisted_charge(gateway: Gateway) -> None:
with gateway.scenario() as scenario:
for include in (None, False, True):
identity: Final = "stream-usage-" + uuid.uuid4().hex
with wire_server(lambda request, identity=identity: Reply(content_type="text/event-stream", chunks=text_stream(identity))) as wire:
model: Final = scenario.model(api_base=wire.url + "/v1", input_cost_per_token=0.001, output_cost_per_token=0.002)
with OpenAI(api_key=gateway.key, base_url=str(gateway.client.base_url), timeout=5, max_retries=0) as client:
stream: Final = client.chat.completions.create(model=model, messages=[{"role": "user", "content": identity}], stream=True, **({} if include is None else {"stream_options": {"include_usage": include}}))
with stream:
chunks: Final = tuple(stream)
assert "".join(choice.delta.content or "" for chunk in chunks for choice in chunk.choices) == "Hello 雪 café"
assert {chunk.id for chunk in chunks} == {identity}
usages: Final = tuple(chunk.usage for chunk in chunks if chunk.usage is not None)
assert len(usages) == (1 if include else 0)
if include:
assert usages[0].prompt_tokens == 11 and usages[0].completion_tokens == 4
requests: Final = wire.drain()
assert len(requests) == 1
assert json.loads(requests[0].body)["stream_options"]["include_usage"] is True
rows: Final = eventually(lambda identity=identity: read_rows('SELECT spend, prompt_tokens, completion_tokens FROM "LiteLLM_SpendLogs" WHERE request_id=%s', (identity,)), lambda values: len(values) == 1, seconds=70)
assert rows[0]["prompt_tokens"] == 11 and rows[0]["completion_tokens"] == 4
assert float(rows[0]["spend"]) == pytest.approx(0.019)
@pytest.mark.covers("other.streaming.failure.truncated_transport_raises_and_control_recovers")
def test_truncated_http_stream_is_an_error_and_next_stream_succeeds() -> None:
import litellm
for truncated in (True, False):
with wire_server(lambda request, truncated=truncated: Reply(content_type="text/event-stream", chunks=text_stream("stream-truncated"), abort_after=1 if truncated else None)) as wire:
stream: Final = litellm.completion(model="openai/gpt-4o-mini", api_base=wire.url + "/v1", api_key="synthetic-stream-key", messages=[{"role": "user", "content": "truncation control"}], stream=True, timeout=5, num_retries=0)
try:
if truncated:
with pytest.raises(litellm.exceptions.MidStreamFallbackError, match="incomplete chunked read") as failure:
tuple(stream)
assert isinstance(failure.value.original_exception, litellm.APIConnectionError)
assert failure.value.generated_content == "Hello "
assert failure.value.is_pre_first_chunk is False
else:
chunks: Final = tuple(stream)
assert "".join(choice.delta.content or "" for chunk in chunks for choice in chunk.choices) == "Hello 雪 café"
assert any(choice.finish_reason == "stop" for chunk in chunks for choice in chunk.choices)
finally:
asyncio.run(stream.aclose())
assert len(wire.drain()) == 1
@pytest.mark.covers("other.streaming.cancellation.closes_actual_provider_connection")
def test_client_cancellation_releases_the_actual_provider_connection() -> None:
import litellm
gate: Final = threading.Event()
frames: Final = (frame("stream-cancel", {"role": "assistant", "content": "first"}), b":" + b"x" * 4_000_000 + b"\n\n", b"data: [DONE]\n\n")
with wire_server(lambda request: Reply(content_type="text/event-stream", chunks=frames, gate_after_first=gate)) as wire:
stream: Final = litellm.completion(model="openai/gpt-4o-mini", api_base=wire.url + "/v1", api_key="synthetic-stream-key", messages=[{"role": "user", "content": "cancellation control"}], stream=True, timeout=5, num_retries=0)
try:
first: Final = next(stream)
assert first.choices[0].delta.content == "first"
finally:
try:
asyncio.run(stream.aclose())
finally:
gate.set()
assert wire.disconnected.get(timeout=5) == "/v1/chat/completions"
assert len(wire.drain()) == 1