aiverify-moonshot 0.4.0__py3-none-any.whl

moonshot/integrations/web_api/routes/runner.py

@@ -0,0 +1,239 @@
+from dependency_injector.wiring import Provide, inject
+from fastapi import APIRouter, Depends, HTTPException
+
+from ..container import Container
+from ..services.runner_service import RunnerService
+from ..services.utils.exceptions_handler import ServiceException
+
+router = APIRouter(tags=["Runner"])
+
+
+@router.get("/api/v1/runners")
+@inject
+def get_all_runners(
+    runner_service: RunnerService = Depends(Provide[Container.runner_service]),
+) -> list:
+    """
+    Retrieve all runners from the database.
+
+    Args:
+        runner_service (RunnerService): The service responsible for retrieving runners.
+
+    Returns:
+        list: A list of all runners.
+
+    Raises:
+        HTTPException: An error with status code 404 if no runners are found.
+                       An error with status code 400 if there is a validation error.
+                       An error with status code 500 for any other server-side error.
+    """
+    try:
+        return runner_service.get_all_runner()
+    except ServiceException as e:
+        if e.error_code == "FileNotFound":
+            raise HTTPException(
+                status_code=404, detail=f"Failed to retrieve runners: {e.msg}"
+            )
+        elif e.error_code == "ValidationError":
+            raise HTTPException(
+                status_code=400, detail=f"Failed to retrieve runners: {e.msg}"
+            )
+        else:
+            raise HTTPException(
+                status_code=500, detail=f"Failed to retrieve runners: {e.msg}"
+            )
+
+
+@router.get("/api/v1/runners/name")
+@inject
+def get_all_runner_name(
+    runner_service: RunnerService = Depends(Provide[Container.runner_service]),
+) -> list:
+    """
+    Retrieve all runner names from the database.
+
+    Args:
+        runner_service (RunnerService): The service responsible for retrieving runner names.
+
+    Returns:
+        list: A list of all runner names.
+
+    Raises:
+        HTTPException: An error with status code 404 if no runner names are found.
+                       An error with status code 400 if there is a validation error.
+                       An error with status code 500 for any other server-side error.
+    """
+    try:
+        return runner_service.get_all_runner_name()
+    except ServiceException as e:
+        if e.error_code == "FileNotFound":
+            raise HTTPException(
+                status_code=404, detail=f"Failed to retrieve runner names: {e.msg}"
+            )
+        elif e.error_code == "ValidationError":
+            raise HTTPException(
+                status_code=400, detail=f"Failed to retrieve runner names: {e.msg}"
+            )
+        else:
+            raise HTTPException(
+                status_code=500, detail=f"Failed to retrieve runner names: {e.msg}"
+            )
+
+
+@router.get("/api/v1/runners/{runner_id}")
+@inject
+def get_runner_by_id(
+    runner_id: str,
+    runner_service: RunnerService = Depends(Provide[Container.runner_service]),
+) -> dict:
+    """
+    Retrieve a specific runner by their ID from the database.
+
+    Args:
+        runner_id (str): The unique identifier of the runner to retrieve.
+        runner_service (RunnerService): The service responsible for retrieving the runner.
+
+    Returns:
+        dict: The runner information if found.
+
+    Raises:
+        HTTPException: An error with status code 404 if the runner is not found.
+                       An error with status code 400 if there is a validation error.
+                       An error with status code 500 for any other server-side error.
+    """
+    try:
+        run = runner_service.get_runner_by_id(runner_id)
+        return run
+    except ServiceException as e:
+        if e.error_code == "FileNotFound":
+            raise HTTPException(
+                status_code=404, detail=f"Failed to retrieve runner: {e.msg}"
+            )
+        elif e.error_code == "ValidationError":
+            raise HTTPException(
+                status_code=400, detail=f"Failed to retrieve runner: {e.msg}"
+            )
+        else:
+            raise HTTPException(
+                status_code=500, detail=f"Failed to retrieve runner: {e.msg}"
+            )
+
+
+@router.delete("/api/v1/runners/{runner_id}")
+@inject
+def delete_runner(
+    runner_id: str,
+    runner_service: RunnerService = Depends(Provide[Container.runner_service]),
+) -> dict:
+    """
+    Delete a specific runner by their ID from the database.
+
+    Args:
+        runner_id (str): The unique identifier of the runner to delete.
+        runner_service (RunnerService): The service responsible for deleting the runner.
+
+    Returns:
+        dict: A message indicating the successful deletion of the runner.
+
+    Raises:
+        HTTPException: An error with status code 404 if the runner is not found.
+                       An error with status code 400 if there is a validation error.
+                       An error with status code 500 for any other server-side error.
+    """
+    try:
+        runner_service.delete_run(runner_id)
+        return {"message": "Runner deleted successfully"}
+    except ServiceException as e:
+        if e.error_code == "FileNotFound":
+            raise HTTPException(
+                status_code=404, detail=f"Failed to delete runner: {e.msg}"
+            )
+        elif e.error_code == "ValidationError":
+            raise HTTPException(
+                status_code=400, detail=f"Failed to delete runner: {e.msg}"
+            )
+        else:
+            raise HTTPException(
+                status_code=500, detail=f"Failed to delete runner: {e.msg}"
+            )
+
+
+@router.get("/api/v1/runners/{runner_id}/runs/{run_id}")
+@inject
+def get_run_details_by_runner(
+    runner_id: str,
+    run_id: str,
+    runner_service: RunnerService = Depends(Provide[Container.runner_service]),
+) -> dict:
+    """
+    Retrieve the details of a specific run by a runner.
+
+    Args:
+        runner_id (str): The unique identifier of the runner.
+        run_id (str): The unique identifier of the run.
+        runner_service (RunnerService): The service responsible for retrieving run details.
+
+    Returns:
+        dict: The details of the run.
+
+    Raises:
+        HTTPException: An error with status code 404 if the run details are not found.
+                       An error with status code 400 if there is a validation error.
+                       An error with status code 500 for any other server-side error.
+    """
+    try:
+        return runner_service.get_run_details_by_runner(runner_id, run_id)
+    except ServiceException as e:
+        if e.error_code == "FileNotFound":
+            raise HTTPException(
+                status_code=404,
+                detail=f"Failed to get run details from runner: {e.msg}",
+            )
+        elif e.error_code == "ValidationError":
+            raise HTTPException(
+                status_code=400,
+                detail=f"Failed to get run details from runner: {e.msg}",
+            )
+        else:
+            raise HTTPException(
+                status_code=500,
+                detail=f"Failed to get run details from runner: {e.msg}",
+            )
+
+
+@router.get("/api/v1/runners/{runner_id}/runs")
+@inject
+def get_runs_id_in_runner(
+    runner_id: str,
+    runner_service: RunnerService = Depends(Provide[Container.runner_service]),
+) -> list[int]:
+    """
+    Retrieve a list of run identifiers associated with a specific runner.
+
+    Args:
+        runner_id (str): The unique identifier of the runner.
+        runner_service (RunnerService): The service responsible for retrieving the list of runs.
+
+    Returns:
+        List[str]: A list of run identifiers.
+
+    Raises:
+        HTTPException: An error with status code 404 if no runs are found for the runner.
+                       An error with status code 400 if there is a validation error.
+                       An error with status code 500 for any other server-side error.
+    """
+    try:
+        return runner_service.get_runs_id_in_runner(runner_id)
+    except ServiceException as e:
+        if e.error_code == "FileNotFound":
+            raise HTTPException(
+                status_code=404, detail=f"Failed to get runs from runner: {e.msg}"
+            )
+        elif e.error_code == "ValidationError":
+            raise HTTPException(
+                status_code=400, detail=f"Failed to get runs from runner: {e.msg}"
+            )
+        else:
+            raise HTTPException(
+                status_code=500, detail=f"Failed to get runs from runner: {e.msg}"
+            )

