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
+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