aiverify-moonshot 0.4.0__py3-none-any.whl

moonshot/integrations/web_api/services/utils/results_formatter.py

@@ -0,0 +1,47 @@
+def transform_web_format(original_dict: dict[str, dict[str, list[dict[str, dict[str, float]]]]]) -> dict[str, dict[str, list[dict[str, list[dict[str, list[dict[str, float]]]]]]]]:
+    transformed = {"metadata": original_dict["metadata"], "results": {"cookbooks": []}}
+    
+    for cookbook_name, cookbook_data in original_dict.items():
+        if cookbook_name == "metadata":
+            continue  # Skip the metadata entry
+        
+        cookbook_entry = {"id": cookbook_name, "recipes": []}
+        for recipe_name, recipe_data in cookbook_data.items():
+            recipe_entry = {"id": recipe_name, "models": []}
+            
+            for model_data_str, model_data in recipe_data.items():
+                # Extract model, recipe, dataset, prompt_template from the string
+                model, _, dataset, prompt_template = eval(model_data_str)
+                
+                model_entry = {
+                    "id": model,
+                    "datasets": [
+                        {
+                            "id": dataset,
+                            "prompt_templates": [
+                                {
+                                    "id": prompt_template,
+                                    "metrics": model_data["results"]
+                                }
+                            ]
+                        }
+                    ]
+                }
+                
+                # Check if the model already exists in the recipe_entry
+                existing_model = next((m for m in recipe_entry["models"] if m["id"] == model), None)
+                if existing_model:
+                    # Check if the dataset already exists in the model
+                    existing_dataset = next((d for d in existing_model["datasets"] if d["id"] == dataset), None)
+                    if existing_dataset:
+                        existing_dataset["prompt_templates"].append(model_entry["datasets"][0]["prompt_templates"][0])
+                    else:
+                        existing_model["datasets"].append(model_entry["datasets"][0])
+                else:
+                    recipe_entry["models"].append(model_entry)
+            
+            cookbook_entry["recipes"].append(recipe_entry)
+        
+        transformed["results"]["cookbooks"].append(cookbook_entry)
+    
+    return transformed

moonshot/integrations/web_api/status_updater/interface/benchmark_progress_callback.py

@@ -0,0 +1,14 @@
+from abc import ABC, abstractmethod
+from ...types.types import TestRunProgress
+
+class InterfaceBenchmarkProgressCallback(ABC):
+    """
+    The interface to implement for handling of benchmark progress updates from moonshot library
+
+    """
+    @abstractmethod
+    def on_progress_update(self, progress_data: TestRunProgress) -> None:
+        pass
+
+
+

moonshot/integrations/web_api/status_updater/interface/redteam_progress_callback.py

@@ -0,0 +1,14 @@
+from abc import ABC, abstractmethod
+from ...types.types import TestRunProgress
+
+class InterfaceRedTeamProgressCallback(ABC):
+    """
+    The interface to implement for handling of redteam progress updates from moonshot library
+
+    """
+    @abstractmethod
+    def on_art_progress_update(self, progress_data: TestRunProgress) -> None:
+        pass
+
+
+

moonshot/integrations/web_api/status_updater/moonshot_ui_webhook.py

