diff --git a/litellm/litellm_core_utils/cloud_storage_security.py b/litellm/litellm_core_utils/cloud_storage_security.py index f462421c2c3..ac4362f2b78 100644 --- a/litellm/litellm_core_utils/cloud_storage_security.py +++ b/litellm/litellm_core_utils/cloud_storage_security.py @@ -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: diff --git a/litellm/litellm_core_utils/prompt_templates/factory.py b/litellm/litellm_core_utils/prompt_templates/factory.py index 9a8e511295b..f531d4902f7 100644 --- a/litellm/litellm_core_utils/prompt_templates/factory.py +++ b/litellm/litellm_core_utils/prompt_templates/factory.py @@ -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}.", diff --git a/litellm/llms/bedrock/files/transformation.py b/litellm/llms/bedrock/files/transformation.py index d8c5b49d17a..fa4404375e7 100644 --- a/litellm/llms/bedrock/files/transformation.py +++ b/litellm/llms/bedrock/files/transformation.py @@ -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}" diff --git a/litellm/responses/litellm_completion_transformation/transformation.py b/litellm/responses/litellm_completion_transformation/transformation.py index 7de3cb19ec5..55b7c2f43d0 100644 --- a/litellm/responses/litellm_completion_transformation/transformation.py +++ b/litellm/responses/litellm_completion_transformation/transformation.py @@ -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 {} diff --git a/tests/unit/litellm_core_utils/test_cloud_storage_security.py b/tests/unit/litellm_core_utils/test_cloud_storage_security.py index 106f221a886..ba24233433c 100644 --- a/tests/unit/litellm_core_utils/test_cloud_storage_security.py +++ b/tests/unit/litellm_core_utils/test_cloud_storage_security.py @@ -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"), ( diff --git a/tests/unit/llms/bedrock/files/test_bedrock_files_transformation.py b/tests/unit/llms/bedrock/files/test_bedrock_files_transformation.py index f05c03b484a..44a0bd4a963 100644 --- a/tests/unit/llms/bedrock/files/test_bedrock_files_transformation.py +++ b/tests/unit/llms/bedrock/files/test_bedrock_files_transformation.py @@ -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: diff --git a/tests/unit/proxy/test_common_request_processing.py b/tests/unit/proxy/test_common_request_processing.py index ca6100f54c1..a4033385eb3 100644 --- a/tests/unit/proxy/test_common_request_processing.py +++ b/tests/unit/proxy/test_common_request_processing.py @@ -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"),