aiverify-moonshot 0.4.0__py3-none-any.whl

moonshot/integrations/web_api/container.py

@@ -0,0 +1,146 @@
+import importlib.resources
+
+from dependency_injector import containers, providers
+
+from .services.attack_module_service import AttackModuleService
+from .services.auto_red_team_test_manager import AutoRedTeamTestManager
+from .services.auto_red_team_test_state import AutoRedTeamTestState
+from .services.benchmark_result_service import BenchmarkResultService
+from .services.benchmark_test_manager import BenchmarkTestManager
+from .services.benchmark_test_state import BenchmarkTestState
+from .services.benchmarking_service import BenchmarkingService
+from .services.context_strategy_service import ContextStrategyService
+from .services.cookbook_service import CookbookService
+from .services.dataset_service import DatasetService
+from .services.endpoint_service import EndpointService
+from .services.metric_service import MetricService
+from .services.prompt_template_service import PromptTemplateService
+from .services.recipe_service import RecipeService
+from .services.runner_service import RunnerService
+from .services.session_service import SessionService
+from .status_updater.moonshot_ui_webhook import MoonshotUIWebhook
+
+
+class Container(containers.DeclarativeContainer):
+    config = providers.Configuration("config")
+    config.from_dict(
+        {
+            "app_environment": "DEV",
+            "asyncio": {
+                "monitor_task": False,
+            },
+            "ssl": {
+                "enabled": False,
+                "file_path": str(
+                    importlib.resources.files("moonshot").joinpath(
+                        "integrations/web_api/certs"
+                    )
+                ),
+                "cert_filename": "cert.pem",
+                "key_filename": "key.pem",
+            },
+            "cors": {
+                "enabled": False,
+                "allowed_origins": "http://localhost:3000",
+            },
+            "log": {
+                "logging": True,
+                "level": "DEBUG",
+                "format": "[%(asctime)s] [%(levelname)s] [%(name)s]: %(message)s",
+                "log_file_path": str(
+                    importlib.resources.files("moonshot").joinpath(
+                        "integrations/web_api/log"
+                    )
+                ),
+                "log_file_max_size": 5242880,
+                "log_file_backup_count": 3,
+            },
+        }
+    )
+
+    benchmark_test_state: providers.Singleton[BenchmarkTestState] = providers.Singleton(
+        BenchmarkTestState
+    )
+    auto_red_team_test_state: providers.Singleton[
+        AutoRedTeamTestState
+    ] = providers.Singleton(AutoRedTeamTestState)
+    webhook: providers.Singleton[MoonshotUIWebhook] = providers.Singleton(
+        MoonshotUIWebhook,
+        benchmark_test_state=benchmark_test_state,
+        auto_red_team_test_state=auto_red_team_test_state,
+    )
+    runner_service: providers.Singleton[RunnerService] = providers.Singleton(
+        RunnerService
+    )
+    auto_red_team_test_manager: providers.Singleton[
+        AutoRedTeamTestManager
+    ] = providers.Singleton(
+        AutoRedTeamTestManager,
+        auto_red_team_test_state=auto_red_team_test_state,
+        progress_status_updater=webhook,
+        runner_service=runner_service,
+    )
+    benchmark_test_manager: providers.Singleton[
+        BenchmarkTestManager
+    ] = providers.Singleton(
+        BenchmarkTestManager,
+        benchmark_test_state=benchmark_test_state,
+        progress_status_updater=webhook,
+        runner_service=runner_service,
+    )
+    session_service: providers.Singleton[SessionService] = providers.Singleton(
+        SessionService,
+        auto_red_team_test_manager=auto_red_team_test_manager,
+        progress_status_updater=webhook,
+        runner_service=runner_service,
+    )
+    prompt_template_service: providers.Singleton[
+        PromptTemplateService
+    ] = providers.Singleton(PromptTemplateService)
+    context_strategy_service: providers.Singleton[ContextStrategyService] = providers.Singleton(
+        ContextStrategyService
+    )
+    benchmarking_service: providers.Singleton[
+        BenchmarkingService
+    ] = providers.Singleton(
+        BenchmarkingService, benchmark_test_manager=benchmark_test_manager
+    )
+    endpoint_service: providers.Singleton[EndpointService] = providers.Singleton(
+        EndpointService
+    )
+    recipe_service: providers.Singleton[RecipeService] = providers.Singleton(
+        RecipeService
+    )
+    cookbook_service: providers.Singleton[CookbookService] = providers.Singleton(
+        CookbookService
+    )
+    benchmark_result_service: providers.Singleton[
+        BenchmarkResultService
+    ] = providers.Singleton(BenchmarkResultService)
+    metric_service: providers.Singleton[MetricService] = providers.Singleton(
+        MetricService
+    )
+
+    dataset_service: providers.Singleton[DatasetService] = providers.Singleton(
+        DatasetService,
+    )
+    am_service: providers.Singleton[AttackModuleService] = providers.Singleton(
+        AttackModuleService,
+    )
+    wiring_config = containers.WiringConfiguration(
+        modules=[
+            ".routes.redteam",
+            ".routes.prompt_template",
+            ".routes.context_strategy",
+            ".routes.benchmark",
+            ".routes.endpoint",
+            ".routes.recipe",
+            ".routes.cookbook",
+            ".routes.benchmark_result",
+            ".routes.metric",
+            ".routes.runner",
+            ".routes.dataset",
+            ".routes.attack_modules",
+            ".services.benchmarking_service",
+        ]
+    )