@@ -0,0 +1,72 @@
+import json
+import logging
+
+import requests
+from dependency_injector.wiring import inject
+from dotenv import dotenv_values
+
+from ..services.auto_red_team_test_state import AutoRedTeamTestState
+from ..services.benchmark_test_state import BenchmarkTestState
+from ..types.types import RedTeamTestProgress, TestRunProgress
+from .interface.benchmark_progress_callback import InterfaceBenchmarkProgressCallback
+from .interface.redteam_progress_callback import InterfaceRedTeamProgressCallback
+
+
+class MoonshotUIWebhook(
+    InterfaceBenchmarkProgressCallback, InterfaceRedTeamProgressCallback
+):
+    """
+    Implementation of the benchmark progress callback.
+    This class is responsible for sending the progress data to the webhook exposed by moonshot-ui server.
+    """
+
+    @inject
+    def __init__(
+        self,
+        benchmark_test_state: BenchmarkTestState,
+        auto_red_team_test_state: AutoRedTeamTestState,
+    ) -> None:
+        self.benchmark_test_state = benchmark_test_state
+        self.auto_red_team_test_state = auto_red_team_test_state
+        self.url = str(
+            dotenv_values().get(
+                "MOONSHOT_UI_CALLBACK_URL",
+                "http://localhost:3000/api/v1/benchmarks/status",
+            )
+        )
+        self.art_url = str(
+            dotenv_values().get(
+                "MOONSHOT_UI_ART_CALLBACK_URL",
+                "http://localhost:3000/api/v1/redteaming/status",
+            )
+        )
+
+    def on_progress_update(self, progress_data: TestRunProgress) -> None:
+        logger = logging.getLogger()
+        logger.debug(json.dumps(progress_data, indent=2))
+        self.benchmark_test_state.update_progress_status(progress_data)
+
+        try:
+            response = requests.post(self.url, json=progress_data)
+            response.raise_for_status()
+            logger.log(level=logging.DEBUG, msg=response.json())
+            logger.log(
+                level=logging.INFO, msg="Progress data successfully sent to the server."
+            )
+        except requests.RequestException as e:
+            logger.critical(msg=f"Failed to send progress data: {e}")
+
+    def on_art_progress_update(self, progress_data: RedTeamTestProgress) -> None:
+        logger = logging.getLogger()
+        logger.debug(json.dumps(progress_data, indent=2))
+        self.auto_red_team_test_state.update_progress_status(progress_data)
+        print("Calling Auto Red Team Callback")
+        try:
+            response = requests.post(self.art_url, json=progress_data)
+            response.raise_for_status()
+            logger.log(level=logging.DEBUG, msg=response.json())
+            logger.log(
+                level=logging.INFO, msg="Progress data successfully sent to the server."
+            )
+        except requests.RequestException as e:
+            logger.critical(msg=f"Failed to send progress data: {e}")

moonshot/integrations/web_api/types/types.py

@@ -0,0 +1,99 @@
+from enum import Enum
+from typing import Any, List, NotRequired
+
+from typing_extensions import TypedDict
+
+
+class PromptDetails(TypedDict):
+    chat_record_id: int
+    conn_id: str
+    context_strategy: str
+    prompt_template: str
+    prompt: str
+    prepared_prompt: str
+    predicted_result: str
+    duration: str
+    prompt_time: str
+
+
+class EndpointChatHistory(TypedDict):
+    chat_id: str
+    endpoint: str
+    chat_history: list[PromptDetails]
+
+
+SessionChats = list[EndpointChatHistory]
+SessionChatsGroupedBySessionId = dict[str, list[PromptDetails]]
+
+ExecutiResultItem = dict[str, float]
+ExecutionResults = dict[str, list[ExecutiResultItem]]
+
+
+class TestRunProgress(TypedDict):
+    current_runner_id: str
+    current_runner_name: str
+    current_runner_type: str
+    current_duration: int
+    current_status: str
+    current_cookbook_index: int
+    current_cookbook_name: str
+    current_cookbook_total: int
+    current_recipe_index: int
+    current_recipe_name: str
+    current_recipe_total: int
+    current_progress: int
+    current_error_messages: List[str]
+
+
+class RedTeamTestProgress(TypedDict):
+    current_runner_id: str
+    current_am_id: str
+    current_pt_id: str
+    current_cs_id: str
+    current_chats: list[dict]
+    current_batch_size: str
+    current_status: str
+
+
+class UvicornLoggingConfig(TypedDict):
+    version: int
+    formatters: dict[str, dict[str, Any]]
+    handlers: dict[str, dict[str, Any]]
+    root: dict[str, dict[str, Any] | list[str]]
+    disable_existing_loggers: bool
+
+
+class UvicornRunArgs(TypedDict, total=False):
+    host: NotRequired[str]
+    port: NotRequired[int]
+    ssl_keyfile: NotRequired[str]
+    ssl_certfile: NotRequired[str]
+    log_config: NotRequired[UvicornLoggingConfig]
+
+
+class BenchmarkCollectionType(Enum):
+    COOKBOOK = "cookbook"
+    RECIPE = "recipe"
+
+
+class ResultMetadata(TypedDict):
+    id: str
+    name: str
+    start_time: str
+    end_time: str
+    duration: int
+    recipes: List[str]
+    cookbooks: List[str]
+    endpoints: List[str]
+    num_of_prompts: int
+    status: str
+
+
+class RequiredMetadata(TypedDict):
+    metadata: ResultMetadata
+
+
+class BenchmarkResult(TypedDict, RequiredMetadata, total=False):
+    # This indicates that any other keys should map to dictionaries, but this is not enforced by static type checkers.
+    # It serves more as documentation for developers.
+    additional_properties: dict[str, Any]

