Fixed s3 vectors memory initialization issue from configuration (#3481)

This commit is contained in:
Abdullah Irfan
2025-09-20 12:43:18 +05:30
committed by GitHub
parent ead210ffe4
commit 9df392b26a
4 changed files with 94 additions and 38 deletions
+17 -17
View File
@@ -28,7 +28,7 @@ config = {
"provider": "s3_vectors",
"config": {
"vector_bucket_name": "my-mem0-vector-bucket",
"index_name": "my-memories-index",
"collection_name": "my-memories-index",
"embedding_model_dims": 1536,
"distance_metric": "cosine",
"region_name": "us-east-1"
@@ -50,13 +50,13 @@ m.add(messages, user_id="alice", metadata={"category": "movies"})
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) |
| 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 |
| `collection_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
@@ -64,15 +64,15 @@ Your AWS identity (user or role) needs permissions to perform actions on S3 Vect
```json
{
"Version": "2012-10-17",
"Statement": [
{
"Effect": "Allow",
"Action": "s3vectors:*",
"Resource": "*"
}
]
"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.
For production, it is recommended to scope down the resource ARN to your specific buckets and indexes.
+3 -7
View File
@@ -5,17 +5,13 @@ from pydantic import BaseModel, ConfigDict, 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"
)
collection_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"
)
region_name: Optional[str] = Field(None, description="AWS region for the S3 Vectors client")
@model_validator(mode="before")
@classmethod
+2 -2
View File
@@ -25,14 +25,14 @@ class S3Vectors(VectorStoreBase):
def __init__(
self,
vector_bucket_name: str,
index_name: str,
collection_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.collection_name = collection_name
self.embedding_model_dims = embedding_model_dims
self.distance_metric = distance_metric
+72 -12
View File
@@ -1,6 +1,8 @@
from mem0.configs.vector_stores.s3_vectors import S3VectorsConfig
import pytest
from botocore.exceptions import ClientError
from mem0.memory.main import Memory
from mem0.vector_stores.s3_vectors import S3Vectors
BUCKET_NAME = "test-bucket"
@@ -17,20 +19,42 @@ def mock_boto_client(mocker):
return mock_client
@pytest.fixture
def mock_embedder(mocker):
mock_embedder = mocker.MagicMock()
mock_embedder.return_value.embed.return_value = [0.1, 0.2, 0.3]
mocker.patch("mem0.utils.factory.EmbedderFactory.create", mock_embedder)
return mock_embedder
@pytest.fixture
def mock_llm(mocker):
mock_llm = mocker.MagicMock()
mocker.patch("mem0.utils.factory.LlmFactory.create", mock_llm)
mocker.patch("mem0.memory.storage.SQLiteManager", mocker.MagicMock())
return mock_llm
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")
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,
collection_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_vector_bucket.assert_called_once_with(
vectorBucketName=BUCKET_NAME
)
mock_boto_client.create_index.assert_called_once_with(
vectorBucketName=BUCKET_NAME,
indexName=INDEX_NAME,
@@ -47,7 +71,7 @@ def test_initialization_uses_existing_resources(mock_boto_client):
S3Vectors(
vector_bucket_name=BUCKET_NAME,
index_name=INDEX_NAME,
collection_name=INDEX_NAME,
embedding_model_dims=EMBEDDING_DIMS,
region_name=REGION,
)
@@ -56,11 +80,41 @@ def test_initialization_uses_existing_resources(mock_boto_client):
mock_boto_client.create_index.assert_not_called()
def test_memory_initialization_with_config(mock_boto_client, mock_llm, mock_embedder):
"""Test Memory initialization with S3Vectors from config."""
# check that Attribute error is not raised
mock_boto_client.get_vector_bucket.return_value = {}
mock_boto_client.get_index.return_value = {}
config = {
"vector_store": {
"provider": "s3_vectors",
"config": {
"vector_bucket_name": BUCKET_NAME,
"collection_name": INDEX_NAME,
"embedding_model_dims": EMBEDDING_DIMS,
"distance_metric": "cosine",
"region_name": REGION,
},
}
}
try:
memory = Memory.from_config(config)
assert memory.vector_store is not None
assert isinstance(memory.vector_store, S3Vectors)
assert isinstance(memory.config.vector_store.config, S3VectorsConfig)
except AttributeError:
pytest.fail("Memory initialization failed")
def test_insert(mock_boto_client):
"""Test inserting vectors."""
store = S3Vectors(
vector_bucket_name=BUCKET_NAME,
index_name=INDEX_NAME,
collection_name=INDEX_NAME,
embedding_model_dims=EMBEDDING_DIMS,
)
vectors = [[0.1, 0.2], [0.3, 0.4]]
@@ -94,7 +148,7 @@ def test_search(mock_boto_client):
}
store = S3Vectors(
vector_bucket_name=BUCKET_NAME,
index_name=INDEX_NAME,
collection_name=INDEX_NAME,
embedding_model_dims=EMBEDDING_DIMS,
)
query_vector = [0.1, 0.2]
@@ -108,10 +162,12 @@ def test_search(mock_boto_client):
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"}}]}
mock_boto_client.get_vectors.return_value = {
"vectors": [{"key": "id1", "metadata": {"meta": "data1"}}]
}
store = S3Vectors(
vector_bucket_name=BUCKET_NAME,
index_name=INDEX_NAME,
collection_name=INDEX_NAME,
embedding_model_dims=EMBEDDING_DIMS,
)
result = store.get("id1")
@@ -131,7 +187,7 @@ def test_delete(mock_boto_client):
"""Test deleting a vector."""
store = S3Vectors(
vector_bucket_name=BUCKET_NAME,
index_name=INDEX_NAME,
collection_name=INDEX_NAME,
embedding_model_dims=EMBEDDING_DIMS,
)
store.delete("id1")
@@ -144,13 +200,15 @@ def test_delete(mock_boto_client):
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")
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,
collection_name=INDEX_NAME,
embedding_model_dims=EMBEDDING_DIMS,
)
@@ -161,5 +219,7 @@ def test_reset(mock_boto_client):
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)
mock_boto_client.delete_index.assert_called_once_with(
vectorBucketName=BUCKET_NAME, indexName=INDEX_NAME
)
assert mock_boto_client.create_index.call_count == 2