mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-16 04:29:31 +00:00
test(integration): cover files_batches endpoint contracts (#44703)
* test(integration): cover files_batches endpoint contracts Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(integration): split raw cloud-storage delete test by caller Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(integration): drop mutable polling state in content download test Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(integration): dispatch request spellings without rebinding Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(integration): import assert_never from typing_extensions for python 3.10 Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(integration): tighten files_batches wire and access assertions Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(integration): tidy files_batches access tests Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(integration): cover repeated target models, upload expiry, batch aliases Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(integration): prove files streaming path, SDK model denial and admin GCS delete Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(integration): give the streaming head the wire barrier's full budget Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test: strengthen files and batches wire coverage Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test: drop duplicate parametrize from the model-routed upload helper Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test: pass the SDK base url in the skipped repeated-plain upload test Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(integration): allow batch output metadata lookup Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test: avoid draining file content requests before wait Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --------- Co-authored-by: mateo <mateo@berri.ai> Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
This commit is contained in:
parent
e226016aee
commit
0ca9f8185c
2 changed files with 1733 additions and 0 deletions
882
tests/integration/providers/test_files_batches_wire.py
Normal file
882
tests/integration/providers/test_files_batches_wire.py
Normal file
|
|
@ -0,0 +1,882 @@
|
|||
import base64
|
||||
import json
|
||||
import re
|
||||
import threading
|
||||
import uuid
|
||||
from collections import deque
|
||||
from collections.abc import Callable, Iterator, Mapping
|
||||
from concurrent.futures import ThreadPoolExecutor
|
||||
from email.message import Message
|
||||
from email.parser import BytesParser
|
||||
from email.policy import HTTP
|
||||
from functools import partial
|
||||
from queue import Queue
|
||||
from typing import Final, Literal, cast
|
||||
from urllib.parse import urlsplit
|
||||
|
||||
import httpx
|
||||
import pytest
|
||||
from integration._support.client import Gateway, eventually, object_value, string_value
|
||||
from integration._support.database import read_rows
|
||||
from integration._support.wire import Reply, Request, Wire, wire_server
|
||||
from openai import OpenAI
|
||||
from pydantic import BaseModel, JsonValue, TypeAdapter
|
||||
from typing_extensions import assert_never
|
||||
|
||||
from litellm.proxy.openai_files_endpoints.common_utils import encode_file_id_with_model
|
||||
|
||||
INPUT_FILENAME: Final = "in.jsonl"
|
||||
OUTPUT_BYTES: Final = b'{"custom_id": "r1", "result": {"response": {"body": "ok"}}}\n'
|
||||
STREAM_CHUNK_BYTES: Final = 1024 * 1024
|
||||
_HEAD_OPEN: Final = b'{"custom_id": "r0", "padding": "'
|
||||
_HEAD_CLOSE: Final = b'"}\n'
|
||||
OUTPUT_HEAD: Final = _HEAD_OPEN + b"h" * (STREAM_CHUNK_BYTES - len(_HEAD_OPEN) - len(_HEAD_CLOSE)) + _HEAD_CLOSE
|
||||
MANAGED_PREFIX: Final = "litellm_proxy:"
|
||||
JSON_OBJECT: Final = TypeAdapter(dict[str, JsonValue])
|
||||
MANAGED_FILE_ROW: Final = (
|
||||
'SELECT model_mappings, flat_model_file_ids FROM "LiteLLM_ManagedFileTable" WHERE unified_file_id = %s'
|
||||
)
|
||||
_BatchUploadSpelling = Literal["sdk_list", "repeated_bracket", "indexed", "comma_joined", "repeated_plain"]
|
||||
_ModelRoutedSpelling = Literal[
|
||||
"sdk_extra_body",
|
||||
"sdk_proxy_fields",
|
||||
"multipart_field",
|
||||
"query_param",
|
||||
"model_header",
|
||||
"multipart_proxy_fields",
|
||||
"sdk_expires_after",
|
||||
"multipart_expires_after",
|
||||
]
|
||||
_ContentRoute = Literal["v1_encoded", "files_encoded", "files_raw_query", "v1_raw_header"]
|
||||
_BatchCreateSpelling = Literal["encoded_input", "body_model", "query_model", "header_model", "sanitized_metadata"]
|
||||
|
||||
|
||||
def _jsonl(model: str) -> bytes:
|
||||
line: Final = {
|
||||
"custom_id": "r1",
|
||||
"method": "POST",
|
||||
"url": "/v1/chat/completions",
|
||||
"body": {"model": model, "messages": [{"role": "user", "content": "one"}]},
|
||||
}
|
||||
return json.dumps(line, separators=(",", ":")).encode() + b"\n"
|
||||
|
||||
|
||||
def _rewritten_jsonl(model: str) -> bytes:
|
||||
line: Final = {
|
||||
"custom_id": "r1",
|
||||
"method": "POST",
|
||||
"url": "/v1/chat/completions",
|
||||
"body": {"model": model, "messages": [{"role": "user", "content": "one"}]},
|
||||
}
|
||||
return json.dumps(line).encode()
|
||||
|
||||
|
||||
class _FileObject(BaseModel):
|
||||
id: str
|
||||
object: str
|
||||
bytes: int
|
||||
created_at: int
|
||||
filename: str
|
||||
purpose: str
|
||||
|
||||
|
||||
class _Batch(BaseModel):
|
||||
id: str
|
||||
object: str
|
||||
input_file_id: str | None = None
|
||||
output_file_id: str | None = None
|
||||
error_file_id: str | None = None
|
||||
status: str
|
||||
|
||||
|
||||
def _json(response: httpx.Response) -> dict[str, JsonValue]:
|
||||
return JSON_OBJECT.validate_json(response.content)
|
||||
|
||||
|
||||
def _json_reply(body: Mapping[str, JsonValue], status: int = 200) -> Reply:
|
||||
return Reply(status=status, body=json.dumps(body).encode())
|
||||
|
||||
|
||||
def _file_object(file_id: str, byte_count: int, filename: str = INPUT_FILENAME) -> dict[str, JsonValue]:
|
||||
return {
|
||||
"id": file_id,
|
||||
"object": "file",
|
||||
"bytes": byte_count,
|
||||
"created_at": 1700000000,
|
||||
"filename": filename,
|
||||
"purpose": "batch",
|
||||
"status": "processed",
|
||||
}
|
||||
|
||||
|
||||
def _expected_file_object(file_id: str, byte_count: int, filename: str = INPUT_FILENAME) -> _FileObject:
|
||||
return _FileObject(
|
||||
id=file_id,
|
||||
object="file",
|
||||
bytes=byte_count,
|
||||
created_at=1700000000,
|
||||
filename=filename,
|
||||
purpose="batch",
|
||||
)
|
||||
|
||||
|
||||
def _batch_object(
|
||||
batch_id: str, input_file_id: str, status: str, output_file_id: str | None = None, error_file_id: str | None = None
|
||||
) -> dict[str, JsonValue]:
|
||||
return {
|
||||
"id": batch_id,
|
||||
"object": "batch",
|
||||
"endpoint": "/v1/chat/completions",
|
||||
"input_file_id": input_file_id,
|
||||
"output_file_id": output_file_id,
|
||||
"error_file_id": error_file_id,
|
||||
"status": status,
|
||||
"created_at": 1700000000,
|
||||
"completion_window": "24h",
|
||||
}
|
||||
|
||||
|
||||
def _multipart_parts(request: Request) -> tuple[Message, ...]:
|
||||
envelope: Final = f"content-type: {request.headers['content-type']}\r\n\r\n".encode() + request.body
|
||||
parsed: Final = BytesParser(policy=HTTP).parsebytes(envelope)
|
||||
assert parsed.is_multipart(), request.headers["content-type"]
|
||||
return tuple(parsed.iter_parts())
|
||||
|
||||
|
||||
def _text_fields(parts: tuple[Message, ...]) -> dict[str, str]:
|
||||
return {
|
||||
part.get_param("name", header="content-disposition"): part.get_payload(decode=True).decode()
|
||||
for part in parts
|
||||
if part.get_filename() is None
|
||||
}
|
||||
|
||||
|
||||
def _file_fields(parts: tuple[Message, ...]) -> dict[str, tuple[str, str, bytes]]:
|
||||
return {
|
||||
part.get_param("name", header="content-disposition") or "": (
|
||||
part.get_filename() or "",
|
||||
part.get_content_type(),
|
||||
part.get_payload(decode=True) or b"",
|
||||
)
|
||||
for part in parts
|
||||
if part.get_filename() is not None
|
||||
}
|
||||
|
||||
|
||||
def _drained_other_than_model_list_probes(wire: Wire) -> tuple[Request, ...]:
|
||||
return tuple(
|
||||
request for request in wire.drain() if urlsplit(request.target).path not in {"/v1/models", "/openai/models"}
|
||||
)
|
||||
|
||||
|
||||
def _bearer(request: Request) -> str:
|
||||
return request.headers.get("authorization", request.headers.get("api-key", "")).removeprefix("Bearer ")
|
||||
|
||||
|
||||
def _sdk_base_url(gateway: Gateway) -> str:
|
||||
return str(gateway.client.base_url).rstrip("/") + "/v1"
|
||||
|
||||
|
||||
def _decoded_unified(managed_file_id: str) -> str:
|
||||
decoded: Final = base64.urlsafe_b64decode(managed_file_id + "=" * (-len(managed_file_id) % 4)).decode()
|
||||
assert decoded.startswith(MANAGED_PREFIX), decoded
|
||||
return decoded
|
||||
|
||||
|
||||
def _models_over_a_fresh_connection(gateway: Gateway, _: int) -> frozenset[str]:
|
||||
with httpx.Client(base_url=gateway.client.base_url, timeout=15, trust_env=False) as client:
|
||||
listed: Final = client.get("/v1/models", headers={"Authorization": f"Bearer {gateway.key}"})
|
||||
assert listed.status_code == 200, listed.text
|
||||
data: Final = _json(listed)["data"]
|
||||
assert isinstance(data, list), listed.text
|
||||
return frozenset(string_value(object_value(entry)["id"]) for entry in data)
|
||||
|
||||
|
||||
def _every_worker_serves(gateway: Gateway, *models: str) -> bool:
|
||||
with ThreadPoolExecutor(max_workers=16) as pool:
|
||||
rounds: Final = tuple(
|
||||
tuple(pool.map(partial(_models_over_a_fresh_connection, gateway), range(16))) for _ in range(2)
|
||||
)
|
||||
return all(model in seen for model in models for round_ in rounds for seen in round_)
|
||||
|
||||
|
||||
def _wait_until_every_worker_serves(gateway: Gateway, *models: str) -> None:
|
||||
eventually(lambda: _every_worker_serves(gateway, *models), lambda served: served, seconds=90)
|
||||
|
||||
|
||||
def _deployment_ids(gateway: Gateway, *model_names: str) -> dict[str, str]:
|
||||
entries: Final = gateway.get("/model/info")["data"]
|
||||
assert isinstance(entries, list)
|
||||
by_name: Final = {
|
||||
string_value(entry["model_name"]): string_value(object_value(entry["model_info"])["id"])
|
||||
for entry in entries
|
||||
if isinstance(entry, dict) and entry.get("model_name") in model_names
|
||||
}
|
||||
return by_name
|
||||
|
||||
|
||||
def _multiplexed_openai_backends(
|
||||
created_batch: Mapping[str, JsonValue] | None = None,
|
||||
content_gate: threading.Event | None = None,
|
||||
) -> Callable[[Request], Reply]:
|
||||
def respond(request: Request) -> Reply:
|
||||
raw_path: Final = urlsplit(request.target).path
|
||||
path: Final = "/v1/" + raw_path.removeprefix("/openai/") if raw_path.startswith("/openai/") else raw_path
|
||||
if request.method == "POST" and path == "/v1/files":
|
||||
file_part: Final = _file_fields(_multipart_parts(request))["file"]
|
||||
return _json_reply(_file_object("file-" + _bearer(request), len(file_part[2]), file_part[0]))
|
||||
if request.method == "POST" and path == "/v1/batches":
|
||||
if created_batch is not None:
|
||||
return _json_reply(created_batch)
|
||||
return _json_reply(_batch_object("batch-" + _bearer(request), "file-pending", "in_progress"))
|
||||
if request.method == "POST" and re.fullmatch(r"/v1/batches/[^/]+/cancel", path):
|
||||
return _json_reply(_batch_object("batch-" + _bearer(request), "file-pending", "cancelled"))
|
||||
if request.method == "GET" and path.startswith("/v1/batches/"):
|
||||
return _json_reply(
|
||||
_batch_object(
|
||||
path.rsplit("/", 1)[1], "file-in", "completed", output_file_id="file-out-" + _bearer(request)
|
||||
)
|
||||
)
|
||||
if request.method == "GET" and path.startswith("/v1/files/") and path.endswith("/content"):
|
||||
if content_gate is not None:
|
||||
return Reply(
|
||||
content_type="application/octet-stream",
|
||||
chunks=(OUTPUT_HEAD, OUTPUT_BYTES),
|
||||
gate_after_first=content_gate,
|
||||
)
|
||||
return Reply(body=OUTPUT_BYTES, content_type="application/octet-stream")
|
||||
if request.method == "GET" and path.startswith("/v1/files/"):
|
||||
return _json_reply(_file_object(path.rsplit("/", 1)[1], len(_jsonl("uploaded-model"))))
|
||||
if request.method == "DELETE" and path.startswith("/v1/files/"):
|
||||
return _json_reply({"id": path.rsplit("/", 1)[1], "object": "file", "deleted": True})
|
||||
return _json_reply({"error": {"message": f"unscripted {request.method} {request.target}"}}, 404)
|
||||
|
||||
return respond
|
||||
|
||||
|
||||
def _managed_upload(
|
||||
gateway: Gateway, spelling: _BatchUploadSpelling, fields: Mapping[str, object], sdk_base_url: str
|
||||
) -> _FileObject:
|
||||
key: Final = gateway.key
|
||||
match spelling:
|
||||
case "sdk_list":
|
||||
with OpenAI(base_url=sdk_base_url, api_key=key, max_retries=0) as client:
|
||||
response: Final = client.files.with_raw_response.create(
|
||||
file=(INPUT_FILENAME, _jsonl("uploaded-model"), "application/jsonl"),
|
||||
purpose="batch",
|
||||
extra_body=fields,
|
||||
)
|
||||
assert response.status_code == 200, response.text
|
||||
return _FileObject.model_validate_json(response.content)
|
||||
case "repeated_bracket" | "indexed" | "repeated_plain":
|
||||
posted: Final = gateway.client.post(
|
||||
"/v1/files",
|
||||
data={"purpose": "batch", **fields},
|
||||
files={"file": (INPUT_FILENAME, _jsonl("uploaded-model"), "application/jsonl")},
|
||||
headers={"Authorization": f"Bearer {key}"},
|
||||
)
|
||||
assert posted.status_code == 200, posted.text
|
||||
return _FileObject.model_validate_json(posted.content)
|
||||
case "comma_joined":
|
||||
created: Final = gateway.request_multipart(
|
||||
"/v1/files",
|
||||
{"purpose": "batch", **fields},
|
||||
{"file": (INPUT_FILENAME, _jsonl("uploaded-model"), "application/jsonl")},
|
||||
)
|
||||
assert created.status_code == 200, created.text
|
||||
return _FileObject.model_validate_json(created.content)
|
||||
case _:
|
||||
assert_never(spelling)
|
||||
|
||||
|
||||
def _assert_upload_reaches_both_target_models(
|
||||
gateway: Gateway, spelling: _BatchUploadSpelling, sdk_base_url: str
|
||||
) -> None:
|
||||
key_a: Final = "provider-key-a-" + uuid.uuid4().hex[:8]
|
||||
key_b: Final = "provider-key-b-" + uuid.uuid4().hex[:8]
|
||||
with wire_server(_multiplexed_openai_backends()) as wire, gateway.scenario() as scenario:
|
||||
model_a: Final = scenario.model(model="openai/gpt-4o-mini", api_base=wire.url + "/v1", api_key=key_a)
|
||||
model_b: Final = scenario.model(model="openai/gpt-4.1-mini", api_base=wire.url + "/v1", api_key=key_b)
|
||||
_wait_until_every_worker_serves(gateway, model_a, model_b)
|
||||
fields: Final[Mapping[str, object]] = {
|
||||
"sdk_list": {"target_model_names": [model_a, model_b]},
|
||||
"repeated_bracket": {"target_model_names[]": [model_a, model_b]},
|
||||
"indexed": {"target_model_names[0]": model_a, "target_model_names[1]": model_b},
|
||||
"comma_joined": {"target_model_names": f"{model_a},{model_b}"},
|
||||
"repeated_plain": {"target_model_names": [model_a, model_b]},
|
||||
}[spelling]
|
||||
managed_file: Final = _managed_upload(gateway, spelling, fields, sdk_base_url)
|
||||
assert managed_file.created_at > 0, managed_file.model_dump()
|
||||
managed: Final = managed_file.id
|
||||
|
||||
requests: Final = _drained_other_than_model_list_probes(wire)
|
||||
assert sorted((r.method, r.target, r.headers["authorization"]) for r in requests) == [
|
||||
("POST", "/v1/files", f"Bearer {key_a}"),
|
||||
("POST", "/v1/files", f"Bearer {key_b}"),
|
||||
], requests
|
||||
uploads: Final = {request.headers["authorization"]: request for request in requests}
|
||||
for bearer, underlying in ((key_a, "gpt-4o-mini"), (key_b, "gpt-4.1-mini")):
|
||||
parts: Final = _multipart_parts(uploads[f"Bearer {bearer}"])
|
||||
assert len(parts) == 2, [part.get_param("name", header="content-disposition") for part in parts]
|
||||
assert _text_fields(parts) == {"purpose": "batch"}
|
||||
file_part: Final = _file_fields(parts)["file"]
|
||||
assert (file_part[0], file_part[2]) == ("modified_file.jsonl", _rewritten_jsonl(underlying)), file_part
|
||||
|
||||
managed_file_part: Final = _file_fields(_multipart_parts(uploads[f"Bearer {key_a}"]))["file"]
|
||||
expected_file: Final = _expected_file_object(
|
||||
managed_file.id,
|
||||
len(managed_file_part[2]),
|
||||
managed_file_part[0],
|
||||
)
|
||||
assert managed_file == expected_file, managed_file.model_dump()
|
||||
|
||||
decoded: Final = _decoded_unified(managed)
|
||||
assert f"target_model_names,{model_a},{model_b}" in decoded, decoded
|
||||
deployment_ids: Final = _deployment_ids(gateway, model_a, model_b)
|
||||
rows: Final = eventually(
|
||||
lambda: read_rows(MANAGED_FILE_ROW, (managed,)),
|
||||
lambda found: len(found) == 1,
|
||||
)
|
||||
assert rows[0]["model_mappings"] == {
|
||||
deployment_ids[model_a]: f"file-{key_a}",
|
||||
deployment_ids[model_b]: f"file-{key_b}",
|
||||
}, rows
|
||||
assert sorted(rows[0]["flat_model_file_ids"]) == sorted([f"file-{key_a}", f"file-{key_b}"]), rows
|
||||
|
||||
readback: Final = gateway.request("GET", f"/v1/files/{managed}")
|
||||
assert readback.status_code == 200, readback.text
|
||||
readback_file: Final = _FileObject.model_validate_json(readback.content)
|
||||
assert readback_file == expected_file, readback.text
|
||||
assert _drained_other_than_model_list_probes(wire) == ()
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"spelling",
|
||||
["sdk_list", "repeated_bracket", "indexed", "comma_joined"],
|
||||
ids=["sdk-list-extra-body", "multipart-repeated-bracket", "multipart-indexed", "multipart-comma-joined"],
|
||||
)
|
||||
def test_batch_upload_to_two_target_models_reaches_each_deployment(
|
||||
gateway: Gateway, spelling: _BatchUploadSpelling
|
||||
) -> None:
|
||||
_assert_upload_reaches_both_target_models(gateway, spelling, _sdk_base_url(gateway))
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"route",
|
||||
["/files", "/openai/v1/files"],
|
||||
ids=["files-alias", "openai-v1-files-alias"],
|
||||
)
|
||||
def test_batch_upload_to_two_target_models_through_a_route_alias_reaches_each_deployment(
|
||||
gateway: Gateway, route: str
|
||||
) -> None:
|
||||
sdk_base_url: Final = str(gateway.client.base_url).rstrip("/") + route.removesuffix("/files")
|
||||
_assert_upload_reaches_both_target_models(gateway, "sdk_list", sdk_base_url)
|
||||
|
||||
|
||||
def test_batch_upload_with_repeated_plain_target_model_names_reaches_each_deployment(gateway: Gateway) -> None:
|
||||
pytest.skip("BUG: repeated plain target_model_names keeps only the last value")
|
||||
_assert_upload_reaches_both_target_models(gateway, "repeated_plain", _sdk_base_url(gateway))
|
||||
|
||||
|
||||
def test_batch_upload_to_two_target_models_keeps_the_jsonl_content_type(gateway: Gateway) -> None:
|
||||
pytest.skip(
|
||||
"BUG: managed multi-target upload sends the rewritten batch file as application/octet-stream instead of "
|
||||
"application/jsonl"
|
||||
)
|
||||
key_a: Final = "provider-key-a-" + uuid.uuid4().hex[:8]
|
||||
key_b: Final = "provider-key-b-" + uuid.uuid4().hex[:8]
|
||||
with wire_server(_multiplexed_openai_backends()) as wire, gateway.scenario() as scenario:
|
||||
model_a: Final = scenario.model(model="openai/gpt-4o-mini", api_base=wire.url + "/v1", api_key=key_a)
|
||||
model_b: Final = scenario.model(model="openai/gpt-4.1-mini", api_base=wire.url + "/v1", api_key=key_b)
|
||||
_wait_until_every_worker_serves(gateway, model_a, model_b)
|
||||
_managed_upload(
|
||||
gateway,
|
||||
"repeated_bracket",
|
||||
{"target_model_names[]": [model_a, model_b]},
|
||||
_sdk_base_url(gateway),
|
||||
)
|
||||
uploads: Final = {
|
||||
request.headers["authorization"]: request for request in _drained_other_than_model_list_probes(wire)
|
||||
}
|
||||
content_types: Final = {
|
||||
f"Bearer {key}": _file_fields(_multipart_parts(uploads[f"Bearer {key}"]))["file"][1]
|
||||
for key in (key_a, key_b)
|
||||
}
|
||||
assert content_types == {
|
||||
f"Bearer {key_a}": "application/jsonl",
|
||||
f"Bearer {key_b}": "application/jsonl",
|
||||
}, content_types
|
||||
|
||||
|
||||
def _model_routed_upload(
|
||||
gateway: Gateway, spelling: _ModelRoutedSpelling, alias: str, uploaded: bytes, sdk_base_url: str
|
||||
) -> _FileObject:
|
||||
match spelling:
|
||||
case "sdk_extra_body":
|
||||
with OpenAI(base_url=sdk_base_url, api_key=gateway.key, max_retries=0) as client:
|
||||
response: Final = client.files.with_raw_response.create(
|
||||
file=(INPUT_FILENAME, uploaded, "application/jsonl"),
|
||||
purpose="batch",
|
||||
extra_body={"model": alias},
|
||||
)
|
||||
assert response.status_code == 200, response.text
|
||||
return _FileObject.model_validate_json(response.content)
|
||||
case "sdk_proxy_fields":
|
||||
with OpenAI(base_url=sdk_base_url, api_key=gateway.key, max_retries=0) as client:
|
||||
response: Final = client.files.with_raw_response.create(
|
||||
file=(INPUT_FILENAME, uploaded, "application/jsonl"),
|
||||
purpose="batch",
|
||||
extra_body={
|
||||
"model": alias,
|
||||
"custom_llm_provider": "openai",
|
||||
"target_model_names": alias,
|
||||
},
|
||||
)
|
||||
assert response.status_code == 200, response.text
|
||||
return _FileObject.model_validate_json(response.content)
|
||||
case "multipart_field":
|
||||
created: Final = gateway.request_multipart(
|
||||
"/v1/files",
|
||||
{"purpose": "batch", "model": alias},
|
||||
{"file": (INPUT_FILENAME, uploaded, "application/jsonl")},
|
||||
)
|
||||
assert created.status_code == 200, created.text
|
||||
return _FileObject.model_validate_json(created.content)
|
||||
case "query_param":
|
||||
queried: Final = gateway.client.post(
|
||||
"/v1/files",
|
||||
params={"model": alias},
|
||||
data={"purpose": "batch"},
|
||||
files={"file": (INPUT_FILENAME, uploaded, "application/jsonl")},
|
||||
headers={"Authorization": f"Bearer {gateway.key}"},
|
||||
)
|
||||
assert queried.status_code == 200, queried.text
|
||||
return _FileObject.model_validate_json(queried.content)
|
||||
case "model_header":
|
||||
headed: Final = gateway.client.post(
|
||||
"/v1/files",
|
||||
data={"purpose": "batch"},
|
||||
files={"file": (INPUT_FILENAME, uploaded, "application/jsonl")},
|
||||
headers={"Authorization": f"Bearer {gateway.key}", "x-litellm-model": alias},
|
||||
)
|
||||
assert headed.status_code == 200, headed.text
|
||||
return _FileObject.model_validate_json(headed.content)
|
||||
case "multipart_proxy_fields":
|
||||
posted: Final = gateway.client.post(
|
||||
"/v1/files",
|
||||
data={
|
||||
"purpose": "batch",
|
||||
"model": alias,
|
||||
"target_model_names": alias,
|
||||
"custom_llm_provider": "openai",
|
||||
"litellm_metadata[tag]": "x",
|
||||
},
|
||||
files={"file": (INPUT_FILENAME, uploaded, "application/jsonl")},
|
||||
headers={"Authorization": f"Bearer {gateway.key}"},
|
||||
)
|
||||
assert posted.status_code == 200, posted.text
|
||||
return _FileObject.model_validate_json(posted.content)
|
||||
case "sdk_expires_after":
|
||||
with OpenAI(base_url=sdk_base_url, api_key=gateway.key, max_retries=0) as client:
|
||||
response: Final = client.files.with_raw_response.create(
|
||||
file=(INPUT_FILENAME, uploaded, "application/jsonl"),
|
||||
purpose="batch",
|
||||
expires_after={"anchor": "created_at", "seconds": 2592000},
|
||||
extra_body={"model": alias},
|
||||
)
|
||||
assert response.status_code == 200, response.text
|
||||
return _FileObject.model_validate_json(response.content)
|
||||
case "multipart_expires_after":
|
||||
created: Final = gateway.request_multipart(
|
||||
"/v1/files",
|
||||
{
|
||||
"purpose": "batch",
|
||||
"model": alias,
|
||||
"expires_after[anchor]": "created_at",
|
||||
"expires_after[seconds]": "2592000",
|
||||
},
|
||||
{"file": (INPUT_FILENAME, uploaded, "application/jsonl")},
|
||||
)
|
||||
assert created.status_code == 200, created.text
|
||||
return _FileObject.model_validate_json(created.content)
|
||||
case _:
|
||||
assert_never(spelling)
|
||||
|
||||
|
||||
def _assert_model_routed_upload_reaches_one_deployment_and_the_encoded_id_reads_back(
|
||||
gateway: Gateway, spelling: _ModelRoutedSpelling, sdk_base_url: str
|
||||
) -> None:
|
||||
key: Final = "provider-key-" + uuid.uuid4().hex[:8]
|
||||
uploaded: Final = _jsonl("uploaded-model")
|
||||
with wire_server(_multiplexed_openai_backends()) as wire, gateway.scenario() as scenario:
|
||||
alias: Final = scenario.model(model="openai/gpt-4o-mini", api_base=wire.url + "/v1", api_key=key)
|
||||
_wait_until_every_worker_serves(gateway, alias)
|
||||
uploaded_file: Final = _model_routed_upload(gateway, spelling, alias, uploaded, sdk_base_url)
|
||||
|
||||
expected_encoded: Final = encode_file_id_with_model(file_id=f"file-{key}", model=alias)
|
||||
assert uploaded_file == _expected_file_object(expected_encoded, len(uploaded)), uploaded_file.model_dump()
|
||||
encoded: Final = uploaded_file.id
|
||||
(upload,) = _drained_other_than_model_list_probes(wire)
|
||||
assert (upload.method, upload.target) == ("POST", "/v1/files"), upload.target
|
||||
assert upload.headers["authorization"] == f"Bearer {key}", upload.headers
|
||||
parts: Final = _multipart_parts(upload)
|
||||
expected_text: Final = {
|
||||
"sdk_extra_body": {"purpose": "batch"},
|
||||
"multipart_field": {"purpose": "batch"},
|
||||
"query_param": {"purpose": "batch"},
|
||||
"model_header": {"purpose": "batch"},
|
||||
"multipart_proxy_fields": {"purpose": "batch"},
|
||||
"sdk_proxy_fields": {"purpose": "batch"},
|
||||
"sdk_expires_after": {
|
||||
"purpose": "batch",
|
||||
"expires_after[anchor]": "created_at",
|
||||
"expires_after[seconds]": "2592000",
|
||||
},
|
||||
"multipart_expires_after": {
|
||||
"purpose": "batch",
|
||||
"expires_after[anchor]": "created_at",
|
||||
"expires_after[seconds]": "2592000",
|
||||
},
|
||||
}[spelling]
|
||||
assert len(parts) == len(expected_text) + 1, [
|
||||
part.get_param("name", header="content-disposition") for part in parts
|
||||
]
|
||||
assert _text_fields(parts) == expected_text
|
||||
assert _file_fields(parts) == {"file": (INPUT_FILENAME, "application/jsonl", uploaded)}
|
||||
|
||||
readback: Final = gateway.request("GET", f"/v1/files/{encoded}")
|
||||
assert readback.status_code == 200, readback.text
|
||||
readback_file: Final = _FileObject.model_validate_json(readback.content)
|
||||
assert readback_file == _expected_file_object(encoded, len(uploaded)), readback.text
|
||||
(fetched,) = _drained_other_than_model_list_probes(wire)
|
||||
assert (fetched.method, fetched.target) == ("GET", f"/v1/files/file-{key}"), fetched.target
|
||||
assert fetched.headers["authorization"] == f"Bearer {key}", fetched.headers
|
||||
assert fetched.body == b""
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"spelling",
|
||||
[
|
||||
"sdk_extra_body",
|
||||
"sdk_proxy_fields",
|
||||
"multipart_field",
|
||||
"query_param",
|
||||
"model_header",
|
||||
"multipart_proxy_fields",
|
||||
"sdk_expires_after",
|
||||
"multipart_expires_after",
|
||||
],
|
||||
ids=[
|
||||
"sdk-extra-body-model",
|
||||
"sdk-extra-body-with-proxy-fields",
|
||||
"multipart-model-field",
|
||||
"query-model",
|
||||
"x-litellm-model-header",
|
||||
"multipart-model-with-proxy-fields",
|
||||
"sdk-expires-after",
|
||||
"multipart-bracketed-expires-after",
|
||||
],
|
||||
)
|
||||
def test_model_routed_upload_reaches_one_deployment_and_the_encoded_id_reads_back(
|
||||
gateway: Gateway, spelling: _ModelRoutedSpelling
|
||||
) -> None:
|
||||
_assert_model_routed_upload_reaches_one_deployment_and_the_encoded_id_reads_back(
|
||||
gateway, spelling, _sdk_base_url(gateway)
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"route",
|
||||
["/files", "/openai/v1/files"],
|
||||
ids=["files-alias", "openai-v1-files-alias"],
|
||||
)
|
||||
def test_model_routed_upload_through_a_route_alias_reaches_one_deployment(
|
||||
gateway: Gateway, route: str
|
||||
) -> None:
|
||||
sdk_base_url: Final = str(gateway.client.base_url).rstrip("/") + route.removesuffix("/files")
|
||||
_assert_model_routed_upload_reaches_one_deployment_and_the_encoded_id_reads_back(
|
||||
gateway, "sdk_extra_body", sdk_base_url
|
||||
)
|
||||
|
||||
|
||||
AZURE_API_VERSION: Final = "2024-10-21"
|
||||
|
||||
|
||||
def _encoded_output_file_id(gateway: Gateway, alias: str, bearer_key: str) -> tuple[str, str]:
|
||||
uploaded_bytes: Final = _jsonl("uploaded-model")
|
||||
uploaded: Final = gateway.request_multipart(
|
||||
"/v1/files",
|
||||
{"purpose": "batch", "model": alias},
|
||||
{"file": (INPUT_FILENAME, uploaded_bytes, "application/jsonl")},
|
||||
)
|
||||
assert uploaded.status_code == 200, uploaded.text
|
||||
encoded_file: Final = encode_file_id_with_model(file_id=f"file-{bearer_key}", model=alias)
|
||||
uploaded_file: Final = _FileObject.model_validate_json(uploaded.content)
|
||||
assert uploaded_file == _expected_file_object(encoded_file, len(uploaded_bytes)), uploaded.text
|
||||
encoded_input: Final = uploaded_file.id
|
||||
created: Final = gateway.request(
|
||||
"POST",
|
||||
"/v1/batches",
|
||||
{"input_file_id": encoded_input, "endpoint": "/v1/chat/completions", "completion_window": "24h"},
|
||||
)
|
||||
assert created.status_code == 200, created.text
|
||||
encoded_batch: Final = string_value(_json(created)["id"])
|
||||
retrieved: Final = gateway.request("GET", f"/v1/batches/{encoded_batch}")
|
||||
assert retrieved.status_code == 200, retrieved.text
|
||||
batch: Final = _Batch.model_validate_json(retrieved.content)
|
||||
assert batch.output_file_id is not None, retrieved.text
|
||||
return batch.output_file_id, f"file-out-{bearer_key}"
|
||||
|
||||
|
||||
def _content_request(
|
||||
gateway: Gateway, route: _ContentRoute, encoded_output: str, raw_output: str, alias: str
|
||||
) -> httpx.Request:
|
||||
authorization: Final = {"Authorization": f"Bearer {gateway.key}"}
|
||||
match route:
|
||||
case "v1_encoded":
|
||||
return gateway.client.build_request("GET", f"/v1/files/{encoded_output}/content", headers=authorization)
|
||||
case "files_encoded":
|
||||
return gateway.client.build_request("GET", f"/files/{encoded_output}/content", headers=authorization)
|
||||
case "files_raw_query":
|
||||
return gateway.client.build_request(
|
||||
"GET", f"/files/{raw_output}/content", params={"model": alias}, headers=authorization
|
||||
)
|
||||
case "v1_raw_header":
|
||||
return gateway.client.build_request(
|
||||
"GET", f"/v1/files/{raw_output}/content", headers={**authorization, "x-litellm-model": alias}
|
||||
)
|
||||
case _:
|
||||
assert_never(route)
|
||||
|
||||
|
||||
def _open_and_read_head(
|
||||
client: httpx.Client, request: httpx.Request
|
||||
) -> tuple[httpx.Response, Iterator[bytes], bytes]:
|
||||
response: Final = client.send(request, stream=True)
|
||||
chunks: Final = response.iter_bytes()
|
||||
return response, chunks, next(chunks, b"")
|
||||
|
||||
|
||||
def _gated_download(gateway: Gateway, request: httpx.Request, gate: threading.Event) -> tuple[int, bytes, bytes]:
|
||||
with ThreadPoolExecutor(max_workers=1) as reader:
|
||||
head_read: Final = reader.submit(_open_and_read_head, gateway.client, request)
|
||||
try:
|
||||
head_arrived: Final = eventually(head_read.done, bool, seconds=4.5, return_last_on_timeout=True)
|
||||
finally:
|
||||
gate.set()
|
||||
response, chunks, head = head_read.result()
|
||||
try:
|
||||
rest: Final = b"".join(chunks)
|
||||
finally:
|
||||
response.close()
|
||||
assert head_arrived, (
|
||||
f"client got no content bytes while the upstream held the rest of the file: {response.status_code}"
|
||||
)
|
||||
return response.status_code, head, rest
|
||||
|
||||
|
||||
@pytest.mark.parametrize("provider", ["openai", "azure"], ids=["openai-streaming", "azure-non-streaming"])
|
||||
@pytest.mark.parametrize(
|
||||
"route",
|
||||
["v1_encoded", "files_encoded", "files_raw_query", "v1_raw_header"],
|
||||
ids=["v1-files-encoded", "files-encoded", "files-raw-model-query", "v1-files-raw-model-header"],
|
||||
)
|
||||
def test_batch_output_file_content_downloads_through_the_model_encoded_id(
|
||||
gateway: Gateway, provider: str, route: _ContentRoute
|
||||
) -> None:
|
||||
key: Final = f"provider-key-{provider}-" + uuid.uuid4().hex[:8]
|
||||
gate: Final = threading.Event()
|
||||
gate.set()
|
||||
received: Final[Queue[Request]] = Queue()
|
||||
backend: Final = _multiplexed_openai_backends(content_gate=gate)
|
||||
|
||||
def respond(request: Request) -> Reply:
|
||||
received.put(request)
|
||||
return backend(request)
|
||||
|
||||
with wire_server(respond) as wire, gateway.scenario() as scenario:
|
||||
deployment: Final = (
|
||||
{"model": "openai/gpt-4o-mini", "api_base": wire.url + "/v1", "api_key": key}
|
||||
if provider == "openai"
|
||||
else {
|
||||
"model": "azure/gpt-4o-mini",
|
||||
"api_base": wire.url,
|
||||
"api_version": AZURE_API_VERSION,
|
||||
"api_key": key,
|
||||
}
|
||||
)
|
||||
alias: Final = scenario.model(**deployment)
|
||||
_wait_until_every_worker_serves(gateway, alias)
|
||||
encoded_output, raw_output = _encoded_output_file_id(gateway, alias, key)
|
||||
assert encoded_output == encode_file_id_with_model(file_id=raw_output, model=alias), encoded_output
|
||||
|
||||
eventually(
|
||||
lambda: tuple(cast(deque[Request], received.queue)),
|
||||
lambda pending: any(urlsplit(request.target).path.endswith("/content") for request in pending),
|
||||
)
|
||||
settled: Final = _drained_other_than_model_list_probes(wire)
|
||||
base_prefix: Final = "/v1/" if provider == "openai" else "/openai/"
|
||||
api_suffix: Final = "" if provider == "openai" else f"?api-version={AZURE_API_VERSION}"
|
||||
expected_content_target: Final = f"{base_prefix}files/{raw_output}/content{api_suffix}"
|
||||
expected_requests: Final = (
|
||||
("POST", f"{base_prefix}files{api_suffix}"),
|
||||
("POST", f"{base_prefix}batches{api_suffix}"),
|
||||
("GET", f"{base_prefix}batches/batch-{key}{api_suffix}"),
|
||||
("GET", expected_content_target),
|
||||
)
|
||||
observed_requests: Final = tuple((r.method, r.target) for r in settled)
|
||||
expected_metadata_request: Final = ("GET", f"{base_prefix}files/{raw_output}{api_suffix}")
|
||||
assert observed_requests.count(expected_metadata_request) <= 1, observed_requests
|
||||
assert tuple(request for request in observed_requests if request != expected_metadata_request) == (
|
||||
expected_requests
|
||||
), observed_requests
|
||||
|
||||
if provider == "openai":
|
||||
gate.clear()
|
||||
status, head, rest = _gated_download(
|
||||
gateway, _content_request(gateway, route, encoded_output, raw_output, alias), gate
|
||||
)
|
||||
assert status == 200, (head + rest).decode()
|
||||
assert head + rest == OUTPUT_HEAD + OUTPUT_BYTES, f"{len(head)} head bytes, {len(rest)} tail bytes"
|
||||
(content_request,) = _drained_other_than_model_list_probes(wire)
|
||||
assert (content_request.method, content_request.target) == ("GET", expected_content_target), (
|
||||
content_request.target
|
||||
)
|
||||
assert content_request.body == b"", content_request.body
|
||||
if provider == "openai":
|
||||
assert content_request.headers["authorization"] == f"Bearer {key}", content_request.headers
|
||||
else:
|
||||
assert content_request.headers["api-key"] == key, content_request.headers
|
||||
|
||||
|
||||
def _encoded_batch_input(
|
||||
gateway: Gateway, spelling: _BatchCreateSpelling, alias: str, wire: Wire, raw_input: str, bearer_key: str
|
||||
) -> str:
|
||||
if spelling == "encoded_input":
|
||||
uploaded_bytes: Final = _jsonl("uploaded-model")
|
||||
uploaded: Final = gateway.request_multipart(
|
||||
"/v1/files",
|
||||
{"purpose": "batch", "model": alias},
|
||||
{"file": (INPUT_FILENAME, uploaded_bytes, "application/jsonl")},
|
||||
)
|
||||
assert uploaded.status_code == 200, uploaded.text
|
||||
encoded_file: Final = encode_file_id_with_model(file_id=f"file-{bearer_key}", model=alias)
|
||||
uploaded_file: Final = _FileObject.model_validate_json(uploaded.content)
|
||||
assert uploaded_file == _expected_file_object(encoded_file, len(uploaded_bytes)), uploaded.text
|
||||
wire.drain()
|
||||
return uploaded_file.id
|
||||
return encode_file_id_with_model(file_id=raw_input, model=alias)
|
||||
|
||||
|
||||
def _created_batch(
|
||||
gateway: Gateway,
|
||||
spelling: _BatchCreateSpelling,
|
||||
encoded_input: str,
|
||||
raw_input: str,
|
||||
alias: str,
|
||||
sdk_base_url: str,
|
||||
) -> _Batch:
|
||||
if spelling in {"query_model", "header_model"}:
|
||||
created: Final = gateway.request(
|
||||
"POST",
|
||||
"/v1/batches",
|
||||
{
|
||||
"input_file_id": raw_input,
|
||||
"endpoint": "/v1/chat/completions",
|
||||
"completion_window": "24h",
|
||||
"metadata": {"job": "x"},
|
||||
},
|
||||
params={"model": alias} if spelling == "query_model" else None,
|
||||
headers={"x-litellm-model": alias} if spelling == "header_model" else None,
|
||||
)
|
||||
assert created.status_code == 200, created.text
|
||||
return _Batch.model_validate_json(created.content)
|
||||
if spelling == "sanitized_metadata":
|
||||
created: Final = gateway.request(
|
||||
"POST",
|
||||
"/v1/batches",
|
||||
{
|
||||
"input_file_id": raw_input,
|
||||
"endpoint": "/v1/chat/completions",
|
||||
"completion_window": "24h",
|
||||
"metadata": {"job": "x", "attempt": 2},
|
||||
"model": alias,
|
||||
},
|
||||
)
|
||||
assert created.status_code == 200, created.text
|
||||
return _Batch.model_validate_json(created.content)
|
||||
with OpenAI(base_url=sdk_base_url, api_key=gateway.key, max_retries=0) as client:
|
||||
extra: Final = {} if spelling == "encoded_input" else {"extra_body": {"model": alias}}
|
||||
response: Final = client.batches.with_raw_response.create(
|
||||
input_file_id=encoded_input if spelling == "encoded_input" else raw_input,
|
||||
endpoint="/v1/chat/completions",
|
||||
completion_window="24h",
|
||||
metadata={"job": "x"},
|
||||
**extra,
|
||||
)
|
||||
assert response.status_code == 200, response.text
|
||||
return _Batch.model_validate_json(response.content)
|
||||
|
||||
|
||||
def _assert_create_batch_reaches_one_deployment_and_returns_model_encoded_ids(
|
||||
gateway: Gateway, spelling: _BatchCreateSpelling, sdk_base_url: str
|
||||
) -> None:
|
||||
key: Final = "provider-key-" + uuid.uuid4().hex[:8]
|
||||
raw_input: Final = f"file-in-{key}"
|
||||
raw_batch: Final = f"batch-{key}"
|
||||
created_batch: Final = _batch_object(
|
||||
raw_batch, raw_input, "validating", output_file_id="file-out", error_file_id="file-err"
|
||||
)
|
||||
with wire_server(_multiplexed_openai_backends(created_batch)) as wire, gateway.scenario() as scenario:
|
||||
alias: Final = scenario.model(model="openai/gpt-4o-mini", api_base=wire.url + "/v1", api_key=key)
|
||||
_wait_until_every_worker_serves(gateway, alias)
|
||||
encoded_input: Final = _encoded_batch_input(gateway, spelling, alias, wire, raw_input, key)
|
||||
created_batch_payload: Final = _created_batch(
|
||||
gateway, spelling, encoded_input, raw_input, alias, sdk_base_url
|
||||
)
|
||||
|
||||
expected_raw_input: Final = f"file-{key}" if spelling == "encoded_input" else raw_input
|
||||
expected_input: Final = (
|
||||
encoded_input if spelling == "encoded_input" else encode_file_id_with_model(file_id=raw_input, model=alias)
|
||||
)
|
||||
assert created_batch_payload.id == encode_file_id_with_model(file_id=raw_batch, model=alias, id_type="batch"), (
|
||||
created_batch_payload.id
|
||||
)
|
||||
assert created_batch_payload.output_file_id == encode_file_id_with_model(file_id="file-out", model=alias), (
|
||||
created_batch_payload.output_file_id
|
||||
)
|
||||
assert created_batch_payload.error_file_id == encode_file_id_with_model(file_id="file-err", model=alias), (
|
||||
created_batch_payload.error_file_id
|
||||
)
|
||||
assert created_batch_payload.input_file_id == expected_input, created_batch_payload.input_file_id
|
||||
|
||||
requests: Final = _drained_other_than_model_list_probes(wire)
|
||||
assert [(r.method, r.target) for r in requests] == [("POST", "/v1/batches")], requests
|
||||
(create_request,) = requests
|
||||
assert create_request.headers["authorization"] == f"Bearer {key}", create_request.headers
|
||||
expected_body: Final = {
|
||||
"input_file_id": expected_raw_input,
|
||||
"endpoint": "/v1/chat/completions",
|
||||
"completion_window": "24h",
|
||||
"metadata": {"job": "x"},
|
||||
}
|
||||
assert JSON_OBJECT.validate_json(create_request.body) == expected_body, create_request.body
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"spelling",
|
||||
["encoded_input", "body_model", "query_model", "header_model", "sanitized_metadata"],
|
||||
ids=["encoded-input-id", "body-model", "query-model", "x-litellm-model-header", "non-string-metadata-dropped"],
|
||||
)
|
||||
def test_create_batch_reaches_one_deployment_and_returns_model_encoded_ids(
|
||||
gateway: Gateway, spelling: _BatchCreateSpelling
|
||||
) -> None:
|
||||
_assert_create_batch_reaches_one_deployment_and_returns_model_encoded_ids(
|
||||
gateway, spelling, _sdk_base_url(gateway)
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("route", ["/batches", "/openai/v1/batches"], ids=["batches-alias", "openai-v1-batches-alias"])
|
||||
def test_create_batch_through_a_route_alias_reaches_one_deployment(gateway: Gateway, route: str) -> None:
|
||||
sdk_base_url: Final = str(gateway.client.base_url).rstrip("/") + route.removesuffix("/batches")
|
||||
_assert_create_batch_reaches_one_deployment_and_returns_model_encoded_ids(gateway, "encoded_input", sdk_base_url)
|
||||
851
tests/integration/security/test_files_batches_access.py
Normal file
851
tests/integration/security/test_files_batches_access.py
Normal file
|
|
@ -0,0 +1,851 @@
|
|||
import base64
|
||||
import contextlib
|
||||
import hashlib
|
||||
import json
|
||||
import socket
|
||||
import socketserver
|
||||
import threading
|
||||
import uuid
|
||||
from collections.abc import Callable, Generator
|
||||
from concurrent.futures import ThreadPoolExecutor
|
||||
from contextlib import contextmanager
|
||||
from dataclasses import dataclass
|
||||
from functools import partial
|
||||
from pathlib import Path
|
||||
from queue import SimpleQueue
|
||||
from typing import Final, Literal
|
||||
from urllib.parse import quote, urlsplit
|
||||
|
||||
import httpx
|
||||
import pytest
|
||||
from botocore.auth import S3SigV4Auth
|
||||
from botocore.awsrequest import AWSRequest
|
||||
from botocore.credentials import Credentials
|
||||
from cryptography.hazmat.primitives import serialization
|
||||
from cryptography.hazmat.primitives.asymmetric import rsa
|
||||
from integration._support.client import Gateway, Scenario, eventually, object_value, string_value
|
||||
from integration._support.database import read_rows
|
||||
from integration._support.process import owned_proxy
|
||||
from integration._support.tls import server_context, write_self_signed_cert
|
||||
from integration._support.wire import Reply, Request, Wire, wire_server
|
||||
from openai import OpenAI, PermissionDeniedError
|
||||
from pydantic import BaseModel, JsonValue, TypeAdapter
|
||||
from typing_extensions import assert_never
|
||||
|
||||
from litellm.proxy.openai_files_endpoints.common_utils import encode_file_id_with_model
|
||||
|
||||
JSON_OBJECT: Final = TypeAdapter(dict[str, JsonValue])
|
||||
MANAGED_FILE_ROW: Final = 'SELECT model_mappings, flat_model_file_ids, created_by, team_id FROM "LiteLLM_ManagedFileTable" WHERE unified_file_id = %s'
|
||||
INPUT_FILENAME: Final = "in.jsonl"
|
||||
MANAGED_PREFIX: Final = "litellm_proxy:"
|
||||
BUCKET: Final = "integration-delete-bucket"
|
||||
VERTEX_PROJECT: Final = "files-batches-delete-project"
|
||||
VERTEX_LOCATION: Final = "us-central1"
|
||||
VERTEX_MODEL: Final = "vertex_ai/gemini-2.5-flash"
|
||||
GCS_HOST: Final = "storage.googleapis.com"
|
||||
GCS_AUTHORITY: Final = f"{GCS_HOST}:443"
|
||||
VERTEX_ACCESS_TOKEN: Final = "ya29.scripted-delete-token"
|
||||
AWS_ACCESS_KEY_ID: Final = "AKIAIOSFODNN7EXAMPLE"
|
||||
AWS_SECRET_ACCESS_KEY: Final = "wJalrXUtnFEMI/K7MDENG/bPxRfiCYEXAMPLEKEY"
|
||||
AWS_REGION: Final = "us-east-1"
|
||||
_ForbiddenSurface = Literal[
|
||||
"upload_sdk_extra_body",
|
||||
"upload_form",
|
||||
"upload_query",
|
||||
"upload_header",
|
||||
"batch_encoded_input",
|
||||
"batch_body_model",
|
||||
"batch_query_model",
|
||||
"batch_header_model",
|
||||
"file_get_encoded",
|
||||
"file_get_query",
|
||||
"file_get_header",
|
||||
"file_content_encoded",
|
||||
"file_content_query",
|
||||
"file_content_header",
|
||||
"batch_get_encoded",
|
||||
"batch_get_query",
|
||||
"batch_get_header",
|
||||
]
|
||||
|
||||
|
||||
def _jsonl(model: str) -> bytes:
|
||||
line: Final = {
|
||||
"custom_id": "r1",
|
||||
"method": "POST",
|
||||
"url": "/v1/chat/completions",
|
||||
"body": {"model": model, "messages": [{"role": "user", "content": "one"}]},
|
||||
}
|
||||
return json.dumps(line, separators=(",", ":")).encode() + b"\n"
|
||||
|
||||
|
||||
class _FileObject(BaseModel):
|
||||
id: str
|
||||
object: str
|
||||
purpose: str
|
||||
|
||||
|
||||
class _Batch(BaseModel):
|
||||
id: str
|
||||
object: str
|
||||
status: str
|
||||
input_file_id: str | None = None
|
||||
|
||||
|
||||
class _ErrorEnvelope(BaseModel):
|
||||
error: dict[str, JsonValue]
|
||||
|
||||
|
||||
class _FileDeleted(BaseModel):
|
||||
id: str
|
||||
object: str
|
||||
deleted: bool
|
||||
|
||||
|
||||
def _json(response: httpx.Response) -> dict[str, JsonValue]:
|
||||
return JSON_OBJECT.validate_json(response.content)
|
||||
|
||||
|
||||
def _json_reply(body: dict[str, JsonValue], status: int = 200) -> Reply:
|
||||
return Reply(status=status, body=json.dumps(body).encode())
|
||||
|
||||
|
||||
def _drained_other_than_model_list_probes(wire: Wire) -> tuple[Request, ...]:
|
||||
return tuple(
|
||||
request for request in wire.drain() if urlsplit(request.target).path not in {"/v1/models", "/openai/models"}
|
||||
)
|
||||
|
||||
|
||||
def _bearer(request: Request) -> str:
|
||||
return request.headers.get("authorization", request.headers.get("api-key", "")).removeprefix("Bearer ")
|
||||
|
||||
|
||||
def _file_object(file_id: str) -> dict[str, JsonValue]:
|
||||
return {
|
||||
"id": file_id,
|
||||
"object": "file",
|
||||
"bytes": 123,
|
||||
"created_at": 1700000000,
|
||||
"filename": INPUT_FILENAME,
|
||||
"purpose": "batch",
|
||||
"status": "processed",
|
||||
}
|
||||
|
||||
|
||||
def _batch_object(batch_id: str, status: str) -> dict[str, JsonValue]:
|
||||
return {
|
||||
"id": batch_id,
|
||||
"object": "batch",
|
||||
"endpoint": "/v1/chat/completions",
|
||||
"input_file_id": "file-pending",
|
||||
"status": status,
|
||||
"created_at": 1700000000,
|
||||
"completion_window": "24h",
|
||||
}
|
||||
|
||||
|
||||
def _openai_backend() -> Callable[[Request], Reply]:
|
||||
def respond(request: Request) -> Reply:
|
||||
path: Final = urlsplit(request.target).path
|
||||
if request.method == "POST" and path == "/v1/files":
|
||||
return _json_reply(_file_object("file-" + _bearer(request)))
|
||||
if request.method == "POST" and path == "/v1/batches":
|
||||
return _json_reply(_batch_object("batch-" + _bearer(request), "in_progress"))
|
||||
if request.method == "POST" and path.endswith("/cancel"):
|
||||
return _json_reply(_batch_object(path.rsplit("/", 2)[1], "cancelled"))
|
||||
if request.method == "GET" and path.startswith("/v1/batches/"):
|
||||
return _json_reply(_batch_object(path.rsplit("/", 1)[1], "in_progress"))
|
||||
if request.method == "GET" and path.startswith("/v1/files/") and path.endswith("/content"):
|
||||
return Reply(body=b"file-content\n", content_type="application/octet-stream")
|
||||
if request.method == "GET" and path.startswith("/v1/files/"):
|
||||
return _json_reply(_file_object(path.rsplit("/", 1)[1]))
|
||||
if request.method == "DELETE" and path.startswith("/v1/files/"):
|
||||
return _json_reply({"id": path.rsplit("/", 1)[1], "object": "file", "deleted": True})
|
||||
return _json_reply({"error": {"message": f"unscripted {request.method} {request.target}"}}, 404)
|
||||
|
||||
return respond
|
||||
|
||||
|
||||
def _decoded_unified(managed_file_id: str) -> str:
|
||||
decoded: Final = base64.urlsafe_b64decode(managed_file_id + "=" * (-len(managed_file_id) % 4)).decode()
|
||||
assert decoded.startswith(MANAGED_PREFIX), decoded
|
||||
return decoded
|
||||
|
||||
|
||||
def _models_over_a_fresh_connection(gateway: Gateway, _: int) -> frozenset[str]:
|
||||
with httpx.Client(base_url=gateway.client.base_url, timeout=15, trust_env=False) as client:
|
||||
listed: Final = client.get("/v1/models", headers={"Authorization": f"Bearer {gateway.key}"})
|
||||
assert listed.status_code == 200, listed.text
|
||||
data: Final = _json(listed)["data"]
|
||||
assert isinstance(data, list), listed.text
|
||||
return frozenset(string_value(object_value(entry)["id"]) for entry in data)
|
||||
|
||||
|
||||
def _every_worker_serves(gateway: Gateway, *models: str) -> bool:
|
||||
with ThreadPoolExecutor(max_workers=16) as pool:
|
||||
rounds: Final = tuple(
|
||||
tuple(pool.map(partial(_models_over_a_fresh_connection, gateway), range(16))) for _ in range(2)
|
||||
)
|
||||
return all(model in seen for model in models for round_ in rounds for seen in round_)
|
||||
|
||||
|
||||
def _wait_until_every_worker_serves(gateway: Gateway, *models: str) -> None:
|
||||
eventually(lambda: _every_worker_serves(gateway, *models), lambda served: served, seconds=90)
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class _Member:
|
||||
team: str
|
||||
user: str
|
||||
key: str
|
||||
|
||||
|
||||
def _member(scenario: Scenario, *models: str) -> _Member:
|
||||
team: Final = scenario.team(models=list(models))
|
||||
user: Final = scenario.member(team)
|
||||
return _Member(team, user, scenario.key(team_id=team, user_id=user))
|
||||
|
||||
|
||||
def _error_message(response: httpx.Response, status: int) -> str:
|
||||
assert response.status_code == status, response.text
|
||||
error: Final = _ErrorEnvelope.model_validate_json(response.content)
|
||||
return str(error.error["message"])
|
||||
|
||||
|
||||
def test_stranger_cannot_read_download_or_delete_another_teams_managed_file(gateway: Gateway) -> None:
|
||||
key: Final = "provider-key-" + uuid.uuid4().hex[:8]
|
||||
with wire_server(_openai_backend()) as wire, gateway.scenario() as scenario:
|
||||
model: Final = scenario.model(model="openai/gpt-4o-mini", api_base=wire.url + "/v1", api_key=key)
|
||||
_wait_until_every_worker_serves(gateway, model)
|
||||
owner: Final = _member(scenario, model)
|
||||
stranger: Final = _member(scenario, model)
|
||||
|
||||
uploaded: Final = gateway.request_multipart(
|
||||
"/v1/files",
|
||||
{"purpose": "batch", "target_model_names": model},
|
||||
{"file": (INPUT_FILENAME, _jsonl(model), "application/jsonl")},
|
||||
key=owner.key,
|
||||
)
|
||||
assert uploaded.status_code == 200, uploaded.text
|
||||
managed: Final = string_value(_json(uploaded)["id"])
|
||||
wire.drain()
|
||||
row: Final = eventually(lambda: read_rows(MANAGED_FILE_ROW, (managed,)), lambda rows: len(rows) == 1)[0]
|
||||
|
||||
denied: Final = (
|
||||
gateway.request("GET", f"/v1/files/{managed}", key=stranger.key),
|
||||
gateway.request("GET", f"/v1/files/{managed}/content", key=stranger.key),
|
||||
gateway.request("DELETE", f"/v1/files/{managed}", key=stranger.key),
|
||||
)
|
||||
for response in denied:
|
||||
assert (
|
||||
_error_message(response, 403) == f"User {stranger.user} does not have access to the file {managed}"
|
||||
), response.text
|
||||
assert _drained_other_than_model_list_probes(wire) == (), "stranger calls must not reach the provider"
|
||||
assert read_rows(MANAGED_FILE_ROW, (managed,)) == [row], read_rows(MANAGED_FILE_ROW, (managed,))
|
||||
|
||||
readback: Final = gateway.request("GET", f"/v1/files/{managed}", key=owner.key)
|
||||
assert readback.status_code == 200, readback.text
|
||||
assert string_value(_json(readback)["id"]) == managed, readback.text
|
||||
assert _drained_other_than_model_list_probes(wire) == ()
|
||||
|
||||
|
||||
def test_stranger_cannot_retrieve_or_cancel_another_teams_managed_batch(gateway: Gateway) -> None:
|
||||
key: Final = "provider-key-" + uuid.uuid4().hex[:8]
|
||||
with wire_server(_openai_backend()) as wire, gateway.scenario() as scenario:
|
||||
model: Final = scenario.model(model="openai/gpt-4o-mini", api_base=wire.url + "/v1", api_key=key)
|
||||
_wait_until_every_worker_serves(gateway, model)
|
||||
owner: Final = _member(scenario, model)
|
||||
stranger: Final = _member(scenario, model)
|
||||
|
||||
uploaded: Final = gateway.request_multipart(
|
||||
"/v1/files",
|
||||
{"purpose": "batch", "target_model_names": model},
|
||||
{"file": (INPUT_FILENAME, _jsonl(model), "application/jsonl")},
|
||||
key=owner.key,
|
||||
)
|
||||
assert uploaded.status_code == 200, uploaded.text
|
||||
managed_input: Final = string_value(_json(uploaded)["id"])
|
||||
created: Final = gateway.request(
|
||||
"POST",
|
||||
"/v1/batches",
|
||||
{"input_file_id": managed_input, "endpoint": "/v1/chat/completions", "completion_window": "24h"},
|
||||
key=owner.key,
|
||||
)
|
||||
assert created.status_code == 200, created.text
|
||||
unified_batch: Final = string_value(_json(created)["id"])
|
||||
wire.drain()
|
||||
|
||||
expected_message: Final = f"User {stranger.user} does not have access to the object {unified_batch}"
|
||||
stranger_get: Final = gateway.request("GET", f"/v1/batches/{unified_batch}", key=stranger.key)
|
||||
assert _error_message(stranger_get, 403) == expected_message, stranger_get.text
|
||||
stranger_cancel: Final = gateway.request("POST", f"/v1/batches/{unified_batch}/cancel", {}, key=stranger.key)
|
||||
assert _error_message(stranger_cancel, 403) == expected_message, stranger_cancel.text
|
||||
assert _drained_other_than_model_list_probes(wire) == (), "stranger batch calls must not reach the provider"
|
||||
|
||||
owner_cancel: Final = gateway.request("POST", f"/v1/batches/{unified_batch}/cancel", {}, key=owner.key)
|
||||
assert owner_cancel.status_code == 200, owner_cancel.text
|
||||
assert string_value(_json(owner_cancel)["id"]) == unified_batch, owner_cancel.text
|
||||
(cancel_request,) = _drained_other_than_model_list_probes(wire)
|
||||
assert (cancel_request.method, cancel_request.target) == ("POST", f"/v1/batches/batch-{key}/cancel"), (
|
||||
cancel_request.target
|
||||
)
|
||||
assert cancel_request.headers["authorization"] == f"Bearer {key}", cancel_request.headers
|
||||
|
||||
|
||||
def _sdk_upload_refusal(gateway: Gateway, forbidden: str, caller_key: str) -> httpx.Response:
|
||||
sdk_base_url: Final = str(gateway.client.base_url).rstrip("/") + "/v1"
|
||||
with OpenAI(base_url=sdk_base_url, api_key=caller_key, max_retries=0) as client:
|
||||
try:
|
||||
uploaded: Final = client.files.with_raw_response.create(
|
||||
file=(INPUT_FILENAME, _jsonl(forbidden), "application/jsonl"),
|
||||
purpose="batch",
|
||||
extra_body={"model": forbidden},
|
||||
)
|
||||
except PermissionDeniedError as refused:
|
||||
return refused.response
|
||||
raise AssertionError(f"SDK upload with a forbidden model was accepted: {uploaded.status_code} {uploaded.text}")
|
||||
|
||||
|
||||
def _forbidden_surface_request(
|
||||
gateway: Gateway,
|
||||
surface: _ForbiddenSurface,
|
||||
forbidden: str,
|
||||
caller_key: str,
|
||||
raw_file: str,
|
||||
raw_batch: str,
|
||||
encoded_file: str,
|
||||
encoded_batch: str,
|
||||
) -> httpx.Response:
|
||||
match surface:
|
||||
case "upload_sdk_extra_body":
|
||||
return _sdk_upload_refusal(gateway, forbidden, caller_key)
|
||||
case "upload_form":
|
||||
return gateway.request_multipart(
|
||||
"/v1/files",
|
||||
{"purpose": "batch", "model": forbidden},
|
||||
{"file": (INPUT_FILENAME, _jsonl(forbidden), "application/jsonl")},
|
||||
key=caller_key,
|
||||
)
|
||||
case "upload_query":
|
||||
return gateway.client.post(
|
||||
"/v1/files",
|
||||
params={"model": forbidden},
|
||||
data={"purpose": "batch"},
|
||||
files={"file": (INPUT_FILENAME, _jsonl(forbidden), "application/jsonl")},
|
||||
headers={"Authorization": f"Bearer {caller_key}"},
|
||||
)
|
||||
case "upload_header":
|
||||
return gateway.client.post(
|
||||
"/v1/files",
|
||||
data={"purpose": "batch"},
|
||||
files={"file": (INPUT_FILENAME, _jsonl(forbidden), "application/jsonl")},
|
||||
headers={"Authorization": f"Bearer {caller_key}", "x-litellm-model": forbidden},
|
||||
)
|
||||
case "batch_encoded_input":
|
||||
return gateway.request(
|
||||
"POST",
|
||||
"/v1/batches",
|
||||
{
|
||||
"input_file_id": encoded_file,
|
||||
"endpoint": "/v1/chat/completions",
|
||||
"completion_window": "24h",
|
||||
},
|
||||
key=caller_key,
|
||||
)
|
||||
case "batch_body_model":
|
||||
return gateway.request(
|
||||
"POST",
|
||||
"/v1/batches",
|
||||
{
|
||||
"input_file_id": raw_file,
|
||||
"endpoint": "/v1/chat/completions",
|
||||
"completion_window": "24h",
|
||||
"model": forbidden,
|
||||
},
|
||||
key=caller_key,
|
||||
)
|
||||
case "batch_query_model":
|
||||
return gateway.request(
|
||||
"POST",
|
||||
"/v1/batches",
|
||||
{"input_file_id": raw_file, "endpoint": "/v1/chat/completions", "completion_window": "24h"},
|
||||
key=caller_key,
|
||||
params={"model": forbidden},
|
||||
)
|
||||
case "batch_header_model":
|
||||
return gateway.request(
|
||||
"POST",
|
||||
"/v1/batches",
|
||||
{"input_file_id": raw_file, "endpoint": "/v1/chat/completions", "completion_window": "24h"},
|
||||
key=caller_key,
|
||||
headers={"x-litellm-model": forbidden},
|
||||
)
|
||||
case "file_get_encoded":
|
||||
return gateway.request("GET", f"/v1/files/{encoded_file}", key=caller_key)
|
||||
case "file_get_query":
|
||||
return gateway.request("GET", f"/v1/files/{raw_file}", key=caller_key, params={"model": forbidden})
|
||||
case "file_get_header":
|
||||
return gateway.request(
|
||||
"GET", f"/v1/files/{raw_file}", key=caller_key, headers={"x-litellm-model": forbidden}
|
||||
)
|
||||
case "file_content_encoded":
|
||||
return gateway.request("GET", f"/v1/files/{encoded_file}/content", key=caller_key)
|
||||
case "file_content_query":
|
||||
return gateway.request(
|
||||
"GET", f"/v1/files/{raw_file}/content", key=caller_key, params={"model": forbidden}
|
||||
)
|
||||
case "file_content_header":
|
||||
return gateway.request(
|
||||
"GET", f"/v1/files/{raw_file}/content", key=caller_key, headers={"x-litellm-model": forbidden}
|
||||
)
|
||||
case "batch_get_encoded":
|
||||
return gateway.request("GET", f"/v1/batches/{encoded_batch}", key=caller_key)
|
||||
case "batch_get_query":
|
||||
return gateway.request("GET", f"/v1/batches/{raw_batch}", key=caller_key, params={"model": forbidden})
|
||||
case "batch_get_header":
|
||||
return gateway.request(
|
||||
"GET", f"/v1/batches/{raw_batch}", key=caller_key, headers={"x-litellm-model": forbidden}
|
||||
)
|
||||
case _:
|
||||
assert_never(surface)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"surface",
|
||||
[
|
||||
"upload_sdk_extra_body",
|
||||
"upload_form",
|
||||
"upload_query",
|
||||
"upload_header",
|
||||
"batch_encoded_input",
|
||||
"batch_body_model",
|
||||
"batch_query_model",
|
||||
"batch_header_model",
|
||||
"file_get_encoded",
|
||||
"file_get_query",
|
||||
"file_get_header",
|
||||
"file_content_encoded",
|
||||
"file_content_query",
|
||||
"file_content_header",
|
||||
"batch_get_encoded",
|
||||
"batch_get_query",
|
||||
"batch_get_header",
|
||||
],
|
||||
ids=[
|
||||
"upload-sdk-extra-body-model",
|
||||
"upload-form-model",
|
||||
"upload-query-model",
|
||||
"upload-header-model",
|
||||
"batch-create-encoded-input",
|
||||
"batch-create-body-model",
|
||||
"batch-create-query-model",
|
||||
"batch-create-model-header",
|
||||
"file-retrieve-encoded-id",
|
||||
"file-retrieve-model-query",
|
||||
"file-retrieve-model-header",
|
||||
"file-content-encoded-id",
|
||||
"file-content-model-query",
|
||||
"file-content-model-header",
|
||||
"batch-retrieve-encoded-id",
|
||||
"batch-retrieve-model-query",
|
||||
"batch-retrieve-model-header",
|
||||
],
|
||||
)
|
||||
def test_a_key_granted_one_model_cannot_spend_another_deployments_files_or_batches(
|
||||
gateway: Gateway, surface: _ForbiddenSurface
|
||||
) -> None:
|
||||
key_allowed: Final = "provider-key-allowed-" + uuid.uuid4().hex[:8]
|
||||
key_forbidden: Final = "provider-key-forbidden-" + uuid.uuid4().hex[:8]
|
||||
with wire_server(_openai_backend()) as wire, gateway.scenario() as scenario:
|
||||
allowed: Final = scenario.model(model="openai/gpt-4o-mini", api_base=wire.url + "/v1", api_key=key_allowed)
|
||||
forbidden: Final = scenario.model(model="openai/gpt-4.1-mini", api_base=wire.url + "/v1", api_key=key_forbidden)
|
||||
_wait_until_every_worker_serves(gateway, allowed, forbidden)
|
||||
team: Final = scenario.team(models=[allowed, forbidden])
|
||||
user: Final = scenario.member(team)
|
||||
caller_key: Final = scenario.key(team_id=team, user_id=user, models=[allowed])
|
||||
caller: Final = _Member(team, user, caller_key)
|
||||
provider_client: Final = httpx.Client(base_url=wire.url, timeout=15, trust_env=False)
|
||||
with provider_client:
|
||||
seeded_file: Final = provider_client.post(
|
||||
"/v1/files",
|
||||
data={"purpose": "batch"},
|
||||
files={"file": (INPUT_FILENAME, _jsonl(forbidden), "application/jsonl")},
|
||||
headers={"Authorization": f"Bearer {key_forbidden}"},
|
||||
)
|
||||
assert seeded_file.status_code == 200, seeded_file.text
|
||||
raw_file_object: Final = _FileObject.model_validate_json(seeded_file.content)
|
||||
assert raw_file_object.id == f"file-{key_forbidden}", seeded_file.text
|
||||
seeded_batch: Final = provider_client.post(
|
||||
"/v1/batches",
|
||||
json={
|
||||
"input_file_id": raw_file_object.id,
|
||||
"endpoint": "/v1/chat/completions",
|
||||
"completion_window": "24h",
|
||||
},
|
||||
headers={"Authorization": f"Bearer {key_forbidden}"},
|
||||
)
|
||||
assert seeded_batch.status_code == 200, seeded_batch.text
|
||||
raw_batch_object: Final = _Batch.model_validate_json(seeded_batch.content)
|
||||
assert raw_batch_object.id == f"batch-{key_forbidden}", seeded_batch.text
|
||||
wire.drain()
|
||||
|
||||
encoded_forbidden_file: Final = encode_file_id_with_model(file_id=raw_file_object.id, model=forbidden)
|
||||
encoded_forbidden_batch: Final = encode_file_id_with_model(
|
||||
file_id=raw_batch_object.id, model=forbidden, id_type="batch"
|
||||
)
|
||||
response: Final = _forbidden_surface_request(
|
||||
gateway,
|
||||
surface,
|
||||
forbidden,
|
||||
caller.key,
|
||||
raw_file_object.id,
|
||||
raw_batch_object.id,
|
||||
encoded_forbidden_file,
|
||||
encoded_forbidden_batch,
|
||||
)
|
||||
provider_requests: Final = _drained_other_than_model_list_probes(wire)
|
||||
assert response.status_code == 403, f"{response.text}; provider requests={provider_requests}"
|
||||
assert _error_message(response, 403) == (
|
||||
f"The requested model '{forbidden}' is not available for this API key, or the model name is invalid. "
|
||||
"Check the models available to you and try again."
|
||||
), response.text
|
||||
assert provider_requests == (), f"{surface} reached the provider: {provider_requests}"
|
||||
|
||||
control: Final = gateway.request_multipart(
|
||||
"/v1/files",
|
||||
{"purpose": "batch", "model": allowed},
|
||||
{"file": (INPUT_FILENAME, _jsonl(allowed), "application/jsonl")},
|
||||
key=caller.key,
|
||||
)
|
||||
assert control.status_code == 200, control.text
|
||||
(control_request,) = _drained_other_than_model_list_probes(wire)
|
||||
assert control_request.headers["authorization"] == f"Bearer {key_allowed}", control_request.headers
|
||||
|
||||
|
||||
def test_indexed_target_model_upload_cannot_use_a_forbidden_deployment(gateway: Gateway) -> None:
|
||||
pytest.skip("BUG: target_model_names[0] upload bypasses the key's model restriction and reaches the deployment")
|
||||
key_allowed: Final = "provider-key-allowed-" + uuid.uuid4().hex[:8]
|
||||
key_forbidden: Final = "provider-key-forbidden-" + uuid.uuid4().hex[:8]
|
||||
with wire_server(_openai_backend()) as wire, gateway.scenario() as scenario:
|
||||
allowed: Final = scenario.model(model="openai/gpt-4o-mini", api_base=wire.url + "/v1", api_key=key_allowed)
|
||||
forbidden: Final = scenario.model(model="openai/gpt-4.1-mini", api_base=wire.url + "/v1", api_key=key_forbidden)
|
||||
_wait_until_every_worker_serves(gateway, allowed, forbidden)
|
||||
team: Final = scenario.team(models=[allowed, forbidden])
|
||||
user: Final = scenario.member(team)
|
||||
caller_key: Final = scenario.key(team_id=team, user_id=user, models=[allowed])
|
||||
|
||||
response: Final = gateway.client.post(
|
||||
"/v1/files",
|
||||
data={"purpose": "batch", "target_model_names[0]": forbidden},
|
||||
files={"file": (INPUT_FILENAME, _jsonl(forbidden), "application/jsonl")},
|
||||
headers={"Authorization": f"Bearer {caller_key}"},
|
||||
)
|
||||
provider_requests: Final = _drained_other_than_model_list_probes(wire)
|
||||
assert response.status_code == 403, f"{response.text}; provider requests={provider_requests}"
|
||||
assert _error_message(response, 403) == (
|
||||
f"The requested model '{forbidden}' is not available for this API key, or the model name is invalid. "
|
||||
"Check the models available to you and try again."
|
||||
), response.text
|
||||
assert provider_requests == (), f"indexed target-model upload reached the provider: {provider_requests}"
|
||||
|
||||
|
||||
def test_repeated_target_model_upload_cannot_use_a_forbidden_deployment(gateway: Gateway) -> None:
|
||||
pytest.skip("BUG: repeated target_model_names[] upload bypasses the key's model restriction and reaches the deployment")
|
||||
key_allowed: Final = "provider-key-allowed-" + uuid.uuid4().hex[:8]
|
||||
key_forbidden: Final = "provider-key-forbidden-" + uuid.uuid4().hex[:8]
|
||||
with wire_server(_openai_backend()) as wire, gateway.scenario() as scenario:
|
||||
allowed: Final = scenario.model(model="openai/gpt-4o-mini", api_base=wire.url + "/v1", api_key=key_allowed)
|
||||
forbidden: Final = scenario.model(model="openai/gpt-4.1-mini", api_base=wire.url + "/v1", api_key=key_forbidden)
|
||||
_wait_until_every_worker_serves(gateway, allowed, forbidden)
|
||||
team: Final = scenario.team(models=[allowed, forbidden])
|
||||
user: Final = scenario.member(team)
|
||||
caller_key: Final = scenario.key(team_id=team, user_id=user, models=[allowed])
|
||||
|
||||
response: Final = gateway.client.post(
|
||||
"/v1/files",
|
||||
data={"purpose": "batch", "target_model_names[]": [allowed, forbidden]},
|
||||
files={"file": (INPUT_FILENAME, _jsonl(forbidden), "application/jsonl")},
|
||||
headers={"Authorization": f"Bearer {caller_key}"},
|
||||
)
|
||||
provider_requests: Final = _drained_other_than_model_list_probes(wire)
|
||||
assert response.status_code == 403, f"{response.text}; provider requests={provider_requests}"
|
||||
assert _error_message(response, 403) == (
|
||||
f"The requested model '{forbidden}' is not available for this API key, or the model name is invalid. "
|
||||
"Check the models available to you and try again."
|
||||
), response.text
|
||||
assert provider_requests == (), f"repeated target-model upload reached the provider: {provider_requests}"
|
||||
|
||||
|
||||
def test_plain_target_model_upload_cannot_use_a_forbidden_deployment(gateway: Gateway) -> None:
|
||||
key_allowed: Final = "provider-key-allowed-" + uuid.uuid4().hex[:8]
|
||||
key_forbidden: Final = "provider-key-forbidden-" + uuid.uuid4().hex[:8]
|
||||
with wire_server(_openai_backend()) as wire, gateway.scenario() as scenario:
|
||||
allowed: Final = scenario.model(model="openai/gpt-4o-mini", api_base=wire.url + "/v1", api_key=key_allowed)
|
||||
forbidden: Final = scenario.model(model="openai/gpt-4.1-mini", api_base=wire.url + "/v1", api_key=key_forbidden)
|
||||
_wait_until_every_worker_serves(gateway, allowed, forbidden)
|
||||
team: Final = scenario.team(models=[allowed, forbidden])
|
||||
user: Final = scenario.member(team)
|
||||
caller_key: Final = scenario.key(team_id=team, user_id=user, models=[allowed])
|
||||
|
||||
response: Final = gateway.request_multipart(
|
||||
"/v1/files",
|
||||
{"purpose": "batch", "target_model_names": forbidden},
|
||||
{"file": (INPUT_FILENAME, _jsonl(forbidden), "application/jsonl")},
|
||||
key=caller_key,
|
||||
)
|
||||
provider_requests: Final = _drained_other_than_model_list_probes(wire)
|
||||
assert response.status_code == 403, f"{response.text}; provider requests={provider_requests}"
|
||||
assert _error_message(response, 403) == (
|
||||
f"The requested model '{forbidden}' is not available for this API key, or the model name is invalid. "
|
||||
"Check the models available to you and try again."
|
||||
), response.text
|
||||
assert provider_requests == (), f"plain target-model upload reached the provider: {provider_requests}"
|
||||
|
||||
|
||||
def _s3_backend() -> Callable[[Request], Reply]:
|
||||
def respond(request: Request) -> Reply:
|
||||
if request.method in {"DELETE", "GET"}:
|
||||
return Reply(status=204, body=b"")
|
||||
return _json_reply({"error": {"message": f"unscripted {request.method} {request.target}"}}, 404)
|
||||
|
||||
return respond
|
||||
|
||||
|
||||
def _bedrock_model(scenario: Scenario, wire: Wire) -> str:
|
||||
return scenario.model(
|
||||
model="bedrock/anthropic.claude-3-haiku-20240307-v1:0",
|
||||
api_key=None,
|
||||
api_base=None,
|
||||
aws_access_key_id=AWS_ACCESS_KEY_ID,
|
||||
aws_secret_access_key=AWS_SECRET_ACCESS_KEY,
|
||||
aws_region_name=AWS_REGION,
|
||||
s3_bucket_name=BUCKET,
|
||||
s3_endpoint_url=wire.url,
|
||||
)
|
||||
|
||||
|
||||
def _vertex_credentials(token_url: str) -> str:
|
||||
private_key: Final = (
|
||||
rsa.generate_private_key(public_exponent=65537, key_size=2048)
|
||||
.private_bytes(
|
||||
serialization.Encoding.PEM,
|
||||
serialization.PrivateFormat.PKCS8,
|
||||
serialization.NoEncryption(),
|
||||
)
|
||||
.decode()
|
||||
)
|
||||
return json.dumps(
|
||||
{
|
||||
"type": "service_account",
|
||||
"project_id": VERTEX_PROJECT,
|
||||
"private_key_id": "scripted",
|
||||
"private_key": private_key,
|
||||
"client_email": f"scripted@{VERTEX_PROJECT}.iam.gserviceaccount.com",
|
||||
"client_id": "0",
|
||||
"auth_uri": f"{token_url}/_oauth/authorize",
|
||||
"token_uri": f"{token_url}/_oauth/token",
|
||||
}
|
||||
)
|
||||
|
||||
|
||||
def _vertex_token_backend(request: Request) -> Reply:
|
||||
if request.method == "POST" and request.target == "/_oauth/token":
|
||||
return Reply(status=400, body=b'{"error":"invalid_grant"}', content_type="application/json")
|
||||
return _json_reply({"error": {"message": f"unscripted {request.method} {request.target}"}}, 404)
|
||||
|
||||
|
||||
def test_non_admin_cannot_delete_a_raw_s3_uri(gateway: Gateway) -> None:
|
||||
with wire_server(_s3_backend()) as wire, gateway.scenario() as scenario:
|
||||
model: Final = _bedrock_model(scenario, wire)
|
||||
_wait_until_every_worker_serves(gateway, model)
|
||||
caller_user: Final = scenario.user(user_role="internal_user")
|
||||
caller_key: Final = scenario.key(user_id=caller_user, models=[model])
|
||||
object_uri: Final = f"s3://{BUCKET}/litellm-bedrock-files-obj.jsonl"
|
||||
provider_path: Final = "/bedrock/v1/files/" + quote(object_uri, safe="")
|
||||
|
||||
denied: Final = gateway.request("DELETE", provider_path, key=caller_key, params={"model": model})
|
||||
assert denied.status_code == 403, denied.text
|
||||
assert _error_message(denied, 403) == (
|
||||
"Raw cloud storage file ids can only be deleted by a proxy admin key. "
|
||||
"Use the LiteLLM managed file id returned when the file was created."
|
||||
), denied.text
|
||||
assert _drained_other_than_model_list_probes(wire) == (), (
|
||||
"non-admin cloud-storage delete reached the storage backend"
|
||||
)
|
||||
|
||||
|
||||
def test_non_admin_cannot_delete_a_raw_gcs_uri_before_vertex_auth_or_storage(gateway: Gateway) -> None:
|
||||
with (
|
||||
wire_server(_vertex_token_backend) as token_wire,
|
||||
gateway.scenario() as scenario,
|
||||
):
|
||||
model: Final = scenario.model(
|
||||
model=VERTEX_MODEL,
|
||||
api_key=None,
|
||||
vertex_project=VERTEX_PROJECT,
|
||||
vertex_location=VERTEX_LOCATION,
|
||||
vertex_credentials=_vertex_credentials(token_wire.url),
|
||||
gcs_bucket_name=BUCKET,
|
||||
)
|
||||
_wait_until_every_worker_serves(gateway, model)
|
||||
caller_user: Final = scenario.user(user_role="internal_user")
|
||||
caller_key: Final = scenario.key(user_id=caller_user, models=[model])
|
||||
object_uri: Final = f"gs://{BUCKET}/litellm-vertex-files/litellm-vertex-files-obj.jsonl"
|
||||
token_wire.drain()
|
||||
|
||||
denied: Final = gateway.request(
|
||||
"DELETE",
|
||||
"/vertex_ai/v1/files/" + quote(object_uri, safe=""),
|
||||
key=caller_key,
|
||||
params={"model": model},
|
||||
)
|
||||
token_requests: Final = token_wire.drain()
|
||||
assert denied.status_code == 403, f"{denied.text}; token requests={token_requests}"
|
||||
assert _error_message(denied, 403) == (
|
||||
"Raw cloud storage file ids can only be deleted by a proxy admin key. "
|
||||
"Use the LiteLLM managed file id returned when the file was created."
|
||||
), denied.text
|
||||
assert token_requests == (), "non-admin raw GCS delete reached the Vertex token endpoint"
|
||||
|
||||
|
||||
def test_proxy_admin_deletes_a_raw_s3_uri_with_one_signed_delete(gateway: Gateway) -> None:
|
||||
with wire_server(_s3_backend()) as wire, gateway.scenario() as scenario:
|
||||
model: Final = _bedrock_model(scenario, wire)
|
||||
_wait_until_every_worker_serves(gateway, model)
|
||||
object_uri: Final = f"s3://{BUCKET}/litellm-bedrock-files-obj.jsonl"
|
||||
encoded_path: Final = "/v1/files/" + quote(object_uri, safe="")
|
||||
|
||||
deleted: Final = gateway.request("DELETE", encoded_path, params={"model": model})
|
||||
assert deleted.status_code == 200, deleted.text
|
||||
parsed: Final = _FileDeleted.model_validate_json(deleted.content)
|
||||
assert parsed.deleted is True and parsed.id == object_uri, deleted.text
|
||||
(request,) = _drained_other_than_model_list_probes(wire)
|
||||
assert request.target == f"/{BUCKET}/litellm-bedrock-files-obj.jsonl", request.target
|
||||
request_date: Final = request.headers["x-amz-date"]
|
||||
content_hash: Final = hashlib.sha256(b"").hexdigest()
|
||||
assert request.headers["x-amz-content-sha256"] == content_hash, request.headers
|
||||
aws_request: Final = AWSRequest(
|
||||
method="DELETE",
|
||||
url=wire.url + request.target,
|
||||
headers={
|
||||
"host": request.headers["host"],
|
||||
"x-amz-date": request_date,
|
||||
"x-amz-content-sha256": content_hash,
|
||||
},
|
||||
)
|
||||
aws_request.context["timestamp"] = request_date
|
||||
auth: Final = S3SigV4Auth(Credentials(AWS_ACCESS_KEY_ID, AWS_SECRET_ACCESS_KEY), "s3", AWS_REGION)
|
||||
signature: Final = auth.signature(
|
||||
auth.string_to_sign(aws_request, auth.canonical_request(aws_request)), aws_request
|
||||
)
|
||||
expected_authorization: Final = (
|
||||
f"AWS4-HMAC-SHA256 Credential={AWS_ACCESS_KEY_ID}/{request_date[:8]}/{AWS_REGION}/s3/aws4_request, "
|
||||
f"SignedHeaders={auth.signed_headers(auth.headers_to_sign(aws_request))}, Signature={signature}"
|
||||
)
|
||||
assert request.headers["authorization"] == expected_authorization, request.headers["authorization"]
|
||||
assert request.method == "DELETE", request.method
|
||||
|
||||
|
||||
def _granting_vertex_token_backend(request: Request) -> Reply:
|
||||
if request.method == "POST" and request.target == "/_oauth/token":
|
||||
return _json_reply({"access_token": VERTEX_ACCESS_TOKEN, "expires_in": 3600, "token_type": "Bearer"})
|
||||
return _json_reply({"error": {"message": f"unscripted {request.method} {request.target}"}}, 404)
|
||||
|
||||
|
||||
def _gcs_backend(request: Request) -> Reply:
|
||||
if request.method == "DELETE":
|
||||
return Reply(status=204, body=b"")
|
||||
return _json_reply({"error": {"message": f"unscripted {request.method} {request.target}"}}, 404)
|
||||
|
||||
|
||||
def _pipe(source: socket.socket, sink: socket.socket) -> None:
|
||||
with contextlib.suppress(OSError):
|
||||
for chunk in iter(lambda: source.recv(65536), b""):
|
||||
sink.sendall(chunk)
|
||||
with contextlib.suppress(OSError):
|
||||
sink.shutdown(socket.SHUT_WR)
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class _ConnectTunnel:
|
||||
url: str
|
||||
authorities: SimpleQueue[str]
|
||||
|
||||
|
||||
@contextmanager
|
||||
def _gcs_tunnel(destination: Wire) -> Generator[_ConnectTunnel, None, None]:
|
||||
authorities: Final[SimpleQueue[str]] = SimpleQueue()
|
||||
destination_port: Final = int(destination.url.rsplit(":", 1)[1])
|
||||
|
||||
class Tunnel(socketserver.StreamRequestHandler):
|
||||
rbufsize = 0
|
||||
request: socket.socket
|
||||
|
||||
def handle(self) -> None:
|
||||
authority: Final = self.rfile.readline().decode().split()[1]
|
||||
while self.rfile.readline() not in (b"\r\n", b""):
|
||||
pass
|
||||
authorities.put(authority)
|
||||
if authority != GCS_AUTHORITY:
|
||||
self.wfile.write(b"HTTP/1.1 403 Forbidden\r\ncontent-length: 0\r\n\r\n")
|
||||
return
|
||||
self.wfile.write(b"HTTP/1.1 200 Connection established\r\n\r\n")
|
||||
self.request.settimeout(10)
|
||||
with socket.create_connection(("127.0.0.1", destination_port), timeout=10) as upstream:
|
||||
outbound: Final = threading.Thread(target=_pipe, args=(self.request, upstream))
|
||||
outbound.start()
|
||||
_pipe(upstream, self.request)
|
||||
outbound.join(timeout=12)
|
||||
|
||||
with socketserver.ThreadingTCPServer(("127.0.0.1", 0), Tunnel) as server:
|
||||
thread: Final = threading.Thread(target=server.serve_forever, kwargs={"poll_interval": 0.05})
|
||||
thread.start()
|
||||
try:
|
||||
yield _ConnectTunnel(f"http://127.0.0.1:{server.server_address[1]}", authorities)
|
||||
finally:
|
||||
server.shutdown()
|
||||
thread.join(timeout=6)
|
||||
|
||||
|
||||
def _queued(queue: SimpleQueue[str]) -> tuple[str, ...]:
|
||||
return tuple(queue.get_nowait() for _ in range(queue.qsize()))
|
||||
|
||||
|
||||
@pytest.mark.timeout(180)
|
||||
def test_proxy_admin_deletes_a_raw_gcs_uri_with_one_authorized_storage_delete(
|
||||
gateway: Gateway, tmp_path: Path
|
||||
) -> None:
|
||||
cert_file, key_file = write_self_signed_cert(tmp_path, (GCS_HOST,))
|
||||
object_name: Final = f"litellm-vertex-files/admin-{uuid.uuid4().hex}.jsonl"
|
||||
object_uri: Final = f"gs://{BUCKET}/{object_name}"
|
||||
with (
|
||||
wire_server(_granting_vertex_token_backend) as token_wire,
|
||||
wire_server(_gcs_backend, tls=server_context(cert_file, key_file)) as storage,
|
||||
_gcs_tunnel(storage) as tunnel,
|
||||
owned_proxy(gateway, tmp_path, {"HTTPS_PROXY": tunnel.url, "SSL_VERIFY": "False"}) as candidate,
|
||||
candidate.scenario() as scenario,
|
||||
):
|
||||
model: Final = scenario.model(
|
||||
model=VERTEX_MODEL,
|
||||
api_key=None,
|
||||
vertex_project=VERTEX_PROJECT,
|
||||
vertex_location=VERTEX_LOCATION,
|
||||
vertex_credentials=_vertex_credentials(token_wire.url),
|
||||
gcs_bucket_name=BUCKET,
|
||||
)
|
||||
token_wire.drain()
|
||||
|
||||
deleted: Final = candidate.request(
|
||||
"DELETE", "/v1/files/" + quote(object_uri, safe=""), params={"model": model}
|
||||
)
|
||||
storage_requests: Final = storage.drain()
|
||||
assert deleted.status_code == 200, f"{deleted.text}; storage requests={storage_requests}"
|
||||
parsed: Final = _FileDeleted.model_validate_json(deleted.content)
|
||||
assert parsed.deleted is True and parsed.id == object_uri, deleted.text
|
||||
assert [(r.method, r.target) for r in storage_requests] == [
|
||||
("DELETE", f"/storage/v1/b/{BUCKET}/o/{quote(object_name, safe='')}")
|
||||
], storage_requests
|
||||
assert storage_requests[0].headers["authorization"] == f"Bearer {VERTEX_ACCESS_TOKEN}", (
|
||||
storage_requests[0].headers
|
||||
)
|
||||
assert _queued(tunnel.authorities) == (GCS_AUTHORITY,), "raw GCS delete took an unexpected route"
|
||||
Loading…
Add table
Reference in a new issue