aiverify-moonshot 0.4.0__py3-none-any.whl

moonshot/src/results/result.py

@@ -0,0 +1,119 @@
+from __future__ import annotations
+
+from pathlib import Path
+
+from pydantic import validate_call
+
+from moonshot.src.configs.env_variables import EnvVariables
+from moonshot.src.storage.storage import Storage
+
+
+class Result:
+    @staticmethod
+    @validate_call
+    def read(result_id: str) -> dict:
+        """
+        Reads the result data from storage for a given result ID.
+
+        This method attempts to retrieve the result data associated with the specified result ID from storage.
+        If the data is found, it is returned as a dictionary. If no data is found, an exception is raised and
+        an error message is printed.
+
+        Args:
+            result_id (str): The unique identifier of the result to be read.
+
+        Returns:
+            dict: A dictionary containing the result data if found.
+
+        Raises:
+            Exception: If no result data is found or if an error occurs during the read operation.
+        """
+        try:
+            if result_id:
+                return Result._read_result(result_id)
+            else:
+                raise RuntimeError("Result ID is empty")
+
+        except Exception as e:
+            print(f"Failed to read result: {str(e)}")
+            raise e
+
+    @staticmethod
+    def _read_result(result_id: str) -> dict:
+        """
+        Reads the result data from storage for a given result ID.
+
+        This method attempts to retrieve the result data associated with the specified result ID from storage.
+        If the data is found, it is returned as a dictionary. If no data is found, a RuntimeError is raised.
+
+        Args:
+            result_id (str): The unique identifier of the result to be read.
+
+        Returns:
+            dict: A dictionary containing the result data.
+
+        Raises:
+            RuntimeError: If no result data is found for the given result ID.
+        """
+        obj_results = Storage.read_object(EnvVariables.RESULTS.name, result_id, "json")
+        if obj_results:
+            return obj_results
+        else:
+            raise RuntimeError(f"Unable to get results for {result_id}.")
+
+    @staticmethod
+    @validate_call
+    def delete(result_id: str) -> bool:
+        """
+        Deletes the result data associated with the given result ID from storage.
+
+        This method attempts to delete the result data identified by the specified result ID from storage.
+        If the deletion is successful, it returns True. If an exception occurs during the deletion process,
+        an error message is printed and the exception is re-raised.
+
+        Args:
+            result_id (str): The unique identifier of the result to be deleted.
+
+        Returns:
+            bool: True if the result data was successfully deleted.
+
+        Raises:
+            Exception: If an error occurs during the deletion process.
+        """
+        try:
+            Storage.delete_object(EnvVariables.RESULTS.name, result_id, "json")
+            return True
+
+        except Exception as e:
+            print(f"Failed to delete result: {str(e)}")
+            raise e
+
+    @staticmethod
+    def get_available_items() -> tuple[list[str], list[dict]]:
+        """
+        Retrieves the list of available result IDs and their corresponding result data.
+
+        This method queries the storage to obtain all the result objects, filters out any that are not relevant
+        (e.g., internal use objects with "__" in their name), and then reads the result data for each remaining
+        result ID. It returns a tuple containing a list of result IDs and a list of dictionaries with the result
+        data.
+
+        Returns:
+            tuple[list[str], list[dict]]: A tuple with the first element being a list of result IDs and the
+            second element being a list of dictionaries containing the result data for each ID.
+        """
+        try:
+            retn_results = []
+            retn_results_ids = []
+
+            for result in Storage.get_objects(EnvVariables.RESULTS.name, "json"):
+                if "__" in result:
+                    continue
+                result_info = Result._read_result(Path(result).stem)
+                retn_results.append(result_info)
+                retn_results_ids.append(Path(result).stem)
+            return retn_results_ids, retn_results
+
+        except Exception as e:
+            print(f"Failed to get available results: {str(e)}")
+            raise e

moonshot/src/results/result_arguments.py

