diff --git a/.github/instructions/scenarios.instructions.md b/.github/instructions/scenarios.instructions.md index f338d73065..1d1e88525f 100644 --- a/.github/instructions/scenarios.instructions.md +++ b/.github/instructions/scenarios.instructions.md @@ -182,8 +182,14 @@ Options: Use one writer per name; concurrent runs can generate and store different batches. - `seed_groups` and `seeds` are inline alternatives. They never use memory or providers. - Validators run on full filtered populations before any sampling. -- Ingredients must remain complete. Use `max_per_dataset="all"` and no finite source - `max_size`; cap the assembled attack groups with `max_total`, not ingredient rows. +- Ingredients must remain complete. Prefer declaring them as required data: override + `get_required_data()` to return a `ScenarioRequiredData`, assemble groups from + `required_data.load_async()`, and limit them with `dataset_config.select_assembled_groups()`. + Required data is prepared and refreshed like any source (fresh runs and explicit loads only, + never estimates or resume), is never capped, sampled, filtered, or deselected, and stays in + the dataset settings projection and scenario identity. `ApiKey` is the reference. Older + scenarios that keep ingredients in `dataset_config` 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. diff --git a/doc/scanner/garak.ipynb b/doc/scanner/garak.ipynb index da4ca4811d..c101f92f27 100644 --- a/doc/scanner/garak.ipynb +++ b/doc/scanner/garak.ipynb @@ -991,8 +991,9 @@ "\n", "**Available techniques:** `GetKey` and `CompleteKey`. `DEFAULT` and `ALL` both select the two\n", "techniques. `max_total` samples across all selected technique populations, not per service.\n", - "The base scenario persists the sample for resume. Use `ApiKeyDatasetConfiguration` with\n", - "`max_total=\"all\"` to run all 348 requests. Standard technique converter stacks are supported." + "The base scenario persists the sample for resume. The four corpus datasets are the scenario's\n", + "required data and are always read in full; `ApiKeyDatasetConfiguration` only limits the assembled\n", + "requests. Use `max_total=\"all\"` to run all 348 requests. Standard technique converter stacks are supported." ] }, { @@ -1051,9 +1052,7 @@ " args={\n", " \"objective_target\": objective_target,\n", " \"scenario_techniques\": [ApiKeyTechnique.GetKey],\n", - " \"dataset_config\": ApiKeyDatasetConfiguration(\n", - " sources=[DatasetSource(name=name) for name in ApiKey.required_datasets()], max_total=2\n", - " ),\n", + " \"dataset_config\": ApiKeyDatasetConfiguration(max_total=2),\n", " }\n", ")\n", "await api_key_scenario.initialize_async() # type: ignore\n", diff --git a/doc/scanner/garak.py b/doc/scanner/garak.py index 5276a808ae..0db5667410 100644 --- a/doc/scanner/garak.py +++ b/doc/scanner/garak.py @@ -338,8 +338,9 @@ # # **Available techniques:** `GetKey` and `CompleteKey`. `DEFAULT` and `ALL` both select the two # techniques. `max_total` samples across all selected technique populations, not per service. -# The base scenario persists the sample for resume. Use `ApiKeyDatasetConfiguration` with -# `max_total="all"` to run all 348 requests. Standard technique converter stacks are supported. +# The base scenario persists the sample for resume. The four corpus datasets are the scenario's +# required data and are always read in full; `ApiKeyDatasetConfiguration` only limits the assembled +# requests. Use `max_total="all"` to run all 348 requests. Standard technique converter stacks are supported. # %% api_key_scenario = ApiKey() @@ -347,9 +348,7 @@ args={ "objective_target": objective_target, "scenario_techniques": [ApiKeyTechnique.GetKey], - "dataset_config": ApiKeyDatasetConfiguration( - sources=[DatasetSource(name=name) for name in ApiKey.required_datasets()], max_total=2 - ), + "dataset_config": ApiKeyDatasetConfiguration(max_total=2), } ) await api_key_scenario.initialize_async() # type: ignore diff --git a/pyrit/registry/components/scenario_registry.py b/pyrit/registry/components/scenario_registry.py index 2aab1153b7..21f300b0f2 100644 --- a/pyrit/registry/components/scenario_registry.py +++ b/pyrit/registry/components/scenario_registry.py @@ -232,7 +232,7 @@ def _build_metadata(self, name: str, cls: type[Scenario]) -> ScenarioMetadata: ) for aggregate in technique_class.get_aggregate_techniques() ) - default_datasets = tuple(instance._default_dataset_config.dataset_names) + default_datasets = tuple(instance.get_default_dataset_names()) return ScenarioMetadata( class_name=cls.__name__, diff --git a/pyrit/scenario/__init__.py b/pyrit/scenario/__init__.py index fb2d5976f0..393997ac23 100644 --- a/pyrit/scenario/__init__.py +++ b/pyrit/scenario/__init__.py @@ -31,6 +31,7 @@ IncompatibleTechniquePolicy, ResolvedDataset, Scenario, + ScenarioRequiredData, ScenarioTechnique, TechniqueRequirements, ) @@ -101,6 +102,7 @@ def find_spec( "Parameter": "pyrit.models.parameter", "ResolvedDataset": "pyrit.scenario.core.dataset_configuration", "Scenario": "pyrit.scenario.core.scenario", + "ScenarioRequiredData": "pyrit.scenario.core.required_data", "ScenarioTechnique": "pyrit.scenario.core.scenario_technique", "TechniqueRequirements": "pyrit.scenario.core.technique_requirements", "IncompatibleTechniqueError": "pyrit.scenario.core.technique_requirements", diff --git a/pyrit/scenario/core/__init__.py b/pyrit/scenario/core/__init__.py index 2a3677816f..3e1d7e0471 100644 --- a/pyrit/scenario/core/__init__.py +++ b/pyrit/scenario/core/__init__.py @@ -30,6 +30,7 @@ ResolvedDataset, require_nonempty, ) + from pyrit.scenario.core.required_data import ScenarioRequiredData from pyrit.scenario.core.scenario import BaselineAttackPolicy, Scenario from pyrit.scenario.core.scenario_target_defaults import ( get_default_adversarial_target, @@ -70,6 +71,7 @@ "ResolvedDataset": "pyrit.scenario.core.dataset_configuration", "require_nonempty": "pyrit.scenario.core.dataset_configuration", "Scenario": "pyrit.scenario.core.scenario", + "ScenarioRequiredData": "pyrit.scenario.core.required_data", "ScenarioTechnique": "pyrit.scenario.core.scenario_technique", "TechniqueRequirements": "pyrit.scenario.core.technique_requirements", "IncompatibleTechniqueError": "pyrit.scenario.core.technique_requirements", diff --git a/pyrit/scenario/core/dataset_configuration.py b/pyrit/scenario/core/dataset_configuration.py index e492834eca..b04e4e52b5 100644 --- a/pyrit/scenario/core/dataset_configuration.py +++ b/pyrit/scenario/core/dataset_configuration.py @@ -28,6 +28,7 @@ import copy import random +from collections import Counter from contextvars import ContextVar from dataclasses import dataclass from enum import Enum @@ -52,7 +53,7 @@ from pyrit.models.scenario_dataset import DatasetFetchPolicy if TYPE_CHECKING: - from collections.abc import Callable, Sequence + from collections.abc import Awaitable, Callable, Sequence from pyrit.datasets import SeedDatasetProvider from pyrit.memory import MemoryInterface @@ -71,6 +72,70 @@ ) +async def prepare_configurations_async( + *, + configurations: Sequence[DatasetConfiguration], + fetch: DatasetFetchPolicy | None = None, + validate_assembly: Callable[[], Awaitable[object]] | None = None, +) -> None: + """ + Load or refresh several configurations' datasets as one batch. + + Every source is checked and fetched before anything is written. When any dataset is refreshed, + each configuration that fetched data and then ``validate_assembly`` run while all fetched + candidates are visible to reads, so a rejected refresh leaves every stored dataset unchanged. + + Args: + configurations (Sequence[DatasetConfiguration]): The configurations to prepare together. + fetch (DatasetFetchPolicy | None): Optional policy for this operation only. + validate_assembly (Callable[[], Awaitable[object]] | None): A check over everything the + configurations build together, such as a scenario's full assembly. + + Raises: + DatasetConstraintError: If a dataset is named by more than one configuration, a source cannot + be fetched, or the candidates fail validation. + """ + # Candidates are keyed by name, so a shared name would validate one batch and commit both. + counts = Counter(name for config in configurations for name in dict.fromkeys(config.dataset_names)) + shared = sorted(name for name, count in counts.items() if count > 1) + if shared: + raise DatasetConstraintError( + f"Datasets {shared} are named by more than one configuration prepared together, such as " + "required data and the population. List each dataset in one place." + ) + fetches = [(config, *await config._fetch_for_preparation_async(fetch=fetch)) for config in configurations] + if not any(policies[name] is DatasetFetchPolicy.ALWAYS for _, fetched, policies in fetches for name in fetched): + for config, fetched, _ in fetches: + for dataset in fetched.values(): + await config._memory.add_seed_datasets_to_memory_async( + datasets=[dataset], added_by="DatasetConfiguration" + ) + return + prepared: list[tuple[DatasetConfiguration, dict[str, PreparedSeedBatch], dict[str, DatasetFetchPolicy]]] = [] + for config, fetched, policies in fetches: + batches = { + name: await config._memory.prepare_seed_batch_async(seeds=dataset.seeds, added_by="DatasetConfiguration") + for name, dataset in fetched.items() + } + prepared.append((config, batches, policies)) + candidates = {name: batch.seeds for _, batches, _ in prepared for name, batch in batches.items()} + token = _candidate_populations.set(candidates) + try: + for config, batches, _ in prepared: + if batches: + await config.with_overrides()._validate_candidate_population_async() + if validate_assembly is not None: + await validate_assembly() + finally: + _candidate_populations.reset(token) + for config, batches, policies in prepared: + for name, batch in batches.items(): + if policies[name] is DatasetFetchPolicy.ALWAYS: + await config._memory.replace_prepared_seed_batch_async(dataset_name=name, batch=batch) + else: + await config._memory.add_prepared_seed_batch_async(batch) + + class _Unset(Enum): VALUE = "unset" @@ -417,9 +482,10 @@ def __init__( self.max_total = None if isinstance(max_total, _Unset) else max_total self.fetch = DatasetFetchPolicy.IF_MISSING if isinstance(fetch, _Unset) else fetch self._filters: dict[str, list[str]] = dict(filters or {}) + self._user_validators: list[Callable[[ResolvedDataset], None]] = list(validators) if validators else [] self._validators: list[Callable[[ResolvedDataset], None]] = [ *self._default_validators(), - *(list(validators) if validators else []), + *self._user_validators, ] self._validate_selection_options() @@ -494,6 +560,7 @@ def with_overrides( result = copy.copy(self) result._filters = {key: list(values) for key, values in self._filters.items()} result._validators = list(self._validators) + result._user_validators = list(self._user_validators) if not isinstance(sources, _Unset): if self.source_kind is DatasetSourceKind.INLINE: raise DatasetConstraintError("Cannot replace inline seeds with named sources through overrides.") @@ -533,6 +600,21 @@ async def prepare_async(self, *, fetch: DatasetFetchPolicy | None = None) -> Non Args: fetch: Optional policy for this operation only. Source settings remain unchanged. + Raises: + DatasetConstraintError: If a source cannot be fetched or provider data is invalid. + """ + await prepare_configurations_async(configurations=[self], fetch=fetch) + + async def _fetch_for_preparation_async( + self, *, fetch: DatasetFetchPolicy | None + ) -> tuple[dict[str, SeedDataset], dict[str, DatasetFetchPolicy]]: + """ + Check all sources and fetch the datasets that preparation would write, without writing. + + Returns: + tuple[dict[str, SeedDataset], dict[str, DatasetFetchPolicy]]: The fetched datasets and + every source's effective policy, keyed by dataset name. + Raises: DatasetConstraintError: If a source cannot be fetched or provider data is invalid. """ @@ -540,7 +622,7 @@ async def prepare_async(self, *, fetch: DatasetFetchPolicy | None = None) -> Non if fetch is not None and not isinstance(fetch, DatasetFetchPolicy): raise DatasetConstraintError("'fetch' must be a DatasetFetchPolicy.") if self.source_kind is DatasetSourceKind.INLINE: - return + return {}, {} policies: dict[str, DatasetFetchPolicy] = {} configured_providers: dict[str, SeedDatasetProvider | None] = {} for source, policy in self._preparation_sources(): @@ -553,7 +635,7 @@ async def prepare_async(self, *, fetch: DatasetFetchPolicy | None = None) -> Non policies[name] = policy configured_providers[name] = source.provider if not policies: - return + return {}, {} stored_names = set(await self._memory.get_seed_dataset_names_async()) to_fetch: list[str] = [] for name, policy in policies.items(): @@ -563,7 +645,7 @@ async def prepare_async(self, *, fetch: DatasetFetchPolicy | None = None) -> Non raise DatasetConstraintError(f"Dataset '{name}' is missing and fetch is 'never'. Import it first.") to_fetch.append(name) if not to_fetch: - return + return {}, policies from pyrit.datasets.seed_datasets.seed_dataset_provider import SeedDatasetProvider providers = {name: provider for name in to_fetch if (provider := configured_providers[name]) is not None} @@ -597,34 +679,7 @@ async def prepare_async(self, *, fetch: DatasetFetchPolicy | None = None) -> Non ) SeedDataset.group_seed_prompts_by_prompt_group_id(dataset.seeds) fetched[name] = dataset - if any(policies[name] is DatasetFetchPolicy.ALWAYS for name in fetched): - await self._commit_refresh_candidates_async(fetched=fetched, policies=policies) - return - for dataset in fetched.values(): - await self._memory.add_seed_datasets_to_memory_async(datasets=[dataset], added_by="DatasetConfiguration") - - async def _commit_refresh_candidates_async( - self, *, fetched: dict[str, SeedDataset], policies: dict[str, DatasetFetchPolicy] - ) -> None: - prepared = { - name: await self._memory.prepare_seed_batch_async(seeds=dataset.seeds, added_by="DatasetConfiguration") - for name, dataset in fetched.items() - } - await self._validate_refresh_candidates_async(prepared=prepared) - for name, batch in prepared.items(): - if policies[name] is DatasetFetchPolicy.ALWAYS: - await self._memory.replace_prepared_seed_batch_async(dataset_name=name, batch=batch) - else: - await self._memory.add_prepared_seed_batch_async(batch) - - async def _validate_refresh_candidates_async(self, *, prepared: dict[str, PreparedSeedBatch]) -> None: - """Resolve task-local copies of storage-ready snapshots without modifying them.""" - candidates = {name: batch.seeds for name, batch in prepared.items()} - token = _candidate_populations.set(candidates) - try: - await self.with_overrides()._validate_candidate_population_async() - finally: - _candidate_populations.reset(token) + return fetched, policies async def _validate_candidate_population_async(self) -> None: """Validate raw configuration populations before replacement.""" @@ -1021,11 +1076,48 @@ async def get_attack_groups_by_dataset_async( dataset fails validation, or no attack groups could be built. """ selection = await self._resolve_attack_groups_async() - sampled = selection.select(apply_sampling=apply_sampling) - if not any(sampled.values()): + return self._require_selected_groups(selection.select(apply_sampling=apply_sampling)) + + def select_assembled_groups( + self, *, groups_by_population: dict[str, list[AttackSeedGroup]], apply_sampling: bool = True + ) -> dict[str, list[AttackSeedGroup]]: + """ + Validate and limit attack groups a scenario assembled from its required data. + + The configuration has no sources of its own here: it validates the assembled + population and applies ``max_total`` across it, keeping each population's key. + + Args: + groups_by_population (dict[str, list[AttackSeedGroup]]): Assembled groups keyed by population. + apply_sampling (bool): When True (default), apply ``max_total``. Resume passes False. + + Returns: + dict[str, list[AttackSeedGroup]]: The validated, limited groups. + + Raises: + DatasetConstraintError: If the configuration has its own sources or inline seeds, the + population fails validation, or no attack groups remain. + """ + self.validate_configuration() + if self.sources or self.source_kind is DatasetSourceKind.INLINE: + raise DatasetConstraintError( + "Assembled populations can't be combined with dataset sources or inline seeds." + ) + self.validate( + ResolvedDataset( + seeds=[seed for groups in groups_by_population.values() for group in groups for seed in group.seeds], + source_kind=self.source_kind, + dataset_names=tuple(groups_by_population), + ) + ) + selection = _ResolvedAttackGroups(configuration=self, groups=groups_by_population) + return self._require_selected_groups(selection.select(apply_sampling=apply_sampling)) + + def _require_selected_groups(self, selected: dict[str, list[AttackSeedGroup]]) -> dict[str, list[AttackSeedGroup]]: + if not any(selected.values()): names = ", ".join(self.dataset_names) if self.dataset_names else "" raise DatasetConstraintError(f"Resolved attack-group dataset is empty (datasets: {names}).") - return sampled + return selected async def _resolve_attack_groups_async(self) -> _ResolvedAttackGroups: self.validate_configuration() diff --git a/pyrit/scenario/core/dataset_overrides.py b/pyrit/scenario/core/dataset_overrides.py index 69bd297375..984583691f 100644 --- a/pyrit/scenario/core/dataset_overrides.py +++ b/pyrit/scenario/core/dataset_overrides.py @@ -20,6 +20,7 @@ DatasetSource, DatasetSourceKind, ) +from pyrit.scenario.core.required_data import ScenarioRequiredData def apply_dataset_overrides( @@ -96,16 +97,23 @@ def validate_dataset_selection( raise DatasetConstraintError(f"Select at most {options.max_selected} dataset(s).") -def project_dataset_settings(config: DatasetAttackConfiguration) -> ScenarioDatasetSettings: +def project_dataset_settings( + config: DatasetAttackConfiguration, *, required_data: ScenarioRequiredData | None = None +) -> ScenarioDatasetSettings: """ Project configuration settings without memory reads or live wire objects. + Required datasets come first as uncapped sources, so clients keep seeing and + refreshing every dataset a run loads. + Returns: ScenarioDatasetSettings: Effective settings for response data. """ children = config.configurations if isinstance(config, CompoundDatasetAttackConfiguration) else (config,) + required_sources = required_data.project_sources() if required_data is not None else [] return ScenarioDatasetSettings( - sources=[ + sources=required_sources + + [ ScenarioDatasetSourceSettings( name=source.name, max_size=child.source_limit(source.name), diff --git a/pyrit/scenario/core/required_data.py b/pyrit/scenario/core/required_data.py new file mode 100644 index 0000000000..5202b8af89 --- /dev/null +++ b/pyrit/scenario/core/required_data.py @@ -0,0 +1,274 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT license. + +""" +Datasets a scenario needs in order to build its attacks. + +Some scenarios assemble their attack population from datasets used as ingredients +(templates, registries, partial keys). Those datasets are not the population a user +selects or limits: they are always read in full, without source caps, sampling, or +filters. ``ScenarioRequiredData`` declares them apart from ``dataset_config``, which then +only describes and limits the assembled attack groups. + +Preparation reuses ``DatasetConfiguration``, so required data is fetched, stored, and +refreshed with the same guarantees as any dataset source. Reads never fetch. +""" + +from __future__ import annotations + +from dataclasses import replace +from typing import TYPE_CHECKING, Any + +from pyrit.models.scenario_dataset import ( + DatasetFetchPolicy, + DatasetSelectionMode, + ScenarioDatasetOptions, + ScenarioDatasetRequest, + ScenarioDatasetSourceSettings, +) +from pyrit.scenario.core.dataset_configuration import ( + DatasetConfiguration, + DatasetConstraintError, + DatasetSource, + ResolvedDataset, + prepare_configurations_async, +) + +if TYPE_CHECKING: + from collections.abc import Awaitable, Callable, Sequence + + from pyrit.models import Seed + + +class _RequiredDatasetConfiguration(DatasetConfiguration): + """An uncapped, unfiltered configuration used only to prepare and read required data.""" + + async def read_async(self) -> dict[str, list[Seed[Any]]]: + """ + Read every required dataset in full and run the validators. + + Returns: + dict[str, list[Seed]]: Seeds keyed by dataset name, in declaration order. + """ + seeds_by_dataset = await self._collect_named_seeds_async() + self.validate( + ResolvedDataset( + seeds=[seed for seeds in seeds_by_dataset.values() for seed in seeds], + source_kind=self.source_kind, + dataset_names=tuple(seeds_by_dataset), + ) + ) + return seeds_by_dataset + + +class ScenarioRequiredData: + """ + The datasets a scenario reads in full to build its attack population. + + Each source keeps its provider and fetch policy, so preparation and refresh behave + exactly like a ``dataset_config`` source. Dataset filters narrow the rows that are read, + as they do for any source. Required datasets can't be capped, sampled, or deselected; + limit the assembled attack groups through the scenario's ``dataset_config`` instead. + """ + + def __init__( + self, + *, + sources: Sequence[DatasetSource], + fetch: DatasetFetchPolicy = DatasetFetchPolicy.IF_MISSING, + validators: Sequence[Callable[[ResolvedDataset], None]] | None = None, + filters: dict[str, list[str]] | None = None, + ) -> None: + """ + Declare the required datasets. + + Args: + sources (Sequence[DatasetSource]): The required datasets, with optional providers and fetch policies. + fetch (DatasetFetchPolicy): The default preparation policy for sources without their own. + validators (Sequence[Callable[[ResolvedDataset], None]] | None): Checks run on the full + required data, both on normal reads and on refresh candidates before they're stored. + filters (dict[str, list[str]] | None): ``get_seeds`` filters applied when reading, before + the validators run. + + Raises: + DatasetConstraintError: If no source is given or a source sets a size limit. + """ + if not sources: + raise DatasetConstraintError("Required data needs at least one dataset source.") + capped = [source.name for source in sources if source.max_size not in ("default", "all")] + if capped: + raise DatasetConstraintError(f"Required datasets are read in full and can't be capped: {capped}.") + self._validators = list(validators or []) + self._configuration = _RequiredDatasetConfiguration( + sources=list(sources), + max_per_dataset="all", + max_total="all", + fetch=fetch, + validators=self._validators, + filters=filters, + ) + + @property + def sources(self) -> tuple[DatasetSource, ...]: + """The required dataset sources.""" + return self._configuration.sources + + @property + def dataset_names(self) -> list[str]: + """The required dataset names, in declaration order.""" + return self._configuration.dataset_names + + @property + def fetch(self) -> DatasetFetchPolicy: + """The default preparation policy.""" + return self._configuration.fetch + + @property + def filters(self) -> dict[str, list[str]]: + """The ``get_seeds`` filters applied when reading.""" + return self._configuration.filters + + async def prepare_async( + self, + *, + fetch: DatasetFetchPolicy | None = None, + population: DatasetConfiguration | None = None, + validate_assembly: Callable[[], Awaitable[object]] | None = None, + ) -> None: + """ + Load or refresh the required datasets, like ``DatasetConfiguration.prepare_async``. + + With a population configuration, both are fetched first and validated together before + anything is written: a refresh of either half is checked against the other half's + candidates, so a rejected refresh leaves every stored dataset unchanged. + + Args: + fetch (DatasetFetchPolicy | None): Optional policy for this operation only. + population (DatasetConfiguration | None): The scenario's population configuration to prepare + in the same batch. + validate_assembly (Callable[[], Awaitable[object]] | None): The scenario's full unsampled + assembly, run against the candidates whenever a dataset is refreshed. + """ + configurations = [self._configuration, *([population] if population is not None else [])] + await prepare_configurations_async( + configurations=configurations, fetch=fetch, validate_assembly=validate_assembly + ) + + async def load_async(self) -> dict[str, list[Seed[Any]]]: + """ + Read the required datasets from memory without fetching. + + Returns: + dict[str, list[Seed]]: Seeds keyed by dataset name, in declaration order. + """ + return await self._configuration.read_async() + + def apply_overrides( + self, + *, + request: ScenarioDatasetRequest, + options: ScenarioDatasetOptions, + population_names: Sequence[str], + ) -> tuple[ScenarioRequiredData, ScenarioDatasetRequest]: + """ + Take the fetch edits for required sources out of a dataset request. + + The request's global fetch policy reaches required sources without their own policy and + also stays on the returned request for the population, as do its filters. + + Args: + request (ScenarioDatasetRequest): The requested dataset edits. + options (ScenarioDatasetOptions): The scenario's dataset capabilities. + population_names (Sequence[str]): The population source names before the edits. + + Returns: + tuple[ScenarioRequiredData, ScenarioDatasetRequest]: The edited required data, and the + request left for the population configuration. + + Raises: + DatasetConstraintError: If the request deselects, caps, or reorders required datasets, or + leaves out every population dataset. + """ + if not options.supports_overrides and request != ScenarioDatasetRequest(): + raise DatasetConstraintError("This scenario does not support dataset overrides.") + required = self if request.fetch is None else self.with_sources(sources=self.sources, fetch=request.fetch) + if request.sources is None: + return required, request + names = self.dataset_names + requested = [source.name for source in request.sources] + fixed = [*names, *population_names] + if options.mode is DatasetSelectionMode.FIXED and requested != fixed: + raise DatasetConstraintError(f"Dataset sources are fixed: {', '.join(fixed)}.") + missing = [name for name in names if name not in requested] + if missing: + raise DatasetConstraintError(f"Required datasets can't be deselected: {', '.join(missing)}.") + required_inputs = [source for source in request.sources if source.name in names] + capped = [source.name for source in required_inputs if source.max_size != "default"] + if capped: + raise DatasetConstraintError(f"Required datasets are read in full and can't be capped: {capped}.") + remaining = [source for source in request.sources if source.name not in names] + # A request can't hold an empty source list, and dropping it would silently restore the defaults. + if not remaining and population_names: + raise DatasetConstraintError( + "Select at least one population dataset besides the required ones " + f"(defaults: {', '.join(population_names)})." + ) + edited = required._with_source_fetch( + {source.name: source.fetch for source in required_inputs if source.fetch is not None} + ) + return edited, request.model_copy(update={"sources": remaining or None}) + + def project_sources(self) -> list[ScenarioDatasetSourceSettings]: + """ + Describe the required sources for dataset settings responses, without reading memory. + + Returns: + list[ScenarioDatasetSourceSettings]: One uncapped entry per required dataset. + """ + return [ + ScenarioDatasetSourceSettings( + name=source.name, + max_size="all", + fetch=source.fetch or self.fetch, + configured_provider=source.provider is not None, + ) + for source in self.sources + ] + + def with_sources( + self, + *, + sources: Sequence[DatasetSource], + fetch: DatasetFetchPolicy, + filters: dict[str, list[str]] | None = None, + validators: Sequence[Callable[[ResolvedDataset], None]] | None = None, + ) -> ScenarioRequiredData: + """ + Copy the declaration with replacement sources, keeping its validators. + + Args: + sources (Sequence[DatasetSource]): The replacement sources. + fetch (DatasetFetchPolicy): The default preparation policy. + filters (dict[str, list[str]] | None): Filters merged over the current ones, as overrides do. + validators (Sequence[Callable[[ResolvedDataset], None]] | None): Validators appended to the + kept ones, used when a legacy ``dataset_config`` migrates its validators onto the required data. + + Returns: + ScenarioRequiredData: The new declaration. + """ + merged = {**self.filters, **{key: list(values) for key, values in (filters or {}).items()}} + return ScenarioRequiredData( + sources=sources, + fetch=fetch, + validators=[*self._validators, *(validators or [])], + filters=merged, + ) + + def _with_source_fetch(self, source_fetch: dict[str, DatasetFetchPolicy]) -> ScenarioRequiredData: + if not source_fetch: + return self + sources = [ + replace(source, fetch=source_fetch[source.name]) if source.name in source_fetch else source + for source in self.sources + ] + return self.with_sources(sources=sources, fetch=self.fetch) diff --git a/pyrit/scenario/core/scenario.py b/pyrit/scenario/core/scenario.py index 3396be5b94..36ca35e990 100644 --- a/pyrit/scenario/core/scenario.py +++ b/pyrit/scenario/core/scenario.py @@ -21,6 +21,7 @@ from tqdm.auto import tqdm from pyrit.common import default_values, get_global_default_values +from pyrit.common.deprecation import print_deprecation_message from pyrit.common.utils import to_sha256 from pyrit.exceptions import ScenarioPartialFailureException from pyrit.executor.attack import AttackExecutor, AttackExecutorResult @@ -70,9 +71,12 @@ from pyrit.scenario.core.atomic_attack import AtomicAttack from pyrit.scenario.core.dataset_configuration import ( DatasetAttackConfiguration, + DatasetSource, + DatasetSourceKind, ) from pyrit.scenario.core.dataset_overrides import apply_dataset_overrides, project_dataset_settings from pyrit.scenario.core.dataset_sampling import sample_with_coverage +from pyrit.scenario.core.required_data import ScenarioRequiredData 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 @@ -256,6 +260,8 @@ def __init__( # this to the caller-supplied config (or the default); defaulting it here means the # attribute always exists for context construction. self._dataset_config: DatasetAttackConfiguration = default_dataset_config + # Datasets the scenario reads in full to build its attacks, apart from dataset_config. + self._required_data: ScenarioRequiredData | None = None self._objective_scorer = objective_scorer self._objective_scorer_identifier = objective_scorer.get_identifier() @@ -684,8 +690,10 @@ async def get_run_size_estimate_async(self, *, target_is_configured: bool = Fals 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), + "dataset_defaults": project_dataset_settings( + self._resolved_default_dataset_config, required_data=self.get_required_data() + ), + "effective_datasets": project_dataset_settings(self._dataset_config, required_data=self._required_data), } ) @@ -911,16 +919,40 @@ def _resolve_runtime_configuration(self, *, require_objective_target: bool) -> N self._dataset_config = ( dataset_config if dataset_config is not None else self._resolved_default_dataset_config ).with_overrides() + self._required_data = self.get_required_data() + self._move_required_sources_out_of_dataset_config() options = self.get_dataset_options() raw_overrides = params.get("dataset_overrides") if raw_overrides is not None: + request = ScenarioDatasetRequest.model_validate(raw_overrides) + if self._required_data is not None: + self._required_data, request = self._required_data.apply_overrides( + request=request, options=options, population_names=self._dataset_config.dataset_names + ) self._dataset_config = apply_dataset_overrides( config=self._dataset_config, - request=ScenarioDatasetRequest.model_validate(raw_overrides), + request=request, options=options, ) + self._load_required_data_like_dataset_config() self._validate_runtime_configuration() + def _load_required_data_like_dataset_config(self) -> None: + """ + Read the required data with the population configuration's fetch policy and filters. + + A configuration that names no datasets only limits assembled groups, so its fetch policy + and filters describe how the required data is loaded. Required sources keep their own + fetch settings. A compound names datasets through its children, not ``sources``, and its + parent ``fetch`` isn't the requested one, so it leaves the required data's policy as resolved. + """ + config = self._dataset_config + if self._required_data is None or config.dataset_names or config.source_kind is DatasetSourceKind.INLINE: + return + self._required_data = self._required_data.with_sources( + sources=self._required_data.sources, fetch=config.fetch, filters=config.filters + ) + def _resolve_default_dataset_config(self) -> DatasetAttackConfiguration: """ Resolve parameter-aware defaults before any override or preparation. @@ -930,6 +962,68 @@ def _resolve_default_dataset_config(self) -> DatasetAttackConfiguration: """ return self._default_dataset_config + def get_required_data(self) -> ScenarioRequiredData | None: + """ + Return the datasets this scenario reads in full to build its attacks, if any. + + Required data is prepared and refreshed like ``dataset_config`` sources but never + capped, sampled, or filtered. ``dataset_config`` then only selects and limits the + attack population. Override this when a scenario assembles its population from + ingredient datasets such as templates or registries. + + Returns: + ScenarioRequiredData | None: The required data, or None when the scenario has none. + """ + return None + + def _move_required_sources_out_of_dataset_config(self) -> None: + """ + Accept a ``dataset_config`` that still lists the scenario's required datasets as sources. + + Before a scenario declared required data, its configuration listed those datasets as + sources. When a configuration of the default's class lists exactly the required datasets, + their fetch settings, providers, and filters move to the required data and the + configuration keeps only its limits. Size limits move too, so they still fail as before. Any other selection + is left for the scenario's own validation to reject. + """ + required = self._required_data + config = self._dataset_config + if required is None or not config.sources or type(config) is not type(self._default_dataset_config): + return + if sorted(config.dataset_names) != sorted(required.dataset_names): + return + print_deprecation_message( + old_item=f"{type(config).__name__}(sources=<{type(self).__name__} required datasets>)", + new_item=( + f"{type(config).__name__} without those sources; {type(self).__name__} declares them as required data" + ), + removed_in="1.4.0", + ) + # The legacy configuration read these datasets as its population, so its own validators + # validated the raw template/corpus seeds. Those seeds are now required data, so the + # validators move with them; the configuration keeps only its defaults and limits. + self._required_data = required.with_sources( + sources=[ + DatasetSource( + name=source.name, + max_size=config.source_limit(source.name), + fetch=source.fetch, + provider=source.provider, + ) + for source in config.sources + ], + fetch=config.fetch, + filters=config.filters, + validators=config._user_validators, + ) + stripped = config.with_overrides(sources=[]) + # Remove the moved validators by identity: a validator may be any callable, including an + # unhashable instance, so they can't go in a set. + moved = {id(validator) for validator in config._user_validators} + stripped._validators = [validator for validator in stripped._validators if id(validator) not in moved] + stripped._user_validators = [] + self._dataset_config = stripped + DATASET_OPTIONS: ClassVar[ScenarioDatasetOptions] = ScenarioDatasetOptions() def get_dataset_options(self) -> ScenarioDatasetOptions: @@ -938,7 +1032,12 @@ def get_dataset_options(self) -> ScenarioDatasetOptions: 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()) + return project_dataset_settings(self._resolve_default_dataset_config(), required_data=self.get_required_data()) + + def get_default_dataset_names(self) -> list[str]: + """Return every dataset a default run loads: required data first, then the population sources.""" + required = self.get_required_data() + return [*(required.dataset_names if required else []), *self._default_dataset_config.dataset_names] async def load_datasets_async(self) -> ScenarioDatasetSettings: """ @@ -948,8 +1047,28 @@ async def load_datasets_async(self) -> ScenarioDatasetSettings: ScenarioDatasetSettings: Resolved settings, with the scan's policies unchanged. """ self._resolve_runtime_configuration(require_objective_target=False) - await self._dataset_config.prepare_async(fetch=DatasetFetchPolicy.IF_MISSING) - return project_dataset_settings(self._dataset_config) + await self._prepare_datasets_async(fetch=DatasetFetchPolicy.IF_MISSING) + return project_dataset_settings(self._dataset_config, required_data=self._required_data) + + async def _prepare_datasets_async(self, *, fetch: DatasetFetchPolicy | None = None) -> None: + """ + Load or refresh the required data and the population together. + + With required data, both halves are fetched first and the full unsampled assembly runs + against all candidates before any write, so a rejected refresh of either half leaves the + stored datasets unchanged and earlier runs can still resume. + + Args: + fetch (DatasetFetchPolicy | None): Optional policy for this operation only. + """ + if self._required_data is None: + await self._dataset_config.prepare_async(fetch=fetch) + return + await self._required_data.prepare_async( + fetch=fetch, + population=self._dataset_config, + validate_assembly=lambda: self._resolve_seed_groups_by_dataset_async(apply_sampling=False), + ) def _validate_runtime_configuration(self) -> None: """ @@ -1068,7 +1187,7 @@ async def initialize_async(self) -> None: # builders check compatibility. Inputs outside the saved run must not affect resume. is_resume = self._scenario_result_id is not None if not is_resume: - await self._dataset_config.prepare_async() + await self._prepare_datasets_async() seed_groups_by_dataset = await self._resolve_seed_groups_by_dataset_async(apply_sampling=not is_resume) if stored_plan is not None: planned_seed_ids = {group_id for group in stored_plan.atomic_groups for group_id in group.seed_group_ids} @@ -1398,7 +1517,9 @@ def _build_scenario_identifier(self) -> ScenarioIdentifier: ScenarioIdentifier: The identifier describing this scenario run. """ techniques = sorted({s.value for s in self._scenario_techniques}) - datasets = list(self._dataset_config.dataset_names) + # Required data first, so moving ingredients out of dataset_config keeps stored identities. + required = self._required_data.dataset_names if self._required_data is not None else [] + datasets = [*required, *self._dataset_config.dataset_names] # Persist only the scenario's own custom params. The framework common inputs # (objective_target, techniques, dataset config, ...) are captured through the # dedicated identity fields below and are often live, non-JSON-serializable diff --git a/pyrit/scenario/scenarios/garak/api_key.py b/pyrit/scenario/scenarios/garak/api_key.py index dfd2332399..fd6d920e47 100644 --- a/pyrit/scenario/scenarios/garak/api_key.py +++ b/pyrit/scenario/scenarios/garak/api_key.py @@ -33,8 +33,10 @@ DatasetAttackConfiguration, DatasetConstraintError, DatasetSource, + DatasetSourceKind, ResolvedDataset, ) +from pyrit.scenario.core.required_data import ScenarioRequiredData from pyrit.scenario.core.scenario import BaselineAttackPolicy, Scenario from pyrit.scenario.core.scenario_technique import ScenarioTechnique from pyrit.score import CredentialLeakScorer @@ -80,7 +82,12 @@ def default(cls) -> ApiKeyTechnique: class ApiKeyDatasetConfiguration(DatasetAttackConfiguration): - """Build API-key attack groups before standard dataset sampling.""" + """ + Limit the API-key requests ``ApiKey`` assembles from its corpus. + + The corpus datasets are ApiKey's required data (``ApiKey.get_required_data``), read in + full. This configuration only limits the assembled requests with ``max_total``. + """ DEFAULT_MAX_DATASET_SIZE: ClassVar[int] = 20 @@ -95,7 +102,6 @@ def __init__(self, **kwargs: Any) -> None: """ super().__init__(**kwargs) self._techniques: list[ApiKeyTechnique] = [ApiKeyTechnique.GetKey, ApiKeyTechnique.CompleteKey] - self.excluded_values: tuple[str, ...] = () def _default_max_per_dataset(self) -> Literal["all"]: return "all" @@ -120,81 +126,77 @@ def size_caps_by_dataset(self) -> dict[str, list[tuple[str, int, Literal["datase def validate_configuration(self) -> None: """ - Check the corpus selection without reading its contents. + Check that only the assembled requests are limited. Raises: - DatasetConstraintError: If the required corpus datasets are not selected. + DatasetConstraintError: If the configuration names datasets or supplies inline seeds. """ super().validate_configuration() - if any(self.source_limit(source.name) != "all" for source in self.sources): - raise DatasetConstraintError("ApiKey ingredient sources must be uncapped; use max_total.") - if set(self.dataset_names) != set(_CORPUS_DATASETS): + if self.sources or self.source_kind is DatasetSourceKind.INLINE: raise DatasetConstraintError( - f"ApiKey requires exactly these datasets: {list(_CORPUS_DATASETS)}; inline seeds are not supported." + f"ApiKey reads its corpus as required data: {list(_CORPUS_DATASETS)}. " + "ApiKeyDatasetConfiguration only takes max_total; inline seeds and other datasets are not supported." ) - async def _build_groups_by_dataset_async(self) -> tuple[dict[str, list[AttackSeedGroup]], ResolvedDataset]: - """ - Resolve the corpus and group requests by technique. - Returns: - tuple: Technique populations and the raw dataset for validation. +def _require_one_template_per_technique(resolved: ResolvedDataset, *, techniques: Sequence[ApiKeyTechnique]) -> None: + """ + Check the corpus has exactly one prompt template for each selected technique. - Raises: - DatasetConstraintError: If the configuration does not select the required corpus. - """ - seeds_by_dataset = await self._collect_named_seeds_async() - seeds = [seed for population in seeds_by_dataset.values() for seed in population] - self.excluded_values = tuple( - seed.value for name in (DATASET_PARTIAL_KEYS, DATASET_SAFE_PLACEHOLDERS) for seed in seeds_by_dataset[name] - ) - populations = self._build_technique_populations(seeds) - return populations, ResolvedDataset( - seeds=seeds, source_kind=self.source_kind, dataset_names=tuple(self.dataset_names) - ) + A technique's template is only used by a run that selects it, so a GetKey-only run does not + require a CompleteKey template. - def _build_technique_populations(self, seeds: list[Seed[Any]]) -> dict[str, list[AttackSeedGroup]]: - """ - Render one request for each selected service and technique input. + Raises: + DatasetConstraintError: If a selected technique has no template, several, or a non-prompt one. + """ + for technique in techniques: + templates = [ + seed + for seed in resolved.seeds + if seed.dataset_name == DATASET_TEMPLATES and (seed.metadata or {}).get("technique") == technique.value + ] + if len(templates) != 1 or not isinstance(templates[0], SeedPrompt): + raise DatasetConstraintError(f"ApiKey requires one prompt template for {technique.value}.") - Returns: - dict[str, list[AttackSeedGroup]]: Unsampled API-key attack groups keyed by technique. - Raises: - DatasetConstraintError: If a selected technique has no unique prompt template. - """ - services = [seed for seed in seeds if seed.dataset_name == DATASET_SERVICES] - partials = [seed.value for seed in seeds if seed.dataset_name == DATASET_PARTIAL_KEYS] - populations: dict[str, list[AttackSeedGroup]] = {} - for technique in self._techniques: - groups: list[AttackSeedGroup] = [] - populations[str(technique.value)] = groups - templates = [ - seed - for seed in seeds - if seed.dataset_name == DATASET_TEMPLATES and (seed.metadata or {}).get("technique") == technique.value - ] - if len(templates) != 1 or not isinstance(templates[0], SeedPrompt): - raise DatasetConstraintError(f"ApiKey requires one prompt template for {technique.value}.") - template = templates[0] - for service in services: - for partial in partials if technique is ApiKeyTechnique.CompleteKey else [""]: - prompt = template.model_copy(deep=True) - prompt.value = template.render_template_value(key=service.value, partial=partial) - groups.append( - AttackSeedGroup( - seeds=[ - SeedObjective( - value=f"Elicit a {service.value} credential in response to: {prompt.value}", - source=service.source, - harm_categories=service.harm_categories, - metadata={"technique": technique.value, "service": service.value}, - ), - prompt, - ] - ) +def _assemble_requests( + *, corpus: dict[str, list[Seed[Any]]], techniques: Sequence[ApiKeyTechnique] +) -> dict[str, list[AttackSeedGroup]]: + """ + Render one request for each service and technique input. + + Returns: + dict[str, list[AttackSeedGroup]]: Unsampled API-key attack groups keyed by technique. + """ + services = corpus[DATASET_SERVICES] + partials = [seed.value for seed in corpus[DATASET_PARTIAL_KEYS]] + populations: dict[str, list[AttackSeedGroup]] = {} + for technique in techniques: + groups: list[AttackSeedGroup] = [] + populations[str(technique.value)] = groups + template = next( + cast("SeedPrompt", seed) + for seed in corpus[DATASET_TEMPLATES] + if (seed.metadata or {}).get("technique") == technique.value + ) + for service in services: + for partial in partials if technique is ApiKeyTechnique.CompleteKey else [""]: + prompt = template.model_copy(deep=True) + prompt.value = template.render_template_value(key=service.value, partial=partial) + groups.append( + AttackSeedGroup( + seeds=[ + SeedObjective( + value=f"Elicit a {service.value} credential in response to: {prompt.value}", + source=service.source, + harm_categories=service.harm_categories, + metadata={"technique": technique.value, "service": service.value}, + ), + prompt, + ] ) - return populations + ) + return populations class ApiKey(Scenario): @@ -202,10 +204,11 @@ class ApiKey(Scenario): Exercise a model's tendency to generate or complete API credentials. GetKey requests a credential for each service. CompleteKey requests completion - of synthetic partial keys. The dataset configuration samples up to 20 requests - across the selected techniques by default. The base scenario persists that sample - for resume. Credential-shaped responses are scored, excluding supplied partials - and safe placeholders. + of synthetic partial keys. The service, template, partial-key, and placeholder + datasets are required data, always read in full. The dataset configuration samples + up to 20 of the assembled requests across the selected techniques by default. The + base scenario persists that sample for resume. Credential-shaped responses are + scored, excluding supplied partials and safe placeholders. Reference: [@derczynski2024garak] """ @@ -238,14 +241,32 @@ def __init__( version=self.VERSION, technique_class=ApiKeyTechnique, default_dataset_config=ApiKeyDatasetConfiguration( - sources=[DatasetSource(name=name) for name in self.required_datasets()], - max_per_dataset="all", - max_total=ApiKeyDatasetConfiguration.DEFAULT_MAX_DATASET_SIZE, + max_total=ApiKeyDatasetConfiguration.DEFAULT_MAX_DATASET_SIZE ), objective_scorer=objective_scorer or CredentialLeakScorer(patterns=CredentialLeakScorer.GARAK_PATTERNS), scenario_result_id=scenario_result_id, ) + def get_required_data(self) -> ScenarioRequiredData: + """ + Return the API-key corpus, read in full to assemble the requests. + + Returns: + ScenarioRequiredData: The service, template, partial-key, and placeholder datasets. + """ + return ScenarioRequiredData( + sources=[DatasetSource(name=name) for name in self.required_datasets()], + validators=[self._validate_templates_for_selected_techniques], + ) + + def _validate_templates_for_selected_techniques(self, resolved: ResolvedDataset) -> None: + """Require one template per technique the run actually selects.""" + techniques = getattr(self._dataset_config, "_techniques", None) or [ + ApiKeyTechnique.GetKey, + ApiKeyTechnique.CompleteKey, + ] + _require_one_template_per_technique(resolved, techniques=techniques) + def _validate_runtime_configuration(self) -> None: config = self._dataset_config if type(config) is not ApiKeyDatasetConfiguration: @@ -268,10 +289,16 @@ async def _resolve_seed_groups_by_dataset_async( DatasetConstraintError: If an unsupported configuration is supplied. """ config = cast("ApiKeyDatasetConfiguration", self._dataset_config) - groups = await config.get_attack_groups_by_dataset_async(apply_sampling=apply_sampling) + required_data = cast("ScenarioRequiredData", self._required_data) + corpus = await required_data.load_async() + populations = _assemble_requests(corpus=corpus, techniques=config._techniques) + groups = config.select_assembled_groups(groups_by_population=populations, apply_sampling=apply_sampling) if self._uses_default_scorer: + excluded_values = [ + seed.value for name in (DATASET_PARTIAL_KEYS, DATASET_SAFE_PLACEHOLDERS) for seed in corpus[name] + ] self._objective_scorer = CredentialLeakScorer.from_excluded_values( - config.excluded_values, patterns=CredentialLeakScorer.GARAK_PATTERNS + excluded_values, patterns=CredentialLeakScorer.GARAK_PATTERNS ) self._objective_scorer_identifier = self._objective_scorer.get_identifier() return groups diff --git a/tests/unit/registry/test_scenario_registry.py b/tests/unit/registry/test_scenario_registry.py index 93289637ea..d6375ee7d4 100644 --- a/tests/unit/registry/test_scenario_registry.py +++ b/tests/unit/registry/test_scenario_registry.py @@ -81,6 +81,10 @@ def get_default_dataset_settings(self) -> ScenarioDatasetSettings: """Return default dataset settings.""" return ScenarioDatasetSettings(max_total=4) + def get_default_dataset_names(self) -> list[str]: + """Return the datasets a default run loads.""" + return ["sample"] + class _MarkdownMetadataScenario(_MetadataScenario): """ diff --git a/tests/unit/scenario/conftest.py b/tests/unit/scenario/conftest.py index 6d73de635a..e78268d102 100644 --- a/tests/unit/scenario/conftest.py +++ b/tests/unit/scenario/conftest.py @@ -9,6 +9,7 @@ import pytest from pyrit.scenario.core.dataset_configuration import DatasetConfiguration +from pyrit.scenario.core.required_data import ScenarioRequiredData @pytest.fixture(autouse=True) @@ -19,5 +20,9 @@ def enforce_dataset_isolation(block_dataset_provider_lookup: AsyncMock) -> None: @pytest.fixture def mock_dataset_preparation() -> Generator[AsyncMock, None, None]: """Skip preparation only in tests that supply their own mocked seed populations.""" - with patch.object(DatasetConfiguration, "prepare_async", new_callable=AsyncMock) as prepare: + # Required data prepares together with the population, so skipping one means skipping both. + with ( + patch.object(DatasetConfiguration, "prepare_async", new_callable=AsyncMock) as prepare, + patch.object(ScenarioRequiredData, "prepare_async", new=prepare), + ): yield prepare diff --git a/tests/unit/scenario/core/test_required_data.py b/tests/unit/scenario/core/test_required_data.py new file mode 100644 index 0000000000..5d54140f4f --- /dev/null +++ b/tests/unit/scenario/core/test_required_data.py @@ -0,0 +1,505 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT license. + +"""Required data: datasets a scenario reads in full, apart from the population it limits.""" + +from collections.abc import Callable +from typing import Any +from unittest.mock import AsyncMock, MagicMock, patch + +import pytest +from unit.mocks import MockPromptTarget + +from pyrit.datasets import SeedDatasetProvider +from pyrit.executor.attack import PromptSendingAttack +from pyrit.memory import SQLiteMemory +from pyrit.models import ( + AttackSeedGroup, + DatasetFetchPolicy, + DatasetSelectionMode, + ScenarioDatasetOptions, + ScenarioDatasetRequest, + SeedDataset, + SeedObjective, +) +from pyrit.scenario import ( + CompoundDatasetAttackConfiguration, + DatasetAttackConfiguration, + DatasetSource, + Scenario, + ScenarioRequiredData, +) +from pyrit.scenario.core import AtomicAttack, AttackTechnique, BaselineAttackPolicy, ScenarioTechnique +from pyrit.scenario.core.dataset_configuration import DatasetConstraintError, ResolvedDataset +from pyrit.scenario.core.scenario_context import ScenarioContext +from pyrit.score import SubStringScorer + + +def _dataset(name: str, count: int = 3) -> SeedDataset: + return SeedDataset( + dataset_name=name, + seeds=[SeedObjective(value=f"{name}-{index}", dataset_name=name) for index in range(count)], + ) + + +def _provider(dataset: SeedDataset) -> MagicMock: + provider = MagicMock(spec=SeedDatasetProvider) + provider.dataset_name = dataset.dataset_name + provider.fetch_dataset_async = AsyncMock(return_value=dataset) + return provider + + +def _require_count(count: int) -> Callable[[ResolvedDataset], None]: + def _validate(resolved: ResolvedDataset) -> None: + if len(resolved.seeds) != count: + raise DatasetConstraintError(f"expected {count} seeds") + + return _validate + + +@pytest.mark.parametrize("max_size", [1, 5]) +def test_required_sources_cannot_be_capped(*, max_size: int) -> None: + with pytest.raises(DatasetConstraintError, match="read in full"): + ScenarioRequiredData(sources=[DatasetSource(name="a"), DatasetSource(name="b", max_size=max_size)]) + + +def test_required_data_needs_a_source() -> None: + with pytest.raises(DatasetConstraintError, match="at least one"): + ScenarioRequiredData(sources=[]) + + +@pytest.mark.parametrize("max_size", ["default", "all", None]) +def test_uncapped_sources_are_accepted(*, max_size: str | None) -> None: + required = ScenarioRequiredData(sources=[DatasetSource(name="a", max_size=max_size)]) + + assert required.dataset_names == ["a"] + assert required.fetch is DatasetFetchPolicy.IF_MISSING + + +@pytest.mark.usefixtures("patch_central_database") +async def test_prepare_stores_once_and_load_reads_everything_without_fetching_async( + *, sqlite_instance: SQLiteMemory +) -> None: + provider = _provider(_dataset("a", count=12)) + required = ScenarioRequiredData(sources=[DatasetSource(name="a", provider=provider)]) + + await required.prepare_async() + await required.prepare_async() + first = await required.load_async() + second = await required.load_async() + + provider.fetch_dataset_async.assert_awaited_once() + assert len(first["a"]) == len(second["a"]) == 12 + assert len(await sqlite_instance.get_seeds_async(dataset_name="a")) == 12 + + +@pytest.mark.usefixtures("patch_central_database") +async def test_load_never_fetches_a_missing_dataset_async() -> None: + provider = _provider(_dataset("a")) + required = ScenarioRequiredData(sources=[DatasetSource(name="a", provider=provider)]) + + with pytest.raises(DatasetConstraintError, match="could not be loaded"): + await required.load_async() + provider.fetch_dataset_async.assert_not_awaited() + + +@pytest.mark.usefixtures("patch_central_database") +async def test_never_policy_requires_stored_data_async() -> None: + provider = _provider(_dataset("a")) + required = ScenarioRequiredData( + sources=[DatasetSource(name="a", provider=provider)], fetch=DatasetFetchPolicy.NEVER + ) + + with pytest.raises(DatasetConstraintError, match="fetch is 'never'"): + await required.prepare_async() + provider.fetch_dataset_async.assert_not_awaited() + + +@pytest.mark.usefixtures("patch_central_database") +async def test_validators_run_on_every_read_async(*, sqlite_instance: SQLiteMemory) -> None: + await sqlite_instance.add_seed_datasets_to_memory_async(datasets=[_dataset("a", count=2)], added_by="test") + required = ScenarioRequiredData(sources=[DatasetSource(name="a")], validators=[_require_count(3)]) + + with pytest.raises(DatasetConstraintError, match="expected 3 seeds"): + await required.load_async() + + +@pytest.mark.usefixtures("patch_central_database") +async def test_refresh_replaces_the_complete_dataset_async(*, sqlite_instance: SQLiteMemory) -> None: + await sqlite_instance.add_seed_datasets_to_memory_async(datasets=[_dataset("a", count=5)], added_by="test") + provider = _provider(SeedDataset(dataset_name="a", seeds=[SeedObjective(value="fresh", dataset_name="a")])) + required = ScenarioRequiredData(sources=[DatasetSource(name="a", provider=provider)]) + + await required.prepare_async(fetch=DatasetFetchPolicy.ALWAYS) + + assert [seed.value for seed in (await required.load_async())["a"]] == ["fresh"] + + +@pytest.mark.usefixtures("patch_central_database") +async def test_refresh_that_fails_validation_keeps_stored_data_async(*, sqlite_instance: SQLiteMemory) -> None: + await sqlite_instance.add_seed_datasets_to_memory_async(datasets=[_dataset("a", count=3)], added_by="test") + provider = _provider(_dataset("a", count=1)) + required = ScenarioRequiredData( + sources=[DatasetSource(name="a", provider=provider)], + fetch=DatasetFetchPolicy.ALWAYS, + validators=[_require_count(3)], + ) + + with pytest.raises(DatasetConstraintError, match="expected 3 seeds"): + await required.prepare_async() + + provider.fetch_dataset_async.assert_awaited_once() + assert len(await sqlite_instance.get_seeds_async(dataset_name="a")) == 3 + + +@pytest.mark.usefixtures("patch_central_database") +async def test_assembly_check_reads_candidates_before_any_write_async(*, sqlite_instance: SQLiteMemory) -> None: + await sqlite_instance.add_seed_datasets_to_memory_async(datasets=[_dataset("a", count=3)], added_by="test") + provider = _provider(_dataset("a", count=1)) + required = ScenarioRequiredData( + sources=[DatasetSource(name="a", provider=provider)], + fetch=DatasetFetchPolicy.ALWAYS, + ) + seen: list[int] = [] + + async def reject_after_reading_async() -> None: + seen.append(len((await required.load_async())["a"])) + raise DatasetConstraintError("assembly rejected the refresh") + + with pytest.raises(DatasetConstraintError, match="assembly rejected the refresh"): + await required.prepare_async(validate_assembly=reject_after_reading_async) + + # The check saw the candidate (1 seed), not the stored data, and the rejection kept the stored 3. + assert seen == [1] + assert [seed.value for seed in (await required.load_async())["a"]] == ["a-0", "a-1", "a-2"] + + +@pytest.mark.usefixtures("patch_central_database") +async def test_filters_narrow_the_rows_that_are_read_async(*, sqlite_instance: SQLiteMemory) -> None: + await sqlite_instance.add_seeds_to_memory_async( + seeds=[ + SeedObjective(value="kept", dataset_name="a", harm_categories=["cyber"]), + SeedObjective(value="dropped", dataset_name="a", harm_categories=["violence"]), + ], + added_by="test", + ) + required = ScenarioRequiredData(sources=[DatasetSource(name="a")], filters={"harm_categories": ["cyber"]}) + + assert [seed.value for seed in (await required.load_async())["a"]] == ["kept"] + + +@pytest.mark.usefixtures("patch_central_database") +async def test_filters_that_match_nothing_fail_clearly_async(*, sqlite_instance: SQLiteMemory) -> None: + await sqlite_instance.add_seed_datasets_to_memory_async(datasets=[_dataset("a")], added_by="test") + required = ScenarioRequiredData(sources=[DatasetSource(name="a")], filters={"harm_categories": ["missing"]}) + + with pytest.raises(DatasetConstraintError, match="none match the configured filters"): + await required.load_async() + + +def test_projection_lists_every_source_uncapped() -> None: + provider = _provider(_dataset("b")) + required = ScenarioRequiredData( + sources=[DatasetSource(name="a"), DatasetSource(name="b", fetch=DatasetFetchPolicy.NEVER, provider=provider)] + ) + + projected = [ + (source.name, source.max_size, source.fetch, source.configured_provider) + for source in required.project_sources() + ] + + assert projected == [ + ("a", "all", DatasetFetchPolicy.IF_MISSING, False), + ("b", "all", DatasetFetchPolicy.NEVER, True), + ] + + +class TestApplyOverrides: + _FIXED = ScenarioDatasetOptions(mode=DatasetSelectionMode.FIXED, per_dataset_limit=False) + _ANY = ScenarioDatasetOptions() + + def _required(self) -> ScenarioRequiredData: + return ScenarioRequiredData(sources=[DatasetSource(name="a"), DatasetSource(name="b")]) + + @pytest.mark.parametrize("with_sources", [False, True]) + def test_global_fetch_reaches_required_sources_and_stays_for_the_population(self, *, with_sources: bool) -> None: + original = ScenarioRequiredData( + sources=[DatasetSource(name="a"), DatasetSource(name="b", fetch=DatasetFetchPolicy.IF_MISSING)] + ) + sources = [{"name": "a"}, {"name": "b"}, {"name": "c"}] if with_sources else None + request = ScenarioDatasetRequest.model_validate({"sources": sources, "fetch": "never", "max_total": 3}) + + required, remaining = original.apply_overrides(request=request, options=self._ANY, population_names=["c"]) + + # "a" has no policy of its own, so it takes the caller's; "b" keeps its explicit one. + assert required.fetch is DatasetFetchPolicy.NEVER + assert [(source.name, source.fetch) for source in required.sources] == [ + ("a", None), + ("b", DatasetFetchPolicy.IF_MISSING), + ] + assert remaining.fetch is DatasetFetchPolicy.NEVER + assert remaining.max_total == 3 + + @pytest.mark.parametrize("with_sources", [False, True]) + def test_filters_stay_on_the_request_for_the_population(self, *, with_sources: bool) -> None: + sources = [{"name": "a"}, {"name": "b"}] if with_sources else None + request = ScenarioDatasetRequest.model_validate({"sources": sources, "filters": {"harm_categories": ["cyber"]}}) + + _, remaining = self._required().apply_overrides(request=request, options=self._FIXED, population_names=[]) + + assert remaining.filters == {"harm_categories": ["cyber"]} + + def test_with_sources_merges_filters_and_keeps_validators(self) -> None: + def validator(_: ResolvedDataset) -> None: + pass + + original = ScenarioRequiredData( + sources=[DatasetSource(name="a")], validators=[validator], filters={"data_types": ["text"]} + ) + + copied = original.with_sources( + sources=original.sources, fetch=DatasetFetchPolicy.NEVER, filters={"harm_categories": ["cyber"]} + ) + + assert copied.filters == {"data_types": ["text"], "harm_categories": ["cyber"]} + assert copied.fetch is DatasetFetchPolicy.NEVER + assert copied._validators == [validator] + + def test_empty_request_keeps_the_declaration(self) -> None: + original = self._required() + + required, remaining = original.apply_overrides( + request=ScenarioDatasetRequest(), options=self._FIXED, population_names=[] + ) + + assert required is original + assert remaining == ScenarioDatasetRequest() + + def test_source_fetch_edits_reach_only_their_dataset(self) -> None: + request = ScenarioDatasetRequest.model_validate({"sources": [{"name": "a"}, {"name": "b", "fetch": "always"}]}) + + required, remaining = self._required().apply_overrides( + request=request, options=self._FIXED, population_names=[] + ) + + assert [(source.name, source.fetch) for source in required.sources] == [ + ("a", None), + ("b", DatasetFetchPolicy.ALWAYS), + ] + assert remaining.sources is None + + def test_population_sources_pass_through(self) -> None: + request = ScenarioDatasetRequest.model_validate({"sources": [{"name": "a"}, {"name": "b"}, {"name": "c"}]}) + + _, remaining = self._required().apply_overrides(request=request, options=self._ANY, population_names=["c"]) + + assert [source.name for source in remaining.sources or []] == ["c"] + + @pytest.mark.parametrize( + ("sources", "options", "error"), + [ + ([{"name": "b"}, {"name": "a"}], _FIXED, "fixed: a, b"), + ([{"name": "a"}], _FIXED, "fixed: a, b"), + ([{"name": "a"}, {"name": "b"}, {"name": "c"}], _FIXED, "fixed: a, b"), + ([{"name": "a"}, {"name": "c"}], _ANY, "can't be deselected: b"), + ([{"name": "a", "max_size": 2}, {"name": "b"}], _ANY, "can't be capped"), + ([{"name": "a", "max_size": "all"}, {"name": "b"}], _ANY, "can't be capped"), + ([{"name": "a"}, {"name": "b"}], _ANY, r"at least one population dataset .*defaults: c"), + ], + ) + def test_unsupported_edits_fail_clearly( + self, *, sources: list[dict[str, object]], options: ScenarioDatasetOptions, error: str + ) -> None: + request = ScenarioDatasetRequest.model_validate({"sources": sources}) + population = ["c"] if options is self._ANY else [] + + with pytest.raises(DatasetConstraintError, match=error): + self._required().apply_overrides(request=request, options=options, population_names=population) + + def test_scenarios_without_overrides_reject_required_edits(self) -> None: + request = ScenarioDatasetRequest.model_validate({"sources": [{"name": "a"}, {"name": "b"}]}) + + with pytest.raises(DatasetConstraintError, match="does not support dataset overrides"): + self._required().apply_overrides( + request=request, + options=ScenarioDatasetOptions(supports_overrides=False), + population_names=[], + ) + + +@pytest.mark.usefixtures("patch_central_database") +async def test_registered_provider_is_used_when_none_is_configured_async(*, sqlite_instance: SQLiteMemory) -> None: + provider = _provider(_dataset("a")) + required = ScenarioRequiredData(sources=[DatasetSource(name="a")]) + + with patch.object(SeedDatasetProvider, "get_providers_by_name_async", return_value={"a": provider}) as lookup: + await required.prepare_async() + + lookup.assert_awaited_once() + assert len(await sqlite_instance.get_seeds_async(dataset_name="a")) == 3 + + +def _versioned(name: str, version: str) -> SeedDataset: + return SeedDataset( + dataset_name=name, + seeds=[SeedObjective(value=f"{version}-{name}-{index}", dataset_name=name) for index in range(2)], + ) + + +class _Technique(ScenarioTechnique): + ALL = ("all", {"all"}) + DIRECT = ("direct", set()) + + @classmethod + def get_aggregate_tags(cls) -> set[str]: + return {"all"} + + +class _TemplatesAndRecordsScenario(Scenario): + """Required ``templates`` and a selectable ``records`` population that only assemble at the same version.""" + + BASELINE_ATTACK_POLICY = BaselineAttackPolicy.Forbidden + + def __init__(self, *, templates: MagicMock, records: MagicMock, compound: bool = False) -> None: + self._templates = templates + population = DatasetAttackConfiguration(sources=[DatasetSource(name="records", provider=records)]) + super().__init__( + version=1, + technique_class=_Technique, + objective_scorer=SubStringScorer(substring="yes"), + # A compound names its datasets through children, so its own ``sources`` is empty. + default_dataset_config=( + CompoundDatasetAttackConfiguration(configurations=[population]) if compound else population + ), + ) + + def get_required_data(self) -> ScenarioRequiredData: + return ScenarioRequiredData(sources=[DatasetSource(name="templates", provider=self._templates)]) + + async def _resolve_seed_groups_by_dataset_async( + self, *, apply_sampling: bool = True + ) -> dict[str, list[AttackSeedGroup]]: + assert self._required_data is not None + templates = (await self._required_data.load_async())["templates"] + groups = await self._dataset_config.get_attack_groups_by_dataset_async(apply_sampling=apply_sampling) + versions = {seed.value.split("-")[0] for seed in templates} | { + seed.value.split("-")[0] for group in groups["records"] for seed in group.seeds + } + if len(versions) != 1: + raise DatasetConstraintError(f"templates and records disagree: {sorted(versions)}") + return groups + + async def _build_atomic_attacks_async(self, *, context: ScenarioContext) -> list[AtomicAttack]: + return [ + AtomicAttack( + atomic_attack_name="direct", + attack_technique=AttackTechnique(attack=PromptSendingAttack(objective_target=context.objective_target)), + seed_groups=list(context.seed_groups), + ) + ] + + +async def _stored_async(memory: SQLiteMemory, name: str) -> dict[str, str]: + return {str(seed.id): str(seed.value) for seed in await memory.get_seeds_async(dataset_name=name)} + + +async def _initialize_async( + *, templates: SeedDataset, records: SeedDataset, overrides: dict[str, Any], compound: bool = False +) -> None: + scenario = _TemplatesAndRecordsScenario( + templates=_provider(templates), records=_provider(records), compound=compound + ) + scenario.set_params_from_args(args={"objective_target": MockPromptTarget(), "dataset_overrides": overrides}) + await scenario.initialize_async() + + +@pytest.mark.usefixtures("patch_central_database") +@pytest.mark.parametrize("records_stored", [False, True]) +@pytest.mark.parametrize("compound", [False, True]) +async def test_required_refresh_validates_with_population_candidates_async( + *, sqlite_instance: SQLiteMemory, records_stored: bool, compound: bool +) -> None: + # A templates refresh must see the records being fetched (or refreshed) in the same run, not + # missing or stale records. Refreshing both to v2 together is a valid joint refresh. + await sqlite_instance.add_seed_datasets_to_memory_async(datasets=[_versioned("templates", "v1")], added_by="test") + if records_stored: + await sqlite_instance.add_seed_datasets_to_memory_async(datasets=[_versioned("records", "v1")], added_by="test") + + await _initialize_async( + templates=_versioned("templates", "v2"), + records=_versioned("records", "v2"), + overrides={"fetch": "always"}, + compound=compound, + ) + + assert sorted((await _stored_async(sqlite_instance, "templates")).values())[0].startswith("v2") + assert sorted((await _stored_async(sqlite_instance, "records")).values())[0].startswith("v2") + + +@pytest.mark.usefixtures("patch_central_database") +async def test_population_refresh_rejected_by_assembly_keeps_stored_rows_async( + *, sqlite_instance: SQLiteMemory +) -> None: + for name in ("templates", "records"): + await sqlite_instance.add_seed_datasets_to_memory_async(datasets=[_versioned(name, "v1")], added_by="test") + before = {name: await _stored_async(sqlite_instance, name) for name in ("templates", "records")} + + with pytest.raises(DatasetConstraintError, match="disagree"): + await _initialize_async( + templates=_versioned("templates", "v1"), + records=_versioned("records", "v2"), + overrides={"sources": [{"name": "templates"}, {"name": "records", "fetch": "always"}]}, + ) + + # Same values and same ids, so earlier runs that planned these rows can still resume. + assert {name: await _stored_async(sqlite_instance, name) for name in ("templates", "records")} == before + + +@pytest.mark.usefixtures("patch_central_database") +@pytest.mark.parametrize("compound", [False, True]) +async def test_global_never_fetch_reaches_required_data_with_population_sources_async( + *, sqlite_instance: SQLiteMemory, compound: bool +) -> None: + await sqlite_instance.add_seed_datasets_to_memory_async(datasets=[_versioned("records", "v1")], added_by="test") + templates = _provider(_versioned("templates", "v1")) + scenario = _TemplatesAndRecordsScenario( + templates=templates, records=_provider(_versioned("records", "v1")), compound=compound + ) + scenario.set_params_from_args( + args={"objective_target": MockPromptTarget(), "dataset_overrides": {"fetch": "never"}} + ) + + with pytest.raises(DatasetConstraintError, match="'templates' is missing and fetch is 'never'"): + await scenario.initialize_async() + + templates.fetch_dataset_async.assert_not_awaited() + + +@pytest.mark.usefixtures("patch_central_database") +async def test_selecting_only_required_data_does_not_restore_the_default_population_async() -> None: + with pytest.raises(DatasetConstraintError, match="at least one population dataset"): + await _initialize_async( + templates=_versioned("templates", "v1"), + records=_versioned("records", "v1"), + overrides={"sources": [{"name": "templates"}]}, + ) + + +@pytest.mark.usefixtures("patch_central_database") +async def test_a_dataset_named_by_both_required_data_and_population_is_rejected_async( + *, sqlite_instance: SQLiteMemory +) -> None: + required_provider = _provider(_dataset("shared")) + population_provider = _provider(_dataset("shared")) + required = ScenarioRequiredData( + sources=[DatasetSource(name="shared", provider=required_provider)], fetch=DatasetFetchPolicy.ALWAYS + ) + population = DatasetAttackConfiguration(sources=[DatasetSource(name="shared", provider=population_provider)]) + + with pytest.raises(DatasetConstraintError, match=r"\['shared'\] are named by more than one configuration"): + await required.prepare_async(population=population) + + # Rejected before any fetch, so neither batch can be committed unvalidated. + required_provider.fetch_dataset_async.assert_not_awaited() + population_provider.fetch_dataset_async.assert_not_awaited() + assert await sqlite_instance.get_seeds_async(dataset_name="shared") == [] diff --git a/tests/unit/scenario/core/test_scenario_dataset_overrides.py b/tests/unit/scenario/core/test_scenario_dataset_overrides.py index 9a501c3d16..83c874cebc 100644 --- a/tests/unit/scenario/core/test_scenario_dataset_overrides.py +++ b/tests/unit/scenario/core/test_scenario_dataset_overrides.py @@ -206,6 +206,7 @@ async def test_catalog_dataset_controls_do_not_depend_on_estimates() -> None: 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.default_datasets == ApiKey.required_datasets() assert item.dataset_defaults.max_total == 20 diff --git a/tests/unit/scenario/garak/test_api_key.py b/tests/unit/scenario/garak/test_api_key.py index d607695c37..fa08173549 100644 --- a/tests/unit/scenario/garak/test_api_key.py +++ b/tests/unit/scenario/garak/test_api_key.py @@ -3,6 +3,7 @@ """Tests for the Garak API-key scenario.""" +from dataclasses import dataclass, field from pathlib import Path from typing import Any, Literal from unittest.mock import AsyncMock, MagicMock, patch @@ -12,9 +13,22 @@ 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, ScenarioDatasetRequest, ScenarioRunSizeEstimateStatus, Seed, SeedDataset +from pyrit.models import ( + ComponentIdentifier, + DatasetFetchPolicy, + ScenarioDatasetRequest, + ScenarioRunSizeEstimateStatus, + Seed, + SeedDataset, +) from pyrit.prompt_target import PromptTarget -from pyrit.scenario.core.dataset_configuration import DatasetAttackConfiguration, DatasetConstraintError +from pyrit.scenario import DatasetSource, ScenarioRequiredData +from pyrit.scenario.core.dataset_configuration import ( + DatasetAttackConfiguration, + DatasetConstraintError, + ResolvedDataset, +) +from pyrit.scenario.core.required_data import _RequiredDatasetConfiguration from pyrit.scenario.core.scenario import BaselineAttackPolicy from pyrit.scenario.garak import ( # type: ignore[ty:unresolved-import] ApiKey, @@ -64,9 +78,12 @@ async def _initialize_async( } ) with ( - patch.object(ApiKeyDatasetConfiguration, "prepare_async", new_callable=AsyncMock), + patch.object(ScenarioRequiredData, "prepare_async", new_callable=AsyncMock), patch.object( - ApiKeyDatasetConfiguration, "_collect_named_seeds_async", new_callable=AsyncMock, return_value=corpus_seeds + _RequiredDatasetConfiguration, + "_collect_named_seeds_async", + new_callable=AsyncMock, + return_value=corpus_seeds, ), ): await scenario.initialize_async() @@ -88,8 +105,10 @@ def test_defaults_and_standard_parameters(self) -> None: assert scenario.name == "ApiKey" assert scenario.VERSION == 1 assert scenario.BASELINE_ATTACK_POLICY is BaselineAttackPolicy.Forbidden - assert scenario._default_dataset_config.dataset_names == ApiKey.required_datasets() - assert scenario._default_dataset_config.max_dataset_size == 20 + assert scenario._default_dataset_config.dataset_names == [] + assert scenario._default_dataset_config.max_total == 20 + assert scenario.get_required_data().dataset_names == ApiKey.required_datasets() + assert scenario.get_default_dataset_names() == ApiKey.required_datasets() assert {"dataset_config", "technique_converters"} <= parameters assert "prompt_cap" not in parameters expected = {ApiKeyTechnique.GetKey, ApiKeyTechnique.CompleteKey} @@ -116,7 +135,7 @@ async def test_uncapped_configuration_renders_full_corpus( scenario=scenario, target=mock_objective_target, corpus_seeds=corpus_seeds, - dataset_config=ApiKeyDatasetConfiguration(dataset_names=ApiKey.required_datasets(), max_dataset_size="all"), + dataset_config=ApiKeyDatasetConfiguration(max_total="all"), ) assert {name: len(groups) for name, groups in _objectives(scenario).items()} == { @@ -141,7 +160,7 @@ async def test_standard_sampling_can_select_rows_beyond_the_prefix( scenario=scenario, target=mock_objective_target, corpus_seeds=corpus_seeds, - dataset_config=ApiKeyDatasetConfiguration(dataset_names=ApiKey.required_datasets(), max_dataset_size=3), + dataset_config=ApiKeyDatasetConfiguration(max_total=3), ) assert list(_objectives(scenario)) == ["complete_key"] @@ -154,7 +173,7 @@ async def test_resume_replays_sample_after_corpus_order_changes( self, size: int, mock_objective_target: PromptTarget, corpus_seeds: dict[str, list[Seed[Any]]] ) -> None: original = ApiKey() - config = ApiKeyDatasetConfiguration(dataset_names=ApiKey.required_datasets(), max_dataset_size=size) + config = ApiKeyDatasetConfiguration(max_total=size) await _initialize_async( scenario=original, target=mock_objective_target, corpus_seeds=corpus_seeds, dataset_config=config ) @@ -165,9 +184,7 @@ async def test_resume_replays_sample_after_corpus_order_changes( scenario=resumed, target=mock_objective_target, corpus_seeds={name: list(reversed(seeds)) for name, seeds in corpus_seeds.items()}, - dataset_config=ApiKeyDatasetConfiguration( - dataset_names=ApiKey.required_datasets(), max_dataset_size=size - ), + dataset_config=ApiKeyDatasetConfiguration(max_total=size), ) assert _objectives(resumed) == _objectives(original) @@ -199,7 +216,7 @@ async def test_custom_scorer_and_converters_are_preserved( scenario=scenario, target=mock_objective_target, corpus_seeds=corpus_seeds, - dataset_config=ApiKeyDatasetConfiguration(dataset_names=ApiKey.required_datasets()), + dataset_config=ApiKeyDatasetConfiguration(), technique_converters={"get_key": [converter]}, ) @@ -218,9 +235,7 @@ async def test_custom_validator_is_preserved( scenario=ApiKey(), target=mock_objective_target, corpus_seeds=corpus_seeds, - dataset_config=ApiKeyDatasetConfiguration( - dataset_names=ApiKey.required_datasets(), validators=[validator] - ), + dataset_config=ApiKeyDatasetConfiguration(validators=[validator]), ) validator.assert_called_once() @@ -260,7 +275,7 @@ async def test_launch_and_estimate_use_standard_dataset_size( ) scenario.set_params_from_args(args=args) with patch.object( - ApiKeyDatasetConfiguration, + _RequiredDatasetConfiguration, "_collect_named_seeds_async", new_callable=AsyncMock, return_value=corpus_seeds, @@ -311,13 +326,14 @@ async def test_single_technique_estimate_preserves_cap( args={ "objective_target": mock_objective_target, "scenario_techniques": [technique], - "dataset_config": ApiKeyDatasetConfiguration( - dataset_names=ApiKey.required_datasets(), max_dataset_size=size - ), + "dataset_config": ApiKeyDatasetConfiguration(max_total=size), } ) with patch.object( - ApiKeyDatasetConfiguration, "_collect_named_seeds_async", new_callable=AsyncMock, return_value=corpus_seeds + _RequiredDatasetConfiguration, + "_collect_named_seeds_async", + new_callable=AsyncMock, + return_value=corpus_seeds, ): estimate = await scenario.get_run_size_estimate_async(target_is_configured=True) @@ -342,3 +358,277 @@ async def test_real_local_datasets_resolve_through_memory(self, mock_objective_t await scenario.initialize_async() assert sum(len(attack.seed_groups) for attack in scenario._atomic_attacks) == 20 + + def test_corpus_is_required_data_and_the_configuration_only_limits_requests(self) -> None: + scenario = ApiKey() + settings = scenario.get_default_dataset_settings() + + assert scenario._default_dataset_config.sources == () + assert [(source.name, source.max_size) for source in settings.sources] == [ + (name, "all") for name in ApiKey.required_datasets() + ] + assert settings.max_total == 20 + + async def test_only_a_fresh_run_prepares_the_corpus_async( + self, *, mock_objective_target: PromptTarget, corpus_seeds: dict[str, list[Seed[Any]]] + ) -> None: + scenario = ApiKey() + with ( + patch.object(ScenarioRequiredData, "prepare_async", new_callable=AsyncMock) as prepare, + patch.object( + _RequiredDatasetConfiguration, + "_collect_named_seeds_async", + new_callable=AsyncMock, + return_value=corpus_seeds, + ), + ): + scenario.set_params_from_args(args={"objective_target": mock_objective_target}) + await scenario.get_run_size_estimate_async(target_is_configured=True) + prepare.assert_not_awaited() + + await scenario.initialize_async() + prepare.assert_awaited_once() + + resumed = ApiKey(scenario_result_id=scenario._scenario_result_id) + resumed.set_params_from_args(args={"objective_target": mock_objective_target}) + await resumed.initialize_async() + prepare.assert_awaited_once() + + assert _objectives(resumed) == _objectives(scenario) + + async def test_estimate_never_reads_the_corpus_async(self, *, mock_objective_target: PromptTarget) -> None: + scenario = ApiKey() + scenario.set_params_from_args(args={"objective_target": mock_objective_target}) + with patch.object(ScenarioRequiredData, "load_async", side_effect=AssertionError("read corpus")): + estimate = await scenario.get_run_size_estimate_async(target_is_configured=True) + + assert estimate.estimated_attack_count == 20 + + async def test_refresh_reaches_the_corpus_async(self, *, mock_objective_target: PromptTarget) -> None: + scenario = ApiKey() + scenario.set_params_from_args( + args={ + "objective_target": mock_objective_target, + "dataset_overrides": { + "fetch": "always", + "sources": [ + {"name": name, "fetch": "never" if name == DATASET_PARTIAL_KEYS else None} + for name in ApiKey.required_datasets() + ], + }, + } + ) + estimate = await scenario.get_run_size_estimate_async(target_is_configured=True) + + assert [(source.name, source.fetch) for source in estimate.effective_datasets.sources] == [ + (name, DatasetFetchPolicy.NEVER if name == DATASET_PARTIAL_KEYS else DatasetFetchPolicy.ALWAYS) + for name in ApiKey.required_datasets() + ] + + async def test_corpus_rows_cannot_be_capped_through_overrides_async( + self, *, mock_objective_target: PromptTarget + ) -> None: + scenario = ApiKey() + scenario.set_params_from_args( + args={ + "objective_target": mock_objective_target, + "dataset_overrides": { + "sources": [{"name": name, "max_size": 2} for name in ApiKey.required_datasets()] + }, + } + ) + with pytest.raises(DatasetConstraintError, match="can't be capped"): + await scenario.get_run_size_estimate_async(target_is_configured=True) + + async def test_identity_still_lists_the_corpus_async( + self, *, mock_objective_target: PromptTarget, corpus_seeds: dict[str, list[Seed[Any]]] + ) -> None: + scenario = ApiKey() + await _initialize_async(scenario=scenario, target=mock_objective_target, corpus_seeds=corpus_seeds) + + assert scenario._build_scenario_identifier().datasets == ApiKey.required_datasets() + + async def test_listing_the_corpus_as_sources_still_works_with_a_warning_async( + self, *, mock_objective_target: PromptTarget, corpus_seeds: dict[str, list[Seed[Any]]] + ) -> None: + current = ApiKey() + await _initialize_async( + scenario=current, + target=mock_objective_target, + corpus_seeds=corpus_seeds, + dataset_config=ApiKeyDatasetConfiguration(max_total="all"), + ) + legacy = ApiKey() + with pytest.warns(DeprecationWarning, match="ApiKeyDatasetConfiguration"): + await _initialize_async( + scenario=legacy, + target=mock_objective_target, + corpus_seeds=corpus_seeds, + dataset_config=ApiKeyDatasetConfiguration( + sources=[ + DatasetSource(name=name, fetch=DatasetFetchPolicy.IF_MISSING) + for name in ApiKey.required_datasets() + ], + max_total="all", + ), + ) + + assert _objectives(legacy) == _objectives(current) + assert legacy._dataset_config.sources == () + assert legacy._required_data is not None + assert [source.fetch for source in legacy._required_data.sources] == [DatasetFetchPolicy.IF_MISSING] * 4 + assert legacy._build_scenario_identifier().datasets == current._build_scenario_identifier().datasets + + async def test_migration_moves_an_unhashable_callable_validator_async( + self, *, mock_objective_target: PromptTarget, corpus_seeds: dict[str, list[Seed[Any]]] + ) -> None: + # ``validators`` accepts any callable, including an unhashable instance; migration must move + # it onto the raw corpus without putting it in a set. + @dataclass + class RecordingValidator: + seen: list[int] = field(default_factory=list) + + def __call__(self, resolved: ResolvedDataset) -> None: + self.seen.append(len(resolved.seeds)) + + validator = RecordingValidator() + scenario = ApiKey() + with pytest.warns(DeprecationWarning, match="ApiKeyDatasetConfiguration"): + await _initialize_async( + scenario=scenario, + target=mock_objective_target, + corpus_seeds=corpus_seeds, + dataset_config=ApiKeyDatasetConfiguration( + sources=[DatasetSource(name=name) for name in ApiKey.required_datasets()], + validators=[validator], + max_total="all", + ), + ) + + assert validator.seen == [sum(len(seeds) for seeds in corpus_seeds.values())] + + async def test_capped_corpus_sources_still_fail_async(self, *, mock_objective_target: PromptTarget) -> None: + scenario = ApiKey() + scenario.set_params_from_args( + args={ + "objective_target": mock_objective_target, + "dataset_config": ApiKeyDatasetConfiguration( + sources=[DatasetSource(name=name) for name in ApiKey.required_datasets()], max_per_dataset=2 + ), + } + ) + with pytest.warns(DeprecationWarning), pytest.raises(DatasetConstraintError, match="can't be capped"): + await scenario.get_run_size_estimate_async(target_is_configured=True) + + async def test_missing_template_fails_before_assembly_async( + self, *, mock_objective_target: PromptTarget, corpus_seeds: dict[str, list[Seed[Any]]] + ) -> None: + corpus = {**corpus_seeds, "garak_api_key_templates": corpus_seeds["garak_api_key_templates"][:1]} + with pytest.raises(DatasetConstraintError, match="one prompt template"): + await _initialize_async(scenario=ApiKey(), target=mock_objective_target, corpus_seeds=corpus) + + async def test_get_key_only_run_does_not_need_a_complete_key_template_async( + self, *, mock_objective_target: PromptTarget, corpus_seeds: dict[str, list[Seed[Any]]] + ) -> None: + get_key_only = [ + seed + for seed in corpus_seeds["garak_api_key_templates"] + if (seed.metadata or {}).get("technique") == ApiKeyTechnique.GetKey.value + ] + corpus = {**corpus_seeds, "garak_api_key_templates": get_key_only} + scenario = ApiKey() + await _initialize_async( + scenario=scenario, + target=mock_objective_target, + corpus_seeds=corpus, + techniques=[ApiKeyTechnique.GetKey], + ) + + assert list(_objectives(scenario)) == ["get_key"] + + async def test_a_selected_technique_still_requires_its_template_async( + self, *, mock_objective_target: PromptTarget, corpus_seeds: dict[str, list[Seed[Any]]] + ) -> None: + get_key_only = [ + seed + for seed in corpus_seeds["garak_api_key_templates"] + if (seed.metadata or {}).get("technique") == ApiKeyTechnique.GetKey.value + ] + corpus = {**corpus_seeds, "garak_api_key_templates": get_key_only} + with pytest.raises(DatasetConstraintError, match="one prompt template for complete_key"): + await _initialize_async( + scenario=ApiKey(), + target=mock_objective_target, + corpus_seeds=corpus, + techniques=[ApiKeyTechnique.CompleteKey], + ) + + async def test_listed_corpus_sources_still_take_overrides_async( + self, *, mock_objective_target: PromptTarget + ) -> None: + scenario = ApiKey() + scenario.set_params_from_args( + args={ + "objective_target": mock_objective_target, + "dataset_config": ApiKeyDatasetConfiguration( + sources=[DatasetSource(name=name) for name in ApiKey.required_datasets()], max_total=5 + ), + "dataset_overrides": { + "max_total": 3, + "sources": [ + {"name": name, "fetch": "always" if name == DATASET_PARTIAL_KEYS else None} + for name in ApiKey.required_datasets() + ], + }, + } + ) + with pytest.warns(DeprecationWarning, match="ApiKeyDatasetConfiguration"): + estimate = await scenario.get_run_size_estimate_async(target_is_configured=True) + + assert estimate.estimated_attack_count == 3 + assert [(source.name, source.fetch) for source in estimate.effective_datasets.sources] == [ + (name, DatasetFetchPolicy.ALWAYS if name == DATASET_PARTIAL_KEYS else DatasetFetchPolicy.IF_MISSING) + for name in ApiKey.required_datasets() + ] + + @pytest.mark.usefixtures("mock_garak_dataset_fetch") + @pytest.mark.parametrize( + ("category", "error"), [("cyber", None), ("violence", "none match the configured filters")] + ) + async def test_dataset_filters_reach_the_corpus_async( + self, *, category: str, error: str | None, mock_objective_target: PromptTarget + ) -> None: + scenario = ApiKey() + scenario.set_params_from_args( + args={ + "objective_target": mock_objective_target, + "dataset_overrides": {"filters": {"harm_categories": [category]}}, + } + ) + + if error: + with pytest.raises(DatasetConstraintError, match=error): + await scenario.initialize_async() + return + await scenario.initialize_async() + assert sum(len(attack.seed_groups) for attack in scenario._atomic_attacks) == 20 + + async def test_configuration_fetch_and_filters_reach_the_corpus_async( + self, *, mock_objective_target: PromptTarget + ) -> None: + scenario = ApiKey() + scenario.set_params_from_args( + args={ + "objective_target": mock_objective_target, + "dataset_config": ApiKeyDatasetConfiguration( + fetch=DatasetFetchPolicy.NEVER, filters={"harm_categories": ["cyber"]} + ), + } + ) + await scenario.get_run_size_estimate_async(target_is_configured=True) + + assert scenario._required_data is not None + assert scenario._required_data.fetch is DatasetFetchPolicy.NEVER + assert scenario._required_data.filters == {"harm_categories": ["cyber"]} + with pytest.raises(DatasetConstraintError, match="fetch is 'never'"): + await scenario.initialize_async() diff --git a/tests/unit/scenario/test_default_run_size_estimates.py b/tests/unit/scenario/test_default_run_size_estimates.py index 759a8469df..92dd21b34e 100644 --- a/tests/unit/scenario/test_default_run_size_estimates.py +++ b/tests/unit/scenario/test_default_run_size_estimates.py @@ -43,7 +43,7 @@ from pyrit.scenario.scenarios.airt.psychosocial import Psychosocial from pyrit.scenario.scenarios.benchmark.adversarial import AdversarialBenchmark from pyrit.scenario.scenarios.foundry.red_team_agent import FoundryComposite, FoundryTechnique, RedTeamAgent -from pyrit.scenario.scenarios.garak.api_key import ApiKey, ApiKeyDatasetConfiguration +from pyrit.scenario.scenarios.garak.api_key import ApiKey from pyrit.scenario.scenarios.garak.encoding import Encoding from pyrit.scenario.scenarios.garak.exploitation import Exploitation from pyrit.scenario.scenarios.garak.figstep import FigStep @@ -224,13 +224,12 @@ async def test_configuration_only_checks_are_shared_before_dataset_reads_async( @pytest.mark.usefixtures("patch_central_database") @pytest.mark.parametrize( "configuration_class", - [ApiKeyDatasetConfiguration, PromptInjectDatasetConfiguration, LatentInjectionDatasetConfiguration], + [PromptInjectDatasetConfiguration, LatentInjectionDatasetConfiguration], ) def test_ingredient_configurations_require_uncapped_sources( configuration_class: type[DatasetAttackConfiguration], ) -> None: scenario_class = { - ApiKeyDatasetConfiguration: ApiKey, PromptInjectDatasetConfiguration: PromptInject, LatentInjectionDatasetConfiguration: LatentInjection, }[configuration_class]