aiverify-moonshot 0.4.0__py3-none-any.whl

moonshot/integrations/cli/common/connectors.py

@@ -0,0 +1,325 @@
+from ast import literal_eval
+
+import cmd2
+from rich.console import Console
+from rich.markup import escape
+from rich.table import Table
+
+from moonshot.api import (
+    api_create_endpoint,
+    api_delete_endpoint,
+    api_get_all_connector_type,
+    api_get_all_endpoint,
+    api_read_endpoint,
+    api_update_endpoint,
+)
+
+console = Console()
+
+
+# ------------------------------------------------------------------------------
+# CLI Functions
+# ------------------------------------------------------------------------------
+def add_endpoint(args) -> None:
+    """
+    Add a new endpoint.
+
+    This function adds a new endpoint by calling the api_create_endpoint function from the
+    moonshot.api module using the endpoint details provided in the args.
+
+    Args:
+        args: A namespace object from argparse. It should have the following attributes:
+            name (str): The name of the new endpoint.
+            connector_type (str): The type of the connector (e.g., 'GPT-3', 'Bert', etc.).
+            uri (str): The URI for the connector's API.
+            token (str): The access token for the API.
+            max_calls_per_second (int): The maximum number of API calls allowed per second.
+            max_concurrency (int): The maximum number of concurrent API calls.
+            params (str): A string representation of a dictionary containing additional parameters for the connector.
+
+    Returns:
+        None
+    """
+    try:
+        params_dict = literal_eval(args.params)
+
+        new_endpoint_id = api_create_endpoint(
+            args.name,
+            args.connector_type,
+            args.uri,
+            args.token,
+            args.max_calls_per_second,
+            args.max_concurrency,
+            params_dict,
+        )
+        print(f"[add_endpoint]: Endpoint ({new_endpoint_id}) created.")
+    except Exception as e:
+        print(f"[add_endpoint]: {str(e)}")
+
+
+def list_endpoints() -> 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.
+
+    Returns:
+        None
+    """
+    try:
+        endpoint_list = api_get_all_endpoint()
+        display_endpoints(endpoint_list)
+    except Exception as e:
+        print(f"[list_endpoints]: {str(e)}")
+
+
+def list_connector_types() -> 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.
+
+    Returns:
+        None
+    """
+    try:
+        connector_type_list = api_get_all_connector_type()
+        display_connector_types(connector_type_list)
+    except Exception as e:
+        print(f"[list_connector_types]: {str(e)}")
+
+
+def view_endpoint(args) -> None:
+    """
+    View a specific endpoint.
+
+    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.
+
+    Args:
+        args: A namespace object from argparse. It should have the following attribute:
+            endpoint (str): The name of the endpoint to view.
+
+    Returns:
+        None
+    """
+    try:
+        endpoint_info = api_read_endpoint(args.endpoint)
+        display_endpoints([endpoint_info])
+    except Exception as e:
+        print(f"[view_endpoint]: {str(e)}")
+
+
+def update_endpoint(args) -> None:
+    """
+    Update a specific endpoint.
+
+    This function updates a specific endpoint by calling the api_update_endpoint function from the
+    moonshot.api module using the endpoint name and update values provided in the args.
+
+    Args:
+        args: A namespace object from argparse. It should have the following attributes:
+            endpoint (str): The name of the endpoint to update.
+            update_kwargs (str): A string representation of a dictionary. Each key-value pair in the dictionary
+            represents a field to update in the endpoint and the new value for that field.
+
+    Returns:
+        None
+    """
+    try:
+        endpoint = args.endpoint
+        update_values = dict(literal_eval(args.update_kwargs))
+        api_update_endpoint(endpoint, **update_values)
+        print("[update_endpoint]: Endpoint updated.")
+    except Exception as e:
+        print(f"[update_endpoint]: {str(e)}")
+
+
+def delete_endpoint(args) -> None:
+    """
+    Delete a specific endpoint.
+
+    This function deletes a specific endpoint by calling the api_delete_endpoint function from the
+    moonshot.api module using the endpoint name provided in the args. Before deletion, it asks for
+    user confirmation. If the user confirms, the endpoint is deleted; otherwise, the deletion is cancelled.
+
+    Args:
+        args: A namespace object from argparse. It should have the following attribute:
+            endpoint (str): The name of the endpoint to delete.
+
+    Returns:
+        None
+    """
+    # Confirm with the user before deleting an endpoint
+    confirmation = console.input(
+        "[bold red]Are you sure you want to delete the endpoint (y/N)? [/]"
+    )
+    if confirmation.lower() != "y":
+        console.print("[bold yellow]Endpoint deletion cancelled.[/]")
+        return
+    try:
+        api_delete_endpoint(args.endpoint)
+        print("[delete_endpoint]: Endpoint deleted.")
+    except Exception as e:
+        print(f"[delete_endpoint]: {str(e)}")
+
+
+# ------------------------------------------------------------------------------
+# Helper functions: Display on cli
+# ------------------------------------------------------------------------------
+def display_connector_types(connector_types):
+    """
+    Display a list of connector types.
+
+    This function takes a list of connector types and displays them in a table format. If the list is empty, it prints a
+    message indicating that no connector types were found.
+
+    Args:
+        connector_types (list): A list of connector types.
+
+    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]")
+
+
+def display_endpoints(endpoints_list):
+    """
+    Display a list of endpoints.
+
+    This function takes a list of endpoints and displays them in a table format. If the list is empty, it prints a
+    message indicating that no endpoints were found.
+
+    Args:
+        endpoints_list (list): A list of endpoints. Each endpoint is a dictionary with keys 'id', 'name',
+        'connector_type', 'uri', 'token', 'max_calls_per_second', 'max_concurrency', 'params', and 'created_date'.
+
+    Returns:
+        None
+    """
+    if endpoints_list:
+        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 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]")
+
+
+# ------------------------------------------------------------------------------
+# Cmd2 Arguments Parsers
+# ------------------------------------------------------------------------------
+# Add endpoint arguments
+add_endpoint_args = cmd2.Cmd2ArgumentParser(
+    description="Add a new endpoint. The 'name' argument will be slugified to create a unique identifier.",
+    epilog="Example:\n add_endpoint openai-connector 'OpenAI GPT3.5 Turbo 1106' "
+    "MY_URI ADD_YOUR_TOKEN_HERE 1 1 \"{'temperature': 0.5, 'model': 'gpt-3.5-turbo-1106'}\"",
+)
+add_endpoint_args.add_argument(
+    "connector_type",
+    type=str,
+    help="Type of connection for the endpoint",
+)
+add_endpoint_args.add_argument("name", type=str, help="Name of the new endpoint")
+add_endpoint_args.add_argument("uri", type=str, help="URI of the new endpoint")
+add_endpoint_args.add_argument("token", type=str, help="Token of the new endpoint")
+add_endpoint_args.add_argument(
+    "max_calls_per_second",
+    type=int,
+    help="Max calls per second of the new endpoint",
+)
+add_endpoint_args.add_argument(
+    "max_concurrency", type=int, help="Max concurrency of the new endpoint"
+)
+add_endpoint_args.add_argument("params", type=str, help="Params of the new endpoint")
+
+# Update endpoint arguments
+update_endpoint_args = cmd2.Cmd2ArgumentParser(
+    description="Update an endpoint.",
+    epilog=(
+        "Available keys:\n"
+        "  name: name of the endpoint\n"
+        "  uri: URI of the endpoint\n"
+        "  token: Token of the endpoint\n"
+        "  max_calls_per_second: Rate limit for max calls per second\n"
+        "  max_concurrency: Rate limit for max concurrency\n"
+        "  params: Extra arguments for the endpoint\n\n"
+        "Example:\n"
+        "  update_endpoint openai-gpt4 \"[('name', 'my-special-openai-endpoint'), "
+        "('uri', 'my-uri-loc'), ('token', 'my-token-here')]\""
+    ),
+)
+update_endpoint_args.add_argument("endpoint", type=str, help="ID of the endpoint. This field is not editable via CLI after creation.")
+update_endpoint_args.add_argument(
+    "update_kwargs", type=str, help="Update endpoint key/value"
+)
+
+# View endpoint arguments
+view_endpoint_args = cmd2.Cmd2ArgumentParser(
+    description="View an endpoint.",
+    epilog="Example:\n view_endpoint openai-gpt4",
+)
+view_endpoint_args.add_argument("endpoint", type=str, help="ID of the endpoint")
+
+# Delete endpoint arguments
+delete_endpoint_args = cmd2.Cmd2ArgumentParser(
+    description="Delete an endpoint.",
+    epilog="Example:\n delete_endpoint openai-gpt4",
+)
+delete_endpoint_args.add_argument("endpoint", type=str, help="ID of the endpoint")

moonshot/integrations/cli/common/display_helper.py

@@ -0,0 +1,42 @@
+def display_view_list_format(title: str, items: list) -> str:
+    """
+    Format a list of items for display.
+
+    This function takes a title and a list of items and formats them into a string suitable for display.
+    Each item in the list is displayed on a new line with an index number. If the list is empty, it returns
+    the title with 'nil'.
+
+    Args:
+        title (str): The title to display above the list.
+        items (list): A list of items to be formatted.
+
+    Returns:
+        str: The formatted list as a string.
+    """
+    if items:
+        return f"[blue]{title}[/blue]:" + "".join(
+            f"\n{i + 1}. {item}" for i, item in enumerate(items)
+        )
+    else:
+        return f"[blue]{title}[/blue]: nil"
+
+
+def display_view_str_format(title: str, item: str) -> str:
+    """
+    Format a string item for display with a title.
+
+    This function takes a title and a string item and formats them into a string suitable for display.
+    If the item is not an empty string, it is displayed next to the title with a blue color for the title.
+    If the item is an empty string, 'nil' is displayed next to the title.
+
+    Args:
+        title (str): The title to display next to the item.
+        item (str): The string item to be formatted.
+
+    Returns:
+        str: The formatted string with the title and item.
+    """
+    if item:
+        return f"[blue]{title}[/blue]: {item}"
+    else:
+        return f"[blue]{title}[/blue]: nil"

moonshot/integrations/cli/common/prompt_template.py

@@ -0,0 +1,94 @@
+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
+
+console = Console()
+
+
+# ------------------------------------------------------------------------------
+# CLI Functions
+# ------------------------------------------------------------------------------
+def list_prompt_templates() -> None:
+    """
+    List all prompt templates available.
+    """
+    try:
+        prompt_templates = api_get_all_prompt_template_detail()
+        display_prompt_templates(prompt_templates)
+    except Exception as e:
+        print(f"[list_prompt_templates]: {str(e)}")
+
+
+def delete_prompt_template(args) -> None:
+    """
+    Deletes a prompt_template after confirming with the user.
+
+    Args:
+        args (object): The arguments object. It should have a 'prompt_template' attribute
+                       which is the ID of the prompt template to delete.
+    """
+    # Confirm with the user before deleting a prompt template
+    confirmation = console.input(
+        "[bold red]Are you sure you want to delete the prompt template (y/N)? [/]"
+    )
+    if confirmation.lower() != "y":
+        console.print("[bold yellow]Prompt template deletion cancelled.[/]")
+        return
+    try:
+        api_delete_prompt_template(args.prompt_template)
+        print("[delete_prompt_template]: Prompt template deleted.")
+    except Exception as e:
+        print(f"[delete_prompt_template]: {str(e)}")
+
+
+# ------------------------------------------------------------------------------
+# Helper functions: Display on cli
+# ------------------------------------------------------------------------------
+def display_prompt_templates(prompt_templates) -> None:
+    """
+    Display the list of prompt templates in a formatted table.
+
+    This function takes a list of prompt templates and displays them in a formatted table.
+    Each row in the table represents a prompt template with its ID, name, description, and contents.
+    If the list of prompt templates is empty, it prints a message indicating that no prompt templates were found.
+
+    Args:
+        prompt_templates (list): A list of dictionaries, each representing a prompt template.
+    """
+    table = Table(
+        title="List of Prompt Templates",
+        show_lines=True,
+        expand=True,
+        header_style="bold",
+    )
+    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]")
+
+
+# Delete prompt template arguments
+delete_prompt_template_args = cmd2.Cmd2ArgumentParser(
+    description="Delete a prompt template.",
+    epilog="Example:\n delete_prompt_template squad-shifts",
+)
+
+delete_prompt_template_args.add_argument(
+    "prompt_template", type=str, help="The ID of the prompt template to delete"
+)

moonshot/integrations/cli/initialisation/initialisation.py

@@ -0,0 +1,14 @@
+import cmd2
+
+
+@cmd2.with_default_category("Initialisation")
+class InitialisationCommandSet(cmd2.CommandSet):
+    def __init__(self):
+        super().__init__()
+
+    def do_interactive(self, _: cmd2.Statement) -> None:
+        """
+        Run the interactive shell.
+        """
+        # To prevent 'interactive is not a recognized command, alias, or macro' from triggering.
+        pass

moonshot/integrations/cli/redteam/attack_module.py

@@ -0,0 +1,70 @@
+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
+
+console = Console()
+
+
+def list_attack_modules() -> None:
+    """
+    Retrieves and prints the metadata of all attack modules.
+    """
+    print("Listing attack modules may take a while...")
+    attack_module_metadata_list = api_get_all_attack_module_metadata()
+
+    if attack_module_metadata_list:
+        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_module_metadata_list, 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)
+
+        console.print(table)
+    else:
+        console.print("[red]There are no attack modules found.[/red]", style="bold")
+
+
+def delete_attack_module(args) -> None:
+    """
+    Deletes an attack module after confirming with the user.
+
+    Args:
+        args (object): The arguments object. It should have a 'attack_modulee' attribute
+                       which is the ID of the attack module to delete.
+    """
+    # Confirm with the user before deleting an attack module
+    confirmation = console.input(
+        "[bold red]Are you sure you want to delete the attack module (y/N)? [/]"
+    )
+    if confirmation.lower() != "y":
+        console.print("[bold yellow]Attack module deletion cancelled.[/]")
+        return
+    try:
+        api_delete_attack_module(args.attack_module)
+        print("[delete_attack_module]: Attack module deleted.")
+    except Exception as e:
+        print(f"[delete_attack_module]: {str(e)}")
+
+
+# Delete attack module arguments
+delete_attack_module_args = cmd2.Cmd2ArgumentParser(
+    description="Delete an attack module.",
+    epilog="Example:\n delete_attack_module sample_attack_module",
+)
+
+delete_attack_module_args.add_argument(
+    "attack_module", type=str, help="The ID of the attack module to delete"
+)