Added mulit id filters support for all vectorstores (#3269)

2025-08-04 14:52:57 -07:00
parent 57a16aeb4b
commit 42e60d6724
8 changed files with 728 additions and 45 deletions
@@ -3,7 +3,13 @@ import uuid
 from unittest.mock import MagicMock

 from qdrant_client import QdrantClient
-from qdrant_client.models import Distance, PointIdsList, PointStruct, VectorParams
+from qdrant_client.models import (
+    Distance,
+    Filter,
+    PointIdsList,
+    PointStruct,
+    VectorParams,
+)

 from mem0.vector_stores.qdrant import Qdrant

@@ -64,6 +70,121 @@ class TestQdrant(unittest.TestCase):
        self.assertEqual(results[0].payload, {"key": "value"})
        self.assertEqual(results[0].score, 0.95)

+    def test_search_with_filters(self):
+        """Test search with agent_id and run_id filters."""
+        vectors = [[0.1, 0.2]]
+        mock_point = MagicMock(
+            id=str(uuid.uuid4()), 
+            score=0.95, 
+            payload={"user_id": "alice", "agent_id": "agent1", "run_id": "run1"}
+        )
+        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)
+
+        # Verify that _create_filter was called and query_filter was passed
+        self.client_mock.query_points.assert_called_once()
+        call_args = self.client_mock.query_points.call_args[1]
+        self.assertEqual(call_args["collection_name"], "test_collection")
+        self.assertEqual(call_args["query"], vectors)
+        self.assertEqual(call_args["limit"], 1)
+        
+        # Verify that a Filter object was created
+        query_filter = call_args["query_filter"]
+        self.assertIsInstance(query_filter, Filter)
+        self.assertEqual(len(query_filter.must), 3)  # user_id, agent_id, run_id
+
+        self.assertEqual(len(results), 1)
+        self.assertEqual(results[0].payload["user_id"], "alice")
+        self.assertEqual(results[0].payload["agent_id"], "agent1")
+        self.assertEqual(results[0].payload["run_id"], "run1")
+
+    def test_search_with_single_filter(self):
+        """Test search with single filter."""
+        vectors = [[0.1, 0.2]]
+        mock_point = MagicMock(
+            id=str(uuid.uuid4()), 
+            score=0.95, 
+            payload={"user_id": "alice"}
+        )
+        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)
+
+        # Verify that a Filter object was created with single condition
+        call_args = self.client_mock.query_points.call_args[1]
+        query_filter = call_args["query_filter"]
+        self.assertIsInstance(query_filter, Filter)
+        self.assertEqual(len(query_filter.must), 1)  # Only user_id
+
+        self.assertEqual(len(results), 1)
+        self.assertEqual(results[0].payload["user_id"], "alice")
+
+    def test_search_with_no_filters(self):
+        """Test search with no filters."""
+        vectors = [[0.1, 0.2]]
+        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)
+
+        call_args = self.client_mock.query_points.call_args[1]
+        self.assertIsNone(call_args["query_filter"])
+
+        self.assertEqual(len(results), 1)
+
+    def test_create_filter_multiple_filters(self):
+        """Test _create_filter with multiple filters."""
+        filters = {"user_id": "alice", "agent_id": "agent1", "run_id": "run1"}
+        result = self.qdrant._create_filter(filters)
+        
+        self.assertIsInstance(result, Filter)
+        self.assertEqual(len(result.must), 3)
+        
+        # Check that all conditions are present
+        conditions = [cond.key for cond in result.must]
+        self.assertIn("user_id", conditions)
+        self.assertIn("agent_id", conditions)
+        self.assertIn("run_id", conditions)
+
+    def test_create_filter_single_filter(self):
+        """Test _create_filter with single filter."""
+        filters = {"user_id": "alice"}
+        result = self.qdrant._create_filter(filters)
+        
+        self.assertIsInstance(result, Filter)
+        self.assertEqual(len(result.must), 1)
+        self.assertEqual(result.must[0].key, "user_id")
+        self.assertEqual(result.must[0].match.value, "alice")
+
+    def test_create_filter_no_filters(self):
+        """Test _create_filter with no filters."""
+        result = self.qdrant._create_filter(None)
+        self.assertIsNone(result)
+        
+        result = self.qdrant._create_filter({})
+        self.assertIsNone(result)
+
+    def test_create_filter_with_range_values(self):
+        """Test _create_filter with range values."""
+        filters = {"user_id": "alice", "count": {"gte": 5, "lte": 10}}
+        result = self.qdrant._create_filter(filters)
+        
+        self.assertIsInstance(result, Filter)
+        self.assertEqual(len(result.must), 2)
+        
+        # Check that range condition is created
+        range_conditions = [cond for cond in result.must if hasattr(cond, 'range') and cond.range is not None]
+        self.assertEqual(len(range_conditions), 1)
+        self.assertEqual(range_conditions[0].key, "count")
+        
+        # Check that string condition is created
+        string_conditions = [cond for cond in result.must if hasattr(cond, 'match') and cond.match is not None]
+        self.assertEqual(len(string_conditions), 1)
+        self.assertEqual(string_conditions[0].key, "user_id")
+
    def test_delete(self):
        vector_id = str(uuid.uuid4())
        self.qdrant.delete(vector_id=vector_id)
