diff --git a/cookbook/litellm_proxy_server/batch_api/bedrock/README.md b/cookbook/litellm_proxy_server/batch_api/bedrock/README.md new file mode 100644 index 00000000000..1410d76658e --- /dev/null +++ b/cookbook/litellm_proxy_server/batch_api/bedrock/README.md @@ -0,0 +1,41 @@ +# Bedrock batch endpoints + +Set `aws_bedrock_control_plane_endpoint` on a deployment when batch create, retrieve, and cancel requests must use a custom Bedrock control-plane endpoint. The proxy signs these requests for the configured URL. Without this setting, batch requests use the standard regional Bedrock endpoint + +```yaml +model_list: + - model_name: bedrock-batch + litellm_params: + model: bedrock/us.anthropic.claude-haiku-4-5-20251001-v1:0 + aws_region_name: us-east-1 + aws_access_key_id: os.environ/AWS_ACCESS_KEY_ID + aws_secret_access_key: os.environ/AWS_SECRET_ACCESS_KEY + aws_batch_role_arn: os.environ/AWS_BATCH_ROLE_ARN + s3_output_bucket_name: os.environ/AWS_S3_OUTPUT_BUCKET_NAME + aws_bedrock_control_plane_endpoint: os.environ/BEDROCK_CONTROL_PLANE_ENDPOINT +``` + +Set `BEDROCK_CONTROL_PLANE_ENDPOINT` to your Bedrock control-plane VPC endpoint or compatible gateway base URL, without `/model-invocation-job`. The proxy adds the operation path. The endpoint must be reachable from the proxy, and the configured AWS region must match the endpoint's signing region + +`api_base` and `aws_bedrock_runtime_endpoint` configure runtime inference separately. They do not override the batch control-plane endpoint. `s3_endpoint_url` configures S3 file access separately + +The control-plane setting belongs in deployment configuration. Sending `aws_bedrock_control_plane_endpoint` in an API request body is rejected. There is no automatic `AWS_BEDROCK_CONTROL_PLANE_ENDPOINT` environment fallback, so reference the desired environment variable explicitly as shown above + +Start the proxy with this configuration and set `LITELLM_API_BASE` to its URL. Set `BEDROCK_BATCH_INPUT_URI` to a valid input JSONL file in S3 that the batch service role can read + +```bash +curl -sS "$LITELLM_API_BASE/v1/batches" \ + -H "Authorization: Bearer $LITELLM_MASTER_KEY" \ + -H 'Content-Type: application/json' \ + -d "{\"model\":\"bedrock-batch\",\"input_file_id\":\"$BEDROCK_BATCH_INPUT_URI\",\"endpoint\":\"/v1/chat/completions\",\"completion_window\":\"24h\"}" +``` + +Use the returned `id` as `BATCH_ID` to retrieve or cancel the same job through the deployment's configured endpoint + +```bash +curl -sS "$LITELLM_API_BASE/v1/batches/$BATCH_ID?model=bedrock-batch" \ + -H "Authorization: Bearer $LITELLM_MASTER_KEY" + +curl -sS -X POST "$LITELLM_API_BASE/v1/batches/$BATCH_ID/cancel?model=bedrock-batch" \ + -H "Authorization: Bearer $LITELLM_MASTER_KEY" +``` diff --git a/litellm/litellm_core_utils/get_litellm_params.py b/litellm/litellm_core_utils/get_litellm_params.py index 6ca02cbbe64..657fd3abfbc 100644 --- a/litellm/litellm_core_utils/get_litellm_params.py +++ b/litellm/litellm_core_utils/get_litellm_params.py @@ -27,6 +27,7 @@ AWS_CREDENTIAL_KWARGS_KEYS: Final = frozenset( "aws_external_id", "aws_session_tags", "aws_bedrock_runtime_endpoint", + "aws_bedrock_control_plane_endpoint", "aws_bedrock_project_id", } ) diff --git a/litellm/llms/bedrock/batches/handler.py b/litellm/llms/bedrock/batches/handler.py index 81057a1ee23..48e79d39b8e 100644 --- a/litellm/llms/bedrock/batches/handler.py +++ b/litellm/llms/bedrock/batches/handler.py @@ -5,7 +5,10 @@ from typing import TYPE_CHECKING, Final, Literal from openai.types.batch import BatchRequestCounts from openai.types.batch import Metadata as OpenAIBatchMetadata -from litellm.litellm_core_utils.aws_partition import get_aws_dns_suffix +from litellm.llms.bedrock.common_utils import ( + bedrock_control_plane_base_url, + configured_bedrock_control_plane_endpoint, +) from litellm.types.llms.bedrock import AwsAuthParams, AwsSessionTag from litellm.types.utils import LiteLLMBatch @@ -156,10 +159,14 @@ class BedrockBatchesHandler: aws_session_tags=aws_session_tags, ) creds: Final = BedrockBatchesConfig().resolve_credentials(auth_params, region) + control_plane_endpoint: Final = configured_bedrock_control_plane_endpoint( + kwargs.get("aws_bedrock_control_plane_endpoint") + ) client: Final = boto3.client( "bedrock", region_name=region, + endpoint_url=control_plane_endpoint, aws_access_key_id=creds.access_key, aws_secret_access_key=creds.secret_key, aws_session_token=creds.token, @@ -171,6 +178,7 @@ class BedrockBatchesHandler: batch_id=batch_id, aws_region_name=region, logging_obj=logging_obj, + aws_bedrock_control_plane_endpoint=control_plane_endpoint, **auth_params.model_dump(), ) @@ -270,7 +278,7 @@ class BedrockBatchesHandler: batch_id: str, aws_region_name: str | None = None, logging_obj=None, - **kwargs, + **kwargs: object, ) -> "LiteLLMBatch": """ Handle ``GetModelInvocationJob`` status check for AWS Bedrock bulk batch @@ -294,7 +302,10 @@ class BedrockBatchesHandler: ``aws_session_token``, ``aws_profile_name``, ``aws_role_name``, ``aws_session_name``, ``aws_web_identity_token``, ``aws_sts_endpoint``, - ``aws_external_id``, ``aws_session_tags``). Unknown keys are ignored. + ``aws_external_id``, ``aws_session_tags``), plus + ``aws_bedrock_control_plane_endpoint`` to send the call to a VPC + endpoint or gateway instead of the regional ``bedrock`` host. + Unknown keys are ignored. Returns: ``LiteLLMBatch`` shaped like an OpenAI Batch resource. @@ -318,10 +329,14 @@ class BedrockBatchesHandler: from litellm.llms.bedrock.batches.transformation import BedrockBatchesConfig creds: Final = BedrockBatchesConfig().resolve_credentials(AwsAuthParams.model_validate(kwargs), region) + control_plane_endpoint: Final = configured_bedrock_control_plane_endpoint( + kwargs.get("aws_bedrock_control_plane_endpoint") + ) client: Final = boto3.client( "bedrock", region_name=region, + endpoint_url=control_plane_endpoint, aws_access_key_id=creds.access_key, aws_secret_access_key=creds.secret_key, aws_session_token=creds.token, @@ -334,14 +349,13 @@ class BedrockBatchesHandler: # `GetModelInvocationJob` accepts either form, but only the bare id # produces a sensible-looking URL in logs. url_path_id: Final = _extract_job_id_from_arn(batch_id) or batch_id + control_plane_base: Final = bedrock_control_plane_base_url(control_plane_endpoint, region) logging_obj.pre_call( input=batch_id, api_key="", additional_args={ "complete_input_dict": {"jobIdentifier": batch_id}, - "api_base": ( - f"https://bedrock.{region}.{get_aws_dns_suffix(region)}/model-invocation-job/{url_path_id}" - ), + "api_base": f"{control_plane_base}/model-invocation-job/{url_path_id}", }, ) diff --git a/litellm/llms/bedrock/batches/transformation.py b/litellm/llms/bedrock/batches/transformation.py index 50fb911709a..84f40bec9ca 100644 --- a/litellm/llms/bedrock/batches/transformation.py +++ b/litellm/llms/bedrock/batches/transformation.py @@ -7,7 +7,7 @@ from typing import TYPE_CHECKING, Final, Literal, cast from httpx import Headers, Response from pydantic import TypeAdapter, ValidationError -from litellm.litellm_core_utils.aws_partition import get_aws_dns_suffix, is_bedrock_arn +from litellm.litellm_core_utils.aws_partition import is_bedrock_arn from litellm.litellm_core_utils.cloud_storage_security import ( BEDROCK_MANAGED_S3_BATCH_PREFIX, ) @@ -32,6 +32,7 @@ from litellm.types.utils import LiteLLMBatch, LlmProviders, Usage from ..base_aws_llm import BaseAWSLLM from ..common_utils import ( CommonBatchFilesUtils, + bedrock_control_plane_base_url, merge_bedrock_aws_request_params, resolve_s3_bucket_owner, resolve_s3_encryption_key_id, @@ -177,11 +178,10 @@ class BedrockBatchesConfig(BaseAWSLLM, BaseBatchesConfig): request_params: Final = merge_bedrock_aws_request_params(litellm_params, optional_params) aws_region_name: Final = self._get_aws_region_name(request_params, model) - # Bedrock model invocation job endpoint - # Format: https://bedrock.{region}.{partition dns suffix}/model-invocation-job - bedrock_endpoint: Final = ( - f"https://bedrock.{aws_region_name}.{get_aws_dns_suffix(aws_region_name)}/model-invocation-job" + control_plane_base: Final = bedrock_control_plane_base_url( + request_params.get("aws_bedrock_control_plane_endpoint"), aws_region_name ) + bedrock_endpoint: Final = f"{control_plane_base}/model-invocation-job" return bedrock_endpoint @@ -279,9 +279,10 @@ class BedrockBatchesConfig(BaseAWSLLM, BaseBatchesConfig): # Use common utility for AWS signing request_params: Final = merge_bedrock_aws_request_params(litellm_params, optional_params) aws_region_name: Final = self._get_aws_region_name(request_params, model) - endpoint_url: Final = ( - f"https://bedrock.{aws_region_name}.{get_aws_dns_suffix(aws_region_name)}/model-invocation-job" + control_plane_base: Final = bedrock_control_plane_base_url( + request_params.get("aws_bedrock_control_plane_endpoint"), aws_region_name ) + endpoint_url: Final = f"{control_plane_base}/model-invocation-job" signed_headers, signed_data = self.common_utils.sign_aws_request( service_name="bedrock", data=bedrock_request, @@ -431,12 +432,13 @@ class BedrockBatchesConfig(BaseAWSLLM, BaseBatchesConfig): import urllib.parse as _ul encoded_arn: Final = _ul.quote(batch_id, safe="") - endpoint_url: Final = ( - f"https://bedrock.{region}.{get_aws_dns_suffix(region)}/model-invocation-job/{encoded_arn}" + request_params: Final = merge_bedrock_aws_request_params(litellm_params, optional_params) + control_plane_base: Final = bedrock_control_plane_base_url( + request_params.get("aws_bedrock_control_plane_endpoint"), region ) + endpoint_url: Final = f"{control_plane_base}/model-invocation-job/{encoded_arn}" # Use common utility for AWS signing - request_params: Final = merge_bedrock_aws_request_params(litellm_params, optional_params) signed_headers, _ = self.common_utils.sign_aws_request( service_name="bedrock", data={}, # GET request has no body diff --git a/litellm/llms/bedrock/common_utils.py b/litellm/llms/bedrock/common_utils.py index bbbc929fad3..0cf3e895ad2 100644 --- a/litellm/llms/bedrock/common_utils.py +++ b/litellm/llms/bedrock/common_utils.py @@ -138,6 +138,25 @@ def merge_bedrock_aws_request_params( return request_params +def configured_bedrock_control_plane_endpoint(endpoint: object) -> str | None: + """An ``aws_bedrock_control_plane_endpoint`` value without a trailing slash, or None when unset.""" + if isinstance(endpoint, str) and endpoint.strip(): + return endpoint.strip().rstrip("/") + return None + + +def bedrock_control_plane_base_url(endpoint: object, aws_region_name: str) -> str: + """Base URL for Bedrock control-plane calls such as model invocation jobs. + + A configured ``aws_bedrock_control_plane_endpoint`` (a VPC endpoint or a compatible gateway) + wins over the regional ``bedrock`` host. ``api_base`` is deliberately not used here, since on a + Bedrock deployment it points at the runtime host, which does not serve these APIs. + """ + return configured_bedrock_control_plane_endpoint(endpoint) or ( + f"https://bedrock.{aws_region_name}.{get_aws_dns_suffix(aws_region_name)}" + ) + + def s3_static_key_pair(params: Mapping[str, object]) -> tuple[str, str] | None: """The s3_access_key_id / s3_secret_access_key pair when both are set, otherwise None.""" s3_access_key_id: Final = params.get("s3_access_key_id") diff --git a/litellm/proxy/auth/auth_utils.py b/litellm/proxy/auth/auth_utils.py index 73d20ae632c..996c28bebaa 100644 --- a/litellm/proxy/auth/auth_utils.py +++ b/litellm/proxy/auth/auth_utils.py @@ -419,6 +419,7 @@ _BANNED_REQUEST_BODY_PARAMS: Final[tuple[str, ...]] = ( # tokens) to the attacker's host, or coerces the proxy into # authenticating against the attacker's host with admin secrets. "aws_bedrock_runtime_endpoint", + "aws_bedrock_control_plane_endpoint", # Bedrock project/workspace association. Deployments pin this to # enforce a data-retention policy, so a caller-supplied value would # re-route the request's retention and accounting to any project diff --git a/litellm/types/litellm_params.py b/litellm/types/litellm_params.py index 978eb9c4557..6b08920e042 100644 --- a/litellm/types/litellm_params.py +++ b/litellm/types/litellm_params.py @@ -103,6 +103,7 @@ class ProviderConnection: class BedrockBatchConnection: # Bedrock rejects these names in request bodies, so register them as LiteLLM-owned aws_batch_role_arn: str | None = None + aws_bedrock_control_plane_endpoint: str | None = None s3_bucket_name: str | None = None s3_region_name: str | None = None s3_endpoint_url: str | None = None diff --git a/litellm/types/router.py b/litellm/types/router.py index 1f159cb8b16..2618bf698ee 100644 --- a/litellm/types/router.py +++ b/litellm/types/router.py @@ -413,6 +413,7 @@ class CredentialLiteLLMParams(LiteLLMBaseModel): aws_external_id: str | None = None aws_session_tags: Sequence[AwsSessionTag] | None = None aws_bedrock_runtime_endpoint: str | None = None + aws_bedrock_control_plane_endpoint: str | None = None aws_bedrock_project_id: str | None = None s3_bucket_name: str | None = None s3_endpoint_url: str | None = None diff --git a/tests/unit/llms/bedrock/batches/test_handler.py b/tests/unit/llms/bedrock/batches/test_handler.py index d920eabcdf6..602504f0c49 100644 --- a/tests/unit/llms/bedrock/batches/test_handler.py +++ b/tests/unit/llms/bedrock/batches/test_handler.py @@ -257,6 +257,25 @@ def test_explicit_region_overrides_arn(patched_boto3): assert kwargs["region_name"] == "eu-central-1" +def test_status_client_uses_regional_host_when_no_control_plane_endpoint( + patched_boto3: tuple[MagicMock, MagicMock], +) -> None: + _, boto_client_factory = patched_boto3 + BedrockBatchesHandler.handle_model_invocation_job_status(batch_id=JOB_ARN) + _, kwargs = boto_client_factory.call_args + assert kwargs["endpoint_url"] is None + + +def test_status_client_uses_configured_control_plane_endpoint(patched_boto3: tuple[MagicMock, MagicMock]) -> None: + _, boto_client_factory = patched_boto3 + BedrockBatchesHandler.handle_model_invocation_job_status( + batch_id=JOB_ARN, + aws_bedrock_control_plane_endpoint="https://bedrock-gateway.internal/", + ) + _, kwargs = boto_client_factory.call_args + assert kwargs["endpoint_url"] == "https://bedrock-gateway.internal" + + def test_failure_message_propagates(patched_boto3): fake_client, _ = patched_boto3 failed_response = _fake_boto3_response(status="Failed") @@ -479,6 +498,29 @@ def test_litellm_cancel_batch_dispatches_to_bedrock(patched_boto3): assert batch.status == "cancelled" +def test_litellm_batch_calls_send_every_bedrock_client_to_the_configured_endpoint( + patched_boto3: tuple[MagicMock, MagicMock], +) -> None: + import litellm + + fake_client, boto_client_factory = patched_boto3 + fake_client.get_model_invocation_job.return_value = _fake_boto3_response(status="Stopped") + + litellm.cancel_batch( + batch_id=JOB_ARN, + custom_llm_provider="bedrock", + aws_bedrock_control_plane_endpoint="https://bedrock-gateway.internal", + ) + litellm.retrieve_batch( + batch_id=JOB_ARN, + custom_llm_provider="bedrock", + aws_bedrock_control_plane_endpoint="https://bedrock-gateway.internal", + ) + + endpoints: Final = tuple(call.kwargs["endpoint_url"] for call in boto_client_factory.call_args_list) + assert endpoints == ("https://bedrock-gateway.internal",) * 3 + + class _TagGatedSTSClient: """Stands in for STS behind a trust policy that only admits sessions carrying ``tags``.""" diff --git a/tests/unit/llms/bedrock/batches/test_transformation.py b/tests/unit/llms/bedrock/batches/test_transformation.py index 5e987239c54..5360ac00018 100644 --- a/tests/unit/llms/bedrock/batches/test_transformation.py +++ b/tests/unit/llms/bedrock/batches/test_transformation.py @@ -14,10 +14,15 @@ URL/ARN handling, and the error class. AWS auth/sigv4 is the only external seam we mock; everything else runs for real. """ +from datetime import datetime, timezone +from typing import Final from unittest.mock import MagicMock, patch import httpx import pytest +from botocore.auth import SigV4Auth +from botocore.awsrequest import AWSRequest +from botocore.credentials import Credentials from litellm.llms.bedrock.batches.transformation import BedrockBatchesConfig from litellm.types.utils import LlmProviders @@ -41,7 +46,7 @@ ARN = "arn:aws:bedrock:us-west-2:123456789012:model-invocation-job/abc1234567" @pytest.fixture -def config(): +def config() -> BedrockBatchesConfig: return BedrockBatchesConfig() @@ -66,6 +71,90 @@ def test_get_complete_batch_url_uses_region(config): assert url == "https://bedrock.eu-central-1.amazonaws.com/model-invocation-job" +CONTROL_PLANE_ENDPOINT: Final = "https://vpce-0abc.bedrock.us-west-2.vpce.amazonaws.com/" + + +def test_get_complete_batch_url_uses_configured_control_plane_endpoint(config: BedrockBatchesConfig) -> None: + url: Final = config.get_complete_batch_url( + api_base=None, + api_key=None, + model="anthropic.claude-3", + optional_params={"aws_region_name": "us-west-2"}, + litellm_params={"aws_bedrock_control_plane_endpoint": CONTROL_PLANE_ENDPOINT}, + data={"input_file_id": "s3://b/k"}, + ) + assert url == "https://vpce-0abc.bedrock.us-west-2.vpce.amazonaws.com/model-invocation-job" + + +def test_proxy_deployment_credentials_carry_the_control_plane_endpoint(config: BedrockBatchesConfig) -> None: + from litellm import Router + + router: Final = Router( + model_list=[ + { + "model_name": "claude-batch", + "litellm_params": { + "model": "bedrock/anthropic.claude-3", + "aws_region_name": "us-west-2", + "aws_bedrock_control_plane_endpoint": CONTROL_PLANE_ENDPOINT, + }, + } + ] + ) + credentials: Final = router.get_deployment_credentials_with_provider("claude-batch") + + url: Final = config.get_complete_batch_url( + api_base=None, + api_key=None, + model="anthropic.claude-3", + optional_params={}, + litellm_params=credentials, + data={"input_file_id": "s3://b/k"}, + ) + assert url == "https://vpce-0abc.bedrock.us-west-2.vpce.amazonaws.com/model-invocation-job" + + +def test_get_complete_batch_url_ignores_runtime_api_base(config: BedrockBatchesConfig) -> None: + url: Final = config.get_complete_batch_url( + api_base="https://bedrock-runtime.us-west-2.amazonaws.com", + api_key=None, + model="anthropic.claude-3", + optional_params={"aws_region_name": "us-west-2"}, + litellm_params={"api_base": "https://bedrock-runtime.us-west-2.amazonaws.com"}, + data={"input_file_id": "s3://b/k"}, + ) + assert url == "https://bedrock.us-west-2.amazonaws.com/model-invocation-job" + + +def test_create_request_sends_and_signs_for_configured_control_plane_endpoint(config: BedrockBatchesConfig) -> None: + expected_url: Final = "https://vpce-0abc.bedrock.us-west-2.vpce.amazonaws.com/model-invocation-job" + with patch("botocore.auth.get_current_datetime", return_value=datetime(2026, 1, 1, tzinfo=timezone.utc)): + result: Final = config.transform_create_batch_request( + model="anthropic.claude-haiku-4-5-20251001-v1:0", + create_batch_data={ + "input_file_id": "s3://in-bucket/input.jsonl", + "endpoint": "/v1/chat/completions", + "completion_window": "24h", + }, + optional_params={"aws_region_name": "us-west-2"}, + litellm_params={ + "aws_batch_role_arn": "arn:aws:iam::123:role/my-batch-role", + "aws_bedrock_control_plane_endpoint": CONTROL_PLANE_ENDPOINT, + "aws_access_key_id": "endpoint-test-key", + "aws_secret_access_key": "endpoint-test-secret", + }, + ) + expected_request: Final = AWSRequest( + method="POST", url=expected_url, data=result["data"], headers={"Content-Type": "application/json"} + ) + SigV4Auth(Credentials("endpoint-test-key", "endpoint-test-secret"), "bedrock", "us-west-2").add_auth( + expected_request + ) + + assert result["url"] == expected_url + assert result["headers"]["Authorization"] == expected_request.headers["Authorization"] + + # --------------------------------------------------------------------------- # # transform_create_batch_request - request construction (sign_aws_request mocked) # --------------------------------------------------------------------------- # @@ -608,6 +697,30 @@ def test_retrieve_request_builds_encoded_arn_url(config): assert mock_sign.call_args.kwargs["data"] == {} +def test_retrieve_request_sends_and_signs_for_configured_control_plane_endpoint(config: BedrockBatchesConfig) -> None: + expected_url: Final = ( + "https://vpce-0abc.bedrock.us-west-2.vpce.amazonaws.com/model-invocation-job/" + "arn%3Aaws%3Abedrock%3Aus-west-2%3A123456789012%3Amodel-invocation-job%2Fabc1234567" + ) + with patch("botocore.auth.get_current_datetime", return_value=datetime(2026, 1, 1, tzinfo=timezone.utc)): + result: Final = config.transform_retrieve_batch_request( + batch_id=ARN, + optional_params={}, + litellm_params={ + "aws_bedrock_control_plane_endpoint": CONTROL_PLANE_ENDPOINT, + "aws_access_key_id": "endpoint-test-key", + "aws_secret_access_key": "endpoint-test-secret", + }, + ) + expected_request: Final = AWSRequest(method="GET", url=expected_url, data="", headers={}) + SigV4Auth(Credentials("endpoint-test-key", "endpoint-test-secret"), "bedrock", "us-west-2").add_auth( + expected_request + ) + + assert result["url"] == expected_url + assert result["headers"]["Authorization"] == expected_request.headers["Authorization"] + + def test_retrieve_request_rejects_non_arn(config): with pytest.raises(ValueError, match="Expected ARN"): config.transform_retrieve_batch_request( diff --git a/tests/unit/llms/bedrock/chat/test_converse_transformation.py b/tests/unit/llms/bedrock/chat/test_converse_transformation.py index 1fb5542c964..154797bab5a 100644 --- a/tests/unit/llms/bedrock/chat/test_converse_transformation.py +++ b/tests/unit/llms/bedrock/chat/test_converse_transformation.py @@ -10739,3 +10739,53 @@ def test_a_refused_native_schema_is_sent_again_only_when_the_retry_drops_native_ assert route.call_count == expected_calls assert all("outputConfig" in json.loads(call.request.content) for call in route.calls[:1]) assert all("outputConfig" not in json.loads(call.request.content) for call in route.calls[1:]) + + +CONTROL_PLANE_ENDPOINT: Final = "https://bedrock-fips.us-east-1.amazonaws.com" + + +def _bedrock_chat_answer(request: httpx.Request) -> httpx.Response: + if request.url.path.endswith("/invoke"): + return httpx.Response( + 200, + json={ + "id": "msg_control_plane", + "type": "message", + "role": "assistant", + "model": "claude-haiku-4-5", + "content": [{"type": "text", "text": "ok"}], + "stop_reason": "end_turn", + "usage": {"input_tokens": 1, "output_tokens": 1}, + }, + ) + return httpx.Response( + 200, + json={ + "output": {"message": {"role": "assistant", "content": [{"text": "ok"}]}}, + "stopReason": "end_turn", + "usage": {"inputTokens": 1, "outputTokens": 1, "totalTokens": 2}, + }, + ) + + +@pytest.mark.parametrize("route_prefix", ["bedrock/", "bedrock/converse/", "bedrock/invoke/"]) +@pytest.mark.respx(assert_all_called=True) +def test_a_deployment_with_a_batch_control_plane_endpoint_keeps_it_out_of_chat_requests( + respx_mock: respx.Router, route_prefix: str +) -> None: + route: Final = respx_mock.post( + url__regex=r"^https://bedrock-runtime\.us-east-1\.amazonaws\.com/model/.+/(converse|invoke)$" + ).mock(side_effect=_bedrock_chat_answer) + + response: Final = litellm.completion( + model=f"{route_prefix}us.anthropic.claude-haiku-4-5-20251001-v1:0", + messages=[{"role": "user", "content": "hi"}], + aws_access_key_id="AKIACONTROLPLANE", + aws_secret_access_key="control-plane-secret", + aws_region_name="us-east-1", + aws_bedrock_control_plane_endpoint=CONTROL_PLANE_ENDPOINT, + client=HTTPHandler(), + ) + + assert response.choices[0].message.content == "ok" + assert "aws_bedrock_control_plane_endpoint" not in route.calls[0].request.content.decode() diff --git a/tests/unit/proxy/auth/test_auth_utils.py b/tests/unit/proxy/auth/test_auth_utils.py index f81ba155097..67ff7dd9344 100644 --- a/tests/unit/proxy/auth/test_auth_utils.py +++ b/tests/unit/proxy/auth/test_auth_utils.py @@ -3566,6 +3566,7 @@ class TestIsRequestBodySafeBlocksEndpointTargetingFields: "field", [ "aws_bedrock_runtime_endpoint", + "aws_bedrock_control_plane_endpoint", "langsmith_base_url", "langfuse_host", "posthog_host", diff --git a/tests/unit/types/test_litellm_params.py b/tests/unit/types/test_litellm_params.py index e88d757519c..84ba85f4a68 100644 --- a/tests/unit/types/test_litellm_params.py +++ b/tests/unit/types/test_litellm_params.py @@ -95,6 +95,7 @@ CONNECTION_NAMES: Final = ( "use_xai_oauth", "fireworks_forward_user_id", "aws_batch_role_arn", + "aws_bedrock_control_plane_endpoint", "s3_bucket_name", "s3_region_name", "s3_endpoint_url", @@ -274,6 +275,7 @@ INTERNAL_STATE_NAMES: Final = ( BEDROCK_BATCH_NAMES: Final = ( "aws_batch_role_arn", + "aws_bedrock_control_plane_endpoint", "s3_bucket_name", "s3_region_name", "s3_endpoint_url", @@ -686,6 +688,7 @@ NAMES_SHARED_WITH_TYPED_MODELS: Final[Mapping[str, tuple[str, ...]]] = MappingPr "api_key", "api_version", "aws_batch_role_arn", + "aws_bedrock_control_plane_endpoint", "azure_password", "azure_scope", "azure_username", diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index b93f7dce2ca..e432c947116 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -35048,6 +35048,8 @@ export interface components { aws_access_key_id?: string | null; /** Aws Batch Role Arn */ aws_batch_role_arn?: string | null; + /** Aws Bedrock Control Plane Endpoint */ + aws_bedrock_control_plane_endpoint?: string | null; /** Aws Bedrock Project Id */ aws_bedrock_project_id?: string | null; /** Aws Bedrock Runtime Endpoint */ @@ -50336,6 +50338,8 @@ export interface components { aws_access_key_id?: string | null; /** Aws Batch Role Arn */ aws_batch_role_arn?: string | null; + /** Aws Bedrock Control Plane Endpoint */ + aws_bedrock_control_plane_endpoint?: string | null; /** Aws Bedrock Project Id */ aws_bedrock_project_id?: string | null; /** Aws Bedrock Runtime Endpoint */