mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-19 00:01:29 +00:00
Merge pull request #41075 from BerriAI/litellm_integration_providers
test: provider wire contracts, streaming and recovery
This commit is contained in:
commit
545df49374
14 changed files with 967 additions and 4 deletions
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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",
|
||||
|
|
|
|||
116
tests/integration/_support/redis_process.py
Normal file
116
tests/integration/_support/redis_process.py
Normal 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()
|
||||
25
tests/integration/_support/sigv4.py
Normal file
25
tests/integration/_support/sigv4.py
Normal 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()
|
||||
|
|
@ -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)
|
||||
|
|
|
|||
113
tests/integration/_support/wire.py
Normal file
113
tests/integration/_support/wire.py
Normal 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}"
|
||||
|
|
@ -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"
|
||||
]
|
||||
}
|
||||
}
|
||||
|
|
|
|||
61
tests/integration/providers/test_anthropic_wire.py
Normal file
61
tests/integration/providers/test_anthropic_wire.py
Normal 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)
|
||||
99
tests/integration/providers/test_bedrock_auth_wire.py
Normal file
99
tests/integration/providers/test_bedrock_auth_wire.py
Normal 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
|
||||
|
|
@ -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)
|
||||
111
tests/integration/providers/test_s3_wire.py
Normal file
111
tests/integration/providers/test_s3_wire.py
Normal 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)
|
||||
98
tests/integration/routing/test_observed_routing.py
Normal file
98
tests/integration/routing/test_observed_routing.py
Normal 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"]
|
||||
59
tests/integration/routing/test_redis_recovery.py
Normal file
59
tests/integration/routing/test_redis_recovery.py
Normal 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() == []
|
||||
149
tests/integration/streaming/test_stream_contracts.py
Normal file
149
tests/integration/streaming/test_stream_contracts.py
Normal 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
|
||||
Loading…
Add table
Reference in a new issue