aiverify-moonshot 0.4.0__py3-none-any.whl

moonshot/src/api/api_session.py

@@ -0,0 +1,290 @@
+from operator import itemgetter
+from typing import Any
+
+from pydantic import validate_call
+
+from moonshot.src.api.api_runner import api_get_all_runner, api_load_runner
+from moonshot.src.configs.env_variables import EnvVariables
+from moonshot.src.redteaming.session.session import Session
+from moonshot.src.runners.runner_type import RunnerType
+from moonshot.src.storage.db_interface import DBInterface
+from moonshot.src.storage.storage import Storage
+
+
+# ------------------------------------------------------------------------------
+# Session and Chat APIs
+# ------------------------------------------------------------------------------
+@validate_call
+def api_load_session(runner_id: str) -> dict | None:
+    """
+    Loads the session details for a specific runner.
+
+    This function calls the `Session.load` method to retrieve the session details associated with the
+    specified runner ID.
+
+    Args:
+        runner_id (str): The unique identifier of the runner for which the session details are to be loaded.
+
+    Returns:
+        dict | None: A dictionary containing the session details if available, otherwise None.
+    """
+    return Session.load(api_load_runner(runner_id).database_instance)
+
+
+@validate_call
+def api_create_session(
+    runner_id: str, database_instance: Any, endpoints: list[str], runner_args: dict
+) -> Session:
+    """
+    Creates a new session for a specific runner.
+
+    This function creates a new session by calling the `Session` constructor with the provided arguments.
+
+    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.
+        runner_args (dict): A dictionary of arguments for the runner.
+
+    Returns:
+        Session
+    """
+    if isinstance(database_instance, DBInterface):
+        if runner_id:
+            session_instance = Session(
+                runner_id,
+                RunnerType.REDTEAM,
+                {**runner_args},
+                database_instance,
+                endpoints,
+                Storage.get_filepath(
+                    EnvVariables.RESULTS.name, runner_id, "json", True
+                ),
+            )
+            return session_instance
+        else:
+            raise RuntimeError(
+                "[Session] Failed to initialise Session. String should have at least 1 character."
+            )
+    else:
+        raise RuntimeError(
+            "[Session] Failed to initialise Session. No database instance provided."
+        )
+
+
+def api_get_all_session_names() -> list[str]:
+    """
+    Retrieves a list of all session names.
+
+    This function calls the `api_get_available_session_info` method to obtain the available session information
+    and returns a list of session names.
+
+    Returns:
+        list[str]: A list of strings, each denoting a session name.
+    """
+    session_names, _ = api_get_available_session_info()
+    return session_names
+
+
+def api_get_available_session_info() -> tuple[list, list]:
+    """
+    Retrieves the IDs and database instances of runners with active sessions.
+
+    This function retrieves the IDs and database instances 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
+        metadata for runners with active sessions.
+    """
+    runners_info = api_get_all_runner()
+    runner_instances = [
+        api_load_runner(str(runner_info.get("id"))) for runner_info in runners_info
+    ]
+
+    runner_ids = []
+    session_metadata_list = []
+
+    for runner_instance in runner_instances:
+        if runner_instance.database_instance:
+            session_metadata = Session.load(runner_instance.database_instance)
+            if session_metadata is not None:
+                runner_ids.append(runner_instance.id)
+                session_metadata_list.append(session_metadata)
+    return runner_ids, session_metadata_list
+
+
+def api_get_all_session_metadata() -> list:
+    """
+    Retrieves metadata for all sessions.
+
+    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.
+    """
+    _, session_metadata_list = api_get_available_session_info()
+    return sorted(
+        session_metadata_list, key=itemgetter("created_datetime"), reverse=True
+    )
+
+
+@validate_call
+def api_update_context_strategy(runner_id: str, context_strategy: str) -> bool:
+    """
+    Updates the context strategy for a specific runner.
+
+    This function updates the context strategy for a specific runner identified by the given runner_id. It calls the
+    `Session.update_context_strategy` method with the runner's database instance,
+    runner_id, and the new context_strategy.
+
+    Args:
+        runner_id (str): The ID of the runner for which the context strategy needs to be updated.
+        context_strategy (str): The new context strategy to be set for the runner.
+
+    Returns:
+        bool: The status on whether the context strategy is updated successfully.
+    """
+    return Session.update_context_strategy(
+        api_load_runner(runner_id).database_instance, runner_id, context_strategy
+    )
+
+
+@validate_call
+def api_update_cs_num_of_prev_prompts(runner_id: str, num_of_prev_prompts: int) -> bool:
+    """
+    Updates the number of previous prompts used in a context strategy for a specific runner.
+
+    This function updates the number of previous prompts used in a context strategy for a specific runner identified by
+    the given runner_id. It calls the `Session.update_cs_num_of_prev_prompts` method with the runner's database
+    instance, runner_id, and the new num_of_prev_prompts.
+
+    Args:
+        runner_id (str): The ID of the runner for which the number of previous prompts needs to be updated.
+        num_of_prev_prompts (int): The new number of previous prompts to be set for the runner.
+
+    Returns:
+        bool: The status on whether the number of prompts for context strategy is updated successfully.
+    """
+    return Session.update_cs_num_of_prev_prompts(
+        api_load_runner(runner_id).database_instance, runner_id, num_of_prev_prompts
+    )
+
+
+@validate_call
+def api_update_prompt_template(runner_id: str, prompt_template: str) -> bool:
+    """
+    Updates the prompt template for a specific runner.
+
+    This function updates the prompt template for a specific runner identified by the given runner_id. It calls the
+    `Session.update_prompt_template` method with the runner's database instance,
+    runner_id, and the new prompt_template.
+
+    Args:
+        runner_id (str): The ID of the runner for which the prompt template needs to be updated.
+        prompt_template (str): The new prompt template to be set for the runner.
+
+    Returns:
+        bool: The status on whether the prompt template is updated successfully.
+    """
+    return Session.update_prompt_template(
+        api_load_runner(runner_id).database_instance, runner_id, prompt_template
+    )
+
+
+@validate_call
+def api_update_metric(runner_id: str, metric_id: str) -> bool:
+    """
+    Updates the metric for a specific runner.
+
+    This function updates the metric for a specific runner identified by the given runner_id. It calls the
+    `Session.update_metric` method with the runner's database instance,
+    runner_id, and the new metric_id.
+
+    Args:
+        runner_id (str): The ID of the runner for which the metric needs to be updated.
+        metric_id (str): The new metric to be set for the runner.
+
+    Returns:
+        bool: The status on whether the metric is updated successfully.
+    """
+    return Session.update_metric(
+        api_load_runner(runner_id).database_instance, runner_id, metric_id
+    )
+
+
+@validate_call
+def api_update_system_prompt(runner_id: str, system_prompt: str) -> bool:
+    """
+    Updates the system prompt for a specific runner.
+
+    This function updates the system prompt for a specific runner identified by the given runner_id. It calls the
+    `Session.update_system_prompt` method with the runner's database instance,
+    runner_id, and the new system_prompt.
+
+    Args:
+        runner_id (str): The ID of the runner for which the system prompt needs to be updated.
+        system_prompt (str): The new system prompt to be set for the runner.
+
+    Returns:
+        bool: The status on whether the system prompt is updated successfully.
+    """
+    return Session.update_system_prompt(
+        api_load_runner(runner_id).database_instance, runner_id, system_prompt
+    )
+
+
+@validate_call
+def api_update_attack_module(runner_id: str, attack_module_id: str) -> bool:
+    """
+    Updates the attack module for a specific runner.
+
+    This function updates the attack module for a specific runner identified by the given runner_id. It calls the
+    `Session.update_attack_module` method with the runner's database instance,
+    runner_id, and the new attack_module_id.
+
+    Args:
+        runner_id (str): The ID of the runner for which the attack module needs to be updated.
+        attack_module_id (str): The new attack module to be set for the runner.
+
+    Returns:
+        bool: The status on whether the attack module is updated successfully.
+    """
+    return Session.update_attack_module(
+        api_load_runner(runner_id).database_instance, runner_id, attack_module_id
+    )
+
+
+@validate_call
+def api_delete_session(runner_id: str) -> bool:
+    """
+    Deletes the session for a specific runner.
+
+    This function deletes the session for the runner identified by the given runner_id. It calls the
+    `Session.delete` method with the runner's database instance.
+
+    Args:
+        runner_id (str): The ID of the runner for which the session needs to be deleted.
+
+    Returns:
+        bool: The status on whether the session is deleted successfully.
+    """
+    return Session.delete(api_load_runner(runner_id).database_instance)
+
+
+@validate_call
+def api_get_all_chats_from_session(runner_id: str) -> dict | None:
+    """
+    Retrieves all chat messages from a specific session.
+
+    This function retrieves all chat messages from the session associated with the specified runner ID.
+    It calls the `Session.get_session_chats` method with the runner's database instance.
+
+    Args:
+        runner_id (str): The unique identifier of the runner for which the chat messages are to be retrieved.
+
+    Returns:
+        dict | None: A dictionary containing all chat messages if available, otherwise None.
+    """
+    return Session.get_session_chats(api_load_runner(runner_id).database_instance)

