tiennm99/litellm
1
0
Fork 0
mirror of https://github.com/tiennm99/litellm.git synced 2026-07-10 19:04:29 +00:00
Code Issues Packages Projects Releases Wiki Activity

Merge pull request #14582 from timelfrink/fix/issue-14573-aws-external-id-support

Browse Source 
Add AWS external ID parameter support for Bedrock authentication
This commit is contained in:
Krish Dholakia
2025-09-15 17:32:00 -07:00
committed by GitHub
parent eb3e159b7c 5bd94cccb9
commit e1a6b9f858
4 changed files with 209 additions and 26 deletions
Download Patch File Download Diff File Expand all files Collapse all files
litellm/llms/bedrock/base_aws_llm.py
+65 -22
View File
@@ -66,6 +66,7 @@ class BaseAWSLLM:
"aws_web_identity_token",
"aws_sts_endpoint",
"aws_bedrock_runtime_endpoint",
"aws_external_id",
]
def get_cache_key(self, credential_args: Dict[str, Optional[str]]) -> str:
@@ -88,6 +89,7 @@ class BaseAWSLLM:
aws_role_name: Optional[str] = None,
aws_web_identity_token: Optional[str] = None,
aws_sts_endpoint: Optional[str] = None,
aws_external_id: Optional[str] = None,
):
"""
Return a boto3.Credentials object
@@ -103,6 +105,7 @@ class BaseAWSLLM:
aws_role_name,
aws_web_identity_token,
aws_sts_endpoint,
aws_external_id,
]
# Iterate over parameters and update if needed
@@ -127,6 +130,7 @@ class BaseAWSLLM:
aws_role_name,
aws_web_identity_token,
aws_sts_endpoint,
aws_external_id,
) = params_to_check
verbose_logger.debug(
@@ -139,7 +143,8 @@ class BaseAWSLLM:
"aws_profile_name=%s\n"
"aws_role_name=%s\n"
"aws_web_identity_token=%s\n"
"aws_sts_endpoint=%s",
"aws_sts_endpoint=%s\n"
"aws_external_id=%s",
aws_access_key_id,
aws_secret_access_key,
aws_session_token,
@@ -149,6 +154,7 @@ class BaseAWSLLM:
aws_role_name,
aws_web_identity_token,
aws_sts_endpoint,
aws_external_id,
)
# create cache key for non-expiring auth flows
@@ -177,6 +183,7 @@ class BaseAWSLLM:
aws_session_name=aws_session_name,
aws_region_name=aws_region_name,
aws_sts_endpoint=aws_sts_endpoint,
aws_external_id=aws_external_id,
)
elif aws_role_name is not None:
# Check if we're in IRSA and trying to assume the same role we already have
@@ -205,6 +212,7 @@ class BaseAWSLLM:
aws_session_token=aws_session_token,
aws_role_name=aws_role_name,
aws_session_name=aws_session_name,
aws_external_id=aws_external_id,
)
elif aws_profile_name is not None: ### CHECK SESSION ###
@@ -406,6 +414,7 @@ class BaseAWSLLM:
aws_session_name: str,
aws_region_name: Optional[str],
aws_sts_endpoint: Optional[str],
aws_external_id: Optional[str] = None,
) -> Tuple[Credentials, Optional[int]]:
"""
Authenticate with AWS Web Identity Token
@@ -438,13 +447,19 @@ class BaseAWSLLM:
# https://docs.aws.amazon.com/STS/latest/APIReference/API_AssumeRoleWithWebIdentity.html
# https://boto3.amazonaws.com/v1/documentation/api/latest/reference/services/sts/client/assume_role_with_web_identity.html
sts_response = sts_client.assume_role_with_web_identity(
RoleArn=aws_role_name,
RoleSessionName=aws_session_name,
WebIdentityToken=oidc_token,
DurationSeconds=3600,
Policy='{"Version":"2012-10-17","Statement":[{"Sid":"BedrockLiteLLM","Effect":"Allow","Action":["bedrock:InvokeModel","bedrock:InvokeModelWithResponseStream"],"Resource":"*","Condition":{"Bool":{"aws:SecureTransport":"true"},"StringLike":{"aws:UserAgent":"litellm/*"}}}]}',
)
assume_role_params = {
"RoleArn": aws_role_name,
"RoleSessionName": aws_session_name,
"WebIdentityToken": oidc_token,
"DurationSeconds": 3600,
"Policy": '{"Version":"2012-10-17","Statement":[{"Sid":"BedrockLiteLLM","Effect":"Allow","Action":["bedrock:InvokeModel","bedrock:InvokeModelWithResponseStream"],"Resource":"*","Condition":{"Bool":{"aws:SecureTransport":"true"},"StringLike":{"aws:UserAgent":"litellm/*"}}}]}',
}
# Add ExternalId parameter if provided
if aws_external_id is not None:
assume_role_params["ExternalId"] = aws_external_id
sts_response = sts_client.assume_role_with_web_identity(**assume_role_params)
iam_creds_dict = {
"aws_access_key_id": sts_response["Credentials"]["AccessKeyId"],
@@ -464,8 +479,9 @@ class BaseAWSLLM:
iam_creds = session.get_credentials()
return iam_creds, self._get_default_ttl_for_boto3_credentials()
def _handle_irsa_cross_account(self, irsa_role_arn: str, aws_role_name: str,
aws_session_name: str, region: str, web_identity_token_file: str) -> dict:
def _handle_irsa_cross_account(self, irsa_role_arn: str, aws_role_name: str,
aws_session_name: str, region: str, web_identity_token_file: str,
aws_external_id: Optional[str] = None) -> dict:
"""Handle cross-account role assumption for IRSA."""
import boto3
@@ -509,11 +525,19 @@ class BaseAWSLLM:
# Now assume the target role
verbose_logger.debug(f"Attempting to assume target role: {aws_role_name} with session: {aws_session_name}")
return sts_client_with_creds.assume_role(
RoleArn=aws_role_name, RoleSessionName=aws_session_name
)
assume_role_params = {
"RoleArn": aws_role_name,
"RoleSessionName": aws_session_name
}
def _handle_irsa_same_account(self, aws_role_name: str, aws_session_name: str, region: str) -> dict:
# Add ExternalId parameter if provided
if aws_external_id is not None:
assume_role_params["ExternalId"] = aws_external_id
return sts_client_with_creds.assume_role(**assume_role_params)
def _handle_irsa_same_account(self, aws_role_name: str, aws_session_name: str, region: str,
aws_external_id: Optional[str] = None) -> dict:
"""Handle same-account role assumption for IRSA."""
import boto3
@@ -530,9 +554,16 @@ class BaseAWSLLM:
# Assume the role
verbose_logger.debug(f"Attempting to assume role: {aws_role_name} with session: {aws_session_name}")
return sts_client.assume_role(
RoleArn=aws_role_name, RoleSessionName=aws_session_name
)
assume_role_params = {
"RoleArn": aws_role_name,
"RoleSessionName": aws_session_name
}
# Add ExternalId parameter if provided
if aws_external_id is not None:
assume_role_params["ExternalId"] = aws_external_id
return sts_client.assume_role(**assume_role_params)
def _extract_credentials_and_ttl(self, sts_response: dict) -> Tuple[Credentials, Optional[int]]:
"""Extract credentials and TTL from STS response."""
@@ -558,6 +589,7 @@ class BaseAWSLLM:
aws_session_token: Optional[str],
aws_role_name: str,
aws_session_name: str,
aws_external_id: Optional[str] = None,
) -> Tuple[Credentials, Optional[int]]:
"""
Authenticate with AWS Role
@@ -584,11 +616,11 @@ class BaseAWSLLM:
# Check if we need to do cross-account role assumption
if aws_role_name != irsa_role_arn:
sts_response = self._handle_irsa_cross_account(
irsa_role_arn, aws_role_name, aws_session_name, region, web_identity_token_file
irsa_role_arn, aws_role_name, aws_session_name, region, web_identity_token_file, aws_external_id
)
else:
sts_response = self._handle_irsa_same_account(
aws_role_name, aws_session_name, region
aws_role_name, aws_session_name, region, aws_external_id
)
return self._extract_credentials_and_ttl(sts_response)
@@ -619,9 +651,16 @@ class BaseAWSLLM:
aws_session_token=aws_session_token,
)
sts_response = sts_client.assume_role(
RoleArn=aws_role_name, RoleSessionName=aws_session_name
)
assume_role_params = {
"RoleArn": aws_role_name,
"RoleSessionName": aws_session_name
}
# Add ExternalId parameter if provided
if aws_external_id is not None:
assume_role_params["ExternalId"] = aws_external_id
sts_response = sts_client.assume_role(**assume_role_params)
# Extract the credentials from the response and convert to Session Credentials
sts_credentials = sts_response["Credentials"]
@@ -800,6 +839,7 @@ class BaseAWSLLM:
aws_bedrock_runtime_endpoint = optional_params.pop(
"aws_bedrock_runtime_endpoint", None
) # https://bedrock-runtime.{region_name}.amazonaws.com
aws_external_id = optional_params.pop("aws_external_id", None)
credentials: Credentials = self.get_credentials(
aws_access_key_id=aws_access_key_id,
@@ -811,6 +851,7 @@ class BaseAWSLLM:
aws_role_name=aws_role_name,
aws_web_identity_token=aws_web_identity_token,
aws_sts_endpoint=aws_sts_endpoint,
aws_external_id=aws_external_id,
)
return Boto3CredentialsInfo(
@@ -915,6 +956,7 @@ class BaseAWSLLM:
aws_profile_name = optional_params.get("aws_profile_name", None)
aws_web_identity_token = optional_params.get("aws_web_identity_token", None)
aws_sts_endpoint = optional_params.get("aws_sts_endpoint", None)
aws_external_id = optional_params.get("aws_external_id", None)
aws_region_name = self._get_aws_region_name(
optional_params=optional_params, model=model
)
@@ -929,6 +971,7 @@ class BaseAWSLLM:
aws_role_name=aws_role_name,
aws_web_identity_token=aws_web_identity_token,
aws_sts_endpoint=aws_sts_endpoint,
aws_external_id=aws_external_id,
)
sigv4 = SigV4Auth(credentials, service_name, aws_region_name)
litellm/llms/bedrock/chat/converse_handler.py
+2
View File
@@ -307,6 +307,7 @@ class BedrockConverseLLM(BaseAWSLLM):
) # https://bedrock-runtime.{region_name}.amazonaws.com
aws_web_identity_token = optional_params.pop("aws_web_identity_token", None)
aws_sts_endpoint = optional_params.pop("aws_sts_endpoint", None)
aws_external_id = optional_params.pop("aws_external_id", None)
optional_params.pop("aws_region_name", None)
litellm_params[
@@ -323,6 +324,7 @@ class BedrockConverseLLM(BaseAWSLLM):
aws_role_name=aws_role_name,
aws_web_identity_token=aws_web_identity_token,
aws_sts_endpoint=aws_sts_endpoint,
aws_external_id=aws_external_id,
)
### SET RUNTIME ENDPOINT ###
tests/llm_translation/test_bedrock_dynamic_auth_params_unit_tests.py
+1
View File
@@ -207,6 +207,7 @@ class DummyCredentials:
("aws_role_name", "dummy_role_name"),
("aws_web_identity_token", "dummy_web_identity_token"),
("aws_sts_endpoint", "dummy_sts_endpoint"),
("aws_external_id", "dummy_external_id"),
],
)
def test_dynamic_aws_params_propagation(model, param_name, param_value):
tests/test_litellm/llms/bedrock/test_base_aws_llm.py
+141 -4
View File
@@ -1026,7 +1026,7 @@ def test_auth_with_aws_role_irsa_environment():
def test_auth_with_aws_role_same_role_irsa():
"""Test that when IRSA role matches the requested role, we skip assumption"""
base_llm = BaseAWSLLM()
# Set IRSA environment variables
with patch.dict(os.environ, {
'AWS_ROLE_ARN': 'arn:aws:iam::111111111111:role/LitellmRole',
@@ -1037,7 +1037,7 @@ def test_auth_with_aws_role_same_role_irsa():
mock_creds.access_key = 'irsa-access-key'
mock_creds.secret_key = 'irsa-secret-key'
mock_creds.token = 'irsa-session-token'
with patch.object(base_llm, '_auth_with_env_vars', return_value=(mock_creds, None)) as mock_env_auth:
# Call get_credentials instead of _auth_with_aws_role directly
# This tests the full flow
@@ -1048,9 +1048,146 @@ def test_auth_with_aws_role_same_role_irsa():
aws_session_name='test-session',
aws_region_name='us-east-1'
)
# Verify it used the env vars auth (no role assumption)
mock_env_auth.assert_called_once()
# Verify the returned credentials
assert creds.access_key == 'irsa-access-key'
def test_assume_role_with_external_id():
"""Test that assume_role STS call includes ExternalId parameter when provided"""
base_aws_llm = BaseAWSLLM()
# Mock the boto3 STS client
mock_sts_client = MagicMock()
mock_expiry = datetime.now(timezone.utc) + timedelta(hours=1)
mock_sts_response = {
"Credentials": {
"AccessKeyId": "test-access-key",
"SecretAccessKey": "test-secret-key",
"SessionToken": "test-session-token",
"Expiration": mock_expiry,
}
}
mock_sts_client.assume_role.return_value = mock_sts_response
with patch("boto3.client", return_value=mock_sts_client):
# Call _auth_with_aws_role with external ID
credentials, ttl = base_aws_llm._auth_with_aws_role(
aws_access_key_id=None,
aws_secret_access_key=None,
aws_session_token=None,
aws_role_name="arn:aws:iam::123456789012:role/ExampleRole",
aws_session_name="test-session",
aws_external_id="UniqueExternalID123"
)
# Verify assume_role was called with ExternalId
mock_sts_client.assume_role.assert_called_once_with(
RoleArn="arn:aws:iam::123456789012:role/ExampleRole",
RoleSessionName="test-session",
ExternalId="UniqueExternalID123"
)
def test_assume_role_without_external_id():
"""Test that assume_role STS call excludes ExternalId parameter when not provided"""
base_aws_llm = BaseAWSLLM()
# Mock the boto3 STS client
mock_sts_client = MagicMock()
mock_expiry = datetime.now(timezone.utc) + timedelta(hours=1)
mock_sts_response = {
"Credentials": {
"AccessKeyId": "test-access-key",
"SecretAccessKey": "test-secret-key",
"SessionToken": "test-session-token",
"Expiration": mock_expiry,
}
}
mock_sts_client.assume_role.return_value = mock_sts_response
with patch("boto3.client", return_value=mock_sts_client):
# Call _auth_with_aws_role without external ID
credentials, ttl = base_aws_llm._auth_with_aws_role(
aws_access_key_id=None,
aws_secret_access_key=None,
aws_session_token=None,
aws_role_name="arn:aws:iam::123456789012:role/ExampleRole",
aws_session_name="test-session"
)
# Verify assume_role was called without ExternalId
mock_sts_client.assume_role.assert_called_once_with(
RoleArn="arn:aws:iam::123456789012:role/ExampleRole",
RoleSessionName="test-session"
)
def test_converse_handler_external_id_extraction():
"""Test that BedrockConverseLLM properly extracts and passes aws_external_id parameter"""
from litellm.llms.bedrock.chat.converse_handler import BedrockConverseLLM
converse_llm = BedrockConverseLLM()
# Mock get_credentials to capture parameters
def mock_get_credentials(**kwargs):
mock_get_credentials.called_kwargs = kwargs
mock_credentials = MagicMock()
mock_credentials.access_key = "test-access-key"
mock_credentials.secret_key = "test-secret-key"
mock_credentials.token = "test-session-token"
return mock_credentials
with patch.object(converse_llm, 'get_credentials', side_effect=mock_get_credentials):
with patch.object(converse_llm, '_get_aws_region_name', return_value="us-west-2"):
with patch.object(converse_llm, 'get_runtime_endpoint', return_value=("https://test", "https://test")):
with patch('litellm.AmazonConverseConfig') as mock_config:
mock_config.return_value._transform_request.return_value = {"test": "data"}
with patch.object(converse_llm, 'get_request_headers') as mock_headers:
mock_headers.return_value = MagicMock()
mock_headers.return_value.headers = {"Authorization": "test"}
with patch('litellm.llms.custom_httpx.http_handler._get_httpx_client') as mock_client:
mock_http_client = MagicMock()
mock_response = MagicMock()
mock_response.raise_for_status.return_value = None
mock_http_client.post.return_value = mock_response
mock_client.return_value = mock_http_client
# Mock the transform_response method
mock_config.return_value._transform_response.return_value = MagicMock()
# Call completion with aws_external_id in optional_params
optional_params = {
"aws_role_name": "arn:aws:iam::123456789012:role/ExampleRole",
"aws_session_name": "test-session",
"aws_external_id": "TestExternalID123"
}
try:
converse_llm.completion(
model="anthropic.claude-3-sonnet-20240229-v1:0",
messages=[{"role": "user", "content": "Hello"}],
api_base=None,
custom_prompt_dict={},
model_response=MagicMock(),
encoding="utf-8",
logging_obj=MagicMock(),
optional_params=optional_params,
acompletion=False,
timeout=None,
litellm_params={}
)
except Exception:
# We expect this to fail due to mocking, but that's OK
# We just want to verify the parameter extraction
pass
# Verify aws_external_id was extracted and passed to get_credentials
assert hasattr(mock_get_credentials, 'called_kwargs')
assert "aws_external_id" in mock_get_credentials.called_kwargs
assert mock_get_credentials.called_kwargs["aws_external_id"] == "TestExternalID123"