aiverify-moonshot 0.4.0__py3-none-any.whl

moonshot/integrations/cli/redteam/context_strategy.py

@@ -0,0 +1,147 @@
+import argparse
+
+import cmd2
+from rich.console import Console
+from rich.table import Table
+
+from moonshot.api import (
+    api_delete_context_strategy,
+    api_get_all_context_strategy_metadata,
+    api_update_context_strategy,
+)
+from moonshot.integrations.cli.active_session_cfg import active_session
+from moonshot.src.redteaming.session.session import Session
+
+console = Console()
+
+
+def use_context_strategy(args: argparse.Namespace) -> None:
+    """
+    Use a context strategy for process the user's prompt (i.e. summarise past 3 prompts and add
+    it to the current user's prompt)
+
+    Args:
+        args: A namespace with the context strategy parameters. Expected to have 'context_strategy'.
+    """
+    new_context_strategy_name = args.context_strategy
+    num_of_prev_prompts = (
+        args.num_of_prev_prompts
+        if args.num_of_prev_prompts
+        else Session.DEFAULT_CONTEXT_STRATEGY_PROMPT
+    )
+
+    # Check if current session exists. If it does, update context strategy and number of previous prompts
+    if active_session:
+        active_session["context_strategy"] = new_context_strategy_name
+        active_session["cs_num_of_prev_prompts"] = num_of_prev_prompts
+        try:
+            api_update_context_strategy(
+                active_session["session_id"], new_context_strategy_name
+            )
+            print(
+                f"Updated session: {active_session['session_id']}. "
+                f"Context Strategy: {active_session['context_strategy']}."
+                f"No. of previous prompts for Context Strategy: {active_session['cs_num_of_prev_prompts']}."
+            )
+        except Exception as e:
+            print(f"[use_context_strategy]: {str(e)}")
+
+    else:
+        print(
+            "There is no active session. Activate a session to send a prompt with a context strategy."
+        )
+
+
+def list_context_strategies() -> None:
+    """
+    List all context strategies available.
+    """
+    context_strategy_metadata_list = api_get_all_context_strategy_metadata()
+    if context_strategy_metadata_list:
+        table = Table(
+            title="Context Strategy List",
+            show_lines=True,
+            expand=True,
+            header_style="bold",
+        )
+        table.add_column("No.", justify="left", width=2)
+        table.add_column("Context Strategy Information", justify="left", width=98)
+        for context_strategy_index, context_strategy_data in enumerate(
+            context_strategy_metadata_list, 1
+        ):
+            context_strategy_data_str = ""
+            for k, v in context_strategy_data.items():
+                context_strategy_data_str += f"[blue]{k.capitalize()}:[/blue] {v}\n\n"
+            table.add_row(str(context_strategy_index), context_strategy_data_str)
+        console.print(table)
+    else:
+        console.print("[red]There are no context strategies found.[/red]", style="bold")
+
+
+def clear_context_strategy() -> None:
+    """
+    Resets the context in a session.
+    """
+    # Check if current session exists
+    if active_session:
+        try:
+            api_update_context_strategy(active_session["session_id"], "")
+            active_session["context_strategy"] = ""
+            print("Cleared context strategy.")
+        except Exception as e:
+            print(f"[clear_context_strategy: {str(e)}]")
+    else:
+        print(
+            "There is no active session. Activate a session to send a prompt with a context strategy."
+        )
+
+
+def delete_context_strategy(args) -> None:
+    """
+    Deletes a context strategy after confirming with the user.
+
+    Args:
+        args (object): The arguments object. It should have a 'context_strategy' attribute
+                       which is the ID of the context strategy to delete.
+    """
+    # Confirm with the user before deleting a context strategy
+    confirmation = console.input(
+        "[bold red]Are you sure you want to delete the context strategy (y/N)? [/]"
+    )
+    if confirmation.lower() != "y":
+        console.print("[bold yellow]Context strategy deletion cancelled.[/]")
+        return
+    try:
+        api_delete_context_strategy(args.context_strategy)
+        print("[delete_context_strategy]: Context strategy deleted.")
+    except Exception as e:
+        print(f"[delete_context_strategy]: {str(e)}")
+
+
+# Use context strategy arguments
+use_context_strategy_args = cmd2.Cmd2ArgumentParser(
+    description="Use a context strategy.",
+    epilog="Example:\n use_context_strategy my_strategy_one",
+)
+use_context_strategy_args.add_argument(
+    "context_strategy",
+    type=str,
+    help="The ID of the context strategy to use",
+)
+use_context_strategy_args.add_argument(
+    "-n",
+    "--num_of_prev_prompts",
+    type=int,
+    help="The number of previous prompts to use with the context strategy",
+    nargs="?",
+)
+
+# Delete context strategy arguments
+delete_context_strategy_args = cmd2.Cmd2ArgumentParser(
+    description="Delete a context strategy.",
+    epilog="Example:\n delete_context_strategy add_previous_prompt",
+)
+
+delete_context_strategy_args.add_argument(
+    "context_strategy", type=str, help="The ID of the context strategy to delete"
+)

moonshot/integrations/cli/redteam/prompt_template.py

