aiverify-moonshot 0.4.1__py3-none-any.whl → 0.4.3__py3-none-any.whl

moonshot/integrations/cli/utils/process_data.py

@@ -0,0 +1,52 @@
+from moonshot.src.utils.find_feature import find_keyword
+from moonshot.src.utils.pagination import get_paginated_lists
+
+
+def filter_data(
+    list_of_data: list, keyword: str = "", pagination: tuple = ()
+) -> list | None:
+    """
+    Filters and paginates a list of data.
+
+    Args:
+        list_of_data (list): The list of data to be filtered and paginated.
+        keyword (str, optional): The keyword to filter the data. Defaults to "".
+        pagination (tuple, optional): A tuple containing the page number and page size. Defaults to ().
+
+    Returns:
+        list | None: The filtered and paginated list of data, or None if no data matches the criteria.
+    """
+    # if there is a find keyword
+    if keyword:
+        list_of_data = find_keyword(keyword, list_of_data)
+        if not list_of_data:
+            return
+
+    # if pagination is required
+    if pagination:
+        # add index to every dictionary in the list
+        if all(isinstance(item, dict) for item in list_of_data):
+            for index, item in enumerate(list_of_data, 1):
+                if isinstance(item, dict):
+                    item["idx"] = index
+
+        # get paginated data
+        page_number = pagination[0]
+        page_size = pagination[1]
+
+        if page_number <= 0 or page_size <= 0:
+            # print("Invalid page number or page size. Page number and page size should start from 1.")
+            raise RuntimeError(
+                "Invalid page number or page size. Page number and page size should start from 1."
+            )
+
+        paginated_data = get_paginated_lists(page_size, list_of_data)
+
+        # perform index checks
+        paginated_data_size = len(paginated_data)
+        if page_number > paginated_data_size:
+            list_of_data = paginated_data[(paginated_data_size - 1)]
+        else:
+            list_of_data = paginated_data[(page_number - 1)]
+
+    return list_of_data

moonshot/integrations/web_api/__main__.py

@@ -7,6 +7,7 @@ from dotenv import dotenv_values, load_dotenv
 from .app import create_app
 from .container import Container
 from .logging_conf import configure_app_logging, create_uvicorn_log_config
+from .services.utils.file_manager import create_temp_dir
 from .types.types import UvicornRunArgs
 
 
@@ -20,6 +21,7 @@ def start_app():
     else:
         container.config.from_yaml(f"{config_file}", required=True)
 
+    create_temp_dir(container.config.temp_folder())
     configure_app_logging(container.config)
     logging.info(f"Environment: {container.config.app_environment()}")
     ENABLE_SSL = container.config.ssl.enabled()

moonshot/integrations/web_api/app.py

@@ -14,12 +14,13 @@ from .routes import (
     attack_modules,
     benchmark,
     benchmark_result,
+    bookmark,
+    context_strategy,
     cookbook,
     dataset,
     endpoint,
     metric,
     prompt_template,
-    context_strategy,
     recipe,
     runner,
 )
@@ -66,13 +67,11 @@ def create_app(cfg: providers.Configuration) -> CustomFastAPI:
 
     app_kwargs["swagger_ui_parameters"] = {
         "defaultModelsExpandDepth": -1,
-        "docExpansion": None
-        }
+        "docExpansion": None,
+    }
 
     app: CustomFastAPI = CustomFastAPI(
-        title="Project Moonshot",
-        version="0.4.1",
-        **app_kwargs
+        title="Project Moonshot", version="0.4.3", **app_kwargs
     )
 
     if cfg.cors.enabled():
@@ -104,6 +103,7 @@ def create_app(cfg: providers.Configuration) -> CustomFastAPI:
     app.include_router(runner.router)
     app.include_router(dataset.router)
     app.include_router(attack_modules.router)
+    app.include_router(bookmark.router)
 
     @app.exception_handler(RequestValidationError)
     async def validation_exception_handler(

moonshot/integrations/web_api/container.py

@@ -9,6 +9,7 @@ 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.bookmark_service import BookmarkService
 from .services.context_strategy_service import ContextStrategyService
 from .services.cookbook_service import CookbookService
 from .services.dataset_service import DatasetService
@@ -55,6 +56,11 @@ class Container(containers.DeclarativeContainer):
                 "log_file_max_size": 5242880,
                 "log_file_backup_count": 3,
             },
+            "temp_folder": str(
+                    importlib.resources.files("moonshot").joinpath(
+                        "integrations/web_api/temp"
+                    )
+                ),
         }
     )
 
