From 6a6080a1529038fa38d692dd482d25fac5a2a8e3 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Fri, 4 Sep 2026 18:30:52 -0700 Subject: [PATCH] feat(bedrock): follow S3 continuation tokens when listing managed files --- litellm/llms/base_llm/files/transformation.py | 11 +- litellm/llms/bedrock/files/transformation.py | 43 ++++- litellm/llms/custom_httpx/llm_http_handler.py | 90 ++++++++++- .../test_bedrock_files_transformation.py | 148 ++++++++++++++++++ 4 files changed, 276 insertions(+), 16 deletions(-) diff --git a/litellm/llms/base_llm/files/transformation.py b/litellm/llms/base_llm/files/transformation.py index 7a7088c2fb5..3f8fec354b7 100644 --- a/litellm/llms/base_llm/files/transformation.py +++ b/litellm/llms/base_llm/files/transformation.py @@ -1,5 +1,5 @@ from abc import ABC, abstractmethod -from collections.abc import Iterator +from collections.abc import Iterator, Mapping from typing import TYPE_CHECKING, Any, Union import httpx @@ -160,6 +160,15 @@ class BaseFilesConfig(BaseConfig): ) -> tuple[str, dict]: """Transform file list request into provider-specific format.""" + def transform_list_files_next_request( + self, + raw_response: httpx.Response, + optional_params: Mapping[str, object], + litellm_params: dict, # mutable-ok: carries provider stashes from the request transform to the response one + ) -> tuple[str, dict[str, str]] | None: + """Request for the page after `raw_response`, or None once the listing is complete.""" + return None + @abstractmethod def transform_list_files_response( self, diff --git a/litellm/llms/bedrock/files/transformation.py b/litellm/llms/bedrock/files/transformation.py index 7d133b411a7..1266397636a 100644 --- a/litellm/llms/bedrock/files/transformation.py +++ b/litellm/llms/bedrock/files/transformation.py @@ -281,6 +281,11 @@ def _managed_listing_prefix(configured_prefix: str, purpose: str | None) -> str: return f"{configured_prefix}/{managed_prefix}" if configured_prefix else managed_prefix +def _requested_listing_purpose(litellm_params: Mapping[str, object]) -> str | None: + requested_purpose: Final = litellm_params.get(LIST_FILES_PURPOSE_PARAM) + return requested_purpose if isinstance(requested_purpose, str) else None + + 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": @@ -1310,16 +1315,42 @@ class BedrockFilesConfig(BaseAWSLLM, BaseFilesConfig): purpose: str | None, optional_params: Mapping[str, object], litellm_params: MutableMapping[str, object], + ) -> tuple[str, dict[str, str]]: + litellm_params[LIST_FILES_PURPOSE_PARAM] = purpose # rebind-ok: handed to the response transform + return self._signed_listing_request(purpose, optional_params, litellm_params, continuation_token=None) + + def transform_list_files_next_request( + self, + raw_response: httpx.Response, + optional_params: Mapping[str, object], + litellm_params: MutableMapping[str, object], + ) -> tuple[str, dict[str, str]] | None: + if raw_response.status_code >= 400: + return None + continuation_token: Final = ET.fromstring(raw_response.content).findtext("{*}NextContinuationToken") + if not continuation_token: + return None + return self._signed_listing_request( + _requested_listing_purpose(litellm_params), optional_params, litellm_params, continuation_token + ) + + def _signed_listing_request( + self, + purpose: str | None, + optional_params: Mapping[str, object], + litellm_params: MutableMapping[str, object], + continuation_token: str | None, ) -> tuple[str, dict[str, str]]: bucket_name, configured_prefix = split_configured_cloud_bucket_name( _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, purpose), - } + listing_query: Final = (("list-type", "2"), ("prefix", _managed_listing_prefix(configured_prefix, purpose))) + continuation_query: Final = (("continuation-token", continuation_token),) if continuation_token else () + query: Final[dict[str, str]] = dict( # mutable-ok: the base files contract returns the query as a dict + listing_query + continuation_query + ) signed_headers: Final = self._sign_s3_empty_body_request( method="GET", api_base=f"{url}?{urlencode(query, quote_via=quote, safe='')}", @@ -1327,7 +1358,6 @@ class BedrockFilesConfig(BaseAWSLLM, BaseFilesConfig): request_params=target.request_params, ) litellm_params[S3_SIGNED_REQUEST_HEADERS_PARAM] = signed_headers # rebind-ok: handed to validate_environment - litellm_params[LIST_FILES_PURPOSE_PARAM] = purpose # rebind-ok: handed to the response transform return url, query def transform_list_files_response( @@ -1342,8 +1372,7 @@ class BedrockFilesConfig(BaseAWSLLM, BaseFilesConfig): message=raw_response.text, headers=raw_response.headers, ) - requested_purpose: Final = litellm_params.get(LIST_FILES_PURPOSE_PARAM) - purpose: Final = requested_purpose if isinstance(requested_purpose, str) else None + purpose: Final = _requested_listing_purpose(litellm_params) 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) diff --git a/litellm/llms/custom_httpx/llm_http_handler.py b/litellm/llms/custom_httpx/llm_http_handler.py index f281c249c72..e5389d0e0b7 100644 --- a/litellm/llms/custom_httpx/llm_http_handler.py +++ b/litellm/llms/custom_httpx/llm_http_handler.py @@ -4928,11 +4928,11 @@ class BaseLLMHTTPHandler: except Exception as e: raise self._handle_error(e=e, provider_config=provider_config) - return provider_config.transform_list_files_response( - raw_response=response, - logging_obj=logging_obj, - litellm_params=litellm_params, + pages: Final = ( + response, + *self._following_list_files_pages(response, provider_config, litellm_params, headers, sync_httpx_client), ) + return self._listed_files_across_pages(pages, provider_config, logging_obj, litellm_params) async def async_list_files( self, @@ -4984,11 +4984,85 @@ class BaseLLMHTTPHandler: except Exception as e: raise self._handle_error(e=e, provider_config=provider_config) - return provider_config.transform_list_files_response( - raw_response=response, - logging_obj=logging_obj, - litellm_params=litellm_params, + following_pages: Final = self._following_async_list_files_pages( + response, provider_config, litellm_params, headers, async_httpx_client ) + pages: Final = ( + response, + *[page async for page in following_pages], # mutable-ok: an async comprehension is spelled as a list + ) + return self._listed_files_across_pages(pages, provider_config, logging_obj, litellm_params) + + def _listed_files_across_pages( + self, + pages: Sequence[httpx.Response], + provider_config: BaseFilesConfig, + logging_obj: LiteLLMLoggingObj, + litellm_params: dict, # mutable-ok: handed to the files contract, which types it as a dict + ) -> list[OpenAIFileObject]: + return [ # mutable-ok: the base files contract returns a list + listed_file + for page in pages + for listed_file in provider_config.transform_list_files_response( + raw_response=page, + logging_obj=logging_obj, + litellm_params=litellm_params, + ) + ] + + def _following_list_files_pages( + self, + page: httpx.Response, + provider_config: BaseFilesConfig, + litellm_params: dict, # mutable-ok: handed to the files contract, which types it as a dict + headers: dict, # mutable-ok: handed to validate_environment, which types it as a dict + client: HTTPHandler, + ) -> Iterator[httpx.Response]: + latest_page = page # rebind-ok: advances one page per loop turn + while next_request := provider_config.transform_list_files_next_request( + raw_response=latest_page, optional_params={}, litellm_params=litellm_params + ): + url, params = next_request + next_headers = provider_config.validate_environment( + api_key=litellm_params.get("api_key"), + headers=headers, + model="", + messages=[], + optional_params={}, + litellm_params=litellm_params, + ) + try: + latest_page = client.get(url=url, headers=next_headers, params=params) + except Exception as e: # noqa: BLE001 # _handle_error maps every failure kind, like the first page's fetch + raise self._handle_error(e=e, provider_config=provider_config) + yield latest_page + + async def _following_async_list_files_pages( + self, + page: httpx.Response, + provider_config: BaseFilesConfig, + litellm_params: dict, # mutable-ok: handed to the files contract, which types it as a dict + headers: dict, # mutable-ok: handed to validate_environment, which types it as a dict + client: AsyncHTTPHandler, + ) -> AsyncIterator[httpx.Response]: + latest_page = page # rebind-ok: advances one page per loop turn + while next_request := provider_config.transform_list_files_next_request( + raw_response=latest_page, optional_params={}, litellm_params=litellm_params + ): + url, params = next_request + next_headers = provider_config.validate_environment( + api_key=litellm_params.get("api_key"), + headers=headers, + model="", + messages=[], + optional_params={}, + litellm_params=litellm_params, + ) + try: + latest_page = await client.get(url=url, headers=next_headers, params=params) + except Exception as e: # noqa: BLE001 # _handle_error maps every failure kind, like the first page's fetch + raise self._handle_error(e=e, provider_config=provider_config) + yield latest_page def retrieve_file_content( self, 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 ceac641bbf0..8e564280bf7 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 @@ -2780,6 +2780,39 @@ class TestBedrockFileListTransformation: "s3://my-bucket/litellm-bedrock-files/job-123/input.jsonl", ) OUTPUT_ID = "s3://my-bucket/litellm-batch-outputs/job-123/input.jsonl.out" + CONTINUATION_TOKEN = "1ueGcxLPRx1Tr/XYExHnhbYLgveDs2J/wm36Hy4vbOwM=" + FIRST_PAGE = b""" + + my-bucket + litellm-bedrock-files + 1 + 1 + true + 1ueGcxLPRx1Tr/XYExHnhbYLgveDs2J/wm36Hy4vbOwM= + + litellm-bedrock-files/job-1/input.jsonl + 2026-09-01T10:00:00.000Z + 10 + +""" + LAST_PAGE = b""" + + my-bucket + litellm-bedrock-files + 1 + 1 + false + 1ueGcxLPRx1Tr/XYExHnhbYLgveDs2J/wm36Hy4vbOwM= + + litellm-bedrock-files/job-2/input.jsonl + 2026-09-02T10:00:00.000Z + 20 + +""" + PAGED_IDS = ( + "s3://my-bucket/litellm-bedrock-files/job-1/input.jsonl", + "s3://my-bucket/litellm-bedrock-files/job-2/input.jsonl", + ) def test_transform_list_files_request_signs_managed_prefix_listing(self, monkeypatch): from litellm.llms.bedrock.files.transformation import ( @@ -3075,3 +3108,118 @@ 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_BUCKET_ID] + + def test_transform_list_files_next_request_signs_the_continuation_page(self, monkeypatch): + import httpx + + from litellm.llms.bedrock.files.transformation import ( + S3_SIGNED_REQUEST_HEADERS_PARAM, + BedrockFilesConfig, + ) + + monkeypatch.setenv("AWS_S3_BUCKET_NAME", "my-bucket") + litellm_params = _bedrock_s3_params() + config = BedrockFilesConfig() + config.transform_list_files_request(purpose="batch", optional_params={}, litellm_params=litellm_params) + first_signature = _sent_signature(litellm_params.pop(S3_SIGNED_REQUEST_HEADERS_PARAM)) + + next_request = config.transform_list_files_next_request( + raw_response=httpx.Response(200, content=self.FIRST_PAGE), + optional_params={}, + litellm_params=litellm_params, + ) + + assert next_request == (self.BUCKET_URL, {**self.BATCH_QUERY, "continuation-token": self.CONTINUATION_TOKEN}) + signed_headers = litellm_params[S3_SIGNED_REQUEST_HEADERS_PARAM] + signed_url = ( + f"{self.BUCKET_URL}?list-type=2&prefix=litellm-bedrock-files" + "&continuation-token=1ueGcxLPRx1Tr%2FXYExHnhbYLgveDs2J%2Fwm36Hy4vbOwM%3D" + ) + assert _sent_signature(signed_headers) == _s3_signature_for("GET", signed_url, signed_headers) + assert _sent_signature(signed_headers) != first_signature + + @pytest.mark.parametrize( + ("status_code", "content"), + [ + pytest.param(200, LAST_PAGE, id="last-page"), + pytest.param(403, b"AccessDenied", id="error-page"), + ], + ) + def test_transform_list_files_next_request_stops_after_the_last_page(self, monkeypatch, status_code, content): + import httpx + + from litellm.llms.bedrock.files.transformation import ( + S3_SIGNED_REQUEST_HEADERS_PARAM, + BedrockFilesConfig, + ) + + monkeypatch.setenv("AWS_S3_BUCKET_NAME", "my-bucket") + litellm_params = _bedrock_s3_params() + + next_request = BedrockFilesConfig().transform_list_files_next_request( + raw_response=httpx.Response(status_code, content=content), + optional_params={}, + litellm_params=litellm_params, + ) + + assert next_request is None + assert S3_SIGNED_REQUEST_HEADERS_PARAM not in litellm_params + + def _mock_paged_listing(self, respx_module): + import httpx + + last_page = respx_module.get( + self.BUCKET_URL, params__contains={"continuation-token": self.CONTINUATION_TOKEN} + ).mock(return_value=httpx.Response(200, content=self.LAST_PAGE)) + first_page = respx_module.get(self.BUCKET_URL, params__contains=self.BATCH_QUERY).mock( + return_value=httpx.Response(200, content=self.FIRST_PAGE) + ) + return first_page, last_page + + def _assert_paged_listing(self, first_page, last_page, files): + assert (first_page.call_count, last_page.call_count) == (1, 1) + assert "continuation-token" not in str(first_page.calls[0].request.url) + last_request = last_page.calls[0].request + assert _sent_signature(last_request.headers) == _s3_signature_for( + "GET", str(last_request.url), last_request.headers + ) + assert [file.id for file in files] == list(self.PAGED_IDS) + + def test_file_list_follows_continuation_tokens_across_pages(self, monkeypatch): + import respx + + import litellm + + monkeypatch.delenv("AWS_S3_BUCKET_NAME", raising=False) + monkeypatch.delenv("AWS_S3_OUTPUT_BUCKET_NAME", raising=False) + + with respx.mock: + first_page, last_page = self._mock_paged_listing(respx) + files = litellm.file_list( + custom_llm_provider="bedrock", + purpose="batch", + **_trusted_bucket_snapshot(s3_bucket_name="my-bucket"), + ) + + self._assert_paged_listing(first_page, last_page, files) + + @pytest.mark.asyncio + async def test_afile_list_follows_continuation_tokens_across_pages(self, monkeypatch): + import respx + + import litellm + + monkeypatch.delenv("AWS_S3_BUCKET_NAME", raising=False) + monkeypatch.delenv("AWS_S3_OUTPUT_BUCKET_NAME", raising=False) + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + litellm.in_memory_llm_clients_cache.flush_cache() + + with respx.mock: + first_page, last_page = self._mock_paged_listing(respx) + files = await litellm.afile_list( + custom_llm_provider="bedrock", + purpose="batch", + **_trusted_bucket_snapshot(s3_bucket_name="my-bucket"), + ) + + self._assert_paged_listing(first_page, last_page, files)