aiverify-moonshot 0.4.0__py3-none-any.whl

moonshot/__main__.py

@@ -0,0 +1,198 @@
+import sys
+import warnings
+import argparse
+import platform
+import subprocess
+import os
+import threading 
+from dotenv import dotenv_values
+from moonshot.api import api_set_environment_variables
+"""
+Run the Moonshot application
+"""
+def run_subprocess(*args, **kwargs):
+    """
+    Run a subprocess with the option to use shell=True on Windows.
+    """
+    if platform.system() == "Windows":
+        kwargs["shell"] = True
+    return subprocess.run(*args, **kwargs)
+
+def ms_lib_env_file(data_repo_name):
+    """
+    Writes the env file to be used for moonshot library
+    """
+    env_content_data = f"""
+    # For Data
+    ATTACK_MODULES="./{data_repo_name}/attack-modules"
+    CONNECTORS="./{data_repo_name}/connectors"
+    CONNECTORS_ENDPOINTS="./{data_repo_name}/connectors-endpoints"
+    CONTEXT_STRATEGY="./{data_repo_name}/context-strategy"
+    COOKBOOKS="./{data_repo_name}/cookbooks"
+    DATABASES="./{data_repo_name}/generated-outputs/databases"
+    DATABASES_MODULES="./{data_repo_name}/databases-modules"
+    DATASETS="./{data_repo_name}/datasets"
+    IO_MODULES="./{data_repo_name}/io-modules"
+    METRICS="./{data_repo_name}/metrics"
+    PROMPT_TEMPLATES="./{data_repo_name}/prompt-templates"
+    RECIPES="./{data_repo_name}/recipes"
+    RESULTS="./{data_repo_name}/generated-outputs/results"
+    RESULTS_MODULES="./{data_repo_name}/results-modules"
+    RUNNERS="./{data_repo_name}/generated-outputs/runners"
+    RUNNERS_MODULES="./{data_repo_name}/runners-modules"
+    TOKENIZERS_PARALLELISM = false 
+    """
+
+    env_content_web_api = """
+    # For Web API
+    HOST_ADDRESS=127.0.0.1 # The interface the server will bind to
+    HOST_PORT=5000
+
+    # Below is the uri of the Web UI webhook.
+    # In the next section, if Web UI listens on a different port,
+    # update this uri accordingly and restart.
+
+    MOONSHOT_UI_CALLBACK_URL=http://localhost:3000/api/v1/benchmarks/status
+    """
+    with open(".env", "w") as env_file:
+        combined_content = env_content_data + env_content_web_api
+        env_file.write(combined_content.strip())
+
+def ms_ui_env_file(ui_repo):
+    """
+    Write the env file to be used with moonshot ui
+    """
+    env_content = """
+    # This should be the URL of the Moonshot Web Api module which was started in the previous section.
+    # Check the startup logs to determine the hostname and port number.
+    MOONSHOT_API_URL=http://127.0.0.1:5000
+    """
+    # Write the .env content to a file
+    with open(os.path.join(ui_repo, ".env"), "w") as env_file:
+        env_file.write(env_content.strip())
+
+def moonshot_data_installation():
+    # Code for moonshot-data installation
+    print("Installing Moonshot Data from GitHub")
+    repo = "https://github.com/aiverify-foundation/moonshot-data.git"
+    folder_name = repo.split("/")[-1].replace(".git", "")
+    
+    # Check if the directory already exists
+    if not os.path.exists(folder_name):
+        print(f"Cloning {repo}")
+        # Clone the repository
+        run_subprocess(["git", "clone", repo], check=True)
+        
+        # Change directory to the folder
+        os.chdir(folder_name)
+
+        print(f"Installing requirements for {folder_name}")
+        # Install the requirements if they exist
+        if os.path.exists("requirements.txt"):
+            run_subprocess(["pip", "install", "-r", "requirements.txt"], check=True)
+            import nltk
+            nltk.download('punkt')
+            nltk.download('stopwords')
+            nltk.download('averaged_perceptron_tagger')
+            nltk.download('universal_tagset')
+        
+        # Change back to the base directory
+        os.chdir("..")
+
+        # Create .env to point to installed folder
+        ms_lib_env_file(folder_name)
+    else:
+        print(f"Directory {folder_name} already exists, skipping clone.")
+
+def moonshot_ui_installation():
+    # Code for moonshot-ui installation
+    repo = "https://github.com/aiverify-foundation/moonshot-ui.git"
+    folder_name = repo.split("/")[-1].replace(".git", "")
+    
+    # Check if the directory already exists
+    if not os.path.exists(folder_name):
+        print(f"Cloning {repo}")
+        # Clone the repository
+        run_subprocess(["git", "clone", repo], check=True)
+
+        # Change directory to the folder
+        os.chdir(folder_name)
+
+        print(f"Installing requirements for {folder_name}")        
+        # Install the requirements if they exist
+        if os.path.exists("package.json"):
+            run_subprocess(["npm", "install"], check=True)
+            run_subprocess(["npm", "run", "build"], check=True)
+        
+        # Change back to the base directory
+        os.chdir("..")
+
+        # Create .env for ui
+        ms_ui_env_file(folder_name)
+    else:
+        print(f"Directory {folder_name} already exists, skipping installation.")
+
+def run_moonshot_ui_dev():
+    """
+    To start a thread to run the Moonshot UI
+    """
+    base_directory = os.getcwd()
+    ui_dev_dir = os.path.join(base_directory, "moonshot-ui")
+
+    if not os.path.exists(ui_dev_dir):
+        print("moonshot-ui does not exist. Please run with '-i moonshot-ui' to install moonshot-ui first.")
+        sys.exit(1)
+    # ms_ui_env_file(ui_dev_dir)
+    run_subprocess(['npm', 'start'], cwd=ui_dev_dir)
+
+def main():
+    parser = argparse.ArgumentParser(description="Run the Moonshot application")
+    parser.add_argument('mode', nargs='?', choices=['web-api', 'cli', 'web', 'help'], help='Mode to run Moonshot in', default=help)
+    parser.add_argument('cli_command', nargs='?', help='The CLI command to run (e.g., "interactive")')
+    parser.add_argument('-i', '--install', action='append', choices=['moonshot-data', 'moonshot-ui'], help='Modules to install', default=[])
+    parser.add_argument('-e', '--env', type=str, help='Path to the .env file', default='.env')
+    
+    args = parser.parse_args()
+    
+    # Handle installations based on the -i include arguments
+    if 'moonshot-data' in args.install:
+        moonshot_data_installation()
+
+    if 'moonshot-ui' in args.install:
+        moonshot_ui_installation()
+
+    # If mode is not specified, skip running any modes
+    if args.mode is None:
+        return
+    
+    if args.mode == "help":
+        parser.print_help()
+        sys.exit(1)
+
+    api_set_environment_variables(dotenv_values(args.env))
+
+    if args.mode == "web-api":
+        from moonshot.integrations.web_api import __main__ as web_api
+        web_api.start_app()
+    elif args.mode == "web":
+    # Create and start the UI dev server thread
+        ui_thread = threading.Thread(target=run_moonshot_ui_dev)
+        ui_thread.start()
+        ui_thread.join(timeout=0.1)  # Wait briefly for the thread to become alive
+        if not ui_thread.is_alive():
+            sys.exit(1)
+        from moonshot.integrations.web_api import __main__ as web_api
+        web_api.start_app()
+    elif args.mode == "cli":
+        from moonshot.integrations.cli import __main__ as cli
+        cli.start_app(args.cli_command)
+
+        # Handle CLI mode here, possibly also with additional arguments
+        pass
+    else:
+        parser.print_help()
+        sys.exit(1)
+
+if __name__ == "__main__":
+    warnings.filterwarnings("ignore")
+    main()