moonshot/src/configs/env_variables.py

@@ -0,0 +1,187 @@
+import importlib.resources
+from enum import Enum
+from pathlib import Path
+
+__app_name__ = "moonshot"
+
+
+class EnvVariables(Enum):
+    ATTACK_MODULES = "ATTACK_MODULES"
+    CONNECTORS = "CONNECTORS"
+    CONNECTORS_ENDPOINTS = "CONNECTORS_ENDPOINTS"
+    CONTEXT_STRATEGY = "CONTEXT_STRATEGY"
+    COOKBOOKS = "COOKBOOKS"
+    DATABASES = "DATABASES"
+    DATABASES_MODULES = "DATABASES_MODULES"
+    DATASETS = "DATASETS"
+    IO_MODULES = "IO_MODULES"
+    METRICS = "METRICS"
+    PROMPT_TEMPLATES = "PROMPT_TEMPLATES"
+    RECIPES = "RECIPES"
+    RESULTS = "RESULTS"
+    RESULTS_MODULES = "RESULTS_MODULES"
+    RUNNERS = "RUNNERS"
+    RUNNERS_MODULES = "RUNNERS_MODULES"
+
+
+class EnvironmentVars:
+    env_vars = {}
+
+    ATTACK_MODULES = [
+        env_vars.get(EnvVariables.ATTACK_MODULES.value),
+        str(importlib.resources.files(__app_name__).joinpath("data/attack-modules")),
+    ]
+    CONNECTORS = [
+        env_vars.get(EnvVariables.CONNECTORS.value),
+        str(importlib.resources.files(__app_name__).joinpath("data/connectors")),
+    ]
+    CONNECTORS_ENDPOINTS = [
+        env_vars.get(EnvVariables.CONNECTORS_ENDPOINTS.value),
+        str(
+            importlib.resources.files(__app_name__).joinpath(
+                "data/connectors-endpoints"
+            )
+        ),
+    ]
+    CONTEXT_STRATEGY = [
+        env_vars.get(EnvVariables.CONTEXT_STRATEGY.value),
+        str(importlib.resources.files(__app_name__).joinpath("data/context-strategy")),
+    ]
+    COOKBOOKS = [
+        env_vars.get(EnvVariables.COOKBOOKS.value),
+        str(importlib.resources.files(__app_name__).joinpath("data/cookbooks")),
+    ]
+    DATABASES = [
+        env_vars.get(EnvVariables.DATABASES.value),
+        str(
+            importlib.resources.files(__app_name__).joinpath(
+                "data/generated-outputs/databases"
+            )
+        ),
+    ]
+    DATABASES_MODULES = [
+        env_vars.get(EnvVariables.DATABASES_MODULES.value),
+        str(importlib.resources.files(__app_name__).joinpath("data/databases-modules")),
+    ]
+    DATASETS = [
+        env_vars.get(EnvVariables.DATASETS.value),
+        str(importlib.resources.files(__app_name__).joinpath("data/datasets")),
+    ]
+    IO_MODULES = [
+        env_vars.get(EnvVariables.IO_MODULES.value),
+        str(importlib.resources.files(__app_name__).joinpath("data/io-modules")),
+    ]
+    METRICS = [
+        env_vars.get(EnvVariables.METRICS.value),
+        str(importlib.resources.files(__app_name__).joinpath("data/metrics")),
+    ]
+    PROMPT_TEMPLATES = [
+        env_vars.get(EnvVariables.PROMPT_TEMPLATES.value),
+        str(importlib.resources.files(__app_name__).joinpath("data/prompt-templates")),
+    ]
+    RECIPES = [
+        env_vars.get(EnvVariables.RECIPES.value),
+        str(importlib.resources.files(__app_name__).joinpath("data/recipes")),
+    ]
+    RESULTS = [
+        env_vars.get(EnvVariables.RESULTS.value),
+        str(
+            importlib.resources.files(__app_name__).joinpath(
+                "data/generated-outputs/results"
+            )
+        ),
+    ]
+    RESULTS_MODULES = [
+        env_vars.get(EnvVariables.RESULTS_MODULES.value),
+        str(importlib.resources.files(__app_name__).joinpath("data/results-modules")),
+    ]
+    RUNNERS = [
+        env_vars.get(EnvVariables.RUNNERS.value),
+        str(
+            importlib.resources.files(__app_name__).joinpath(
+                "data/generated-outputs/runners"
+            )
+        ),
+    ]
+    RUNNERS_MODULES = [
+        env_vars.get(EnvVariables.RUNNERS_MODULES.value),
+        str(importlib.resources.files(__app_name__).joinpath("data/runners-modules")),
+    ]
+
+    @staticmethod
+    def load_env(env_dict: dict | None = None) -> None:
+        """
+        This method is used to load environment variables from a provided dictionary.
+        If the dictionary is not provided, it defaults to an empty dictionary.
+
+        The method updates the class attributes with the corresponding values from the dictionary.
+        If a key from the class attributes is not found in the dictionary, it will not be updated and a
+        message will be printed.
+
+        Args:
+            env_dict (dict | None): A dictionary containing the environment variables to be loaded.
+                                    If None, an empty dictionary will be used.
+        """
+        if env_dict is None:
+            env_dict = dict()
+
+        keys = [e.value for e in EnvVariables]
+        unset_keys = []
+        for key in keys:
+            if key in env_dict:
+                given_path = Path(env_dict[key])
+                if given_path.exists():
+                    EnvironmentVars.__dict__[key][0] = str(given_path)
+                else:
+                    print(
+                        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(
+                f"Unable to retrieve the following environment variables: {unset_keys}. The stock set will be used."
+            )
+
+    @staticmethod
+    def get_file_path(file_type: str, file_name: str, ignore_existance: bool) -> str:
+        """
+        This method is used to get the file path for a given file type and file name.
+        If the ignore existance flag is set to True, it returns the file path even if the file does not exist.
+
+        Args:
+            file_type (str): The type of the file (e.g., 'recipe', 'cookbook').
+            file_name (str): The name of the file.
+            ignore_existance (bool): A flag indicating whether to return the file path
+                             even if the file does not exist.
+
+        Returns:
+            str: The file path of the file.
+        """
+        for directory in EnvironmentVars.get_file_directory(file_type):
+            file_path = Path(directory) / file_name
+            if ignore_existance:
+                return str(file_path)
+            else:
+                if Path(file_path).exists():
+                    return str(file_path)
+        return ""
+
+    @staticmethod
+    def get_file_directory(file_type: str) -> list[str]:
+        """
+        This method retrieves the directory paths associated with a specified file type.
+
+        Args:
+            file_type (str): The type of file for which the directory paths are to be retrieved.
+
+        Returns:
+            list[str]: A list containing the directory paths associated with the given file type.
+        """
+        path_from_user, path_from_resource = getattr(EnvironmentVars, file_type)
+
+        if path_from_user is not None:
+            return [path_from_user, path_from_resource]
+        else:
+            return [path_from_resource]