aiverify-moonshot 0.4.1__py3-none-any.whl → 0.4.2__py3-none-any.whl

moonshot/src/redteaming/session/session.py

@@ -17,6 +17,10 @@ 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
+from moonshot.src.utils.log import configure_logger
+
+# Create a logger for this module
+logger = configure_logger(__name__)
 
 
 class SessionMetadata:
@@ -283,13 +287,13 @@ class Session:
 
             # check if the session metadata record already exists
             if session_metadata_records:
-                print("[Session] Session already exists.")
+                logger.info("[Session] Session already exists.")
                 self.session_metadata = SessionMetadata.from_tuple(
                     session_metadata_records[0]
                 )
             # create a new record if session metadata does not exist
             else:
-                print("[Session] Creating new session.")
+                logger.info("[Session] Creating new session.")
 
                 # create chat history table for each endpoint
 
@@ -379,13 +383,13 @@ class Session:
         # ------------------------------------------------------------------------------
         # Part 1: Get asyncio running loop
         # ------------------------------------------------------------------------------
-        print("[Session] Part 1: Loading asyncio running loop...")
+        logger.debug("[Session] Part 1: Loading asyncio running loop...")
         loop = asyncio.get_running_loop()
 
         # ------------------------------------------------------------------------------
         # Part 2: Load runner processing module
         # ------------------------------------------------------------------------------
-        print("[Session] Part 2: Loading runner processing module...")
+        logger.debug("[Session] Part 2: Loading runner processing module...")
         start_time = time.perf_counter()
         runner_module_instance = None
         try:
@@ -414,20 +418,20 @@ class Session:
                 )
 
         except Exception as e:
-            print(
+            logger.error(
                 f"[Session] Failed to load runner processing module in Part 2 due to error: {str(e)}"
             )
             raise e
 
         finally:
-            print(
+            logger.debug(
                 f"[Session] Loading runner processing module took {(time.perf_counter() - start_time):.4f}s"
             )
 
         # ------------------------------------------------------------------------------
         # Part 3: Run runner processing module
         # ------------------------------------------------------------------------------
-        print("[Session] Part 3: Running runner processing module...")
+        logger.debug("[Session] Part 3: Running runner processing module...")
         start_time = time.perf_counter()
         runner_results = {}
 
@@ -446,7 +450,7 @@ class Session:
                 raise RuntimeError("Failed to initialise runner module instance.")
 
         except Exception as e:
-            print(
+            logger.error(
                 f"[Session] Failed to run runner processing module in Part 3 due to error: {str(e)}"
             )
             raise e
@@ -455,14 +459,14 @@ class Session:
             self.red_teaming_progress.status = RunStatus.COMPLETED
             if self.check_redteaming_type() == RedTeamingType.AUTOMATED:
                 self.red_teaming_progress.notify_progress()
-            print(
+            logger.debug(
                 f"[Session] Running runner processing module took {(time.perf_counter() - start_time):.4f}s"
             )
 
         # ------------------------------------------------------------------------------
         # Part 4: Wrap up run
         # ------------------------------------------------------------------------------
-        print("[Session] Part 4: Wrap up run...")
+        logger.debug("[Session] Part 4: Wrap up run...")
         return runner_results
 
     def cancel(self) -> None:
@@ -476,7 +480,7 @@ class Session:
         Returns:
             None
         """
-        print("[Session] Cancelling automated red teaming...")
+        logger.warning("[Session] Cancelling automated red teaming...")
         self.cancel_event.set()
 
     def check_redteaming_type(self) -> RedTeamingType:

moonshot/src/results/result.py

@@ -6,6 +6,10 @@ from pydantic import validate_call
 
 from moonshot.src.configs.env_variables import EnvVariables
 from moonshot.src.storage.storage import Storage
+from moonshot.src.utils.log import configure_logger
+
+# Create a logger for this module
+logger = configure_logger(__name__)
 
 
 class Result:
@@ -35,7 +39,7 @@ class Result:
                 raise RuntimeError("Result ID is empty")
 
         except Exception as e:
-            print(f"Failed to read result: {str(e)}")
+            logger.error(f"Failed to read result: {str(e)}")
             raise e
 
     @staticmethod
@@ -85,7 +89,7 @@ class Result:
             return True
 
         except Exception as e:
-            print(f"Failed to delete result: {str(e)}")
+            logger.error(f"Failed to delete result: {str(e)}")
             raise e
 
     @staticmethod
@@ -115,5 +119,5 @@ class Result:
             return retn_results_ids, retn_results
 
         except Exception as e:
-            print(f"Failed to get available results: {str(e)}")
+            logger.error(f"Failed to get available results: {str(e)}")
             raise e

moonshot/src/runners/runner.py

@@ -13,6 +13,10 @@ 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
+from moonshot.src.utils.log import configure_logger
+
+# Create a logger for this module
+logger = configure_logger(__name__)
 
 
 class Runner:
@@ -82,13 +86,13 @@ class Runner:
                 )
             runner_args = Runner.read(runner_id)
             runner_args.database_instance = Storage.create_database_connection(
-                EnvVariables.DATABASES.name, runner_id, "db"
+                EnvVariables.DATABASES.name, Path(runner_args.database_file).stem, "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)}")
+            logger.error(f"[Runner] Failed to load runner: {str(e)}")
             raise e
 
     @classmethod
@@ -114,6 +118,14 @@ class Runner:
         """
         try:
             runner_id = slugify(runner_args.name, lowercase=True)
+            runner_info = {
+                "name": runner_args.name,
+                "database_file": Storage.get_filepath(
+                    EnvVariables.DATABASES.name, runner_id, "db", True
+                ),
+                "endpoints": runner_args.endpoints,
+                "description": runner_args.description,
+            }
 
             # Check if runner file exists. If it exists, raise an error.
             if Storage.is_object_exists(EnvVariables.RUNNERS.name, runner_id, "json"):
@@ -129,35 +141,27 @@ class Runner:
                         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"
+                EnvVariables.RUNNERS.name, runner_id, runner_info, "json"
+            )
+
+            # Add additional attributes (id, database instance and update progress_callback_func)
+            runner_info["id"] = runner_id
+            runner_info["database_instance"] = Storage.create_database_connection(
+                EnvVariables.DATABASES.name, runner_id, "db"
             )
+            runner_info["progress_callback_func"] = runner_args.progress_callback_func
 
             # Create runner cache table
             Storage.create_database_table(
-                runner_args.database_instance, Runner.sql_create_runner_cache_table
+                runner_info["database_instance"], Runner.sql_create_runner_cache_table
             )
 
-            return cls(runner_args)
+            return cls(RunnerArguments(**runner_info))
 
         except Exception as e:
-            print(f"[Runner] Failed to create runner: {str(e)}")
+            logger.error(f"[Runner] Failed to create runner: {str(e)}")
             raise e
 
     @staticmethod
@@ -179,17 +183,40 @@ class Runner:
             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")
+            if not runner_id:
+                raise RuntimeError("Runner ID is empty.")
+
+            runner_details = Runner._read_runner(runner_id)
+            if not runner_details:
+                raise RuntimeError(f"Runner with ID '{runner_id}' does not exist.")
+
+            return RunnerArguments(**runner_details)
 
         except Exception as e:
-            print(f"[Runner] Failed to read runner: {str(e)}")
+            logger.error(f"[Runner] Failed to read runner: {str(e)}")
             raise e
 
+    @staticmethod
+    def _read_runner(runner_id: str) -> dict:
+        """
+        Retrieves the runner's information from a JSON file.
+
+        This internal method is designed to fetch the details of a specific runner by its ID. It searches for the
+        corresponding JSON file within the directory specified by `EnvVariables.RUNNERS`. The method returns a
+        dictionary containing the runner's information.
+
+        Args:
+            runner_id (str): The unique identifier of the runner whose information is being retrieved.
+
+        Returns:
+            dict: A dictionary with the runner's information.
+        """
+        runner_info = {"id": runner_id}
+        runner_info.update(
+            Storage.read_object(EnvVariables.RUNNERS.name, runner_id, "json")
+        )
+        return runner_info
+
     @staticmethod
     @validate_call
     def delete(runner_id: str) -> bool:
@@ -216,7 +243,7 @@ class Runner:
             return True
 
         except Exception as e:
-            print(f"[Runner] Failed to delete runner: {str(e)}")
+            logger.error(f"[Runner] Failed to delete runner: {str(e)}")
             raise e
 
     @staticmethod
@@ -245,18 +272,14 @@ class Runner:
                 if "__" in runner:
                     continue
 
-                runner_info = RunnerArguments(
-                    **Storage.read_object(
-                        EnvVariables.RUNNERS.name, Path(runner).stem, "json"
-                    )
-                )
+                runner_info = RunnerArguments(**Runner._read_runner(Path(runner).stem))
                 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)}")
+            logger.error(f"[Runner] Failed to get available runners: {str(e)}")
             raise e
 
     def close(self) -> None:
@@ -285,7 +308,7 @@ class Runner:
         """
         async with self.current_operation_lock:
             if self.current_operation:
-                print(f"[Runner] {self.id} - Cancelling current operation...")
+                logger.warning(f"[Runner] {self.id} - Cancelling current operation...")
                 self.current_operation.cancel()
                 self.current_operation = None  # Reset the current operation
 
@@ -328,7 +351,7 @@ class Runner:
         """
         async with self.current_operation_lock:  # Acquire the lock
             # Create new benchmark recipe test run
-            print(f"[Runner] {self.id} - Running benchmark recipe run...")
+            logger.info(f"[Runner] {self.id} - Running benchmark recipe run...")
             self.current_operation = Run(
                 self.id,
                 RunnerType.BENCHMARK,
@@ -354,7 +377,7 @@ class Runner:
         # 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.")
+            logger.info(f"[Runner] {self.id} - Benchmark recipe run completed.")
 
     async def run_cookbooks(
         self,
@@ -396,7 +419,7 @@ class Runner:
         """
         async with self.current_operation_lock:  # Acquire the lock
             # Create new benchmark cookbook test run
-            print(f"[Runner] {self.id} - Running benchmark cookbook run...")
+            logger.info(f"[Runner] {self.id} - Running benchmark cookbook run...")
             self.current_operation = Run(
                 self.id,
                 RunnerType.BENCHMARK,
@@ -422,7 +445,7 @@ class Runner:
         # 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.")
+            logger.info(f"[Runner] {self.id} - Benchmark cookbook run completed.")
 
     async def run_red_teaming(
         self,
@@ -450,7 +473,7 @@ class Runner:
             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...")
+            logger.info(f"[Runner] {self.id} - Running red teaming session...")
             self.current_operation = Session(
                 self.id,
                 RunnerType.REDTEAM,
@@ -471,6 +494,6 @@ class Runner:
         # 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.")
+            logger.info(f"[Runner] {self.id} - Red teaming run completed.")
 
         return red_teaming_results

moonshot/src/runs/run.py

@@ -13,6 +13,10 @@ 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
+from moonshot.src.utils.log import configure_logger
+
+# Create a logger for this module
+logger = configure_logger(__name__)
 
 
 class Run:
@@ -164,7 +168,7 @@ class Run:
         Returns:
             None
         """
-        print("[Run] Cancelling run...")
+        logger.warning("[Run] Cancelling run...")
         self.cancel_event.set()
 
     async def run(self) -> ResultArguments | None:
@@ -184,7 +188,7 @@ class Run:
         # ------------------------------------------------------------------------------
         # Part 0: Initialise
         # ------------------------------------------------------------------------------
-        print("[Run] Part 0: Initialising run...")
+        logger.debug("[Run] Part 0: Initialising run...")
         start_time = time.perf_counter()
         try:
             # Initialise the run
@@ -219,20 +223,20 @@ class Run:
             self.run_progress.notify_error(error_message)
 
         finally:
-            print(
+            logger.debug(
                 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...")
+        logger.debug("[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...")
+        logger.debug("[Run] Part 2: Loading modules...")
         start_time = time.perf_counter()
         runner_module_instance = None
         result_module_instance = None
@@ -246,14 +250,14 @@ class Run:
         except Exception as e:
             self.run_progress.notify_error(f"[Run] Module loading error: {e}")
         finally:
-            print(
+            logger.debug(
                 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...")
+        logger.debug("[Run] Part 3: Running runner processing module...")
         start_time = time.perf_counter()
         runner_results = None
         try:
@@ -274,14 +278,14 @@ class Run:
             self.run_progress.notify_error(error_message)
 
         finally:
-            print(
+            logger.debug(
                 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...")
+        logger.debug("[Run] Part 4: Running result processing module...")
         start_time = time.perf_counter()
         updated_runner_results = None
         try:
@@ -301,14 +305,14 @@ class Run:
             self.run_progress.notify_error(error_message)
 
         finally:
-            print(
+            logger.debug(
                 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...")
+        logger.debug("[Run] Part 5: Wrap up run...")
         return updated_runner_results
 
     def _load_module(self, arg_key: str, env_var: str):

moonshot/src/runs/run_progress.py

@@ -4,6 +4,10 @@ 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
+from moonshot.src.utils.log import configure_logger
+
+# Create a logger for this module
+logger = configure_logger(__name__)
 
 
 class RunProgress:
@@ -40,7 +44,7 @@ class RunProgress:
             error_message (str): The error message to be logged and recorded.
         """
         # Update error message
-        print(error_message)
+        logger.error(error_message)
         self.run_arguments.error_messages.append(error_message)
 
         # Update progress status
@@ -108,8 +112,8 @@ class RunProgress:
                 RunProgress.sql_update_run_record,
             )
         else:
-            print(
-                "[RunProgress] Failed to update run progress: db_instance is not initialised."
+            logger.warning(
+                "[RunProgress] Failed to update run progress: database instance is not initialised."
             )
 
         # If a callback function is provided, call it with the updated run arguments

moonshot/src/storage/db_interface.py

@@ -100,6 +100,20 @@ class DBInterface:
         """
         pass
 
+    @abstractmethod
+    def delete_records_in_table(self, delete_record_sql: str) -> None:
+        """
+        This method is used to delete all records from a specified table in the database. The details of the operation
+        are implementation specific and should be provided by the concrete class that inherits from this abstract class.
+
+        Args:
+            delete_record_sql (str): The SQL query to delete all records from a table.
+
+        Returns:
+            None
+        """
+        pass
+
     @abstractmethod
     def check_database_table_exists(self, table_name: str) -> bool:
         """

moonshot/src/storage/storage.py

@@ -22,7 +22,7 @@ class Storage:
         obj_info: dict,
         obj_extension: str,
         obj_mod_type: str = "jsonio",
-    ) -> bool:
+    ) -> str:
         """
         Writes the object information to a file.
 
@@ -32,6 +32,9 @@ class Storage:
             obj_info (dict): A dictionary containing the object information.
             obj_extension (str): The file extension (e.g., 'json', 'py').
             obj_mod_type (str, optional): The module type for object serialization. Defaults to 'jsonio'.
+
+        Returns:
+            str: A filepath string of the object that has just been created.
         """
         if not hasattr(EnvironmentVars, obj_type):
             raise RuntimeError(
@@ -45,7 +48,11 @@ class Storage:
                 Storage.get_filepath(EnvVariables.IO_MODULES.name, obj_mod_type, "py"),
             )
             if obj_mod_instance and callable(obj_mod_instance):
-                return obj_mod_instance(obj_filepath).create_file(obj_info)
+                try:
+                    obj_mod_instance(obj_filepath).create_file(obj_info)
+                    return obj_filepath
+                except Exception as e:
+                    raise e
             else:
                 raise RuntimeError(
                     f"Unable to get defined object module instance - {obj_mod_instance}"
@@ -479,6 +486,30 @@ class Storage:
         else:
             raise RuntimeError("Database instance is not initialised.")
 
+    @staticmethod
+    def delete_database_record_in_table(
+        database_instance: DBInterface, sql_delete_record: str
+    ) -> None:
+        """
+        Deletes records from a table in the database based on a SQL condition.
+
+        This method is used to delete records from a specific table in the database
+            that meet the condition specified in the SQL delete statement.
+        If the database instance is not initialised, it raises a RuntimeError.
+        Otherwise, it calls the delete_records_in_table method of the database instance with the provided SQL query.
+
+        Args:
+            database_instance (DBInterface): The database accessor instance.
+            sql_delete_record (str): The SQL query to delete records from a table.
+
+        Returns:
+            None
+        """
+        if database_instance:
+            database_instance.delete_records_in_table(sql_delete_record)
+        else:
+            raise RuntimeError("Database instance is not initialised.")
+
     @staticmethod
     def check_database_table_exists(
         database_instance: DBInterface, table_name: str

moonshot/src/utils/find_feature.py

@@ -0,0 +1,45 @@
+from typing import Any, Union
+
+
+def find_keyword(
+    keyword: str, target: Union[str, list[str], list[list[str]], list[dict], list, dict]
+) -> Any:
+    """
+    Find the keyword in the target and return the matching elements.
+
+    Args:
+        keyword (str): The keyword to search for.
+        target: The target to search within.
+
+    Returns:
+        Any: The matching elements or None if no match is found.
+    """
+
+    # find in single string
+    if isinstance(target, str):
+        if target.find(keyword) != -1:
+            return target
+        return ""
+
+    # find in single dict
+    elif isinstance(target, dict):
+        if any(str(keyword) in str(value).lower() for value in target.values()):
+            return target
+        return {}
+
+    # find in list
+    elif isinstance(target, list):
+        # list of string
+        if all(isinstance(item, str) for item in target):
+            return [item for item in target if str(item).lower().find(keyword) != -1]
+
+        # list of dict
+        elif all(isinstance(item, dict) for item in target):
+            return [
+                item
+                for item in target
+                if isinstance(item, dict)
+                and any(str(keyword) in str(value).lower() for value in item.values())
+            ]
+
+    return None