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

moonshot/src/recipes/recipe.py

@@ -9,6 +9,10 @@ from moonshot.src.configs.env_variables import EnvVariables
 from moonshot.src.datasets.dataset import Dataset
 from moonshot.src.recipes.recipe_arguments import RecipeArguments
 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 Recipe:
@@ -21,7 +25,6 @@ class Recipe:
         self.datasets = rec_args.datasets
         self.prompt_templates = rec_args.prompt_templates
         self.metrics = rec_args.metrics
-        self.attack_modules = rec_args.attack_modules
         self.grading_scale = rec_args.grading_scale
         self.stats = rec_args.stats
 
@@ -64,12 +67,7 @@ class Recipe:
         """
         try:
             rec_id = slugify(rec_args.name, lowercase=True)
-            # check if the recipe exists
-            if Storage.is_object_exists(EnvVariables.RECIPES.name, rec_id, "json"):
-                raise RuntimeError(f"Recipe with ID '{rec_id}' already exists.")
-
             rec_info = {
-                "id": rec_id,
                 "name": rec_args.name,
                 "description": rec_args.description,
                 "tags": rec_args.tags,
@@ -77,10 +75,13 @@ class Recipe:
                 "datasets": rec_args.datasets,
                 "prompt_templates": rec_args.prompt_templates,
                 "metrics": rec_args.metrics,
-                "attack_modules": rec_args.attack_modules,
                 "grading_scale": rec_args.grading_scale,
             }
 
+            # check if the recipe exists
+            if Storage.is_object_exists(EnvVariables.RECIPES.name, rec_id, "json"):
+                raise RuntimeError(f"Recipe with ID '{rec_id}' already exists.")
+
             Recipe.check_file_exists(
                 EnvVariables.PROMPT_TEMPLATES.name,
                 rec_args.prompt_templates,
@@ -93,19 +94,13 @@ class Recipe:
             Recipe.check_file_exists(
                 EnvVariables.METRICS.name, rec_args.metrics, "Metric", "py"
             )
-            Recipe.check_file_exists(
-                EnvVariables.ATTACK_MODULES.name,
-                rec_args.attack_modules,
-                "Attack Module",
-                "py",
-            )
 
             # Write as json output
             Storage.create_object(EnvVariables.RECIPES.name, rec_id, rec_info, "json")
             return rec_id
 
         except Exception as e:
-            print(f"Failed to create recipe: {str(e)}")
+            logger.error(f"Failed to create recipe: {str(e)}")
             raise e
 
     @staticmethod
@@ -128,13 +123,17 @@ class Recipe:
             Exception: If there is an issue reading the file or during any other part of the process.
         """
         try:
-            if rec_id:
-                return RecipeArguments(**Recipe._read_recipe(rec_id, {}))
-            else:
-                raise RuntimeError("Recipe ID is empty")
+            if not rec_id:
+                raise RuntimeError("Recipe ID is empty.")
+
+            recipe_details = Recipe._read_recipe(rec_id, {})
+            if not recipe_details:
+                raise RuntimeError(f"Recipe with ID '{rec_id}' does not exist.")
+
+            return RecipeArguments(**recipe_details)
 
         except Exception as e:
-            print(f"Failed to read recipe: {str(e)}")
+            logger.error(f"Failed to read recipe: {str(e)}")
             raise e
 
     @staticmethod
@@ -171,9 +170,12 @@ class Recipe:
         Raises:
             RuntimeError: If the recipe file cannot be read or does not exist.
         """
-        obj_results = Storage.read_object(EnvVariables.RECIPES.name, rec_id, "json")
-        if not obj_results:
+        obj_results = {"id": rec_id}
+        recipe_info = Storage.read_object(EnvVariables.RECIPES.name, rec_id, "json")
+        if not recipe_info:
             raise RuntimeError(f"Unable to get results for {rec_id}.")
+        else:
+            obj_results.update(recipe_info)
 
         # Calculate statistics for the recipe and update the results dictionary with them
         stats = {
@@ -181,7 +183,6 @@ class Recipe:
             "num_of_datasets": len(obj_results["datasets"]),
             "num_of_prompt_templates": len(obj_results["prompt_templates"]),
             "num_of_metrics": len(obj_results["metrics"]),
-            "num_of_attack_modules": len(obj_results["attack_modules"]),
             "num_of_datasets_prompts": {},
         }
 
@@ -203,25 +204,28 @@ class Recipe:
     @staticmethod
     def update(rec_args: RecipeArguments) -> bool:
         """