moonshot/integrations/web_api/logging_conf.py

@@ -0,0 +1,114 @@
+import logging
+import os
+import sys
+from logging.handlers import RotatingFileHandler
+from typing import Literal
+
+from dependency_injector import providers
+
+from .types.types import UvicornLoggingConfig
+
+COLORS = {
+    "HEADER": "\033[95m",
+    "OKBLUE": "\033[94m",
+    "OKGREEN": "\033[92m",
+    "WARNING": "\033[93m",
+    "FAIL": "\033[91m",
+    "ENDC": "\033[0m",
+    "BOLD": "\033[1m",
+    "UNDERLINE": "\033[4m",
+    "WHITE": "\033[97m",
+}
+
+
+class ColorizedFormatter(logging.Formatter):
+    LEVEL_COLORS = {
+        logging.DEBUG: COLORS["OKBLUE"],
+        logging.INFO: COLORS["OKGREEN"],
+        logging.WARNING: COLORS["WARNING"],
+        logging.ERROR: COLORS["FAIL"],
+        logging.CRITICAL: COLORS["HEADER"],
+    }
+
+    def __init__(
+        self,
+        fmt: str,
+        datefmt: str | None = None,
+        style: Literal["%"] = "%",
+        disableColor: bool = False,
+    ):
+        super().__init__(fmt, datefmt, style)
+        self.disableColor = disableColor
+
+    def format(self, record: logging.LogRecord):
+        if self.disableColor:
+            return super().format(record)
+        else:
+            color = str(self.LEVEL_COLORS.get(record.levelno))
+            message = super().format(record)
+            return color + message + COLORS["ENDC"]
+
+
+def create_logging_dir(log_file_path: str):
+    if not os.path.exists(log_file_path):
+        os.makedirs(log_file_path)
+
+
+def configure_app_logging(cfg: providers.Configuration):
+    if cfg.log.logging():
+        create_logging_dir(cfg.log.log_file_path())
+
+    file_handler = RotatingFileHandler(
+        filename=cfg.log.log_file_path() + "/web_api.log",
+        maxBytes=cfg.log.log_file_max_size(),
+        backupCount=cfg.log.log_file_backup_count(),
+    )
+    stream_handler = logging.StreamHandler(sys.stdout)
+    stream_handler.setFormatter(ColorizedFormatter(cfg.log.format()))
+
+    logging.basicConfig(
+        handlers=[file_handler, stream_handler],
+        level=cfg.log.level(),
+        format=cfg.log.format(),
+    )
+
+    logging.info("Logging is configured.")
+
+
+def create_uvicorn_log_config(cfg: providers.Configuration) -> UvicornLoggingConfig:
+    if cfg.log.logging():
+        create_logging_dir(cfg.log.log_file_path())
+
+    return {
+        "version": 1,
+        "disable_existing_loggers": False,
+        "formatters": {
+            "default": {
+                "()": "moonshot.integrations.web_api.logging_conf.ColorizedFormatter",
+                "format": cfg.log.format(),
+            },
+            "file_formatter": {
+                "()": "moonshot.integrations.web_api.logging_conf.ColorizedFormatter",
+                "format": cfg.log.format(),
+                "disableColor": True,
+            },
+        },
+        "handlers": {
+            "file": {
+                "class": "logging.handlers.RotatingFileHandler",
+                "filename": cfg.log.log_file_path() + "/web_api.log",
+                "maxBytes": cfg.log.log_file_max_size(),
+                "backupCount": cfg.log.log_file_backup_count(),
+                "formatter": "file_formatter",
+            },
+            "console": {
+                "class": "logging.StreamHandler",
+                "stream": "ext://sys.stdout",
+                "formatter": "default",
+            },
+        },
+        "root": {
+            "level": cfg.log.level(),
+            "handlers": ["file", "console"],
+        },
+    }