moonshot/integrations/web_api/schemas/benchmark_runner_dto.py

@@ -0,0 +1,13 @@
+from pydantic import BaseModel, ConfigDict
+
+
+class BenchmarkRunnerDTO(BaseModel):
+    model_config = ConfigDict(from_attributes=True)
+    run_name: str
+    description: str
+    endpoints: list[str]
+    inputs: list[str]
+    num_of_prompts: int
+    random_seed: int
+    system_prompt: str
+    runner_processing_module: str

moonshot/integrations/web_api/schemas/cookbook_create_dto.py

@@ -0,0 +1,19 @@
+from typing import Optional
+
+from moonshot.src.cookbooks.cookbook_arguments import (
+    CookbookArguments as CookbookPydanticModel,
+)
+from pydantic import Field
+
+class CookbookCreateDTO(CookbookPydanticModel):
+    id: Optional[str] = None
+    name: str = Field(..., min_length=1)
+    description: Optional[str] = Field(default="", min_length=1)
+    recipes: list[str] = Field(..., min_length=1)
+
+
+class CookbookUpdateDTO(CookbookPydanticModel):
+    id: Optional[str] = None
+    name: Optional[str] = Field(default=None, min_length=1)
+    description: Optional[str] = Field(default=None, min_length=1)
+    recipes: Optional[list[str]] = Field(default=None, min_length=1)

