JDSS/mem0
2
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Add Amazon S3 Vectors Support (#3237)

Browse Source
This commit is contained in:
Enzo Biondo
2025-08-23 08:14:42 -03:00
committed by GitHub
parent ee00bd5731
commit c72bfc3285
10 changed files with 461 additions and 2 deletions
Download Patch File Download Diff File Expand all files Collapse all files
LLM.md
+3 -2
View File
@@ -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
+78
View File
@@ -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
+1
View File
@@ -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
+1
View File
@@ -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
+1
View File
@@ -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
+34
View File
@@ -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
+1
View File
@@ -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
+1
View File
@@ -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
+176
View File
@@ -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
+165
View File
@@ -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