From 20fcd187b48594cef318f2bad29c02c3833948e0 Mon Sep 17 00:00:00 2001 From: Sameer Kankute Date: Tue, 5 May 2026 13:28:24 +0530 Subject: [PATCH] feat(batches): request-level Bedrock batch S3 bucket overrides via kwargs - Proxy: attach _litellm_batch_s3_bucket_overrides from POST body for s3_bucket_name / s3_output_bucket_name - create_batch: merge overrides into litellm_params after GenericLiteLLMParams (proxy-safe, provider-agnostic) - GenericLiteLLMParams: s3_output_bucket_name - TypedDict LiteLLMBatchCreateRequest documents optional bucket fields Tests: tests/test_litellm/proxy/test_batch_s3_bucket_request_overrides.py Co-authored-by: Cursor --- litellm/batches/main.py | 6 ++ litellm/proxy/batches_endpoints/endpoints.py | 6 ++ litellm/types/llms/openai.py | 3 + litellm/types/router.py | 1 + .../test_batch_s3_bucket_request_overrides.py | 91 +++++++++++++++++++ 5 files changed, 107 insertions(+) create mode 100644 tests/test_litellm/proxy/test_batch_s3_bucket_request_overrides.py diff --git a/litellm/batches/main.py b/litellm/batches/main.py index 259439d4d09..f4e3bc48198 100644 --- a/litellm/batches/main.py +++ b/litellm/batches/main.py @@ -194,6 +194,12 @@ def create_batch( # noqa: PLR0915 _is_async = kwargs.pop("acreate_batch", False) is True litellm_params = dict(GenericLiteLLMParams(**kwargs)) + batch_s3_bucket_overrides = kwargs.get("_litellm_batch_s3_bucket_overrides") + if isinstance(batch_s3_bucket_overrides, dict): + for key in ("s3_bucket_name", "s3_output_bucket_name"): + value = batch_s3_bucket_overrides.get(key) + if isinstance(value, str) and value.strip(): + litellm_params[key] = value.strip() litellm_logging_obj: LiteLLMLoggingObj = cast( LiteLLMLoggingObj, kwargs.get("litellm_logging_obj", None) ) diff --git a/litellm/proxy/batches_endpoints/endpoints.py b/litellm/proxy/batches_endpoints/endpoints.py index 935b96a0e39..85b225cf65a 100644 --- a/litellm/proxy/batches_endpoints/endpoints.py +++ b/litellm/proxy/batches_endpoints/endpoints.py @@ -121,6 +121,12 @@ async def create_batch( # noqa: PLR0915 or "openai" ) _create_batch_data = LiteLLMBatchCreateRequest(**data) + _create_batch_data["_litellm_batch_s3_bucket_overrides"] = { + key: value.strip() + for key in ("s3_bucket_name", "s3_output_bucket_name") + for value in [_create_batch_data.get(key)] + if isinstance(value, str) and value.strip() + } # Apply team-level batch output expiry enforcement team_metadata = user_api_key_dict.team_metadata or {} diff --git a/litellm/types/llms/openai.py b/litellm/types/llms/openai.py index 986ec39f3bb..ed0d14d2d2c 100644 --- a/litellm/types/llms/openai.py +++ b/litellm/types/llms/openai.py @@ -434,6 +434,9 @@ class CreateBatchRequest(TypedDict, total=False): class LiteLLMBatchCreateRequest(CreateBatchRequest, total=False): model: str + # Bedrock batch / proxy: optional request-level buckets (see _litellm_batch_s3_bucket_overrides) + s3_bucket_name: Optional[str] + s3_output_bucket_name: Optional[str] class RetrieveBatchRequest(TypedDict, total=False): diff --git a/litellm/types/router.py b/litellm/types/router.py index 926815ba317..2b14a259047 100644 --- a/litellm/types/router.py +++ b/litellm/types/router.py @@ -246,6 +246,7 @@ class GenericLiteLLMParams(CredentialLiteLLMParams, CustomPricingLiteLLMParams): # Batch/File API Params s3_bucket_name: Optional[str] = None + s3_output_bucket_name: Optional[str] = None s3_encryption_key_id: Optional[str] = None gcs_bucket_name: Optional[str] = None diff --git a/tests/test_litellm/proxy/test_batch_s3_bucket_request_overrides.py b/tests/test_litellm/proxy/test_batch_s3_bucket_request_overrides.py new file mode 100644 index 00000000000..3e936384cbd --- /dev/null +++ b/tests/test_litellm/proxy/test_batch_s3_bucket_request_overrides.py @@ -0,0 +1,91 @@ +"""Bedrock batch: request-level S3 buckets flow through create_batch kwargs into litellm_params.""" + +import json +from unittest.mock import MagicMock, patch + +import litellm + + +def test_bedrock_batch_s3_output_bucket_name_kwarg_in_payload(): + captured_request_body = None + + def mock_post(*args, **kwargs): + nonlocal captured_request_body + if "data" in kwargs: + captured_request_body = kwargs["data"] + + mock_response = MagicMock() + mock_response.json.return_value = { + "jobArn": ( + "arn:aws:bedrock:us-west-2:123456789012:" + "model-invocation-job/test-job" + ), + "jobName": "test-job", + "status": "Submitted", + } + mock_response.status_code = 200 + mock_response.raise_for_status.return_value = None + return mock_response + + with patch( + "litellm.llms.custom_httpx.http_handler.HTTPHandler.post", + side_effect=mock_post, + ): + litellm.create_batch( + completion_window="24h", + endpoint="/v1/chat/completions", + input_file_id="s3://input-bucket/input/test.jsonl", + custom_llm_provider="bedrock", + model="us.anthropic.claude-haiku-4-5-20251001-v1:0", + aws_batch_role_arn="arn:aws:iam::123456789012:role/test-role", + s3_output_bucket_name="dedicated-output-bucket", + ) + + assert captured_request_body is not None + request_data = json.loads(captured_request_body) + s3_uri = request_data["outputDataConfig"]["s3OutputDataConfig"]["s3Uri"] + assert s3_uri.startswith("s3://dedicated-output-bucket/") + + +def test_bedrock_batch_s3_bucket_override_dict_wins_in_litellm_params(): + captured_request_body = None + + def mock_post(*args, **kwargs): + nonlocal captured_request_body + if "data" in kwargs: + captured_request_body = kwargs["data"] + + mock_response = MagicMock() + mock_response.json.return_value = { + "jobArn": ( + "arn:aws:bedrock:us-west-2:123456789012:" + "model-invocation-job/test-job" + ), + "jobName": "test-job", + "status": "Submitted", + } + mock_response.status_code = 200 + mock_response.raise_for_status.return_value = None + return mock_response + + with patch( + "litellm.llms.custom_httpx.http_handler.HTTPHandler.post", + side_effect=mock_post, + ): + litellm.create_batch( + completion_window="24h", + endpoint="/v1/chat/completions", + input_file_id="s3://input-bucket/input/test.jsonl", + custom_llm_provider="bedrock", + model="us.anthropic.claude-haiku-4-5-20251001-v1:0", + aws_batch_role_arn="arn:aws:iam::123456789012:role/test-role", + s3_output_bucket_name="model-default-output-bucket", + _litellm_batch_s3_bucket_overrides={ + "s3_output_bucket_name": "request-level-output-bucket" + }, + ) + + assert captured_request_body is not None + request_data = json.loads(captured_request_body) + s3_uri = request_data["outputDataConfig"]["s3OutputDataConfig"]["s3Uri"] + assert s3_uri.startswith("s3://request-level-output-bucket/")