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

moonshot/src/api/api_recipe.py

@@ -79,7 +79,7 @@ def api_read_recipes(rec_ids: conlist(str, min_length=1)) -> list[dict]:
     and returns a list of dictionaries containing each recipe's information.
 
     Args:
-        rec_ids (list[str]): The IDs of the recipes.
+        rec_ids (conlist(str, min_length=1)): The IDs of the recipes.
 
     Returns:
         list[dict]: A list of dictionaries, each containing a recipe's information.
@@ -155,7 +155,7 @@ def api_get_all_recipe() -> list[dict]:
     """
     Retrieves all available recipes.
 
-    This function calls the get_available_recipes method to retrieve all available recipes. It then converts each
+    This function calls the get_available_items method to retrieve all available recipes. It then converts each
     recipe into a dictionary using the to_dict method and returns a list of these dictionaries.
 
     Returns:
@@ -169,7 +169,7 @@ def api_get_all_recipe_name() -> list[str]:
     """
     Retrieves all available recipe names.
 
-    This function calls the get_available_recipes method to retrieve all available recipes. It then extracts the names
+    This function calls the get_available_items method to retrieve all available recipes. It then extracts the names
     of each recipe and returns a list of these names.
 
     Returns:

moonshot/src/api/api_red_teaming.py

@@ -11,9 +11,7 @@ def api_get_all_attack_modules() -> list[str]:
     Retrieves all available attack module IDs.
 
     This function calls the `get_available_items` method from the `AttackModule` class to retrieve all available
-    attack modules.
-
-    It then extracts the IDs of each attack module and returns a list of these IDs.
+    attack modules. It then extracts the IDs of each attack module and returns a list of these IDs.
 
     Returns:
         list[str]: A list of strings, each representing an attack module ID.
@@ -27,10 +25,8 @@ def api_get_all_attack_module_metadata() -> list[dict]:
     Retrieves metadata for all available attack modules.
 
     This function calls the `get_available_items` method from the `AttackModule` class to retrieve all available
-    attack modules metadata.
-
-    It then extracts the metadata for each attack module and returns a list of dictionaries, each containing the
-    metadata of an attack module.
+    attack modules. It then extracts the metadata for each attack module and returns a list of dictionaries,
+    each containing the metadata of an attack module.
 
     Returns:
         list[dict]: A list of dictionaries, each representing the metadata of an attack module.
@@ -44,7 +40,7 @@ def api_delete_attack_module(am_id: str) -> bool:
     """
     Deletes an attack module by its identifier.
 
-    This function takes an attack module ID as input and calls the delete method from the AttackModule class
+    This function takes an attack module ID as input and calls the `delete` method from the `AttackModule` class
     to remove the specified attack module from storage.
 
     Args:

moonshot/src/api/api_result.py

@@ -11,7 +11,7 @@ def api_read_result(res_id: str) -> dict:
     """
     Reads a result and returns its information.
 
-    This function takes a result ID as input, reads the corresponding database file from the storage manager,
+    This function takes a result ID as input, reads the corresponding result from the storage manager,
     and returns a dictionary containing the result's information.
 
     Args:
@@ -28,7 +28,7 @@ def api_read_results(res_ids: conlist(str, min_length=1)) -> list[dict]:
     """
     Reads multiple results and returns their information.
 
-    This function takes a list of result IDs as input, reads the corresponding database files from the storage manager,
+    This function takes a list of result IDs as input, reads the corresponding results from the storage manager,
     and returns a list of dictionaries, each containing a result's information.
 
     Args:
@@ -37,7 +37,6 @@ def api_read_results(res_ids: conlist(str, min_length=1)) -> list[dict]:
     Returns:
         list[dict]: A list of dictionaries, each containing a result's information.
     """
-
     return [Result.read(res_id) for res_id in res_ids]
 
 
