scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-10-26 06:03:19 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 35

Merge 21e06501e6 into c6c303a618

Browse source
This commit is contained in:
Priyansh Nandwana 2026-10-20 01:52:11 +00:00 • committed by GitHub
commit 040d12f689
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
10 changed files with 187 additions and 16 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

1
litellm/litellm_core_utils/get_litellm_params.py
View file

@ -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",
}
)

18
litellm/llms/bedrock/batches/handler.py
View file

@ -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,6 +278,7 @@ class BedrockBatchesHandler:
batch_id: str,
aws_region_name: str | None = None,
logging_obj=None,
aws_bedrock_control_plane_endpoint: str | None = None,
**kwargs,
) -> "LiteLLMBatch":
"""
@ -289,6 +298,8 @@ class BedrockBatchesHandler:
aws_region_name: Region for the boto3 ``bedrock`` client. If omitted,
we fall back to parsing the region out of ``batch_id`` itself.
logging_obj: Optional litellm logging object.
aws_bedrock_control_plane_endpoint: Sends the call to a VPC endpoint
or gateway instead of the regional ``bedrock`` host.
**kwargs: Optional AWS credential overrides
(``aws_access_key_id``, ``aws_secret_access_key``,
``aws_session_token``, ``aws_profile_name``,
@ -322,6 +333,7 @@ class BedrockBatchesHandler:
client: Final = boto3.client(
"bedrock",
region_name=region,
endpoint_url=configured_bedrock_control_plane_endpoint(aws_bedrock_control_plane_endpoint),
aws_access_key_id=creds.access_key,
aws_secret_access_key=creds.secret_key,
aws_session_token=creds.token,
@ -339,9 +351,7 @@ class BedrockBatchesHandler:
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"{bedrock_control_plane_base_url(aws_bedrock_control_plane_endpoint, region)}/model-invocation-job/{url_path_id}",
},
)

19
litellm/llms/bedrock/batches/transformation.py
View file

@ -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,
@ -178,10 +179,8 @@ class BedrockBatchesConfig(BaseAWSLLM, BaseBatchesConfig):
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"
)
# Format: {control plane base}/model-invocation-job
bedrock_endpoint: Final = f"{bedrock_control_plane_base_url(request_params.get('aws_bedrock_control_plane_endpoint'), aws_region_name)}/model-invocation-job"
return bedrock_endpoint
@ -279,9 +278,7 @@ 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"
)
endpoint_url: Final = f"{bedrock_control_plane_base_url(request_params.get('aws_bedrock_control_plane_endpoint'), aws_region_name)}/model-invocation-job"
signed_headers, signed_data = self.common_utils.sign_aws_request(
service_name="bedrock",
data=bedrock_request,
@ -431,12 +428,10 @@ 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)
endpoint_url: Final = f"{bedrock_control_plane_base_url(request_params.get('aws_bedrock_control_plane_endpoint'), region)}/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

19
litellm/llms/bedrock/common_utils.py
View file

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

1
litellm/proxy/auth/auth_utils.py
View file

@ -417,6 +417,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

1
litellm/types/router.py
View file

@ -411,6 +411,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

39
tests/unit/llms/bedrock/batches/test_handler.py
View file

@ -257,6 +257,23 @@ 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):
_, 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):
_, 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 +496,28 @@ 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):
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",
)
# cancel builds a client to stop the job and another to read its status, retrieve builds one more
endpoints = [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``."""

100
tests/unit/llms/bedrock/batches/test_transformation.py
View file

@ -66,6 +66,89 @@ def test_get_complete_batch_url_uses_region(config):
assert url == "https://bedrock.eu-central-1.amazonaws.com/model-invocation-job"
CONTROL_PLANE_ENDPOINT = "https://vpce-0abc.bedrock.us-west-2.vpce.amazonaws.com/"
def test_get_complete_batch_url_uses_configured_control_plane_endpoint(config):
url = 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):
# the proxy hands batch calls the deployment's credentials through the strict CredentialLiteLLMParams filter
from litellm import Router
router = 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 = router.get_deployment_credentials_with_provider("claude-batch")
url = 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):
# api_base on a Bedrock deployment points at the runtime host, which does not serve batch jobs
url = 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):
with patch.object(
config.common_utils,
"generate_unique_job_name",
return_value="litellm-batch-deadbeef",
), patch.object(config.common_utils, "sign_aws_request") as mock_sign:
mock_sign.return_value = ({"Authorization": "signed"}, b"{}")
result = config.transform_create_batch_request(
model="anthropic.claude-3-5-sonnet",
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,
},
)
expected = "https://vpce-0abc.bedrock.us-west-2.vpce.amazonaws.com/model-invocation-job"
assert result["url"] == expected
assert mock_sign.call_args.kwargs["endpoint_url"] == expected
# --------------------------------------------------------------------------- #
# transform_create_batch_request - request construction (sign_aws_request mocked)
# --------------------------------------------------------------------------- #
@ -608,6 +691,23 @@ 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):
with patch.object(config.common_utils, "sign_aws_request") as mock_sign:
mock_sign.return_value = ({"Authorization": "signed"}, b"")
result = config.transform_retrieve_batch_request(
batch_id=ARN,
optional_params={},
litellm_params={"aws_bedrock_control_plane_endpoint": CONTROL_PLANE_ENDPOINT},
)
expected = (
"https://vpce-0abc.bedrock.us-west-2.vpce.amazonaws.com/model-invocation-job/"
"arn%3Aaws%3Abedrock%3Aus-west-2%3A123456789012%3Amodel-invocation-job%2Fabc1234567"
)
assert result["url"] == expected
assert mock_sign.call_args.kwargs["endpoint_url"] == expected
def test_retrieve_request_rejects_non_arn(config):
with pytest.raises(ValueError, match="Expected ARN"):
config.transform_retrieve_batch_request(

1
tests/unit/proxy/auth/test_auth_utils.py
View file

@ -3357,6 +3357,7 @@ class TestIsRequestBodySafeBlocksEndpointTargetingFields:
"field",
[
"aws_bedrock_runtime_endpoint",
"aws_bedrock_control_plane_endpoint",
"langsmith_base_url",
"langfuse_host",
"posthog_host",

4
ui/litellm-dashboard/src/lib/http/schema.d.ts generated vendored
View file

@ -35028,6 +35028,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 */
@ -50278,6 +50280,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 */