moonshot/integrations/web_api/routes/attack_modules.py

@@ -0,0 +1,66 @@
+from dependency_injector.wiring import Provide, inject
+from fastapi import APIRouter, Depends, HTTPException
+
+from ..container import Container
+from ..services.attack_module_service import AttackModuleService
+from ..services.utils.exceptions_handler import ServiceException
+
+router = APIRouter(tags=["Attack Modules"])
+
+
+@router.get("/api/v1/attack-modules")
+@inject
+def get_all_attack_module(
+    am_service: AttackModuleService = Depends(Provide[Container.am_service]),
+) -> list[str]:
+    """
+    Retrieve all attack modules from the database.
+
+    Args:
+        am_service (AttackModuleService): The service responsible for fetching attack modules.
+
+    Returns:
+        list: A list of attack modules if successful.
+
+    Raises:
+        HTTPException: An error with status code 404 if attack modules file is not found.
+        HTTPException: An error with status code 400 if there is a validation error with the request.
+        HTTPException: An error with status code 500 for any other type of server-side error.
+    """
+    try:
+        return am_service.get_all_attack_module()
+    except ServiceException as e:
+        if e.error_code == "FileNotFound":
+            raise HTTPException(
+                status_code=404, detail=f"Failed to retrieve attack modules: {e.msg}"
+            )
+        elif e.error_code == "ValidationError":
+            raise HTTPException(
+                status_code=400, detail=f"Failed to retrieve attack modules: {e.msg}"
+            )
+        else:
+            raise HTTPException(
+                status_code=500, detail=f"Failed to retrieve attack modules: {e.msg}"
+            )
+
+
+@router.get("/api/v1/attack-modules/metadata")
+@inject
+def get_all_attack_module_metadata(
+    am_service: AttackModuleService = Depends(Provide[Container.am_service]),
+) -> list:
+    try:
+        return am_service.get_all_attack_module_metadata()
+    except ServiceException as e:
+        if e.error_code == "FileNotFound":
+            raise HTTPException(
+                status_code=404, detail=f"Failed to retrieve attack modules: {e.msg}"
+            )
+        elif e.error_code == "ValidationError":
+            raise HTTPException(
+                status_code=400, detail=f"Failed to retrieve attack modules: {e.msg}"
+            )
+        else:
+            raise HTTPException(
+                status_code=500, detail=f"Failed to retrieve attack modules: {e.msg}"
+            )

moonshot/integrations/web_api/routes/benchmark.py

