diff --git a/litellm/litellm_core_utils/cloud_storage_security.py b/litellm/litellm_core_utils/cloud_storage_security.py index 88ddb14f498..e06cadcaa93 100644 --- a/litellm/litellm_core_utils/cloud_storage_security.py +++ b/litellm/litellm_core_utils/cloud_storage_security.py @@ -44,8 +44,7 @@ def is_managed_cloud_storage_uri(file_id: str) -> bool: _S3_SCHEME: Final = "s3://" _FIRST_CHARACTERS_THAT_CAN_DECODE_TO_S3: Final = ("s", "%") -_CLIENT_MEDIA_FIELDS: Final = ("messages", "input", "system") -_MEDIA_URI_KEYS: Final = frozenset({"url", "image_url", "file_id", "file_url", "file_data"}) +_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]: @@ -89,8 +88,7 @@ def _is_s3_uri(uri: str) -> bool: def s3_media_uris(request_body: Mapping[str, object]) -> tuple[str, ...]: - media_uris: Final = chain.from_iterable(_media_uris(request_body.get(field)) for field in _CLIENT_MEDIA_FIELDS) - return tuple(filter(_is_s3_uri, media_uris)) + return tuple(filter(_is_s3_uri, _media_uris(request_body))) def is_managed_s3_uri(uri: str) -> bool: diff --git a/litellm/llms/bedrock/files/transformation.py b/litellm/llms/bedrock/files/transformation.py index 4d40e0dec85..3ab84c6d58c 100644 --- a/litellm/llms/bedrock/files/transformation.py +++ b/litellm/llms/bedrock/files/transformation.py @@ -32,6 +32,7 @@ from litellm.litellm_core_utils.cloud_storage_security import ( BEDROCK_MANAGED_S3_UPLOAD_PREFIX, build_managed_cloud_object_name, encode_s3_object_key_for_url, + s3_media_uris, sanitize_cloud_object_component, should_allow_legacy_cloud_file_ids, split_configured_cloud_bucket_name, @@ -321,8 +322,7 @@ def _validate_file_id_against_configured_buckets( _REJECTED_FILE_ID_REQUEST_URL: Final = "https://litellm.ai" -def _rejected_file_id(reason: ValueError) -> BedrockError: - message: Final = str(reason) +def _rejected_request(message: str) -> BedrockError: return BedrockError( status_code=400, message=message, @@ -334,6 +334,29 @@ def _rejected_file_id(reason: ValueError) -> BedrockError: ) +def _rejected_file_id(reason: ValueError) -> BedrockError: + return _rejected_request(str(reason)) + + +def _batch_line_error(line_number: int, problem: str) -> BedrockError: + return _rejected_request(f"Batch input file line {line_number} {problem}") + + +def _json_object(value: object) -> Mapping[str, object] | None: + return value if isinstance(value, dict) else None + + +def _batch_body_problem(body: object) -> str | None: + json_body: Final = _json_object(body) + s3_uris: Final = s3_media_uris(json_body) if json_body is not None else () + if not s3_uris: + return None + 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" + ) + + def _resolve_managed_s3_object(file_id: str, litellm_params: Mapping[str, object]) -> tuple[str, str]: configured_bucket_names: Final = get_configured_s3_bucket_names(litellm_params) allow_legacy_cloud_file_ids: Final = should_allow_legacy_cloud_file_ids(litellm_params) @@ -1083,6 +1106,8 @@ class BedrockFilesConfig(BaseAWSLLM, BaseFilesConfig): for idx, _openai_jsonl_content in enumerate(openai_jsonl_content): # Extract the request body from OpenAI format openai_body = _openai_jsonl_content.get("body", {}) + if (body_problem := _batch_body_problem(openai_body)) is not None: + raise _batch_line_error(idx + 1, body_problem) record_model = openai_body.get("model", "") resolved_model = litellm.model_alias_map.get(record_model, record_model) model_for_transform, provider = self._resolve_batch_record_model_and_provider( 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 ea4953cb7c8..e476650c0c8 100644 --- a/tests/unit/litellm_core_utils/test_cloud_storage_security.py +++ b/tests/unit/litellm_core_utils/test_cloud_storage_security.py @@ -44,6 +44,10 @@ def _responses_turn(*parts: object) -> dict[str, object]: return {"input": [{"role": "user", "content": list(parts)}]} +def _gemini_turn(*parts: object) -> dict[str, object]: + return {"contents": [{"role": "user", "parts": list(parts)}]} + + @pytest.mark.parametrize( "request_body", ( @@ -89,6 +93,19 @@ def _responses_turn(*parts: object) -> dict[str, object]: {"system": [{"type": "image", "source": {"type": "url", "url": _MANAGED_UPLOAD}}], "messages": []}, id="anthropic-system", ), + pytest.param( + _gemini_turn({"file_data": {"mime_type": "image/png", "file_uri": _MANAGED_UPLOAD}}), + id="gemini-file-data-snake-case", + ), + pytest.param( + _gemini_turn({"fileData": {"mimeType": "image/png", "fileUri": _MANAGED_UPLOAD}}), + id="gemini-file-data-camel-case", + ), + pytest.param( + {"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"), ), ) def test_s3_media_uris_finds_an_s3_object_on_every_media_surface(request_body: dict[str, object]): @@ -151,7 +168,7 @@ def test_is_managed_s3_uri_leaves_other_objects_alone(uri: str): pytest.param({"messages": [{"role": "user", "content": _MANAGED_UPLOAD}]}, id="string-content"), pytest.param({"input": _MANAGED_UPLOAD}, id="responses-string-input"), pytest.param({"system": _MANAGED_UPLOAD, "messages": []}, id="string-system"), - pytest.param({"metadata": {"url": _MANAGED_UPLOAD}}, id="non-media-field"), + pytest.param({"metadata": {"note": _MANAGED_UPLOAD}}, id="non-media-key"), ), ) def test_s3_media_uris_ignores_what_never_becomes_an_s3_media_source(request_body: dict[str, object]): 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 425beefac8a..0ec966dfb05 100644 --- a/tests/unit/llms/bedrock/files/test_bedrock_files_transformation.py +++ b/tests/unit/llms/bedrock/files/test_bedrock_files_transformation.py @@ -22,6 +22,7 @@ from botocore.credentials import Credentials import litellm from litellm.constants import DEFAULT_REASONING_EFFORT_LOW_THINKING_BUDGET +from litellm.llms.bedrock.common_utils import BedrockError from litellm.llms.bedrock.files.transformation import BedrockJsonlFilesTransformation from tests.unit.proxy.conftest import httpx_transport @@ -4334,3 +4335,74 @@ async def test_async_create_file( "s3://litellm-proxy-941277531214/litellm-bedrock-files-" ) assert file_obj.filename.endswith(".jsonl") + + +_VALID_BATCH_LINE: Final = json.dumps( + { + "custom_id": "req-1", + "method": "POST", + "url": "/v1/chat/completions", + "body": { + "model": "bedrock/amazon.nova-pro-v1:0", + "messages": [{"role": "user", "content": "Hello"}], + "max_tokens": 10, + }, + } +) +_BATCH_S3_IMAGE: Final = "s3://media-bucket/dir/cat.png" +_BATCH_S3_MEDIA_REFUSAL: Final = ( + f"Batch input file line 2 has s3:// media URLs ({_BATCH_S3_IMAGE}). " + "LiteLLM does not accept s3:// media URLs in Bedrock batch files, so send the media inline as base64 data" +) +_CHAT_BATCH_LINE_WITH_AN_S3_IMAGE: Final = json.dumps( + { + "custom_id": "req-2", + "method": "POST", + "url": "/v1/chat/completions", + "body": { + "model": "bedrock/amazon.nova-pro-v1:0", + "messages": [{"role": "user", "content": [{"type": "image_url", "image_url": {"url": _BATCH_S3_IMAGE}}]}], + }, + } +) +_RESPONSES_BATCH_LINE_WITH_AN_S3_IMAGE: Final = json.dumps( + { + "custom_id": "req-2", + "method": "POST", + "url": "/v1/responses", + "body": { + "model": "bedrock/amazon.nova-pro-v1:0", + "input": [{"role": "user", "content": [{"type": "input_image", "image_url": _BATCH_S3_IMAGE}]}], + }, + } +) + + +@pytest.mark.usefixtures(httpx_transport.__name__) +@pytest.mark.asyncio +@pytest.mark.parametrize( + "s3_media_line", + ( + pytest.param(_CHAT_BATCH_LINE_WITH_AN_S3_IMAGE, id="chat-s3-image"), + pytest.param(_RESPONSES_BATCH_LINE_WITH_AN_S3_IMAGE, id="responses-s3-image"), + ), +) +async def test_create_file_refuses_a_batch_line_with_s3_media_with_a_400_before_uploading( + respx_mock: respx.MockRouter, s3_media_line: str +) -> None: + put_route: Final = respx_mock.put(url__regex=r"https://s3\..+").respond(status_code=200) + + with pytest.raises(BedrockError) as exc_info: + await litellm.acreate_file( + file=("batch.jsonl", f"{_VALID_BATCH_LINE}\n{s3_media_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", + ) + + assert exc_info.value.status_code == 400 + assert exc_info.value.message == _BATCH_S3_MEDIA_REFUSAL + assert not put_route.called diff --git a/tests/unit/proxy/test_common_request_processing.py b/tests/unit/proxy/test_common_request_processing.py index fd0aec87b2d..5c07835b3f9 100644 --- a/tests/unit/proxy/test_common_request_processing.py +++ b/tests/unit/proxy/test_common_request_processing.py @@ -4,7 +4,7 @@ import datetime import json from collections.abc import AsyncIterable, Mapping from types import MappingProxyType, SimpleNamespace -from typing import AsyncGenerator, Callable, Final, Iterator, Literal, Optional, Sequence +from typing import AsyncGenerator, Callable, Final, Iterator, Literal, Optional, Sequence, TypeAlias from urllib.parse import unquote_plus from unittest.mock import AsyncMock, MagicMock, patch @@ -220,15 +220,27 @@ def _anthropic_image(uri: str) -> dict[str, object]: } +def _gemini_snake_case_image(uri: str) -> dict[str, object]: + return {"contents": [{"role": "user", "parts": [{"file_data": {"mime_type": "image/png", "file_uri": uri}}]}]} + + +def _gemini_camel_case_image(uri: str) -> dict[str, object]: + return {"contents": [{"role": "user", "parts": [{"fileData": {"mimeType": "image/png", "fileUri": uri}}]}]} + + +_S3MediaRoute: TypeAlias = Literal["acompletion", "aresponses", "anthropic_messages", "agenerate_content"] +_GEMINI_GENERATE_CONTENT_PATH: Final = "/v1beta/models/nova:generateContent" _S3_MEDIA_SURFACES: Final = ( pytest.param("acompletion", "/v1/chat/completions", _chat_image, id="chat-completions"), pytest.param("aresponses", "/v1/responses", _responses_image, id="responses"), pytest.param("anthropic_messages", "/v1/messages", _anthropic_image, id="anthropic-messages"), + pytest.param("agenerate_content", _GEMINI_GENERATE_CONTENT_PATH, _gemini_snake_case_image, id="gemini-snake-case"), + pytest.param("agenerate_content", _GEMINI_GENERATE_CONTENT_PATH, _gemini_camel_case_image, id="gemini-camel-case"), ) async def _send_s3_media( - route_type: Literal["acompletion", "aresponses", "anthropic_messages"], + route_type: _S3MediaRoute, path: str, body: dict[str, object], user_role: LitellmUserRoles, @@ -596,7 +608,7 @@ class TestProxyBaseLLMRequestProcessing: ) async def test_non_admin_s3_media_is_refused_while_a_proxy_admin_key_passes( self, - route_type: Literal["acompletion", "aresponses", "anthropic_messages"], + route_type: _S3MediaRoute, path: str, body_for: Callable[[str], dict[str, object]], uri: str, @@ -623,13 +635,17 @@ class TestProxyBaseLLMRequestProcessing: ) async def test_non_admin_s3_media_passes_in_a_bucket_the_admin_allowed( self, - route_type: Literal["acompletion", "aresponses", "anthropic_messages"], + route_type: _S3MediaRoute, path: str, body_for: Callable[[str], dict[str, object]], general_settings: dict[str, object], ) -> None: response: Final = await _send_s3_media( - route_type, path, body_for("s3://media-bucket/dir/cat.png"), LitellmUserRoles.INTERNAL_USER, general_settings + route_type, + path, + body_for("s3://media-bucket/dir/cat.png"), + LitellmUserRoles.INTERNAL_USER, + general_settings, ) assert "a cat" in str(response)