@@ -0,0 +1,44 @@
+from __future__ import annotations
+
+from pydantic import BaseModel
+
+from moonshot.src.runs.run_status import RunStatus
+
+
+class ResultArguments(BaseModel):
+    id: str  # The ID of the Runner.
+
+    start_time: float  # The start time of the Run.
+
+    end_time: float  # The end time of the Run.
+
+    duration: int  # The duration of the Run.
+
+    status: RunStatus  # Status of the Run.
+
+    raw_results: dict = {}  # Raw Results of the Run from runners-modules.
+
+    results: dict = {}  # Results of the Run from results-modules.
+
+    params: dict = {}  # Other information required for results module
+
+    def to_dict(self) -> dict:
+        """
+        Transforms the ResultArguments instance into a dictionary format.
+
+        This method serializes the ResultArguments instance into a dictionary where attribute names become keys
+        and their corresponding values are the dictionary values.
+
+        Returns:
+            A dictionary representation of the ResultArguments instance.
+        """
+        return {
+            "id": self.id,
+            "start_time": self.start_time,
+            "end_time": self.end_time,
+            "duration": self.duration,
+            "status": self.status.name,
+            "raw_results": self.raw_results,
+            "results": self.results,
+            "params": self.params,
+        }

moonshot/src/runners/runner.py

