aiverify-moonshot 0.4.0__py3-none-any.whl

moonshot/integrations/web_api/services/recipe_service.py

@@ -0,0 +1,155 @@
+from moonshot.src.recipes.recipe_arguments import RecipeArguments as Recipe
+
+from .... import api as moonshot_api
+from ..schemas.recipe_create_dto import RecipeCreateDTO, RecipeUpdateDTO
+from ..schemas.recipe_response_model import RecipeResponseModel
+from ..services.base_service import BaseService
+from ..services.utils.exceptions_handler import exception_handler
+
+
+class RecipeService(BaseService):
+    @exception_handler
+    def create_recipe(self, recipe_data: RecipeCreateDTO) -> None:
+        """
+        Create a new recipe with the given data.
+
+        Args:
+            recipe_data (RecipeCreateDTO): Data transfer object containing recipe details.
+        """
+        moonshot_api.api_create_recipe(
+            name=recipe_data.name,
+            description=recipe_data.description,
+            tags=recipe_data.tags,
+            categories=recipe_data.categories,
+            datasets=recipe_data.datasets,
+            prompt_templates=recipe_data.prompt_templates,
+            metrics=recipe_data.metrics,
+            attack_modules=recipe_data.attack_modules,
+            grading_scale=recipe_data.grading_scale,
+        )
+
+    @exception_handler
+    def get_all_recipes(
+        self,
+        tags: str,
+        categories: str,
+        sort_by: str,
+        count: bool,
+        ids: str | None = None,
+    ) -> list[RecipeResponseModel]:
+        """
+        Retrieve all recipes, with optional filters for tags, categories, sorting, and including prompt counts.
+
+        Args:
+            ids (str, optional): Filter recipes by IDs. If None, no ID-based filtering is applied.
+            tags (str, optional): Filter recipes by tags. If None, no tag-based filtering is applied.
+            categories (str, optional): Filter recipes by categories. If None, no category-based filtering is applied.
+            sort_by (str, optional): Sort recipes by a specified field. If None, no sorting is applied.
+            count (bool, optional): Include the total prompt count in each recipe if True.
+
+        Returns:
+            list[RecipeResponseModel]: A list of recipe, filtered and sorted, with optional prompt counts.
+        """
+        filtered_recipes: list[RecipeResponseModel] = []
+
+        if ids:
+            recipe_ids = ids.split(",")
+            recipes = [moonshot_api.api_read_recipe(id) for id in recipe_ids]
+        else:
+            recipes = moonshot_api.api_get_all_recipe()
+
+        for recipe_dict in recipes:
+            recipe = RecipeResponseModel(**recipe_dict)
+            if count:
+                recipe.total_prompt_in_recipe = get_total_prompt_in_recipe(recipe)
+            filtered_recipes.append(recipe)
+
+        # TODO - do all filtering in 1 pass
+        if tags:
+            filtered_recipes = [
+                recipe for recipe in filtered_recipes if tags in recipe.tags
+            ]
+
+        if categories:
+            categories_list = categories.split(",") if categories else []
+            if categories_list:
+                filtered_recipes = [
+                    recipe
+                    for recipe in filtered_recipes
+                    if any(
+                        category.lower() in (cat.lower() for cat in recipe.categories)
+                        for category in categories_list
+                    )
+                ]
+        if sort_by:
+            if sort_by == "id":
+                filtered_recipes.sort(key=lambda x: x.id)
+
+        return filtered_recipes
+
+    @exception_handler
+    def get_all_recipes_name(self) -> list[str]:
+        """
+        Retrieve the names of all recipes.
+
+        Returns:
+            list[str]: A list of recipe names.
+        """
+        recipes = moonshot_api.api_get_all_recipe_name()
+        return recipes
+
+    @exception_handler
+    def update_recipe(self, recipe_data: RecipeUpdateDTO, recipe_id: str) -> None:
+        """
+        Update an existing recipe with new data.
+
+        Args:
+            recipe_data (RecipeCreateDTO): Data transfer object containing new recipe details.
+            recipe_id (str): The ID of the recipe to update.
+        """
+        update_data = {
+            k: v
+            for k, v in recipe_data.to_dict().items()
+            if v is not None and k not in ["id", "stats"]
+        }
+
+        moonshot_api.api_update_recipe(rec_id=recipe_id, **update_data)
+
+    @exception_handler
+    def delete_recipe(self, recipe_id: str) -> None:
+        """
+        Delete a recipe by its ID.
+
+        Args:
+            recipe_id (str): The ID of the recipe to delete.
+        """
+        moonshot_api.api_delete_recipe(recipe_id)
+
+
+@staticmethod
+def get_total_prompt_in_recipe(recipe: Recipe) -> int:
+    """
+    Calculate the total number of prompts in a recipe.
+
+    This function sums up the number of dataset prompts and then multiplies
+    the result by the number of prompt templates if they exist.
+
+    Args:
+        recipe (Recipe): The recipe object containing the stats and
+                                  prompt templates information.
+
+    Returns:
+        int: The total count of prompts within the recipe.
+    """
+    # Initialize total prompt count
+    total_prompt_count = 0
+
+    # Add counts from dataset prompts if available
+    datasets_prompts = recipe.stats.get("num_of_datasets_prompts", {})
+    total_prompt_count += sum(datasets_prompts.values())
+
+    # If there are prompt templates, scale the total count by the number of templates
+    if recipe.prompt_templates:
+        total_prompt_count *= len(recipe.prompt_templates)
+
+    return total_prompt_count