@@ -0,0 +1,67 @@
+import argparse
+
+import cmd2
+from rich.console import Console
+
+from moonshot.api import api_update_prompt_template
+from moonshot.integrations.cli.active_session_cfg import active_session
+
+console = Console()
+
+
+def use_prompt_template(args: argparse.Namespace) -> None:
+    """
+    Use a prompt template by specifying its name while user is in a session.
+
+    Args:
+        args: A namespace with the prompt template parameters. Expected to have 'prompt_template'.
+    """
+    new_prompt_template_name = args.prompt_template
+
+    # Check if current session exists
+    if active_session:
+        active_session["prompt_template"] = new_prompt_template_name
+        try:
+            api_update_prompt_template(
+                active_session["session_id"], new_prompt_template_name
+            )
+            print(
+                f"Updated session: {active_session['session_id']}. "
+                f"Prompt Template: {active_session['prompt_template']}."
+            )
+        except Exception as e:
+            print(f"[use_prompt_template]: {str(e)}")
+    else:
+        print(
+            "There is no active session. Activate a session to send a prompt with a prompt template."
+        )
+
+
+def clear_prompt_template() -> None:
+    """
+    Resets the prompt template in a session.
+    """
+    # Check if current session exists
+    if active_session:
+        try:
+            api_update_prompt_template(active_session["session_id"], "")
+            active_session["prompt_template"] = ""
+            print("Cleared prompt template.")
+        except Exception as e:
+            print(f"[clear_prompt_template: {str(e)}]")
+    else:
+        print(
+            "There is no active session. Activate a session to send a prompt with a prompt template."
+        )
+
+
+# Use prompt template arguments
+use_prompt_template_args = cmd2.Cmd2ArgumentParser(
+    description="Use a prompt template.",
+    epilog="Example:\n use_prompt_template 'analogical-similarity'",
+)
+use_prompt_template_args.add_argument(
+    "prompt_template",
+    type=str,
+    help="Name of the prompt template",
+)

moonshot/integrations/cli/redteam/redteam.py

@@ -0,0 +1,90 @@
+import argparse
+
+import cmd2
+
+from moonshot.integrations.cli.redteam.attack_module import (
+    delete_attack_module,
+    delete_attack_module_args,
+    list_attack_modules,
+)
+from moonshot.integrations.cli.redteam.context_strategy import (
+    clear_context_strategy,
+    delete_context_strategy,
+    delete_context_strategy_args,
+    list_context_strategies,
+    use_context_strategy,
+    use_context_strategy_args,
+)
+from moonshot.integrations.cli.redteam.prompt_template import (
+    clear_prompt_template,
+    use_prompt_template,
+    use_prompt_template_args,
+)
+from moonshot.integrations.cli.redteam.session import (
+    automated_rt_session_args,
+    delete_session,
+    delete_session_args,
+    end_session,
+    list_sessions,
+    new_session,
+    new_session_args,
+    run_attack_module,
+    use_session,
+    use_session_args,
+)
+
+
+@cmd2.with_default_category("Moonshot RedTeaming")
+class RedTeamCommandSet(cmd2.CommandSet):
+    def __init__(self):
+        super().__init__()
+
+    @cmd2.with_argparser(new_session_args)
+    def do_new_session(self, args: argparse.Namespace) -> None:
+        new_session(args)
+
+    def do_end_session(self, _: cmd2.Statement) -> None:
+        end_session()
+
+    def do_list_sessions(self, _: cmd2.Statement) -> None:
+        list_sessions()
+
+    @cmd2.with_argparser(use_session_args)
+    def do_use_session(self, args: argparse.Namespace) -> None:
+        use_session(args)
+
+    @cmd2.with_argparser(use_prompt_template_args)
+    def do_use_prompt_template(self, args: argparse.Namespace) -> None:
+        use_prompt_template(args)
+
+    def do_clear_prompt_template(self, _: cmd2.Statement) -> None:
+        clear_prompt_template()
+
+    def do_list_context_strategies(self, _: cmd2.Statement) -> None:
+        list_context_strategies()
+
+    @cmd2.with_argparser(use_context_strategy_args)
+    def do_use_context_strategy(self, args: argparse.Namespace) -> None:
+        use_context_strategy(args)
+
+    def do_clear_context_strategy(self, _: cmd2.Statement) -> None:
+        clear_context_strategy()
+
+    @cmd2.with_argparser(automated_rt_session_args)
+    def do_run_attack_module(self, args: argparse.Namespace) -> None:
+        run_attack_module(args)
+
+    def do_list_attack_modules(self, _: cmd2.Statement) -> None:
+        list_attack_modules()
+
+    @cmd2.with_argparser(delete_session_args)
+    def do_delete_session(self, args: argparse.Namespace) -> None:
+        delete_session(args)
+
+    @cmd2.with_argparser(delete_context_strategy_args)
+    def do_delete_context_strategy(self, args: argparse.Namespace) -> None:
+        delete_context_strategy(args)
+
+    @cmd2.with_argparser(delete_attack_module_args)
+    def do_delete_attack_module(self, args: argparse.Namespace) -> None:
+        delete_attack_module(args)