moonshot/integrations/web_api/schemas/cookbook_response_model.py

@@ -0,0 +1,9 @@
+from typing import Optional
+
+from moonshot.src.cookbooks.cookbook_arguments import (
+    CookbookArguments as CookbookPydanticModel,
+)
+
+
+class CookbookResponseModel(CookbookPydanticModel):
+    total_prompt_in_cookbook: Optional[int] = None

moonshot/integrations/web_api/schemas/dataset_response_dto.py

@@ -0,0 +1,9 @@
+from pydantic import BaseModel
+
+
+class DatasetResponseDTO(BaseModel):
+    id: str
+    name: str
+    description: str
+    num_of_dataset_prompts: int
+    created_date: str

moonshot/integrations/web_api/schemas/endpoint_create_dto.py

@@ -0,0 +1,21 @@
+from typing import Optional
+
+from moonshot.src.connectors_endpoints.connector_endpoint_arguments import (
+    ConnectorEndpointArguments as ConnectorEndpointPydanticModel,
+)
+
+
+class EndpointCreateDTO(ConnectorEndpointPydanticModel):
+    id: Optional[str] = None
+
+
+class EndpointUpdateDTO(ConnectorEndpointPydanticModel):
+    id: Optional[str] = None
+    name: Optional[str] = None
+    connector_type: Optional[str] = None
+    uri: Optional[str] = None
+    token: Optional[str] = None
+    max_calls_per_second: Optional[int] = None
+    max_concurrency: Optional[int] = None
+    params: Optional[dict] = None
+    created_date: Optional[str] = None

moonshot/integrations/web_api/schemas/endpoint_response_model.py

@@ -0,0 +1,11 @@
+from moonshot.src.connectors_endpoints.connector_endpoint_arguments import (
+    ConnectorEndpointArguments as ConnectorEndpointPydanticModel,
+)
+
+
+class EndpointDataModel(ConnectorEndpointPydanticModel):
+    pass
+
+    def mask_token(self):
+        if self.token:
+            self.token = "*" * len(self.token)

moonshot/integrations/web_api/schemas/prompt_response_model.py

@@ -0,0 +1,14 @@
+from pydantic import BaseModel, RootModel
+
+from ..schemas.session_response_model import ChatRecord
+
+
+
+class PromptInfo(BaseModel):
+    current_runner_id: str
+    current_chats: dict[str, list[ChatRecord]]
+    current_batch_size: int
+    current_status: str
+
+
+PromptResponseModel = RootModel[PromptInfo]

moonshot/integrations/web_api/schemas/prompt_template_response_model.py

