mirror of
https://github.com/BerriAI/litellm.git
synced 2026-11-14 09:03:16 +00:00
fix(proxy): check s3:// URLs only inside media parts and refuse them only on Converse batch lines
The s3:// bucket check now reads URL fields only inside the content parts that a Bedrock transform turns into an S3 source, and it uses the same s3:// test as the prompt factory. Values in metadata, annotations, tool arguments and other fields pass as before. Each URL is named once in the 403 and 400 messages Bedrock batch files refuse s3:// media only on the lines that go through the Converse transform, so passthrough-provider lines upload unchanged
This commit is contained in:
mateo-berri
parent
b089fb9f90
commit
2c23d2253a
7 changed files with 367 additions and 22 deletions
|
|
@ -50,14 +50,36 @@ def is_managed_cloud_storage_uri(file_id: str) -> bool:
|
|||
|
||||
|
||||
_S3_SCHEME: Final = "s3://"
|
||||
_FIRST_CHARACTERS_THAT_CAN_DECODE_TO_S3: Final = ("s", "%")
|
||||
_S3_MEDIA_PART_TYPES: Final = frozenset(
|
||||
{"image_url", "file", "input_image", "input_file", "computer_screenshot", "image", "document"}
|
||||
)
|
||||
_GEMINI_MEDIA_PART_KEYS: Final = frozenset({"file_data", "fileData"})
|
||||
_MEDIA_URI_KEYS: Final = frozenset({"url", "image_url", "file_id", "file_url", "file_data", "file_uri", "fileUri"})
|
||||
|
||||
|
||||
def _media_uris(value: object, key: str | None = None) -> Iterator[str]:
|
||||
def is_s3_uri(value: str) -> bool:
|
||||
return value.startswith(_S3_SCHEME)
|
||||
|
||||
|
||||
def _media_part_uris(value: object, key: str | None = None) -> Iterator[str]:
|
||||
if isinstance(value, str):
|
||||
return iter((value,) if key in _MEDIA_URI_KEYS else ())
|
||||
fields: Final = as_str_mapping(value)
|
||||
if fields is not None:
|
||||
return chain.from_iterable(_media_part_uris(child, child_key) for child_key, child in fields.items())
|
||||
items: Final = as_object_list(value)
|
||||
return chain.from_iterable(map(_media_part_uris, items)) if items is not None else iter(())
|
||||
|
||||
|
||||
def _is_media_part(fields: Mapping[str, object], key: str | None) -> bool:
|
||||
part_type: Final = fields.get("type")
|
||||
return (isinstance(part_type, str) and part_type in _S3_MEDIA_PART_TYPES) or key in _GEMINI_MEDIA_PART_KEYS
|
||||
|
||||
|
||||
def _media_uris(value: object, key: str | None = None) -> Iterator[str]:
|
||||
fields: Final = as_str_mapping(value)
|
||||
if fields is not None and _is_media_part(fields, key):
|
||||
return _media_part_uris(fields)
|
||||
if fields is not None:
|
||||
return chain.from_iterable(_media_uris(child, child_key) for child_key, child in fields.items())
|
||||
items: Final = as_object_list(value)
|
||||
|
|
@ -88,14 +110,8 @@ def _refused_in_some_decoding(
|
|||
return bool(parts) and (not settled or not all(map(is_allowed, parts)))
|
||||
|
||||
|
||||
def _is_s3_uri(uri: str) -> bool:
|
||||
return uri.startswith(_FIRST_CHARACTERS_THAT_CAN_DECODE_TO_S3) and any(
|
||||
decoding.startswith(_S3_SCHEME) for decoding in _decodings(uri)
|
||||
)
|
||||
|
||||
|
||||
def s3_media_uris(request_body: Mapping[str, object]) -> tuple[str, ...]:
|
||||
return tuple(filter(_is_s3_uri, _media_uris(request_body)))
|
||||
return tuple(dict.fromkeys(filter(is_s3_uri, _media_uris(request_body))))
|
||||
|
||||
|
||||
def is_managed_s3_uri(uri: str) -> bool:
|
||||
|
|
|
|||
|
|
@ -23,6 +23,7 @@ import litellm.types.llms
|
|||
from litellm import verbose_logger
|
||||
from litellm._uuid import uuid
|
||||
from litellm.constants import REDACTED_BY_LITELLM
|
||||
from litellm.litellm_core_utils.cloud_storage_security import is_s3_uri
|
||||
from litellm.litellm_core_utils.core_helpers import as_str_mapping, attribute_or_none
|
||||
from litellm.litellm_core_utils.prompt_templates.mid_conversation_system import (
|
||||
anthropic_system_messages,
|
||||
|
|
@ -4260,7 +4261,7 @@ class BedrockImageProcessor:
|
|||
) -> BedrockContentBlock:
|
||||
"""Synchronous image processing."""
|
||||
|
||||
if image_url.startswith("s3://"):
|
||||
if is_s3_uri(image_url):
|
||||
return cls._create_s3_block(image_url, format, bucket_owner)
|
||||
if "base64" in image_url:
|
||||
img_bytes, mime_type, image_format = cls._parse_base64_image(image_url)
|
||||
|
|
@ -4283,7 +4284,7 @@ class BedrockImageProcessor:
|
|||
) -> BedrockContentBlock:
|
||||
"""Asynchronous image processing."""
|
||||
|
||||
if image_url.startswith("s3://"):
|
||||
if is_s3_uri(image_url):
|
||||
return cls._create_s3_block(image_url, format, bucket_owner)
|
||||
if "base64" in image_url:
|
||||
img_bytes, mime_type, image_format = cls._parse_base64_image(image_url)
|
||||
|
|
@ -4452,7 +4453,7 @@ def _s3_bucket_owner(media: object, url: str) -> str | None:
|
|||
owner: Final = fields.get("bucket_owner") if fields is not None else None
|
||||
if owner is None or (isinstance(owner, str) and owner):
|
||||
return owner
|
||||
if not url.startswith("s3://"):
|
||||
if not is_s3_uri(url):
|
||||
return None
|
||||
raise litellm.BadRequestError(
|
||||
message=f"`bucket_owner` for {url} must be the AWS account id as a non-empty string, not {owner!r}.",
|
||||
|
|
|
|||
|
|
@ -357,6 +357,14 @@ def _batch_line_error(line_number: int, problem: str) -> BedrockError:
|
|||
return _client_error(f"Batch input file line {line_number} {problem}")
|
||||
|
||||
|
||||
class _S3MediaInBatchRecord(ValueError):
|
||||
def __init__(self, s3_uris: tuple[str, ...]) -> None:
|
||||
super().__init__(
|
||||
f"has s3:// media URLs ({', '.join(s3_uris)}). "
|
||||
"LiteLLM does not accept s3:// media URLs in Bedrock batch files, so send the media inline as base64 data"
|
||||
)
|
||||
|
||||
|
||||
def _decode_batch_input(content: str | bytes) -> str:
|
||||
if isinstance(content, str):
|
||||
return content.removeprefix("\ufeff")
|
||||
|
|
@ -392,12 +400,6 @@ def _batch_body_problem(record: Mapping[str, object]) -> str | None:
|
|||
return "has a body that is missing or not a JSON object"
|
||||
if "model" in body and not isinstance(body["model"], str):
|
||||
return "has a body.model that is not a string"
|
||||
s3_uris: Final = s3_media_uris(body)
|
||||
if s3_uris:
|
||||
return (
|
||||
f"has s3:// media URLs ({', '.join(s3_uris)}). "
|
||||
"LiteLLM does not accept s3:// media URLs in Bedrock batch files, so send the media inline as base64 data"
|
||||
)
|
||||
return None
|
||||
|
||||
|
||||
|
|
@ -1094,6 +1096,9 @@ class BedrockFilesConfig(BaseAWSLLM, BaseFilesConfig):
|
|||
AmazonConverseConfig,
|
||||
)
|
||||
|
||||
s3_uris: Final = s3_media_uris(openai_request_body)
|
||||
if s3_uris:
|
||||
raise _S3MediaInBatchRecord(s3_uris)
|
||||
converse_config: Final = AmazonConverseConfig()
|
||||
mapped_params = converse_config.map_openai_params(
|
||||
non_default_params={k: v for k, v in optional_params.items() if k != "requestMetadata"},
|
||||
|
|
@ -1169,6 +1174,8 @@ class BedrockFilesConfig(BaseAWSLLM, BaseFilesConfig):
|
|||
) -> _BedrockBatchRecord:
|
||||
try:
|
||||
return self._transform_openai_batch_record(idx, batch_input_line.record, target_model)
|
||||
except _S3MediaInBatchRecord as e:
|
||||
raise _batch_line_error(batch_input_line.line_number, str(e)) from e
|
||||
except _BATCH_RECORD_CONVERSION_ERRORS as e:
|
||||
raise _batch_line_error(
|
||||
batch_input_line.line_number, f"could not be converted to a Bedrock batch record: {e}"
|
||||
|
|
|
|||
|
|
@ -56,6 +56,7 @@ from litellm.integrations.anthropic_cache_control_hook import (
|
|||
whole_messages_within_cache_control_limit,
|
||||
with_cache_control,
|
||||
)
|
||||
from litellm.litellm_core_utils.cloud_storage_security import is_s3_uri
|
||||
from litellm.litellm_core_utils.core_helpers import as_str_mapping, drops_unsupported_params, normalize_drop_params
|
||||
from litellm.litellm_core_utils.get_supported_openai_params import (
|
||||
get_supported_openai_params,
|
||||
|
|
@ -295,7 +296,7 @@ def _validate_responses_media(adapter: TypeAdapter[_MediaT], value: object) -> _
|
|||
|
||||
|
||||
def _s3_bucket_owner_field(url: object, bucket_owner: object) -> Mapping[str, object]:
|
||||
is_s3_url: Final = isinstance(url, str) and url.startswith("s3://")
|
||||
is_s3_url: Final = isinstance(url, str) and is_s3_uri(url)
|
||||
return {"bucket_owner": bucket_owner} if is_s3_url and bucket_owner is not None else {}
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -132,7 +132,18 @@ def _gemini_turn(*parts: object) -> dict[str, object]:
|
|||
{"system_instruction": {"parts": [{"file_data": {"file_uri": _MANAGED_UPLOAD}}]}, "contents": []},
|
||||
id="gemini-system-instruction",
|
||||
),
|
||||
pytest.param({"metadata": {"url": _MANAGED_UPLOAD}}, id="media-key-in-any-field"),
|
||||
pytest.param(
|
||||
{
|
||||
"input": [
|
||||
{
|
||||
"type": "computer_call_output",
|
||||
"call_id": "call_1",
|
||||
"output": {"type": "computer_screenshot", "image_url": _MANAGED_UPLOAD},
|
||||
}
|
||||
]
|
||||
},
|
||||
id="responses-computer-screenshot",
|
||||
),
|
||||
),
|
||||
)
|
||||
def test_s3_media_uris_finds_an_s3_object_on_every_media_surface(request_body: dict[str, object]):
|
||||
|
|
@ -151,8 +162,6 @@ def test_s3_media_uris_finds_an_s3_object_on_every_media_surface(request_body: d
|
|||
"s3://media-bucket/litellm%2Dbedrock-files/x.pdf",
|
||||
"s3://media-bucket/litellm%252Dbedrock-files/x.pdf",
|
||||
"s3://media-bucket%2Flitellm-bedrock-files/x.pdf",
|
||||
"s3%3A%2F%2Fmedia-bucket%2Flitellm-bedrock-files%2Fx.pdf",
|
||||
"%73%33%3A%2F%2Fmedia-bucket%2Flitellm-bedrock-files%2Fx.pdf",
|
||||
"s3:///litellm-bedrock-files/x.pdf",
|
||||
),
|
||||
)
|
||||
|
|
@ -196,12 +205,98 @@ def test_is_managed_s3_uri_leaves_other_objects_alone(uri: str):
|
|||
pytest.param({"input": _MANAGED_UPLOAD}, id="responses-string-input"),
|
||||
pytest.param({"system": _MANAGED_UPLOAD, "messages": []}, id="string-system"),
|
||||
pytest.param({"metadata": {"note": _MANAGED_UPLOAD}}, id="non-media-key"),
|
||||
pytest.param({"metadata": {"url": _MANAGED_UPLOAD}}, id="metadata-url"),
|
||||
pytest.param({"extra_body": {"url": _MANAGED_UPLOAD}}, id="extra-body-url"),
|
||||
pytest.param({"name": "fetch", "arguments": {"url": _MANAGED_UPLOAD}}, id="mcp-tool-arguments"),
|
||||
pytest.param(
|
||||
{
|
||||
"messages": [
|
||||
{
|
||||
"role": "assistant",
|
||||
"content": "See the source.",
|
||||
"annotations": [{"type": "url_citation", "url_citation": {"url": _MANAGED_UPLOAD}}],
|
||||
}
|
||||
]
|
||||
},
|
||||
id="chat-url-citation",
|
||||
),
|
||||
pytest.param(
|
||||
{
|
||||
"input": [
|
||||
{
|
||||
"type": "message",
|
||||
"role": "assistant",
|
||||
"content": [
|
||||
{
|
||||
"type": "output_text",
|
||||
"text": "See the file.",
|
||||
"annotations": [{"type": "file_citation", "file_id": _MANAGED_UPLOAD, "index": 0}],
|
||||
}
|
||||
],
|
||||
}
|
||||
]
|
||||
},
|
||||
id="responses-file-citation",
|
||||
),
|
||||
pytest.param(
|
||||
_gemini_turn({"functionCall": {"name": "fetch", "args": {"url": _MANAGED_UPLOAD}}}),
|
||||
id="gemini-function-call-camel-case",
|
||||
),
|
||||
pytest.param(
|
||||
_gemini_turn({"function_call": {"name": "fetch", "args": {"url": _MANAGED_UPLOAD}}}),
|
||||
id="gemini-function-call-snake-case",
|
||||
),
|
||||
pytest.param(
|
||||
{
|
||||
"tools": [
|
||||
{
|
||||
"type": "function",
|
||||
"function": {
|
||||
"name": "fetch",
|
||||
"parameters": {
|
||||
"type": "object",
|
||||
"properties": {"url": {"type": "string", "enum": [_MANAGED_UPLOAD]}},
|
||||
"default": {"url": _MANAGED_UPLOAD},
|
||||
},
|
||||
},
|
||||
}
|
||||
]
|
||||
},
|
||||
id="tool-schema",
|
||||
),
|
||||
pytest.param(
|
||||
_user_turn({"type": "image_url", "image_url": {"url": "S3://media-bucket/litellm-bedrock-files/x.png"}}),
|
||||
id="uppercase-scheme",
|
||||
),
|
||||
pytest.param(
|
||||
_user_turn(
|
||||
{"type": "file", "file": {"file_id": "s3%3A%2F%2Fmedia-bucket%2Flitellm-bedrock-files%2Fx.pdf"}}
|
||||
),
|
||||
id="encoded-scheme",
|
||||
),
|
||||
pytest.param(
|
||||
_user_turn(
|
||||
{"type": "file", "file": {"file_id": "%73%33%3A%2F%2Fmedia-bucket%2Flitellm-bedrock-files%2Fx.pdf"}}
|
||||
),
|
||||
id="fully-encoded-scheme",
|
||||
),
|
||||
),
|
||||
)
|
||||
def test_s3_media_uris_ignores_what_never_becomes_an_s3_media_source(request_body: dict[str, object]):
|
||||
assert s3_media_uris(request_body) == ()
|
||||
|
||||
|
||||
def test_s3_media_uris_names_each_s3_object_once_in_the_order_it_first_appears():
|
||||
other_upload: Final = "s3://media-bucket/dir/cat.png"
|
||||
request_body: Final = _user_turn(
|
||||
{"type": "image_url", "image_url": {"url": other_upload}},
|
||||
{"type": "file", "file": {"file_id": _MANAGED_UPLOAD}},
|
||||
{"type": "image_url", "image_url": {"url": other_upload}},
|
||||
)
|
||||
|
||||
assert s3_media_uris(request_body) == (other_upload, _MANAGED_UPLOAD)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("uri", "names_an_unlisted_bucket"),
|
||||
(
|
||||
|
|
|
|||
|
|
@ -4487,6 +4487,10 @@ _CHAT_BATCH_LINE_WITH_AN_S3_IMAGE: Final = _batch_line_with_body(
|
|||
"messages",
|
||||
[{"role": "user", "content": [{"type": "image_url", "image_url": {"url": _BATCH_S3_IMAGE}}]}],
|
||||
)
|
||||
_S3_IMAGE_PART: Final = {"type": "image_url", "image_url": {"url": _BATCH_S3_IMAGE}}
|
||||
_CHAT_BATCH_LINE_WITH_THE_SAME_S3_IMAGE_TWICE: Final = _batch_line_with_body(
|
||||
"messages", [{"role": "user", "content": [_S3_IMAGE_PART, _S3_IMAGE_PART]}]
|
||||
)
|
||||
_RESPONSES_BATCH_LINE_WITH_AN_S3_IMAGE: Final = json.dumps(
|
||||
{
|
||||
"custom_id": "req-2",
|
||||
|
|
@ -4551,6 +4555,11 @@ _RESPONSES_BATCH_LINE_WITH_AN_S3_IMAGE: Final = json.dumps(
|
|||
_BATCH_S3_MEDIA_REFUSAL,
|
||||
id="responses-s3-image",
|
||||
),
|
||||
pytest.param(
|
||||
f"{_VALID_BATCH_LINE}\n{_CHAT_BATCH_LINE_WITH_THE_SAME_S3_IMAGE_TWICE}\n".encode(),
|
||||
_BATCH_S3_MEDIA_REFUSAL,
|
||||
id="chat-repeated-s3-image",
|
||||
),
|
||||
pytest.param(
|
||||
f"{_VALID_BATCH_LINE}\n{'[' * 1_000_000}{']' * 1_000_000}\n".encode(),
|
||||
"Batch input file line 2 could not be parsed",
|
||||
|
|
@ -4584,6 +4593,35 @@ async def test_create_file_rejects_a_malformed_batch_line_with_a_400_before_uplo
|
|||
assert not put_route.called
|
||||
|
||||
|
||||
@pytest.mark.usefixtures(httpx_transport.__name__)
|
||||
@pytest.mark.asyncio
|
||||
async def test_create_file_uploads_s3_values_that_no_converse_record_sends_as_an_s3_source(
|
||||
respx_mock: respx.MockRouter,
|
||||
) -> None:
|
||||
put_route: Final = respx_mock.put(url__regex=r"https://s3\..+").respond(status_code=200)
|
||||
passthrough_messages: Final = [{"role": "user", "content": [_S3_IMAGE_PART]}]
|
||||
passthrough_line: Final = _batch_line(
|
||||
{"model": "bedrock/openai.gpt-oss-120b-1:0", "max_tokens": 10, "messages": passthrough_messages}
|
||||
)
|
||||
metadata_only_line: Final = _batch_line_with_body("metadata", {"url": _BATCH_S3_IMAGE})
|
||||
|
||||
await litellm.acreate_file(
|
||||
file=("batch.jsonl", f"{passthrough_line}\n{metadata_only_line}\n".encode(), "application/jsonl"),
|
||||
purpose="batch",
|
||||
custom_llm_provider="bedrock",
|
||||
s3_bucket_name="litellm-batch-bucket",
|
||||
aws_region_name="us-west-2",
|
||||
aws_access_key_id="test-access-key",
|
||||
aws_secret_access_key="test-secret-key",
|
||||
)
|
||||
|
||||
passthrough_record, metadata_only_record = (
|
||||
json.loads(line) for line in put_route.calls.last.request.content.decode().split("\n")
|
||||
)
|
||||
assert passthrough_record["modelInput"]["messages"] == passthrough_messages
|
||||
assert metadata_only_record["recordId"] == _VALID_BATCH_RECORD["custom_id"]
|
||||
|
||||
|
||||
@pytest.mark.usefixtures(httpx_transport.__name__)
|
||||
@pytest.mark.asyncio
|
||||
async def test_create_file_accepts_a_batch_file_with_a_utf8_byte_order_mark(respx_mock: respx.MockRouter) -> None:
|
||||
|
|
|
|||
|
|
@ -253,6 +253,142 @@ _S3_MEDIA_SURFACES: Final = (
|
|||
)
|
||||
|
||||
|
||||
_S3_OBJECT: Final = "s3://media-bucket/dir/cat.png"
|
||||
_HELLO: Final = {"role": "user", "content": "hello"}
|
||||
_S3_VALUES_OUTSIDE_MEDIA_PARTS: Final = (
|
||||
pytest.param(
|
||||
"acompletion",
|
||||
"/v1/chat/completions",
|
||||
{"messages": [_HELLO], "metadata": {"url": _S3_OBJECT}},
|
||||
id="chat-metadata-url",
|
||||
),
|
||||
pytest.param(
|
||||
"acompletion",
|
||||
"/v1/chat/completions",
|
||||
{"messages": [_HELLO], "extra_body": {"url": _S3_OBJECT}},
|
||||
id="chat-extra-body-url",
|
||||
),
|
||||
pytest.param(
|
||||
"acompletion",
|
||||
"/v1/chat/completions",
|
||||
{
|
||||
"messages": [
|
||||
_HELLO,
|
||||
{
|
||||
"role": "assistant",
|
||||
"content": "See the source.",
|
||||
"annotations": [{"type": "url_citation", "url_citation": {"url": _S3_OBJECT}}],
|
||||
},
|
||||
_HELLO,
|
||||
]
|
||||
},
|
||||
id="chat-url-citation",
|
||||
),
|
||||
pytest.param(
|
||||
"acompletion",
|
||||
"/v1/chat/completions",
|
||||
{
|
||||
"messages": [_HELLO],
|
||||
"tools": [
|
||||
{
|
||||
"type": "function",
|
||||
"function": {
|
||||
"name": "fetch",
|
||||
"parameters": {
|
||||
"type": "object",
|
||||
"properties": {"url": {"type": "string", "enum": [_S3_OBJECT]}},
|
||||
"default": {"url": _S3_OBJECT},
|
||||
},
|
||||
},
|
||||
}
|
||||
],
|
||||
},
|
||||
id="chat-tool-schema",
|
||||
),
|
||||
pytest.param(
|
||||
"aresponses",
|
||||
"/v1/responses",
|
||||
{
|
||||
"input": [
|
||||
{
|
||||
"type": "message",
|
||||
"role": "assistant",
|
||||
"content": [
|
||||
{
|
||||
"type": "output_text",
|
||||
"text": "See the file.",
|
||||
"annotations": [{"type": "file_citation", "file_id": _S3_OBJECT, "index": 0}],
|
||||
}
|
||||
],
|
||||
},
|
||||
_HELLO,
|
||||
]
|
||||
},
|
||||
id="responses-file-citation",
|
||||
),
|
||||
pytest.param(
|
||||
"agenerate_content",
|
||||
_GEMINI_GENERATE_CONTENT_PATH,
|
||||
{
|
||||
"contents": [
|
||||
{"role": "user", "parts": [{"text": "hello"}]},
|
||||
{"role": "model", "parts": [{"functionCall": {"name": "fetch", "args": {"url": _S3_OBJECT}}}]},
|
||||
{"role": "user", "parts": [{"functionResponse": {"name": "fetch", "response": {"ok": True}}}]},
|
||||
]
|
||||
},
|
||||
id="gemini-function-call-camel-case",
|
||||
),
|
||||
pytest.param(
|
||||
"agenerate_content",
|
||||
_GEMINI_GENERATE_CONTENT_PATH,
|
||||
{
|
||||
"contents": [
|
||||
{"role": "user", "parts": [{"text": "hello"}]},
|
||||
{"role": "model", "parts": [{"function_call": {"name": "fetch", "args": {"url": _S3_OBJECT}}}]},
|
||||
{"role": "user", "parts": [{"function_response": {"name": "fetch", "response": {"ok": True}}}]},
|
||||
]
|
||||
},
|
||||
id="gemini-function-call-snake-case",
|
||||
),
|
||||
)
|
||||
_S3_VALUES_ON_NON_MEDIA_ROUTES: Final = (
|
||||
pytest.param(
|
||||
"call_mcp_tool",
|
||||
"/mcp-rest/tools/call",
|
||||
{"name": "fetch", "arguments": {"url": _S3_OBJECT}},
|
||||
id="mcp-tool-arguments",
|
||||
),
|
||||
pytest.param("afile_retrieve", "/v1/files/s3-object", {"file_id": _S3_OBJECT}, id="file-retrieve"),
|
||||
pytest.param("afile_content", "/v1/files/s3-object/content", {"file_id": _S3_OBJECT}, id="file-content"),
|
||||
pytest.param(
|
||||
"avector_store_file_create",
|
||||
"/v1/vector_stores/vs_1/files",
|
||||
{"vector_store_id": "vs_1", "file_id": _S3_OBJECT},
|
||||
id="vector-store-file-create",
|
||||
),
|
||||
pytest.param(
|
||||
"asend_message",
|
||||
"/a2a/agent-1/message/send",
|
||||
{
|
||||
"jsonrpc": "2.0",
|
||||
"id": "1",
|
||||
"method": "message/send",
|
||||
"params": {
|
||||
"message": {
|
||||
"role": "user",
|
||||
"messageId": "m1",
|
||||
"parts": [
|
||||
{"kind": "data", "data": {"file_url": _S3_OBJECT}},
|
||||
{"kind": "file", "file": {"uri": _S3_OBJECT, "mimeType": "image/png"}},
|
||||
],
|
||||
}
|
||||
},
|
||||
},
|
||||
id="a2a-data-and-file-parts",
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
async def _send_s3_media(
|
||||
route_type: _S3MediaRoute,
|
||||
path: str,
|
||||
|
|
@ -650,6 +786,57 @@ class TestProxyBaseLLMRequestProcessing:
|
|||
|
||||
assert "a cat" in str(response)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize(("route_type", "path", "body"), _S3_VALUES_OUTSIDE_MEDIA_PARTS)
|
||||
async def test_non_admin_s3_values_outside_media_parts_pass_with_no_bucket_listed(
|
||||
self, route_type: _S3MediaRoute, path: str, body: dict[str, object]
|
||||
) -> None:
|
||||
response: Final = await _send_s3_media(route_type, path, body, LitellmUserRoles.INTERNAL_USER, {})
|
||||
|
||||
assert "a cat" in str(response)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize(("route_type", "path", "body"), _S3_VALUES_ON_NON_MEDIA_ROUTES)
|
||||
async def test_non_admin_s3_values_on_routes_without_media_parts_reach_the_route(
|
||||
self, route_type: str, path: str, body: dict[str, object]
|
||||
) -> None:
|
||||
from litellm.proxy.proxy_server import ProxyConfig as RealProxyConfig
|
||||
|
||||
data, _ = await ProxyBaseLLMRequestProcessing(data=copy.deepcopy(body)).common_processing_pre_call_logic(
|
||||
request=Request({"type": "http", "method": "POST", "path": path, "query_string": b"", "headers": []}),
|
||||
general_settings={},
|
||||
user_api_key_dict=ProxyUserAPIKeyAuth(api_key="sk-test", user_role=LitellmUserRoles.INTERNAL_USER),
|
||||
proxy_logging_obj=ProxyLogging(user_api_key_cache=UserApiKeyCache()),
|
||||
proxy_config=RealProxyConfig(),
|
||||
route_type=route_type,
|
||||
)
|
||||
|
||||
assert {key: data[key] for key in body} == body
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_non_admin_s3_media_refusal_names_a_repeated_object_once(self) -> None:
|
||||
with pytest.raises(HTTPException) as refused:
|
||||
await _send_s3_media(
|
||||
"acompletion",
|
||||
"/v1/chat/completions",
|
||||
{
|
||||
"messages": [
|
||||
{
|
||||
"role": "user",
|
||||
"content": [
|
||||
{"type": "image_url", "image_url": {"url": _S3_OBJECT}},
|
||||
{"type": "image_url", "image_url": {"url": _S3_OBJECT}},
|
||||
],
|
||||
}
|
||||
]
|
||||
},
|
||||
LitellmUserRoles.INTERNAL_USER,
|
||||
{},
|
||||
)
|
||||
|
||||
assert refused.value.status_code == status.HTTP_403_FORBIDDEN
|
||||
assert refused.value.detail.count(_S3_OBJECT) == 1
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize(
|
||||
("path", "route_type", "has_router", "priority_header", "body_priority", "expected_priority"),
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue