From e84e5d03bdf9ea6fefd64110f45bbc939faca121 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Fri, 4 Sep 2026 17:52:52 -0700 Subject: [PATCH] fix(bedrock): list the configured output bucket for purpose=batch_output --- litellm/llms/bedrock/files/transformation.py | 38 ++++-- .../test_bedrock_files_transformation.py | 124 +++++++++++++++++- 2 files changed, 150 insertions(+), 12 deletions(-) diff --git a/litellm/llms/bedrock/files/transformation.py b/litellm/llms/bedrock/files/transformation.py index 41ea206a250..7d133b411a7 100644 --- a/litellm/llms/bedrock/files/transformation.py +++ b/litellm/llms/bedrock/files/transformation.py @@ -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( diff --git a/tests/test_litellm/llms/bedrock/files/test_bedrock_files_transformation.py b/tests/test_litellm/llms/bedrock/files/test_bedrock_files_transformation.py index 37195771e0d..ceac641bbf0 100644 --- a/tests/test_litellm/llms/bedrock/files/test_bedrock_files_transformation.py +++ b/tests/test_litellm/llms/bedrock/files/test_bedrock_files_transformation.py @@ -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""" + + my-output-bucket + litellm-batch-outputs/ + + litellm-batch-outputs/job-9/input.jsonl.out + 2026-09-04T08:00:00.000Z + 70 + +""" + OUTPUT_BUCKET_ID = "s3://my-output-bucket/litellm-batch-outputs/job-9/input.jsonl.out" LISTING = b""" my-bucket @@ -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]