-        Updates the recipe information based on the provided RecipeArguments.
+        Updates the recipe information in the storage based on the provided RecipeArguments object.
 
-        This method takes RecipeArguments, converts it to a dictionary, and writes the updated
-        recipe information to the storage. If the operation is successful, it returns True.
-        If an exception occurs, it prints an error message and re-raises the exception.
+        This method processes an instance of RecipeArguments, transforming it into a dictionary format. It ensures
+        that certain required files exist before proceeding to update the recipe information in the storage system.
+        Upon successful update, it returns True. If any error is encountered during the process, the error message is
+        logged, and the exception is propagated upwards.
 
         Args:
-            rec_args (RecipeArguments): The recipe arguments containing updated values.
+            rec_args (RecipeArguments): An object containing the updated recipe information.
 
         Returns:
-            bool: True if the recipe was successfully updated.
+            bool: Indicates whether the recipe update was successful (True) or not (False).
 
         Raises:
-            Exception: If an error occurs during the update process.
+            Exception: Propagates any exceptions that occur during the update process.
         """
         try:
-            # Convert the recipe arguments to a dictionary
+            # Transform RecipeArguments into a dictionary, excluding the 'id' field
             rec_info = rec_args.to_dict()
+            rec_info.pop("id", None)  # The 'id' is derived and should not be written
 
+            # Verify existence of related files in storage
             Recipe.check_file_exists(
                 EnvVariables.PROMPT_TEMPLATES.name,
                 rec_args.prompt_templates,
@@ -234,21 +238,15 @@ class Recipe:
             Recipe.check_file_exists(
                 EnvVariables.METRICS.name, rec_args.metrics, "Metric", "py"
             )
-            Recipe.check_file_exists(
-                EnvVariables.ATTACK_MODULES.name,
-                rec_args.attack_modules,
-                "Attack Module",
-                "py",
-            )
 
-            # Write the updated recipe information to the file
+            # Persist the updated recipe information to storage
             Storage.create_object(
                 EnvVariables.RECIPES.name, rec_args.id, rec_info, "json"
             )
             return True
 
         except Exception as e:
-            print(f"Failed to update recipe: {str(e)}")
+            logger.error(f"Failed to update recipe: {str(e)}")
             raise e
 
     @staticmethod
@@ -275,7 +273,7 @@ class Recipe:
             return True
 
         except Exception as e:
-            print(f"Failed to delete recipe: {str(e)}")
+            logger.error(f"Failed to delete recipe: {str(e)}")
             raise e
 
     @staticmethod
@@ -314,7 +312,7 @@ class Recipe:
             return retn_recs_ids, retn_recs
 
         except Exception as e:
-            print(f"Failed to get available recipes: {str(e)}")
+            logger.error(f"Failed to get available recipes: {str(e)}")
             raise e
 
     @staticmethod

moonshot/src/recipes/recipe_arguments.py

@@ -27,9 +27,6 @@ class RecipeArguments(BaseModel):
     # metrics (list): The list of metrics in the recipe.
     metrics: list[str] = Field(min_length=1)
 
-    # attack_modules (list): The list of attack modules in the recipe.
-    attack_modules: list[str]
-
     # grading_scale (dict): A dictionary where keys are grading categories and values are lists of grading scale.
     grading_scale: dict[str, list[int]]
 
@@ -105,7 +102,6 @@ class RecipeArguments(BaseModel):
             "datasets": self.datasets,
             "prompt_templates": self.prompt_templates,
             "metrics": self.metrics,
-            "attack_modules": self.attack_modules,
             "grading_scale": self.grading_scale,
             "stats": self.stats,
         }

moonshot/src/redteaming/attack/attack_module.py

@@ -18,6 +18,10 @@ from moonshot.src.redteaming.attack.context_strategy import ContextStrategy
 from moonshot.src.runs.run_status import RunStatus
 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 AttackModule:
@@ -173,7 +177,7 @@ class AttackModule:
             for generator in generator_list:
                 async for result in generator:
                     if self.cancel_event.is_set():
-                        print(
+                        logger.warning(
                             "[Red Teaming] Cancellation flag is set. Cancelling task..."
                         )
                         break
@@ -197,7 +201,9 @@ class AttackModule:
         for prepared_prompt in list_of_prompts:
             for target_llm_connector in self.connector_instances:
                 if self.cancel_event.is_set():
-                    print("[Red Teaming] Cancellation flag is set. Cancelling task...")
+                    logger.warning(
+                        "[Red Teaming] Cancellation flag is set. Cancelling task..."
+                    )
                     break
 
                 if self.red_teaming_progress:
@@ -256,7 +262,9 @@ class AttackModule:
         consolidated_responses = []
         for prepared_prompt in list_of_prompts:
             if self.cancel_event.is_set():
-                print("[Red Teaming] Cancellation flag is set. Cancelling task...")
+                logger.warning(
+                    "[Red Teaming] Cancellation flag is set. Cancelling task..."
+                )
                 break
 
             if self.red_teaming_progress:
@@ -344,7 +352,9 @@ class AttackModule:
         """
         async for prompt_info in gen_prompts_generator:
             if self.cancel_event.is_set():
