diff --git a/mem0/configs/embeddings/base.py b/mem0/configs/embeddings/base.py index a5a174b7b..0737088fc 100644 --- a/mem0/configs/embeddings/base.py +++ b/mem0/configs/embeddings/base.py @@ -1,3 +1,4 @@ +import os from abc import ABC from typing import Dict, Optional, Union @@ -38,7 +39,7 @@ class BaseEmbedderConfig(ABC): # AWS Bedrock specific aws_access_key_id: Optional[str] = None, aws_secret_access_key: Optional[str] = None, - aws_region: Optional[str] = "us-west-2", + aws_region: Optional[str] = None, ): """ Initializes a configuration class instance for the Embeddings. @@ -105,4 +106,5 @@ 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_region = aws_region + self.aws_region = aws_region or os.environ.get("AWS_REGION") or "us-west-2" + diff --git a/mem0/configs/llms/aws_bedrock.py b/mem0/configs/llms/aws_bedrock.py index bbdebef0a..a285f9074 100644 --- a/mem0/configs/llms/aws_bedrock.py +++ b/mem0/configs/llms/aws_bedrock.py @@ -1,6 +1,7 @@ -from typing import Optional, Dict, Any, List -from mem0.configs.llms.base import BaseLlmConfig import os +from typing import Any, Dict, List, Optional + +from mem0.configs.llms.base import BaseLlmConfig class AWSBedrockConfig(BaseLlmConfig): @@ -19,7 +20,7 @@ class AWSBedrockConfig(BaseLlmConfig): top_k: int = 1, aws_access_key_id: Optional[str] = None, aws_secret_access_key: Optional[str] = None, - aws_region: str = "us-west-2", + aws_region: str = "", aws_session_token: Optional[str] = None, aws_profile: Optional[str] = None, model_kwargs: Optional[Dict[str, Any]] = None, @@ -53,7 +54,7 @@ class AWSBedrockConfig(BaseLlmConfig): self.aws_access_key_id = aws_access_key_id self.aws_secret_access_key = aws_secret_access_key - self.aws_region = aws_region + self.aws_region = aws_region or os.getenv("AWS_REGION", "us-west-2") self.aws_session_token = aws_session_token self.aws_profile = aws_profile self.model_kwargs = model_kwargs or {} diff --git a/mem0/embeddings/aws_bedrock.py b/mem0/embeddings/aws_bedrock.py index d5ce21db6..5c3c1acd1 100644 --- a/mem0/embeddings/aws_bedrock.py +++ b/mem0/embeddings/aws_bedrock.py @@ -28,15 +28,15 @@ class AWSBedrockEmbedding(EmbeddingBase): aws_access_key = os.environ.get("AWS_ACCESS_KEY_ID", "") aws_secret_key = os.environ.get("AWS_SECRET_ACCESS_KEY", "") aws_session_token = os.environ.get("AWS_SESSION_TOKEN", "") - aws_region = os.environ.get("AWS_REGION", "us-west-2") # Check if AWS config is provided in the config if hasattr(self.config, "aws_access_key_id"): 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 - if hasattr(self.config, "aws_region"): - aws_region = self.config.aws_region + + # AWS region is always set in config - see BaseEmbedderConfig + aws_region = self.config.aws_region or "us-west-2" self.client = boto3.client( "bedrock-runtime",