aiverify-moonshot 0.4.0__py3-none-any.whl

moonshot/src/redteaming/attack/context_strategy.py

@@ -0,0 +1,131 @@
+from __future__ import annotations
+
+from pathlib import Path
+
+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
+
+
+class ContextStrategy:
+    def __init__(self, cs_id: str) -> None:
+        self.id = cs_id
+
+    @classmethod
+    def load(cls, cs_id: str) -> ContextStrategy:
+        """
+        Retrieves a context strategy module instance by its ID.
+
+        This method attempts to load a context strategy instance using the provided ID. If the context strategy instance
+        is found, it is returned. If the context strategy instance does not exist, a RuntimeError is raised.
+
+        Args:
+            cs_id (str): The unique identifier of the context strategy to be retrieved.
+
+        Returns:
+            ContextStrategy: The retrieved context strategy instance.
+
+        Raises:
+            RuntimeError: If the context strategy instance does not exist.
+        """
+        context_strategy_inst = get_instance(
+            cs_id,
+            Storage.get_filepath(EnvVariables.CONTEXT_STRATEGY.name, cs_id, "py"),
+        )
+        if context_strategy_inst:
+            return context_strategy_inst(cs_id)
+        else:
+            raise RuntimeError(
+                f"Unable to get defined context strategy instance - {cs_id}"
+            )
+
+    @staticmethod
+    def get_all_context_strategies() -> list[str]:
+        """
+        Retrieves the names of all context strategy files.
+
+        This method fetches the names of all context strategy files by scanning the directory specified in the
+        EnvironmentVars. It filters out filenames containing double underscores before returning the list
+        of context strategy names.
+
+        Returns:
+            list: A list of context strategy file names.
+        """
+        filepaths = []
+        context_strategy_files = Storage.get_objects(
+            EnvVariables.CONTEXT_STRATEGY.name, "py"
+        )
+        for context_strategy in context_strategy_files:
+            filepaths.append(Path(context_strategy).stem)
+        return filepaths
+
+    @staticmethod
+    def delete(cs_id: str) -> bool:
+        """
+        Deletes a context strategy module instance by its ID.
+
+        This method attempts to delete a context strategy instance using the provided ID.
+        If the deletion is successful, it returns True.
+
+        If an error occurs during the deletion process, it prints an error message and re-raises the exception.
+
+        Args:
+            cs_id (str): The unique identifier of the context strategy to be deleted.
+
+        Returns:
+            bool: True if the deletion was successful.
+
+        Raises:
+            Exception: If an error occurs during the deletion process.
+        """
+        try:
+            Storage.delete_object(EnvVariables.CONTEXT_STRATEGY.name, cs_id, "py")
+            return True
+
+        except Exception as e:
+            print(f"Failed to delete context strategy: {str(e)}")
+            raise e
+
+    @staticmethod
+    def process_prompt_cs(
+        user_prompt: str,
+        context_strategy_name: str,
+        db_instance: DBInterface,
+        endpoint_id,
+        num_of_previous_chats: int,
+    ) -> str:
+        """
+        Process the user prompt using the specified context strategy.
+
+        Args:
+            user_prompt (str): The user prompt to process.
+            context_strategy_name (str): The name of the context strategy to use.
+            db_instance (DBInterface): The database interface instance.
+            endpoint_id: The ID of the endpoint.
+            num_of_previous_chats (int): The number of previous chats to add.
+
+        Returns:
+            str: The processed user prompt based on the context strategy.
+        """
+        from moonshot.src.redteaming.session.chat import Chat
+
+        context_strategy_instance = get_instance(
+            context_strategy_name,
+            Storage.get_filepath(
+                EnvVariables.CONTEXT_STRATEGY.name, context_strategy_name, "py"
+            ),
+        )
+
+        if context_strategy_instance:
+            # get the last n chats
+            list_of_chats = Chat.get_n_chat_history(
+                db_instance, endpoint_id, num_of_previous_chats
+            )
+            context_strategy_instance = context_strategy_instance(context_strategy_name)
+            return context_strategy_instance.add_in_context(user_prompt, list_of_chats)
+        else:
+            print(
+                "Cannot load context strategy. Make sure the name of the context strategy is correct."
+            )
+            return ""

moonshot/src/redteaming/context_strategy/context_strategy_interface.py

