Adding Native OpenSearch support for Mem0 (#2211)

tests/vector_stores/test_opensearch.py

@@ -0,0 +1,150 @@
+import os
+import unittest
+from unittest.mock import MagicMock, patch
+
+import dotenv
+
+try:
+    from opensearchpy import OpenSearch
+except ImportError:
+    raise ImportError(
+        "OpenSearch requires extra dependencies. Install with `pip install opensearch-py`"
+    ) from None
+
+from mem0.vector_stores.opensearch import OpenSearchDB
+
+
+class TestOpenSearchDB(unittest.TestCase):
+    @classmethod
+    def setUpClass(cls):
+        dotenv.load_dotenv()
+        cls.original_env = {
+            'OS_URL': os.getenv('OS_URL', 'http://localhost:9200'),
+            'OS_USERNAME': os.getenv('OS_USERNAME', 'test_user'),
+            'OS_PASSWORD': os.getenv('OS_PASSWORD', 'test_password')
+        }
+        os.environ['OS_URL'] = 'http://localhost'
+        os.environ['OS_USERNAME'] = 'test_user'
+        os.environ['OS_PASSWORD'] = 'test_password'
+
+    def setUp(self):
+        self.client_mock = MagicMock(spec=OpenSearch)
+        self.client_mock.indices = MagicMock()
+        self.client_mock.indices.exists = MagicMock(return_value=False)
+        self.client_mock.indices.create = MagicMock()
+        self.client_mock.indices.delete = MagicMock()
+        self.client_mock.indices.get_alias = MagicMock()
+        self.client_mock.get = MagicMock()
+        self.client_mock.update = MagicMock()
+        self.client_mock.delete = MagicMock()
+        self.client_mock.search = MagicMock()
+
+        patcher = patch('mem0.vector_stores.opensearch.OpenSearch', return_value=self.client_mock)
+        self.mock_os = patcher.start()
+        self.addCleanup(patcher.stop)
+
+        self.os_db = OpenSearchDB(
+            host=os.getenv('OS_URL'),
+            port=9200,
+            collection_name="test_collection",
+            embedding_model_dims=1536,
+            user=os.getenv('OS_USERNAME'),
+            password=os.getenv('OS_PASSWORD'),
+            verify_certs=False,
+            use_ssl=False,
+            auto_create_index=False
+        )
+        self.client_mock.reset_mock()
+
+    @classmethod
+    def tearDownClass(cls):
+        for key, value in cls.original_env.items():
+            if value is not None:
+                os.environ[key] = value
+            else:
+                os.environ.pop(key, None)
+
+    def tearDown(self):
+        self.client_mock.reset_mock()
+
+    def test_create_index(self):
+        self.client_mock.indices.exists.return_value = False
+        self.os_db.create_index()
+        self.client_mock.indices.create.assert_called_once()
+        create_args = self.client_mock.indices.create.call_args[1]
+        self.assertEqual(create_args["index"], "test_collection")
+        mappings = create_args["body"]["mappings"]["properties"]
+        self.assertEqual(mappings["vector"]["type"], "knn_vector")
+        self.assertEqual(mappings["vector"]["dimension"], 1536)
+        self.client_mock.reset_mock()
+        self.client_mock.indices.exists.return_value = True
+        self.os_db.create_index()
+        self.client_mock.indices.create.assert_not_called()
+
+    def test_insert(self):
+        vectors = [[0.1] * 1536, [0.2] * 1536]
+        payloads = [{"key1": "value1"}, {"key2": "value2"}]
+        ids = ["id1", "id2"]
+        with patch('mem0.vector_stores.opensearch.bulk') as mock_bulk:
+            mock_bulk.return_value = (2, [])
+            results = self.os_db.insert(vectors=vectors, payloads=payloads, ids=ids)
+            mock_bulk.assert_called_once()
+            actions = mock_bulk.call_args[0][1]
+            self.assertEqual(actions[0]["_index"], "test_collection")
+            self.assertEqual(actions[0]["_id"], "id1")
+            self.assertEqual(actions[0]["_source"]["vector"], vectors[0])
+            self.assertEqual(actions[0]["_source"]["metadata"], payloads[0])
+            self.assertEqual(len(results), 2)
+            self.assertEqual(results[0].id, "id1")
+            self.assertEqual(results[0].payload, payloads[0])
+    
+    def test_get(self):
+        mock_response = {"_id": "id1", "_source": {"metadata": {"key1": "value1"}}}
+        self.client_mock.get.return_value = mock_response
+        result = self.os_db.get("id1")
+        self.client_mock.get.assert_called_once_with(index="test_collection", id="id1")
+        self.assertIsNotNone(result)
+        self.assertEqual(result.id, "id1")
+        self.assertEqual(result.payload, {"key1": "value1"})
+    
+    def test_update(self):
+        vector = [0.3] * 1536
+        payload = {"key3": "value3"}
+        self.os_db.update("id1", vector=vector, payload=payload)
+        self.client_mock.update.assert_called_once()
+        update_args = self.client_mock.update.call_args[1]
+        self.assertEqual(update_args["index"], "test_collection")
+        self.assertEqual(update_args["id"], "id1")
+        self.assertEqual(update_args["body"], {"doc": {"vector": vector, "metadata": payload}})
+
+    def test_list_cols(self):
+        self.client_mock.indices.get_alias.return_value = {"test_collection": {}}
+        result = self.os_db.list_cols()
+        self.client_mock.indices.get_alias.assert_called_once()
+        self.assertEqual(result, ["test_collection"])
+
+    def test_search(self):
+        mock_response = {"hits": {"hits": [{"_id": "id1", "_score": 0.8, "_source": {"vector": [0.1] * 1536, "metadata": {"key1": "value1"}}}]}}
+        self.client_mock.search.return_value = mock_response
+        query_vector = [0.1] * 1536
+        results = self.os_db.search(query=query_vector, limit=5)
+        self.client_mock.search.assert_called_once()
+        search_args = self.client_mock.search.call_args[1]
+        self.assertEqual(search_args["index"], "test_collection")
+        body = search_args["body"]
+        self.assertIn("knn", body["query"])
+        self.assertIn("vector", body["query"]["knn"])
+        self.assertEqual(body["query"]["knn"]["vector"]["vector"], query_vector)
+        self.assertEqual(body["query"]["knn"]["vector"]["k"], 5)
+        self.assertEqual(len(results), 1)
+        self.assertEqual(results[0].id, "id1")
+        self.assertEqual(results[0].score, 0.8)
+        self.assertEqual(results[0].payload, {"key1": "value1"})
+
+    def test_delete(self):
+        self.os_db.delete(vector_id="id1")
+        self.client_mock.delete.assert_called_once_with(index="test_collection", id="id1")
+
+    def test_delete_col(self):
+        self.os_db.delete_col()
+        self.client_mock.indices.delete.assert_called_once_with(index="test_collection")