moonshot/integrations/web_api/services/runner_service.py

@@ -0,0 +1,147 @@
+from typing import Callable
+
+from moonshot.src.runners.runner import Runner
+
+from .... import api as moonshot_api
+from ..services.base_service import BaseService
+from ..services.utils.exceptions_handler import exception_handler
+from ..types.types import TestRunProgress
+
+
+class RunnerService(BaseService):
+    @exception_handler
+    def get_all_runner(self) -> list[dict]:
+        """
+        Retrieves all runners.
+
+        Returns:
+            A list of dictionaries, where each dictionary contains the details of a runner.
+        """
+        runners = moonshot_api.api_get_all_runner()
+        return runners
+
+    @exception_handler
+    def get_all_runner_name(self) -> list[str]:
+        """
+        Retrieves the names of all runners.
+
+        Returns:
+            A list of strings, where each string is the name of a runner.
+        """
+        runners = moonshot_api.api_get_all_runner_name()
+        return runners
+
+    @exception_handler
+    def get_runner_by_id(self, runner_id: str) -> dict | None:
+        """
+        Retrieves a runner by its ID.
+
+        Args:
+            runner_id: The unique identifier of the runner.
+
+        Returns:
+            A dictionary containing the runner's details if found, otherwise None.
+        """
+        runner = moonshot_api.api_read_runner(runner_id)
+        return runner
+
+    @exception_handler
+    def delete_run(self, runner_id: str) -> None:
+        """
+        Deletes a runner by its ID.
+
+        Args:
+            runner_id: The unique identifier of the runner to be deleted.
+        """
+        moonshot_api.api_delete_runner(runner_id)
+
+    @exception_handler
+    def create_runner(
+        self,
+        runner_name: str,
+        endpoints: list[str],
+        description: str,
+        progress_callback_func: Callable[[TestRunProgress], None] = None,
+    ) -> Runner:
+        """
+        Creates a new runner with the specified details.
+
+        Args:
+            runner_name: The name of the runner.
+            endpoints: A list of endpoints the runner will interact with.
+            description: A brief description of the runner.
+            progress_callback_func: An optional callback function for progress updates.
+
+        Returns:
+            An instance of Runner.
+        """
+        return moonshot_api.api_create_runner(
+            name=runner_name,
+            endpoints=endpoints,
+            description=description,
+            progress_callback_func=progress_callback_func,
+        )
+
+    @exception_handler
+    def load_runner(
+        self,
+        runner_id: str,
+        progress_callback_func: Callable[[TestRunProgress], None] = None,
+    ) -> Runner:
+        """
+        Loads a runner by its ID.
+
+        Args:
+            runner_id: The unique identifier of the runner to be loaded.
+            progress_callback_func: An optional callback function for progress updates.
+
+        Returns:
+            An instance of Runner.
+        """
+        return moonshot_api.api_load_runner(
+            runner_id, progress_callback_func=progress_callback_func
+        )
+
+    @exception_handler
+    def get_run_details_by_runner(self, runner_id: str, run_id: int):
+        """
+        Retrieves the details of a specific run by the runner ID and run ID.
+
+        Args:
+            runner_id: The unique identifier of the runner.
+            run_id: The unique identifier of the run.
+
+        Returns:
+            A dictionary containing the details of the run if found, otherwise an empty dictionary.
+        """
+        runs = moonshot_api.api_get_all_run(runner_id)
+        retn = {}
+        run_id_int = int(run_id)
+        for run in runs:
+            run_id_from_run = int(run.get("run_id"))
+            if run_id_from_run == run_id_int:
+                retn = {
+                    "run_id": run.get("run_id"),
+                    "runner_id": run.get("runner_id"),
+                    "runner_args": run.get("runner_args"),
+                    "endpoints": run.get("endpoints"),
+                    "start_time": run.get("start_time"),
+                }
+        return retn
+
+    @exception_handler
+    def get_runs_id_in_runner(self, runner_id) -> list[int]:
+        """
+        Retrieves a list of run IDs for a given runner.
+
+        Args:
+            runner_id: The unique identifier of the runner.
+
+        Returns:
+            A list of integers, where each integer is the ID of a run associated with the runner.
+        """
+        runs = moonshot_api.api_get_all_run(runner_id)
+        retn = []
+        for run in runs:
+            retn.append(run.get("run_id"))
+        return retn

