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

moonshot/src/api/api_bookmark.py

@@ -0,0 +1,95 @@
+from moonshot.src.bookmark.bookmark import Bookmark
+from moonshot.src.bookmark.bookmark_arguments import BookmarkArguments
+
+
+def api_insert_bookmark(
+    name: str,
+    prompt: str,
+    prepared_prompt: str,
+    response: str,
+    context_strategy: str = "",
+    prompt_template: str = "",
+    attack_module: str = "",
+    metric: str = "",
+) -> dict:
+    """
+    Inserts a new bookmark into the database.
+
+    This function constructs a BookmarkArguments object with the provided details and
+    invokes the add_bookmark method of a Bookmark instance to persist the new bookmark.
+
+    Args:
+        name (str): The unique name of the bookmark.
+        prompt (str): The associated prompt text for the bookmark.
+        response (str): The corresponding response text for the bookmark.
+        context_strategy (str): The strategy used for context management in the bookmark.
+        prompt_template (str): The template used for generating the prompt.
+        attack_module (str): The attack module linked with the bookmark.
+    """
+    # Create a new BookmarkArguments object
+    bookmark_args = BookmarkArguments(
+        id=0,  # id will be auto-generated by the database
+        name=name,
+        prompt=prompt,
+        prepared_prompt=prepared_prompt,
+        response=response,
+        context_strategy=context_strategy,
+        prompt_template=prompt_template,
+        attack_module=attack_module,
+        metric=metric,
+        bookmark_time="",  # bookmark_time will be set to current time in add_bookmark method
+    )
+    return Bookmark.get_instance().add_bookmark(bookmark_args)
+
+
+def api_get_all_bookmarks() -> list[dict]:
+    """
+    Retrieves a list of all bookmarks from the database.
+
+    Returns:
+        list[dict]: A list of bookmarks, each represented as a dictionary.
+    """
+    return Bookmark.get_instance().get_all_bookmarks()
+
+
+def api_get_bookmark(bookmark_name: str) -> dict:
+    """
+    Retrieves the details of a specific bookmark by its name.
+
+    Args:
+        bookmark_name (int): The name of the bookmark to retrieve.
+
+    Returns:
+        dict: The bookmark details corresponding to the provided ID.
+    """
+    return Bookmark.get_instance().get_bookmark(bookmark_name)
+
+
+def api_delete_bookmark(bookmark_name: str) -> dict:
+    """
+    Removes a bookmark from the database using its name.
+
+    Args:
+        bookmark_name (str): The name of the bookmark to be removed.
+    """
+    return Bookmark.get_instance().delete_bookmark(bookmark_name)
+
+
+def api_delete_all_bookmark() -> dict:
+    """
+    Removes all bookmarks from the database.
+    """
+    return Bookmark.get_instance().delete_all_bookmark()
+
+
+def api_export_bookmarks(export_file_name: str = "bookmarks") -> str:
+    """
+    Exports bookmarks to a specified file.
+
+    Args:
+        export_file_name (str): The name of the file to export the bookmarks to.
+
+    Returns:
+        str: The filepath of where the file is written.
+    """
+    return Bookmark.get_instance().export_bookmarks(export_file_name)

moonshot/src/api/api_connector_endpoint.py

@@ -98,7 +98,7 @@ def api_update_endpoint(ep_id: str, **kwargs) -> bool:
         existing_endpoint = ConnectorEndpoint.read(ep_id)
     except Exception:
         raise RuntimeError(
-            f"[api_connector_endpoint]: Endpoint with ID '{ep_id}' does not exist"
+            f"[api_update_endpoint]: Endpoint with ID '{ep_id}' does not exist"
         )
 
     # Update the fields of the existing endpoint with the provided kwargs

moonshot/src/api/api_context_strategy.py

@@ -40,7 +40,7 @@ def api_delete_context_strategy(cs_id: str) -> bool:
     return ContextStrategy.delete(cs_id)
 
 