moonshot/src/api/api_connector.py

@@ -0,0 +1,58 @@
+from pydantic import conlist, validate_call
+
+from moonshot.src.connectors.connector import Connector
+from moonshot.src.connectors_endpoints.connector_endpoint import ConnectorEndpoint
+
+
+# ------------------------------------------------------------------------------
+# Connector APIs
+# ------------------------------------------------------------------------------
+@validate_call
+def api_create_connector_from_endpoint(ep_id: str) -> Connector:
+    """
+    Creates a connector based on the provided endpoint ID.
+
+    This function retrieves the endpoint arguments using the provided endpoint ID and then creates a connector
+    based on those arguments. It utilizes the ConnectorManager's read_endpoint method to fetch the endpoint
+    arguments and then calls the create_connector method to initialize and return the connector.
+
+    Args:
+        ep_id (str): The ID of the endpoint for which to create a connector.
+
+    Returns:
+        Connector: An initialized Connector object.
+    """
+    return Connector.create(ConnectorEndpoint.read(ep_id))
+
+
+@validate_call
+def api_create_connectors_from_endpoints(
+    ep_ids: conlist(str, min_length=1)
+) -> list[Connector]:
+    """
+    Creates connectors for multiple endpoints based on their IDs.
+
+    This function takes a list of endpoint IDs, retrieves the corresponding endpoint arguments for each ID,
+    and then creates a connector for each set of arguments. It utilizes the ConnectorEndpoint's read method
+    to fetch the endpoint arguments and then calls the Connector's create method to initialize the connectors.
+
+    Args:
+        ep_ids (conlist(str, min_length=1)): A list of endpoint IDs for which to create connectors.
+
+    Returns:
+        list[Connector]: A list of initialized Connector objects.
+    """
+    return [Connector.create(ConnectorEndpoint.read(ep_id)) for ep_id in ep_ids]
+
+
+def api_get_all_connector_type() -> list[str]:
+    """
+    Retrieves a list of all available connector types.
+
+    This function calls the ConnectorManager's get_available_connector_types method to retrieve a list of all available
+    connector types. It returns the list of connector types.
+
+    Returns:
+        list[str]: A list of connector types.
+    """
+    return Connector.get_available_items()

moonshot/src/api/api_connector_endpoint.py