@@ -0,0 +1,476 @@
+from __future__ import annotations
+
+import asyncio
+from pathlib import Path
+from typing import Callable
+
+from pydantic import validate_call
+from slugify import slugify
+
+from moonshot.src.configs.env_variables import EnvVariables
+from moonshot.src.redteaming.session.session import Session
+from moonshot.src.runners.runner_arguments import RunnerArguments
+from moonshot.src.runners.runner_type import RunnerType
+from moonshot.src.runs.run import Run
+from moonshot.src.storage.storage import Storage
+
+
+class Runner:
+    sql_create_runner_cache_table = """
+        CREATE TABLE IF NOT EXISTS runner_cache_table (
+        id INTEGER PRIMARY KEY AUTOINCREMENT,
+        connection_id text NOT NULL,
+        recipe_id text,
+        dataset_id text,
+        prompt_template_id text,
+        context_strategy_id text,
+        attack_module_id text,
+        prompt_index INTEGER,
+        prompt text NOT NULL,
+        target text NOT NULL,
+        predicted_results text NOT NULL,
+        duration text NOT NULL,
+        random_seed INTEGER,
+        system_prompt text
+        );
+    """
+
+    def __init__(self, runner_args: RunnerArguments) -> None:
+        self.id = runner_args.id
+        self.name = runner_args.name
+        self.description = runner_args.description
+        self.endpoints = runner_args.endpoints
+        self.database_instance = runner_args.database_instance
+        self.database_file = runner_args.database_file
+        self.progress_callback_func = runner_args.progress_callback_func
+
+        # Set current run
+        self.current_operation = None
+        self.current_operation_lock = asyncio.Lock()  # Mutex lock for current operation
+
+    @classmethod
+    def load(
+        cls, runner_id: str, progress_callback_func: Callable | None = None
+    ) -> Runner:
+        """
+        This method is responsible for loading an existing runner.
+
+        It accepts a runner_id and an optional progress_callback_func as arguments. The method first verifies the
+        existence of the runner file corresponding to the provided runner_id. If the runner file does not exist,
+        it raises a RuntimeError.
+        If the runner file is found, the method reads the file and establishes a database connection. It then assigns
+        the progress_callback_func (if provided) to the runner and returns a new Runner instance.
+
+        Args:
+            runner_id (str): The unique identifier of the runner to be loaded.
+            progress_callback_func (Callable | None): An optional callback function for tracking the progress of
+            the runner.
+
+        Returns:
+            Runner: An instance of the Runner class, initialized with the data loaded from the runner file.
+
+        Raises:
+            RuntimeError: If the runner file corresponding to the provided runner_id does not exist.
+        """
+        try:
+            # Check if runner file exists. If it does not exists, raise an error.
+            if not Storage.is_object_exists(
+                EnvVariables.RUNNERS.name, runner_id, "json"
+            ):
+                raise RuntimeError(
+                    "[Runner] Unable to load runner because the runner file does not exist."
+                )
+            runner_args = Runner.read(runner_id)
+            runner_args.database_instance = Storage.create_database_connection(
+                EnvVariables.DATABASES.name, runner_id, "db"
+            )
+            runner_args.progress_callback_func = progress_callback_func
+            return cls(runner_args)
+
+        except Exception as e:
+            print(f"[Runner] Failed to load runner: {str(e)}")
+            raise e
+
+    @classmethod
+    def create(cls, runner_args: RunnerArguments) -> Runner:
+        """
+        Creates a new runner instance.
+
+        This method takes a RunnerArguments object to generate a unique runner_id from the runner's name.
+        It checks for the existence of a runner file with the same id.
+        If found, a RuntimeError is raised to indicate the conflict.
+        Otherwise, it proceeds to create a new runner file and sets up a database connection for the runner.
+        A new Runner class instance, initialized with the provided arguments, is then returned.
+
+        Args:
+            runner_args (RunnerArguments): The configuration parameters for creating the runner.
+
+        Returns:
+            Runner: An instance of the Runner class, newly created with the specified arguments.
+
+        Raises:
+            RuntimeError: Raised if a runner file with the generated runner_id already exists,
+            indicating a duplicate runner.
+        """
+        try:
+            runner_id = slugify(runner_args.name, lowercase=True)
+
+            # Check if runner file exists. If it exists, raise an error.
+            if Storage.is_object_exists(EnvVariables.RUNNERS.name, runner_id, "json"):
+                raise RuntimeError(
+                    "[Runner] Unable to create runner because the runner file exists."
+                )
+            # Check if all endpoint configuration files exist. If not, raise an error.
+            for endpoint in runner_args.endpoints:
+                if not Storage.is_object_exists(
+                    EnvVariables.CONNECTORS_ENDPOINTS.name, endpoint, "json"
+                ):
+                    raise RuntimeError(
+                        f"[Runner] Connector endpoint {endpoint} does not exist."
+                    )
+
+            runner_info = {
+                "id": runner_id,
+                "name": runner_args.name,
+                "endpoints": runner_args.endpoints,
+                "database_file": Storage.get_filepath(
+                    EnvVariables.DATABASES.name, runner_id, "db", True
+                ),
+                "progress_callback_func": runner_args.progress_callback_func,
+                "description": runner_args.description,
+            }
+            runner_args = RunnerArguments(**runner_info)
+            runner_args.database_instance = Storage.create_database_connection(
+                EnvVariables.DATABASES.name, runner_id, "db"
+            )
+
+            # Create runner file
+            Storage.create_object(
+                EnvVariables.RUNNERS.name, runner_id, runner_args.to_dict(), "json"
+            )
+
+            # Create runner cache table
+            Storage.create_database_table(
+                runner_args.database_instance, Runner.sql_create_runner_cache_table
+            )
+
+            return cls(runner_args)
+
+        except Exception as e:
+            print(f"[Runner] Failed to create runner: {str(e)}")
+            raise e
+
+    @staticmethod
+    @validate_call
+    def read(runner_id: str) -> RunnerArguments:
+        """
+        Retrieves the runner data and constructs a RunnerArguments object.
+
+        This method accepts a runner_id as an input and utilizes the StorageManager to fetch the runner data.
+        It subsequently builds a RunnerArguments object using the fetched data and returns this object.
+
+        Args:
+            runner_id (str): The unique identifier of the runner.
+
+        Returns:
+            RunnerArguments: An object of RunnerArguments constructed with the runner's data.
+
+        Raises:
+            Exception: If an error occurs during the data retrieval or any other operation within the method.
+        """
+        try:
+            if runner_id:
+                return RunnerArguments(
+                    **Storage.read_object(EnvVariables.RUNNERS.name, runner_id, "json")
+                )
+            else:
+                raise RuntimeError("Runner ID is empty")
+
+        except Exception as e:
+            print(f"[Runner] Failed to read runner: {str(e)}")
+            raise e
+
+    @staticmethod
+    @validate_call
+    def delete(runner_id: str) -> bool:
+        """
+        Deletes the runner and its associated database instance.
+
+        This method attempts to delete the runner identified by the provided runner_id from storage.
+        It also attempts to delete the associated database instance. If both deletions are successful,
+        it returns True. If an exception occurs during the deletion process, an error message is printed
+        and the exception is re-raised.
+
+        Args:
+            runner_id (str): The unique identifier of the runner to be deleted.
+
+        Returns:
+            bool: True if the runner and its associated database instance were successfully deleted.
+
+        Raises:
+            Exception: If an error occurs during the deletion process.
+        """
+        try:
+            Storage.delete_object(EnvVariables.RUNNERS.name, runner_id, "json")
+            Storage.delete_object(EnvVariables.DATABASES.name, runner_id, "db")
+            return True
+
+        except Exception as e:
+            print(f"[Runner] Failed to delete runner: {str(e)}")
+            raise e
+
+    @staticmethod
+    def get_available_items() -> tuple[list[str], list[RunnerArguments]]:
+        """
+        Retrieves and returns a list of available runners.
+
+        This method scans the directory specified by `EnvironmentVars.RUNNERS` and collects all stored runner files.
+        It excludes any files that contain "__" in their names. For each valid runner file, the method reads the file
+        content and constructs a RunnerArguments object encapsulating the runner's details. Both the RunnerArguments
+        object and the runner ID are then appended to their respective lists.
+
+        Returns:
+            tuple[list[str], list[RunnerArguments]]: A tuple where the first element is a list of runner IDs and
+            the second element is a list of RunnerArguments objects representing the details of each runner.
+
+        Raises:
+            Exception: If an error is encountered during the file reading process or any other operation within
+            the method.
+        """
+        try:
+            retn_runners = []
+            retn_runners_ids = []
+            runners = Storage.get_objects(EnvVariables.RUNNERS.name, "json")
+            for runner in runners:
+                if "__" in runner:
+                    continue
+
+                runner_info = RunnerArguments(
+                    **Storage.read_object(
+                        EnvVariables.RUNNERS.name, Path(runner).stem, "json"
+                    )
+                )
+                retn_runners.append(runner_info)
+                retn_runners_ids.append(runner_info.id)
+
+            return retn_runners_ids, retn_runners
+
+        except Exception as e:
+            print(f"[Runner] Failed to get available runners: {str(e)}")
+            raise e
+
+    def close(self) -> None:
+        """
+        Closes the runner instance.
+
+        This method is responsible for closing the runner instance. If a database instance is associated with the
+        runner, it also closes the database connection using the StorageManager's close_database_connection
+        method.
+
+        Raises:
+            Exception: If any error occurs while closing the runner or the database connection.
+        """
+        if self.database_instance:
+            Storage.close_database_connection(self.database_instance)
+
+    async def cancel(self) -> None:
+        """
+        Cancels the runner instance.
+
+        This method is responsible for cancelling the runner instance. If a run is currently in progress,
+        it stops the run and releases any resources associated with it.
+
+        Raises:
+            Exception: If any error occurs while cancelling the runner or releasing the resources.
+        """
+        async with self.current_operation_lock:
+            if self.current_operation:
+                print(f"[Runner] {self.id} - Cancelling current operation...")
+                self.current_operation.cancel()
+                self.current_operation = None  # Reset the current operation
+
+    async def run_recipes(
+        self,
+        recipes: list[str],
+        num_of_prompts: int = 0,
+        random_seed: int = 0,
+        system_prompt: str = "",
+        runner_processing_module: str = "benchmarking",
+        result_processing_module: str = "benchmarking-result",
+    ) -> None:
+        """
+        Initiates an asynchronous benchmark run using a set of recipes.
+
+        This method sets up and starts a benchmark run tailored for recipes. It instantiates a benchmark run object,
+        applies the configuration based on the provided recipes, number of prompts, random seed, system prompt, and
+        the specified runner and result processing modules, and then commences the run asynchronously.
+
+        Args:
+            recipes (list[str]): The recipes to be included in the benchmark run.
+
+            num_of_prompts (int, optional): The count of prompts to utilize during the benchmark.
+            Defaults to 0.
+
+            random_seed (int, optional): The seed for random number generation to ensure reproducibility.
+            Defaults to 0.
+
+            system_prompt (str, optional): The system prompt to be used during the benchmark.
+            Defaults to an empty string.
+
+            runner_processing_module (str, optional): The module responsible for processing the runner.
+            Defaults to "benchmarking".
+
+            result_processing_module (str, optional): The module responsible for processing the results.
+            Defaults to "benchmarking-result".
+
+        Raises:
+            Exception: If any error occurs during the setup or execution of the benchmark run.
+        """
+        async with self.current_operation_lock:  # Acquire the lock
+            # Create new benchmark recipe test run
+            print(f"[Runner] {self.id} - Running benchmark recipe run...")
+            self.current_operation = Run(
+                self.id,
+                RunnerType.BENCHMARK,
+                {
+                    "recipes": recipes,
+                    "num_of_prompts": num_of_prompts,
+                    "random_seed": random_seed,
+                    "system_prompt": system_prompt,
+                    "runner_processing_module": runner_processing_module,
+                    "result_processing_module": result_processing_module,
+                },
+                self.database_instance,
+                self.endpoints,
+                Storage.get_filepath(EnvVariables.RESULTS.name, self.id, "json", True),
+                self.progress_callback_func,
+            )
+            # Note: The lock is held during setup but should be released before long-running operations
+
+        # Execute the long-running operation outside of the lock
+        # Run new benchmark recipe test run
+        await self.current_operation.run()
+
+        # After completion, reset current_operation to None within the lock
+        async with self.current_operation_lock:
+            self.current_operation = None
+            print(f"[Runner] {self.id} - Benchmark recipe run completed and reset.")
+
+    async def run_cookbooks(
+        self,
+        cookbooks: list[str],
+        num_of_prompts: int = 0,
+        random_seed: int = 0,
+        system_prompt: str = "",
+        runner_processing_module: str = "benchmarking",
+        result_processing_module: str = "benchmarking-result",
+    ) -> None:
+        """
+        Asynchronously runs a set of cookbooks with the provided parameters.
+
+        This method is responsible for initiating a benchmark cookbook run with the specified cookbooks and parameters.
+        It creates a new benchmark cookbook run instance, configures it with the provided cookbook names,
+        number of prompts, random seed, system prompt, runner processing module, and result processing module,
+        and then starts the run asynchronously.
+
+        Args:
+            cookbooks (list[str]): A list of cookbook names to be run in the benchmark.
+
+            num_of_prompts (int, optional): The number of prompts to be used in the benchmark run.
+            Defaults to 0.
+
+            random_seed (int, optional): The seed for random number generation to ensure reproducibility.
+            Defaults to 0.
+
+            system_prompt (str, optional): A system prompt to be used in the benchmark run.
+            Defaults to an empty string.
+
+            runner_processing_module (str, optional): The module responsible for processing the runner.
+            Defaults to "benchmarking".
+
+            result_processing_module (str, optional): The module responsible for processing the results.
+            Defaults to "benchmarking-result".
+
+        Raises:
+            Exception: If any error occurs during the setup or execution of the benchmark run.
+        """
+        async with self.current_operation_lock:  # Acquire the lock
+            # Create new benchmark cookbook test run
+            print(f"[Runner] {self.id} - Running benchmark cookbook run...")
+            self.current_operation = Run(
+                self.id,
+                RunnerType.BENCHMARK,
+                {
+                    "cookbooks": cookbooks,
+                    "num_of_prompts": num_of_prompts,
+                    "random_seed": random_seed,
+                    "system_prompt": system_prompt,
+                    "runner_processing_module": runner_processing_module,
+                    "result_processing_module": result_processing_module,
+                },
+                self.database_instance,
+                self.endpoints,
+                Storage.get_filepath(EnvVariables.RESULTS.name, self.id, "json", True),
+                self.progress_callback_func,
+            )
+            # Note: The lock is held during setup but should be released before long-running operations
+
+        # Execute the long-running operation outside of the lock
+        # Run new benchmark cookbook test run
+        await self.current_operation.run()
+
+        # After completion, reset current_operation to None within the lock
+        async with self.current_operation_lock:
+            self.current_operation = None
+            print(f"[Runner] {self.id} - Benchmark cookbook run completed and reset.")
+
+    async def run_red_teaming(
+        self,
+        red_team_args: dict,
+        system_prompt: str = "",
+        runner_processing_module: str = "redteaming",
+    ) -> list | None:
+        """
+        Asynchronously runs a red teaming session with the provided arguments.
+
+        This method is responsible for initiating a red teaming session with the specified arguments. It creates a new
+        red teaming session instance, configures it with the provided red teaming arguments, system prompt, and
+        runner processing module, and then starts the session asynchronously.
+
+        Args:
+            red_team_args (dict): A dictionary of arguments for the red teaming session.
+
+            system_prompt (str, optional): A system prompt to be used in the red teaming session.
+            Defaults to an empty string.
+
+            runner_processing_module (str, optional): The processing module to be used for the session.
+            Defaults to "redteaming".
+
+        Raises:
+            Exception: If any error occurs during the setup or execution of the red teaming session.
+        """
+        async with self.current_operation_lock:  # Acquire the lock
+            print(f"[Runner] {self.id} - Running red teaming session...")
+            self.current_operation = Session(
+                self.id,
+                RunnerType.REDTEAM,
+                {
+                    **red_team_args,
+                    "runner_processing_module": runner_processing_module,
+                },
+                self.database_instance,
+                self.endpoints,
+                Storage.get_filepath(EnvVariables.RESULTS.name, self.id, "json", True),
+                self.progress_callback_func,
+            )
+
+        # Note: The lock is held during setup but should be released before long-running operations
+        # Execute the long-running operation outside of the lock
+        red_teaming_results = await self.current_operation.run()
+
+        # After completion, reset current_operation to None within the lock
+        async with self.current_operation_lock:
+            self.current_operation = None
+            print(f"[Runner] {self.id} - Red teaming run completed.")
+
+        return red_teaming_results

