diff --git a/litellm/files/main.py b/litellm/files/main.py index 669d50dde414..6786615b19b3 100644 --- a/litellm/files/main.py +++ b/litellm/files/main.py @@ -39,7 +39,6 @@ from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.llms.azure.common_utils import get_azure_credentials from litellm.llms.azure.files.handler import AzureOpenAIFilesAPI -from litellm.llms.bedrock.files.handler import BedrockFilesHandler from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler from litellm.llms.custom_httpx.llm_http_handler import BaseLLMHTTPHandler from litellm.llms.openai.common_utils import get_openai_credentials @@ -82,7 +81,6 @@ def _should_sdk_support_streaming( openai_files_instance = OpenAIFilesAPI() azure_files_instance = AzureOpenAIFilesAPI() vertex_ai_files_instance = VertexAIFilesHandler() -bedrock_files_instance = BedrockFilesHandler() ################################################# @@ -1018,15 +1016,6 @@ def file_content( max_retries=optional_params.max_retries, litellm_params=litellm_params_dict, ) - elif custom_llm_provider == "bedrock": - response = bedrock_files_instance.file_content( - _is_async=_is_async, - file_content_request=_file_content_request, - api_base=optional_params.api_base, - optional_params=litellm_params_dict, - timeout=timeout, - max_retries=optional_params.max_retries, - ) else: raise litellm.exceptions.BadRequestError( message="LiteLLM doesn't support {} for 'file_content'. Supported providers are 'openai', 'azure', 'vertex_ai', 'bedrock', 'manus', 'anthropic'.".format( diff --git a/litellm/llms/bedrock/files/handler.py b/litellm/llms/bedrock/files/handler.py deleted file mode 100644 index ecf157e12eef..000000000000 --- a/litellm/llms/bedrock/files/handler.py +++ /dev/null @@ -1,242 +0,0 @@ -import asyncio -import base64 -import os -from types import MappingProxyType -from typing import Any, Coroutine, Mapping, Optional, Tuple, Union, cast - -import httpx - -from litellm import LlmProviders -from litellm.litellm_core_utils.cloud_storage_security import ( - BEDROCK_MANAGED_S3_PREFIXES, - should_allow_legacy_cloud_file_ids, - validate_managed_cloud_file_id, -) -from litellm.llms.custom_httpx.http_handler import get_async_httpx_client -from litellm.types.llms.openai import ( - FileContentRequest, - HttpxBinaryResponseContent, -) -from litellm.types.utils import SpecialEnums - -from ..base_aws_llm import BaseAWSLLM - - -class BedrockFilesHandler(BaseAWSLLM): - """ - Handles downloading files from S3 for Bedrock batch processing. - - This implementation downloads files from S3 buckets where Bedrock - stores batch output files. - """ - - def __init__(self): - super().__init__() - self.async_httpx_client = get_async_httpx_client( - llm_provider=LlmProviders.BEDROCK, - ) - - def _extract_s3_uri_from_file_id(self, file_id: str) -> str: - """ - Extract S3 URI from encoded file ID. - - The file ID can be in two formats: - 1. Base64-encoded unified file ID containing: llm_output_file_id,s3://bucket/path - 2. Direct S3 URI: s3://bucket/litellm-managed-prefix/path - - Args: - file_id: Encoded file ID or direct S3 URI - - Returns: - S3 URI (e.g., "s3://bucket-name/path/to/file") - """ - # First, try to decode if it's a base64-encoded unified file ID - try: - # Add padding if needed - padded = file_id + "=" * (-len(file_id) % 4) - decoded = base64.urlsafe_b64decode(padded).decode() - - # Check if it's a unified file ID format - if decoded.startswith(SpecialEnums.LITELM_MANAGED_FILE_ID_PREFIX.value): - # Extract llm_output_file_id from the decoded string - if "llm_output_file_id," in decoded: - s3_uri = decoded.split("llm_output_file_id,")[1].split(";")[0] - return s3_uri - except Exception: - pass - - # If not base64 encoded or doesn't contain llm_output_file_id, accept only - # explicit S3 URIs. Bucket and key validation happens before any S3 call. - if file_id.startswith("s3://"): - return file_id - - raise ValueError("file_id must be a managed LiteLLM S3 file id") - - def _parse_s3_uri( - self, - s3_uri: str, - configured_bucket_name: str, - allow_legacy_cloud_file_ids: bool = False, - ) -> Tuple[str, str]: - """ - Parse S3 URI to extract bucket name and object key. - - Args: - s3_uri: S3 URI (e.g., "s3://bucket-name/path/to/file") - - Returns: - Tuple of (bucket_name, object_key) - """ - return validate_managed_cloud_file_id( - file_id=s3_uri, - scheme="s3://", - configured_bucket_name=configured_bucket_name, - allowed_object_prefixes=BEDROCK_MANAGED_S3_PREFIXES, - allow_legacy_cloud_file_ids=allow_legacy_cloud_file_ids, - ) - - def _get_configured_s3_bucket_name(self, litellm_params: dict) -> str: - trusted_model_credentials = litellm_params.get( - "_litellm_internal_model_credentials" - ) - bucket_name = None - if isinstance(trusted_model_credentials, type(MappingProxyType({}))): - trusted_model_credentials_mapping = cast( - Mapping[str, Any], trusted_model_credentials - ) - candidate_bucket_name = trusted_model_credentials_mapping.get( - "s3_bucket_name" - ) - if isinstance(candidate_bucket_name, str): - bucket_name = candidate_bucket_name - bucket_name = bucket_name or os.getenv("AWS_S3_BUCKET_NAME") - if not bucket_name: - raise ValueError( - "S3 bucket_name is required. Set 's3_bucket_name' in proxy config or AWS_S3_BUCKET_NAME for Bedrock file content retrieval." - ) - return bucket_name - - async def afile_content( - self, - file_content_request: FileContentRequest, - optional_params: dict, - timeout: Union[float, httpx.Timeout], - max_retries: Optional[int], - ) -> HttpxBinaryResponseContent: - """ - Download file content from S3 bucket for Bedrock files. - - Args: - file_content_request: Contains file_id (encoded or S3 URI) - optional_params: Optional parameters containing AWS credentials - timeout: Request timeout - max_retries: Max retry attempts - - Returns: - HttpxBinaryResponseContent: Binary content wrapped in compatible response format - """ - import boto3 - from botocore.credentials import Credentials - - file_id = file_content_request.get("file_id") - if not file_id: - raise ValueError("file_id is required in file_content_request") - - # Extract S3 URI from file ID - s3_uri = self._extract_s3_uri_from_file_id(file_id) - configured_bucket_name = self._get_configured_s3_bucket_name(optional_params) - bucket_name, object_key = self._parse_s3_uri( - s3_uri=s3_uri, - configured_bucket_name=configured_bucket_name, - allow_legacy_cloud_file_ids=should_allow_legacy_cloud_file_ids( - optional_params - ), - ) - - # Get AWS credentials - aws_region_name = self._get_aws_region_name( - optional_params=optional_params, model="" - ) - credentials: Credentials = self.get_credentials( - aws_access_key_id=optional_params.get("aws_access_key_id"), - aws_secret_access_key=optional_params.get("aws_secret_access_key"), - aws_session_token=optional_params.get("aws_session_token"), - aws_region_name=aws_region_name, - aws_session_name=optional_params.get("aws_session_name"), - aws_profile_name=optional_params.get("aws_profile_name"), - aws_role_name=optional_params.get("aws_role_name"), - aws_web_identity_token=optional_params.get("aws_web_identity_token"), - aws_sts_endpoint=optional_params.get("aws_sts_endpoint"), - ) - - # Create S3 client - s3_client = boto3.client( - "s3", - aws_access_key_id=credentials.access_key, - aws_secret_access_key=credentials.secret_key, - aws_session_token=credentials.token, - region_name=aws_region_name, - verify=self._get_ssl_verify(), - ) - - # Download file from S3 - try: - response = s3_client.get_object(Bucket=bucket_name, Key=object_key) - file_content = response["Body"].read() - except Exception as e: - raise ValueError( - f"Failed to download file from S3: {s3_uri}. Error: {str(e)}" - ) - - # Create mock HTTP response - mock_response = httpx.Response( - status_code=200, - content=file_content, - headers={"content-type": "application/octet-stream"}, - request=httpx.Request(method="GET", url=s3_uri), - ) - - return HttpxBinaryResponseContent(response=mock_response) - - def file_content( - self, - _is_async: bool, - file_content_request: FileContentRequest, - api_base: Optional[str], - optional_params: dict, - timeout: Union[float, httpx.Timeout], - max_retries: Optional[int], - ) -> Union[ - HttpxBinaryResponseContent, Coroutine[Any, Any, HttpxBinaryResponseContent] - ]: - """ - Download file content from S3 bucket for Bedrock files. - Supports both sync and async operations. - - Args: - _is_async: Whether to run asynchronously - file_content_request: Contains file_id (encoded or S3 URI) - api_base: API base (unused for S3 operations) - optional_params: Optional parameters containing AWS credentials - timeout: Request timeout - max_retries: Max retry attempts - - Returns: - HttpxBinaryResponseContent or Coroutine: Binary content wrapped in compatible response format - """ - if _is_async: - return self.afile_content( - file_content_request=file_content_request, - optional_params=optional_params, - timeout=timeout, - max_retries=max_retries, - ) - else: - return asyncio.run( - self.afile_content( - file_content_request=file_content_request, - optional_params=optional_params, - timeout=timeout, - max_retries=max_retries, - ) - ) diff --git a/litellm/llms/bedrock/files/transformation.py b/litellm/llms/bedrock/files/transformation.py index cec2e934af85..1636ec9cd034 100644 --- a/litellm/llms/bedrock/files/transformation.py +++ b/litellm/llms/bedrock/files/transformation.py @@ -1,7 +1,9 @@ +import base64 import json import os import time -from typing import Any, Dict, List, Optional, Tuple, Union +from types import MappingProxyType +from typing import Any, Dict, List, Mapping, Optional, Tuple, Union, cast from urllib.parse import unquote import httpx @@ -13,11 +15,14 @@ from litellm.files.utils import FilesAPIUtils from litellm.litellm_core_utils.cloud_storage_security import ( BEDROCK_MANAGED_S3_BATCH_PREFIX, + BEDROCK_MANAGED_S3_PREFIXES, BEDROCK_MANAGED_S3_UPLOAD_PREFIX, build_managed_cloud_object_name, encode_s3_object_key_for_url, sanitize_cloud_object_component, + should_allow_legacy_cloud_file_ids, split_configured_cloud_bucket_name, + validate_managed_cloud_file_id, ) from litellm.litellm_core_utils.prompt_templates.common_utils import extract_file_data from litellm.llms.base_llm.chat.transformation import BaseLLMException @@ -34,7 +39,7 @@ OpenAIFileObject, PathLike, ) -from litellm.types.utils import ExtractedFileData, LlmProviders +from litellm.types.utils import ExtractedFileData, LlmProviders, SpecialEnums from litellm.utils import get_llm_provider from ..base_aws_llm import BaseAWSLLM @@ -877,6 +882,64 @@ def get_error_class( status_code=status_code, message=error_message, headers=headers ) + def _get_trusted_credentials(self, litellm_params: dict) -> Mapping[str, Any]: + snapshot = litellm_params.get("_litellm_internal_model_credentials") + if isinstance(snapshot, type(MappingProxyType({}))): + return cast(Mapping[str, Any], snapshot) + return MappingProxyType({}) + + def _extract_s3_uri_from_file_id(self, file_id: str) -> str: + try: + padded = file_id + "=" * (-len(file_id) % 4) + decoded = base64.urlsafe_b64decode(padded).decode() + if decoded.startswith(SpecialEnums.LITELM_MANAGED_FILE_ID_PREFIX.value): + if "llm_output_file_id," in decoded: + return decoded.split("llm_output_file_id,")[1].split(";")[0] + except Exception: + pass + + if file_id.startswith("s3://"): + return file_id + + raise ValueError("file_id must be a managed LiteLLM S3 file id") + + def _get_configured_s3_buckets(self, litellm_params: dict) -> tuple[str, ...]: + trusted = self._get_trusted_credentials(litellm_params) + + input_candidate = trusted.get("s3_bucket_name") + input_bucket = ( + input_candidate + if isinstance(input_candidate, str) and input_candidate + else os.getenv("AWS_S3_BUCKET_NAME") + ) + + output_candidate = trusted.get("s3_output_bucket_name") + output_bucket = ( + output_candidate + if isinstance(output_candidate, str) and output_candidate + else os.getenv("AWS_S3_OUTPUT_BUCKET_NAME") + ) + + buckets = tuple( + dict.fromkeys(bucket for bucket in (input_bucket, output_bucket) if bucket) + ) + if not buckets: + raise ValueError( + "S3 bucket_name is required for Bedrock file content retrieval. " + "Set AWS_S3_BUCKET_NAME (and AWS_S3_OUTPUT_BUCKET_NAME if batch " + "outputs use a separate bucket), or retrieve via a model so the " + "deployment's s3_bucket_name is resolved through model-based routing." + ) + return buckets + + def _resolve_s3_region(self, litellm_params: dict) -> str: + trusted = self._get_trusted_credentials(litellm_params) + for key in ("s3_region_name", "aws_region_name"): + trusted_region = trusted.get(key) + if isinstance(trusted_region, str) and trusted_region: + return trusted_region + return self._get_aws_region_name(optional_params=litellm_params, model="") + def transform_retrieve_file_request( self, file_id: str, @@ -925,15 +988,82 @@ def transform_list_files_response( ) -> List[OpenAIFileObject]: raise NotImplementedError("BedrockFilesConfig does not support file listing") + def _validate_against_configured_buckets( + self, s3_uri: str, litellm_params: dict + ) -> tuple[str, str]: + configured_buckets = self._get_configured_s3_buckets(litellm_params) + allow_legacy = should_allow_legacy_cloud_file_ids(litellm_params) + for index, configured_bucket in enumerate(configured_buckets): + is_last_bucket = index == len(configured_buckets) - 1 + try: + return validate_managed_cloud_file_id( + file_id=s3_uri, + scheme="s3://", + configured_bucket_name=configured_bucket, + allowed_object_prefixes=BEDROCK_MANAGED_S3_PREFIXES, + allow_legacy_cloud_file_ids=allow_legacy, + ) + except ValueError: + if is_last_bucket: + raise + raise ValueError("file_id must reference a LiteLLM-managed storage object") + + def _generate_presigned_s3_get_url( + self, bucket_name: str, object_key: str, litellm_params: dict + ) -> str: + try: + import boto3 + from botocore.config import Config + except ImportError: + raise ImportError("Missing boto3 to call bedrock. Run 'pip install boto3'.") + + region = self._resolve_s3_region(litellm_params) + credentials = self.get_credentials( + aws_access_key_id=litellm_params.get("aws_access_key_id"), + aws_secret_access_key=litellm_params.get("aws_secret_access_key"), + aws_session_token=litellm_params.get("aws_session_token"), + aws_region_name=region, + aws_session_name=litellm_params.get("aws_session_name"), + aws_profile_name=litellm_params.get("aws_profile_name"), + aws_role_name=litellm_params.get("aws_role_name"), + aws_web_identity_token=litellm_params.get("aws_web_identity_token"), + aws_sts_endpoint=litellm_params.get("aws_sts_endpoint"), + aws_external_id=litellm_params.get("aws_external_id"), + ) + + s3_client = boto3.client( + "s3", + aws_access_key_id=credentials.access_key, + aws_secret_access_key=credentials.secret_key, + aws_session_token=credentials.token, + region_name=region, + config=Config(signature_version="s3v4", s3={"addressing_style": "path"}), + verify=self._get_ssl_verify(), + ) + + return s3_client.generate_presigned_url( + "get_object", + Params={"Bucket": bucket_name, "Key": object_key}, + ExpiresIn=300, + ) + def transform_file_content_request( self, file_content_request, optional_params: dict, litellm_params: dict, ) -> tuple[str, dict]: - raise NotImplementedError( - "BedrockFilesConfig does not support file content retrieval" + file_id = file_content_request.get("file_id") or "" + s3_uri = self._extract_s3_uri_from_file_id(file_id) + bucket_name, object_key = self._validate_against_configured_buckets( + s3_uri=s3_uri, litellm_params=litellm_params + ) + url = self._generate_presigned_s3_get_url( + bucket_name=bucket_name, + object_key=object_key, + litellm_params=litellm_params, ) + return url, {} def transform_file_content_response( self, @@ -941,9 +1071,7 @@ def transform_file_content_response( logging_obj: LiteLLMLoggingObj, litellm_params: dict, ) -> HttpxBinaryResponseContent: - raise NotImplementedError( - "BedrockFilesConfig does not support file content retrieval" - ) + return HttpxBinaryResponseContent(response=raw_response) class BedrockJsonlFilesTransformation: diff --git a/tests/test_litellm/llms/bedrock/files/test_bedrock_files_handler.py b/tests/test_litellm/llms/bedrock/files/test_bedrock_files_handler.py index 7f91b49a6f5d..9ef68bf4a532 100644 --- a/tests/test_litellm/llms/bedrock/files/test_bedrock_files_handler.py +++ b/tests/test_litellm/llms/bedrock/files/test_bedrock_files_handler.py @@ -1,168 +1,120 @@ import base64 -import os from types import MappingProxyType from unittest.mock import MagicMock, patch import pytest import litellm.files.main as files_main -from litellm.llms.bedrock.files.handler import BedrockFilesHandler +from litellm.litellm_core_utils.cloud_storage_security import ( + BEDROCK_MANAGED_S3_PREFIXES, + validate_managed_cloud_file_id, +) +from litellm.llms.bedrock.files.transformation import BedrockFilesConfig from litellm.types.utils import SpecialEnums def _encode_unified_file_id(s3_uri: str) -> str: unified_file_id = SpecialEnums.LITELLM_MANAGED_FILE_COMPLETE_STR.value.format( - "application/json", - "unified-id", - "", - s3_uri, - "model-id", + "application/json", "unified-id", "", s3_uri, "model-id" ) return base64.urlsafe_b64encode(unified_file_id.encode()).decode().rstrip("=") -class TestBedrockFilesHandler: +def _parse(s3_uri, configured_bucket_name, allow_legacy_cloud_file_ids=False): + return validate_managed_cloud_file_id( + file_id=s3_uri, + scheme="s3://", + configured_bucket_name=configured_bucket_name, + allowed_object_prefixes=BEDROCK_MANAGED_S3_PREFIXES, + allow_legacy_cloud_file_ids=allow_legacy_cloud_file_ids, + ) + + +class TestBedrockFilesConfigS3Validation: def setup_method(self): - self.handler = BedrockFilesHandler() + self.config = BedrockFilesConfig() def test_should_parse_direct_managed_s3_uri(self): - bucket, key = self.handler._parse_s3_uri( - s3_uri="s3://safe-bucket/litellm-bedrock-files-model-id-abc.jsonl", - configured_bucket_name="safe-bucket", + bucket, key = _parse( + "s3://safe-bucket/litellm-bedrock-files-model-id-abc.jsonl", "safe-bucket" ) - assert bucket == "safe-bucket" assert key == "litellm-bedrock-files-model-id-abc.jsonl" def test_should_parse_managed_batch_output_uri(self): - bucket, key = self.handler._parse_s3_uri( - s3_uri="s3://safe-bucket/litellm-batch-outputs/job/", - configured_bucket_name="safe-bucket", + bucket, key = _parse( + "s3://safe-bucket/litellm-batch-outputs/job/", "safe-bucket" ) - assert bucket == "safe-bucket" assert key == "litellm-batch-outputs/job/" def test_should_reject_arbitrary_bucket(self): with pytest.raises(ValueError, match="configured storage bucket"): - self.handler._parse_s3_uri( - s3_uri="s3://other-bucket/litellm-bedrock-files-model-id-abc.jsonl", - configured_bucket_name="safe-bucket", + _parse( + "s3://other-bucket/litellm-bedrock-files-model-id-abc.jsonl", + "safe-bucket", ) def test_should_reject_unmanaged_same_bucket_key(self): with pytest.raises(ValueError, match="LiteLLM-managed"): - self.handler._parse_s3_uri( - s3_uri="s3://safe-bucket/private/output.jsonl", - configured_bucket_name="safe-bucket", - ) + _parse("s3://safe-bucket/private/output.jsonl", "safe-bucket") def test_should_allow_legacy_same_bucket_key_when_server_flag_enabled(self): - bucket, key = self.handler._parse_s3_uri( - s3_uri="s3://safe-bucket/private/output.jsonl", - configured_bucket_name="safe-bucket", + bucket, key = _parse( + "s3://safe-bucket/private/output.jsonl", + "safe-bucket", allow_legacy_cloud_file_ids=True, ) - assert bucket == "safe-bucket" assert key == "private/output.jsonl" def test_should_keep_configured_prefix_for_legacy_keys(self): - bucket, key = self.handler._parse_s3_uri( - s3_uri="s3://safe-bucket/team-a/private/output.jsonl", - configured_bucket_name="safe-bucket/team-a", + bucket, key = _parse( + "s3://safe-bucket/team-a/private/output.jsonl", + "safe-bucket/team-a", allow_legacy_cloud_file_ids=True, ) - assert bucket == "safe-bucket" assert key == "team-a/private/output.jsonl" def test_should_reject_legacy_key_outside_configured_prefix(self): with pytest.raises(ValueError, match="configured storage prefix"): - self.handler._parse_s3_uri( - s3_uri="s3://safe-bucket/team-b/private/output.jsonl", - configured_bucket_name="safe-bucket/team-a", + _parse( + "s3://safe-bucket/team-b/private/output.jsonl", + "safe-bucket/team-a", allow_legacy_cloud_file_ids=True, ) def test_should_reject_dot_segment_key(self): with pytest.raises(ValueError, match="invalid path segment"): - self.handler._parse_s3_uri( - s3_uri="s3://safe-bucket/litellm-bedrock-files/../secret.jsonl", - configured_bucket_name="safe-bucket", + _parse( + "s3://safe-bucket/litellm-bedrock-files/../secret.jsonl", "safe-bucket" ) def test_should_reject_empty_middle_path_segment(self): with pytest.raises(ValueError, match="invalid path segment"): - self.handler._parse_s3_uri( - s3_uri="s3://safe-bucket/litellm-bedrock-files//secret.jsonl", - configured_bucket_name="safe-bucket", + _parse( + "s3://safe-bucket/litellm-bedrock-files//secret.jsonl", "safe-bucket" ) def test_should_extract_unified_managed_s3_uri(self): file_id = _encode_unified_file_id( "s3://safe-bucket/litellm-batch-outputs/job/output.jsonl" ) - assert ( - self.handler._extract_s3_uri_from_file_id(file_id) + self.config._extract_s3_uri_from_file_id(file_id) == "s3://safe-bucket/litellm-batch-outputs/job/output.jsonl" ) def test_should_reject_file_id_without_s3_scheme(self): with pytest.raises(ValueError, match="managed LiteLLM S3 file id"): - self.handler._extract_s3_uri_from_file_id("safe-bucket/private.jsonl") + self.config._extract_s3_uri_from_file_id("safe-bucket/private.jsonl") def test_should_reject_unified_unmanaged_s3_uri(self): file_id = _encode_unified_file_id("s3://safe-bucket/private/output.jsonl") - s3_uri = self.handler._extract_s3_uri_from_file_id(file_id) - + s3_uri = self.config._extract_s3_uri_from_file_id(file_id) with pytest.raises(ValueError, match="LiteLLM-managed"): - self.handler._parse_s3_uri( - s3_uri=s3_uri, - configured_bucket_name="safe-bucket", - ) - - def test_should_not_trust_request_s3_bucket_name_for_expected_bucket(self): - with patch.dict(os.environ, {"AWS_S3_BUCKET_NAME": "safe-bucket"}): - assert ( - self.handler._get_configured_s3_bucket_name( - {"s3_bucket_name": "attacker-bucket"} - ) - == "safe-bucket" - ) - - def test_should_trust_proxy_config_s3_bucket_name_for_expected_bucket(self): - trusted_credentials = MappingProxyType({"s3_bucket_name": "safe-bucket"}) - - with patch.dict(os.environ, {}, clear=True): - assert ( - self.handler._get_configured_s3_bucket_name( - { - "s3_bucket_name": "attacker-bucket", - "_litellm_internal_model_credentials": trusted_credentials, - } - ) - == "safe-bucket" - ) - - def test_should_not_trust_user_supplied_internal_credentials_dict(self): - with patch.dict(os.environ, {}, clear=True): - with pytest.raises(ValueError, match="S3 bucket_name is required"): - self.handler._get_configured_s3_bucket_name( - { - "_litellm_internal_model_credentials": { - "s3_bucket_name": "attacker-bucket" - } - } - ) - - def test_should_require_server_s3_bucket_name(self): - with patch.dict(os.environ, {}, clear=True): - with pytest.raises(ValueError, match="S3 bucket_name is required"): - self.handler._get_configured_s3_bucket_name( - {"s3_bucket_name": "attacker-bucket"} - ) + _parse(s3_uri, "safe-bucket") def test_should_forward_trusted_model_credentials_to_bedrock_provider_config(): diff --git a/tests/test_litellm/llms/bedrock/files/test_bedrock_files_transformation.py b/tests/test_litellm/llms/bedrock/files/test_bedrock_files_transformation.py index 4731be13e78e..1d8e23ee7905 100644 --- a/tests/test_litellm/llms/bedrock/files/test_bedrock_files_transformation.py +++ b/tests/test_litellm/llms/bedrock/files/test_bedrock_files_transformation.py @@ -2,11 +2,23 @@ Test bedrock files transformation functionality """ +import base64 import json import os -from urllib.parse import unquote, urlparse +from types import MappingProxyType +from unittest.mock import patch +from urllib.parse import parse_qs, unquote, urlparse, urlsplit -from litellm.llms.bedrock.files.transformation import BedrockJsonlFilesTransformation +import httpx +import pytest +from botocore.credentials import Credentials + +from litellm.llms.bedrock.files.transformation import ( + BedrockFilesConfig, + BedrockJsonlFilesTransformation, +) +from litellm.types.llms.openai import HttpxBinaryResponseContent +from litellm.types.utils import SpecialEnums class TestBedrockFilesTransformation: @@ -1173,3 +1185,289 @@ def test_other_non_embedding_urls_route_to_chat(self): assert not BedrockFilesConfig._is_embedding_record( {"url": "/v1/responses", "body": {"input": "x"}} ) + + +def _encode_unified_file_id(s3_uri: str) -> str: + unified_file_id = SpecialEnums.LITELLM_MANAGED_FILE_COMPLETE_STR.value.format( + "application/json", + "unified-id", + "", + s3_uri, + "model-id", + ) + return base64.urlsafe_b64encode(unified_file_id.encode()).decode().rstrip("=") + + +def _trusted(**creds) -> dict: + return {"_litellm_internal_model_credentials": MappingProxyType(dict(creds))} + + +class TestBedrockFilesConfigResolution: + def setup_method(self): + self.config = BedrockFilesConfig() + + def test_extract_direct_s3_uri(self): + assert ( + self.config._extract_s3_uri_from_file_id( + "s3://b/litellm-batch-outputs/job/output.jsonl" + ) + == "s3://b/litellm-batch-outputs/job/output.jsonl" + ) + + def test_extract_unified_managed_s3_uri(self): + file_id = _encode_unified_file_id( + "s3://b/litellm-batch-outputs/job/output.jsonl" + ) + assert ( + self.config._extract_s3_uri_from_file_id(file_id) + == "s3://b/litellm-batch-outputs/job/output.jsonl" + ) + + def test_extract_rejects_non_s3_file_id(self): + with pytest.raises(ValueError, match="managed LiteLLM S3 file id"): + self.config._extract_s3_uri_from_file_id("b/private.jsonl") + + def test_buckets_prefer_trusted_snapshot_over_request(self): + params = { + "s3_bucket_name": "attacker-bucket", + **_trusted(s3_bucket_name="safe-bucket"), + } + with patch.dict(os.environ, {}, clear=True): + assert self.config._get_configured_s3_buckets(params) == ("safe-bucket",) + + def test_buckets_include_output_bucket(self): + params = _trusted( + s3_bucket_name="in-bucket", s3_output_bucket_name="out-bucket" + ) + with patch.dict(os.environ, {}, clear=True): + assert self.config._get_configured_s3_buckets(params) == ( + "in-bucket", + "out-bucket", + ) + + def test_buckets_fall_back_to_env(self): + with patch.dict( + os.environ, + {"AWS_S3_BUCKET_NAME": "env-in", "AWS_S3_OUTPUT_BUCKET_NAME": "env-out"}, + clear=True, + ): + assert self.config._get_configured_s3_buckets({}) == ("env-in", "env-out") + + def test_buckets_require_at_least_one(self): + with patch.dict(os.environ, {}, clear=True): + with pytest.raises(ValueError, match="S3 bucket_name is required"): + self.config._get_configured_s3_buckets({"s3_bucket_name": "ignored"}) + + def test_region_prefers_trusted_s3_region_name(self): + params = { + "aws_region_name": "us-east-1", + **_trusted(s3_region_name="eu-central-1"), + } + assert self.config._resolve_s3_region(params) == "eu-central-1" + + def test_region_falls_back_to_aws_region_name(self): + assert self.config._resolve_s3_region( + {"aws_region_name": "ap-southeast-2"} + ) == ("ap-southeast-2") + + +class TestBedrockFilesContentRetrieval: + def setup_method(self): + self.config = BedrockFilesConfig() + self.fake_creds = Credentials( + access_key="AKIAIOSFODNN7EXAMPLE", + secret_key="wJalrXUtnFEMI/K7MDENG/bPxRfiCYEXAMPLEKEY", + token=None, + ) + + def _request(self, file_id, litellm_params, capture=None): + def spy(**kwargs): + if capture is not None: + capture.update(kwargs) + return self.fake_creds + + with patch.object(self.config, "get_credentials", side_effect=spy): + return self.config.transform_file_content_request( + file_content_request={"file_id": file_id}, + optional_params={}, + litellm_params=litellm_params, + ) + + def test_returns_presigned_get_url_for_input_bucket(self): + url, params = self._request( + "s3://in-bucket/litellm-batch-outputs/job/in.jsonl.out", + _trusted(s3_bucket_name="in-bucket", aws_region_name="us-west-2"), + ) + assert params == {} + parts = urlsplit(url) + query = parse_qs(parts.query) + assert parts.path == "/in-bucket/litellm-batch-outputs/job/in.jsonl.out" + assert query["X-Amz-Algorithm"] == ["AWS4-HMAC-SHA256"] + assert "X-Amz-Signature" in query + + def test_region_from_trusted_snapshot_wins(self): + url, _ = self._request( + "s3://in-bucket/litellm-batch-outputs/job/in.jsonl.out", + { + "aws_region_name": "us-east-1", + **_trusted( + s3_bucket_name="in-bucket", + s3_region_name="eu-central-1", + aws_region_name="us-east-1", + ), + }, + ) + credential = parse_qs(urlsplit(url).query)["X-Amz-Credential"][0] + assert "/eu-central-1/s3/aws4_request" in credential + assert urlsplit(url).netloc == "s3.eu-central-1.amazonaws.com" + + def test_china_partition_uses_correct_endpoint_suffix(self): + url, _ = self._request( + "s3://in-bucket/litellm-batch-outputs/job/in.jsonl.out", + _trusted(s3_bucket_name="in-bucket", aws_region_name="cn-north-1"), + ) + assert urlsplit(url).netloc == "s3.cn-north-1.amazonaws.com.cn" + assert "X-Amz-Signature" in parse_qs(urlsplit(url).query) + + def test_uses_sigv4_not_deprecated_sigv2(self): + url, _ = self._request( + "s3://in-bucket/litellm-batch-outputs/job/in.jsonl.out", + _trusted(s3_bucket_name="in-bucket", aws_region_name="us-west-2"), + ) + query = parse_qs(urlsplit(url).query) + assert query["X-Amz-Algorithm"] == ["AWS4-HMAC-SHA256"] + assert "Signature" not in query and "AWSAccessKeyId" not in query + + def test_output_bucket_distinct_from_input_validates(self): + url, _ = self._request( + "s3://out-bucket/litellm-batch-outputs/job/in.jsonl.out", + _trusted( + s3_bucket_name="in-bucket", + s3_output_bucket_name="out-bucket", + aws_region_name="us-west-2", + ), + ) + assert ( + urlsplit(url).path == "/out-bucket/litellm-batch-outputs/job/in.jsonl.out" + ) + + def test_batch_output_at_root_validates_with_prefixed_input_bucket(self): + # input bucket is prefix-scoped; batch output lands at bucket root + url, _ = self._request( + "s3://shared/litellm-batch-outputs/job/in.jsonl.out", + _trusted( + s3_bucket_name="shared/team-a", + s3_output_bucket_name="shared", + aws_region_name="us-west-2", + ), + ) + assert urlsplit(url).path == "/shared/litellm-batch-outputs/job/in.jsonl.out" + + def test_unified_file_id_is_decoded_and_presigned(self): + file_id = _encode_unified_file_id( + "s3://in-bucket/litellm-batch-outputs/job/in.jsonl.out" + ) + url, _ = self._request( + file_id, _trusted(s3_bucket_name="in-bucket", aws_region_name="us-west-2") + ) + assert urlsplit(url).path == "/in-bucket/litellm-batch-outputs/job/in.jsonl.out" + assert "X-Amz-Signature" in parse_qs(urlsplit(url).query) + + def test_rejects_bucket_outside_configured_set(self): + with pytest.raises(ValueError): + self._request( + "s3://other-bucket/litellm-batch-outputs/job/in.jsonl.out", + _trusted( + s3_bucket_name="in-bucket", + s3_output_bucket_name="out-bucket", + aws_region_name="us-west-2", + ), + ) + + def test_rejects_unmanaged_prefix_in_configured_bucket(self): + with pytest.raises(ValueError): + self._request( + "s3://in-bucket/private/secret.jsonl", + _trusted(s3_bucket_name="in-bucket", aws_region_name="us-west-2"), + ) + + def test_ignores_plain_request_bucket(self): + with pytest.raises(ValueError, match="S3 bucket_name is required"): + self._request( + "s3://attacker-bucket/litellm-batch-outputs/job/in.jsonl.out", + {"s3_bucket_name": "attacker-bucket", "aws_region_name": "us-west-2"}, + ) + + def test_forwards_aws_external_id_to_get_credentials(self): + capture: dict = {} + self._request( + "s3://in-bucket/litellm-batch-outputs/job/in.jsonl.out", + { + "aws_external_id": "ext-123", + **_trusted(s3_bucket_name="in-bucket", aws_region_name="us-west-2"), + }, + capture=capture, + ) + assert capture.get("aws_external_id") == "ext-123" + + def test_response_returns_raw_bytes_unchanged(self): + body = b'{"recordId":"req-1","modelOutput":{"foo":"bar"}}\n' + raw_response = httpx.Response( + status_code=200, + content=body, + headers={"content-type": "application/octet-stream"}, + request=httpx.Request( + "GET", "https://s3.us-west-2.amazonaws.com/in-bucket/x" + ), + ) + result = self.config.transform_file_content_response( + raw_response=raw_response, logging_obj=None, litellm_params={} + ) + assert isinstance(result, HttpxBinaryResponseContent) + assert result.content == body + + def test_retrieve_file_content_through_generic_handler(self): + import time as _time + + from litellm.files.main import base_llm_http_handler + from litellm.litellm_core_utils.litellm_logging import Logging + from litellm.llms.custom_httpx.http_handler import HTTPHandler + + body = b'{"recordId":"req-1","modelOutput":{"ok":true}}\n' + seen = {} + + def handler(request: httpx.Request) -> httpx.Response: + seen["url"] = str(request.url) + return httpx.Response(status_code=200, content=body) + + sync_client = HTTPHandler() + sync_client.client = httpx.Client(transport=httpx.MockTransport(handler)) + + logging_obj = Logging( + model="", + messages=[], + stream=False, + call_type="file_content", + start_time=_time.time(), + litellm_call_id="test-call", + function_id="", + ) + + with patch.object(self.config, "get_credentials", return_value=self.fake_creds): + result = base_llm_http_handler.retrieve_file_content( + file_content_request={ + "file_id": "s3://in-bucket/litellm-batch-outputs/job/in.jsonl.out" + }, + provider_config=self.config, + litellm_params=_trusted( + s3_bucket_name="in-bucket", aws_region_name="us-west-2" + ), + headers={}, + logging_obj=logging_obj, + _is_async=False, + client=sync_client, + ) + + assert "X-Amz-Signature=" in seen["url"] + assert "/in-bucket/litellm-batch-outputs/job/in.jsonl.out" in seen["url"] + assert result.content == body