Search task added

opensearch_py_ml/ml_commons/ml_commons_client.py

@@ -6,6 +6,7 @@
 # GitHub history for details.
 
 
+import json
 import time
 from typing import Any, List, Union
 
@@ -376,6 +377,34 @@ def _get_task_info(self, task_id: str):
             url=API_URL,
         )
 
+    def search_task(self, json_input) -> object:
+        """
+        This method searches a task from opensearch cluster (using ml commons api)
+        :param json: json input for the search request
+        :type json: string or dict
+        :return: returns a json object, with detailed information about the searched task
+        :rtype: object
+        """
+
+        API_URL = f"{ML_BASE_URI}/tasks/_search"
+
+        if isinstance(json_input, str):
+            try:
+                json_obj = json.loads(json_input)
+                API_BODY = json.dumps(json_obj)
+            except json.JSONDecodeError:
+                return "Invalid JSON string passed as argument."
+        elif isinstance(json_input, dict):
+            API_BODY = json.dumps(json_input)
+        else:
+            return "Invalid JSON object passed as argument."
+
+        return self._client.transport.perform_request(
+            method="POST",
+            url=API_URL,
+            body=API_BODY,
+        )
+
     def get_model_info(self, model_id: str) -> object:
         """
         This method return information about a model registered in the opensearch cluster (using ml commons api)

tests/ml_commons/test_ml_commons_client.py

@@ -372,6 +372,26 @@ def test_integration_model_train_register_full_cycle():
                     print("Model Task Status:", ml_task_status)
                     raised = True
                 assert raised == False, "Raised Exception in pulling task info"
+
+                try:
+                    search_task_obj = ml_client.search_task(json="{'_id':'{task_id}'}")
+                    assert search_task_obj["hits"]["hits"][0]["_id"] == task_id
+                except:  # noqa: E722
+                    raised = True
+                assert raised == False, "Raised Exception in searching task"
+
+                try:
+                    search_task_obj = ml_client.search_task(json={"_id": task_id})
+                except:  # noqa: E722
+                    raised = True
+                assert raised == False, "Raised Exception in searching task"
+
+                try:
+                    search_task_obj = ml_client.search_task(json=15)
+                    assert search_task_obj == "Invalid JSON object passed as argument."
+                except:  # noqa: E722
+                    raised = True
+                assert raised == False, "Raised Exception in searching task"
                 # This is test is being flaky. Sometimes the test is passing and sometimes showing 500 error
                 # due to memory circuit breaker.
                 # Todo: We need to revisit this test.