diff --git a/docs/components/embedders/models/aws_bedrock.mdx b/docs/components/embedders/models/aws_bedrock.mdx new file mode 100644 index 00000000..470963ef --- /dev/null +++ b/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 + + +```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") +``` + + +### Config + +Here are the parameters available for configuring AWS Bedrock embedder: + + + +| Parameter | Description | Default Value | +| --- | --- | --- | +| `model` | The name of the embedding model to use | `amazon.titan-embed-text-v1` | + + diff --git a/docs/components/embedders/overview.mdx b/docs/components/embedders/overview.mdx index 3ae4d592..22f7308e 100644 --- a/docs/components/embedders/overview.mdx +++ b/docs/components/embedders/overview.mdx @@ -26,6 +26,7 @@ See the list of supported embedders below. + ## Usage diff --git a/docs/docs.json b/docs/docs.json index edccbca8..58880951 100644 --- a/docs/docs.json +++ b/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" ] } ] diff --git a/mem0/embeddings/aws_bedrock.py b/mem0/embeddings/aws_bedrock.py new file mode 100644 index 00000000..2fcf6df3 --- /dev/null +++ b/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) diff --git a/mem0/embeddings/configs.py b/mem0/embeddings/configs.py index 9b1be04a..b4fadd68 100644 --- a/mem0/embeddings/configs.py +++ b/mem0/embeddings/configs.py @@ -23,6 +23,7 @@ class EmbedderConfig(BaseModel): "together", "lmstudio", "langchain", + "aws_bedrock", ]: return v else: diff --git a/mem0/utils/factory.py b/mem0/utils/factory.py index 03fc0d2c..ce75cb7e 100644 --- a/mem0/utils/factory.py +++ b/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 diff --git a/pyproject.toml b/pyproject.toml index 658e118c..d161670f 100644 --- a/pyproject.toml +++ b/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 "] exclude = [