moonshot/integrations/web_api/services/session_service.py

@@ -0,0 +1,350 @@
+from dependency_injector.wiring import inject
+
+from moonshot.src.api.api_session import api_create_session
+from moonshot.src.runners.runner import Runner
+
+from .... import api as moonshot_api
+from ..schemas.prompt_response_model import PromptResponseModel
+from ..schemas.session_create_dto import SessionCreateDTO
+from ..schemas.session_prompt_dto import SessionPromptDTO
+from ..schemas.session_response_model import SessionMetadataModel, SessionResponseModel
+from ..services.auto_red_team_test_manager import AutoRedTeamTestManager
+from ..services.runner_service import RunnerService
+from ..services.utils.exceptions_handler import exception_handler
+from ..status_updater.interface.redteam_progress_callback import (
+    InterfaceRedTeamProgressCallback,
+)
+from .base_service import BaseService
+
+
+class SessionService(BaseService):
+    active_runner: Runner
+
+    @inject
+    def __init__(
+        self,
+        auto_red_team_test_manager: AutoRedTeamTestManager,
+        progress_status_updater: InterfaceRedTeamProgressCallback,
+        runner_service: RunnerService,
+    ) -> None:
+        """
+        Initialize the SessionService with dependencies.
+
+        Args:
+            progress_status_updater (InterfaceRedTeamProgressCallback): The callback interface for progress updates.
+            runner_service (RunnerService): The service for managing runners.
+        """
+        self.auto_red_team_test_manager = auto_red_team_test_manager
+        self.progress_status_updater = progress_status_updater
+        self.runner_service = runner_service
+        super().__init__()
+
+    @exception_handler
+    def create_new_session(
+        self, session_create_dto: SessionCreateDTO
+    ) -> SessionResponseModel:
+        """
+        Create a new session with a new runner and return the session metadata.
+
+        Args:
+            session_create_dto (SessionCreateDTO): Data transfer object containing session creation details.
+
+        Returns:
+            SessionMetadataModel: The metadata of the newly created session.
+        """
+        # Create a new runner instance using the runner service
+        runner = self.runner_service.create_runner(
+            runner_name=session_create_dto.name,
+            endpoints=session_create_dto.endpoints,
+            description=session_create_dto.description,
+            progress_callback_func=self.progress_status_updater.on_art_progress_update,
+        )
+        self.active_runner = runner
+
+        # Prepare runner arguments for session creation
+        runner_args = {
+            "context_strategy": session_create_dto.context_strategy,
+            "prompt_template": session_create_dto.prompt_template,
+            "cs_num_of_prev_prompts": session_create_dto.cs_num_of_prev_prompts,
+            "attack_module": session_create_dto.attack_module,
+            "metric": session_create_dto.metric,
+            "system_prompt": session_create_dto.system_prompt,
+        }
+
+        # Create a new session if the runner has a database instance
+        if runner.database_instance:
+            api_create_session(
+                runner.id, runner.database_instance, runner.endpoints, runner_args
+            )
+
+        # Load and return the session metadata
+        session_metadata_dict = moonshot_api.api_load_session(runner.id)
+        if session_metadata_dict is None:
+            raise ValueError(f"No session metadata found for runner ID {runner.id}")
+
+        return SessionResponseModel(
+            session_name=runner.name,
+            session_description=runner.description,
+            session=SessionMetadataModel(**session_metadata_dict),
+            chat_records=None,
+        )
+
+    @exception_handler
+    def get_all_session(self) -> list[SessionMetadataModel]:
+        """
+        Retrieve all session metadata.
+
+        Returns:
+            list[SessionMetadataModel]: A list of session metadata models for all sessions.
+        """
+        retn_session = []
+        runners_with_session = moonshot_api.api_get_all_runner()
+        sessions_metadata_dicts = moonshot_api.api_get_all_session_metadata()
+        
+        runners_dict = {runner['id']: runner for runner in runners_with_session}
+        
+        for session in sessions_metadata_dicts:
+            sess_id = session.get("session_id")
+            if sess_id in runners_dict:
+                session['description'] = runners_dict[sess_id]['description']
+                retn_session.append(session)
+        
+        return [
+            SessionMetadataModel(**metadata) for metadata in sessions_metadata_dicts
+        ]
+
+    @exception_handler
+    def get_all_sessions_names(self) -> list[str]:
+        """
+        Retrieve a list of all session names.
+
+        This method calls the moonshot API to get the names of all sessions and returns them.
+
+        Returns:
+            list[str]: A list of session names.
+        """
+        sessions = moonshot_api.api_get_all_session_names()
+        return sessions
+
+    @exception_handler
+    def get_session_by_runner_id(
+        self, runner_id: str, include_history: bool
+    ) -> SessionResponseModel:
+        """
+        Retrieve session information by runner ID, optionally including chat history.
+
+        Args:
+            runner_id (str): The unique identifier of the runner.
+            include_history (bool): Flag to determine if chat history should be included.
+
+        Returns:
+            SessionResponseModel: An object containing session metadata and optionally chat records.
+        """
+        runner: Runner = self.runner_service.load_runner(
+            runner_id, self.progress_status_updater.on_art_progress_update
+        )
+        self.active_runner = runner
+        session_metadata_dict = moonshot_api.api_load_session(self.active_runner.id)
+
+        if not isinstance(session_metadata_dict, dict):
+            raise ValueError(
+                f"Session metadata for runner ID {runner.id} must be a dictionary."
+            )
+
+        session_chat = (
+            moonshot_api.api_get_all_chats_from_session(runner.id)
+            if include_history
+            else None
+        )
+
+        return SessionResponseModel(
+            session_name=runner.name,
+            session_description=runner.description,
+            session=SessionMetadataModel(**session_metadata_dict),
+            chat_records=session_chat,
+        )
+
+    @exception_handler
+    def update_session_chat(self, runner_id: str):
+        """
+        Update the chat for the current session.
+
+        Args:
+            runner_id (str): The unique identifier of the runner.
+
+        Returns:
+            The updated chat records for the session.
+        """
+        if self.active_runner.id != runner_id:
+            raise RuntimeError("Active session and requested session do not match.")
+        session_chat = moonshot_api.api_get_all_chats_from_session(
+            self.active_runner.id
+        )
+        return session_chat
+
+    @exception_handler
+    def delete_session(self, runner_id: str) -> None:
+        """
+        Delete a session by runner ID.
+
+        This method calls the moonshot API to delete the session associated with the given runner ID.
+
+        Args:
+            runner_id (str): The unique identifier of the runner whose session is to be deleted.
+        """
+        moonshot_api.api_delete_session(runner_id)
+
+    @exception_handler
+    def select_prompt_template(self, runner_id: str, prompt_template_name: str = ""):
+        """
+        Select a prompt template for the current session.
+
+        Args:
+            runner_id (str): The unique identifier of the runner.
+            prompt_template_name (str): The name of the prompt template to be selected.
+        """
+        if self.active_runner.id != runner_id:
+            raise RuntimeError("Active session and requested session do not match.")
+        moonshot_api.api_update_prompt_template(
+            self.active_runner.id, prompt_template_name
+        )
+
+    @exception_handler
+    def select_ctx_strategy(
+        self, runner_id: str, ctx_strategy_name: str = "", num_of_prompt: int = 5
+    ) -> None:
+        """
+        Select a context strategy for the current session.
+
+        Args:
+            runner_id (str): The unique identifier of the runner.
+            ctx_strategy_name (str): The name of the context strategy to be selected.
+            num_of_prompt (int): The number of previous prompts to consider in the strategy.
+        """
+        if self.active_runner.id != runner_id:
+            raise RuntimeError("Active session and requested session do not match.")
+        moonshot_api.api_update_context_strategy(
+            self.active_runner.id, ctx_strategy_name
+        )
+        moonshot_api.api_update_cs_num_of_prev_prompts(
+            self.active_runner.id, num_of_prompt
+        )
+
+    @exception_handler
+    def select_attack_module(self, runner_id: str, attack_module_name: str = ""):
+        """
+        Select an attack module for the current session.
+
+        Args:
+            runner_id (str): The unique identifier of the runner.
+            attack_module_name (str): The name of the attack module to be selected.
+        """
+        if self.active_runner.id != runner_id:
+            raise RuntimeError("Active session and requested session do not match.")
+        moonshot_api.api_update_attack_module(self.active_runner.id, attack_module_name)
+
+    @exception_handler
+    def select_metric(self, runner_id: str, metric_name: str = ""):
+        """
+        Select a metric for the current session.
+
+        Args:
+            runner_id (str): The unique identifier of the runner.
+            metric_name (str): The name of the metric to be selected.
+        """
+        if self.active_runner.id != runner_id:
+            raise RuntimeError("Active session and requested session do not match.")
+        moonshot_api.api_update_metric(self.active_runner.id, metric_name)
+
+    @exception_handler
+    def update_system_prompt(self, runner_id: str, system_prompt: str = ""):
+        """
+        Update the system prompt for the current session.
+
+        Args:
+            runner_id (str): The unique identifier of the runner.
+            system_prompt (str): The new system prompt to be set.
+        """
+        if self.active_runner.id != runner_id:
+            raise RuntimeError("Active session and requested session do not match.")
+        moonshot_api.api_update_system_prompt(self.active_runner.id, system_prompt)
+
+    @exception_handler
+    async def send_prompt(
+        self, runner_id: str, prompt: SessionPromptDTO, batch_size: int = 5
+    ) -> PromptResponseModel | str:
+        """
+        Send a prompt to the runner for processing.
+
+        Args:
+            runner_id (str): The unique identifier of the runner.
+            prompt (SessionPromptDTO): The prompt data transfer object containing user prompt details.
+
+        Returns:
+            The result of the red teaming operation initiated by the prompt.
+        """
+        if self.active_runner.id != runner_id:
+            raise RuntimeError("Active session and requested session do not match.")
+
+        session_metadata = moonshot_api.api_load_session(self.active_runner.id)
+        if session_metadata is None:
+            raise RuntimeError(
+                f"No session metadata found for runner ID: {self.active_runner.id}"
+            )
+
+        def get_metadata_value(key, default=""):
+            return session_metadata.get(key, default)
+
+        prompt_template = get_metadata_value("prompt_template")
+        context_strategy = get_metadata_value("context_strategy")
+        num_of_prev_prompts = get_metadata_value("cs_num_of_prev_prompts")
+        attack_module = get_metadata_value("attack_module")
+        system_prompt = get_metadata_value("system_prompt")
+        metric = get_metadata_value("metric")
+
+        rt_args = {
+            "prompt": prompt.user_prompt,
+            "system_prompt": system_prompt,
+            "context_strategy_info": [
+                {
+                    "context_strategy_id": context_strategy,
+                    "num_of_prev_prompts": num_of_prev_prompts,
+                }
+            ]
+            if context_strategy
+            else [],
+            "prompt_template_ids": [prompt_template] if prompt_template else [],
+        }
+
+        if attack_module:
+            rt_args["attack_module_id"] = attack_module
+            rt_args["metric_ids"] = [metric] if metric else []
+            id = await self.auto_red_team_test_manager.schedule_art_task(
+                rt_args, self.active_runner, batch_size
+            )
+            return id
+        else:
+            response = await self.active_runner.run_red_teaming(
+                {"manual_rt_args": rt_args}
+            )
+            return PromptResponseModel.model_validate(response)
+
+    @exception_handler
+    async def cancel_auto_redteam(self, runner_id: str):
+        if self.active_runner.id != runner_id:
+            raise RuntimeError("Active session and requested session do not match.")
+
+        await self.auto_red_team_test_manager.cancel_task(self.active_runner.id)
+
+    @exception_handler
+    async def end_session(self, runner_id: str):
+        """
+        End the current session.
+
+        Args:
+            runner_id (str): The unique identifier of the runner whose session is to be ended.
+        """
+        if self.active_runner.id != runner_id:
+            raise RuntimeError("Active session and requested session do not match.")
+
+        self.active_runner.close()

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