@@ -0,0 +1,116 @@
+from dependency_injector.wiring import Provide, inject
+from fastapi import APIRouter, Depends, HTTPException
+
+from ..container import Container
+from ..schemas.benchmark_runner_dto import BenchmarkRunnerDTO
+from ..services.benchmark_test_state import BenchmarkTestState
+from ..services.benchmarking_service import BenchmarkingService
+from ..services.utils.exceptions_handler import ServiceException
+from ..types.types import BenchmarkCollectionType
+
+router = APIRouter(tags=["Benchmarking"])
+
+
+@router.post("/api/v1/benchmarks")
+@inject
+async def benchmark_executor(
+    type: BenchmarkCollectionType,
+    data: BenchmarkRunnerDTO,
+    benchmarking_service: BenchmarkingService = Depends(
+        Provide[Container.benchmarking_service]
+    ),
+) -> dict:
+    """
+    Execute a benchmark test.
+
+    Args:
+        type (BenchmarkCollectionType): The type of benchmark to execute.
+        data (BenchmarkRunnerDTO): The data required to execute the benchmark.
+        benchmarking_service (BenchmarkingService, optional): The service that will execute the benchmark.
+
+    Returns:
+        dict: A dictionary with the 'id' key containing the ID of the created execution task.
+
+    Raises:
+        HTTPException: If the provided type is invalid (status code 400) or if the service fails to create
+        and execute the benchmark (status code 500).
+    """
+    try:
+        if type is BenchmarkCollectionType.COOKBOOK:
+            id = await benchmarking_service.execute_cookbook(data)
+            return {"id": id}
+        elif type is BenchmarkCollectionType.RECIPE:
+            id = await benchmarking_service.execute_recipe(data)
+            return {"id": id}
+        else:
+            raise HTTPException(status_code=400, detail="Invalid query parameter: type")
+    except ServiceException as e:
+        raise HTTPException(
+            status_code=500, detail=f"Unable to create and execute benchmark: {e}"
+        )
+
+
+@router.get("/api/v1/benchmarks/status")
+@inject
+def get_benchmark_progress(
+    benchmark_state: BenchmarkTestState = Depends(
+        Provide[Container.benchmark_test_state]
+    ),
+):
+    """
+    Retrieve the progress status of all benchmarks.
+
+    Args:
+        benchmark_state (BenchmarkTestState, optional): The state service that tracks benchmark progress.
+
+    Returns:
+        The progress status of all benchmarks.
+
+    Raises:
+        HTTPException: If there is an error retrieving the progress status, with a status code indicating the
+        nature of the error (404 for file not found, 400 for validation error).
+    """
+    try:
+        all_status = benchmark_state.get_all_progress_status()
+        return all_status
+    except ServiceException as e:
+        if e.error_code == "FileNotFound":
+            raise HTTPException(
+                status_code=404, detail=f"Failed to retrieve progress status: {e.msg}"
+            )
+        elif e.error_code == "ValidationError":
+            raise HTTPException(
+                status_code=400, detail=f"Failed to retrieve progress status: {e.msg}"
+            )
+        else:
+            raise HTTPException(
+                status_code=500, detail=f"Failed to retrieve progress status: {e.msg}"
+            )
+
+
+@router.post("/api/v1/benchmarks/cancel/{runner_id}")
+@inject
+async def cancel_benchmark_executor(
+    runner_id: str,
+    benchmarking_service: BenchmarkingService = Depends(
+        Provide[Container.benchmarking_service]
+    ),
+):
+    """
+    Cancel a benchmark execution task.
+
+    Args:
+        runner_id (str): The ID of the runner executing the benchmark.
+        benchmarking_service (BenchmarkingService): The service that will cancel the benchmark execution.
+
+    Returns:
+        None
+
+    Raises:
+        HTTPException: If the service is unable to cancel the benchmark, with a status code
+        500 indicating an internal server error.
+    """
+    try:
+        await benchmarking_service.cancel_executor(runner_id)
+    except ServiceException as e:
+        raise HTTPException(status_code=500, detail=f"Unable to cancel benchmark: {e}")

moonshot/integrations/web_api/routes/benchmark_result.py

