REFACTOR Breaking: Migrate CLI to use pyrit.models (#1997)

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
Co-authored-by: hannahwestra25 <hannahwestra@microsoft.com>

Author: 3 people

pyrit/backend/mappers/target_mappers.py

@@ -12,8 +12,8 @@
 ``TargetIdentifier`` instead of poking ``identifier.params`` by string key.
 """
 
-from pyrit.backend.models import TargetCapabilitiesInfo, TargetInstance
 from pyrit.models import TargetIdentifier
+from pyrit.models.catalog.target import TargetCapabilitiesInfo, TargetInstance
 from pyrit.prompt_target import PromptTarget
 from pyrit.prompt_target.common.target_capabilities import CapabilityName, TargetCapabilities
 from pyrit.prompt_target.round_robin_target import RoundRobinTarget
@@ -46,6 +46,7 @@ def _target_capabilities_to_info(
         supports_json_output=capabilities.supports_json_output,
         supports_editable_history=capabilities.supports_editable_history,
         supports_system_prompt=capabilities.supports_system_prompt,
+        supports_streaming_audio=capabilities.supports_streaming_audio,
         supported_input_modalities=sorted({str(t) for combo in capabilities.input_modalities for t in combo}),
         supported_output_modalities=sorted({str(t) for combo in capabilities.output_modalities for t in combo}),
     )

pyrit/backend/models/__init__.py

@@ -49,20 +49,15 @@
     PreviewStep,
 )
 from pyrit.backend.models.initializers import (
-    InitializerParameterSummary,
     ListRegisteredInitializersResponse,
-    RegisteredInitializer,
     RegisterInitializerRequest,
 )
 from pyrit.backend.models.scenarios import (
     ListRegisteredScenariosResponse,
-    RegisteredScenario,
-    ScenarioParameterSummary,
+    ScenarioRunListResponse,
 )
 from pyrit.backend.models.targets import (
     CreateTargetRequest,
-    TargetCapabilitiesInfo,
-    TargetInstance,
     TargetListResponse,
 )
 
@@ -108,16 +103,11 @@
     "PreviewStep",
     # Scenarios
     "ListRegisteredScenariosResponse",
-    "RegisteredScenario",
-    "ScenarioParameterSummary",
+    "ScenarioRunListResponse",
     # Initializers
-    "InitializerParameterSummary",
     "ListRegisteredInitializersResponse",
-    "RegisteredInitializer",
     "RegisterInitializerRequest",
     # Targets
     "CreateTargetRequest",
-    "TargetCapabilitiesInfo",
-    "TargetInstance",
     "TargetListResponse",
 ]

pyrit/backend/models/initializers.py

@@ -2,38 +2,23 @@
 # Licensed under the MIT license.
 
 """
-Initializer API response models.
+REST envelopes for the initializer endpoints.
 
-Initializers configure the PyRIT environment (targets, datasets, env vars)
-before scenario execution. These models represent initializer metadata.
+Canonical initializer catalog types (``RegisteredInitializer``,
+``InitializerParameterSummary``) live in ``pyrit.models.catalog.initializer``
+and should be imported from there directly.
 """
 
 from pydantic import BaseModel, Field
 
 from pyrit.backend.models.common import PaginationInfo
 from pyrit.models import REGISTRY_NAME_PATTERN
+from pyrit.models.catalog.initializer import RegisteredInitializer
 
-
-class InitializerParameterSummary(BaseModel):
-    """Summary of an initializer-declared parameter."""
-
-    name: str = Field(..., description="Parameter name")
-    description: str = Field(..., description="Human-readable description of the parameter")
-    default: list[str] | None = Field(None, description="Default value(s), or None if required")
-
-
-class RegisteredInitializer(BaseModel):
-    """Summary of a registered initializer."""
-
-    initializer_name: str = Field(..., description="Initializer registry name (e.g., 'target')")
-    initializer_type: str = Field(..., description="Initializer class name (e.g., 'TargetInitializer')")
-    description: str = Field("", description="Human-readable description of the initializer")
-    required_env_vars: list[str] = Field(
-        default_factory=list, description="Environment variables required by this initializer"
-    )
-    supported_parameters: list[InitializerParameterSummary] = Field(
-        default_factory=list, description="Parameters accepted by this initializer"
-    )
+__all__ = [
+    "ListRegisteredInitializersResponse",
+    "RegisterInitializerRequest",
+]
 
 
 class ListRegisteredInitializersResponse(BaseModel):

pyrit/backend/models/scenarios.py

@@ -2,48 +2,23 @@
 # Licensed under the MIT license.
 
 """
-Scenario API response models.
+REST envelopes for the scenario endpoints.
 
-Scenarios are multi-attack security testing campaigns. These models represent
-the metadata about available scenarios (listing) and scenario execution (runs).
+Canonical scenario catalog/run types (``RegisteredScenario``,
+``ScenarioParameterSummary``, ``ScenarioRunSummary``, ``RunScenarioRequest``)
+live in ``pyrit.models.catalog.scenario`` and should be imported from there
+directly.
 """
 
-from datetime import datetime
-from enum import Enum
-from typing import Any
-
 from pydantic import BaseModel, Field
 
 from pyrit.backend.models.common import PaginationInfo
+from pyrit.models.catalog.scenario import RegisteredScenario, ScenarioRunSummary
 
-
-class ScenarioParameterSummary(BaseModel):
-    """Summary of a scenario-declared parameter."""
-
-    name: str = Field(..., description="Parameter name (e.g., 'max_turns')")
-    description: str = Field(..., description="Human-readable description of the parameter")
-    default: str | None = Field(None, description="Default value as a display string, or None if required")
-    param_type: str = Field(..., description="Type of the parameter as a display string (e.g., 'int', 'str')")
-    choices: list[str] | None = Field(None, description="Allowed values as strings, or None if unconstrained")
-    is_list: bool = Field(False, description="True when the parameter accepts a list of values (e.g., list[str])")
-
-
-class RegisteredScenario(BaseModel):
-    """Summary of a registered scenario."""
-
-    scenario_name: str = Field(..., description="Scenario name  (e.g., 'foundry.red_team_agent')")
-    scenario_type: str = Field(..., description="Scenario type identifier (e.g., 'RedTeamAgentScenario')")
-    description: str = Field(..., description="Human-readable description of the scenario")
-    default_strategy: str = Field(..., description="Default strategy name used when none specified")
-    aggregate_strategies: list[str] = Field(
-        ..., description="Aggregate strategies that combine multiple attack approaches"
-    )
-    all_strategies: list[str] = Field(..., description="All available concrete strategy names")
-    default_datasets: list[str] = Field(..., description="Default dataset names used by the scenario")
-    max_dataset_size: int | None = Field(None, description="Maximum items per dataset (None means unlimited)")
-    supported_parameters: list[ScenarioParameterSummary] = Field(
-        default_factory=list, description="Scenario-declared custom parameters"
-    )
+__all__ = [
+    "ListRegisteredScenariosResponse",
+    "ScenarioRunListResponse",
+]
 
 
 class ListRegisteredScenariosResponse(BaseModel):
@@ -53,73 +28,6 @@ class ListRegisteredScenariosResponse(BaseModel):
     pagination: PaginationInfo = Field(..., description="Pagination metadata")
 
 
-# ============================================================================
-# Scenario Run Models
-# ============================================================================
-
-
-class ScenarioRunStatus(str, Enum):
-    """Status of a scenario run, aligned with core ScenarioRunState."""
-
-    CREATED = "CREATED"
-    INITIALIZING = "INITIALIZING"
-    IN_PROGRESS = "IN_PROGRESS"
-    COMPLETED = "COMPLETED"
-    FAILED = "FAILED"
-    CANCELLED = "CANCELLED"
-
-
-class RunScenarioRequest(BaseModel):
-    """Request body for starting a scenario run."""
-
-    scenario_name: str = Field(..., description="Scenario name (e.g., 'foundry.red_team_agent')")
-    target_name: str = Field(..., description="Name of a registered target from the TargetRegistry")
-    initializers: list[str] | None = Field(
-        None, description="Initializer names to run before scenario (e.g., ['target', 'load_default_datasets'])"
-    )
-    strategies: list[str] | None = Field(None, description="Strategy names to use (uses scenario default if omitted)")
-    dataset_names: list[str] | None = Field(None, description="Dataset names to use (uses scenario default if omitted)")
-    max_dataset_size: int | None = Field(None, ge=1, description="Maximum items per dataset")
-    max_concurrency: int = Field(10, ge=1, le=100, description="Maximum concurrent operations")
-    max_retries: int = Field(0, ge=0, le=20, description="Maximum retry attempts on failure")
-    labels: dict[str, str] | None = Field(None, description="Labels to attach to memory entries")
-    scenario_params: dict[str, Any] | None = Field(
-        None,
-        description="Custom parameters for the scenario (passed to scenario.set_params_from_args). "
-        "Keys are parameter names declared by the scenario's supported_parameters().",
-    )
-    initializer_args: dict[str, dict[str, Any]] | None = Field(
-        None,
-        description="Per-initializer arguments keyed by initializer name. "
-        "Each value is a dict of args passed to that initializer's set_params_from_args(). "
-        "Example: {'target': {'endpoint': 'https://...'}}.",
-    )
-    scenario_result_id: str | None = Field(
-        None,
-        description="Optional ID of an existing ScenarioResult to resume. "
-        "If provided, the scenario will resume from prior progress instead of starting fresh.",
-    )
-
-
-class ScenarioRunSummary(BaseModel):
-    """Response for a scenario run (status + result details)."""
-
-    scenario_result_id: str = Field(..., description="UUID of the ScenarioResult in memory")
-    scenario_name: str = Field(..., description="Registry key of the scenario being run")
-    scenario_version: int = Field(0, ge=0, description="Version of the scenario")
-    status: ScenarioRunStatus = Field(..., description="Current run status")
-    created_at: datetime = Field(..., description="When the run was created")
-    updated_at: datetime = Field(..., description="When the run status last changed")
-    error: str | None = Field(None, description="Error message if status is FAILED")
-    error_type: str | None = Field(None, description="Exception class name if status is FAILED")
-    strategies_used: list[str] = Field(default_factory=list, description="Strategy names that were executed")
-    total_attacks: int = Field(0, ge=0, description="Total number of attack results persisted for this run")
-    completed_attacks: int = Field(0, ge=0, description="Number of attacks that reached a terminal outcome")
-    objective_achieved_rate: int = Field(0, ge=0, le=100, description="Success rate as percentage (0-100)")
-    labels: dict[str, str] = Field(default_factory=dict, description="Labels attached to this run")
-    completed_at: datetime | None = Field(None, description="When the scenario finished")
-
-
 class ScenarioRunListResponse(BaseModel):
     """Response for listing scenario runs."""
 

pyrit/backend/models/targets.py

@@ -2,72 +2,24 @@
 # Licensed under the MIT license.
 
 """
-Target instance models.
+REST envelopes and write-request types for the target endpoints.
 
-Targets have two concepts:
-- Types: Static metadata bundled with frontend (from registry)
-- Instances: Runtime objects created via API with specific configuration
-
-This module defines the Instance models for runtime target management.
+Canonical target catalog types (``TargetInstance``, ``TargetCapabilitiesInfo``)
+live in ``pyrit.models.catalog.target`` and should be imported from there
+directly.
 """
 
 from typing import Any, Literal
 
 from pydantic import BaseModel, Field
 
 from pyrit.backend.models.common import PaginationInfo
+from pyrit.models.catalog.target import TargetInstance
 
-
-class TargetCapabilitiesInfo(BaseModel):
-    """
-    Wire-format snapshot of a target's capabilities.
-
-    Mirrors the domain ``TargetCapabilities`` dataclass for API consumers
-    (notably the GUI). Modality combinations (``frozenset[frozenset[...]]``)
-    are flattened into sorted unique modality lists since the frontend uses
-    them only for per-piece modality checks.
-    """
-
-    supports_multi_turn: bool = Field(False, description="Target natively supports multi-turn conversations")
-    supports_multi_message_pieces: bool = Field(
-        False, description="Target supports multiple message pieces in a single request"
-    )
-    supports_json_schema: bool = Field(False, description="Target can constrain output to a provided JSON schema")
-    supports_json_output: bool = Field(False, description="Target supports JSON output mode")
-    supports_editable_history: bool = Field(False, description="Target allows attack history to be modified")
-    supports_system_prompt: bool = Field(False, description="Target natively supports system prompts")
-    supported_input_modalities: list[str] = Field(
-        default_factory=lambda: ["text"],
-        description="Sorted unique input modality data types the target accepts (e.g., ['image_path', 'text'])",
-    )
-    supported_output_modalities: list[str] = Field(
-        default_factory=lambda: ["text"],
-        description="Sorted unique output modality data types the target produces (e.g., ['audio_path', 'text'])",
-    )
-
-
-class TargetInstance(BaseModel):
-    """
-    A runtime target instance.
-
-    Created either by an initializer (at startup) or by user (via API).
-    Also used as the create-target response (same shape as GET).
-    """
-
-    target_registry_name: str = Field(..., description="Target registry key (e.g., 'azure_openai_chat')")
-    target_type: str = Field(..., description="Target class name (e.g., 'OpenAIChatTarget')")
-    endpoint: str | None = Field(None, description="Target endpoint URL")
-    model_name: str | None = Field(None, description="Model or deployment name used in API calls")
-    underlying_model_name: str | None = Field(None, description="Underlying model name if different (e.g., 'gpt-4o')")
-    temperature: float | None = Field(None, description="Temperature parameter for generation")
-    top_p: float | None = Field(None, description="Top-p parameter for generation")
-    max_requests_per_minute: int | None = Field(None, description="Maximum requests per minute")
-    capabilities: TargetCapabilitiesInfo = Field(..., description="Structured snapshot of target capabilities")
-    target_specific_params: dict[str, Any] | None = Field(None, description="Additional target-specific parameters")
-    inner_targets: list["TargetInstance"] | None = Field(
-        None, description="Inner targets for composite targets like RoundRobinTarget"
-    )
-    identifier_hash: str | None = Field(None, description="ComponentIdentifier content hash for duplicate detection")
+__all__ = [
+    "CreateTargetRequest",
+    "TargetListResponse",
+]
 
 
 class TargetListResponse(BaseModel):

pyrit/backend/routes/initializers.py

@@ -18,10 +18,10 @@
 from pyrit.backend.models.common import ProblemDetail
 from pyrit.backend.models.initializers import (
     ListRegisteredInitializersResponse,
-    RegisteredInitializer,
     RegisterInitializerRequest,
 )
 from pyrit.backend.services.initializer_service import get_initializer_service
+from pyrit.models.catalog.initializer import RegisteredInitializer
 
 router = APIRouter(prefix="/initializers", tags=["initializers"])
 

pyrit/backend/routes/scenarios.py

@@ -17,14 +17,16 @@
 from pyrit.backend.models.common import ProblemDetail
 from pyrit.backend.models.scenarios import (
     ListRegisteredScenariosResponse,
-    RegisteredScenario,
-    RunScenarioRequest,
     ScenarioRunListResponse,
-    ScenarioRunSummary,
 )
 from pyrit.backend.services.scenario_run_service import get_scenario_run_service
 from pyrit.backend.services.scenario_service import get_scenario_service
 from pyrit.models import ScenarioResult
+from pyrit.models.catalog.scenario import (
+    RegisteredScenario,
+    RunScenarioRequest,
+    ScenarioRunSummary,
+)
 
 router = APIRouter(prefix="/scenarios", tags=["scenarios"])
 

pyrit/backend/routes/targets.py

@@ -13,10 +13,10 @@
 from pyrit.backend.models.common import ProblemDetail
 from pyrit.backend.models.targets import (
     CreateTargetRequest,
-    TargetInstance,
     TargetListResponse,
 )
 from pyrit.backend.services.target_service import get_target_service
+from pyrit.models.catalog.target import TargetInstance
 
 router = APIRouter(prefix="/targets", tags=["targets"])
 

pyrit/backend/services/initializer_service.py

@@ -13,8 +13,10 @@
 
 from pyrit.backend.models.common import PaginationInfo
 from pyrit.backend.models.initializers import (
-    InitializerParameterSummary,
     ListRegisteredInitializersResponse,
+)
+from pyrit.models.catalog.initializer import (
+    InitializerParameterSummary,
     RegisteredInitializer,
 )
 from pyrit.registry import InitializerMetadata, InitializerRegistry