aiverify-moonshot 0.4.0__py3-none-any.whl

moonshot/src/runners/runner_type.py

@@ -0,0 +1,6 @@
+from enum import Enum
+
+
+class RunnerType(Enum):
+    BENCHMARK = "benchmark"
+    REDTEAM = "redteam"

moonshot/src/runs/run.py

@@ -0,0 +1,344 @@
+from __future__ import annotations
+
+import asyncio
+import time
+from typing import Any, Callable
+
+from moonshot.src.configs.env_variables import EnvVariables
+from moonshot.src.results.result_arguments import ResultArguments
+from moonshot.src.runners.runner_type import RunnerType
+from moonshot.src.runs.run_arguments import RunArguments
+from moonshot.src.runs.run_progress import RunProgress
+from moonshot.src.runs.run_status import RunStatus
+from moonshot.src.storage.db_interface import DBInterface
+from moonshot.src.storage.storage import Storage
+from moonshot.src.utils.import_modules import get_instance
+
+
+class Run:
+    sql_create_run_table = """
+        CREATE TABLE IF NOT EXISTS run_table (
+        run_id INTEGER PRIMARY KEY AUTOINCREMENT,
+        runner_id text NOT NULL,
+        runner_type text NOT NULL,
+        runner_args text NOT NULL,
+        endpoints text NOT NULL,
+        results_file text NOT NULL,
+        start_time INTEGER NOT NULL,
+        end_time INTEGER NOT NULL,
+        duration INTEGER NOT NULL,
+        error_messages text NOT NULL,
+        raw_results text NOT NULL,
+        results text NOT NULL,
+        status text NOT NULL
+        );
+    """
+    sql_create_run_record = """
+        INSERT INTO run_table (
+        runner_id,runner_type,runner_args,endpoints,results_file,start_time,end_time,duration,error_messages,raw_results,results,status)
+        VALUES(?,?,?,?,?,?,?,?,?,?,?,?)
+    """
+    sql_read_run_record = """
+        SELECT * from run_table WHERE run_id=?
+    """
+    sql_read_latest_run_record = """
+        SELECT * FROM run_table WHERE run_id=(SELECT MAX(run_id) FROM run_table)
+    """
+    sql_read_all_run_records = """
+        SELECT * FROM run_table
+    """
+
+    def __init__(
+        self,
+        runner_id: str,
+        runner_type: RunnerType,
+        runner_args: dict,
+        database_instance: Any | None,
+        endpoints: list[str],
+        results_file: str,
+        progress_callback_func: Callable | None = None,
+    ) -> None:
+        # Create run arguments
+        self.run_arguments = RunArguments(
+            runner_id=runner_id,
+            runner_type=runner_type,
+            runner_args=runner_args,
+            database_instance=database_instance,
+            endpoints=endpoints,
+            results_file=results_file,
+            start_time=0.0,
+            end_time=0.0,
+            duration=0,
+            error_messages=[],
+            raw_results={},
+            results={},
+            status=RunStatus.PENDING,
+        )
+        # Pass the reference of run_arguments to RunProgress
+        self.run_progress = RunProgress(
+            self.run_arguments,
+            progress_callback_func,
+        )
+        # Create a cancellation asyncio event
+        self.cancel_event = asyncio.Event()
+        # Create run table
+        if database_instance:
+            Storage.create_database_table(database_instance, Run.sql_create_run_table)
+
+    @staticmethod
+    def load(database_instance: DBInterface | None, run_id: int | None) -> RunArguments:
+        """
+        Loads run data for a given run_id from the database, or the latest run if run_id is None.
+
+        This method retrieves run data for the specified run_id from the database, or if run_id is None,
+        it retrieves the latest run. If the database instance is not provided, it raises a RuntimeError.
+        If the database instance is provided, it invokes the read_record method of the database instance
+        with the given run_id or the latest run and returns a RunArguments object created from the retrieved record.
+
+        Parameters:
+            database_instance (DBAccessor | None): The database accessor instance.
+            run_id (int | None): The ID of the run to retrieve, or None to retrieve the latest run.
+
+        Returns:
+            RunArguments: An object containing the details of the run with the given run_id or the latest run.
+        """
+        if not database_instance:
+            raise RuntimeError("[Run] Database instance not provided.")
+
+        if run_id is not None:
+            run_arguments_info = Storage.read_database_record(
+                database_instance,
+                (run_id,),
+                Run.sql_read_run_record,
+            )
+        else:
+            run_arguments_info = Storage.read_database_record(
+                database_instance,
+                (),
+                Run.sql_read_latest_run_record,
+            )
+
+        if run_arguments_info:
+            return RunArguments.from_tuple(run_arguments_info)
+        else:
+            raise RuntimeError(
+                f"[Run] Failed to get database record for run_id {run_id}: {database_instance}"
+            )
+
+    @staticmethod
+    def get_all_runs(database_instance: DBInterface) -> list[RunArguments]:
+        """
+        Retrieves all run records from the database.
+
+        This method fetches all the run records from the database and converts them into a list of RunArguments objects.
+        If the database instance is not provided, it raises a RuntimeError. If no records are found, it also raises
+        a RuntimeError.
+
+        Parameters:
+            database_instance (DBInterface): The database interface to fetch records from.
+
+        Returns:
+            list[RunArguments]: A list of RunArguments objects representing each run record.
+        """
+        if not database_instance:
+            raise RuntimeError("[Run] Database instance not provided.")
+
+        all_run_arguments_info = Storage.read_database_records(
+            database_instance,
+            Run.sql_read_all_run_records,
+        )
+
+        if all_run_arguments_info:
+            output = [RunArguments.from_tuple(info) for info in all_run_arguments_info]
+            return output
+        else:
+            return []
+
+    def cancel(self) -> None:
+        """
+        Sets the cancel event to stop the run process.
+
+        This method is used to signal that the run process should be cancelled. It sets the cancel_event
+        which can be checked in various points of the asynchronous run process to gracefully stop the execution.
+
+        Returns:
+            None
+        """
+        print("[Run] Cancelling run...")
+        self.cancel_event.set()
+
+    async def run(self) -> ResultArguments | None:
+        """
+        Executes the run process asynchronously.
+
+        This method is the main entry point for running the process. It performs the run operation
+        asynchronously and returns a ResultArguments object if the run completes successfully, or None
+        if the run is cancelled or fails to complete.
+
+        Returns:
+            ResultArguments | None: The result of the run operation if successful, otherwise None.
+
+        Raises:
+            RuntimeError: If any error occurs during the run process.
+        """
+        # ------------------------------------------------------------------------------
+        # Part 0: Initialise
+        # ------------------------------------------------------------------------------
+        print("[Run] Part 0: Initialising run...")
+        start_time = time.perf_counter()
+        try:
+            # Initialise the run
+            self.run_arguments.start_time = time.time()
+            self.run_arguments.end_time = time.time()
+
+            # Create a new run record in database
+            if self.run_arguments.database_instance:
+                inserted_record = Storage.create_database_record(
+                    self.run_arguments.database_instance,
+                    self.run_arguments.to_create_tuple(),
+                    Run.sql_create_run_record,
+                )
+                if inserted_record:
+                    self.run_arguments.run_id = inserted_record[0]
+                else:
+                    raise RuntimeError(
+                        "[Run] Failed to create record: record not inserted."
+                    )
+            else:
+                raise RuntimeError(
+                    "[Run] Failed to create record: db_instance is not initialised."
+                )
+
+            # Set status to running
+            self.run_progress.notify_progress(status=RunStatus.RUNNING)
+
+        except Exception as e:
+            error_message = (
+                f"[Run] Failed to initialise run in Part 0 due to error: {str(e)}"
+            )
+            self.run_progress.notify_error(error_message)
+
+        finally:
+            print(
+                f"[Run] Initialise run took {(time.perf_counter() - start_time):.4f}s"
+            )
+
+        # ------------------------------------------------------------------------------
+        # Part 1: Get asyncio running loop
+        # ------------------------------------------------------------------------------
+        print("[Run] Part 1: Loading asyncio running loop...")
+        loop = asyncio.get_running_loop()
+
+        # ------------------------------------------------------------------------------
+        # Part 2: Load runner and result processing module
+        # ------------------------------------------------------------------------------
+        print("[Run] Part 2: Loading modules...")
+        start_time = time.perf_counter()
+        runner_module_instance = None
+        result_module_instance = None
+        try:
+            runner_module_instance = self._load_module(
+                "runner_processing_module", EnvVariables.RUNNERS_MODULES.name
+            )
+            result_module_instance = self._load_module(
+                "result_processing_module", EnvVariables.RESULTS_MODULES.name
+            )
+        except Exception as e:
+            self.run_progress.notify_error(f"[Run] Module loading error: {e}")
+        finally:
+            print(
+                f"[Run] Module loading took {(time.perf_counter() - start_time):.4f}s"
+            )
+
+        # ------------------------------------------------------------------------------
+        # Part 3: Run runner processing module
+        # ------------------------------------------------------------------------------
+        print("[Run] Part 3: Running runner processing module...")
+        start_time = time.perf_counter()
+        runner_results = None
+        try:
+            if runner_module_instance:
+                runner_results = await runner_module_instance.generate(  # type: ignore ; ducktyping
+                    loop,
+                    self.run_arguments.runner_args,
+                    self.run_arguments.database_instance,
+                    self.run_arguments.endpoints,
+                    self.run_progress,
+                    self.cancel_event,
+                )
+            else:
+                raise RuntimeError("Failed to initialise runner module instance.")
+
+        except Exception as e:
+            error_message = f"[Run] Failed to run runner processing module in Part 3 due to error: {str(e)}"
+            self.run_progress.notify_error(error_message)
+
+        finally:
+            print(
+                f"[Run] Running runner processing module took {(time.perf_counter() - start_time):.4f}s"
+            )
+
+        # ------------------------------------------------------------------------------
+        # Part 4: Run result processing module
+        # ------------------------------------------------------------------------------
+        print("[Run] Part 4: Running result processing module...")
+        start_time = time.perf_counter()
+        updated_runner_results = None
+        try:
+            if result_module_instance:
+                updated_runner_results = result_module_instance.generate(  # type: ignore ; ducktyping
+                    runner_results
+                )
+                if updated_runner_results:
+                    self.run_progress.notify_progress(
+                        results=updated_runner_results.results
+                    )
+            else:
+                raise RuntimeError("Failed to initialise result module instance.")
+
+        except Exception as e:
+            error_message = f"[Run] Failed to run result processing module in Part 4 due to error: {str(e)}"
+            self.run_progress.notify_error(error_message)
+
+        finally:
+            print(
+                f"[Run] Running result processing module took {(time.perf_counter() - start_time):.4f}s"
+            )
+
+        # ------------------------------------------------------------------------------
+        # Part 5: Wrap up run
+        # ------------------------------------------------------------------------------
+        print("[Run] Part 5: Wrap up run...")
+        return updated_runner_results
+
+    def _load_module(self, arg_key: str, env_var: str):
+        """
+        Load a module based on the argument key and environment variable.
+
+        This method retrieves the module name from the runner arguments using the provided
+        argument key. It then attempts to load the module using the name and the file path
+        obtained from the environment variable. If the module name is not provided or the
+        module instance cannot be created, a RuntimeError is raised.
+
+        Args:
+            arg_key (str): The key to look up the module name in the runner arguments.
+            env_var (str): The environment variable used to obtain the file path of the module.
+
+        Returns:
+            An instance of the loaded module.
+
+        Raises:
+            RuntimeError: If the module name is not provided or the module instance cannot be created.
+        """
+        module_name = self.run_arguments.runner_args.get(arg_key)
+        if not module_name:
+            raise RuntimeError(f"[Run] Module name for '{arg_key}' not provided.")
+        module_instance = get_instance(
+            module_name,
+            Storage.get_filepath(env_var, module_name, "py"),
+        )
+        if not module_instance:
+            raise RuntimeError(
+                f"[Run] Unable to get instance for module '{module_name}'."
+            )
+        return module_instance()