@@ -63,11 +62,13 @@ def api_delete_result(res_id: str) -> bool:
 
 def api_get_all_result() -> list[dict]:
     """
-    This function retrieves all available results and returns them as a list of dictionaries. Each dictionary
-    represents a result and contains its information.
+    This function retrieves all available results and returns them as a list of dictionaries.
+
+    This function calls the get_available_items method from the Result class to retrieve all available results.
+    It then returns a list of dictionaries, each containing the details of a result.
 
     Returns:
-        list[dict]: A list of dictionaries, each representing a result.
+        list[dict]: A list of dictionaries, each representing a result's details.
     """
     _, results = Result.get_available_items()
     return [result for result in results]
@@ -77,6 +78,9 @@ def api_get_all_result_name() -> list[str]:
     """
     This function retrieves all available result names and returns them as a list.
 
+    This function calls the get_available_items method from the Result class to retrieve all available results.
+    It then extracts the names of each result and returns a list of these names.
+
     Returns:
         list[str]: A list of result names.
     """

moonshot/src/api/api_run.py

@@ -14,7 +14,7 @@ def api_get_all_run(runner_id: str = "") -> list[dict]:
 
     Args:
         runner_id (str, optional): The ID of the runner to retrieve runs for. If empty, runs for all runners
-        are retrieved.
+                                   are retrieved.
 
     Returns:
         list[dict]: A list of dictionaries, each representing a run's data.
