FEAT Add dataset info to backend and expose --list-datasets in pyrit_scan (microsoft#2106)

Co-authored-by: Behnam Ousat <behnamousat@microsoft.com>

Author: behnam-o

pyrit/backend/main.py

@@ -25,6 +25,7 @@
     attacks,
     auth,
     converters,
+    datasets,
     health,
     initializers,
     labels,
@@ -120,6 +121,7 @@ async def lifespan(app: FastAPI) -> AsyncGenerator[None, None]:
 app.include_router(attacks.router, prefix="/api", tags=["attacks"])
 app.include_router(targets.router, prefix="/api", tags=["targets"])
 app.include_router(converters.router, prefix="/api", tags=["converters"])
+app.include_router(datasets.router, prefix="/api", tags=["datasets"])
 app.include_router(scenarios.router, prefix="/api", tags=["scenarios"])
 app.include_router(initializers.router, prefix="/api", tags=["initializers"])
 app.include_router(labels.router, prefix="/api", tags=["labels"])

pyrit/backend/models/__init__.py

@@ -48,6 +48,10 @@
     CreateConverterResponse,
     PreviewStep,
 )
+from pyrit.backend.models.datasets import (
+    DatasetInfo,
+    DatasetListResponse,
+)
 from pyrit.backend.models.initializers import (
     ListRegisteredInitializersResponse,
     RegisterInitializerRequest,
@@ -101,6 +105,9 @@
     "CreateConverterRequest",
     "CreateConverterResponse",
     "PreviewStep",
+    # Datasets
+    "DatasetInfo",
+    "DatasetListResponse",
     # Scenarios
     "ListRegisteredScenariosResponse",
     "ScenarioRunListResponse",

pyrit/backend/models/datasets.py

@@ -0,0 +1,24 @@
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT license.
+
+"""
+Dataset models for the PyRIT API.
+
+Datasets are seed prompt/objective collections provided by
+``SeedDatasetProvider`` subclasses. These models describe the wire format for
+listing available datasets.
+"""
+
+from pydantic import BaseModel, Field
+
+
+class DatasetInfo(BaseModel):
+    """Metadata about a single available dataset."""
+
+    name: str = Field(..., description="Dataset name (e.g., 'harmbench')")
+
+
+class DatasetListResponse(BaseModel):
+    """Response for listing available datasets."""
+
+    items: list[DatasetInfo] = Field(..., description="List of available datasets")

pyrit/backend/routes/__init__.py

@@ -5,11 +5,23 @@
 API route handlers.
 """
 
-from pyrit.backend.routes import attacks, converters, health, initializers, labels, media, scenarios, targets, version
+from pyrit.backend.routes import (
+    attacks,
+    converters,
+    datasets,
+    health,
+    initializers,
+    labels,
+    media,
+    scenarios,
+    targets,
+    version,
+)
 
 __all__ = [
     "attacks",
     "converters",
+    "datasets",
     "health",
     "initializers",
     "labels",

pyrit/backend/routes/datasets.py

@@ -0,0 +1,37 @@
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT license.
+
+"""
+Dataset API routes.
+
+Provides an endpoint for listing available seed datasets. Datasets are
+discovered from registered ``SeedDatasetProvider`` subclasses.
+"""
+
+from fastapi import APIRouter
+
+from pyrit.backend.models.common import ProblemDetail
+from pyrit.backend.models.datasets import (
+    DatasetListResponse,
+)
+from pyrit.backend.services.dataset_service import get_dataset_service
+
+router = APIRouter(prefix="/datasets", tags=["datasets"])
+
+
+@router.get(
+    "",
+    response_model=DatasetListResponse,
+    responses={
+        500: {"model": ProblemDetail, "description": "Internal server error"},
+    },
+)
+async def list_datasets() -> DatasetListResponse:  # pyrit-async-suffix-exempt
+    """
+    List all available datasets.
+
+    Returns:
+        DatasetListResponse: Available datasets.
+    """
+    service = get_dataset_service()
+    return await service.list_datasets_async()

pyrit/backend/services/__init__.py

@@ -15,6 +15,10 @@
     ConverterService,
     get_converter_service,
 )
+from pyrit.backend.services.dataset_service import (
+    DatasetService,
+    get_dataset_service,
+)
 from pyrit.backend.services.initializer_service import (
     InitializerService,
     get_initializer_service,
@@ -37,6 +41,8 @@
     "get_attack_service",
     "ConverterService",
     "get_converter_service",
+    "DatasetService",
+    "get_dataset_service",
     "InitializerService",
     "get_initializer_service",
     "ScenarioService",

pyrit/backend/services/dataset_service.py

@@ -0,0 +1,55 @@
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT license.
+
+"""
+Dataset service for listing seed datasets.
+
+Wraps ``SeedDatasetProvider`` discovery and memory to list available datasets.
+"""
+
+import logging
+from functools import lru_cache
+
+from pyrit.backend.models.datasets import (
+    DatasetInfo,
+    DatasetListResponse,
+)
+from pyrit.datasets import SeedDatasetProvider
+from pyrit.memory import CentralMemory
+
+logger = logging.getLogger(__name__)
+
+
+class DatasetService:
+    """Service for listing seed datasets."""
+
+    def __init__(self) -> None:
+        """Initialize the dataset service."""
+        self._memory = CentralMemory.get_memory_instance()
+
+    async def list_datasets_async(self) -> DatasetListResponse:
+        """
+        List all available datasets.
+
+        Combines datasets discoverable via registered providers with those
+        already loaded into memory, since both are available for use.
+
+        Returns:
+            DatasetListResponse: Available datasets.
+        """
+        provider_names = await SeedDatasetProvider.get_all_dataset_names_async()
+        memory_names = self._memory.get_seed_dataset_names()
+        available = sorted(set(provider_names) | set(memory_names))
+        items = [DatasetInfo(name=name) for name in available]
+        return DatasetListResponse(items=items)
+
+
+@lru_cache(maxsize=1)
+def get_dataset_service() -> DatasetService:
+    """
+    Get the global dataset service instance.
+
+    Returns:
+        The singleton DatasetService instance.
+    """
+    return DatasetService()

pyrit/cli/_output.py

@@ -13,7 +13,7 @@
 from __future__ import annotations
 
 import sys
-from typing import TYPE_CHECKING
+from typing import TYPE_CHECKING, Any
 
 if TYPE_CHECKING:
     from pyrit.models import ScenarioResult
@@ -193,6 +193,31 @@ def print_target_list(*, items: list[TargetInstance]) -> None:
     print(f"\nTotal targets: {len(items)}")
 
 
+# ---------------------------------------------------------------------------
+# Dataset listing
+# ---------------------------------------------------------------------------
+
+
+def print_dataset_list(*, items: list[dict[str, Any]]) -> None:
+    """
+    Print a formatted list of available datasets.
+
+    Args:
+        items: List of dataset dicts from ``GET /api/datasets``.
+    """
+    if not items:
+        print("No datasets found.")
+        return
+
+    print("\nAvailable Datasets:")
+    print("=" * 80)
+    for ds in items:
+        name = ds.get("name", "unknown")
+        print(f"    {name}")
+    print("=" * 80)
+    print(f"\nTotal datasets: {len(items)}")
+
+
 # ---------------------------------------------------------------------------
 # Scenario run progress & summary
 # ---------------------------------------------------------------------------

pyrit/cli/api_client.py

@@ -196,6 +196,19 @@ async def list_targets_async(self, *, limit: int = 200) -> list[TargetInstance]:
         payload = await self._get_json_async(path="/api/targets", params={"limit": limit})
         return [TargetInstance.model_validate(item) for item in payload.get("items", [])]
 
+    # ------------------------------------------------------------------
+    # Datasets
+    # ------------------------------------------------------------------
+
+    async def list_datasets_async(self) -> dict[str, Any]:
+        """
+        List all available datasets.
+
+        Returns:
+            dict: ``DatasetListResponse`` payload.
+        """
+        return await self._get_json_async(path="/api/datasets")
+
     # ------------------------------------------------------------------
     # Scenario runs
     # ------------------------------------------------------------------

pyrit/cli/pyrit_scan.py

@@ -91,6 +91,9 @@ def _print_cli_exception(*, exc: BaseException) -> None:
   pyrit_scan --list-initializers
   pyrit_scan --list-targets
 
+  # List available datasets
+  pyrit_scan --list-datasets
+
   # Run single-turn cyber attacks against a target
   pyrit_scan airt.cyber --target openai_chat --strategies single_turn
 
@@ -188,6 +191,11 @@ def _build_base_parser(*, add_help: bool = True) -> ArgumentParser:
         action="store_true",
         help="List all available targets and exit",
     )
+    discovery_group.add_argument(
+        "--list-datasets",
+        action="store_true",
+        help="List all available datasets and exit",
+    )
     discovery_group.add_argument(
         "--add-initializer",
         type=str,
@@ -451,6 +459,7 @@ def _is_command_specified(*, parsed_args: Namespace) -> bool:
         parsed_args.list_scenarios
         or parsed_args.list_initializers
         or parsed_args.list_targets
+        or parsed_args.list_datasets
         or parsed_args.add_initializer
         or parsed_args.scenario_name
     )
@@ -514,6 +523,10 @@ async def _handle_list_commands_async(*, client: Any, parsed_args: Namespace) ->
         targets = await client.list_targets_async()
         _output.print_target_list(items=targets)
         return 0
+    if parsed_args.list_datasets:
+        resp = await client.list_datasets_async()
+        _output.print_dataset_list(items=resp.get("items", []))
+        return 0
     return None