moonshot/api.py

@@ -0,0 +1,155 @@
+from moonshot.src.api.api_connector import (
+    api_create_connector_from_endpoint,
+    api_create_connectors_from_endpoints,
+    api_get_all_connector_type,
+)
+from moonshot.src.api.api_connector_endpoint import (
+    api_create_endpoint,
+    api_delete_endpoint,
+    api_get_all_endpoint,
+    api_get_all_endpoint_name,
+    api_read_endpoint,
+    api_update_endpoint,
+)
+from moonshot.src.api.api_context_strategy import (
+    api_delete_context_strategy,
+    api_get_all_context_strategies,
+    api_get_all_context_strategy_metadata,
+)
+from moonshot.src.api.api_cookbook import (
+    api_create_cookbook,
+    api_delete_cookbook,
+    api_get_all_cookbook,
+    api_get_all_cookbook_name,
+    api_read_cookbook,
+    api_read_cookbooks,
+    api_update_cookbook,
+)
+from moonshot.src.api.api_dataset import (
+    api_delete_dataset,
+    api_get_all_datasets,
+    api_get_all_datasets_name,
+)
+from moonshot.src.api.api_environment_variables import api_set_environment_variables
+from moonshot.src.api.api_metrics import (
+    api_delete_metric,
+    api_get_all_metric,
+    api_get_all_metric_name,
+)
+from moonshot.src.api.api_prompt_template import (
+    api_delete_prompt_template,
+    api_get_all_prompt_template_detail,
+    api_get_all_prompt_template_name,
+)
+from moonshot.src.api.api_recipe import (
+    api_create_recipe,
+    api_delete_recipe,
+    api_get_all_recipe,
+    api_get_all_recipe_name,
+    api_read_recipe,
+    api_read_recipes,
+    api_update_recipe,
+)
+from moonshot.src.api.api_red_teaming import (
+    api_delete_attack_module,
+    api_get_all_attack_module_metadata,
+    api_get_all_attack_modules,
+)
+from moonshot.src.api.api_result import (
+    api_delete_result,
+    api_get_all_result,
+    api_get_all_result_name,
+    api_read_result,
+    api_read_results,
+)
+from moonshot.src.api.api_run import api_get_all_run
+from moonshot.src.api.api_runner import (
+    api_create_runner,
+    api_delete_runner,
+    api_get_all_runner,
+    api_get_all_runner_name,
+    api_load_runner,
+    api_read_runner,
+)
+from moonshot.src.api.api_session import (
+    api_create_session,
+    api_delete_session,
+    api_get_all_chats_from_session,
+    api_get_all_session_metadata,
+    api_get_all_session_names,
+    api_get_available_session_info,
+    api_load_session,
+    api_update_attack_module,
+    api_update_context_strategy,
+    api_update_cs_num_of_prev_prompts,
+    api_update_metric,
+    api_update_prompt_template,
+    api_update_system_prompt,
+)
+
+__all__ = [
+    "api_create_connector_from_endpoint",
+    "api_create_connectors_from_endpoints",
+    "api_get_all_connector_type",
+    "api_create_endpoint",
+    "api_delete_endpoint",
+    "api_get_all_endpoint",
+    "api_get_all_endpoint_name",
+    "api_read_endpoint",
+    "api_update_endpoint",
+    "api_delete_context_strategy",
+    "api_get_all_context_strategies",
+    "api_get_all_context_strategy_metadata",
+    "api_create_cookbook",
+    "api_delete_cookbook",
+    "api_get_all_cookbook",
+    "api_get_all_cookbook_name",
+    "api_read_cookbook",
+    "api_read_cookbooks",
+    "api_update_cookbook",
+    "api_delete_dataset",
+    "api_get_all_datasets",
+    "api_get_all_datasets_name",
+    "api_set_environment_variables",
+    "api_delete_metric",
+    "api_get_all_metric",
+    "api_get_all_metric_name",
+    "api_get_all_prompt_template_detail",
+    "api_get_all_prompt_template_name",
+    "api_delete_prompt_template",
+    "api_create_recipe",
+    "api_delete_recipe",
+    "api_get_all_recipe",
+    "api_get_all_recipe_name",
+    "api_read_recipe",
+    "api_read_recipes",
+    "api_update_recipe",
+    "api_get_all_attack_module_metadata",
+    "api_get_all_attack_modules",
+    "api_delete_attack_module",
+    "api_delete_result",
+    "api_get_all_result",
+    "api_get_all_result_name",
+    "api_read_result",
+    "api_read_results",
+    "api_get_all_run",
+    "api_create_runner",
+    "api_delete_runner",
+    "api_get_all_runner",
+    "api_get_all_runner_name",
+    "api_load_runner",
+    "api_read_runner",
+    "api_create_session",
+    "api_delete_session",
+    "api_get_all_chats_from_session",
+    "api_get_all_session_metadata",
+    "api_get_all_session_names",
+    "api_get_available_session_info",
+    "api_load_session",
+    "api_update_attack_module",
+    "api_update_context_strategy",
+    "api_update_cs_num_of_prev_prompts",
+    "api_update_metric",
+    "api_update_prompt_template",
+    "api_update_system_prompt",
+]

