Support for AWS Bedrock Embeddings (#2660)

docs/components/embedders/models/aws_bedrock.mdx

@@ -0,0 +1,62 @@
+---
+title: AWS Bedrock
+---
+
+To use AWS Bedrock embedding models, you need to have the appropriate AWS credentials and permissions. The embeddings implementation relies on the `boto3` library.
+
+### Setup
+- Ensure you have model access from the [AWS Bedrock Console](https://us-east-1.console.aws.amazon.com/bedrock/home?region=us-east-1#/modelaccess)
+- Authenticate the boto3 client using a method described in the [AWS documentation](https://boto3.amazonaws.com/v1/documentation/api/latest/guide/credentials.html)
+- Set up environment variables for authentication:
+  ```bash
+  export AWS_REGION=us-east-1
+  export AWS_ACCESS_KEY_ID=your-access-key
+  export AWS_SECRET_ACCESS_KEY=your-secret-key
+  ```
+
+### Usage
+
+<CodeGroup>
+```python Python
+import os
+from mem0 import Memory
+
+# For LLM if needed
+os.environ["OPENAI_API_KEY"] = "your-openai-api-key"
+
+# AWS credentials
+os.environ["AWS_REGION"] = "us-east-1"
+os.environ["AWS_ACCESS_KEY_ID"] = "your-access-key"
+os.environ["AWS_SECRET_ACCESS_KEY"] = "your-secret-key"
+
+config = {
+    "embedder": {
+        "provider": "aws_bedrock",
+        "config": {
+            "model": "amazon.titan-embed-text-v1"
+        }
+    }
+}
+
+m = Memory.from_config(config)
+messages = [
+    {"role": "user", "content": "I'm planning to watch a movie tonight. Any recommendations?"},
+    {"role": "assistant", "content": "How about a thriller movies? They can be quite engaging."},
+    {"role": "user", "content": "I'm not a big fan of thriller movies but I love sci-fi movies."},
+    {"role": "assistant", "content": "Got it! I'll avoid thriller recommendations and suggest sci-fi movies in the future."}
+]
+m.add(messages, user_id="alice")
+```
+</CodeGroup>
+
+### Config
+
+Here are the parameters available for configuring AWS Bedrock embedder:
+
+<Tabs>
+<Tab title="Python">
+| Parameter | Description | Default Value |
+| --- | --- | --- |
+| `model` | The name of the embedding model to use | `amazon.titan-embed-text-v1` |
+</Tab>
+</Tabs>

docs/components/embedders/overview.mdx

@@ -26,6 +26,7 @@ See the list of supported embedders below.
   <Card title="Together" href="/components/embedders/models/together"></Card>
   <Card title="LM Studio" href="/components/embedders/models/lmstudio"></Card>
   <Card title="Langchain" href="/components/embedders/models/langchain"></Card>
+  <Card title="AWS Bedrock" href="/components/embedders/models/aws_bedrock"></Card>
 </CardGroup>
 
 ## Usage

docs/docs.json

@@ -166,7 +166,8 @@
                           "components/embedders/models/gemini",
                           "components/embedders/models/lmstudio",
                           "components/embedders/models/together",
-                          "components/embedders/models/langchain"
+                          "components/embedders/models/langchain",
+                          "components/embedders/models/aws_bedrock"
                         ]
                       }
                     ]

mem0/embeddings/aws_bedrock.py

@@ -0,0 +1,78 @@
+import json
+from typing import Literal, Optional
+
+try:
+    import boto3
+except ImportError:
+    raise ImportError("The 'boto3' library is required. Please install it using 'pip install boto3'.")
+
+import numpy as np
+
+from mem0.configs.embeddings.base import BaseEmbedderConfig
+from mem0.embeddings.base import EmbeddingBase
+
+
+class AWSBedrockEmbedding(EmbeddingBase):
+    """AWS Bedrock embedding implementation.
+
+    This class uses AWS Bedrock's embedding models.
+    """
+
+    def __init__(self, config: Optional[BaseEmbedderConfig] = None):
+        super().__init__(config)
+
+        self.config.model = self.config.model or "amazon.titan-embed-text-v1"
+        self.client = boto3.client("bedrock-runtime")
+
+    def _normalize_vector(self, embeddings):
+        """Normalize the embedding to a unit vector."""
+        emb = np.array(embeddings)
+        norm_emb = emb / np.linalg.norm(emb)
+        return norm_emb.tolist()
+
+    def _get_embedding(self, text):
+        """Call out to Bedrock embedding endpoint."""
+
+        # Format input body based on the provider
+        provider = self.config.model.split(".")[0]
+        input_body = {}
+
+        if provider == "cohere":
+            input_body["input_type"] = "search_document"
+            input_body["texts"] = [text]
+        else:
+            # Amazon and other providers
+            input_body["inputText"] = text
+
+        body = json.dumps(input_body)
+
+        try:
+            response = self.client.invoke_model(
+                body=body,
+                modelId=self.config.model,
+                accept="application/json",
+                contentType="application/json",
+            )
+
+            response_body = json.loads(response.get("body").read())
+
+            if provider == "cohere":
+                embeddings = response_body.get("embeddings")[0]
+            else:
+                embeddings = response_body.get("embedding")
+
+            return embeddings
+        except Exception as e:
+            raise ValueError(f"Error getting embedding from AWS Bedrock: {e}")
+
+    def embed(self, text, memory_action: Optional[Literal["add", "search", "update"]] = None):
+        """
+        Get the embedding for the given text using AWS Bedrock.
+
+        Args:
+            text (str): The text to embed.
+            memory_action (optional): The type of embedding to use. Must be one of "add", "search", or "update". Defaults to None.
+        Returns:
+            list: The embedding vector.
+        """
+        return self._get_embedding(text)

mem0/embeddings/configs.py

@@ -23,6 +23,7 @@ class EmbedderConfig(BaseModel):
             "together",
             "lmstudio",
             "langchain",
+            "aws_bedrock",
         ]:
             return v
         else:

mem0/utils/factory.py

@@ -53,6 +53,7 @@ class EmbedderFactory:
         "together": "mem0.embeddings.together.TogetherEmbedding",
         "lmstudio": "mem0.embeddings.lmstudio.LMStudioEmbedding",
         "langchain": "mem0.embeddings.langchain.LangchainEmbedding",
+        "aws_bedrock": "mem0.embeddings.aws_bedrock.AWSBedrockEmbedding",
     }
 
     @classmethod

pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "mem0ai"
-version = "0.1.97"
+version = "0.1.98"
 description = "Long-term memory for AI Agents"
 authors = ["Mem0 <founders@mem0.ai>"]
 exclude = [