@@ -0,0 +1,162 @@
+from pydantic import validate_call
+
+from moonshot.src.connectors_endpoints.connector_endpoint import ConnectorEndpoint
+from moonshot.src.connectors_endpoints.connector_endpoint_arguments import (
+    ConnectorEndpointArguments,
+)
+
+
+# ------------------------------------------------------------------------------
+# Connector Endpoint APIs
+# ------------------------------------------------------------------------------
+@validate_call
+def api_create_endpoint(
+    name: str,
+    connector_type: str,
+    uri: str,
+    token: str,
+    max_calls_per_second: int,
+    max_concurrency: int,
+    params: dict,
+) -> str:
+    """
+    Creates a new connector endpoint.
+
+    This function creates a new connector endpoint with the specified parameters. It initializes
+    a ConnectorEndpointArguments instance and then uses the ConnectorEndpoint class to create
+    the endpoint.
+
+    Args:
+        name (str): The name of the endpoint.
+        connector_type (str): The type of the connector.
+        uri (str): The URI for the connector.
+        token (str): The token for authentication with the connector.
+        max_calls_per_second (int): The maximum number of calls allowed per second.
+        max_concurrency (int): The maximum number of concurrent calls allowed.
+        params (dict): Additional parameters for the connector.
+
+    Returns:
+        str: The ID of the newly created connector endpoint.
+    """
+    # Create a new connector endpoint arguments instance.
+    # We do not need to provide id and created_date.
+    # This is because during creation:
+    #   1. the id is slugify from the name and stored as id.
+    #   2. the created_date is based on the os file created date and time.
+    connector_endpoint_args = ConnectorEndpointArguments(
+        id="",
+        name=name,
+        connector_type=connector_type,
+        uri=uri,
+        token=token,
+        max_calls_per_second=max_calls_per_second,
+        max_concurrency=max_concurrency,
+        params=params,
+        created_date="",
+    )
+    return ConnectorEndpoint.create(connector_endpoint_args)
+
+
+@validate_call
+def api_read_endpoint(ep_id: str) -> dict:
+    """
+    Reads an endpoint from the connector manager.
+
+    This function reads an endpoint from the connector manager using the provided endpoint ID.
+
+    Args:
+        ep_id (str): The ID of the endpoint to read.
+
+    Returns:
+        dict: A dictionary containing the endpoint information.
+    """
+    return ConnectorEndpoint.read(ep_id).to_dict()
+
+
+@validate_call
+def api_update_endpoint(ep_id: str, **kwargs) -> bool:
+    """
+    Updates an existing endpoint with new values.
+
+    This function updates an existing endpoint in the connector manager using the provided endpoint ID and
+    keyword arguments.
+
+    Each keyword argument corresponds to an attribute of the endpoint that should be updated.
+
+    Args:
+        ep_id (str): The ID of the endpoint to update.
+        **kwargs: Arbitrary keyword arguments representing the attributes to update.
+
+    Returns:
+        bool: True if the update was successful.
+
+    Raises:
+        RuntimeError: If the endpoint with the given ID does not exist or the update failed.
+    """
+    # Check if the endpoint exists
+    try:
+        existing_endpoint = ConnectorEndpoint.read(ep_id)
+    except Exception:
+        raise RuntimeError(
+            f"[api_connector_endpoint]: Endpoint with ID '{ep_id}' does not exist"
+        )
+
+    # Update the fields of the existing endpoint with the provided kwargs
+    for key, value in kwargs.items():
+        if hasattr(existing_endpoint, key):
+            setattr(existing_endpoint, key, value)
+
+    # Perform pydantic check on the updated existing endpoint
+    ConnectorEndpointArguments.model_validate(existing_endpoint.to_dict())
+
+    # Update the endpoint
+    return ConnectorEndpoint.update(existing_endpoint)
+
+
+@validate_call
+def api_delete_endpoint(ep_id: str) -> bool:
+    """
+    Deletes an endpoint from the connector manager.
+
+    This function deletes an endpoint from the connector manager using the provided endpoint ID.
+
+    Args:
+        ep_id (str): The ID of the endpoint to delete.
+
+    Returns:
+        bool: True if the deletion was successful.
+
+    Raises:
+        RuntimeError: If the endpoint with the given ID does not exist or the deletion failed.
+    """
+    return ConnectorEndpoint.delete(ep_id)
+
+
+def api_get_all_endpoint() -> list[dict]:
+    """
+    Retrieves a list of all available endpoints.
+
+    This function calls the ConnectorManager's get_available_endpoints method to retrieve a list of all available
+    endpoints and their details. It then converts each ConnectorEndpointArguments object into a dictionary for easier
+    consumption by the caller.
+
+    Returns:
+        list[dict]: A list of dictionaries, each representing an endpoint's details.
+    """
+    _, endpoints = ConnectorEndpoint.get_available_items()
+    return [endpoint.to_dict() for endpoint in endpoints]
+
+
+def api_get_all_endpoint_name() -> list[str]:
+    """
+    Retrieves a list of all endpoint names.
+
+    This function calls the ConnectorManager's get_available_endpoints method to retrieve a list of all available
+    endpoint names. It extracts the names from the tuple returned by get_available_endpoints, which contains a list
+    of endpoint names and a list of ConnectorEndpointArguments objects.
+
+    Returns:
+        list[str]: A list of endpoint names.
+    """
+    endpoints_names, _ = ConnectorEndpoint.get_available_items()
+    return endpoints_names

