scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-10-07 02:59:05 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 103

feat(bedrock): support file content retrieval for batch output files (#30595)

Browse source 
Implements transform_file_content_request and transform_file_content_response
in BedrockFilesConfig so GET /v1/files/{id}/content works for Bedrock batch
files. The request transform resolves the file id (direct s3:// URI or base64
unified id) to its S3 object, validates bucket and key prefix against the
server-configured bucket, and SigV4-signs an S3 GetObject using the same
credential and region resolution as the existing upload path. The credential
and region params are validated into a typed model at the boundary, so the only
untyped values left are the botocore signing primitives.

Also fixes the proxy managed-files path: CredentialLiteLLMParams now carries
s3_bucket_name (previously dropped when building deployment credentials) and
the managed-files hook passes the deployment credential snapshot when routing
afile_content, so unified-id content retrieval works with per-model bucket
config instead of only the AWS_S3_BUCKET_NAME env var.

Preserves managed-file access control: the proxy file-content endpoint now
rejects raw cloud-storage ids (s3://, gs://), which would otherwise skip the
owner/team check that only runs for unified ids and let a caller read another
tenant's batch output by its object key. Managed outputs are reachable only
through their unified file id. The afile_content "not found" error now reports
the caller's unified id rather than the resolved internal S3 URI.

Fixes #16186, #15563
This commit is contained in:
AD Mohanraj 2026-06-18 05:34:25 -07:00 • committed by GitHub
parent 85828da695
commit df029adeb4
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
10 changed files with 741 additions and 80 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

23
enterprise/litellm_enterprise/proxy/hooks/managed_files.py
View file

@ -1,9 +1,9 @@
# What is this?
## This hook is used to check for LiteLLM managed files in the request body, and replace them with model-specific file id
import asyncio
import base64
import json
from types import MappingProxyType
from typing import TYPE_CHECKING, Any, Dict, List, Literal, Optional, Union, cast
from fastapi import HTTPException
@ -1472,8 +1472,8 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
error_message += f" (showing {MAX_BATCHES_IN_ERROR} most recent): {', '.join(batch_statuses)}. "
error_message += (
f"To delete this file before complete cost tracking, please delete or cancel the referencing batch(es) first. "
f"Alternatively, wait for all batches to complete and for cost to be computed (batch_processed=true)."
"To delete this file before complete cost tracking, please delete or cancel the referencing batch(es) first. "
"Alternatively, wait for all batches to complete and for cost to be computed (batch_processed=true)."
)
# Record blocked deletion metric
@ -1550,9 +1550,22 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
if specific_model_file_id_mapping:
exception_dict = {}
for model_id, file_id in specific_model_file_id_mapping.items():
for model_id, provider_file_id in specific_model_file_id_mapping.items():
try:
return await llm_router.afile_content(model=model_id, file_id=file_id, **data) # type: ignore
# Cloud-storage providers (e.g. Bedrock S3) validate file ids
# against the deployment's configured bucket, which they only
# trust from this immutable server-side snapshot, never from
# request params.
credentials = llm_router.get_deployment_credentials_with_provider(
model_id=model_id
)
if credentials is not None:
data["_litellm_internal_model_credentials"] = cast(
Dict, MappingProxyType(dict(credentials))
)
else:
data.pop("_litellm_internal_model_credentials", None)
return await llm_router.afile_content(model=model_id, file_id=provider_file_id, **data) # type: ignore
except Exception as e:
exception_dict[model_id] = str(e)
raise Exception(

15
litellm/litellm_core_utils/cloud_storage_security.py
View file

@ -15,8 +15,23 @@ BEDROCK_MANAGED_S3_PREFIXES = (
BEDROCK_MANAGED_S3_UPLOAD_PREFIX,
BEDROCK_MANAGED_S3_OUTPUT_PREFIX,
)
MANAGED_CLOUD_STORAGE_SCHEMES = ("s3://", "gs://")
_MAPPING_PROXY_TYPE: type = type(MappingProxyType({}))
def is_managed_cloud_storage_uri(file_id: str) -> bool:
"""
True if file_id is a raw cloud-storage object URI (e.g. ``s3://bucket/key``).
These are internal provider artifacts. On the multi-tenant proxy they must be
retrieved through their managed unified file id so owner/team access is enforced;
a raw URI supplied by a caller bypasses that check.
"""
return isinstance(file_id, str) and file_id.startswith(
MANAGED_CLOUD_STORAGE_SCHEMES
)
_SAFE_OBJECT_COMPONENT_PATTERN = re.compile(r"[^A-Za-z0-9._-]+")

68
litellm/llms/bedrock/files/handler.py
View file

@ -1,8 +1,6 @@
import asyncio
import base64
import os
from types import MappingProxyType
from typing import Any, Coroutine, Mapping, Optional, Tuple, Union, cast
from collections.abc import Mapping
from typing import Any, Coroutine, Optional, Tuple, Union
import httpx
@ -17,7 +15,6 @@ from litellm.types.llms.openai import (
FileContentRequest,
HttpxBinaryResponseContent,
)
from litellm.types.utils import SpecialEnums
from ..base_aws_llm import BaseAWSLLM
@ -37,40 +34,9 @@ class BedrockFilesHandler(BaseAWSLLM):
)
def _extract_s3_uri_from_file_id(self, file_id: str) -> str:
"""
Extract S3 URI from encoded file ID.
from .transformation import extract_s3_uri_from_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")
return extract_s3_uri_from_file_id(file_id)
def _parse_s3_uri(
self,
@ -95,26 +61,12 @@ class BedrockFilesHandler(BaseAWSLLM):
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
def _get_configured_s3_bucket_name(
self, litellm_params: Mapping[str, object]
) -> str:
from .transformation import get_configured_s3_bucket_name
return get_configured_s3_bucket_name(litellm_params)
async def afile_content(
self,

224
litellm/llms/bedrock/files/transformation.py
View file

@ -1,23 +1,37 @@
import base64
import json
import os
import time
from typing import Any, Dict, List, Optional, Tuple, Union
from collections.abc import Mapping, MutableMapping
from types import MappingProxyType
from typing import (
Any,
Dict,
List,
Optional,
Tuple,
Union,
)
from urllib.parse import unquote
import httpx
from httpx import Headers, Response
from openai.types.file_deleted import FileDeleted
from pydantic import BaseModel, ConfigDict
from litellm._logging import verbose_logger
from litellm._uuid import uuid
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
@ -28,18 +42,98 @@ from litellm.llms.base_llm.files.transformation import (
from litellm.types.llms.openai import (
AllMessageValues,
CreateFileRequest,
FileContentRequest,
FileTypes,
HttpxBinaryResponseContent,
OpenAICreateFileRequestOptionalParams,
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
from ..common_utils import BedrockError
# 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).
# Same pattern as the `upload_url` handoff in `transform_create_file_request`.
S3_SIGNED_GET_HEADERS_PARAM = "_s3_signed_get_headers"
class _BedrockS3RequestParams(BaseModel):
"""Typed view of the credential/region params the S3 GetObject path reads."""
model_config = ConfigDict(extra="ignore")
aws_access_key_id: str | None = None
aws_secret_access_key: str | None = None
aws_session_token: str | None = None
aws_region_name: str | None = None
aws_session_name: str | None = None
aws_profile_name: str | None = None
aws_role_name: str | None = None
aws_web_identity_token: str | None = None
aws_sts_endpoint: str | None = None
s3_region_name: str | None = None
s3_endpoint_url: str | None = None
class _TrustedS3ModelCredentials(BaseModel):
"""The S3 bucket the server trusts file ids against, from the deployment snapshot."""
model_config = ConfigDict(extra="ignore")
s3_bucket_name: str | None = None
def extract_s3_uri_from_file_id(file_id: str) -> str:
"""
Resolve a Bedrock file id to its S3 URI.
Accepts either a base64-encoded LiteLLM unified file id (whose decoded
form carries `llm_output_file_id,s3://...`) or a direct `s3://` URI.
"""
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_bucket_name(litellm_params: Mapping[str, object]) -> str:
"""
Resolve the server-configured S3 bucket for Bedrock file operations.
Only trusts the immutable server-side credential snapshot or the
environment; never a request-supplied param, since the bucket is what
`validate_managed_cloud_file_id` checks file ids against.
"""
trusted_model_credentials = litellm_params.get(
"_litellm_internal_model_credentials"
)
bucket_name: str | None = None
if isinstance(trusted_model_credentials, MappingProxyType):
snapshot: dict[str, object] = {}
snapshot.update(trusted_model_credentials) # any-ok: untyped snapshot
bucket_name = _TrustedS3ModelCredentials.model_validate(snapshot).s3_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
class BedrockFilesConfig(BaseAWSLLM, BaseFilesConfig):
"""
@ -63,16 +157,21 @@ class BedrockFilesConfig(BaseAWSLLM, BaseFilesConfig):
def validate_environment(
self,
headers: dict,
headers: MutableMapping[str, object],
model: str,
messages: List[AllMessageValues],
optional_params: dict,
litellm_params: dict,
api_key: Optional[str] = None,
api_base: Optional[str] = None,
litellm_params: MutableMapping[str, object],
api_key: str | None = None,
api_base: str | None = None,
) -> dict:
# No additional headers needed for S3 uploads - AWS credentials handled by BaseAWSLLM
return headers
result: dict[str, object] = {}
result.update(headers)
signed_headers = litellm_params.pop(S3_SIGNED_GET_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
return result
def _get_content_from_openai_file(self, openai_file_content: FileTypes) -> str:
"""
@ -927,23 +1026,114 @@ class BedrockFilesConfig(BaseAWSLLM, BaseFilesConfig):
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_content_request: FileContentRequest,
optional_params: Mapping[str, object],
litellm_params: MutableMapping[str, object],
) -> tuple[str, dict[str, str]]:
"""
Build a SigV4-signed S3 GetObject request for a Bedrock batch file.
Bedrock batch file ids are `s3://bucket/key` URIs (or unified ids
that decode to one); the bucket and key are validated against the
server-configured bucket before any request is signed.
"""
file_id = file_content_request.get("file_id")
if not file_id:
raise ValueError("file_id is required for Bedrock file content retrieval")
s3_uri = extract_s3_uri_from_file_id(file_id)
bucket_name, object_key = validate_managed_cloud_file_id(
file_id=s3_uri,
scheme="s3://",
configured_bucket_name=get_configured_s3_bucket_name(litellm_params),
allowed_object_prefixes=BEDROCK_MANAGED_S3_PREFIXES,
allow_legacy_cloud_file_ids=should_allow_legacy_cloud_file_ids(
litellm_params
),
)
# 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: dict[str, object] = {}
merged_params.update(litellm_params)
merged_params.update(optional_params)
request_params = _BedrockS3RequestParams.model_validate(merged_params)
region_preference = (
request_params.s3_region_name or request_params.aws_region_name
)
region_params: dict[str, str | None] = {"aws_region_name": region_preference}
aws_region_name = 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}.amazonaws.com"
).rstrip("/")
url = 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(
api_base=url,
aws_region_name=aws_region_name,
request_params=request_params,
)
return url, {}
def _sign_s3_get_request(
self,
api_base: str,
aws_region_name: str,
request_params: _BedrockS3RequestParams,
) -> dict[str, str]:
"""
SigV4-sign an S3 GetObject request, mirroring `_sign_s3_request` (PUT).
"""
try:
import hashlib
from botocore.auth import SigV4Auth
from botocore.awsrequest import AWSRequest
except ImportError:
raise ImportError("Missing boto3 to call bedrock. Run 'pip install boto3'.")
credentials = self.get_credentials( # any-ok: boto3 Credentials is untyped
aws_access_key_id=request_params.aws_access_key_id,
aws_secret_access_key=request_params.aws_secret_access_key,
aws_session_token=request_params.aws_session_token,
aws_region_name=aws_region_name,
aws_session_name=request_params.aws_session_name,
aws_profile_name=request_params.aws_profile_name,
aws_role_name=request_params.aws_role_name,
aws_web_identity_token=request_params.aws_web_identity_token,
aws_sts_endpoint=request_params.aws_sts_endpoint,
)
empty_body_hash = hashlib.sha256(b"").hexdigest()
aws_request = AWSRequest( # any-ok: botocore AWSRequest is untyped
method="GET",
url=api_base,
headers={"x-amz-content-sha256": empty_body_hash},
)
auth = SigV4Auth(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
def transform_file_content_response(
self,
raw_response: httpx.Response,
logging_obj: LiteLLMLoggingObj,
litellm_params: dict,
) -> HttpxBinaryResponseContent:
raise NotImplementedError(
"BedrockFilesConfig does not support file content retrieval"
)
if raw_response.status_code >= 400:
raise BedrockError(
status_code=raw_response.status_code,
message=raw_response.text,
headers=raw_response.headers,
)
return HttpxBinaryResponseContent(response=raw_response)
class BedrockJsonlFilesTransformation:

12
litellm/proxy/openai_files_endpoints/files_endpoints.py
View file

@ -38,6 +38,9 @@ from litellm.proxy.common_utils.openai_endpoint_utils import (
get_custom_llm_provider_from_request_headers,
get_custom_llm_provider_from_request_query,
)
from litellm.litellm_core_utils.cloud_storage_security import (
is_managed_cloud_storage_uri,
)
from litellm.proxy.openai_files_endpoints.common_utils import (
_is_base64_encoded_unified_file_id,
encode_file_id_with_model,
@ -726,6 +729,15 @@ async def get_file_content(
}
)
else:
# A raw cloud-storage URI (s3://, gs://) supplied here would skip the
# managed-file owner/team check that only runs for unified ids, letting
# a caller read another tenant's object by its key. Such objects are only
# reachable through their managed unified id.
if is_managed_cloud_storage_uri(file_id):
raise HTTPException(
status_code=400,
detail="Raw cloud storage file ids cannot be retrieved directly. Use the LiteLLM managed file id returned when the file was created.",
)
# Check for model-based credential routing
(
should_route,

1
litellm/types/router.py
View file

@ -186,6 +186,7 @@ class CredentialLiteLLMParams(BaseModel):
aws_region_name: Optional[str] = None
aws_bedrock_runtime_endpoint: Optional[str] = None
aws_bedrock_project_id: Optional[str] = None
s3_bucket_name: Optional[str] = None
## IBM WATSONX ##
watsonx_region_name: Optional[str] = None

130
tests/test_litellm/enterprise/proxy/test_managed_files_hook.py
View file

@ -255,3 +255,133 @@ async def test_should_skip_non_file_unified_id_on_output_file_id():
assert batch_response.output_file_id == batch_unified
mock_afile_retrieve.assert_not_called()
managed_files.store_unified_file_id.assert_not_awaited()
@pytest.mark.asyncio
async def test_afile_content_passes_trusted_model_credentials_to_router():
"""
afile_content must hand the deployment's credential snapshot to the router
call as an immutable server-side mapping. Cloud-storage providers (Bedrock
S3) validate file ids against the bucket in that snapshot, so without it
unified-id content retrieval only works when AWS_S3_BUCKET_NAME is set.
"""
from types import MappingProxyType
managed_files = _make_managed_files_instance()
unified_file_id = "unified-file-id"
s3_uri = "s3://my-bucket/litellm-batch-outputs/job-123/input.jsonl.out"
managed_files.get_model_file_id_mapping = AsyncMock(
return_value={unified_file_id: {"model-123": s3_uri}}
)
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_content = AsyncMock(return_value=MagicMock())
await managed_files.afile_content(
file_id=unified_file_id,
litellm_parent_otel_span=None,
llm_router=mock_router,
)
call_kwargs = mock_router.afile_content.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_content_bedrock_unified_id_end_to_end(monkeypatch):
"""
Proxy repro for Bedrock batch output retrieval: a unified file id that
resolves to an s3:// output object must be fetched via a SigV4-signed S3
GET using the deployment's s3_bucket_name (no AWS_S3_BUCKET_NAME env).
Regression test for "BedrockFilesConfig does not support file content
retrieval" 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-batch-outputs/job-123/input.jsonl.out"
managed_files.get_model_file_id_mapping = AsyncMock(
return_value={unified_file_id: {"model-123": s3_uri}}
)
expected_url = "https://s3.us-west-2.amazonaws.com/my-bucket/litellm-batch-outputs/job-123/input.jsonl.out"
with respx.mock:
route = respx.get(expected_url).mock(
return_value=httpx.Response(200, content=b'{"recordId": "x"}')
)
response = await managed_files.afile_content(
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.content == b'{"recordId": "x"}'
@pytest.mark.asyncio
async def test_afile_content_error_reports_unified_id_not_provider_uri():
"""When every model attempt fails, the error must name the caller's unified
file id, never the resolved internal s3:// URI (no internal-path leak)."""
managed_files = _make_managed_files_instance()
unified_file_id = "litellm_proxy_unified_id_abc"
s3_uri = "s3://my-bucket/litellm-batch-outputs/job-123/input.jsonl.out"
managed_files.get_model_file_id_mapping = AsyncMock(
return_value={unified_file_id: {"model-123": s3_uri}}
)
mock_router = MagicMock()
mock_router.get_deployment_credentials_with_provider = MagicMock(return_value=None)
mock_router.afile_content = AsyncMock(side_effect=Exception("deployment failed"))
with pytest.raises(Exception) as exc_info:
await managed_files.afile_content(
file_id=unified_file_id,
litellm_parent_otel_span=None,
llm_router=mock_router,
)
message = str(exc_info.value)
assert unified_file_id in message
assert s3_uri not in message

15
tests/test_litellm/litellm_core_utils/test_cloud_storage_security.py Normal file
View file

@ -0,0 +1,15 @@
from litellm.litellm_core_utils.cloud_storage_security import (
is_managed_cloud_storage_uri,
)
def test_is_managed_cloud_storage_uri_detects_raw_object_uris():
assert is_managed_cloud_storage_uri("s3://bucket/litellm-batch-outputs/x.jsonl.out")
assert is_managed_cloud_storage_uri("gs://bucket/litellm-vertex-files/x")
def test_is_managed_cloud_storage_uri_ignores_provider_and_unified_ids():
# Plain provider ids and base64 unified ids carry no storage scheme.
assert not is_managed_cloud_storage_uri("file-abc123")
assert not is_managed_cloud_storage_uri("bGl0ZWxsbV9wcm94eQ==")
assert not is_managed_cloud_storage_uri("")

314
tests/test_litellm/llms/bedrock/files/test_bedrock_files_transformation.py
View file

@ -4,8 +4,11 @@ Test bedrock files transformation functionality
import json
import os
from unittest.mock import MagicMock
from urllib.parse import unquote, urlparse
import pytest
from litellm.llms.bedrock.files.transformation import BedrockJsonlFilesTransformation
@ -1173,3 +1176,314 @@ class TestBedrockFilesEmbeddingTransformation:
assert not BedrockFilesConfig._is_embedding_record(
{"url": "/v1/responses", "body": {"input": "x"}}
)
class TestBedrockFileContentTransformation:
"""SigV4-signed S3 GetObject retrieval of Bedrock batch output files."""
S3_URI = "s3://my-bucket/litellm-batch-outputs/job-123/input.jsonl.out"
EXPECTED_URL = "https://s3.us-west-2.amazonaws.com/my-bucket/litellm-batch-outputs/job-123/input.jsonl.out"
def _litellm_params(self) -> dict:
return {
"aws_access_key_id": "AKIAEXAMPLE",
"aws_secret_access_key": "secret",
"aws_region_name": "us-west-2",
}
def test_transform_file_content_request_signs_s3_get(self, monkeypatch):
"""The request transform must produce the S3 object URL plus SigV4 GET headers."""
import hashlib
from litellm.llms.bedrock.files.transformation import (
S3_SIGNED_GET_HEADERS_PARAM,
BedrockFilesConfig,
)
monkeypatch.setenv("AWS_S3_BUCKET_NAME", "my-bucket")
litellm_params = self._litellm_params()
url, params = BedrockFilesConfig().transform_file_content_request(
file_content_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_GET_HEADERS_PARAM]
assert (
signed_headers["x-amz-content-sha256"] == hashlib.sha256(b"").hexdigest()
), "GET has no payload, so the content hash must be the empty-body hash"
authorization = signed_headers["Authorization"]
assert authorization.startswith("AWS4-HMAC-SHA256 Credential=AKIAEXAMPLE/")
assert "/us-west-2/s3/aws4_request" in authorization
assert "x-amz-content-sha256" in authorization
assert "X-Amz-Date" in signed_headers
def test_transform_file_content_request_decodes_unified_file_id(self, monkeypatch):
"""Base64 unified ids carrying llm_output_file_id must resolve to their S3 object."""
import base64
from litellm.llms.bedrock.files.transformation import 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("=")
)
url, _ = BedrockFilesConfig().transform_file_content_request(
file_content_request={"file_id": encoded_file_id},
optional_params={},
litellm_params=self._litellm_params(),
)
assert url == self.EXPECTED_URL
def test_transform_file_content_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_file_content_request(
file_content_request={
"file_id": "s3://other-bucket/litellm-batch-outputs/job/x.jsonl.out"
},
optional_params={},
litellm_params=self._litellm_params(),
)
def test_transform_file_content_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_file_content_request(
file_content_request={"file_id": "s3://my-bucket/private/x.jsonl"},
optional_params={},
litellm_params=self._litellm_params(),
)
def test_extract_s3_uri_rejects_non_managed_file_id(self):
"""A file id that is neither an s3:// URI nor a unified id must be rejected."""
from litellm.llms.bedrock.files.transformation import (
extract_s3_uri_from_file_id,
)
with pytest.raises(ValueError, match="managed LiteLLM S3 file id"):
extract_s3_uri_from_file_id("file-1234567890")
def test_transform_file_content_request_requires_configured_bucket(
self, monkeypatch
):
"""Without a server-configured bucket (env or snapshot), the request must fail
before any S3 call rather than guessing a bucket from the file id."""
from litellm.llms.bedrock.files.transformation import BedrockFilesConfig
monkeypatch.delenv("AWS_S3_BUCKET_NAME", raising=False)
with pytest.raises(ValueError, match="S3 bucket_name is required"):
BedrockFilesConfig().transform_file_content_request(
file_content_request={"file_id": self.S3_URI},
optional_params={},
litellm_params=self._litellm_params(),
)
def test_transform_file_content_request_requires_file_id(self, monkeypatch):
from litellm.llms.bedrock.files.transformation import BedrockFilesConfig
monkeypatch.setenv("AWS_S3_BUCKET_NAME", "my-bucket")
with pytest.raises(ValueError, match="file_id is required"):
BedrockFilesConfig().transform_file_content_request(
file_content_request={},
optional_params={},
litellm_params=self._litellm_params(),
)
def test_sign_request_without_botocore_raises_helpful_error(self, monkeypatch):
"""A missing botocore must surface an actionable 'install boto3' error
rather than a raw import failure."""
import sys
from litellm.llms.bedrock.files.transformation import BedrockFilesConfig
monkeypatch.setenv("AWS_S3_BUCKET_NAME", "my-bucket")
monkeypatch.setitem(sys.modules, "botocore.auth", None)
with pytest.raises(ImportError, match="boto3"):
BedrockFilesConfig().transform_file_content_request(
file_content_request={"file_id": self.S3_URI},
optional_params={},
litellm_params=self._litellm_params(),
)
def test_bucket_resolved_from_trusted_model_credentials(self, monkeypatch):
"""Per-model s3_bucket_name must be honored via the server-side credential snapshot."""
from types import MappingProxyType
from litellm.llms.bedrock.files.transformation import BedrockFilesConfig
monkeypatch.delenv("AWS_S3_BUCKET_NAME", raising=False)
litellm_params = self._litellm_params()
litellm_params["_litellm_internal_model_credentials"] = MappingProxyType(
{"s3_bucket_name": "my-bucket"}
)
url, _ = BedrockFilesConfig().transform_file_content_request(
file_content_request={"file_id": self.S3_URI},
optional_params={},
litellm_params=litellm_params,
)
assert url == self.EXPECTED_URL
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,
BedrockFilesConfig,
)
monkeypatch.setenv("AWS_S3_BUCKET_NAME", "my-bucket")
litellm_params = self._litellm_params()
litellm_params["s3_region_name"] = "eu-west-1"
url, _ = BedrockFilesConfig().transform_file_content_request(
file_content_request={"file_id": self.S3_URI},
optional_params={},
litellm_params=litellm_params,
)
assert url.startswith("https://s3.eu-west-1.amazonaws.com/")
authorization = litellm_params[S3_SIGNED_GET_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,
BedrockFilesConfig,
)
litellm_params = {
S3_SIGNED_GET_HEADERS_PARAM: {"Authorization": "AWS4-HMAC-SHA256 test"}
}
headers = BedrockFilesConfig().validate_environment(
headers={"x-custom": "kept"},
model="",
messages=[],
optional_params={},
litellm_params=litellm_params,
)
assert headers == {
"x-custom": "kept",
"Authorization": "AWS4-HMAC-SHA256 test",
}
assert S3_SIGNED_GET_HEADERS_PARAM not in litellm_params
def test_transform_file_content_response_wraps_binary_content(self):
import httpx
from litellm.llms.bedrock.files.transformation import BedrockFilesConfig
from litellm.types.llms.openai import HttpxBinaryResponseContent
raw_response = httpx.Response(
status_code=200,
content=b'{"recordId": "CALL0000001"}',
request=httpx.Request("GET", self.EXPECTED_URL),
)
result = BedrockFilesConfig().transform_file_content_response(
raw_response=raw_response,
logging_obj=MagicMock(),
litellm_params={},
)
assert isinstance(result, HttpxBinaryResponseContent)
assert result.response.content == b'{"recordId": "CALL0000001"}'
def test_transform_file_content_response_raises_on_s3_error(self):
import httpx
from litellm.llms.bedrock.common_utils import BedrockError
from litellm.llms.bedrock.files.transformation import BedrockFilesConfig
raw_response = httpx.Response(
status_code=403,
content=b"<Error><Code>AccessDenied</Code></Error>",
request=httpx.Request("GET", self.EXPECTED_URL),
)
with pytest.raises(BedrockError, match="AccessDenied"):
BedrockFilesConfig().transform_file_content_response(
raw_response=raw_response,
logging_obj=MagicMock(),
litellm_params={},
)
def test_file_content_end_to_end_sends_signed_get(self, monkeypatch):
"""litellm.file_content must issue a SigV4-signed GET and return the S3 object bytes."""
import httpx
import respx
import litellm
monkeypatch.setenv("AWS_S3_BUCKET_NAME", "my-bucket")
with respx.mock:
route = respx.get(self.EXPECTED_URL).mock(
return_value=httpx.Response(200, content=b'{"recordId": "x"}')
)
response = litellm.file_content(
file_id=self.S3_URI,
custom_llm_provider="bedrock",
**self._litellm_params(),
)
assert route.called
request = route.calls[0].request
assert request.headers["Authorization"].startswith("AWS4-HMAC-SHA256")
assert "x-amz-content-sha256" in request.headers
assert response.content == b'{"recordId": "x"}'
@pytest.mark.asyncio
async def test_afile_content_end_to_end_sends_signed_get(self, monkeypatch):
"""Async variant: litellm.afile_content over the same signed GET path."""
import httpx
import respx
import litellm
monkeypatch.setenv("AWS_S3_BUCKET_NAME", "my-bucket")
# respx can only intercept httpx transports
monkeypatch.setattr(litellm, "disable_aiohttp_transport", True)
litellm.in_memory_llm_clients_cache.flush_cache()
with respx.mock:
route = respx.get(self.EXPECTED_URL).mock(
return_value=httpx.Response(200, content=b'{"recordId": "x"}')
)
response = await litellm.afile_content(
file_id=self.S3_URI,
custom_llm_provider="bedrock",
**self._litellm_params(),
)
assert route.called
assert (
route.calls[0]
.request.headers["Authorization"]
.startswith("AWS4-HMAC-SHA256")
)
assert response.content == b'{"recordId": "x"}'

19
tests/test_litellm/proxy/openai_files_endpoint/test_files_endpoint.py
View file

@ -228,6 +228,25 @@ def test_invalid_purpose(mocker: MockerFixture, monkeypatch, llm_router: Router)
assert "Invalid purpose: my-bad-purpose" in response.json()["error"]["message"]
def test_get_file_content_rejects_raw_cloud_storage_uri(llm_router: Router):
"""A raw s3:// file id must be rejected on the proxy content endpoint.
Such an id is not a managed unified id, so it would otherwise skip the
owner/team access check and let a caller read another tenant's batch output
object by its key. Callers must use the managed unified file id.
"""
from urllib.parse import quote
s3_file_id = "s3://my-bucket/litellm-batch-outputs/job-123/input.jsonl.out"
response = client.get(
f"/v1/files/{quote(s3_file_id, safe='')}/content?provider=bedrock",
headers={"Authorization": "Bearer test-key"},
)
assert response.status_code == 400
assert "managed file id" in response.json()["error"]["message"].lower()
def test_mock_create_audio_file(mocker: MockerFixture, monkeypatch, llm_router: Router):
"""
Asserts 'create_file' is called with the correct arguments