@@ -0,0 +1,10 @@
+from pydantic import BaseModel, RootModel
+
+
+class PromptTemplate(BaseModel):
+    name: str
+    description: str
+    template: str
+
+
+PromptTemplatesResponseModel = RootModel[list[PromptTemplate]]

moonshot/integrations/web_api/schemas/recipe_create_dto.py

@@ -0,0 +1,32 @@
+from typing import Optional
+
+from moonshot.src.recipes.recipe_arguments import RecipeArguments as RecipePydanticModel
+
+from pydantic import Field
+
+class RecipeCreateDTO(RecipePydanticModel):
+    id: Optional[str] = None
+    name: str = Field(..., min_length=1)
+    description: Optional[str] = Field(default="", min_length=1)
+    tags: Optional[list[str]] = None
+    categories: Optional[list[str]] = None
+    datasets: list[str] = Field(..., min_length=1)
+    metrics: list[str] = Field(..., min_length=1)
+    prompt_templates: Optional[list[str]] = None
+    attack_modules: Optional[list[str]] = None
+    grading_scale: Optional[dict[str, list[int]]] = None
+    stats: Optional[dict] = None
+
+
+class RecipeUpdateDTO(RecipePydanticModel):
+    id: Optional[str] = None
+    name: Optional[str] = None
+    description: Optional[str] = Field(default="", min_length=1)
+    tags: Optional[list[str]] = None
+    categories: Optional[list[str]] = None
+    datasets: Optional[list[str]] = None
+    prompt_templates: Optional[list[str]] = None
+    metrics: Optional[list[str]] = None
+    attack_modules: Optional[list[str]] = None
+    grading_scale: Optional[dict[str, list[int]]] = None
+    stats: Optional[dict] = None

moonshot/integrations/web_api/schemas/recipe_response_model.py

@@ -0,0 +1,7 @@
+from typing import Optional
+
+from moonshot.src.recipes.recipe_arguments import RecipeArguments as RecipePydanticModel
+
+
+class RecipeResponseModel(RecipePydanticModel):
+    total_prompt_in_recipe: Optional[int] = None

moonshot/integrations/web_api/schemas/session_create_dto.py

@@ -0,0 +1,16 @@
+from typing import Optional
+
+from pydantic import BaseModel, ConfigDict, Field
+
+
+class SessionCreateDTO(BaseModel):
+    model_config = ConfigDict(from_attributes=True)
+    name: str = Field(min_length=3)
+    description: Optional[str] = Field(default="", max_length=1000)
+    endpoints: list[str] = Field(min_length=1)
+    context_strategy: str = ""
+    prompt_template: str = ""
+    attack_module: str = ""
+    metric: str = ""
+    system_prompt: str = ""
+    cs_num_of_prev_prompts: int = 5

moonshot/integrations/web_api/schemas/session_prompt_dto.py

@@ -0,0 +1,7 @@
+from pydantic import BaseModel, ConfigDict, Field
+
+
+class SessionPromptDTO(BaseModel):
+    model_config = ConfigDict(from_attributes=True)
+    history_length: int | None = 10
+    user_prompt: str = Field(min_length=1)

moonshot/integrations/web_api/schemas/session_response_model.py

@@ -0,0 +1,38 @@
+from typing import Dict, List, Optional
+
+from pydantic import BaseModel, Field
+
+
+class ChatRecord(BaseModel):
+    conn_id: str
+    context_strategy: Optional[str]
+    prompt_template: Optional[str]
+    attack_module: Optional[str]
+    metric: Optional[str]
+    prompt: str
+    prepared_prompt: str
+    system_prompt: Optional[str]
+    predicted_result: str
+    duration: str
+    prompt_time: str
+
+
+class SessionMetadataModel(BaseModel):
+    session_id: str = Field(min_length=1)
+    description: Optional[str] = ""
+    endpoints: List[str] = Field(min_length=1)
+    created_epoch: str
+    created_datetime: str = Field(min_length=1)
+    prompt_template: Optional[str]
+    context_strategy: Optional[str]
+    cs_num_of_prev_prompts: int | None = 5
+    attack_module: Optional[str]
+    metric: Optional[str]
+    system_prompt: Optional[str]
+
+
+class SessionResponseModel(BaseModel):
+    session_name: str
+    session_description: str
+    session: SessionMetadataModel
+    chat_records: Optional[Dict[str, List[ChatRecord]]]

moonshot/integrations/web_api/services/attack_module_service.py

@@ -0,0 +1,34 @@
+from .... import api as moonshot_api
+from .base_service import BaseService
+from .utils.exceptions_handler import exception_handler
+
+
+class AttackModuleService(BaseService):
+    @exception_handler
+    def get_all_attack_module(self) -> list[str]:
+        """
+        Retrieves a list of all available attack modules.
+
+        This method wraps the `api_get_all_attack_modules` function from the moonshot API and returns
+        a list of attack module names.
+
+        Returns:
+            list[str]: A list of strings, each denoting the name of an attack module.
+        """
+        attack_modules = moonshot_api.api_get_all_attack_modules()
+        return attack_modules
+
+    @exception_handler
+    def get_all_attack_module_metadata(self) -> list[dict]:
+        """
+        Retrieves metadata for all available attack modules.
+
+        This method calls the `api_get_all_attack_module_metadata` function from the moonshot API. It collects
+        metadata for each attack module, which includes details such as the module's name, description, and
+        any other relevant information provided by the module's `get_metadata` method.
+
+        Returns:
+            list[dict]: A list of dictionaries, each containing metadata of an attack module.
+        """
+        am_metadata = moonshot_api.api_get_all_attack_module_metadata()
+        return am_metadata

moonshot/integrations/web_api/services/auto_red_team_test_manager.py