moonshot/src/runners/runner_arguments.py

@@ -0,0 +1,46 @@
+from __future__ import annotations
+
+from typing import Any, Callable
+
+from pydantic import BaseModel, Field
+
+
+class RunnerArguments(BaseModel):
+    id: str  # The ID of the Runner.
+
+    name: str = Field(min_length=1)  # The name of the Runner.
+
+    database_file: str = ""  # The database file associated with the Runner.
+
+    endpoints: list[str] = Field(min_length=1)  # List of endpoints for the Runner.
+
+    description: str = ""  # A brief description of the Runner.
+
+    # ------------------------------------------------------------------------------
+    # These attributes are not exported to dict
+    # ------------------------------------------------------------------------------
+    database_instance: Any | None = (
+        None  # The database instance associated with the Runner.
+    )
+
+    progress_callback_func: Callable | None = (
+        None  # The progress callback function for the Runner.
+    )
+
+    def to_dict(self) -> dict:
+        """
+        Transforms the RunnerArguments instance into a dictionary format.
+
+        This method serializes the RunnerArguments instance, excluding 'database_instance' and 'progress_callback_func',
+        into a dictionary where attribute names become keys and their corresponding values are the dictionary values.
+
+        Returns:
+            A dictionary representation of the RunnerArguments instance, excluding non-serializable attributes.
+        """
+        return {
+            "id": self.id,
+            "name": self.name,
+            "database_file": self.database_file,
+            "endpoints": self.endpoints,
+            "description": self.description,
+        }