moonshot/src/api/api_context_strategy.py

@@ -0,0 +1,57 @@
+from pydantic import validate_call
+
+from moonshot.src.redteaming.attack.context_strategy import ContextStrategy
+
+
+# ------------------------------------------------------------------------------
+# Context Strategy APIs
+# ------------------------------------------------------------------------------
+def api_get_all_context_strategies() -> list[str]:
+    """
+    Retrieves and returns the names of all context strategies currently available.
+
+    This API endpoint interfaces with the `ContextStrategy.get_all_context_strategy_names` method to fetch a list
+    of all context strategy names. It's designed for clients that need to know what context strategies are available for
+    use in sessions or other components of the system.
+
+    Returns:
+        list[str]: A list of strings, each representing the name of a context strategy.
+    """
+    return ContextStrategy.get_all_context_strategies()
+
+
+@validate_call
+def api_delete_context_strategy(cs_id: str) -> bool:
+    """
+    Deletes a context strategy identified by its ID.
+
+    This API endpoint interfaces with the `ContextStrategy.delete` method to remove a context strategy from the system.
+    It is used to manage the available context strategies by allowing for their removal when they are no longer needed.
+
+    Args:
+        cs_id (str): The unique identifier of the context strategy to be deleted.
+
+    Returns:
+        bool: True if the context strategy was successfully deleted.
+
+    Raises:
+        Exception: If the deletion process encounters an error.
+    """
+    return ContextStrategy.delete(cs_id)
+
+
+def api_get_all_context_strategy_metadata() -> list:
+    """
+    Retrieves metadata for all context strategy modules.
+
+    This function retrieves the metadata for all available context strategies and
+    returns a list of metadata dictionaries.
+
+    Returns:
+        list: A list of attack module metadata.
+    """
+
+    return [
+        ContextStrategy.load(context_strategy_name).get_metadata()  # type: ignore ; ducktyping
+        for context_strategy_name in ContextStrategy.get_all_context_strategies()
+    ]

moonshot/src/api/api_cookbook.py

