Add Amazon S3 Vectors Support (#3237)

LLM.md

@@ -304,7 +304,7 @@ config = MemoryConfig(
 - **langchain** - LangChain embeddings
 - **aws_bedrock** - AWS Bedrock embeddings
 
-#### Vector Store Providers (17 supported)
+#### Vector Store Providers (19 supported)
 - **qdrant** - Qdrant vector database (default)
 - **chroma** - ChromaDB
 - **pinecone** - Pinecone vector database
@@ -322,9 +322,10 @@ config = MemoryConfig(
 - **supabase** - Supabase vector
 - **baidu** - Baidu vector database
 - **langchain** - LangChain vector stores
+- **s3_vectors** - Amazon S3 Vectors
 - **databricks** - Databricks vector stores
 
-#### Graph Store Providers (3 supported)
+#### Graph Store Providers (4 supported)
 - **neo4j** - Neo4j graph database
 - **memgraph** - Memgraph
 - **neptune** - AWS Neptune Analytics

docs/components/vectordbs/dbs/s3_vectors.mdx

@@ -0,0 +1,78 @@
+---
+title: Amazon S3 Vectors
+---
+
+[Amazon S3 Vectors](https://aws.amazon.com/s3/features/vectors/) is a purpose-built, cost-optimized vector storage and query service for semantic search and AI applications. It provides S3-level elasticity and durability with sub-second query performance.
+
+### Installation
+
+S3 Vectors support requires additional dependencies. Install them with:
+
+```bash
+pip install boto3
+```
+
+### Usage
+
+To use Amazon S3 Vectors with Mem0, you need to have an AWS account and the necessary IAM permissions (`s3vectors:*`). Ensure your environment is configured with AWS credentials (e.g., via `~/.aws/credentials` or environment variables).
+
+```python
+import os
+from mem0 import Memory
+
+# Ensure your AWS credentials are configured in your environment
+# e.g., by setting AWS_ACCESS_KEY_ID, AWS_SECRET_ACCESS_KEY, and AWS_DEFAULT_REGION
+
+config = {
+    "vector_store": {
+        "provider": "s3_vectors",
+        "config": {
+            "vector_bucket_name": "my-mem0-vector-bucket",
+            "index_name": "my-memories-index",
+            "embedding_model_dims": 1536,
+            "distance_metric": "cosine",
+            "region_name": "us-east-1"
+        }
+    }
+}
+
+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 movie? 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", metadata={"category": "movies"})
+```
+
+### Config
+
+Here are the available parameters for the `s3_vectors` config:
+
+| Parameter              | Description                                                          | Default Value |
+| ---------------------- | -------------------------------------------------------------------- | ------------- |
+| `vector_bucket_name`   | The name of the S3 Vector bucket to use. It will be created if it doesn't exist. | Required      |
+| `index_name`           | The name of the vector index within the bucket.                        | `mem0`        |
+| `embedding_model_dims` | Dimensions of the embedding model. Must match your embedder.         | `1536`        |
+| `distance_metric`      | Distance metric for similarity search. Options: `cosine`, `euclidean`. | `cosine`      |
+| `region_name`          | The AWS region where the bucket and index reside.                    | `None` (uses default from AWS config) |
+
+### IAM Permissions
+
+Your AWS identity (user or role) needs permissions to perform actions on S3 Vectors. A minimal policy would look like this:
+
+```json
+{
+    "Version": "2012-10-17",
+    "Statement": [
+        {
+            "Effect": "Allow",
+            "Action": "s3vectors:*",
+            "Resource": "*"
+        }
+    ]
+}
+```
+
+For production, it is recommended to scope down the resource ARN to your specific buckets and indexes.

docs/components/vectordbs/overview.mdx

@@ -31,6 +31,7 @@ See the list of supported vector databases below.
   <Card title="Weaviate" href="/components/vectordbs/dbs/weaviate"></Card>
   <Card title="FAISS" href="/components/vectordbs/dbs/faiss"></Card>
   <Card title="LangChain" href="/components/vectordbs/dbs/langchain"></Card>
+  <Card title="Amazon S3 Vectors" href="/components/vectordbs/dbs/s3_vectors"></Card>
   <Card title="Databricks" href="/components/vectordbs/dbs/databricks"></Card>
 </CardGroup>
 

docs/docs.json

@@ -161,6 +161,7 @@
                           "components/vectordbs/dbs/faiss",
                           "components/vectordbs/dbs/langchain",
                           "components/vectordbs/dbs/baidu",
+                          "components/vectordbs/dbs/s3_vectors",
                           "components/vectordbs/dbs/databricks"
                         ]
                       }

docs/llms.txt

@@ -93,6 +93,7 @@ Key differentiators:
 - [Databricks](https://docs.mem0.ai/components/vectordbs/dbs/databricks): Delta Lake integration for vector search
 - [Baidu](https://docs.mem0.ai/components/vectordbs/dbs/baidu): Baidu vector database integration
 - [LangChain](https://docs.mem0.ai/components/vectordbs/dbs/langchain): LangChain vector store integration
+- [S3 Vectors](https://docs.mem0.ai/components/vectordbs/dbs/s3_vectors): Amazon S3 Vectors integration
 
 ### Supported Embeddings
 

mem0/configs/vector_stores/s3_vectors.py

@@ -0,0 +1,34 @@
+from typing import Any, Dict, Optional
+
+from pydantic import BaseModel, Field, model_validator
+
+
+class S3VectorsConfig(BaseModel):
+    vector_bucket_name: str = Field(description="Name of the S3 Vector bucket")
+    index_name: str = Field("mem0", description="Name of the vector index")
+    embedding_model_dims: int = Field(
+        1536, description="Dimension of the embedding vector"
+    )
+    distance_metric: str = Field(
+        "cosine",
+        description="Distance metric for similarity search. Options: 'cosine', 'euclidean'",
+    )
+    region_name: Optional[str] = Field(
+        None, description="AWS region for the S3 Vectors client"
+    )
+
+    @model_validator(mode="before")
+    @classmethod
+    def validate_extra_fields(cls, values: Dict[str, Any]) -> Dict[str, Any]:
+        allowed_fields = set(cls.model_fields.keys())
+        input_fields = set(values.keys())
+        extra_fields = input_fields - allowed_fields
+        if extra_fields:
+            raise ValueError(
+                f"Extra fields not allowed: {', '.join(extra_fields)}. Please input only the following fields: {', '.join(allowed_fields)}"
+            )
+        return values
+
+    model_config = {
+        "arbitrary_types_allowed": True,
+    }

mem0/utils/factory.py

@@ -173,6 +173,7 @@ class VectorStoreFactory:
         "weaviate": "mem0.vector_stores.weaviate.Weaviate",
         "faiss": "mem0.vector_stores.faiss.FAISS",
         "langchain": "mem0.vector_stores.langchain.Langchain",
+        "s3_vectors": "mem0.vector_stores.s3_vectors.S3Vectors",
     }
 
     @classmethod

mem0/vector_stores/configs.py

@@ -29,6 +29,7 @@ class VectorStoreConfig(BaseModel):
         "weaviate": "WeaviateConfig",
         "faiss": "FAISSConfig",
         "langchain": "LangchainConfig",
+        "s3_vectors": "S3VectorsConfig",
     }
 
     @model_validator(mode="after")

mem0/vector_stores/s3_vectors.py

@@ -0,0 +1,176 @@
+import json
+import logging
+from typing import Dict, List, Optional
+
+from pydantic import BaseModel
+
+from mem0.vector_stores.base import VectorStoreBase
+
+try:
+    import boto3
+    from botocore.exceptions import ClientError
+except ImportError:
+    raise ImportError("The 'boto3' library is required. Please install it using 'pip install boto3'.")
+
+logger = logging.getLogger(__name__)
+
+
+class OutputData(BaseModel):
+    id: Optional[str]
+    score: Optional[float]
+    payload: Optional[Dict]
+
+
+class S3Vectors(VectorStoreBase):
+    def __init__(
+        self,
+        vector_bucket_name: str,
+        index_name: str,
+        embedding_model_dims: int,
+        distance_metric: str = "cosine",
+        region_name: Optional[str] = None,
+    ):
+        self.client = boto3.client("s3vectors", region_name=region_name)
+        self.vector_bucket_name = vector_bucket_name
+        self.collection_name = index_name
+        self.embedding_model_dims = embedding_model_dims
+        self.distance_metric = distance_metric
+
+        self._ensure_bucket_exists()
+        self.create_col(self.collection_name, self.embedding_model_dims, self.distance_metric)
+
+    def _ensure_bucket_exists(self):
+        try:
+            self.client.get_vector_bucket(vectorBucketName=self.vector_bucket_name)
+            logger.info(f"Vector bucket '{self.vector_bucket_name}' already exists.")
+        except ClientError as e:
+            if e.response["Error"]["Code"] == "NotFoundException":
+                logger.info(f"Vector bucket '{self.vector_bucket_name}' not found. Creating it.")
+                self.client.create_vector_bucket(vectorBucketName=self.vector_bucket_name)
+                logger.info(f"Vector bucket '{self.vector_bucket_name}' created.")
+            else:
+                raise
+
+    def create_col(self, name, vector_size, distance="cosine"):
+        try:
+            self.client.get_index(vectorBucketName=self.vector_bucket_name, indexName=name)
+            logger.info(f"Index '{name}' already exists in bucket '{self.vector_bucket_name}'.")
+        except ClientError as e:
+            if e.response["Error"]["Code"] == "NotFoundException":
+                logger.info(f"Index '{name}' not found in bucket '{self.vector_bucket_name}'. Creating it.")
+                self.client.create_index(
+                    vectorBucketName=self.vector_bucket_name,
+                    indexName=name,
+                    dataType="float32",
+                    dimension=vector_size,
+                    distanceMetric=distance,
+                )
+                logger.info(f"Index '{name}' created.")
+            else:
+                raise
+
+    def _parse_output(self, vectors: List[Dict]) -> List[OutputData]:
+        results = []
+        for v in vectors:
+            payload = v.get("metadata", {})
+            # Boto3 might return metadata as a JSON string
+            if isinstance(payload, str):
+                try:
+                    payload = json.loads(payload)
+                except json.JSONDecodeError:
+                    logger.warning(f"Failed to parse metadata for key {v.get('key')}")
+                    payload = {}
+            results.append(OutputData(id=v.get("key"), score=v.get("distance"), payload=payload))
+        return results
+
+    def insert(self, vectors, payloads=None, ids=None):
+        vectors_to_put = []
+        for i, vec in enumerate(vectors):
+            vectors_to_put.append(
+                {
+                    "key": ids[i],
+                    "data": {"float32": vec},
+                    "metadata": payloads[i] if payloads else {},
+                }
+            )
+        self.client.put_vectors(
+            vectorBucketName=self.vector_bucket_name,
+            indexName=self.collection_name,
+            vectors=vectors_to_put,
+        )
+
+    def search(self, query, vectors, limit=5, filters=None):
+        params = {
+            "vectorBucketName": self.vector_bucket_name,
+            "indexName": self.collection_name,
+            "queryVector": {"float32": vectors},
+            "topK": limit,
+            "returnMetadata": True,
+            "returnDistance": True,
+        }
+        if filters:
+            params["filter"] = filters
+
+        response = self.client.query_vectors(**params)
+        return self._parse_output(response.get("vectors", []))
+
+    def delete(self, vector_id):
+        self.client.delete_vectors(
+            vectorBucketName=self.vector_bucket_name,
+            indexName=self.collection_name,
+            keys=[vector_id],
+        )
+
+    def update(self, vector_id, vector=None, payload=None):
+        # S3 Vectors uses put_vectors for updates (overwrite)
+        self.insert(vectors=[vector], payloads=[payload], ids=[vector_id])
+
+    def get(self, vector_id) -> Optional[OutputData]:
+        response = self.client.get_vectors(
+            vectorBucketName=self.vector_bucket_name,
+            indexName=self.collection_name,
+            keys=[vector_id],
+            returnData=False,
+            returnMetadata=True,
+        )
+        vectors = response.get("vectors", [])
+        if not vectors:
+            return None
+        return self._parse_output(vectors)[0]
+
+    def list_cols(self):
+        response = self.client.list_indexes(vectorBucketName=self.vector_bucket_name)
+        return [idx["indexName"] for idx in response.get("indexes", [])]
+
+    def delete_col(self):
+        self.client.delete_index(vectorBucketName=self.vector_bucket_name, indexName=self.collection_name)
+
+    def col_info(self):
+        response = self.client.get_index(vectorBucketName=self.vector_bucket_name, indexName=self.collection_name)
+        return response.get("index", {})
+
+    def list(self, filters=None, limit=None):
+        # Note: list_vectors does not support metadata filtering.
+        if filters:
+            logger.warning("S3 Vectors `list` does not support metadata filtering. Ignoring filters.")
+
+        params = {
+            "vectorBucketName": self.vector_bucket_name,
+            "indexName": self.collection_name,
+            "returnData": False,
+            "returnMetadata": True,
+        }
+        if limit:
+            params["maxResults"] = limit
+
+        paginator = self.client.get_paginator("list_vectors")
+        pages = paginator.paginate(**params)
+        all_vectors = []
+        for page in pages:
+            all_vectors.extend(page.get("vectors", []))
+        return [self._parse_output(all_vectors)]
+
+    def reset(self):
+        logger.warning(f"Resetting index {self.collection_name}...")
+        self.delete_col()
+        self.create_col(self.collection_name, self.embedding_model_dims, self.distance_metric)

tests/vector_stores/test_s3_vectors.py

@@ -0,0 +1,165 @@
+import pytest
+from botocore.exceptions import ClientError
+
+from mem0.vector_stores.s3_vectors import S3Vectors
+
+BUCKET_NAME = "test-bucket"
+INDEX_NAME = "test-index"
+EMBEDDING_DIMS = 1536
+REGION = "us-east-1"
+
+
+@pytest.fixture
+def mock_boto_client(mocker):
+    """Fixture to mock the boto3 S3Vectors client."""
+    mock_client = mocker.MagicMock()
+    mocker.patch("boto3.client", return_value=mock_client)
+    return mock_client
+
+
+def test_initialization_creates_resources(mock_boto_client):
+    """Test that bucket and index are created if they don't exist."""
+    not_found_error = ClientError({"Error": {"Code": "NotFoundException"}}, "OperationName")
+    mock_boto_client.get_vector_bucket.side_effect = not_found_error
+    mock_boto_client.get_index.side_effect = not_found_error
+
+    S3Vectors(
+        vector_bucket_name=BUCKET_NAME,
+        index_name=INDEX_NAME,
+        embedding_model_dims=EMBEDDING_DIMS,
+        region_name=REGION,
+    )
+
+    mock_boto_client.create_vector_bucket.assert_called_once_with(vectorBucketName=BUCKET_NAME)
+    mock_boto_client.create_index.assert_called_once_with(
+        vectorBucketName=BUCKET_NAME,
+        indexName=INDEX_NAME,
+        dataType="float32",
+        dimension=EMBEDDING_DIMS,
+        distanceMetric="cosine",
+    )
+
+
+def test_initialization_uses_existing_resources(mock_boto_client):
+    """Test that existing bucket and index are used if found."""
+    mock_boto_client.get_vector_bucket.return_value = {}
+    mock_boto_client.get_index.return_value = {}
+
+    S3Vectors(
+        vector_bucket_name=BUCKET_NAME,
+        index_name=INDEX_NAME,
+        embedding_model_dims=EMBEDDING_DIMS,
+        region_name=REGION,
+    )
+
+    mock_boto_client.create_vector_bucket.assert_not_called()
+    mock_boto_client.create_index.assert_not_called()
+
+
+def test_insert(mock_boto_client):
+    """Test inserting vectors."""
+    store = S3Vectors(
+        vector_bucket_name=BUCKET_NAME,
+        index_name=INDEX_NAME,
+        embedding_model_dims=EMBEDDING_DIMS,
+    )
+    vectors = [[0.1, 0.2], [0.3, 0.4]]
+    payloads = [{"meta": "data1"}, {"meta": "data2"}]
+    ids = ["id1", "id2"]
+
+    store.insert(vectors, payloads, ids)
+
+    mock_boto_client.put_vectors.assert_called_once_with(
+        vectorBucketName=BUCKET_NAME,
+        indexName=INDEX_NAME,
+        vectors=[
+            {
+                "key": "id1",
+                "data": {"float32": [0.1, 0.2]},
+                "metadata": {"meta": "data1"},
+            },
+            {
+                "key": "id2",
+                "data": {"float32": [0.3, 0.4]},
+                "metadata": {"meta": "data2"},
+            },
+        ],
+    )
+
+
+def test_search(mock_boto_client):
+    """Test searching for vectors."""
+    mock_boto_client.query_vectors.return_value = {
+        "vectors": [{"key": "id1", "distance": 0.9, "metadata": {"meta": "data1"}}]
+    }
+    store = S3Vectors(
+        vector_bucket_name=BUCKET_NAME,
+        index_name=INDEX_NAME,
+        embedding_model_dims=EMBEDDING_DIMS,
+    )
+    query_vector = [0.1, 0.2]
+    results = store.search(query="test", vectors=query_vector, limit=1)
+
+    mock_boto_client.query_vectors.assert_called_once()
+    assert len(results) == 1
+    assert results[0].id == "id1"
+    assert results[0].score == 0.9
+
+
+def test_get(mock_boto_client):
+    """Test retrieving a vector by ID."""
+    mock_boto_client.get_vectors.return_value = {"vectors": [{"key": "id1", "metadata": {"meta": "data1"}}]}
+    store = S3Vectors(
+        vector_bucket_name=BUCKET_NAME,
+        index_name=INDEX_NAME,
+        embedding_model_dims=EMBEDDING_DIMS,
+    )
+    result = store.get("id1")
+
+    mock_boto_client.get_vectors.assert_called_once_with(
+        vectorBucketName=BUCKET_NAME,
+        indexName=INDEX_NAME,
+        keys=["id1"],
+        returnData=False,
+        returnMetadata=True,
+    )
+    assert result.id == "id1"
+    assert result.payload["meta"] == "data1"
+
+
+def test_delete(mock_boto_client):
+    """Test deleting a vector."""
+    store = S3Vectors(
+        vector_bucket_name=BUCKET_NAME,
+        index_name=INDEX_NAME,
+        embedding_model_dims=EMBEDDING_DIMS,
+    )
+    store.delete("id1")
+
+    mock_boto_client.delete_vectors.assert_called_once_with(
+        vectorBucketName=BUCKET_NAME, indexName=INDEX_NAME, keys=["id1"]
+    )
+
+
+def test_reset(mock_boto_client):
+    """Test resetting the vector index."""
+    # GIVEN: The index does not exist, so it gets created on init and reset
+    not_found_error = ClientError({"Error": {"Code": "NotFoundException"}}, "OperationName")
+    mock_boto_client.get_index.side_effect = not_found_error
+
+    # WHEN: The store is initialized
+    store = S3Vectors(
+        vector_bucket_name=BUCKET_NAME,
+        index_name=INDEX_NAME,
+        embedding_model_dims=EMBEDDING_DIMS,
+    )
+
+    # THEN: The index is created once during initialization
+    assert mock_boto_client.create_index.call_count == 1
+
+    # WHEN: The store is reset
+    store.reset()
+
+    # THEN: The index is deleted and then created again
+    mock_boto_client.delete_index.assert_called_once_with(vectorBucketName=BUCKET_NAME, indexName=INDEX_NAME)
+    assert mock_boto_client.create_index.call_count == 2