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