Fixed s3 vectors memory initialization issue from configuration (#3481)
This commit is contained in:
@@ -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.
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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
|
||||
|
||||
|
||||
@@ -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
|
||||
|
||||
Reference in New Issue
Block a user