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

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:

moonshot/src/connectors_endpoints/connector_endpoint.py

@@ -8,33 +8,39 @@ from moonshot.src.connectors_endpoints.connector_endpoint_arguments import (
     ConnectorEndpointArguments,
 )
 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 ConnectorEndpoint:
     @staticmethod
     def create(ep_args: ConnectorEndpointArguments) -> str:
         """
-        Creates a new connector endpoint.
+        Creates a new connector endpoint and stores its details as a JSON object.
+
+        This method accepts a ConnectorEndpointArguments object, generates a unique slugified ID from the endpoint's
+        name, and stores the endpoint's details in a JSON file within a specified directory.
 
-        This method takes a ConnectorEndpointArguments object as input, generates a unique slugified ID based on the
-        endpoint's name, and then creates a new endpoint with the provided details. The endpoint information is stored
-        as a JSON object in the directory specified by `EnvVariables.CONNECTORS_ENDPOINTS`. If the operation is
-        successful, the unique ID of the new endpoint is returned. If any error arises during the process, an exception
-        is raised and the error message is logged.
+        The directory path is determined by the `EnvVariables.CONNECTORS_ENDPOINTS` environment variable.
+        Upon successful creation, the method returns the unique ID of the endpoint.
+        If an error occurs during the creation process, the method raises an exception and logs the error message.
 
         Args:
-            ep_args (ConnectorEndpointArguments): An object containing the details of the endpoint to be created.
+            ep_args (ConnectorEndpointArguments): The details of the endpoint to be created,
+            encapsulated in a ConnectorEndpointArguments object.
 
         Returns:
-            str: The unique ID of the newly created endpoint.
+            str: The unique ID of the newly created endpoint, derived from slugifying the endpoint's name.
 
         Raises:
-            Exception: If there's an error during the endpoint creation process.
+            Exception: If an error occurs during the creation process, including issues with storing the endpoint's
+            details.
         """
         try:
             ep_id = slugify(ep_args.name, lowercase=True)
             ep_info = {
-                "id": ep_id,
                 "name": ep_args.name,
                 "connector_type": ep_args.connector_type,
                 "uri": ep_args.uri,
@@ -51,59 +57,63 @@ class ConnectorEndpoint:
             return ep_id
 
         except Exception as e:
-            print(f"Failed to create endpoint: {str(e)}")
+            logger.error(f"Failed to create endpoint: {str(e)}")
             raise e
 
     @staticmethod
     @validate_call
     def read(ep_id: str) -> ConnectorEndpointArguments:
         """
-        Fetches the details of a given endpoint.
+        Retrieves the details of a specified endpoint by its ID.
 
-        This method takes an endpoint ID as input, finds the corresponding JSON file in the directory
-        specified by `EnvironmentVars.CONNECTORS_ENDPOINTS`, and returns a ConnectorEndpointArguments object
-        that contains the endpoint's details. If any error arises during the process, an exception is raised and the
-        error message is logged.
+        This method searches for the endpoint's corresponding JSON file within the directory defined by the
+        `EnvVariables.CONNECTORS_ENDPOINTS` environment variable. It then constructs and returns a
+        ConnectorEndpointArguments object populated with the endpoint's details. If the endpoint ID is not found or
+        any other error occurs, an exception is raised with an appropriate error message.
 
         Args:
-            ep_id (str): The unique ID of the endpoint to be fetched.
+            ep_id (str): The unique identifier of the endpoint whose details are to be retrieved.
 
         Returns:
-            ConnectorEndpointArguments: An object encapsulating the details of the fetched endpoint.
+            ConnectorEndpointArguments: An instance filled with the endpoint's details.
 
         Raises:
-            Exception: If there's an error during the file reading process or any other operation within the method.
+            RuntimeError: If the endpoint ID is empty or the specified endpoint does not exist.
+            Exception: For any issues encountered during the file reading or data parsing process.
         """
         try:
-            if ep_id:
-                return ConnectorEndpointArguments(
-                    **ConnectorEndpoint._read_endpoint(ep_id)
-                )
-            else:
-                raise RuntimeError("Connector Endpoint ID is empty")
+            if not ep_id:
+                raise RuntimeError("Connector Endpoint ID is empty.")
+
+            endpoint_details = ConnectorEndpoint._read_endpoint(ep_id)
+            if not endpoint_details:
+                raise RuntimeError(f"Endpoint with ID '{ep_id}' does not exist.")
+
+            return ConnectorEndpointArguments(**endpoint_details)
 
         except Exception as e:
-            print(f"Failed to read endpoint: {str(e)}")
+            logger.error(f"Failed to read endpoint: {str(e)}")
             raise e
 
     @staticmethod
     def _read_endpoint(ep_id: str) -> dict:
         """
-        Reads the endpoint information from a JSON file and adds the creation datetime.
+        Retrieves the endpoint's information from a JSON file, including its creation datetime.
 
-        This method accepts an endpoint ID as an argument, locates the corresponding JSON file in the directory
-        defined by `EnvironmentVars.CONNECTORS_ENDPOINTS`, and returns a dictionary that encapsulates the endpoint's
-        details along with its creation datetime. If any error occurs during the process, it is handled by the calling
-        method.
+        This internal method is designed to fetch the details of a specific endpoint by its ID. It searches for the
+        corresponding JSON file within the directory specified by `EnvVariables.CONNECTORS_ENDPOINTS`. The method
+        returns a dictionary containing the endpoint's information, enriched with the creation datetime. Errors
+        encountered during this process are managed by the method that invokes this one.
 
         Args:
-            ep_id (str): The unique identifier of the endpoint to be retrieved.
+            ep_id (str): The unique identifier of the endpoint whose information is being retrieved.
 
         Returns:
-            dict: A dictionary containing the details of the retrieved endpoint along with its creation datetime.
+            dict: A dictionary with the endpoint's information, including its creation datetime.
         """
-        connector_endpoint_info = Storage.read_object(
-            EnvVariables.CONNECTORS_ENDPOINTS.name, ep_id, "json"
+        connector_endpoint_info = {"id": ep_id}
+        connector_endpoint_info.update(
+            Storage.read_object(EnvVariables.CONNECTORS_ENDPOINTS.name, ep_id, "json")
         )
         creation_datetime = Storage.get_creation_datetime(
             EnvVariables.CONNECTORS_ENDPOINTS.name, ep_id, "json"
@@ -116,35 +126,41 @@ class ConnectorEndpoint:
     @staticmethod
     def update(ep_args: ConnectorEndpointArguments) -> bool:
         """
-        Updates the endpoint information based on the provided arguments.
+        Updates the endpoint information in the storage based on the provided ConnectorEndpointArguments object.
+
+        This method serializes the provided ConnectorEndpointArguments object into a dictionary, excluding the 'id' and
+        'created_date' keys. It then persists the updated information to the corresponding JSON file within the
+        directory defined by `EnvVariables.CONNECTORS_ENDPOINTS`.
 
-        This method takes a ConnectorEndpointArguments object, converts it to a dictionary, and removes the
-        'created_date' key if it exists. It then writes the updated information to the corresponding JSON file
-        in the directory specified by `EnvVariables.CONNECTORS_ENDPOINTS`.
+        This operation ensures that the endpoint's mutable attributes are updated according to the provided arguments.
 
         Args:
-            ep_args (ConnectorEndpointArguments): An object containing the updated details of the endpoint.
+            ep_args (ConnectorEndpointArguments): The object encapsulating the updated attributes of the endpoint.
 
         Returns:
-            bool: True if the update operation was successful.
+            bool: Indicates whether the update operation was successful. Returns True if the update was successfully
+            persisted to the storage; otherwise, an exception is raised.
 
         Raises:
-            Exception: If there's an error during the update process.
+            Exception: Signifies a failure in the update process, potentially due to issues with data serialization or
+            storage access.
         """
         try:
-            # Convert the endpoint arguments to a dictionary
-            # Remove created_date if it exists
+            # Serialize the ConnectorEndpointArguments object to a dictionary and remove derived properties
             ep_info = ep_args.to_dict()
-            ep_info.pop("created_date", None)
+            ep_info.pop("id", None)  # The 'id' is derived and should not be written
+            ep_info.pop(
+                "created_date", None
+            )  # The 'created_date' is derived and should not be written
 
-            # Write the updated endpoint information to the file
+            # Write the updated endpoint information to the storage
             Storage.create_object(
                 EnvVariables.CONNECTORS_ENDPOINTS.name, ep_args.id, ep_info, "json"
             )
             return True
 
         except Exception as e:
-            print(f"Failed to update endpoint: {str(e)}")
+            logger.error(f"Failed to update endpoint: {str(e)}")
             raise e
 
     @staticmethod
@@ -171,7 +187,7 @@ class ConnectorEndpoint:
             return True
 
         except Exception as e:
-            print(f"Failed to delete endpoint: {str(e)}")
+            logger.error(f"Failed to delete endpoint: {str(e)}")
             raise e
 
     @staticmethod
@@ -207,5 +223,5 @@ class ConnectorEndpoint:
             return retn_eps_ids, retn_eps
 
         except Exception as e:
-            print(f"Failed to get available endpoints: {str(e)}")
+            logger.error(f"Failed to get available endpoints: {str(e)}")
             raise e