scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-14 23:21:35 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 102

fix(bedrock): list the configured output bucket for purpose=batch_output

Browse source
This commit is contained in:
mateo-berri 2026-09-04 17:52:52 -07:00
parent fc978aec21
commit e84e5d03bd
2 changed files with 150 additions and 12 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

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

@ -263,9 +263,30 @@ 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
_ANY_MANAGED_LISTING_PREFIX: Final = os.path.commonprefix(BEDROCK_MANAGED_S3_PREFIXES)
_MANAGED_LISTING_PREFIX_BY_PURPOSE: Final = MappingProxyType(
{
"batch": os.path.commonprefix((BEDROCK_MANAGED_S3_BATCH_PREFIX, BEDROCK_MANAGED_S3_UPLOAD_PREFIX)),
"batch_output": BEDROCK_MANAGED_S3_OUTPUT_PREFIX,
}
)
def _managed_listing_prefix(configured_prefix: str, purpose: str | None) -> str:
managed_prefix: Final = (
_MANAGED_LISTING_PREFIX_BY_PURPOSE.get(purpose, _ANY_MANAGED_LISTING_PREFIX)
if purpose
else _ANY_MANAGED_LISTING_PREFIX
)
return f"{configured_prefix}/{managed_prefix}" if configured_prefix else managed_prefix
def _listing_bucket_name(litellm_params: Mapping[str, object], purpose: str | None) -> str:
input_bucket_name: Final = get_configured_s3_bucket_name(litellm_params)
if purpose != "batch_output":
return input_bucket_name
trusted: Final = _trusted_s3_model_credentials(litellm_params)
return trusted.s3_output_bucket_name or os.getenv("AWS_S3_OUTPUT_BUCKET_NAME") or input_bucket_name
def _listed_object_created_at(entry: ET.Element) -> int:
@ -1291,13 +1312,13 @@ class BedrockFilesConfig(BaseAWSLLM, BaseFilesConfig):
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)
_listing_bucket_name(litellm_params, purpose)
)
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),
"prefix": _managed_listing_prefix(configured_prefix, purpose),
}
signed_headers: Final = self._sign_s3_empty_body_request(
method="GET",
@ -1321,9 +1342,10 @@ class BedrockFilesConfig(BaseAWSLLM, BaseFilesConfig):
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)
purpose: Final = requested_purpose if isinstance(requested_purpose, str) else None
configured_bucket_name: Final = _listing_bucket_name(litellm_params, purpose)
allow_legacy_cloud_file_ids: Final = should_allow_legacy_cloud_file_ids(litellm_params)
listing: Final = ET.fromstring(raw_response.content)
bucket_name: Final = (
listing.findtext("{*}Name") or split_configured_cloud_bucket_name(configured_bucket_name)[0]
@ -1335,7 +1357,7 @@ class BedrockFilesConfig(BaseAWSLLM, BaseFilesConfig):
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)
if listed_file is not None and (purpose is None or listed_file.purpose == purpose)
]
def transform_file_content_request(

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

@ -2734,6 +2734,20 @@ class TestBedrockFileListTransformation:
BUCKET_URL = "https://s3.us-west-2.amazonaws.com/my-bucket/"
MANAGED_QUERY = {"list-type": "2", "prefix": "litellm-b"}
BATCH_QUERY = {"list-type": "2", "prefix": "litellm-bedrock-files"}
OUTPUT_QUERY = {"list-type": "2", "prefix": "litellm-batch-outputs/"}
OUTPUT_BUCKET_URL = "https://s3.us-west-2.amazonaws.com/my-output-bucket/"
OUTPUT_BUCKET_LISTING = b"""<?xml version="1.0" encoding="UTF-8"?>
<ListBucketResult xmlns="http://s3.amazonaws.com/doc/2006-03-01/">
<Name>my-output-bucket</Name>
<Prefix>litellm-batch-outputs/</Prefix>
<Contents>
<Key>litellm-batch-outputs/job-9/input.jsonl.out</Key>
<LastModified>2026-09-04T08:00:00.000Z</LastModified>
<Size>70</Size>
</Contents>
</ListBucketResult>"""
OUTPUT_BUCKET_ID = "s3://my-output-bucket/litellm-batch-outputs/job-9/input.jsonl.out"
LISTING = b"""<?xml version="1.0" encoding="UTF-8"?>
<ListBucketResult xmlns="http://s3.amazonaws.com/doc/2006-03-01/">
<Name>my-bucket</Name>
@ -2784,10 +2798,10 @@ class TestBedrockFileListTransformation:
)
assert url == self.BUCKET_URL
assert params == self.MANAGED_QUERY
assert params == self.BATCH_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
"GET", f"{url}?list-type=2&prefix=litellm-bedrock-files", signed_headers
)
assert litellm_params[LIST_FILES_PURPOSE_PARAM] == "batch"
@ -2902,7 +2916,7 @@ class TestBedrockFileListTransformation:
monkeypatch.setenv("AWS_S3_BUCKET_NAME", "my-bucket")
with respx.mock:
route = respx.get(self.BUCKET_URL, params__contains=self.MANAGED_QUERY).mock(
route = respx.get(self.BUCKET_URL, params__contains=self.BATCH_QUERY).mock(
return_value=httpx.Response(200, content=self.LISTING)
)
@ -2947,7 +2961,7 @@ class TestBedrockFileListTransformation:
litellm.in_memory_llm_clients_cache.flush_cache()
with respx.mock:
route = respx.get(self.BUCKET_URL, params__contains=self.MANAGED_QUERY).mock(
route = respx.get(self.BUCKET_URL, params__contains=self.OUTPUT_QUERY).mock(
return_value=httpx.Response(200, content=self.LISTING)
)
@ -2959,3 +2973,105 @@ class TestBedrockFileListTransformation:
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]
def test_transform_list_files_request_narrows_prefix_to_requested_purpose(self, monkeypatch):
from litellm.llms.bedrock.files.transformation import BedrockFilesConfig
monkeypatch.setenv("AWS_S3_BUCKET_NAME", "my-bucket")
monkeypatch.delenv("AWS_S3_OUTPUT_BUCKET_NAME", raising=False)
batch_url, batch_params = BedrockFilesConfig().transform_list_files_request(
purpose="batch", optional_params={}, litellm_params=_bedrock_s3_params()
)
output_url, output_params = BedrockFilesConfig().transform_list_files_request(
purpose="batch_output", optional_params={}, litellm_params=_bedrock_s3_params()
)
assert (batch_url, batch_params) == (self.BUCKET_URL, self.BATCH_QUERY)
assert (output_url, output_params) == (self.BUCKET_URL, self.OUTPUT_QUERY)
def test_transform_list_files_request_lists_configured_output_bucket_for_batch_output(self, monkeypatch):
from litellm.llms.bedrock.files.transformation import (
S3_SIGNED_REQUEST_HEADERS_PARAM,
BedrockFilesConfig,
)
monkeypatch.delenv("AWS_S3_BUCKET_NAME", raising=False)
monkeypatch.delenv("AWS_S3_OUTPUT_BUCKET_NAME", raising=False)
litellm_params = _trusted_bucket_snapshot(
s3_bucket_name="my-bucket", s3_output_bucket_name="my-output-bucket/team-a"
)
url, params = BedrockFilesConfig().transform_list_files_request(
purpose="batch_output", optional_params={}, litellm_params=litellm_params
)
signed_headers = litellm_params[S3_SIGNED_REQUEST_HEADERS_PARAM]
input_url, input_params = BedrockFilesConfig().transform_list_files_request(
purpose="batch", optional_params={}, litellm_params=dict(litellm_params)
)
assert url == self.OUTPUT_BUCKET_URL
assert params == {"list-type": "2", "prefix": "team-a/litellm-batch-outputs/"}
assert _sent_signature(signed_headers) == _s3_signature_for(
"GET", f"{url}?list-type=2&prefix=team-a%2Flitellm-batch-outputs%2F", signed_headers
)
assert (input_url, input_params) == (self.BUCKET_URL, self.BATCH_QUERY)
def test_transform_list_files_request_reads_output_bucket_from_env(self, monkeypatch):
from litellm.llms.bedrock.files.transformation import BedrockFilesConfig
monkeypatch.setenv("AWS_S3_BUCKET_NAME", "my-bucket")
monkeypatch.setenv("AWS_S3_OUTPUT_BUCKET_NAME", "my-output-bucket")
url, params = BedrockFilesConfig().transform_list_files_request(
purpose="batch_output", optional_params={}, litellm_params=_bedrock_s3_params()
)
assert (url, params) == (self.OUTPUT_BUCKET_URL, self.OUTPUT_QUERY)
def test_transform_list_files_response_accepts_output_bucket_objects(self, monkeypatch):
import httpx
from litellm.llms.bedrock.files.transformation import (
LIST_FILES_PURPOSE_PARAM,
BedrockFilesConfig,
)
monkeypatch.delenv("AWS_S3_BUCKET_NAME", raising=False)
monkeypatch.delenv("AWS_S3_OUTPUT_BUCKET_NAME", raising=False)
files = BedrockFilesConfig().transform_list_files_response(
raw_response=httpx.Response(200, content=self.OUTPUT_BUCKET_LISTING),
logging_obj=MagicMock(),
litellm_params={
**_trusted_bucket_snapshot(s3_bucket_name="my-bucket", s3_output_bucket_name="my-output-bucket"),
LIST_FILES_PURPOSE_PARAM: "batch_output",
},
)
assert [(file.id, file.purpose, file.bytes) for file in files] == [(self.OUTPUT_BUCKET_ID, "batch_output", 70)]
def test_file_list_batch_output_end_to_end_lists_output_bucket(self, monkeypatch):
import httpx
import respx
import litellm
monkeypatch.delenv("AWS_S3_BUCKET_NAME", raising=False)
monkeypatch.delenv("AWS_S3_OUTPUT_BUCKET_NAME", raising=False)
with respx.mock:
route = respx.get(self.OUTPUT_BUCKET_URL, params__contains=self.OUTPUT_QUERY).mock(
return_value=httpx.Response(200, content=self.OUTPUT_BUCKET_LISTING)
)
files = litellm.file_list(
custom_llm_provider="bedrock",
purpose="batch_output",
**_trusted_bucket_snapshot(s3_bucket_name="my-bucket", s3_output_bucket_name="my-output-bucket"),
)
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_BUCKET_ID]