moonshot/integrations/cli/__main__.py

@@ -0,0 +1,25 @@
+import sys
+import warnings
+
+from moonshot.integrations.cli.cli import CommandLineInterface
+
+
+def start_app(cli_command = None):
+    """
+    Run the Moonshot application
+    """
+    # Setting the warnings to be ignored
+    warnings.filterwarnings("ignore")
+
+    cli_instance = CommandLineInterface()
+    if cli_command == "interactive":
+        # Run in interactive mode
+        cli_instance.debug = False
+        cli_instance.cmdloop("Starting moonshot interactive prompt...")
+    elif cli_command:
+        # Run a specific command
+        cli_instance.onecmd(cli_command)
+    else:
+        # Show help if no command is provided
+        cli_instance.onecmd("help")
+              

moonshot/integrations/cli/active_session_cfg.py

@@ -0,0 +1 @@
+active_session = {}

moonshot/integrations/cli/benchmark/benchmark.py

@@ -0,0 +1,186 @@
+import argparse
+
+import cmd2
+
+from moonshot.integrations.cli.benchmark.cookbook import (
+    add_cookbook,
+    add_cookbook_args,
+    delete_cookbook,
+    delete_cookbook_args,
+    list_cookbooks,
+    run_cookbook,
+    run_cookbook_args,
+    update_cookbook,
+    update_cookbook_args,
+    view_cookbook,
+    view_cookbook_args,
+)
+from moonshot.integrations.cli.benchmark.datasets import (
+    delete_dataset,
+    delete_dataset_args,
+    list_datasets,
+    view_dataset,
+    view_dataset_args,
+)
+from moonshot.integrations.cli.benchmark.metrics import (
+    delete_metric,
+    delete_metric_args,
+    list_metrics,
+    view_metric,
+    view_metric_args,
+)
+from moonshot.integrations.cli.benchmark.recipe import (
+    add_recipe,
+    add_recipe_args,
+    delete_recipe,
+    delete_recipe_args,
+    list_recipes,
+    run_recipe,
+    run_recipe_args,
+    update_recipe,
+    update_recipe_args,
+    view_recipe,
+    view_recipe_args,
+)
+from moonshot.integrations.cli.benchmark.result import (
+    delete_result,
+    delete_result_args,
+    list_results,
+    view_result,
+    view_result_args,
+)
+from moonshot.integrations.cli.benchmark.run import list_runs, view_run, view_run_args
+from moonshot.integrations.cli.benchmark.runner import (
+    delete_runner,
+    delete_runner_args,
+    list_runners,
+    view_runner,
+    view_runner_args,
+)
+
+
+@cmd2.with_default_category("Moonshot Benchmarking")
+class BenchmarkCommandSet(cmd2.CommandSet):
+    def __init__(self):
+        super().__init__()
+
+    # ------------------------------------------------------------------------------
+    # List contents
+    # ------------------------------------------------------------------------------
+
+    def do_list_cookbooks(self, _: cmd2.Statement) -> None:
+        list_cookbooks()
+
+    def do_list_recipes(self, _: cmd2.Statement) -> None:
+        list_recipes()
+
+    def do_list_results(self, _: cmd2.Statement) -> None:
+        list_results()
+
+    def do_list_runners(self, _: cmd2.Statement) -> None:
+        list_runners()
+
+    def do_list_runs(self, _: cmd2.Statement) -> None:
+        list_runs()
+
+    def do_list_metrics(self, _: cmd2.Statement) -> None:
+        list_metrics()
+
+    def do_list_datasets(self, _: cmd2.Statement) -> None:
+        list_datasets()
+
+    # ------------------------------------------------------------------------------
+    # Add contents
+    # ------------------------------------------------------------------------------
+
+    @cmd2.with_argparser(add_cookbook_args)
+    def do_add_cookbook(self, args: argparse.Namespace) -> None:
+        add_cookbook(args)
+
+    @cmd2.with_argparser(add_recipe_args)
+    def do_add_recipe(self, args: argparse.Namespace) -> None:
+        add_recipe(args)
+
+    # ------------------------------------------------------------------------------
+    # Delete contents
+    # ------------------------------------------------------------------------------
+
+    @cmd2.with_argparser(delete_cookbook_args)
+    def do_delete_cookbook(self, args: argparse.Namespace) -> None:
+        delete_cookbook(args)
+
+    @cmd2.with_argparser(delete_recipe_args)
+    def do_delete_recipe(self, args: argparse.Namespace) -> None:
+        delete_recipe(args)
+
+    @cmd2.with_argparser(delete_result_args)
+    def do_delete_result(self, args: argparse.Namespace) -> None:
+        delete_result(args)
+
+    @cmd2.with_argparser(delete_runner_args)
+    def do_delete_runner(self, args: argparse.Namespace) -> None:
+        delete_runner(args)
+
+    @cmd2.with_argparser(delete_metric_args)
+    def do_delete_metric(self, args: argparse.Namespace) -> None:
+        delete_metric(args)
+
+    @cmd2.with_argparser(delete_dataset_args)
+    def do_delete_dataset(self, args: argparse.Namespace) -> None:
+        delete_dataset(args)
+
+    # ------------------------------------------------------------------------------
+    # Update contents
+    # ------------------------------------------------------------------------------
+
+    @cmd2.with_argparser(update_cookbook_args)
+    def do_update_cookbook(self, args: argparse.Namespace) -> None:
+        update_cookbook(args)
+
+    @cmd2.with_argparser(update_recipe_args)
+    def do_update_recipe(self, args: argparse.Namespace) -> None:
+        update_recipe(args)
+
+    # ------------------------------------------------------------------------------
+    # Run contents
+    # ------------------------------------------------------------------------------
+
+    @cmd2.with_argparser(run_cookbook_args)
+    def do_run_cookbook(self, args: argparse.Namespace) -> None:
+        run_cookbook(args)
+
+    @cmd2.with_argparser(run_recipe_args)
+    def do_run_recipe(self, args: argparse.Namespace) -> None:
+        run_recipe(args)
+
+    # ------------------------------------------------------------------------------
+    # View contents
+    # ------------------------------------------------------------------------------
+
+    @cmd2.with_argparser(view_cookbook_args)
+    def do_view_cookbook(self, args: argparse.Namespace) -> None:
+        view_cookbook(args)
+
+    @cmd2.with_argparser(view_recipe_args)
+    def do_view_recipe(self, args: argparse.Namespace) -> None:
+        view_recipe(args)
+
+    @cmd2.with_argparser(view_result_args)
+    def do_view_result(self, args: argparse.Namespace) -> None:
+        view_result(args)
+
+    @cmd2.with_argparser(view_runner_args)
+    def do_view_runner(self, args: argparse.Namespace) -> None:
+        view_runner(args)
+
+    @cmd2.with_argparser(view_run_args)
+    def do_view_run(self, args: argparse.Namespace) -> None:
+        view_run(args)
+
+    @cmd2.with_argparser(view_metric_args)
+    def do_view_metric(self, args: argparse.Namespace) -> None:
+        view_metric(args)
+
+    @cmd2.with_argparser(view_dataset_args)
+    def do_view_dataset(self, args: argparse.Namespace) -> None:
+        view_dataset(args)