From 2c7b965ff4db4dd7e64b67107c704ea019e44009 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Thu, 5 Nov 2026 22:42:35 -0800 Subject: [PATCH] fix(proxy): limit non-admin s3:// media to allowed buckets and refuse managed prefixes A non-admin key may only send an s3:// media URL whose bucket is listed in general_settings.allowed_client_side_s3_buckets (or when client-side credentials are allowed), and never one under litellm's managed prefixes. Proxy admins are not limited. A bucket_owner on an s3:// URL that is not a non-empty string is a 400, the Responses bridge forwards the owner only for s3:// URLs, audio on S3 is a 400, and video in a tool result is sent to Bedrock. --- .../cloud_storage_security.py | 64 +++++++- .../prompt_templates/factory.py | 107 +++++++------ litellm/proxy/_lazy_openapi_snapshot.json | 8 - litellm/proxy/_types.py | 2 +- litellm/proxy/auth/auth_utils.py | 4 +- litellm/proxy/common_request_processing.py | 48 +++++- .../transformation.py | 42 ++--- litellm/types/llms/bedrock.py | 11 +- litellm/types/llms/openai.py | 2 - ...llm_core_utils_prompt_templates_factory.py | 92 ++++++++++- .../test_cloud_storage_security.py | 150 ++++++++++++++++++ .../chat/test_converse_transformation.py | 37 +++++ .../test_generic_guardrail_api.py | 84 ++++++++++ .../proxy/test_common_request_processing.py | 126 ++++++++++++++- .../test_litellm_completion_responses.py | 53 +++++++ ui/litellm-dashboard/src/lib/http/schema.d.ts | 6 +- 16 files changed, 743 insertions(+), 93 deletions(-) diff --git a/litellm/litellm_core_utils/cloud_storage_security.py b/litellm/litellm_core_utils/cloud_storage_security.py index 477775bb423..37c8efa0ac3 100644 --- a/litellm/litellm_core_utils/cloud_storage_security.py +++ b/litellm/litellm_core_utils/cloud_storage_security.py @@ -1,11 +1,12 @@ import posixpath import re -from collections.abc import Mapping, Sequence -from itertools import accumulate, repeat +from collections.abc import Callable, Iterator, Mapping, Sequence +from itertools import accumulate, chain, repeat from typing import Final from urllib.parse import quote, unquote from litellm._uuid import uuid +from litellm.litellm_core_utils.core_helpers import as_object_list, as_str_mapping VERTEX_AI_MANAGED_GCS_PREFIX: Final = "litellm-vertex-files/" BEDROCK_MANAGED_S3_BATCH_PREFIX: Final = "litellm-bedrock-files-" @@ -48,6 +49,65 @@ def is_managed_cloud_storage_uri(file_id: str) -> bool: return decodings[-1].startswith(MANAGED_CLOUD_STORAGE_SCHEMES) or not settled +_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"}) + + +def _media_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_uris(child, child_key) for child_key, child in fields.items()) + items: Final = as_object_list(value) + return chain.from_iterable(map(_media_uris, items)) if items is not None else iter(()) + + +def _s3_bucket(uri: str) -> str | None: + return uri.removeprefix(_S3_SCHEME).partition("/")[0] if uri.startswith(_S3_SCHEME) else None + + +def _s3_object_key(uri: str) -> str | None: + if not uri.startswith(_S3_SCHEME): + return None + _, separator, object_key = uri.removeprefix(_S3_SCHEME).partition("/") + return object_key if separator else None + + +def _outside_managed_prefixes(object_key: str) -> bool: + return not any(prefix in object_key for prefix in BEDROCK_MANAGED_S3_PREFIXES) + + +def _refused_in_some_decoding( + uri: str, s3_part: Callable[[str], str | None], is_allowed: Callable[[str], bool] +) -> bool: + decodings: Final = _decodings(uri) + parts: Final = tuple(part for part in map(s3_part, decodings) if part is not None) + settled: Final = decodings[-1] == decodings[-2] + 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, ...]: + 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)) + + +def is_managed_s3_uri(uri: str) -> bool: + return _refused_in_some_decoding(uri, _s3_object_key, _outside_managed_prefixes) + + +def names_an_unlisted_s3_bucket(uri: str, is_listed: Callable[[str], bool]) -> bool: + return _refused_in_some_decoding(uri, _s3_bucket, is_listed) + + _SAFE_OBJECT_COMPONENT_PATTERN: Final = re.compile(r"[^A-Za-z0-9._-]+") diff --git a/litellm/litellm_core_utils/prompt_templates/factory.py b/litellm/litellm_core_utils/prompt_templates/factory.py index d56175ee888..9a8e511295b 100644 --- a/litellm/litellm_core_utils/prompt_templates/factory.py +++ b/litellm/litellm_core_utils/prompt_templates/factory.py @@ -11,12 +11,11 @@ from dataclasses import dataclass, replace from enum import Enum from pathlib import PurePosixPath from types import MappingProxyType -from typing import Annotated, Any, Final, TypeAlias, TypeVar, cast, overload +from typing import Annotated, Any, Final, TypeAlias, TypedDict, TypeVar, cast, overload from jinja2.sandbox import ImmutableSandboxedEnvironment from pydantic import BaseModel, BeforeValidator, ConfigDict, JsonValue, TypeAdapter, ValidationError, model_validator from pydantic_core import ErrorDetails, PydanticCustomError -from typing_extensions import ReadOnly, TypedDict import litellm import litellm.types @@ -4215,12 +4214,11 @@ class BedrockImageProcessor: extension: Final = PurePosixPath(uri).suffix.removeprefix(".").lower() if not extension: return None - bedrock_format: Final = "jpeg" if extension == "jpg" else extension config: Final = litellm.AmazonConverseConfig() - if bedrock_format in config.get_supported_image_types(): - return f"image/{bedrock_format}" - if bedrock_format in config.get_supported_video_types(): - return f"video/{bedrock_format}" + if extension in config.get_supported_image_types(): + return f"image/{extension}" + if extension in config.get_supported_video_types(): + return f"video/{extension}" try: return get_file_mime_type_from_extension(extension) except ValueError: @@ -4236,6 +4234,13 @@ class BedrockImageProcessor: model="", llm_provider="bedrock", ) + if mime_type.startswith("audio/"): + raise litellm.BadRequestError( + message=f"Bedrock Converse does not accept {uri} as {mime_type}. " + "Audio is not one of its image, document, or video formats.", + model="", + llm_provider="bedrock", + ) try: image_format: Final = cls.validate_format(mime_type, mime_type.partition("/")[2]) except ValueError as e: @@ -4424,13 +4429,15 @@ def _convert_to_bedrock_tool_call_invoke( def _append_bedrock_tool_result_media_block( tool_result_content_blocks: list[BedrockToolResultContentBlock], processed_block: BedrockContentBlock, - content: dict, + content: Mapping[str, object], content_type: str, ) -> None: if "image" in processed_block: tool_result_content_blocks.append(BedrockToolResultContentBlock(image=processed_block["image"])) elif "document" in processed_block: tool_result_content_blocks.append(BedrockToolResultContentBlock(document=processed_block["document"])) + elif "video" in processed_block: + tool_result_content_blocks.append(BedrockToolResultContentBlock(video=processed_block["video"])) else: verbose_logger.warning( "Bedrock Converse: unrecognized BedrockContentBlock keys %s for %s tool-result block %s; dropping.", @@ -4440,21 +4447,23 @@ def _append_bedrock_tool_result_media_block( ) -class _S3MediaOwner(TypedDict, total=False): - bucket_owner: ReadOnly[object] - - -_S3_MEDIA_OWNER: Final = TypeAdapter(_S3MediaOwner) - - -def _s3_bucket_owner(media: _S3MediaOwner) -> str | None: - owner: Final = media.get("bucket_owner") - return owner if isinstance(owner, str) else None +def _s3_bucket_owner(media: object, url: str) -> str | None: + fields: Final = as_str_mapping(media) + 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://"): + 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}.", + model="", + llm_provider="bedrock", + ) def _append_bedrock_tool_result_image_url_block( tool_result_content_blocks: list[BedrockToolResultContentBlock], - content: dict, + content: ChatCompletionImageObject, ) -> None: format: str | None = None if isinstance(content["image_url"], dict): @@ -4462,26 +4471,21 @@ def _append_bedrock_tool_result_image_url_block( format = content["image_url"].get("format") else: image_url = content["image_url"] - bucket_owner: Final = ( - _s3_bucket_owner(_S3_MEDIA_OWNER.validate_python(content["image_url"])) - if isinstance(content["image_url"], dict) - else None - ) processed_block: Final = BedrockImageProcessor.process_image_sync( image_url=image_url, format=format, - bucket_owner=bucket_owner, + bucket_owner=_s3_bucket_owner(content["image_url"], image_url), ) _append_bedrock_tool_result_media_block(tool_result_content_blocks, processed_block, content, "image_url") def _append_bedrock_tool_result_file_block( tool_result_content_blocks: list[BedrockToolResultContentBlock], - content: dict, + content: ChatCompletionFileObject, ) -> None: # Match the user-message path (_process_file_message): accept either # file_data (base64 data URI) or file_id (server-side reference / URL). - file_obj: Final = content.get("file") or {} + file_obj: Final = content.get("file") or ChatCompletionFileObjectFile() file_data: Final = file_obj.get("file_data") file_id: Final = file_obj.get("file_id") if file_data is None and file_id is None: @@ -4490,10 +4494,11 @@ def _append_bedrock_tool_result_file_block( model="", llm_provider="bedrock", ) + url: Final = cast(str, file_id or file_data) processed_block: Final = BedrockImageProcessor.process_image_sync( - image_url=cast(str, file_id or file_data), + image_url=url, format=file_obj.get("format"), - bucket_owner=_s3_bucket_owner(_S3_MEDIA_OWNER.validate_python(file_obj)), + bucket_owner=_s3_bucket_owner(file_obj, url), ) _append_bedrock_tool_result_media_block(tool_result_content_blocks, processed_block, content, "file") @@ -5143,13 +5148,10 @@ class BedrockConverseMessagesProcessor: format = element["image_url"].get("format") else: image_url = element["image_url"] - bucket_owner: str | None = ( - element["image_url"].get("bucket_owner") - if isinstance(element["image_url"], dict) - else None - ) _part = await BedrockImageProcessor.process_image_async( - image_url=image_url, format=format, bucket_owner=bucket_owner + image_url=image_url, + format=format, + bucket_owner=_s3_bucket_owner(element["image_url"], image_url), ) _parts.append(_part) elif element["type"] == "file": @@ -5286,7 +5288,15 @@ class BedrockConverseMessagesProcessor: image_url = element["image_url"]["url"] else: image_url = element["image_url"] - assistants_part = await BedrockImageProcessor.process_image_async(image_url=image_url) + assistants_part = await BedrockImageProcessor.process_image_async( + image_url=image_url, + format=( + element["image_url"].get("format") + if isinstance(element["image_url"], dict) + else None + ), + bucket_owner=_s3_bucket_owner(element["image_url"], image_url), + ) assistants_parts.append(assistants_part) # Add cache point block for assistant content elements _cache_point_block = litellm.AmazonConverseConfig().get_cache_point_block( @@ -5416,8 +5426,9 @@ class BedrockConverseMessagesProcessor: llm_provider="bedrock", ) format: Final = file_message.get("format") + url: Final = cast(str, file_id or file_data) return BedrockImageProcessor.process_image_sync( - image_url=cast(str, file_id or file_data), format=format, bucket_owner=file_message.get("bucket_owner") + image_url=url, format=format, bucket_owner=_s3_bucket_owner(file_message, url) ) _process_file_message = process_file_message @@ -5442,8 +5453,9 @@ class BedrockConverseMessagesProcessor: model="", llm_provider="bedrock", ) + url: Final = cast(str, file_id or file_data) return await BedrockImageProcessor.process_image_async( - image_url=cast(str, file_id or file_data), format=format, bucket_owner=file_message.get("bucket_owner") + image_url=url, format=format, bucket_owner=_s3_bucket_owner(file_message, url) ) @staticmethod @@ -5573,13 +5585,10 @@ def _bedrock_converse_messages_pt( format = element["image_url"].get("format") else: image_url = element["image_url"] - bucket_owner: str | None = ( - element["image_url"].get("bucket_owner") - if isinstance(element["image_url"], dict) - else None - ) _part = BedrockImageProcessor.process_image_sync( - image_url=image_url, format=format, bucket_owner=bucket_owner + image_url=image_url, + format=format, + bucket_owner=_s3_bucket_owner(element["image_url"], image_url), ) _parts.append(_part) elif element["type"] == "file": @@ -5717,7 +5726,15 @@ def _bedrock_converse_messages_pt( image_url = element["image_url"]["url"] else: image_url = element["image_url"] - assistants_part = BedrockImageProcessor.process_image_sync(image_url=image_url) + assistants_part = BedrockImageProcessor.process_image_sync( + image_url=image_url, + format=( + element["image_url"].get("format") + if isinstance(element["image_url"], dict) + else None + ), + bucket_owner=_s3_bucket_owner(element["image_url"], image_url), + ) assistants_parts.append(assistants_part) # Add cache point block for assistant content elements _cache_point_block = litellm.AmazonConverseConfig().get_cache_point_block( diff --git a/litellm/proxy/_lazy_openapi_snapshot.json b/litellm/proxy/_lazy_openapi_snapshot.json index ac2a6179d8d..5cb63c98d56 100644 --- a/litellm/proxy/_lazy_openapi_snapshot.json +++ b/litellm/proxy/_lazy_openapi_snapshot.json @@ -49500,10 +49500,6 @@ }, "ChatCompletionFileObjectFile": { "properties": { - "bucket_owner": { - "title": "Bucket Owner", - "type": "string" - }, "detail": { "title": "Detail", "type": "string" @@ -49610,10 +49606,6 @@ }, "ChatCompletionImageUrlObject": { "properties": { - "bucket_owner": { - "title": "Bucket Owner", - "type": "string" - }, "detail": { "title": "Detail", "type": "string" diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index e53b5d4ad52..9ac6e40994c 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -2815,7 +2815,7 @@ class ConfigGeneralSettings(LiteLLMPydanticObjectBase): allowed_routes: list | None = Field(None, description="Proxy API Endpoints you want users to be able to access") allowed_client_side_s3_buckets: tuple[str, ...] | None = Field( None, - description="S3 bucket names a client may send as `s3_bucket_name` or `s3_output_bucket_name` in a request body, e.g. on /v1/batches. Any other bucket is rejected unless client-side credentials are allowed", + description="S3 bucket names a client may send as `s3_bucket_name` or `s3_output_bucket_name` in a request body, e.g. on /v1/batches, and the buckets of the s3:// image and file URLs a non-admin key sends. Any other bucket is rejected unless client-side credentials are allowed", ) allowed_client_side_s3_endpoint_urls: tuple[str, ...] | None = Field( None, diff --git a/litellm/proxy/auth/auth_utils.py b/litellm/proxy/auth/auth_utils.py index 2b5fc6d877e..bd8ba46ad84 100644 --- a/litellm/proxy/auth/auth_utils.py +++ b/litellm/proxy/auth/auth_utils.py @@ -495,7 +495,7 @@ _PROXY_ALLOWLIST_SETTING_BY_PARAM: Final[Mapping[str, _ClientSideAllowlistSettin ) -class _ClientSideStorageAllowlist(BaseModel): +class ClientSideStorageAllowlist(BaseModel): model_config = ConfigDict(frozen=True) allowed_client_side_s3_buckets: tuple[str, ...] | None = None @@ -555,7 +555,7 @@ def _check_banned_params( return if ( setting := _PROXY_ALLOWLIST_SETTING_BY_PARAM.get(param) - ) is not None and _ClientSideStorageAllowlist.model_validate(general_settings).admits(setting, body[param]): + ) is not None and ClientSideStorageAllowlist.model_validate(general_settings).admits(setting, body[param]): continue if ( _allow_model_level_clientside_configurable_parameters( diff --git a/litellm/proxy/common_request_processing.py b/litellm/proxy/common_request_processing.py index e812b03385d..8ebefb74327 100644 --- a/litellm/proxy/common_request_processing.py +++ b/litellm/proxy/common_request_processing.py @@ -6,7 +6,7 @@ import math from collections.abc import AsyncGenerator, AsyncIterator, Awaitable, Callable, Coroutine, Mapping, Sequence from dataclasses import dataclass from datetime import datetime, timezone -from functools import lru_cache +from functools import lru_cache, partial from itertools import chain from types import MappingProxyType from typing import ( @@ -62,6 +62,11 @@ from litellm.litellm_core_utils.bug_report import ( should_report_bug, strip_bug_report_notice, ) +from litellm.litellm_core_utils.cloud_storage_security import ( + is_managed_s3_uri, + names_an_unlisted_s3_bucket, + s3_media_uris, +) from litellm.litellm_core_utils.core_helpers import ( finite_cost, get_or_create_metadata_bucket, @@ -94,7 +99,7 @@ from litellm.litellm_core_utils.streaming_handler import ( backfill_missing_cache_usage_fields, ) from litellm.litellm_core_utils.tool_fields import DropPolicy, MalformedField, drop_policy, resolve_tool_fields -from litellm.proxy._types import LiteLLMRoutes, ProxyErrorTypes, ProxyException, UserAPIKeyAuth +from litellm.proxy._types import LiteLLMRoutes, LitellmUserRoles, ProxyErrorTypes, ProxyException, UserAPIKeyAuth from litellm.proxy.anthropic_endpoints.answered_status import is_provider_overload, with_openai_overload_status from litellm.proxy.auth.auth_checks import ( can_key_call_resolved_model, @@ -102,7 +107,7 @@ from litellm.proxy.auth.auth_checks import ( request_skips_budget_checks, tag_max_budget_check_for_tags, ) -from litellm.proxy.auth.auth_utils import check_response_size_is_safe, get_request_route +from litellm.proxy.auth.auth_utils import ClientSideStorageAllowlist, check_response_size_is_safe, get_request_route from litellm.proxy.bug_report_config import build_proxy_bug_report from litellm.proxy.common_utils.callback_utils import ( get_logging_caching_headers, @@ -981,6 +986,37 @@ def _request_model(data: Mapping[str, object]) -> str | list[str] | None: return None +def _refuse_s3_media_the_key_may_not_read( + data: Mapping[str, object], + user_api_key_dict: UserAPIKeyAuth, + general_settings: Mapping[str, object], +) -> None: + if user_api_key_dict.user_role == LitellmUserRoles.PROXY_ADMIN: + return + s3_uris: Final = s3_media_uris(data) + if not s3_uris: + return + managed_uris: Final = tuple(filter(is_managed_s3_uri, s3_uris)) + if managed_uris: + raise HTTPException( + status_code=status.HTTP_403_FORBIDDEN, + detail=f"Raw s3:// URIs under LiteLLM-managed storage can only be sent by a proxy admin key " + f"({', '.join(managed_uris)}). Use the LiteLLM managed file id returned when the file was created.", + ) + if general_settings.get("allow_client_side_credentials") is True: + return + is_listed: Final = partial( + ClientSideStorageAllowlist.model_validate(general_settings).admits, "allowed_client_side_s3_buckets" + ) + unlisted_uris: Final = tuple(uri for uri in s3_uris if names_an_unlisted_s3_bucket(uri, is_listed)) + if unlisted_uris: + raise HTTPException( + status_code=status.HTTP_403_FORBIDDEN, + detail=f"Raw s3:// URIs can only name a bucket listed in `general_settings.allowed_client_side_s3_buckets` " + f"({', '.join(unlisted_uris)}). Ask the proxy admin to list the bucket there.", + ) + + async def _enforce_guardrail_added_tag_budgets( data: Mapping[str, object], tags_before_guardrails: frozenset[str], @@ -2492,7 +2528,8 @@ class ProxyBaseLLMRequestProcessing: ) -> tuple[dict, LiteLLMLoggingObj]: start_time: Final = datetime.now() # start before calling guardrail hooks - requested_model: Final = self.data.get("model") + raw_request_fields: Final = _REQUEST_FIELDS.validate_python(self.data) + requested_model: Final = raw_request_fields.get("model") if requested_model is not None and not isinstance(requested_model, str): raise ProxyException( message="'model' must be a string.", @@ -2500,6 +2537,9 @@ class ProxyBaseLLMRequestProcessing: param="model", code=status.HTTP_400_BAD_REQUEST, ) + _refuse_s3_media_the_key_may_not_read( + raw_request_fields, user_api_key_dict, _REQUEST_FIELDS.validate_python(general_settings) + ) self.data = await add_litellm_data_to_request( data=self.data, request=request, diff --git a/litellm/responses/litellm_completion_transformation/transformation.py b/litellm/responses/litellm_completion_transformation/transformation.py index 1e87dfa0906..7de3cb19ec5 100644 --- a/litellm/responses/litellm_completion_transformation/transformation.py +++ b/litellm/responses/litellm_completion_transformation/transformation.py @@ -86,6 +86,7 @@ from litellm.responses.litellm_completion_transformation.session_handler import ) from litellm.types.llms.anthropic import AnthropicToolSearchToolRegex from litellm.types.llms.base import CachedTokensDetails +from litellm.types.llms.bedrock import S3MediaFileObjectFile, S3MediaImageUrlObject from litellm.types.llms.openai import ( AllMessageValues, ChatCompletionAssistantMessage, @@ -93,7 +94,6 @@ from litellm.types.llms.openai import ( ChatCompletionCachedContent, ChatCompletionFileObject, ChatCompletionImageObject, - ChatCompletionImageUrlObject, ChatCompletionRedactedThinkingBlock, ChatCompletionResponseMessage, ChatCompletionSystemMessage, @@ -272,8 +272,8 @@ _ANY_KEY_DICT_ADAPTER: Final = TypeAdapter(dict[object, object]) _STR_KEY_DICT_ADAPTER: Final = TypeAdapter(dict[str, object]) _OBJECT_LIST_ADAPTER: Final = TypeAdapter(list[object]) _AUDIO_PART_ADAPTER: Final = TypeAdapter(ChatCompletionAudioObject) -_IMAGE_URL_ADAPTER: Final = TypeAdapter(ChatCompletionImageUrlObject) -_FILE_PART_ADAPTER: Final = TypeAdapter(_ResponsesFilePart) +_IMAGE_URL_ADAPTER: Final = TypeAdapter(S3MediaImageUrlObject) +_FILE_ADAPTER: Final = TypeAdapter(S3MediaFileObjectFile) _MessageT = TypeVar("_MessageT") _MediaT: Final = TypeVar("_MediaT") _InputItemT: Final = TypeVar("_InputItemT") @@ -294,6 +294,11 @@ def _validate_responses_media(adapter: TypeAdapter[_MediaT], value: object) -> _ ) from exc +def _s3_bucket_owner_field(url: object, bucket_owner: object) -> Mapping[str, object]: + is_s3_url: Final = isinstance(url, str) and url.startswith("s3://") + return {"bucket_owner": bucket_owner} if is_s3_url and bucket_owner is not None else {} + + def _responses_image_part(item: Mapping[str, object]) -> _ResponsesImagePart: source: Final = item.get("image_url") or item.get("url") fields: Final = _validate_responses_media(_STR_KEY_DICT_ADAPTER, source) if isinstance(source, Mapping) else {} @@ -307,15 +312,15 @@ def _responses_image_part(item: Mapping[str, object]) -> _ResponsesImagePart: ) detail: Final = fields.get("detail") or item.get("detail") video_metadata: Final = fields.get("video_metadata") or item.get("video_metadata") - bucket_owner: Final = fields.get("bucket_owner") or item.get("bucket_owner") + url: Final = (fields.get("url") if isinstance(source, Mapping) else source) or "" image_url: Final = _validate_responses_media( _IMAGE_URL_ADAPTER, { - "url": (fields.get("url") if isinstance(source, Mapping) else source) or "", + "url": url, "detail": detail if isinstance(detail, str) else "auto", **({"format": media_format} if media_format is not None else {}), **({"video_metadata": video_metadata} if video_metadata is not None else {}), - **({"bucket_owner": bucket_owner} if bucket_owner is not None else {}), + **_s3_bucket_owner_field(url, fields.get("bucket_owner", item.get("bucket_owner"))), }, ) return ( @@ -330,23 +335,20 @@ def _responses_file_part(item: Mapping[str, object]) -> _ResponsesFilePart: fields: Final = {**nested, **{key: value for key, value in item.items() if key != "file" and value is not None}} file_id: Final = fields.get("file_id") or fields.get("file_url") media_format: Final = fields.get("format") or fields.get("mime_type") or fields.get("content_type") - file: Final = { - **{ - key: fields[key] - for key in ("file_data", "filename", "detail", "video_metadata", "bucket_owner") - if fields.get(key) - }, - **({"file_id": file_id} if file_id else {}), - **({"format": media_format} if media_format is not None else {}), - } - return _validate_responses_media( - _FILE_PART_ADAPTER, + file: Final = _validate_responses_media( + _FILE_ADAPTER, { - "type": "file", - "file": file, - **({"cache_control": item["cache_control"]} if "cache_control" in item else {}), + **{key: fields[key] for key in ("file_data", "filename", "detail", "video_metadata") if fields.get(key)}, + **({"file_id": file_id} if file_id else {}), + **({"format": media_format} if media_format is not None else {}), + **_s3_bucket_owner_field(file_id or fields.get("file_data"), fields.get("bucket_owner")), }, ) + return ( + _ResponsesFilePart(type="file", file=file, cache_control=item["cache_control"]) + if "cache_control" in item + else _ResponsesFilePart(type="file", file=file) + ) def decode_responses_media(item: Mapping[str, object]) -> _ResponsesMediaPart | None: diff --git a/litellm/types/llms/bedrock.py b/litellm/types/llms/bedrock.py index 5425391ee64..04233f37d8a 100644 --- a/litellm/types/llms/bedrock.py +++ b/litellm/types/llms/bedrock.py @@ -9,7 +9,7 @@ from typing_extensions import NotRequired, ReadOnly, Required, TypedDict, overri from litellm.types.llms.base import LiteLLMBaseModel -from .openai import ChatCompletionToolCallChunk +from .openai import ChatCompletionFileObjectFile, ChatCompletionImageUrlObject, ChatCompletionToolCallChunk class CachePointBlock(TypedDict, total=False): @@ -32,6 +32,14 @@ class SourceBlock(TypedDict, total=False): s3Location: ReadOnly[S3Location] +class S3MediaImageUrlObject(ChatCompletionImageUrlObject, total=False): + bucket_owner: ReadOnly[object] + + +class S3MediaFileObjectFile(ChatCompletionFileObjectFile, total=False): + bucket_owner: ReadOnly[object] + + BedrockImageTypes = Literal["png", "jpeg", "gif", "webp"] @@ -74,6 +82,7 @@ class SearchResultBlock(TypedDict, total=False): class ToolResultContentBlock(TypedDict, total=False): image: ImageBlock document: DocumentBlock + video: ReadOnly[VideoBlock] searchResult: SearchResultBlock json: dict text: str diff --git a/litellm/types/llms/openai.py b/litellm/types/llms/openai.py index 0db3010022e..c9739eee91d 100644 --- a/litellm/types/llms/openai.py +++ b/litellm/types/llms/openai.py @@ -822,7 +822,6 @@ class ChatCompletionImageUrlObject(TypedDict, total=False): detail: str format: str video_metadata: ReadOnly[Mapping[str, object]] - bucket_owner: ReadOnly[str] class ChatCompletionImageObject(TypedDict): @@ -876,7 +875,6 @@ class ChatCompletionFileObjectFile(TypedDict, total=False): format: str detail: str # For video/image resolution control (low, medium, high, ultra_high) video_metadata: dict[str, Any] # For video-specific metadata (fps, start_offset, end_offset) - bucket_owner: ReadOnly[str] class ChatCompletionFileObject(TypedDict): diff --git a/tests/unit/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_factory.py b/tests/unit/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_factory.py index f716278a7c7..3c2532aba7f 100644 --- a/tests/unit/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_factory.py +++ b/tests/unit/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_factory.py @@ -3119,6 +3119,21 @@ def _s3_source(uri: str, bucket_owner: str | None = None) -> Mapping[str, object {"image": {"format": "jpeg", "source": _s3_source("s3://media-bucket/dir/CAT.JPG")}}, id="string-url-uppercase-extension", ), + pytest.param( + {"type": "image_url", "image_url": {"url": "s3://media-bucket/dir/cat.jpg"}}, + {"image": {"format": "jpeg", "source": _s3_source("s3://media-bucket/dir/cat.jpg")}}, + id="jpg-is-jpeg", + ), + pytest.param( + {"type": "image_url", "image_url": {"url": "s3://media-bucket/dir/CAT.PNG"}}, + {"image": {"format": "png", "source": _s3_source("s3://media-bucket/dir/CAT.PNG")}}, + id="uppercase-image-extension", + ), + pytest.param( + {"type": "image_url", "image_url": {"url": "s3://media-bucket/clips/DEMO.MOV"}}, + {"video": {"format": "mov", "source": _s3_source("s3://media-bucket/clips/DEMO.MOV")}}, + id="uppercase-video-extension", + ), pytest.param( { "type": "image_url", @@ -3196,6 +3211,10 @@ async def test_bedrock_converse_s3_media_without_a_known_extension_or_format_is_ pytest.param( {"url": "s3://media-bucket/scans/page.png", "format": "image/tiff"}, "image/tiff", id="tiff-format" ), + pytest.param({"url": "s3://media-bucket/voice/note.mp3"}, "audio/mpeg", id="mp3-key"), + pytest.param( + {"url": "s3://media-bucket/voice/note.bin", "format": "audio/mp4"}, "audio/mp4", id="audio-format" + ), ], ) async def test_bedrock_converse_s3_media_in_a_format_bedrock_does_not_take_is_a_bad_request( @@ -3208,6 +3227,42 @@ async def test_bedrock_converse_s3_media_in_a_format_bedrock_does_not_take_is_a_ assert named_format in excinfo.value.message +_PNG_BASE64: Final = "iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAQAAAC1HAwCAAAAC0lEQVR42mNkYAAAAAYAAjCB0C8AAAAASUVORK5CYII=" + + +def _owned_image_url_part(url: str, bucket_owner: object) -> Mapping[str, object]: + return {"type": "image_url", "image_url": {"url": url, "bucket_owner": bucket_owner}} + + +def _owned_file_part(url: str, bucket_owner: object) -> Mapping[str, object]: + return {"type": "file", "file": {"file_id": url, "bucket_owner": bucket_owner}} + + +@pytest.mark.parametrize("run_async", [False, True], ids=["sync", "async"]) +@pytest.mark.parametrize("bucket_owner", [111122223333, ["111122223333"], ""], ids=["int", "list", "empty-string"]) +@pytest.mark.parametrize("part_for", [_owned_image_url_part, _owned_file_part], ids=["image-url", "file"]) +async def test_bedrock_converse_s3_media_with_a_bucket_owner_that_is_not_an_account_id_is_a_bad_request( + part_for: Callable[[str, object], Mapping[str, object]], bucket_owner: object, run_async: bool +): + with pytest.raises(litellm.BadRequestError) as excinfo: + await _converse_user_content(part_for("s3://media-bucket/dir/cat.png", bucket_owner), run_async) + + assert "`bucket_owner`" in excinfo.value.message + assert "s3://media-bucket/dir/cat.png" in excinfo.value.message + + +@pytest.mark.parametrize("run_async", [False, True], ids=["sync", "async"]) +@pytest.mark.parametrize("part_for", [_owned_image_url_part, _owned_file_part], ids=["image-url", "file"]) +async def test_bedrock_converse_ignores_a_bucket_owner_on_media_that_is_not_on_s3( + part_for: Callable[[str, object], Mapping[str, object]], run_async: bool +): + content: Final = await _converse_user_content( + part_for(f"data:image/png;base64,{_PNG_BASE64}", 111122223333), run_async + ) + + assert content == [{"image": {"format": "png", "source": {"bytes": _PNG_BASE64}}}] + + @pytest.mark.parametrize( "part", [ @@ -3240,11 +3295,44 @@ def test_bedrock_converse_tool_result_s3_media_keeps_its_bucket_owner(part: Mapp ] +@pytest.mark.parametrize("bucket_owner", [111122223333, ""], ids=["int", "empty-string"]) +@pytest.mark.parametrize("part_for", [_owned_image_url_part, _owned_file_part], ids=["image-url", "file"]) +def test_bedrock_converse_tool_result_s3_media_with_a_bucket_owner_that_is_not_an_account_id_is_a_bad_request( + part_for: Callable[[str, object], Mapping[str, object]], bucket_owner: object +): + messages: Final = [ + {"role": "user", "content": "Describe the screenshot."}, + { + "tool_call_id": "tooluse_s3_1", + "role": "tool", + "name": "fetch_screenshot", + "content": [part_for("s3://media-bucket/dir/cat.png", bucket_owner)], + }, + ] + + with respx.mock, pytest.raises(litellm.BadRequestError) as excinfo: + _bedrock_converse_messages_pt(messages=messages, model=_S3_MEDIA_MODEL, llm_provider="bedrock") + + assert "`bucket_owner`" in excinfo.value.message + + @pytest.mark.parametrize("run_async", [False, True], ids=["sync", "async"]) async def test_bedrock_converse_assistant_s3_image_is_sent_as_an_s3_location(run_async: bool): messages: Final = [ {"role": "user", "content": "Draw a cat."}, - {"role": "assistant", "content": [{"type": "image_url", "image_url": {"url": "s3://media-bucket/dir/cat.png"}}]}, + { + "role": "assistant", + "content": [ + { + "type": "image_url", + "image_url": { + "url": "s3://media-bucket/dir/cat.png", + "format": "image/webp", + "bucket_owner": _S3_BUCKET_OWNER, + }, + } + ], + }, {"role": "user", "content": "Make it bigger."}, ] @@ -3258,7 +3346,7 @@ async def test_bedrock_converse_assistant_s3_image_is_sent_as_an_s3_location(run ) assert converted[1]["content"] == [ - {"image": {"format": "png", "source": _s3_source("s3://media-bucket/dir/cat.png")}} + {"image": {"format": "webp", "source": _s3_source("s3://media-bucket/dir/cat.png", _S3_BUCKET_OWNER)}} ] 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 46131aa64e5..e1bd97368c9 100644 --- a/tests/unit/litellm_core_utils/test_cloud_storage_security.py +++ b/tests/unit/litellm_core_utils/test_cloud_storage_security.py @@ -7,6 +7,9 @@ import pytest from litellm.litellm_core_utils.cloud_storage_security import ( MAX_FILE_ID_DECODE_PASSES, is_managed_cloud_storage_uri, + is_managed_s3_uri, + names_an_unlisted_s3_bucket, + s3_media_uris, ) @@ -55,3 +58,150 @@ def test_is_managed_cloud_storage_uri_fails_closed_on_an_id_encoded_past_the_pas assert is_managed_cloud_storage_uri(_quoted_times("file-abc/x", MAX_FILE_ID_DECODE_PASSES + 1)) assert is_managed_cloud_storage_uri(nested_gs_id) assert is_managed_cloud_storage_uri(nested_plain_id) + + +_MANAGED_UPLOAD: Final = "s3://media-bucket/litellm-bedrock-files/other-team.pdf" + + +def _user_turn(*parts: object) -> dict[str, object]: + return {"messages": [{"role": "user", "content": list(parts)}]} + + +def _responses_turn(*parts: object) -> dict[str, object]: + return {"input": [{"role": "user", "content": list(parts)}]} + + +@pytest.mark.parametrize( + "request_body", + ( + pytest.param(_user_turn({"type": "image_url", "image_url": {"url": _MANAGED_UPLOAD}}), id="chat-image-url"), + pytest.param(_user_turn({"type": "image_url", "image_url": _MANAGED_UPLOAD}), id="chat-image-url-string"), + pytest.param(_user_turn({"type": "file", "file": {"file_id": _MANAGED_UPLOAD}}), id="chat-file-id"), + pytest.param(_user_turn({"type": "file", "file": {"file_data": _MANAGED_UPLOAD}}), id="chat-file-data"), + pytest.param( + { + "messages": [ + { + "role": "tool", + "tool_call_id": "call_1", + "content": [{"type": "image_url", "image_url": {"url": _MANAGED_UPLOAD}}], + } + ] + }, + id="chat-tool-result", + ), + pytest.param(_responses_turn({"type": "input_image", "image_url": _MANAGED_UPLOAD}), id="responses-image"), + pytest.param(_responses_turn({"type": "input_file", "file_url": _MANAGED_UPLOAD}), id="responses-file-url"), + pytest.param(_responses_turn({"type": "input_file", "file_id": _MANAGED_UPLOAD}), id="responses-file-id"), + pytest.param( + { + "input": [ + { + "type": "function_call_output", + "call_id": "call_1", + "output": [{"type": "input_image", "image_url": _MANAGED_UPLOAD}], + } + ] + }, + id="responses-function-call-output", + ), + pytest.param( + _user_turn({"type": "image", "source": {"type": "url", "url": _MANAGED_UPLOAD}}), id="anthropic-image" + ), + pytest.param( + _user_turn({"type": "document", "source": {"type": "url", "url": _MANAGED_UPLOAD}}), + id="anthropic-document", + ), + pytest.param( + {"system": [{"type": "image", "source": {"type": "url", "url": _MANAGED_UPLOAD}}], "messages": []}, + id="anthropic-system", + ), + ), +) +def test_s3_media_uris_finds_an_s3_object_on_every_media_surface(request_body: dict[str, object]): + assert s3_media_uris(request_body) == (_MANAGED_UPLOAD,) + + +@pytest.mark.parametrize( + "uri", + ( + "s3://media-bucket/litellm-bedrock-files/x.pdf", + "s3://media-bucket/litellm-bedrock-files-batch-1/input.jsonl", + "s3://media-bucket/litellm-batch-outputs/x.jsonl.out", + "s3://media-bucket/tenant-a/litellm-bedrock-files/x.pdf", + "s3://media-bucket//litellm-bedrock-files/x.pdf", + "s3://media-bucket/./litellm-bedrock-files/x.pdf", + "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", + ), +) +def test_is_managed_s3_uri_sees_a_managed_prefix_anywhere_in_the_decoded_key(uri: str): + assert s3_media_uris(_user_turn({"type": "file", "file": {"file_id": uri}})) == (uri,) + assert is_managed_s3_uri(uri) + + +def test_is_managed_s3_uri_fails_closed_on_a_key_encoded_past_the_pass_cap(): + uri: Final = "s3://media-bucket/" + _quoted_times("litellm-bedrock-files/x.pdf", MAX_FILE_ID_DECODE_PASSES + 2) + + assert s3_media_uris(_user_turn({"type": "image_url", "image_url": {"url": uri}})) == (uri,) + assert is_managed_s3_uri(uri) + + +@pytest.mark.parametrize( + "uri", + ( + pytest.param("s3://media-bucket/dir/cat.png", id="unmanaged-key"), + pytest.param("s3://litellm-bedrock-files-bucket/cat.png", id="managed-prefix-in-bucket-name-only"), + pytest.param("s3://media-bucket", id="no-key"), + ), +) +def test_is_managed_s3_uri_leaves_other_objects_alone(uri: str): + assert not is_managed_s3_uri(uri) + + +@pytest.mark.parametrize( + "request_body", + ( + pytest.param( + _user_turn({"type": "image_url", "image_url": {"url": "https://example.test/litellm-bedrock-files/x.png"}}), + id="https-url", + ), + pytest.param( + _user_turn({"type": "image_url", "image_url": {"url": "data:image/png;base64,czM6Ly9idWNrZXQva2V5"}}), + id="data-url", + ), + pytest.param(_user_turn({"type": "text", "text": _MANAGED_UPLOAD}), id="text-part"), + 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"), + ), +) +def test_s3_media_uris_ignores_what_never_becomes_an_s3_media_source(request_body: dict[str, object]): + assert s3_media_uris(request_body) == () + + +@pytest.mark.parametrize( + ("uri", "names_an_unlisted_bucket"), + ( + pytest.param("s3://listed-bucket/dir/cat.png", False, id="listed"), + pytest.param("s3://listed-bucket/dir/a%20cat.png", False, id="listed-with-an-encoded-key"), + pytest.param("s3://other-bucket/dir/cat.png", True, id="unlisted"), + pytest.param("s3://listed-bucket.example/cat.png", True, id="listed-name-as-a-prefix"), + pytest.param("s3://listed%2Dbucket/cat.png", True, id="encoded-bucket"), + pytest.param("s3://other-bucket%2Flisted-bucket/cat.png", True, id="encoded-slash-in-the-bucket"), + pytest.param("s3:///listed-bucket/cat.png", True, id="empty-bucket"), + pytest.param("s3%3A%2F%2Fother-bucket%2Fcat.png", True, id="encoded-scheme"), + pytest.param( + "s3://listed-bucket/" + _quoted_times("a cat.png", MAX_FILE_ID_DECODE_PASSES + 2), + True, + id="encoded-past-the-pass-cap", + ), + ), +) +def test_names_an_unlisted_s3_bucket_checks_the_bucket_of_every_decoding(uri: str, names_an_unlisted_bucket: bool): + assert names_an_unlisted_s3_bucket(uri, frozenset({"listed-bucket"}).__contains__) is names_an_unlisted_bucket diff --git a/tests/unit/llms/bedrock/chat/test_converse_transformation.py b/tests/unit/llms/bedrock/chat/test_converse_transformation.py index b37d6977c10..bcdcd378ff4 100644 --- a/tests/unit/llms/bedrock/chat/test_converse_transformation.py +++ b/tests/unit/llms/bedrock/chat/test_converse_transformation.py @@ -367,6 +367,43 @@ def test_bedrock_system_prompt_variants_keep_continue_message( ] +def test_bedrock_tool_result_s3_video_reaches_the_converse_request_as_a_video_block() -> None: + clip: Final = "s3://media-bucket/clips/demo.mp4" + request: Final = AmazonConverseConfig()._transform_request( + model="us.amazon.nova-2-lite-v1:0", + messages=[ + {"role": "user", "content": "What happens in the clip?"}, + { + "role": "assistant", + "content": None, + "tool_calls": [ + {"id": "call_1", "type": "function", "function": {"name": "fetch_clip", "arguments": "{}"}} + ], + }, + { + "role": "tool", + "tool_call_id": "call_1", + "content": [{"type": "text", "text": "video attached"}, {"type": "file", "file": {"file_id": clip}}], + }, + ], + optional_params={ + "tools": [ + { + "type": "function", + "function": {"name": "fetch_clip", "parameters": {"type": "object", "properties": {}}}, + } + ] + }, + litellm_params={}, + headers={}, + ) + + assert request["messages"][-1]["content"][0]["toolResult"]["content"] == [ + {"text": "video attached"}, + {"video": {"format": "mp4", "source": {"s3Location": {"uri": clip}}}}, + ] + + @pytest.mark.respx(assert_all_called=True) def test_bedrock_json_schema_description_is_sent_in_request(respx_mock: respx.Router) -> None: description: Final = "Find the meaning inside a poem" diff --git a/tests/unit/proxy/guardrails/guardrail_hooks/test_generic_guardrail_api.py b/tests/unit/proxy/guardrails/guardrail_hooks/test_generic_guardrail_api.py index b5f794cb136..8f3697e9b24 100644 --- a/tests/unit/proxy/guardrails/guardrail_hooks/test_generic_guardrail_api.py +++ b/tests/unit/proxy/guardrails/guardrail_hooks/test_generic_guardrail_api.py @@ -33,6 +33,7 @@ from litellm.proxy.guardrails.guardrail_hooks.generic_guardrail_api import ( from litellm.proxy.guardrails.guardrail_hooks.generic_guardrail_api.generic_guardrail_api import ( _HEADER_PRESENT_PLACEHOLDER, ) +from litellm.litellm_core_utils.prompt_templates.factory import _bedrock_converse_messages_pt from litellm.proxy.guardrails.guardrail_hooks.unified_guardrail.unified_guardrail import UnifiedLLMGuardrails from litellm.types.guardrails import GuardrailEventHooks from litellm.types.utils import Choices, GenericGuardrailAPIInputs, Message @@ -949,6 +950,89 @@ class TestStructuredMessagesInResponse: assert guardrailed_inputs["structured_messages"] == [assistant_row, redacted_row] + @pytest.mark.asyncio + async def test_an_echo_that_drops_bucket_owner_still_sends_bucket_owner_to_bedrock( + self, raw_rows_guardrail, mock_request_data_input + ): + image_row = { + "role": "user", + "content": [ + {"type": "text", "text": "Describe it."}, + { + "type": "image_url", + "image_url": {"url": "s3://media-bucket/dir/cat.png", "bucket_owner": "111122223333"}, + }, + ], + } + redacted_row = {"role": "user", "content": "my ssn is [SSN]"} + + def echo_without_bucket_owner_and_redact_one_row(url, json, headers, **_kwargs): + sent_image_row = json["structured_messages"][0] + answer = MagicMock() + answer.json.return_value = { + "action": "GUARDRAIL_INTERVENED", + "structured_messages": [ + { + "role": "user", + "content": [ + sent_image_row["content"][0], + {"type": "image_url", "image_url": {"url": sent_image_row["content"][1]["image_url"]["url"]}}, + ], + }, + redacted_row, + ], + } + answer.raise_for_status = MagicMock() + return answer + + with patch.object( + raw_rows_guardrail.async_handler, "post", side_effect=echo_without_bucket_owner_and_redact_one_row + ): + guardrailed_inputs = await raw_rows_guardrail.apply_guardrail( + inputs={ + "texts": ["Describe it.", "my ssn is 123-45-6789"], + "structured_messages": [image_row, {"role": "user", "content": "my ssn is 123-45-6789"}], + }, + request_data=mock_request_data_input, + input_type="request", + ) + bedrock_messages = _bedrock_converse_messages_pt( + messages=guardrailed_inputs["structured_messages"], model="us.amazon.nova-2-lite-v1:0", llm_provider="bedrock" + ) + + assert [block for block in bedrock_messages[0]["content"] if "image" in block] == [ + { + "image": { + "format": "png", + "source": {"s3Location": {"uri": "s3://media-bucket/dir/cat.png", "bucketOwner": "111122223333"}}, + } + } + ] + + @pytest.mark.asyncio + async def test_an_image_with_a_bucket_owner_that_is_not_a_string_keeps_its_turn_for_the_guardrail( + self, generic_guardrail + ): + row = { + "role": "user", + "content": [ + {"type": "text", "text": "What is written here?"}, + {"type": "image_url", "image_url": {"url": "data:image/png;base64,AA", "bucket_owner": 111122223333}}, + ], + } + + assert await _structured_messages_sent(generic_guardrail, [row]) == [ + [ + { + "role": "user", + "content": [ + {"type": "text", "text": "What is written here?"}, + {"type": "image_url", "image_url": {"url": "data:image/png;base64,AA"}}, + ], + } + ] + ] + @pytest.mark.asyncio @pytest.mark.parametrize( "returned_row", diff --git a/tests/unit/proxy/test_common_request_processing.py b/tests/unit/proxy/test_common_request_processing.py index 5c27ad8ce77..605a2b9497f 100644 --- a/tests/unit/proxy/test_common_request_processing.py +++ b/tests/unit/proxy/test_common_request_processing.py @@ -76,7 +76,7 @@ from litellm.proxy.common_utils.sse_keepalive import ANTHROPIC_PING_SSE_CHUNK from litellm.proxy.dd_span_tagger import DDSpanTagger from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.llms.anthropic.chat.transformation import AnthropicConfig -from litellm.proxy._types import Litellm_EntityType, ProxyErrorTypes, ProxyException +from litellm.proxy._types import Litellm_EntityType, LitellmUserRoles, ProxyErrorTypes, ProxyException from litellm.proxy._types import UserAPIKeyAuth as ProxyUserAPIKeyAuth from litellm.proxy.auth.fallback_budget import RouterFallbackBudgetCheck from litellm.proxy.hooks.model_max_budget_limiter import ( @@ -219,6 +219,65 @@ def test_include_guardrail_response_requested_is_false_without_exact_true(): assert include_guardrail_response_requested({}) is False +def _chat_image(uri: str) -> dict[str, object]: + return {"messages": [{"role": "user", "content": [{"type": "image_url", "image_url": {"url": uri}}]}]} + + +def _responses_image(uri: str) -> dict[str, object]: + return {"input": [{"role": "user", "content": [{"type": "input_image", "image_url": uri}]}]} + + +def _anthropic_image(uri: str) -> dict[str, object]: + return { + "max_tokens": 16, + "messages": [{"role": "user", "content": [{"type": "image", "source": {"type": "url", "url": uri}}]}], + } + + +_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"), +) + + +async def _send_s3_media( + route_type: Literal["acompletion", "aresponses", "anthropic_messages"], + path: str, + body: dict[str, object], + user_role: LitellmUserRoles, + general_settings: dict[str, object], +) -> object: + from litellm.proxy.proxy_server import ProxyConfig as RealProxyConfig + + router: Final = Router( + model_list=[ + { + "model_name": "nova", + "litellm_params": {"model": "bedrock/us.amazon.nova-2-lite-v1:0", "mock_response": "a cat"}, + } + ] + ) + return await ProxyBaseLLMRequestProcessing(data={"model": "nova", **body}).base_process_llm_request( + request=Request( + { + "type": "http", + "method": "POST", + "path": path, + "query_string": b"", + "headers": [(b"content-type", b"application/json")], + } + ), + fastapi_response=Response(), + user_api_key_dict=ProxyUserAPIKeyAuth(api_key="sk-test", user_role=user_role), + route_type=route_type, + proxy_logging_obj=ProxyLogging(user_api_key_cache=UserApiKeyCache()), + general_settings=general_settings, + proxy_config=RealProxyConfig(), + llm_router=router, + ) + + class TestProxyBaseLLMRequestProcessing: @pytest.mark.asyncio async def test_base_passthrough_process_llm_request_preserves_litellm_headers_for_non_streaming_response( @@ -510,6 +569,71 @@ class TestProxyBaseLLMRequestProcessing: assert exc_info.value.param == "model" add_litellm_data_to_request.assert_not_awaited() + @pytest.mark.asyncio + @pytest.mark.parametrize(("route_type", "path", "body_for"), _S3_MEDIA_SURFACES) + @pytest.mark.parametrize( + ("uri", "general_settings", "named_in_refusal"), + [ + pytest.param( + "s3://other-bucket/cat.png", + {"allowed_client_side_s3_buckets": ["media-bucket"]}, + "general_settings.allowed_client_side_s3_buckets", + id="unlisted-bucket", + ), + pytest.param( + "s3://media-bucket/cat.png", + {}, + "general_settings.allowed_client_side_s3_buckets", + id="no-bucket-listed", + ), + pytest.param( + "s3://media-bucket/litellm-bedrock-files/other-team.png", + {"allowed_client_side_s3_buckets": ["media-bucket"], "allow_client_side_credentials": True}, + "LiteLLM-managed storage", + id="managed-key-in-a-listed-bucket", + ), + ], + ) + async def test_non_admin_s3_media_is_refused_while_a_proxy_admin_key_passes( + self, + route_type: Literal["acompletion", "aresponses", "anthropic_messages"], + path: str, + body_for: Callable[[str], dict[str, object]], + uri: str, + general_settings: dict[str, object], + named_in_refusal: str, + ) -> None: + with pytest.raises(HTTPException) as refused: + await _send_s3_media(route_type, path, body_for(uri), LitellmUserRoles.INTERNAL_USER, general_settings) + admin_response: Final = await _send_s3_media(route_type, path, body_for(uri), LitellmUserRoles.PROXY_ADMIN, {}) + + assert refused.value.status_code == status.HTTP_403_FORBIDDEN, refused.value.detail + assert uri in refused.value.detail + assert named_in_refusal in refused.value.detail + assert "a cat" in str(admin_response) + + @pytest.mark.asyncio + @pytest.mark.parametrize(("route_type", "path", "body_for"), _S3_MEDIA_SURFACES) + @pytest.mark.parametrize( + "general_settings", + [ + pytest.param({"allowed_client_side_s3_buckets": ["media-bucket"]}, id="listed-bucket"), + pytest.param({"allow_client_side_credentials": True}, id="client-side-credentials-allowed"), + ], + ) + async def test_non_admin_s3_media_passes_in_a_bucket_the_admin_allowed( + self, + route_type: Literal["acompletion", "aresponses", "anthropic_messages"], + 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 + ) + + assert "a cat" in str(response) + @pytest.mark.asyncio @pytest.mark.parametrize( ("path", "route_type", "has_router", "priority_header", "body_priority", "expected_priority"), diff --git a/tests/unit/responses/litellm_completion_transformation/test_litellm_completion_responses.py b/tests/unit/responses/litellm_completion_transformation/test_litellm_completion_responses.py index 9e2659a8781..8bacb03fac4 100644 --- a/tests/unit/responses/litellm_completion_transformation/test_litellm_completion_responses.py +++ b/tests/unit/responses/litellm_completion_transformation/test_litellm_completion_responses.py @@ -307,6 +307,30 @@ class TestLiteLLMCompletionResponsesConfig: }, id="input-file", ), + pytest.param( + {"type": "input_image", "image_url": "s3://media-bucket/dir/cat.png", "bucket_owner": 111122223333}, + { + "type": "image_url", + "image_url": { + "url": "s3://media-bucket/dir/cat.png", + "detail": "auto", + "bucket_owner": 111122223333, + }, + }, + id="input-image-owner-left-for-bedrock-to-check", + ), + pytest.param( + { + "type": "input_image", + "image_url": {"url": "s3://media-bucket/dir/cat.png", "bucket_owner": ""}, + "bucket_owner": "111122223333", + }, + { + "type": "image_url", + "image_url": {"url": "s3://media-bucket/dir/cat.png", "detail": "auto", "bucket_owner": ""}, + }, + id="input-image-empty-owner-left-for-bedrock-to-check", + ), ], ) def test_transform_input_media_item_keeps_bucket_owner(self, item: Mapping[str, object], expected: object): @@ -314,6 +338,35 @@ class TestLiteLLMCompletionResponsesConfig: expected ] + @pytest.mark.parametrize( + ("item", "expected"), + [ + pytest.param( + {"type": "input_image", "image_url": "data:image/png;base64,AA", "bucket_owner": "111122223333"}, + {"type": "image_url", "image_url": {"url": "data:image/png;base64,AA", "detail": "auto"}}, + id="data-url-image", + ), + pytest.param( + {"type": "input_image", "image_url": "https://example.test/cat.png", "bucket_owner": 111122223333}, + {"type": "image_url", "image_url": {"url": "https://example.test/cat.png", "detail": "auto"}}, + id="https-image-with-an-int-owner", + ), + pytest.param( + { + "type": "input_file", + "file_data": "data:application/pdf;base64,AA", + "bucket_owner": ["111122223333"], + }, + {"type": "file", "file": {"file_data": "data:application/pdf;base64,AA"}}, + id="data-url-file-with-a-list-owner", + ), + ], + ) + def test_transform_input_media_item_drops_bucket_owner_off_s3(self, item: Mapping[str, object], expected: object): + assert LiteLLMCompletionResponsesConfig._transform_responses_api_content_to_chat_completion_content([item]) == [ + expected + ] + def test_transform_chat_completion_response_with_reasoning_content(self): """Test that reasoning content is preserved in the full transformation pipeline""" # Setup diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index 127735a2974..e4cb9fccdb7 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -29783,8 +29783,6 @@ export interface components { }; /** ChatCompletionFileObjectFile */ ChatCompletionFileObjectFile: { - /** Bucket Owner */ - bucket_owner?: string; /** Detail */ detail?: string; /** File Data */ @@ -29826,8 +29824,6 @@ export interface components { }; /** ChatCompletionImageUrlObject */ ChatCompletionImageUrlObject: { - /** Bucket Owner */ - bucket_owner?: string; /** Detail */ detail?: string; /** Format */ @@ -30797,7 +30793,7 @@ export interface components { allow_unmanaged_response_ids?: boolean | null; /** * Allowed Client Side S3 Buckets - * @description S3 bucket names a client may send as `s3_bucket_name` or `s3_output_bucket_name` in a request body, e.g. on /v1/batches. Any other bucket is rejected unless client-side credentials are allowed + * @description S3 bucket names a client may send as `s3_bucket_name` or `s3_output_bucket_name` in a request body, e.g. on /v1/batches, and the buckets of the s3:// image and file URLs a non-admin key sends. Any other bucket is rejected unless client-side credentials are allowed */ allowed_client_side_s3_buckets?: string[] | null; /**