-def api_get_all_context_strategy_metadata() -> list:
+def api_get_all_context_strategy_metadata() -> list[dict]:
     """
     Retrieves metadata for all context strategy modules.
 
@@ -48,7 +48,7 @@ def api_get_all_context_strategy_metadata() -> list:
     returns a list of metadata dictionaries.
 
     Returns:
-        list: A list of attack module metadata.
+        list[dict]: A list of dictionaries, each representing the details of a context strategy metadata.
     """
 
     return [

moonshot/src/api/api_dataset.py

@@ -1,6 +1,7 @@
 from pydantic import validate_call
 
 from moonshot.src.datasets.dataset import Dataset
+from moonshot.src.datasets.dataset_arguments import DatasetArguments
 
 
 # ------------------------------------------------------------------------------
@@ -44,3 +45,37 @@ def api_get_all_datasets_name() -> list[str]:
     """
     datasets_name, _ = Dataset.get_available_items()
     return datasets_name
+
+
+def api_create_datasets(
+    name: str, description: str, reference: str, license: str, method: str, **kwargs
+) -> str:
+    """
+    This function creates a new dataset.
+
+    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.
+
+    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.
+
+    Returns:
+        str: The ID of the newly created dataset.
+    """
+    ds_args = DatasetArguments(
+        id="",
+        name=name,
+        description=description,
+        reference=reference,
+        license=license,
+        examples=None,
+    )
+
+    return Dataset.create(ds_args, method, **kwargs)

moonshot/src/api/api_recipe.py

@@ -16,7 +16,6 @@ def api_create_recipe(
     datasets: list[str],
     prompt_templates: list[str],
     metrics: list[str],
-    attack_modules: list[str],
     grading_scale: dict[str, list[int]],
 ) -> str:
     """
@@ -34,7 +33,6 @@ def api_create_recipe(
         datasets (list[str]): A list of datasets used in the recipe.
         prompt_templates (list[str]): A list of prompt templates for the recipe.
         metrics (list[str]): A list of metrics to evaluate the recipe.
-        attack_modules (list[str]): A list of attack modules used in the recipe.
         grading_scale (dict[str, list[int]]): A grading scale dictionary where the key is the grade and the
         value is a list of integers representing the scale.
 
@@ -50,7 +48,6 @@ def api_create_recipe(
         datasets=datasets,
         prompt_templates=prompt_templates,
         metrics=metrics,
-        attack_modules=attack_modules,
         grading_scale=grading_scale,
     )
     return Recipe.create(rec_args)

moonshot/src/api/api_session.py

@@ -115,7 +115,7 @@ def api_get_available_session_info() -> tuple[list, list]:
     return runner_ids, session_metadata_list
 
 
-def api_get_all_session_metadata() -> list:
+def api_get_all_session_metadata() -> list[dict]:
     """
     Retrieves metadata for all sessions.
 

moonshot/src/bookmark/bookmark.py

@@ -0,0 +1,257 @@
+from datetime import datetime
+
+from moonshot.src.bookmark.bookmark_arguments import BookmarkArguments
+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 Bookmark:
+    _instance = None
+
+    def __new__(cls, db_name="bookmark"):
+        """
+        Create a new instance of the Bookmark class or return the existing instance.
+
+        Args:
+            db_name (str): The name of the database.
+
+        Returns:
+            Bookmark: The singleton instance of the Bookmark class.
+        """
+        if cls._instance is None:
+            cls._instance = super(Bookmark, cls).__new__(cls)
+            cls._instance.__init_instance(db_name)
+        return cls._instance
+
+    @classmethod
+    def get_instance(cls, db_name="bookmark"):
+        """
+        Get the singleton instance of the Bookmark class.
+
+        Args:
+            db_name (str): The name of the database.
+
+        Returns:
+            Bookmark: The singleton instance of the Bookmark class.
+        """
+        if cls._instance is None:
+            cls._instance = super(Bookmark, cls).__new__(cls)
+            cls._instance.__init_instance(db_name)
+        return cls._instance
+
+    sql_create_bookmark_table = """
+        CREATE TABLE IF NOT EXISTS bookmark (
+        id INTEGER PRIMARY KEY AUTOINCREMENT,
+        name TEXT NOT NULL UNIQUE,
+        prompt TEXT NOT NULL,
+        prepared_prompt TEXT NOT NULL,
+        response TEXT NOT NULL,
+        context_strategy TEXT,
+        prompt_template TEXT,
+        attack_module TEXT,
+        metric TEXT,
+        bookmark_time DATETIME NOT NULL DEFAULT CURRENT_TIMESTAMP
+        );
+    """
+
+    sql_insert_bookmark_record = """
+        INSERT INTO bookmark (
+        name, prompt, prepared_prompt, response, context_strategy, prompt_template, attack_module,
+        metric, bookmark_time)
+        VALUES (?,?,?,?,?,?,?,?,?);
+    """
+
+    sql_select_bookmarks_record = """
+        SELECT * FROM bookmark;
+    """
+
+    sql_select_bookmark_record = """
+        SELECT * FROM bookmark WHERE name = ? ;
+    """
+
+    sql_delete_bookmark_records = """
+        DELETE FROM bookmark;
+    """
+
+    def __init_instance(self, db_name) -> None:
+        """
+        Initialize the database instance for the Bookmark class.
+
+        Args:
+            db_name (str): The name of the database.
+        """
+        self.db_instance = Storage.create_database_connection(
+            EnvVariables.BOOKMARKS.name, db_name, "db"
+        )
+        Storage.create_database_table(
+            self.db_instance, Bookmark.sql_create_bookmark_table
+        )
+
+    def add_bookmark(self, bookmark: BookmarkArguments) -> dict:
+        """
+        Add a new bookmark to the database and return the success status.
+
+        Args:
+            bookmark (BookmarkArguments): The bookmark data to add.
+
+        Returns:
+            bool: True if the bookmark was added successfully, False otherwise.
+        """
+        bookmark.bookmark_time = datetime.now().replace(microsecond=0).isoformat(" ")
+
+        data = (
+            bookmark.name,
+            bookmark.prompt,
+            bookmark.prepared_prompt,
+            bookmark.response,
+            bookmark.context_strategy,
+            bookmark.prompt_template,
+            bookmark.attack_module,
+            bookmark.metric,
+            bookmark.bookmark_time,
+        )
+        try:
+            results = Storage.create_database_record(
+                self.db_instance, data, Bookmark.sql_insert_bookmark_record
+            )
+            if results is not None:
+                return {"success": True, "message": "Bookmark added successfully."}
+            else:
+                raise Exception("Error inserting record into database.")
+        except Exception as e:
+            error_message = f"Failed to add bookmark record: {e}"
+            return {"success": False, "message": error_message}
+
+    def get_all_bookmarks(self) -> list[dict]:
+        """
+        Retrieve all bookmarks from the database.
+
+        Returns:
+            list[dict]: A list of all bookmarks as dictionaries.
+        """
+        list_of_bookmarks_tuples = Storage.read_database_records(
+            self.db_instance,
+            Bookmark.sql_select_bookmarks_record,
+        )
+        if list_of_bookmarks_tuples:
+            list_of_bookmarks = [
+                BookmarkArguments.from_tuple_to_dict(bookmark_tuple)
+                for bookmark_tuple in list_of_bookmarks_tuples
+            ]
+        else:
+            list_of_bookmarks = []
+        return list_of_bookmarks
+
+    def get_bookmark(self, bookmark_name: str) -> dict:
+        """
+        Retrieve a bookmark by its unique name.
+
+        Args:
+            bookmark_name (int): The unique name for the bookmark.
+
+        Returns:
+            dict: The bookmark information as a dictionary.
+
+        Raises:
+            RuntimeError: If the bookmark cannot be found.
+        """
+        if bookmark_name is not None:
+            bookmark_info = Storage.read_database_record(
+                self.db_instance, (bookmark_name,), Bookmark.sql_select_bookmark_record
+            )
+            if bookmark_info is not None:
+                return BookmarkArguments.from_tuple_to_dict(bookmark_info)
+            else:
+                raise RuntimeError(
+                    f"[Bookmark] No record found for bookmark name {bookmark_name}"
+                )
+        else:
+            raise RuntimeError(f"[Bookmark] Invalid bookmark name: {bookmark_name}")
+
+    def delete_bookmark(self, bookmark_name: str) -> dict:
+        """
+        Delete a bookmark by its unique name.
+
+        Args:
+            bookmark_name (str): The unique name for the bookmark to be deleted.
+        """
+        if bookmark_name is not None:
+            try:
+                sql_delete_bookmark_record = 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 record deleted."}
+            except Exception as e:
+                error_message = f"Failed to delete bookmark record: {e}"
+                return {"success": False, "message": error_message}
+        else:
+            error_message = f"[Bookmark] Invalid bookmark name: {bookmark_name}"
+            return {"success": False, "message": error_message}
+
+    def delete_all_bookmark(self) -> dict:
+        """
+        Delete all bookmarks from the database and return the operation result.
+
+        Returns:
+            dict: A dictionary with 'success' status and 'message' containing an error message if failed.
+        """
+        try:
+            Storage.delete_database_record_in_table(
+                self.db_instance, Bookmark.sql_delete_bookmark_records
+            )
+            return {"success": True, "message": "All bookmark records deleted."}
+        except Exception as e:
+            error_message = f"Failed to delete all bookmark records: {e}"
+            return {"success": False, "message": error_message}
+
+    def export_bookmarks(self, export_file_name: str = "bookmarks") -> str:
+        """
+        Export all bookmarks to a JSON file.
+
+        This method retrieves all bookmarks from the database, converts them to a JSON format,
+        and writes them to a file in the 'moonshot-data/bookmark' directory with the provided file name.
+
+        Args:
+            export_file_name (str): The base name of the file to export the bookmarks to.
+                                    The '.json' extension will be appended to this base name.
+
+        Returns:
+            str: The path to the exported JSON file containing the bookmarks.
+        """
+        list_of_bookmarks_tuples = Storage.read_database_records(
+            self.db_instance,
+            Bookmark.sql_select_bookmarks_record,
+        )
+
+        if list_of_bookmarks_tuples is not None:
+            bookmarks_json = [
+                BookmarkArguments.from_tuple_to_dict(bookmark_tuple)
+                for bookmark_tuple in list_of_bookmarks_tuples
+            ]
+        else:
+            bookmarks_json = []
+
+        try:
+            return Storage.create_object(
+                EnvVariables.BOOKMARKS.name,
+                export_file_name,
+                {"bookmarks": bookmarks_json},
+                "json",
+            )
+        except Exception as e:
+            logger.error(f"Failed to export bookmarks - {str(e)}.")
+            raise e
+
+    def close(self) -> None:
+        """
+        Close the database connection.
+        """
+        if self.db_instance:
+            Storage.close_database_connection(self.db_instance)

moonshot/src/bookmark/bookmark_arguments.py

@@ -0,0 +1,38 @@
+from __future__ import annotations
+
+from pydantic import BaseModel, Field
+
+
+class BookmarkArguments(BaseModel):
+    name: str = Field(min_length=1)
+    prompt: str = Field(min_length=1)
+    prepared_prompt: str = Field(min_length=1)
+    response: str = Field(min_length=1)
+    context_strategy: str
+    prompt_template: str
+    attack_module: str
+    metric: str
+    bookmark_time: str
+
+    @classmethod
+    def from_tuple_to_dict(cls, values: tuple) -> dict:
+        """
+        Creates a dictionary from a tuple of values representing BookmarkArguments fields.
+
+        Args:
+            values (tuple): A tuple containing the fields of BookmarkArguments in order.
+
+        Returns:
+            dict: A dictionary representing the BookmarkArguments.
+        """
+        return {
+            "name": values[1],
+            "prompt": values[2],
+            "prepared_prompt": values[3],
+            "response": values[4],
+            "context_strategy": values[5],
+            "prompt_template": values[6],
+            "attack_module": values[7],
+            "metric": values[8],
+            "bookmark_time": values[9],
+        }

moonshot/src/configs/env_variables.py

@@ -2,11 +2,17 @@ import importlib.resources
 from enum import Enum
 from pathlib import Path
 
+from moonshot.src.utils.log import configure_logger
+
+# Create a logger for this module
+logger = configure_logger(__name__)
+
 __app_name__ = "moonshot"
 
 
 class EnvVariables(Enum):
     ATTACK_MODULES = "ATTACK_MODULES"
+    BOOKMARKS = "BOOKMARKS"
     CONNECTORS = "CONNECTORS"
     CONNECTORS_ENDPOINTS = "CONNECTORS_ENDPOINTS"
     CONTEXT_STRATEGY = "CONTEXT_STRATEGY"
@@ -31,6 +37,10 @@ class EnvironmentVars:
         env_vars.get(EnvVariables.ATTACK_MODULES.value),
         str(importlib.resources.files(__app_name__).joinpath("data/attack-modules")),
     ]
+    BOOKMARKS = [
+        env_vars.get(EnvVariables.BOOKMARKS.value),
+        str(importlib.resources.files(__app_name__).joinpath("data/generated-outputs/bookmarks")),
+    ]
     CONNECTORS = [
         env_vars.get(EnvVariables.CONNECTORS.value),
         str(importlib.resources.files(__app_name__).joinpath("data/connectors")),
@@ -133,14 +143,14 @@ class EnvironmentVars:
                 if given_path.exists():
                     EnvironmentVars.__dict__[key][0] = str(given_path)
                 else:
-                    print(
+                    logger.warning(
                         f"Unable to set {key}. The provided path {given_path} does not exist. ",
                         "The stock set will be used.",
                     )
             else:
                 unset_keys.append(key)
         if unset_keys:
-            print(
+            logger.warning(
                 f"Unable to retrieve the following environment variables: {unset_keys}. The stock set will be used."
             )
 

moonshot/src/connectors/connector.py

@@ -15,6 +15,10 @@ from moonshot.src.connectors_endpoints.connector_endpoint_arguments import (
 )
 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__)
 
 
 def perform_retry(func):
@@ -42,13 +46,15 @@ def perform_retry(func):
                 try:
                     return await func(self, *args, **kwargs)
                 except Exception as exc:
-                    print(f"Operation failed. {str(exc)} - Retrying...")
+                    logger.warning(f"Operation failed. {str(exc)} - Retrying...")
 
                 # Perform retry
                 retry_count += 1
                 if retry_count <= self.retries_times:
                     delay = base_delay * (2**retry_count)
-                    print(f"Attempt {retry_count}, Retrying in {delay} seconds...")
+                    logger.warning(
+                        f"Attempt {retry_count}, Retrying in {delay} seconds..."
+                    )
                     await sleep(delay)
             # Raise an exception
             raise ConnectionError("Failed to get response.")
@@ -228,7 +234,7 @@ class Connector:
             return Connector.load(ep_args)
 
         except Exception as e:
-            print(f"Failed to create connector: {str(e)}")
+            logger.error(f"Failed to create connector: {str(e)}")
             raise e
 
     @staticmethod
@@ -255,7 +261,7 @@ class Connector:
             ]
 
         except Exception as e:
-            print(f"Failed to get available connectors: {str(e)}")
+            logger.error(f"Failed to get available connectors: {str(e)}")
             raise e
 
     @staticmethod
@@ -292,14 +298,16 @@ class Connector:
             Exception: If there is an error during prediction.
         """
         try:
-            print(f"Predicting prompt {generated_prompt.prompt_index} [{connector.id}]")
+            logger.info(
+                f"Predicting prompt {generated_prompt.prompt_index} [{connector.id}]"
+            )
 
             start_time = time.perf_counter()
             generated_prompt.predicted_results = await connector.get_response(
                 generated_prompt.prompt
             )
             generated_prompt.duration = time.perf_counter() - start_time
-            print(
+            logger.debug(
                 f"[Prompt {generated_prompt.prompt_index}] took {generated_prompt.duration:.4f}s"
             )
 
@@ -311,7 +319,7 @@ class Connector:
             return generated_prompt
 
         except Exception as e:
-            print(f"Failed to get prediction: {str(e)}")
+            logger.error(f"Failed to get prediction: {str(e)}")
             raise e
 
     def set_system_prompt(self, system_prompt: str) -> None: