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

feat(bedrock): follow S3 continuation tokens when listing managed files

Browse source
This commit is contained in:
mateo-berri 2026-09-04 18:30:52 -07:00
parent e84e5d03bd
commit 6a6080a152
4 changed files with 276 additions and 16 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

11
litellm/llms/base_llm/files/transformation.py
View file

@ -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,

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

@ -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)

90
litellm/llms/custom_httpx/llm_http_handler.py
View file

@ -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,

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

@ -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"""<?xml version="1.0" encoding="UTF-8"?>
<ListBucketResult xmlns="http://s3.amazonaws.com/doc/2006-03-01/">
<Name>my-bucket</Name>
<Prefix>litellm-bedrock-files</Prefix>
<KeyCount>1</KeyCount>
<MaxKeys>1</MaxKeys>
<IsTruncated>true</IsTruncated>
<NextContinuationToken>1ueGcxLPRx1Tr/XYExHnhbYLgveDs2J/wm36Hy4vbOwM=</NextContinuationToken>
<Contents>
<Key>litellm-bedrock-files/job-1/input.jsonl</Key>
<LastModified>2026-09-01T10:00:00.000Z</LastModified>
<Size>10</Size>
</Contents>
</ListBucketResult>"""
LAST_PAGE = b"""<?xml version="1.0" encoding="UTF-8"?>
<ListBucketResult xmlns="http://s3.amazonaws.com/doc/2006-03-01/">
<Name>my-bucket</Name>
<Prefix>litellm-bedrock-files</Prefix>
<KeyCount>1</KeyCount>
<MaxKeys>1</MaxKeys>
<IsTruncated>false</IsTruncated>
<ContinuationToken>1ueGcxLPRx1Tr/XYExHnhbYLgveDs2J/wm36Hy4vbOwM=</ContinuationToken>
<Contents>
<Key>litellm-bedrock-files/job-2/input.jsonl</Key>
<LastModified>2026-09-02T10:00:00.000Z</LastModified>
<Size>20</Size>
</Contents>
</ListBucketResult>"""
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"<Error><Code>AccessDenied</Code></Error>", 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)