@@ -103,6 +224,67 @@ class TestQdrant(unittest.TestCase):
        result = self.qdrant.list_cols()
        self.assertEqual(result.collections[0]["name"], "test_collection")

+    def test_list_with_filters(self):
+        """Test list with agent_id and run_id filters."""
+        mock_point = MagicMock(
+            id=str(uuid.uuid4()), 
+            score=0.95, 
+            payload={"user_id": "alice", "agent_id": "agent1", "run_id": "run1"}
+        )
+        self.client_mock.scroll.return_value = MagicMock(points=[mock_point])
+
+        filters = {"user_id": "alice", "agent_id": "agent1", "run_id": "run1"}
+        results = self.qdrant.list(filters=filters, limit=10)
+
+        # Verify that _create_filter was called and scroll_filter was passed
+        self.client_mock.scroll.assert_called_once()
+        call_args = self.client_mock.scroll.call_args[1]
+        self.assertEqual(call_args["collection_name"], "test_collection")
+        self.assertEqual(call_args["limit"], 10)
+        
+        # Verify that a Filter object was created
+        scroll_filter = call_args["scroll_filter"]
+        self.assertIsInstance(scroll_filter, Filter)
+        self.assertEqual(len(scroll_filter.must), 3)  # user_id, agent_id, run_id
+
+        self.assertEqual(len(results), 1)
+        self.assertEqual(results[0].payload["user_id"], "alice")
+        self.assertEqual(results[0].payload["agent_id"], "agent1")
+        self.assertEqual(results[0].payload["run_id"], "run1")
+
+    def test_list_with_single_filter(self):
+        """Test list with single filter."""
+        mock_point = MagicMock(
+            id=str(uuid.uuid4()), 
+            score=0.95, 
+            payload={"user_id": "alice"}
+        )
+        self.client_mock.scroll.return_value = MagicMock(points=[mock_point])
+
+        filters = {"user_id": "alice"}
+        results = self.qdrant.list(filters=filters, limit=10)
+
+        # Verify that a Filter object was created with single condition
+        call_args = self.client_mock.scroll.call_args[1]
+        scroll_filter = call_args["scroll_filter"]
+        self.assertIsInstance(scroll_filter, Filter)
+        self.assertEqual(len(scroll_filter.must), 1)  # Only user_id
+
+        self.assertEqual(len(results), 1)
+        self.assertEqual(results[0].payload["user_id"], "alice")
+
+    def test_list_with_no_filters(self):
+        """Test list with no filters."""
+        mock_point = MagicMock(id=str(uuid.uuid4()), score=0.95, payload={"key": "value"})
+        self.client_mock.scroll.return_value = MagicMock(points=[mock_point])
+
+        results = self.qdrant.list(filters=None, limit=10)
+
+        call_args = self.client_mock.scroll.call_args[1]
+        self.assertIsNone(call_args["scroll_filter"])
+
+        self.assertEqual(len(results), 1)
+
    def test_delete_col(self):
        self.qdrant.delete_col()
        self.client_mock.delete_collection.assert_called_once_with(collection_name="test_collection")