feat: support opensearch approximate k-NN (#5322)

api/tests/integration_tests/vdb/opensearch/test_opensearch.py

@@ -0,0 +1,186 @@
+from unittest.mock import MagicMock, patch
+
+import pytest
+
+from core.rag.datasource.vdb.field import Field
+from core.rag.datasource.vdb.opensearch.opensearch_vector import OpenSearchConfig, OpenSearchVector
+from core.rag.models.document import Document
+from extensions import ext_redis
+
+
+def get_example_text() -> str:
+    return "This is a sample text for testing purposes."
+
+
+@pytest.fixture(scope="module")
+def setup_mock_redis():
+    ext_redis.redis_client.get = MagicMock(return_value=None)
+    ext_redis.redis_client.set = MagicMock(return_value=None)
+
+    mock_redis_lock = MagicMock()
+    mock_redis_lock.__enter__ = MagicMock()
+    mock_redis_lock.__exit__ = MagicMock()
+    ext_redis.redis_client.lock = MagicMock(return_value=mock_redis_lock)
+
+
+class TestOpenSearchVector:
+    def setup_method(self):
+        self.collection_name = "test_collection"
+        self.example_doc_id = "example_doc_id"
+        self.vector = OpenSearchVector(
+            collection_name=self.collection_name,
+            config=OpenSearchConfig(
+                host='localhost',
+                port=9200,
+                user='admin',
+                password='password',
+                secure=False
+            )
+        )
+        self.vector._client = MagicMock()
+
+    @pytest.mark.parametrize("search_response, expected_length, expected_doc_id", [
+        ({
+            'hits': {
+                'total': {'value': 1},
+                'hits': [
+                    {'_source': {'page_content': get_example_text(), 'metadata': {"document_id": "example_doc_id"}}}
+                ]
+            }
+        }, 1, "example_doc_id"),
+        ({
+            'hits': {
+                'total': {'value': 0},
+                'hits': []
+            }
+        }, 0, None)
+    ])
+    def test_search_by_full_text(self, search_response, expected_length, expected_doc_id):
+        self.vector._client.search.return_value = search_response
+
+        hits_by_full_text = self.vector.search_by_full_text(query=get_example_text())
+        assert len(hits_by_full_text) == expected_length
+        if expected_length > 0:
+            assert hits_by_full_text[0].metadata['document_id'] == expected_doc_id
+
+    def test_search_by_vector(self):
+        vector = [0.1] * 128
+        mock_response = {
+            'hits': {
+                'total': {'value': 1},
+                'hits': [
+                    {
+                        '_source': {
+                            Field.CONTENT_KEY.value: get_example_text(),
+                            Field.METADATA_KEY.value: {"document_id": self.example_doc_id}
+                        },
+                        '_score': 1.0
+                    }
+                ]
+            }
+        }
+        self.vector._client.search.return_value = mock_response
+
+        hits_by_vector = self.vector.search_by_vector(query_vector=vector)
+
+        print("Hits by vector:", hits_by_vector)
+        print("Expected document ID:", self.example_doc_id)
+        print("Actual document ID:", hits_by_vector[0].metadata['document_id'] if hits_by_vector else "No hits")
+
+        assert len(hits_by_vector) > 0, f"Expected at least one hit, got {len(hits_by_vector)}"
+        assert hits_by_vector[0].metadata['document_id'] == self.example_doc_id, \
+            f"Expected document ID {self.example_doc_id}, got {hits_by_vector[0].metadata['document_id']}"
+
+    def test_delete_by_document_id(self):
+        self.vector._client.delete_by_query.return_value = {'deleted': 1}
+
+        doc = Document(page_content="Test content to delete", metadata={"document_id": self.example_doc_id})
+        embedding = [0.1] * 128
+
+        with patch('opensearchpy.helpers.bulk') as mock_bulk:
+            mock_bulk.return_value = ([], [])
+            self.vector.add_texts([doc], [embedding])
+
+        self.vector.delete_by_document_id(document_id=self.example_doc_id)
+
+        self.vector._client.search.return_value = {'hits': {'total': {'value': 0}, 'hits': []}}
+
+        ids = self.vector.get_ids_by_metadata_field(key='document_id', value=self.example_doc_id)
+        assert ids is None or len(ids) == 0
+
+    def test_get_ids_by_metadata_field(self):
+        mock_response = {
+            'hits': {
+                'total': {'value': 1},
+                'hits': [{'_id': 'mock_id'}]
+            }
+        }
+        self.vector._client.search.return_value = mock_response
+
+        doc = Document(page_content="Test content", metadata={"document_id": self.example_doc_id})
+        embedding = [0.1] * 128
+
+        with patch('opensearchpy.helpers.bulk') as mock_bulk:
+            mock_bulk.return_value = ([], [])
+            self.vector.add_texts([doc], [embedding])
+
+        ids = self.vector.get_ids_by_metadata_field(key='document_id', value=self.example_doc_id)
+        assert len(ids) == 1
+        assert ids[0] == 'mock_id'
+
+    def test_add_texts(self):
+        self.vector._client.index.return_value = {'result': 'created'}
+
+        doc = Document(page_content="Test content", metadata={"document_id": self.example_doc_id})
+        embedding = [0.1] * 128
+
+        with patch('opensearchpy.helpers.bulk') as mock_bulk:
+            mock_bulk.return_value = ([], [])
+            self.vector.add_texts([doc], [embedding])
+
+        mock_response = {
+            'hits': {
+                'total': {'value': 1},
+                'hits': [{'_id': 'mock_id'}]
+            }
+        }
+        self.vector._client.search.return_value = mock_response
+
+        ids = self.vector.get_ids_by_metadata_field(key='document_id', value=self.example_doc_id)
+        assert len(ids) == 1
+        assert ids[0] == 'mock_id'
+
+@pytest.mark.usefixtures("setup_mock_redis")
+class TestOpenSearchVectorWithRedis:
+    def setup_method(self):
+        self.tester = TestOpenSearchVector()
+
+    def test_search_by_full_text(self):
+        self.tester.setup_method()
+        search_response = {
+            'hits': {
+                'total': {'value': 1},
+                'hits': [
+                    {'_source': {'page_content': get_example_text(), 'metadata': {"document_id": "example_doc_id"}}}
+                ]
+            }
+        }
+        expected_length = 1
+        expected_doc_id = "example_doc_id"
+        self.tester.test_search_by_full_text(search_response, expected_length, expected_doc_id)
+
+    def test_delete_by_document_id(self):
+        self.tester.setup_method()
+        self.tester.test_delete_by_document_id()
+
+    def test_get_ids_by_metadata_field(self):
+        self.tester.setup_method()
+        self.tester.test_get_ids_by_metadata_field()
+
+    def test_add_texts(self):
+        self.tester.setup_method()
+        self.tester.test_add_texts()
+
+    def test_search_by_vector(self):
+        self.tester.setup_method()
+        self.tester.test_search_by_vector()