@@ -0,0 +1,86 @@
+import asyncio
+import uuid
+from typing import Any
+
+from dependency_injector.wiring import inject
+
+from moonshot.src.runners.runner import Runner
+
+from ..services.auto_red_team_test_state import AutoRedTeamTestState
+from ..services.base_service import BaseService
+from ..services.runner_service import RunnerService
+from ..status_updater.interface.redteam_progress_callback import (
+    InterfaceRedTeamProgressCallback,
+)
+from ..types.types import RedTeamTestProgress
+
+
+class AutoRedTeamTestManager(BaseService):
+    @inject
+    def __init__(
+        self,
+        auto_red_team_test_state: AutoRedTeamTestState,
+        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_state = auto_red_team_test_state
+        self.progress_status_updater = progress_status_updater
+        self.runner_service = runner_service
+        super().__init__()
+
+    def generate_unique_task_id(self) -> str:
+        unique_id = str(uuid.uuid4())
+        return f"task_{unique_id}"
+
+    def add_task(
+        self, executor_id: str, task: asyncio.Task[Any], moonshot_runner: Runner
+    ) -> None:
+        self.auto_red_team_test_state.add_task(executor_id, task, moonshot_runner)
+
+    def remove_task(self, executor_id: str) -> None:
+        self.auto_red_team_test_state.remove_task(executor_id)
+
+    async def cancel_task(self, executor_id: str) -> None:
+        await self.auto_red_team_test_state.cancel_task(executor_id)
+
+    def update_progress_status(self, updates: RedTeamTestProgress):
+        self.auto_red_team_test_state.update_progress_status(updates)
+
+    def on_task_completed(self, task: asyncio.Task[Any]) -> None:
+        self.logger.debug(f"Task {task.get_name()} has completed")
+
+    async def schedule_art_task(
+        self, art_args: dict, active_runner: Runner, batch_size: int
+    ) -> str:
+        task_id = self.generate_unique_task_id()
+
+        auto_redteam_coroutine = self.send_art_prompt(
+            art_args, active_runner, batch_size
+        )
+
+        task = asyncio.create_task(auto_redteam_coroutine, name=task_id)
+
+        def on_executor_completion(task: asyncio.Task[Any]):
+            if task.exception():
+                self.logger.error(
+                    f"Executor {task.get_name()} has failed - {task.exception()}"
+                )
+            else:
+                self.logger.debug(f"Executor {task.get_name()} has completed")
+
+        task.add_done_callback(on_executor_completion)
+
+        self.add_task(active_runner.id, task, active_runner)
+        return active_runner.id
+
+    async def send_art_prompt(self, art_args, active_runner: Runner, batch_size: int):
+        await active_runner.run_red_teaming(
+            {"attack_strategies": [art_args], "chat_batch_size": batch_size}
+        )

moonshot/integrations/web_api/services/auto_red_team_test_state.py

@@ -0,0 +1,57 @@
+import asyncio
+from typing import Any, TypedDict
+
+from moonshot.integrations.web_api.services.base_service import BaseService
+from moonshot.integrations.web_api.types.types import RedTeamTestProgress
+from moonshot.src.runners.runner import Runner
+
+
+class AutoRedTeamTaskInfo(TypedDict):
+    async_task: asyncio.Task[Any]
+    moonshot_runner: Runner
+    status: RedTeamTestProgress | None
+
+
+class AutoRedTeamTestState(BaseService):
+    state: dict[str, AutoRedTeamTaskInfo] = {}
+
+    def add_task(
+        self, id: str, task: asyncio.Task[Any], moonshot_runner: Runner
+    ) -> None:
+        self.state[id] = {
+            "async_task": task,
+            "moonshot_runner": moonshot_runner,
+            "status": None,
+        }
+
+    def remove_task(self, id: str) -> AutoRedTeamTaskInfo:
+        return self.state.pop(id)
+
+    async def cancel_task(self, id: str) -> None:
+        task_info = self.state[id]
+        if task_info and not task_info["async_task"].done():
+            if task_info["moonshot_runner"]:
+                await task_info["moonshot_runner"].cancel()
+            task_info["async_task"].cancel()
+            self.logger.debug(f"Task {id} has been cancelled")
+        else:
+            self.logger.debug(f"Task {id} is already completed or does not exist.")
+
+    def update_progress_status(self, updates: RedTeamTestProgress) -> None:
+        self.logger.debug(updates)
+        self.state[updates["current_runner_id"]]["status"] = updates
+
+    def get_state(self) -> dict[str, AutoRedTeamTaskInfo]:
+        return self.state
+
+    def get_progress_status(self, id: str) -> RedTeamTestProgress | None:
+        if id in self.state:
+            return self.state[id]["status"]
+        return None
+
+    def get_all_progress_status(self) -> dict[str, RedTeamTestProgress]:
+        return {
+            id: task_info["status"]
+            for id, task_info in self.state.items()
+            if task_info["status"] is not None
+        }

moonshot/integrations/web_api/services/base_service.py

@@ -0,0 +1,8 @@
+import logging
+
+
+class BaseService:
+    def __init__(self) -> None:
+        self.logger = logging.getLogger(
+            f"{__name__}.{self.__class__.__name__}",
+        )

moonshot/integrations/web_api/services/benchmark_result_service.py

@@ -0,0 +1,25 @@
+from .... import api as moonshot_api
+from ..services.base_service import BaseService
+from ..services.utils.exceptions_handler import exception_handler
+from ..types.types import BenchmarkResult
+
+
+class BenchmarkResultService(BaseService):
+    @exception_handler
+    def get_all_results(self) -> list[BenchmarkResult]:
+        results = moonshot_api.api_get_all_result()
+        return [BenchmarkResult(**result) for result in results]
+
+    @exception_handler
+    def get_all_result_name(self) -> list[str] | None:
+        results = moonshot_api.api_get_all_result_name()
+        return results
+
+    @exception_handler
+    def get_result_by_id(self, result_id: str) -> BenchmarkResult:
+        result = moonshot_api.api_read_result(result_id)
+        return BenchmarkResult(**result)
+
+    @exception_handler
+    def delete_result(self, result_id: str) -> None:
+        moonshot_api.api_delete_result(result_id)