diff --git a/litellm/litellm_core_utils/cloud_storage_security.py b/litellm/litellm_core_utils/cloud_storage_security.py index c913e4de02e..411759059fe 100644 --- a/litellm/litellm_core_utils/cloud_storage_security.py +++ b/litellm/litellm_core_utils/cloud_storage_security.py @@ -1,11 +1,18 @@ +import base64 import posixpath import re -from collections.abc import Mapping, Sequence +from collections.abc import Callable, Iterable, Iterator, Mapping, Sequence +from dataclasses import dataclass +from functools import partial +from itertools import accumulate, chain, repeat, takewhile from types import MappingProxyType from typing import Any, Final, cast from urllib.parse import quote, unquote +from pydantic import TypeAdapter + 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-" @@ -18,6 +25,15 @@ BEDROCK_MANAGED_S3_PREFIXES: Final = ( ) MANAGED_CLOUD_STORAGE_SCHEMES: Final = ("s3://", "gs://") _MAPPING_PROXY_TYPE: Final[type] = type(MappingProxyType({})) +MAX_FILE_ID_DECODE_PASSES: Final = 8 + + +def _unquote_once(current: str, _: str) -> str: + return unquote(current) + + +def _decodings(value: str) -> tuple[str, ...]: + return tuple(accumulate(repeat(value, MAX_FILE_ID_DECODE_PASSES + 2), _unquote_once)) def is_managed_cloud_storage_uri(file_id: str) -> bool: @@ -31,6 +47,182 @@ def is_managed_cloud_storage_uri(file_id: str) -> bool: return isinstance(file_id, str) and file_id.startswith(MANAGED_CLOUD_STORAGE_SCHEMES) +_S3_SCHEME: Final = "s3://" +_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"}) +_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_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} +_SCALAR_TYPES: Final = frozenset({int, float, bool, type(None)}) +_OBJECT_TUPLE: Final = TypeAdapter(tuple[object, ...]) + + +def is_s3_uri(value: str) -> bool: + return value.startswith(_S3_SCHEME) + + +def _base64url_text(value: str) -> str | None: + try: + return base64.urlsafe_b64decode(value + "=" * (-len(value) % 4)).decode() + except ValueError: + return None + + +def _s3_uris_in(text: str) -> Iterator[str]: + return (match.group(0) for match in _EMBEDDED_S3_URI.finditer(text)) + + +def _s3_uris_embedded_in_file_id(file_id: str) -> Iterator[str]: + encodings: Final = dict.fromkeys((file_id, *(file_id.removeprefix(prefix) for prefix in _ENCODED_FILE_ID_PREFIXES))) + return chain.from_iterable(map(_s3_uris_in, filter(None, map(_base64url_text, encodings)))) + + +def _with_embedded_s3_uris(file_id: str) -> Iterator[str]: + return chain((file_id,), _s3_uris_embedded_in_file_id(file_id)) + + +def _non_scalar(values: Iterable[object]) -> Iterator[object]: + return (value for value in values if type(value) not in _SCALAR_TYPES) + + +def _sequence_items(value: object) -> Sequence[object] | None: + if isinstance(value, tuple): + return _OBJECT_TUPLE.validate_python(value) + return as_object_list(value) + + +@dataclass(frozen=True, slots=True) +class _BodyNode: + value: object + key: str | None = None + inside_media_part: bool = False + + +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 _child_nodes(node: _BodyNode) -> Iterator[_BodyNode]: + fields: Final = as_str_mapping(node.value) + if fields is not None: + inside_media_part: Final = node.inside_media_part or _is_media_part(fields, node.key) + return ( + _BodyNode(child, child_key, inside_media_part) + for child_key, child in fields.items() + if type(child) not in _SCALAR_TYPES + ) + items: Final = _sequence_items(node.value) + if items is None: + return iter(()) + return (_BodyNode(item, None, node.inside_media_part) for item in _non_scalar(items)) + + +def _next_level(level: tuple[_BodyNode, ...], _: None) -> tuple[_BodyNode, ...]: + return tuple(chain.from_iterable(map(_child_nodes, level))) + + +def _body_nodes(request_body: object) -> Iterator[_BodyNode]: + levels: Final = accumulate(repeat(None), _next_level, initial=(_BodyNode(request_body),)) + return chain.from_iterable(takewhile(bool, levels)) + + +def _string_leaf(node: _BodyNode) -> str | None: + return node.value if isinstance(node.value, str) else None + + +def _media_file_id(node: _BodyNode) -> str | 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 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 ()) + + +def _mapped_provider_file_ids(request_body: Mapping[str, object]) -> Iterator[str]: + mapping_leaves: Final = filter(None, map(_string_leaf, _body_nodes(request_body.get(_MODEL_FILE_ID_MAPPING_KEY)))) + return chain.from_iterable(map(_with_embedded_s3_uris, mapping_leaves)) + + +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 _may_name_an_s3_uri(value: object) -> bool: + if isinstance(value, str): + return is_s3_uri(value) + fields: Final = as_str_mapping(value) + if fields is not None: + return not _KEYS_THAT_MAY_HIDE_AN_S3_URI.isdisjoint(fields) or any( + map(_may_name_an_s3_uri, _non_scalar(fields.values())) + ) + items: Final = _sequence_items(value) + return items is not None and any(map(_may_name_an_s3_uri, _non_scalar(items))) + + +def _may_carry_s3_media(request_body: Mapping[str, object]) -> bool: + try: + return _may_name_an_s3_uri(request_body) + except RecursionError: + return True + + +def media_file_ids(request_body: Mapping[str, object]) -> tuple[str, ...]: + return tuple(dict.fromkeys(filter(None, map(_media_file_id, _body_nodes(request_body))))) + + +def file_id_embeds_s3_uri_matching(file_id: str, matches: Callable[[str], bool]) -> bool: + return any(map(matches, _s3_uris_embedded_in_file_id(file_id))) + + +def s3_media_uris( + request_body: Mapping[str, object], trusted_file_ids: frozenset[str] = frozenset() +) -> tuple[str, ...]: + if not _may_carry_s3_media(request_body): + return () + media_uris: Final = chain.from_iterable( + map(partial(_media_uri_candidates, trusted_file_ids), _body_nodes(request_body)) + ) + return tuple(dict.fromkeys(filter(is_s3_uri, chain(media_uris, _mapped_provider_file_ids(request_body))))) + + +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/core_helpers.py b/litellm/litellm_core_utils/core_helpers.py index 7ab3dd7687d..5f209b3de49 100644 --- a/litellm/litellm_core_utils/core_helpers.py +++ b/litellm/litellm_core_utils/core_helpers.py @@ -8,7 +8,7 @@ from types import MappingProxyType from typing import TYPE_CHECKING, Any, Final, Literal, Protocol, cast, runtime_checkable import httpx -from pydantic import InstanceOf, TypeAdapter, ValidationError +from pydantic import TypeAdapter, ValidationError import litellm from litellm._logging import verbose_logger @@ -57,14 +57,8 @@ def as_str_mapping(value: object) -> Mapping[str, object] | None: return value if isinstance(value, Mapping) else None -_OBJECT_LIST: Final = TypeAdapter(InstanceOf[list[object]]) - - def as_object_list(value: object) -> Sequence[object] | None: - try: - return _OBJECT_LIST.validate_python(value) - except ValidationError: - return None + return value if isinstance(value, list) else None @runtime_checkable diff --git a/litellm/litellm_core_utils/prompt_templates/factory.py b/litellm/litellm_core_utils/prompt_templates/factory.py index 767f3b4c30d..2e11ef2c775 100644 --- a/litellm/litellm_core_utils/prompt_templates/factory.py +++ b/litellm/litellm_core_utils/prompt_templates/factory.py @@ -8,6 +8,7 @@ import re import xml.etree.ElementTree as ET from collections.abc import Container, Iterable, Iterator, Mapping, Sequence from enum import Enum +from pathlib import PurePosixPath from types import MappingProxyType from typing import Any, Final, TypeAlias, TypedDict, cast, overload @@ -20,6 +21,8 @@ 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 from litellm.litellm_core_utils.prompt_templates.mid_conversation_system import ( anthropic_system_messages, find_anthropic_server_tool_result, @@ -32,7 +35,7 @@ from litellm.litellm_core_utils.prompt_templates.mid_conversation_system import from litellm.litellm_core_utils.tool_fields import tools_without_hosted_mcp from litellm.litellm_core_utils.url_utils import async_safe_get, safe_get from litellm.llms.custom_httpx.http_handler import HTTPHandler, get_async_httpx_client, header_value -from litellm.types.files import get_file_extension_from_mime_type +from litellm.types.files import get_file_extension_from_mime_type, get_file_mime_type_from_extension from litellm.types.llms.anthropic import * from litellm.types.llms.bedrock import CachePointBlock from litellm.types.llms.bedrock import MessageBlock as BedrockMessageBlock @@ -2190,6 +2193,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]" @@ -3639,6 +3646,7 @@ from litellm.types.llms.bedrock import ( from litellm.types.llms.bedrock import ContentBlock as BedrockContentBlock from litellm.types.llms.bedrock import DocumentBlock as BedrockDocumentBlock from litellm.types.llms.bedrock import ImageBlock as BedrockImageBlock +from litellm.types.llms.bedrock import S3Location as BedrockS3Location from litellm.types.llms.bedrock import SourceBlock as BedrockSourceBlock from litellm.types.llms.bedrock import ToolBlock as BedrockToolBlock from litellm.types.llms.bedrock import ToolResultBlock as BedrockToolResultBlock @@ -3811,10 +3819,8 @@ class BedrockImageProcessor: return valid_extensions[0] @staticmethod - def _create_bedrock_block(image_bytes: str, mime_type: str, image_format: str) -> BedrockContentBlock: + def _create_bedrock_block(source: BedrockSourceBlock, mime_type: str, image_format: str) -> BedrockContentBlock: """Create appropriate Bedrock content block based on mime type.""" - _blob: Final = BedrockSourceBlock(bytes=image_bytes) - document_types: Final = ["application", "text"] is_document: Final = any(mime_type.startswith(doc_type) for doc_type in document_types) @@ -3824,12 +3830,10 @@ class BedrockImageProcessor: HASH_SAMPLE_BYTES: Final = 64 * 1024 # hash up to 64 KB of data if is_document: - # --- Prepare normalized bytes for hashing (without modifying original) --- - if isinstance(image_bytes, str): - # Remove whitespace/newlines so base64 variations hash identically - normalized = "".join(image_bytes.split()).encode("utf-8") - else: - normalized = image_bytes + s3_location: Final = source.get("s3Location") + hashed_content: Final = s3_location["uri"] if s3_location is not None else source.get("bytes") or "" + # Remove whitespace/newlines so base64 variations hash identically + normalized: Final = "".join(hashed_content.split()).encode("utf-8") # --- Use only the first 64 KB for speed --- if len(normalized) <= HASH_SAMPLE_BYTES: @@ -3849,20 +3853,75 @@ class BedrockImageProcessor: # --- Return content block --- return BedrockContentBlock( document=BedrockDocumentBlock( - source=_blob, + source=source, format=image_format, name=document_name, ) ) elif is_video: - return BedrockContentBlock(video=BedrockVideoBlock(source=_blob, 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=_blob, format=image_format)) + return BedrockContentBlock(image=BedrockImageBlock(source=source, format=image_format)) + + @staticmethod + def _s3_key_mime_type(uri: str) -> str | None: + extension: Final = PurePosixPath(uri).suffix.removeprefix(".").lower() + if not extension: + return None + config: Final = litellm.AmazonConverseConfig() + 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: + return None @classmethod - def process_image_sync(cls, image_url: str, format: str | None = None) -> BedrockContentBlock: + def _create_s3_block(cls, uri: str, format: str | None, bucket_owner: str | None) -> BedrockContentBlock: + 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}. " + "Audio is not one of its image, document, or video formats.", + model="", + llm_provider="bedrock", + ) + try: + 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}", + model="", + llm_provider="bedrock", + ) from e + s3_location: Final = ( + BedrockS3Location(uri=uri) if bucket_owner is None else BedrockS3Location(uri=uri, bucketOwner=bucket_owner) + ) + return cls._create_bedrock_block(BedrockSourceBlock(s3Location=s3_location), mime_type, image_format) + + @classmethod + def process_image_sync( + cls, image_url: str, format: str | None = None, bucket_owner: str | None = None + ) -> BedrockContentBlock: """Synchronous image processing.""" + 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) elif "http://" in image_url or "https://" in image_url: @@ -3876,12 +3935,16 @@ class BedrockImageProcessor: image_format = mime_type.split("/")[1] image_format = cls.validate_format(mime_type, image_format) - return cls._create_bedrock_block(img_bytes, mime_type, image_format) + return cls._create_bedrock_block(BedrockSourceBlock(bytes=img_bytes), mime_type, image_format) @classmethod - async def process_image_async(cls, image_url: str, format: str | None) -> BedrockContentBlock: + async def process_image_async( + cls, image_url: str, format: str | None = None, bucket_owner: str | None = None + ) -> BedrockContentBlock: """Asynchronous image processing.""" + 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) elif "http://" in image_url or "https://" in image_url: @@ -3895,7 +3958,7 @@ class BedrockImageProcessor: image_format = mime_type.split("/")[1] image_format = cls.validate_format(mime_type, image_format) - return cls._create_bedrock_block(img_bytes, mime_type, image_format) + return cls._create_bedrock_block(BedrockSourceBlock(bytes=img_bytes), mime_type, image_format) def _convert_to_bedrock_tool_call_invoke( @@ -4021,13 +4084,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 and "s3Location" in processed_block["video"]["source"]: + 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.", @@ -4037,9 +4102,23 @@ def _append_bedrock_tool_result_media_block( ) +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 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}.", + 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): @@ -4050,17 +4129,18 @@ def _append_bedrock_tool_result_image_url_block( processed_block: Final = BedrockImageProcessor.process_image_sync( image_url=image_url, format=format, + 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: @@ -4069,9 +4149,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(file_obj, url), ) _append_bedrock_tool_result_media_block(tool_result_content_blocks, processed_block, content, "file") @@ -4706,7 +4788,9 @@ class BedrockConverseMessagesProcessor: else: image_url = element["image_url"] _part = await BedrockImageProcessor.process_image_async( - image_url=image_url, format=format + image_url=image_url, + format=format, + bucket_owner=_s3_bucket_owner(element["image_url"], image_url), ) _parts.append(_part) elif element["type"] == "file": @@ -4849,7 +4933,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( @@ -4986,7 +5078,10 @@ class BedrockConverseMessagesProcessor: llm_provider="bedrock", ) format: Final = file_message.get("format") - return BedrockImageProcessor.process_image_sync(image_url=cast(str, file_id or file_data), format=format) + url: Final = cast(str, file_id or file_data) + return BedrockImageProcessor.process_image_sync( + image_url=url, format=format, bucket_owner=_s3_bucket_owner(file_message, url) + ) _process_file_message = process_file_message @@ -5010,7 +5105,10 @@ class BedrockConverseMessagesProcessor: model="", llm_provider="bedrock", ) - return await BedrockImageProcessor.process_image_async(image_url=cast(str, file_id or file_data), format=format) + url: Final = cast(str, file_id or file_data) + return await BedrockImageProcessor.process_image_async( + image_url=url, format=format, bucket_owner=_s3_bucket_owner(file_message, url) + ) @staticmethod def _process_document_message(element: dict) -> BedrockContentBlock: @@ -5144,6 +5242,7 @@ def _bedrock_converse_messages_pt( _part = BedrockImageProcessor.process_image_sync( image_url=image_url, format=format, + bucket_owner=_s3_bucket_owner(element["image_url"], image_url), ) _parts.append(_part) elif element["type"] == "file": @@ -5289,7 +5388,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/llms/bedrock/files/transformation.py b/litellm/llms/bedrock/files/transformation.py index 1b3a1b206f5..a1808134c60 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, @@ -323,8 +324,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, @@ -336,6 +336,22 @@ 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}") + + +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 _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) @@ -1027,6 +1043,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=optional_params, @@ -1108,11 +1127,14 @@ class BedrockFilesConfig(BaseAWSLLM, BaseFilesConfig): openai_request_body=openai_body, model=model_for_transform ) else: - model_input = self._map_openai_to_bedrock_params( - openai_request_body=self._transform_batch_body_to_chat_body(openai_body, record_kind), - model=model_for_transform, - provider=provider, - ) + try: + model_input = self._map_openai_to_bedrock_params( + openai_request_body=self._transform_batch_body_to_chat_body(openai_body, record_kind), + model=model_for_transform, + provider=provider, + ) + except _S3MediaInBatchRecord as e: + raise _batch_line_error(idx + 1, str(e)) from e # Create Bedrock batch record record_id = _openai_jsonl_content.get("custom_id", f"CALL{str(idx).zfill(7)}") diff --git a/litellm/llms/custom_httpx/llm_http_handler.py b/litellm/llms/custom_httpx/llm_http_handler.py index 6737905192e..6685ca1c6b1 100644 --- a/litellm/llms/custom_httpx/llm_http_handler.py +++ b/litellm/llms/custom_httpx/llm_http_handler.py @@ -126,6 +126,7 @@ from litellm.responses.streaming_iterator import ( ProjectQuotaCallback, ResponsesAPIStreamingIterator, ResponsesWebSocketStreaming, + ResponsesWsFrameCheck, SyncResponsesAPIStreamingIterator, ) from litellm.types.containers.main import ( @@ -6933,6 +6934,7 @@ class BaseLLMHTTPHandler: custom_llm_provider: str | None = None, first_message: str | None = None, request_defaults: ResponsesWebSocketRequestDefaults | None = None, + frame_check: ResponsesWsFrameCheck | None = None, **kwargs: object, ) -> Exception | None: """ @@ -6965,6 +6967,7 @@ class BaseLLMHTTPHandler: custom_llm_provider=custom_llm_provider, first_message=first_message, quota_callbacks=_ws_quota_callbacks, + frame_check=frame_check, **kwargs, ) await handler.run() diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index 5e812eec477..f29dab18e33 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -2736,6 +2736,10 @@ class ConfigGeneralSettings(LiteLLMPydanticObjectBase): description="How long a blocked source address, or source address and username, stays blocked. Every attempt from a blocked key, right or wrong, is refused with 429 before the password is checked; the block is not extended by refused attempts. Set under `general_settings` in config.yaml. Defaults to 300", ) 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 that keys other than proxy admins may name in `s3://` image and file URLs. Any other bucket is rejected unless client-side credentials are allowed", + ) reject_clientside_metadata_tags: bool | None = Field( None, description="When set to True, rejects requests that contain client-side 'metadata.tags' to prevent users from influencing budgets by sending different tags. Tags can only be inherited from the API key metadata.", diff --git a/litellm/proxy/auth/auth_utils.py b/litellm/proxy/auth/auth_utils.py index f3d7805693b..dd57fffc6de 100644 --- a/litellm/proxy/auth/auth_utils.py +++ b/litellm/proxy/auth/auth_utils.py @@ -2,7 +2,7 @@ import importlib.util import os import re import sys -from collections.abc import Collection, Iterator, Mapping +from collections.abc import Callable, Collection, Iterator, Mapping from functools import lru_cache from itertools import accumulate, chain, repeat from logging import Logger @@ -10,7 +10,7 @@ from types import MappingProxyType from typing import Any, Final, Protocol from fastapi import HTTPException, Request, status -from pydantic import PositiveInt, TypeAdapter, ValidationError +from pydantic import BaseModel, ConfigDict, PositiveInt, TypeAdapter, ValidationError import litellm from litellm import Router, constants, provider_list @@ -41,7 +41,11 @@ from litellm.router_utils.clientside_credential_handler import ( from litellm.types.passthrough_endpoints.pass_through_endpoints import ( LITELLM_PASS_THROUGH_ENDPOINT_MARKER, ) -from litellm.types.router import CONFIGURABLE_CLIENTSIDE_AUTH_PARAMS, Deployment, server_owned_wif_fields_present +from litellm.types.router import ( + CONFIGURABLE_CLIENTSIDE_AUTH_PARAMS, + Deployment, + server_owned_wif_fields_present, +) from litellm.types.router import reject_server_owned_wif_params as _reject_server_owned_wif_params from litellm.types.utils import CustomPricingLiteLLMParams, oauth_token_exchange_litellm_params from litellm.types.videos.main import DecodedVideoId @@ -528,6 +532,19 @@ _BANNED_METADATA_PARAMS: Final[tuple[str, ...]] = tuple( ) +class ClientSideStorageAllowlist(BaseModel): + model_config = ConfigDict(frozen=True) + + allowed_client_side_s3_buckets: tuple[str, ...] | None = None + + def admits_s3_bucket(self, bucket: str) -> bool: + return bucket in (self.allowed_client_side_s3_buckets or ()) + + +def client_side_s3_bucket_check(general_settings: Mapping[str, object]) -> Callable[[str], bool]: + return ClientSideStorageAllowlist.model_validate(general_settings).admits_s3_bucket + + def _check_banned_params( body: Mapping[str, object], general_settings: dict, diff --git a/litellm/proxy/common_request_processing.py b/litellm/proxy/common_request_processing.py index c6562dc2273..ddf2917d881 100644 --- a/litellm/proxy/common_request_processing.py +++ b/litellm/proxy/common_request_processing.py @@ -3,9 +3,18 @@ import contextlib import json import logging import math -from collections.abc import AsyncGenerator, AsyncIterator, Awaitable, Callable, Coroutine, Mapping, Sequence +from collections.abc import ( + AsyncGenerator, + AsyncIterator, + Awaitable, + Callable, + Coroutine, + Iterable, + Mapping, + Sequence, +) from datetime import datetime -from functools import lru_cache +from functools import lru_cache, partial from itertools import chain from types import MappingProxyType from typing import ( @@ -58,6 +67,13 @@ from litellm.litellm_core_utils.bug_report import ( should_report_bug, strip_bug_report_notice, ) +from litellm.litellm_core_utils.cloud_storage_security import ( + file_id_embeds_s3_uri_matching, + is_managed_s3_uri, + media_file_ids, + names_an_unlisted_s3_bucket, + s3_media_uris, +) from litellm.litellm_core_utils.core_helpers import ( get_or_create_metadata_bucket, independent_snapshot, @@ -87,14 +103,14 @@ 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.auth.auth_checks import ( can_key_call_resolved_model, is_model_cost_zero, 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 check_response_size_is_safe, client_side_s3_bucket_check, 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, @@ -125,6 +141,10 @@ from litellm.proxy.common_utils.sse_keepalive import ( from litellm.proxy.dd_span_tagger import DDSpanTagger from litellm.proxy.guardrails.auto_router_compression import arm_pre_call as _arm_auto_router_compression from litellm.proxy.native_compaction import with_proxy_compaction_executor +from litellm.proxy.openai_files_endpoints.common_utils import ( + ManagedResourceAccessChecker, + is_base64_encoded_unified_file_id, +) from litellm.proxy.route_llm_request import ( route_request, ) @@ -855,6 +875,81 @@ def _request_model(data: Mapping[str, object]) -> str | list[str] | None: return None +async def _key_may_read_unified_file_id( + managed_files: ManagedResourceAccessChecker, user_api_key_dict: UserAPIKeyAuth, unified_file_id: str +) -> bool: + try: + return await managed_files.can_user_call_unified_file_id(unified_file_id, user_api_key_dict) + except HTTPException: + return False + + +async def _leading_readable_file_ids( + managed_files: ManagedResourceAccessChecker, user_api_key_dict: UserAPIKeyAuth, unified_file_ids: Iterable[str] +) -> AsyncIterator[str]: + for file_id in unified_file_ids: + if not await _key_may_read_unified_file_id(managed_files, user_api_key_dict, file_id): + return + yield file_id + + +async def _unified_file_ids_the_key_may_read( + data: Mapping[str, object], + user_api_key_dict: UserAPIKeyAuth, + managed_files: object, + is_refused: Callable[[str], bool], +) -> frozenset[str]: + if not isinstance(managed_files, ManagedResourceAccessChecker): + return frozenset() + file_ids_hiding_a_refused_uri: Final = filter( + partial(file_id_embeds_s3_uri_matching, matches=is_refused), + filter(is_base64_encoded_unified_file_id, media_file_ids(data)), + ) + return frozenset( + [ + file_id + async for file_id in _leading_readable_file_ids( + managed_files, user_api_key_dict, file_ids_hiding_a_refused_uri + ) + ] + ) + + +def _unlisted_s3_uri_check(general_settings: Mapping[str, object]) -> Callable[[str], bool]: + if general_settings.get("allow_client_side_credentials") is True: + return lambda _: False + return partial(names_an_unlisted_s3_bucket, is_listed=client_side_s3_bucket_check(general_settings)) + + +async def refuse_s3_media_the_key_may_not_read( + data: Mapping[str, object], + user_api_key_dict: UserAPIKeyAuth, + general_settings: Mapping[str, object], + managed_files: object, +) -> None: + if not s3_media_uris(data) or user_api_key_dict.user_role == LitellmUserRoles.PROXY_ADMIN: + return + is_unlisted: Final = _unlisted_s3_uri_check(general_settings) + readable_file_ids: Final = await _unified_file_ids_the_key_may_read( + data, user_api_key_dict, managed_files, lambda uri: is_managed_s3_uri(uri) or is_unlisted(uri) + ) + s3_uris: Final = s3_media_uris(data, readable_file_ids) + 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 a LiteLLM managed file id that this key's user or team created.", + ) + unlisted_uris: Final = tuple(filter(is_unlisted, s3_uris)) + 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], @@ -2277,6 +2372,13 @@ class ProxyBaseLLMRequestProcessing: param="model", code=status.HTTP_400_BAD_REQUEST, ) + if route_type != "call_mcp_tool": + await refuse_s3_media_the_key_may_not_read( + self.data, + user_api_key_dict, + general_settings, + proxy_logging_obj.get_proxy_hook("managed_files"), + ) self.data = await add_litellm_data_to_request( data=self.data, request=request, diff --git a/litellm/proxy/guardrails/guardrail_hooks/model_armor/file_scanning.py b/litellm/proxy/guardrails/guardrail_hooks/model_armor/file_scanning.py index 366ea7326ae..931aacbe8be 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/model_armor/file_scanning.py +++ b/litellm/proxy/guardrails/guardrail_hooks/model_armor/file_scanning.py @@ -7,7 +7,7 @@ maps its MIME type to a Model Armor ``byteDataType``, and returns the decoded by guardrail hooks can submit them. ``plan_file_scans`` classifies each block: blocks with no inline bytes (``file_id`` or remote -``gs://`` / ``http(s)`` references) and supported documents whose base64 will not decode are +``gs://`` / ``s3://`` / ``http(s)`` references) and supported documents whose base64 will not decode are reported as unscannable so the guardrail hook can fail closed (blocking unless ``fail_on_error`` is false) rather than letting an unscanned document reach the model. """ @@ -27,7 +27,7 @@ from litellm.types.llms.openai import AllMessageValues MODEL_ARMOR_MAX_FILE_SIZE_BYTES: Final = 4 * 1024 * 1024 -_REMOTE_URI_SCHEMES: Final = ("gs://", "http://", "https://") +_REMOTE_URI_SCHEMES: Final = ("gs://", "s3://", "http://", "https://") ModelArmorByteDataType = Literal["PDF", "WORD_DOCUMENT", "EXCEL_DOCUMENT", "POWERPOINT_DOCUMENT", "CSV", "TXT"] diff --git a/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py b/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py index 5a41e2ea068..16001d5bf43 100644 --- a/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py +++ b/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py @@ -93,6 +93,7 @@ from litellm.proxy.common_request_processing import ( ProxyBaseLLMRequestProcessing, log_llm_api_exception, open_sse_before_first_byte, + refuse_s3_media_the_key_may_not_read, resolve_litellm_call_id, ) from litellm.proxy.common_utils.error_body_call_id import JSON_OBJECT, error_body_call_id, with_call_id @@ -253,6 +254,10 @@ async def chat_completion_pass_through_endpoint( data["adapter_id"] = adapter_id + await refuse_s3_media_the_key_may_not_read( + data, user_api_key_dict, general_settings, proxy_logging_obj.get_proxy_hook("managed_files") + ) + verbose_proxy_logger.debug("Request received by LiteLLM:\n%s", data) data["model"] = ( general_settings.get("completion_model", None) # server default diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 80970cba90b..e44af2aaf6d 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -445,6 +445,7 @@ from litellm.proxy.common_request_processing import ( # noqa: F401, RUF100 # l log_llm_api_exception, open_sse_before_first_byte, proxy_exception_from_route_error, + refuse_s3_media_the_key_may_not_read, request_litellm_call_id, resolve_litellm_call_id, should_return_raw_model_name, @@ -17679,6 +17680,9 @@ async def async_queue_request( try: data = await request.json() data.pop("_litellm_strip_stream_usage", None) + await refuse_s3_media_the_key_may_not_read( + data, user_api_key_dict, general_settings, proxy_logging_obj.get_proxy_hook("managed_files") + ) # Include original request and headers in the data data["proxy_server_request"] = { diff --git a/litellm/proxy/rag_endpoints/endpoints.py b/litellm/proxy/rag_endpoints/endpoints.py index a1a7d62899f..ea3c9036d20 100644 --- a/litellm/proxy/rag_endpoints/endpoints.py +++ b/litellm/proxy/rag_endpoints/endpoints.py @@ -29,6 +29,7 @@ from litellm.proxy.auth.user_api_key_auth import UserAPIKeyAuth, user_api_key_au from litellm.proxy.common_request_processing import ( ProxyBaseLLMRequestProcessing, open_sse_before_first_byte, + refuse_s3_media_the_key_may_not_read, ttft_keepalive_interval, ) from litellm.proxy.common_utils.http_parsing_utils import ( # noqa: F401 # legacy module exports @@ -766,6 +767,7 @@ async def rag_query( general_settings, llm_router, proxy_config, + proxy_logging_obj, select_data_generator, version, ) @@ -828,6 +830,10 @@ async def rag_query( **store_data, } + await refuse_s3_media_the_key_may_not_read( + data, user_api_key_dict, general_settings, proxy_logging_obj.get_proxy_hook("managed_files") + ) + # Add litellm data request_data: dict[str, object] = {} request_data = await add_litellm_data_to_request( diff --git a/litellm/proxy/response_api_endpoints/endpoints.py b/litellm/proxy/response_api_endpoints/endpoints.py index 63062569af4..0d77c9750d2 100644 --- a/litellm/proxy/response_api_endpoints/endpoints.py +++ b/litellm/proxy/response_api_endpoints/endpoints.py @@ -51,6 +51,7 @@ from litellm.proxy.common_request_processing import ( ProxyBaseLLMRequestProcessing, cost_headers, create_response, + refuse_s3_media_the_key_may_not_read, resolve_counted_tool_fields, ) from litellm.proxy.common_utils.error_body_call_id import error_body_call_id, with_call_id @@ -71,6 +72,7 @@ from litellm.proxy.pass_through_endpoints.llm_provider_handlers.batch_attributio request_tags_from_request_data, ) from litellm.proxy.route_llm_request import raise_if_required_body_param_missing +from litellm.responses.streaming_iterator import ResponsesWsFrameCheck from litellm.router import Router from litellm.types.llms.base import LiteLLMBaseModel from litellm.types.llms.openai import ( @@ -1904,6 +1906,24 @@ async def _enforce_responses_ws_first_frame_model_auth( ) +def _s3_media_frame_check( + admitted_model: object, + user_api_key_dict: UserAPIKeyAuth, + general_settings: Mapping[str, object], + managed_files: object, +) -> ResponsesWsFrameCheck: + async def refusal(params: Mapping[str, object], /) -> str | None: + try: + await refuse_s3_media_the_key_may_not_read( + {**params, "model": admitted_model}, user_api_key_dict, general_settings, managed_files + ) + except HTTPException as refused: + return str(refused.detail) + return None + + return refusal + + async def _responses_websocket_session( websocket: WebSocket, model: str | None, @@ -2000,7 +2020,15 @@ async def _responses_websocket_session( return routed_data: Final = dict( - data, user_api_key_dict=user_api_key_dict, **_routing_hints_from_first_ws_frame(first_message) + data, + user_api_key_dict=user_api_key_dict, + responses_ws_frame_check=_s3_media_frame_check( + data.get("model"), + user_api_key_dict, + general_settings, + proxy_logging_obj.get_proxy_hook("managed_files"), + ), + **_routing_hints_from_first_ws_frame(first_message), ) # Phase 2: route to upstream provider try: diff --git a/litellm/responses/litellm_completion_transformation/transformation.py b/litellm/responses/litellm_completion_transformation/transformation.py index 31cd78c7517..b8e60d40e29 100644 --- a/litellm/responses/litellm_completion_transformation/transformation.py +++ b/litellm/responses/litellm_completion_transformation/transformation.py @@ -41,6 +41,7 @@ from litellm._logging import verbose_logger from litellm.caching import InMemoryCache from litellm.constants import REDACTED_BY_LITELLM, REDACTED_TOOL_CALL_ARGUMENTS_PLACEHOLDER from litellm.integrations.anthropic_cache_control_hook import MAX_CACHE_CONTROL_BLOCKS, AnthropicCacheControlHook +from litellm.litellm_core_utils.cloud_storage_security import is_s3_uri from litellm.litellm_core_utils.core_helpers import drops_unsupported_params, normalize_drop_params from litellm.litellm_core_utils.get_supported_openai_params import ( get_supported_openai_params, @@ -62,6 +63,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 S3MediaImageUrlObject from litellm.types.llms.openai import ( AllMessageValues, ChatCompletionAssistantMessage, @@ -226,6 +228,10 @@ _TEXT_ADAPTER: Final = TypeAdapter(str) _RESPONSES_API_TOOL_CHOICE_ADAPTER: Final = TypeAdapter(ToolChoice) +def _is_s3_url(url: object) -> bool: + return isinstance(url, str) and is_s3_uri(url) + + @runtime_checkable class _SupportsIter(Protocol): def __iter__(self) -> Iterator[object]: ... @@ -2253,6 +2259,12 @@ class LiteLLMCompletionResponsesConfig: file_dict["file_data"] = item["file_data"] if item.get("filename"): file_dict["filename"] = item["filename"] + media_format: Final = item.get("format") + if media_format is not None: + file_dict["format"] = media_format + bucket_owner: Final = item.get("bucket_owner") + if bucket_owner is not None and _is_s3_url(file_id or item.get("file_data")): + file_dict["bucket_owner"] = bucket_owner new_item: Final[dict[str, object]] = {"type": "file", "file": file_dict} if "cache_control" in item: @@ -2266,9 +2278,17 @@ class LiteLLMCompletionResponsesConfig: """ Transform a Responses API input_image item to a Chat Completion image item """ - image_url_obj: Final = ChatCompletionImageUrlObject( - url=item.get("image_url") or "", detail=item.get("detail") or "auto" + url: Final = item.get("image_url") or "" + detail: Final = item.get("detail") or "auto" + bucket_owner: Final = item.get("bucket_owner") + image_url_obj: Final = ( + S3MediaImageUrlObject(url=url, detail=detail, bucket_owner=bucket_owner) + if bucket_owner is not None and _is_s3_url(url) + else ChatCompletionImageUrlObject(url=url, detail=detail) ) + media_format: Final = item.get("format") + if media_format is not None: + image_url_obj["format"] = media_format return ChatCompletionImageObject(type="image_url", image_url=image_url_obj) diff --git a/litellm/responses/main.py b/litellm/responses/main.py index e73b4668e27..a4196ca0a2d 100644 --- a/litellm/responses/main.py +++ b/litellm/responses/main.py @@ -84,7 +84,7 @@ if TYPE_CHECKING: else: MCPTool = Any -from .streaming_iterator import BaseResponsesAPIStreamingIterator +from .streaming_iterator import BaseResponsesAPIStreamingIterator, ResponsesWsFrameCheck __all__ = ( "acancel_responses", @@ -2555,6 +2555,7 @@ async def _aresponses_websocket( api_base: str | None = None, api_key: str | None = None, timeout: float | None = None, + responses_ws_frame_check: ResponsesWsFrameCheck | None = None, **kwargs, ) -> Exception | None: """ @@ -2647,5 +2648,6 @@ async def _aresponses_websocket( litellm_metadata=_build_litellm_metadata_for_ws(kwargs), custom_llm_provider=_custom_llm_provider, request_defaults=_build_responses_websocket_request_defaults(deployment_kwargs), + frame_check=responses_ws_frame_check, **remaining_kwargs, ) diff --git a/litellm/responses/streaming_iterator.py b/litellm/responses/streaming_iterator.py index 7a6685ae824..7d00bea5ba8 100644 --- a/litellm/responses/streaming_iterator.py +++ b/litellm/responses/streaming_iterator.py @@ -66,6 +66,10 @@ if TYPE_CHECKING: from litellm.types.router import LiteLLM_Params +class ResponsesWsFrameCheck(Protocol): + async def __call__(self, params: Mapping[str, object], /) -> str | None: ... + + class ProjectQuotaCallback(Protocol): async def enforce_project_io_token_quota_for_frame( self, @@ -2719,6 +2723,7 @@ class ManagedResponsesWebSocketHandler: custom_llm_provider: str | None = None, first_message: str | None = None, quota_callbacks: Sequence[ProjectQuotaCallback] | None = None, + frame_check: ResponsesWsFrameCheck | None = None, **kwargs: object, ) -> None: self.websocket = websocket @@ -2737,6 +2742,7 @@ class ManagedResponsesWebSocketHandler: self._connection_provider = self._resolve_provider(model) or custom_llm_provider self.first_message = first_message self.quota_callbacks: tuple[ProjectQuotaCallback, ...] = tuple(quota_callbacks) if quota_callbacks else () + self.frame_check = frame_check # Carry through safe pass-through kwargs (e.g. extra_headers) self.extra_kwargs: dict[str, object] = {k: v for k, v in kwargs.items() if k not in _MANAGED_WS_SKIP_KWARGS} # In-memory session history: response_id → full accumulated message list. @@ -3140,6 +3146,12 @@ class ManagedResponsesWebSocketHandler: verbose_logger.debug("ManagedResponsesWS: error sending warmup ack: %s", exc) return + call_kwargs: Final = self._build_base_call_kwargs(msg_obj) + refusal: Final = None if self.frame_check is None else await self.frame_check(call_kwargs) + if refusal is not None: + await self._send_error(refusal, "invalid_request_error") + return + try: await _enforce_frame_project_quota( self.quota_callbacks, self.user_api_key_dict, self.model_group or self.model, raw_message @@ -3148,7 +3160,6 @@ class ManagedResponsesWebSocketHandler: await self._send_error(str(e), error_type="rate_limit_exceeded") return - call_kwargs: Final = self._build_base_call_kwargs(msg_obj) call_kwargs["stream"] = True # A frame that repeats the connection's public alias (model_group) must diff --git a/litellm/types/files.py b/litellm/types/files.py index 259a836d9ad..083f5bf77d1 100644 --- a/litellm/types/files.py +++ b/litellm/types/files.py @@ -29,6 +29,7 @@ class FileType(Enum): JSON = "JSON" M4A = "M4A" M4V = "M4V" + MD = "MD" MOV = "MOV" MP3 = "MP3" MP4 = "MP4" @@ -75,6 +76,7 @@ FILE_EXTENSIONS: Final[Mapping[FileType, list[str]]] = MappingProxyType( FileType.JSON: ["json"], FileType.M4A: ["m4a"], FileType.M4V: ["m4v"], + FileType.MD: ["md", "markdown"], FileType.MOV: ["mov"], FileType.MP3: ["mp3"], FileType.MP4: ["mp4"], @@ -122,6 +124,7 @@ FILE_MIME_TYPES: Final[Mapping[FileType, str]] = MappingProxyType( FileType.JSON: "application/json", FileType.M4A: "audio/x-m4a", FileType.M4V: "video/x-m4v", + FileType.MD: "text/markdown", FileType.MOV: "video/quicktime", FileType.MP3: "audio/mpeg", FileType.MP4: "video/mp4", @@ -238,7 +241,7 @@ def is_audio_file_type(file_type): # Text -TEXT_FILE_TYPES: Final = {FileType.CSV, FileType.HTML, FileType.RTF, FileType.TXT} +TEXT_FILE_TYPES: Final = {FileType.CSV, FileType.HTML, FileType.MD, FileType.RTF, FileType.TXT} def is_text_file_type(file_type): diff --git a/litellm/types/llms/bedrock.py b/litellm/types/llms/bedrock.py index e7bef9b3946..20de3a40e8a 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 ChatCompletionImageUrlObject, ChatCompletionToolCallChunk class CachePointBlock(TypedDict, total=False): @@ -22,8 +22,18 @@ class SystemContentBlock(TypedDict, total=False): cachePoint: CachePointBlock -class SourceBlock(TypedDict): - bytes: str | None # base 64 encoded string +class S3Location(TypedDict): + uri: ReadOnly[str] + bucketOwner: NotRequired[ReadOnly[str]] + + +class SourceBlock(TypedDict, total=False): + bytes: ReadOnly[str | None] # base 64 encoded string + s3Location: ReadOnly[S3Location] + + +class S3MediaImageUrlObject(ChatCompletionImageUrlObject, total=False): + bucket_owner: ReadOnly[object] BedrockImageTypes = Literal["png", "jpeg", "gif", "webp"] @@ -68,6 +78,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/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 b31f45e506e..ddc37b80705 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 @@ -68,6 +68,7 @@ from litellm.litellm_core_utils.prompt_templates.common_utils import ( ) from litellm.llms.vertex_ai.gemini.transformation import gemini_convert_messages_with_history from litellm.types.llms.bedrock import BedrockConverseReasoningContentBlock +from litellm.types.llms.bedrock import ContentBlock as BedrockContentBlock from litellm.types.llms.openai import AllMessageValues, ChatCompletionUserMessage @@ -2481,10 +2482,10 @@ def test_bedrock_create_bedrock_block_deterministic_document_hash(): # Create two blocks with same content block1 = BedrockImageProcessor._create_bedrock_block( - image_bytes=base64_content, mime_type="application/pdf", image_format="pdf" + source={"bytes": base64_content}, mime_type="application/pdf", image_format="pdf" ) block2 = BedrockImageProcessor._create_bedrock_block( - image_bytes=base64_content, mime_type="application/pdf", image_format="pdf" + source={"bytes": base64_content}, mime_type="application/pdf", image_format="pdf" ) # Both should have the same document name @@ -2509,10 +2510,10 @@ def test_bedrock_create_bedrock_block_different_content_different_hash(): # Create blocks block1 = BedrockImageProcessor._create_bedrock_block( - image_bytes=base64_content1, mime_type="application/pdf", image_format="pdf" + source={"bytes": base64_content1}, mime_type="application/pdf", image_format="pdf" ) block2 = BedrockImageProcessor._create_bedrock_block( - image_bytes=base64_content2, mime_type="application/pdf", image_format="pdf" + source={"bytes": base64_content2}, mime_type="application/pdf", image_format="pdf" ) # Should have different document names @@ -2535,15 +2536,15 @@ def test_bedrock_create_bedrock_block_normalized_base64(): # Create blocks block1 = BedrockImageProcessor._create_bedrock_block( - image_bytes=base64_content, mime_type="application/pdf", image_format="pdf" + source={"bytes": base64_content}, mime_type="application/pdf", image_format="pdf" ) block2 = BedrockImageProcessor._create_bedrock_block( - image_bytes=base64_with_newlines, + source={"bytes": base64_with_newlines}, mime_type="application/pdf", image_format="pdf", ) block3 = BedrockImageProcessor._create_bedrock_block( - image_bytes=base64_with_spaces, + source={"bytes": base64_with_spaces}, mime_type="application/pdf", image_format="pdf", ) @@ -2570,10 +2571,10 @@ def test_bedrock_create_bedrock_block_large_file_sampling(): # Create blocks block1 = BedrockImageProcessor._create_bedrock_block( - image_bytes=base64_content1, mime_type="application/pdf", image_format="pdf" + source={"bytes": base64_content1}, mime_type="application/pdf", image_format="pdf" ) block2 = BedrockImageProcessor._create_bedrock_block( - image_bytes=base64_content2, mime_type="application/pdf", image_format="pdf" + source={"bytes": base64_content2}, mime_type="application/pdf", image_format="pdf" ) # Should have different names because total length is different @@ -2592,7 +2593,7 @@ def test_bedrock_create_bedrock_block_very_large_file(): # Create block block = BedrockImageProcessor._create_bedrock_block( - image_bytes=base64_content, mime_type="application/pdf", image_format="pdf" + source={"bytes": base64_content}, mime_type="application/pdf", image_format="pdf" ) # Should have a valid document name @@ -2614,7 +2615,7 @@ def test_bedrock_create_bedrock_block_image_type(): # Create block block = BedrockImageProcessor._create_bedrock_block( - image_bytes=base64_content, mime_type="image/png", image_format="png" + source={"bytes": base64_content}, mime_type="image/png", image_format="png" ) # Should be an image block, not a document block @@ -2635,7 +2636,7 @@ def test_bedrock_create_bedrock_block_video_type(): # Create block block = BedrockImageProcessor._create_bedrock_block( - image_bytes=base64_content, mime_type="video/mp4", image_format="mp4" + source={"bytes": base64_content}, mime_type="video/mp4", image_format="mp4" ) # Should be a video block, not a document block @@ -2656,7 +2657,7 @@ def test_bedrock_create_bedrock_block_document_name_format(): base64_content = base64.b64encode(pdf_content).decode("utf-8") block = BedrockImageProcessor._create_bedrock_block( - image_bytes=base64_content, mime_type="application/pdf", image_format="pdf" + source={"bytes": base64_content}, mime_type="application/pdf", image_format="pdf" ) document_name = block["document"]["name"] @@ -2685,7 +2686,7 @@ def test_bedrock_create_bedrock_block_different_document_formats(): for content, mime_type, format_type in test_cases: base64_content = base64.b64encode(content + b"\x00" * 100).decode("utf-8") block = BedrockImageProcessor._create_bedrock_block( - image_bytes=base64_content, mime_type=mime_type, image_format=format_type + source={"bytes": base64_content}, mime_type=mime_type, image_format=format_type ) assert block.get("document") is not None @@ -2694,6 +2695,365 @@ def test_bedrock_create_bedrock_block_different_document_formats(): assert block["document"]["format"] == format_type +_S3_MEDIA_MODEL: Final = "us.amazon.nova-2-lite-v1:0" +_S3_BUCKET_OWNER: Final = "111122223333" + + +async def _converse_user_content(part: Mapping[str, object], run_async: bool) -> Sequence[BedrockContentBlock]: + messages: Final = [{"role": "user", "content": [{"type": "text", "text": "Describe this."}, part]}] + with respx.mock: + converted: Final = ( + await BedrockConverseMessagesProcessor._bedrock_converse_messages_pt_async( + messages=messages, model=_S3_MEDIA_MODEL, llm_provider="bedrock" + ) + if run_async + else _bedrock_converse_messages_pt(messages=messages, model=_S3_MEDIA_MODEL, llm_provider="bedrock") + ) + return converted[0]["content"][1:] + + +def _s3_source(uri: str, bucket_owner: str | None = None) -> Mapping[str, object]: + return {"s3Location": {"uri": uri} if bucket_owner is None else {"uri": uri, "bucketOwner": bucket_owner}} + + +@pytest.mark.parametrize("run_async", [False, True], ids=["sync", "async"]) +@pytest.mark.parametrize( + ("part", "expected"), + [ + 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="image-format-from-extension", + ), + pytest.param( + {"type": "image_url", "image_url": "s3://media-bucket/dir/CAT.JPG"}, + {"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", + "image_url": {"url": "s3://media-bucket/dir/cat.png", "bucket_owner": _S3_BUCKET_OWNER}, + }, + {"image": {"format": "png", "source": _s3_source("s3://media-bucket/dir/cat.png", _S3_BUCKET_OWNER)}}, + id="bucket-owner", + ), + pytest.param( + {"type": "image_url", "image_url": {"url": "s3://media-bucket/dir/cat.png", "format": "image/webp"}}, + {"image": {"format": "webp", "source": _s3_source("s3://media-bucket/dir/cat.png")}}, + id="explicit-format-wins", + ), + pytest.param( + {"type": "image_url", "image_url": {"url": "s3://media-bucket/base64/cat.png"}}, + {"image": {"format": "png", "source": _s3_source("s3://media-bucket/base64/cat.png")}}, + id="base64-in-key", + ), + pytest.param( + {"type": "image_url", "image_url": {"url": "s3://media-bucket/clips/demo.mp4"}}, + {"video": {"format": "mp4", "source": _s3_source("s3://media-bucket/clips/demo.mp4")}}, + id="video", + ), + 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="video-extension-is-the-bedrock-format", + ), + pytest.param( + {"type": "file", "file": {"file_id": "s3://media-bucket/dir/cat.png", "bucket_owner": _S3_BUCKET_OWNER}}, + {"image": {"format": "png", "source": _s3_source("s3://media-bucket/dir/cat.png", _S3_BUCKET_OWNER)}}, + id="file-id", + ), + ], +) +async def test_bedrock_converse_s3_media_is_sent_as_an_s3_location_without_a_download( + part: Mapping[str, object], expected: Mapping[str, object], run_async: bool +): + 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" + invoice: Final = "s3://media-bucket/docs/invoice.pdf" + + report_content: Final = await _converse_user_content({"type": "file", "file": {"file_id": report}}, run_async) + report_again: Final = await _converse_user_content({"type": "file", "file": {"file_id": report}}, run_async) + invoice_content: Final = await _converse_user_content({"type": "file", "file": {"file_id": invoice}}, run_async) + + report_name: Final = report_content[0]["document"]["name"] + assert report_content == [{"document": {"format": "pdf", "name": report_name, "source": _s3_source(report)}}] + assert re.fullmatch(r"DocumentPDFmessages_[0-9a-f]{16}_pdf", report_name) + assert report_again == report_content + assert invoice_content[0]["document"]["name"] != report_name + + +@pytest.mark.parametrize("run_async", [False, True], ids=["sync", "async"]) +@pytest.mark.parametrize("uri", ["s3://media-bucket/docs/notes.md", "s3://media-bucket/docs/NOTES.markdown"]) +async def test_bedrock_converse_s3_markdown_without_a_format_is_an_md_document(uri: str, run_async: bool): + content: Final = await _converse_user_content({"type": "file", "file": {"file_id": uri}}, run_async) + + assert [{**block, "document": {**block["document"], "name": "-"}} for block in content] == [ + {"document": {"format": "md", "name": "-", "source": _s3_source(uri)}} + ] + + +@pytest.mark.parametrize("run_async", [False, True], ids=["sync", "async"]) +@pytest.mark.parametrize("uri", ["s3://media-bucket/uploads/7f3a9c", "s3://media-bucket/uploads/blob.unknownext"]) +async def test_bedrock_converse_s3_media_without_a_known_extension_or_format_is_a_bad_request( + uri: str, run_async: bool +): + with pytest.raises(litellm.BadRequestError) as excinfo: + await _converse_user_content({"type": "image_url", "image_url": {"url": uri}}, run_async) + + assert uri in excinfo.value.message + assert "`format`" in excinfo.value.message + + +@pytest.mark.parametrize("run_async", [False, True], ids=["sync", "async"]) +@pytest.mark.parametrize( + ("image_url", "named_format"), + [ + pytest.param({"url": "s3://media-bucket/scans/page.tiff"}, "page.tiff", id="tiff-key"), + 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( + image_url: Mapping[str, str], named_format: str, run_async: bool +): + with pytest.raises(litellm.BadRequestError) as excinfo: + await _converse_user_content({"type": "image_url", "image_url": image_url}, run_async) + + assert image_url["url"] in excinfo.value.message + 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", + [ + pytest.param( + { + "type": "image_url", + "image_url": {"url": "s3://media-bucket/dir/cat.png", "bucket_owner": _S3_BUCKET_OWNER}, + }, + id="image-url", + ), + pytest.param( + {"type": "file", "file": {"file_id": "s3://media-bucket/dir/cat.png", "bucket_owner": _S3_BUCKET_OWNER}}, + id="file", + ), + ], +) +def test_bedrock_converse_tool_result_s3_media_keeps_its_bucket_owner(part: Mapping[str, object]): + messages: Final = [ + {"role": "user", "content": "Describe the screenshot."}, + {"tool_call_id": "tooluse_s3_1", "role": "tool", "name": "fetch_screenshot", "content": [part]}, + ] + + with respx.mock: + converted: Final = _bedrock_converse_messages_pt( + messages=messages, model=_S3_MEDIA_MODEL, llm_provider="bedrock" + ) + + assert converted[-1]["content"][-1]["toolResult"]["content"] == [ + {"image": {"format": "png", "source": _s3_source("s3://media-bucket/dir/cat.png", _S3_BUCKET_OWNER)}} + ] + + +@pytest.mark.parametrize( + "file, expected_media", + [ + pytest.param( + {"file_id": "s3://media-bucket/clips/demo.mp4"}, + [{"video": {"format": "mp4", "source": _s3_source("s3://media-bucket/clips/demo.mp4")}}], + id="s3-video-is-forwarded", + ), + pytest.param( + {"file_data": "data:video/mp4;base64,AAAAGGZ0eXBtcDQy"}, + [], + id="inline-video-is-dropped", + ), + ], +) +def test_bedrock_converse_tool_result_video_is_forwarded_only_from_s3( + file: Mapping[str, object], expected_media: list[Mapping[str, object]] +): + messages: Final = [ + {"role": "user", "content": "Describe the clip."}, + { + "tool_call_id": "tooluse_video_1", + "role": "tool", + "name": "fetch_clip", + "content": [{"type": "text", "text": "clip:"}, {"type": "file", "file": file}], + }, + ] + + with respx.mock: + converted: Final = _bedrock_converse_messages_pt( + messages=messages, model=_S3_MEDIA_MODEL, llm_provider="bedrock" + ) + + assert converted[-1]["content"][-1]["toolResult"]["content"] == [{"text": "clip:"}, *expected_media] + + +@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", + "format": "image/webp", + "bucket_owner": _S3_BUCKET_OWNER, + }, + } + ], + }, + {"role": "user", "content": "Make it bigger."}, + ] + + with respx.mock: + converted: Final = ( + await BedrockConverseMessagesProcessor._bedrock_converse_messages_pt_async( + messages=messages, model=_S3_MEDIA_MODEL, llm_provider="bedrock" + ) + if run_async + else _bedrock_converse_messages_pt(messages=messages, model=_S3_MEDIA_MODEL, llm_provider="bedrock") + ) + + assert converted[1]["content"] == [ + {"image": {"format": "webp", "source": _s3_source("s3://media-bucket/dir/cat.png", _S3_BUCKET_OWNER)}} + ] + + def test_bedrock_nova_web_search_options_mapping(): """ Test that web_search_options is correctly mapped to Nova grounding. 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 c3a2511a263..5e711d58d28 100644 --- a/tests/unit/litellm_core_utils/test_cloud_storage_security.py +++ b/tests/unit/litellm_core_utils/test_cloud_storage_security.py @@ -1,5 +1,18 @@ +import sys +from base64 import urlsafe_b64encode +from functools import reduce +from typing import Final, cast +from urllib.parse import quote + +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, + media_file_ids, + names_an_unlisted_s3_bucket, + s3_media_uris, ) @@ -13,3 +26,407 @@ def test_is_managed_cloud_storage_uri_ignores_provider_and_unified_ids(): assert not is_managed_cloud_storage_uri("file-abc123") assert not is_managed_cloud_storage_uri("bGl0ZWxsbV9wcm94eQ==") assert not is_managed_cloud_storage_uri("") + + +def _quote_once(current: str, _: int) -> str: + return quote(current, safe="") + + +def _quoted_times(value: str, times: int) -> str: + return reduce(_quote_once, range(times), value) + + +_MANAGED_UPLOAD: Final = "s3://media-bucket/litellm-bedrock-files/other-team.pdf" + + +def _unpadded_base64url(text: str) -> str: + return urlsafe_b64encode(text.encode()).decode().rstrip("=") + + +_MODEL_EMBEDDED_UPLOAD_ID: Final = "file-" + _unpadded_base64url(f"litellm:{_MANAGED_UPLOAD};model,nova") +_UNIFIED_BATCH_OUTPUT_ID: Final = _unpadded_base64url( + f"litellm_proxy:application/json;unified_id,u1;target_model_names,nova;llm_output_file_id,{_MANAGED_UPLOAD};" +) + + +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)}]} + + +def _gemini_turn(*parts: object) -> dict[str, object]: + return {"contents": [{"role": "user", "parts": 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( + _user_turn({"type": "file", "file": {"file_id": _MODEL_EMBEDDED_UPLOAD_ID}}), + id="chat-model-embedded-file-id", + ), + pytest.param( + _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"}}), + "model_file_id_mapping": {"alias": {"deployment-1": _MANAGED_UPLOAD}}, + }, + id="model-file-id-mapping", + ), + pytest.param( + { + **_responses_turn({"type": "input_file", "file_id": "alias"}), + "model_file_id_mapping": {"alias": {"deployment-1": _MODEL_EMBEDDED_UPLOAD_ID}}, + }, + id="model-file-id-mapping-to-a-model-embedded-file-id", + ), + 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", + ), + 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( + { + "input": [ + { + "type": "computer_call_output", + "call_id": "call_1", + "output": {"type": "computer_screenshot", "image_url": _MANAGED_UPLOAD}, + } + ] + }, + id="responses-computer-screenshot", + ), + pytest.param( + { + "messages": ( + {"role": "user", "content": ({"type": "image_url", "image_url": {"url": _MANAGED_UPLOAD}},)}, + ) + }, + id="tuples", + ), + pytest.param( + _user_turn(1, 2.5, True, None, {"type": "image_url", "image_url": {"url": _MANAGED_UPLOAD}}), + id="numbers-before-the-media-part", + ), + ), +) +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:///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": "file", + "file": { + "file_id": _unpadded_base64url( + "litellm_proxy:application/pdf;unified_id,u2;target_model_names,nova;" + "llm_output_file_id,file-abc123;" + ) + }, + } + ), + id="unified-file-id-of-a-provider-file", + ), + pytest.param( + _user_turn( + { + "type": "file", + "file": {"file_data": "data:application/pdf;base64," + _unpadded_base64url(_MANAGED_UPLOAD)}, + } + ), + id="file-data-payload-naming-an-s3-object", + ), + pytest.param( + _user_turn({"type": "text", "text": _MODEL_EMBEDDED_UPLOAD_ID}), id="encoded-file-id-outside-a-media-part" + ), + 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": {"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": "file", + "file": {"file_data": "data:application/pdf;base64,JVBERi0=", "filename": _MANAGED_UPLOAD}, + } + ), + id="media-part-filename", + ), + 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_finds_an_s3_object_nested_deeper_than_the_recursion_limit(): + nested: Final = reduce( + lambda inner, _: [inner], + range(sys.getrecursionlimit() + 100), + cast(object, [{"type": "image_url", "image_url": {"url": _MANAGED_UPLOAD}}]), + ) + + assert s3_media_uris({"messages": [{"role": "user", "content": nested}]}) == (_MANAGED_UPLOAD,) + + +def test_s3_media_uris_skips_only_the_objects_a_trusted_file_id_embeds(): + listed_image: Final = "s3://media-bucket/dir/cat.png" + request_body: Final = _responses_turn( + {"type": "input_file", "file_id": _UNIFIED_BATCH_OUTPUT_ID}, + {"type": "input_file", "file_id": _MODEL_EMBEDDED_UPLOAD_ID}, + {"type": "input_image", "image_url": listed_image}, + ) + + assert s3_media_uris(request_body, frozenset({_UNIFIED_BATCH_OUTPUT_ID})) == (_MANAGED_UPLOAD, listed_image) + assert s3_media_uris(request_body, frozenset({_UNIFIED_BATCH_OUTPUT_ID, _MODEL_EMBEDDED_UPLOAD_ID})) == ( + listed_image, + ) + + +def test_media_file_ids_lists_the_file_ids_of_media_parts_only(): + request_body: Final = { + **_user_turn( + {"type": "file", "file": {"file_id": "file-a"}}, + {"type": "text", "text": "file-b", "file_id": "file-b"}, + {"type": "file", "file": {"file_id": "file-a"}}, + ), + "metadata": {"file_id": "file-c"}, + } + + 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( + {"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"), + ( + 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 69269e31e0b..955e752d924 100644 --- a/tests/unit/llms/bedrock/chat/test_converse_transformation.py +++ b/tests/unit/llms/bedrock/chat/test_converse_transformation.py @@ -362,6 +362,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/llms/bedrock/files/test_bedrock_files_transformation.py b/tests/unit/llms/bedrock/files/test_bedrock_files_transformation.py index 2410a594e56..ec63e73e1f1 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 BedrockFilesConfig, BedrockJsonlFilesTransformation from tests.unit.proxy.conftest import httpx_transport @@ -4402,3 +4403,139 @@ 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}}]}], + }, + } +) +_S3_IMAGE_PART: Final = {"type": "image_url", "image_url": {"url": _BATCH_S3_IMAGE}} +_OTHER_BATCH_S3_IMAGE: Final = "s3://media-bucket/dir/dog.png" +_BATCH_TWO_S3_MEDIA_REFUSAL: Final = _BATCH_S3_MEDIA_REFUSAL.replace( + f"({_BATCH_S3_IMAGE})", f"({_BATCH_S3_IMAGE}, {_OTHER_BATCH_S3_IMAGE})" +) + + +def _nova_chat_batch_line(custom_id: str, **body: object) -> str: + return json.dumps( + { + "custom_id": custom_id, + "method": "POST", + "url": "/v1/chat/completions", + "body": {"model": "bedrock/amazon.nova-pro-v1:0", **body}, + } + ) + + +_CHAT_BATCH_LINE_WITH_THE_SAME_S3_IMAGE_TWICE: Final = _nova_chat_batch_line( + "req-2", messages=[{"role": "user", "content": [_S3_IMAGE_PART, _S3_IMAGE_PART]}] +) +_CHAT_BATCH_LINE_WITH_TWO_S3_IMAGES: Final = _nova_chat_batch_line( + "req-2", + messages=[ + { + "role": "user", + "content": [_S3_IMAGE_PART, {"type": "image_url", "image_url": {"url": _OTHER_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", "expected_message"), + ( + pytest.param(_CHAT_BATCH_LINE_WITH_AN_S3_IMAGE, _BATCH_S3_MEDIA_REFUSAL, id="chat-s3-image"), + pytest.param(_RESPONSES_BATCH_LINE_WITH_AN_S3_IMAGE, _BATCH_S3_MEDIA_REFUSAL, id="responses-s3-image"), + pytest.param( + _CHAT_BATCH_LINE_WITH_THE_SAME_S3_IMAGE_TWICE, _BATCH_S3_MEDIA_REFUSAL, id="chat-repeated-s3-image" + ), + pytest.param(_CHAT_BATCH_LINE_WITH_TWO_S3_IMAGES, _BATCH_TWO_S3_MEDIA_REFUSAL, id="chat-two-s3-images"), + ), +) +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, expected_message: 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 == expected_message + 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 = _nova_chat_batch_line( + "req-1", model="bedrock/openai.gpt-oss-120b-1:0", max_tokens=10, messages=passthrough_messages + ) + metadata_only_line: Final = _nova_chat_batch_line( + "req-2", messages=[{"role": "user", "content": "Hello"}], max_tokens=10, 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().splitlines() + ) + assert passthrough_record["modelInput"]["messages"] == passthrough_messages + assert metadata_only_record["recordId"] == "req-2" 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 19eb3c71a5b..45d217dd5da 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 @@ -5,6 +5,7 @@ This test file tests the Generic Guardrail API implementation, specifically focusing on metadata extraction and passing. """ +import json import os from typing import Final from unittest.mock import AsyncMock, MagicMock, patch @@ -16,6 +17,7 @@ import litellm from litellm import ModelResponse from litellm._version import version as litellm_version from litellm.exceptions import GuardrailRaisedException, Timeout +from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler from litellm.proxy._types import UserAPIKeyAuth from litellm.proxy.guardrails.guardrail_hooks.generic_guardrail_api import ( GenericGuardrailAPI, @@ -26,6 +28,26 @@ from litellm.proxy.guardrails.guardrail_hooks.generic_guardrail_api.generic_guar from litellm.types.utils import Choices, Message +async def _structured_messages_sent(guardrail: GenericGuardrailAPI, rows: list[dict]) -> list: + sent_bodies: Final[list[bytes]] = [] # mutable-ok: records every body the guardrail endpoint received + + async def answer(request: httpx.Request) -> httpx.Response: + sent_bodies.append(request.content) + return httpx.Response(200, json={"action": "NONE"}, request=request) + + endpoint: Final = AsyncHTTPHandler(transport=httpx.MockTransport(answer)) + try: + with patch.object(guardrail, "async_handler", endpoint): + await guardrail.apply_guardrail( + inputs={"texts": ["What is written here?"], "structured_messages": rows}, + request_data={}, + input_type="request", + ) + finally: + await endpoint.close() + return [json.loads(body)["structured_messages"] for body in sent_bodies] + + @pytest.fixture def generic_guardrail(): """Create a GenericGuardrailAPI instance for testing""" @@ -672,6 +694,30 @@ class TestStructuredMessagesInResponse: assert "structured_messages" not in guardrailed_inputs assert guardrailed_inputs["texts"] == ["Never repeat an SSN.", "Look up [REDACTED] for me."] + @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: Final = { + "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( "structured_messages", diff --git a/tests/unit/proxy/guardrails/guardrail_hooks/test_model_armor.py b/tests/unit/proxy/guardrails/guardrail_hooks/test_model_armor.py index 5aa3c7dfc17..c0901f83853 100644 --- a/tests/unit/proxy/guardrails/guardrail_hooks/test_model_armor.py +++ b/tests/unit/proxy/guardrails/guardrail_hooks/test_model_armor.py @@ -3,13 +3,11 @@ import base64 import io import json from collections.abc import Iterator, Sequence -from typing import cast +from typing import Final, cast from unittest.mock import AsyncMock, MagicMock, Mock, patch -import pytest - - import httpx +import pytest from fastapi import HTTPException import litellm @@ -17,13 +15,18 @@ import litellm.types.utils from litellm._logging import verbose_proxy_logger from litellm.caching import DualCache from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, MaskedHTTPStatusError -from litellm.proxy.guardrails.anthropic_sse import anthropic_sse_error_frames from litellm.proxy._types import UserAPIKeyAuth +from litellm.proxy.guardrails.anthropic_sse import anthropic_sse_error_frames from litellm.proxy.guardrails.guardrail_hooks.model_armor import ModelArmorGuardrail +from litellm.proxy.guardrails.guardrail_hooks.model_armor.file_scanning import ( + FileScanPlan, + plan_file_scans, +) from litellm.proxy.guardrails.guardrail_hooks.model_armor.model_armor import ( ModelArmorAPIError, ) from litellm.types.guardrails import GuardrailEventHooks +from litellm.types.llms.openai import AllMessageValues @pytest.mark.asyncio @@ -2670,6 +2673,44 @@ async def test_pre_call_blocks_unresolvable_file_id_reference(): assert "could not scan" in str(exc_info.value.detail) +@pytest.mark.asyncio +@pytest.mark.parametrize( + "uri", + [ + "s3://listed-bucket/report.pdf", + " S3://listed-bucket/report.pdf ", + "gs://listed-bucket/report.pdf", + "https://example.com/report.pdf", + "http://example.com/report.pdf", + ], +) +async def test_pre_call_blocks_remote_document_without_format_or_filename(uri: str) -> None: + messages: Final[list[AllMessageValues]] = [ + {"role": "user", "content": [{"type": "file", "file": {"file_data": uri}}]} + ] + assert plan_file_scans(messages) == FileScanPlan(attachments=(), unscannable_count=1) + + guardrail: Final = ModelArmorGuardrail( + template_id="test-template", + project_id="test-project", + guardrail_name="model-armor-test", + fail_on_error=True, + ) + data: Final = {"messages": messages, "metadata": {"guardrails": ["model-armor-test"]}} + + with pytest.raises(HTTPException) as exc_info: + await guardrail.async_pre_call_hook( + user_api_key_dict=UserAPIKeyAuth(), + cache=DualCache(), + data=data, + call_type="completion", + ) + + assert exc_info.value.status_code == 400 + assert "could not scan" in str(exc_info.value.detail) + assert data["metadata"]["_model_armor_status"] == "blocked" + + @pytest.mark.asyncio async def test_pre_call_blocks_remote_url_document_reference(): """A remote (https) document reference cannot be fetched here, so it fails closed by default.""" diff --git a/tests/unit/proxy/pass_through_endpoints/test_pass_through_endpoints.py b/tests/unit/proxy/pass_through_endpoints/test_pass_through_endpoints.py index 53bff9dd30c..e8c4161c6b1 100644 --- a/tests/unit/proxy/pass_through_endpoints/test_pass_through_endpoints.py +++ b/tests/unit/proxy/pass_through_endpoints/test_pass_through_endpoints.py @@ -30,6 +30,7 @@ from starlette.datastructures import UploadFile as StarletteUploadFile import litellm from litellm._logging import verbose_proxy_logger +from litellm.caching.dual_cache import DualCache from litellm.constants import DEFAULT_REQUEST_TIMEOUT_SECONDS from litellm.integrations.custom_guardrail import CustomGuardrail from litellm.integrations.custom_logger import CustomLogger @@ -61,6 +62,7 @@ from litellm.proxy.pass_through_endpoints.success_handler import ( PassThroughEndpointLogging, ) from litellm.proxy.route_llm_request import ProxyModelNotFoundError +from litellm.proxy.utils import ProxyLogging from litellm.types import utils as types_utils from litellm.types.guardrails import GuardrailEventHooks from litellm.types.passthrough_endpoints.pass_through_endpoints import ( @@ -10328,3 +10330,101 @@ def test_init_kwargs_no_authenticated_identity_field_is_client_settable(mock_req metadata: Final = result["litellm_params"]["metadata"] survived: Final = {key: metadata.get(key) for key in expected if metadata.get(key) != expected[key]} assert survived == {}, f"client-supplied values survived for: {sorted(survived)}" + + +def _stub_adapter_proxy(monkeypatch: pytest.MonkeyPatch) -> MagicMock: + proxy_logging: Final = MagicMock() + proxy_logging.pre_call_hook = AsyncMock(side_effect=lambda **kwargs: kwargs["data"]) + proxy_logging.post_call_failure_hook = AsyncMock() + + async def fake_add_litellm_data_to_request(**kwargs: object) -> object: + return kwargs["data"] + + monkeypatch.setattr("litellm.proxy.proxy_server.proxy_logging_obj", proxy_logging) + monkeypatch.setattr("litellm.proxy.proxy_server.add_litellm_data_to_request", fake_add_litellm_data_to_request) + monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", None) + monkeypatch.setattr("litellm.proxy.proxy_server.user_model", None) + monkeypatch.setattr("litellm.proxy.proxy_server.general_settings", {}) + return proxy_logging + +def _request_with_body(content_type: str, body: bytes) -> Request: + async def receive() -> dict[str, object]: + return {"type": "http.request", "body": body, "more_body": False} + + return Request( + { + "type": "http", + "method": "POST", + "path": "/pass-through", + "query_string": b"", + "headers": [(b"content-type", content_type.encode())], + }, + receive, + ) + +_ADAPTER_BODY_WITH_S3_IMAGE: Final = json.dumps( + { + "model": "unknown-model", + "messages": [ + {"role": "user", "content": [{"type": "image_url", "image_url": {"url": "s3://private-bucket/cat.png"}}]} + ], + } +).encode() +_ADAPTER_PYTHON_LITERAL_BODY_WITH_S3_IMAGE_IN_TUPLES: Final = repr( + { + "model": "unknown-model", + "messages": ( + {"role": "user", "content": ({"type": "image_url", "image_url": {"url": "s3://private-bucket/cat.png"}},)}, + ), + } +).encode() + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "body", + [_ADAPTER_BODY_WITH_S3_IMAGE, _ADAPTER_PYTHON_LITERAL_BODY_WITH_S3_IMAGE_IN_TUPLES], + ids=["json", "python_literal_with_tuples"], +) +async def test_chat_completion_pass_through_endpoint_refuses_an_s3_image_in_an_unlisted_bucket( + monkeypatch: pytest.MonkeyPatch, + body: bytes, +): + proxy_logging: Final = _stub_adapter_proxy(monkeypatch) + request: Final = MagicMock(spec=Request) + request.body = AsyncMock(return_value=body) + + with pytest.raises(ProxyException) as raised: + await chat_completion_pass_through_endpoint( + fastapi_response=Response(), + request=request, + adapter_id="anthropic", + user_api_key_dict=UserAPIKeyAuth(api_key="sk-test"), + ) + + assert raised.value.code == "403" + proxy_logging.pre_call_hook.assert_not_called() + + +@pytest.mark.asyncio +async def test_chat_completion_pass_through_endpoint_routes_an_s3_image_in_a_listed_bucket( + monkeypatch: pytest.MonkeyPatch, +): + monkeypatch.setattr("litellm.proxy.proxy_server.proxy_logging_obj", ProxyLogging(user_api_key_cache=DualCache())) + monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", None) + monkeypatch.setattr("litellm.proxy.proxy_server.user_model", None) + monkeypatch.setattr( + "litellm.proxy.proxy_server.general_settings", {"allowed_client_side_s3_buckets": ["private-bucket"]} + ) + request: Final = _request_with_body("application/json", _ADAPTER_BODY_WITH_S3_IMAGE) + + with pytest.raises(ProxyException) as raised: + await chat_completion_pass_through_endpoint( + fastapi_response=Response(), + request=request, + adapter_id="anthropic", + user_api_key_dict=UserAPIKeyAuth(api_key="sk-test"), + ) + + assert raised.value.code == "400" + assert raised.value.message == str(ProxyModelNotFoundError(route="completion", model_name="unknown-model")) diff --git a/tests/unit/proxy/rag_endpoints/test_rag_endpoints.py b/tests/unit/proxy/rag_endpoints/test_rag_endpoints.py index 2d654ea28ec..680813bcc3e 100644 --- a/tests/unit/proxy/rag_endpoints/test_rag_endpoints.py +++ b/tests/unit/proxy/rag_endpoints/test_rag_endpoints.py @@ -7,9 +7,11 @@ Covers: import io import json +from typing import Final from unittest.mock import AsyncMock, MagicMock, patch import pytest +import respx from fastapi.testclient import TestClient import litellm @@ -1174,3 +1176,64 @@ class TestVectorStoreUploadControls: assert server_filename.endswith(".txt") assert secured_content_type == "text/plain" assert content_bytes == b"benign document text\n" + + +@pytest.mark.usefixtures("httpx_transport") +@pytest.mark.parametrize("listed, expected_status", [(False, 403), (True, 200)]) +def test_rag_query_checks_the_s3_image_bucket_before_search( + client_internal_user: TestClient, + respx_mock: respx.MockRouter, + monkeypatch: pytest.MonkeyPatch, + listed: bool, + expected_status: int, +) -> None: + monkeypatch.setenv("OPENAI_API_KEY", "sk-test") + search: Final = respx_mock.post("https://api.openai.com/v1/vector_stores/provider-native-store/search").respond( + json={ + "object": "vector_store.search_results.page", + "search_query": ["hello"], + "data": [], + "has_more": False, + "next_page": None, + } + ) + completion: Final = respx_mock.post("https://api.openai.com/v1/chat/completions").respond( + json={ + "id": "chatcmpl-test", + "object": "chat.completion", + "created": 1, + "model": "gpt-4o-mini", + "choices": [{"index": 0, "message": {"role": "assistant", "content": "hello"}, "finish_reason": "stop"}], + "usage": {"prompt_tokens": 1, "completion_tokens": 1, "total_tokens": 2}, + } + ) + with ( + patch( + "litellm.proxy.proxy_server.general_settings", + {"allowed_client_side_s3_buckets": ["private-bucket"]} if listed else {}, + ), + patch("litellm.proxy.proxy_server.llm_router", None), + patch("litellm.proxy.proxy_server.prisma_client", None), + patch("litellm.vector_store_registry", None), + ): + response: Final = client_internal_user.post( + "/v1/rag/query", + json={ + "model": "gpt-4o-mini", + "messages": [ + { + "role": "user", + "content": [ + {"type": "text", "text": "hello"}, + {"type": "image_url", "image_url": {"url": "s3://private-bucket/cat.png"}}, + ], + } + ], + "retrieval_config": {"vector_store_id": "provider-native-store", "custom_llm_provider": "openai"}, + }, + ) + assert response.status_code == expected_status, response.text + assert search.call_count == int(listed) + assert completion.call_count == int(listed) + if not listed: + assert "allowed_client_side_s3_buckets" in response.json()["detail"] diff --git a/tests/unit/proxy/response_api_endpoints/test_endpoints.py b/tests/unit/proxy/response_api_endpoints/test_endpoints.py index bde570a8f61..077d7e2aa3c 100644 --- a/tests/unit/proxy/response_api_endpoints/test_endpoints.py +++ b/tests/unit/proxy/response_api_endpoints/test_endpoints.py @@ -1612,6 +1612,81 @@ class TestManagedResponsesWSFirstMessage: assert processed == [subsequent] +class TestResponsesWSS3MediaFrameCheck: + @pytest.mark.asyncio + async def test_managed_socket_refuses_a_non_admin_s3_image_on_every_frame_and_forwards_a_proxy_admin_frame( + self, monkeypatch: pytest.MonkeyPatch + ): + from litellm import Router + from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth + from litellm.proxy.response_api_endpoints.endpoints import _responses_websocket_session + + image_uri: Final = "s3://media-bucket/dir/cat.png" + frame: Final = json.dumps( + { + "type": "response.create", + "input": [{"role": "user", "content": [{"type": "input_image", "image_url": image_uri}]}], + } + ) + forwarded_event: Final = {"type": "response.created", "response": {"id": "resp_admin"}} + provider_inputs: Final[list[object]] = [] + + async def fake_aresponses(**kwargs: object): + provider_inputs.append(kwargs["input"]) + + async def events(): + yield forwarded_event + + return events() + + monkeypatch.setattr(litellm, "aresponses", fake_aresponses) + test_router: Final = Router( + model_list=[{"model_name": "nova", "litellm_params": {"model": "bedrock/us.amazon.nova-2-lite-v1:0"}}] + ) + + async def route_to_test_router(*, data: dict[str, Any], **_kwargs: object): + return test_router._aresponses_websocket(**data) + + async def frames_sent_back(user_role: LitellmUserRoles) -> list[object]: + ws = MagicMock() + ws.scope = {"headers": []} + ws.url = "ws://testserver/v1/responses" + ws.receive_text = AsyncMock(side_effect=[frame, frame, Exception("disconnect")]) + ws.send_text = AsyncMock() + ws.close = AsyncMock() + with ( + patch( # test-quality-ok: the pre-call processor needs a live proxy and is not what this test covers + "litellm.proxy.response_api_endpoints.endpoints.ProxyBaseLLMRequestProcessing", + side_effect=lambda data: MagicMock( + common_processing_pre_call_logic=AsyncMock(return_value=(data, MagicMock())) + ), + ), + patch( # test-quality-ok: routes to a router holding one Bedrock deployment instead of the proxy's + "litellm.proxy.route_llm_request.route_request", + side_effect=route_to_test_router, + ), + ): + await _responses_websocket_session( + websocket=ws, + model="nova", + user_api_key_dict=UserAPIKeyAuth(api_key="sk-test", user_id="user-a", user_role=user_role), + ) + return [json.loads(call.args[0]) for call in ws.send_text.await_args_list] + + refused: Final = await frames_sent_back(LitellmUserRoles.INTERNAL_USER) + + assert provider_inputs == [] + assert [event["type"] for event in refused] == ["error", "error"] + assert all("allowed_client_side_s3_buckets" in event["error"]["message"] for event in refused) + assert all(image_uri in event["error"]["message"] for event in refused) + + admitted: Final = await frames_sent_back(LitellmUserRoles.PROXY_ADMIN) + + assert admitted == [forwarded_event, forwarded_event] + assert len(provider_inputs) == 2 + assert all(image_uri in json.dumps(provider_input) for provider_input in provider_inputs) + + class TestResponsesWSStreamingFirstMessage: @pytest.mark.asyncio async def test_client_to_backend_replays_first_message(self): diff --git a/tests/unit/proxy/test_common_request_processing.py b/tests/unit/proxy/test_common_request_processing.py index 8da89598eff..e48624cd6b2 100644 --- a/tests/unit/proxy/test_common_request_processing.py +++ b/tests/unit/proxy/test_common_request_processing.py @@ -2,9 +2,10 @@ import asyncio import copy import datetime import json +from base64 import urlsafe_b64encode 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 @@ -25,6 +26,7 @@ from litellm.litellm_core_utils.bug_report import ( build_bug_report, ) from litellm.types.router import RouterRateLimitError, RouterRateLimitErrorBasic +from litellm.types.utils import SpecialEnums from litellm.constants import ( CLIENT_REQUESTED_MODEL_SCOPE_KEY, MAX_LITELLM_CALL_ID_LENGTH, @@ -63,9 +65,10 @@ from litellm.proxy.common_request_processing import ( from litellm.proxy.common_utils.callback_utils import add_guardrail_to_applied_guardrails_header from litellm.proxy.common_utils.proxy_rate_limit_error import ProxyRateLimitError from litellm.proxy.common_utils.sse_keepalive import ANTHROPIC_PING_SSE_CHUNK +from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache from litellm.proxy.dd_span_tagger import DDSpanTagger from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj -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 ( @@ -204,6 +207,309 @@ 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 _chat_image_after_numbers(uri: str) -> dict[str, object]: + image_part: Final = {"type": "image_url", "image_url": {"url": uri}} + return {"messages": [{"role": "user", "content": [1, 2.5, True, None, image_part]}]} + + +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}}]}], + } + + +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("acompletion", "/v1/chat/completions", _chat_image_after_numbers, id="chat-completions-numbers-first"), + 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"), +) + + +def _chat_file(file_id: str) -> dict[str, object]: + return {"messages": [{"role": "user", "content": [{"type": "file", "file": {"file_id": file_id}}]}]} + + +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 = ( + urlsafe_b64encode( + SpecialEnums.LITELLM_MANAGED_FILE_COMPLETE_STR.value.format( + "application/pdf", "u1", "nova", _UPLOADED_OBJECT, "deployment-1" + ).encode() + ) + .decode() + .rstrip("=") +) + + +class _ManagedFilesCreatedBy(CustomLogger): + def __init__(self, creators: Mapping[str, str]) -> None: + super().__init__() + self.creators: Final = creators + + async def can_user_call_unified_file_id( + self, unified_file_id: str, user_api_key_dict: ProxyUserAPIKeyAuth + ) -> bool: + if unified_file_id not in self.creators: + raise HTTPException(status_code=status.HTTP_404_NOT_FOUND, detail=f"File not found: {unified_file_id}") + return self.creators[unified_file_id] == user_api_key_dict.user_id + + async def can_user_call_unified_object_id( + self, unified_object_id: str, user_api_key_dict: ProxyUserAPIKeyAuth + ) -> bool: + return False + + +class _ManagedFilesCountingLookups(_ManagedFilesCreatedBy): + def __init__(self, creators: Mapping[str, str]) -> None: + super().__init__(creators) + self.looked_up: Final[list[str]] = [] # mutable-ok: the fake records the lookups the gate makes + + async def can_user_call_unified_file_id( + self, unified_file_id: str, user_api_key_dict: ProxyUserAPIKeyAuth + ) -> bool: + self.looked_up.append(unified_file_id) + return await super().can_user_call_unified_file_id(unified_file_id, user_api_key_dict) + + +def _unified_file_id_for(s3_uri: str) -> str: + managed_id: Final = SpecialEnums.LITELLM_MANAGED_FILE_COMPLETE_STR.value.format( + "application/pdf", "u1", "nova", s3_uri, "deployment-1" + ) + return urlsafe_b64encode(managed_id.encode()).decode().rstrip("=") + + +def _chat_files(file_ids: Sequence[str]) -> dict[str, object]: + parts: Final = [{"type": "file", "file": {"file_id": file_id}} for file_id in file_ids] + return {"messages": [{"role": "user", "content": parts}]} + + +_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, + body: dict[str, object], + user_role: LitellmUserRoles, + general_settings: dict[str, object], + managed_files: CustomLogger | None = None, + client_auth_params: tuple[str, ...] = (), + deployment_aws_params: Mapping[str, str] = MappingProxyType({}), +) -> object: + from litellm.proxy.proxy_server import ProxyConfig as RealProxyConfig + + proxy_logging_obj: Final = ProxyLogging(user_api_key_cache=UserApiKeyCache()) + if managed_files is not None: + proxy_logging_obj.proxy_hook_mapping["managed_files"] = managed_files + + router: Final = Router( + model_list=[ + { + "model_name": "nova", + "litellm_params": { + "model": "bedrock/us.amazon.nova-2-lite-v1:0", + "mock_response": "a cat", + "configurable_clientside_auth_params": list(client_auth_params), + **deployment_aws_params, + }, + } + ] + ) + 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_id="user-a", user_role=user_role), + route_type=route_type, + proxy_logging_obj=proxy_logging_obj, + 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( @@ -480,6 +786,37 @@ 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, 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 = {} + + pre_call: Final = 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, + ) + + if not refused: + await pre_call + return + with pytest.raises(HTTPException) as refusal: + await pre_call + assert refusal.value.status_code == 403 + assert "s3://other-bucket/cat.png" in str(refusal.value.detail) + @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( @@ -509,6 +846,304 @@ 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://media-bucket/cat.png", + {}, + "general_settings.allowed_client_side_s3_buckets", + id="no-bucket-listed", + ), + 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/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: _S3MediaRoute, + 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: _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, + ) + + assert "a cat" in str(response) + + @pytest.mark.asyncio + @pytest.mark.parametrize( + ("identity", "deployment_aws_params"), + [ + pytest.param({"aws_role_name": "arn:aws:iam::111122223333:role/reader"}, {}, id="role"), + pytest.param( + {"aws_web_identity_token": "oidc-token"}, + {"aws_role_name": "arn:aws:iam::111122223333:role/web", "aws_session_name": "proxy"}, + id="web-identity-token-with-the-deployment-role-and-session", + ), + pytest.param({"aws_profile_name": "reader"}, {}, id="profile"), + ], + ) + async def test_non_admin_s3_media_needs_a_listed_bucket_even_with_a_client_aws_identity_the_deployment_admits( + self, identity: dict[str, str], deployment_aws_params: dict[str, str] + ) -> None: + with pytest.raises(HTTPException) as refused: + await _send_s3_media( + "acompletion", + "/v1/chat/completions", + {**_chat_image("s3://media-bucket/dir/cat.png"), **identity}, + LitellmUserRoles.INTERNAL_USER, + {}, + client_auth_params=tuple(identity), + deployment_aws_params=deployment_aws_params, + ) + + assert refused.value.status_code == status.HTTP_403_FORBIDDEN, refused.value.detail + assert "general_settings.allowed_client_side_s3_buckets" in str(refused.value.detail) + + @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 + @pytest.mark.parametrize(("route_type", "path", "body_for"), _S3_FILE_ID_SURFACES) + async def test_non_admin_managed_file_id_its_user_created_passes( + self, route_type: _S3MediaRoute, path: str, body_for: Callable[[str], dict[str, object]] + ) -> None: + response: Final = await _send_s3_media( + route_type, + path, + body_for(_UPLOADED_FILE_ID), + LitellmUserRoles.INTERNAL_USER, + {}, + _ManagedFilesCreatedBy({_UPLOADED_FILE_ID: "user-a"}), + ) + + assert "a cat" in str(response) + + @pytest.mark.asyncio + @pytest.mark.parametrize(("route_type", "path", "body_for"), _S3_FILE_ID_SURFACES) + @pytest.mark.parametrize( + "managed_files", + [ + pytest.param(_ManagedFilesCreatedBy({_UPLOADED_FILE_ID: "user-b"}), id="another-user-created-it"), + pytest.param(_ManagedFilesCreatedBy({}), id="forged-id"), + pytest.param(None, id="no-managed-files-hook"), + ], + ) + async def test_non_admin_managed_file_id_the_key_may_not_read_is_refused( + self, + route_type: _S3MediaRoute, + path: str, + body_for: Callable[[str], dict[str, object]], + managed_files: CustomLogger | None, + ) -> None: + with pytest.raises(HTTPException) as refused: + await _send_s3_media( + route_type, path, body_for(_UPLOADED_FILE_ID), LitellmUserRoles.INTERNAL_USER, {}, managed_files + ) + + assert refused.value.status_code == status.HTTP_403_FORBIDDEN, refused.value.detail + assert _UPLOADED_OBJECT in refused.value.detail + assert "LiteLLM-managed storage" in refused.value.detail + + @pytest.mark.asyncio + async def test_non_admin_forged_managed_file_ids_stop_at_the_first_failed_lookup(self) -> None: + forged_ids: Final = tuple( + _unified_file_id_for(f"s3://media-bucket/litellm-bedrock-files/forged-{index}.pdf") for index in range(2000) + ) + managed_files: Final = _ManagedFilesCountingLookups({}) + + with pytest.raises(HTTPException) as refused: + await _send_s3_media( + "acompletion", + "/v1/chat/completions", + _chat_files(forged_ids), + LitellmUserRoles.INTERNAL_USER, + {}, + managed_files, + ) + + assert refused.value.status_code == status.HTTP_403_FORBIDDEN, refused.value.detail + assert managed_files.looked_up == [forged_ids[0]] + + @pytest.mark.asyncio + async def test_non_admin_file_ids_that_cannot_change_the_outcome_cost_no_lookup(self) -> None: + listed_object_id: Final = _unified_file_id_for("s3://media-bucket/dir/report.pdf") + id_without_s3_uri: Final = _unified_file_id_for("gs://other/report.pdf") + managed_files: Final = _ManagedFilesCountingLookups({_UPLOADED_FILE_ID: "user-a"}) + + response: Final = await _send_s3_media( + "acompletion", + "/v1/chat/completions", + _chat_files((listed_object_id, id_without_s3_uri, _UPLOADED_FILE_ID)), + LitellmUserRoles.INTERNAL_USER, + {"allowed_client_side_s3_buckets": ["media-bucket"]}, + managed_files, + ) + + assert "a cat" in str(response) + assert managed_files.looked_up == [_UPLOADED_FILE_ID] + + @pytest.mark.asyncio + async def test_non_admin_managed_file_id_its_user_created_passes_from_an_unlisted_bucket(self) -> None: + owned_id: Final = _unified_file_id_for("s3://other-bucket/report.pdf") + + response: Final = await _send_s3_media( + "acompletion", + "/v1/chat/completions", + _chat_files((owned_id,)), + LitellmUserRoles.INTERNAL_USER, + {}, + _ManagedFilesCreatedBy({owned_id: "user-a"}), + ) + + assert "a cat" in str(response) + + @pytest.mark.asyncio + async def test_non_admin_unlisted_file_id_costs_no_lookup_when_client_side_credentials_are_allowed(self) -> None: + unlisted_object_id: Final = _unified_file_id_for("s3://other-bucket/report.pdf") + managed_files: Final = _ManagedFilesCountingLookups({}) + + response: Final = await _send_s3_media( + "acompletion", + "/v1/chat/completions", + _chat_files((unlisted_object_id,)), + LitellmUserRoles.INTERNAL_USER, + {"allow_client_side_credentials": True}, + managed_files, + ) + + assert "a cat" in str(response) + assert managed_files.looked_up == [] + + @pytest.mark.asyncio + @pytest.mark.parametrize( + ("mapped_source", "named_in_refusal"), + [ + pytest.param( + "s3://other-bucket/private.pdf", + "general_settings.allowed_client_side_s3_buckets", + id="unlisted-bucket", + ), + pytest.param( + "s3://media-bucket/litellm-bedrock-files/other-team.pdf", "LiteLLM-managed storage", id="managed-key" + ), + ], + ) + async def test_non_admin_s3_source_behind_a_model_file_id_mapping_is_refused_while_a_proxy_admin_key_passes( + self, mapped_source: str, named_in_refusal: str + ) -> None: + body: Final = {**_chat_file("alias"), "model_file_id_mapping": {"alias": {"deployment-1": mapped_source}}} + + with pytest.raises(HTTPException) as refused: + await _send_s3_media( + "acompletion", + "/v1/chat/completions", + body, + LitellmUserRoles.INTERNAL_USER, + {"allowed_client_side_s3_buckets": ["media-bucket"]}, + ) + admin_response: Final = await _send_s3_media( + "acompletion", "/v1/chat/completions", body, LitellmUserRoles.PROXY_ADMIN, {} + ) + + assert refused.value.status_code == status.HTTP_403_FORBIDDEN, refused.value.detail + assert mapped_source in refused.value.detail + assert named_in_refusal in refused.value.detail + assert "a cat" in str(admin_response) + + @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("safe_memory_mode", [False, True]) @pytest.mark.parametrize( diff --git a/tests/unit/proxy/test_proxy_server.py b/tests/unit/proxy/test_proxy_server.py index 80bc3ad1a34..c3a7a9c0925 100644 --- a/tests/unit/proxy/test_proxy_server.py +++ b/tests/unit/proxy/test_proxy_server.py @@ -3789,3 +3789,64 @@ def test_live_session_preserves_pretty_printed_start_frame( assert captured.route_type == "_arealtime" assert routed_data["model"] == "gpt-live-1" assert routed_data["live_session_start"] == expected_start + + +@pytest.mark.asyncio +async def test_queue_request_refuses_a_non_admin_s3_image_outside_the_allowed_buckets( + monkeypatch: pytest.MonkeyPatch, +) -> None: + from starlette.requests import Request + from starlette.responses import Response + + from litellm.proxy import proxy_server + from litellm.proxy._types import ProxyException + from litellm.proxy.proxy_server import async_queue_request + + body: Final = json.dumps( + { + "model": "nova-lite", + "messages": [ + { + "role": "user", + "content": [{"type": "image_url", "image_url": {"url": "s3://another-teams-bucket/cat.png"}}], + } + ], + } + ).encode() + + async def receive() -> dict[str, object]: + return {"type": "http.request", "body": body, "more_body": False} + + request: Final = Request( + { + "type": "http", + "method": "POST", + "path": "/queue/chat/completions", + "headers": [(b"content-type", b"application/json")], + "query_string": b"", + "scheme": "http", + "server": ("testserver", 80), + "root_path": "", + }, + receive, + ) + router: Final = litellm.Router( + model_list=[ + { + "model_name": "nova-lite", + "litellm_params": {"model": "openai/gpt-5.4-mini", "api_key": "sk-test", "mock_response": "ok"}, + } + ], + num_retries=0, + ) + monkeypatch.setattr(proxy_server, "llm_router", router) + monkeypatch.setattr(proxy_server, "general_settings", {"allowed_client_side_s3_buckets": ["media-bucket"]}) + + with pytest.raises(ProxyException) as exc_info: + await async_queue_request( + request=request, + fastapi_response=Response(), + user_api_key_dict=UserAPIKeyAuth(user_id="user-1", user_role=LitellmUserRoles.INTERNAL_USER), + ) + + assert exc_info.value.code == "403" diff --git a/tests/unit/proxy/test_proxy_server_endpoints_and_startup.py b/tests/unit/proxy/test_proxy_server_endpoints_and_startup.py index 114f5123dd6..f929a8c730d 100644 --- a/tests/unit/proxy/test_proxy_server_endpoints_and_startup.py +++ b/tests/unit/proxy/test_proxy_server_endpoints_and_startup.py @@ -12519,6 +12519,30 @@ async def test_update_config_general_settings_emits_audit_log(monkeypatch): assert before["some_api_key"] != "sk-stored-secret" +@pytest.mark.asyncio +async def test_update_config_general_settings_stores_a_client_side_s3_bucket_the_gate_admits(monkeypatch): + import litellm.proxy.proxy_server as proxy_server_module + from litellm.proxy._types import ConfigFieldUpdate + from litellm.proxy.auth.auth_utils import ClientSideStorageAllowlist + from litellm.proxy.proxy_server import update_config_general_settings + + fake: Final = _fake_prisma_with_config({}) + monkeypatch.setattr(proxy_server_module, "prisma_client", fake) + monkeypatch.setattr(litellm, "store_audit_logs", False) + + await update_config_general_settings( + data=ConfigFieldUpdate( + field_name="allowed_client_side_s3_buckets", + field_value=["media-bucket"], + config_type="general_settings", + ), + user_api_key_dict=UserAPIKeyAuth(api_key="hashed-admin", user_id="admin-1", user_role=LitellmUserRoles.PROXY_ADMIN), + ) + + stored: Final = json.loads(fake.db.litellm_config.upsert.call_args.kwargs["data"]["update"]["param_value"]) + assert ClientSideStorageAllowlist.model_validate(stored).admits_s3_bucket("media-bucket") + + @pytest.mark.asyncio async def test_delete_config_general_settings_is_visible_to_the_next_read(monkeypatch): import litellm.proxy.proxy_server as proxy_server_module 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 4fb2a4af571..c17946a2f6b 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 @@ -44,6 +44,68 @@ from litellm.types.utils import ( class TestLiteLLMCompletionResponsesConfig: + @pytest.mark.parametrize("part_type", ["input_image", "input_file"]) + @pytest.mark.parametrize( + ("key", "media_format", "expected_format"), + [("key", "image/png", "png"), ("key.jpeg", "image/png", "png"), ("key.png", None, "png")], + ) + def test_explicit_responses_media_format_reaches_the_converse_s3_image( + self, part_type: str, key: str, media_format: str | None, expected_format: str + ) -> None: + uri: Final = f"s3://bucket/{key}" + messages: Final = LiteLLMCompletionResponsesConfig.transform_responses_api_input_to_messages( + input=[ + { + "role": "user", + "content": [ + { + "type": part_type, + "image_url" if part_type == "input_image" else "file_id": uri, + **({"format": media_format} if media_format is not None else {}), + } + ], + } + ], + responses_api_request={}, + ) + request: Final = litellm.AmazonConverseConfig()._transform_request( + model="us.amazon.nova-2-lite-v1:0", + messages=messages, + optional_params={}, + litellm_params={}, + headers={}, + ) + + assert request["messages"] == [ + { + "role": "user", + "content": [{"image": {"format": expected_format, "source": {"s3Location": {"uri": uri}}}}], + } + ] + + @pytest.mark.parametrize("part_type", ["input_image", "input_file"]) + def test_extensionless_responses_s3_media_without_format_is_a_bad_request(self, part_type: str) -> None: + messages: Final = LiteLLMCompletionResponsesConfig.transform_responses_api_input_to_messages( + input=[ + { + "role": "user", + "content": [ + {"type": part_type, "image_url" if part_type == "input_image" else "file_id": "s3://bucket/key"} + ], + } + ], + responses_api_request={}, + ) + + with pytest.raises(litellm.BadRequestError, match="Cannot tell the media type"): + litellm.AmazonConverseConfig()._transform_request( + model="us.amazon.nova-2-lite-v1:0", + messages=messages, + optional_params={}, + litellm_params={}, + headers={}, + ) + def test_transform_input_file_item_to_file_item_with_file_id(self): """Test transformation of input_file item with file_id to Chat Completion file format""" # Setup @@ -263,6 +325,81 @@ class TestLiteLLMCompletionResponsesConfig: assert "extra_field" not in result assert "another_field" not in result + @pytest.mark.parametrize( + ("item", "expected"), + [ + 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", + ), + pytest.param( + {"type": "input_file", "file_id": "s3://media-bucket/docs/report.pdf", "bucket_owner": "111122223333"}, + { + "type": "file", + "file": {"file_id": "s3://media-bucket/docs/report.pdf", "bucket_owner": "111122223333"}, + }, + 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": "s3://media-bucket/dir/cat.png", "bucket_owner": ""}, + { + "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): + assert LiteLLMCompletionResponsesConfig._transform_responses_api_content_to_chat_completion_content([item]) == [ + 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 13b84896a84..bae8cea7153 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -29715,6 +29715,11 @@ export interface components { * @description If True, lets keys address Responses API ids that this proxy did not issue (raw provider ids, or ids issued before response-id encryption was configured). Such an id carries no owner, so no ownership check can run on it; ids this proxy did issue keep full ownership enforcement. Off by default, in which case an unrecognized response id is rejected with 403 */ allow_unmanaged_response_ids?: boolean | null; + /** + * Allowed Client Side S3 Buckets + * @description S3 bucket names that keys other than proxy admins may name in `s3://` image and file URLs. Any other bucket is rejected unless client-side credentials are allowed + */ + allowed_client_side_s3_buckets?: string[] | null; /** * Allowed File Extensions * @description the only file extensions (e.g. ['.jsonl', '.pdf', '.txt']) accepted on /v1/files uploads, for any purpose, matched case-insensitively against the uploaded filename. Files with any other extension, or none, are rejected. An empty list rejects every upload. Unset means no allowlist is applied