-                print("[Red Teaming] Cancellation flag is set. Cancelling task...")
+                logger.warning(
+                    "[Red Teaming] Cancellation flag is set. Cancelling task..."
+                )
                 break
             new_prompt_info = RedTeamingPromptArguments(
                 conn_id=prompt_info.conn_id,
@@ -415,7 +425,7 @@ class AttackModule:
             )
             return cache_info if cache_info else {}
         except Exception as e:
-            print(f"No previous cache information: {str(e)}")
+            logger.error(f"No previous cache information: {str(e)}")
             return {}
 
     @staticmethod
@@ -434,7 +444,7 @@ class AttackModule:
                 obj_extension=AttackModule.cache_extension,
             )
         except Exception as e:
-            print(f"Failed to write cache information: {str(e)}")
+            logger.error(f"Failed to write cache information: {str(e)}")
             raise e
 
     @staticmethod
@@ -479,7 +489,7 @@ class AttackModule:
             return retn_am_ids, retn_ams
 
         except Exception as e:
-            print(f"Failed to get available attack modules: {str(e)}")
+            logger.error(f"Failed to get available attack modules: {str(e)}")
             raise e
 
     @staticmethod
@@ -542,7 +552,7 @@ class AttackModule:
             return True
 
         except Exception as e:
-            print(f"Failed to delete attack module: {str(e)}")
+            logger.error(f"Failed to delete attack module: {str(e)}")
             raise e
 
 

moonshot/src/redteaming/attack/context_strategy.py

@@ -6,6 +6,10 @@ from moonshot.src.configs.env_variables import EnvVariables
 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 ContextStrategy:
@@ -84,7 +88,7 @@ class ContextStrategy:
             return True
 
         except Exception as e:
-            print(f"Failed to delete context strategy: {str(e)}")
+            logger.error(f"Failed to delete context strategy: {str(e)}")
             raise e
 
     @staticmethod
@@ -125,7 +129,7 @@ class ContextStrategy:
             context_strategy_instance = context_strategy_instance(context_strategy_name)
             return context_strategy_instance.add_in_context(user_prompt, list_of_chats)
         else:
-            print(
+            logger.error(
                 "Cannot load context strategy. Make sure the name of the context strategy is correct."
             )
             return ""

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