@@ -97,9 +103,9 @@ class Container(containers.DeclarativeContainer):
     prompt_template_service: providers.Singleton[
         PromptTemplateService
     ] = providers.Singleton(PromptTemplateService)
-    context_strategy_service: providers.Singleton[ContextStrategyService] = providers.Singleton(
+    context_strategy_service: providers.Singleton[
         ContextStrategyService
-    )
+    ] = providers.Singleton(ContextStrategyService)
     benchmarking_service: providers.Singleton[
         BenchmarkingService
     ] = providers.Singleton(
@@ -127,6 +133,9 @@ class Container(containers.DeclarativeContainer):
     am_service: providers.Singleton[AttackModuleService] = providers.Singleton(
         AttackModuleService,
     )
+    bookmark_service: providers.Singleton[BookmarkService] = providers.Singleton(
+        BookmarkService,
+    )
     wiring_config = containers.WiringConfiguration(
         modules=[
             ".routes.redteam",
@@ -141,6 +150,7 @@ class Container(containers.DeclarativeContainer):
             ".routes.runner",
             ".routes.dataset",
             ".routes.attack_modules",
+            ".routes.bookmark",
             ".services.benchmarking_service",
         ]
     )

moonshot/integrations/web_api/routes/bookmark.py

@@ -0,0 +1,173 @@
+from typing import Optional
+
+from dependency_injector.wiring import Provide, inject
+from fastapi import APIRouter, Depends, HTTPException, Query
+
+from ..container import Container
+from ..schemas.bookmark_create_dto import BookmarkCreateDTO, BookmarkPydanticModel
+from ..services.bookmark_service import BookmarkService
+from ..services.utils.exceptions_handler import ServiceException
+
+router = APIRouter(tags=["Bookmark"])
+
+
+@router.post(
+    "/api/v1/bookmarks", response_description="Bookmark data added to the database"
+)
+@inject
+def insert_bookmark(
+    bookmark_data: BookmarkCreateDTO,
+    bookmark_service: BookmarkService = Depends(Provide[Container.bookmark_service]),
+) -> dict:
+    """
+    Insert a new bookmark into the database.
+
+    Args:
+        bookmark_data: The data of the bookmark to be added.
+        bookmark_service: The service responsible for bookmark operations.
+
+    Returns:
+        A dictionary with a message indicating successful insertion.
+
+    Raises:
+        HTTPException: An error occurred while inserting the bookmark.
+    """
+    try:
+        result = bookmark_service.insert_bookmark(bookmark_data)
+        return result
+    except ServiceException as e:
+        if e.error_code == "FileNotFound":
+            raise HTTPException(
+                status_code=404, detail=f"Failed to insert bookmark: {e.msg}"
+            )
+        elif e.error_code == "ValidationError":
+            raise HTTPException(
+                status_code=400, detail=f"Failed to insert bookmark: {e.msg}"
+            )
+        else:
+            raise HTTPException(
+                status_code=500, detail=f"Failed to insert bookmark: {e.msg}"
+            )
+
+
+@router.get(
+    "/api/v1/bookmarks",
+    response_description="List of all bookmarks or a specific bookmark by name",
+)
+@inject
+def get_all_bookmarks(
+    name: Optional[str] = Query(None, description="Name of the bookmark to query"),
+    bookmark_service: BookmarkService = Depends(Provide[Container.bookmark_service]),
+) -> list[BookmarkPydanticModel]:
+    """
+    Retrieve all bookmarks or a specific bookmark by name from the database.
+
+    Args:
+        name: The name of the bookmark to retrieve. If None, all bookmarks are retrieved.
+        bookmark_service: The service responsible for bookmark operations.
+
+    Returns:
+        A list of bookmarks or a single bookmark if a name is provided.
+
+    Raises:
+        HTTPException: An error occurred while retrieving bookmarks.
+    """
+    try:
+        return bookmark_service.get_all_bookmarks(
+            name=name,
+        )
+    except ServiceException as e:
+        if e.error_code == "FileNotFound":
+            raise HTTPException(
+                status_code=404, detail=f"Failed to retrieve bookmarks: {e.msg}"
+            )
+        elif e.error_code == "ValidationError":
+            raise HTTPException(
+                status_code=400, detail=f"Failed to retrieve bookmarks: {e.msg}"
+            )
+        else:
+            raise HTTPException(
+                status_code=500, detail=f"Failed to retrieve bookmarks: {e.msg}"
+            )
+
+
+@router.delete(
+    "/api/v1/bookmarks", response_description="Bookmark data deleted from the database"
+)
+@inject
+def delete_bookmark(
+    all: bool = Query(False, description="Flag to delete all bookmarks"),
+    name: Optional[str] = Query(None, description="Name of the bookmark to delete"),
+    bookmark_service: BookmarkService = Depends(Provide[Container.bookmark_service]),
+) -> dict:
+    """
+    Delete a specific bookmark by name or all bookmarks from the database.
+
+    Args:
+        all: A flag indicating whether to delete all bookmarks.
+        name: The name of the bookmark to delete. If 'all' is False and 'name' is None, no bookmark will be deleted.
+        bookmark_service: The service responsible for bookmark operations.
+
+    Returns:
+        A dictionary with a message indicating successful deletion or an error message.
+
+    Raises:
+        HTTPException: An error occurred while deleting the bookmark(s).
+    """
+    try:
+        if all:
+            return bookmark_service.delete_bookmarks(all=True)
+        elif name is not None:
+            return bookmark_service.delete_bookmarks(all=False, name=name)
+        else:
+            raise HTTPException(
+                status_code=400, detail="Must specify 'all' or 'name' parameter"
+            )
+    except ServiceException as e:
+        if e.error_code == "FileNotFound":
+            raise HTTPException(
+                status_code=404, detail=f"Failed to delete bookmark: {e.msg}"
+            )
+        elif e.error_code == "ValidationError":
+            raise HTTPException(
+                status_code=400, detail=f"Failed to delete bookmark: {e.msg}"
+            )
+        else:
+            raise HTTPException(
+                status_code=500, detail=f"Failed to delete bookmark: {e.msg}"
+            )
+@router.post(
+    "/api/v1/bookmarks/export", response_description="Exporting Bookmark to JSON file"
+)
+@inject
+def export_bookbookmarks(
+    export_file_name: Optional[str] = Query(
+        "bookmarks", description="Name of the exported file"
+    ),
+    bookmark_service: BookmarkService = Depends(Provide[Container.bookmark_service]),
+) -> str:
+    """
+    Export bookmarks to a JSON file with a given file name.
+
+    Args:
+        export_file_name: The name of the file to export the bookmarks to.z
+        bookmark_service: The service responsible for bookmark operations.
+
+    Returns:
+        A string with the path to the exported file or an error message.
+    """
+    try:
+        return bookmark_service.export_bookmarks(export_file_name)
+    except ServiceException as e:
+        if e.error_code == "FileNotFound":
+            raise HTTPException(
+                status_code=404, detail=f"Failed to export bookmark: {e.msg}"
+            )
+        elif e.error_code == "ValidationError":
+            raise HTTPException(
+                status_code=400, detail=f"Failed to export bookmark: {e.msg}"
+            )
+        else:
+            raise HTTPException(
+                status_code=500, detail=f"Failed to export bookmark: {e.msg}"
+            )

moonshot/integrations/web_api/routes/dataset.py

@@ -1,7 +1,8 @@
 from dependency_injector.wiring import Provide, inject
-from fastapi import APIRouter, Depends, HTTPException
+from fastapi import APIRouter, Depends, HTTPException, Query
 
 from ..container import Container
+from ..schemas.dataset_create_dto import DatasetCreateDTO
 from ..schemas.dataset_response_dto import DatasetResponseDTO
 from ..services.dataset_service import DatasetService
 from ..services.utils.exceptions_handler import ServiceException
@@ -9,6 +10,50 @@ from ..services.utils.exceptions_handler import ServiceException
 router = APIRouter(tags=["Datasets"])
 
 
+@router.post("/api/v1/datasets")
+@inject
+def create_dataset(
+    dataset_data: DatasetCreateDTO,
+    method: str = Query(
+        ...,
+        description="The method to use for creating the dataset. Supported methods are 'hf' and 'csv'.",
+    ),
+    dataset_service: DatasetService = Depends(Provide[Container.dataset_service]),
+) -> str:
+    """
+    Create a new dataset using the specified method.
+
+    Args:
+        dataset_data (DatasetCreateDTO): The data required to create the dataset.
+        method (str): The method to use for creating the dataset. Supported methods are "hf" and "csv".
+        dataset_service (DatasetService, optional): The service responsible for creating the dataset.
+        Defaults to Depends(Provide[Container.dataset_service]).
+
+    Returns:
+        dict: A message indicating the dataset was created successfully.
+
+    Raises:
+        HTTPException: An error with status code 404 if the dataset file 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:
+        return dataset_service.create_dataset(dataset_data, method)
+    except ServiceException as e:
+        if e.error_code == "FileNotFound":
+            raise HTTPException(
+                status_code=404, detail=f"Failed to retrieve datasets: {e.msg}"
+            )
+        elif e.error_code == "ValidationError":
+            raise HTTPException(
+                status_code=400, detail=f"Failed to retrieve datasets: {e.msg}"
+            )
+        else:
+            raise HTTPException(
+                status_code=500, detail=f"Failed to retrieve datasets: {e.msg}"
+            )
+
+
 @router.get("/api/v1/datasets")
 @inject
 def get_all_datasets(

moonshot/integrations/web_api/schemas/bookmark_create_dto.py

@@ -0,0 +1,13 @@
+from typing import Optional
+
+from moonshot.src.bookmark.bookmark_arguments import (
+    BookmarkArguments as BookmarkPydanticModel,
+)
+
+
+class BookmarkCreateDTO(BookmarkPydanticModel):
+    prompt_template: Optional[str] = ""
+    context_strategy: Optional[str] = ""
+    attack_module: Optional[str] = ""
+    metric: Optional[str] = ""
+    bookmark_time: Optional[str] = None

moonshot/integrations/web_api/schemas/dataset_create_dto.py

@@ -0,0 +1,18 @@
+from typing import Optional
+
+from pydantic import Field
+from pyparsing import Iterator
+
+from moonshot.src.datasets.dataset_arguments import (
+    DatasetArguments as DatasetPydanticModel,
+)
+
+
+class DatasetCreateDTO(DatasetPydanticModel):
+    id: Optional[str] = None
+    examples: Iterator[dict] = None
+    name: str = Field(..., min_length=1)
+    description: str = Field(default="", min_length=1)
+    license: Optional[str] = ""
+    reference: Optional[str] = ""
+    params: dict

moonshot/integrations/web_api/schemas/recipe_create_dto.py

@@ -13,7 +13,6 @@ class RecipeCreateDTO(RecipePydanticModel):
     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
 
@@ -27,6 +26,5 @@ class RecipeUpdateDTO(RecipePydanticModel):
     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/services/bookmark_service.py

@@ -0,0 +1,94 @@
+from .... import api as moonshot_api
+from ..schemas.bookmark_create_dto import BookmarkCreateDTO, BookmarkPydanticModel
+from ..services.base_service import BaseService
+from ..services.utils.exceptions_handler import exception_handler
+from .utils.file_manager import copy_file
+
+
+class BookmarkService(BaseService):
+    @exception_handler
+    def insert_bookmark(self, bookmark_data: BookmarkCreateDTO) -> dict:
+        """
+        Inserts a new bookmark into the system.
+
+        Args:
+            bookmark_data (BookmarkCreateDTO): The data transfer object containing bookmark details.
+        """
+        result = moonshot_api.api_insert_bookmark(
+            name=bookmark_data.name,
+            prompt=bookmark_data.prompt,
+            prepared_prompt=bookmark_data.prepared_prompt,
+            response=bookmark_data.response,
+            context_strategy=bookmark_data.context_strategy,
+            prompt_template=bookmark_data.prompt_template,
+            attack_module=bookmark_data.attack_module,
+            metric=bookmark_data.metric,
+        )
+
+        return result
+
+    @exception_handler
+    def get_all_bookmarks(self, name: str | None = None) -> list[BookmarkPydanticModel]:
+        """
+        Retrieves all bookmarks or a specific bookmark by its name.
+
+        Args:
+            name (str | None, optional): The name of the bookmark to retrieve. If None, all bookmarks are retrieved.
+
+        Returns:
+            list[BookmarkPydanticModel]: A list of bookmark models.
+        """
+        retn_bookmark: list[BookmarkPydanticModel] = []
+
+        if name:
+            bookmarks = [moonshot_api.api_get_bookmark(name)]
+        else:
+            bookmarks = moonshot_api.api_get_all_bookmarks()
+
+        for bookmark in bookmarks:
+            retn_bookmark.append(BookmarkPydanticModel(**bookmark))
+        return retn_bookmark
+
+    @exception_handler
+    def delete_bookmarks(self, all: bool = False, name: str | None = None) -> dict:
+        """
+        Deletes a single bookmark by its name or all bookmarks if the 'all' flag is set to True and returns
+        a boolean indicating the success of the operation.
+
+        Args:
+            all (bool, optional): If True, all bookmarks will be deleted. Defaults to False.
+            name (str | None, optional): The name of the bookmark to delete. If 'all' is False, 'name' must be provided.
+
+        Returns:
+            dict: True if the deletion was successful, False otherwise.
+        """
+        if all:
+            result = moonshot_api.api_delete_all_bookmark()
+        elif name is not None:
+            result = moonshot_api.api_delete_bookmark(name)
+        else:
+            raise ValueError("Either 'all' must be True or 'name' must be provided.")
+
+        if not result["success"]:
+            raise Exception(
+                result["message"], "delete_bookmarks", "DeleteBookmarkError"
+            )
+
+        return result
+
+    @exception_handler
+    def export_bookmarks(self, export_file_name: str = "bookmarks") -> str:
+        """
+        Exports bookmarks to a file or returns them as a list of dictionaries.
+
+        Args:
+            write_file (bool, optional): If True, bookmarks will be written to a file. Defaults to False.
+            export_file_name (str, optional): The name of the file to write bookmarks to. Defaults to "bookmarks".
+
+        Returns:
+            list[dict]: A list of bookmarks as dictionaries.
+        """
+
+        new_file_path = moonshot_api.api_export_bookmarks(export_file_name)
+
+        return copy_file(new_file_path)

moonshot/integrations/web_api/services/dataset_service.py

@@ -1,10 +1,35 @@
 from .... import api as moonshot_api
+from ..schemas.dataset_create_dto import DatasetCreateDTO
 from ..schemas.dataset_response_dto import DatasetResponseDTO
 from ..services.base_service import BaseService
 from ..services.utils.exceptions_handler import exception_handler
+from .utils.file_manager import copy_file
 
 
 class DatasetService(BaseService):
+    @exception_handler
+    def create_dataset(self, dataset_data: DatasetCreateDTO, method: str) -> str:
+        """
+        Create a dataset using the specified method.
+
+        Args:
+            dataset_data (DatasetCreateDTO): The data required to create the dataset.
+            method (str): The method to use for creating the dataset.
+                          Supported methods are "hf" and "csv".
+
+        Raises:
+            Exception: If an error occurs during dataset creation.
+        """
+        new_ds_path = moonshot_api.api_create_datasets(
+            name=dataset_data.name,
+            description=dataset_data.description,
+            reference=dataset_data.reference,
+            license=dataset_data.license,
+            method=method,
+            **dataset_data.params,
+        )
+        return copy_file(new_ds_path)
+
     @exception_handler
     def get_all_datasets(self) -> list[DatasetResponseDTO]:
         datasets = moonshot_api.api_get_all_datasets()

moonshot/integrations/web_api/services/recipe_service.py

@@ -24,7 +24,6 @@ class RecipeService(BaseService):
             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,
         )
 

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

@@ -0,0 +1,52 @@
+import os
+import shutil
+
+# Global variable to store the path to the temporary folder
+TEMP_FOLDER_PATH = None
+
+
+def create_temp_dir(temp_file_path: str):
+    """
+    Create a temporary directory at the specified path.
+
+    This function sets a global variable with the path to the temporary directory
+    and creates the directory if it does not already exist.
+
+    Args:
+        temp_file_path (str): The file path where the temporary directory will be created.
+    """
+    global TEMP_FOLDER_PATH
+    TEMP_FOLDER_PATH = temp_file_path
+    if not os.path.exists(TEMP_FOLDER_PATH):
+        os.makedirs(TEMP_FOLDER_PATH)
+
+
+def copy_file(file_path: str) -> str:
+    """
+    Copy a file to the temporary directory.
+
+    This function copies a file from the given file path to the temporary directory
+    set by the create_temp_dir function. It raises an error if the temporary directory
+    path is not set before calling this function.
+
+    Args:
+        file_path (str): The path of the file to be copied.
+
+    Returns:
+        str: The path to the copied file in the temporary directory.
+
+    Raises:
+        ValueError: If the temporary folder path is not set.
+    """
+    if TEMP_FOLDER_PATH is None:
+        raise ValueError(
+            "Temporary folder path is not set. Call create_temp_dir first."
+        )
+
+    # Define the destination path for the copied file
+    dest_path = os.path.join(TEMP_FOLDER_PATH, os.path.basename(file_path))
+
+    # Copy the file to the destination
+    shutil.copy(file_path, dest_path)
+
+    return dest_path

moonshot/integrations/web_api/status_updater/moonshot_ui_webhook.py

@@ -60,7 +60,6 @@ class MoonshotUIWebhook(
         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()