moonshot/src/runs/run_arguments.py

@@ -0,0 +1,162 @@
+from __future__ import annotations
+
+import ast
+
+from pydantic import BaseModel
+
+from moonshot.src.runners.runner_type import RunnerType
+from moonshot.src.runs.run_status import RunStatus
+from moonshot.src.storage.db_interface import DBInterface
+
+
+class RunArguments(BaseModel):
+    class Config:
+        arbitrary_types_allowed = True
+
+    run_id: int = 0  # The id of the run
+
+    runner_id: str  # The id of the runner
+
+    runner_type: RunnerType  # Run type for the Run.
+
+    runner_args: dict  # Dictionary containing arguments for the runner.
+
+    database_instance: DBInterface | None  # Database instance for the Run.
+
+    endpoints: list[str]  # List of endpoints for the Run.
+
+    results_file: str  # The results file associated with the Run.
+
+    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.
+
+    error_messages: list[str]  # The error messages associated with the Run.
+
+    raw_results: dict  # Results of the Run by runners-module.
+
+    results: dict  # Generated Results of the Run by results-module.
+
+    status: RunStatus  # Status of the Run.
+
+    def to_dict(self) -> dict:
+        """
+        Converts the RunArguments object into a dictionary format.
+
+        This method transforms the RunArguments instance into a dictionary, encapsulating all the critical attributes
+        associated with the run. The resulting dictionary includes keys for runner_id, runner_type, runner_args,
+        database_instance, endpoints, results_file, start_time, end_time, duration, error_messages, raw_results,
+        results, and status offering a comprehensive snapshot of the run's parameters for straightforward access
+        and further processing.
+
+        Returns:
+            dict: A dictionary containing all the significant attributes of the RunArguments instance.
+        """
+        return {
+            "run_id": self.run_id,
+            "runner_id": self.runner_id,
+            "runner_type": self.runner_type,
+            "runner_args": self.runner_args,
+            "endpoints": self.endpoints,
+            "results_file": self.results_file,
+            "start_time": self.start_time,
+            "end_time": self.end_time,
+            "duration": self.duration,
+            "error_messages": self.error_messages,
+            "raw_results": self.raw_results,
+            "results": self.results,
+            "status": self.status,
+        }
+
+    def to_create_tuple(self) -> tuple:
+        """
+        Creates a tuple of run arguments for database insertion.
+
+        This method prepares a tuple of run arguments that can be used to insert a new record into the run_table
+        in the database. The tuple includes the runner_id, runner_type in lowercase, string representation of
+        runner_args, string representation of endpoints, results_file, start_time, end_time, duration,
+        string representation of error_messages, string representation of raw_results, string representation of results,
+        and the run status in lowercase.
+
+        Returns:
+            tuple: A tuple containing the run arguments ready for database insertion.
+        """
+        return (
+            self.runner_id,
+            self.runner_type.name.lower(),
+            str(self.runner_args),
+            str(self.endpoints),
+            self.results_file,
+            self.start_time,
+            self.end_time,
+            self.duration,
+            str(self.error_messages),
+            str(self.raw_results),
+            str(self.results),
+            self.status.name.lower(),
+        )
+
+    def to_tuple(self) -> tuple:
+        """
+        Serializes the RunArguments object to a tuple format.
+
+        This method serializes the RunArguments instance to a tuple, encapsulating the primary attributes of the run.
+        The resulting tuple contains the runner ID, runner type in lowercase, string representation of runner arguments,
+        string representation of endpoints, results file path, start time, end time, run duration, string representation
+        of error messages, string representation of raw results, string representation of results,
+        and the run status in lowercase. This provides a complete summary of the run's parameters for easy storage
+        or transmission.
+
+        Returns:
+            tuple: A tuple containing the serialized attributes of the RunArguments instance.
+        """
+        return (
+            self.runner_id,
+            self.runner_type.name.lower(),
+            str(self.runner_args),
+            str(self.endpoints),
+            self.results_file,
+            self.start_time,
+            self.end_time,
+            self.duration,
+            str(self.error_messages),
+            str(self.raw_results),
+            str(self.results),
+            self.status.name.lower(),
+            self.run_id,
+        )
+
+    @classmethod
+    def from_tuple(cls, run_record: tuple) -> RunArguments:
+        """
+        Reconstructs a RunArguments object from a serialized tuple.
+
+        This class method takes a tuple that contains the serialized form of a RunArguments object and reconstructs it
+        into a new RunArguments instance. The tuple is expected to contain the following elements in order: runner ID,
+        runner type, runner arguments, endpoints, results file, start time, end time, duration, error messages, results,
+        and status. These elements collectively represent the state of a run at a specific point in time.
+
+        Args:
+            run_record (tuple): A tuple containing the serialized state of a RunArguments object.
+
+        Returns:
+            RunArguments: An instance of RunArguments initialized with the data extracted from the tuple.
+        """
+        return cls(
+            run_id=run_record[0],
+            runner_id=run_record[1],
+            runner_type=RunnerType(run_record[2]),
+            runner_args=ast.literal_eval(run_record[3]),
+            database_instance=None,
+            endpoints=ast.literal_eval(run_record[4]),
+            results_file=run_record[5],
+            start_time=float(run_record[6]),
+            end_time=float(run_record[7]),
+            duration=run_record[8],
+            error_messages=ast.literal_eval(run_record[9]),
+            raw_results=ast.literal_eval(run_record[10]),
+            results=ast.literal_eval(run_record[11]),
+            status=RunStatus(run_record[12]),
+        )