@@ -0,0 +1,160 @@
+from pydantic import conlist, validate_call
+
+from moonshot.src.cookbooks.cookbook import Cookbook
+from moonshot.src.cookbooks.cookbook_arguments import CookbookArguments
+
+
+# ------------------------------------------------------------------------------
+# Cookbook APIs
+# ------------------------------------------------------------------------------
+@validate_call
+def api_create_cookbook(name: str, description: str, recipes: list[str]) -> str:
+    """
+    Creates a new cookbook.
+
+    This function takes the name, description, and recipes for a new cookbook as input. It then creates a new
+    CookbookArguments 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 Cookbook's create_cookbook method to
+    create the new cookbook.
+
+    Args:
+        name (str): The name of the new cookbook.
+        description (str): A brief description of the new cookbook.
+        recipes (list[str]): A list of recipes to be included in the new cookbook.
+
+    Returns:
+        str: The ID of the newly created cookbook.
+    """
+    # Create a new cookbook
+    # We do not need to provide the id.
+    # This is because during creation:
+    # 1. the id is slugify from the name and stored as id.
+    cb_args = CookbookArguments(
+        id="",
+        name=name,
+        description=description,
+        recipes=recipes,
+    )
+    return Cookbook.create(cb_args)
+
+
+@validate_call
+def api_read_cookbook(cb_id: str) -> dict:
+    """
+    Retrieves a cookbook based on the provided cookbook ID.
+
+    This function reads a cookbook using the `read_cookbook` method
+    of the `Cookbook` class, and converts the returned `Cookbook` object to a dictionary using its `to_dict` method.
+
+    Args:
+        cb_id (str): A cookbook ID.
+
+    Returns:
+        dict: A dictionary representing a cookbook.
+    """
+    return Cookbook.read(cb_id).to_dict()
+
+
+@validate_call
+def api_read_cookbooks(cb_ids: conlist(str, min_length=1)) -> list[dict]:
+    """
+    Retrieves a list of cookbooks based on the provided list of cookbook IDs.
+
+    This function iterates over the list of provided cookbook IDs, reads each cookbook using the `read_cookbook` method
+    of the `Cookbook` class, and converts the returned `Cookbook` objects to dictionaries using their `to_dict` method.
+    It then returns a list of these dictionary representations.
+
+    Args:
+        cb_ids (conlist(str, min_length=1)): A list of cookbook IDs.
+
+    Returns:
+        list[dict]: A list of dictionaries representing the cookbooks.
+    """
+    return [Cookbook.read(cb_id).to_dict() for cb_id in cb_ids]
+
+
+@validate_call
+def api_update_cookbook(cb_id: str, **kwargs) -> bool:
+    """
+    Updates the fields of an existing cookbook with the provided keyword arguments.
+
+    This function first checks if the cookbook with the given ID exists. If it does, it updates the fields
+    of the cookbook with the provided keyword arguments. If a field does not exist on the cookbook, it is ignored.
+    After updating the fields, it persists the changes to the cookbook.
+
+    Args:
+        cb_id (str): The ID of the cookbook to update.
+        **kwargs: Arbitrary keyword arguments representing the fields to update and their new values.
+
+    Returns:
+        bool: True if the cookbook was successfully updated.
+
+    Raises:
+        Exception: If there's an error during the update process.
+    """
+    # Check if the cookbook exists
+    try:
+        existing_cookbook = Cookbook.read(cb_id)
+    except Exception:
+        raise RuntimeError(f"Cookbook with ID '{cb_id}' does not exist")
+
+    # Update the fields of the existing cookbook with the provided kwargs
+    for key, value in kwargs.items():
+        if hasattr(existing_cookbook, key):
+            setattr(existing_cookbook, key, value)
+
+    # Perform pydantic check on the updated existing cookbook
+    CookbookArguments.model_validate(existing_cookbook.to_dict())
+
+    # Update the cookbook
+    return Cookbook.update(existing_cookbook)
+
+
+@validate_call
+def api_delete_cookbook(cb_id: str) -> bool:
+    """
+    Deletes a cookbook based on the provided cookbook ID.
+
+    This function calls the `delete` method of the `Cookbook` class with the given cookbook ID. If the cookbook
+    is successfully deleted, the method returns True, otherwise it returns False.
+
+    Args:
+        cb_id (str): The ID of the cookbook to delete.
+
+    Returns:
+        bool: True if the cookbook was successfully deleted.
+
+    Raises:
+        Exception: If the deletion process encounters an error.
+    """
+    return Cookbook.delete(cb_id)
+
+
+def api_get_all_cookbook() -> list[dict]:
+    """
+    Retrieves all available cookbooks.
+
+    This function calls the `get_available_cookbooks` method of the `Cookbook` class, which returns a tuple
+    containing a list of cookbook IDs and a list of `CookbookArguments` objects. The function then returns a list
+    of dictionaries, each representing a cookbook.
+
+    Returns:
+        list[dict]: A list of dictionaries, each representing a cookbook.
+    """
+    _, cookbooks = Cookbook.get_available_items()
+    return [cookbook.to_dict() for cookbook in cookbooks]
+
+
+def api_get_all_cookbook_name() -> list[str]:
+    """
+    Retrieves the names of all available cookbooks.
+
+    This function calls the `get_available_cookbooks` method of the `Cookbook` class, which returns a tuple
+    containing a list of cookbook IDs and a list of `CookbookArguments` objects. The function then returns the
+    list of cookbook IDs, which are the names of the cookbooks.
+
+    Returns:
+        list[str]: A list of cookbook names.
+    """
+    cookbooks_names, _ = Cookbook.get_available_items()
+    return cookbooks_names