@@ -0,0 +1,46 @@
+from abc import abstractmethod
+
+from moonshot.src.utils.timeit import timeit
+
+
+class ContextStrategyInterface:
+    @abstractmethod
+    @timeit
+    def get_metadata(self) -> dict | None:
+        """
+        Abstract method to retrieve metadata from the context strategy.
+
+        Returns:
+            dict | None: Returns a dictionary containing the metadata of the context strategy, or None if the
+            operation was unsuccessful.
+        """
+        pass
+
+    @abstractmethod
+    def add_in_context(
+        self, user_prompt: str, list_of_previous_prompts: list[dict] = []
+    ) -> None:
+        """
+        Abstract method to add a user prompt and list of previous prompts to the context strategy.
+
+        Args:
+            user_prompt (str): The user prompt to be added to the context strategy.
+            list_of_previous_prompts (list[dict], optional): List of previous prompts. Defaults to [].
+
+        Returns:
+            None
+        """
+        pass
+
+    @abstractmethod
+    def get_number_of_prev_prompts(self, no_prev_prompts: int) -> int:
+        """
+        Abstract method to get the number of previous prompts to be retrieved from the context strategy.
+
+        Args:
+            no_prev_prompts (int): The number of previous prompts to retrieve.
+
+        Returns:
+            int: The number of previous prompts to be retrieved.
+        """
+        pass

moonshot/src/redteaming/session/chat.py