@@ -0,0 +1,175 @@
+from dependency_injector.wiring import Provide, inject
+from fastapi import APIRouter, Depends, HTTPException
+
+from ..container import Container
+from ..services.benchmark_result_service import BenchmarkResultService
+from ..services.utils.exceptions_handler import ServiceException
+
+router = APIRouter(tags=["Benchmark Results"])
+
+
+@router.get("/api/v1/benchmarks/results")
+@inject
+async def get_all_results(
+    benchmark_result_service: BenchmarkResultService = Depends(
+        Provide[Container.benchmark_result_service]
+    ),
+) -> list[dict]:
+    """
+    Retrieve all benchmark results.
+
+    This endpoint retrieves a list of all benchmark results from the database. Each benchmark result is
+    represented as a dictionary containing its associated data.
+
+    Args:
+        benchmark_result_service (BenchmarkResultService): The service responsible for fetching benchmark results.
+
+    Returns:
+        list[dict]: A list of dictionaries, each representing a single benchmark result.
+
+    Raises:
+        HTTPException: Raised if the results file cannot be found (404) or if an unspecified error occurs (500).
+    """
+    try:
+        results = benchmark_result_service.get_all_results()
+        return results
+    except ServiceException as e:
+        if e.error_code == "FileNotFound":
+            raise HTTPException(
+                status_code=404, detail=f"Failed to retrieve results: {e.msg}"
+            )
+        elif e.error_code == "ValidationError":
+            raise HTTPException(
+                status_code=400, detail=f"Failed to retrieve results: {e.msg}"
+            )
+        else:
+            raise HTTPException(
+                status_code=500, detail=f"Failed to retrieve results: {e.msg}"
+            )
+
+
+@router.get("/api/v1/benchmarks/results/name")
+@inject
+async def get_all_results_name(
+    benchmark_result_service: BenchmarkResultService = Depends(
+        Provide[Container.benchmark_result_service]
+    ),
+):
+    """
+    Get all benchmark result names from the database.
+
+    This endpoint retrieves the names of all benchmark results stored in the database.
+
+    Args:
+        benchmark_result_service (BenchmarkResultService): The service responsible for fetching
+        the names of the benchmark results.
+
+    Returns:
+        A list of all benchmark result names.
+
+    Raises:
+        HTTPException: An error occurred while trying to find the result names file (404),
+                       a validation error occurred (400), or
+                       an unspecified error occurred (500).
+    """
+    try:
+        results = benchmark_result_service.get_all_result_name()
+        return results
+    except ServiceException as e:
+        if e.error_code == "FileNotFound":
+            raise HTTPException(
+                status_code=404, detail=f"Failed to retrieve result name: {e.msg}"
+            )
+        elif e.error_code == "ValidationError":
+            raise HTTPException(
+                status_code=400, detail=f"Failed to retrieve result name: {e.msg}"
+            )
+        else:
+            raise HTTPException(
+                status_code=500, detail=f"Failed to retrieve result name: {e.msg}"
+            )
+
+
+@router.get("/api/v1/benchmarks/results/{result_id}")
+@inject
+async def get_one_results(
+    result_id: str,
+    benchmark_result_service: BenchmarkResultService = Depends(
+        Provide[Container.benchmark_result_service]
+    ),
+):
+    """
+    Retrieve a single benchmark result by its ID.
+
+    This endpoint fetches the details of a specific benchmark result identified by the provided result_id.
+
+    Args:
+        result_id (str): The unique identifier of the benchmark result to retrieve.
+        benchmark_result_service (BenchmarkResultService): The service responsible for fetching the benchmark result.
+
+    Returns:
+        dict: A dictionary containing the details of the benchmark result.
+
+    Raises:
+        HTTPException: An error occurred while trying to find the results file (404) or
+                       an unspecified error occurred (500).
+    """
+    try:
+        results = benchmark_result_service.get_result_by_id(result_id)
+        return results
+    except ServiceException as e:
+        if e.error_code == "FileNotFound":
+            raise HTTPException(
+                status_code=404, detail=f"Failed to retrieve result: {e.msg}"
+            )
+        elif e.error_code == "ValidationError":
+            raise HTTPException(
+                status_code=400, detail=f"Failed to retrieve result: {e.msg}"
+            )
+        else:
+            raise HTTPException(
+                status_code=500, detail=f"Failed to retrieve result: {e.msg}"
+            )
+
+@router.delete("/api/v1/benchmarks/results/{result_id}")
+@inject
+def delete_result(
+    result_id: str,
+    benchmark_result_service: BenchmarkResultService = Depends(
+        Provide[Container.benchmark_result_service]
+    ),
+) -> dict[str, str] | tuple[dict[str, str], int]:
+    """
+    Delete a benchmark result by its ID.
+
+    This endpoint deletes a specific benchmark result identified by the provided result_id.
+
+    Args:
+        result_id (str): The unique identifier of the benchmark result to delete.
+        benchmark_result_service (BenchmarkResultService): The service responsible for deleting the benchmark result.
+
+    Returns:
+        dict[str, str] | tuple[dict[str, str], int]: A message indicating successful deletion,
+        or an HTTPException with an appropriate status code.
+
+    Raises:
+        HTTPException: An error occurred while trying to delete the result due to the result not being found (404),
+                       a validation error occurred (400), or
+                       an unspecified error occurred (500).
+    """
+    try:
+        benchmark_result_service.delete_result(result_id)
+        return {"message": "Result deleted successfully"}
+    except ServiceException as e:
+        if e.error_code == "FileNotFound":
+            raise HTTPException(
+                status_code=404, detail=f"Failed to delete result: {e.msg}"
+            )
+        elif e.error_code == "ValidationError":
+            raise HTTPException(
+                status_code=400, detail=f"Failed to delete result: {e.msg}"
+            )
+        else:
+            raise HTTPException(
+                status_code=500, detail=f"Failed to delete result: {e.msg}"
+            )