diff --git a/litellm/llms/bedrock/base_aws_llm.py b/litellm/llms/bedrock/base_aws_llm.py index ce196757f9..0ddf8896fd 100644 --- a/litellm/llms/bedrock/base_aws_llm.py +++ b/litellm/llms/bedrock/base_aws_llm.py @@ -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) diff --git a/litellm/llms/bedrock/chat/converse_handler.py b/litellm/llms/bedrock/chat/converse_handler.py index 15a5002f0e..54c603e596 100644 --- a/litellm/llms/bedrock/chat/converse_handler.py +++ b/litellm/llms/bedrock/chat/converse_handler.py @@ -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 ### diff --git a/tests/llm_translation/test_bedrock_dynamic_auth_params_unit_tests.py b/tests/llm_translation/test_bedrock_dynamic_auth_params_unit_tests.py index 7220ffbb2c..06a3086857 100644 --- a/tests/llm_translation/test_bedrock_dynamic_auth_params_unit_tests.py +++ b/tests/llm_translation/test_bedrock_dynamic_auth_params_unit_tests.py @@ -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): diff --git a/tests/test_litellm/llms/bedrock/test_base_aws_llm.py b/tests/test_litellm/llms/bedrock/test_base_aws_llm.py index 5effa6fa01..f5856cd12d 100644 --- a/tests/test_litellm/llms/bedrock/test_base_aws_llm.py +++ b/tests/test_litellm/llms/bedrock/test_base_aws_llm.py @@ -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"