From b4919d9bd78386db9a65c9c381b45c0aecce3a24 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Wed, 9 Sep 2026 19:53:57 -0700 Subject: [PATCH] fix(bedrock): walk the output location on an unfiltered files list A list without a purpose covered the input bucket only, so a deployment with a separate s3_output_bucket_name never saw its batch outputs unless the caller passed purpose=batch_output. The listing now follows the input location to its last page and then walks the output location whenever it differs from the input one, in bucket or in prefix, so the unfiltered list matches what OpenAI returns --- litellm/llms/bedrock/files/transformation.py | 27 +++- .../test_bedrock_files_transformation.py | 127 ++++++++++++++++++ 2 files changed, 149 insertions(+), 5 deletions(-) diff --git a/litellm/llms/bedrock/files/transformation.py b/litellm/llms/bedrock/files/transformation.py index 29732f8afe2..6e2b0c12090 100644 --- a/litellm/llms/bedrock/files/transformation.py +++ b/litellm/llms/bedrock/files/transformation.py @@ -69,6 +69,8 @@ S3_SIGNED_REQUEST_HEADERS_PARAM: Final = "_s3_signed_request_headers" LIST_FILES_PURPOSE_PARAM: Final = "_s3_list_files_purpose" +LIST_FILES_LOCATION_PARAM: Final = "_s3_list_files_location" + class _S3DeleteContext(BaseModel): file_id: str = Field(min_length=1) @@ -322,6 +324,17 @@ def _requested_listing_purpose(litellm_params: Mapping[str, object]) -> str | No return requested_purpose if isinstance(requested_purpose, str) else None +def _walked_listing_purpose(litellm_params: Mapping[str, object]) -> str | None: + walked_purpose: Final = litellm_params.get(LIST_FILES_LOCATION_PARAM) + return walked_purpose if isinstance(walked_purpose, str) else _requested_listing_purpose(litellm_params) + + +def _output_location_still_unlisted(litellm_params: Mapping[str, object]) -> bool: + if _walked_listing_purpose(litellm_params) is not None: + return False + return _listing_bucket_name(litellm_params, "batch_output") != _listing_bucket_name(litellm_params, None) + + def _listing_bucket_name(litellm_params: Mapping[str, object], purpose: str | None) -> str: if purpose != "batch_output": return get_configured_s3_bucket_name(litellm_params) @@ -1342,6 +1355,7 @@ class BedrockFilesConfig(BaseAWSLLM, BaseFilesConfig): litellm_params: MutableMapping[str, object], ) -> tuple[str, dict[str, str]]: litellm_params[LIST_FILES_PURPOSE_PARAM] = purpose # rebind-ok: handed to the response transform + litellm_params[LIST_FILES_LOCATION_PARAM] = purpose # rebind-ok: names the location the next page walks return self._signed_listing_request(purpose, optional_params, litellm_params, continuation_token=None) def transform_list_files_next_request( @@ -1353,11 +1367,14 @@ class BedrockFilesConfig(BaseAWSLLM, BaseFilesConfig): if raw_response.status_code >= 400: return None continuation_token: Final = ET.fromstring(raw_response.content).findtext("{*}NextContinuationToken") - if not continuation_token: + if continuation_token: + return self._signed_listing_request( + _walked_listing_purpose(litellm_params), optional_params, litellm_params, continuation_token + ) + if not _output_location_still_unlisted(litellm_params): return None - return self._signed_listing_request( - _requested_listing_purpose(litellm_params), optional_params, litellm_params, continuation_token - ) + litellm_params[LIST_FILES_LOCATION_PARAM] = "batch_output" # rebind-ok: the input location is fully listed + return self._signed_listing_request("batch_output", optional_params, litellm_params, continuation_token=None) def _signed_listing_request( self, @@ -1399,7 +1416,7 @@ class BedrockFilesConfig(BaseAWSLLM, BaseFilesConfig): response=raw_response, ) purpose: Final = _requested_listing_purpose(litellm_params) - configured_bucket_name: Final = _listing_bucket_name(litellm_params, purpose) + configured_bucket_name: Final = _listing_bucket_name(litellm_params, _walked_listing_purpose(litellm_params)) allow_legacy_cloud_file_ids: Final = should_allow_legacy_cloud_file_ids(litellm_params) listing: Final = ET.fromstring(raw_response.content) bucket_name: Final = ( 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 84433ea79a9..c609455f3d8 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 @@ -3340,6 +3340,133 @@ class TestBedrockFileListTransformation: 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] + def test_file_list_without_purpose_also_walks_a_separate_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: + input_route = respx.get(self.BUCKET_URL, params__contains=self.MANAGED_QUERY).mock( + return_value=httpx.Response(200, content=self.LISTING) + ) + output_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", + **_trusted_bucket_snapshot(s3_bucket_name="my-bucket", s3_output_bucket_name="my-output-bucket"), + ) + + assert (input_route.call_count, output_route.call_count) == (1, 1) + output_request = output_route.calls[0].request + assert _sent_signature(output_request.headers) == _s3_signature_for( + "GET", str(output_request.url), output_request.headers + ) + assert [file.id for file in files] == [*self.BATCH_IDS, self.OUTPUT_ID, self.OUTPUT_BUCKET_ID] + + def test_file_list_without_purpose_walks_the_output_bucket_after_the_last_input_page(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: + respx.get(self.BUCKET_URL, params__contains={"continuation-token": self.CONTINUATION_TOKEN}).mock( + return_value=httpx.Response(200, content=self.LAST_PAGE) + ) + respx.get(self.BUCKET_URL, params__contains=self.MANAGED_QUERY).mock( + return_value=httpx.Response(200, content=self.FIRST_PAGE) + ) + 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", + **_trusted_bucket_snapshot(s3_bucket_name="my-bucket", s3_output_bucket_name="my-output-bucket"), + ) + requested_urls = [str(call.request.url) for call in respx.calls] + + assert requested_urls == [ + f"{self.BUCKET_URL}?list-type=2&prefix=litellm-b", + f"{self.BUCKET_URL}?list-type=2&prefix=litellm-b" + "&continuation-token=1ueGcxLPRx1Tr%2FXYExHnhbYLgveDs2J%2Fwm36Hy4vbOwM%3D", + f"{self.OUTPUT_BUCKET_URL}?list-type=2&prefix=litellm-batch-outputs%2F", + ] + assert [file.id for file in files] == [*self.PAGED_IDS, self.OUTPUT_BUCKET_ID] + + @pytest.mark.parametrize( + ("purpose", "bucket_snapshot"), + [ + pytest.param(None, {"s3_bucket_name": "my-bucket"}, id="outputs-share-the-input-bucket"), + pytest.param( + "batch", + {"s3_bucket_name": "my-bucket", "s3_output_bucket_name": "my-output-bucket"}, + id="input-purpose-requested", + ), + ], + ) + def test_file_list_leaves_the_output_bucket_alone_unless_an_unfiltered_list_needs_it( + self, monkeypatch, purpose, bucket_snapshot + ): + 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: + input_route = respx.get(self.BUCKET_URL).mock(return_value=httpx.Response(200, content=self.LISTING)) + output_route = respx.get(self.OUTPUT_BUCKET_URL).mock( + return_value=httpx.Response(200, content=self.OUTPUT_BUCKET_LISTING) + ) + + files = litellm.file_list( + custom_llm_provider="bedrock", purpose=purpose, **_trusted_bucket_snapshot(**bucket_snapshot) + ) + + assert (input_route.call_count, output_route.call_count) == (1, 0) + assert [file.id for file in files] == [*self.BATCH_IDS, *(() if purpose else (self.OUTPUT_ID,))] + + def test_transform_list_files_next_request_walks_an_output_prefix_inside_the_input_bucket(self, monkeypatch): + import httpx + + 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-bucket/out") + config = BedrockFilesConfig() + config.transform_list_files_request(purpose=None, optional_params={}, litellm_params=litellm_params) + litellm_params.pop(S3_SIGNED_REQUEST_HEADERS_PARAM) + + output_request = config.transform_list_files_next_request( + raw_response=httpx.Response(200, content=self.LISTING), optional_params={}, litellm_params=litellm_params + ) + after_output_request = config.transform_list_files_next_request( + raw_response=httpx.Response(200, content=self.LISTING), optional_params={}, litellm_params=litellm_params + ) + + assert output_request == (self.BUCKET_URL, {"list-type": "2", "prefix": "out/litellm-batch-outputs/"}) + signed_headers = litellm_params[S3_SIGNED_REQUEST_HEADERS_PARAM] + assert _sent_signature(signed_headers) == _s3_signature_for( + "GET", f"{self.BUCKET_URL}?list-type=2&prefix=out%2Flitellm-batch-outputs%2F", signed_headers + ) + assert after_output_request is None + def test_transform_list_files_next_request_signs_the_continuation_page(self, monkeypatch): import httpx