Modified generator to generate plugins

Signed-off-by: saimedhi <[email protected]>

opensearchpy/client/plugins.py

@@ -12,6 +12,7 @@
 
 from ..plugins.alerting import AlertingClient
 from ..plugins.index_management import IndexManagementClient
+from ..plugins.knn import KnnClient
 from .client import Client
 from .utils import NamespacedClient
 
@@ -29,6 +30,7 @@ def __init__(self, client: Client) -> None:
         # self.anomaly_detection = AnomalyDetectionClient(client)
         # self.trace_analytics = TraceAnalyticsClient(client)
         self.index_management = IndexManagementClient(client)
+        self.knn = KnnClient(client)
 
         self._dynamic_lookup(client)
 

test_opensearchpy/test_server/test_plugins/test_knn.py

@@ -0,0 +1,103 @@
+# SPDX-License-Identifier: Apache-2.0
+#
+# The OpenSearch Contributors require contributions made to
+# this file be licensed under the Apache-2.0 license or a
+# compatible open source license.
+#
+# Modifications Copyright OpenSearch Contributors. See
+# GitHub history for details.
+
+
+from __future__ import unicode_literals
+from math import e
+
+import random
+import time
+from urllib import response
+from opensearchpy.exceptions import RequestError
+from opensearchpy import helpers
+
+from .. import OpenSearchTestCase
+
+
+class TestKNNPlugin(OpenSearchTestCase):
+    def setUp(self) -> None: 
+        index_name = "train-index"
+        if self.client.indices.exists(index=index_name):
+            self.client.indices.delete(index=index_name)
+        try:
+            self.client.plugins.knn.delete_model(model_id= 'my-model')
+        except Exception:
+            pass
+
+    def tearDown(self) -> None:
+        index_name = "train-index"
+        if not self.client.indices.exists(self.index_name):
+            self.client.indices.delete(index=self.index_name)
+        try:
+            self.client.plugins.knn.delete_model(model_id= 'my-model')
+        except Exception:
+            pass
+
+    def test_knn_model(self) -> None:
+
+        index_name = "train-index"
+        # Create an index with training data
+        self.client.indices.create(
+            index=index_name,
+            body={ "settings": {"number_of_shards": 3,"number_of_replicas": 0 },
+                "mappings": {"properties": {"train-field": {"type": "knn_vector", "dimension": 4}}}
+                })
+
+
+        # Adding data to the index  
+        docs = '''
+        { "index": { "_index": "train-index", "_id": "1" } }
+        { "train-field": [1.5, 5.5, 4.5, 6.4]}
+        { "index": { "_index": "train-index", "_id": "2" } }
+        { "train-field": [2.5, 3.5, 5.6, 6.7]}
+        { "index": { "_index": "train-index", "_id": "3" } }
+        { "train-field": [4.5, 5.5, 6.7, 3.7]}
+        { "index": { "_index": "train-index", "_id": "4" } }
+        { "train-field": [1.5, 5.5, 4.5, 6.4]}
+        '''
+        self.client.bulk(docs)
+
+        self.client.plugins.knn.train_model(model_id= 'my-model', body = {
+            "training_index": "train-index",
+            "training_field": "train-field",
+            "dimension": 4,
+            "description": "My model description",
+            "method": {
+                "name": "ivf",
+                "engine": "faiss",
+                "space_type": "l2",
+                "parameters": {
+                "nlist": 4,
+                "nprobes": 2
+                }
+            }
+            })
+
+        time_out = time.time() + 900  # Waiting for model to train before searching it. 
+
+        while time.time() < time_out:
+
+            # Fetch the model state
+            try:
+                response = self.client.plugins.knn.get_model(model_id=model_id, filter_path='state')['state']
+                if response['state'] != 'training':
+                    pass 
+            except Exception as e:
+                print(f"Error fetching knn model: {e}")
+
+            # Sleep for 1 minute before the next attempt
+            time.sleep(10)
+
+
+
+
+        #self.client.plugins.knn.get_model(model_id= 'my-model')
+
+        print("response", response)
+        self.client.plugins.knn.delete_model(model_id= 'my-model')