diff --git a/litellm/llms/bedrock/embed/embedding.py b/litellm/llms/bedrock/embed/embedding.py index 27dc785bf57..f58a54e5054 100644 --- a/litellm/llms/bedrock/embed/embedding.py +++ b/litellm/llms/bedrock/embed/embedding.py @@ -57,6 +57,7 @@ class BedrockEmbedding(BaseAWSLLM): aws_profile_name = optional_params.pop("aws_profile_name", None) 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) ### SET REGION NAME ### if aws_region_name is None: @@ -87,6 +88,7 @@ class BedrockEmbedding(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 credentials, aws_region_name diff --git a/tests/test_litellm/llms/bedrock/embed/test_bedrock_embedding.py b/tests/test_litellm/llms/bedrock/embed/test_bedrock_embedding.py index c67a8712340..309b709770e 100644 --- a/tests/test_litellm/llms/bedrock/embed/test_bedrock_embedding.py +++ b/tests/test_litellm/llms/bedrock/embed/test_bedrock_embedding.py @@ -957,3 +957,44 @@ def test_titan_image_embedding_cost_uses_per_image_rate(): assert response.usage is not None assert response.usage.prompt_tokens_details is not None assert response.usage.prompt_tokens_details.image_count == 1 + + +def test_bedrock_embedding_with_aws_external_id(): + """Test that aws_external_id is correctly passed to get_credentials for Bedrock embeddings""" + litellm.set_verbose = True + client = HTTPHandler() + test_api_key = "test-bearer-token-12345" + model = "bedrock/amazon.titan-embed-text-v1" + embed_response = {"embedding": [0.1, 0.2, 0.3], "inputTextTokenCount": 10} + + with patch.object(client, "post") as mock_post: + mock_response = Mock() + mock_response.status_code = 200 + mock_response.text = json.dumps(embed_response) + mock_response.json = lambda: json.loads(mock_response.text) + mock_post.return_value = mock_response + + with patch( + "litellm.llms.bedrock.embed.embedding.BedrockEmbedding.get_credentials" + ) as mock_get_creds: + mock_get_creds.return_value = (Mock(), "us-east-1") + + response = litellm.embedding( + model=model, + input=test_input, + client=client, + aws_region_name="us-east-1", + aws_bedrock_runtime_endpoint="https://bedrock-runtime.us-east-1.amazonaws.com", + api_key=test_api_key, + aws_role_name="arn:aws:iam::123456789012:role/test-role", + aws_external_id="test-external-id", + ) + + assert isinstance(response, litellm.EmbeddingResponse) + mock_get_creds.assert_called_once() + _, kwargs = mock_get_creds.call_args + assert kwargs.get("aws_external_id") == "test-external-id" + assert ( + kwargs.get("aws_role_name") + == "arn:aws:iam::123456789012:role/test-role" + )