From fd5a995f5d4447993c218cb4984abb1943f254b5 Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Thu, 8 Oct 2026 19:55:02 +0000 Subject: [PATCH] fix(bedrock): classify retrieved file purpose by its own bucket prefix Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/llms/bedrock/files/transformation.py | 34 +++++++++++-- .../test_bedrock_files_transformation.py | 51 +++++++++++++++++++ 2 files changed, 80 insertions(+), 5 deletions(-) diff --git a/litellm/llms/bedrock/files/transformation.py b/litellm/llms/bedrock/files/transformation.py index 00515150050..60a1d98d6a8 100644 --- a/litellm/llms/bedrock/files/transformation.py +++ b/litellm/llms/bedrock/files/transformation.py @@ -73,6 +73,7 @@ from ..common_utils import ( S3_SIGNED_REQUEST_HEADERS_PARAM: Final = "_s3_signed_request_headers" S3_RETRIEVE_FILE_ID_PARAM: Final = "_s3_retrieve_file_id" S3_RETRIEVE_FILE_KEY_PARAM: Final = "_s3_retrieve_file_key" +S3_RETRIEVE_FILE_RELATIVE_KEY_PARAM: Final = "_s3_retrieve_file_relative_key" LIST_FILES_PURPOSE_PARAM: Final = "_s3_list_files_purpose" @@ -335,6 +336,25 @@ def _resolve_managed_s3_object(file_id: str, litellm_params: Mapping[str, object raise _rejected_file_id(reason) from reason +def _relative_s3_object_key( + bucket_name: str, + object_key: str, + litellm_params: Mapping[str, object], +) -> str: + configured_bucket_prefixes: Final = tuple( + split_configured_cloud_bucket_name(configured_bucket_name) + for configured_bucket_name in get_configured_s3_bucket_names(litellm_params) + ) + matching_prefixes: Final = tuple( + configured_prefix + for configured_bucket, configured_prefix in configured_bucket_prefixes + if configured_bucket == bucket_name + and (not configured_prefix or object_key.startswith(f"{configured_prefix}/")) + ) + configured_prefix: Final = max(matching_prefixes, key=len, default="") + return object_key[len(configured_prefix) + 1 :] if configured_prefix else object_key + + _ANY_MANAGED_LISTING_PREFIX: Final = os.path.commonprefix(BEDROCK_MANAGED_S3_PREFIXES) _MANAGED_LISTING_PREFIX_BY_PURPOSE: Final = MappingProxyType( { @@ -1338,7 +1358,12 @@ class BedrockFilesConfig(BaseAWSLLM, BaseFilesConfig): litellm_params: MutableMapping[str, object], ) -> tuple[str, dict[str, str]]: """Prepare a ranged S3 GET for file retrieval.""" - _, object_key = _resolve_managed_s3_object(file_id=file_id, litellm_params=litellm_params) + bucket_name, object_key = _resolve_managed_s3_object(file_id=file_id, litellm_params=litellm_params) + relative_key: Final = _relative_s3_object_key( + bucket_name=bucket_name, + object_key=object_key, + litellm_params=litellm_params, + ) url, params = self._transform_s3_file_request( file_id=file_id, method="GET", @@ -1352,6 +1377,7 @@ class BedrockFilesConfig(BaseAWSLLM, BaseFilesConfig): litellm_params[S3_SIGNED_REQUEST_HEADERS_PARAM] = MappingProxyType({**signed_headers, "Range": "bytes=0-0"}) litellm_params[S3_RETRIEVE_FILE_ID_PARAM] = file_id litellm_params[S3_RETRIEVE_FILE_KEY_PARAM] = object_key + litellm_params[S3_RETRIEVE_FILE_RELATIVE_KEY_PARAM] = relative_key return url, params def transform_retrieve_file_response( @@ -1363,14 +1389,12 @@ class BedrockFilesConfig(BaseAWSLLM, BaseFilesConfig): """Build file metadata, accepting 416 only when S3 proves the object is empty.""" file_id: Final = litellm_params.get(S3_RETRIEVE_FILE_ID_PARAM) object_key: Final = litellm_params.get(S3_RETRIEVE_FILE_KEY_PARAM) - if not isinstance(file_id, str) or not isinstance(object_key, str): + relative_key: Final = litellm_params.get(S3_RETRIEVE_FILE_RELATIVE_KEY_PARAM) + if not isinstance(file_id, str) or not isinstance(object_key, str) or not isinstance(relative_key, str): raise ValueError("S3 retrieve response is missing request context") file_size: Final = _retrieved_s3_file_size(raw_response) - configured_bucket_name: Final = _listing_bucket_name(litellm_params, "batch_output") - _, configured_prefix = split_configured_cloud_bucket_name(configured_bucket_name) - relative_key: Final = object_key[len(configured_prefix) + 1 :] if configured_prefix else object_key last_modified: Final = raw_response.headers.get("Last-Modified", "") created_at: Final = int(parsedate_to_datetime(last_modified).timestamp()) if last_modified else 0 return OpenAIFileObject( diff --git a/tests/unit/llms/bedrock/files/test_bedrock_files_transformation.py b/tests/unit/llms/bedrock/files/test_bedrock_files_transformation.py index 1a919e7b6ad..6685deccf65 100644 --- a/tests/unit/llms/bedrock/files/test_bedrock_files_transformation.py +++ b/tests/unit/llms/bedrock/files/test_bedrock_files_transformation.py @@ -2276,6 +2276,57 @@ class TestBedrockFileContentTransformation: assert response.status == "processed" assert response.object == "file" + @pytest.mark.parametrize( + ("file_id", "purpose"), + [ + pytest.param( + "s3://out-bucket/outpfx/litellm-batch-outputs/job-123/x.jsonl.out", + "batch_output", + id="output-bucket", + ), + pytest.param( + "s3://in-bucket/pfx/litellm-bedrock-files/job-123/input.jsonl", + "batch", + id="input-bucket-upload", + ), + pytest.param( + "s3://in-bucket/pfx/litellm-batch-outputs/job-123/x.jsonl.out", + "batch_output", + id="input-bucket-output", + ), + ], + ) + def test_transform_retrieve_file_response_uses_the_retrieved_bucket_prefix( + self, file_id: str, purpose: str, monkeypatch: pytest.MonkeyPatch + ) -> None: + import httpx + + from litellm.llms.bedrock.files.transformation import BedrockFilesConfig + + monkeypatch.delenv("AWS_S3_BUCKET_NAME", raising=False) + monkeypatch.delenv("AWS_S3_OUTPUT_BUCKET_NAME", raising=False) + litellm_params: Final = _trusted_bucket_snapshot( + s3_bucket_name="in-bucket/pfx", + s3_output_bucket_name="out-bucket/outpfx", + ) + config: Final = BedrockFilesConfig() + config.transform_retrieve_file_request( + file_id=file_id, + optional_params={}, + litellm_params=litellm_params, + ) + response: Final = config.transform_retrieve_file_response( + raw_response=httpx.Response( + 206, + headers={"Content-Range": "bytes 0-0/1"}, + request=httpx.Request("GET", file_id), + ), + logging_obj=MagicMock(), + litellm_params=litellm_params, + ) + + assert response.purpose == purpose + def test_transform_retrieve_file_response_accepts_verified_empty_object(self, monkeypatch): import httpx