diff --git a/docs/components/vectordbs/dbs/s3_vectors.mdx b/docs/components/vectordbs/dbs/s3_vectors.mdx index 8faf09b46..47be4fb83 100644 --- a/docs/components/vectordbs/dbs/s3_vectors.mdx +++ b/docs/components/vectordbs/dbs/s3_vectors.mdx @@ -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. \ No newline at end of file +For production, it is recommended to scope down the resource ARN to your specific buckets and indexes. diff --git a/mem0/configs/vector_stores/s3_vectors.py b/mem0/configs/vector_stores/s3_vectors.py index 95c50f675..4118a4086 100644 --- a/mem0/configs/vector_stores/s3_vectors.py +++ b/mem0/configs/vector_stores/s3_vectors.py @@ -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 diff --git a/mem0/vector_stores/s3_vectors.py b/mem0/vector_stores/s3_vectors.py index 37b80cb9c..f6504c379 100644 --- a/mem0/vector_stores/s3_vectors.py +++ b/mem0/vector_stores/s3_vectors.py @@ -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 diff --git a/tests/vector_stores/test_s3_vectors.py b/tests/vector_stores/test_s3_vectors.py index 29b877c0c..e8141e2f5 100644 --- a/tests/vector_stores/test_s3_vectors.py +++ b/tests/vector_stores/test_s3_vectors.py @@ -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