@@ -0,0 +1,41 @@
+import asyncio
+from typing import Any, Callable
+from pydantic import ValidationError
+
+class ServiceException(Exception):
+    error_code: str
+    msg: str
+
+    def __init__(self, msg: str, method_name: str, error_code: str = 'UnknownSessionError'):
+        self.error_code = error_code
+        message = f"[ServiceException] {error_code} in {method_name} - {msg}"
+        self.msg = message 
+        super().__init__(message)
+
+def exception_handler(func: Callable[..., Any]) -> Callable[..., Any]:
+    if asyncio.iscoroutinefunction(func):
+        async def async_wrapper(*args: Any, **kwargs: Any) -> Any:
+            try:
+                return await func(*args, **kwargs)
+            except FileNotFoundError as e:
+                raise ServiceException(f"A file not found error occurred: {e}", func.__name__, "FileNotFound")
+            except ValidationError as e:
+                raise ServiceException(f"A validation error occurred: {e}", func.__name__, "ValidationError")
+            except ValueError as e:
+                raise ServiceException(f"An value error occurred: {e}", func.__name__, "ValueError")
+            except Exception as e:
+                raise ServiceException(f"An unexpected error occurred: {e}", func.__name__, "UnexpectedError")
+        return async_wrapper
+    else:
+        def wrapper(*args: Any, **kwargs: Any) -> Any:
+            try:
+                return func(*args, **kwargs)
+            except FileNotFoundError as e:
+                raise ServiceException(f"A file not found error occurred: {e}", func.__name__, "FileNotFound")
+            except ValidationError as e:
+                raise ServiceException(f"A validation error occurred: {e}", func.__name__, "ValidationError")
+            except ValueError as e:
+                raise ServiceException(f"An value error occurred: {e}", func.__name__, "ValueError")
+            except Exception as e:
+                raise ServiceException(f"An unexpected error occurred: {e}", func.__name__, "UnexpectedError")
+        return wrapper