moonshot/src/runs/run_progress.py

@@ -0,0 +1,145 @@
+import time
+from typing import Callable
+
+from moonshot.src.runs.run_arguments import RunArguments
+from moonshot.src.runs.run_status import RunStatus
+from moonshot.src.storage.storage import Storage
+
+
+class RunProgress:
+    sql_update_run_record = """
+        UPDATE run_table SET runner_id=?,runner_type=?,runner_args=?,endpoints=?,results_file=?,start_time=?,end_time=?,
+        duration=?,error_messages=?,raw_results=?,results=?,status=?
+        WHERE run_id=?
+    """
+
+    def __init__(
+        self, run_arguments: RunArguments, run_progress_callback_func: Callable | None
+    ):
+        # Information on the run and callback for progress updating
+        self.run_arguments = run_arguments
+        self.run_progress_callback_func = run_progress_callback_func
+
+        # Information to be sent back through callback
+        self.cookbook_index: int = -1
+        self.cookbook_name: str = ""
+        self.cookbook_total: int = -1
+        self.recipe_index: int = -1
+        self.recipe_name: str = ""
+        self.recipe_total: int = -1
+        self.progress: int = 0
+
+    def notify_error(self, error_message: str) -> None:
+        """
+        Notifies about an error that occurred during the run process.
+
+        This method logs the error message, appends it to the run arguments' error messages list,
+        and updates the run progress status to indicate that the run is continuing with errors.
+
+        Args:
+            error_message (str): The error message to be logged and recorded.
+        """
+        # Update error message
+        print(error_message)
+        self.run_arguments.error_messages.append(error_message)
+
+        # Update progress status
+        self.notify_progress(status=RunStatus.RUNNING_WITH_ERRORS)
+
+    def notify_progress(self, **kwargs) -> None:
+        """
+        Updates the run progress information and the run status.
+
+        This method is responsible for updating the run status, setting the end time, calculating the run duration,
+        and persisting these changes to the database. Additionally, if a callback function for run progress has been
+        set, this method will invoke it with the current state of run arguments.
+
+        The method accepts arbitrary keyword arguments which are used to update specific attributes of the run_arguments
+        object. It ensures that the run progress is accurately reflected in both the object's state and the
+        corresponding database record.
+
+        Args:
+            **kwargs: Keyword arguments that correspond to the attributes of run_arguments which should be updated.
+
+        Returns:
+            None
+        """
+        # Update run arguments values
+        if self.run_arguments.start_time > 0.0:
+            self.run_arguments.end_time = time.time()
+            self.run_arguments.duration = int(
+                self.run_arguments.end_time - self.run_arguments.start_time
+            )
+
+        # Update run arguments values with provided key-value pairs
+        for key, value in kwargs.items():
+            # Update self values
+            if hasattr(self, key):
+                setattr(self, key, value)
+
+            # Update self.run_arguments
+            if hasattr(self.run_arguments, key):
+                setattr(self.run_arguments, key, value)
+
+        # Calculate percentage
+        if self.cookbook_total > 0:
+            if self.recipe_total > 0 and self.cookbook_index != self.cookbook_total:
+                # Calculate percentage with cookbook and recipe defined
+                per_recipe_percentage = (100 / self.cookbook_total) / self.recipe_total
+                self.progress = int(
+                    self.cookbook_index * (100 / self.cookbook_total)
+                ) + int(self.recipe_index * per_recipe_percentage)
+            else:
+                # Calculate percentage with cookbook defined and no recipes defined
+                # Or cookbook index and total is same.
+                self.progress = int(self.cookbook_index * (100 / self.cookbook_total))
+        elif self.recipe_total > 0:
+            # There is no cookbook, calculate for recipes defined
+            self.progress = int(self.recipe_index * (100 / self.recipe_total))
+        else:
+            # Initialization: set 0
+            self.progress = 0
+
+        # Update database record
+        if self.run_arguments.database_instance:
+            Storage.update_database_record(
+                self.run_arguments.database_instance,
+                self.run_arguments.to_tuple(),
+                RunProgress.sql_update_run_record,
+            )
+        else:
+            print(
+                "[RunProgress] Failed to update run progress: db_instance is not initialised."
+            )
+
+        # If a callback function is provided, call it with the updated run arguments
+        if self.run_progress_callback_func:
+            self.run_progress_callback_func(self.get_dict())
+
+    def get_dict(self) -> dict:
+        """
+        Constructs and returns a dictionary with the current state of the benchmark execution.
+
+        This method assembles a dictionary encapsulating the current progress of the benchmark execution.
+        The resulting dictionary is composed of the following keys: execution identifier, name, type, current duration,
+        status, cookbook index, cookbook name, cookbook total, recipe index, recipe name, recipe total,
+        progress percentage, and a list of error messages.
+
+        Returns:
+            dict: A dictionary representing the current benchmark execution state, with keys for various progress
+            metrics and details.
+        """
+        return {
+            "current_runner_id": self.run_arguments.runner_id,
+            "current_runner_type": self.run_arguments.runner_type.value,
+            "current_duration": self.run_arguments.duration,
+            "current_status": self.run_arguments.status.value,
+            "current_cookbook_index": self.cookbook_index,
+            "current_cookbook_name": self.cookbook_name,
+            "current_cookbook_total": self.cookbook_total,
+            "current_recipe_index": self.recipe_index,
+            "current_recipe_name": self.recipe_name,
+            "current_recipe_total": self.recipe_total,
+            "current_progress": self.progress,
+            "current_error_messages": list(set(self.run_arguments.error_messages)),
+        }

moonshot/src/runs/run_status.py

@@ -0,0 +1,10 @@
+from enum import Enum
+
+
+class RunStatus(Enum):
+    PENDING = "pending"
+    CANCELLED = "cancelled"
+    RUNNING = "running"
+    RUNNING_WITH_ERRORS = "running_with_errors"
+    COMPLETED = "completed"
+    COMPLETED_WITH_ERRORS = "completed_with_errors"