scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-11-14 09:03:16 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 214

fix(bedrock): check every file id field, skip MCP tool calls, accept registered S3 media types

Browse source
This commit is contained in:
mateo-berri 2026-11-10 20:08:21 -08:00
parent f923c4da26
commit 2b13898be4
6 changed files with 137 additions and 13 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

7
litellm/litellm_core_utils/cloud_storage_security.py
View file

@ -83,7 +83,8 @@ _MEDIA_URI_KEYS: Final = frozenset({"url", "image_url", "file_id", "file_url", "
_ENCODED_FILE_ID_PREFIXES: Final = ("file-", "batch_")
_EMBEDDED_S3_URI: Final = re.compile(r"s3://[^;]*")
_MODEL_FILE_ID_MAPPING_KEY: Final = "model_file_id_mapping"
_KEYS_THAT_MAY_HIDE_AN_S3_URI: Final = frozenset({"file_id", _MODEL_FILE_ID_MAPPING_KEY})
_KEYS_THAT_MAY_HOLD_A_FILE_ID: Final = frozenset({"file_id", "file_url", "file_uri", "fileUri"})
_KEYS_THAT_MAY_HIDE_AN_S3_URI: Final = _KEYS_THAT_MAY_HOLD_A_FILE_ID | {_MODEL_FILE_ID_MAPPING_KEY}
def is_s3_uri(value: str) -> bool:
@ -145,13 +146,13 @@ def _string_leaf(node: _BodyNode) -> str | None:
def _media_file_id(node: _BodyNode) -> str | None:
return _string_leaf(node) if node.inside_media_part and node.key == "file_id" else None
return _string_leaf(node) if node.inside_media_part and node.key in _KEYS_THAT_MAY_HOLD_A_FILE_ID else None
def _media_uri_candidates(trusted_file_ids: frozenset[str], node: _BodyNode) -> Iterator[str]:
if not node.inside_media_part or not isinstance(node.value, str):
return iter(())
if node.key == "file_id":
if node.key in _KEYS_THAT_MAY_HOLD_A_FILE_ID:
return iter((node.value,)) if node.value in trusted_file_ids else _with_embedded_s3_uris(node.value)
return iter((node.value,) if node.key in _MEDIA_URI_KEYS else ())

18
litellm/litellm_core_utils/prompt_templates/factory.py
View file

@ -2411,6 +2411,10 @@ def anthropic_process_openai_file_message(
raise Exception(f"Either file_data or file_id must be present in the file message: {message}")
_S3_MIME_SUBTYPE_FORMATS: Final = MappingProxyType(
{"jpg": "jpeg", "quicktime": "mov", "x-matroska": "mkv", "x-flv": "flv", "x-ms-wmv": "wmv", "3gpp": "3gp"}
)
_CONVERSE_VIDEO_FORMAT_NAMES: Final = MappingProxyType({"3gp": "three_gp"})
_EMPTY_TEXT_PLACEHOLDER: Final = "[System: Empty message content sanitised to satisfy protocol]"
@ -4238,7 +4242,11 @@ class BedrockImageProcessor:
)
)
elif is_video:
return BedrockContentBlock(video=BedrockVideoBlock(source=source, format=image_format))
return BedrockContentBlock(
video=BedrockVideoBlock(
source=source, format=_CONVERSE_VIDEO_FORMAT_NAMES.get(image_format, image_format)
)
)
else:
return BedrockContentBlock(image=BedrockImageBlock(source=source, format=image_format))
@ -4259,14 +4267,15 @@ class BedrockImageProcessor:
@classmethod
def _create_s3_block(cls, uri: str, format: str | None, bucket_owner: str | None) -> BedrockContentBlock:
mime_type: Final = format or cls._s3_key_mime_type(uri)
if mime_type is None:
declared_type: Final = format or cls._s3_key_mime_type(uri)
if declared_type is None:
raise litellm.BadRequestError(
message=f"Cannot tell the media type of {uri} from its file extension. "
"Set the part's `format` to its MIME type, e.g. `image/png`.",
model="",
llm_provider="bedrock",
)
mime_type: Final = declared_type.partition(";")[0].strip().lower()
if mime_type.startswith("audio/"):
raise litellm.BadRequestError(
message=f"Bedrock Converse does not accept {uri} as {mime_type}. "
@ -4275,7 +4284,8 @@ class BedrockImageProcessor:
llm_provider="bedrock",
)
try:
image_format: Final = cls.validate_format(mime_type, mime_type.partition("/")[2])
subtype: Final = mime_type.partition("/")[2]
image_format: Final = cls.validate_format(mime_type, _S3_MIME_SUBTYPE_FORMATS.get(subtype, subtype))
except ValueError as e:
raise litellm.BadRequestError(
message=f"Bedrock Converse does not accept {uri} as {mime_type}. {e}",

13
litellm/proxy/common_request_processing.py
View file

@ -2734,12 +2734,13 @@ class ProxyBaseLLMRequestProcessing:
param="model",
code=status.HTTP_400_BAD_REQUEST,
)
await refuse_s3_media_the_key_may_not_read(
raw_request_fields,
user_api_key_dict,
_REQUEST_FIELDS.validate_python(general_settings),
proxy_logging_obj.get_proxy_hook("managed_files"),
)
if route_type != "call_mcp_tool":
await refuse_s3_media_the_key_may_not_read(
raw_request_fields,
user_api_key_dict,
_REQUEST_FIELDS.validate_python(general_settings),
proxy_logging_obj.get_proxy_hook("managed_files"),
)
self.data = await add_litellm_data_to_request(
data=self.data,
request=request,

47
tests/unit/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_factory.py
View file

@ -3406,6 +3406,53 @@ async def test_bedrock_converse_s3_media_is_sent_as_an_s3_location_without_a_dow
assert await _converse_user_content(part, run_async) == [expected]
@pytest.mark.parametrize("run_async", [False, True], ids=["sync", "async"])
@pytest.mark.parametrize(
("declared", "block", "converse_format"),
[
pytest.param("image/jpg", "image", "jpeg", id="image-jpg"),
pytest.param("IMAGE/PNG", "image", "png", id="upper-case"),
pytest.param("image/png; charset=binary", "image", "png", id="image-with-parameter"),
pytest.param("video/quicktime", "video", "mov", id="quicktime"),
pytest.param("video/x-matroska", "video", "mkv", id="matroska"),
pytest.param("video/x-flv", "video", "flv", id="flv"),
pytest.param("video/x-ms-wmv", "video", "wmv", id="wmv"),
pytest.param("video/3gpp", "video", "three_gp", id="3gpp"),
],
)
async def test_bedrock_converse_s3_media_format_accepts_a_registered_mime_type(
declared: str, block: str, converse_format: str, run_async: bool
):
uri: Final = "s3://media-bucket/uploads/7f3a9c"
content: Final = await _converse_user_content(
{"type": "image_url", "image_url": {"url": uri, "format": declared}}, run_async
)
assert content == [{block: {"format": converse_format, "source": _s3_source(uri)}}]
@pytest.mark.parametrize("run_async", [False, True], ids=["sync", "async"])
async def test_bedrock_converse_s3_3gp_key_uses_the_converse_format_name(run_async: bool):
uri: Final = "s3://media-bucket/clips/demo.3gp"
content: Final = await _converse_user_content({"type": "image_url", "image_url": {"url": uri}}, run_async)
assert content == [{"video": {"format": "three_gp", "source": _s3_source(uri)}}]
@pytest.mark.parametrize("run_async", [False, True], ids=["sync", "async"])
async def test_bedrock_converse_s3_text_with_a_charset_is_a_txt_document(run_async: bool):
uri: Final = "s3://media-bucket/docs/7f3a9c"
content: Final = await _converse_user_content(
{"type": "file", "file": {"file_id": uri, "format": "text/plain; charset=utf-8"}}, run_async
)
assert content[0]["document"]["format"] == "txt"
assert content[0]["document"]["source"] == _s3_source(uri)
@pytest.mark.parametrize("run_async", [False, True], ids=["sync", "async"])
async def test_bedrock_converse_s3_pdf_is_a_document_named_after_its_uri(run_async: bool):
report: Final = "s3://media-bucket/docs/report.pdf"

21
tests/unit/litellm_core_utils/test_cloud_storage_security.py
View file

@ -115,6 +115,18 @@ def _gemini_turn(*parts: object) -> dict[str, object]:
_responses_turn({"type": "input_file", "file_id": _UNIFIED_BATCH_OUTPUT_ID}),
id="responses-unified-file-id",
),
pytest.param(
_responses_turn({"type": "input_file", "file_url": _UNIFIED_BATCH_OUTPUT_ID}),
id="responses-unified-file-id-in-file-url",
),
pytest.param(
_gemini_turn({"file_data": {"mime_type": "application/pdf", "file_uri": _UNIFIED_BATCH_OUTPUT_ID}}),
id="gemini-unified-file-id-in-file-uri",
),
pytest.param(
_gemini_turn({"fileData": {"mimeType": "application/pdf", "fileUri": _MODEL_EMBEDDED_UPLOAD_ID}}),
id="gemini-model-embedded-file-id-in-camel-case-file-uri",
),
pytest.param(
{
**_user_turn({"type": "file", "file": {"file_id": "alias"}}),
@ -393,6 +405,15 @@ def test_media_file_ids_lists_the_file_ids_of_media_parts_only():
assert media_file_ids(request_body) == ("file-a",)
def test_media_file_ids_lists_every_field_a_bridge_reads_as_a_file_id():
request_body: Final = {
**_responses_turn({"type": "input_file", "file_url": "file-a"}),
**_gemini_turn({"file_data": {"file_uri": "file-b"}}, {"fileData": {"fileUri": "file-c"}}),
}
assert set(media_file_ids(request_body)) == {"file-a", "file-b", "file-c"}
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(

44
tests/unit/proxy/test_common_request_processing.py
View file

@ -270,9 +270,21 @@ def _responses_file(file_id: str) -> dict[str, object]:
return {"input": [{"role": "user", "content": [{"type": "input_file", "file_id": file_id}]}]}
def _responses_file_url(file_id: str) -> dict[str, object]:
return {"input": [{"role": "user", "content": [{"type": "input_file", "file_url": file_id}]}]}
def _gemini_file_uri(file_id: str) -> dict[str, object]:
return {
"contents": [{"role": "user", "parts": [{"file_data": {"mime_type": "application/pdf", "file_uri": file_id}}]}]
}
_S3_FILE_ID_SURFACES: Final = (
pytest.param("acompletion", "/v1/chat/completions", _chat_file, id="chat-completions"),
pytest.param("aresponses", "/v1/responses", _responses_file, id="responses"),
pytest.param("aresponses", "/v1/responses", _responses_file_url, id="responses-file-url"),
pytest.param("agenerate_content", _GEMINI_GENERATE_CONTENT_PATH, _gemini_file_uri, id="gemini-file-uri"),
)
_UPLOADED_OBJECT: Final = "s3://media-bucket/litellm-bedrock-files/upload.pdf"
_UPLOADED_FILE_ID: Final = (
@ -513,6 +525,10 @@ async def _send_s3_media(
)
class _StopAfterTheS3Check(Exception):
pass
class TestProxyBaseLLMRequestProcessing:
@pytest.mark.asyncio
async def test_base_passthrough_process_llm_request_preserves_litellm_headers_for_non_streaming_response(
@ -775,6 +791,34 @@ class TestProxyBaseLLMRequestProcessing:
pytest.fail("litellm_call_id is not a valid UUID")
assert data_passed["litellm_call_id"] == returned_data["litellm_call_id"]
@pytest.mark.asyncio
@pytest.mark.parametrize(
("route_type", "refused"), [("call_mcp_tool", False), ("acompletion", True)], ids=["mcp-tool", "chat"]
)
async def test_s3_media_check_covers_model_requests_and_leaves_mcp_tool_arguments_alone(
self, monkeypatch, route_type: Literal["call_mcp_tool", "acompletion"], refused: bool
):
tool_argument: Final = {"type": "image_url", "image_url": {"url": "s3://other-bucket/cat.png"}}
processing_obj: Final = ProxyBaseLLMRequestProcessing(data={"name": "describe", "arguments": tool_argument})
mock_request: Final = MagicMock(spec=Request)
mock_request.headers = {}
reached_request_setup: Final = AsyncMock(side_effect=_StopAfterTheS3Check)
monkeypatch.setattr(
litellm.proxy.common_request_processing, "add_litellm_data_to_request", reached_request_setup
)
with pytest.raises(HTTPException if refused else _StopAfterTheS3Check):
await processing_obj.common_processing_pre_call_logic(
request=mock_request,
general_settings={},
user_api_key_dict=ProxyUserAPIKeyAuth(
api_key="sk-test", user_id="user-a", user_role=LitellmUserRoles.INTERNAL_USER
),
proxy_logging_obj=ProxyLogging(user_api_key_cache=UserApiKeyCache()),
proxy_config=MagicMock(spec=ProxyConfig),
route_type=route_type,
)
@pytest.mark.asyncio
@pytest.mark.parametrize("requested_model", [{"bad": "value"}, ["gpt-5.2"], 1])
async def test_common_processing_pre_call_logic_rejects_a_non_string_model_with_400(