aiverify-moonshot 0.4.2__py3-none-any.whl → 0.4.4__py3-none-any.whl

moonshot/integrations/cli/redteam/context_strategy.py

@@ -1,4 +1,5 @@
 import argparse
+from ast import literal_eval
 
 import cmd2
 from rich.console import Console
@@ -10,8 +11,8 @@ from moonshot.api import (
     api_update_context_strategy,
 )
 from moonshot.integrations.cli.active_session_cfg import active_session
+from moonshot.integrations.cli.utils.process_data import filter_data
 from moonshot.src.redteaming.session.session import Session
-from moonshot.src.utils.find_feature import find_keyword
 
 console = Console()
 
@@ -60,6 +61,7 @@ def list_context_strategies(args) -> list | None:
     Args:
         args: A namespace object from argparse. It should have an optional attribute:
         find (str): Optional field to find context strategies with a keyword.
+        pagination (str): Optional field to paginate context strategies.
 
     Returns:
         list | None: A list of ContextStrategy or None if there is no result.
@@ -67,19 +69,19 @@ def list_context_strategies(args) -> list | None:
     try:
         context_strategy_metadata_list = api_get_all_context_strategy_metadata()
         keyword = args.find.lower() if args.find else ""
-        if keyword:
-            filtered_context_strategies_list = find_keyword(
-                keyword, context_strategy_metadata_list
+        pagination = literal_eval(args.pagination) if args.pagination else ()
+
+        if context_strategy_metadata_list:
+            filtered_context_strategies_list = filter_data(
+                context_strategy_metadata_list, keyword, pagination
             )
             if filtered_context_strategies_list:
-                display_context_strategies(filtered_context_strategies_list)
+                _display_context_strategies(filtered_context_strategies_list)
                 return filtered_context_strategies_list
-            else:
-                print("No context strategies containing keyword found.")
-                return None
-        else:
-            display_context_strategies(context_strategy_metadata_list)
-            return context_strategy_metadata_list
+
+        console.print("[red]There are no context strategies found.[/red]")
+        return None
+
     except Exception as e:
         print(f"[list_context_strategies]: {str(e)}")
 
@@ -124,7 +126,7 @@ def delete_context_strategy(args) -> None:
         print(f"[delete_context_strategy]: {str(e)}")
 
 
-def display_context_strategies(context_strategies: list) -> None:
+def _display_context_strategies(context_strategies: list) -> None:
     """
     Display a list of context strategies.
 
@@ -137,25 +139,21 @@ def display_context_strategies(context_strategies: list) -> None:
     Returns:
         None
     """
-    if context_strategies:
-        table = Table(
-            title="Context Strategy List",
-            show_lines=True,
-            expand=True,
-            header_style="bold",
-        )
-        table.add_column("No.", justify="left", width=2)
-        table.add_column("Context Strategy Information", justify="left", width=98)
-        for context_strategy_index, context_strategy_data in enumerate(
-            context_strategies, 1
-        ):
-            context_strategy_data_str = ""
-            for k, v in context_strategy_data.items():
+    table = Table(
+        title="Context Strategy List",
+        show_lines=True,
+        expand=True,
+        header_style="bold",
+    )
+    table.add_column("No.", justify="left", width=2)
+    table.add_column("Context Strategy Information", justify="left", width=98)
+    for idx, context_strategy_data in enumerate(context_strategies, 1):
+        context_strategy_data_str = ""
+        for k, v in context_strategy_data.items():
+            if k != "idx":
                 context_strategy_data_str += f"[blue]{k.capitalize()}:[/blue] {v}\n\n"
-            table.add_row(str(context_strategy_index), context_strategy_data_str)
-        console.print(table)
-    else:
-        console.print("[red]There are no context strategies found.[/red]", style="bold")
+        table.add_row(str(idx), context_strategy_data_str)
+    console.print(table)
 
 
 # Use context strategy arguments
@@ -199,3 +197,11 @@ list_context_strategies_args.add_argument(
     help="Optional field to find context strategies with keyword",
     nargs="?",
 )
+
+list_context_strategies_args.add_argument(
+    "-p",
+    "--pagination",
+    type=str,
+    help="Optional tuple to paginate context strategies(s). E.g. (2,10) returns 2nd page with 10 items in each page.",
+    nargs="?",
+)

moonshot/integrations/cli/redteam/session.py

@@ -22,8 +22,8 @@ from moonshot.api import (
     api_load_session,
 )
 from moonshot.integrations.cli.active_session_cfg import active_session
+from moonshot.integrations.cli.utils.process_data import filter_data
 from moonshot.src.redteaming.session.session import Session
-from moonshot.src.utils.find_feature import find_keyword
 
 console = Console()
 
@@ -139,6 +139,7 @@ def list_sessions(args) -> list | None:
     Args:
         args: A namespace object from argparse. It should have an optional attribute:
         find (str): Optional field to find session(s) with a keyword.
+        pagination (str): Optional field to paginate sessions.
 
     Returns:
         list | None: A list of Session or None if there is no result.
@@ -146,19 +147,19 @@ def list_sessions(args) -> list | None:
     try:
         session_metadata_list = api_get_all_session_metadata()
         keyword = args.find.lower() if args.find else ""
-        if keyword:
-            filtered_session_metadata_list = find_keyword(
-                keyword, session_metadata_list
+        pagination = literal_eval(args.pagination) if args.pagination else ()
+
+        if session_metadata_list:
+            filtered_session_metadata_list = filter_data(
+                session_metadata_list, keyword, pagination
             )
             if filtered_session_metadata_list:
-                display_sessions(filtered_session_metadata_list)
+                _display_sessions(filtered_session_metadata_list)
                 return filtered_session_metadata_list
-            else:
-                print("No sessions containing keyword found.")
-                return None
-        else:
-            display_sessions(session_metadata_list)
-            return session_metadata_list
+
+        console.print("[red]There are no sessions found.[/red]")
+        return None
+
     except Exception as e:
         print(f"[list_sessions]: {str(e)}")
 
@@ -363,12 +364,13 @@ def list_bookmarks(args) -> list | None:
 
     This function retrieves all available bookmarks by calling the api_get_all_bookmarks function from the
     moonshot.api module.
-    It then displays the retrieved bookmarks using the display_bookmarks function.
+    It then displays the retrieved bookmarks using the _display_bookmarks function.
     If no bookmarks are found, a message is printed to the console.
 
     Args:
         args: A namespace object from argparse. It should have an optional attribute:
         find (str): Optional field to find bookmark(s) with a keyword.
+        pagination (str): Optional field to paginate bookmarks.
 
     Returns:
         list | None: A list of Bookmark or None if there is no result.
@@ -376,22 +378,22 @@ def list_bookmarks(args) -> list | None:
     try:
         bookmarks_list = api_get_all_bookmarks()
         keyword = args.find.lower() if args.find else ""
-        if keyword:
-            filtered_bookmarks_list = find_keyword(keyword, bookmarks_list)
+        pagination = literal_eval(args.pagination) if args.pagination else ()
+
+        if bookmarks_list:
+            filtered_bookmarks_list = filter_data(bookmarks_list, keyword, pagination)
             if filtered_bookmarks_list:
-                display_bookmarks(filtered_bookmarks_list)
+                _display_bookmarks(filtered_bookmarks_list)
                 return filtered_bookmarks_list
-            else:
-                print("No bookmarks containing keyword found.")
-                return None
-        else:
-            display_bookmarks(bookmarks_list)
-            return bookmarks_list
+
+        console.print("[red]There are no bookmarks found.[/red]")
+        return None
+
     except Exception as e:
         print(f"[list_bookmarks]: {str(e)}")
 
 
-def display_bookmarks(bookmarks_list) -> None:
+def _display_bookmarks(bookmarks_list) -> None:
     """
     Display the list of bookmarks in a tabular format.
 
@@ -402,38 +404,39 @@ def display_bookmarks(bookmarks_list) -> None:
     Args:
         bookmarks_list (list): A list of dictionaries, where each dictionary contains the details of a bookmark.
     """
-    if bookmarks_list:
-        table = Table(
-            title="Bookmark List", show_lines=True, expand=True, header_style="bold"
+
+    table = Table(
+        title="Bookmark List", show_lines=True, expand=True, header_style="bold"
+    )
+    table.add_column("ID.", justify="left", width=5)
+    table.add_column("Name", justify="left", width=20)
+    table.add_column("Prepared Prompt", justify="left", width=50)
+    table.add_column("Predicted Response", justify="left", width=50)
+    table.add_column("Bookmark Time", justify="left", width=20)
+    for idx, bookmark in enumerate(bookmarks_list, 1):
+        (
+            name,
+            prompt,
+            prepared_prompt,
+            response,
+            context_strategy,
+            prompt_template,
+            attack_module,
+            metric,
+            bookmark_time,
+            *other_args,
+        ) = bookmark.values()
+        idx = bookmark.get("idx", idx)
+
+        table.add_section()
+        table.add_row(
+            str(idx),
+            name,
+            prepared_prompt,
+            response,
+            bookmark_time,
         )
-        table.add_column("ID.", justify="left", width=5)
-        table.add_column("Name", justify="left", width=20)
-        table.add_column("Prepared Prompt", justify="left", width=50)
-        table.add_column("Predicted Response", justify="left", width=50)
-        table.add_column("Bookmark Time", justify="left", width=20)
-        for idx, bookmark in enumerate(bookmarks_list, 1):
-            (
-                name,
-                prompt,
-                prepared_prompt,
-                response,
-                context_strategy,
-                prompt_template,
-                attack_module,
-                metric,
-                bookmark_time,
-            ) = bookmark.values()
-            table.add_section()
-            table.add_row(
-                str(idx),
-                name,
-                prepared_prompt,
-                response,
-                bookmark_time,
-            )
-        console.print(table)
-    else:
-        console.print("[red]There are no bookmarks found.[/red]")
+    console.print(table)
 
 
 def view_bookmark(args) -> None:
@@ -446,12 +449,12 @@ def view_bookmark(args) -> None:
 
     try:
         bookmark_info = api_get_bookmark(args.bookmark_name)
-        display_bookmark(bookmark_info)
+        _display_bookmark(bookmark_info)
     except Exception as e:
         print(f"[view_bookmark]: {str(e)}")
 
 
-def display_bookmark(bookmark_info: dict) -> None:
+def _display_bookmark(bookmark_info: dict) -> None:
     """
     Display the filtered bookmark in a tabular format.
 
@@ -600,7 +603,7 @@ def run_attack_module(args):
         if args.prompt_template:
             prompt_template = [args.prompt_template]
         elif active_session["prompt_template"]:
-            prompt_template = [args.prompt_template]
+            prompt_template = [active_session["prompt_template"]]
         else:
             prompt_template = []
 
@@ -703,7 +706,7 @@ def delete_session(args) -> None:
         print(f"[delete_session]: {str(e)}")
 
 
-def display_sessions(sessions: list) -> None:
+def _display_sessions(sessions: list) -> None:
     """
     Display a list of sessions.
 
@@ -717,25 +720,32 @@ def display_sessions(sessions: list) -> None:
         None
     """
 
-    if sessions:
-        table = Table(
-            title="Session List", show_lines=True, expand=True, header_style="bold"
-        )
-        table.add_column("No.", justify="left", width=2)
-        table.add_column("Session ID", justify="left", width=20)
-        table.add_column("Contains", justify="left", width=78)
-
-        for session_index, session_data in enumerate(sessions, 1):
-            session_id = session_data.get("session_id", "")
-            endpoints = ", ".join(session_data.get("endpoints", []))
-            created_datetime = session_data.get("created_datetime", "")
-
-            session_info = f"[red]id: {session_id}[/red]\n\nCreated: {created_datetime}"
-            contains_info = f"[blue]Endpoints:[/blue] {endpoints}\n\n"
-            table.add_row(str(session_index), session_info, contains_info)
-        console.print(table)
-    else:
-        console.print("[red]There are no sessions found.[/red]", style="bold")
+    table = Table(
+        title="Session List", show_lines=True, expand=True, header_style="bold"
+    )
+    table.add_column("No.", justify="left", width=2)
+    table.add_column("Session ID", justify="left", width=20)
+    table.add_column("Contains", justify="left", width=78)
+
+    for idx, session_data in enumerate(sessions, 1):
+        (
+            session_id,
+            endpoints,
+            created_epoch,
+            created_datetime,
+            prompt_template,
+            context_strategy,
+            cs_num_of_prev_prompts,
+            attack_module,
+            metric,
+            system_prompt,
+            *other_args,
+        ) = session_data.values()
+        idx = session_data.get("idx", idx)
+        session_info = f"[red]id: {session_id}[/red]\n\nCreated: {created_datetime}"
+        contains_info = f"[blue]Endpoints:[/blue] {endpoints}\n\n"
+        table.add_row(str(idx), session_info, contains_info)
+    console.print(table)
 
 
 # use session arguments
@@ -885,6 +895,13 @@ list_sessions_args.add_argument(
     nargs="?",
 )
 
+list_sessions_args.add_argument(
+    "-p",
+    "--pagination",
+    type=str,
+    help="Optional tuple to paginate session(s). E.g. (2,10) returns 2nd page with 10 items in each page.",
+    nargs="?",
+)
 
 # Add bookmark arguments
 add_bookmark_args = cmd2.Cmd2ArgumentParser(
@@ -973,3 +990,11 @@ list_bookmarks_args.add_argument(
     help="Optional field to find bookmark(s) with keyword",
     nargs="?",
 )
+
+list_bookmarks_args.add_argument(
+    "-p",
+    "--pagination",
+    type=str,
+    help="Optional tuple to paginate bookmark(s). E.g. (2,10) returns 2nd page with 10 items in each page.",
+    nargs="?",
+)

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/app.py

@@ -71,7 +71,7 @@ def create_app(cfg: providers.Configuration) -> CustomFastAPI:
     }
 
     app: CustomFastAPI = CustomFastAPI(
-        title="Project Moonshot", version="0.4.2", **app_kwargs
+        title="Project Moonshot", version="0.4.4", **app_kwargs
     )
 
     if cfg.cors.enabled():

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/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

@@ -52,18 +52,22 @@ class BookmarkService(BaseService):
     @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.
+        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 'id' must be provided.")
+            raise ValueError("Either 'all' must be True or 'name' must be provided.")
 
         if not result["success"]:
             raise Exception(

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/src/api/api_dataset.py

@@ -1,6 +1,7 @@
 from pydantic import validate_call
 
 from moonshot.src.datasets.dataset import Dataset
+from moonshot.src.datasets.dataset_arguments import DatasetArguments
 
 
 # ------------------------------------------------------------------------------
@@ -44,3 +45,37 @@ def api_get_all_datasets_name() -> list[str]:
     """
     datasets_name, _ = Dataset.get_available_items()
     return datasets_name
+
+
+def api_create_datasets(
+    name: str, description: str, reference: str, license: str, method: str, **kwargs
+) -> str:
+    """
+    This function creates a new dataset.
+
+    This function takes the name, description, reference, and license for a new dataset as input. It then creates a new
+    DatasetArguments object with these details and an empty id. The id is left empty because it will be generated
+    from the name during the creation process. The function then calls the Dataset's create method to
+    create the new dataset.
+
+    Args:
+        name (str): The name of the new dataset.
+        description (str): A brief description of the new dataset.
+        reference (str): A reference link for the new dataset.
+        license (str): The license of the new dataset.
+        method (str): The method to create new dataset. (csv/hf)
+        kwargs: Additional keyword arguments for the Dataset's create method.
+
+    Returns:
+        str: The ID of the newly created dataset.
+    """
+    ds_args = DatasetArguments(
+        id="",
+        name=name,
+        description=description,
+        reference=reference,
+        license=license,
+        examples=None,
+    )
+
+    return Dataset.create(ds_args, method, **kwargs)