aiverify-moonshot 0.4.8__py3-none-any.whl → 0.4.9__py3-none-any.whl

moonshot/integrations/web_api/services/cookbook_service.py

@@ -5,7 +5,7 @@ from .... import api as moonshot_api
 from ..schemas.cookbook_create_dto import CookbookCreateDTO, CookbookUpdateDTO
 from ..schemas.cookbook_response_model import CookbookResponseModel
 from ..services.base_service import BaseService
-from ..services.recipe_service import get_total_prompt_in_recipe
+from ..services.recipe_service import get_total_prompt_in_recipe, get_endpoint_dependency_in_recipe
 from ..services.utils.exceptions_handler import exception_handler
 
 
@@ -63,24 +63,24 @@ class CookbookService(BaseService):
                 if cookbook not in cookbooks_list:
                     cookbooks_list.append(cookbook)
                     if count:
-                        cookbook.total_prompt_in_cookbook = (
-                            get_total_prompt_in_cookbook(cookbook)
+                        cookbook.total_prompt_in_cookbook, cookbook.total_dataset_in_cookbook = (
+                            get_total_prompt_and_dataset_in_cookbook(cookbook)
                         )
 
             if tags and cookbooks_recipe_has_tags(tags, cookbook):
                 if cookbook not in cookbooks_list:
                     cookbooks_list.append(cookbook)
                     if count:
-                        cookbook.total_prompt_in_cookbook = (
-                            get_total_prompt_in_cookbook(cookbook)
+                        cookbook.total_prompt_in_cookbook, cookbook.total_dataset_in_cookbook = (
+                            get_total_prompt_and_dataset_in_cookbook(cookbook)
                         )
 
             if categories and cookbooks_recipe_has_categories(categories, cookbook):
                 if cookbook not in cookbooks_list:
                     cookbooks_list.append(cookbook)
                     if count:
-                        cookbook.total_prompt_in_cookbook = (
-                            get_total_prompt_in_cookbook(cookbook)
+                        cookbook.total_prompt_in_cookbook, cookbook.total_dataset_in_cookbook = (
+                            get_total_prompt_and_dataset_in_cookbook(cookbook)
                         )
 
             if categories_excluded and cookbooks_recipe_has_categories(
@@ -88,6 +88,9 @@ class CookbookService(BaseService):
             ):
                 cookbooks_list.remove(cookbook)
 
+        for cookbook in cookbooks_list:
+            cookbook.endpoint_required = cookbook_endpoint_dependency(cookbook)
+
         return cookbooks_list
 
     @exception_handler
@@ -131,7 +134,7 @@ class CookbookService(BaseService):
 
 
 @staticmethod
-def get_total_prompt_in_cookbook(cookbook: Cookbook) -> int:
+def get_total_prompt_and_dataset_in_cookbook(cookbook: Cookbook) -> tuple[int, int]:
     """
     Calculate the total number of prompts in a cookbook.
 
@@ -144,7 +147,8 @@ def get_total_prompt_in_cookbook(cookbook: Cookbook) -> int:
         int: The total count of prompts within the cookbook.
     """
     recipes = moonshot_api.api_read_recipes(cookbook.recipes)
-    return sum(get_total_prompt_in_recipe(Recipe(**recipe)) for recipe in recipes)
+    total_prompts, total_datasets = zip(*(get_total_prompt_in_recipe(Recipe(**recipe)) for recipe in recipes))
+    return sum(total_prompts), sum(total_datasets)
 
 
 @staticmethod
@@ -192,3 +196,26 @@ def cookbooks_recipe_has_categories(categories: str, cookbook: Cookbook) -> bool
         ):
             return True
     return False
+
+@staticmethod
+def cookbook_endpoint_dependency(cookbook: Cookbook) -> list[str] | None:
+    """
+    Retrieve a list of endpoint dependencies for all recipes in a given cookbook.
+
+    Args:
+        cookbook (Cookbook): The cookbook object containing the recipe IDs.
+
+    Returns:
+        list[str] | None: A list of endpoint dependencies if any are found, otherwise None.
+    """
+    recipes_in_cookbook = cookbook.recipes
+    recipes = moonshot_api.api_read_recipes(recipes_in_cookbook)
+    list_of_endpoints = set()
+
+    for recipe in recipes:
+        recipe = Recipe(**recipe)
+        endpoints = get_endpoint_dependency_in_recipe(recipe)
+        if endpoints:
+            list_of_endpoints.update(endpoints)
+
+    return list(list_of_endpoints) if list_of_endpoints else None

moonshot/integrations/web_api/services/dataset_service.py

@@ -1,5 +1,5 @@
 from .... import api as moonshot_api
-from ..schemas.dataset_create_dto import DatasetCreateDTO
+from ..schemas.dataset_create_dto import CSV_Dataset_DTO, HF_Dataset_DTO
 from ..schemas.dataset_response_dto import DatasetResponseDTO
 from ..services.base_service import BaseService
 from ..services.utils.exceptions_handler import exception_handler
@@ -8,24 +8,49 @@ from .utils.file_manager import copy_file
 
 class DatasetService(BaseService):
     @exception_handler
-    def create_dataset(self, dataset_data: DatasetCreateDTO, method: str) -> str:
+    def convert_dataset(self, dataset_data: CSV_Dataset_DTO) -> str:
         """
-        Create a dataset using the specified method.
+        Convert a dataset using the provided dataset data.
 
         Args:
-            dataset_data (DatasetCreateDTO): The data required to create the dataset.
-            method (str): The method to use for creating the dataset.
-                          Supported methods are "hf" and "csv".
+            dataset_data (CSV_Dataset_DTO): The data required to convert the dataset.
+
+        Returns:
+            str: The path to the newly created dataset.
+
+        Raises:
+            Exception: If an error occurs during dataset conversion.
+        """
+
+        new_ds_path = moonshot_api.api_convert_dataset(
+            name=dataset_data.name,
+            description=dataset_data.description,
+            reference=dataset_data.reference,
+            license=dataset_data.license,
+            csv_file_path=dataset_data.csv_file_path,
+        )
+        return copy_file(new_ds_path)
+
+    @exception_handler
+    def download_dataset(self, dataset_data: HF_Dataset_DTO) -> str:
+        """
+        Download a dataset using the provided dataset data.
+
+        Args:
+            dataset_data (HF_Dataset_DTO): The data required to download the dataset.
+
+        Returns:
+            str: The path to the newly downloaded dataset.
 
         Raises:
-            Exception: If an error occurs during dataset creation.
+            Exception: If an error occurs during dataset download.
         """
-        new_ds_path = moonshot_api.api_create_datasets(
+
+        new_ds_path = moonshot_api.api_download_dataset(
             name=dataset_data.name,
             description=dataset_data.description,
             reference=dataset_data.reference,
             license=dataset_data.license,
-            method=method,
             **dataset_data.params,
         )
         return copy_file(new_ds_path)

moonshot/integrations/web_api/services/recipe_service.py

@@ -60,7 +60,7 @@ class RecipeService(BaseService):
         for recipe_dict in recipes:
             recipe = RecipeResponseModel(**recipe_dict)
             if count:
-                recipe.total_prompt_in_recipe = get_total_prompt_in_recipe(recipe)
+                recipe.total_prompt_in_recipe, _ = get_total_prompt_in_recipe(recipe)
             filtered_recipes.append(recipe)
 
         # TODO - do all filtering in 1 pass
@@ -84,6 +84,9 @@ class RecipeService(BaseService):
             if sort_by == "id":
                 filtered_recipes.sort(key=lambda x: x.id)
 
+        for recipe in filtered_recipes:
+            recipe.endpoint_required = get_endpoint_dependency_in_recipe(recipe)
+
         return filtered_recipes
 
     @exception_handler
@@ -126,7 +129,7 @@ class RecipeService(BaseService):
 
 
 @staticmethod
-def get_total_prompt_in_recipe(recipe: Recipe) -> int:
+def get_total_prompt_in_recipe(recipe: Recipe) -> tuple[int, int]:
     """
     Calculate the total number of prompts in a recipe.
 
@@ -139,6 +142,7 @@ def get_total_prompt_in_recipe(recipe: Recipe) -> int:
 
     Returns:
         int: The total count of prompts within the recipe.
+        int: The total count of datasets within the recipe.
     """
     # Initialize total prompt count
     total_prompt_count = 0
@@ -151,4 +155,30 @@ def get_total_prompt_in_recipe(recipe: Recipe) -> int:
     if recipe.prompt_templates:
         total_prompt_count *= len(recipe.prompt_templates)
 
-    return total_prompt_count
+    return total_prompt_count, int(recipe.stats.get("num_of_datasets", 0))
+
+@staticmethod
+def get_endpoint_dependency_in_recipe(recipe: Recipe) -> list[str] | None:
+    """
+    Retrieve the list of endpoint dependencies for a given recipe.
+
+    This function fetches all metrics and their associated endpoints, then
+    matches the metrics in the provided recipe to find and compile a list
+    of endpoint dependencies.
+
+    Args:
+        recipe (Recipe): The recipe object containing the metrics information.
+
+    Returns:
+        list[str] | None: A list of endpoint dependencies if found, otherwise None.
+    """
+    metrics = recipe.metrics
+    all_metrics = moonshot_api.api_get_all_metric()
+    
+    endpoints = set()
+    for metric in metrics:
+        for m in all_metrics:
+            if m['id'] == metric:
+                endpoints.update(m['endpoints'])
+    
+    return list(endpoints) if endpoints else None

moonshot/src/api/api_dataset.py

@@ -47,35 +47,67 @@ def api_get_all_datasets_name() -> list[str]:
     return datasets_name
 
 
-def api_create_datasets(
-    name: str, description: str, reference: str, license: str, method: str, **kwargs
+def api_download_dataset(
+    name: str, description: str, reference: str, license: str, **kwargs
 ) -> str:
     """
-    This function creates a new dataset.
+    Downloads a dataset from Hugging Face and creates a new dataset with the provided details.
 
-    This function takes the name, description, reference, and license for a new dataset as input. It then creates a new
-    DatasetArguments object with these details and an empty id. The id is left empty because it will be generated
-    from the name during the creation process. The function then calls the Dataset's create method to
-    create the new dataset.
+    This function takes the name, description, reference, and license for a new dataset as input, along with additional
+    keyword arguments for downloading the dataset from Hugging Face. It then creates a new DatasetArguments object with
+    these details and an empty id. The id is left empty because it will be generated from the name during the creation
+    process. The function then calls the Dataset's create method to create the new dataset.
 
     Args:
         name (str): The name of the new dataset.
         description (str): A brief description of the new dataset.
         reference (str): A reference link for the new dataset.
         license (str): The license of the new dataset.
-        method (str): The method to create new dataset. (csv/hf)
-        kwargs: Additional keyword arguments for the Dataset's create method.
+        kwargs: Additional keyword arguments for downloading the dataset from Hugging Face.
 
     Returns:
         str: The ID of the newly created dataset.
     """
+    examples = Dataset.download_hf(**kwargs)
     ds_args = DatasetArguments(
         id="",
         name=name,
         description=description,
         reference=reference,
         license=license,
-        examples=None,
+        examples=examples,
     )
+    return Dataset.create(ds_args)
 
-    return Dataset.create(ds_args, method, **kwargs)
+
+def api_convert_dataset(
+    name: str, description: str, reference: str, license: str, csv_file_path: str
+) -> str:
+    """
+    Converts a CSV file to a dataset and creates a new dataset with the provided details.
+
+    This function takes the name, description, reference, and license for a new dataset as input, along with the file
+    path to a CSV file. It then creates a new DatasetArguments object with these details and an empty id. The id is left
+    empty because it will be generated from the name during the creation process. The function then calls the Dataset's
+    create method to create the new dataset.
+
+    Args:
+        name (str): The name of the new dataset.
+        description (str): A brief description of the new dataset.
+        reference (str): A reference link for the new dataset.
+        license (str): The license of the new dataset.
+        csv_file_path (str): The file path to the CSV file.
+
+    Returns:
+        str: The ID of the newly created dataset.
+    """
+    examples = Dataset.convert_data(csv_file_path)
+    ds_args = DatasetArguments(
+        id="",
+        name=name,
+        description=description,
+        reference=reference,
+        license=license,
+        examples=examples,
+    )
+    return Dataset.create(ds_args)

moonshot/src/bookmark/bookmark.py

@@ -11,6 +11,7 @@ from moonshot.src.messages_constants import (
     BOOKMARK_ADD_BOOKMARK_VALIDATION_ERROR,
     BOOKMARK_DELETE_ALL_BOOKMARK_ERROR,
     BOOKMARK_DELETE_ALL_BOOKMARK_SUCCESS,
+    BOOKMARK_DELETE_BOOKMARK_FAIL,
     BOOKMARK_DELETE_BOOKMARK_ERROR,
     BOOKMARK_DELETE_BOOKMARK_ERROR_1,
     BOOKMARK_DELETE_BOOKMARK_SUCCESS,
@@ -209,15 +210,21 @@ class Bookmark:
         """
         if isinstance(bookmark_name, str) and bookmark_name:
             try:
-                sql_delete_bookmark_record = textwrap.dedent(
-                    f"""
-                    DELETE FROM bookmark WHERE name = '{bookmark_name}';
-                """
-                )
-                Storage.delete_database_record_in_table(
-                    self.db_instance, sql_delete_bookmark_record
-                )
-                return {"success": True, "message": BOOKMARK_DELETE_BOOKMARK_SUCCESS}
+
+                bookmark_info = Storage.read_database_record(
+                self.db_instance, (bookmark_name,), Bookmark.sql_select_bookmark_record)
+                if bookmark_info is not None:
+                    sql_delete_bookmark_record = textwrap.dedent(
+                        f"""
+                        DELETE FROM bookmark WHERE name = '{bookmark_name}';
+                    """
+                    )
+                    Storage.delete_database_record_in_table(
+                        self.db_instance, sql_delete_bookmark_record
+                    )
+                    return {"success": True, "message": BOOKMARK_DELETE_BOOKMARK_SUCCESS}
+                else:
+                    return {"success": False, "message": BOOKMARK_DELETE_BOOKMARK_FAIL}
             except Exception as e:
                 return {
                     "success": False,

moonshot/src/datasets/dataset.py

@@ -1,6 +1,7 @@
 from __future__ import annotations
 
 from pathlib import Path
+from typing import Iterator
 
 import pandas as pd
 from datasets import load_dataset
@@ -22,34 +23,21 @@ class Dataset:
 
     @staticmethod
     @validate_call
-    def create(ds_args: DatasetArguments, method: str, **kwargs) -> str:
+    def create(ds_args: DatasetArguments) -> str:
         """
         Creates a new dataset based on the provided arguments and method.
 
         This method generates a unique dataset ID using the dataset name,
         checks if a dataset with the same ID already exists, and then
-        creates the dataset using the specified method (either 'csv' or
-        'hf'). The dataset information is then stored as a JSON object.
+        creates the dataset using the specified method. The dataset information
+        is then stored as a JSON object.
 
         Args:
             ds_args (DatasetArguments): The arguments containing dataset
-                details such as name, description, reference, and license.
-            method (str): The method to create the dataset. It can be either
-                'csv' or 'hf'.
-            **kwargs: Additional keyword arguments required for the specified
-                method.
-                - For 'csv' method: 'csv_file_path' (str): The file path to
-                    the CSV file.
-                - For 'hf' method: 'dataset_name' (str): The name of the
-                    Hugging Face dataset.
-                    'dataset_config' (str): The configuration of the Hugging
-                    Face dataset.
-                    'split' (str): The split of the dataset to load.
-                    'input_col' (list[str]): The list of input columns.
-                    'target_col' (str): The target column.
+                details such as name, description, reference, license, and examples.
 
         Returns:
-            str: The unique ID of the created dataset.
+            str: The file path of the created dataset JSON object.
 
         Raises:
             RuntimeError: If a dataset with the same ID already exists.
@@ -63,56 +51,58 @@ class Dataset:
             if Storage.is_object_exists(EnvVariables.DATASETS.name, ds_id, "json"):
                 raise RuntimeError(f"Dataset with ID '{ds_id}' already exists.")
 
-            examples = [{}]
-            if method == "csv":
-                examples = Dataset._convert_csv(kwargs["csv_file_path"])
-            elif method == "hf":
-                examples = Dataset._download_hf(kwargs)
-
             ds_info = {
                 "id": ds_id,
                 "name": ds_args.name,
                 "description": ds_args.description,
                 "reference": ds_args.reference,
                 "license": ds_args.license,
-                "examples": examples,
             }
 
+            examples = ds_args.examples
+
             # Write as JSON output
-            file_path = Storage.create_object(
-                EnvVariables.DATASETS.name, ds_id, ds_info, "json"
+            file_path = Storage.create_object_with_iterator(
+                EnvVariables.DATASETS.name,
+                ds_id,
+                ds_info,
+                "json",
+                iterator_keys=["examples"],
+                iterator_data=examples,
             )
-            return file_path
 
+            return file_path
         except Exception as e:
             logger.error(f"Failed to create dataset: {str(e)}")
             raise e
 
     @staticmethod
-    def _convert_csv(csv_file: str) -> list[dict]:
+    @validate_call
+    def convert_data(csv_file_path: str) -> Iterator[dict]:
         """
-        Converts a CSV file to a list of dictionaries.
+        Converts a CSV file to an iterator of dictionaries.
 
-        This method reads a CSV file and converts its contents into a list of dictionaries,
+        This method reads a CSV file and converts its contents into an iterator of dictionaries,
         where each dictionary represents a row in the CSV file.
 
         Args:
-            csv_file (str): The file path to the CSV file.
+            csv_file_path (str): The file path to the CSV file.
 
         Returns:
-            list[dict]: A list of dictionaries representing the CSV data.
+            Iterator[dict]: An iterator of dictionaries representing the CSV data.
         """
-        df = pd.read_csv(csv_file)
-        data = df.to_dict("records")
-        return data
+        df = pd.read_csv(csv_file_path, chunksize=1)
+        for chunk in df:
+            yield chunk.to_dict("records")[0]
 
     @staticmethod
-    def _download_hf(hf_args) -> list[dict]:
+    @validate_call
+    def download_hf(**hf_args) -> Iterator[dict]:
         """
-        Downloads a dataset from Hugging Face and converts it to a list of dictionaries.
+        Downloads a dataset from Hugging Face and converts it to an iterator of dictionaries.
 
         This method loads a dataset from Hugging Face based on the provided arguments and converts
-        its contents into a list of dictionaries, where each dictionary contains 'input' and 'target' keys.
+        its contents into an iterator of dictionaries, where each dictionary contains 'input' and 'target' keys.
 
         Args:
             hf_args (dict): A dictionary containing the following keys:
@@ -123,15 +113,17 @@ class Dataset:
                 - 'target_col' (str): The target column.
 
         Returns:
-            list[dict]: A list of dictionaries representing the dataset.
+            Iterator[dict]: An iterator of dictionaries representing the dataset.
         """
-        dataset = load_dataset(hf_args["dataset_name"], hf_args["dataset_config"])
-        data = []
-        for example in dataset[hf_args["split"]]:
+
+        dataset = load_dataset(
+            hf_args["dataset_name"], hf_args["dataset_config"], split=hf_args["split"]
+        )
+
+        for example in dataset:
             input_data = " ".join([str(example[col]) for col in hf_args["input_col"]])
             target_data = str(example[hf_args["target_col"]])
-            data.append({"input": input_data, "target": target_data})
-        return data
+            yield {"input": input_data, "target": target_data}
 
     @staticmethod
     @validate_call

moonshot/src/datasets/dataset_arguments.py

@@ -1,5 +1,6 @@
+from typing import Iterator
+
 from pydantic import BaseModel
-from pyparsing import Iterator
 
 
 class DatasetArguments(BaseModel):

moonshot/src/messages_constants.py

@@ -15,6 +15,7 @@ BOOKMARK_GET_BOOKMARK_ERROR_1 = "[Bookmark] Invalid bookmark name: {message}"
 # BOOKMARK - delete_bookmark
 # ------------------------------------------------------------------------------
 BOOKMARK_DELETE_BOOKMARK_SUCCESS = "[Bookmark] Bookmark record deleted."
+BOOKMARK_DELETE_BOOKMARK_FAIL = "[Bookmark] Bookmark record not found. Unable to delete."
 BOOKMARK_DELETE_BOOKMARK_ERROR = (
     "[Bookmark] Failed to delete bookmark record: {message}"
 )

moonshot/src/redteaming/attack/attack_module.py

@@ -34,6 +34,7 @@ class AttackModule:
 
     def __init__(self, am_id: str, am_arguments: AttackModuleArguments | None = None):
         self.id = am_id
+        self.req_and_config = self.get_attack_module_req_and_config()
         if am_arguments is not None:
             self.connector_ids = am_arguments.connector_ids
             self.prompt_templates = am_arguments.prompt_templates
@@ -529,6 +530,45 @@ class AttackModule:
 
         return am_metadata, cache_updated
 
+    def get_attack_module_req_and_config(self) -> dict:
+        """
+        Retrieves the configuration for a specific attack module by its identifier.
+
+        Returns:
+            dict: The attack module configuration as a dictionary. Returns an empty dict if the configuration
+            is not found.
+
+        Raises:
+            Exception: If reading the attack module configuration fails or if the configuration cannot be created.
+        """
+        attack_module_config = "attack_modules_config"
+        try:
+            obj_results = Storage.read_object(
+                EnvVariables.ATTACK_MODULES.name, attack_module_config, "json"
+            )
+            return obj_results.get(self.id, {})
+        except Exception as e:
+            logger.warning(
+                f"[AttackModule] Failed to read attack module configuration: {str(e)}"
+            )
+            logger.info("Attempting to create empty attack module configuration...")
+            try:
+                Storage.create_object(
+                    obj_type=EnvVariables.ATTACK_MODULES.name,
+                    obj_id=attack_module_config,
+                    obj_info={},
+                    obj_extension="json",
+                )
+                # After creation, attempt to read it again to ensure it was created successfully
+                obj_results = Storage.read_object(
+                    EnvVariables.ATTACK_MODULES.name, attack_module_config, "json"
+                )
+                return obj_results.get(self.id, {})
+            except Exception as e:
+                raise Exception(
+                    f"[AttackModule] Failed to retrieve attack modules configuration: {str(e)}"
+                )
+
     @staticmethod
     def delete(am_id: str) -> bool:
         """

moonshot/src/storage/io_interface.py

@@ -1,5 +1,5 @@
 from abc import abstractmethod
-from typing import Any
+from typing import Any, Iterator
 
 
 class IOInterface:
@@ -16,6 +16,23 @@ class IOInterface:
         """
         pass
 
+    @abstractmethod
+    def create_file_with_iterator(
+        self, data: dict, iterator_keys: list[str], iterator_data: Iterator[dict]
+    ) -> bool:
+        """
+        Creates a file using an iterator to provide the data for specified keys.
+
+        Args:
+            data (dict): The data to be serialized into JSON and written to the file.
+            iterator_keys (list[str]): A list of keys for which the values will be written using iterators.
+            iterator_data (Iterator[dict]): An iterator for the data to be written for the specified keys.
+
+        Returns:
+            bool: Always returns True to indicate the operation was executed without raising an exception.
+        """
+        pass
+
     @abstractmethod
     def read_file(self, filepath: str) -> dict | None:
         """