@@ -0,0 +1,209 @@
+from __future__ import annotations
+
+from typing import Union
+
+from slugify import slugify
+
+from moonshot.src.storage.db_interface import DBInterface
+from moonshot.src.storage.storage import Storage
+
+
+class ChatRecord:
+    def __init__(
+        self,
+        chat_record_id: str,
+        conn_id: str,
+        context_strategy: str,
+        prompt_template: str,
+        attack_module: str,
+        metric: str,
+        prompt: str,
+        prepared_prompt: str,
+        system_prompt: str,
+        predicted_result: str,
+        duration: str,
+        prompt_time: str,
+    ):
+        self.chat_record_id = chat_record_id
+        self.conn_id = conn_id
+        self.context_strategy = context_strategy
+        self.prompt_template = prompt_template
+        self.attack_module = attack_module
+        self.metric = metric
+        self.prompt = prompt
+        self.prepared_prompt = prepared_prompt
+        self.system_prompt = system_prompt
+        self.predicted_result = predicted_result
+        self.duration = duration
+        self.prompt_time = prompt_time
+
+    def to_dict(self) -> dict[str, str]:
+        """
+        Converts the ChatRecord instance into a dictionary.
+
+        Returns:
+            dict: A dictionary representation of the ChatRecord instance.
+        """
+        return {
+            "chat_record_id": self.chat_record_id,
+            "conn_id": self.conn_id,
+            "context_strategy": self.context_strategy,
+            "prompt_template": self.prompt_template,
+            "attack_module": self.attack_module,
+            "metric": self.metric,
+            "prompt": self.prompt,
+            "prepared_prompt": self.prepared_prompt,
+            "system_prompt": self.system_prompt,
+            "predicted_result": self.predicted_result,
+            "duration": self.duration,
+            "prompt_time": self.prompt_time,
+        }
+
+
+class Chat:
+    sql_create_chat_metadata_record = """
+            INSERT INTO chat_metadata_table (
+            chat_id,endpoint,created_epoch,created_datetime)
+            VALUES(?,?,?,?)
+    """
+
+    sql_select_n_chat_from_chat_table = (
+        """SELECT * FROM {} order by prompt_time desc limit {}"""
+    )
+
+    def __init__(
+        self,
+        session_db_instance: DBInterface,
+        endpoint: str = "",
+        created_epoch: float = 0.0,
+        created_datetime: str = "",
+        chat_id: str = "",
+    ):
+        self.chat_history = []
+        if chat_id:
+            db_chat_id = chat_id.replace("-", "_")
+            # There is an existing chat
+            self.chat_id = db_chat_id
+            self.endpoint = endpoint
+            self.chat_history = self.load_chat_history(session_db_instance, db_chat_id)
+        else:
+            # No existing chat, create new chat
+            created_datetime = str(created_datetime).replace("-", "_")
+            chat_id = f"{slugify(endpoint)}_{created_datetime}"
+            db_chat_id = chat_id.replace("-", "_")
+
+            sql_create_chat_history_table = f"""
+                CREATE TABLE IF NOT EXISTS {db_chat_id} (
+                id INTEGER PRIMARY KEY AUTOINCREMENT,
+                connection_id text NOT NULL,
+                context_strategy int,
+                prompt_template text,
+                prompt text NOT NULL,
+                prepared_prompt text NOT NULL,
+                predicted_result text NOT NULL,
+                duration text NOT NULL,
+                prompt_time text NOT NULL
+                );
+            """
+            Storage.create_database_table(
+                session_db_instance, sql_create_chat_history_table
+            )
+
+            chat_meta_tuple = (
+                chat_id,
+                endpoint,
+                created_epoch,
+                created_datetime,
+            )
+
+            Storage.create_database_record(
+                session_db_instance,
+                chat_meta_tuple,
+                Chat.sql_create_chat_metadata_record,
+            )
+
+            self.chat_id = db_chat_id
+            self.endpoint = endpoint
+            self.chat_history: list[ChatRecord] = []
+
+    def to_dict(self) -> dict[str, Union[str, list[dict[str, str]]]]:
+        """
+        Converts the Chat instance into a dictionary.
+
+        This method iterates over the chat history, converting each ChatRecord instance into a dictionary
+        using its `to_dict` method. It then constructs a dictionary that includes the chat ID, endpoint,
+        and the list of chat history dictionaries.
+
+        Returns:
+            dict: A dictionary representation of the Chat instance, including chat ID, endpoint, and chat history.
+        """
+        list_of_chat_history_dict = [
+            chat_record.to_dict() for chat_record in self.chat_history
+        ]
+        return {
+            "chat_id": self.chat_id,
+            "endpoint": self.endpoint,
+            "chat_history": list_of_chat_history_dict,
+        }
+
+    @staticmethod
+    def load_chat_history(session_db_instance: DBInterface, chat_id: str) -> list:
+        """
+        Loads the chat history for a specific chat ID.
+
+        This method retrieves the chat history for a given chat ID by calling the StorageManager's method
+        to get the chat history for one endpoint. It then converts the chat record tuples into ChatRecord instances
+        and returns a list of ChatRecord objects.
+
+        Args:
+            db_instance: The database instance associated with the chat session.
+            chat_id (str): The unique identifier for the chat session.
+
+        Returns:
+            list[ChatRecord]: A list of ChatRecord instances representing the chat history.
+        """
+        sql_read_chat_history_for_one_endpoint = f"""SELECT * FROM {chat_id}"""
+        list_of_chat_record_tuples = Storage.read_database_records(
+            session_db_instance, sql_read_chat_history_for_one_endpoint
+        )
+
+        list_of_chat_records = []
+        if list_of_chat_record_tuples:
+            list_of_chat_records = [
+                ChatRecord(*chat_record_tuple).to_dict()
+                for chat_record_tuple in list_of_chat_record_tuples
+            ]
+        return list_of_chat_records
+
+    @staticmethod
+    def get_n_chat_history(
+        session_db_instance: DBInterface, endpoint_id: str, num_of_previous_chats: int
+    ) -> list[dict]:
+        """
+        Loads the chat history for a specific chat ID.
+
+        This method retrieves the chat history for a given chat ID by calling the StorageManager's method
+        to get the chat history for one endpoint. It then converts the chat record tuples into ChatRecord instances
+        and returns a list of ChatRecord objects.
+
+        Args:
+            db_instance: The database instance associated with the chat session.
+            chat_id (str): The unique identifier for the chat session.
+
+        Returns:
+            list[ChatRecord]: A list of ChatRecord instances representing the chat history.
+        """
+        endpoint_id = endpoint_id.replace("-", "_")
+        list_of_chat_record_tuples = Storage.read_database_records(
+            session_db_instance,
+            Chat.sql_select_n_chat_from_chat_table.format(
+                endpoint_id, num_of_previous_chats
+            ),
+        )
+        list_of_chat_records = []
+        if list_of_chat_record_tuples:
+            list_of_chat_records = [
+                ChatRecord(*chat_record_tuple).to_dict()
+                for chat_record_tuple in list_of_chat_record_tuples
+            ]
+        return list_of_chat_records

moonshot/src/redteaming/session/red_teaming_progress.py

