fix

Signed-off-by: kalyanr <[email protected]>
rawwar · Dec 19, 2023 · 0d7b096 · 0d7b096
1 parent fefca75
commit 0d7b096
Showing 1 changed file with 13 additions and 2 deletions.
diff --git a/opensearch_py_ml/ml_commons/ml_commons_client.py b/opensearch_py_ml/ml_commons/ml_commons_client.py
@@ -669,22 +669,33 @@ def _get_tasks_profile(
     def get_profile(
         self,
         profile_type: str = "all",
-        ids: Optional[List[str]] = None,
+        ids: Optional[Union[str, List[str]]] = None,
         request_body: Optional[dict] = None,
     ) -> dict:
         """
         Get profile information based on the profile type.
 
         Args:
             profile_type: The type of profile to retrieve. Valid values are 'all', 'model', or 'task'. Default is 'all'.
-            ids: A list of profile IDs to retrieve. Default is None.
+            ids: Either a single profile ID as a string, or a list of profile IDs to retrieve. Default is None.
             request_body: The request body containing additional information. Default is None.
 
         Returns:
             The profile information.
 
         Raises:
             ValueError: If the profile_type is not 'all', 'model', or 'task'.
+
+        Example:
+            get_profile()
+
+            get_profile(profile_type='model', ids='model1')
+
+            get_profile(profile_type='model', ids=['model1', 'model2'])
+
+            get_profile(profile_type='task', ids='task1', request_body={"node_ids": ["KzONM8c8T4Od-NoUANQNGg"],"return_all_tasks": true,"return_all_models": true})
+
+            get_profile(profile_type='task', ids=['task1', 'task2'], request_body={'additional': 'info'})
         """
 
         if profile_type == "all":