From 78f57417653fe8468f88d8c892cde266c556de7d Mon Sep 17 00:00:00 2001 From: Richard Lundeen Date: Sun, 11 Oct 2026 15:16:19 -0700 Subject: [PATCH] [BREAKING] FEAT: Add scenario dataset controls across application surfaces Expose scenario-owned dataset options, shared requests, and source-level controls in Python, REST, CLI, shell, and GUI. Preserve read-only estimates and resume behavior. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> Copilot-Session: 16cbdc1b-2bf6-4286-be66-c9f30c7ed72b --- .../instructions/scenarios.instructions.md | 14 ++ doc/code/scenarios/0_scenarios.ipynb | 3 +- doc/code/scenarios/0_scenarios.py | 3 +- .../1_common_scenario_parameters.ipynb | 2 +- .../scenarios/1_common_scenario_parameters.py | 2 +- doc/code/scenarios/3_adaptive_scenarios.ipynb | 2 +- doc/code/scenarios/3_adaptive_scenarios.py | 2 +- doc/gui/0_gui.md | 24 ++- doc/scanner/1_pyrit_scan.ipynb | 10 +- doc/scanner/1_pyrit_scan.py | 10 +- doc/scanner/adaptive.ipynb | 4 +- doc/scanner/adaptive.py | 4 +- doc/scanner/airt.ipynb | 22 +- doc/scanner/airt.py | 22 +- doc/scanner/benchmark.ipynb | 2 +- doc/scanner/benchmark.py | 2 +- doc/scanner/foundry.ipynb | 2 +- doc/scanner/foundry.py | 2 +- doc/scanner/garak.ipynb | 28 +-- doc/scanner/garak.py | 28 +-- .../Scenarios/ScenarioDatasetTable.styles.ts | 13 ++ .../Scenarios/ScenarioDatasetTable.test.tsx | 109 ++++++++++ .../Scenarios/ScenarioDatasetTable.tsx | 194 +++++++++++++++++ .../Scenarios/ScenarioDetail.test.tsx | 113 +++++++--- .../components/Scenarios/ScenarioDetail.tsx | 177 +++++++--------- .../src/hooks/useScenarioDatasets.test.tsx | 58 +++++ frontend/src/hooks/useScenarioDatasets.ts | 50 +++++ frontend/src/services/api.ts | 10 + frontend/src/types/index.ts | 59 ++++++ pyrit/backend/models/datasets.py | 1 + pyrit/backend/services/dataset_service.py | 3 + .../scenario_configuration_resolver.py | 31 +-- .../backend/services/scenario_run_service.py | 22 +- pyrit/backend/services/scenario_service.py | 6 +- pyrit/cli/_cli_args.py | 88 +++++++- pyrit/cli/_output.py | 11 + pyrit/cli/pyrit_scan.py | 29 ++- pyrit/cli/pyrit_shell.py | 27 ++- pyrit/memory/memory_interface.py | 12 +- pyrit/models/__init__.py | 16 ++ pyrit/models/catalog/scenario.py | 142 ++++++------- pyrit/models/scenario_dataset.py | 130 ++++++++++++ pyrit/models/seeds/seed_dataset_summary.py | 1 + .../registry/components/scenario_registry.py | 15 +- pyrit/scenario/core/dataset_configuration.py | 13 +- pyrit/scenario/core/dataset_overrides.py | 122 +++++++++++ pyrit/scenario/core/scenario.py | 55 ++++- pyrit/scenario/scenarios/airt/psychosocial.py | 42 +++- pyrit/scenario/scenarios/garak/api_key.py | 2 + pyrit/scenario/scenarios/garak/encoding.py | 4 + .../scenario/scenarios/garak/exploitation.py | 6 +- pyrit/scenario/scenarios/garak/figstep.py | 4 + .../scenarios/garak/latent_injection.py | 2 + .../scenarios/garak/package_hallucination.py | 15 +- .../scenario/scenarios/garak/prompt_inject.py | 2 + .../garak/system_prompt_extraction.py | 4 + .../scenario/scenarios/garak/web_injection.py | 16 ++ tests/unit/backend/test_dataset_service.py | 3 + .../test_scenario_configuration_resolver.py | 46 ++-- tests/unit/backend/test_scenario_resume.py | 30 ++- .../unit/backend/test_scenario_run_service.py | 44 ++-- tests/unit/backend/test_scenario_service.py | 19 +- tests/unit/cli/test_dataset_request.py | 54 +++++ ...st_interface_seed_dataset_summary_query.py | 16 +- tests/unit/models/test_scenario_catalog.py | 3 + tests/unit/models/test_scenario_request.py | 99 +++++++++ tests/unit/registry/test_scenario_registry.py | 11 + tests/unit/scenario/airt/test_leakage.py | 4 + tests/unit/scenario/airt/test_psychosocial.py | 11 +- tests/unit/scenario/airt/test_scam.py | 4 + .../core/test_scenario_dataset_overrides.py | 198 ++++++++++++++++++ .../scenario/foundry/test_red_team_agent.py | 4 + tests/unit/scenario/garak/test_api_key.py | 4 +- tests/unit/scenario/garak/test_divergence.py | 3 +- tests/unit/scenario/garak/test_encoding.py | 4 + 75 files changed, 1935 insertions(+), 414 deletions(-) create mode 100644 frontend/src/components/Scenarios/ScenarioDatasetTable.styles.ts create mode 100644 frontend/src/components/Scenarios/ScenarioDatasetTable.test.tsx create mode 100644 frontend/src/components/Scenarios/ScenarioDatasetTable.tsx create mode 100644 frontend/src/hooks/useScenarioDatasets.test.tsx create mode 100644 frontend/src/hooks/useScenarioDatasets.ts create mode 100644 pyrit/models/scenario_dataset.py create mode 100644 pyrit/scenario/core/dataset_overrides.py create mode 100644 tests/unit/cli/test_dataset_request.py create mode 100644 tests/unit/scenario/core/test_scenario_dataset_overrides.py diff --git a/.github/instructions/scenarios.instructions.md b/.github/instructions/scenarios.instructions.md index bef7214bd7..20e5f4c80b 100644 --- a/.github/instructions/scenarios.instructions.md +++ b/.github/instructions/scenarios.instructions.md @@ -163,6 +163,20 @@ Options: - 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. +- Clients and Python use `ScenarioDatasetRequest` through the common `dataset_overrides` + parameter. `dataset_config` still replaces the default; overrides apply on top. +- Resolve parameter/technique-dependent names in `_resolve_default_dataset_config()`, + before overrides, validation, and preparation. Do not rewrite user selections in validators. +- Declare `DATASET_OPTIONS` (or parameter-aware `get_dataset_options()`): `ANY` only for + generic named populations, `SUBSET` for approved alternatives, `FIXED` for ingredients. + Disable limits that would truncate ingredients. Compound source edits are ambiguous; + preserve child shaping/caps and allow only supported global edits. +- Estimates expose `dataset_defaults`, `effective_datasets`, and `dataset_options`. + These are response projections, not incoming edit state. Never read seeds for projections. + The catalog reads options and defaults from registry metadata, never from an estimate. +- Resume replays the saved launch request; the stored run plan fixes the selected seed groups, + so caps and fetch policies cannot change a resumed run. Dataset names are checked by the + stored scenario identity. Resume never prepares datasets. - 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. diff --git a/doc/code/scenarios/0_scenarios.ipynb b/doc/code/scenarios/0_scenarios.ipynb index b088db7741..7b0b110a57 100644 --- a/doc/code/scenarios/0_scenarios.ipynb +++ b/doc/code/scenarios/0_scenarios.ipynb @@ -66,10 +66,11 @@ "\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 `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", + " - Users select sources with `--datasets name[:N|:all]`, or pass `dataset_overrides=ScenarioDatasetRequest(...)` in `set_params_from_args`. A custom `dataset_config` replaces the default; overrides apply on top.\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", + " - Scenarios declare allowed source selections and limit controls. Estimates return resolved `dataset_defaults` and `effective_datasets` separately from requested edits. Name-only selection preserves existing source caps and fetch policies.\n", "\n", "4. **Constructor**: Use `@apply_defaults` decorator and call `super().__init__()` with scenario metadata:\n", " - `name`: Descriptive name for your scenario\n", diff --git a/doc/code/scenarios/0_scenarios.py b/doc/code/scenarios/0_scenarios.py index 3a951368a4..a9407131da 100644 --- a/doc/code/scenarios/0_scenarios.py +++ b/doc/code/scenarios/0_scenarios.py @@ -68,10 +68,11 @@ # # 3. **Default Dataset**: Pass `default_dataset_config=` to `super().__init__()` to specify the datasets your scenario uses out of the box. # - 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 +# - Users select sources with `--datasets name[:N|:all]`, or pass `dataset_overrides=ScenarioDatasetRequest(...)` in `set_params_from_args`. A custom `dataset_config` replaces the default; overrides apply on top. # - 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 +# - Scenarios declare allowed source selections and limit controls. Estimates return resolved `dataset_defaults` and `effective_datasets` separately from requested edits. Name-only selection preserves existing source caps and fetch policies. # # 4. **Constructor**: Use `@apply_defaults` decorator and call `super().__init__()` with scenario metadata: # - `name`: Descriptive name for your scenario diff --git a/doc/code/scenarios/1_common_scenario_parameters.ipynb b/doc/code/scenarios/1_common_scenario_parameters.ipynb index e6081ccad4..5950f87f7c 100644 --- a/doc/code/scenarios/1_common_scenario_parameters.ipynb +++ b/doc/code/scenarios/1_common_scenario_parameters.ipynb @@ -15,7 +15,7 @@ "\n", "> **Two selection axes**: *Techniques* select attack techniques (*how* attacks run — e.g., prompt\n", "> sending, role play, TAP). *Datasets* select objectives (*what* is tested — e.g., harm categories,\n", - "> compliance topics). Use `--dataset-names` on the CLI to filter by content category.\n", + "> compliance topics). Use `--datasets` on the CLI to filter by content category.\n", "\n", "> **Running scenarios from the command line?** See the [Scanner documentation](../../scanner/0_scanner.md).\n", "\n", diff --git a/doc/code/scenarios/1_common_scenario_parameters.py b/doc/code/scenarios/1_common_scenario_parameters.py index 354df1410e..0ced660ed1 100644 --- a/doc/code/scenarios/1_common_scenario_parameters.py +++ b/doc/code/scenarios/1_common_scenario_parameters.py @@ -19,7 +19,7 @@ # # > **Two selection axes**: *Techniques* select attack techniques (*how* attacks run — e.g., prompt # > sending, role play, TAP). *Datasets* select objectives (*what* is tested — e.g., harm categories, -# > compliance topics). Use `--dataset-names` on the CLI to filter by content category. +# > compliance topics). Use `--datasets` on the CLI to filter by content category. # # > **Running scenarios from the command line?** See the [Scanner documentation](../../scanner/0_scanner.md). # diff --git a/doc/code/scenarios/3_adaptive_scenarios.ipynb b/doc/code/scenarios/3_adaptive_scenarios.ipynb index 2c9b2bc2f5..436b55b09a 100644 --- a/doc/code/scenarios/3_adaptive_scenarios.ipynb +++ b/doc/code/scenarios/3_adaptive_scenarios.ipynb @@ -804,7 +804,7 @@ "# Use specific datasets and limit size\n", "pyrit_scan --scenario TextAdaptive --target openai_chat \\\n", " --datasets airt_hate airt_violence \\\n", - " --max-dataset-size 10\n", + " --max-total 10\n", "```" ] } diff --git a/doc/code/scenarios/3_adaptive_scenarios.py b/doc/code/scenarios/3_adaptive_scenarios.py index 247b24b786..2614eb1d17 100644 --- a/doc/code/scenarios/3_adaptive_scenarios.py +++ b/doc/code/scenarios/3_adaptive_scenarios.py @@ -301,5 +301,5 @@ def _technique_label(result) -> str: # # Use specific datasets and limit size # pyrit_scan --scenario TextAdaptive --target openai_chat \ # --datasets airt_hate airt_violence \ -# --max-dataset-size 10 +# --max-total 10 # ``` diff --git a/doc/gui/0_gui.md b/doc/gui/0_gui.md index d23fe60a8f..7707351f71 100644 --- a/doc/gui/0_gui.md +++ b/doc/gui/0_gui.md @@ -490,7 +490,29 @@ Defaults are saved in this browser, separately for each signed-in account. They A missing or changed default is shown as unavailable. Select a new default or clear it; CoPyRIT does not silently substitute a different target. If browser storage is unavailable, a warning states that the choice applies only in the current session. -Scanner forms let you override either selection for one run. **Objective Target** is the target the scenario runs against. **Adversarial Target**, under **Parameters** above **Dataset override**, is the target used to generate attacks. The adversarial selector appears only for scenarios whose available techniques use the shared default; scenarios with no such use, including the explicit-target adversarial benchmark, do not show it or send an override. **Use server default** clears the per-run adversarial override. Changes to your defaults do not change existing chats or queued/running scans. Explicit adversarial targets in a scenario or technique still take priority. Scorer targets are unchanged. +Scanner forms let you override either selection for one run. **Objective Target** is the target the scenario runs against. **Adversarial Target**, under **Parameters** above **Datasets**, is the target used to generate attacks. The adversarial selector appears only for scenarios whose available techniques use the shared default; scenarios with no such use, including the explicit-target adversarial benchmark, do not show it or send an override. **Use server default** clears the per-run adversarial override. Changes to your defaults do not change existing chats or queued/running scans. Explicit adversarial targets in a scenario or technique still take priority. Scorer targets are unchanged. + +### Scenario datasets + +The **Datasets** table shows each selected name, stored seed origins (including mixed +origins), stored example count, source limit, and fetch policy. Unknown counts for +provider-only datasets are not zero. Provider availability is separate from stored origin. +**Reload dataset metadata** refreshes the page-scoped catalog without fetching seed data. +Runtime changes and launches also invalidate this catalog. + +Use **Add dataset** for allowed alternatives or named memory datasets, **Remove** to +change the selection, and **Use default sources** to restore parameter-aware defaults. +FigStep allows one of `figstep` or `figstep_pro`; ingredient and compound sources are fixed. +Unavailable limit/source controls are not submitted. Parameter and technique changes +refresh resolved defaults without discarding edits; invalid selections display an error. + +**Max per dataset**, source limits, and **Max total** accept positive integers, `all`, +or `default`/blank. Source caps apply before the total; `all` removes only that cap. +Name-only selection retains existing explicit source settings. Fetch is `never` or +`if_missing`: only a new run may prepare missing registered datasets. Preview and launch +send the same edit-only dataset object, with harm/category filters preserved. Resume uses +the saved request and the original seed selection; a different dataset selection is +rejected. Resume never prepares datasets. #### Core Adversarial Default Override diff --git a/doc/scanner/1_pyrit_scan.ipynb b/doc/scanner/1_pyrit_scan.ipynb index ee28f8b9e7..670f985d40 100644 --- a/doc/scanner/1_pyrit_scan.ipynb +++ b/doc/scanner/1_pyrit_scan.ipynb @@ -1194,9 +1194,17 @@ "on demand, so a full dataset preload is not required:\n", "\n", "```shell\n", - "pyrit_scan run garak.figstep --target openai_chat --dataset-names figstep_pro --max-dataset-size 1\n", + "pyrit_scan run garak.figstep --target openai_chat --datasets figstep_pro --max-total 1 --dataset-fetch if_missing\n", "```\n", "\n", + "Scanner and shell share `--datasets name[:N|:all]`, `--max-per-dataset`, `--max-total`,\n", + "and `--dataset-fetch never|if_missing`. Source limits apply before the combined total.\n", + "Omitted, empty, or `default` keeps the scenario setting; `all` removes only its own cap.\n", + "`list-scenarios` displays allowed sources and defaults. Ingredient/compound scenarios\n", + "restrict controls to preserve their shaping. `--dataset-filters` works with either form.\n", + "Deprecated `--dataset-names` and `--max-dataset-size` cannot be mixed with the new flags.\n", + "Discovery and estimates never fetch; only a new run prepares absent registered datasets.\n", + "\n", "Custom initialization scripts are loaded by the backend at startup: list them in the\n", "`initialization_scripts` section of the config the server loads (paths are relative to your\n", "working directory, but full paths avoid confusion). Once the server is running with that\n", diff --git a/doc/scanner/1_pyrit_scan.py b/doc/scanner/1_pyrit_scan.py index 0aa18722aa..1e5dce2700 100644 --- a/doc/scanner/1_pyrit_scan.py +++ b/doc/scanner/1_pyrit_scan.py @@ -125,9 +125,17 @@ # on demand, so a full dataset preload is not required: # # ```shell -# pyrit_scan run garak.figstep --target openai_chat --dataset-names figstep_pro --max-dataset-size 1 +# pyrit_scan run garak.figstep --target openai_chat --datasets figstep_pro --max-total 1 --dataset-fetch if_missing # ``` # +# Scanner and shell share `--datasets name[:N|:all]`, `--max-per-dataset`, `--max-total`, +# and `--dataset-fetch never|if_missing`. Source limits apply before the combined total. +# Omitted, empty, or `default` keeps the scenario setting; `all` removes only its own cap. +# `list-scenarios` displays allowed sources and defaults. Ingredient/compound scenarios +# restrict controls to preserve their shaping. `--dataset-filters` works with either form. +# Deprecated `--dataset-names` and `--max-dataset-size` cannot be mixed with the new flags. +# Discovery and estimates never fetch; only a new run prepares absent registered datasets. +# # Custom initialization scripts are loaded by the backend at startup: list them in the # `initialization_scripts` section of the config the server loads (paths are relative to your # working directory, but full paths avoid confusion). Once the server is running with that diff --git a/doc/scanner/adaptive.ipynb b/doc/scanner/adaptive.ipynb index a4bc5849a4..1f7c72fd64 100644 --- a/doc/scanner/adaptive.ipynb +++ b/doc/scanner/adaptive.ipynb @@ -26,8 +26,8 @@ "pyrit_scan run adaptive.text_adaptive \\\n", " --initializers target \\\n", " --target openai_chat \\\n", - " --dataset-names airt_hate \\\n", - " --max-dataset-size 2 \\\n", + " --datasets airt_hate \\\n", + " --max-total 2 \\\n", " --max-attempts-per-objective 2\n", "```" ] diff --git a/doc/scanner/adaptive.py b/doc/scanner/adaptive.py index b229e0e55d..0eda326f0b 100644 --- a/doc/scanner/adaptive.py +++ b/doc/scanner/adaptive.py @@ -25,8 +25,8 @@ # pyrit_scan run adaptive.text_adaptive \ # --initializers target \ # --target openai_chat \ -# --dataset-names airt_hate \ -# --max-dataset-size 2 \ +# --datasets airt_hate \ +# --max-total 2 \ # --max-attempts-per-objective 2 # ``` diff --git a/doc/scanner/airt.ipynb b/doc/scanner/airt.ipynb index 09cf12e75e..04866a9005 100644 --- a/doc/scanner/airt.ipynb +++ b/doc/scanner/airt.ipynb @@ -102,8 +102,8 @@ " --initializers target \\\n", " --target openai_chat \\\n", " --techniques role_play_movie_script \\\n", - " --dataset-names airt_hate \\\n", - " --max-dataset-size 1\n", + " --datasets airt_hate \\\n", + " --max-total 1\n", "```\n", "\n", "**Available techniques:** ALL, DEFAULT, SINGLE_TURN, MULTI_TURN, role_play_movie_script, many_shot, tap" @@ -255,8 +255,8 @@ "\n", "Each sub-harm escalates a simulated multi-turn conversation toward the objective, then layers the\n", "selected converter techniques on top (natural-language rewrites that preserve emotional framing;\n", - "obfuscation converters like base64/morse are excluded). Datasets are bound to the sub-harms, so\n", - "`--dataset-names` is ignored (`--max-dataset-size` still applies).\n", + "obfuscation converters like base64/morse are excluded). `--datasets` can select a subset of the\n", + "selected sub-harms' datasets; `--max-total` must cover each selected sub-harm.\n", "\n", "**Available techniques:** ALL, DEFAULT, tone, language, persuasion, deterministic, crescendo" ] @@ -413,7 +413,7 @@ " --initializers target \\\n", " --target openai_chat \\\n", " --techniques role_play_movie_script \\\n", - " --max-dataset-size 1\n", + " --max-total 1\n", "```\n", "\n", "**Available techniques:** Use `pyrit_scan run airt.cyber --list-scenario-parameters` to inspect\n", @@ -570,8 +570,8 @@ "pyrit_scan run airt.jailbreak \\\n", " --initializers target \\\n", " --target openai_chat \\\n", - " --dataset-names harmbench \\\n", - " --max-dataset-size 1\n", + " --datasets harmbench \\\n", + " --max-total 1\n", "```\n", "\n", "**Available technique selectors:** ALL, DEFAULT, and SINGLE_TURN currently select both\n", @@ -730,8 +730,8 @@ "pyrit_scan airt.multilingual \\\n", " --initializers target \\\n", " --target openai_chat \\\n", - " --dataset-names harmbench \\\n", - " --max-dataset-size 1\n", + " --datasets harmbench \\\n", + " --max-total 1\n", "```\n", "\n", "**Available techniques:** `prompt_sending` is the default. Every registry technique (`role_play_*`,\n", @@ -870,7 +870,7 @@ "plagiarism detection.\n", "\n", "```bash\n", - "pyrit_scan run airt.leakage --target openai_chat --techniques first_letter --max-dataset-size 1\n", + "pyrit_scan run airt.leakage --target openai_chat --techniques first_letter --max-total 1\n", "```\n", "\n", "**Available techniques:** ALL, SINGLE_TURN, MULTI_TURN, IP, SENSITIVE_DATA, FirstLetter, Image, RolePlay, Crescendo\n", @@ -1049,7 +1049,7 @@ " --initializers target \\\n", " --target openai_chat \\\n", " --techniques context_compliance \\\n", - " --max-dataset-size 1\n", + " --max-total 1\n", "```\n", "\n", "**Available techniques:** ALL, DEFAULT, SINGLE_TURN, MULTI_TURN, ContextCompliance, RolePlay,\n", diff --git a/doc/scanner/airt.py b/doc/scanner/airt.py index cb804f14e7..a7fba16ebc 100644 --- a/doc/scanner/airt.py +++ b/doc/scanner/airt.py @@ -43,8 +43,8 @@ # --initializers target \ # --target openai_chat \ # --techniques role_play_movie_script \ -# --dataset-names airt_hate \ -# --max-dataset-size 1 +# --datasets airt_hate \ +# --max-total 1 # ``` # # **Available techniques:** ALL, DEFAULT, SINGLE_TURN, MULTI_TURN, role_play_movie_script, many_shot, tap @@ -84,8 +84,8 @@ # # Each sub-harm escalates a simulated multi-turn conversation toward the objective, then layers the # selected converter techniques on top (natural-language rewrites that preserve emotional framing; -# obfuscation converters like base64/morse are excluded). Datasets are bound to the sub-harms, so -# `--dataset-names` is ignored (`--max-dataset-size` still applies). +# obfuscation converters like base64/morse are excluded). `--datasets` can select a subset of the +# selected sub-harms' datasets; `--max-total` must cover each selected sub-harm. # # **Available techniques:** ALL, DEFAULT, tone, language, persuasion, deterministic, crescendo @@ -124,7 +124,7 @@ # --initializers target \ # --target openai_chat \ # --techniques role_play_movie_script \ -# --max-dataset-size 1 +# --max-total 1 # ``` # # **Available techniques:** Use `pyrit_scan run airt.cyber --list-scenario-parameters` to inspect @@ -169,8 +169,8 @@ # pyrit_scan run airt.jailbreak \ # --initializers target \ # --target openai_chat \ -# --dataset-names harmbench \ -# --max-dataset-size 1 +# --datasets harmbench \ +# --max-total 1 # ``` # # **Available technique selectors:** ALL, DEFAULT, and SINGLE_TURN currently select both @@ -213,8 +213,8 @@ # pyrit_scan airt.multilingual \ # --initializers target \ # --target openai_chat \ -# --dataset-names harmbench \ -# --max-dataset-size 1 +# --datasets harmbench \ +# --max-total 1 # ``` # # **Available techniques:** `prompt_sending` is the default. Every registry technique (`role_play_*`, @@ -254,7 +254,7 @@ # plagiarism detection. # # ```bash -# pyrit_scan run airt.leakage --target openai_chat --techniques first_letter --max-dataset-size 1 +# pyrit_scan run airt.leakage --target openai_chat --techniques first_letter --max-total 1 # ``` # # **Available techniques:** ALL, SINGLE_TURN, MULTI_TURN, IP, SENSITIVE_DATA, FirstLetter, Image, RolePlay, Crescendo @@ -308,7 +308,7 @@ # --initializers target \ # --target openai_chat \ # --techniques context_compliance \ -# --max-dataset-size 1 +# --max-total 1 # ``` # # **Available techniques:** ALL, DEFAULT, SINGLE_TURN, MULTI_TURN, ContextCompliance, RolePlay, diff --git a/doc/scanner/benchmark.ipynb b/doc/scanner/benchmark.ipynb index 6abd678fe9..2ed6d84aeb 100644 --- a/doc/scanner/benchmark.ipynb +++ b/doc/scanner/benchmark.ipynb @@ -35,7 +35,7 @@ " --target openai_chat \\\n", " --adversarial-targets adversarial_chat \\\n", " --techniques role_play_video_game \\\n", - " --max-dataset-size 1\n", + " --max-total 1\n", "```\n", "\n", "Pass multiple `--adversarial-targets` values to compare across models in a single run.\n", diff --git a/doc/scanner/benchmark.py b/doc/scanner/benchmark.py index 8adddffbc1..92588c50db 100644 --- a/doc/scanner/benchmark.py +++ b/doc/scanner/benchmark.py @@ -34,7 +34,7 @@ # --target openai_chat \ # --adversarial-targets adversarial_chat \ # --techniques role_play_video_game \ -# --max-dataset-size 1 +# --max-total 1 # ``` # # Pass multiple `--adversarial-targets` values to compare across models in a single run. diff --git a/doc/scanner/foundry.ipynb b/doc/scanner/foundry.ipynb index 78a5dbeaab..dca0cb3342 100644 --- a/doc/scanner/foundry.ipynb +++ b/doc/scanner/foundry.ipynb @@ -73,7 +73,7 @@ "**CLI example:**\n", "\n", "```bash\n", - "pyrit_scan run foundry.red_team_agent --target openai_chat --techniques base64 --max-dataset-size 1\n", + "pyrit_scan run foundry.red_team_agent --target openai_chat --techniques base64 --max-total 1\n", "```\n", "\n", "**Available techniques by difficulty:**\n", diff --git a/doc/scanner/foundry.py b/doc/scanner/foundry.py index bbcc9b8bb6..e73747b8ca 100644 --- a/doc/scanner/foundry.py +++ b/doc/scanner/foundry.py @@ -40,7 +40,7 @@ # **CLI example:** # # ```bash -# pyrit_scan run foundry.red_team_agent --target openai_chat --techniques base64 --max-dataset-size 1 +# pyrit_scan run foundry.red_team_agent --target openai_chat --techniques base64 --max-total 1 # ``` # # **Available techniques by difficulty:** diff --git a/doc/scanner/garak.ipynb b/doc/scanner/garak.ipynb index 01f9dc85fb..e4be9c7077 100644 --- a/doc/scanner/garak.ipynb +++ b/doc/scanner/garak.ipynb @@ -36,7 +36,7 @@ "PromptInject tests whether a target follows injected instructions hidden in benign tasks.\n", "`--techniques` selects the injection wording, and `--goal-texts` selects the requested output.\n", "The default run uses all five techniques, the three Garak goal texts, and up to 60 attempts.\n", - "`--max-dataset-size` must be at least the number of selected goal texts.\n", + "`--max-total` must be at least the number of selected goal texts.\n", "Sampling selects at least one context for every goal, then fills the remaining budget\n", "without replacement. All techniques use the same sample. Injection wording is applied\n", "through request converters, before any caller-supplied converters.\n", @@ -48,7 +48,7 @@ "pyrit_scan run garak.prompt_inject --target openai_chat \\\n", " --techniques ignore_print \\\n", " --goal-texts \"Custom trigger text\" \\\n", - " --max-dataset-size 16\n", + " --max-total 16\n", "```\n", "\n", "**Techniques:** `ignore_print`, `ignore_say`, `ignore_print_upper`, `nevermind`, and\n", @@ -295,7 +295,7 @@ "**CLI example:**\n", "\n", "```bash\n", - "pyrit_scan run garak.encoding --target openai_chat --techniques base64 --max-dataset-size 1\n", + "pyrit_scan run garak.encoding --target openai_chat --techniques base64 --max-total 1\n", "```\n", "\n", "**Available techniques** (17 encodings): Base64, Base2048, Base16, Base32, ASCII85, Hex,\n", @@ -453,8 +453,8 @@ "**CLI examples:**\n", "\n", "```bash\n", - "pyrit_scan garak.figstep --target openai_chat --dataset-names figstep --max-dataset-size 1\n", - "pyrit_scan garak.figstep --target openai_chat --dataset-names figstep_pro --max-dataset-size 1\n", + "pyrit_scan garak.figstep --target openai_chat --datasets figstep --max-total 1\n", + "pyrit_scan garak.figstep --target openai_chat --datasets figstep_pro --max-total 1\n", "```\n", "\n", "> **Note:** The objective target must natively support multi-piece user messages and accept text\n", @@ -616,7 +616,7 @@ "**CLI example:**\n", "\n", "```bash\n", - "pyrit_scan run garak.web_injection --target openai_chat --techniques xss --max-dataset-size 1\n", + "pyrit_scan run garak.web_injection --target openai_chat --techniques xss\n", "```\n", "\n", "**Available techniques** (8 probes): MarkdownImageExfil, ColabAIDataLeakage,\n", @@ -626,7 +626,7 @@ "**Aggregate techniques:** `ALL` (all 8), `DEFAULT` (excludes the two combinatorial extended\n", "probes), `EXFIL` (the 6 markdown-exfil probes), and `XSS` (TaskXSS + MarkdownXSS).\n", "\n", - "**Dataset requirements:** `--dataset-names` replaces the default datasets, so it must include\n", + "**Dataset requirements:** `--datasets` replaces the default datasets, so it must include\n", "every dataset the selected techniques read. Otherwise the scenario stops before running and names\n", "each missing dataset. Datasets already loaded into memory never fill the gap.\n", "\n", @@ -641,7 +641,7 @@ "| StringAssemblyDataExfil | None (built-in seeds) |\n", "\n", "```bash\n", - "pyrit_scan run garak.web_injection --target openai_chat --techniques task_xss --dataset-names garak_xss_normal_instructions garak_web_html_js\n", + "pyrit_scan run garak.web_injection --target openai_chat --techniques task_xss --datasets garak_xss_normal_instructions garak_web_html_js\n", "```" ] }, @@ -982,7 +982,7 @@ "pyrit_scan run garak.api_key --target openai_chat\n", "\n", "# Run only GetKey with a smaller sample.\n", - "pyrit_scan run garak.api_key --target openai_chat --techniques get_key --max-dataset-size 2\n", + "pyrit_scan run garak.api_key --target openai_chat --techniques get_key --max-total 2\n", "```\n", "\n", "**Available techniques:** `GetKey` and `CompleteKey`. `DEFAULT` and `ALL` both select the two\n", @@ -1157,7 +1157,7 @@ "\n", "```bash\n", "pyrit_scan run garak.latent_injection --target openai_chat --techniques bare \\\n", - " --families whois --max-dataset-size 1\n", + " --families whois --max-total 1\n", "```\n", "\n", "**Available techniques** (14 separator styles): Bare, Newline, Blockquote, BlockquoteInline,\n", @@ -1349,7 +1349,7 @@ "**CLI example:**\n", "\n", "```bash\n", - "pyrit_scan run garak.doctor --target openai_chat --techniques policy_puppetry --max-dataset-size 1\n", + "pyrit_scan run garak.doctor --target openai_chat --techniques policy_puppetry --max-total 1\n", "```\n", "\n", "**Available techniques** (2 probes): `PolicyPuppetry` (wraps the objective in the Dr House\n", @@ -1841,13 +1841,13 @@ "**CLI example:**\n", "\n", "```bash\n", - "pyrit_scan garak.audio_achilles_heel --target azure_openai_realtime --max-dataset-size 2\n", + "pyrit_scan garak.audio_achilles_heel --target azure_openai_realtime --max-total 2\n", "```\n", "\n", "> **Note:** The objective target must accept `audio_path` input (i.e. be multimodal). The example\n", "> below uses the registered Azure OpenAI Realtime target; non-audio targets such as the default\n", "> `openai_chat` will error when the audio request is sent. The full dataset holds ~350\n", - "> clips, so a default run samples a small subset to finish quickly — raise `--max-dataset-size`\n", + "> clips, so a default run samples a small subset to finish quickly — raise `--max-total`\n", "> for broader coverage." ] }, @@ -2005,7 +2005,7 @@ "**CLI example:**\n", "\n", "```bash\n", - "pyrit_scan run garak.divergence --target openai_chat --techniques repeat --max-dataset-size 2\n", + "pyrit_scan run garak.divergence --target openai_chat --techniques repeat --max-total 2\n", "```\n", "\n", "**Available techniques:** `Repeat`, `DEFAULT`, and `ALL` all select the same probe.\n", diff --git a/doc/scanner/garak.py b/doc/scanner/garak.py index 766c8a9f83..1d3426f5dd 100644 --- a/doc/scanner/garak.py +++ b/doc/scanner/garak.py @@ -35,7 +35,7 @@ # PromptInject tests whether a target follows injected instructions hidden in benign tasks. # `--techniques` selects the injection wording, and `--goal-texts` selects the requested output. # The default run uses all five techniques, the three Garak goal texts, and up to 60 attempts. -# `--max-dataset-size` must be at least the number of selected goal texts. +# `--max-total` must be at least the number of selected goal texts. # Sampling selects at least one context for every goal, then fills the remaining budget # without replacement. All techniques use the same sample. Injection wording is applied # through request converters, before any caller-supplied converters. @@ -47,7 +47,7 @@ # pyrit_scan run garak.prompt_inject --target openai_chat \ # --techniques ignore_print \ # --goal-texts "Custom trigger text" \ -# --max-dataset-size 16 +# --max-total 16 # ``` # # **Techniques:** `ignore_print`, `ignore_say`, `ignore_print_upper`, `nevermind`, and @@ -128,7 +128,7 @@ # **CLI example:** # # ```bash -# pyrit_scan run garak.encoding --target openai_chat --techniques base64 --max-dataset-size 1 +# pyrit_scan run garak.encoding --target openai_chat --techniques base64 --max-total 1 # ``` # # **Available techniques** (17 encodings): Base64, Base2048, Base16, Base32, ASCII85, Hex, @@ -176,8 +176,8 @@ # **CLI examples:** # # ```bash -# pyrit_scan garak.figstep --target openai_chat --dataset-names figstep --max-dataset-size 1 -# pyrit_scan garak.figstep --target openai_chat --dataset-names figstep_pro --max-dataset-size 1 +# pyrit_scan garak.figstep --target openai_chat --datasets figstep --max-total 1 +# pyrit_scan garak.figstep --target openai_chat --datasets figstep_pro --max-total 1 # ``` # # > **Note:** The objective target must natively support multi-piece user messages and accept text @@ -217,7 +217,7 @@ # **CLI example:** # # ```bash -# pyrit_scan run garak.web_injection --target openai_chat --techniques xss --max-dataset-size 1 +# pyrit_scan run garak.web_injection --target openai_chat --techniques xss # ``` # # **Available techniques** (8 probes): MarkdownImageExfil, ColabAIDataLeakage, @@ -227,7 +227,7 @@ # **Aggregate techniques:** `ALL` (all 8), `DEFAULT` (excludes the two combinatorial extended # probes), `EXFIL` (the 6 markdown-exfil probes), and `XSS` (TaskXSS + MarkdownXSS). # -# **Dataset requirements:** `--dataset-names` replaces the default datasets, so it must include +# **Dataset requirements:** `--datasets` replaces the default datasets, so it must include # every dataset the selected techniques read. Otherwise the scenario stops before running and names # each missing dataset. Datasets already loaded into memory never fill the gap. # @@ -242,7 +242,7 @@ # | StringAssemblyDataExfil | None (built-in seeds) | # # ```bash -# pyrit_scan run garak.web_injection --target openai_chat --techniques task_xss --dataset-names garak_xss_normal_instructions garak_web_html_js +# pyrit_scan run garak.web_injection --target openai_chat --techniques task_xss --datasets garak_xss_normal_instructions garak_web_html_js # ``` # %% @@ -329,7 +329,7 @@ # pyrit_scan run garak.api_key --target openai_chat # # # Run only GetKey with a smaller sample. -# pyrit_scan run garak.api_key --target openai_chat --techniques get_key --max-dataset-size 2 +# pyrit_scan run garak.api_key --target openai_chat --techniques get_key --max-total 2 # ``` # # **Available techniques:** `GetKey` and `CompleteKey`. `DEFAULT` and `ALL` both select the two @@ -377,7 +377,7 @@ # # ```bash # pyrit_scan run garak.latent_injection --target openai_chat --techniques bare \ -# --families whois --max-dataset-size 1 +# --families whois --max-total 1 # ``` # # **Available techniques** (14 separator styles): Bare, Newline, Blockquote, BlockquoteInline, @@ -443,7 +443,7 @@ # **CLI example:** # # ```bash -# pyrit_scan run garak.doctor --target openai_chat --techniques policy_puppetry --max-dataset-size 1 +# pyrit_scan run garak.doctor --target openai_chat --techniques policy_puppetry --max-total 1 # ``` # # **Available techniques** (2 probes): `PolicyPuppetry` (wraps the objective in the Dr House @@ -571,13 +571,13 @@ # **CLI example:** # # ```bash -# pyrit_scan garak.audio_achilles_heel --target azure_openai_realtime --max-dataset-size 2 +# pyrit_scan garak.audio_achilles_heel --target azure_openai_realtime --max-total 2 # ``` # # > **Note:** The objective target must accept `audio_path` input (i.e. be multimodal). The example # > below uses the registered Azure OpenAI Realtime target; non-audio targets such as the default # > `openai_chat` will error when the audio request is sent. The full dataset holds ~350 -# > clips, so a default run samples a small subset to finish quickly — raise `--max-dataset-size` +# > clips, so a default run samples a small subset to finish quickly — raise `--max-total` # > for broader coverage. # %% @@ -616,7 +616,7 @@ # **CLI example:** # # ```bash -# pyrit_scan run garak.divergence --target openai_chat --techniques repeat --max-dataset-size 2 +# pyrit_scan run garak.divergence --target openai_chat --techniques repeat --max-total 2 # ``` # # **Available techniques:** `Repeat`, `DEFAULT`, and `ALL` all select the same probe. diff --git a/frontend/src/components/Scenarios/ScenarioDatasetTable.styles.ts b/frontend/src/components/Scenarios/ScenarioDatasetTable.styles.ts new file mode 100644 index 0000000000..e40aea247c --- /dev/null +++ b/frontend/src/components/Scenarios/ScenarioDatasetTable.styles.ts @@ -0,0 +1,13 @@ +import { makeStyles, tokens } from '@fluentui/react-components' + +import { mobileTouchTarget } from '@/styles/touchTargets' + +export const useScenarioDatasetTableStyles = makeStyles({ + root: { display: 'flex', flexDirection: 'column', gap: tokens.spacingVerticalM }, + controls: { display: 'flex', flexWrap: 'wrap', gap: tokens.spacingHorizontalM }, + scroll: { overflowX: 'auto' }, + table: { width: '100%', borderCollapse: 'collapse' }, + cell: { padding: tokens.spacingHorizontalS, textAlign: 'left' }, + input: { maxWidth: '12em' }, + button: { ...mobileTouchTarget }, +}) diff --git a/frontend/src/components/Scenarios/ScenarioDatasetTable.test.tsx b/frontend/src/components/Scenarios/ScenarioDatasetTable.test.tsx new file mode 100644 index 0000000000..5eea65b2dc --- /dev/null +++ b/frontend/src/components/Scenarios/ScenarioDatasetTable.test.tsx @@ -0,0 +1,109 @@ +import { useState } from 'react' + +import { render, screen, within } from '@testing-library/react' +import userEvent from '@testing-library/user-event' +import { FluentProvider, webLightTheme } from '@fluentui/react-components' + +import type { ScenarioDatasetOptions, ScenarioDatasetRequest, ScenarioDatasetSettings } from '@/types' + +import ScenarioDatasetTable from './ScenarioDatasetTable' + +const DEFAULTS: ScenarioDatasetSettings = { + sources: [{ name: 'stored', max_size: 2, fetch: 'never', configured_provider: false }], + max_per_dataset: 5, max_total: 'all', fetch: 'if_missing', filters: {}, +} +const OPTIONS: ScenarioDatasetOptions = { + mode: 'any', allowed_names: [], per_dataset_limit: true, total_limit: true, supports_overrides: true, filters: true, +} + +interface HarnessProps { + options?: ScenarioDatasetOptions + defaults?: ScenarioDatasetSettings + onChange: jest.Mock +} + +function Harness({ options = OPTIONS, defaults = DEFAULTS, onChange }: HarnessProps) { + const [request, setRequest] = useState({}) + return { setRequest(next); onChange(next) }} + /> +} + +describe('ScenarioDatasetTable', () => { + beforeEach(() => jest.clearAllMocks()) + + it('shows mixed origins, stored counts, and unknown provider-only counts without sending defaults', async () => { + const user = userEvent.setup() + const onChange = jest.fn() + render() + expect(screen.getByText('local, unknown')).toBeInTheDocument() + expect(screen.getByText('8')).toBeInTheDocument() + expect(onChange).not.toHaveBeenCalled() + await user.type(screen.getByRole('combobox', { name: 'Add dataset' }), 'remote') + await user.click(screen.getByRole('button', { name: 'Add source' })) + const row = screen.getByRole('rowheader', { name: 'remote' }).closest('tr') + expect(row).not.toBeNull() + if (row) { + expect(within(row).getByText('Provider available')).toBeInTheDocument() + expect(within(row).getAllByText('Unknown')).toHaveLength(2) + } + expect(onChange.mock.calls.at(-1)?.[0]).toEqual({ sources: [{ name: 'stored' }, { name: 'remote' }] }) + }) + + it('edits source caps and policies independently of global limits with keyboard-accessible controls', async () => { + const user = userEvent.setup() + const onChange = jest.fn() + render() + await user.type(screen.getByRole('textbox', { name: 'stored limit' }), 'all') + await user.selectOptions(screen.getByRole('combobox', { name: 'stored fetch' }), 'if_missing') + await user.type(screen.getByRole('textbox', { name: 'Max total' }), '12') + await user.selectOptions(screen.getByRole('combobox', { name: 'Default dataset fetch' }), 'never') + expect(onChange.mock.calls.at(-1)?.[0]).toEqual({ + sources: [{ name: 'stored', max_size: 'all', fetch: 'if_missing' }], max_total: 12, fetch: 'never', + }) + await user.clear(screen.getByRole('textbox', { name: 'Max total' })) + expect(onChange.mock.calls.at(-1)?.[0]).not.toHaveProperty('max_total') + screen.getByRole('textbox', { name: 'Max per dataset' }).focus() + await user.keyboard('3') + expect(onChange.mock.calls.at(-1)?.[0]).toHaveProperty('max_per_dataset', 3) + }) + + it('does not send fixed or unsupported inputs and supports one-source replacement', async () => { + const user = userEvent.setup() + const onChange = jest.fn() + const view = render() + expect(screen.queryByRole('button', { name: 'Add source' })).not.toBeInTheDocument() + expect(screen.queryByRole('textbox', { name: 'stored limit' })).not.toBeInTheDocument() + expect(screen.queryByRole('textbox', { name: 'Max total' })).not.toBeInTheDocument() + expect(onChange).not.toHaveBeenCalled() + view.rerender() + await user.click(screen.getByRole('combobox', { name: 'Add dataset' })) + await user.click(screen.getByRole('option', { name: 'remote' })) + await user.click(screen.getByRole('button', { name: 'Add source' })) + expect(onChange.mock.calls.at(-1)?.[0]).toEqual({ sources: [{ name: 'remote' }] }) + expect(screen.queryByRole('rowheader', { name: 'stored' })).not.toBeInTheDocument() + }) + + it('retains edits and displays invalid selections when parameter-dependent options change', async () => { + const user = userEvent.setup() + const onChange = jest.fn() + const view = render() + await user.type(screen.getByRole('textbox', { name: 'stored limit' }), '1') + view.rerender() + expect(screen.getByRole('alert')).toHaveTextContent('Invalid selection for current parameters: stored') + expect(screen.getByRole('textbox', { name: 'stored limit' })).toHaveValue('1') + await user.click(screen.getByRole('button', { name: 'Use default sources' })) + expect(screen.queryByRole('rowheader', { name: 'stored' })).not.toBeInTheDocument() + expect(onChange.mock.calls.at(-1)?.[0]).toEqual({}) + }) +}) diff --git a/frontend/src/components/Scenarios/ScenarioDatasetTable.tsx b/frontend/src/components/Scenarios/ScenarioDatasetTable.tsx new file mode 100644 index 0000000000..51078acc07 --- /dev/null +++ b/frontend/src/components/Scenarios/ScenarioDatasetTable.tsx @@ -0,0 +1,194 @@ +import { useState } from 'react' + +import { Button, Combobox, Field, Input, MessageBar, MessageBarBody, Option, Select, Text } from '@fluentui/react-components' + +import type { + DatasetInfo, ScenarioDatasetLimit, ScenarioDatasetOptions, ScenarioDatasetRequest, + ScenarioDatasetSettings, ScenarioDatasetSourceInput, +} from '@/types' + +import { useScenarioDatasetTableStyles } from './ScenarioDatasetTable.styles' + +interface ScenarioDatasetTableProps { + request: ScenarioDatasetRequest + defaults: ScenarioDatasetSettings + effective: ScenarioDatasetSettings + options: ScenarioDatasetOptions + datasets: DatasetInfo[] + disabled: boolean + loading: boolean + metadataError: string | null + onReload: () => void + onChange: (request: ScenarioDatasetRequest) => void + onValidationChange: (error: string | null) => void +} + +interface DatasetLimitFieldProps { + label: string + value: ScenarioDatasetLimit | null | undefined + effective: number | 'all' + disabled: boolean + onChange: (value: ScenarioDatasetLimit) => void + onError: (error: string | null) => void +} + +function DatasetLimitField({ label, value, effective, disabled, onChange, onError }: DatasetLimitFieldProps) { + const [invalidText, setInvalidText] = useState(null) + const styles = useScenarioDatasetTableStyles() + return ( + + { + const text = data.value.trim() + const parsed = Number(text) + if (text !== '' && text !== 'default' && text !== 'all' && (!Number.isSafeInteger(parsed) || parsed < 1)) { + setInvalidText(data.value) + onError(`${label} must be a positive integer, all, or default.`) + return + } + setInvalidText(null) + onError(null) + onChange(text === '' || text === 'default' ? 'default' : text === 'all' ? 'all' : parsed) + }} + /> + + ) +} + +/** Displays effective settings while emitting only user edits, never catalog defaults. */ +export default function ScenarioDatasetTable({ + request, defaults, effective, options, datasets, disabled, loading, metadataError, + onReload, onChange, onValidationChange, +}: ScenarioDatasetTableProps) { + const styles = useScenarioDatasetTableStyles() + const [newName, setNewName] = useState('') + const [errors, setErrors] = useState>({}) + const selected: ScenarioDatasetSourceInput[] = request.sources ?? effective.sources.map((source) => ({ name: source.name })) + const canSelect = options.supports_overrides && options.mode !== 'fixed' + const canEdit = options.supports_overrides && options.source_settings !== false + const candidates = datasets.filter((dataset) => options.mode === 'any' || options.allowed_names.includes(dataset.name)) + const setError = (key: string, error: string | null): void => { + const next = { ...errors } + if (error) next[key] = error + else delete next[key] + setErrors(next) + onValidationChange(Object.values(next)[0] ?? null) + } + const updateSource = (name: string, edits: Partial): void => { + onChange({ ...request, sources: selected.map((source) => source.name === name ? { ...source, ...edits } : source) }) + } + const updateGlobal = (key: 'max_per_dataset' | 'max_total', value: ScenarioDatasetLimit): void => { + const edits = { ...request } + if (value === 'default') delete edits[key] + else edits[key] = value + onChange(edits) + } + const invalidNames = selected.filter((source) => + options.mode === 'subset' && !options.allowed_names.includes(source.name), + ) + return ( +
+ Datasets + + {metadataError && {metadataError}} + {invalidNames.length > 0 && + Invalid selection for current parameters: {invalidNames.map((source) => source.name).join(', ')} + } + {Object.values(errors).map((error) => {error})} + {options.supports_overrides &&
+ {options.per_dataset_limit && setError('perDataset', error)} + onChange={(value) => updateGlobal('max_per_dataset', value)} />} + {options.total_limit && setError('total', error)} + onChange={(value) => updateGlobal('max_total', value)} />} + + + +
} +
+ + {['Name', 'Stored origins', 'Stored examples', 'Limit', 'Fetch policy', 'Selection'].map((label) => + )} + {selected.map((source) => { + const metadata = datasets.find((dataset) => dataset.name === source.name) + const settings = effective.sources.find((item) => item.name === source.name) + ?? defaults.sources.find((item) => item.name === source.name) + const hasProvider = metadata?.provider_available || settings?.configured_provider + return + + + + + + + + })} +
{label}
{source.name}{metadata?.origins?.join(', ') || 'Unknown'} +
{hasProvider ? 'Provider available' : metadata ? 'Memory only' : 'Availability unknown'}
{metadata?.logical_examples ?? 'Unknown'}{canEdit && options.per_dataset_limit + ? setError(source.name, error)} + onChange={(value) => updateSource(source.name, { max_size: value })} /> + : settings?.max_size ?? 'all'}{canEdit + ? : settings?.fetch ?? effective.fetch}{canSelect && }
+
+ {canSelect &&
+ + setNewName(event.target.value)} + onOptionSelect={(_, data) => setNewName(data.optionValue ?? '')}> + {candidates.map((dataset) => )} + + + + +
} +
+ ) +} diff --git a/frontend/src/components/Scenarios/ScenarioDetail.test.tsx b/frontend/src/components/Scenarios/ScenarioDetail.test.tsx index 477c78c8fb..d6f47508c1 100644 --- a/frontend/src/components/Scenarios/ScenarioDetail.test.tsx +++ b/frontend/src/components/Scenarios/ScenarioDetail.test.tsx @@ -6,6 +6,7 @@ import { MemoryRouter, Route, Routes } from 'react-router' import { scenariosApi, targetsApi } from '@/services/api' import type { RegisteredScenario, + ScenarioRunSizeEstimateRequest, ScenarioRunSizeEstimateResponse, TargetInstance, } from '@/types' @@ -13,6 +14,7 @@ import type { import ScenarioDetail from './ScenarioDetail' jest.mock('@/services/api', () => ({ + datasetsApi: { list: jest.fn().mockResolvedValue([]) }, scenariosApi: { estimateRun: jest.fn(), getScenario: jest.fn(), @@ -335,7 +337,7 @@ describe('ScenarioDetail', () => { expect(within(screen.getByRole('region', { name: 'Objective Target' })).queryByRole( 'combobox', { name: 'Adversarial Target' }, )).not.toBeInTheDocument() - expect(fallback.compareDocumentPosition(within(parameters).getByRole('textbox', { name: 'Dataset override' }))) + expect(fallback.compareDocumentPosition(within(parameters).getByRole('table', { name: 'Selected datasets' }))) .toBe(Node.DOCUMENT_POSITION_FOLLOWING) expect(fallback).toHaveAccessibleDescription('The registered target this scenario uses to generate attacks.') expect(fallback).toHaveValue('adversary') @@ -942,10 +944,8 @@ describe('ScenarioDetail', () => { expect(await screen.findByRole('heading', { name: 'Parameters' })).toBeInTheDocument() expect(screen.queryByText('Advanced options')).not.toBeInTheDocument() - expect(screen.getByTestId('max-dataset-size-input')).toHaveValue(8) - expect(screen.getByText( - 'The scenario default is 8. Edit it to override the default.', - )).toBeInTheDocument() + expect(screen.getByRole('textbox', { name: 'Max total' })).toHaveValue('') + expect(screen.getByText('Default/effective: 8. Blank or default retains the scenario setting.')).toBeInTheDocument() const estimate = screen.getByTestId('run-estimate') expect(within(estimate).getByText('Dataset size')).toBeInTheDocument() expect(within(estimate).getByText('8')).toBeInTheDocument() @@ -975,7 +975,8 @@ describe('ScenarioDetail', () => { renderDetail('/scanner/foundry.red_team_agent') - expect(await screen.findByTestId('max-dataset-size-input')).toHaveValue(5) + expect(await screen.findByRole('textbox', { name: 'Max total' })).toHaveValue('') + expect(screen.getByText('Default/effective: 5. Blank or default retains the scenario setting.')).toBeInTheDocument() expect(await screen.findByText('Up to 10')).toBeInTheDocument() expect(mockEstimateRun.mock.calls.at(-1)?.[1]).not.toHaveProperty('max_dataset_size') }) @@ -989,7 +990,7 @@ describe('ScenarioDetail', () => { }, })) renderDetail('/scanner/foundry.red_team_agent') - expect(await screen.findByRole('spinbutton', { name: 'Max dataset size' })).toHaveValue(null) + expect(await screen.findByRole('textbox', { name: 'Max total' })).toHaveValue('') await waitFor(() => expect(mockEstimateRun).toHaveBeenCalled()) expect(mockEstimateRun.mock.calls.at(-1)?.[1]).not.toHaveProperty('max_dataset_size') }) @@ -1003,9 +1004,8 @@ describe('ScenarioDetail', () => { }, })) renderDetail('/scanner/foundry.red_team_agent') - const input = await screen.findByRole('spinbutton', { name: 'Max dataset size' }) - expect(input).toBeDisabled() - expect(input).toHaveValue(null) + await screen.findByRole('table', { name: 'Selected datasets' }) + expect(screen.queryByRole('textbox', { name: 'Max total' })).not.toBeInTheDocument() expect(screen.queryByRole('button', { name: 'Use all data' })).not.toBeInTheDocument() }) @@ -1019,8 +1019,8 @@ describe('ScenarioDetail', () => { }, })) renderDetail('/scanner/foundry.red_team_agent') - const input = await screen.findByRole('spinbutton', { name: 'Max dataset size' }) - expect(input).toHaveValue(5) + const input = await screen.findByRole('textbox', { name: 'Max total' }) + await user.type(input, '2') await user.clear(input) expect(screen.queryByRole('button', { name: 'Use all data' })).not.toBeInTheDocument() await waitFor(() => expect(mockEstimateRun).toHaveBeenCalled()) @@ -1033,13 +1033,12 @@ describe('ScenarioDetail', () => { it('sends an explicit all total for both estimation and launch', async () => { const user = userEvent.setup() renderDetail('/scanner/foundry.red_team_agent') - const unlimited = await screen.findByRole('checkbox', { name: 'No total dataset limit' }) - await user.click(unlimited) - expect(screen.getByRole('spinbutton', { name: 'Max dataset size' })).toBeDisabled() - await waitFor(() => expect(mockEstimateRun.mock.calls.at(-1)?.[1]).toHaveProperty('max_dataset_size', 'all')) + const input = await screen.findByRole('textbox', { name: 'Max total' }) + await user.type(input, 'all') + await waitFor(() => expect(mockEstimateRun.mock.calls.at(-1)?.[1]).toHaveProperty('datasets.max_total', 'all')) await confirmRunPreview(user) await waitFor(() => expect(mockStartRun).toHaveBeenCalled()) - expect(mockStartRun.mock.calls[0][0]).toHaveProperty('max_dataset_size', 'all') + expect(mockStartRun.mock.calls[0][0]).toHaveProperty('datasets.max_total', 'all') }) it('restores the scenario default when no-total-limit is unchecked', async () => { @@ -1051,11 +1050,11 @@ describe('ScenarioDetail', () => { }, })) renderDetail('/scanner/foundry.red_team_agent') - const unlimited = await screen.findByRole('checkbox', { name: 'No total dataset limit' }) - await user.click(unlimited) - await waitFor(() => expect(mockEstimateRun.mock.calls.at(-1)?.[1]).toHaveProperty('max_dataset_size', 'all')) - await user.click(unlimited) - expect(screen.getByRole('spinbutton', { name: 'Max dataset size' })).toHaveValue(5) + const input = await screen.findByRole('textbox', { name: 'Max total' }) + await user.type(input, 'all') + await waitFor(() => expect(mockEstimateRun.mock.calls.at(-1)?.[1]).toHaveProperty('datasets.max_total', 'all')) + await user.clear(input) + expect(input).toHaveValue('') await waitFor(() => expect(mockEstimateRun.mock.calls.at(-1)?.[1]).not.toHaveProperty('max_dataset_size')) await confirmRunPreview(user) await waitFor(() => expect(mockStartRun).toHaveBeenCalled()) @@ -1067,17 +1066,22 @@ describe('ScenarioDetail', () => { renderDetail('/scanner/foundry.red_team_agent') await screen.findByTestId('scenario-target-select') - await user.type(screen.getByTestId('dataset-override-input'), 'ds_a, ds_b') - await user.type(screen.getByTestId('max-dataset-size-input'), '25') + const addDataset = screen.getByRole('combobox', { name: 'Add dataset' }) + await user.type(addDataset, 'ds_a') + await user.click(screen.getByRole('button', { name: 'Add source' })) + await user.type(addDataset, 'ds_b') + await user.click(screen.getByRole('button', { name: 'Add source' })) + await user.click(screen.getByRole('button', { name: 'Remove harmbench' })) + await user.type(screen.getByRole('textbox', { name: 'Max total' }), '25') await user.type(screen.getByTestId('harm-categories-filter-input'), 'cyber, violence') await user.type(screen.getByTestId('data-types-filter-input'), 'text, image_path') await confirmRunPreview(user) await waitFor(() => expect(mockStartRun).toHaveBeenCalled()) const request = mockStartRun.mock.calls[0][0] - expect(request.dataset_names).toEqual(['ds_a', 'ds_b']) - expect(request.max_dataset_size).toBe(25) - expect(request.dataset_filters).toEqual({ + expect(request.datasets.sources).toEqual([{ name: 'ds_a' }, { name: 'ds_b' }]) + expect(request.datasets.max_total).toBe(25) + expect(request.datasets.filters).toEqual({ harm_categories: ['cyber', 'violence'], data_types: ['text', 'image_path'], }) @@ -1086,11 +1090,13 @@ describe('ScenarioDetail', () => { expect.objectContaining({ target_name: 'target-a', techniques: ['default_technique'], - dataset_names: ['ds_a', 'ds_b'], - max_dataset_size: 25, - dataset_filters: { + datasets: { + sources: [{ name: 'ds_a' }, { name: 'ds_b' }], + max_total: 25, + filters: { harm_categories: ['cyber', 'violence'], data_types: ['text', 'image_path'], + }, }, include_baseline: true, }), @@ -1099,17 +1105,56 @@ describe('ScenarioDetail', () => { expect(mockEstimateRun.mock.calls.at(-1)?.[1]).not.toHaveProperty('labels') }) + it('refreshes parameter-aware defaults after invalid source edits without silently dropping them', async () => { + const user = userEvent.setup() + mockGetScenario.mockResolvedValueOnce(makeScenario({ + supported_parameters: [ + { name: 'scope', type_name: 'str', required: false, default: 'all', choices: null, is_list: false }, + ], + })) + mockEstimateRun.mockImplementation(async (_name: string, request: ScenarioRunSizeEstimateRequest) => { + const narrowed = request.scenario_params?.scope === 'one' + if (narrowed && request.datasets?.sources) { + throw { isAxiosError: true, response: { status: 400, data: { detail: 'Dataset is outside the selected scope.' } } } + } + const name = narrowed ? 'remote' : 'harmbench' + const settings = { + sources: [{ name, max_size: 5, fetch: 'never', configured_provider: false }], + max_per_dataset: 5, max_total: 5, fetch: 'never', filters: {}, + } + return { + ...makeEstimate(10), dataset_defaults: settings, effective_datasets: settings, + dataset_options: { + mode: 'subset', allowed_names: [name], per_dataset_limit: true, total_limit: true, + supports_overrides: true, filters: true, + }, + } + }) + renderDetail('/scanner/foundry.red_team_agent') + await screen.findByRole('textbox', { name: 'harmbench limit' }) + await user.type(screen.getByRole('textbox', { name: 'harmbench limit' }), '1') + await user.clear(screen.getByTestId('scenario-param-scope')) + await user.type(screen.getByTestId('scenario-param-scope'), 'one') + await screen.findByText('Invalid selection for current parameters: harmbench') + expect(screen.getByRole('textbox', { name: 'harmbench limit' })).toHaveValue('1') + expect(mockEstimateRun.mock.calls.some((call) => + call[1].scenario_params?.scope === 'one' && call[1].datasets === undefined, + )).toBe(true) + await user.click(screen.getByRole('button', { name: 'Use default sources' })) + await screen.findByRole('rowheader', { name: 'remote' }) + }) + it('rejects a non-positive-integer max dataset size', async () => { const user = userEvent.setup() renderDetail('/scanner/foundry.red_team_agent') await screen.findByTestId('scenario-target-select') - await user.type(screen.getByTestId('max-dataset-size-input'), '0') + await user.type(screen.getByRole('textbox', { name: 'Max total' }), '0') await user.click(screen.getByTestId('launch-scenario-btn')) - expect(await screen.findByRole('alert')).toHaveTextContent( - 'Max dataset size must be a positive integer.', - ) + expect(await screen.findAllByRole('alert')).toEqual(expect.arrayContaining([ + expect.objectContaining({ textContent: expect.stringContaining('Max total must be a positive integer, all, or default.') }), + ])) expect(mockStartRun).not.toHaveBeenCalled() }) diff --git a/frontend/src/components/Scenarios/ScenarioDetail.tsx b/frontend/src/components/Scenarios/ScenarioDetail.tsx index f4378c4f3c..096636a64e 100644 --- a/frontend/src/components/Scenarios/ScenarioDetail.tsx +++ b/frontend/src/components/Scenarios/ScenarioDetail.tsx @@ -32,6 +32,7 @@ import { Link, useNavigate, useParams } from 'react-router' import MarkdownContent from '@/components/Markdown/MarkdownContent' import TargetSelect from '@/components/Config/TargetSelect' import { useRuntime } from '@/hooks/useRuntime' +import { useScenarioDatasets } from '@/hooks/useScenarioDatasets' import ParameterField from '@/components/Parameters/ParameterField' import SingleStepSpinButton from '@/components/Parameters/SingleStepSpinButton' import { @@ -51,6 +52,10 @@ import type { ScenarioRunSizeEstimateRequest, ScenarioRunEstimateState, ScenarioTechniqueSummary, + ScenarioDatasetRequest, + ScenarioDatasetSettings, + ScenarioDatasetOptions, + ScenarioRunSizeEstimateResponse, TargetInstance, } from '@/types' import { routerPathParamValue, scenarioRunRoutePath } from '@/utils/routeParams' @@ -61,6 +66,7 @@ import { ScenarioRunEstimateDetails } from './ScenarioRunEstimate' import { normalizeScenarioMarkdown } from './scenarioMarkdown' import { mapScenarioRunEstimate } from './scenarioRunEstimateAdapter' import { techniqueSetName } from './scenarioTechniqueSets' +import ScenarioDatasetTable from './ScenarioDatasetTable' /** * Common/opaque parameters every scenario declares via @@ -74,6 +80,7 @@ const COMMON_SCENARIO_PARAMETER_NAMES = new Set([ 'scenario_techniques', 'technique_converters', 'dataset_config', + 'dataset_overrides', 'memory_labels', 'max_concurrency', 'max_retries', @@ -95,11 +102,6 @@ function targetOptionLabel(target: TargetInstance): string { : target.target_registry_name } -function defaultMaxDatasetSize(scenario: RegisteredScenario): string { - const limit = scenario.default_run_size.dataset_limit - return limit.state === 'value' ? String(limit.value) : '' -} - /** Resolves a Fluent `SpinButton` change event to a numeric value, preferring the parsed `value` over the raw `displayValue`. */ function resolveSpinButtonValue(data: { value?: number | null; displayValue?: string }, previous: number): number { if (typeof data.value === 'number') { @@ -234,8 +236,8 @@ interface BuildEstimateRequestInput { techniques: string[] dynamicParameters: Parameter[] scenarioParamValues: Record - datasetOverride: string - maxDatasetSize: string + datasetEdits: ScenarioDatasetRequest + datasetError: string | null harmCategoriesFilter: string dataTypesFilter: string includeBaseline: boolean @@ -291,8 +293,8 @@ function buildEstimateRequest({ techniques, dynamicParameters, scenarioParamValues, - datasetOverride, - maxDatasetSize, + datasetEdits, + datasetError, harmCategoriesFilter, dataTypesFilter, includeBaseline, @@ -310,18 +312,7 @@ function buildEstimateRequest({ scenarioParams = result.parameters } - let maxDatasetSizeValue: number | 'all' | undefined - const trimmedMaxDatasetSize = maxDatasetSize.trim() - if (trimmedMaxDatasetSize === 'all') { - maxDatasetSizeValue = 'all' - } else if (trimmedMaxDatasetSize.length > 0) { - const parsed = Number(trimmedMaxDatasetSize) - if (!Number.isInteger(parsed) || parsed < 1) { - return { ok: false, error: 'Max dataset size must be a positive integer.' } - } - maxDatasetSizeValue = parsed - } - const datasetNames = parseDatasetNames(datasetOverride) + if (datasetError) return { ok: false, error: datasetError } const request: ScenarioRunSizeEstimateRequest = { techniques, include_baseline: includeBaseline, @@ -332,19 +323,15 @@ function buildEstimateRequest({ if (scenario.uses_default_adversarial_target && adversarialTargetName) { request.adversarial_target_name = adversarialTargetName } - if (datasetNames.length > 0) { - request.dataset_names = datasetNames - } - if (maxDatasetSizeValue !== undefined) { - request.max_dataset_size = maxDatasetSizeValue - } const harmCategories = parseDatasetNames(harmCategoriesFilter) const dataTypes = parseDatasetNames(dataTypesFilter) if (harmCategories.length > 0 || dataTypes.length > 0) { - request.dataset_filters = { + request.datasets = { ...datasetEdits, filters: { ...(harmCategories.length > 0 ? { harm_categories: harmCategories } : {}), ...(dataTypes.length > 0 ? { data_types: dataTypes } : {}), - } + } } + } else if (Object.keys(datasetEdits).length > 0) { + request.datasets = datasetEdits } if (scenarioParams) { request.scenario_params = scenarioParams @@ -394,15 +381,7 @@ function buildRunRequest(input: BuildRunRequestInput): BuildRunRequestResult { if (estimateRequest.adversarial_target_name !== undefined) { request.adversarial_target_name = estimateRequest.adversarial_target_name } - if (estimateRequest.dataset_names !== undefined) { - request.dataset_names = estimateRequest.dataset_names - } - if (estimateRequest.max_dataset_size !== undefined) { - request.max_dataset_size = estimateRequest.max_dataset_size - } - if (estimateRequest.dataset_filters !== undefined) { - request.dataset_filters = estimateRequest.dataset_filters - } + if (estimateRequest.datasets !== undefined) request.datasets = estimateRequest.datasets if (estimateRequest.scenario_params !== undefined) { request.scenario_params = estimateRequest.scenario_params } @@ -612,12 +591,21 @@ function ScenarioLaunchForm({ const [baselineChecked, setBaselineChecked] = useState( () => !isBaselineForbidden && scenario.include_baseline_by_default, ) - const [datasetOverride, setDatasetOverride] = useState('') - const configuredDefaultMaxDatasetSize = useMemo( - () => defaultMaxDatasetSize(scenario), - [scenario], - ) - const [maxDatasetSize, setMaxDatasetSize] = useState(configuredDefaultMaxDatasetSize) + const [datasetEdits, setDatasetEdits] = useState({}) + const [datasetError, setDatasetError] = useState(null) + const datasetMetadata = useScenarioDatasets() + const [datasetProjection, setDatasetProjection] = useState(null) + const datasetDefaults: ScenarioDatasetSettings = datasetProjection?.dataset_defaults ?? scenario.dataset_defaults ?? { + sources: scenario.default_datasets.map((name) => ({ name, max_size: 'all', fetch: 'if_missing', configured_provider: false })), + max_per_dataset: 'all', max_total: scenario.default_run_size.dataset_limit.value ?? 'all', + fetch: 'if_missing', filters: {}, + } + const effectiveDatasets = datasetProjection?.effective_datasets ?? datasetDefaults + const datasetOptions: ScenarioDatasetOptions = datasetProjection?.dataset_options ?? scenario.dataset_options ?? { + mode: 'any', allowed_names: [], per_dataset_limit: true, + total_limit: scenario.default_run_size.dataset_limit.state !== 'not_applicable', + supports_overrides: true, filters: true, + } const [harmCategoriesFilter, setHarmCategoriesFilter] = useState('') const [dataTypesFilter, setDataTypesFilter] = useState('') const [maxConcurrency, setMaxConcurrency] = useState(DEFAULT_MAX_CONCURRENCY) @@ -670,15 +658,8 @@ function ScenarioLaunchForm({ [isBaselineForbidden, techniqueOptions], ) const techniques = selectedTechniques - const maxDatasetSizeOverride = maxDatasetSize.trim() - && maxDatasetSize !== configuredDefaultMaxDatasetSize - ? maxDatasetSize - : '' - const datasetSizeLabel = scenario.default_run_size.dataset_limit.state === 'not_applicable' - ? 'Not applicable' - : maxDatasetSize === 'all' - ? 'No total limit' - : maxDatasetSize.trim() || configuredDefaultMaxDatasetSize || 'Scenario default' + const datasetSizeLabel = !datasetOptions.total_limit ? 'Not applicable' + : effectiveDatasets.max_total === 'all' ? 'No total limit' : String(effectiveDatasets.max_total) const estimateResult = useMemo( () => buildEstimateRequest({ scenario, @@ -687,8 +668,8 @@ function ScenarioLaunchForm({ techniques, dynamicParameters, scenarioParamValues, - datasetOverride, - maxDatasetSize: maxDatasetSizeOverride, + datasetEdits, + datasetError, harmCategoriesFilter, dataTypesFilter, includeBaseline: isBaselineForbidden ? false : baselineChecked, @@ -696,12 +677,12 @@ function ScenarioLaunchForm({ [ adversarialTargetName, baselineChecked, - datasetOverride, + datasetEdits, + datasetError, dataTypesFilter, dynamicParameters, harmCategoriesFilter, isBaselineForbidden, - maxDatasetSizeOverride, scenario, scenarioParamValues, targetName, @@ -716,8 +697,8 @@ function ScenarioLaunchForm({ techniques, dynamicParameters, scenarioParamValues, - datasetOverride, - maxDatasetSize: maxDatasetSizeOverride, + datasetEdits, + datasetError, harmCategoriesFilter, dataTypesFilter, maxConcurrency, @@ -728,14 +709,14 @@ function ScenarioLaunchForm({ [ adversarialTargetName, baselineChecked, - datasetOverride, + datasetEdits, + datasetError, dataTypesFilter, dynamicParameters, harmCategoriesFilter, isBaselineForbidden, labels, maxConcurrency, - maxDatasetSizeOverride, maxRetries, scenario, scenarioParamValues, @@ -774,6 +755,7 @@ function ScenarioLaunchForm({ return } const result = mapScenarioRunEstimate(response, 'request') + setDatasetProjection(response) setEstimateRequestState({ status: 'resolved', requestKey: estimateRequestKey, @@ -790,11 +772,26 @@ function ScenarioLaunchForm({ ) { return } + const error = toApiError(err) setEstimateRequestState({ status: 'error', requestKey: estimateRequestKey, - error: toApiError(err).detail, + error: error.detail, }) + if (estimateRequest.datasets && (error.status === 400 || error.status === 422)) { + const { datasets, ...defaultRequest } = estimateRequest + void datasets + // Keep invalid edits visible while refreshing parameter-aware allowed sources. + void scenariosApi.estimateRun(scenario.scenario_name, defaultRequest, controller.signal) + .then((response) => { + if (controller.signal.aborted || requestSequence !== estimateSequenceRef.current) return + setDatasetProjection((previous) => ({ + ...response, + effective_datasets: previous?.effective_datasets ?? response.effective_datasets, + })) + }) + .catch(() => { /* The original validation error remains authoritative. */ }) + } }) }, ESTIMATE_DEBOUNCE_MS) @@ -908,6 +905,7 @@ function ScenarioLaunchForm({ try { const summary = await scenariosApi.startRun(requestResult.request) + datasetMetadata.reload() setPreviewOpen(false) navigate(scenarioRunRoutePath(summary.scenario_result_id), { state: { scenarioName: scenario.scenario_name }, @@ -935,8 +933,8 @@ function ScenarioLaunchForm({ const displayedEstimateNotes = estimateNotes(estimateState) const displayedEstimate = estimateFromState(estimateState) const selectedTechniqueCount = selectedTechniques.length + (baselineChecked ? 1 : 0) - const previewDatasets = parseDatasetNames(datasetOverride) - const effectiveDatasets = previewDatasets.length > 0 ? previewDatasets : scenario.default_datasets + const previewDatasets = (datasetEdits.sources ?? effectiveDatasets.sources).map((source) => source.name) + const previewDatasetNames = previewDatasets.length > 0 ? previewDatasets : scenario.default_datasets const previewHarmCategories = parseDatasetNames(harmCategoriesFilter) const previewDataTypes = parseDatasetNames(dataTypesFilter) @@ -1130,45 +1128,12 @@ function ScenarioLaunchForm({ disabled={submitting} /> )} - - setDatasetOverride(data.value)} - placeholder={scenario.default_datasets.join(', ') || undefined} - data-testid="dataset-override-input" - /> - - - setMaxDatasetSize(data.value)} - data-testid="max-dataset-size-input" - /> - - setMaxDatasetSize(data.checked === true ? 'all' : configuredDefaultMaxDatasetSize)} + - Per-dataset limits still apply. An empty input uses the scenario default total. setHarmCategoriesFilter(data.value)} data-testid="harm-categories-filter-input" @@ -1189,7 +1154,7 @@ function ScenarioLaunchForm({ setDataTypesFilter(data.value)} data-testid="data-types-filter-input" @@ -1341,8 +1306,8 @@ function ScenarioLaunchForm({
- {effectiveDatasets.length > 0 - ? effectiveDatasets.join(', ') + {previewDatasetNames.length > 0 + ? previewDatasetNames.join(', ') : 'No datasets declared'} diff --git a/frontend/src/hooks/useScenarioDatasets.test.tsx b/frontend/src/hooks/useScenarioDatasets.test.tsx new file mode 100644 index 0000000000..c522c620f8 --- /dev/null +++ b/frontend/src/hooks/useScenarioDatasets.test.tsx @@ -0,0 +1,58 @@ +import { act, renderHook, waitFor } from '@testing-library/react' + +import { datasetsApi } from '@/services/api' +import type { DatasetInfo } from '@/types' + +import { useScenarioDatasets } from './useScenarioDatasets' + +jest.mock('@/services/api', () => ({ datasetsApi: { list: jest.fn() } })) +let generation = 'one' +jest.mock('./useRuntime', () => ({ useRuntime: () => ({ ready: true, generation }) })) +const list = datasetsApi.list as jest.Mock +const NAMED: DatasetInfo = { + name: 'named', selection_key: 'dataset:named:named', loaded: true, provider_available: false, logical_examples: 1, origins: ['local'], +} + +describe('useScenarioDatasets', () => { + beforeEach(() => { jest.clearAllMocks(); generation = 'one' }) + + it('caches metadata for the page, excludes unnamed populations, and explicitly reloads', async () => { + list.mockResolvedValue([NAMED, { ...NAMED, name: '(unnamed)', selection_key: 'dataset:unnamed' }]) + const view = renderHook(() => useScenarioDatasets()) + await waitFor(() => expect(view.result.current.items).toEqual([NAMED])) + view.rerender() + expect(list).toHaveBeenCalledTimes(1) + act(() => { view.result.current.reload() }) + await waitFor(() => expect(list).toHaveBeenCalledTimes(2)) + generation = 'two' + view.rerender() + await waitFor(() => expect(list).toHaveBeenCalledTimes(3)) + }) + + it('ignores old-generation responses and surfaces reload failures without generating datasets', async () => { + let completeOld: (value: DatasetInfo[]) => void = () => {} + list.mockReturnValueOnce(new Promise((resolve) => { completeOld = resolve })) + list.mockResolvedValueOnce([{ ...NAMED, name: 'new' }]) + const view = renderHook(() => useScenarioDatasets()) + generation = 'two' + view.rerender() + await waitFor(() => expect(view.result.current.items[0]?.name).toBe('new')) + await act(async () => { completeOld([NAMED]) }) + expect(view.result.current.items[0]?.name).toBe('new') + list.mockRejectedValueOnce(new Error('offline')) + act(() => { view.result.current.reload() }) + await waitFor(() => expect(view.result.current.error).not.toBeNull()) + }) + + it('hides loaded metadata immediately when the runtime changes', async () => { + list.mockResolvedValueOnce([NAMED]) + const view = renderHook(() => useScenarioDatasets()) + await waitFor(() => expect(view.result.current.items).toEqual([NAMED])) + list.mockReturnValueOnce(new Promise(() => {})) + generation = 'two' + view.rerender() + expect(view.result.current.items).toEqual([]) + expect(view.result.current.loading).toBe(true) + view.unmount() + }) +}) diff --git a/frontend/src/hooks/useScenarioDatasets.ts b/frontend/src/hooks/useScenarioDatasets.ts new file mode 100644 index 0000000000..08e1b7503c --- /dev/null +++ b/frontend/src/hooks/useScenarioDatasets.ts @@ -0,0 +1,50 @@ +import { useCallback, useEffect, useState } from 'react' + +import { datasetsApi } from '@/services/api' +import { toApiError } from '@/services/errors' +import type { DatasetInfo } from '@/types' + +import { useRuntime } from './useRuntime' + +interface DatasetSnapshot { + key: string + items: DatasetInfo[] + error: string | null +} + +/** Page-scoped metadata cache; no dataset loading or generation occurs here. */ +export function useScenarioDatasets() { + const runtime = useRuntime() + const [snapshot, setSnapshot] = useState(null) + const [reloadCount, setReloadCount] = useState(0) + const reload = useCallback(() => { setReloadCount((count) => count + 1) }, []) + const key = JSON.stringify([runtime.generation, runtime.ready, reloadCount]) + + useEffect(() => { + if (!runtime.ready) return + const controller = new AbortController() + const load = async (): Promise => { + try { + const data = await datasetsApi.list(controller.signal) + if (controller.signal.aborted) return + setSnapshot({ + key, + items: data.filter((item) => item.selection_key !== 'dataset:unnamed'), + error: null, + }) + } catch (err: unknown) { + if (!controller.signal.aborted) setSnapshot({ key, items: [], error: toApiError(err).detail }) + } + } + void load() + return () => { controller.abort() } + }, [key, runtime.ready]) + + const current = runtime.ready && snapshot?.key === key ? snapshot : null + return { + items: current?.items ?? [], + error: current?.error ?? null, + loading: runtime.ready && current === null, + reload, + } +} diff --git a/frontend/src/services/api.ts b/frontend/src/services/api.ts index 4e4cde69df..f910c41d28 100644 --- a/frontend/src/services/api.ts +++ b/frontend/src/services/api.ts @@ -5,6 +5,7 @@ import { toApiError } from './errors' import { compatibility, COMPATIBILITY_HEADER } from './compatibility' import { getGraphScopes } from '../auth/msalConfig' import type { + DatasetInfo, TargetInstance, TargetListResponse, TargetTypeListResponse, @@ -462,6 +463,15 @@ export const labelsApi = { }, } +export const datasetsApi = { + list: async (signal?: AbortSignal): Promise => { + const response = await apiClient.get<{ items: DatasetInfo[] }>('/datasets', { + params: { loaded_only: false }, signal, + }) + return response.data.items + }, +} + export const scenariosApi = { /** * Lists one page of the scenario catalog. Callers that need the full diff --git a/frontend/src/types/index.ts b/frontend/src/types/index.ts index 938631dc8f..186af95382 100644 --- a/frontend/src/types/index.ts +++ b/frontend/src/types/index.ts @@ -770,6 +770,58 @@ export interface ChangeMainConversationResponse { // --- Scenarios --- +export type ScenarioDatasetLimit = number | 'all' | 'default' +export type DatasetFetchPolicy = 'never' | 'if_missing' + +export interface ScenarioDatasetSourceInput { + name: string + max_size?: ScenarioDatasetLimit | null + fetch?: DatasetFetchPolicy | 'default' | null +} + +export interface ScenarioDatasetRequest { + sources?: ScenarioDatasetSourceInput[] | null + max_per_dataset?: ScenarioDatasetLimit | null + max_total?: ScenarioDatasetLimit | null + fetch?: DatasetFetchPolicy | 'default' | null + filters?: Record | null +} + +export interface ScenarioDatasetOptions { + mode: 'any' | 'subset' | 'fixed' + allowed_names: string[] + max_selected?: number | null + per_dataset_limit: boolean + total_limit: boolean + supports_overrides: boolean + filters: boolean + source_settings?: boolean +} + +export interface ScenarioDatasetSourceSettings { + name: string + max_size: number | 'all' + fetch: DatasetFetchPolicy + configured_provider: boolean +} + +export interface ScenarioDatasetSettings { + sources: ScenarioDatasetSourceSettings[] + max_per_dataset: number | 'all' + max_total: number | 'all' + fetch: DatasetFetchPolicy + filters: Record +} + +export interface DatasetInfo { + name: string + selection_key: string + loaded: boolean + provider_available: boolean + logical_examples: number | null + origins: string[] +} + export interface RegisteredScenario { scenario_name: string scenario_type: string @@ -783,6 +835,8 @@ export interface RegisteredScenario { all_techniques: string[] technique_summaries: ScenarioTechniqueSummary[] default_datasets: string[] + dataset_options?: ScenarioDatasetOptions + dataset_defaults?: ScenarioDatasetSettings baseline_policy: 'enabled' | 'disabled' | 'forbidden' include_baseline_by_default: boolean uses_default_adversarial_target: boolean @@ -810,6 +864,7 @@ export interface RunScenarioRequest { dataset_names?: string[] | null max_dataset_size?: number | 'all' | 'default' | '' | null dataset_filters?: Record | null + datasets?: ScenarioDatasetRequest max_concurrency?: number max_retries?: number include_baseline?: boolean | null @@ -855,6 +910,9 @@ export interface ScenarioRunSizeEstimateResponse { status?: 'exact' | 'approximate' | 'conditional' | 'unavailable' dataset_size: ScenarioDatasetSizeEstimate dataset_limit: DatasetLimitInput + dataset_options?: ScenarioDatasetOptions + dataset_defaults?: ScenarioDatasetSettings + effective_datasets?: ScenarioDatasetSettings estimated_attack_count: number | null minimum_attack_count?: number | null maximum_attack_count?: number | null @@ -871,6 +929,7 @@ export interface ScenarioRunSizeEstimateRequest { dataset_names?: string[] | null max_dataset_size?: number | 'all' | 'default' | '' | null dataset_filters?: Record | null + datasets?: ScenarioDatasetRequest include_baseline?: boolean | null scenario_params?: Record | null } diff --git a/pyrit/backend/models/datasets.py b/pyrit/backend/models/datasets.py index f3566c340a..97bf5ef246 100644 --- a/pyrit/backend/models/datasets.py +++ b/pyrit/backend/models/datasets.py @@ -30,6 +30,7 @@ class DatasetInfo(BaseModel): ) loaded: bool = Field(..., description="Whether this dataset currently has seeds in memory") provider_available: bool = Field(..., description="Whether a registered provider can load this dataset") + origins: list[str] = Field(default_factory=list, description="Distinct origins of stored seeds, not providers") logical_examples: int | None = Field( None, description="Number of logical examples represented by stored seed groups" ) diff --git a/pyrit/backend/services/dataset_service.py b/pyrit/backend/services/dataset_service.py index 84920ffb85..c41c6fe6cc 100644 --- a/pyrit/backend/services/dataset_service.py +++ b/pyrit/backend/services/dataset_service.py @@ -83,6 +83,7 @@ async def list_datasets_async(self, *, loaded_only: bool = False) -> DatasetList selection_key=f"dataset:named:{dataset_name}", loaded=summary is not None, provider_available=dataset_name in provider_names, + origins=list(summary.origins) if summary else [], logical_examples=summary.logical_examples if summary else None, seed_pieces=summary.seed_pieces if summary else None, objectives=summary.objectives if summary else None, @@ -100,6 +101,7 @@ async def list_datasets_async(self, *, loaded_only: bool = False) -> DatasetList selection_key="dataset:unnamed", loaded=True, provider_available=False, + origins=list(unnamed_summary.origins), logical_examples=unnamed_summary.logical_examples, seed_pieces=unnamed_summary.seed_pieces, objectives=unnamed_summary.objectives, @@ -290,6 +292,7 @@ def _merge_unnamed_summaries(summaries: Sequence[SeedDatasetSummary]) -> SeedDat modalities=tuple(sorted({modality for summary in unnamed for modality in summary.modalities})), harm_categories=tuple(sorted({category for summary in unnamed for category in summary.harm_categories})), has_unlabeled_harm_categories=any(summary.has_unlabeled_harm_categories for summary in unnamed), + origins=tuple(sorted({origin for summary in unnamed for origin in summary.origins})), ) diff --git a/pyrit/backend/services/scenario_configuration_resolver.py b/pyrit/backend/services/scenario_configuration_resolver.py index b0c6e19256..5933763248 100644 --- a/pyrit/backend/services/scenario_configuration_resolver.py +++ b/pyrit/backend/services/scenario_configuration_resolver.py @@ -7,13 +7,12 @@ from typing import TYPE_CHECKING, Any -from pyrit.models.dataset_limit import DatasetLimit, normalize_dataset_limit from pyrit.registry import ConverterRegistry, ScenarioRegistry, TargetRegistry -from pyrit.scenario import DatasetSource from pyrit.scenario.core.scenario_target_defaults import validate_default_adversarial_target if TYPE_CHECKING: from pyrit.converter import Converter + from pyrit.models.scenario_dataset import ScenarioDatasetRequest from pyrit.prompt_target import PromptTarget from pyrit.scenario import Scenario @@ -91,9 +90,7 @@ def resolve_configuration( scenario_class: type[Scenario], objective_target: Any | None = None, techniques: list[str] | None = None, - dataset_names: list[str] | None = None, - max_dataset_size: DatasetLimit = "default", - dataset_filters: dict[str, list[str]] | None = None, + datasets: ScenarioDatasetRequest | None = None, include_baseline: bool | None = None, max_concurrency: int | None = None, max_retries: int | None = None, @@ -102,7 +99,7 @@ def resolve_configuration( """ Resolve shared launch and estimate fields into scenario parameters. - Omitted, null, empty, and "default" limits retain the default. "all" removes only the total limit. + Dataset edits pass through unchanged; the scenario applies and validates them. Returns: dict[str, Any]: Values accepted by ``Scenario.set_params_from_args``. @@ -122,11 +119,9 @@ def resolve_configuration( if memory_labels: resolved["memory_labels"] = memory_labels - filters = dataset_filters or {} - total_limit = normalize_dataset_limit(max_dataset_size) - has_total_override = total_limit != "default" - needs_introspection = bool(techniques) or bool(dataset_names) or has_total_override or bool(filters) - if not needs_introspection: + if datasets is not None: + resolved["dataset_overrides"] = datasets + if not techniques: return resolved try: @@ -147,20 +142,6 @@ def resolve_configuration( if technique_converters: resolved["technique_converters"] = technique_converters - if dataset_names or has_total_override or filters: - default_config = introspection_instance._default_dataset_config - config = default_config.with_overrides(filters=filters) - if dataset_names: - existing_sources = {source.name: source for source in config.sources} - config = config.with_overrides( - sources=[ - existing_sources[name] if name in existing_sources else DatasetSource(name=name) - for name in dataset_names - ] - ) - config = config.with_overrides(max_total=total_limit) - resolved["dataset_config"] = config - return resolved @classmethod diff --git a/pyrit/backend/services/scenario_run_service.py b/pyrit/backend/services/scenario_run_service.py index a187b94e8e..f800cf25ba 100644 --- a/pyrit/backend/services/scenario_run_service.py +++ b/pyrit/backend/services/scenario_run_service.py @@ -103,6 +103,7 @@ "target_name", "adversarial_target_name", "techniques", + "datasets", "dataset_names", "max_dataset_size", "dataset_filters", @@ -110,6 +111,8 @@ "max_retries", "include_baseline", ) +# Fields added after the first saved launch records; older records omit them. +_OPTIONAL_LAUNCH_FIELDS = frozenset({"adversarial_target_name", "datasets", "max_dataset_size"}) _SAFE_SCENARIO_PARAMETER_NAMES = frozenset( { @@ -324,13 +327,15 @@ def _restore_launch_request(self, *, stored: ScenarioResult) -> RunScenarioReque raw_request = {"adversarial_target_name": None, **raw_request} if ( not isinstance(raw_request, dict) - or any(name not in raw_request for name in _LAUNCH_REQUEST_FIELDS if name != "max_dataset_size") + or any(name not in raw_request for name in _LAUNCH_REQUEST_FIELDS if name not in _OPTIONAL_LAUNCH_FIELDS) or raw_request["include_baseline"] is None ): raise ScenarioRunConflictError("The saved launch configuration is incomplete; resume was not started.") try: - request = RunScenarioRequest.model_validate( - {name: raw_request[name] for name in _LAUNCH_REQUEST_FIELDS if name in raw_request}, strict=True + # JSON validation accepts saved enum values as strings while keeping strict type checks. + request = RunScenarioRequest.model_validate_json( + json.dumps({name: raw_request[name] for name in _LAUNCH_REQUEST_FIELDS if name in raw_request}), + strict=True, ) except ValidationError as exc: raise ScenarioRunConflictError( @@ -340,7 +345,7 @@ def _restore_launch_request(self, *, stored: ScenarioResult) -> RunScenarioReque raise ScenarioRunConflictError("The saved scenario or target registration name is empty.") identifier = stored.scenario_identifier if (request.techniques is None and identifier.techniques is None) or ( - request.dataset_names is None and identifier.datasets is None + request.dataset_request is None and identifier.datasets is None ): raise ScenarioRunConflictError("The saved scenario identity is missing techniques or datasets.") custom_params = { @@ -348,6 +353,8 @@ def _restore_launch_request(self, *, stored: ScenarioResult) -> RunScenarioReque for name, value in identifier.params.items() if name not in {"version", "techniques", "datasets"} } + # Dataset names are not copied from the identity: the saved request reproduces them, and the + # stored-identity check on initialization rejects any difference. return request.model_copy( update={ "scenario_result_id": str(stored.id), @@ -355,7 +362,6 @@ def _restore_launch_request(self, *, stored: ScenarioResult) -> RunScenarioReque "initializer_args": None, "scenario_params": custom_params, "techniques": request.techniques if request.techniques is not None else identifier.techniques, - "dataset_names": request.dataset_names if request.dataset_names is not None else identifier.datasets, "labels": dict(stored.labels), }, deep=True, @@ -643,9 +649,7 @@ async def _prepare_run_async(self, *, request: RunScenarioRequest) -> _PreparedR scenario_class=scenario_class, objective_target=objective_target, techniques=request.techniques, - dataset_names=request.dataset_names, - max_dataset_size=request.max_dataset_size, - dataset_filters=request.dataset_filters, + datasets=request.dataset_request, include_baseline=request.include_baseline, max_concurrency=request.max_concurrency, max_retries=request.max_retries, @@ -1231,7 +1235,7 @@ async def _initialize_scenario_async(self, *, request: RunScenarioRequest, init_ The fully initialized Scenario instance ready for run_async. """ scenario_registry = ScenarioRegistry.get_registry_singleton() - launch_request = {name: getattr(request, name) for name in _LAUNCH_REQUEST_FIELDS} + launch_request = request.model_dump(mode="json", include=set(_LAUNCH_REQUEST_FIELDS)) if launch_request["include_baseline"] is None: scenario_class = scenario_registry.get_class(request.scenario_name) baseline_parameter = next( diff --git a/pyrit/backend/services/scenario_service.py b/pyrit/backend/services/scenario_service.py index a1b1a620b7..d6e95d3e2e 100644 --- a/pyrit/backend/services/scenario_service.py +++ b/pyrit/backend/services/scenario_service.py @@ -68,6 +68,8 @@ def _metadata_to_registered_scenario( all_techniques=list(metadata.all_techniques), technique_summaries=list(metadata.technique_summaries), default_datasets=list(metadata.default_datasets), + dataset_options=metadata.dataset_options, + dataset_defaults=metadata.dataset_defaults, supported_parameters=[ parameter.for_external_catalog() for parameter in metadata.supported_parameters @@ -498,9 +500,7 @@ async def _estimate_configured_run_size_async( scenario_class=scenario_class, objective_target=objective_target, techniques=request.techniques, - dataset_names=request.dataset_names, - max_dataset_size=request.max_dataset_size, - dataset_filters=request.dataset_filters, + datasets=request.dataset_request, include_baseline=request.include_baseline, ) return await self._registry.create_and_estimate_async( diff --git a/pyrit/cli/_cli_args.py b/pyrit/cli/_cli_args.py index 3d92084889..3cd2c7b0ea 100644 --- a/pyrit/cli/_cli_args.py +++ b/pyrit/cli/_cli_args.py @@ -173,6 +173,84 @@ def parse_dataset_limit(value: str) -> int | Literal["all", "default"]: return normalize_dataset_limit(value) +def parse_dataset_source(value: str) -> Any: + """ + Parse a named source with an optional source cap. + + Returns: + DatasetSourceInput: Validated named source. + + Raises: + ValueError: If the token is malformed. + """ + from pyrit.models.scenario_dataset import DatasetSourceInput + + parts = value.split(":") + if len(parts) > 2 or not parts[0].strip(): + raise ValueError("Expected dataset name[:positive-limit|all|default].") + return DatasetSourceInput(name=parts[0], max_size=parse_dataset_limit(parts[1]) if len(parts) == 2 else "default") + + +def parse_dataset_fetch(value: str) -> str: + """ + Validate the global preparation policy without importing dataset providers. + + Returns: + str: Validated preparation policy. + + Raises: + ValueError: If the policy is unsupported. + """ + if value not in {"never", "if_missing"}: + raise ValueError("Dataset fetch must be 'never' or 'if_missing'.") + return value + + +def build_dataset_request_fields(args: dict[str, Any]) -> dict[str, Any]: + """ + Build the same edit-only dataset request for scanner and shell. + + Returns: + dict[str, Any]: Either the new request or legacy dataset fields. + + Raises: + ValueError: If new and deprecated flags are mixed. + """ + from pyrit.models.scenario_dataset import ScenarioDatasetRequest + + new_keys = ("datasets", "max_per_dataset", "max_total", "dataset_fetch") + new_form = any(args.get(key) is not None for key in new_keys) + if new_form and any(args.get(key) is not None for key in ("dataset_names", "max_dataset_size")): + raise ValueError( + "Do not mix --datasets/--max-per-dataset/--max-total/--dataset-fetch with deprecated dataset flags." + ) + filters = collapse_dataset_filters(args["dataset_filters"]) if args.get("dataset_filters") else None + if new_form: + return { + "datasets": ScenarioDatasetRequest( + sources=args.get("datasets"), + max_per_dataset=args.get("max_per_dataset") or "default", + max_total=args.get("max_total") or "default", + fetch=args.get("dataset_fetch"), + filters=filters, + ) + } + fields: dict[str, Any] = {} + for key in ("dataset_names", "max_dataset_size"): + if args.get(key) is not None: + from pyrit.common.deprecation import print_deprecation_message + + print_deprecation_message( + old_item="--" + key.replace("_", "-"), + new_item="--datasets" if key == "dataset_names" else "--max-total", + removed_in="1.4.0", + ) + fields[key] = args[key] + if filters: + fields["dataset_filters"] = filters + return fields + + # --------------------------------------------------------------------------- # Argparse adapter # --------------------------------------------------------------------------- @@ -426,9 +504,11 @@ def _coerce_filter_values(value: str) -> list[str]: "memory_labels": 'Additional labels as JSON string (e.g., \'{"experiment": "test1"}\')', "database": "Database type to use for memory storage", "log_level": "Logging level", - "dataset_names": "List of dataset names to use instead of scenario defaults (e.g., harmbench advbench). " + "dataset_names": "Deprecated; use --datasets. List of dataset names instead of scenario defaults " + "(e.g., harmbench advbench). " "Retains source settings for names already selected by the scenario", - "max_dataset_size": "Total attack-group limit (positive integer, 'default', or 'all' for no total limit). " + "max_dataset_size": "Deprecated; use --max-total. Total attack-group limit " + "(positive integer, 'default', or 'all'). " "Omit to retain the scenario default. Per-dataset limits still apply", "dataset_filters": "Dataset seed filters as KEY=VALUE tokens " "(e.g., harm_categories=cyber data_types=text). Keys filter seeds before sizing. " @@ -662,6 +742,10 @@ class _ArgSpec: _DATASET_NAMES_ARG, _MAX_DATASET_SIZE_ARG, _DATASET_FILTERS_ARG, + _ArgSpec(flags=["--datasets"], result_key="datasets", multi_value=True, parser=parse_dataset_source), + _ArgSpec(flags=["--max-per-dataset"], result_key="max_per_dataset", parser=parse_dataset_limit), + _ArgSpec(flags=["--max-total"], result_key="max_total", parser=parse_dataset_limit), + _ArgSpec(flags=["--dataset-fetch"], result_key="dataset_fetch", parser=parse_dataset_fetch), _TARGET_ARG, ] diff --git a/pyrit/cli/_output.py b/pyrit/cli/_output.py index 3cdc8c3291..8b7f015680 100644 --- a/pyrit/cli/_output.py +++ b/pyrit/cli/_output.py @@ -123,6 +123,17 @@ def print_scenario_list(*, items: list[RegisteredScenario]) -> None: if sc.default_datasets: print(f" Default Datasets ({len(sc.default_datasets)}):") print(_wrap(text=", ".join(sc.default_datasets), indent=" ")) + options = sc.dataset_options + print(f" Dataset selection: {options.mode.value}; overrides: {options.supports_overrides}") + if options.allowed_names: + print(_wrap(text="Options: " + ", ".join(options.allowed_names), indent=" ")) + defaults = sc.dataset_defaults + print( + f" Dataset defaults: per-source={defaults.max_per_dataset}, " + f"total={defaults.max_total}, fetch={defaults.fetch.value}" + ) + for source in defaults.sources: + print(f" {source.name}: limit={source.max_size}, fetch={source.fetch.value}") if sc.supported_parameters: print(" Supported Parameters:") for p in sc.supported_parameters: diff --git a/pyrit/cli/pyrit_scan.py b/pyrit/cli/pyrit_scan.py index f912bb5bb5..3f2cb8a7cf 100644 --- a/pyrit/cli/pyrit_scan.py +++ b/pyrit/cli/pyrit_scan.py @@ -27,11 +27,12 @@ ARG_HELP, _parse_initializer_arg, add_results_arguments, + build_dataset_request_fields, build_parameters_from_api, - collapse_dataset_filters, non_negative_int, parse_dataset_filter, parse_dataset_limit, + parse_dataset_source, positive_int, validate_log_level_argparse, ) @@ -313,6 +314,19 @@ def _add_run_arguments(*, parser: ArgumentParser, scenario_params: list[Paramete parser.add_argument("--memory-labels", type=str, help=ARG_HELP["memory_labels"]) parser.add_argument("--dataset-names", type=str, nargs="+", help=ARG_HELP["dataset_names"]) parser.add_argument("--max-dataset-size", type=parse_dataset_limit, help=ARG_HELP["max_dataset_size"]) + parser.add_argument( + "--datasets", + type=parse_dataset_source, + nargs="+", + help="Ordered dataset names with optional :N or :all source caps", + ) + parser.add_argument( + "--max-per-dataset", type=parse_dataset_limit, help="Default source cap; use all for no source cap" + ) + parser.add_argument("--max-total", type=parse_dataset_limit, help="Combined attack-group cap") + parser.add_argument( + "--dataset-fetch", choices=["never", "if_missing"], help="Preparation policy for missing datasets" + ) parser.add_argument( "--dataset-filters", type=parse_dataset_filter, @@ -1004,12 +1018,7 @@ def _build_run_request(*, parsed_args: Namespace, scenario_name: str) -> RunScen kwargs["max_concurrency"] = parsed_args.max_concurrency if parsed_args.max_retries is not None: kwargs["max_retries"] = parsed_args.max_retries - if parsed_args.dataset_names: - kwargs["dataset_names"] = parsed_args.dataset_names - if parsed_args.max_dataset_size is not None: - kwargs["max_dataset_size"] = parsed_args.max_dataset_size - if parsed_args.dataset_filters: - kwargs["dataset_filters"] = collapse_dataset_filters(parsed_args.dataset_filters) + kwargs.update(build_dataset_request_fields(vars(parsed_args))) if parsed_args.memory_labels: kwargs["labels"] = parse_memory_labels(json_string=parsed_args.memory_labels) @@ -1062,7 +1071,11 @@ async def _run_scenario_async( from pyrit.models import ScenarioRunState scenario_name = parsed_args.scenario_name - request = _build_run_request(parsed_args=parsed_args, scenario_name=scenario_name) + try: + request = _build_run_request(parsed_args=parsed_args, scenario_name=scenario_name) + except ValueError as exc: + print(f"Error: {exc}") + return 1 print(f"\nRunning scenario: {scenario_name}") sys.stdout.flush() diff --git a/pyrit/cli/pyrit_shell.py b/pyrit/cli/pyrit_shell.py index 8110dada98..746f36fb1f 100644 --- a/pyrit/cli/pyrit_shell.py +++ b/pyrit/cli/pyrit_shell.py @@ -434,8 +434,12 @@ def do_run(self, line: str) -> None: --max-concurrency Maximum concurrent operations --max-retries Maximum retry attempts --memory-labels JSON string of labels - --dataset-names ... Override default dataset names - --max-dataset-size Maximum items per dataset + --datasets ... Select sources and optional caps + --max-per-dataset Default per-source cap + --max-total Combined attack-group cap + --dataset-fetch never or if_missing + --dataset-names ... Deprecated source selection + --max-dataset-size Deprecated total cap -- Scenario-declared parameters (see list-scenarios) Notes: @@ -452,8 +456,8 @@ def do_run(self, line: str) -> None: return from pyrit.cli._cli_args import ( + build_dataset_request_fields, build_parameters_from_api, - collapse_dataset_filters, extract_scenario_args, parse_run_arguments, ) @@ -517,12 +521,11 @@ def do_run(self, line: str) -> None: request_kwargs["max_concurrency"] = args["max_concurrency"] if args.get("max_retries") is not None: request_kwargs["max_retries"] = args["max_retries"] - if args.get("dataset_names"): - request_kwargs["dataset_names"] = args["dataset_names"] - if args.get("max_dataset_size") is not None: - request_kwargs["max_dataset_size"] = args["max_dataset_size"] - if args.get("dataset_filters"): - request_kwargs["dataset_filters"] = collapse_dataset_filters(args["dataset_filters"]) + try: + request_kwargs.update(build_dataset_request_fields(args)) + except ValueError as exc: + print(f"Error: {exc}") + return if args.get("memory_labels"): request_kwargs["labels"] = args["memory_labels"] @@ -530,7 +533,11 @@ def do_run(self, line: str) -> None: if scenario_params: request_kwargs["scenario_params"] = scenario_params - request = RunScenarioRequest(**request_kwargs) + try: + request = RunScenarioRequest(**request_kwargs) + except ValueError as exc: + print(f"Error: {exc}") + return # Start run print(f"\nRunning scenario: {scenario_name}") diff --git a/pyrit/memory/memory_interface.py b/pyrit/memory/memory_interface.py index ca0997e73d..037e2f84c6 100644 --- a/pyrit/memory/memory_interface.py +++ b/pyrit/memory/memory_interface.py @@ -4295,6 +4295,7 @@ def _execute_get_seed_dataset_summaries(self) -> Sequence[SeedDatasetSummary]: SeedEntry.dataset_name.label("dataset_name"), SeedEntry.data_type, SeedEntry.harm_categories, + SeedEntry.origin, ) .where(named_condition) .distinct() @@ -4307,6 +4308,7 @@ def _execute_get_seed_dataset_summaries(self) -> Sequence[SeedDatasetSummary]: named_aggregate.c.objectives, named_metadata.c.data_type, named_metadata.c.harm_categories, + named_metadata.c.origin, ).join( named_metadata, named_aggregate.c.dataset_name == named_metadata.c.dataset_name, @@ -4323,7 +4325,10 @@ def _execute_get_seed_dataset_summaries(self) -> Sequence[SeedDatasetSummary]: .subquery() ) unnamed_metadata = ( - select(SeedEntry.data_type, SeedEntry.harm_categories).where(unnamed_condition).distinct().subquery() + select(SeedEntry.data_type, SeedEntry.harm_categories, SeedEntry.origin) + .where(unnamed_condition) + .distinct() + .subquery() ) unnamed_statement = select( unnamed_aggregate.c.dataset_name, @@ -4332,6 +4337,7 @@ def _execute_get_seed_dataset_summaries(self) -> Sequence[SeedDatasetSummary]: unnamed_aggregate.c.objectives, unnamed_metadata.c.data_type, unnamed_metadata.c.harm_categories, + unnamed_metadata.c.origin, ).select_from(unnamed_aggregate.join(unnamed_metadata, literal(True))) combined_statement = named_statement.union_all(unnamed_statement) @@ -4351,9 +4357,12 @@ def _execute_get_seed_dataset_summaries(self) -> Sequence[SeedDatasetSummary]: "modalities": set(), "harm_categories": set(), "has_unlabeled_harm_categories": False, + "origins": set(), } dataset_order.append(dataset_name) summary = summaries_by_dataset[dataset_name] + if row.origin: + summary["origins"].add(row.origin) if row.data_type: summary["modalities"].add(row.data_type) categories = row.harm_categories or [] @@ -4371,6 +4380,7 @@ def _execute_get_seed_dataset_summaries(self) -> Sequence[SeedDatasetSummary]: modalities=tuple(sorted(summaries_by_dataset[dataset_name]["modalities"])), harm_categories=tuple(sorted(summaries_by_dataset[dataset_name]["harm_categories"])), has_unlabeled_harm_categories=summaries_by_dataset[dataset_name]["has_unlabeled_harm_categories"], + origins=tuple(sorted(summaries_by_dataset[dataset_name]["origins"])), ) for dataset_name in dataset_order ] diff --git a/pyrit/models/__init__.py b/pyrit/models/__init__.py index 3cfc87e1aa..3bec545932 100644 --- a/pyrit/models/__init__.py +++ b/pyrit/models/__init__.py @@ -134,6 +134,15 @@ from pyrit.models.results.scenario_result import ScenarioResult, ScenarioRunState from pyrit.models.results.strategy_result import StrategyResult, StrategyResultT from pyrit.models.retry_event import RetryEvent + from pyrit.models.scenario_dataset import ( + DatasetFetchPolicy, + DatasetSelectionMode, + DatasetSourceInput, + ScenarioDatasetOptions, + ScenarioDatasetRequest, + ScenarioDatasetSettings, + ScenarioDatasetSourceSettings, + ) from pyrit.models.scenario_dataset_size_estimate import ( AllAvailableDatasetSize, BoundedDatasetSize, @@ -264,6 +273,13 @@ "ConversationObservationPayload": "pyrit.models.score", "AllAvailableDatasetSize": "pyrit.models.scenario_dataset_size_estimate", "BoundedDatasetSize": "pyrit.models.scenario_dataset_size_estimate", + "DatasetFetchPolicy": "pyrit.models.scenario_dataset", + "DatasetSelectionMode": "pyrit.models.scenario_dataset", + "DatasetSourceInput": "pyrit.models.scenario_dataset", + "ScenarioDatasetOptions": "pyrit.models.scenario_dataset", + "ScenarioDatasetRequest": "pyrit.models.scenario_dataset", + "ScenarioDatasetSettings": "pyrit.models.scenario_dataset", + "ScenarioDatasetSourceSettings": "pyrit.models.scenario_dataset", "DatasetLimitInput": "pyrit.models.scenario_dataset_size_estimate", "DatasetLimitState": "pyrit.models.scenario_dataset_size_estimate", "IndeterminateDatasetSize": "pyrit.models.scenario_dataset_size_estimate", diff --git a/pyrit/models/catalog/scenario.py b/pyrit/models/catalog/scenario.py index 9a9f269a25..8350f7e1a5 100644 --- a/pyrit/models/catalog/scenario.py +++ b/pyrit/models/catalog/scenario.py @@ -16,15 +16,29 @@ from datetime import datetime from enum import Enum from math import prod -from typing import Annotated, Any, Literal - -from pydantic import AliasChoices, BaseModel, BeforeValidator, Field, computed_field, field_validator, model_validator +from typing import Annotated, Any, Literal, Self + +from pydantic import ( + AliasChoices, + BaseModel, + Field, + computed_field, + field_validator, + model_validator, +) -from pyrit.models.dataset_limit import normalize_dataset_limit from pyrit.models.parameter import Parameter from pyrit.models.request_limits import MAX_IDENTIFIER_LENGTH, MAX_ITEMS, MAX_LABEL_KEY_LENGTH, MAX_LABEL_VALUE_LENGTH from pyrit.models.results.scenario_result import ScenarioRunState from pyrit.models.retry_event import RetryEvent +from pyrit.models.scenario_dataset import ( + DATASET_FILTERS, + DatasetLimitInputValue, + DatasetSourceInput, + ScenarioDatasetOptions, + ScenarioDatasetRequest, + ScenarioDatasetSettings, +) from pyrit.models.scenario_dataset_size_estimate import ( DatasetLimitInput, IndeterminateDatasetSize, @@ -48,28 +62,6 @@ Field(max_length=MAX_ITEMS), ] -# Authoritative set of dataset seed filters exposed over the run request surface. Each entry -# is used verbatim as a ``MemoryInterface.get_seeds`` keyword argument, so a filter key IS the -# get_seeds kwarg. Every exposed filter must be a list-valued (Sequence) get_seeds parameter. -# Adding a filterable field is a one-line change here; the CLI ``--dataset-filters`` help text -# describes these keys, and this request model validates them server-side (covering the GUI too). -# -# Comma-list semantics differ per key because ``get_seeds`` treats each field differently, and -# that behavior lives in ``pyrit.memory`` (this layer cannot import it). As of today -# (see ``MemoryInterface.get_seeds`` / ``_add_list_conditions``): -# - harm_categories -> AND + substring: a seed must be tagged with EVERY value, and each value -# is a substring match (``cyber`` matches ``cyber_harm``). So ``harm_categories=cyber,violence`` -# is an intersection, not a union. -# - data_types -> OR + exact: a seed matches ANY value, compared for exact equality. So -# ``data_types=text,image_path`` is a union. -DATASET_FILTERS: frozenset[str] = frozenset({"harm_categories", "data_types"}) - - -_RequestDatasetLimit = Annotated[ - Annotated[int, Field(ge=1, strict=True)] | Literal["all", "default"] | None, - BeforeValidator(normalize_dataset_limit), -] - def _validate_dataset_filter_mapping( value: dict[str, list[str]] | None, @@ -98,6 +90,45 @@ class ScenarioRunSizeEstimateStatus(str, Enum): Unavailable = "unavailable" +class _DatasetRequestModel(BaseModel): + """Dataset selection shared by launch and estimate requests.""" + + datasets: ScenarioDatasetRequest | None = Field( + None, description="Dataset sources, limits, fetch policy, and filters. Omit to use scenario defaults." + ) + dataset_names: _RequestNames | None = Field(None, description="Deprecated; use datasets.sources.") + max_dataset_size: DatasetLimitInputValue = Field("default", description="Deprecated; use datasets.max_total.") + dataset_filters: _RequestFilters | None = Field( + None, description="Deprecated; use datasets.filters. Accepted keys: harm_categories, data_types." + ) + + @property + def dataset_request(self) -> ScenarioDatasetRequest | None: + """The dataset edits, converted from the deprecated fields when they are used.""" + if self.datasets is not None or not self._uses_deprecated_dataset_fields(): + return self.datasets + return ScenarioDatasetRequest( + sources=[DatasetSourceInput(name=name) for name in self.dataset_names] if self.dataset_names else None, + max_total=self.max_dataset_size, + filters=self.dataset_filters or None, + ) + + @field_validator("dataset_filters") + @classmethod + def _validate_dataset_filters(cls, value: dict[str, list[str]] | None) -> dict[str, list[str]] | None: + return _validate_dataset_filter_mapping(value) + + @model_validator(mode="after") + def _validate_dataset_form(self) -> Self: + if self.datasets is not None and self._uses_deprecated_dataset_fields(): + raise ValueError("Use either 'datasets' or the deprecated dataset fields, not both.") + self.dataset_request # noqa: B018 - converts deprecated fields so invalid values fail here + return self + + def _uses_deprecated_dataset_fields(self) -> bool: + return bool(self.dataset_names) or self.max_dataset_size != "default" or bool(self.dataset_filters) + + class ScenarioRunSizeEstimateCondition(str, Enum): """Reason an estimate remains conditional until launch.""" @@ -194,6 +225,9 @@ class ScenarioRunSizeEstimate(BaseModel): components: list[ScenarioRunSizeComponent] = Field(default_factory=list) datasets: list[ScenarioDatasetSummary] = Field(default_factory=list) dataset_size: ScenarioDatasetSizeEstimate = Field(default_factory=IndeterminateDatasetSize) + dataset_options: ScenarioDatasetOptions = Field(default_factory=ScenarioDatasetOptions) + dataset_defaults: ScenarioDatasetSettings = Field(default_factory=ScenarioDatasetSettings) + effective_datasets: ScenarioDatasetSettings = Field(default_factory=ScenarioDatasetSettings) dataset_limit: DatasetLimitInput = Field(default_factory=DatasetLimitInput) effective_parameters: dict[str, bool | int | float | str | list[str]] = Field( default_factory=dict, @@ -365,6 +399,8 @@ class RegisteredScenario(BaseModel): description="Descriptions and tags for the available concrete techniques", ) default_datasets: list[str] = Field(..., description="Default dataset names used by the scenario") + dataset_options: ScenarioDatasetOptions = Field(default_factory=ScenarioDatasetOptions) + dataset_defaults: ScenarioDatasetSettings = Field(default_factory=ScenarioDatasetSettings) baseline_policy: Literal["enabled", "disabled", "forbidden"] = Field( "enabled", description="Whether baseline execution is enabled, disabled, or forbidden" ) @@ -381,7 +417,7 @@ class RegisteredScenario(BaseModel): ) -class ScenarioRunSizeEstimateRequest(BaseModel): +class ScenarioRunSizeEstimateRequest(_DatasetRequestModel): """Request-specific scenario run-size configuration.""" adversarial_target_name: _RequestName | None = Field( @@ -396,19 +432,6 @@ class ScenarioRunSizeEstimateRequest(BaseModel): techniques: _RequestTechniques | None = Field( None, description="Technique names to estimate (uses scenario default if omitted)" ) - dataset_names: _RequestNames | None = Field( - None, description="Dataset names to estimate (uses scenario default if omitted)" - ) - max_dataset_size: _RequestDatasetLimit = Field( - "default", - description="Total selected logical seed-group limit. Omitted, null, empty, or 'default': scenario default. " - "'all': no total limit. " - "Per-dataset limits remain in effect.", - ) - dataset_filters: _RequestFilters | None = Field( - None, - description="Dataset seed filters keyed by field. Accepted keys: harm_categories, data_types.", - ) include_baseline: bool | None = Field( None, description="Override the scenario baseline default; forbidden scenarios reject true", @@ -418,19 +441,8 @@ class ScenarioRunSizeEstimateRequest(BaseModel): description="Scenario-declared parameters such as Jailbreak template and attempt counts", ) - @field_validator("dataset_filters") - @classmethod - def _validate_dataset_filters(cls, value: dict[str, list[str]] | None) -> dict[str, list[str]] | None: - """ - Validate estimate dataset filters against the shared allow-list. - Returns: - dict[str, list[str]] | None: Validated filters. - """ - return _validate_dataset_filter_mapping(value) - - -class RunScenarioRequest(BaseModel): +class RunScenarioRequest(_DatasetRequestModel): """Request body for starting a scenario run.""" scenario_name: _RequestName = Field(..., description="Scenario name (e.g., 'foundry.red_team_agent')") @@ -446,21 +458,6 @@ class RunScenarioRequest(BaseModel): techniques: _RequestTechniques | None = Field( None, description="Technique names to use (uses scenario default if omitted)" ) - dataset_names: _RequestNames | None = Field( - None, description="Dataset names to use (uses scenario default if omitted)" - ) - max_dataset_size: _RequestDatasetLimit = Field( - "default", - description="Total selected logical seed-group limit. Omitted, null, empty, or 'default': scenario default. " - "'all': no total limit. " - "Per-dataset limits remain in effect.", - ) - dataset_filters: _RequestFilters | None = Field( - None, - description=( - "Dataset seed filters keyed by field, applied before sampling. Accepted keys: harm_categories, data_types." - ), - ) 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") include_baseline: bool | None = Field( @@ -485,17 +482,6 @@ class RunScenarioRequest(BaseModel): "If provided, the scenario will resume from prior progress instead of starting fresh.", ) - @field_validator("dataset_filters") - @classmethod - def _validate_dataset_filters(cls, value: dict[str, list[str]] | None) -> dict[str, list[str]] | None: - """ - Reject any dataset-filter key not in the exposed ``DATASET_FILTERS`` allow-list. - - Returns: - dict[str, list[str]] | None: The validated filters, unchanged. - """ - return _validate_dataset_filter_mapping(value) - class AttackErrorSummary(BaseModel): """A single errored attack result surfaced in a run summary.""" diff --git a/pyrit/models/scenario_dataset.py b/pyrit/models/scenario_dataset.py new file mode 100644 index 0000000000..45d93ac303 --- /dev/null +++ b/pyrit/models/scenario_dataset.py @@ -0,0 +1,130 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT license. + +"""Shared scenario dataset inputs and selection capabilities.""" + +from enum import Enum +from typing import Annotated, Literal + +from pydantic import BaseModel, BeforeValidator, ConfigDict, Field, field_validator + +from pyrit.models.dataset_limit import normalize_dataset_limit +from pyrit.models.request_limits import MAX_IDENTIFIER_LENGTH, MAX_ITEMS + +# Dataset seed filters exposed to clients. Each key is a list-valued ``MemoryInterface.get_seeds`` +# keyword argument. Comma-list semantics differ per key because ``get_seeds`` treats each field +# differently (see ``MemoryInterface.get_seeds`` / ``_add_list_conditions``): +# - harm_categories -> AND + substring: a seed must be tagged with EVERY value, and each value +# is a substring match (``cyber`` matches ``cyber_harm``). +# - data_types -> OR + exact: a seed matches ANY value, compared for exact equality. +DATASET_FILTERS: frozenset[str] = frozenset({"harm_categories", "data_types"}) +DatasetLimitInputValue = Annotated[ + Annotated[int, Field(ge=1, strict=True)] | Literal["all", "default"], + BeforeValidator(normalize_dataset_limit), +] + + +class DatasetFetchPolicy(str, Enum): + """When scenario preparation may load a dataset.""" + + NEVER = "never" + IF_MISSING = "if_missing" + + +def _normalize_fetch(value: object) -> object: + if isinstance(value, str): + value = value.strip().lower() + return None if value in (None, "", "default") else value + + +class DatasetSelectionMode(str, Enum): + """The permitted named-dataset selection.""" + + ANY = "any" + SUBSET = "subset" + FIXED = "fixed" + + +class ScenarioDatasetOptions(BaseModel): + """Scenario-owned capabilities for dataset overrides.""" + + model_config = ConfigDict(extra="forbid", frozen=True) + + mode: DatasetSelectionMode = DatasetSelectionMode.ANY + allowed_names: list[str] = Field(default_factory=list) + max_selected: int | None = Field(None, ge=1) + per_dataset_limit: bool = True + total_limit: bool = True + supports_overrides: bool = True + filters: bool = True + source_settings: bool = True + + +class DatasetSourceInput(BaseModel): + """A named source and optional limit and preparation overrides.""" + + model_config = ConfigDict(extra="forbid") + + name: str = Field(min_length=1, max_length=MAX_IDENTIFIER_LENGTH, strict=True) + max_size: DatasetLimitInputValue = "default" + fetch: Annotated[DatasetFetchPolicy | None, BeforeValidator(_normalize_fetch)] = None + + @field_validator("name") + @classmethod + def _validate_name(cls, value: str) -> str: + if not value.strip(): + raise ValueError("Dataset names must be non-empty.") + return value + + +class ScenarioDatasetRequest(BaseModel): + """Serializable dataset overrides shared by all application surfaces.""" + + model_config = ConfigDict(extra="forbid") + + sources: list[DatasetSourceInput] | None = Field(None, min_length=1, max_length=MAX_ITEMS) + max_per_dataset: DatasetLimitInputValue = "default" + max_total: DatasetLimitInputValue = "default" + fetch: Annotated[DatasetFetchPolicy | None, BeforeValidator(_normalize_fetch)] = None + filters: dict[str, list[str]] | None = Field(None, max_length=MAX_ITEMS) + + @field_validator("sources") + @classmethod + def _validate_sources(cls, value: list[DatasetSourceInput] | None) -> list[DatasetSourceInput] | None: + names = [source.name for source in value or []] + if len(names) != len(set(names)): + raise ValueError("Duplicate dataset source names are not allowed.") + return value + + @field_validator("filters") + @classmethod + def _validate_filters(cls, value: dict[str, list[str]] | None) -> dict[str, list[str]] | None: + for key, values in (value or {}).items(): + if key not in DATASET_FILTERS: + raise ValueError(f"Unknown dataset filter '{key}'. Allowed: {', '.join(sorted(DATASET_FILTERS))}.") + if len(values) > MAX_ITEMS or any(len(item) > MAX_IDENTIFIER_LENGTH for item in values): + raise ValueError("Dataset filter values exceed the request limits.") + return value + + +class ScenarioDatasetSourceSettings(BaseModel): + """Effective source settings, not editable request fields.""" + + model_config = ConfigDict(extra="forbid") + + name: str + max_size: Annotated[int, Field(ge=1, strict=True)] | Literal["all"] + fetch: DatasetFetchPolicy + configured_provider: bool = False + + +class ScenarioDatasetSettings(BaseModel): + """Resolved defaults or effective settings for display and persistence.""" + + model_config = ConfigDict(extra="forbid") + + sources: list[ScenarioDatasetSourceSettings] = Field(default_factory=list) + max_per_dataset: Annotated[int, Field(ge=1, strict=True)] | Literal["all"] = "all" + max_total: Annotated[int, Field(ge=1, strict=True)] | Literal["all"] = "all" + fetch: DatasetFetchPolicy = DatasetFetchPolicy.IF_MISSING + filters: dict[str, list[str]] = Field(default_factory=dict) diff --git a/pyrit/models/seeds/seed_dataset_summary.py b/pyrit/models/seeds/seed_dataset_summary.py index bcfb8e85a1..0b5b882303 100644 --- a/pyrit/models/seeds/seed_dataset_summary.py +++ b/pyrit/models/seeds/seed_dataset_summary.py @@ -15,3 +15,4 @@ class SeedDatasetSummary: modalities: tuple[str, ...] harm_categories: tuple[str, ...] has_unlabeled_harm_categories: bool + origins: tuple[str, ...] = () diff --git a/pyrit/registry/components/scenario_registry.py b/pyrit/registry/components/scenario_registry.py index a11afd7828..8a15dcc360 100644 --- a/pyrit/registry/components/scenario_registry.py +++ b/pyrit/registry/components/scenario_registry.py @@ -18,7 +18,13 @@ from dataclasses import dataclass, field from typing import TYPE_CHECKING, Any, Literal -from pyrit.models import ScenarioRunSizeEstimate, ScenarioTechniqueSummary, class_name_to_snake_case +from pyrit.models import ( + ScenarioDatasetOptions, + ScenarioDatasetSettings, + ScenarioRunSizeEstimate, + ScenarioTechniqueSummary, + class_name_to_snake_case, +) from pyrit.models.identifiers.scenario_identifier import ScenarioIdentifier from pyrit.registry.registry import ParamBagRegistry from pyrit.registry.registry_metadata import RegistryMetadata @@ -66,6 +72,11 @@ class ScenarioMetadata(RegistryMetadata): # Default dataset names used by this scenario. default_datasets: tuple[str, ...] = field(kw_only=True) + # Dataset changes the scenario accepts, and its default dataset settings. Both are static: + # they never depend on memory contents or on a run-size estimate succeeding. + dataset_options: ScenarioDatasetOptions = field(kw_only=True, default_factory=ScenarioDatasetOptions) + dataset_defaults: ScenarioDatasetSettings = field(kw_only=True, default_factory=ScenarioDatasetSettings) + # Scenario-declared custom parameters. supported_parameters: tuple[Parameter, ...] = field(kw_only=True, default=()) @@ -210,6 +221,8 @@ def _build_metadata(self, name: str, cls: type[Scenario]) -> ScenarioMetadata: aggregate_techniques=aggregate_techniques, aggregate_technique_expansions=aggregate_technique_expansions, default_datasets=default_datasets, + dataset_options=instance.get_dataset_options(), + dataset_defaults=instance.get_default_dataset_settings(), supported_parameters=supported_parameters, baseline_policy=instance.BASELINE_ATTACK_POLICY.value, include_baseline_by_default=instance.BASELINE_ATTACK_POLICY.value == "enabled", diff --git a/pyrit/scenario/core/dataset_configuration.py b/pyrit/scenario/core/dataset_configuration.py index 16e6c246b1..f31fe6449a 100644 --- a/pyrit/scenario/core/dataset_configuration.py +++ b/pyrit/scenario/core/dataset_configuration.py @@ -47,6 +47,7 @@ scenario_dataset_size_from_limit, ) from pyrit.models.dataset_limit import DatasetLimit, ResolvedDatasetLimit, normalize_dataset_limit +from pyrit.models.scenario_dataset import DatasetFetchPolicy if TYPE_CHECKING: from collections.abc import Callable, Sequence @@ -66,13 +67,6 @@ class _Unset(Enum): VALUE = "unset" -class DatasetFetchPolicy(str, Enum): - """When preparation may load a registered dataset into memory.""" - - NEVER = "never" - IF_MISSING = "if_missing" - - def _resolve_limit(*, name: str, value: DatasetLimit | _Unset, default: ResolvedDatasetLimit) -> ResolvedDatasetLimit: try: limit = normalize_dataset_limit(None if isinstance(value, _Unset) else value) @@ -997,6 +991,11 @@ def __init__( ) self._configurations = list(configurations) + @property + def configurations(self) -> tuple[DatasetAttackConfiguration, ...]: + """The child configurations, in order.""" + return tuple(self._configurations) + @property def has_sampling_limits(self) -> bool: """Whether this compound or any child can select a subset.""" diff --git a/pyrit/scenario/core/dataset_overrides.py b/pyrit/scenario/core/dataset_overrides.py new file mode 100644 index 0000000000..69bd297375 --- /dev/null +++ b/pyrit/scenario/core/dataset_overrides.py @@ -0,0 +1,122 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT license. + +"""Scenario-owned override validation and lightweight dataset projections.""" + +from dataclasses import replace +from typing import Any + +from pyrit.models.scenario_dataset import ( + DatasetSelectionMode, + ScenarioDatasetOptions, + ScenarioDatasetRequest, + ScenarioDatasetSettings, + ScenarioDatasetSourceSettings, +) +from pyrit.scenario.core.dataset_configuration import ( + CompoundDatasetAttackConfiguration, + DatasetAttackConfiguration, + DatasetConstraintError, + DatasetSource, + DatasetSourceKind, +) + + +def apply_dataset_overrides( + *, + config: DatasetAttackConfiguration, + request: ScenarioDatasetRequest, + options: ScenarioDatasetOptions, +) -> DatasetAttackConfiguration: + """ + Apply named-source edits without rebuilding custom configurations. + + Returns: + DatasetAttackConfiguration: Shallow copy preserving live object identity. + + Raises: + DatasetConstraintError: If a requested edit is unsupported. + """ + if not options.supports_overrides and request != ScenarioDatasetRequest(): + raise DatasetConstraintError("This scenario does not support dataset overrides.") + if not options.per_dataset_limit and ( + request.max_per_dataset != "default" or any(source.max_size != "default" for source in request.sources or []) + ): + raise DatasetConstraintError("This scenario does not support per-dataset limits.") + if not options.total_limit and request.max_total != "default": + raise DatasetConstraintError("This scenario does not support a total dataset limit.") + if not options.filters and request.filters: + raise DatasetConstraintError("This scenario does not support dataset filters.") + kwargs: dict[str, Any] = {} + if request.sources is not None: + if isinstance(config, CompoundDatasetAttackConfiguration): + raise DatasetConstraintError("Compound sources are fixed; source edits are ambiguous.") + existing = {source.name: source for source in config.sources} + selected = [] + for source in request.sources: + base = existing.get(source.name, DatasetSource(name=source.name)) + edits = { + **({"max_size": source.max_size} if source.max_size != "default" else {}), + **({"fetch": source.fetch} if source.fetch is not None else {}), + } + selected.append(replace(base, **edits) if edits else base) + kwargs["sources"] = selected + if request.fetch is not None: + kwargs["fetch"] = request.fetch + effective = config.with_overrides( + **kwargs, + max_per_dataset=request.max_per_dataset, + max_total=request.max_total, + filters=request.filters, + ) + validate_dataset_selection(config=effective, defaults=config, options=options) + return effective + + +def validate_dataset_selection( + *, + config: DatasetAttackConfiguration, + defaults: DatasetAttackConfiguration, + options: ScenarioDatasetOptions, +) -> None: + """ + Validate names before preparation while preserving inline Python inputs. + + Raises: + DatasetConstraintError: If the source selection violates scenario capabilities. + """ + if config.source_kind is DatasetSourceKind.INLINE: + return + names = config.dataset_names + if options.mode is DatasetSelectionMode.FIXED and names != defaults.dataset_names: + raise DatasetConstraintError(f"Dataset sources are fixed: {', '.join(defaults.dataset_names)}.") + if options.mode is DatasetSelectionMode.SUBSET and (not names or set(names) - set(options.allowed_names)): + raise DatasetConstraintError(f"Select datasets from: {', '.join(options.allowed_names)}.") + if options.max_selected is not None and len(names) > options.max_selected: + raise DatasetConstraintError(f"Select at most {options.max_selected} dataset(s).") + + +def project_dataset_settings(config: DatasetAttackConfiguration) -> ScenarioDatasetSettings: + """ + Project configuration settings without memory reads or live wire objects. + + Returns: + ScenarioDatasetSettings: Effective settings for response data. + """ + children = config.configurations if isinstance(config, CompoundDatasetAttackConfiguration) else (config,) + return ScenarioDatasetSettings( + sources=[ + ScenarioDatasetSourceSettings( + name=source.name, + max_size=child.source_limit(source.name), + fetch=source.fetch or child.fetch, + configured_provider=getattr(source, "provider", None) is not None, + ) + for child in children + for source in child.sources + ], + max_per_dataset=config.max_per_dataset, + max_total=config.max_total, + fetch=children[0].fetch, + filters=config.filters, + ) diff --git a/pyrit/scenario/core/scenario.py b/pyrit/scenario/core/scenario.py index 2b2ed3fb73..b49f249a42 100644 --- a/pyrit/scenario/core/scenario.py +++ b/pyrit/scenario/core/scenario.py @@ -53,6 +53,7 @@ config_hash, ) from pyrit.models.parameter import ComponentType, Parameter, RegistryReference +from pyrit.models.scenario_dataset import ScenarioDatasetOptions, ScenarioDatasetRequest, ScenarioDatasetSettings from pyrit.prompt_target import PromptTarget from pyrit.prompt_target.common.target_requirements import TargetRequirements from pyrit.registry import ScorerRegistry @@ -61,6 +62,7 @@ from pyrit.scenario.core.dataset_configuration import ( DatasetAttackConfiguration, ) +from pyrit.scenario.core.dataset_overrides import apply_dataset_overrides, project_dataset_settings from pyrit.scenario.core.scenario_context import ScenarioContext from pyrit.scenario.core.scenario_target_defaults import get_default_scorer_target from pyrit.scenario.core.scenario_technique import ScenarioTechnique @@ -348,6 +350,11 @@ def _common_scenario_parameters(cls) -> list[Parameter]: description="Techniques to execute; defaults to the scenario's default aggregate when omitted.", opaque=True, ), + Parameter( + name="dataset_overrides", + description="Serializable dataset edits applied on top of the dataset configuration.", + opaque=True, + ), Parameter( name="technique_converters", description="Mapping of concrete technique name to extra request converters to append.", @@ -636,6 +643,22 @@ async def get_default_run_size_estimate_async(self) -> ScenarioRunSizeEstimate: @final async def get_run_size_estimate_async(self, *, target_is_configured: bool = False) -> ScenarioRunSizeEstimate: + """ + Estimate without preparation and expose resolved defaults separately from edits. + + Returns: + ScenarioRunSizeEstimate: Estimate with dataset defaults and effective settings. + """ + estimate = await self._get_run_size_estimate_async(target_is_configured=target_is_configured) + return estimate.model_copy( + update={ + "dataset_options": self.get_dataset_options(), + "dataset_defaults": project_dataset_settings(self._resolved_default_dataset_config), + "effective_datasets": project_dataset_settings(self._dataset_config), + } + ) + + async def _get_run_size_estimate_async(self, *, target_is_configured: bool = False) -> ScenarioRunSizeEstimate: """ Estimate the currently configured run without creating or persisting it. @@ -818,7 +841,6 @@ def _resolve_runtime_configuration(self, *, require_objective_target: bool) -> N dataset_config = params.get("dataset_config") self._dataset_config_provided = dataset_config is not None - self._dataset_config = (dataset_config if dataset_config else self._default_dataset_config).with_overrides() self._max_concurrency = params.get("max_concurrency", 4) self._max_retries = params.get("max_retries", 0) self._memory_labels = params.get("memory_labels") or {} @@ -839,8 +861,39 @@ def _resolve_runtime_configuration(self, *, require_objective_target: bool) -> N scenario_techniques=params.get("scenario_techniques") ) self._technique_converters = params.get("technique_converters") or {} + self._resolved_default_dataset_config = self._resolve_default_dataset_config().with_overrides() + self._dataset_config = ( + dataset_config if dataset_config is not None else self._resolved_default_dataset_config + ).with_overrides() + options = self.get_dataset_options() + raw_overrides = params.get("dataset_overrides") + if raw_overrides is not None: + self._dataset_config = apply_dataset_overrides( + config=self._dataset_config, + request=ScenarioDatasetRequest.model_validate(raw_overrides), + options=options, + ) self._validate_runtime_configuration() + def _resolve_default_dataset_config(self) -> DatasetAttackConfiguration: + """ + Resolve parameter-aware defaults before any override or preparation. + + Returns: + DatasetAttackConfiguration: Scenario-owned resolved default configuration. + """ + return self._default_dataset_config + + DATASET_OPTIONS: ClassVar[ScenarioDatasetOptions] = ScenarioDatasetOptions() + + def get_dataset_options(self) -> ScenarioDatasetOptions: + """Return the scenario's named-dataset capabilities for the current parameters.""" + return self.DATASET_OPTIONS + + def get_default_dataset_settings(self) -> ScenarioDatasetSettings: + """Return the default dataset settings for the current parameters, without reading memory.""" + return project_dataset_settings(self._resolve_default_dataset_config()) + def _validate_runtime_configuration(self) -> None: """ Check resolved parameters for preview and launch without reading datasets. diff --git a/pyrit/scenario/scenarios/airt/psychosocial.py b/pyrit/scenario/scenarios/airt/psychosocial.py index 1f54d65918..0a006f7f95 100644 --- a/pyrit/scenario/scenarios/airt/psychosocial.py +++ b/pyrit/scenario/scenarios/airt/psychosocial.py @@ -40,11 +40,13 @@ SeedPrompt, ) from pyrit.models.parameter import Parameter +from pyrit.models.scenario_dataset import DatasetSelectionMode, ScenarioDatasetOptions from pyrit.prompt_normalizer.converter_configuration import ConverterConfiguration from pyrit.scenario.core.atomic_attack import AtomicAttack from pyrit.scenario.core.attack_technique import AttackTechnique from pyrit.scenario.core.attack_technique_factory import AttackTechniqueFactory from pyrit.scenario.core.dataset_configuration import ( + CompoundDatasetAttackConfiguration, DatasetAttackConfiguration, DatasetConstraintError, DatasetSource, @@ -478,25 +480,43 @@ def _selected_sub_harms(self) -> list[_SubHarm]: ) return [_SUB_HARMS_BY_NAME[name]] - def _validate_runtime_configuration(self) -> None: - config = self._dataset_config - allowed = {harm.dataset_name for harm in _SUB_HARMS} - if set(config.dataset_names) - allowed: - raise DatasetConstraintError("Psychosocial sources must match its sub-harm datasets.") + def _resolve_default_dataset_config(self) -> DatasetAttackConfiguration: + config = self._default_dataset_config by_name = {source.name: source for source in config.sources} - self._dataset_config = config.with_overrides( + return config.with_overrides( sources=[ by_name.get(harm.dataset_name, DatasetSource(name=harm.dataset_name)) for harm in self._selected_sub_harms() ] ) + + def get_dataset_options(self) -> ScenarioDatasetOptions: + """ + Describe named datasets allowed for the selected sub-harm parameter. + + Returns: + ScenarioDatasetOptions: Parameter-aware source selection capabilities. + """ + return ScenarioDatasetOptions( + mode=DatasetSelectionMode.SUBSET, + allowed_names=[harm.dataset_name for harm in self._selected_sub_harms()], + ) + + def _effective_sub_harms(self) -> list[_SubHarm]: + return [harm for harm in self._selected_sub_harms() if harm.dataset_name in self._dataset_config.dataset_names] + + def _validate_runtime_configuration(self) -> None: + if isinstance(self._dataset_config, CompoundDatasetAttackConfiguration): + raise DatasetConstraintError("Psychosocial requires flat sources, not compound children.") + if set(self._dataset_config.dataset_names) - set(self.get_dataset_options().allowed_names): + raise DatasetConstraintError("Psychosocial sources must match the selected sub-harm datasets.") super()._validate_runtime_configuration() cap = self._dataset_config.max_total - if cap != "all" and cap < len(self._selected_sub_harms()): + if cap != "all" and cap < len(self._effective_sub_harms()): raise DatasetConstraintError( f"Psychosocial max_total ({cap}) must cover every selected sub-harm " - f"({len(self._selected_sub_harms())}). Increase the total limit to at least " - f"{len(self._selected_sub_harms())}, or select one sub-harm." + f"({len(self._effective_sub_harms())}). Increase the total limit to at least " + f"{len(self._effective_sub_harms())}, or select one sub-harm." ) async def _resolve_seed_groups_by_dataset_async( @@ -515,7 +535,7 @@ async def _resolve_seed_groups_by_dataset_async( harm.dataset_name: self._dataset_config._sample_source_groups( name=harm.dataset_name, groups=populations.get(harm.dataset_name, []) ) - for harm in self._selected_sub_harms() + for harm in self._effective_sub_harms() } return sample_with_coverage( groups_by_dataset=selected, @@ -591,7 +611,7 @@ async def _build_atomic_attacks_async(self, *, context: ScenarioContext) -> list # Resolved lazily so a no-arg ``Psychosocial()`` works for registry metadata introspection. adversarial_chat = self._adversarial_chat or get_default_adversarial_target() - sub_harms = self._selected_sub_harms() + sub_harms = self._effective_sub_harms() techniques = [t for t in context.scenario_techniques if isinstance(t, PsychosocialTechnique)] max_turns = int(self.params.get("max_turns", 5)) seed_groups_by_dataset = context.seed_groups_by_dataset diff --git a/pyrit/scenario/scenarios/garak/api_key.py b/pyrit/scenario/scenarios/garak/api_key.py index 648de88d0a..95ff038900 100644 --- a/pyrit/scenario/scenarios/garak/api_key.py +++ b/pyrit/scenario/scenarios/garak/api_key.py @@ -25,6 +25,7 @@ SeedObjective, SeedPrompt, ) +from pyrit.models.scenario_dataset import DatasetSelectionMode, ScenarioDatasetOptions from pyrit.prompt_normalizer import ConverterConfiguration from pyrit.scenario.core.atomic_attack import AtomicAttack from pyrit.scenario.core.attack_technique import AttackTechnique @@ -210,6 +211,7 @@ class ApiKey(Scenario): """ VERSION: int = 1 + DATASET_OPTIONS = ScenarioDatasetOptions(mode=DatasetSelectionMode.FIXED, per_dataset_limit=False) BASELINE_ATTACK_POLICY: ClassVar[BaselineAttackPolicy] = BaselineAttackPolicy.Forbidden @classmethod diff --git a/pyrit/scenario/scenarios/garak/encoding.py b/pyrit/scenario/scenarios/garak/encoding.py index d777cc415d..e70188bd7f 100644 --- a/pyrit/scenario/scenarios/garak/encoding.py +++ b/pyrit/scenario/scenarios/garak/encoding.py @@ -34,6 +34,7 @@ SeedObjective, SeedPrompt, ) +from pyrit.models.scenario_dataset import DatasetSelectionMode, ScenarioDatasetOptions from pyrit.prompt_normalizer.converter_configuration import ConverterConfiguration from pyrit.scenario.core.atomic_attack import AtomicAttack from pyrit.scenario.core.attack_technique import AttackTechnique @@ -168,6 +169,9 @@ class Encoding(Scenario): """ VERSION: int = 2 + DATASET_OPTIONS = ScenarioDatasetOptions( + mode=DatasetSelectionMode.FIXED, per_dataset_limit=False, source_settings=False + ) SUPPORTS_TECHNIQUE_CONVERTERS: ClassVar[bool] = False @apply_defaults diff --git a/pyrit/scenario/scenarios/garak/exploitation.py b/pyrit/scenario/scenarios/garak/exploitation.py index c6d81968b6..4444089eec 100644 --- a/pyrit/scenario/scenarios/garak/exploitation.py +++ b/pyrit/scenario/scenarios/garak/exploitation.py @@ -28,6 +28,7 @@ SeedObjective, SeedPrompt, ) +from pyrit.models.scenario_dataset import DatasetSelectionMode, ScenarioDatasetOptions from pyrit.prompt_normalizer import ConverterConfiguration from pyrit.scenario.core.atomic_attack import AtomicAttack from pyrit.scenario.core.attack_technique_factory import AttackTechniqueFactory @@ -189,6 +190,9 @@ class Exploitation(Scenario): """ VERSION: int = 1 + DATASET_OPTIONS = ScenarioDatasetOptions( + mode=DatasetSelectionMode.FIXED, per_dataset_limit=False, total_limit=False, supports_overrides=False + ) BASELINE_ATTACK_POLICY: ClassVar[BaselineAttackPolicy] = BaselineAttackPolicy.Forbidden @classmethod @@ -222,7 +226,7 @@ def additional_parameters(cls) -> list[Parameter]: @classmethod def supported_parameters(cls) -> list[Parameter]: """Return parameters, excluding overrides of the fixed reference corpus.""" - unsupported = {"dataset_config"} + unsupported = {"dataset_config", "dataset_overrides"} return [parameter for parameter in super().supported_parameters() if parameter.name not in unsupported] @apply_defaults diff --git a/pyrit/scenario/scenarios/garak/figstep.py b/pyrit/scenario/scenarios/garak/figstep.py index 3dea0f7a3b..a9d8109c2e 100644 --- a/pyrit/scenario/scenarios/garak/figstep.py +++ b/pyrit/scenario/scenarios/garak/figstep.py @@ -11,6 +11,7 @@ from pyrit.common.path import SCORER_SEED_PROMPT_PATH from pyrit.executor.attack import AttackScoringConfig, PromptSendingAttack from pyrit.models import AttackSeedGroup, PromptDataType +from pyrit.models.scenario_dataset import DatasetSelectionMode, ScenarioDatasetOptions from pyrit.prompt_normalizer import ConverterConfiguration from pyrit.prompt_target import CapabilityName, TargetRequirements from pyrit.scenario.core.atomic_attack import AtomicAttack @@ -68,6 +69,9 @@ class FigStep(Scenario): """ VERSION: int = 1 + DATASET_OPTIONS = ScenarioDatasetOptions( + mode=DatasetSelectionMode.SUBSET, allowed_names=["figstep", "figstep_pro"], max_selected=1 + ) TARGET_REQUIREMENTS: ClassVar[TargetRequirements] = TargetRequirements( native_required=frozenset({CapabilityName.MULTI_MESSAGE_PIECES}), required_input_modalities=_FIGSTEP_INPUT_MODALITIES, diff --git a/pyrit/scenario/scenarios/garak/latent_injection.py b/pyrit/scenario/scenarios/garak/latent_injection.py index 9615946c51..e62ca764a5 100644 --- a/pyrit/scenario/scenarios/garak/latent_injection.py +++ b/pyrit/scenario/scenarios/garak/latent_injection.py @@ -18,6 +18,7 @@ from pyrit.converter import SearchReplaceConverter from pyrit.executor.attack import AttackConverterConfig, AttackScoringConfig, PromptSendingAttack from pyrit.models import AttackSeedGroup, Parameter, Seed, SeedObjective, SeedPrompt +from pyrit.models.scenario_dataset import DatasetSelectionMode, ScenarioDatasetOptions from pyrit.prompt_normalizer import ConverterConfiguration from pyrit.scenario.core.atomic_attack import AtomicAttack from pyrit.scenario.core.attack_technique import AttackTechnique @@ -314,6 +315,7 @@ class LatentInjection(Scenario): """Apply each selected separator to the same sample of document injections.""" VERSION: int = 2 + DATASET_OPTIONS = ScenarioDatasetOptions(mode=DatasetSelectionMode.FIXED, per_dataset_limit=False) BASELINE_ATTACK_POLICY: ClassVar[BaselineAttackPolicy] = BaselineAttackPolicy.Forbidden SEPARATORS: ClassVar[dict[str, tuple[str, str]]] = { "bare": ("", ""), diff --git a/pyrit/scenario/scenarios/garak/package_hallucination.py b/pyrit/scenario/scenarios/garak/package_hallucination.py index e3a502e286..aa821b54ff 100644 --- a/pyrit/scenario/scenarios/garak/package_hallucination.py +++ b/pyrit/scenario/scenarios/garak/package_hallucination.py @@ -23,6 +23,7 @@ SeedObjective, SeedPrompt, ) +from pyrit.models.scenario_dataset import DatasetSelectionMode, ScenarioDatasetOptions from pyrit.scenario.core.atomic_attack import AtomicAttack from pyrit.scenario.core.attack_technique import AttackTechnique from pyrit.scenario.core.dataset_configuration import DatasetAttackConfiguration, DatasetSource @@ -139,6 +140,9 @@ class PackageHallucination(Scenario): """ VERSION: int = 3 + DATASET_OPTIONS = ScenarioDatasetOptions( + mode=DatasetSelectionMode.FIXED, per_dataset_limit=False, total_limit=False + ) SUPPORTS_TECHNIQUE_CONVERTERS: ClassVar[bool] = False # The plain code request is not an adversarial baseline to compare against, so no baseline. @@ -203,15 +207,18 @@ def __init__( USES_DATASET_SIZE_LIMIT: ClassVar[bool] = False - def _validate_runtime_configuration(self) -> None: + def _resolve_default_dataset_config(self) -> DatasetAttackConfiguration: + techniques = self._resolve_scenario_techniques(scenario_techniques=self.params.get("scenario_techniques")) names = [ *_CORPUS_DATASETS, - *(_LANGUAGE_SPECS[technique.value].dataset_name for technique in self._scenario_techniques), + *(_LANGUAGE_SPECS[technique.value].dataset_name for technique in techniques), ] - by_name = {source.name: source for source in self._dataset_config.sources} - self._dataset_config = self._dataset_config.with_overrides( + by_name = {source.name: source for source in self._default_dataset_config.sources} + return self._default_dataset_config.with_overrides( sources=[by_name.get(name, DatasetSource(name=name)) for name in dict.fromkeys(names)] ) + + def _validate_runtime_configuration(self) -> None: super()._validate_runtime_configuration() if self._max_prompts_per_language < 1: raise ValueError("max_prompts_per_language must be greater than zero") diff --git a/pyrit/scenario/scenarios/garak/prompt_inject.py b/pyrit/scenario/scenarios/garak/prompt_inject.py index 7608f5cb7d..533a3e5a57 100644 --- a/pyrit/scenario/scenarios/garak/prompt_inject.py +++ b/pyrit/scenario/scenarios/garak/prompt_inject.py @@ -17,6 +17,7 @@ from pyrit.executor.attack import AttackConverterConfig, AttackScoringConfig, PromptSendingAttack from pyrit.memory import CentralMemory from pyrit.models import AttackSeedGroup, Parameter, Seed, SeedObjective, SeedPrompt +from pyrit.models.scenario_dataset import DatasetSelectionMode, ScenarioDatasetOptions from pyrit.prompt_normalizer import ConverterConfiguration from pyrit.scenario.core.atomic_attack import AtomicAttack from pyrit.scenario.core.attack_technique import AttackTechnique @@ -217,6 +218,7 @@ class PromptInject(Scenario): """ VERSION: int = 3 + DATASET_OPTIONS = ScenarioDatasetOptions(mode=DatasetSelectionMode.FIXED, per_dataset_limit=False) BASELINE_ATTACK_POLICY: ClassVar[BaselineAttackPolicy] = BaselineAttackPolicy.Forbidden DEFAULT_GOAL_TEXTS: ClassVar[tuple[str, str, str]] = _DEFAULT_GOAL_TEXTS TECHNIQUE_DATASET_NAME: ClassVar[str] = "prompt_inject_techniques" diff --git a/pyrit/scenario/scenarios/garak/system_prompt_extraction.py b/pyrit/scenario/scenarios/garak/system_prompt_extraction.py index ffc53cd1ce..2eb4e5b65a 100644 --- a/pyrit/scenario/scenarios/garak/system_prompt_extraction.py +++ b/pyrit/scenario/scenarios/garak/system_prompt_extraction.py @@ -22,6 +22,7 @@ scenario_dataset_size_from_limit, ) from pyrit.models.dataset_limit import DatasetLimit, normalize_dataset_limit +from pyrit.models.scenario_dataset import DatasetSelectionMode, ScenarioDatasetOptions from pyrit.scenario.core.atomic_attack import AtomicAttack from pyrit.scenario.core.attack_technique import AttackTechnique from pyrit.scenario.core.dataset_configuration import DatasetAttackConfiguration, DatasetSource @@ -98,6 +99,9 @@ class SystemPromptExtraction(Scenario): """ VERSION: int = 1 + DATASET_OPTIONS = ScenarioDatasetOptions( + mode=DatasetSelectionMode.FIXED, per_dataset_limit=False, total_limit=False + ) SUPPORTS_TECHNIQUE_CONVERTERS: ClassVar[bool] = False # Template-dominated like the Doctor/Jailbreak scenarios: the bare system prompt with no diff --git a/pyrit/scenario/scenarios/garak/web_injection.py b/pyrit/scenario/scenarios/garak/web_injection.py index 6ad61d6e65..26f64a8b78 100644 --- a/pyrit/scenario/scenarios/garak/web_injection.py +++ b/pyrit/scenario/scenarios/garak/web_injection.py @@ -23,6 +23,7 @@ SeedObjective, SeedPrompt, ) +from pyrit.models.scenario_dataset import DatasetSelectionMode, ScenarioDatasetOptions from pyrit.scenario.core.atomic_attack import AtomicAttack from pyrit.scenario.core.attack_technique import AttackTechnique from pyrit.scenario.core.dataset_configuration import ( @@ -104,6 +105,21 @@ class WebInjection(Scenario): """ VERSION: int = 2 + + def get_dataset_options(self) -> ScenarioDatasetOptions: + """ + Describe approved ingredient sources subject to technique coverage validation. + + Returns: + ScenarioDatasetOptions: Approved sources with ingredient limits disabled. + """ + return ScenarioDatasetOptions( + mode=DatasetSelectionMode.SUBSET, + allowed_names=self._default_dataset_config.dataset_names, + per_dataset_limit=False, + total_limit=False, + ) + SUPPORTS_TECHNIQUE_CONVERTERS: ClassVar[bool] = False USES_DATASET_SIZE_LIMIT: ClassVar[bool] = False BASELINE_ATTACK_POLICY: ClassVar[BaselineAttackPolicy] = BaselineAttackPolicy.Enabled diff --git a/tests/unit/backend/test_dataset_service.py b/tests/unit/backend/test_dataset_service.py index 86cc51f3a8..19da7f43e5 100644 --- a/tests/unit/backend/test_dataset_service.py +++ b/tests/unit/backend/test_dataset_service.py @@ -58,6 +58,7 @@ async def test_list_datasets_includes_memory_summaries(self, dataset_service, mo modalities=("image_path", "text"), harm_categories=("hate", "violence"), has_unlabeled_harm_categories=True, + origins=("local", "unknown"), ), SeedDatasetSummary( dataset_name=None, @@ -87,11 +88,13 @@ async def test_list_datasets_includes_memory_summaries(self, dataset_service, mo assert loaded.modalities == ["image_path", "text"] assert loaded.harm_categories == ["hate", "violence"] assert loaded.has_unlabeled_harm_categories is True + assert loaded.origins == ["local", "unknown"] provider_only = next(item for item in result.items if item.name == "provider_only") assert provider_only.loaded is False assert provider_only.provider_available is True assert provider_only.logical_examples is None + assert provider_only.origins == [] unnamed = next(item for item in result.items if item.name == "(unnamed)") assert unnamed.selection_key == "dataset:unnamed" diff --git a/tests/unit/backend/test_scenario_configuration_resolver.py b/tests/unit/backend/test_scenario_configuration_resolver.py index 22ff1e2c0d..7710f991fb 100644 --- a/tests/unit/backend/test_scenario_configuration_resolver.py +++ b/tests/unit/backend/test_scenario_configuration_resolver.py @@ -11,7 +11,9 @@ from pyrit.backend.services.scenario_configuration_resolver import ScenarioConfigurationResolver from pyrit.memory import CentralMemory, MemoryInterface from pyrit.models import SeedObjective +from pyrit.models.catalog.scenario import RunScenarioRequest from pyrit.models.dataset_limit import DatasetLimit, ResolvedDatasetLimit +from pyrit.models.scenario_dataset import ScenarioDatasetOptions, ScenarioDatasetRequest from pyrit.prompt_target.common.target_capabilities import TargetCapabilities from pyrit.registry import TargetRegistry from pyrit.scenario import DatasetAttackConfiguration, DatasetFetchPolicy, DatasetSource, Scenario @@ -20,6 +22,7 @@ override_default_adversarial_target, scenario_target_defaults, ) +from pyrit.scenario.core.dataset_overrides import apply_dataset_overrides from pyrit.scenario.scenarios.adaptive.text_adaptive import TextAdaptive from pyrit.scenario.scenarios.airt.rapid_response import RapidResponse from pyrit.scenario.scenarios.garak.api_key import ApiKey @@ -28,16 +31,22 @@ from unit.mocks import MockPromptTarget +def _datasets(**fields: object) -> ScenarioDatasetRequest | None: + return RunScenarioRequest.model_validate({"scenario_name": "s", "target_name": "t", **fields}).dataset_request + + @pytest.mark.usefixtures("patch_central_database") @pytest.mark.parametrize(("limit", "expected"), [(None, 20), ("", 20), ("default", 20), ("all", "all"), (7, 7)]) def test_dataset_name_override_applies_explicit_total(*, limit: DatasetLimit, expected: ResolvedDatasetLimit) -> None: resolved = ScenarioConfigurationResolver.resolve_configuration( scenario_name="garak.api_key", scenario_class=ApiKey, - dataset_names=ApiKey.required_datasets(), - max_dataset_size=limit, + datasets=_datasets(dataset_names=ApiKey.required_datasets(), max_dataset_size=limit), ) - assert resolved["dataset_config"].max_dataset_size == expected + scenario = ApiKey() + scenario.set_params_from_args(args=resolved) + scenario._resolve_runtime_configuration(require_objective_target=False) + assert scenario._dataset_config.max_total == expected @pytest.mark.usefixtures("patch_central_database") @@ -51,10 +60,13 @@ def test_dataset_overrides_preserve_subclass_state_and_default() -> None: resolved = ScenarioConfigurationResolver.resolve_configuration( scenario_name="garak.prompt_inject", scenario_class=MagicMock(return_value=scenario), - dataset_names=PromptInject.required_datasets(), - max_dataset_size=7, - dataset_filters={"data_types": ["text"]}, - )["dataset_config"] + datasets=_datasets( + dataset_names=PromptInject.required_datasets(), + max_dataset_size=7, + dataset_filters={"data_types": ["text"]}, + ), + )["dataset_overrides"] + resolved = apply_dataset_overrides(config=original, request=resolved, options=ScenarioDatasetOptions()) assert type(resolved) is PromptInjectDatasetConfiguration assert resolved._goal_texts == ["first goal", "second goal"] assert resolved.max_total == 7 @@ -83,8 +95,9 @@ async def test_total_override_preserves_source_limits_async( resolved = ScenarioConfigurationResolver.resolve_configuration( scenario_name="test", scenario_class=MagicMock(return_value=scenario), - max_dataset_size=total, - )["dataset_config"] + datasets=_datasets(max_dataset_size=total), + )["dataset_overrides"] + resolved = apply_dataset_overrides(config=original, request=resolved, options=ScenarioDatasetOptions()) populations = { name: [SeedObjective(value=f"{name}-{index}", dataset_name=name) for index in range(20)] for name in original.dataset_names @@ -115,9 +128,9 @@ def test_name_override_retains_source_options_and_adds_defaults_for_new_names(to resolved = ScenarioConfigurationResolver.resolve_configuration( scenario_name="test", scenario_class=MagicMock(return_value=scenario), - dataset_names=["new", "retained"], - max_dataset_size=total, - )["dataset_config"] + datasets=_datasets(dataset_names=["new", "retained"], max_dataset_size=total), + )["dataset_overrides"] + resolved = apply_dataset_overrides(config=original, request=resolved, options=ScenarioDatasetOptions()) assert resolved.sources == (DatasetSource(name="new"), retained) assert resolved.sources[1] is retained assert resolved.source_limit("new") == 4 @@ -130,9 +143,14 @@ def test_name_override_retains_source_options_and_adds_defaults_for_new_names(to @pytest.mark.usefixtures("patch_central_database") def test_omitted_total_retains_default() -> None: resolved = ScenarioConfigurationResolver.resolve_configuration( - scenario_name="garak.api_key", scenario_class=ApiKey, dataset_names=ApiKey.required_datasets() + scenario_name="garak.api_key", + scenario_class=ApiKey, + datasets=_datasets(dataset_names=ApiKey.required_datasets()), ) - assert resolved["dataset_config"].max_total == 20 + scenario = ApiKey() + scenario.set_params_from_args(args=resolved) + scenario._resolve_runtime_configuration(require_objective_target=False) + assert scenario._dataset_config.max_total == 20 @pytest.mark.usefixtures("patch_central_database") diff --git a/tests/unit/backend/test_scenario_resume.py b/tests/unit/backend/test_scenario_resume.py index a85af70e57..aa821b0a45 100644 --- a/tests/unit/backend/test_scenario_resume.py +++ b/tests/unit/backend/test_scenario_resume.py @@ -13,6 +13,7 @@ from pyrit.backend.services.scenario_run_service import ( _LAUNCH_REQUEST_FIELDS, _LAUNCH_REQUEST_METADATA_KEY, + _OPTIONAL_LAUNCH_FIELDS, ScenarioRunConflictError, ScenarioRunNotFoundError, ScenarioRunService, @@ -24,8 +25,11 @@ from pyrit.models import ( SCENARIO_RUN_PLAN_METADATA_KEY, AttackOutcome, + DatasetFetchPolicy, + DatasetSourceInput, Message, Parameter, + ScenarioDatasetRequest, ScenarioResult, ScenarioRunState, SeedObjective, @@ -448,9 +452,7 @@ async def test_resume_missing_scenario_registration_is_explicit_async( await service.resume_run_async(scenario_result_id=str(stored.id)) -@pytest.mark.parametrize( - "missing", [name for name in _LAUNCH_REQUEST_FIELDS if name not in {"adversarial_target_name", "max_dataset_size"}] -) +@pytest.mark.parametrize("missing", [name for name in _LAUNCH_REQUEST_FIELDS if name not in _OPTIONAL_LAUNCH_FIELDS]) async def test_resume_incomplete_saved_configuration_never_uses_defaults_async( *, resume_environment: tuple[ScenarioRunService, MockPromptTarget], missing: str ) -> None: @@ -495,6 +497,28 @@ async def test_saved_total_limit_round_trips_async( assert restored.max_dataset_size == (limit_args.get("max_dataset_size") or "default") +async def test_saved_dataset_request_round_trips_async( + resume_environment: tuple[ScenarioRunService, MockPromptTarget], +) -> None: + service, _ = resume_environment + datasets = ScenarioDatasetRequest( + sources=[DatasetSourceInput(name=_DATASET_NAME, max_size=1, fetch=DatasetFetchPolicy.NEVER)], + max_total=1, + filters={"data_types": ["text"]}, + ) + prepared = await service._prepare_run_async( + request=RunScenarioRequest(scenario_name=_SCENARIO_NAME, target_name=_TARGET_NAME, datasets=datasets) + ) + stored = ( + await CentralMemory.get_memory_instance().get_scenario_results_async( + scenario_result_ids=[prepared.scenario._scenario_result_id] + ) + )[0] + restored = service._restore_launch_request(stored=stored) + assert restored.dataset_request == datasets + assert restored.dataset_names is None + + async def test_resume_older_launch_record_without_adversarial_selection_async( resume_environment: tuple[ScenarioRunService, MockPromptTarget], ) -> None: diff --git a/tests/unit/backend/test_scenario_run_service.py b/tests/unit/backend/test_scenario_run_service.py index a3d04ae1e0..86d21bf21d 100644 --- a/tests/unit/backend/test_scenario_run_service.py +++ b/tests/unit/backend/test_scenario_run_service.py @@ -56,6 +56,7 @@ config_hash, ) from pyrit.models.catalog.scenario import RunScenarioRequest, ScenarioTechniqueSummary +from pyrit.models.scenario_dataset import ScenarioDatasetOptions, ScenarioDatasetRequest from pyrit.prompt_target.common.target_capabilities import TargetCapabilities from pyrit.registry import ScenarioMetadata, ScenarioRegistry from pyrit.scenario import Scenario @@ -64,6 +65,7 @@ DatasetConfiguration, get_default_adversarial_target, ) +from pyrit.scenario.core.dataset_overrides import apply_dataset_overrides from pyrit.scenario.core.scenario_technique import ScenarioTechnique from pyrit.scenario.scenarios.airt.jailbreak import Jailbreak from pyrit.scenario.scenarios.airt.scam import Scam @@ -74,6 +76,16 @@ from pyrit.prompt_target import PromptTarget +def _apply_mapped_datasets(registries: dict[str, Any]) -> DatasetAttackConfiguration: + call = registries["scenario_registry"].create_and_initialize_async.await_args + assert "dataset_config" not in call.kwargs + return apply_dataset_overrides( + config=registries["scenario_instance"]._default_dataset_config, + request=call.kwargs.get("dataset_overrides", ScenarioDatasetRequest()), + options=ScenarioDatasetOptions(), + ) + + class _StubTechnique(ScenarioTechnique): """Minimal concrete ScenarioTechnique used to exercise converter-token parsing.""" @@ -423,7 +435,8 @@ def update_state(*, scenario_result_id: str, scenario_run_state: ScenarioRunStat release_first.set() await asyncio.wait_for(finished.wait(), timeout=5) await asyncio.wait_for(first_task, timeout=5) - assert introspected == initialized == [first, second, fallback, second] + assert introspected == [] + assert initialized == [first, second, fallback, second] assert executed == [first, second, fallback] assert get_default_adversarial_target() is fallback finally: @@ -826,8 +839,8 @@ async def test_start_run_max_dataset_size_uses_default_config(self, mock_all_reg assert default_config.max_total == 100 init_call = mock_all_registries["scenario_registry"].create_and_initialize_async.await_args - assert init_call.kwargs["dataset_config"] is not default_config - assert init_call.kwargs["dataset_config"].max_total == 5 + assert init_call.kwargs["dataset_overrides"].max_total == 5 + assert "dataset_config" not in init_call.kwargs @pytest.mark.parametrize( ("limit_args", "expected"), @@ -850,7 +863,7 @@ async def test_start_run_resolves_total_limit_async( request=RunScenarioRequest(scenario_name="test", target_name="my_target", **limit_args) ) init_call = mock_all_registries["scenario_registry"].create_and_initialize_async.await_args - config = init_call.kwargs.get("dataset_config", default_config) + config = _apply_mapped_datasets(mock_all_registries) assert config.max_total == expected saved = init_call.kwargs["initial_metadata"][_svc_mod._LAUNCH_REQUEST_METADATA_KEY] assert saved["max_dataset_size"] == (limit_args.get("max_dataset_size") or "default") @@ -876,7 +889,7 @@ class _MarkerDatasetConfiguration(DatasetConfiguration): await service.start_run_async(request=_make_request(dataset_names=["custom_a", "custom_b"], max_dataset_size=3)) init_call = mock_all_registries["scenario_registry"].create_and_initialize_async.await_args - built_config = init_call.kwargs["dataset_config"] + built_config = _apply_mapped_datasets(mock_all_registries) # Type is preserved (this is the regression assertion) assert type(built_config) is _MarkerDatasetConfiguration @@ -902,7 +915,7 @@ class _MarkerDatasetConfiguration(DatasetConfiguration): await service.start_run_async(request=_make_request(dataset_names=["only_this"])) init_call = mock_all_registries["scenario_registry"].create_and_initialize_async.await_args - built_config = init_call.kwargs["dataset_config"] + built_config = _apply_mapped_datasets(mock_all_registries) assert type(built_config) is _MarkerDatasetConfiguration assert built_config.dataset_names == ["only_this"] assert built_config.max_dataset_size == "all" @@ -926,7 +939,7 @@ def __init__(self, *, required_extra: str, **kwargs: Any) -> None: service = ScenarioRunService() await service.start_run_async(request=_make_request(dataset_names=["custom"])) init_call = mock_all_registries["scenario_registry"].create_and_initialize_async.await_args - config = init_call.kwargs["dataset_config"] + config = _apply_mapped_datasets(mock_all_registries) assert isinstance(config, _RequiresExtraArgConfiguration) assert config._required_extra == "seeded" assert config.dataset_names == ["custom"] @@ -950,7 +963,7 @@ class _MarkerDatasetConfiguration(DatasetConfiguration): ) init_call = mock_all_registries["scenario_registry"].create_and_initialize_async.await_args - built_config = init_call.kwargs["dataset_config"] + built_config = _apply_mapped_datasets(mock_all_registries) assert built_config.dataset_names == ["custom"] assert built_config.max_dataset_size == 7 assert built_config.filters == {"harm_categories": ["cyber"]} @@ -965,7 +978,7 @@ async def test_start_run_dataset_filters_updates_default_config(self, mock_all_r await service.start_run_async(request=_make_request(dataset_filters={"harm_categories": ["cyber"]})) init_call = mock_all_registries["scenario_registry"].create_and_initialize_async.await_args - built_config = init_call.kwargs["dataset_config"] + built_config = _apply_mapped_datasets(mock_all_registries) assert built_config is not default_config assert default_config.filters == {} assert built_config.filters == {"harm_categories": ["cyber"]} @@ -994,8 +1007,13 @@ async def test_start_run_dataset_names_introspection_failure_raises(self, mock_m patch(f"{_REGISTRY_PATCH_BASE}.TargetRegistry.get_registry_singleton", return_value=mock_tr), patch(f"{_REGISTRY_PATCH_BASE}.InitializerRegistry.get_registry_singleton", return_value=mock_ir), ): - with pytest.raises(ValueError, match="not instantiable without arguments"): - await service.start_run_async(request=_make_request(dataset_names=["custom"])) + mapped = _svc_mod.ScenarioConfigurationResolver.resolve_configuration( + scenario_name="test", + scenario_class=mock_scenario_class, + datasets=_make_request(dataset_names=["custom"]).dataset_request, + ) + assert mapped["dataset_overrides"].sources[0].name == "custom" + mock_scenario_class.assert_not_called() async def test_start_run_max_dataset_size_with_dataset_names_uses_subclass_with_both( self, mock_all_registries @@ -1013,9 +1031,7 @@ class _MarkerDatasetConfiguration(DatasetConfiguration): service = ScenarioRunService() await service.start_run_async(request=_make_request(dataset_names=["a", "b"], max_dataset_size=7)) - built_config = mock_all_registries["scenario_registry"].create_and_initialize_async.await_args.kwargs[ - "dataset_config" - ] + built_config = _apply_mapped_datasets(mock_all_registries) assert type(built_config) is _MarkerDatasetConfiguration assert built_config.dataset_names == ["a", "b"] assert built_config.max_dataset_size == 7 diff --git a/tests/unit/backend/test_scenario_service.py b/tests/unit/backend/test_scenario_service.py index 17c0505402..e99fce1b1f 100644 --- a/tests/unit/backend/test_scenario_service.py +++ b/tests/unit/backend/test_scenario_service.py @@ -397,7 +397,8 @@ async def estimate_async(**kwargs: object) -> ScenarioRunSizeEstimate: for name in ("first", "second") ) ) - assert introspected == estimated == [first, second] + assert introspected == [] + assert estimated == [first, second] assert get_default_adversarial_target() is fallback assert dict(service._estimate_cache) == cache_before assert await service.get_scenario_async(scenario_name=metadata.registry_name) == before @@ -1541,7 +1542,11 @@ async def test_estimate_resolves_total_limit_async( scenario_name="test.scenario", request=ScenarioRunSizeEstimateRequest(dataset_names=["harmbench"], **limit_args), ) - config = service._registry.create_and_estimate_async.await_args.kwargs["dataset_config"] + from pyrit.models.scenario_dataset import ScenarioDatasetOptions + from pyrit.scenario.core.dataset_overrides import apply_dataset_overrides + + mapped = service._registry.create_and_estimate_async.await_args.kwargs["dataset_overrides"] + config = apply_dataset_overrides(config=original, request=mapped, options=ScenarioDatasetOptions()) assert config.max_total == expected assert config.max_per_dataset == 5 assert original.max_total == 20 @@ -1619,11 +1624,11 @@ async def test_configured_estimate_uses_shared_launch_resolution(self) -> None: assert call.kwargs["scenario_techniques"] == [_EstimateTechnique.PROMPT_SENDING] assert call.kwargs["include_baseline"] is True assert call.kwargs["objective_target"] is objective_target - dataset_config = call.kwargs["dataset_config"] - assert type(dataset_config) is DatasetAttackConfiguration - assert dataset_config.dataset_names == ["harmbench"] - assert dataset_config.max_dataset_size == 3 - assert dataset_config.filters == {"harm_categories": ["violence"]} + overrides = call.kwargs["dataset_overrides"] + assert "dataset_config" not in call.kwargs + assert [source.name for source in overrides.sources] == ["harmbench"] + assert overrides.max_total == 3 + assert overrides.filters == {"harm_categories": ["violence"]} async def test_configured_estimate_rejects_incompatible_v4_jailbreak_technique(self) -> None: """Request previews reject techniques omitted by Jailbreak's v4 compatibility policy.""" diff --git a/tests/unit/cli/test_dataset_request.py b/tests/unit/cli/test_dataset_request.py new file mode 100644 index 0000000000..58cdbd63b4 --- /dev/null +++ b/tests/unit/cli/test_dataset_request.py @@ -0,0 +1,54 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT license. + +"""Scanner and shell share strict edit-only dataset parsing.""" + +from contextlib import nullcontext + +import pytest + +from pyrit.cli._cli_args import build_dataset_request_fields, parse_run_arguments +from pyrit.cli.pyrit_scan import _build_run_request, parse_args + + +@pytest.mark.parametrize( + "flags", + [ + "--datasets harmbench --dataset-names other", + "--max-per-dataset 2 --max-dataset-size 4", + "--max-total all --dataset-names harmbench", + "--dataset-fetch never --max-dataset-size default", + ], +) +def test_scanner_and_shell_reject_mixed_dataset_flags(flags: str) -> None: + parsed = parse_args(["run", "test", "--target", "target", *flags.split()]) + with pytest.raises(ValueError, match="Do not mix"): + _build_run_request(parsed_args=parsed, scenario_name="test") + with pytest.raises(ValueError, match="Do not mix"): + build_dataset_request_fields(parse_run_arguments(args_string="test " + flags)) + + +@pytest.mark.parametrize( + "flags", + [ + "--datasets name:0", + "--datasets name:4:all", + "--datasets name name", + "--dataset-fetch always", + "--max-total -1", + ], +) +def test_shell_rejects_malformed_datasets(flags: str) -> None: + with pytest.raises(ValueError): + build_dataset_request_fields(parse_run_arguments(args_string="test " + flags)) + + +@pytest.mark.parametrize("source_flags", ["--datasets harmbench:all", "--dataset-names harmbench"]) +def test_common_filters_work_with_either_form(source_flags: str) -> None: + flags = source_flags + " --dataset-filters data_types=text,image_path" + with pytest.warns(DeprecationWarning) if "--dataset-names" in flags else nullcontext(): + fields = build_dataset_request_fields(parse_run_arguments(args_string="test " + flags)) + if "datasets" in fields: + assert fields["datasets"].filters == {"data_types": ["text", "image_path"]} + else: + assert fields["dataset_filters"] == {"data_types": ["text", "image_path"]} diff --git a/tests/unit/memory/memory_interface/test_interface_seed_dataset_summary_query.py b/tests/unit/memory/memory_interface/test_interface_seed_dataset_summary_query.py index 820494cd22..24e9bdeede 100644 --- a/tests/unit/memory/memory_interface/test_interface_seed_dataset_summary_query.py +++ b/tests/unit/memory/memory_interface/test_interface_seed_dataset_summary_query.py @@ -5,7 +5,21 @@ from pyrit.memory import MemoryInterface from pyrit.memory.memory_session import MemorySession -from pyrit.models import SeedPrompt +from pyrit.models import SeedOrigin, SeedPrompt + + +async def test_dataset_origins_aggregate_without_seed_hydration(sqlite_instance: MemoryInterface) -> None: + await sqlite_instance.add_seeds_to_memory_async( + seeds=[ + SeedPrompt(value="local", dataset_name="mixed", origin=SeedOrigin.LOCAL), + SeedPrompt(value="unknown", dataset_name="mixed", origin=SeedOrigin.UNKNOWN), + ], + added_by="test", + ) + with patch.object(sqlite_instance, "get_seeds_async", side_effect=AssertionError("No hydration")): + summaries = await sqlite_instance.get_seed_dataset_summaries_async() + assert summaries[0].origins == ("local", "unknown") + assert summaries[0].logical_examples == summaries[0].seed_pieces == 2 async def test_get_seed_dataset_summaries_avoids_metadata_row_multiplication( diff --git a/tests/unit/models/test_scenario_catalog.py b/tests/unit/models/test_scenario_catalog.py index 6910656e1a..01debc3841 100644 --- a/tests/unit/models/test_scenario_catalog.py +++ b/tests/unit/models/test_scenario_catalog.py @@ -66,6 +66,9 @@ def test_run_size_estimate_preserves_legacy_total_and_serializes_additively() -> assert estimate.minimum_attack_count == 6 assert estimate.maximum_attack_count == 6 payload = estimate.model_dump(mode="json") + assert payload.pop("dataset_options")["supports_overrides"] is True + assert payload.pop("dataset_defaults")["sources"] == [] + assert payload.pop("effective_datasets")["sources"] == [] assert list(payload["effective_parameters"]) == ["include_baseline", "techniques"] assert payload == { "status": "exact", diff --git a/tests/unit/models/test_scenario_request.py b/tests/unit/models/test_scenario_request.py index e51c384c7d..4105756438 100644 --- a/tests/unit/models/test_scenario_request.py +++ b/tests/unit/models/test_scenario_request.py @@ -6,9 +6,108 @@ import pytest from pydantic import ValidationError +from pyrit.models import DatasetSourceInput, ScenarioDatasetRequest from pyrit.models.catalog.scenario import DATASET_FILTERS, RunScenarioRequest, ScenarioRunSizeEstimateRequest +@pytest.mark.parametrize("model", [RunScenarioRequest, ScenarioRunSizeEstimateRequest]) +def test_new_dataset_request_transport_round_trip(model: type) -> None: + data = { + "scenario_name": "s", + "target_name": "t", + "datasets": { + "sources": [{"name": "harmbench"}, {"name": "custom", "max_size": "all", "fetch": "never"}], + "max_per_dataset": 4, + "max_total": 12, + "fetch": "if_missing", + "filters": {"data_types": ["text"]}, + }, + } + request = model.model_validate(data) + assert request.dataset_request == request.datasets + assert model.model_validate(request.model_dump(mode="json", exclude_none=True)) == request + assert model.model_validate_json(request.model_dump_json()) == request + + +@pytest.mark.parametrize("model", [RunScenarioRequest, ScenarioRunSizeEstimateRequest]) +@pytest.mark.parametrize( + "legacy", + [{"dataset_names": ["a"]}, {"max_dataset_size": 3}, {"dataset_filters": {"data_types": ["text"]}}], +) +def test_new_dataset_request_rejects_deprecated_values(*, model: type, legacy: dict[str, object]) -> None: + with pytest.raises(ValidationError, match="not both"): + model.model_validate({"scenario_name": "s", "target_name": "t", "datasets": {"max_total": 2}, **legacy}) + + +@pytest.mark.parametrize("model", [RunScenarioRequest, ScenarioRunSizeEstimateRequest]) +def test_new_dataset_request_accepts_deprecated_defaults(model: type) -> None: + request = model.model_validate( + { + "scenario_name": "s", + "target_name": "t", + "datasets": {"max_total": 2}, + "dataset_names": None, + "max_dataset_size": "default", + "dataset_filters": None, + } + ) + assert request.dataset_request == ScenarioDatasetRequest(max_total=2) + + +@pytest.mark.parametrize("model", [RunScenarioRequest, ScenarioRunSizeEstimateRequest]) +def test_deprecated_dataset_fields_convert_to_one_request(model: type) -> None: + request = model.model_validate( + { + "scenario_name": "s", + "target_name": "t", + "dataset_names": ["a", "b"], + "max_dataset_size": 3, + "dataset_filters": {"data_types": ["text"]}, + } + ) + assert request.datasets is None + assert request.dataset_request == ScenarioDatasetRequest( + sources=[DatasetSourceInput(name="a"), DatasetSourceInput(name="b")], + max_total=3, + filters={"data_types": ["text"]}, + ) + assert model.model_validate({"scenario_name": "s", "target_name": "t"}).dataset_request is None + + +@pytest.mark.parametrize("model", [RunScenarioRequest, ScenarioRunSizeEstimateRequest]) +def test_invalid_deprecated_dataset_names_fail_validation(model: type) -> None: + with pytest.raises(ValidationError, match="Duplicate"): + model.model_validate({"scenario_name": "s", "target_name": "t", "dataset_names": ["a", "a"]}) + + +@pytest.mark.parametrize( + "data", + [ + {"sources": []}, + {"sources": [{"name": "a"}, {"name": "a"}]}, + {"sources": [{"name": " "}]}, + {"sources": [{"name": "a", "provider": {}}]}, + {"max_total": True}, + {"max_per_dataset": 0}, + {"max_total": 1.5}, + {"fetch": "always"}, + {"filters": {"authors": ["me"]}}, + {"filters": {"data_types": "text"}}, + {"unexpected": True}, + ], +) +def test_shared_dataset_request_rejects_invalid_inputs(data: dict[str, object]) -> None: + with pytest.raises(ValidationError): + ScenarioDatasetRequest.model_validate(data) + + +@pytest.mark.parametrize("value", [None, "", "default"]) +def test_shared_dataset_request_default_semantics(value: object) -> None: + request = ScenarioDatasetRequest.model_validate({"max_total": value, "max_per_dataset": value, "fetch": value}) + assert request.max_total == request.max_per_dataset == "default" + assert request.fetch is None + + @pytest.mark.parametrize("model", [RunScenarioRequest, ScenarioRunSizeEstimateRequest]) @pytest.mark.parametrize( ("limit", "expected"), diff --git a/tests/unit/registry/test_scenario_registry.py b/tests/unit/registry/test_scenario_registry.py index 6722782574..10d2af2782 100644 --- a/tests/unit/registry/test_scenario_registry.py +++ b/tests/unit/registry/test_scenario_registry.py @@ -7,6 +7,7 @@ import pytest +from pyrit.models import DatasetSelectionMode, ScenarioDatasetOptions, ScenarioDatasetSettings from pyrit.registry.components.scenario_registry import ScenarioRegistry from pyrit.scenario.core import BaselineAttackPolicy, ScenarioTechnique @@ -62,6 +63,14 @@ def _resolve_scenario_techniques(self, *, scenario_techniques): """Resolve the concrete defaults.""" return _MetadataTechnique.resolve(scenario_techniques, default=self._default_technique) + def get_dataset_options(self) -> ScenarioDatasetOptions: + """Return fixed dataset options.""" + return ScenarioDatasetOptions(mode=DatasetSelectionMode.FIXED) + + def get_default_dataset_settings(self) -> ScenarioDatasetSettings: + """Return default dataset settings.""" + return ScenarioDatasetSettings(max_total=4) + class _MarkdownMetadataScenario(_MetadataScenario): """ @@ -104,6 +113,8 @@ def test_build_metadata_expands_ordered_default_techniques() -> None: "all": ("one", "two"), "default": ("one", "two"), } + assert metadata.dataset_options.mode is DatasetSelectionMode.FIXED + assert metadata.dataset_defaults.max_total == 4 def test_build_metadata_reports_scenario_owned_adversarial_usage() -> None: diff --git a/tests/unit/scenario/airt/test_leakage.py b/tests/unit/scenario/airt/test_leakage.py index 1888d3be4f..903538b642 100644 --- a/tests/unit/scenario/airt/test_leakage.py +++ b/tests/unit/scenario/airt/test_leakage.py @@ -10,6 +10,7 @@ from pyrit.common.path import DATASETS_PATH from pyrit.models import AttackSeedGroup, ComponentIdentifier, SeedDataset, SeedObjective +from pyrit.models.scenario_dataset import DatasetFetchPolicy from pyrit.prompt_target import PromptTarget from pyrit.registry import TargetRegistry from pyrit.registry.components.attack_technique_registry import AttackTechniqueRegistry @@ -52,6 +53,9 @@ def mock_dataset_config(mock_memory_seeds): mock_config.with_overrides.return_value = mock_config mock_config.sources = () mock_config.max_total = "all" + mock_config.max_per_dataset = "all" + mock_config.fetch = DatasetFetchPolicy.IF_MISSING + mock_config._filters = {} mock_config.get_attack_groups_by_dataset_async = AsyncMock(return_value={"airt_leakage": seed_groups}) mock_config.dataset_names = ["airt_leakage"] return mock_config diff --git a/tests/unit/scenario/airt/test_psychosocial.py b/tests/unit/scenario/airt/test_psychosocial.py index 816fb42e0d..da5ada8929 100644 --- a/tests/unit/scenario/airt/test_psychosocial.py +++ b/tests/unit/scenario/airt/test_psychosocial.py @@ -24,6 +24,7 @@ scenario_dataset_size_from_limit, ) from pyrit.models.dataset_limit import ResolvedDatasetLimit +from pyrit.models.scenario_dataset import ScenarioDatasetRequest from pyrit.prompt_target import PromptTarget from pyrit.registry import TargetRegistry from pyrit.scenario.core.dataset_configuration import ( @@ -188,7 +189,9 @@ def get_seeds(*, dataset_name: str, **_: object) -> list[SeedObjective]: return list(seeds_by_dataset[dataset_name]) config = DatasetAttackConfiguration( - sources=[DatasetSource(name=name) for name in seeds_by_dataset], + sources=[ + DatasetSource(name=harm.dataset_name) for harm in _SUB_HARMS if sub_harm == "all" or harm.name == sub_harm + ], max_per_dataset="all", max_total=outer_limit, ) @@ -253,9 +256,9 @@ async def test_backend_total_one_rejects_both_sub_harms_before_reads_async( scenario_name="airt.psychosocial", scenario_class=MagicMock(return_value=scenario), objective_target=mock_objective_target, - max_dataset_size=1, + datasets=ScenarioDatasetRequest(max_total=1), ) - assert args["dataset_config"].max_total == 1 + assert args["dataset_overrides"].max_total == 1 scenario.set_params_from_args(args=args) with ( patch.object(DatasetAttackConfiguration, "prepare_async", side_effect=AssertionError("No preparation")), @@ -310,7 +313,7 @@ async def test_backend_total_one_allows_one_selected_sub_harm_async(mock_objecti scenario_name="airt.psychosocial", scenario_class=MagicMock(return_value=scenario), objective_target=mock_objective_target, - max_dataset_size=1, + datasets=ScenarioDatasetRequest(max_total=1), ) scenario.set_params_from_args( args={**args, "sub_harm": "imminent_crisis", "scenario_techniques": [PsychosocialTechnique.NoConverter]} diff --git a/tests/unit/scenario/airt/test_scam.py b/tests/unit/scenario/airt/test_scam.py index 5c4a99043e..f28390de81 100644 --- a/tests/unit/scenario/airt/test_scam.py +++ b/tests/unit/scenario/airt/test_scam.py @@ -12,6 +12,7 @@ from pyrit.executor.attack import PromptSendingAttack, RedTeamingAttack from pyrit.executor.attack.core.attack_config import AttackScoringConfig from pyrit.models import AttackSeedGroup, ComponentIdentifier, SeedDataset, SeedObjective +from pyrit.models.scenario_dataset import DatasetFetchPolicy from pyrit.prompt_target import OpenAIChatTarget, PromptTarget from pyrit.scenario import DatasetAttackConfiguration, DatasetConfiguration from pyrit.scenario.scenarios.airt.scam import Scam, ScamTechnique @@ -57,6 +58,9 @@ def mock_dataset_config(mock_memory_seed_groups): mock_config.with_overrides.return_value = mock_config mock_config.sources = () mock_config.max_total = "all" + mock_config.max_per_dataset = "all" + mock_config.fetch = DatasetFetchPolicy.IF_MISSING + mock_config._filters = {} mock_config.get_attack_seed_groups_async = AsyncMock(return_value=attack_seed_groups) mock_config.get_attack_groups_by_dataset_async = AsyncMock(return_value={"airt_scam": attack_seed_groups}) mock_config.dataset_names = ["airt_scam"] diff --git a/tests/unit/scenario/core/test_scenario_dataset_overrides.py b/tests/unit/scenario/core/test_scenario_dataset_overrides.py new file mode 100644 index 0000000000..6b54762149 --- /dev/null +++ b/tests/unit/scenario/core/test_scenario_dataset_overrides.py @@ -0,0 +1,198 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT license. + +"""Offline dataset contract coverage across Python, REST, scanner, and shell.""" + +from unittest.mock import AsyncMock, patch + +import pytest +from unit.mocks import MockPromptTarget + +from pyrit.backend.services.scenario_configuration_resolver import ScenarioConfigurationResolver +from pyrit.backend.services.scenario_service import ScenarioService +from pyrit.cli._cli_args import ( + build_dataset_request_fields, + parse_run_arguments, +) +from pyrit.cli.pyrit_scan import _build_run_request, parse_args +from pyrit.memory import CentralMemory +from pyrit.models import ( + DatasetFetchPolicy, + DatasetSelectionMode, + DatasetSourceInput, + ScenarioDatasetRequest, + SeedObjective, +) +from pyrit.models.catalog import RunScenarioRequest +from pyrit.registry import ScenarioRegistry +from pyrit.scenario import DatasetAttackConfiguration, DatasetSource, Scenario +from pyrit.scenario.core.atomic_attack import AtomicAttack +from pyrit.scenario.core.matrix_atomic_attack_builder import build_baseline_atomic_attack +from pyrit.scenario.core.scenario_context import ScenarioContext +from pyrit.scenario.core.scenario_technique import ScenarioTechnique +from pyrit.scenario.scenarios.garak.api_key import ApiKey +from pyrit.scenario.scenarios.garak.encoding import Encoding +from pyrit.scenario.scenarios.garak.figstep import FigStep +from pyrit.scenario.scenarios.garak.package_hallucination import PackageHallucination +from pyrit.score import SubStringScorer + + +class _Technique(ScenarioTechnique): + ALL = ("all", {"all"}) + DIRECT = ("direct", {"direct"}) + + @classmethod + def get_aggregate_tags(cls) -> set[str]: + return {"all"} + + +class _Scenario(Scenario): + def __init__(self, *, scenario_result_id: str | None = None) -> None: + super().__init__( + version=1, + technique_class=_Technique, + default_dataset_config=DatasetAttackConfiguration( + sources=[DatasetSource(name="a", max_size=2, fetch=DatasetFetchPolicy.NEVER), DatasetSource(name="b")], + max_per_dataset=5, + max_total=7, + ), + objective_scorer=SubStringScorer(substring="yes"), + scenario_result_id=scenario_result_id, + ) + + async def _build_atomic_attacks_async(self, *, context: ScenarioContext) -> list[AtomicAttack]: + return [ + build_baseline_atomic_attack( + objective_target=context.objective_target, + objective_scorer=self._objective_scorer, + seed_groups=list(context.seed_groups), + memory_labels=context.memory_labels, + ) + ] + + +@pytest.mark.usefixtures("patch_central_database") +async def test_python_rest_scanner_shell_resolve_identical_settings_without_fetching() -> None: + tokens = "--datasets a b:all --max-per-dataset 3 --max-total 4 --dataset-fetch never" + scanner_request = _build_run_request( + parsed_args=parse_args(["run", "test", "--target", "test", *tokens.split()]), scenario_name="test" + ) + scanner = scanner_request.datasets + shell = build_dataset_request_fields(parse_run_arguments(args_string="test " + tokens))["datasets"] + python = ScenarioDatasetRequest( + sources=[DatasetSourceInput(name="a"), DatasetSourceInput(name="b", max_size="all")], + max_per_dataset=3, + max_total=4, + fetch=DatasetFetchPolicy.NEVER, + ) + assert scanner == shell == python + request = RunScenarioRequest(scenario_name="test", target_name="test", datasets=python) + transported = RunScenarioRequest.model_validate_json(request.model_dump_json(exclude_none=True)) + mapped = ScenarioConfigurationResolver.resolve_configuration( + scenario_name="test", + scenario_class=_Scenario, + datasets=transported.dataset_request, + ) + assert "dataset_config" not in mapped + estimates = [] + memory = CentralMemory.get_memory_instance() + with ( + patch.object(DatasetAttackConfiguration, "prepare_async", side_effect=AssertionError("No preparation")), + patch.object(memory, "get_seeds_async", side_effect=AssertionError("No seed reads")), + ): + for overrides in [python, scanner, shell, mapped["dataset_overrides"]]: + scenario = _Scenario() + scenario.set_params_from_args(args={"dataset_overrides": overrides}) + estimates.append(await scenario.get_run_size_estimate_async()) + assert all(estimate.effective_datasets == estimates[0].effective_datasets for estimate in estimates) + estimate = estimates[0] + assert estimate.dataset_defaults.max_total == 7 + assert estimate.effective_datasets.max_total == 4 + assert [source.max_size for source in estimate.effective_datasets.sources] == [2, "all"] + assert estimate.dataset_defaults.sources[0].fetch is DatasetFetchPolicy.NEVER + + +@pytest.mark.usefixtures("patch_central_database") +async def test_resume_replays_saved_selection_without_preparation() -> None: + memory = CentralMemory.get_memory_instance() + await memory.add_seeds_to_memory_async( + seeds=[SeedObjective(value=f"{name}-{index}", dataset_name=name) for name in ["a", "b"] for index in range(6)], + added_by="test", + ) + target = MockPromptTarget() + overrides = ScenarioDatasetRequest(max_total=4, fetch=DatasetFetchPolicy.NEVER) + scenario = _Scenario() + scenario.set_params_from_args(args={"objective_target": target, "dataset_overrides": overrides}) + with patch.object(DatasetAttackConfiguration, "prepare_async", new_callable=AsyncMock) as prepare: + await scenario.initialize_async() + prepare.assert_awaited_once() + selected = [group.logical_id for group in scenario._atomic_attacks[0].seed_groups] + assert len(selected) == 4 + + # Changed caps do not matter on resume: the stored run plan fixes the selected seed groups. + resumed = _Scenario(scenario_result_id=scenario._scenario_result_id) + resumed._default_dataset_config = resumed._default_dataset_config.with_overrides(max_total=1, max_per_dataset=1) + resumed.set_params_from_args(args={"objective_target": target, "dataset_overrides": overrides}) + with patch.object(DatasetAttackConfiguration, "prepare_async", side_effect=AssertionError("No resume preparation")): + await resumed.initialize_async() + assert [group.logical_id for group in resumed._atomic_attacks[0].seed_groups] == selected + + # A different dataset selection is an error, never a silent replacement. + diverged = _Scenario(scenario_result_id=scenario._scenario_result_id) + diverged.set_params_from_args( + args={ + "objective_target": target, + "dataset_overrides": ScenarioDatasetRequest(sources=[DatasetSourceInput(name="a")]), + } + ) + with pytest.raises(ValueError, match="does not match"): + await diverged.initialize_async() + + +@pytest.mark.usefixtures("patch_central_database") +@pytest.mark.parametrize( + ("scenario_class", "overrides", "error"), + [ + (FigStep, {"sources": [{"name": "other"}]}, "Select datasets"), + (FigStep, {"sources": [{"name": "figstep"}, {"name": "figstep_pro"}]}, "at most"), + (ApiKey, {"max_per_dataset": 1}, "per-dataset"), + (ApiKey, {"sources": [{"name": "harmbench"}]}, "fixed"), + (Encoding, {"sources": [{"name": "garak_slur_terms_en"}]}, "ambiguous"), + (PackageHallucination, {"max_total": 1}, "total dataset"), + ], +) +async def test_constrained_scenarios_reject_unsupported_overrides_before_preparation( + *, + scenario_class: type[Scenario], + overrides: dict[str, object], + error: str, +) -> None: + scenario = scenario_class(objective_scorer=SubStringScorer(substring="yes")) + scenario.set_params_from_args(args={"dataset_overrides": overrides}) + with ( + patch.object(DatasetAttackConfiguration, "prepare_async", side_effect=AssertionError("No preparation")), + pytest.raises(ValueError, match=error), + ): + await scenario.get_run_size_estimate_async() + + +@pytest.mark.usefixtures("patch_central_database") +async def test_encoding_global_overrides_preserve_child_caps() -> None: + scenario = Encoding() + scenario.set_params_from_args(args={"dataset_overrides": {"max_total": 3, "fetch": "never"}}) + estimate = await scenario.get_run_size_estimate_async() + assert [child.max_total for child in scenario._dataset_config.configurations] == [10, 10] + assert estimate.effective_datasets.max_total == 3 + + +@pytest.mark.usefixtures("patch_central_database") +async def test_catalog_dataset_controls_do_not_depend_on_estimates() -> None: + with patch.object(ScenarioRegistry, "_discover"): + registry = ScenarioRegistry() + registry.register_class(ApiKey, name="garak.api_key") + with patch.object(ScenarioRegistry, "get_registry_singleton", return_value=registry): + [item] = (await ScenarioService().list_scenarios_async(include_estimates=False)).items + assert item.dataset_options.mode is DatasetSelectionMode.FIXED + assert item.dataset_options.per_dataset_limit is False + assert [source.name for source in item.dataset_defaults.sources] == ApiKey.required_datasets() + assert item.dataset_defaults.max_total == 20 diff --git a/tests/unit/scenario/foundry/test_red_team_agent.py b/tests/unit/scenario/foundry/test_red_team_agent.py index ad15602a47..cc00251ff3 100644 --- a/tests/unit/scenario/foundry/test_red_team_agent.py +++ b/tests/unit/scenario/foundry/test_red_team_agent.py @@ -41,6 +41,7 @@ from pyrit.executor.attack.core.attack_config import AttackScoringConfig from pyrit.executor.attack.core.attack_strategy import AttackStrategy from pyrit.models import AttackSeedGroup, ComponentIdentifier, SeedObjective +from pyrit.models.scenario_dataset import DatasetFetchPolicy from pyrit.prompt_target import PromptTarget from pyrit.scenario import AtomicAttack, DatasetAttackConfiguration from pyrit.scenario.foundry import ( # type: ignore[ty:unresolved-import] @@ -133,6 +134,9 @@ def mock_dataset_config(mock_memory_seed_groups): mock_config.with_overrides.return_value = mock_config mock_config.sources = () mock_config.max_total = "all" + mock_config.max_per_dataset = "all" + mock_config.fetch = DatasetFetchPolicy.IF_MISSING + mock_config._filters = {} mock_config.get_attack_seed_groups_async = AsyncMock(return_value=mock_memory_seed_groups) mock_config.dataset_names = ["foundry_red_team"] return mock_config diff --git a/tests/unit/scenario/garak/test_api_key.py b/tests/unit/scenario/garak/test_api_key.py index 0a93b2389c..21175a6600 100644 --- a/tests/unit/scenario/garak/test_api_key.py +++ b/tests/unit/scenario/garak/test_api_key.py @@ -12,7 +12,7 @@ from pyrit.backend.services.scenario_configuration_resolver import ScenarioConfigurationResolver from pyrit.converter import Base64Converter, Converter from pyrit.executor.attack import PromptSendingAttack -from pyrit.models import ComponentIdentifier, ScenarioRunSizeEstimateStatus, Seed, SeedDataset +from pyrit.models import ComponentIdentifier, ScenarioDatasetRequest, ScenarioRunSizeEstimateStatus, Seed, SeedDataset from pyrit.prompt_target import PromptTarget from pyrit.scenario.core.dataset_configuration import DatasetAttackConfiguration, DatasetConstraintError from pyrit.scenario.core.scenario import BaselineAttackPolicy @@ -247,7 +247,7 @@ async def test_launch_and_estimate_use_standard_dataset_size( scenario_name="garak.api_key", scenario_class=ApiKey, objective_target=mock_objective_target, - max_dataset_size=size, + datasets=ScenarioDatasetRequest(max_total=size), ) scenario.set_params_from_args(args=args) with patch.object( diff --git a/tests/unit/scenario/garak/test_divergence.py b/tests/unit/scenario/garak/test_divergence.py index bd07e847fa..49730b0fc0 100644 --- a/tests/unit/scenario/garak/test_divergence.py +++ b/tests/unit/scenario/garak/test_divergence.py @@ -18,6 +18,7 @@ DivergesFromRepetition, Message, MessagePiece, + ScenarioDatasetRequest, Seed, SeedDataset, SeedObjective, @@ -154,7 +155,7 @@ async def test_runtime_budget_matches_estimate(self, corpus: list[Seed], size: i scenario_name="garak.divergence", scenario_class=Divergence, objective_target=MockPromptTarget(), - max_dataset_size=size, + datasets=ScenarioDatasetRequest(max_total=size), ) scenario.set_params_from_args(args=args) with patch.object( diff --git a/tests/unit/scenario/garak/test_encoding.py b/tests/unit/scenario/garak/test_encoding.py index f6ef5ea8f1..3668873cc6 100644 --- a/tests/unit/scenario/garak/test_encoding.py +++ b/tests/unit/scenario/garak/test_encoding.py @@ -10,6 +10,7 @@ from pyrit.converter import Base64Converter from pyrit.executor.attack import PromptSendingAttack from pyrit.models import AttackSeedGroup, ComponentIdentifier, SeedObjective, SeedPrompt +from pyrit.models.scenario_dataset import DatasetFetchPolicy from pyrit.prompt_target import PromptTarget from pyrit.scenario import CompoundDatasetAttackConfiguration, DatasetAttackConfiguration, DatasetConfiguration from pyrit.scenario.garak import Encoding, EncodingTechnique # type: ignore[ty:unresolved-import] @@ -65,6 +66,9 @@ def mock_dataset_config(mock_attack_seed_groups): mock_config.with_overrides.return_value = mock_config mock_config.sources = () mock_config.max_total = "all" + mock_config.max_per_dataset = "all" + mock_config.fetch = DatasetFetchPolicy.IF_MISSING + mock_config._filters = {} mock_config.get_attack_seed_groups_async = AsyncMock(return_value=mock_attack_seed_groups) mock_config.dataset_names = ["garak_slur_terms_en", "garak_web_html_js"] return mock_config