@@ -0,0 +1,128 @@
+from typing import Callable
+
+from moonshot.src.runs.run_status import RunStatus
+
+
+class RedTeamingProgress:
+    DEFAULT_CHAT_BATCH_SIZE = 5
+
+    def __init__(
+        self,
+        runner_id: str,
+        red_teaming_arguments: dict,
+        run_progress_callback_func: Callable | None,
+    ):
+        # Information on the run and callback for progress updating
+        self.runner_id = runner_id
+        self.chat_batch_size = red_teaming_arguments.get(
+            "chat_batch_size", RedTeamingProgress.DEFAULT_CHAT_BATCH_SIZE
+        )
+        self.chats = red_teaming_arguments.get("chats", {})
+        self.current_count = 0
+        self.run_progress_callback_func = run_progress_callback_func
+        self.status = RunStatus.PENDING
+
+    def update_red_teaming_chats(
+        self, red_teaming_prompt_arguments: dict, run_status: RunStatus
+    ) -> None:
+        """
+        This method updates the red teaming chats with the provided arguments and run status.
+
+        It calculates the response time by adding the duration to the start time. Then, it creates a dictionary
+        with the prompt, response, prompt time, and response time. This dictionary is then added to the chats.
+
+        Args:
+            red_teaming_prompt_arguments (dict): A dictionary containing the arguments for the red teaming prompt.
+            It should contain all keys in AttackModule.RedTeamingPromptArguments, which are mainly:
+                - conn_id (str): The connection ID of the chat
+                - cs_id (str): The context strategy ID (if any) or ""
+                - pt_id (str): The prompt template ID (if any) or ""
+                - am_id (str): The attack module ID
+                - me_id (str): The metric ID (if any) or ""
+                - original_prompt (str): The original prompt entered by the user
+                - prepared_prompt(str): The modified and final prompt that was sent to the LLM
+                - system_prompt (str): The system prompt entered by the user (if any)
+                - response (str): The response from the LLM
+                - duration (str): The amount of time it takes to get back the response from the LLM in seconds
+                (in string)
+                - start_time (str): The datetime of the prompt (in string)
+
+            - run_status (RunStatus): The current status of the run.
+
+        Returns:
+            None
+        """
+        prompt_response_dict = {
+            "conn_id": red_teaming_prompt_arguments["conn_id"],
+            "context_strategy": red_teaming_prompt_arguments["cs_id"],
+            "prompt_template": red_teaming_prompt_arguments["pt_id"],
+            "attack_module": red_teaming_prompt_arguments["am_id"],
+            "metric": red_teaming_prompt_arguments["me_id"],
+            "prompt": red_teaming_prompt_arguments["original_prompt"],
+            "prepared_prompt": red_teaming_prompt_arguments["prepared_prompt"],
+            "system_prompt": red_teaming_prompt_arguments["system_prompt"],
+            "predicted_result": red_teaming_prompt_arguments["response"],
+            "duration": red_teaming_prompt_arguments["duration"],
+            "prompt_time": red_teaming_prompt_arguments["start_time"],
+        }
+
+        if red_teaming_prompt_arguments["conn_id"] not in self.chats:
+            self.chats[red_teaming_prompt_arguments["conn_id"]] = []
+        self.chats[red_teaming_prompt_arguments["conn_id"]].append(prompt_response_dict)
+        self.status = run_status
+
+    def reset_chats(self) -> None:
+        """
+        This method clears all the chat data stored in the 'chats' attribute.
+        """
+        self.chats.clear()
+
+    def update_red_teaming_progress(self) -> None:
+        """
+        This method updates the progress of the red teaming session.
+
+        It checks if the current count of chats is equal to or greater than the batch size. If it is, it triggers
+        a callback to notify the progress, resets the chats for the next batch, and resets the current count to zero.
+
+        Regardless of the condition, it increments the current count by one.
+
+        Args:
+            None
+
+        Returns:
+            None
+        """
+        if self.current_count >= self.chat_batch_size:
+            self.notify_progress()
+            self.reset_chats()
+            self.current_count = 0
+
+        self.current_count += 1
+
+    def notify_progress(self) -> None:
+        """
+        This method checks if a callback function for run progress exists and if so,
+        it calls the function with the current state of the red teaming progress.
+        """
+        if self.run_progress_callback_func:
+            self.run_progress_callback_func(self.get_dict())
+
+    def get_dict(self) -> dict:
+        """
+        This method returns a dictionary representation of the current state of the red teaming progress.
+
+        The dictionary includes the following keys:
+        - "current_runner_id": The ID of the current runner.
+        - "current_chats": The chats that will be returned during a callback.
+        - "current_batch_size": The current batch size, which indicates the number of chats returned during a callback.
+        - "current_status": The current status of the run.
+
+        Returns:
+            dict: A dictionary representation of the current state of the red teaming progress.
+        """
+        return {
+            "current_runner_id": self.runner_id,
+            "current_chats": self.chats,
+            "current_batch_size": self.chat_batch_size,
+            "current_status": self.status.name,
+        }

moonshot/src/redteaming/session/red_teaming_type.py

@@ -0,0 +1,6 @@
+from enum import Enum
+
+
+class RedTeamingType(Enum):
+    MANUAL = "manual"
+    AUTOMATED = "automated"