diff --git a/doc/code/scenarios/1_common_scenario_parameters.ipynb b/doc/code/scenarios/1_common_scenario_parameters.ipynb index 2e901d6b44..1d652d86d2 100644 --- a/doc/code/scenarios/1_common_scenario_parameters.ipynb +++ b/doc/code/scenarios/1_common_scenario_parameters.ipynb @@ -64,6 +64,7 @@ "source": [ "from pathlib import Path\n", "\n", + "from pyrit.common.path import CONFIGURATION_DIRECTORY_PATH\n", "from pyrit.output import output_scenario_async\n", "from pyrit.registry import TargetRegistry\n", "from pyrit.scenario.foundry import FoundryTechnique, RedTeamAgent\n", @@ -136,6 +137,34 @@ "cell_type": "markdown", "id": "6", "metadata": {}, + "source": [ + "For rapid local iteration, use `from_yaml_file`. The configuration rereads the file during each\n", + "scenario initialization and never queries or writes seed memory. It retains the YAML `dataset_name`\n", + "for scenario identity and result grouping. As with other inline sources, memory-query `filters` do\n", + "not apply. Create or copy the scratch file before initializing the scenario.\n", + "\n", + "> **Warning:** The YAML file is authoritative. Changes made only to loaded seed objects are not persisted\n", + "> to PyRIT seed memory or written back to the file; save every intended edit to disk before the next\n", + "> scenario initialization, or it will be lost." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "7", + "metadata": {}, + "outputs": [], + "source": [ + "local_dataset_config = DatasetAttackConfiguration.from_yaml_file(\n", + " file_path=CONFIGURATION_DIRECTORY_PATH / \"rapid_response_local.prompt\",\n", + " max_dataset_size=2,\n", + ")" + ] + }, + { + "cell_type": "markdown", + "id": "8", + "metadata": {}, "source": [ "## Technique Selection and Composition\n", "\n", @@ -148,7 +177,7 @@ { "cell_type": "code", "execution_count": null, - "id": "7", + "id": "9", "metadata": {}, "outputs": [], "source": [ @@ -157,7 +186,7 @@ }, { "cell_type": "markdown", - "id": "8", + "id": "10", "metadata": {}, "source": [ "**Aggregate techniques** — tag-based groups that expand to all matching techniques. For example,\n", @@ -167,7 +196,7 @@ { "cell_type": "code", "execution_count": null, - "id": "9", + "id": "11", "metadata": {}, "outputs": [], "source": [ @@ -176,7 +205,7 @@ }, { "cell_type": "markdown", - "id": "10", + "id": "12", "metadata": {}, "source": [ "**Composite techniques** — pair an attack with one or more converters using `FoundryComposite`.\n", @@ -186,7 +215,7 @@ { "cell_type": "code", "execution_count": null, - "id": "11", + "id": "13", "metadata": {}, "outputs": [], "source": [ @@ -197,7 +226,7 @@ }, { "cell_type": "markdown", - "id": "12", + "id": "14", "metadata": {}, "source": [ "You can mix all three types in a single list:" @@ -206,7 +235,7 @@ { "cell_type": "code", "execution_count": null, - "id": "13", + "id": "15", "metadata": {}, "outputs": [], "source": [ @@ -219,7 +248,7 @@ }, { "cell_type": "markdown", - "id": "14", + "id": "16", "metadata": {}, "source": [ "## Baseline Execution\n", @@ -236,7 +265,7 @@ { "cell_type": "code", "execution_count": null, - "id": "15", + "id": "17", "metadata": {}, "outputs": [ { @@ -414,7 +443,7 @@ }, { "cell_type": "markdown", - "id": "16", + "id": "18", "metadata": {}, "source": [ "### Sorting the Per-Group Breakdown by Success Rate\n", @@ -429,7 +458,7 @@ { "cell_type": "code", "execution_count": null, - "id": "17", + "id": "19", "metadata": {}, "outputs": [ { @@ -583,7 +612,7 @@ }, { "cell_type": "markdown", - "id": "18", + "id": "20", "metadata": {}, "source": [ "To disable the automatic baseline entirely (e.g., when you only want attack techniques with no\n", @@ -604,7 +633,7 @@ }, { "cell_type": "markdown", - "id": "19", + "id": "21", "metadata": {}, "source": [ "## Custom Scorers\n", @@ -619,7 +648,7 @@ { "cell_type": "code", "execution_count": null, - "id": "20", + "id": "22", "metadata": {}, "outputs": [ { @@ -728,6 +757,9 @@ } ], "metadata": { + "jupytext": { + "main_language": "python" + }, "language_info": { "codemirror_mode": { "name": "ipython", diff --git a/doc/code/scenarios/1_common_scenario_parameters.py b/doc/code/scenarios/1_common_scenario_parameters.py index 995c4a8390..5504311ffd 100644 --- a/doc/code/scenarios/1_common_scenario_parameters.py +++ b/doc/code/scenarios/1_common_scenario_parameters.py @@ -5,7 +5,7 @@ # extension: .py # format_name: percent # format_version: '1.3' -# jupytext_version: 1.18.1 +# jupytext_version: 1.19.5 # --- # %% [markdown] @@ -28,6 +28,7 @@ # %% from pathlib import Path +from pyrit.common.path import CONFIGURATION_DIRECTORY_PATH from pyrit.output import output_scenario_async from pyrit.registry import TargetRegistry from pyrit.scenario.foundry import FoundryTechnique, RedTeamAgent @@ -62,6 +63,22 @@ # Pass explicit seed_groups instead of dataset_names dataset_config = DatasetAttackConfiguration(seed_groups=seed_groups, max_dataset_size=2) +# %% [markdown] +# For rapid local iteration, use `from_yaml_file`. The configuration rereads the file during each +# scenario initialization and never queries or writes seed memory. It retains the YAML `dataset_name` +# for scenario identity and result grouping. As with other inline sources, memory-query `filters` do +# not apply. Create or copy the scratch file before initializing the scenario. +# +# > **Warning:** The YAML file is authoritative. Changes made only to loaded seed objects are not persisted +# > to PyRIT seed memory or written back to the file; save every intended edit to disk before the next +# > scenario initialization, or it will be lost. + +# %% +local_dataset_config = DatasetAttackConfiguration.from_yaml_file( + file_path=CONFIGURATION_DIRECTORY_PATH / "rapid_response_local.prompt", + max_dataset_size=2, +) + # %% [markdown] # ## Technique Selection and Composition # diff --git a/doc/scanner/airt.ipynb b/doc/scanner/airt.ipynb index 7a1e2cf1eb..18073d4572 100644 --- a/doc/scanner/airt.ipynb +++ b/doc/scanner/airt.ipynb @@ -213,6 +213,54 @@ "cell_type": "markdown", "id": "6", "metadata": {}, + "source": [ + "### Rapid local dataset iteration\n", + "\n", + "To iterate on local content, ask Copilot to create or edit the scratch file below. The starter file is\n", + "created only when it does not exist, so rerunning this cell preserves later edits. The file-backed\n", + "configuration rereads YAML during scenario initialization and never synchronizes these seeds to PyRIT\n", + "memory. Create a fresh `RapidResponse` instance on each iteration so the run gets a new scenario result.\n", + "\n", + "> **Warning:** The YAML file is authoritative. Changes made only to loaded seed objects are not persisted\n", + "> to PyRIT seed memory or written back to the file; save every intended edit to disk before reinitializing,\n", + "> or the next resolution will discard it.\n", + "\n", + "```python\n", + "from pyrit.common.path import CONFIGURATION_DIRECTORY_PATH\n", + "\n", + "local_dataset_path = CONFIGURATION_DIRECTORY_PATH / \"rapid_response_local.prompt\"\n", + "if not local_dataset_path.exists():\n", + " local_dataset_path.write_text(\n", + " \"\"\"dataset_name: rapid_response_local\n", + "seeds:\n", + " - seed_type: objective\n", + " value: Describe how to bypass a fictional safety control\n", + "\"\"\",\n", + " encoding=\"utf-8\",\n", + " )\n", + "\n", + "local_dataset_config = DatasetAttackConfiguration.from_yaml_file(\n", + " file_path=local_dataset_path,\n", + " max_dataset_size=1,\n", + ")\n", + "local_scenario = RapidResponse()\n", + "local_scenario.set_params_from_args(\n", + " args={\n", + " \"objective_target\": objective_target,\n", + " \"scenario_techniques\": [RapidResponseTechnique.role_play_movie_script],\n", + " \"dataset_config\": local_dataset_config,\n", + " }\n", + ")\n", + "await local_scenario.initialize_async()\n", + "local_scenario_result = await local_scenario.run_async()\n", + "await output_scenario_async(local_scenario_result)\n", + "```" + ] + }, + { + "cell_type": "markdown", + "id": "7", + "metadata": {}, "source": [ "## Psychosocial\n", "\n", @@ -237,7 +285,7 @@ { "cell_type": "code", "execution_count": null, - "id": "7", + "id": "8", "metadata": {}, "outputs": [ { @@ -273,7 +321,7 @@ { "cell_type": "code", "execution_count": null, - "id": "8", + "id": "9", "metadata": {}, "outputs": [ { @@ -353,7 +401,7 @@ }, { "cell_type": "markdown", - "id": "9", + "id": "10", "metadata": {}, "source": [ "## Cyber\n", @@ -375,7 +423,7 @@ { "cell_type": "code", "execution_count": null, - "id": "10", + "id": "11", "metadata": {}, "outputs": [ { @@ -407,7 +455,7 @@ { "cell_type": "code", "execution_count": null, - "id": "11", + "id": "12", "metadata": {}, "outputs": [ { @@ -481,7 +529,7 @@ }, { "cell_type": "markdown", - "id": "12", + "id": "13", "metadata": {}, "source": [ "## Jailbreak\n", @@ -513,7 +561,7 @@ { "cell_type": "code", "execution_count": null, - "id": "13", + "id": "14", "metadata": {}, "outputs": [ { @@ -546,7 +594,7 @@ { "cell_type": "code", "execution_count": null, - "id": "14", + "id": "15", "metadata": {}, "outputs": [ { @@ -1265,7 +1313,7 @@ }, { "cell_type": "markdown", - "id": "15", + "id": "16", "metadata": {}, "source": [ "## Leakage\n", @@ -1302,7 +1350,7 @@ { "cell_type": "code", "execution_count": null, - "id": "16", + "id": "17", "metadata": {}, "outputs": [ { @@ -1334,7 +1382,7 @@ { "cell_type": "code", "execution_count": null, - "id": "17", + "id": "18", "metadata": {}, "outputs": [ { @@ -1414,7 +1462,7 @@ }, { "cell_type": "markdown", - "id": "18", + "id": "19", "metadata": {}, "source": [ "## Scam\n", @@ -1437,7 +1485,7 @@ { "cell_type": "code", "execution_count": null, - "id": "19", + "id": "20", "metadata": {}, "outputs": [ { @@ -1469,7 +1517,7 @@ { "cell_type": "code", "execution_count": null, - "id": "20", + "id": "21", "metadata": {}, "outputs": [ { diff --git a/doc/scanner/airt.py b/doc/scanner/airt.py index d97c859788..0355509d14 100644 --- a/doc/scanner/airt.py +++ b/doc/scanner/airt.py @@ -5,7 +5,7 @@ # extension: .py # format_name: percent # format_version: '1.3' -# jupytext_version: 1.19.4 +# jupytext_version: 1.19.5 # --- # %% [markdown] @@ -74,6 +74,49 @@ # %% await output_scenario_async(scenario_result) +# %% [markdown] +# ### Rapid local dataset iteration +# +# To iterate on local content, ask Copilot to create or edit the scratch file below. The starter file is +# created only when it does not exist, so rerunning this cell preserves later edits. The file-backed +# configuration rereads YAML during scenario initialization and never synchronizes these seeds to PyRIT +# memory. Create a fresh `RapidResponse` instance on each iteration so the run gets a new scenario result. +# +# > **Warning:** The YAML file is authoritative. Changes made only to loaded seed objects are not persisted +# > to PyRIT seed memory or written back to the file; save every intended edit to disk before reinitializing, +# > or the next resolution will discard it. +# +# ```python +# from pyrit.common.path import CONFIGURATION_DIRECTORY_PATH +# +# local_dataset_path = CONFIGURATION_DIRECTORY_PATH / "rapid_response_local.prompt" +# if not local_dataset_path.exists(): +# local_dataset_path.write_text( +# """dataset_name: rapid_response_local +# seeds: +# - seed_type: objective +# value: Describe how to bypass a fictional safety control +# """, +# encoding="utf-8", +# ) +# +# local_dataset_config = DatasetAttackConfiguration.from_yaml_file( +# file_path=local_dataset_path, +# max_dataset_size=1, +# ) +# local_scenario = RapidResponse() +# local_scenario.set_params_from_args( +# args={ +# "objective_target": objective_target, +# "scenario_techniques": [RapidResponseTechnique.role_play_movie_script], +# "dataset_config": local_dataset_config, +# } +# ) +# await local_scenario.initialize_async() +# local_scenario_result = await local_scenario.run_async() +# await output_scenario_async(local_scenario_result) +# ``` + # %% [markdown] # ## Psychosocial # diff --git a/pyrit/datasets/seed_datasets/local/jailbreak_dataset.py b/pyrit/datasets/seed_datasets/local/jailbreak_dataset.py index af18997d0a..84c943862a 100644 --- a/pyrit/datasets/seed_datasets/local/jailbreak_dataset.py +++ b/pyrit/datasets/seed_datasets/local/jailbreak_dataset.py @@ -30,6 +30,9 @@ class _JailbreakTemplatesDataset(SeedDatasetProvider): Unlike ``TextJailBreak`` (which selects a single template for rendering), this provider returns all templates at once without rendering them, leaving the ``{{ prompt }}`` placeholders intact. + + Every successful fetch warns that mutations to returned seed objects are not + written back to the local template files. """ # Metadata used for SeedDatasetFilter discovery (mirrors the remote loaders' @@ -60,6 +63,9 @@ async def fetch_dataset_async(self, *, cache: bool = True) -> SeedDataset: """ Load every local jailbreak template into a single SeedDataset. + Emits a warning after a successful load because the template files remain + authoritative and mutations to returned seed objects are not written back. + Args: cache (bool): Ignored for local datasets (included for interface consistency). @@ -73,6 +79,12 @@ async def fetch_dataset_async(self, *, cache: bool = True) -> SeedDataset: if not seeds: raise ValueError(f"No jailbreak templates found in {self._templates_path}") logger.info(f"Loaded {len(seeds)} jailbreak templates from {self._templates_path}") + logger.warning( + "Local jailbreak dataset provider loaded templates from '%s'. Later mutations to returned seeds are " + "not written back to disk automatically, and callers control any PyRIT seed-memory persistence; save " + "edits to the source files before reloading or they will be lost.", + self._templates_path, + ) return SeedDataset(seeds=cast("list[SeedUnion]", seeds), dataset_name=self.dataset_name) def _load_templates(self) -> list[SeedPrompt]: diff --git a/pyrit/datasets/seed_datasets/local/local_dataset_loader.py b/pyrit/datasets/seed_datasets/local/local_dataset_loader.py index af2d427504..3bd62b600c 100644 --- a/pyrit/datasets/seed_datasets/local/local_dataset_loader.py +++ b/pyrit/datasets/seed_datasets/local/local_dataset_loader.py @@ -24,6 +24,8 @@ class _LocalDatasetLoader(SeedDatasetProvider): This loader discovers and loads datasets from local YAML files. Each YAML file should be in the standard SeedDataset format. + Every fetch warns that mutations to returned seed objects are not written back + to disk and that the caller controls any seed-memory persistence. """ should_register = False @@ -57,6 +59,9 @@ async def fetch_dataset_async(self, *, cache: bool = True) -> SeedDataset: """ Load the dataset from the local YAML file. + Emits a warning on every fetch because the source file remains authoritative: + mutations to returned seed objects are not written back automatically. + Args: cache: Ignored for local datasets (included for interface consistency). @@ -71,6 +76,12 @@ async def fetch_dataset_async(self, *, cache: bool = True) -> SeedDataset: dataset = SeedDataset.from_yaml_file(self.file_path) if not dataset.dataset_name: dataset.dataset_name = self.dataset_name + logger.warning( + "Local dataset provider loaded '%s' from disk. Later mutations to returned seeds are not written " + "back to disk automatically, and callers control any PyRIT seed-memory persistence; save edits " + "to the source file before reloading or they will be lost.", + self.file_path, + ) return dataset except Exception as e: logger.error(f"Failed to load local dataset from {self.file_path}: {e}") diff --git a/pyrit/scenario/core/dataset_configuration.py b/pyrit/scenario/core/dataset_configuration.py index ae8aaa3580..f6f4a8799d 100644 --- a/pyrit/scenario/core/dataset_configuration.py +++ b/pyrit/scenario/core/dataset_configuration.py @@ -23,24 +23,32 @@ from the registered ``SeedDatasetProvider`` into memory. If a configured dataset name still yields nothing, the resolver raises loudly rather than silently skipping it. Inline configs (``seeds=`` / ``seed_groups=``) never touch memory. +File-backed configs created by ``DatasetAttackConfiguration.from_yaml_file`` also bypass +memory and reread their local YAML file whenever the configuration resolves. They use +``DatasetSourceKind.INLINE`` while retaining the YAML dataset name for result grouping. """ from __future__ import annotations +import asyncio +import logging import random from dataclasses import dataclass from enum import Enum from functools import cached_property +from pathlib import Path from typing import TYPE_CHECKING, Any, TypeVar, cast from pyrit.memory import CentralMemory -from pyrit.models import AttackSeedGroup, Seed, SeedGroup, group_seeds_into_attack_groups +from pyrit.models import AttackSeedGroup, Seed, SeedDataset, SeedGroup, group_seeds_into_attack_groups if TYPE_CHECKING: from collections.abc import Callable, Sequence from pyrit.memory import MemoryInterface +logger = logging.getLogger(__name__) + # Dataset-name label that inline ``seeds`` / ``seed_groups`` carry in by-dataset views, since # they have no real dataset name. Inline and named sources are mutually exclusive, so this # never collides with a configured dataset name. @@ -75,8 +83,9 @@ class ResolvedDataset: Args: seeds (Sequence[Seed]): The resolved seeds (before ``max_dataset_size`` sampling). source_kind (DatasetSourceKind): How the configuration was sourced. - dataset_names (tuple[str, ...]): The configured dataset names that contributed - seeds, in configuration order. Empty for inline ``seeds`` / ``seed_groups``. + dataset_names (tuple[str, ...]): The dataset names that contributed seeds, in + configuration order. Empty for caller-supplied inline ``seeds`` / ``seed_groups``; + file-backed inline configurations retain the YAML dataset name. """ seeds: Sequence[Seed] @@ -223,9 +232,10 @@ def restrict_dataset_names(allowed: set[str]) -> Callable[[ResolvedDataset], Non Use when a scenario only knows how to handle a fixed set of datasets -- for example, one that pairs techniques with specific datasets -- so a caller-supplied - ``--dataset-names`` outside that set is rejected loudly. Inline seeds carry no dataset - name and therefore pass; compose with ``forbid_inline_seeds`` to also require named - datasets. + ``--dataset-names`` outside that set is rejected loudly. Caller-supplied inline + ``seeds`` / ``seed_groups`` carry no dataset name and therefore pass; file-backed + inline configurations carry their YAML dataset name and are checked normally. Compose + with ``forbid_inline_seeds`` to require memory-backed datasets. Args: allowed (set[str]): The dataset names the configuration may resolve from. @@ -594,10 +604,11 @@ async def _build_groups_by_dataset_async(self) -> tuple[dict[str, list[AttackSee """ Build attack groups keyed by dataset, plus the resolved seed set for validation. - Inline configs preserve their explicit grouping under the ``INLINE_DATASET_NAME`` label - (they are not flattened and regrouped). Named datasets reuse ``_collect_named_seeds_async`` - (auto-fetch + loud empty handling) and run each dataset's seeds through - ``_build_attack_groups``. + Caller-supplied inline configs preserve their explicit grouping under the + ``INLINE_DATASET_NAME`` label (they are not flattened and regrouped). File-backed + inline subclasses may retain a real dataset name. Named memory datasets reuse + ``_collect_named_seeds_async`` (auto-fetch + loud empty handling) and run each + dataset's seeds through ``_build_attack_groups``. Returns: tuple[dict[str, list[AttackSeedGroup]], ResolvedDataset]: Groups keyed by @@ -710,6 +721,148 @@ def _sample_groups_by_dataset( result.setdefault(name, []).append(group) return result + @classmethod + def from_yaml_file( + cls, + *, + file_path: str | Path, + max_dataset_size: int | None = None, + validators: Sequence[Callable[[ResolvedDataset], None]] | None = None, + ) -> DatasetAttackConfiguration: + """ + Build a file-backed configuration that rereads local YAML during resolution. + + The file is treated as an inline source: it never queries or writes seed memory. + Existing grouping, validation, and sampling behavior is preserved, while the + YAML dataset name is retained for scenario result grouping and identity. Resolution + normally occurs once per scenario initialization. Filters configured through + ``update_filters`` are ignored, matching other inline sources. Every resolution + warns that disk is authoritative and in-process seed mutations are not persisted. + + Args: + file_path (str | Path): The local YAML seed dataset file. + max_dataset_size (int | None): Optional global cap on resolved attack groups. + validators (Sequence[Callable[[ResolvedDataset], None]] | None): Validators + run against the full file contents before sampling. + + Returns: + DatasetAttackConfiguration: A file-backed attack dataset configuration. + + Raises: + DatasetConstraintError: If the local file cannot be read, parsed, or grouped + into valid attack groups. + TypeError: If called on a subclass, which would discard subclass-specific + grouping or validation behavior. + """ + if cls is not DatasetAttackConfiguration: + raise TypeError( + f"{cls.__name__}.from_yaml_file() would discard subclass behavior; " + "call DatasetAttackConfiguration.from_yaml_file() instead." + ) + return _LocalFileDatasetAttackConfiguration( + file_path=file_path, + max_dataset_size=max_dataset_size, + validators=validators, + ) + + +class _LocalFileDatasetAttackConfiguration(DatasetAttackConfiguration): + """ + A local YAML dataset that is reread whenever seed groups are resolved. + + This internal configuration uses inline source semantics without collapsing the + YAML dataset name to ``INLINE_DATASET_NAME``. + """ + + def __init__( + self, + *, + file_path: str | Path, + max_dataset_size: int | None = None, + validators: Sequence[Callable[[ResolvedDataset], None]] | None = None, + ) -> None: + super().__init__(max_dataset_size=max_dataset_size, validators=validators) + self._file_path = Path(file_path) + self._resolved_dataset_name: str | None = None + + @property + def dataset_names(self) -> list[str]: + """The resolved YAML dataset name, or the file stem before first resolution.""" + return [self._resolved_dataset_name or self._file_path.stem] + + @property + def source_kind(self) -> DatasetSourceKind: + """The inline source kind for the local file.""" + return DatasetSourceKind.INLINE + + async def _build_groups_by_dataset_async(self) -> tuple[dict[str, list[AttackSeedGroup]], ResolvedDataset]: + """ + Load and group the current local file contents. + + Returns: + tuple[dict[str, list[AttackSeedGroup]], ResolvedDataset]: Attack groups keyed + by the YAML dataset name and the full resolved seed set. + + Raises: + DatasetConstraintError: If the local file cannot be read, parsed as a + valid ``SeedDataset``, or grouped into valid attack groups. + """ + dataset = await self._load_dataset_async() + dataset_name = dataset.dataset_name or dataset.name or self._file_path.stem + self._resolved_dataset_name = dataset_name + + seeds = cast("list[Seed]", list(dataset.seeds)) + groups = self._build_file_attack_groups(seeds=seeds) + logger.warning( + "Local file-backed dataset '%s' was loaded from disk. Changes made to loaded seeds are not persisted " + "to PyRIT seed memory or written back to the file; save edits to disk before the next resolution or " + "they will be lost.", + self._file_path, + ) + resolved = ResolvedDataset( + seeds=seeds, + source_kind=self.source_kind, + dataset_names=(dataset_name,), + ) + return {dataset_name: groups}, resolved + + def _build_file_attack_groups(self, *, seeds: list[Seed]) -> list[AttackSeedGroup]: + """ + Group file-backed seeds with actionable local-file error context. + + Args: + seeds (list[Seed]): The current parsed file contents. + + Returns: + list[AttackSeedGroup]: The grouped attack seeds. + + Raises: + DatasetConstraintError: If any group lacks exactly one objective. + """ + try: + return self._build_attack_groups(seeds) + except ValueError as exc: + raise DatasetConstraintError( + f"Local dataset file '{self._file_path}' could not be grouped into attack groups; " + "each group needs exactly one 'seed_type: objective': " + f"{exc}" + ) from exc + + async def _load_dataset_async(self) -> SeedDataset: + """ + Read and validate the local YAML dataset off the event loop. + + Returns: + SeedDataset: The current parsed file contents. + + Raises: + DatasetConstraintError: If the file cannot be read or parsed. + """ + try: + return await asyncio.to_thread(SeedDataset.from_yaml_file, self._file_path) + except (OSError, ValueError) as exc: + raise DatasetConstraintError(f"Local dataset file '{self._file_path}' could not be loaded: {exc}") from exc + class CompoundDatasetAttackConfiguration(DatasetAttackConfiguration): """ diff --git a/tests/unit/datasets/test_jailbreak_dataset.py b/tests/unit/datasets/test_jailbreak_dataset.py index 41d239d57f..8bb74bbe92 100644 --- a/tests/unit/datasets/test_jailbreak_dataset.py +++ b/tests/unit/datasets/test_jailbreak_dataset.py @@ -1,6 +1,9 @@ # Copyright (c) Microsoft Corporation. # Licensed under the MIT license. +import logging +from pathlib import Path + import pytest from pyrit.datasets import SeedDatasetProvider @@ -44,6 +47,25 @@ async def test_fetch_dataset_async_loads_templates_from_path(tmp_path): assert all(seed.dataset_name == dataset.dataset_name for seed in dataset.seeds) +async def test_fetch_dataset_async_warns_that_in_memory_edits_are_not_written_to_disk( + *, + tmp_path: Path, + caplog: pytest.LogCaptureFixture, +) -> None: + (tmp_path / "a.yaml").write_text(_VALID_TEMPLATE, encoding="utf-8") + loader = _JailbreakTemplatesDataset(templates_path=tmp_path) + + with caplog.at_level( + logging.WARNING, + logger="pyrit.datasets.seed_datasets.local.jailbreak_dataset", + ): + await loader.fetch_dataset_async() + + assert str(tmp_path) in caplog.text + assert "not written back to disk" in caplog.text + assert "save edits to the source files before reloading or they will be lost" in caplog.text + + async def test_fetch_dataset_async_skips_invalid_templates(tmp_path): (tmp_path / "valid.yaml").write_text(_VALID_TEMPLATE, encoding="utf-8") (tmp_path / "invalid.yaml").write_text("not: [a, valid: seed", encoding="utf-8") diff --git a/tests/unit/datasets/test_local_dataset_loader.py b/tests/unit/datasets/test_local_dataset_loader.py index 6c51b6b877..ab407523f8 100644 --- a/tests/unit/datasets/test_local_dataset_loader.py +++ b/tests/unit/datasets/test_local_dataset_loader.py @@ -1,6 +1,7 @@ # Copyright (c) Microsoft Corporation. # Licensed under the MIT license. +import logging from pathlib import Path import pytest @@ -49,7 +50,38 @@ async def test_fetch_dataset(self, tmp_path, valid_yaml_content): assert len(dataset.prompts) == 1 assert dataset.prompts[0].value == "test prompt" - async def test_fetch_dataset_file_not_found(self): - loader = _LocalDatasetLoader(file_path=Path("non_existent.yaml")) - with pytest.raises(Exception): + async def test_fetch_dataset_warns_that_in_memory_edits_are_not_written_to_disk( + self, + *, + tmp_path: Path, + valid_yaml_content: str, + caplog: pytest.LogCaptureFixture, + ) -> None: + file_path = tmp_path / "test.yaml" + file_path.write_text(valid_yaml_content, encoding="utf-8") + loader = _LocalDatasetLoader(file_path=file_path) + + with caplog.at_level( + logging.WARNING, + logger="pyrit.datasets.seed_datasets.local.local_dataset_loader", + ): await loader.fetch_dataset_async() + + assert str(file_path) in caplog.text + assert "not written back to disk" in caplog.text + assert "save edits to the source file before reloading or they will be lost" in caplog.text + + async def test_fetch_dataset_file_not_found_does_not_log_success_warning( + self, + *, + caplog: pytest.LogCaptureFixture, + ) -> None: + loader = _LocalDatasetLoader(file_path=Path("non_existent.yaml")) + with caplog.at_level( + logging.WARNING, + logger="pyrit.datasets.seed_datasets.local.local_dataset_loader", + ): + with pytest.raises(Exception): + await loader.fetch_dataset_async() + + assert "Local dataset provider loaded" not in caplog.text diff --git a/tests/unit/scenario/airt/test_rapid_response.py b/tests/unit/scenario/airt/test_rapid_response.py index 66259563e0..d884c4100c 100644 --- a/tests/unit/scenario/airt/test_rapid_response.py +++ b/tests/unit/scenario/airt/test_rapid_response.py @@ -19,7 +19,10 @@ from pyrit.registry import TargetRegistry from pyrit.registry.components.attack_technique_registry import AttackTechniqueRegistry from pyrit.scenario.core.attack_technique_factory import AttackTechniqueFactory -from pyrit.scenario.core.dataset_configuration import CompoundDatasetAttackConfiguration +from pyrit.scenario.core.dataset_configuration import ( + CompoundDatasetAttackConfiguration, + DatasetAttackConfiguration, +) from pyrit.scenario.scenarios.airt.rapid_response import RapidResponse from pyrit.score import TrueFalseScorer from pyrit.setup.initializers.techniques import ( @@ -245,6 +248,39 @@ async def test_initialize_raises_when_no_datasets(self, mock_objective_target, m with pytest.raises(ValueError, match="could not be loaded"): await scenario.initialize_async() + async def test_local_file_dataset_drives_identity_and_display_group( + self, + *, + tmp_path: pathlib.Path, + mock_objective_target, + mock_objective_scorer, + ): + file_path = tmp_path / "rapid-response.prompt" + file_path.write_text( + """dataset_name: copilot_local +seeds: + - seed_type: objective + value: Test local iteration +""", + encoding="utf-8", + ) + config = DatasetAttackConfiguration.from_yaml_file(file_path=file_path) + scenario = RapidResponse(objective_scorer=mock_objective_scorer) + scenario.set_params_from_args( + args={ + "objective_target": mock_objective_target, + "dataset_config": config, + "scenario_techniques": [_technique_class()("role_play_movie_script")], + "include_baseline": False, + } + ) + + await scenario.initialize_async() + + stored_result = scenario._memory.get_scenario_results(scenario_result_ids=[scenario._scenario_result_id])[0] + assert stored_result.scenario_identifier.datasets == ["copilot_local"] + assert {attack.display_group for attack in scenario._atomic_attacks} == {"copilot_local"} + @patch("pyrit.scenario.core.scenario.Scenario._get_default_objective_scorer") @patch.object( CompoundDatasetAttackConfiguration, diff --git a/tests/unit/scenario/core/test_dataset_configuration.py b/tests/unit/scenario/core/test_dataset_configuration.py index 14e914b4d5..d6020ae3ac 100644 --- a/tests/unit/scenario/core/test_dataset_configuration.py +++ b/tests/unit/scenario/core/test_dataset_configuration.py @@ -3,6 +3,8 @@ """Tests for the DatasetConfiguration base class and DatasetAttackConfiguration.""" +import logging +from pathlib import Path from unittest.mock import AsyncMock, MagicMock, patch import pytest @@ -76,6 +78,14 @@ def make_objectives(*values: str) -> list[SeedObjective]: return [SeedObjective(value=v) for v in values] +def write_objective_dataset(*, file_path: Path, dataset_name: str, objectives: list[str]) -> None: + """Write a minimal objective dataset for local-file resolution tests.""" + lines = [f"dataset_name: {dataset_name}", "seeds:"] + for objective in objectives: + lines.extend([" - seed_type: objective", f" value: {objective}"]) + file_path.write_text("\n".join(lines) + "\n", encoding="utf-8") + + class TestDatasetConfigurationInit: """Construction, source-exclusivity, and defensive copying.""" @@ -473,6 +483,202 @@ async def test_restrict_dataset_names_passes_for_allowed_dataset(self, mock_memo assert [g.objective.value for g in groups] == ["a"] +class TestLocalFileDatasetAttackConfiguration: + """Resolution of local YAML files without seed-memory synchronization.""" + + def test_rejects_subclass_factory_call(self, *, tmp_path: Path) -> None: + class CustomDatasetAttackConfiguration(DatasetAttackConfiguration): + pass + + with pytest.raises(TypeError, match="would discard subclass behavior"): + CustomDatasetAttackConfiguration.from_yaml_file(file_path=tmp_path / "dataset.prompt") + + async def test_preserves_dataset_name_and_grouping(self, *, tmp_path: Path, mock_memory: MagicMock) -> None: + file_path = tmp_path / "rapid-response.prompt" + file_path.write_text( + """dataset_name: copilot_generated +seeds: + - seed_type: objective + value: Test the generated objective + prompt_group_alias: group_1 + - seed_type: prompt + value: Use this system context + role: system + sequence: 0 + prompt_group_alias: group_1 +""", + encoding="utf-8", + ) + + config = DatasetAttackConfiguration.from_yaml_file(file_path=file_path) + groups_by_dataset = await config.get_attack_groups_by_dataset_async() + + assert list(groups_by_dataset) == ["copilot_generated"] + assert config.dataset_names == ["copilot_generated"] + assert config.source_kind is DatasetSourceKind.INLINE + assert len(groups_by_dataset["copilot_generated"]) == 1 + assert {seed.value for seed in groups_by_dataset["copilot_generated"][0].seeds} == { + "Test the generated objective", + "Use this system context", + } + mock_memory.get_seeds.assert_not_called() + mock_memory.add_seed_datasets_to_memory_async.assert_not_awaited() + + async def test_warns_on_every_resolution_that_disk_is_authoritative( + self, + *, + tmp_path: Path, + caplog: pytest.LogCaptureFixture, + ) -> None: + file_path = tmp_path / "rapid-response.prompt" + write_objective_dataset( + file_path=file_path, + dataset_name="local_iteration", + objectives=["first objective"], + ) + config = DatasetAttackConfiguration.from_yaml_file(file_path=file_path) + + with caplog.at_level(logging.WARNING, logger="pyrit.scenario.core.dataset_configuration"): + await config.get_attack_groups_by_dataset_async() + await config.get_attack_groups_by_dataset_async() + + warnings = [ + record.message + for record in caplog.records + if record.name == "pyrit.scenario.core.dataset_configuration" + and "Local file-backed dataset" in record.message + ] + assert len(warnings) == 2 + assert all(str(file_path) in message for message in warnings) + assert all("not persisted to PyRIT seed memory" in message for message in warnings) + assert all( + "save edits to disk before the next resolution or they will be lost" in message for message in warnings + ) + + async def test_rereads_file_on_every_resolution(self, *, tmp_path: Path) -> None: + file_path = tmp_path / "rapid-response.prompt" + write_objective_dataset( + file_path=file_path, + dataset_name="local_iteration", + objectives=["first objective", "remove me"], + ) + config = DatasetAttackConfiguration.from_yaml_file(file_path=file_path) + + first = await config.get_attack_groups_by_dataset_async() + write_objective_dataset( + file_path=file_path, + dataset_name="local_iteration", + objectives=["updated objective"], + ) + second = await config.get_attack_groups_by_dataset_async() + + assert sorted(group.objective.value for group in first["local_iteration"]) == [ + "first objective", + "remove me", + ] + assert [group.objective.value for group in second["local_iteration"]] == ["updated objective"] + + async def test_file_backed_inline_source_exposes_dataset_name_to_validators(self, *, tmp_path: Path) -> None: + file_path = tmp_path / "rapid-response.prompt" + write_objective_dataset( + file_path=file_path, + dataset_name="local_iteration", + objectives=["first objective"], + ) + seen: list[ResolvedDataset] = [] + config = DatasetAttackConfiguration.from_yaml_file( + file_path=file_path, + validators=[seen.append, restrict_dataset_names({"local_iteration"})], + ) + + await config.get_attack_seed_groups_async() + + assert seen[0].is_inline + assert seen[0].dataset_names == ("local_iteration",) + + async def test_file_backed_dataset_name_is_rejected_by_validator(self, *, tmp_path: Path) -> None: + file_path = tmp_path / "rapid-response.prompt" + write_objective_dataset( + file_path=file_path, + dataset_name="local_iteration", + objectives=["first objective"], + ) + config = DatasetAttackConfiguration.from_yaml_file( + file_path=file_path, + validators=[restrict_dataset_names({"allowed"})], + ) + + with pytest.raises(DatasetConstraintError, match=r"local_iteration.*not allowed"): + await config.get_attack_seed_groups_async() + + async def test_validates_full_file_before_sampling(self, *, tmp_path: Path) -> None: + file_path = tmp_path / "rapid-response.prompt" + write_objective_dataset( + file_path=file_path, + dataset_name="local_iteration", + objectives=["first objective", "second objective"], + ) + config = DatasetAttackConfiguration.from_yaml_file( + file_path=file_path, + max_dataset_size=1, + validators=[require_min_size(2)], + ) + + groups = await config.get_attack_seed_groups_async() + + assert len(groups) == 1 + + async def test_invalid_file_raises_without_persistence_warning( + self, + *, + tmp_path: Path, + caplog: pytest.LogCaptureFixture, + ) -> None: + file_path = tmp_path / "invalid.prompt" + file_path.write_text("dataset_name: invalid\nseeds: []\n", encoding="utf-8") + config = DatasetAttackConfiguration.from_yaml_file(file_path=file_path) + + with caplog.at_level(logging.WARNING, logger="pyrit.scenario.core.dataset_configuration"): + with pytest.raises(DatasetConstraintError, match=r"invalid\.prompt.*could not be loaded") as exc_info: + await config.get_attack_groups_by_dataset_async() + assert exc_info.value.__cause__ is not None + assert "cannot be empty" in str(exc_info.value.__cause__) + assert "Local file-backed dataset" not in caplog.text + + async def test_deleted_file_raises_dataset_constraint_error(self, *, tmp_path: Path) -> None: + file_path = tmp_path / "deleted.prompt" + write_objective_dataset( + file_path=file_path, + dataset_name="local_iteration", + objectives=["first objective"], + ) + config = DatasetAttackConfiguration.from_yaml_file(file_path=str(file_path)) + await config.get_attack_groups_by_dataset_async() + file_path.unlink() + + with pytest.raises(DatasetConstraintError, match=r"deleted\.prompt.*could not be loaded"): + await config.get_attack_groups_by_dataset_async() + + async def test_missing_objective_raises_grouping_error_with_file_context(self, *, tmp_path: Path) -> None: + file_path = tmp_path / "missing-objective.prompt" + file_path.write_text( + """dataset_name: local_iteration +seeds: + - value: This defaults to a prompt instead of an objective +""", + encoding="utf-8", + ) + config = DatasetAttackConfiguration.from_yaml_file(file_path=file_path) + + with pytest.raises( + DatasetConstraintError, + match=r"missing-objective\.prompt.*exactly one 'seed_type: objective'", + ) as exc_info: + await config.get_attack_groups_by_dataset_async() + assert exc_info.value.__cause__ is not None + assert "Found 0" in str(exc_info.value.__cause__) + + class TestCompoundDatasetAttackConfiguration: """``CompoundDatasetAttackConfiguration`` composes child configs with independent budgets."""