mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-14 23:21:35 +00:00
feat(bedrock): support file delete and list for S3-backed managed files
This commit is contained in:
mateo-berri
parent
b3c867c7b2
commit
fc978aec21
7 changed files with 814 additions and 59 deletions
|
|
@ -1781,7 +1781,13 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
|
|||
# Remove conflicting keys from data to avoid duplicate keyword arguments
|
||||
filtered_data = {k: v for k, v in data.items() if k not in ("model", "file_id")}
|
||||
for model_id, model_file_id in specific_model_file_id_mapping.items():
|
||||
delete_response = await llm_router.afile_delete(model=model_id, file_id=model_file_id, **filtered_data) # type: ignore
|
||||
credentials = llm_router.get_deployment_credentials_with_provider(model_id=model_id)
|
||||
router_kwargs = (
|
||||
{**filtered_data, "_litellm_internal_model_credentials": MappingProxyType(dict(credentials))}
|
||||
if credentials is not None
|
||||
else filtered_data
|
||||
)
|
||||
delete_response = await llm_router.afile_delete(model=model_id, file_id=model_file_id, **router_kwargs)
|
||||
|
||||
stored_file_object = await self.delete_unified_file_id(file_id, litellm_parent_otel_span)
|
||||
|
||||
|
|
|
|||
|
|
@ -681,6 +681,10 @@ def file_list(
|
|||
)
|
||||
if provider_config is not None:
|
||||
litellm_params_dict: Final = get_litellm_params(**kwargs)
|
||||
add_trusted_model_credentials_to_litellm_params(
|
||||
litellm_params_dict=litellm_params_dict,
|
||||
kwargs=kwargs,
|
||||
)
|
||||
litellm_params_dict["api_key"] = optional_params.api_key
|
||||
litellm_params_dict["api_base"] = optional_params.api_base
|
||||
|
||||
|
|
|
|||
|
|
@ -1,14 +1,18 @@
|
|||
import base64
|
||||
import json
|
||||
import os
|
||||
import posixpath
|
||||
import time
|
||||
import xml.etree.ElementTree as ET
|
||||
from collections.abc import Iterable, Mapping, MutableMapping, Sequence
|
||||
from contextlib import suppress
|
||||
from dataclasses import dataclass
|
||||
from datetime import datetime
|
||||
from functools import cache
|
||||
from itertools import chain
|
||||
from types import MappingProxyType
|
||||
from typing import Any, Final, TypeAlias, TypedDict
|
||||
from urllib.parse import unquote
|
||||
from urllib.parse import quote, unquote, urlencode
|
||||
|
||||
import httpx
|
||||
from httpx import Headers, Response
|
||||
|
|
@ -23,6 +27,7 @@ from litellm.files.utils import FilesAPIUtils
|
|||
from litellm.litellm_core_utils.aws_partition import get_aws_dns_suffix
|
||||
from litellm.litellm_core_utils.cloud_storage_security import (
|
||||
BEDROCK_MANAGED_S3_BATCH_PREFIX,
|
||||
BEDROCK_MANAGED_S3_OUTPUT_PREFIX,
|
||||
BEDROCK_MANAGED_S3_PREFIXES,
|
||||
BEDROCK_MANAGED_S3_UPLOAD_PREFIX,
|
||||
build_managed_cloud_object_name,
|
||||
|
|
@ -60,11 +65,15 @@ from litellm.utils import get_llm_provider
|
|||
from ..base_aws_llm import BaseAWSLLM
|
||||
from ..common_utils import BedrockError, merge_bedrock_aws_request_params, resolve_s3_encryption_key_id
|
||||
|
||||
# litellm_params key used to hand the SigV4-signed GET headers from
|
||||
# `transform_file_content_request` to `validate_environment` (the only hook
|
||||
# the shared file-content HTTP handler exposes for setting request headers).
|
||||
# litellm_params key used to hand SigV4-signed request headers from the
|
||||
# content, delete, and list request transforms to `validate_environment` (the
|
||||
# only hook the shared files HTTP handler exposes for setting request headers).
|
||||
# Same pattern as the `upload_url` handoff in `transform_create_file_request`.
|
||||
S3_SIGNED_GET_HEADERS_PARAM: Final = "_s3_signed_get_headers"
|
||||
S3_SIGNED_REQUEST_HEADERS_PARAM: Final = "_s3_signed_request_headers"
|
||||
|
||||
DELETED_FILE_ID_PARAM: Final = "_s3_deleted_file_id"
|
||||
|
||||
LIST_FILES_PURPOSE_PARAM: Final = "_s3_list_files_purpose"
|
||||
|
||||
# litellm_params key carrying the size of the body uploaded to S3, handed from
|
||||
# `transform_create_file_request` to `transform_create_file_response`.
|
||||
|
|
@ -151,6 +160,13 @@ class _BedrockS3RequestParams(BaseModel):
|
|||
s3_endpoint_url: str | None = None
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class _S3RequestTarget:
|
||||
endpoint_url: str
|
||||
aws_region_name: str
|
||||
request_params: _BedrockS3RequestParams
|
||||
|
||||
|
||||
class _TrustedS3ModelCredentials(BaseModel):
|
||||
"""The S3 buckets the server trusts file ids against, from the deployment snapshot."""
|
||||
|
||||
|
|
@ -247,6 +263,51 @@ def _validate_file_id_against_configured_buckets(
|
|||
return validate_against(configured_bucket_names[-1])
|
||||
|
||||
|
||||
def _managed_listing_prefix(configured_prefix: str) -> str:
|
||||
common_prefix: Final = os.path.commonprefix(BEDROCK_MANAGED_S3_PREFIXES)
|
||||
return f"{configured_prefix}/{common_prefix}" if configured_prefix else common_prefix
|
||||
|
||||
|
||||
def _listed_object_created_at(entry: ET.Element) -> int:
|
||||
last_modified: Final = entry.findtext("{*}LastModified")
|
||||
if not last_modified:
|
||||
return 0
|
||||
return int(datetime.fromisoformat(last_modified.replace("Z", "+00:00")).timestamp())
|
||||
|
||||
|
||||
def _listed_managed_file(
|
||||
entry: ET.Element,
|
||||
bucket_name: str,
|
||||
configured_bucket_name: str,
|
||||
allow_legacy_cloud_file_ids: bool,
|
||||
) -> OpenAIFileObject | None:
|
||||
object_key: Final = entry.findtext("{*}Key")
|
||||
if not object_key:
|
||||
return None
|
||||
file_id: Final = f"s3://{bucket_name}/{object_key}"
|
||||
try:
|
||||
validate_managed_cloud_file_id(
|
||||
file_id=file_id,
|
||||
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,
|
||||
)
|
||||
except ValueError:
|
||||
return None
|
||||
_, configured_prefix = split_configured_cloud_bucket_name(configured_bucket_name)
|
||||
relative_key: Final = object_key[len(configured_prefix) + 1 :] if configured_prefix else object_key
|
||||
return OpenAIFileObject(
|
||||
id=file_id,
|
||||
bytes=int(entry.findtext("{*}Size") or 0),
|
||||
created_at=_listed_object_created_at(entry),
|
||||
filename=posixpath.basename(object_key),
|
||||
object="file",
|
||||
purpose="batch_output" if relative_key.startswith(BEDROCK_MANAGED_S3_OUTPUT_PREFIX) else "batch",
|
||||
status="uploaded",
|
||||
)
|
||||
|
||||
|
||||
def _uploaded_object_size(litellm_params: Mapping[str, object], raw_response: Response) -> int:
|
||||
"""
|
||||
S3 answers PutObject with an empty body, so the stored object size comes from the
|
||||
|
|
@ -291,7 +352,7 @@ class BedrockFilesConfig(BaseAWSLLM, BaseFilesConfig):
|
|||
) -> dict:
|
||||
result: Final[dict[str, object]] = {}
|
||||
result.update(headers)
|
||||
signed_headers: Final = litellm_params.pop(S3_SIGNED_GET_HEADERS_PARAM, None)
|
||||
signed_headers: Final = litellm_params.pop(S3_SIGNED_REQUEST_HEADERS_PARAM, None)
|
||||
if isinstance(signed_headers, Mapping):
|
||||
result.update(signed_headers) # any-ok: untyped handoff headers
|
||||
# otherwise no extra headers - AWS credentials are handled by BaseAWSLLM
|
||||
|
|
@ -1187,34 +1248,95 @@ class BedrockFilesConfig(BaseAWSLLM, BaseFilesConfig):
|
|||
def transform_delete_file_request(
|
||||
self,
|
||||
file_id: str,
|
||||
optional_params: dict,
|
||||
litellm_params: dict,
|
||||
) -> tuple[str, dict]:
|
||||
raise NotImplementedError("BedrockFilesConfig does not support file deletion")
|
||||
optional_params: Mapping[str, object],
|
||||
litellm_params: MutableMapping[str, object],
|
||||
) -> tuple[str, dict[str, str]]:
|
||||
if not file_id:
|
||||
raise ValueError("file_id is required for Bedrock file deletion")
|
||||
bucket_name, object_key = _validate_file_id_against_configured_buckets(
|
||||
s3_uri=extract_s3_uri_from_file_id(file_id),
|
||||
configured_bucket_names=get_configured_s3_bucket_names(litellm_params),
|
||||
allow_legacy_cloud_file_ids=should_allow_legacy_cloud_file_ids(litellm_params),
|
||||
)
|
||||
target: Final = self._s3_request_target(optional_params=optional_params, litellm_params=litellm_params)
|
||||
url: Final = f"{target.endpoint_url}/{bucket_name}/{encode_s3_object_key_for_url(object_key)}"
|
||||
signed_headers: Final = self._sign_s3_empty_body_request(
|
||||
method="DELETE",
|
||||
api_base=url,
|
||||
aws_region_name=target.aws_region_name,
|
||||
request_params=target.request_params,
|
||||
)
|
||||
litellm_params[S3_SIGNED_REQUEST_HEADERS_PARAM] = signed_headers # rebind-ok: handed to validate_environment
|
||||
litellm_params[DELETED_FILE_ID_PARAM] = file_id # rebind-ok: S3 DeleteObject answers with an empty body
|
||||
return url, {} # mutable-ok: the base files contract returns the query as a dict
|
||||
|
||||
def transform_delete_file_response(
|
||||
self,
|
||||
raw_response: httpx.Response,
|
||||
logging_obj: LiteLLMLoggingObj,
|
||||
litellm_params: dict,
|
||||
litellm_params: Mapping[str, object],
|
||||
) -> FileDeleted:
|
||||
raise NotImplementedError("BedrockFilesConfig does not support file deletion")
|
||||
if raw_response.status_code >= 400:
|
||||
raise BedrockError(
|
||||
status_code=raw_response.status_code,
|
||||
message=raw_response.text,
|
||||
headers=raw_response.headers,
|
||||
)
|
||||
return FileDeleted(id=str(litellm_params.get(DELETED_FILE_ID_PARAM, "")), deleted=True, object="file")
|
||||
|
||||
def transform_list_files_request(
|
||||
self,
|
||||
purpose: str | None,
|
||||
optional_params: dict,
|
||||
litellm_params: dict,
|
||||
) -> tuple[str, dict]:
|
||||
raise NotImplementedError("BedrockFilesConfig does not support file listing")
|
||||
optional_params: Mapping[str, object],
|
||||
litellm_params: MutableMapping[str, object],
|
||||
) -> tuple[str, dict[str, str]]:
|
||||
bucket_name, configured_prefix = split_configured_cloud_bucket_name(
|
||||
get_configured_s3_bucket_name(litellm_params)
|
||||
)
|
||||
target: Final = self._s3_request_target(optional_params=optional_params, litellm_params=litellm_params)
|
||||
url: Final = f"{target.endpoint_url}/{bucket_name}/"
|
||||
query: Final[dict[str, str]] = { # mutable-ok: the base files contract returns the query as a dict
|
||||
"list-type": "2",
|
||||
"prefix": _managed_listing_prefix(configured_prefix),
|
||||
}
|
||||
signed_headers: Final = self._sign_s3_empty_body_request(
|
||||
method="GET",
|
||||
api_base=f"{url}?{urlencode(query, quote_via=quote, safe='')}",
|
||||
aws_region_name=target.aws_region_name,
|
||||
request_params=target.request_params,
|
||||
)
|
||||
litellm_params[S3_SIGNED_REQUEST_HEADERS_PARAM] = signed_headers # rebind-ok: handed to validate_environment
|
||||
litellm_params[LIST_FILES_PURPOSE_PARAM] = purpose # rebind-ok: handed to the response transform
|
||||
return url, query
|
||||
|
||||
def transform_list_files_response(
|
||||
self,
|
||||
raw_response: httpx.Response,
|
||||
logging_obj: LiteLLMLoggingObj,
|
||||
litellm_params: dict,
|
||||
litellm_params: Mapping[str, object],
|
||||
) -> list[OpenAIFileObject]:
|
||||
raise NotImplementedError("BedrockFilesConfig does not support file listing")
|
||||
if raw_response.status_code >= 400:
|
||||
raise BedrockError(
|
||||
status_code=raw_response.status_code,
|
||||
message=raw_response.text,
|
||||
headers=raw_response.headers,
|
||||
)
|
||||
configured_bucket_name: Final = get_configured_s3_bucket_name(litellm_params)
|
||||
allow_legacy_cloud_file_ids: Final = should_allow_legacy_cloud_file_ids(litellm_params)
|
||||
requested_purpose: Final = litellm_params.get(LIST_FILES_PURPOSE_PARAM)
|
||||
listing: Final = ET.fromstring(raw_response.content)
|
||||
bucket_name: Final = (
|
||||
listing.findtext("{*}Name") or split_configured_cloud_bucket_name(configured_bucket_name)[0]
|
||||
)
|
||||
listed_files: Final = (
|
||||
_listed_managed_file(entry, bucket_name, configured_bucket_name, allow_legacy_cloud_file_ids)
|
||||
for entry in listing.iterfind("{*}Contents")
|
||||
)
|
||||
return [ # mutable-ok: the base files contract returns a list
|
||||
listed_file
|
||||
for listed_file in listed_files
|
||||
if listed_file is not None and (requested_purpose is None or listed_file.purpose == requested_purpose)
|
||||
]
|
||||
|
||||
def transform_file_content_request(
|
||||
self,
|
||||
|
|
@ -1239,40 +1361,53 @@ class BedrockFilesConfig(BaseAWSLLM, BaseFilesConfig):
|
|||
configured_bucket_names=get_configured_s3_bucket_names(litellm_params),
|
||||
allow_legacy_cloud_file_ids=should_allow_legacy_cloud_file_ids(litellm_params),
|
||||
)
|
||||
target: Final = self._s3_request_target(optional_params=optional_params, litellm_params=litellm_params)
|
||||
url: Final = f"{target.endpoint_url}/{bucket_name}/{encode_s3_object_key_for_url(object_key)}"
|
||||
|
||||
# The shared file-content handler passes optional_params={}, so AWS
|
||||
# credentials/region arrive via litellm_params here (unlike the upload
|
||||
# path). s3_region_name wins over aws_region_name, same priority as
|
||||
# get_complete_file_url above.
|
||||
merged_params: Final[dict[str, object]] = {}
|
||||
merged_params.update(litellm_params)
|
||||
merged_params.update(optional_params)
|
||||
request_params: Final = _BedrockS3RequestParams.model_validate(merged_params)
|
||||
|
||||
region_preference: Final = request_params.s3_region_name or request_params.aws_region_name
|
||||
region_params: Final[dict[str, str | None]] = {"aws_region_name": region_preference}
|
||||
aws_region_name: Final = self._get_aws_region_name(optional_params=region_params, model="")
|
||||
|
||||
s3_endpoint_url = (
|
||||
request_params.s3_endpoint_url or f"https://s3.{aws_region_name}.{get_aws_dns_suffix(aws_region_name)}"
|
||||
).rstrip("/")
|
||||
url: Final = f"{s3_endpoint_url}/{bucket_name}/{encode_s3_object_key_for_url(object_key)}"
|
||||
|
||||
litellm_params[S3_SIGNED_GET_HEADERS_PARAM] = self._sign_s3_get_request(
|
||||
signed_headers: Final = self._sign_s3_empty_body_request(
|
||||
method="GET",
|
||||
api_base=url,
|
||||
aws_region_name=aws_region_name,
|
||||
request_params=request_params,
|
||||
aws_region_name=target.aws_region_name,
|
||||
request_params=target.request_params,
|
||||
)
|
||||
litellm_params[S3_SIGNED_REQUEST_HEADERS_PARAM] = signed_headers # rebind-ok: handed to validate_environment
|
||||
return url, {}
|
||||
|
||||
def _sign_s3_get_request(
|
||||
def _s3_request_target(
|
||||
self,
|
||||
optional_params: Mapping[str, object],
|
||||
litellm_params: Mapping[str, object],
|
||||
) -> _S3RequestTarget:
|
||||
"""
|
||||
The shared files handler passes optional_params={}, so AWS credentials and
|
||||
region arrive via litellm_params here (unlike the upload path).
|
||||
s3_region_name wins over aws_region_name, same priority as get_complete_file_url.
|
||||
"""
|
||||
request_params: Final = _BedrockS3RequestParams.model_validate(
|
||||
MappingProxyType({**litellm_params, **optional_params})
|
||||
)
|
||||
region_preference: Final = request_params.s3_region_name or request_params.aws_region_name
|
||||
aws_region_name: Final = self._get_aws_region_name(
|
||||
optional_params={"aws_region_name": region_preference}, # mutable-ok: BaseAWSLLM takes a dict
|
||||
model="",
|
||||
)
|
||||
endpoint_url: Final = (
|
||||
request_params.s3_endpoint_url or f"https://s3.{aws_region_name}.{get_aws_dns_suffix(aws_region_name)}"
|
||||
).rstrip("/")
|
||||
return _S3RequestTarget(
|
||||
endpoint_url=endpoint_url, aws_region_name=aws_region_name, request_params=request_params
|
||||
)
|
||||
|
||||
def _sign_s3_empty_body_request(
|
||||
self,
|
||||
method: str,
|
||||
api_base: str,
|
||||
aws_region_name: str,
|
||||
request_params: _BedrockS3RequestParams,
|
||||
) -> dict[str, str]:
|
||||
) -> Mapping[str, str]:
|
||||
"""
|
||||
SigV4-sign an S3 GetObject request, mirroring `_sign_s3_request` (PUT).
|
||||
SigV4-sign a bodiless S3 request (GetObject, DeleteObject, ListObjectsV2),
|
||||
mirroring `_sign_s3_request` (PUT).
|
||||
"""
|
||||
try:
|
||||
import hashlib
|
||||
|
|
@ -1297,13 +1432,13 @@ class BedrockFilesConfig(BaseAWSLLM, BaseFilesConfig):
|
|||
|
||||
empty_body_hash: Final = hashlib.sha256(b"").hexdigest()
|
||||
aws_request: Final = AWSRequest( # any-ok: botocore AWSRequest is untyped
|
||||
method="GET",
|
||||
method=method,
|
||||
url=api_base,
|
||||
headers={"x-amz-content-sha256": empty_body_hash},
|
||||
headers={"x-amz-content-sha256": empty_body_hash}, # mutable-ok: botocore AWSRequest takes a dict
|
||||
)
|
||||
auth: Final = S3SigV4Auth(credentials, "s3", aws_region_name) # any-ok: botocore untyped
|
||||
auth.add_auth(aws_request) # any-ok: botocore request mutation is untyped
|
||||
return dict(aws_request.headers) # any-ok: botocore headers are untyped
|
||||
return MappingProxyType(dict(aws_request.headers)) # any-ok: botocore headers are untyped
|
||||
|
||||
def transform_file_content_response(
|
||||
self,
|
||||
|
|
|
|||
|
|
@ -1519,7 +1519,7 @@ async def list_files(
|
|||
|
||||
if should_route and credentials is not None:
|
||||
# Use model-based routing with credentials from config
|
||||
prepare_data_with_credentials(data=data, credentials=credentials)
|
||||
prepare_data_with_credentials(data=data, credentials=credentials, include_internal_credentials=True)
|
||||
response = await litellm.afile_list(
|
||||
custom_llm_provider=credentials["custom_llm_provider"],
|
||||
purpose=purpose,
|
||||
|
|
@ -1545,7 +1545,7 @@ async def list_files(
|
|||
model_id=target_model_names_list[0],
|
||||
operation_context="file list",
|
||||
)
|
||||
prepare_data_with_credentials(data=data, credentials=credentials)
|
||||
prepare_data_with_credentials(data=data, credentials=credentials, include_internal_credentials=True)
|
||||
response = await litellm.afile_list(
|
||||
custom_llm_provider=credentials["custom_llm_provider"],
|
||||
purpose=purpose,
|
||||
|
|
|
|||
|
|
@ -1574,3 +1574,99 @@ async def test_batch_retrieve_hook_does_not_claim_attribution():
|
|||
|
||||
managed_files.store_unified_object_id.assert_awaited_once()
|
||||
assert managed_files.store_unified_object_id.await_args.kwargs["persist_attribution"] is False
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_afile_delete_passes_trusted_model_credentials_to_router():
|
||||
"""
|
||||
afile_delete must hand the deployment's credential snapshot to the router
|
||||
call, since Bedrock validates the s3:// file id against the bucket in it.
|
||||
"""
|
||||
from types import MappingProxyType
|
||||
|
||||
managed_files = _make_managed_files_instance()
|
||||
unified_file_id = "unified-file-id"
|
||||
s3_uri = "s3://my-bucket/litellm-bedrock-files/job-123/input.jsonl"
|
||||
managed_files.get_model_file_id_mapping = AsyncMock(return_value={unified_file_id: {"model-123": s3_uri}})
|
||||
managed_files.delete_unified_file_id = AsyncMock(return_value=_make_file_object(unified_file_id))
|
||||
|
||||
mock_router = MagicMock()
|
||||
mock_router.get_deployment_credentials_with_provider = MagicMock(
|
||||
return_value={
|
||||
"custom_llm_provider": "bedrock",
|
||||
"s3_bucket_name": "my-bucket",
|
||||
"aws_region_name": "us-west-2",
|
||||
}
|
||||
)
|
||||
mock_router.afile_delete = AsyncMock(return_value=MagicMock())
|
||||
|
||||
await managed_files.afile_delete(
|
||||
file_id=unified_file_id,
|
||||
litellm_parent_otel_span=None,
|
||||
llm_router=mock_router,
|
||||
)
|
||||
|
||||
call_kwargs = mock_router.afile_delete.call_args.kwargs
|
||||
assert call_kwargs["model"] == "model-123"
|
||||
assert call_kwargs["file_id"] == s3_uri
|
||||
trusted_credentials = call_kwargs["_litellm_internal_model_credentials"]
|
||||
assert isinstance(trusted_credentials, MappingProxyType)
|
||||
assert trusted_credentials["s3_bucket_name"] == "my-bucket"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_afile_delete_bedrock_unified_id_end_to_end(monkeypatch):
|
||||
"""
|
||||
Proxy repro for deleting a Bedrock batch input file by unified id: the
|
||||
s3:// object must be removed via a SigV4-signed S3 DELETE using the
|
||||
deployment's s3_bucket_name (no AWS_S3_BUCKET_NAME env).
|
||||
|
||||
Regression test for "BedrockFilesConfig does not support file deletion"
|
||||
raised on this path.
|
||||
"""
|
||||
import httpx
|
||||
import respx
|
||||
|
||||
import litellm
|
||||
from litellm import Router
|
||||
|
||||
monkeypatch.delenv("AWS_S3_BUCKET_NAME", raising=False)
|
||||
monkeypatch.setattr(litellm, "disable_aiohttp_transport", True)
|
||||
litellm.in_memory_llm_clients_cache.flush_cache()
|
||||
|
||||
router = Router(
|
||||
model_list=[
|
||||
{
|
||||
"model_name": "bedrock-claude",
|
||||
"litellm_params": {
|
||||
"model": "bedrock/anthropic.claude-3-5-sonnet-20240620-v1:0",
|
||||
"aws_access_key_id": "AKIAEXAMPLE",
|
||||
"aws_secret_access_key": "secret",
|
||||
"aws_region_name": "us-west-2",
|
||||
"s3_bucket_name": "my-bucket",
|
||||
},
|
||||
"model_info": {"id": "model-123"},
|
||||
}
|
||||
]
|
||||
)
|
||||
|
||||
managed_files = _make_managed_files_instance()
|
||||
unified_file_id = "unified-file-id"
|
||||
s3_uri = "s3://my-bucket/litellm-bedrock-files/job-123/input.jsonl"
|
||||
managed_files.get_model_file_id_mapping = AsyncMock(return_value={unified_file_id: {"model-123": s3_uri}})
|
||||
managed_files.delete_unified_file_id = AsyncMock(return_value=_make_file_object(unified_file_id))
|
||||
|
||||
expected_url = "https://s3.us-west-2.amazonaws.com/my-bucket/litellm-bedrock-files/job-123/input.jsonl"
|
||||
with respx.mock:
|
||||
route = respx.delete(expected_url).mock(return_value=httpx.Response(204))
|
||||
|
||||
response = await managed_files.afile_delete(
|
||||
file_id=unified_file_id,
|
||||
litellm_parent_otel_span=None,
|
||||
llm_router=router,
|
||||
)
|
||||
|
||||
assert route.called
|
||||
assert route.calls[0].request.headers["Authorization"].startswith("AWS4-HMAC-SHA256")
|
||||
assert response.id == unified_file_id
|
||||
managed_files.delete_unified_file_id.assert_awaited_once_with(unified_file_id, None)
|
||||
|
|
|
|||
|
|
@ -1873,7 +1873,7 @@ class TestBedrockFileContentTransformation:
|
|||
import hashlib
|
||||
|
||||
from litellm.llms.bedrock.files.transformation import (
|
||||
S3_SIGNED_GET_HEADERS_PARAM,
|
||||
S3_SIGNED_REQUEST_HEADERS_PARAM,
|
||||
BedrockFilesConfig,
|
||||
)
|
||||
|
||||
|
|
@ -1889,7 +1889,7 @@ class TestBedrockFileContentTransformation:
|
|||
assert url == self.EXPECTED_URL
|
||||
assert params == {}
|
||||
|
||||
signed_headers = litellm_params[S3_SIGNED_GET_HEADERS_PARAM]
|
||||
signed_headers = litellm_params[S3_SIGNED_REQUEST_HEADERS_PARAM]
|
||||
content_hashes = {
|
||||
value
|
||||
for name, value in signed_headers.items()
|
||||
|
|
@ -2139,7 +2139,7 @@ class TestBedrockFileContentTransformation:
|
|||
def test_s3_region_name_wins_for_content_signing(self, monkeypatch):
|
||||
"""s3_region_name must override aws_region_name for both the URL and the signature."""
|
||||
from litellm.llms.bedrock.files.transformation import (
|
||||
S3_SIGNED_GET_HEADERS_PARAM,
|
||||
S3_SIGNED_REQUEST_HEADERS_PARAM,
|
||||
BedrockFilesConfig,
|
||||
)
|
||||
|
||||
|
|
@ -2154,17 +2154,17 @@ class TestBedrockFileContentTransformation:
|
|||
)
|
||||
|
||||
assert url.startswith("https://s3.eu-west-1.amazonaws.com/")
|
||||
authorization = litellm_params[S3_SIGNED_GET_HEADERS_PARAM]["Authorization"]
|
||||
authorization = litellm_params[S3_SIGNED_REQUEST_HEADERS_PARAM]["Authorization"]
|
||||
assert "/eu-west-1/s3/aws4_request" in authorization
|
||||
|
||||
def test_validate_environment_merges_and_pops_signed_get_headers(self):
|
||||
from litellm.llms.bedrock.files.transformation import (
|
||||
S3_SIGNED_GET_HEADERS_PARAM,
|
||||
S3_SIGNED_REQUEST_HEADERS_PARAM,
|
||||
BedrockFilesConfig,
|
||||
)
|
||||
|
||||
litellm_params = {
|
||||
S3_SIGNED_GET_HEADERS_PARAM: {"Authorization": "AWS4-HMAC-SHA256 test"}
|
||||
S3_SIGNED_REQUEST_HEADERS_PARAM: {"Authorization": "AWS4-HMAC-SHA256 test"}
|
||||
}
|
||||
|
||||
headers = BedrockFilesConfig().validate_environment(
|
||||
|
|
@ -2179,7 +2179,7 @@ class TestBedrockFileContentTransformation:
|
|||
"x-custom": "kept",
|
||||
"Authorization": "AWS4-HMAC-SHA256 test",
|
||||
}
|
||||
assert S3_SIGNED_GET_HEADERS_PARAM not in litellm_params
|
||||
assert S3_SIGNED_REQUEST_HEADERS_PARAM not in litellm_params
|
||||
|
||||
def test_transform_file_content_response_wraps_binary_content(self):
|
||||
import httpx
|
||||
|
|
@ -2379,7 +2379,7 @@ class TestBedrockFilesS3SignatureEncoding:
|
|||
self, monkeypatch: pytest.MonkeyPatch
|
||||
) -> None:
|
||||
from litellm.llms.bedrock.files.transformation import (
|
||||
S3_SIGNED_GET_HEADERS_PARAM,
|
||||
S3_SIGNED_REQUEST_HEADERS_PARAM,
|
||||
BedrockFilesConfig,
|
||||
)
|
||||
|
||||
|
|
@ -2402,7 +2402,7 @@ class TestBedrockFilesS3SignatureEncoding:
|
|||
method="GET",
|
||||
url=url,
|
||||
body=None,
|
||||
headers=litellm_params[S3_SIGNED_GET_HEADERS_PARAM],
|
||||
headers=litellm_params[S3_SIGNED_REQUEST_HEADERS_PARAM],
|
||||
)
|
||||
|
||||
|
||||
|
|
@ -2457,7 +2457,7 @@ def test_sign_s3_request_assumes_role_with_external_id(monkeypatch):
|
|||
assert "ASIAFILESPUTROLE" in authorization
|
||||
|
||||
|
||||
def test_sign_s3_get_request_assumes_role_with_external_id(monkeypatch):
|
||||
def test_sign_s3_empty_body_request_assumes_role_with_external_id(monkeypatch):
|
||||
"""A trust policy requiring sts:ExternalId must be satisfied when signing the S3 download request."""
|
||||
import datetime
|
||||
from unittest.mock import patch
|
||||
|
|
@ -2504,7 +2504,8 @@ def test_sign_s3_get_request_assumes_role_with_external_id(monkeypatch):
|
|||
assert request_params.aws_external_id == "external-id-files-get"
|
||||
|
||||
with patch.object(boto3, "client", return_value=FakeSTSClient()):
|
||||
signed_headers = BedrockFilesConfig()._sign_s3_get_request(
|
||||
signed_headers = BedrockFilesConfig()._sign_s3_empty_body_request(
|
||||
method="GET",
|
||||
api_base="https://s3.us-east-1.amazonaws.com/safe-bucket/litellm-bedrock-files-model-id-abc.jsonl",
|
||||
aws_region_name="us-east-1",
|
||||
request_params=request_params,
|
||||
|
|
@ -2512,3 +2513,449 @@ def test_sign_s3_get_request_assumes_role_with_external_id(monkeypatch):
|
|||
|
||||
authorization = {key.lower(): value for key, value in signed_headers.items()}["authorization"]
|
||||
assert "ASIAFILESGETROLE" in authorization
|
||||
|
||||
|
||||
def _s3_signature_for(method: str, url: str, headers: Mapping[str, str]) -> str:
|
||||
sent = {name.lower(): value for name, value in headers.items()}
|
||||
signed_names = sent["authorization"].split("SignedHeaders=")[1].split(",")[0].split(";")
|
||||
request = AWSRequest(
|
||||
method=method,
|
||||
url=url,
|
||||
headers={name: sent[name] for name in signed_names if name in sent},
|
||||
)
|
||||
request.context["timestamp"] = sent["x-amz-date"]
|
||||
signer = S3SigV4Auth(Credentials("AKIAEXAMPLE", "secret"), "s3", "us-west-2")
|
||||
return signer.signature(signer.string_to_sign(request, signer.canonical_request(request)), request)
|
||||
|
||||
|
||||
def _sent_signature(headers: Mapping[str, str]) -> str:
|
||||
authorization = {name.lower(): value for name, value in headers.items()}["authorization"]
|
||||
return authorization.split("Signature=")[1].strip()
|
||||
|
||||
|
||||
def _bedrock_s3_params() -> dict:
|
||||
return {
|
||||
"aws_access_key_id": "AKIAEXAMPLE",
|
||||
"aws_secret_access_key": "secret",
|
||||
"aws_region_name": "us-west-2",
|
||||
}
|
||||
|
||||
|
||||
def _trusted_bucket_snapshot(**deployment_litellm_params) -> dict:
|
||||
from types import MappingProxyType
|
||||
|
||||
from litellm.types.router import CredentialLiteLLMParams
|
||||
|
||||
snapshot = CredentialLiteLLMParams(**deployment_litellm_params).model_dump(exclude_none=True)
|
||||
return {**_bedrock_s3_params(), "_litellm_internal_model_credentials": MappingProxyType(snapshot)}
|
||||
|
||||
|
||||
class TestBedrockFileDeletionTransformation:
|
||||
"""SigV4-signed S3 DeleteObject for LiteLLM-managed Bedrock batch files."""
|
||||
|
||||
S3_URI = "s3://my-bucket/litellm-bedrock-files/job-123/input.jsonl"
|
||||
EXPECTED_URL = "https://s3.us-west-2.amazonaws.com/my-bucket/litellm-bedrock-files/job-123/input.jsonl"
|
||||
|
||||
def test_transform_delete_file_request_signs_s3_delete(self, monkeypatch):
|
||||
import hashlib
|
||||
|
||||
from litellm.llms.bedrock.files.transformation import (
|
||||
DELETED_FILE_ID_PARAM,
|
||||
S3_SIGNED_REQUEST_HEADERS_PARAM,
|
||||
BedrockFilesConfig,
|
||||
)
|
||||
|
||||
monkeypatch.setenv("AWS_S3_BUCKET_NAME", "my-bucket")
|
||||
litellm_params = _bedrock_s3_params()
|
||||
|
||||
url, params = BedrockFilesConfig().transform_delete_file_request(
|
||||
file_id=self.S3_URI,
|
||||
optional_params={},
|
||||
litellm_params=litellm_params,
|
||||
)
|
||||
|
||||
assert url == self.EXPECTED_URL
|
||||
assert params == {}
|
||||
signed_headers = litellm_params[S3_SIGNED_REQUEST_HEADERS_PARAM]
|
||||
lowered = {name.lower(): value for name, value in signed_headers.items()}
|
||||
assert lowered["x-amz-content-sha256"] == hashlib.sha256(b"").hexdigest()
|
||||
assert "/us-west-2/s3/aws4_request" in lowered["authorization"]
|
||||
assert _sent_signature(signed_headers) == _s3_signature_for("DELETE", url, signed_headers)
|
||||
assert litellm_params[DELETED_FILE_ID_PARAM] == self.S3_URI
|
||||
|
||||
def test_transform_delete_file_request_decodes_unified_file_id(self, monkeypatch):
|
||||
import base64
|
||||
|
||||
from litellm.llms.bedrock.files.transformation import (
|
||||
DELETED_FILE_ID_PARAM,
|
||||
BedrockFilesConfig,
|
||||
)
|
||||
from litellm.types.utils import SpecialEnums
|
||||
|
||||
monkeypatch.setenv("AWS_S3_BUCKET_NAME", "my-bucket")
|
||||
unified_file_id = SpecialEnums.LITELLM_MANAGED_FILE_COMPLETE_STR.value.format(
|
||||
"application/json", "unified-id", "", self.S3_URI, "model-id"
|
||||
)
|
||||
encoded_file_id = base64.urlsafe_b64encode(unified_file_id.encode()).decode().rstrip("=")
|
||||
litellm_params = _bedrock_s3_params()
|
||||
|
||||
url, _ = BedrockFilesConfig().transform_delete_file_request(
|
||||
file_id=encoded_file_id,
|
||||
optional_params={},
|
||||
litellm_params=litellm_params,
|
||||
)
|
||||
|
||||
assert url == self.EXPECTED_URL
|
||||
assert litellm_params[DELETED_FILE_ID_PARAM] == encoded_file_id
|
||||
|
||||
def test_transform_delete_file_request_uses_trusted_snapshot_bucket(self, monkeypatch):
|
||||
from litellm.llms.bedrock.files.transformation import BedrockFilesConfig
|
||||
|
||||
monkeypatch.delenv("AWS_S3_BUCKET_NAME", raising=False)
|
||||
|
||||
url, _ = BedrockFilesConfig().transform_delete_file_request(
|
||||
file_id=self.S3_URI,
|
||||
optional_params={},
|
||||
litellm_params=_trusted_bucket_snapshot(s3_bucket_name="my-bucket"),
|
||||
)
|
||||
|
||||
assert url == self.EXPECTED_URL
|
||||
|
||||
def test_transform_delete_file_request_rejects_foreign_bucket(self, monkeypatch):
|
||||
from litellm.llms.bedrock.files.transformation import BedrockFilesConfig
|
||||
|
||||
monkeypatch.setenv("AWS_S3_BUCKET_NAME", "my-bucket")
|
||||
|
||||
with pytest.raises(ValueError, match="configured storage bucket"):
|
||||
BedrockFilesConfig().transform_delete_file_request(
|
||||
file_id="s3://other-bucket/litellm-bedrock-files/job-123/input.jsonl",
|
||||
optional_params={},
|
||||
litellm_params=_bedrock_s3_params(),
|
||||
)
|
||||
|
||||
def test_transform_delete_file_request_rejects_unmanaged_key(self, monkeypatch):
|
||||
from litellm.llms.bedrock.files.transformation import BedrockFilesConfig
|
||||
|
||||
monkeypatch.setenv("AWS_S3_BUCKET_NAME", "my-bucket")
|
||||
|
||||
with pytest.raises(ValueError, match="LiteLLM-managed"):
|
||||
BedrockFilesConfig().transform_delete_file_request(
|
||||
file_id="s3://my-bucket/private/x.jsonl",
|
||||
optional_params={},
|
||||
litellm_params=_bedrock_s3_params(),
|
||||
)
|
||||
|
||||
def test_transform_delete_file_response_echoes_the_deleted_id(self):
|
||||
import httpx
|
||||
|
||||
from litellm.llms.bedrock.files.transformation import (
|
||||
DELETED_FILE_ID_PARAM,
|
||||
BedrockFilesConfig,
|
||||
)
|
||||
|
||||
deleted = BedrockFilesConfig().transform_delete_file_response(
|
||||
raw_response=httpx.Response(204),
|
||||
logging_obj=MagicMock(),
|
||||
litellm_params={DELETED_FILE_ID_PARAM: self.S3_URI},
|
||||
)
|
||||
|
||||
assert deleted.id == self.S3_URI
|
||||
assert deleted.deleted is True
|
||||
assert deleted.object == "file"
|
||||
|
||||
def test_transform_delete_file_response_raises_on_s3_error(self):
|
||||
import httpx
|
||||
|
||||
from litellm.llms.bedrock.common_utils import BedrockError
|
||||
from litellm.llms.bedrock.files.transformation import BedrockFilesConfig
|
||||
|
||||
with pytest.raises(BedrockError) as excinfo:
|
||||
BedrockFilesConfig().transform_delete_file_response(
|
||||
raw_response=httpx.Response(403, text="<Error><Code>AccessDenied</Code></Error>"),
|
||||
logging_obj=MagicMock(),
|
||||
litellm_params={},
|
||||
)
|
||||
|
||||
assert excinfo.value.status_code == 403
|
||||
|
||||
def test_file_delete_end_to_end_sends_signed_delete(self, monkeypatch):
|
||||
import httpx
|
||||
import respx
|
||||
|
||||
import litellm
|
||||
|
||||
monkeypatch.setenv("AWS_S3_BUCKET_NAME", "my-bucket")
|
||||
|
||||
with respx.mock:
|
||||
route = respx.delete(self.EXPECTED_URL).mock(return_value=httpx.Response(204))
|
||||
|
||||
response = litellm.file_delete(
|
||||
file_id=self.S3_URI,
|
||||
custom_llm_provider="bedrock",
|
||||
**_bedrock_s3_params(),
|
||||
)
|
||||
|
||||
assert route.called
|
||||
request = route.calls[0].request
|
||||
assert request.headers["Authorization"].startswith("AWS4-HMAC-SHA256")
|
||||
assert _sent_signature(request.headers) == _s3_signature_for("DELETE", str(request.url), request.headers)
|
||||
assert response.id == self.S3_URI
|
||||
assert response.deleted is True
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_afile_delete_end_to_end_sends_signed_delete(self, monkeypatch):
|
||||
import httpx
|
||||
import respx
|
||||
|
||||
import litellm
|
||||
|
||||
monkeypatch.setenv("AWS_S3_BUCKET_NAME", "my-bucket")
|
||||
monkeypatch.setattr(litellm, "disable_aiohttp_transport", True)
|
||||
litellm.in_memory_llm_clients_cache.flush_cache()
|
||||
|
||||
with respx.mock:
|
||||
route = respx.delete(self.EXPECTED_URL).mock(return_value=httpx.Response(204))
|
||||
|
||||
response = await litellm.afile_delete(
|
||||
file_id=self.S3_URI,
|
||||
custom_llm_provider="bedrock",
|
||||
**_bedrock_s3_params(),
|
||||
)
|
||||
|
||||
assert route.called
|
||||
request = route.calls[0].request
|
||||
assert _sent_signature(request.headers) == _s3_signature_for("DELETE", str(request.url), request.headers)
|
||||
assert response.id == self.S3_URI
|
||||
assert response.deleted is True
|
||||
|
||||
|
||||
class TestBedrockFileListTransformation:
|
||||
"""SigV4-signed S3 ListObjectsV2 over the LiteLLM-managed key prefixes."""
|
||||
|
||||
BUCKET_URL = "https://s3.us-west-2.amazonaws.com/my-bucket/"
|
||||
MANAGED_QUERY = {"list-type": "2", "prefix": "litellm-b"}
|
||||
LISTING = b"""<?xml version="1.0" encoding="UTF-8"?>
|
||||
<ListBucketResult xmlns="http://s3.amazonaws.com/doc/2006-03-01/">
|
||||
<Name>my-bucket</Name>
|
||||
<Prefix>litellm-b</Prefix>
|
||||
<KeyCount>4</KeyCount>
|
||||
<IsTruncated>false</IsTruncated>
|
||||
<Contents>
|
||||
<Key>litellm-bedrock-files-model-abc.jsonl</Key>
|
||||
<LastModified>2026-09-01T10:00:00.000Z</LastModified>
|
||||
<Size>120</Size>
|
||||
</Contents>
|
||||
<Contents>
|
||||
<Key>litellm-bedrock-files/job-123/input.jsonl</Key>
|
||||
<LastModified>2026-09-02T11:30:00.000Z</LastModified>
|
||||
<Size>340</Size>
|
||||
</Contents>
|
||||
<Contents>
|
||||
<Key>litellm-batch-outputs/job-123/input.jsonl.out</Key>
|
||||
<LastModified>2026-09-03T12:45:00.000Z</LastModified>
|
||||
<Size>560</Size>
|
||||
</Contents>
|
||||
<Contents>
|
||||
<Key>litellm-bogus/other.jsonl</Key>
|
||||
<LastModified>2026-09-03T12:45:00.000Z</LastModified>
|
||||
<Size>1</Size>
|
||||
</Contents>
|
||||
</ListBucketResult>"""
|
||||
BATCH_IDS = (
|
||||
"s3://my-bucket/litellm-bedrock-files-model-abc.jsonl",
|
||||
"s3://my-bucket/litellm-bedrock-files/job-123/input.jsonl",
|
||||
)
|
||||
OUTPUT_ID = "s3://my-bucket/litellm-batch-outputs/job-123/input.jsonl.out"
|
||||
|
||||
def test_transform_list_files_request_signs_managed_prefix_listing(self, monkeypatch):
|
||||
from litellm.llms.bedrock.files.transformation import (
|
||||
LIST_FILES_PURPOSE_PARAM,
|
||||
S3_SIGNED_REQUEST_HEADERS_PARAM,
|
||||
BedrockFilesConfig,
|
||||
)
|
||||
|
||||
monkeypatch.setenv("AWS_S3_BUCKET_NAME", "my-bucket")
|
||||
litellm_params = _bedrock_s3_params()
|
||||
|
||||
url, params = BedrockFilesConfig().transform_list_files_request(
|
||||
purpose="batch",
|
||||
optional_params={},
|
||||
litellm_params=litellm_params,
|
||||
)
|
||||
|
||||
assert url == self.BUCKET_URL
|
||||
assert params == self.MANAGED_QUERY
|
||||
signed_headers = litellm_params[S3_SIGNED_REQUEST_HEADERS_PARAM]
|
||||
assert _sent_signature(signed_headers) == _s3_signature_for(
|
||||
"GET", f"{url}?list-type=2&prefix=litellm-b", signed_headers
|
||||
)
|
||||
assert litellm_params[LIST_FILES_PURPOSE_PARAM] == "batch"
|
||||
|
||||
def test_transform_list_files_request_scopes_to_configured_prefix(self, monkeypatch):
|
||||
from litellm.llms.bedrock.files.transformation import (
|
||||
S3_SIGNED_REQUEST_HEADERS_PARAM,
|
||||
BedrockFilesConfig,
|
||||
)
|
||||
|
||||
monkeypatch.setenv("AWS_S3_BUCKET_NAME", "my-bucket/LLM AI Projects")
|
||||
litellm_params = _bedrock_s3_params()
|
||||
|
||||
url, params = BedrockFilesConfig().transform_list_files_request(
|
||||
purpose=None,
|
||||
optional_params={},
|
||||
litellm_params=litellm_params,
|
||||
)
|
||||
|
||||
assert url == self.BUCKET_URL
|
||||
assert params == {"list-type": "2", "prefix": "LLM AI Projects/litellm-b"}
|
||||
signed_headers = litellm_params[S3_SIGNED_REQUEST_HEADERS_PARAM]
|
||||
assert _sent_signature(signed_headers) == _s3_signature_for(
|
||||
"GET", f"{url}?list-type=2&prefix=LLM%20AI%20Projects%2Flitellm-b", signed_headers
|
||||
)
|
||||
|
||||
def test_transform_list_files_request_uses_trusted_snapshot_bucket(self, monkeypatch):
|
||||
from litellm.llms.bedrock.files.transformation import BedrockFilesConfig
|
||||
|
||||
monkeypatch.delenv("AWS_S3_BUCKET_NAME", raising=False)
|
||||
|
||||
url, params = BedrockFilesConfig().transform_list_files_request(
|
||||
purpose=None,
|
||||
optional_params={},
|
||||
litellm_params=_trusted_bucket_snapshot(s3_bucket_name="my-bucket"),
|
||||
)
|
||||
|
||||
assert url == self.BUCKET_URL
|
||||
assert params == self.MANAGED_QUERY
|
||||
|
||||
def _list_response(self, purpose: str | None, listing: bytes | None = None, status_code: int = 200):
|
||||
import httpx
|
||||
|
||||
from litellm.llms.bedrock.files.transformation import (
|
||||
LIST_FILES_PURPOSE_PARAM,
|
||||
BedrockFilesConfig,
|
||||
)
|
||||
|
||||
return BedrockFilesConfig().transform_list_files_response(
|
||||
raw_response=httpx.Response(status_code, content=listing if listing is not None else self.LISTING),
|
||||
logging_obj=MagicMock(),
|
||||
litellm_params={**_bedrock_s3_params(), LIST_FILES_PURPOSE_PARAM: purpose},
|
||||
)
|
||||
|
||||
def test_transform_list_files_response_maps_managed_objects(self, monkeypatch):
|
||||
from datetime import datetime, timezone
|
||||
|
||||
monkeypatch.setenv("AWS_S3_BUCKET_NAME", "my-bucket")
|
||||
|
||||
files = self._list_response(purpose=None)
|
||||
|
||||
assert [file.id for file in files] == [*self.BATCH_IDS, self.OUTPUT_ID]
|
||||
assert [file.purpose for file in files] == ["batch", "batch", "batch_output"]
|
||||
assert [file.bytes for file in files] == [120, 340, 560]
|
||||
assert [file.filename for file in files] == [
|
||||
"litellm-bedrock-files-model-abc.jsonl",
|
||||
"input.jsonl",
|
||||
"input.jsonl.out",
|
||||
]
|
||||
assert files[1].created_at == int(datetime(2026, 9, 2, 11, 30, tzinfo=timezone.utc).timestamp())
|
||||
assert {file.object for file in files} == {"file"}
|
||||
assert {file.status for file in files} == {"uploaded"}
|
||||
|
||||
def test_transform_list_files_response_filters_by_purpose(self, monkeypatch):
|
||||
monkeypatch.setenv("AWS_S3_BUCKET_NAME", "my-bucket")
|
||||
|
||||
assert [file.id for file in self._list_response(purpose="batch")] == list(self.BATCH_IDS)
|
||||
assert [file.id for file in self._list_response(purpose="batch_output")] == [self.OUTPUT_ID]
|
||||
|
||||
def test_transform_list_files_response_scopes_to_configured_prefix(self, monkeypatch):
|
||||
monkeypatch.setenv("AWS_S3_BUCKET_NAME", "my-bucket/team-a")
|
||||
listing = b"""<?xml version="1.0" encoding="UTF-8"?>
|
||||
<ListBucketResult xmlns="http://s3.amazonaws.com/doc/2006-03-01/">
|
||||
<Name>my-bucket</Name>
|
||||
<Contents><Key>team-a/litellm-bedrock-files/job-1/input.jsonl</Key><Size>10</Size></Contents>
|
||||
<Contents><Key>team-a/litellm-batch-outputs/job-1/input.jsonl.out</Key><Size>20</Size></Contents>
|
||||
<Contents><Key>litellm-bedrock-files/job-2/input.jsonl</Key><Size>30</Size></Contents>
|
||||
</ListBucketResult>"""
|
||||
|
||||
files = self._list_response(purpose=None, listing=listing)
|
||||
|
||||
assert [(file.id, file.purpose) for file in files] == [
|
||||
("s3://my-bucket/team-a/litellm-bedrock-files/job-1/input.jsonl", "batch"),
|
||||
("s3://my-bucket/team-a/litellm-batch-outputs/job-1/input.jsonl.out", "batch_output"),
|
||||
]
|
||||
|
||||
def test_transform_list_files_response_raises_on_s3_error(self, monkeypatch):
|
||||
from litellm.llms.bedrock.common_utils import BedrockError
|
||||
|
||||
monkeypatch.setenv("AWS_S3_BUCKET_NAME", "my-bucket")
|
||||
|
||||
with pytest.raises(BedrockError) as excinfo:
|
||||
self._list_response(purpose=None, listing=b"<Error><Code>AccessDenied</Code></Error>", status_code=403)
|
||||
|
||||
assert excinfo.value.status_code == 403
|
||||
|
||||
def test_file_list_end_to_end_sends_signed_listing(self, monkeypatch):
|
||||
import httpx
|
||||
import respx
|
||||
|
||||
import litellm
|
||||
|
||||
monkeypatch.setenv("AWS_S3_BUCKET_NAME", "my-bucket")
|
||||
|
||||
with respx.mock:
|
||||
route = respx.get(self.BUCKET_URL, params__contains=self.MANAGED_QUERY).mock(
|
||||
return_value=httpx.Response(200, content=self.LISTING)
|
||||
)
|
||||
|
||||
files = litellm.file_list(custom_llm_provider="bedrock", purpose="batch", **_bedrock_s3_params())
|
||||
|
||||
assert route.called
|
||||
request = route.calls[0].request
|
||||
assert request.headers["Authorization"].startswith("AWS4-HMAC-SHA256")
|
||||
assert _sent_signature(request.headers) == _s3_signature_for("GET", str(request.url), request.headers)
|
||||
assert [file.id for file in files] == list(self.BATCH_IDS)
|
||||
|
||||
def test_file_list_uses_trusted_snapshot_bucket_without_env(self, monkeypatch):
|
||||
import httpx
|
||||
import respx
|
||||
|
||||
import litellm
|
||||
|
||||
monkeypatch.delenv("AWS_S3_BUCKET_NAME", raising=False)
|
||||
|
||||
with respx.mock:
|
||||
route = respx.get(self.BUCKET_URL, params__contains=self.MANAGED_QUERY).mock(
|
||||
return_value=httpx.Response(200, content=self.LISTING)
|
||||
)
|
||||
|
||||
files = litellm.file_list(
|
||||
custom_llm_provider="bedrock",
|
||||
**_trusted_bucket_snapshot(s3_bucket_name="my-bucket"),
|
||||
)
|
||||
|
||||
assert route.called
|
||||
assert [file.id for file in files] == [*self.BATCH_IDS, self.OUTPUT_ID]
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_afile_list_end_to_end_sends_signed_listing(self, monkeypatch):
|
||||
import httpx
|
||||
import respx
|
||||
|
||||
import litellm
|
||||
|
||||
monkeypatch.setenv("AWS_S3_BUCKET_NAME", "my-bucket")
|
||||
monkeypatch.setattr(litellm, "disable_aiohttp_transport", True)
|
||||
litellm.in_memory_llm_clients_cache.flush_cache()
|
||||
|
||||
with respx.mock:
|
||||
route = respx.get(self.BUCKET_URL, params__contains=self.MANAGED_QUERY).mock(
|
||||
return_value=httpx.Response(200, content=self.LISTING)
|
||||
)
|
||||
|
||||
files = await litellm.afile_list(
|
||||
custom_llm_provider="bedrock", purpose="batch_output", **_bedrock_s3_params()
|
||||
)
|
||||
|
||||
assert route.called
|
||||
request = route.calls[0].request
|
||||
assert _sent_signature(request.headers) == _s3_signature_for("GET", str(request.url), request.headers)
|
||||
assert [file.id for file in files] == [self.OUTPUT_ID]
|
||||
|
|
|
|||
|
|
@ -4666,3 +4666,70 @@ def test_create_file_path_traversal_filename_rejected_before_forwarding(monkeypa
|
|||
assert error["param"] == "file"
|
||||
assert "traversal" in error["message"].lower()
|
||||
assert forwarded_calls == []
|
||||
|
||||
|
||||
def test_list_files_target_model_names_passes_trusted_bedrock_credentials(
|
||||
mocker: MockerFixture, monkeypatch
|
||||
):
|
||||
"""
|
||||
GET /v1/files?target_model_names=<bedrock model> must hand the deployment's
|
||||
immutable credential snapshot to litellm.afile_list, since Bedrock resolves
|
||||
the S3 bucket to list from that snapshot rather than from request params.
|
||||
"""
|
||||
from types import MappingProxyType
|
||||
|
||||
import litellm.proxy.proxy_server as ps
|
||||
from litellm.proxy._types import LitellmUserRoles
|
||||
|
||||
bedrock_router = Router(
|
||||
model_list=[
|
||||
{
|
||||
"model_name": "bedrock-claude",
|
||||
"litellm_params": {
|
||||
"model": "bedrock/anthropic.claude-3-5-sonnet-20240620-v1:0",
|
||||
"aws_access_key_id": "AKIAEXAMPLE",
|
||||
"aws_secret_access_key": "secret",
|
||||
"aws_region_name": "us-west-2",
|
||||
"s3_bucket_name": "my-bucket",
|
||||
},
|
||||
},
|
||||
]
|
||||
)
|
||||
|
||||
proxy_logging_obj = setup_proxy_logging_object(monkeypatch, bedrock_router)
|
||||
monkeypatch.setattr("litellm.proxy.proxy_server.master_key", None)
|
||||
monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", None)
|
||||
monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", bedrock_router)
|
||||
proxy_logging_obj.update_request_status = mocker.AsyncMock()
|
||||
proxy_logging_obj.post_call_success_hook = mocker.AsyncMock(return_value=[])
|
||||
proxy_logging_obj.post_call_failure_hook = mocker.AsyncMock()
|
||||
|
||||
captured_kwargs: dict = {}
|
||||
|
||||
async def _mock_afile_list(**kwargs):
|
||||
captured_kwargs.update(kwargs)
|
||||
return []
|
||||
|
||||
monkeypatch.setattr(litellm, "afile_list", _mock_afile_list)
|
||||
|
||||
app.dependency_overrides[ps.user_api_key_auth] = lambda: UserAPIKeyAuth(
|
||||
api_key="test-key",
|
||||
user_role=LitellmUserRoles.PROXY_ADMIN,
|
||||
user_id="test-user",
|
||||
)
|
||||
|
||||
try:
|
||||
response = client.get(
|
||||
"/v1/files?target_model_names=bedrock-claude&purpose=batch",
|
||||
headers={"Authorization": "Bearer test-key"},
|
||||
)
|
||||
finally:
|
||||
app.dependency_overrides.pop(ps.user_api_key_auth, None)
|
||||
|
||||
assert response.status_code == 200, response.text
|
||||
assert captured_kwargs["custom_llm_provider"] == "bedrock"
|
||||
assert captured_kwargs["purpose"] == "batch"
|
||||
trusted_credentials = captured_kwargs["_litellm_internal_model_credentials"]
|
||||
assert isinstance(trusted_credentials, MappingProxyType)
|
||||
assert trusted_credentials["s3_bucket_name"] == "my-bucket"
|
||||
proxy_logging_obj.post_call_failure_hook.assert_not_called()
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue