aiverify-moonshot 0.4.2__py3-none-any.whl → 0.4.4__py3-none-any.whl

moonshot/integrations/cli/common/connectors.py

@@ -13,7 +13,7 @@ from moonshot.api import (
     api_read_endpoint,
     api_update_endpoint,
 )
-from moonshot.src.utils.find_feature import find_keyword
+from moonshot.integrations.cli.utils.process_data import filter_data
 
 console = Console()
 
@@ -63,11 +63,12 @@ def list_endpoints(args) -> list | None:
     List all endpoints.
 
     This function retrieves all endpoints by calling the api_get_all_endpoint function from the
-    moonshot.api module. It then displays the endpoints using the display_endpoints function.
+    moonshot.api module. It then displays the endpoints using the _display_endpoints function.
 
     Args:
         args: A namespace object from argparse. It should have an optional attribute:
         find (str): Optional field to find endpoint(s) with a keyword.
+        pagination (str): Optional field to paginate endpoints.
 
     Returns:
         list | None: A list of ConnectorEndpoint or None if there is no result.
@@ -75,17 +76,17 @@ def list_endpoints(args) -> list | None:
     try:
         endpoints_list = api_get_all_endpoint()
         keyword = args.find.lower() if args.find else ""
-        if keyword:
-            filtered_endpoints_list = find_keyword(keyword, endpoints_list)
+        pagination = literal_eval(args.pagination) if args.pagination else ()
+
+        if endpoints_list:
+            filtered_endpoints_list = filter_data(endpoints_list, keyword, pagination)
             if filtered_endpoints_list:
-                display_endpoints(filtered_endpoints_list)
+                _display_endpoints(filtered_endpoints_list)
                 return filtered_endpoints_list
-            else:
-                print("No endpoints containing keyword found.")
-                return None
-        else:
-            display_endpoints(endpoints_list)
-            return endpoints_list
+
+        console.print("[red]There are no endpoints found.[/red]")
+        return None
+
     except Exception as e:
         print(f"[list_endpoints]: {str(e)}")
 
@@ -95,11 +96,12 @@ def list_connector_types(args) -> list | None:
     List all connector types.
 
     This function retrieves all connector types by calling the api_get_all_connector_type function from the
-    moonshot.api module. It then displays the connector types using the display_connector_types function.
+    moonshot.api module. It then displays the connector types using the _display_connector_types function.
 
     Args:
         args: A namespace object from argparse. It should have an optional attribute:
         find (str): Optional field to find connector type(s) with a keyword.
+        pagination (str): Optional field to paginate connector types.
 
     Returns:
         list | None: A list of Connector or None if there is no result.
@@ -107,17 +109,19 @@ def list_connector_types(args) -> list | None:
     try:
         connector_type_list = api_get_all_connector_type()
         keyword = args.find.lower() if args.find else ""
-        if keyword:
-            filtered_connector_type_list = find_keyword(keyword, connector_type_list)
+        pagination = literal_eval(args.pagination) if args.pagination else ()
+
+        if connector_type_list:
+            filtered_connector_type_list = filter_data(
+                connector_type_list, keyword, pagination
+            )
             if filtered_connector_type_list:
-                display_connector_types(filtered_connector_type_list)
+                _display_connector_types(filtered_connector_type_list)
                 return filtered_connector_type_list
-            else:
-                print("No connectors containing keyword found.")
-                return None
-        else:
-            display_connector_types(connector_type_list)
-            return connector_type_list
+
+        console.print("[red]There are no connector types found.[/red]")
+        return None
+
     except Exception as e:
         print(f"[list_connector_types]: {str(e)}")
 
@@ -128,7 +132,7 @@ def view_endpoint(args) -> None:
 
     This function retrieves a specific endpoint by calling the api_read_endpoint function from the
     moonshot.api module using the endpoint name provided in the args. It then displays the endpoint
-    information using the display_endpoints function.
+    information using the _display_endpoints function.
 
     Args:
         args: A namespace object from argparse. It should have the following attribute:
@@ -139,7 +143,7 @@ def view_endpoint(args) -> None:
     """
     try:
         endpoint_info = api_read_endpoint(args.endpoint)
-        display_endpoints([endpoint_info])
+        _display_endpoints([endpoint_info])
     except Exception as e:
         print(f"[view_endpoint]: {str(e)}")
 
@@ -201,7 +205,7 @@ def delete_endpoint(args) -> None:
 # ------------------------------------------------------------------------------
 # Helper functions: Display on cli
 # ------------------------------------------------------------------------------
-def display_connector_types(connector_types: list) -> None:
+def _display_connector_types(connector_types: list) -> None:
     """
     Display a list of connector types.
 
@@ -214,24 +218,22 @@ def display_connector_types(connector_types: list) -> None:
     Returns:
         None
     """
-    if connector_types:
-        table = Table(
-            title="List of Connector Types",
-            show_lines=True,
-            expand=True,
-            header_style="bold",
-        )
-        table.add_column("No.", width=2)
-        table.add_column("Connector Type", justify="left", width=78)
-        for connector_id, connector_type in enumerate(connector_types, 1):
-            table.add_section()
-            table.add_row(str(connector_id), connector_type)
-        console.print(table)
-    else:
-        console.print("[red]There are no connector types found.[/red]")
+    table = Table(
+        title="List of Connector Types",
+        show_lines=True,
+        expand=True,
+        header_style="bold",
+    )
+    table.add_column("No.", width=2)
+    table.add_column("Connector Type", justify="left", width=78)
+
+    for idx, connector_type in enumerate(connector_types, 1):
+        table.add_section()
+        table.add_row(str(idx), connector_type)
+    console.print(table)
 
 
-def display_endpoints(endpoints_list):
+def _display_endpoints(endpoints_list):
     """
     Display a list of endpoints.
 
@@ -245,52 +247,51 @@ def display_endpoints(endpoints_list):
     Returns:
         None
     """
-    if endpoints_list:
-        table = Table(
-            title="List of Connector Endpoints",
-            show_lines=True,
-            expand=True,
-            header_style="bold",
+    table = Table(
+        title="List of Connector Endpoints",
+        show_lines=True,
+        expand=True,
+        header_style="bold",
+    )
+    table.add_column("No.", justify="left", width=2)
+    table.add_column("Id", justify="left", width=10)
+    table.add_column("Name", justify="left", width=10)
+    table.add_column("Connector Type", justify="left", width=10)
+    table.add_column("Uri", justify="left", width=10)
+    table.add_column("Token", justify="left", width=10)
+    table.add_column("Max Calls Per Second", justify="left", width=5)
+    table.add_column("Max concurrency", justify="left", width=5)
+    table.add_column("Params", justify="left", width=30)
+    table.add_column("Created Date", justify="left", width=8)
+
+    for idx, endpoint in enumerate(endpoints_list, 1):
+        (
+            id,
+            name,
+            connector_type,
+            uri,
+            token,
+            max_calls_per_second,
+            max_concurrency,
+            params,
+            created_date,
+            *other_args,
+        ) = endpoint.values()
+        table.add_section()
+        idx = endpoint.get("idx", idx)
+        table.add_row(
+            str(idx),
+            id,
+            name,
+            connector_type,
+            uri,
+            token,
+            str(max_calls_per_second),
+            str(max_concurrency),
+            escape(str(params)),
+            created_date,
         )
-        table.add_column("No.", justify="left", width=2)
-        table.add_column("Id", justify="left", width=10)
-        table.add_column("Name", justify="left", width=10)
-        table.add_column("Connector Type", justify="left", width=10)
-        table.add_column("Uri", justify="left", width=10)
-        table.add_column("Token", justify="left", width=10)
-        table.add_column("Max Calls Per Second", justify="left", width=5)
-        table.add_column("Max concurrency", justify="left", width=5)
-        table.add_column("Params", justify="left", width=30)
-        table.add_column("Created Date", justify="left", width=8)
-
-        for endpoint_id, endpoint in enumerate(endpoints_list, 1):
-            (
-                id,
-                name,
-                connector_type,
-                uri,
-                token,
-                max_calls_per_second,
-                max_concurrency,
-                params,
-                created_date,
-            ) = endpoint.values()
-            table.add_section()
-            table.add_row(
-                str(endpoint_id),
-                id,
-                name,
-                connector_type,
-                uri,
-                token,
-                str(max_calls_per_second),
-                str(max_concurrency),
-                escape(str(params)),
-                created_date,
-            )
-        console.print(table)
-    else:
-        console.print("[red]There are no endpoints found.[/red]")
+    console.print(table)
 
 
 # ------------------------------------------------------------------------------
@@ -373,6 +374,13 @@ list_endpoints_args.add_argument(
     nargs="?",
 )
 
+list_endpoints_args.add_argument(
+    "-p",
+    "--pagination",
+    type=str,
+    help="Optional tuple to paginate endpoint(s). E.g. (2,10) returns 2nd page with 10 items in each page.",
+    nargs="?",
+)
 
 # List connector types arguments
 list_connector_types_args = cmd2.Cmd2ArgumentParser(
@@ -387,3 +395,11 @@ list_connector_types_args.add_argument(
     help="Optional field to find connector type(s) with keyword",
     nargs="?",
 )
+
+list_connector_types_args.add_argument(
+    "-p",
+    "--pagination",
+    type=str,
+    help="Optional tuple to paginate connector type(s). E.g. (2,10) returns 2nd page with 10 items in each page.",
+    nargs="?",
+)

moonshot/integrations/cli/common/dataset.py

@@ -0,0 +1,66 @@
+from ast import literal_eval
+
+import cmd2
+from rich.console import Console
+
+from moonshot.api import (
+    api_create_datasets,
+)
+
+console = Console()
+def add_dataset(args) -> None:
+    """
+    Create a new dataset using the provided arguments and log the result.
+
+    This function attempts to create a new dataset by calling the `api_create_datasets`
+    function with the necessary parameters extracted from `args`. If successful, it logs
+    the creation of the dataset with its ID. If an exception occurs, it logs the error.
+
+    Args:
+        args: An argparse.Namespace object containing the following attributes:
+            - name (str): Name of the new dataset.
+            - description (str): Description of the new dataset.
+            - reference (str): Reference URL for the new dataset.
+            - license (str): License type for the new dataset.
+            - method (str): Method to convert the new dataset ('hf' or 'csv').
+            - params (dict): Additional parameters for dataset creation.
+    """
+    try:
+        new_dataset_id = api_create_datasets(
+            args.name,
+            args.description,
+            args.reference,
+            args.license,
+            args.method,
+            **args.params,
+        )
+        print(f"[add_dataset]: Dataset ({new_dataset_id}) created.")
+    except Exception as e:
+        print(f"[add_dataset]: {str(e)}")
+
+        # ------------------------------------------------------------------------------
+# Cmd2 Arguments Parsers
+# ------------------------------------------------------------------------------
+# Add dataset arguments
+add_dataset_args = cmd2.Cmd2ArgumentParser(
+    description="Add a new dataset. The 'name' argument will be slugified to create a unique identifier.",
+    epilog=(
+        "Examples:\n"
+        "1. add_dataset 'dataset-name' 'A brief description' 'http://reference.com' 'MIT' 'csv' \"{'csv_file_path': '/path/to/your/file.csv'}\"\n"
+        "2. add_dataset 'dataset-name' 'A brief description' 'http://reference.com' 'MIT' 'hf' \"{'dataset_name': 'cais/mmlu', 'dataset_config': 'college_biology', 'split': 'test', 'input_col': ['question','choices'], 'target_col': 'answer'}\""
+    ),
+)
+add_dataset_args.add_argument("name", type=str, help="Name of the new dataset")
+add_dataset_args.add_argument("description", type=str, help="Description of the new dataset")
+add_dataset_args.add_argument("reference", type=str, help="Reference of the new dataset")
+add_dataset_args.add_argument("license", type=str, help="License of the new dataset")
+add_dataset_args.add_argument("method", type=str, choices=['hf', 'csv'], help="Method to convert the new dataset. Choose either 'hf' or 'csv'.")
+add_dataset_args.add_argument(
+    "params", 
+    type=literal_eval, 
+    help=(
+        "Params of the new dataset in dictionary format. For example: \n"
+        "1. For 'csv' method: \"{'csv_file_path': '/path/to/your/file.csv'}\"\n"
+        "2. For 'hf' method: \"{'dataset_name': 'cais_mmlu', 'dataset_config': 'college_biology', 'split': 'test', 'input_col': ['questions','choices'], 'target_col': 'answer'}\""
+    )
+)

moonshot/integrations/cli/common/prompt_template.py

@@ -1,9 +1,11 @@
+from ast import literal_eval
+
 import cmd2
 from rich.console import Console
 from rich.table import Table
 
 from moonshot.api import api_delete_prompt_template, api_get_all_prompt_template_detail
-from moonshot.src.utils.find_feature import find_keyword
+from moonshot.integrations.cli.utils.process_data import filter_data
 
 console = Console()
 
@@ -18,6 +20,7 @@ def list_prompt_templates(args) -> list | None:
     Args:
         args: A namespace object from argparse. It should have an optional attribute:
         find (str): Optional field to find prompt template(s) with a keyword.
+        pagination (str): Optional field to paginate prompt templates.
 
     Returns:
         list | None: A list of PromptTemplate or None if there is no result.
@@ -25,19 +28,19 @@ def list_prompt_templates(args) -> list | None:
     try:
         prompt_templates_list = api_get_all_prompt_template_detail()
         keyword = args.find.lower() if args.find else ""
-        if keyword:
-            filtered_prompt_templates_list = find_keyword(
-                keyword, prompt_templates_list
+        pagination = literal_eval(args.pagination) if args.pagination else ()
+
+        if prompt_templates_list:
+            filtered_prompt_templates_list = filter_data(
+                prompt_templates_list, keyword, pagination
             )
             if filtered_prompt_templates_list:
-                display_prompt_templates(filtered_prompt_templates_list)
+                _display_prompt_templates(filtered_prompt_templates_list)
                 return filtered_prompt_templates_list
-            else:
-                print("No prompt templates containing keyword found.")
-                return None
-        else:
-            display_prompt_templates(prompt_templates_list)
-            return prompt_templates_list
+
+        console.print("[red]There are no prompt templates found.[/red]")
+        return None
+
     except Exception as e:
         print(f"[list_prompt_templates]: {str(e)}")
 
@@ -67,7 +70,7 @@ def delete_prompt_template(args) -> None:
 # ------------------------------------------------------------------------------
 # Helper functions: Display on cli
 # ------------------------------------------------------------------------------
-def display_prompt_templates(prompt_templates) -> None:
+def _display_prompt_templates(prompt_templates) -> None:
     """
     Display the list of prompt templates in a formatted table.
 
@@ -87,21 +90,13 @@ def display_prompt_templates(prompt_templates) -> None:
     table.add_column("No.", width=2)
     table.add_column("Prompt Template", justify="left", width=50)
     table.add_column("Contains", justify="left", width=48, overflow="fold")
-    if prompt_templates:
-        for prompt_index, prompt_template in enumerate(prompt_templates, 1):
-            (
-                id,
-                name,
-                description,
-                contents,
-            ) = prompt_template.values()
-
-            prompt_info = f"[red]id: {id}[/red]\n\n[blue]{name}[/blue]\n{description}"
-            table.add_section()
-            table.add_row(str(prompt_index), prompt_info, contents)
-        console.print(table)
-    else:
-        console.print("[red]There are no prompt templates found.[/red]")
+    for idx, prompt_template in enumerate(prompt_templates, 1):
+        (id, name, description, contents, *other_args) = prompt_template.values()
+        idx = prompt_template.get("idx", idx)
+        prompt_info = f"[red]id: {id}[/red]\n\n[blue]{name}[/blue]\n{description}"
+        table.add_section()
+        table.add_row(str(idx), prompt_info, contents)
+    console.print(table)
 
 
 # Delete prompt template arguments
@@ -127,3 +122,11 @@ list_prompt_templates_args.add_argument(
     help="Optional field to find prompt template(s) with keyword",
     nargs="?",
 )
+
+list_prompt_templates_args.add_argument(
+    "-p",
+    "--pagination",
+    type=str,
+    help="Optional tuple to paginate prompt template(s). E.g. (2,10) returns 2nd page with 10 items in each page.",
+    nargs="?",
+)

moonshot/integrations/cli/redteam/attack_module.py

@@ -1,13 +1,18 @@
+from ast import literal_eval
+
 import cmd2
 from rich.console import Console
 from rich.table import Table
 
 from moonshot.api import api_delete_attack_module, api_get_all_attack_module_metadata
-from moonshot.src.utils.find_feature import find_keyword
+from moonshot.integrations.cli.utils.process_data import filter_data
 
 console = Console()
 
 
+# ------------------------------------------------------------------------------
+# CLI Functions
+# ------------------------------------------------------------------------------
 def list_attack_modules(args) -> list | None:
     """
     Retrieves and prints the metadata of all attack modules.
@@ -15,6 +20,7 @@ def list_attack_modules(args) -> list | None:
     Args:
         args: A namespace object from argparse. It should have an optional attribute:
         find (str): Optional field to find attack module(s) with a keyword.
+        pagination (str): Optional field to paginate attack modules.
 
     Returns:
          list | None: A list of AttackModule or None if there is no result.
@@ -23,19 +29,19 @@ def list_attack_modules(args) -> list | None:
         print("Listing attack modules may take a while...")
         attack_module_metadata_list = api_get_all_attack_module_metadata()
         keyword = args.find.lower() if args.find else ""
-        if keyword:
-            filtered_attack_modules_list = find_keyword(
-                keyword, attack_module_metadata_list
+        pagination = literal_eval(args.pagination) if args.pagination else ()
+
+        if attack_module_metadata_list:
+            filtered_attack_modules_list = filter_data(
+                attack_module_metadata_list, keyword, pagination
             )
             if filtered_attack_modules_list:
-                display_attack_modules(filtered_attack_modules_list)
+                _display_attack_modules(filtered_attack_modules_list)
                 return filtered_attack_modules_list
-            else:
-                print("No attack modules containing keyword found.")
-                return None
-        else:
-            display_attack_modules(attack_module_metadata_list)
-            return attack_module_metadata_list
+
+        console.print("[red]There are no attack modules found.[/red]")
+        return None
+
     except Exception as e:
         print(f"[list_attack_modules]: {str(e)}")
 
@@ -62,7 +68,10 @@ def delete_attack_module(args) -> None:
         print(f"[delete_attack_module]: {str(e)}")
 
 
-def display_attack_modules(attack_modules: list) -> None:
+# ------------------------------------------------------------------------------
+# Helper functions: Display on cli
+# ------------------------------------------------------------------------------
+def _display_attack_modules(attack_modules: list) -> None:
     """
     Display a list of attack modules.
 
@@ -75,25 +84,23 @@ def display_attack_modules(attack_modules: list) -> None:
     Returns:
         None
     """
-    if attack_modules:
-        table = Table(
-            title="Attack Module List",
-            show_lines=True,
-            expand=True,
-            header_style="bold",
-        )
-        table.add_column("No.", width=2)
-        table.add_column("Details", justify="left", width=98)
-
-        for attack_module_index, attack_module_data in enumerate(attack_modules, 1):
-            attack_module_data_str = ""
-            for k, v in attack_module_data.items():
-                attack_module_data_str += f"[blue]{k.capitalize()}:[/blue] {v}\n\n"
-            table.add_row(str(attack_module_index), attack_module_data_str)
+    table = Table(
+        title="Attack Module List",
+        show_lines=True,
+        expand=True,
+        header_style="bold",
+    )
+    table.add_column("No.", width=2)
+    table.add_column("Details", justify="left", width=98)
 
-        console.print(table)
-    else:
-        console.print("[red]There are no attack modules found.[/red]", style="bold")
+    for idx, attack_module_data in enumerate(attack_modules, 1):
+        attack_module_data_str = ""
+        for k, v in attack_module_data.items():
+            if k != "idx":
+                attack_module_data_str += f"[blue]{k.capitalize()}:[/blue] {v}\n\n"
+        idx = attack_module_data.get("idx", idx)
+        table.add_row(str(idx), attack_module_data_str)
+    console.print(table)
 
 
 # Delete attack module arguments
@@ -119,3 +126,11 @@ list_attack_modules_args.add_argument(
     help="Optional field to find attack module(s) with keyword",
     nargs="?",
 )
+
+list_attack_modules_args.add_argument(
+    "-p",
+    "--pagination",
+    type=str,
+    help="Optional tuple to paginate attack module(s). E.g. (2,10) returns 2nd page with 10 items in each page.",
+    nargs="?",
+)