Skip to content
Merged
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
64 changes: 49 additions & 15 deletions .github/instructions/scenarios.instructions.md
Original file line number Diff line number Diff line change
Expand Up @@ -28,7 +28,7 @@ class MyScenario(Scenario):
super().__init__(
version=self.VERSION,
technique_class=MyTechnique,
default_dataset_config=DatasetConfiguration(dataset_names=["my_dataset"]),
default_dataset_config=DatasetAttackConfiguration(sources=[DatasetSource(name="my_dataset")]),
objective_scorer=objective_scorer or self._get_default_objective_scorer(),
scenario_result_id=scenario_result_id,
)
Expand Down Expand Up @@ -68,7 +68,7 @@ def __init__(
super().__init__(
version=self.VERSION,
technique_class=MyTechnique,
default_dataset_config=DatasetConfiguration(dataset_names=["my_dataset"]),
default_dataset_config=DatasetAttackConfiguration(sources=[DatasetSource(name="my_dataset")]),
objective_scorer=objective_scorer,
)
```
Expand Down Expand Up @@ -113,30 +113,64 @@ Dropping a common input is not silent: `set_params_from_args` rejects any value

## Dataset Loading

Datasets are read from `CentralMemory`.
Datasets are read from `CentralMemory`. New runs call `prepare_async()` before
reading. Reads, discovery, estimates, and resume must never fetch or store datasets.
Resolve parameter-dependent source names before preparation, not during seed reads.

### Basic — named datasets:
```python
DatasetConfiguration(
dataset_names=["airt_hate", "airt_violence"],
max_dataset_size=10, # optional: sample up to N per dataset
DatasetAttackConfiguration(
sources=[
DatasetSource(name="harmbench"),
DatasetSource(name="airt_hate", max_size=10),
DatasetSource(name="team_objectives", max_size="all"),
],
max_per_dataset=5,
max_total=12,
)
```

With filtered populations of 100, 40, and 7 groups, selection has two steps:

| Dataset | Full population | After source cap |
| --- | --- | --- |
| `harmbench` | 100 | 5 (inherits `max_per_dataset`) |
| `airt_hate` | 40 | 10 (explicit `max_size`) |
| `team_objectives` | 7 | 7 (`"all"` removes the source cap) |

The total cap then samples 12 groups from these 22 groups. The split can be 3, 6,
and 3, for example; it is not balanced or guaranteed to cover every dataset.
A dataset with no selected groups keeps its key with an empty list.

### Advanced — custom subclass for filtering:
```python
class MyDatasetConfiguration(DatasetConfiguration):
def get_seed_groups(self) -> dict[str, list[SeedGroup]]:
result = super().get_seed_groups()
# Filter by selected techniques via self._scenario_techniques
return filtered_result
class MyDatasetConfiguration(DatasetAttackConfiguration):
def _build_attack_groups(self, seeds: list[Seed]) -> list[AttackSeedGroup]:
return build_custom_groups(seeds)
```

Options:
- `dataset_names` — load by name from memory
- `seed_groups` — pass explicit groups (mutually exclusive with `dataset_names`)
- `max_dataset_size` — cap per dataset
- Override `_load_seed_groups_for_dataset()` for custom loading
- `sources` selects datasets by name. Source `max_size` overrides `max_per_dataset`;
omitted, `None`, empty, and `"default"` mean inherit; `"all"` means no cap.
Comment thread
richlundeen marked this conversation as resolved.
- All dataset limits use the same rule: omitted, `None`, empty, and `"default"` use the
default; `"all"` removes that limit; a positive integer sets a cap.
- `max_per_dataset=5` is the named-objective default. `max_total="all"` leaves the
combined selection uncapped. Apply source limits before the total limit.
- `fetch=DatasetFetchPolicy.IF_MISSING` prepares absent registered datasets.
`NEVER` requires stored data. A filter miss must never fetch.
- `seed_groups` and `seeds` are inline alternatives. They never use memory or providers.
- Validators run on full filtered populations before any sampling.
- Ingredients must remain complete. Use `max_per_dataset="all"` and no finite source
`max_size`; cap the assembled attack groups with `max_total`, not ingredient rows.
- Use `with_overrides()` instead of reconstructing a subclass or mutating its defaults.
- Keep custom shaping in `_build_attack_groups()` or `_build_groups_by_dataset_async()`.
Reads use `_collect_seeds_for_dataset_async()` and must not fetch.
- `dataset_names`, `max_dataset_size`, `auto_fetch`, and `per_dataset()` are deprecated.

`max_dataset_size` is an exact alias for `max_total`; it does not disable source caps.
To keep the old total-only selection of up to 10 groups, use
`max_per_dataset="all", max_total=10`. To remove both caps, set both to `"all"`;
`None` now uses the default, not unlimited selection.

## Technique Enum

Expand Down
4 changes: 2 additions & 2 deletions doc/code/datasets/6_generated_datasets.ipynb
Original file line number Diff line number Diff line change
Expand Up @@ -240,11 +240,11 @@
],
"source": [
"from pyrit.output import output_scenario_async, output_scenario_attacks_async\n",
"from pyrit.scenario import DatasetAttackConfiguration\n",
"from pyrit.scenario import DatasetAttackConfiguration, DatasetFetchPolicy, DatasetSource\n",
"from pyrit.scenario.airt import RapidResponse, RapidResponseTechnique\n",
"from pyrit.score import SelfAskTrueFalseScorer\n",
"\n",
"dataset_config = DatasetAttackConfiguration(dataset_names=[dataset_name], auto_fetch=False)\n",
"dataset_config = DatasetAttackConfiguration(sources=[DatasetSource(name=dataset_name)], fetch=DatasetFetchPolicy.NEVER)\n",
"scenario = RapidResponse(objective_scorer=SelfAskTrueFalseScorer(chat_target=OpenAIChatTarget()))\n",
"scenario.set_params_from_args(\n",
" args={\n",
Expand Down
4 changes: 2 additions & 2 deletions doc/code/datasets/6_generated_datasets.py
Original file line number Diff line number Diff line change
Expand Up @@ -76,11 +76,11 @@

# %%
from pyrit.output import output_scenario_async, output_scenario_attacks_async
from pyrit.scenario import DatasetAttackConfiguration
from pyrit.scenario import DatasetAttackConfiguration, DatasetFetchPolicy, DatasetSource
from pyrit.scenario.airt import RapidResponse, RapidResponseTechnique
from pyrit.score import SelfAskTrueFalseScorer

dataset_config = DatasetAttackConfiguration(dataset_names=[dataset_name], auto_fetch=False)
dataset_config = DatasetAttackConfiguration(sources=[DatasetSource(name=dataset_name)], fetch=DatasetFetchPolicy.NEVER)
scenario = RapidResponse(objective_scorer=SelfAskTrueFalseScorer(chat_target=OpenAIChatTarget()))
scenario.set_params_from_args(
args={
Expand Down
13 changes: 9 additions & 4 deletions doc/code/scenarios/0_scenarios.ipynb
Original file line number Diff line number Diff line change
Expand Up @@ -65,9 +65,11 @@
" Matrix-shaped scenarios delegate to `build_matrix_atomic_attacks(context=...)` in one line.\n",
"\n",
"3. **Default Dataset**: Pass `default_dataset_config=` to `super().__init__()` to specify the datasets your scenario uses out of the box.\n",
" - Returns a `DatasetConfiguration` with one or more named datasets (e.g., `DatasetConfiguration(dataset_names=[\"my_dataset\"])`)\n",
" - Returns a `DatasetAttackConfiguration` with named sources (e.g., `sources=[DatasetSource(name=\"my_dataset\")]`)\n",
" - Users can override this at runtime via `--dataset-names` in the CLI or by passing a custom `dataset_config` programmatically\n",
" - `DatasetAttackConfiguration` selects at most 5 attack groups unless you set `max_dataset_size`; `max_dataset_size=None` uses all groups\n",
" - Named sources select at most 5 attack groups per dataset by default; `max_total` caps the union. For any limit, omitted/`None`/empty/`\"default\"` uses the default; `\"all\"` removes that limit.\n",
" - The deprecated `max_dataset_size` is an exact alias for `max_total`, not a way to disable source caps. To keep the old total-only selection, use `max_per_dataset=\"all\", max_total=10`; to remove both caps, set both to `\"all\"`. `None` now uses the default, not unlimited selection.\n",
" - New runs prepare missing registered datasets once; reads, estimates, and resume never fetch\n",
"\n",
"4. **Constructor**: Use `@apply_defaults` decorator and call `super().__init__()` with scenario metadata:\n",
" - `name`: Descriptive name for your scenario\n",
Expand Down Expand Up @@ -113,7 +115,8 @@
"source": [
"from pyrit.common import apply_defaults\n",
"from pyrit.scenario import (\n",
" DatasetConfiguration,\n",
" DatasetAttackConfiguration,\n",
" DatasetSource,\n",
" Scenario,\n",
" ScenarioTechnique,\n",
")\n",
Expand Down Expand Up @@ -159,7 +162,9 @@
" version=self.VERSION,\n",
" objective_scorer=self._objective_scorer,\n",
" technique_class=MyTechnique,\n",
" default_dataset_config=DatasetConfiguration(dataset_names=[\"dataset_name\"], max_dataset_size=4),\n",
" default_dataset_config=DatasetAttackConfiguration(\n",
" sources=[DatasetSource(name=\"dataset_name\")], max_total=4\n",
" ),\n",
" scenario_result_id=scenario_result_id,\n",
" )\n",
"\n",
Expand Down
13 changes: 9 additions & 4 deletions doc/code/scenarios/0_scenarios.py
Original file line number Diff line number Diff line change
Expand Up @@ -67,9 +67,11 @@
# Matrix-shaped scenarios delegate to `build_matrix_atomic_attacks(context=...)` in one line.
#
# 3. **Default Dataset**: Pass `default_dataset_config=` to `super().__init__()` to specify the datasets your scenario uses out of the box.
# - Returns a `DatasetConfiguration` with one or more named datasets (e.g., `DatasetConfiguration(dataset_names=["my_dataset"])`)
# - Returns a `DatasetAttackConfiguration` with named sources (e.g., `sources=[DatasetSource(name="my_dataset")]`)
# - Users can override this at runtime via `--dataset-names` in the CLI or by passing a custom `dataset_config` programmatically
# - `DatasetAttackConfiguration` selects at most 5 attack groups unless you set `max_dataset_size`; `max_dataset_size=None` uses all groups
# - Named sources select at most 5 attack groups per dataset by default; `max_total` caps the union. For any limit, omitted/`None`/empty/`"default"` uses the default; `"all"` removes that limit.
# - The deprecated `max_dataset_size` is an exact alias for `max_total`, not a way to disable source caps. To keep the old total-only selection, use `max_per_dataset="all", max_total=10`; to remove both caps, set both to `"all"`. `None` now uses the default, not unlimited selection.
# - New runs prepare missing registered datasets once; reads, estimates, and resume never fetch
#
# 4. **Constructor**: Use `@apply_defaults` decorator and call `super().__init__()` with scenario metadata:
# - `name`: Descriptive name for your scenario
Expand Down Expand Up @@ -97,7 +99,8 @@
# %%
from pyrit.common import apply_defaults
from pyrit.scenario import (
DatasetConfiguration,
DatasetAttackConfiguration,
DatasetSource,
Scenario,
ScenarioTechnique,
)
Expand Down Expand Up @@ -143,7 +146,9 @@ def __init__(
version=self.VERSION,
objective_scorer=self._objective_scorer,
technique_class=MyTechnique,
default_dataset_config=DatasetConfiguration(dataset_names=["dataset_name"], max_dataset_size=4),
default_dataset_config=DatasetAttackConfiguration(
sources=[DatasetSource(name="dataset_name")], max_total=4
),
scenario_result_id=scenario_result_id,
)

Expand Down
10 changes: 5 additions & 5 deletions doc/code/scenarios/1_common_scenario_parameters.ipynb
Original file line number Diff line number Diff line change
Expand Up @@ -75,7 +75,7 @@
"## Dataset Configuration\n",
"\n",
"`DatasetAttackConfiguration` controls which prompts (objectives) are sent to the target.\n",
"The simplest approach uses `dataset_names` to load datasets by name from memory.\n",
"Use `sources` to select named datasets from memory.\n",
"By default, `RedTeamAgent` loads four random objectives from HarmBench [@mazeika2024harmbench]."
]
},
Expand All @@ -86,9 +86,9 @@
"metadata": {},
"outputs": [],
"source": [
"from pyrit.scenario import DatasetAttackConfiguration\n",
"from pyrit.scenario import DatasetAttackConfiguration, DatasetSource\n",
"\n",
"dataset_config = DatasetAttackConfiguration(dataset_names=[\"harmbench\"], max_dataset_size=2)"
"dataset_config = DatasetAttackConfiguration(sources=[DatasetSource(name=\"harmbench\")], max_total=2)"
]
},
{
Expand Down Expand Up @@ -129,8 +129,8 @@
" if all(seed.value.isascii() and seed.value.isprintable() for seed in group.seeds)\n",
"]\n",
"\n",
"# Pass explicit seed_groups instead of dataset_names\n",
"dataset_config = DatasetAttackConfiguration(seed_groups=seed_groups, max_dataset_size=2)"
"# Pass explicit seed_groups instead of named sources\n",
"dataset_config = DatasetAttackConfiguration(seed_groups=seed_groups, max_total=2)"
]
},
{
Expand Down
10 changes: 5 additions & 5 deletions doc/code/scenarios/1_common_scenario_parameters.py
Original file line number Diff line number Diff line change
Expand Up @@ -42,13 +42,13 @@
# ## Dataset Configuration
#
# `DatasetAttackConfiguration` controls which prompts (objectives) are sent to the target.
# The simplest approach uses `dataset_names` to load datasets by name from memory.
# Use `sources` to select named datasets from memory.
# By default, `RedTeamAgent` loads four random objectives from HarmBench [@mazeika2024harmbench].

# %%
from pyrit.scenario import DatasetAttackConfiguration
from pyrit.scenario import DatasetAttackConfiguration, DatasetSource

dataset_config = DatasetAttackConfiguration(dataset_names=["harmbench"], max_dataset_size=2)
dataset_config = DatasetAttackConfiguration(sources=[DatasetSource(name="harmbench")], max_total=2)

# %% [markdown]
# For more control, use `SeedDatasetProvider` to fetch datasets and pass explicit `seed_groups`.
Expand All @@ -69,8 +69,8 @@
if all(seed.value.isascii() and seed.value.isprintable() for seed in group.seeds)
]

# Pass explicit seed_groups instead of dataset_names
dataset_config = DatasetAttackConfiguration(seed_groups=seed_groups, max_dataset_size=2)
# Pass explicit seed_groups instead of named sources
dataset_config = DatasetAttackConfiguration(seed_groups=seed_groups, max_total=2)

# %% [markdown]
# ## Technique Selection and Composition
Expand Down
10 changes: 5 additions & 5 deletions doc/code/scenarios/3_adaptive_scenarios.ipynb
Original file line number Diff line number Diff line change
Expand Up @@ -113,7 +113,7 @@
"\n",
"from pyrit.output.scenario_result.pretty import PrettyScenarioResultMemoryPrinter\n",
"from pyrit.registry import TargetRegistry\n",
"from pyrit.scenario import DatasetAttackConfiguration\n",
"from pyrit.scenario import DatasetAttackConfiguration, DatasetSource\n",
"from pyrit.scenario.scenarios.adaptive import TextAdaptive\n",
"from pyrit.setup import initialize_from_config_async\n",
"\n",
Expand Down Expand Up @@ -452,8 +452,8 @@
" \"objective_target\": objective_target,\n",
" \"scenario_techniques\": [technique_class(\"single_turn\")],\n",
" \"dataset_config\": DatasetAttackConfiguration(\n",
" dataset_names=[\"airt_hate\", \"airt_violence\"],\n",
" max_dataset_size=4,\n",
" sources=[DatasetSource(name=\"airt_hate\"), DatasetSource(name=\"airt_violence\")],\n",
" max_total=4,\n",
" ),\n",
" }\n",
")\n",
Expand Down Expand Up @@ -577,8 +577,8 @@
" \"objective_target\": objective_target,\n",
" \"scenario_techniques\": [technique_class(\"single_turn\")],\n",
" \"dataset_config\": DatasetAttackConfiguration(\n",
" dataset_names=[\"airt_hate\", \"airt_violence\"],\n",
" max_dataset_size=4,\n",
" sources=[DatasetSource(name=\"airt_hate\"), DatasetSource(name=\"airt_violence\")],\n",
" max_total=4,\n",
" ),\n",
" }\n",
")\n",
Expand Down
10 changes: 5 additions & 5 deletions doc/code/scenarios/3_adaptive_scenarios.py
Original file line number Diff line number Diff line change
Expand Up @@ -48,7 +48,7 @@

from pyrit.output.scenario_result.pretty import PrettyScenarioResultMemoryPrinter
from pyrit.registry import TargetRegistry
from pyrit.scenario import DatasetAttackConfiguration
from pyrit.scenario import DatasetAttackConfiguration, DatasetSource
from pyrit.scenario.scenarios.adaptive import TextAdaptive
from pyrit.setup import initialize_from_config_async

Expand Down Expand Up @@ -102,8 +102,8 @@
"objective_target": objective_target,
"scenario_techniques": [technique_class("single_turn")],
"dataset_config": DatasetAttackConfiguration(
dataset_names=["airt_hate", "airt_violence"],
max_dataset_size=4,
sources=[DatasetSource(name="airt_hate"), DatasetSource(name="airt_violence")],
max_total=4,
),
}
)
Expand Down Expand Up @@ -132,8 +132,8 @@
"objective_target": objective_target,
"scenario_techniques": [technique_class("single_turn")],
"dataset_config": DatasetAttackConfiguration(
dataset_names=["airt_hate", "airt_violence"],
max_dataset_size=4,
sources=[DatasetSource(name="airt_hate"), DatasetSource(name="airt_violence")],
max_total=4,
),
}
)
Expand Down
6 changes: 3 additions & 3 deletions doc/scanner/1_pyrit_scan.ipynb
Original file line number Diff line number Diff line change
Expand Up @@ -1287,7 +1287,7 @@
"\n",
"from pyrit.common import apply_defaults\n",
"from pyrit.prompt_target.openai.openai_chat_target import OpenAIChatTarget\n",
"from pyrit.scenario import DatasetAttackConfiguration, Scenario, ScenarioTechnique\n",
"from pyrit.scenario import DatasetAttackConfiguration, DatasetSource, Scenario, ScenarioTechnique\n",
"from pyrit.score import SelfAskRefusalScorer, TrueFalseInverterScorer\n",
"from pyrit.setup import initialize_pyrit_async\n",
"\n",
Expand All @@ -1311,8 +1311,8 @@
" version=1,\n",
" objective_scorer=TrueFalseInverterScorer(scorer=SelfAskRefusalScorer(chat_target=OpenAIChatTarget())),\n",
" technique_class=MyCustomTechnique,\n",
" # DatasetAttackConfiguration selects at most 5 attack groups by default; set max_dataset_size to change it.\n",
" default_dataset_config=DatasetAttackConfiguration(dataset_names=[\"harmbench\"]),\n",
" # Named sources default to 5 groups each; use max_per_dataset and max_total to set caps.\n",
" default_dataset_config=DatasetAttackConfiguration(sources=[DatasetSource(name=\"harmbench\")]),\n",
" scenario_result_id=scenario_result_id,\n",
" )\n",
" # ... your scenario-specific initialization code\n",
Expand Down
6 changes: 3 additions & 3 deletions doc/scanner/1_pyrit_scan.py
Original file line number Diff line number Diff line change
Expand Up @@ -182,7 +182,7 @@

