diff --git a/mem0/configs/embeddings/base.py b/mem0/configs/embeddings/base.py index 05c7b5a03..3882b4b2e 100644 --- a/mem0/configs/embeddings/base.py +++ b/mem0/configs/embeddings/base.py @@ -39,6 +39,7 @@ class BaseEmbedderConfig(ABC): # AWS Bedrock specific aws_access_key_id: Optional[str] = None, aws_secret_access_key: Optional[str] = None, + aws_session_token: Optional[str] = None, aws_region: Optional[str] = None, ): """ @@ -106,5 +107,6 @@ class BaseEmbedderConfig(ABC): # AWS Bedrock specific self.aws_access_key_id = aws_access_key_id self.aws_secret_access_key = aws_secret_access_key + self.aws_session_token = aws_session_token self.aws_region = aws_region or os.environ.get("AWS_REGION") or "us-west-2" diff --git a/mem0/embeddings/aws_bedrock.py b/mem0/embeddings/aws_bedrock.py index 5c3c1acd1..a6fe3c224 100644 --- a/mem0/embeddings/aws_bedrock.py +++ b/mem0/embeddings/aws_bedrock.py @@ -34,7 +34,12 @@ class AWSBedrockEmbedding(EmbeddingBase): aws_access_key = self.config.aws_access_key_id if hasattr(self.config, "aws_secret_access_key"): aws_secret_key = self.config.aws_secret_access_key - + # Honor a session token supplied via config (temporary credentials from + # STS / assume-role), falling back to the env var when unset. The LLM + # Bedrock provider already supports this; mirror it here. + if getattr(self.config, "aws_session_token", None): + aws_session_token = self.config.aws_session_token + # AWS region is always set in config - see BaseEmbedderConfig aws_region = self.config.aws_region or "us-west-2" diff --git a/tests/embeddings/test_aws_bedrock_embeddings.py b/tests/embeddings/test_aws_bedrock_embeddings.py new file mode 100644 index 000000000..b1924d43f --- /dev/null +++ b/tests/embeddings/test_aws_bedrock_embeddings.py @@ -0,0 +1,60 @@ +from unittest.mock import Mock, patch + +import pytest + +from mem0.configs.embeddings.base import BaseEmbedderConfig +from mem0.embeddings.aws_bedrock import AWSBedrockEmbedding + + +@pytest.fixture +def mock_boto3_client(): + with patch("mem0.embeddings.aws_bedrock.boto3.client") as mock_client: + mock_client.return_value = Mock() + yield mock_client + + +def test_session_token_from_config_is_passed_to_client(mock_boto3_client): + """A session token supplied via config must reach the bedrock-runtime client. + + Temporary credentials (STS / assume-role) require a session token; the LLM + Bedrock provider already honors it, so the embedding provider should too. + """ + with patch("mem0.embeddings.aws_bedrock.os.environ", {}): + config = BaseEmbedderConfig( + model="amazon.titan-embed-text-v2:0", + aws_access_key_id="AKIA_TEST", + aws_secret_access_key="SECRET_TEST", + aws_session_token="SESSION_TOKEN_TEST", + aws_region="eu-central-1", + ) + AWSBedrockEmbedding(config) + + _, kwargs = mock_boto3_client.call_args + assert kwargs["aws_session_token"] == "SESSION_TOKEN_TEST" + assert kwargs["aws_access_key_id"] == "AKIA_TEST" + assert kwargs["aws_secret_access_key"] == "SECRET_TEST" + assert kwargs["region_name"] == "eu-central-1" + + +def test_session_token_falls_back_to_env(mock_boto3_client): + """When config doesn't set a session token, the env var is still used.""" + with patch.dict( + "mem0.embeddings.aws_bedrock.os.environ", + {"AWS_SESSION_TOKEN": "ENV_SESSION_TOKEN"}, + clear=True, + ): + config = BaseEmbedderConfig(model="amazon.titan-embed-text-v2:0") + AWSBedrockEmbedding(config) + + _, kwargs = mock_boto3_client.call_args + assert kwargs["aws_session_token"] == "ENV_SESSION_TOKEN" + + +def test_no_session_token_passes_none(mock_boto3_client): + """Without a token in config or env, the client receives None (unchanged).""" + with patch("mem0.embeddings.aws_bedrock.os.environ", {}): + config = BaseEmbedderConfig(model="amazon.titan-embed-text-v2:0") + AWSBedrockEmbedding(config) + + _, kwargs = mock_boto3_client.call_args + assert kwargs["aws_session_token"] is None