Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 1 addition & 1 deletion pyrit/backend/mappers/target_mappers.py
Original file line number Diff line number Diff line change
Expand Up @@ -5,7 +5,7 @@
Target mappers – domain → DTO translation for target-related models.
"""

from pyrit.backend.models.targets import TargetCapabilitiesInfo, TargetInstance
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
Expand Down
14 changes: 2 additions & 12 deletions pyrit/backend/models/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -48,20 +48,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,
)

Expand Down Expand Up @@ -105,16 +100,11 @@
"PreviewStep",
# Scenarios
"ListRegisteredScenariosResponse",
"RegisteredScenario",
"ScenarioParameterSummary",
"ScenarioRunListResponse",
# Initializers
"InitializerParameterSummary",
"ListRegisteredInitializersResponse",
"RegisteredInitializer",
"RegisterInitializerRequest",
# Targets
"CreateTargetRequest",
"TargetCapabilitiesInfo",
"TargetInstance",
"TargetListResponse",
]
33 changes: 9 additions & 24 deletions pyrit/backend/models/initializers.py
Original file line number Diff line number Diff line change
Expand Up @@ -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):
Expand Down
112 changes: 10 additions & 102 deletions pyrit/backend/models/scenarios.py
Original file line number Diff line number Diff line change
Expand Up @@ -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):
Expand All @@ -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."""

Expand Down
66 changes: 9 additions & 57 deletions pyrit/backend/models/targets.py
Original file line number Diff line number Diff line change
Expand Up @@ -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):
Expand Down
2 changes: 1 addition & 1 deletion pyrit/backend/routes/initializers.py
Original file line number Diff line number Diff line change
Expand Up @@ -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"])

Expand Down
12 changes: 8 additions & 4 deletions pyrit/backend/routes/scenarios.py
Original file line number Diff line number Diff line change
Expand Up @@ -12,18 +12,22 @@
/api/scenarios/runs — scenario execution lifecycle
"""

from typing import Any

from fastapi import APIRouter, HTTPException, Query, status

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.catalog.scenario import (
RegisteredScenario,
RunScenarioRequest,
ScenarioRunSummary,
)

router = APIRouter(prefix="/scenarios", tags=["scenarios"])

Expand Down Expand Up @@ -199,7 +203,7 @@ async def cancel_scenario_run(scenario_result_id: str) -> ScenarioRunSummary: #
409: {"model": ProblemDetail, "description": "Run not yet completed"},
},
)
async def get_scenario_run_results(scenario_result_id: str) -> dict: # pyrit-async-suffix-exempt
async def get_scenario_run_results(scenario_result_id: str) -> dict[str, Any]: # pyrit-async-suffix-exempt
"""
Get detailed results for a completed scenario run.

Expand Down
2 changes: 1 addition & 1 deletion pyrit/backend/routes/targets.py
Original file line number Diff line number Diff line change
Expand Up @@ -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"])

Expand Down
4 changes: 3 additions & 1 deletion pyrit/backend/services/initializer_service.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
Loading
Loading