from pyrit.common import apply_defaults
from pyrit.prompt_target.openai.openai_chat_target import OpenAIChatTarget
from pyrit.scenario import DatasetAttackConfiguration, Scenario, ScenarioTechnique
from pyrit.scenario import DatasetAttackConfiguration, DatasetSource, Scenario, ScenarioTechnique
from pyrit.score import SelfAskRefusalScorer, TrueFalseInverterScorer
from pyrit.setup import initialize_pyrit_async

Expand All @@ -206,8 +206,8 @@ def __init__(self, *, scenario_result_id=None, **kwargs):
version=1,
objective_scorer=TrueFalseInverterScorer(scorer=SelfAskRefusalScorer(chat_target=OpenAIChatTarget())),
technique_class=MyCustomTechnique,
# DatasetAttackConfiguration selects at most 5 attack groups by default; set max_dataset_size to change it.
default_dataset_config=DatasetAttackConfiguration(dataset_names=["harmbench"]),
# Named sources default to 5 groups each; use max_per_dataset and max_total to set caps.
default_dataset_config=DatasetAttackConfiguration(sources=[DatasetSource(name="harmbench")]),
scenario_result_id=scenario_result_id,
)
# ... your scenario-specific initialization code
Expand Down
4 changes: 2 additions & 2 deletions doc/scanner/adaptive.ipynb
Original file line number Diff line number Diff line change
Expand Up @@ -107,15 +107,15 @@
"\n",
"from pyrit.output import output_scenario_async\n",
"from pyrit.registry import TargetRegistry\n",
"from pyrit.scenario import DatasetAttackConfiguration\n",
"from pyrit.scenario import DatasetAttackConfiguration, DatasetSource\n",
"from pyrit.scenario.adaptive import TextAdaptive\n",
"from pyrit.setup import initialize_from_config_async\n",
"\n",
"await initialize_from_config_async(config_path=Path(\"pyrit_conf.yaml\")) # type: ignore\n",
"\n",
"objective_target = TargetRegistry.get_registry_singleton().instances.get(\"openai_chat\")\n",
"\n",
"dataset_config = DatasetAttackConfiguration(dataset_names=[\"airt_hate\"], max_dataset_size=2)\n",
"dataset_config = DatasetAttackConfiguration(sources=[DatasetSource(name=\"airt_hate\")], max_total=2)\n",
"\n",
"scenario = TextAdaptive()\n",
"scenario.set_params_from_args( # type: ignore\n",
Expand Down
Loading
Loading