aiverify-moonshot 0.4.10__py3-none-any.whl → 0.5.0__py3-none-any.whl

moonshot/src/messages_constants.py

@@ -15,7 +15,9 @@ 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_FAIL = (
+    "[Bookmark] Bookmark record not found. Unable to delete."
+)
 BOOKMARK_DELETE_BOOKMARK_ERROR = (
     "[Bookmark] Failed to delete bookmark record: {message}"
 )
@@ -33,9 +35,89 @@ BOOKMARK_DELETE_ALL_BOOKMARK_ERROR = (
 # BOOKMARK - export_bookmarks
 # ------------------------------------------------------------------------------
 BOOKMARK_EXPORT_BOOKMARK_ERROR = "[Bookmark] Failed to export bookmarks: {message}"
-BOOKMARK_EXPORT_BOOKMARK_VALIDATION_ERROR = "Export filename must be a non-empty string"
+BOOKMARK_EXPORT_BOOKMARK_VALIDATION_ERROR = (
+    "Export filename must be a non-empty string."
+)
 
 # ------------------------------------------------------------------------------
 # BOOKMARK ARGUMENTS - from_tuple_to_dict
 # ------------------------------------------------------------------------------
-BOOKMARK_ARGUMENTS_FROM_TUPLE_TO_DICT_VALIDATION_ERROR = "[BookmarkArguments] Failed to convert to dictionary because of the insufficient number of values"  # noqa: E501
+BOOKMARK_ARGUMENTS_FROM_TUPLE_TO_DICT_VALIDATION_ERROR = "[BookmarkArguments] Failed to convert to dictionary because of the insufficient number of values."  # noqa: E501
+
+# ------------------------------------------------------------------------------
+# CONNECTOR - perform_retry_callback
+# ------------------------------------------------------------------------------
+CONNECTOR_PERFORM_RETRY_CALLBACK_ERROR = "[Connector ID: {connector_id}] Attempt {attempt_no} failed due to error: {message}"  # noqa: E501
+
+# ------------------------------------------------------------------------------
+# CONNECTOR - load
+# ------------------------------------------------------------------------------
+CONNECTOR_LOAD_CONNECTOR_ENDPOINT_ARGUMENTS_VALIDATION_ERROR = "[Connector] The 'ep_args' argument must be an instance of ConnectorEndpointArguments and not None."  # noqa: E501
+CONNECTOR_LOAD_CONNECTOR_INSTANCE_RUNTIME_ERROR = (
+    "[Connector] Failed to get connector instance: {message}"
+)
+
+# ------------------------------------------------------------------------------
+# CONNECTOR - create
+# ------------------------------------------------------------------------------
+CONNECTOR_CREATE_CONNECTOR_ENDPOINT_ARGUMENTS_VALIDATION_ERROR = "[Connector] The 'ep_args' argument must be an instance of ConnectorEndpointArguments and not None."  # noqa: E501
+CONNECTOR_CREATE_ERROR = "[Connector] Failed to create connector: {message}"
+
+# ------------------------------------------------------------------------------
+# CONNECTOR - get_available_items
+# ------------------------------------------------------------------------------
+CONNECTOR_GET_AVAILABLE_ITEMS_ERROR = (
+    "[Connector] Failed to get available connectors: {message}"
+)
+
+# ------------------------------------------------------------------------------
+# CONNECTOR - get_prediction
+# ------------------------------------------------------------------------------
+CONNECTOR_GET_PREDICTION_ARGUMENTS_GENERATED_PROMPT_VALIDATION_ERROR = "[Connector] The 'generated_prompt' argument must be an instance of ConnectorPromptArguments and not None."  # noqa: E501
+CONNECTOR_GET_PREDICTION_ARGUMENTS_CONNECTOR_VALIDATION_ERROR = "[Connector] The 'connector' argument must be an instance of Connector and not None."  # noqa: E501
+CONNECTOR_GET_PREDICTION_INFO = (
+    "[Connector ID: {connector_id}] Predicting Prompt Index {prompt_index}."
+)
+CONNECTOR_GET_PREDICTION_TIME_TAKEN_INFO = "[Connector ID: {connector_id}] Prompt Index {prompt_index} took {prompt_duration}s."  # noqa: E501
+CONNECTOR_GET_PREDICTION_ERROR = "[Connector ID: {connector_id}] Prompt Index {prompt_index} failed to get prediction: {message}"  # noqa: E501
+
+# ------------------------------------------------------------------------------
+# CONNECTOR - set_system_prompt
+# ------------------------------------------------------------------------------
+CONNECTOR_SET_SYSTEM_PROMPT_VALIDATION_ERROR = "[Connector] The 'system_prompt' argument must be an instance of string and not None."  # noqa: E501
+
+# ------------------------------------------------------------------------------
+# CONNECTOR ENDPOINT - create
+# ------------------------------------------------------------------------------
+CONNECTOR_ENDPOINT_CREATE_ERROR = (
+    "[ConnectorEndpoint] Failed to create connector endpoint: {message}"
+)
+
+# ------------------------------------------------------------------------------
+# CONNECTOR ENDPOINT - read
+# ------------------------------------------------------------------------------
+CONNECTOR_ENDPOINT_READ_INVALID = "Invalid connector endpoint id - {ep_id}"
+CONNECTOR_ENDPOINT_READ_ERROR = (
+    "[ConnectorEndpoint] Failed to read connector endpoint: {message}"
+)
+
+# ------------------------------------------------------------------------------
+# CONNECTOR ENDPOINT - update
+# ------------------------------------------------------------------------------
+CONNECTOR_ENDPOINT_UPDATE_ERROR = (
+    "[ConnectorEndpoint] Failed to update connector endpoint: {message}"
+)
+
+# ------------------------------------------------------------------------------
+# CONNECTOR ENDPOINT - delete
+# ------------------------------------------------------------------------------
+CONNECTOR_ENDPOINT_DELETE_ERROR = (
+    "[ConnectorEndpoint] Failed to delete connector endpoint: {message}"
+)
+
+# ------------------------------------------------------------------------------
+# CONNECTOR ENDPOINT - get_available_items
+# ------------------------------------------------------------------------------
+CONNECTOR_ENDPOINT_GET_AVAILABLE_ITEMS_ERROR = (
+    "[ConnectorEndpoint] Failed to get available connector endpoints: {message}"
+)

moonshot/src/redteaming/attack/attack_module.py

@@ -31,6 +31,22 @@ class AttackModule:
         INSERT INTO {} (connection_id,context_strategy,prompt_template,attack_module,
         metric,prompt,prepared_prompt,system_prompt,predicted_result,duration,prompt_time)VALUES(?,?,?,?,?,?,?,?,?,?,?)
         """
+    sql_create_chat_history_table = """
+        CREATE TABLE IF NOT EXISTS {} (
+        id INTEGER PRIMARY KEY AUTOINCREMENT,
+        connection_id text NOT NULL,
+        context_strategy text,
+        prompt_template text,
+        attack_module text,
+        metric text,
+        prompt text NOT NULL,
+        prepared_prompt text NOT NULL,
+        system_prompt text,
+        predicted_result text NOT NULL,
+        duration text NOT NULL,
+        prompt_time text NOT NULL
+        );
+    """
 
     def __init__(self, am_id: str, am_arguments: AttackModuleArguments | None = None):
         self.id = am_id
@@ -261,6 +277,15 @@ class AttackModule:
             list: A list of consolidated responses from the specified LLM connector.
         """
         consolidated_responses = []
+
+        # perform a check to see if the target endpoint has its table in the db. if not, create one
+        endpoint_id = target_llm_connector.id.replace("-", "_")
+        if not Storage.check_database_table_exists(self.db_instance, endpoint_id):
+            Storage.create_database_table(
+                self.db_instance,
+                AttackModule.sql_create_chat_history_table.format(endpoint_id),
+            )
+
         for prepared_prompt in list_of_prompts:
             if self.cancel_event.is_set():
                 logger.warning(
@@ -316,6 +341,7 @@ class AttackModule:
             chat_record_tuple (tuple): A tuple containing the chat record information.
             chat_record_id (str): The ID of the chat record.
         """
+
         endpoint_id = chat_record_id.replace("-", "_")
         Storage.create_database_record(
             self.db_instance,
@@ -623,7 +649,8 @@ class RedTeamingPromptArguments(BaseModel):
 
         This method collects all the attributes of the RedTeamingPromptArguments instance and forms a tuple
         with the attribute values in this specific order: conn_id, cs_id, pt_id, am_id, me_id, original_prompt,
-        connector_prompt.prompt, connector_prompt.predicted_results, connector_prompt.duration, start_time.
+        connector_prompt.prompt, system_prompt, connector_prompt.predicted_results.response, 
+        connector_prompt.duration, start_time.
 
         Returns:
             tuple: A tuple representation of the RedTeamingPromptArguments instance.
@@ -637,21 +664,21 @@ class RedTeamingPromptArguments(BaseModel):
             self.original_prompt,
             self.connector_prompt.prompt,
             self.system_prompt,
-            str(self.connector_prompt.predicted_results),
+            self.connector_prompt.predicted_results.response if self.connector_prompt.predicted_results else "",
             str(self.connector_prompt.duration),
             self.start_time,
         )
 
     def to_dict(self) -> dict:
         """
-        Converts the RedTeamingPromptArguments instance into a dict.
+        Converts the RedTeamingPromptArguments instance into a dictionary.
 
-        This method collects all the attributes of the RedTeamingPromptArguments instance and forms a dict
-        with the keys: conn_id, cs_id, pt_id, am_id, me_id, original_prompt, prepared_prompt, system_prompt
+        This method collects all the attributes of the RedTeamingPromptArguments instance and forms a dictionary
+        with the keys: conn_id, cs_id, pt_id, am_id, me_id, original_prompt, system_prompt, prepared_prompt,
         response, duration, start_time.
 
         Returns:
-            dict: A dict representation of the RedTeamingPromptArguments instance.
+            dict: A dictionary representation of the RedTeamingPromptArguments instance.
         """
         return {
             "conn_id": self.conn_id,
@@ -662,7 +689,7 @@ class RedTeamingPromptArguments(BaseModel):
             "original_prompt": self.original_prompt,
             "prepared_prompt": self.connector_prompt.prompt,
             "system_prompt": self.system_prompt,
-            "response": str(self.connector_prompt.predicted_results),
+            "response": self.connector_prompt.predicted_results.response if self.connector_prompt.predicted_results else "",
             "duration": str(self.connector_prompt.duration),
             "start_time": self.start_time,
         }

moonshot/src/runs/run_arguments.py

@@ -1,6 +1,7 @@
 from __future__ import annotations
 
 import ast
+import json
 
 from pydantic import BaseModel
 
@@ -83,20 +84,7 @@ class RunArguments(BaseModel):
         Returns:
             tuple: A tuple containing the run arguments ready for database insertion.
         """
-        return (
-            self.runner_id,
-            self.runner_type.name.lower(),
-            str(self.runner_args),
-            str(self.endpoints),
-            self.results_file,
-            self.start_time,
-            self.end_time,
-            self.duration,
-            str(self.error_messages),
-            str(self.raw_results),
-            str(self.results),
-            self.status.name.lower(),
-        )
+        return self._to_tuple(include_run_id=False)
 
     def to_tuple(self) -> tuple:
         """
@@ -112,7 +100,42 @@ class RunArguments(BaseModel):
         Returns:
             tuple: A tuple containing the serialized attributes of the RunArguments instance.
         """
-        return (
+        return self._to_tuple(include_run_id=True)
+
+    def _to_tuple(self, include_run_id: bool) -> tuple:
+        """
+        Converts the RunArguments object to a tuple format.
+
+        This method serializes the RunArguments instance into a tuple, including or excluding the run_id based on the
+        include_run_id parameter. It ensures that all dictionary keys are converted to strings and JSON-encodes
+        the raw_results and results attributes.
+
+        Args:
+            include_run_id (bool): A flag indicating whether to include the run_id in the resulting tuple.
+
+        Returns:
+            tuple: A tuple containing the serialized attributes of the RunArguments instance.
+        """
+
+        def convert_keys_to_str(d):
+            """
+            Recursively converts all keys in a dictionary to strings.
+
+            Args:
+                d (dict): The dictionary whose keys need to be converted.
+
+            Returns:
+                dict: A new dictionary with all keys converted to strings.
+            """
+
+            def convert_value(v):
+                if isinstance(v, dict):
+                    return {str(k): convert_value(val) for k, val in v.items()}
+                return v
+
+            return {str(k): convert_value(v) for k, v in d.items()}
+
+        base_tuple = (
             self.runner_id,
             self.runner_type.name.lower(),
             str(self.runner_args),
@@ -122,11 +145,11 @@ class RunArguments(BaseModel):
             self.end_time,
             self.duration,
             str(self.error_messages),
-            str(self.raw_results),
-            str(self.results),
+            json.dumps(convert_keys_to_str(self.raw_results)),
+            json.dumps(self.results),
             self.status.name.lower(),
-            self.run_id,
         )
+        return base_tuple + (self.run_id,) if include_run_id else base_tuple
 
     @classmethod
     def from_tuple(cls, run_record: tuple) -> RunArguments:
@@ -156,7 +179,7 @@ class RunArguments(BaseModel):
             end_time=float(run_record[7]),
             duration=run_record[8],
             error_messages=ast.literal_eval(run_record[9]),
-            raw_results=ast.literal_eval(run_record[10]),
-            results=ast.literal_eval(run_record[11]),
+            raw_results=json.loads(run_record[10]),
+            results=json.loads(run_record[11]),
             status=RunStatus(run_record[12]),
         )