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

fix(sdk): removing deprecating param from our sdk and docs changes with it (#4740)

Browse Source
This commit is contained in:
Kartik
2026-04-12 00:34:58 +05:30
committed by GitHub
parent 3882af7450
commit e44b46ef2e
138 changed files with 1186 additions and 2840 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tests/vector_stores/test_azure_ai_search.py
+1 -1
View File
@@ -533,7 +533,7 @@ def test_search_basic(azure_ai_search_instance):
# Search with a vector
query_text = "test query" # Add a query string
query_vector = [0.1, 0.2, 0.3]
results = instance.search(query_text, query_vector, limit=5) # Pass the query string
results = instance.search(query_text, query_vector, top_k=5) # Pass the query string
# Verify search was called correctly
mock_search_client.search.assert_called_once()
tests/vector_stores/test_azure_mysql.py
+4 -3
View File
@@ -1,7 +1,8 @@
import json
import pytest
from unittest.mock import Mock, patch
import pytest
from mem0.vector_stores.azure_mysql import AzureMySQL, OutputData
@@ -118,7 +119,7 @@ def test_search(azure_mysql_instance):
])
query_vector = [0.2, 0.3, 0.4]
results = azure_mysql_instance.search(query="test", vectors=query_vector, limit=5)
results = azure_mysql_instance.search(query="test", vectors=query_vector, top_k=5)
assert isinstance(results, list)
assert cursor.execute.called
@@ -218,7 +219,7 @@ def test_list(azure_mysql_instance):
}
])
results = azure_mysql_instance.list(limit=10)
results = azure_mysql_instance.list(top_k=10)
assert isinstance(results, list)
assert len(results) > 0
tests/vector_stores/test_baidu.py
+3 -3
View File
@@ -112,7 +112,7 @@ def test_search(mochow_instance, mock_mochow_client):
mochow_instance._table.vector_search.return_value = mock_search_results
vectors = [0.1, 0.2, 0.3]
results = mochow_instance.search(query="test", vectors=vectors, limit=2)
results = mochow_instance.search(query="test", vectors=vectors, top_k=2)
# Verify search was called with correct parameters
mochow_instance._table.vector_search.assert_called_once()
@@ -143,7 +143,7 @@ def test_search_with_filters(mochow_instance, mock_mochow_client):
vectors = [0.1, 0.2, 0.3]
filters = {"user_id": "user123", "agent_id": "agent456"}
mochow_instance.search(query="test", vectors=vectors, limit=2, filters=filters)
mochow_instance.search(query="test", vectors=vectors, top_k=2, filters=filters)
# Verify search was called with filter
call_args = mochow_instance._table.vector_search.call_args
@@ -196,7 +196,7 @@ def test_list(mochow_instance, mock_mochow_client):
mock_result.rows = [{"id": "id1", "metadata": {"name": "vector1"}}, {"id": "id2", "metadata": {"name": "vector2"}}]
mochow_instance._table.select.return_value = mock_result
results = mochow_instance.list(limit=2)
results = mochow_instance.list(top_k=2)
mochow_instance._table.select.assert_called_once_with(filter=None, projections=["id", "metadata"], limit=2)
tests/vector_stores/test_cassandra.py
+5 -4
View File
@@ -1,7 +1,8 @@
import json
import pytest
from unittest.mock import Mock, patch
import pytest
from mem0.vector_stores.cassandra import CassandraDB, OutputData
@@ -106,7 +107,7 @@ def test_search(cassandra_instance):
cassandra_instance.session.execute = Mock(return_value=[mock_row1, mock_row2])
query_vector = [0.2, 0.3, 0.4]
results = cassandra_instance.search(query="test", vectors=query_vector, limit=5)
results = cassandra_instance.search(query="test", vectors=query_vector, top_k=5)
assert isinstance(results, list)
assert len(results) <= 5
@@ -213,7 +214,7 @@ def test_list(cassandra_instance):
cassandra_instance.session.execute = Mock(return_value=[mock_row])
results = cassandra_instance.list(limit=10)
results = cassandra_instance.list(top_k=10)
assert isinstance(results, list)
assert len(results) > 0
@@ -264,7 +265,7 @@ def test_search_with_filters(cassandra_instance):
results = cassandra_instance.search(
query="test",
vectors=query_vector,
limit=5,
top_k=5,
filters={"category": "A"}
)
tests/vector_stores/test_chroma.py
+8 -8
View File
@@ -2,8 +2,8 @@ from unittest.mock import Mock, patch
import pytest
from mem0.vector_stores.chroma import ChromaDB
from mem0.configs.vector_stores.chroma import ChromaDbConfig
from mem0.vector_stores.chroma import ChromaDB
@pytest.fixture
@@ -39,7 +39,7 @@ def test_search_vectors(chromadb_instance, mock_chromadb_client):
chromadb_instance.collection.query.return_value = mock_result
vectors = [[0.1, 0.2, 0.3]]
results = chromadb_instance.search(query="", vectors=vectors, limit=2)
results = chromadb_instance.search(query="", vectors=vectors, top_k=2)
chromadb_instance.collection.query.assert_called_once_with(query_embeddings=vectors, where=None, n_results=2)
@@ -60,7 +60,7 @@ def test_search_vectors_with_filters(chromadb_instance, mock_chromadb_client):
vectors = [[0.1, 0.2, 0.3]]
filters = {"user_id": "alice", "agent_id": "agent1", "run_id": "run1"}
results = chromadb_instance.search(query="", vectors=vectors, limit=2, filters=filters)
results = chromadb_instance.search(query="", vectors=vectors, top_k=2, filters=filters)
# Verify that _generate_where_clause was called with the filters
expected_where = {"$and": [{"user_id": {"$eq": "alice"}}, {"agent_id": {"$eq": "agent1"}}, {"run_id": {"$eq": "run1"}}]}
@@ -86,7 +86,7 @@ def test_search_vectors_with_single_filter(chromadb_instance, mock_chromadb_clie
vectors = [[0.1, 0.2, 0.3]]
filters = {"user_id": "alice"}
results = chromadb_instance.search(query="", vectors=vectors, limit=2, filters=filters)
results = chromadb_instance.search(query="", vectors=vectors, top_k=2, filters=filters)
# Verify that single filter is passed with $eq operator
expected_where = {"user_id": {"$eq": "alice"}}
@@ -108,7 +108,7 @@ def test_search_vectors_with_no_filters(chromadb_instance, mock_chromadb_client)
chromadb_instance.collection.query.return_value = mock_result
vectors = [[0.1, 0.2, 0.3]]
results = chromadb_instance.search(query="", vectors=vectors, limit=2, filters=None)
results = chromadb_instance.search(query="", vectors=vectors, top_k=2, filters=None)
chromadb_instance.collection.query.assert_called_once_with(
query_embeddings=vectors, where=None, n_results=2
@@ -162,7 +162,7 @@ def test_list_vectors(chromadb_instance):
}
chromadb_instance.collection.get.return_value = mock_result
results = chromadb_instance.list(limit=2)
results = chromadb_instance.list(top_k=2)
chromadb_instance.collection.get.assert_called_once_with(where=None, limit=2)
@@ -181,7 +181,7 @@ def test_list_vectors_with_filters(chromadb_instance):
chromadb_instance.collection.get.return_value = mock_result
filters = {"user_id": "alice", "agent_id": "agent1", "run_id": "run1"}
results = chromadb_instance.list(filters=filters, limit=2)
results = chromadb_instance.list(filters=filters, top_k=2)
# Verify that _generate_where_clause was called with the filters
expected_where = {"$and": [{"user_id": {"$eq": "alice"}}, {"agent_id": {"$eq": "agent1"}}, {"run_id": {"$eq": "run1"}}]}
@@ -203,7 +203,7 @@ def test_list_vectors_with_single_filter(chromadb_instance):
chromadb_instance.collection.get.return_value = mock_result
filters = {"user_id": "alice"}
results = chromadb_instance.list(filters=filters, limit=2)
results = chromadb_instance.list(filters=filters, top_k=2)
# Verify that single filter is passed with $eq operator
expected_where = {"user_id": {"$eq": "alice"}}
tests/vector_stores/test_databricks.py
+23 -17
View File
@@ -5,9 +5,15 @@ import pytest
pytest.importorskip("databricks", reason="databricks-sdk package not installed")
from databricks.sdk.service.vectorsearch import VectorIndexType, QueryVectorIndexResponse, ResultManifest, ResultData, ColumnInfo
from mem0.vector_stores.databricks import Databricks
from databricks.sdk.service.vectorsearch import (
ColumnInfo,
QueryVectorIndexResponse,
ResultData,
ResultManifest,
VectorIndexType,
)
from mem0.vector_stores.databricks import Databricks
# ---------------------- Fixtures ---------------------- #
@@ -185,7 +191,7 @@ def test_search_delta_sync_text(db_instance_delta, mock_workspace_client):
mock_workspace_client.vector_search_indexes.query_index.return_value = SimpleNamespace(
result=SimpleNamespace(data_array=[row])
)
results = db_instance_delta.search(query="hello", vectors=None, limit=1)
results = db_instance_delta.search(query="hello", vectors=None, top_k=1)
mock_workspace_client.vector_search_indexes.query_index.assert_called_once()
assert len(results) == 1
assert results[0].id == "id1"
@@ -210,7 +216,7 @@ def test_search_direct_access_vector(db_instance_direct, mock_workspace_client):
mock_workspace_client.vector_search_indexes.query_index.return_value = SimpleNamespace(
result=SimpleNamespace(data_array=[row])
)
results = db_instance_direct.search(query="", vectors=[0.1, 0.2, 0.3, 0.4], limit=1)
results = db_instance_direct.search(query="", vectors=[0.1, 0.2, 0.3, 0.4], top_k=1)
assert len(results) == 1
assert results[0].id == "id2"
assert results[0].score == 0.77
@@ -235,7 +241,7 @@ def test_search_delta_sync_self_managed_vectors(mock_workspace_client):
mock_workspace_client.vector_search_indexes.query_index.return_value = SimpleNamespace(
result=SimpleNamespace(data_array=[])
)
inst.search(query="ignored", vectors=[0.1, 0.2, 0.3, 0.4], limit=5)
inst.search(query="ignored", vectors=[0.1, 0.2, 0.3, 0.4], top_k=5)
call_kwargs = mock_workspace_client.vector_search_indexes.query_index.call_args.kwargs
assert "query_vector" in call_kwargs
assert "query_text" not in call_kwargs
@@ -423,7 +429,7 @@ def test_list_memories(db_instance_delta, mock_workspace_client):
]
)
)
res = db_instance_delta.list(limit=1)
res = db_instance_delta.list(top_k=1)
assert isinstance(res, list)
assert len(res[0]) == 1
assert res[0][0].id == "id-get"
@@ -453,7 +459,7 @@ def test_list_memories_direct_access(db_instance_direct, mock_workspace_client):
]
)
)
res = db_instance_direct.list(limit=5)
res = db_instance_direct.list(top_k=5)
assert isinstance(res, list)
assert len(res[0]) == 1
assert res[0][0].id == "id-da-list"
@@ -516,7 +522,7 @@ def test_list_memories_delta_sync_self_managed(mock_workspace_client):
mock_workspace_client.vector_search_indexes.query_index.return_value = SimpleNamespace(
result=SimpleNamespace(data_array=[])
)
inst.list(limit=5)
inst.list(top_k=5)
call_kwargs = mock_workspace_client.vector_search_indexes.query_index.call_args.kwargs
assert "query_vector" in call_kwargs
assert "query_text" not in call_kwargs
@@ -527,7 +533,7 @@ def test_list_memories_default_limit(db_instance_delta, mock_workspace_client):
mock_workspace_client.vector_search_indexes.query_index.return_value = SimpleNamespace(
result=SimpleNamespace(data_array=[])
)
db_instance_delta.list(limit=None)
db_instance_delta.list(top_k=None)
call_kwargs = mock_workspace_client.vector_search_indexes.query_index.call_args.kwargs
assert call_kwargs["num_results"] == 100
@@ -624,8 +630,8 @@ def test_reset(db_instance_delta, mock_workspace_client):
def test_e2e_config_to_factory_delta_sync(mock_workspace_client):
"""End-to-end: VectorStoreConfig validates docs-correct params, factory creates Databricks instance."""
from mem0.vector_stores.configs import VectorStoreConfig
from mem0.utils.factory import VectorStoreFactory
from mem0.vector_stores.configs import VectorStoreConfig
# Step 1: Config validation (simulates what Memory.from_config does)
vs_config = VectorStoreConfig(
@@ -655,8 +661,8 @@ def test_e2e_config_to_factory_delta_sync(mock_workspace_client):
def test_e2e_config_to_factory_direct_access(mock_workspace_client):
"""End-to-end: DIRECT_ACCESS via config → factory creates correct instance."""
from mem0.vector_stores.configs import VectorStoreConfig
from mem0.utils.factory import VectorStoreFactory
from mem0.vector_stores.configs import VectorStoreConfig
mock_workspace_client.tables.exists.return_value = SimpleNamespace(table_exists=True)
@@ -699,8 +705,8 @@ def test_e2e_old_docs_config_rejected():
def test_e2e_crud_lifecycle_delta_sync(mock_workspace_client):
"""End-to-end CRUD lifecycle: insert → search → get → list → update → delete."""
from mem0.vector_stores.configs import VectorStoreConfig
from mem0.utils.factory import VectorStoreFactory
from mem0.vector_stores.configs import VectorStoreConfig
vs_config = VectorStoreConfig(
provider="databricks",
@@ -736,7 +742,7 @@ def test_e2e_crud_lifecycle_delta_sync(mock_workspace_client):
data_array=[["mem-001", "h1", None, None, "u1", "test memory", None, None, None, 0.95]]
)
)
results = db.search(query="test", vectors=None, limit=5)
results = db.search(query="test", vectors=None, top_k=5)
assert len(results) == 1
assert results[0].id == "mem-001"
assert results[0].payload["data"] == "test memory"
@@ -767,7 +773,7 @@ def test_e2e_crud_lifecycle_delta_sync(mock_workspace_client):
data_array=[["mem-001", "h1", None, None, "u1", "test memory", None, None, None]]
)
)
listed = db.list(filters={"user_id": "u1"}, limit=10)
listed = db.list(filters={"user_id": "u1"}, top_k=10)
assert len(listed[0]) == 1
assert listed[0][0].id == "mem-001"
list_kwargs = mock_workspace_client.vector_search_indexes.query_index.call_args.kwargs
@@ -797,8 +803,8 @@ def test_e2e_crud_lifecycle_delta_sync(mock_workspace_client):
def test_e2e_crud_lifecycle_direct_access(mock_workspace_client):
"""End-to-end CRUD lifecycle for DIRECT_ACCESS: insert → search → get → list."""
from mem0.vector_stores.configs import VectorStoreConfig
from mem0.utils.factory import VectorStoreFactory
from mem0.vector_stores.configs import VectorStoreConfig
mock_workspace_client.tables.exists.return_value = SimpleNamespace(table_exists=True)
@@ -838,7 +844,7 @@ def test_e2e_crud_lifecycle_direct_access(mock_workspace_client):
data_array=[["mem-da-001", "h1", None, None, "u1", "direct memory", None, None, None, [0.1, 0.2, 0.3, 0.4], 0.9]]
)
)
results = db.search(query="", vectors=[0.1, 0.2, 0.3, 0.4], limit=5)
results = db.search(query="", vectors=[0.1, 0.2, 0.3, 0.4], top_k=5)
assert len(results) == 1
search_kwargs = mock_workspace_client.vector_search_indexes.query_index.call_args.kwargs
assert "query_vector" in search_kwargs
@@ -867,7 +873,7 @@ def test_e2e_crud_lifecycle_direct_access(mock_workspace_client):
data_array=[["mem-da-001", "h1", None, None, "u1", "direct memory", None, None, None, [0.1, 0.2, 0.3, 0.4]]]
)
)
db.list(limit=5)
db.list(top_k=5)
list_kwargs = mock_workspace_client.vector_search_indexes.query_index.call_args.kwargs
assert "query_vector" in list_kwargs
assert "query_text" not in list_kwargs
tests/vector_stores/test_elasticsearch.py
+4 -4
View File
@@ -9,8 +9,8 @@ try:
except ImportError:
raise ImportError("Elasticsearch requires extra dependencies. Install with `pip install elasticsearch`") from None
from mem0.vector_stores.elasticsearch import ElasticsearchDB, OutputData
from mem0.configs.vector_stores.elasticsearch import ElasticsearchConfig
from mem0.vector_stores.elasticsearch import ElasticsearchDB, OutputData
class TestElasticsearchDB(unittest.TestCase):
@@ -189,7 +189,7 @@ class TestElasticsearchDB(unittest.TestCase):
# Perform search
vectors = [[0.1] * 1536]
results = self.es_db.search(query="", vectors=vectors, limit=5)
results = self.es_db.search(query="", vectors=vectors, top_k=5)
# Verify search call
self.client_mock.search.assert_called_once()
@@ -221,7 +221,7 @@ class TestElasticsearchDB(unittest.TestCase):
vectors = [[0.1] * 1536]
limit = 5
filters = {"key1": "value1"}
self.es_db.search(query="", vectors=vectors, limit=limit, filters=filters)
self.es_db.search(query="", vectors=vectors, top_k=limit, filters=filters)
# Verify custom search query function was called
self.es_db.custom_search_query.assert_called_once_with(vectors, limit, filters)
@@ -272,7 +272,7 @@ class TestElasticsearchDB(unittest.TestCase):
self.client_mock.search.return_value = mock_response
# Perform list operation
results = self.es_db.list(limit=10)
results = self.es_db.list(top_k=10)
# Verify search call
self.client_mock.search.assert_called_once()
tests/vector_stores/test_faiss.py
+3 -3
View File
@@ -101,7 +101,7 @@ def test_search(faiss_instance, mock_faiss_index):
with patch.object(faiss_instance, "_parse_output", return_value=expected_results):
# Call search
results = faiss_instance.search(query="test query", vectors=query_vector, limit=2)
results = faiss_instance.search(query="test query", vectors=query_vector, top_k=2)
# Verify numpy.array was called (but we don't check exact call arguments since it's complex)
assert mock_np_array.called
@@ -148,7 +148,7 @@ def test_search_with_filters(faiss_instance, mock_faiss_index):
with patch.object(faiss_instance, "_apply_filters", side_effect=lambda p, f: p.get("category") == "A"):
# Call search with filters
results = faiss_instance.search(
query="test query", vectors=query_vector, limit=2, filters={"category": "A"}
query="test query", vectors=query_vector, top_k=2, filters={"category": "A"}
)
# Verify numpy.array was called
@@ -241,7 +241,7 @@ def test_list(faiss_instance):
assert len(results[0]) == 3
# Test listing with a limit
results = faiss_instance.list(limit=2)
results = faiss_instance.list(top_k=2)
assert len(results[0]) == 2
# Test listing with filters
tests/vector_stores/test_langchain_vector_store.py
+9 -9
View File
@@ -48,7 +48,7 @@ def test_search_vectors(langchain_instance):
# Test search without filters
vectors = [[0.1, 0.2, 0.3]]
results = langchain_instance.search(query="", vectors=vectors, limit=2)
results = langchain_instance.search(query="", vectors=vectors, top_k=2)
langchain_instance.client.similarity_search_by_vector.assert_called_once_with(embedding=vectors, k=2)
@@ -60,7 +60,7 @@ def test_search_vectors(langchain_instance):
# Test search with filters
filters = {"name": "vector1"}
langchain_instance.search(query="", vectors=vectors, limit=2, filters=filters)
langchain_instance.search(query="", vectors=vectors, top_k=2, filters=filters)
langchain_instance.client.similarity_search_by_vector.assert_called_with(embedding=vectors, k=2, filter=filters)
@@ -75,7 +75,7 @@ def test_search_vectors_with_agent_id_run_id_filters(langchain_instance):
vectors = [[0.1, 0.2, 0.3]]
filters = {"user_id": "alice", "agent_id": "agent1", "run_id": "run1"}
results = langchain_instance.search(query="", vectors=vectors, limit=2, filters=filters)
results = langchain_instance.search(query="", vectors=vectors, top_k=2, filters=filters)
# Verify that filters were passed to the underlying vector store
langchain_instance.client.similarity_search_by_vector.assert_called_once_with(
@@ -96,7 +96,7 @@ def test_search_vectors_with_single_filter(langchain_instance):
vectors = [[0.1, 0.2, 0.3]]
filters = {"user_id": "alice"}
results = langchain_instance.search(query="", vectors=vectors, limit=2, filters=filters)
results = langchain_instance.search(query="", vectors=vectors, top_k=2, filters=filters)
# Verify that filters were passed to the underlying vector store
langchain_instance.client.similarity_search_by_vector.assert_called_once_with(
@@ -114,7 +114,7 @@ def test_search_vectors_with_no_filters(langchain_instance):
langchain_instance.client.similarity_search_by_vector.return_value = mock_docs
vectors = [[0.1, 0.2, 0.3]]
results = langchain_instance.search(query="", vectors=vectors, limit=2, filters=None)
results = langchain_instance.search(query="", vectors=vectors, top_k=2, filters=None)
# Verify that no filters were passed to the underlying vector store
langchain_instance.client.similarity_search_by_vector.assert_called_once_with(
@@ -155,7 +155,7 @@ def test_list_with_filters(langchain_instance):
langchain_instance.client._collection = mock_collection
filters = {"user_id": "alice", "agent_id": "agent1", "run_id": "run1"}
results = langchain_instance.list(filters=filters, limit=10)
results = langchain_instance.list(filters=filters, top_k=10)
# Verify that the collection.get method was called with the correct filters
mock_collection.get.assert_called_once_with(where=filters, limit=10)
@@ -180,7 +180,7 @@ def test_list_with_single_filter(langchain_instance):
langchain_instance.client._collection = mock_collection
filters = {"user_id": "alice"}
results = langchain_instance.list(filters=filters, limit=10)
results = langchain_instance.list(filters=filters, top_k=10)
# Verify that the collection.get method was called with the correct filter
mock_collection.get.assert_called_once_with(where=filters, limit=10)
@@ -202,7 +202,7 @@ def test_list_with_no_filters(langchain_instance):
}
langchain_instance.client._collection = mock_collection
results = langchain_instance.list(filters=None, limit=10)
results = langchain_instance.list(filters=None, top_k=10)
# Verify that the collection.get method was called with no filters
mock_collection.get.assert_called_once_with(where=None, limit=10)
@@ -220,7 +220,7 @@ def test_list_with_exception(langchain_instance):
mock_collection.get.side_effect = Exception("Test exception")
langchain_instance.client._collection = mock_collection
results = langchain_instance.list(filters={"user_id": "alice"}, limit=10)
results = langchain_instance.list(filters={"user_id": "alice"}, top_k=10)
# Verify that an empty list is returned when an exception occurs
assert results == []
tests/vector_stores/test_milvus.py
+6 -4
View File
@@ -8,10 +8,12 @@ These tests verify:
4. Update/upsert operations
"""
import pytest
from unittest.mock import MagicMock, patch
from mem0.vector_stores.milvus import MilvusDB
import pytest
from mem0.configs.vector_stores.milvus import MetricType
from mem0.vector_stores.milvus import MilvusDB
class TestMilvusDB:
@@ -118,7 +120,7 @@ class TestMilvusDB:
results = milvus_db.search(
query="test query",
vectors=query_vector,
limit=5,
top_k=5,
filters=filters
)
@@ -199,7 +201,7 @@ class TestMilvusDB:
{"id": "mem2", "metadata": {"user_id": "alice"}}
]
results = milvus_db.list(filters={"user_id": "alice"}, limit=10)
results = milvus_db.list(filters={"user_id": "alice"}, top_k=10)
# Verify query was called with filter
call_args = mock_milvus_client.query.call_args
tests/vector_stores/test_mongodb.py
+8 -8
View File
@@ -86,7 +86,7 @@ def test_search(mongo_vector_fixture):
]
mock_collection.list_search_indexes.return_value = ["test_collection_vector_index"]
results = mongo_vector.search("query_str", query_vector, limit=2)
results = mongo_vector.search("query_str", query_vector, top_k=2)
mock_collection.list_search_indexes.assert_called_with(name="test_collection_vector_index")
mock_collection.aggregate.assert_called_once_with(
[
@@ -119,7 +119,7 @@ def test_search_with_filters(mongo_vector_fixture):
mock_collection.list_search_indexes.return_value = ["test_collection_vector_index"]
filters = {"user_id": "alice", "agent_id": "agent1", "run_id": "run1"}
results = mongo_vector.search("query_str", query_vector, limit=2, filters=filters)
results = mongo_vector.search("query_str", query_vector, top_k=2, filters=filters)
# Verify that the aggregation pipeline includes the filter stage
mock_collection.aggregate.assert_called_once()
@@ -153,7 +153,7 @@ def test_search_with_single_filter(mongo_vector_fixture):
mock_collection.list_search_indexes.return_value = ["test_collection_vector_index"]
filters = {"user_id": "alice"}
results = mongo_vector.search("query_str", query_vector, limit=2, filters=filters)
results = mongo_vector.search("query_str", query_vector, top_k=2, filters=filters)
# Verify that the aggregation pipeline includes the filter stage
mock_collection.aggregate.assert_called_once()
@@ -177,7 +177,7 @@ def test_search_with_no_filters(mongo_vector_fixture):
]
mock_collection.list_search_indexes.return_value = ["test_collection_vector_index"]
results = mongo_vector.search("query_str", query_vector, limit=2, filters=None)
results = mongo_vector.search("query_str", query_vector, top_k=2, filters=None)
# Verify that the aggregation pipeline does not include the filter stage
mock_collection.aggregate.assert_called_once()
@@ -315,7 +315,7 @@ def test_list(mongo_vector_fixture):
{"_id": "id2", "payload": {"key": "value2"}},
]
results = mongo_vector.list(limit=2)
results = mongo_vector.list(top_k=2)
mock_collection.find.assert_called_once_with({})
mock_cursor.limit.assert_called_once_with(2)
@@ -334,7 +334,7 @@ def test_list_with_filters(mongo_vector_fixture):
]
filters = {"user_id": "alice", "agent_id": "agent1", "run_id": "run1"}
results = mongo_vector.list(filters=filters, limit=2)
results = mongo_vector.list(filters=filters, top_k=2)
# Verify that the find method was called with the correct query
expected_query = {
@@ -363,7 +363,7 @@ def test_list_with_single_filter(mongo_vector_fixture):
]
filters = {"user_id": "alice"}
results = mongo_vector.list(filters=filters, limit=2)
results = mongo_vector.list(filters=filters, top_k=2)
# Verify that the find method was called with the correct query
expected_query = {
@@ -387,7 +387,7 @@ def test_list_with_no_filters(mongo_vector_fixture):
{"_id": "id1", "payload": {"key": "value1"}},
]
results = mongo_vector.list(filters=None, limit=2)
results = mongo_vector.list(filters=None, top_k=2)
# Verify that the find method was called with empty query
mock_collection.find.assert_called_once_with({})
tests/vector_stores/test_neptune_analytics.py
+1 -1
View File
@@ -123,7 +123,7 @@ class TestNeptuneAnalyticsOperations:
payloads=SAMPLE_PAYLOADS
)
result = na_instance.search(query="", vectors=VECTOR_1, limit=1)
result = na_instance.search(query="", vectors=VECTOR_1, top_k=1)
assert len(result) == 1
assert "label" not in result[0].payload
tests/vector_stores/test_opensearch.py
+2 -2
View File
@@ -229,7 +229,7 @@ class TestOpenSearchDB(unittest.TestCase):
}
self.client_mock.search.return_value = mock_response
vectors = [[0.1] * 1536]
results = self.os_db.search(query="", vectors=vectors, limit=5)
results = self.os_db.search(query="", vectors=vectors, top_k=5)
self.client_mock.search.assert_called_once()
search_args = self.client_mock.search.call_args[1]
self.assertEqual(search_args["index"], "test_collection")
@@ -345,7 +345,7 @@ class TestOpenSearchDB(unittest.TestCase):
def test_search_error_logs_with_exc_info(self, mock_logger):
"""Search error logging should include exc_info for full stack trace."""
self.client_mock.search.side_effect = Exception("Search failed")
results = self.os_db.search(query="", vectors=[[0.1] * 1536], limit=5)
results = self.os_db.search(query="", vectors=[[0.1] * 1536], top_k=5)
self.assertEqual(results, [])
mock_logger.error.assert_called_once()
call_kwargs = mock_logger.error.call_args
tests/vector_stores/test_pgvector.py
+16 -16
View File
@@ -432,7 +432,7 @@ class TestPGVector(unittest.TestCase):
maxconn=4
)
results = pgvector.search("test query", [0.1, 0.2, 0.3], limit=2)
results = pgvector.search("test query", [0.1, 0.2, 0.3], top_k=2)
# Verify the _get_cursor context manager was called
mock_get_cursor.assert_called()
@@ -481,7 +481,7 @@ class TestPGVector(unittest.TestCase):
maxconn=4
)
results = pgvector.search("test query", [0.1, 0.2, 0.3], limit=2)
results = pgvector.search("test query", [0.1, 0.2, 0.3], top_k=2)
# Verify the _get_cursor context manager was called
mock_get_cursor.assert_called()
@@ -1031,7 +1031,7 @@ class TestPGVector(unittest.TestCase):
maxconn=4
)
results = pgvector.list(limit=2)
results = pgvector.list(top_k=2)
# Verify the _get_cursor context manager was called
mock_get_cursor.assert_called()
@@ -1079,7 +1079,7 @@ class TestPGVector(unittest.TestCase):
maxconn=4
)
results = pgvector.list(limit=2)
results = pgvector.list(top_k=2)
# Verify the _get_cursor context manager was called
mock_get_cursor.assert_called()
@@ -1127,7 +1127,7 @@ class TestPGVector(unittest.TestCase):
)
filters = {"user_id": "alice", "agent_id": "agent1", "run_id": "run1"}
results = pgvector.search("test query", [0.1, 0.2, 0.3], limit=2, filters=filters)
results = pgvector.search("test query", [0.1, 0.2, 0.3], top_k=2, filters=filters)
# Verify the _get_cursor context manager was called
mock_get_cursor.assert_called()
@@ -1177,7 +1177,7 @@ class TestPGVector(unittest.TestCase):
)
filters = {"user_id": "alice", "agent_id": "agent1", "run_id": "run1"}
results = pgvector.search("test query", [0.1, 0.2, 0.3], limit=2, filters=filters)
results = pgvector.search("test query", [0.1, 0.2, 0.3], top_k=2, filters=filters)
# Verify the _get_cursor context manager was called
mock_get_cursor.assert_called()
@@ -1227,7 +1227,7 @@ class TestPGVector(unittest.TestCase):
)
filters = {"user_id": "alice"}
results = pgvector.search("test query", [0.1, 0.2, 0.3], limit=2, filters=filters)
results = pgvector.search("test query", [0.1, 0.2, 0.3], top_k=2, filters=filters)
# Verify the _get_cursor context manager was called
mock_get_cursor.assert_called()
@@ -1275,7 +1275,7 @@ class TestPGVector(unittest.TestCase):
)
filters = {"user_id": "alice"}
results = pgvector.search("test query", [0.1, 0.2, 0.3], limit=2, filters=filters)
results = pgvector.search("test query", [0.1, 0.2, 0.3], top_k=2, filters=filters)
# Verify the _get_cursor context manager was called
mock_get_cursor.assert_called()
@@ -1323,7 +1323,7 @@ class TestPGVector(unittest.TestCase):
maxconn=4
)
results = pgvector.search("test query", [0.1, 0.2, 0.3], limit=2, filters=None)
results = pgvector.search("test query", [0.1, 0.2, 0.3], top_k=2, filters=None)
# Verify the _get_cursor context manager was called
mock_get_cursor.assert_called()
@@ -1372,7 +1372,7 @@ class TestPGVector(unittest.TestCase):
maxconn=4
)
results = pgvector.search("test query", [0.1, 0.2, 0.3], limit=2, filters=None)
results = pgvector.search("test query", [0.1, 0.2, 0.3], top_k=2, filters=None)
# Verify the _get_cursor context manager was called
mock_get_cursor.assert_called()
@@ -1421,7 +1421,7 @@ class TestPGVector(unittest.TestCase):
)
filters = {"user_id": "alice", "agent_id": "agent1"}
results = pgvector.list(filters=filters, limit=2)
results = pgvector.list(filters=filters, top_k=2)
# Verify the _get_cursor context manager was called
mock_get_cursor.assert_called()
@@ -1470,7 +1470,7 @@ class TestPGVector(unittest.TestCase):
)
filters = {"user_id": "alice", "agent_id": "agent1"}
results = pgvector.list(filters=filters, limit=2)
results = pgvector.list(filters=filters, top_k=2)
# Verify the _get_cursor context manager was called
mock_get_cursor.assert_called()
@@ -1519,7 +1519,7 @@ class TestPGVector(unittest.TestCase):
)
filters = {"user_id": "alice"}
results = pgvector.list(filters=filters, limit=2)
results = pgvector.list(filters=filters, top_k=2)
# Verify the _get_cursor context manager was called
mock_get_cursor.assert_called()
@@ -1567,7 +1567,7 @@ class TestPGVector(unittest.TestCase):
)
filters = {"user_id": "alice"}
results = pgvector.list(filters=filters, limit=2)
results = pgvector.list(filters=filters, top_k=2)
# Verify the _get_cursor context manager was called
mock_get_cursor.assert_called()
@@ -1615,7 +1615,7 @@ class TestPGVector(unittest.TestCase):
maxconn=4
)
results = pgvector.list(filters=None, limit=2)
results = pgvector.list(filters=None, top_k=2)
# Verify the _get_cursor context manager was called
mock_get_cursor.assert_called()
@@ -1663,7 +1663,7 @@ class TestPGVector(unittest.TestCase):
maxconn=4
)
results = pgvector.list(filters=None, limit=2)
results = pgvector.list(filters=None, top_k=2)
# Verify the _get_cursor context manager was called
mock_get_cursor.assert_called()
tests/vector_stores/test_pinecone.py
+1 -1
View File
@@ -80,7 +80,7 @@ def test_insert_vectors(pinecone_db):
def test_search_vectors(pinecone_db):
pinecone_db.index.query.return_value.matches = [{"id": "id1", "score": 0.9, "metadata": {"name": "vector1"}}]
results = pinecone_db.search("test query", [0.1] * 128, limit=1)
results = pinecone_db.search("test query", [0.1] * 128, top_k=1)
pinecone_db.index.query.assert_called_with(
vector=[0.1] * 128,
top_k=1,
tests/vector_stores/test_qdrant.py
+7 -7
View File
@@ -84,7 +84,7 @@ class TestQdrant(unittest.TestCase):
mock_point = MagicMock(id=str(uuid.uuid4()), score=0.95, payload={"key": "value"})
self.client_mock.query_points.return_value = MagicMock(points=[mock_point])
results = self.qdrant.search(query="", vectors=vectors, limit=1)
results = self.qdrant.search(query="", vectors=vectors, top_k=1)
self.client_mock.query_points.assert_called_once_with(
collection_name="test_collection",
@@ -108,7 +108,7 @@ class TestQdrant(unittest.TestCase):
self.client_mock.query_points.return_value = MagicMock(points=[mock_point])
filters = {"user_id": "alice", "agent_id": "agent1", "run_id": "run1"}
results = self.qdrant.search(query="", vectors=vectors, limit=1, filters=filters)
results = self.qdrant.search(query="", vectors=vectors, top_k=1, filters=filters)
# Verify that _create_filter was called and query_filter was passed
self.client_mock.query_points.assert_called_once()
@@ -138,7 +138,7 @@ class TestQdrant(unittest.TestCase):
self.client_mock.query_points.return_value = MagicMock(points=[mock_point])
filters = {"user_id": "alice"}
results = self.qdrant.search(query="", vectors=vectors, limit=1, filters=filters)
results = self.qdrant.search(query="", vectors=vectors, top_k=1, filters=filters)
# Verify that a Filter object was created with single condition
call_args = self.client_mock.query_points.call_args[1]
@@ -155,7 +155,7 @@ class TestQdrant(unittest.TestCase):
mock_point = MagicMock(id=str(uuid.uuid4()), score=0.95, payload={"key": "value"})
self.client_mock.query_points.return_value = MagicMock(points=[mock_point])
results = self.qdrant.search(query="", vectors=vectors, limit=1, filters=None)
results = self.qdrant.search(query="", vectors=vectors, top_k=1, filters=None)
call_args = self.client_mock.query_points.call_args[1]
self.assertIsNone(call_args["query_filter"])
@@ -298,7 +298,7 @@ class TestQdrant(unittest.TestCase):
self.client_mock.scroll.return_value = [mock_point]
filters = {"user_id": "alice", "agent_id": "agent1", "run_id": "run1"}
results = self.qdrant.list(filters=filters, limit=10)
results = self.qdrant.list(filters=filters, top_k=10)
# Verify that _create_filter was called and scroll_filter was passed
self.client_mock.scroll.assert_called_once()
@@ -327,7 +327,7 @@ class TestQdrant(unittest.TestCase):
self.client_mock.scroll.return_value = [mock_point]
filters = {"user_id": "alice"}
results = self.qdrant.list(filters=filters, limit=10)
results = self.qdrant.list(filters=filters, top_k=10)
# Verify that a Filter object was created with single condition
call_args = self.client_mock.scroll.call_args[1]
@@ -344,7 +344,7 @@ class TestQdrant(unittest.TestCase):
mock_point = MagicMock(id=str(uuid.uuid4()), score=0.95, payload={"key": "value"})
self.client_mock.scroll.return_value = [mock_point]
results = self.qdrant.list(filters=None, limit=10)
results = self.qdrant.list(filters=None, top_k=10)
call_args = self.client_mock.scroll.call_args[1]
self.assertIsNone(call_args["scroll_filter"])
tests/vector_stores/test_s3_vectors.py
+2 -2
View File
@@ -1,7 +1,7 @@
from mem0.configs.vector_stores.s3_vectors import S3VectorsConfig
import pytest
from botocore.exceptions import ClientError
from mem0.configs.vector_stores.s3_vectors import S3VectorsConfig
from mem0.memory.main import Memory
from mem0.vector_stores.s3_vectors import S3Vectors
@@ -152,7 +152,7 @@ def test_search(mock_boto_client):
embedding_model_dims=EMBEDDING_DIMS,
)
query_vector = [0.1, 0.2]
results = store.search(query="test", vectors=query_vector, limit=1)
results = store.search(query="test", vectors=query_vector, top_k=1)
mock_boto_client.query_vectors.assert_called_once()
assert len(results) == 1
tests/vector_stores/test_supabase.py
+2 -2
View File
@@ -67,7 +67,7 @@ def test_search_vectors(supabase_instance, mock_collection):
vectors = [[0.1, 0.2, 0.3]]
filters = {"category": "test"}
results = supabase_instance.search(query="", vectors=vectors, limit=2, filters=filters)
results = supabase_instance.search(query="", vectors=vectors, top_k=2, filters=filters)
mock_collection.query.assert_called_once_with(
data=vectors, limit=2, filters={"category": {"$eq": "test"}}, include_metadata=True, include_value=True
@@ -118,7 +118,7 @@ def test_list_vectors(supabase_instance, mock_collection):
mock_collection.query.return_value = mock_query_results
mock_collection.fetch.return_value = mock_fetch_results
results = supabase_instance.list(limit=2, filters={"category": "test"})
results = supabase_instance.list(top_k=2, filters={"category": "test"})
assert len(results[0]) == 2
assert results[0][0].id == "id1"
tests/vector_stores/test_turbopuffer.py
+3 -3
View File
@@ -289,7 +289,7 @@ class TestSearch:
]
db.namespace.query.return_value = mock_response
results = db.search("test query", [0.1, 0.2, 0.3, 0.4], limit=2)
results = db.search("test query", [0.1, 0.2, 0.3, 0.4], top_k=2)
db.namespace.query.assert_called_once_with(
rank_by=("vector", "ANN", [0.1, 0.2, 0.3, 0.4]),
@@ -306,7 +306,7 @@ class TestSearch:
db.namespace.query.return_value = mock_response
results = db.search(
"query", [0.1, 0.2, 0.3, 0.4], limit=5, filters={"user_id": "u1"}
"query", [0.1, 0.2, 0.3, 0.4], top_k=5, filters={"user_id": "u1"}
)
call_kwargs = db.namespace.query.call_args[1]
@@ -515,7 +515,7 @@ class TestList:
mock_response.rows = [_make_row("id1", dist=0.1, data="hello")]
db.namespace.query.return_value = mock_response
db.list(filters={"user_id": "u1"}, limit=50)
db.list(filters={"user_id": "u1"}, top_k=50)
call_kwargs = db.namespace.query.call_args[1]
assert call_kwargs["filters"] == ("user_id", "Eq", "u1")
tests/vector_stores/test_upstash_vector.py
+4 -4
View File
@@ -60,7 +60,7 @@ def test_search_vectors(upstash_instance, mock_index):
results = upstash_instance.search(
query="hello world",
vectors=vectors,
limit=2,
top_k=2,
filters={"age": 30, "name": "John"},
)
@@ -132,7 +132,7 @@ def test_list_vectors(upstash_instance):
filters = {"age": 30, "name": "John"}
print("filters", filters)
[results] = upstash_instance.list(filters=filters, limit=15)
[results] = upstash_instance.list(filters=filters, top_k=15)
upstash_instance.client.info.return_value = {
"dimension": 10,
@@ -193,7 +193,7 @@ def test_search_vectors_with_embeddings(upstash_instance_with_embeddings, mock_i
results = upstash_instance_with_embeddings.search(
query="hello world",
vectors=[],
limit=2,
top_k=2,
filters={"age": 30, "name": "John"},
)
@@ -302,7 +302,7 @@ def test_search_vectors_empty_filters(upstash_instance):
results = upstash_instance.search(
query="hello world",
vectors=vectors,
limit=1,
top_k=1,
filters=None,
)
tests/vector_stores/test_valkey.py
+7 -7
View File
@@ -59,7 +59,7 @@ def test_search_filter_syntax(valkey_db, mock_valkey_client):
valkey_db.search(
query="test query",
vectors=np.random.rand(1536).tolist(),
limit=5,
top_k=5,
filters={"user_id": "test_user"},
)
@@ -72,7 +72,7 @@ def test_search_filter_syntax(valkey_db, mock_valkey_client):
valkey_db.search(
query="test query",
vectors=np.random.rand(1536).tolist(),
limit=5,
top_k=5,
filters={"user_id": "test_user", "agent_id": "test_agent"},
)
@@ -104,7 +104,7 @@ def test_search_without_filters(valkey_db, mock_valkey_client):
results = valkey_db.search(
query="test query",
vectors=np.random.rand(1536).tolist(),
limit=5,
top_k=5,
)
# Check that the search was called with the correct syntax
@@ -405,7 +405,7 @@ def test_list(valkey_db, mock_valkey_client):
mock_ft.search.return_value = mock_results
# Call list
results = valkey_db.list(filters={"user_id": "test_user"}, limit=10)
results = valkey_db.list(filters={"user_id": "test_user"}, top_k=10)
# Check that search was called with the correct arguments
mock_ft.search.assert_called_once()
@@ -437,7 +437,7 @@ def test_search_error_handling(valkey_db, mock_valkey_client):
valkey_db.search(
query="test query",
vectors=np.random.rand(1536).tolist(),
limit=5,
top_k=5,
filters={"user_id": "test_user"},
)
@@ -776,7 +776,7 @@ def test_search_with_invalid_metadata(valkey_db, mock_valkey_client):
mock_valkey_client.ft.return_value.search.return_value = mock_result
# Should handle invalid JSON gracefully
results = valkey_db.search(query="test query", vectors=np.random.rand(1536).tolist(), limit=5)
results = valkey_db.search(query="test query", vectors=np.random.rand(1536).tolist(), top_k=5)
assert len(results) == 1
@@ -790,7 +790,7 @@ def test_search_with_hnsw_ef_runtime(valkey_db, mock_valkey_client):
mock_result.docs = []
mock_valkey_client.ft.return_value.search.return_value = mock_result
valkey_db.search(query="test query", vectors=np.random.rand(1536).tolist(), limit=5)
valkey_db.search(query="test query", vectors=np.random.rand(1536).tolist(), top_k=5)
# Verify the search was called
assert mock_valkey_client.ft.return_value.search.called
tests/vector_stores/test_vertex_ai_vector_search.py
+1 -1
View File
@@ -100,7 +100,7 @@ def test_search_vectors(vector_store, mock_vertex_ai):
mock_vertex_ai["endpoint"].find_neighbors.return_value = [[mock_neighbor]]
results = vector_store.search(query="", vectors=vectors, filters=filters, limit=1)
results = vector_store.search(query="", vectors=vectors, filters=filters, top_k=1)
mock_vertex_ai["endpoint"].find_neighbors.assert_called_once_with(
deployed_index_id=vector_store.deployment_index_id,
tests/vector_stores/test_weaviate.py
+4 -4
View File
@@ -1,10 +1,10 @@
import os
import uuid
import httpx
import unittest
import uuid
from unittest.mock import MagicMock, patch
import dotenv
import httpx
import weaviate
from weaviate.exceptions import UnexpectedStatusCodeException
@@ -137,7 +137,7 @@ class TestWeaviateDB(unittest.TestCase):
mock_hybrid.return_value = mock_response
vectors = [[0.1] * 1536]
results = self.weaviate_db.search(query="", vectors=vectors, limit=5)
results = self.weaviate_db.search(query="", vectors=vectors, top_k=5)
mock_hybrid.assert_called_once()
@@ -170,7 +170,7 @@ class TestWeaviateDB(unittest.TestCase):
self.client_mock.collections.get.return_value.query.fetch_objects = mock_fetch
mock_fetch.return_value = mock_response
results = self.weaviate_db.list(limit=10)
results = self.weaviate_db.list(top_k=10)
mock_fetch.assert_called_once()