@@ -34,8 +34,8 @@ def _api_get_available_runs(
     RunArguments instances representing the runs.
 
     Args:
-        runner_id (str): The ID of the runner for which to retrieve run information. If empty, information for
-                         all runners is retrieved.
+        runner_id (str, optional): The ID of the runner for which to retrieve run information. If empty, information
+                                   for all runners is retrieved.
 
     Returns:
         tuple[list[str], list[RunArguments]]: A tuple containing a list of runner IDs and a list of RunArguments

moonshot/src/api/api_runner.py

@@ -20,7 +20,7 @@ def api_create_runner(
     Creates a new runner.
 
     This function takes the name, endpoints, and an optional progress callback function to create a new Runner instance.
-    The id of the runner is generated from the name of the runner using the slugify function,
+    The ID of the runner is generated from the name of the runner using the slugify function,
     so it does not need to be provided.
 
     Args:
@@ -59,7 +59,8 @@ def api_load_runner(
 
     Args:
         runner_id (str): The ID of the runner to be loaded.
-        progress_callback_func (Callable | None): The progress callback function to be used by the runner.
+        progress_callback_func (Callable | None, optional): An optional progress callback function for the runner.
+        Defaults to None.
 
     Returns:
         Runner: An initialized Runner object.
@@ -108,8 +109,8 @@ def api_get_all_runner() -> list[dict]:
     """
     Retrieves all available runners.
 
-    This function calls the get_available_items method to retrieve all available runners. It then converts each
-    runner into a dictionary using the to_dict method and returns a list of these dictionaries.
+    This function calls the get_available_items method from the Runner class to retrieve all available runners.
+    It then converts each runner into a dictionary using the to_dict method and returns a list of these dictionaries.
 
     Returns:
         list[dict]: A list of dictionaries, each representing a runner.
@@ -122,8 +123,8 @@ def api_get_all_runner_name() -> list[str]:
     """
     Retrieves all available runner names.
 
-    This function calls the get_available_items method to retrieve all available runners. It then extracts the names of
-    each runner and returns a list of these names.
+    This function calls the get_available_items method from the Runner class to retrieve all available runners.
+    It then extracts the names of each runner and returns a list of these names.
 
     Returns:
         list[str]: A list of runner names.

moonshot/src/api/api_session.py

@@ -43,11 +43,14 @@ def api_create_session(
     Args:
         runner_id (str): The unique identifier of the runner for which the session is to be created.
         database_instance (Any): The database instance to be used for the session.
-        endpoints (list): A list of endpoints for the session.
+        endpoints (list[str]): A list of endpoints for the session.
         runner_args (dict): A dictionary of arguments for the runner.
 
     Returns:
-        Session
+        Session: A new Session object.
+
+    Raises:
+        RuntimeError: If the runner_id is empty or if the database_instance is not provided.
     """
     if isinstance(database_instance, DBInterface):
         if runner_id:
@@ -86,16 +89,16 @@ def api_get_all_session_names() -> list[str]:
     return session_names
 
 
-def api_get_available_session_info() -> tuple[list, list]:
+def api_get_available_session_info() -> tuple[list[str], list[dict]]:
     """
-    Retrieves the IDs and database instances of runners with active sessions.
+    Retrieves the IDs and metadata of runners with active sessions.
 
-    This function retrieves the IDs and database instances of runners with active sessions by querying all runners
+    This function retrieves the IDs and metadata of runners with active sessions by querying all runners
     and checking if each runner has an active session. It returns a tuple containing a list of runner IDs and a list
     of corresponding session metadata for runners with active sessions.
 
     Returns:
-        tuple[list[str], list[str]]: A tuple containing a list of runner IDs and a list of corresponding session
+        tuple[list[str], list[dict]]: A tuple containing a list of runner IDs and a list of corresponding session
         metadata for runners with active sessions.
     """
     runners_info = api_get_all_runner()
@@ -122,7 +125,8 @@ def api_get_all_session_metadata() -> list[dict]:
     This function retrieves the metadata for all active sessions by calling the `api_get_available_session_info` method.
 
     Returns:
-        list: A list containing the metadata for all active sessions, sorted by created datetime in descending order.
+        list[dict]: A list containing the metadata for all active sessions, sorted by created datetime in
+                    descending order.
     """
     _, session_metadata_list = api_get_available_session_info()
     return sorted(

moonshot/src/connectors/connector.py

@@ -3,17 +3,32 @@ from __future__ import annotations
 import asyncio
 import time
 from abc import abstractmethod
-from asyncio import sleep
-from functools import wraps
+from functools import partial, wraps
 from pathlib import Path
 from typing import Callable
 
+from tenacity import RetryCallState, retry, stop_after_attempt, wait_random_exponential
+
 from moonshot.src.configs.env_variables import EnvVariables
 from moonshot.src.connectors.connector_prompt_arguments import ConnectorPromptArguments
 from moonshot.src.connectors.connector_response import ConnectorResponse
 from moonshot.src.connectors_endpoints.connector_endpoint_arguments import (
     ConnectorEndpointArguments,
 )
+from moonshot.src.messages_constants import (
+    CONNECTOR_CREATE_CONNECTOR_ENDPOINT_ARGUMENTS_VALIDATION_ERROR,
+    CONNECTOR_CREATE_ERROR,
+    CONNECTOR_GET_AVAILABLE_ITEMS_ERROR,
+    CONNECTOR_GET_PREDICTION_ARGUMENTS_CONNECTOR_VALIDATION_ERROR,
+    CONNECTOR_GET_PREDICTION_ARGUMENTS_GENERATED_PROMPT_VALIDATION_ERROR,
+    CONNECTOR_GET_PREDICTION_ERROR,
+    CONNECTOR_GET_PREDICTION_INFO,
+    CONNECTOR_GET_PREDICTION_TIME_TAKEN_INFO,
+    CONNECTOR_LOAD_CONNECTOR_ENDPOINT_ARGUMENTS_VALIDATION_ERROR,
+    CONNECTOR_LOAD_CONNECTOR_INSTANCE_RUNTIME_ERROR,
+    CONNECTOR_PERFORM_RETRY_CALLBACK_ERROR,
+    CONNECTOR_SET_SYSTEM_PROMPT_VALIDATION_ERROR,
+)
 from moonshot.src.storage.storage import Storage
 from moonshot.src.utils.import_modules import get_instance
 from moonshot.src.utils.log import configure_logger
@@ -22,43 +37,56 @@ from moonshot.src.utils.log import configure_logger
 logger = configure_logger(__name__)
 
 
+def perform_retry_callback(connector_id: str, retry_state: RetryCallState) -> None:
+    """
+    Callback function to log retry attempts with detailed information.
+
+    This function is called by the tenacity library before each retry attempt.
+    It logs the retry attempt number, the sleep time before the next attempt,
+    the error message that caused the retry, and the connector ID.
+
+    Args:
+        connector_id (str): The ID of the connector.
+        retry_state (RetryCallState): The state of the retry call, which includes
+            information about the current attempt, the exception raised, and the next action.
+    """
+    sleep_time = retry_state.idle_for if retry_state else 0
+    exception = (
+        retry_state.outcome.exception() if retry_state.outcome else "Unknown exception"
+    )
+    logger.error(
+        CONNECTOR_PERFORM_RETRY_CALLBACK_ERROR.format(
+            connector_id=connector_id,
+            attempt_no=retry_state.attempt_number,
+            sleep=f"{sleep_time:.2f}",
+            message=str(exception),
+        )
+    )
+
+
 def perform_retry(func):
     """
-    A decorator to perform retries on a function.
+    A decorator to perform retries on a function using tenacity.
 
-    This decorator wraps a function to enable retrying the function call
-    if it fails. The number of retries and the delay between retries
-    are determined by the `retries_times` and `allow_retries`
-    attributes of the class instance.
+    This decorator wraps an asynchronous function to enable retrying the function call
+    if it fails. The number of attempts and the delay between attempts
+    are determined by the `max_attempts` attribute of the class instance.
 
-    Parameters:
-    - func (Callable): The function to be wrapped and retried.
+    Args:
+        func (Callable): The asynchronous function to be wrapped and retried.
 
     Returns:
-    - Callable: A wrapper function that includes retry logic.
+        Callable: A wrapper function that includes retry logic.
     """
 
     async def wrapper(self, *args, **kwargs):
-        if self.allow_retries:
-            retry_count = 0
-            base_delay = 1
-            while retry_count <= self.retries_times:
-                # Perform the request
-                try:
-                    return await func(self, *args, **kwargs)
-                except Exception as exc:
-                    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)
-                    logger.warning(
-                        f"Attempt {retry_count}, Retrying in {delay} seconds..."
-                    )
-                    await sleep(delay)
-            # Raise an exception
-            raise ConnectionError("Failed to get response.")
+        retry_decorator = retry(
+            wait=wait_random_exponential(min=1, max=60),
+            stop=stop_after_attempt(self.max_attempts),
+            after=partial(perform_retry_callback, self.id),
+            reraise=True,
+        )
+        return await retry_decorator(func)(self, *args, **kwargs)
 
     return wrapper
 
@@ -71,6 +99,7 @@ class Connector:
         self.token = ep_args.token
         self.max_concurrency = ep_args.max_concurrency
         self.max_calls_per_second = ep_args.max_calls_per_second
+        self.model = ep_args.model
         self.params = ep_args.params
 
         # Rate limiting
@@ -80,21 +109,19 @@ class Connector:
         self.updated_at = time.time()
         self.semaphore = asyncio.Semaphore(ep_args.max_concurrency)
 
-        # Set Prompts if they exists
+        # Set Prompts if they exist
         self.pre_prompt = ep_args.params.get("pre_prompt", "")
         self.post_prompt = ep_args.params.get("post_prompt", "")
         self.system_prompt = ep_args.params.get("system_prompt", "")
 
         # Connection timeout
         self.timeout = ep_args.params.get("timeout", 600)
-        self.allow_retries = ep_args.params.get("allow_retries", True)
-        self.retries_times = ep_args.params.get("num_of_retries", 3)
+        self.max_attempts = ep_args.params.get("max_attempts", 3)
 
         # Optional params
         excluded_keys = {
-            "allow_retries",
             "timeout",
-            "num_of_retries",
+            "max_attempts",
             "pre_prompt",
             "post_prompt",
             "system_prompt",
@@ -114,10 +141,6 @@ class Connector:
 
         The method updates the `updated_at` attribute to the current time after tokens are added, ensuring that
         the next token addition will be calculated based on the correct elapsed time.
-
-        Usage:
-            # Assume `self` is an instance of a class with `rate_limiter`, `tokens`, and `updated_at` attributes.
-            await self.add_tokens()
         """
         now = time.time()
         elapsed = now - self.updated_at
@@ -144,11 +167,6 @@ class Connector:
 
         Returns:
             Callable: The decorated function wrapped with rate limiting logic.
-
-        Usage:
-            @rate_limited
-            async def some_async_function(*args, **kwargs):
-                # Function implementation
         """
 
         @wraps(func)
@@ -187,7 +205,7 @@ class Connector:
     @classmethod
     def load(cls, ep_args: ConnectorEndpointArguments) -> Connector:
         """
-        This method dynamically loads a connector instance based on the provided endpoint arguments.
+        Dynamically loads a connector instance based on the provided endpoint arguments.
 
         The connector type specified in the `ep_args` is used to dynamically load the corresponding
         connector class. The connector is then instantiated with the provided endpoint arguments. If the
@@ -203,17 +221,24 @@ class Connector:
         Raises:
             RuntimeError: If the specified connector type does not match any available connector classes.
         """
+        if ep_args is None or not isinstance(ep_args, ConnectorEndpointArguments):
+            raise ValueError(
+                CONNECTOR_LOAD_CONNECTOR_ENDPOINT_ARGUMENTS_VALIDATION_ERROR
+            )
+
         connector_instance = get_instance(
             ep_args.connector_type,
             Storage.get_filepath(
                 EnvVariables.CONNECTORS.name, ep_args.connector_type, "py"
             ),
         )
-        if connector_instance:
+        if connector_instance and isinstance(connector_instance, Callable):
             return connector_instance(ep_args)
         else:
             raise RuntimeError(
-                f"Unable to get defined connector instance - {ep_args.connector_type}"
+                CONNECTOR_LOAD_CONNECTOR_INSTANCE_RUNTIME_ERROR.format(
+                    message=ep_args.connector_type
+                )
             )
 
     @staticmethod
@@ -223,7 +248,7 @@ class Connector:
 
         This method takes a ConnectorEndpointArguments object, which contains the necessary information
         to initialize and return a Connector object. The Connector object is created by calling the
-        `load_connector` method, which dynamically loads and initializes the connector based on the
+        `load` method, which dynamically loads and initializes the connector based on the
         endpoint arguments provided.
 
         Args:
@@ -231,12 +256,20 @@ class Connector:
 
         Returns:
             Connector: An initialized Connector object based on the provided endpoint arguments.
+
+        Raises:
+            ValueError: If the provided endpoint arguments are invalid.
+            Exception: If there is an error during the creation of the connector.
         """
         try:
+            if ep_args is None or not isinstance(ep_args, ConnectorEndpointArguments):
+                raise ValueError(
+                    CONNECTOR_CREATE_CONNECTOR_ENDPOINT_ARGUMENTS_VALIDATION_ERROR
+                )
             return Connector.load(ep_args)
 
         except Exception as e:
-            logger.error(f"Failed to create connector: {str(e)}")
+            logger.error(CONNECTOR_CREATE_ERROR.format(message=str(e)))
             raise e
 
     @staticmethod
@@ -245,7 +278,7 @@ class Connector:
         Fetches a list of all available connector types.
 
         This method employs the `get_connectors` method to locate all Python files in the directory
-        defined by the `EnvironmentVars.CONNECTORS` environment variable. It subsequently excludes any files that are
+        defined by the `EnvVariables.CONNECTORS` environment variable. It subsequently excludes any files that are
         not intended to be exposed as connectors (those containing "__" in their names). The method yields a list of the
         names of these connector types.
 
@@ -263,7 +296,7 @@ class Connector:
             ]
 
         except Exception as e:
-            logger.error(f"Failed to get available connectors: {str(e)}")
+            logger.error(CONNECTOR_GET_AVAILABLE_ITEMS_ERROR.format(message=str(e)))
             raise e
 
     @staticmethod
@@ -299,9 +332,24 @@ class Connector:
         Raises:
             Exception: If there is an error during prediction.
         """
+        if generated_prompt is None or not isinstance(
+            generated_prompt, ConnectorPromptArguments
+        ):
+            raise ValueError(
+                CONNECTOR_GET_PREDICTION_ARGUMENTS_GENERATED_PROMPT_VALIDATION_ERROR
+            )
+
+        if connector is None or not isinstance(connector, Connector):
+            raise ValueError(
+                CONNECTOR_GET_PREDICTION_ARGUMENTS_CONNECTOR_VALIDATION_ERROR
+            )
+
         try:
             logger.info(
-                f"Predicting prompt {generated_prompt.prompt_index} [{connector.id}]"
+                CONNECTOR_GET_PREDICTION_INFO.format(
+                    connector_id=connector.id,
+                    prompt_index=generated_prompt.prompt_index,
+                )
             )
 
             start_time = time.perf_counter()
@@ -310,7 +358,11 @@ class Connector:
             )
             generated_prompt.duration = time.perf_counter() - start_time
             logger.debug(
-                f"[Prompt {generated_prompt.prompt_index}] took {generated_prompt.duration:.4f}s"
+                CONNECTOR_GET_PREDICTION_TIME_TAKEN_INFO.format(
+                    connector_id=connector.id,
+                    prompt_index=generated_prompt.prompt_index,
+                    prompt_duration=f"{generated_prompt.duration:.4f}",
+                )
             )
 
             # Call prompt callback
@@ -321,17 +373,28 @@ class Connector:
             return generated_prompt
 
         except Exception as e:
-            logger.error(f"Failed to get prediction: {str(e)}")
+            logger.error(
+                CONNECTOR_GET_PREDICTION_ERROR.format(
+                    connector_id=connector.id,
+                    prompt_index=generated_prompt.prompt_index,
+                    message=str(e),
+                )
+            )
             raise e
 
     def set_system_prompt(self, system_prompt: str) -> None:
         """
-        Assigns a new system prompt to this connector instance.
+        Sets a new system prompt for this connector instance.
 
-        The system prompt serves as a preconfigured command or message that the connector can use to initiate
-        interactions or execute specific operations.
+        The system prompt is a predefined message or command that the connector can use to start interactions
+        or perform specific tasks.
 
-        Parameters:
+        Args:
             system_prompt (str): The new system prompt to set for this connector.
+
+        Raises:
+            ValueError: If the provided system prompt is not a string or is None.
         """
+        if system_prompt is None or not isinstance(system_prompt, str):
+            raise ValueError(CONNECTOR_SET_SYSTEM_PROMPT_VALIDATION_ERROR)
         self.system_prompt = system_prompt

moonshot/src/connectors/connector_prompt_arguments.py

@@ -1,8 +1,8 @@
 from __future__ import annotations
 
-from typing import Any
+from typing import Annotated, Any
 
-from pydantic import BaseModel
+from pydantic import BaseModel, Field
 
 from moonshot.src.connectors.connector_response import ConnectorResponse
 
@@ -11,7 +11,9 @@ class ConnectorPromptArguments(BaseModel):
     class Config:
         arbitrary_types_allowed = True
 
-    prompt_index: int  # The index of the prompt in the dataset
+    prompt_index: Annotated[
+        int, Field(strict=True, ge=0)
+    ]  # The index of the prompt in the dataset
 
     prompt: str  # The actual prompt text
 
@@ -21,4 +23,5 @@ class ConnectorPromptArguments(BaseModel):
         None  # The predicted results, default is None
     )
 
-    duration: float = 0.0  # The duration it took to get the results, default is 0.0
+    # The duration it took to get the results, must be a positive float
+    duration: Annotated[float, Field(strict=True, ge=0.0)] = 0.0

moonshot/src/connectors/connector_response.py

@@ -1,14 +1,9 @@
-class ConnectorResponse:
-    def __init__(self, response: str, context: list = []):
-        """
-        Initializes a ConnectorResponse instance.
+from pydantic import BaseModel
 
-        Args:
-            response (str): The response text from the connector.
-            context (list, optional): Additional context or metadata related to the response. Defaults to an empty list.
-        """
-        self.response = response
-        self.context = context
+
+class ConnectorResponse(BaseModel):
+    response: str = ""
+    context: list = []
 
     def to_dict(self) -> dict:
         """