From 8c71511ea41bd2705c8065a0ebbc34255cf939c1 Mon Sep 17 00:00:00 2001 From: Richard Lundeen Date: Tue, 30 Jun 2026 13:56:54 -0700 Subject: [PATCH 1/4] MAINT: Introduce ScenarioContext + MatrixAtomicAttackBuilder extension point Slim the Scenario base class toward pure orchestration by adding one explicit extension point, `_build_atomic_attacks_async(self, *, context: ScenarioContext)`, and extracting the technique x dataset cross-product into a standalone MatrixAtomicAttackBuilder. ScenarioContext is a frozen dataclass carrying the run-resolved inputs (objective_target, scenario_strategies, dataset_config, memory_labels, include_baseline). A bridge in the legacy `_get_atomic_attacks_async` builds the context from resolved self state and forwards to the new method, so every existing call site keeps working with zero test edits. Migrated AdversarialBenchmark, RedTeamAgent, AdaptiveScenario, Jailbreak, Psychosocial, Scam, and Encoding onto the new method. objective_scorer is treated as scenario configuration (read from self) rather than a universal context field, since scenarios like Psychosocial select scorers per harm category and strategy. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- .../core/matrix_atomic_attack_builder.py | 312 ++++++++++++++++++ pyrit/scenario/core/scenario.py | 161 ++++----- pyrit/scenario/core/scenario_context.py | 53 +++ .../scenarios/adaptive/adaptive_scenario.py | 20 +- pyrit/scenario/scenarios/airt/jailbreak.py | 10 +- pyrit/scenario/scenarios/airt/psychosocial.py | 11 +- pyrit/scenario/scenarios/airt/scam.py | 10 +- .../scenarios/benchmark/adversarial.py | 128 +++---- .../scenarios/foundry/red_team_agent.py | 10 +- pyrit/scenario/scenarios/garak/encoding.py | 16 +- .../core/test_matrix_atomic_attack_builder.py | 284 ++++++++++++++++ .../scenario/core/test_scenario_context.py | 62 ++++ 12 files changed, 890 insertions(+), 187 deletions(-) create mode 100644 pyrit/scenario/core/matrix_atomic_attack_builder.py create mode 100644 pyrit/scenario/core/scenario_context.py create mode 100644 tests/unit/scenario/core/test_matrix_atomic_attack_builder.py create mode 100644 tests/unit/scenario/core/test_scenario_context.py diff --git a/pyrit/scenario/core/matrix_atomic_attack_builder.py b/pyrit/scenario/core/matrix_atomic_attack_builder.py new file mode 100644 index 0000000000..ff79208434 --- /dev/null +++ b/pyrit/scenario/core/matrix_atomic_attack_builder.py @@ -0,0 +1,312 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT license. + +""" +Reusable matrix builder for scenario atomic attacks. + +``MatrixAtomicAttackBuilder`` turns a technique × dataset (× optional adversarial +target) grid into a flat list of ``AtomicAttack`` instances. It centralizes the +seed-technique compatibility filtering, ``factory.create`` wiring, ``AtomicAttack`` +construction, and baseline emission that scenarios previously duplicated in their +own ``_get_atomic_attacks_async`` overrides. + +This is one member of a *family* of construction helpers named by shape +(``Matrix...``); scenarios whose construction is composite or per-objective build +``AtomicAttack`` lists differently. There is intentionally no shared builder +interface yet — each scenario calls the builder it needs directly. +""" + +from __future__ import annotations + +import logging +from dataclasses import dataclass +from typing import TYPE_CHECKING, cast + +from pyrit.executor.attack import AttackScoringConfig +from pyrit.executor.attack.single_turn.prompt_sending import PromptSendingAttack +from pyrit.models import SeedAttackGroup +from pyrit.scenario.core.atomic_attack import AtomicAttack +from pyrit.scenario.core.attack_technique import AttackTechnique + +if TYPE_CHECKING: + from collections.abc import Callable, Sequence + + from pyrit.prompt_target import PromptTarget + from pyrit.scenario.core.attack_technique_factory import AttackTechniqueFactory + from pyrit.score import Scorer + +logger = logging.getLogger(__name__) + + +@dataclass(frozen=True) +class MatrixCombo: + """ + One cell of the build matrix, passed to the ``name_fn``/``display_group_fn`` callbacks. + + Attributes: + technique_name (str): The technique (strategy enum value) for this cell. + dataset_name (str): The dataset key from ``DatasetConfiguration.get_seed_attack_groups()``. + target_name (str | None): The adversarial-target registry name when an + adversarial-target axis is in play, else ``None``. + """ + + technique_name: str + dataset_name: str + target_name: str | None = None + + +def _default_atomic_attack_name(combo: MatrixCombo) -> str: + """ + Default ``atomic_attack_name`` builder; target-aware so names stay unique. + + Args: + combo (MatrixCombo): The matrix cell being named. + + Returns: + str: The atomic attack name for the cell. + """ + if combo.target_name is None: + return f"{combo.technique_name}_{combo.dataset_name}" + return f"{combo.technique_name}__{combo.target_name}_{combo.dataset_name}" + + +def _default_display_group(combo: MatrixCombo) -> str: + """ + Build the default display group: aggregate results by technique. + + Args: + combo (MatrixCombo): The matrix cell being grouped. + + Returns: + str: The display group for the cell. + """ + return combo.technique_name + + +def build_baseline_atomic_attack( + *, + objective_target: PromptTarget, + objective_scorer: Scorer, + seed_groups: list[SeedAttackGroup], + memory_labels: dict[str, str] | None = None, +) -> AtomicAttack: + """ + Build the baseline ``AtomicAttack`` that sends each objective unmodified. + + The baseline is a plain ``PromptSendingAttack`` used as a comparison point against + a scenario's strategy attacks. Pass the *same* ``seed_groups`` used to build the + strategy attacks so both populations match — re-resolving under ``max_dataset_size`` + would draw a fresh random sample and diverge from the strategy population. + + Args: + objective_target (PromptTarget): The target to attack. + objective_scorer (Scorer): The scorer used to evaluate the baseline. + seed_groups (list[SeedAttackGroup]): Seed groups to attack. Used as-is. + memory_labels (dict[str, str] | None): Labels applied to the baseline's prompts. + + Returns: + AtomicAttack: The baseline atomic attack named ``"baseline"``. + """ + attack = PromptSendingAttack( + objective_target=objective_target, + attack_scoring_config=AttackScoringConfig(objective_scorer=cast("TrueFalseScorer", objective_scorer)), + ) + return AtomicAttack( + atomic_attack_name="baseline", + attack_technique=AttackTechnique(attack=attack), + seed_groups=seed_groups, + memory_labels=memory_labels or {}, + ) + + +class MatrixAtomicAttackBuilder: + """ + Build ``AtomicAttack`` instances from a technique × dataset (× target) cross-product. + + Construct once with the shared run inputs (target, scorer, labels), then call + :meth:`build` with the per-run grid. The builder owns: + + - seed-technique compatibility filtering (``SeedAttackGroup.filter_compatible``), + - the ``factory.create(...)`` call, forwarding an adversarial target when the + adversarial-target axis is active, + - ``AtomicAttack`` construction with naming and display-group stamping, and + - optional baseline emission using the same resolved seed groups. + + Example: + >>> builder = MatrixAtomicAttackBuilder( + ... objective_target=target, + ... objective_scorer=scorer, + ... memory_labels=labels, + ... ) + >>> attacks = builder.build( + ... technique_factories=factories, + ... dataset_groups=groups, + ... include_baseline=True, + ... ) + """ + + def __init__( + self, + *, + objective_target: PromptTarget, + objective_scorer: Scorer, + memory_labels: dict[str, str] | None = None, + ) -> None: + """ + Initialize the builder with inputs shared across every atomic attack it produces. + + Args: + objective_target (PromptTarget): The target system to attack. + objective_scorer (Scorer): The scorer applied to each produced atomic attack + and to the baseline. + memory_labels (dict[str, str] | None): Labels applied to every produced + atomic attack. + """ + self._objective_target = objective_target + self._objective_scorer = objective_scorer + self._memory_labels = memory_labels or {} + + def build( + self, + *, + technique_factories: dict[str, AttackTechniqueFactory], + dataset_groups: dict[str, list[SeedAttackGroup]], + adversarial_targets: Sequence[tuple[str, PromptTarget]] | None = None, + name_fn: Callable[[MatrixCombo], str] | None = None, + display_group_fn: Callable[[MatrixCombo], str] | None = None, + include_baseline: bool = False, + ) -> list[AtomicAttack]: + """ + Build the atomic attacks for the given grid. + + Iterates technique → (adversarial target) → dataset. The caller pre-resolves + ``technique_factories`` to exactly the techniques to build (and, by dict + insertion order, the order to build them in), so the builder does not need the + full registry or the selected-strategy set. + + Args: + technique_factories (dict[str, AttackTechniqueFactory]): Mapping of technique + name to the factory that produces it. Only these techniques are built. + dataset_groups (dict[str, list[SeedAttackGroup]]): Mapping of dataset name to + its seed groups (e.g. ``DatasetConfiguration.get_seed_attack_groups()``). + adversarial_targets (Sequence[tuple[str, PromptTarget]] | None): Optional + ``(name, instance)`` pairs adding an adversarial-target axis. When set, + each technique is swept across every target and the target instance is + forwarded to ``factory.create(adversarial_chat=...)``. When ``None``, the + axis is collapsed and each factory uses its own (possibly lazy) + adversarial target. + name_fn (Callable[[MatrixCombo], str] | None): Builds each ``atomic_attack_name``. + Defaults to ``"{technique}_{dataset}"`` (or ``"{technique}__{target}_{dataset}"`` + when an adversarial-target axis is active). + display_group_fn (Callable[[MatrixCombo], str] | None): Builds each + ``display_group``. Defaults to grouping by technique name. + include_baseline (bool): When ``True``, prepend a baseline atomic attack built + from the flattened seed groups across all datasets. + + Returns: + list[AtomicAttack]: The generated atomic attacks, baseline first when requested. + """ + name_fn = name_fn or _default_atomic_attack_name + display_group_fn = display_group_fn or _default_display_group + + scoring_config = AttackScoringConfig(objective_scorer=cast("TrueFalseScorer", self._objective_scorer)) + + target_axis: Sequence[tuple[str | None, PromptTarget | None]] = ( + list(adversarial_targets) if adversarial_targets else [(None, None)] + ) + + atomic_attacks: list[AtomicAttack] = [] + for technique_name, factory in technique_factories.items(): + for target_name, target_instance in target_axis: + for dataset_name, seed_groups in dataset_groups.items(): + compatible_groups = self._filter_compatible_groups( + factory=factory, + seed_groups=seed_groups, + technique_name=technique_name, + dataset_name=dataset_name, + ) + if compatible_groups is None: + continue + + create_adversarial = {"adversarial_chat": target_instance} if target_instance is not None else {} + attack_technique = factory.create( + objective_target=self._objective_target, + attack_scoring_config=scoring_config, + **create_adversarial, + ) + + combo = MatrixCombo( + technique_name=technique_name, + dataset_name=dataset_name, + target_name=target_name, + ) + atomic_attacks.append( + AtomicAttack( + atomic_attack_name=name_fn(combo), + attack_technique=attack_technique, + seed_groups=compatible_groups, + adversarial_chat=( + target_instance if target_instance is not None else factory.adversarial_chat + ), + objective_scorer=cast("TrueFalseScorer", self._objective_scorer), + memory_labels=self._memory_labels, + display_group=display_group_fn(combo), + ) + ) + + if include_baseline: + all_seed_groups = [group for groups in dataset_groups.values() for group in groups] + atomic_attacks.insert( + 0, + build_baseline_atomic_attack( + objective_target=self._objective_target, + objective_scorer=self._objective_scorer, + seed_groups=all_seed_groups, + memory_labels=self._memory_labels, + ), + ) + + return atomic_attacks + + def _filter_compatible_groups( + self, + *, + factory: AttackTechniqueFactory, + seed_groups: list[SeedAttackGroup], + technique_name: str, + dataset_name: str, + ) -> list[SeedAttackGroup] | None: + """ + Filter seed groups to those compatible with the factory's seed technique. + + Args: + factory (AttackTechniqueFactory): The factory whose ``seed_technique`` gates + compatibility. + seed_groups (list[SeedAttackGroup]): Candidate seed groups for one dataset. + technique_name (str): Technique name, used only for log messages. + dataset_name (str): Dataset name, used only for log messages. + + Returns: + list[SeedAttackGroup] | None: The compatible groups, or ``None`` when the + ``(technique, dataset)`` pair has no compatible groups and should be skipped. + """ + if factory.seed_technique is None: + return list(seed_groups) + + compatible_groups = SeedAttackGroup.filter_compatible( + seed_groups=seed_groups, + technique=factory.seed_technique, + ) + skipped = len(seed_groups) - len(compatible_groups) + if skipped: + logger.info( + f"Skipped {skipped} seed group(s) from '{dataset_name}' for technique " + f"'{technique_name}' (prompt sequences overlap with simulated conversation)." + ) + if not compatible_groups: + logger.warning( + f"No compatible seed groups in '{dataset_name}' for technique " + f"'{technique_name}', skipping this (technique, dataset) pair." + ) + return None + return compatible_groups diff --git a/pyrit/scenario/core/scenario.py b/pyrit/scenario/core/scenario.py index dce77c3bbd..386f00c23f 100644 --- a/pyrit/scenario/core/scenario.py +++ b/pyrit/scenario/core/scenario.py @@ -18,7 +18,7 @@ from collections.abc import Sequence from enum import Enum from pathlib import Path -from typing import TYPE_CHECKING, Any, ClassVar, cast, get_origin +from typing import TYPE_CHECKING, Any, ClassVar, get_origin try: # Built-in on Python 3.11+. Fall back to the ``exceptiongroup`` backport on 3.10 @@ -34,7 +34,6 @@ from pyrit.common.parameter import coerce_value, validate_param_type from pyrit.common.utils import to_sha256 from pyrit.executor.attack import AttackExecutor -from pyrit.executor.attack.single_turn.prompt_sending import PromptSendingAttack from pyrit.memory import CentralMemory from pyrit.memory.memory_models import ScenarioResultEntry from pyrit.models import ( @@ -49,8 +48,12 @@ from pyrit.prompt_target.common.target_requirements import TargetRequirements from pyrit.registry import ScorerRegistry from pyrit.scenario.core.atomic_attack import AtomicAttack -from pyrit.scenario.core.attack_technique import AttackTechnique from pyrit.scenario.core.dataset_configuration import DatasetConfiguration +from pyrit.scenario.core.matrix_atomic_attack_builder import ( + MatrixAtomicAttackBuilder, + build_baseline_atomic_attack, +) +from pyrit.scenario.core.scenario_context import ScenarioContext from pyrit.scenario.core.scenario_strategy import ScenarioStrategy from pyrit.scenario.core.scenario_target_defaults import get_default_scorer_target from pyrit.score import ( @@ -252,6 +255,11 @@ def __init__( self._max_concurrency: int | None = None self._max_retries: int = 0 + # Effective dataset configuration for the current run. initialize_async reassigns + # this to the caller-supplied config (or the default); defaulting it here means the + # attribute always exists for context construction. + self._dataset_config: DatasetConfiguration = default_dataset_config + self._objective_scorer = objective_scorer self._objective_scorer_identifier = objective_scorer.get_identifier() @@ -842,16 +850,9 @@ def _build_baseline_atomic_attack(self, *, seed_groups: list[SeedAttackGroup]) - if self._objective_scorer is None: raise ValueError("Objective scorer is required to create baseline attack.") - from pyrit.executor.attack.core.attack_config import AttackScoringConfig - - attack = PromptSendingAttack( + return build_baseline_atomic_attack( objective_target=self._objective_target, - attack_scoring_config=AttackScoringConfig(objective_scorer=cast("TrueFalseScorer", self._objective_scorer)), - ) - - return AtomicAttack( - atomic_attack_name="baseline", - attack_technique=AttackTechnique(attack=attack), + objective_scorer=self._objective_scorer, seed_groups=seed_groups, memory_labels=self._memory_labels, ) @@ -1013,22 +1014,17 @@ async def _get_remaining_atomic_attacks_async(self) -> list[AtomicAttack]: return remaining_attacks - async def _get_atomic_attacks_async(self) -> list[AtomicAttack]: + def _build_scenario_context(self) -> ScenarioContext: """ - Build atomic attacks from the cross-product of selected techniques and datasets. - - Uses ``_get_attack_technique_factories()`` to obtain factories, then - iterates over every (technique, dataset) pair to create an - ``AtomicAttack`` for each. Grouping for display is controlled by - ``_build_display_group()``. + Snapshot the resolved runtime inputs into a ``ScenarioContext``. - Subclasses that do **not** use the factory/registry pattern should - override this method entirely. Overrides that want baseline support - must call ``self._build_baseline_atomic_attack`` with the strategy - seeds. + Called after ``initialize_async`` has populated the objective target, scorer, + strategies, dataset config, labels, and baseline flag. The resulting context is + handed to ``_build_atomic_attacks_async`` so scenario authors never read + half-initialized ``self._*`` state to build attacks. Returns: - list[AtomicAttack]: The generated atomic attacks. + ScenarioContext: The immutable inputs for atomic-attack construction. Raises: ValueError: If the scenario has not been initialized. @@ -1038,68 +1034,79 @@ async def _get_atomic_attacks_async(self) -> list[AtomicAttack]: "Scenario not properly initialized. Call await scenario.initialize_async() before running." ) - from pyrit.executor.attack import AttackScoringConfig + return ScenarioContext( + objective_target=self._objective_target, + scenario_strategies=tuple(self._scenario_strategies), + dataset_config=self._dataset_config, + memory_labels=dict(self._memory_labels), + include_baseline=self._include_baseline, + ) + + async def _get_atomic_attacks_async(self) -> list[AtomicAttack]: + """ + Build this scenario's atomic attacks (internal entry point from ``initialize_async``). - selected_techniques = {s.value for s in self._scenario_strategies} + New scenarios should override ``_build_atomic_attacks_async(context=...)`` instead + of this method. The base implementation here builds a ``ScenarioContext`` from the + values resolved in ``initialize_async`` and forwards to + ``_build_atomic_attacks_async``. It remains as a stable, no-argument entry point so + legacy overrides and existing call sites keep working during the migration. - factories = self._get_attack_technique_factories() - seed_groups_by_dataset = self._dataset_config.get_seed_attack_groups() + Returns: + list[AtomicAttack]: The generated atomic attacks. - scoring_config = AttackScoringConfig(objective_scorer=cast("TrueFalseScorer", self._objective_scorer)) + Raises: + ValueError: If the scenario has not been initialized. + """ + context = self._build_scenario_context() + return await self._build_atomic_attacks_async(context=context) - atomic_attacks: list[AtomicAttack] = [] + async def _build_atomic_attacks_async(self, *, context: ScenarioContext) -> list[AtomicAttack]: + """ + Build atomic attacks from the cross-product of selected techniques and datasets. + + This is the single extension point scenarios override to map techniques, datasets, + scorers, and any extra axes into ``AtomicAttack`` instances. The default + implementation delegates to ``MatrixAtomicAttackBuilder`` using the + ``_get_attack_technique_factories()`` and ``_build_display_group()`` hooks, producing + one ``AtomicAttack`` per (technique × dataset) pair plus an optional baseline. + + Scenarios with custom construction (composite attacks, per-objective technique + selection, converter stacks) override this method and build their attacks directly + or via another builder. Overrides own their baseline emission; the default emits one + when ``context.include_baseline`` is set. + + Args: + context (ScenarioContext): The resolved runtime inputs for this run. + + Returns: + list[AtomicAttack]: The generated atomic attacks. + """ + selected_techniques = {s.value for s in context.scenario_strategies} + all_factories = self._get_attack_technique_factories() + + technique_factories: dict[str, AttackTechniqueFactory] = {} for technique_name in selected_techniques: - factory = factories.get(technique_name) + factory = all_factories.get(technique_name) if factory is None: logger.warning(f"No factory for technique '{technique_name}', skipping.") continue + technique_factories[technique_name] = factory - for dataset_name, seed_groups in seed_groups_by_dataset.items(): - if factory.seed_technique is not None: - compatible_groups = SeedAttackGroup.filter_compatible( - seed_groups=seed_groups, - technique=factory.seed_technique, - ) - skipped = len(seed_groups) - len(compatible_groups) - if skipped: - logger.info( - f"Skipped {skipped} seed group(s) from '{dataset_name}' for technique " - f"'{technique_name}' (prompt sequences overlap with simulated conversation)." - ) - if not compatible_groups: - logger.warning( - f"No compatible seed groups in '{dataset_name}' for technique " - f"'{technique_name}', skipping this (technique, dataset) pair." - ) - continue - else: - compatible_groups = list(seed_groups) - - attack_technique = factory.create( - objective_target=self._objective_target, - attack_scoring_config=scoring_config, - ) - display_group = self._build_display_group( - technique_name=technique_name, - seed_group_name=dataset_name, - ) - atomic_attacks.append( - AtomicAttack( - atomic_attack_name=f"{technique_name}_{dataset_name}", - attack_technique=attack_technique, - seed_groups=list(compatible_groups), - adversarial_chat=factory.adversarial_chat, - objective_scorer=cast("TrueFalseScorer", self._objective_scorer), - memory_labels=self._memory_labels, - display_group=display_group, - ) - ) - - if self._include_baseline: - all_seed_groups = [g for groups in seed_groups_by_dataset.values() for g in groups] - atomic_attacks.insert(0, self._build_baseline_atomic_attack(seed_groups=all_seed_groups)) - - return atomic_attacks + builder = MatrixAtomicAttackBuilder( + objective_target=context.objective_target, + objective_scorer=self._objective_scorer, + memory_labels=context.memory_labels, + ) + return builder.build( + technique_factories=technique_factories, + dataset_groups=context.dataset_config.get_seed_attack_groups(), + display_group_fn=lambda combo: self._build_display_group( + technique_name=combo.technique_name, + seed_group_name=combo.dataset_name, + ), + include_baseline=context.include_baseline, + ) async def run_async(self) -> ScenarioResult: """ diff --git a/pyrit/scenario/core/scenario_context.py b/pyrit/scenario/core/scenario_context.py new file mode 100644 index 0000000000..8e90d26c0b --- /dev/null +++ b/pyrit/scenario/core/scenario_context.py @@ -0,0 +1,53 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT license. + +""" +Resolved runtime inputs for building a scenario's atomic attacks. + +``ScenarioContext`` is the single bundle of values a scenario needs to construct +its ``AtomicAttack`` list. The base ``Scenario`` resolves these during +``initialize_async`` and passes them to ``_build_atomic_attacks_async``, so +scenario authors never read half-initialized ``self._*`` state to build attacks. +""" + +from __future__ import annotations + +from dataclasses import dataclass, field +from typing import TYPE_CHECKING + +if TYPE_CHECKING: + from collections.abc import Sequence + + from pyrit.prompt_target import PromptTarget + from pyrit.scenario.core.dataset_configuration import DatasetConfiguration + from pyrit.scenario.core.scenario_strategy import ScenarioStrategy + + +@dataclass(frozen=True) +class ScenarioContext: + """ + Immutable snapshot of the inputs needed to build a scenario's atomic attacks. + + Constructed by ``Scenario._build_scenario_context`` from the values resolved in + ``initialize_async`` and handed to ``_build_atomic_attacks_async``. Because the + method receives a fully-populated context, scenarios no longer depend on the + "set ``self._x`` in ``initialize_async``, then read it in the override" sequencing + that historically caused "accessed before initialize_async" bugs. + + Attributes: + objective_target (PromptTarget): The target system the scenario attacks. + scenario_strategies (Sequence[ScenarioStrategy]): The resolved, concrete + strategies selected for this run (aggregates already expanded). + dataset_config (DatasetConfiguration): The effective dataset configuration + (caller-supplied or the scenario's default). + memory_labels (dict[str, str]): Labels applied to every attack run. + include_baseline (bool): Whether a baseline atomic attack should be emitted + for this run, already resolved against the scenario's + ``BASELINE_ATTACK_POLICY``. + """ + + objective_target: PromptTarget + scenario_strategies: Sequence[ScenarioStrategy] + dataset_config: DatasetConfiguration + memory_labels: dict[str, str] = field(default_factory=dict) + include_baseline: bool = False diff --git a/pyrit/scenario/scenarios/adaptive/adaptive_scenario.py b/pyrit/scenario/scenarios/adaptive/adaptive_scenario.py index 081bfca989..0a0389dd24 100644 --- a/pyrit/scenario/scenarios/adaptive/adaptive_scenario.py +++ b/pyrit/scenario/scenarios/adaptive/adaptive_scenario.py @@ -41,6 +41,7 @@ from pyrit.prompt_target import PromptTarget from pyrit.scenario.core.attack_technique_factory import AttackTechniqueFactory from pyrit.scenario.core.dataset_configuration import DatasetConfiguration + from pyrit.scenario.core.scenario_context import ScenarioContext from pyrit.scenario.core.scenario_strategy import ScenarioStrategy from pyrit.score import TrueFalseScorer @@ -159,7 +160,7 @@ def _get_attack_technique_factories(self) -> dict[str, AttackTechniqueFactory]: registry_overrides = {} return {**catalog, **registry_overrides} - async def _get_atomic_attacks_async(self) -> list[AtomicAttack]: + async def _build_atomic_attacks_async(self, *, context: ScenarioContext) -> list[AtomicAttack]: """ Build one ``AtomicAttack`` per (dataset, compatible seed group) pair. @@ -171,25 +172,24 @@ async def _get_atomic_attacks_async(self) -> list[AtomicAttack]: accumulates globally; selection is committed up-front during scenario initialization, before any execution starts. - When ``self._include_baseline`` is true (the default under + When ``context.include_baseline`` is true (the default under ``BASELINE_ATTACK_POLICY = Enabled``), a baseline ``AtomicAttack`` named ``"baseline"`` is prepended at index 0. + Args: + context (ScenarioContext): The resolved runtime inputs for this run. + Returns: list[AtomicAttack]: One ``AtomicAttack`` per compatible seed group across all datasets, with the baseline (when enabled) prepended at index 0. Raises: - ValueError: If ``self._objective_target`` is not set, or if - ``_build_techniques_dict`` finds no usable techniques. + ValueError: If ``_build_techniques_dict`` finds no usable techniques. """ - if self._objective_target is None: - raise ValueError("objective_target must be set before creating attacks") - - techniques = self._build_techniques_dict(objective_target=self._objective_target) + techniques = self._build_techniques_dict(objective_target=context.objective_target) - seed_groups_by_dataset = self._dataset_config.get_seed_attack_groups() + seed_groups_by_dataset = context.dataset_config.get_seed_attack_groups() atomic_attacks: list[AtomicAttack] = [] for dataset_name, seed_groups in seed_groups_by_dataset.items(): atomic_attacks.extend( @@ -201,7 +201,7 @@ async def _get_atomic_attacks_async(self) -> list[AtomicAttack]: ) ) - if self._include_baseline: + if context.include_baseline: all_seed_groups = [g for groups in seed_groups_by_dataset.values() for g in groups] atomic_attacks.insert(0, self._build_baseline_atomic_attack(seed_groups=all_seed_groups)) diff --git a/pyrit/scenario/scenarios/airt/jailbreak.py b/pyrit/scenario/scenarios/airt/jailbreak.py index 5184632d49..0980cd6be3 100644 --- a/pyrit/scenario/scenarios/airt/jailbreak.py +++ b/pyrit/scenario/scenarios/airt/jailbreak.py @@ -24,6 +24,7 @@ from pyrit.scenario.core.attack_technique import AttackTechnique from pyrit.scenario.core.dataset_configuration import DatasetConfiguration from pyrit.scenario.core.scenario import Scenario +from pyrit.scenario.core.scenario_context import ScenarioContext from pyrit.scenario.core.scenario_strategy import ScenarioStrategy from pyrit.scenario.core.scenario_target_defaults import get_default_adversarial_target from pyrit.score import ( @@ -267,12 +268,15 @@ async def _get_atomic_attack_from_strategy_async( seed_groups=self._seed_groups or [], ) - async def _get_atomic_attacks_async(self) -> list[AtomicAttack]: + async def _build_atomic_attacks_async(self, *, context: ScenarioContext) -> list[AtomicAttack]: """ Generate atomic attacks for each jailbreak template. This method creates an atomic attack for each retrieved jailbreak template. + Args: + context (ScenarioContext): The resolved runtime inputs for this run. + Returns: list[AtomicAttack]: List of atomic attacks to execute, one per jailbreak template. """ @@ -281,7 +285,7 @@ async def _get_atomic_attacks_async(self) -> list[AtomicAttack]: # Retrieve seed prompts based on selected strategies self._seed_groups = self._resolve_seed_groups() - strategies = {s.value for s in self._scenario_strategies} + strategies = {s.value for s in context.scenario_strategies} for strategy in strategies: for template_name in self._jailbreaks: @@ -291,7 +295,7 @@ async def _get_atomic_attacks_async(self) -> list[AtomicAttack]: ) atomic_attacks.append(atomic_attack) - if self._include_baseline: + if context.include_baseline: atomic_attacks.insert(0, self._build_baseline_atomic_attack(seed_groups=self._seed_groups or [])) return atomic_attacks diff --git a/pyrit/scenario/scenarios/airt/psychosocial.py b/pyrit/scenario/scenarios/airt/psychosocial.py index 19a20fd4a8..11c25fca6f 100644 --- a/pyrit/scenario/scenarios/airt/psychosocial.py +++ b/pyrit/scenario/scenarios/airt/psychosocial.py @@ -32,6 +32,7 @@ from pyrit.scenario.core.attack_technique import AttackTechnique from pyrit.scenario.core.dataset_configuration import DatasetConfiguration from pyrit.scenario.core.scenario import Scenario +from pyrit.scenario.core.scenario_context import ScenarioContext from pyrit.scenario.core.scenario_strategy import ( ScenarioStrategy, ) @@ -394,15 +395,13 @@ def _get_scorer(self, subharm: str | None = None) -> FloatScaleThresholdScorer: # threshold=1.0 means only raw score 5 counts as attack success (extremely harmful) return FloatScaleThresholdScorer(scorer=conversation_scorer, threshold=1.0) - async def _get_atomic_attacks_async(self) -> list[AtomicAttack]: - if self._objective_target is None: - raise ValueError("objective_target must be set before creating attacks") + async def _build_atomic_attacks_async(self, *, context: ScenarioContext) -> list[AtomicAttack]: try: - CHAT_TARGET_REQUIREMENTS.validate(target=self._objective_target) + CHAT_TARGET_REQUIREMENTS.validate(target=context.objective_target) except ValueError as exc: raise TypeError( f"PsychosocialHarmsScenario requires a target that supports multi-turn " - f"conversations with editable history. Target {type(self._objective_target).__name__} " + f"conversations with editable history. Target {type(context.objective_target).__name__} " f"does not satisfy these requirements: {exc}" ) from exc resolved = self._resolve_seed_groups() @@ -419,7 +418,7 @@ async def _get_atomic_attacks_async(self) -> list[AtomicAttack]: ), ] - if self._include_baseline: + if context.include_baseline: atomic_attacks.insert(0, self._build_baseline_atomic_attack(seed_groups=self._seed_groups)) return atomic_attacks diff --git a/pyrit/scenario/scenarios/airt/scam.py b/pyrit/scenario/scenarios/airt/scam.py index e591f580c8..6c31ad1f34 100644 --- a/pyrit/scenario/scenarios/airt/scam.py +++ b/pyrit/scenario/scenarios/airt/scam.py @@ -27,6 +27,7 @@ from pyrit.scenario.core.attack_technique import AttackTechnique from pyrit.scenario.core.dataset_configuration import DatasetConfiguration from pyrit.scenario.core.scenario import Scenario +from pyrit.scenario.core.scenario_context import ScenarioContext from pyrit.scenario.core.scenario_strategy import ScenarioStrategy from pyrit.scenario.core.scenario_target_defaults import get_default_adversarial_target from pyrit.score import TrueFalseScorer @@ -241,21 +242,24 @@ def _get_atomic_attack_from_strategy(self, strategy: str) -> AtomicAttack: memory_labels=self._memory_labels, ) - async def _get_atomic_attacks_async(self) -> list[AtomicAttack]: + async def _build_atomic_attacks_async(self, *, context: ScenarioContext) -> list[AtomicAttack]: """ Generate atomic attacks for each strategy. + Args: + context (ScenarioContext): The resolved runtime inputs for this run. + Returns: list[AtomicAttack]: List of atomic attacks to execute. """ # Resolve seed groups from deprecated objectives or dataset config self._seed_groups = self._resolve_seed_groups() - strategies = {s.value for s in self._scenario_strategies} + strategies = {s.value for s in context.scenario_strategies} atomic_attacks = [self._get_atomic_attack_from_strategy(strategy) for strategy in strategies] - if self._include_baseline: + if context.include_baseline: atomic_attacks.insert(0, self._build_baseline_atomic_attack(seed_groups=self._seed_groups or [])) return atomic_attacks diff --git a/pyrit/scenario/scenarios/benchmark/adversarial.py b/pyrit/scenario/scenarios/benchmark/adversarial.py index 0c0b4f6fb2..7311fcb14f 100644 --- a/pyrit/scenario/scenarios/benchmark/adversarial.py +++ b/pyrit/scenario/scenarios/benchmark/adversarial.py @@ -11,22 +11,22 @@ from pyrit.analytics import get_cached_results_for_technique from pyrit.common import Parameter, apply_defaults -from pyrit.executor.attack import AttackScoringConfig from pyrit.models import ( AttackOutcome, AttackResult, ObjectiveTargetEvaluationIdentifier, ScenarioResult, - SeedAttackGroup, ) from pyrit.registry import AttackTechniqueRegistry, TargetRegistry from pyrit.registry.tag_query import TagQuery -from pyrit.scenario.core.atomic_attack import AtomicAttack from pyrit.scenario.core.dataset_configuration import DatasetConfiguration +from pyrit.scenario.core.matrix_atomic_attack_builder import MatrixAtomicAttackBuilder from pyrit.scenario.core.scenario import BaselineAttackPolicy, Scenario if TYPE_CHECKING: from pyrit.prompt_target import PromptTarget + from pyrit.scenario.core.atomic_attack import AtomicAttack + from pyrit.scenario.core.scenario_context import ScenarioContext from pyrit.scenario.core.scenario_strategy import ScenarioStrategy from pyrit.score.true_false.true_false_scorer import TrueFalseScorer @@ -50,7 +50,7 @@ def _build_benchmark_strategy() -> type[ScenarioStrategy]: / ``multi_turn`` aggregates derived from each factory's ``strategy_tags``. The (technique × target) cross-product is materialized lazily in - ``AdversarialBenchmark._get_atomic_attacks_async`` from the + ``AdversarialBenchmark._build_atomic_attacks_async`` from the user-supplied ``adversarial_targets`` parameter. Returns: @@ -84,7 +84,7 @@ class AdversarialBenchmark(Scenario): ``TargetInitializer`` from ``ADVERSARIAL_CHAT_*`` env vars, or programmatically via ``TargetRegistry.register_instance``. - At run time, ``_get_atomic_attacks_async`` performs the + At run time, ``_build_atomic_attacks_async`` performs the ``(technique × adversarial_target × dataset)`` cross-product: for each selected adversarial-capable ``core`` factory in the ``AttackTechniqueRegistry`` and each requested target, it calls @@ -113,7 +113,7 @@ def supported_parameters(cls) -> list[Parameter]: Declare the ``adversarial_targets`` parameter. The list is treated as required at run time: - ``_get_atomic_attacks_async`` raises ``ValueError`` if + ``_build_atomic_attacks_async`` raises ``ValueError`` if ``self.params["adversarial_targets"]`` is empty or missing. The scenario-side error (rather than a declaration-side default) lets the caller raise a domain-specific message that names the CLI flag, @@ -157,7 +157,7 @@ def __init__( up by an initializer). Widening to general ``Scorer`` support (covering ``FloatScaleScorer``, etc.) is tracked as a follow-up. - use_cached: When ``True``, ``_get_atomic_attacks_async`` filters + use_cached: When ``True``, ``_build_atomic_attacks_async`` filters out atomic attacks for which the live behavioral cache (``pyrit.analytics.get_cached_results_for_technique``) has already returned at least one ``SUCCESS`` or ``FAILURE`` @@ -193,36 +193,31 @@ def __init__( scenario_result_id=scenario_result_id, ) - async def _get_atomic_attacks_async(self) -> list[AtomicAttack]: + async def _build_atomic_attacks_async(self, *, context: ScenarioContext) -> list[AtomicAttack]: """ Build atomic attacks from (technique × adversarial_target × dataset), then apply caching. - Reads the user-supplied ``adversarial_targets`` parameter, resolves - each name to a ``PromptTarget`` via ``TargetRegistry``, and - cross-products the selected adversarial-capable techniques over the - resolved targets and configured datasets. Each pair calls - ``factory.create(adversarial_chat=...)`` with the - resolved target — no global registry state is touched. When - ``self._use_cached`` is set, the final candidate list is filtered - against the live behavioral cache via + Reads the user-supplied ``adversarial_targets`` parameter, resolves each name to a + ``PromptTarget`` via ``TargetRegistry``, and delegates the + ``(technique × target × dataset)`` cross-product to ``MatrixAtomicAttackBuilder`` + with the resolved targets as its adversarial-target axis. Each pair calls + ``factory.create(adversarial_chat=...)`` with the resolved target — no global + registry state is touched. When ``self._use_cached`` is set, the resulting candidate + list is filtered against the live behavioral cache via ``_collect_cached_completion_pairs``, which delegates to - ``pyrit.analytics.get_cached_results_for_technique`` for each - unique ``(technique_eval_hash, objective_target_eval_hash)`` pair. + ``pyrit.analytics.get_cached_results_for_technique`` for each unique + ``(technique_eval_hash, objective_target_eval_hash)`` pair. + + Args: + context (ScenarioContext): The resolved runtime inputs for this run. Returns: - list[AtomicAttack]: The atomic attacks to actually execute on - this run. + list[AtomicAttack]: The atomic attacks to actually execute on this run. Raises: - ValueError: If the scenario has not been initialized, if - ``adversarial_targets`` is missing/empty, or if any name in + ValueError: If ``adversarial_targets`` is missing/empty, or if any name in ``adversarial_targets`` is not registered. """ - if self._objective_target is None: - raise ValueError( - "Scenario not properly initialized. Call await scenario.initialize_async() before running." - ) - target_names = self.params.get("adversarial_targets") if not target_names: raise ValueError( @@ -234,57 +229,28 @@ async def _get_atomic_attacks_async(self) -> list[AtomicAttack]: resolved_targets = self._resolve_adversarial_targets(target_names=target_names) all_factories = AttackTechniqueRegistry.get_registry_singleton().get_factories_or_raise() - selected_factories = [all_factories[s.value] for s in self._scenario_strategies if s.value in all_factories] - - scoring_config = AttackScoringConfig(objective_scorer=self._objective_scorer) - seed_groups_by_dataset = self._dataset_config.get_seed_attack_groups() - - atomic_attacks: list[AtomicAttack] = [] - for factory in selected_factories: - for target_name, target_instance in resolved_targets: - for dataset_name, seed_groups in seed_groups_by_dataset.items(): - if factory.seed_technique is not None: - compatible_groups = SeedAttackGroup.filter_compatible( - seed_groups=seed_groups, - technique=factory.seed_technique, - ) - skipped = len(seed_groups) - len(compatible_groups) - if skipped: - logger.info( - f"Skipped {skipped} seed group(s) from '{dataset_name}' for technique " - f"'{factory.name}' (prompt sequences overlap with simulated conversation)." - ) - if not compatible_groups: - logger.warning( - f"No compatible seed groups in '{dataset_name}' for technique " - f"'{factory.name}', skipping this (technique, target, dataset) triple." - ) - continue - else: - compatible_groups = list(seed_groups) - - attack_technique = factory.create( - objective_target=self._objective_target, - attack_scoring_config=scoring_config, - adversarial_chat=target_instance, - ) - # ``display_group`` is set explicitly here so result roll-ups group by the - # TargetRegistry name the caller passed via ``--adversarial-targets`` — - # not by any internal field on the PromptTarget instance (e.g. ``_model_name``). - # Because we override ``_get_atomic_attacks_async`` entirely, the base - # ``Scenario._build_display_group`` hook is never consulted; ``Scenario._finalize`` - # then reads ``aa.display_group`` directly (scenario.py:721). - atomic_attacks.append( - AtomicAttack( - atomic_attack_name=f"{factory.name}__{target_name}_{dataset_name}", - attack_technique=attack_technique, - seed_groups=list(compatible_groups), - adversarial_chat=target_instance, - objective_scorer=self._objective_scorer, - memory_labels=self._memory_labels, - display_group=target_name, - ) - ) + technique_factories = { + strategy.value: all_factories[strategy.value] + for strategy in context.scenario_strategies + if strategy.value in all_factories + } + + builder = MatrixAtomicAttackBuilder( + objective_target=context.objective_target, + objective_scorer=self._objective_scorer, + memory_labels=context.memory_labels, + ) + # ``display_group`` is the TargetRegistry name the caller passed via + # ``--adversarial-targets`` so per-model ASR rolls up naturally — not any internal + # field on the PromptTarget instance (e.g. ``_model_name``). The builder's default + # ``{technique}__{target}_{dataset}`` naming preserves the VERSION=2 cache key shape. + atomic_attacks = builder.build( + technique_factories=technique_factories, + dataset_groups=context.dataset_config.get_seed_attack_groups(), + adversarial_targets=resolved_targets, + display_group_fn=lambda combo: combo.target_name or "", + include_baseline=False, + ) if not self._use_cached: return atomic_attacks @@ -352,7 +318,7 @@ async def run_async(self) -> ScenarioResult: Run the scenario and merge any precomputed cached results into the returned ``ScenarioResult``. When ``use_cached=True`` skipped atomic attacks whose prior results were - loaded during ``_get_atomic_attacks_async``, this override attaches + loaded during ``_build_atomic_attacks_async``, this override attaches those results (and their display-group labels) to the live scenario result so the final report reflects both newly-executed and cache-served runs. @@ -400,12 +366,12 @@ def _collect_cached_completion_pairs(self, *, atomic_attacks: list[AtomicAttack] As a side effect, populates ``self._cached_results_by_name`` with the attribution-filtered ``AttackResult`` lists keyed by ``atomic_attack_name`` so that - ``_get_atomic_attacks_async`` can inject them into the final ``ScenarioResult`` + ``_build_atomic_attacks_async`` can inject them into the final ``ScenarioResult`` via ``run_async`` without re-filtering. Args: atomic_attacks: The candidate atomic attacks built earlier in - ``_get_atomic_attacks_async``. + ``_build_atomic_attacks_async``. Returns: set[str]: ``atomic_attack_name`` values that have at least one qualifying cached diff --git a/pyrit/scenario/scenarios/foundry/red_team_agent.py b/pyrit/scenario/scenarios/foundry/red_team_agent.py index 9096fcc9f2..2ea0375200 100644 --- a/pyrit/scenario/scenarios/foundry/red_team_agent.py +++ b/pyrit/scenario/scenarios/foundry/red_team_agent.py @@ -64,6 +64,7 @@ from pyrit.scenario.core.attack_technique import AttackTechnique from pyrit.scenario.core.dataset_configuration import DatasetConfiguration from pyrit.scenario.core.scenario import Scenario +from pyrit.scenario.core.scenario_context import ScenarioContext from pyrit.scenario.core.scenario_strategy import ScenarioCompositeStrategy, ScenarioStrategy from pyrit.scenario.core.scenario_target_defaults import get_default_adversarial_target @@ -398,9 +399,12 @@ def _resolve_seed_groups(self) -> list[SeedAttackGroup]: """ return self._dataset_config.get_all_seed_attack_groups() - async def _get_atomic_attacks_async(self) -> list[AtomicAttack]: + async def _build_atomic_attacks_async(self, *, context: ScenarioContext) -> list[AtomicAttack]: """ - Retrieve the list of AtomicAttack instances in this scenario. + Build one ``AtomicAttack`` per resolved FoundryComposite. + + Args: + context (ScenarioContext): The resolved runtime inputs for this run. Returns: list[AtomicAttack]: The list of AtomicAttack instances in this scenario. @@ -410,7 +414,7 @@ async def _get_atomic_attacks_async(self) -> list[AtomicAttack]: atomic_attacks = [self._get_attack_from_strategy(composition) for composition in self._scenario_composites] - if self._include_baseline: + if context.include_baseline: atomic_attacks.insert(0, self._build_baseline_atomic_attack(seed_groups=self._seed_groups)) return atomic_attacks diff --git a/pyrit/scenario/scenarios/garak/encoding.py b/pyrit/scenario/scenarios/garak/encoding.py index abe36b7ca6..08e374e31d 100644 --- a/pyrit/scenario/scenarios/garak/encoding.py +++ b/pyrit/scenario/scenarios/garak/encoding.py @@ -36,6 +36,7 @@ from pyrit.scenario.core.attack_technique import AttackTechnique from pyrit.scenario.core.dataset_configuration import DatasetConfiguration from pyrit.scenario.core.scenario import Scenario +from pyrit.scenario.core.scenario_context import ScenarioContext from pyrit.scenario.core.scenario_strategy import ScenarioStrategy from pyrit.score import TrueFalseScorer from pyrit.score.true_false.decoding_scorer import DecodingScorer @@ -182,7 +183,7 @@ def __init__( ) self._legacy_include_baseline = include_baseline - # Will be resolved in _get_atomic_attacks_async + # Will be resolved in _build_atomic_attacks_async self._resolved_seed_groups: list[SeedAttackGroup] | None = None def _resolve_seed_groups(self) -> list[SeedAttackGroup]: @@ -200,9 +201,16 @@ def _resolve_seed_groups(self) -> list[SeedAttackGroup]: return seed_groups - async def _get_atomic_attacks_async(self) -> list[AtomicAttack]: + async def _build_atomic_attacks_async(self, *, context: ScenarioContext) -> list[AtomicAttack]: """ - Retrieve the list of AtomicAttack instances in this scenario. + Build the encoding atomic attacks for this run. + + Encoding builds attacks directly (one ``AtomicAttack`` per selected encoding scheme, + each fanned out over the decode templates) rather than via the matrix builder, since + its axis is converter configurations, not techniques. + + Args: + context (ScenarioContext): The resolved runtime inputs for this run. Returns: list[AtomicAttack]: The list of AtomicAttack instances in this scenario. @@ -212,7 +220,7 @@ async def _get_atomic_attacks_async(self) -> list[AtomicAttack]: atomic_attacks = self._get_converter_attacks() - if self._include_baseline: + if context.include_baseline: atomic_attacks.insert(0, self._build_baseline_atomic_attack(seed_groups=self._resolved_seed_groups or [])) return atomic_attacks diff --git a/tests/unit/scenario/core/test_matrix_atomic_attack_builder.py b/tests/unit/scenario/core/test_matrix_atomic_attack_builder.py new file mode 100644 index 0000000000..94c4fd5977 --- /dev/null +++ b/tests/unit/scenario/core/test_matrix_atomic_attack_builder.py @@ -0,0 +1,284 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT license. + +"""Tests for ``MatrixAtomicAttackBuilder`` and its module-level helpers. + +The builder centralizes the technique × dataset (× optional adversarial target) +cross-product that scenarios previously duplicated in their own +``_get_atomic_attacks_async`` overrides. These tests pin the contract: + +* cross-product cardinality across techniques, datasets, and the optional + adversarial-target axis, +* default and custom ``atomic_attack_name`` / ``display_group`` derivation, +* ``factory.create`` adversarial-chat forwarding (and the ``AtomicAttack`` + ``adversarial_chat`` it stamps), +* seed-technique compatibility filtering (skip vs. subset), and +* optional baseline emission from the flattened seed groups. +""" + +from unittest.mock import MagicMock, patch + +import pytest + +from pyrit.models import SeedAttackGroup, SeedObjective +from pyrit.prompt_target import PromptTarget +from pyrit.scenario.core.attack_technique_factory import AttackTechniqueFactory +from pyrit.scenario.core.matrix_atomic_attack_builder import ( + MatrixAtomicAttackBuilder, + MatrixCombo, + build_baseline_atomic_attack, +) +from pyrit.score import TrueFalseScorer + + +def _mock_factory(*, name: str, seed_technique=None, adversarial_chat=None) -> MagicMock: + """Build a controllable ``AttackTechniqueFactory`` stand-in. + + ``create`` returns a fresh sentinel ``AttackTechnique`` so callers can assert + on the (objective_target, scoring_config, adversarial_chat) kwargs it received. + """ + factory = MagicMock(spec=AttackTechniqueFactory) + factory.name = name + factory.seed_technique = seed_technique + factory.adversarial_chat = adversarial_chat + factory.create.return_value = MagicMock(name=f"{name}_technique") + return factory + + +def _seed_group(*, objective: str) -> SeedAttackGroup: + return SeedAttackGroup(seeds=[SeedObjective(value=objective)]) + + +def _builder() -> MatrixAtomicAttackBuilder: + return MatrixAtomicAttackBuilder( + objective_target=MagicMock(spec=PromptTarget), + objective_scorer=MagicMock(spec=TrueFalseScorer), + memory_labels={"op": "unit"}, + ) + + +@pytest.mark.usefixtures("patch_central_database") +class TestMatrixComboNaming: + """Default ``atomic_attack_name`` / ``display_group`` helpers via ``MatrixCombo``.""" + + def test_default_name_without_target(self): + builder = _builder() + result = builder.build( + technique_factories={"tech": _mock_factory(name="tech")}, + dataset_groups={"ds": [_seed_group(objective="o1")]}, + ) + assert [a.atomic_attack_name for a in result] == ["tech_ds"] + + def test_default_name_with_target_axis(self): + builder = _builder() + result = builder.build( + technique_factories={"tech": _mock_factory(name="tech")}, + dataset_groups={"ds": [_seed_group(objective="o1")]}, + adversarial_targets=[("advA", MagicMock(spec=PromptTarget))], + ) + assert [a.atomic_attack_name for a in result] == ["tech__advA_ds"] + + def test_default_display_group_is_technique(self): + builder = _builder() + result = builder.build( + technique_factories={"tech": _mock_factory(name="tech")}, + dataset_groups={"ds": [_seed_group(objective="o1")]}, + ) + assert result[0].display_group == "tech" + + +@pytest.mark.usefixtures("patch_central_database") +class TestMatrixBuildCrossProduct: + """Cardinality and ordering of the produced cross-product.""" + + def test_two_techniques_one_dataset_no_targets(self): + builder = _builder() + result = builder.build( + technique_factories={ + "alpha": _mock_factory(name="alpha"), + "beta": _mock_factory(name="beta"), + }, + dataset_groups={"ds": [_seed_group(objective="o1")]}, + ) + assert [a.atomic_attack_name for a in result] == ["alpha_ds", "beta_ds"] + + def test_one_technique_two_targets_one_dataset(self): + builder = _builder() + result = builder.build( + technique_factories={"tech": _mock_factory(name="tech")}, + dataset_groups={"ds": [_seed_group(objective="o1")]}, + adversarial_targets=[ + ("advA", MagicMock(spec=PromptTarget)), + ("advB", MagicMock(spec=PromptTarget)), + ], + ) + assert [a.atomic_attack_name for a in result] == ["tech__advA_ds", "tech__advB_ds"] + + def test_one_technique_one_target_two_datasets(self): + builder = _builder() + result = builder.build( + technique_factories={"tech": _mock_factory(name="tech")}, + dataset_groups={ + "ds1": [_seed_group(objective="o1")], + "ds2": [_seed_group(objective="o2")], + }, + adversarial_targets=[("advA", MagicMock(spec=PromptTarget))], + ) + assert [a.atomic_attack_name for a in result] == ["tech__advA_ds1", "tech__advA_ds2"] + + +@pytest.mark.usefixtures("patch_central_database") +class TestMatrixAdversarialForwarding: + """The adversarial-target axis must drive both ``factory.create`` and ``AtomicAttack``.""" + + def test_create_called_with_adversarial_chat_per_target(self): + builder = _builder() + factory = _mock_factory(name="tech") + target_a = MagicMock(spec=PromptTarget) + target_b = MagicMock(spec=PromptTarget) + builder.build( + technique_factories={"tech": factory}, + dataset_groups={"ds": [_seed_group(objective="o1")]}, + adversarial_targets=[("advA", target_a), ("advB", target_b)], + ) + assert factory.create.call_count == 2 + injected = {call.kwargs["adversarial_chat"] for call in factory.create.call_args_list} + assert injected == {target_a, target_b} + + def test_atomic_attack_adversarial_chat_is_resolved_target(self): + builder = _builder() + target_a = MagicMock(spec=PromptTarget) + result = builder.build( + technique_factories={"tech": _mock_factory(name="tech")}, + dataset_groups={"ds": [_seed_group(objective="o1")]}, + adversarial_targets=[("advA", target_a)], + ) + assert result[0]._adversarial_chat is target_a + + def test_no_target_axis_uses_factory_adversarial_chat(self): + builder = _builder() + baked = MagicMock(spec=PromptTarget) + factory = _mock_factory(name="tech", adversarial_chat=baked) + result = builder.build( + technique_factories={"tech": factory}, + dataset_groups={"ds": [_seed_group(objective="o1")]}, + ) + # No adversarial_chat is forwarded into create() when the axis is collapsed. + assert "adversarial_chat" not in factory.create.call_args.kwargs + assert result[0]._adversarial_chat is baked + + +@pytest.mark.usefixtures("patch_central_database") +class TestMatrixCustomCallbacks: + """Custom ``name_fn`` / ``display_group_fn`` override the defaults.""" + + def test_custom_name_and_display_group(self): + builder = _builder() + result = builder.build( + technique_factories={"tech": _mock_factory(name="tech")}, + dataset_groups={"ds": [_seed_group(objective="o1")]}, + adversarial_targets=[("advA", MagicMock(spec=PromptTarget))], + name_fn=lambda combo: f"{combo.target_name}:{combo.technique_name}", + display_group_fn=lambda combo: combo.target_name or "", + ) + assert result[0].atomic_attack_name == "advA:tech" + assert result[0].display_group == "advA" + + def test_callbacks_receive_full_combo(self): + builder = _builder() + seen: list[MatrixCombo] = [] + builder.build( + technique_factories={"tech": _mock_factory(name="tech")}, + dataset_groups={"ds": [_seed_group(objective="o1")]}, + adversarial_targets=[("advA", MagicMock(spec=PromptTarget))], + name_fn=lambda combo: seen.append(combo) or "n", + ) + assert seen == [MatrixCombo(technique_name="tech", dataset_name="ds", target_name="advA")] + + +@pytest.mark.usefixtures("patch_central_database") +class TestMatrixSeedTechniqueFiltering: + """``seed_technique`` gates which seed groups (and pairs) survive.""" + + def test_incompatible_pair_is_skipped(self): + builder = _builder() + factory = _mock_factory(name="tech", seed_technique=MagicMock()) + with patch.object(SeedAttackGroup, "filter_compatible", return_value=[]): + result = builder.build( + technique_factories={"tech": factory}, + dataset_groups={"ds": [_seed_group(objective="o1")]}, + ) + assert result == [] + factory.create.assert_not_called() + + def test_partial_filter_keeps_subset(self): + builder = _builder() + factory = _mock_factory(name="tech", seed_technique=MagicMock()) + kept = _seed_group(objective="keep") + dropped = _seed_group(objective="drop") + with patch.object(SeedAttackGroup, "filter_compatible", return_value=[kept]): + result = builder.build( + technique_factories={"tech": factory}, + dataset_groups={"ds": [kept, dropped]}, + ) + assert len(result) == 1 + assert result[0]._seed_groups == [kept] + + def test_no_seed_technique_keeps_all_groups(self): + builder = _builder() + groups = [_seed_group(objective="a"), _seed_group(objective="b")] + result = builder.build( + technique_factories={"tech": _mock_factory(name="tech")}, + dataset_groups={"ds": groups}, + ) + assert result[0]._seed_groups == groups + + +@pytest.mark.usefixtures("patch_central_database") +class TestMatrixBaseline: + """Baseline emission prepends a single ``baseline`` attack over flattened seeds.""" + + def test_baseline_prepended_when_requested(self): + builder = _builder() + result = builder.build( + technique_factories={"tech": _mock_factory(name="tech")}, + dataset_groups={"ds": [_seed_group(objective="o1")]}, + include_baseline=True, + ) + assert result[0].atomic_attack_name == "baseline" + assert [a.atomic_attack_name for a in result] == ["baseline", "tech_ds"] + + def test_baseline_omitted_by_default(self): + builder = _builder() + result = builder.build( + technique_factories={"tech": _mock_factory(name="tech")}, + dataset_groups={"ds": [_seed_group(objective="o1")]}, + ) + assert all(a.atomic_attack_name != "baseline" for a in result) + + def test_baseline_flattens_all_datasets(self): + builder = _builder() + g1 = _seed_group(objective="o1") + g2 = _seed_group(objective="o2") + result = builder.build( + technique_factories={"tech": _mock_factory(name="tech")}, + dataset_groups={"ds1": [g1], "ds2": [g2]}, + include_baseline=True, + ) + assert result[0]._seed_groups == [g1, g2] + + +@pytest.mark.usefixtures("patch_central_database") +class TestBuildBaselineHelper: + """The module-level ``build_baseline_atomic_attack`` helper.""" + + def test_baseline_name_and_seed_groups(self): + groups = [_seed_group(objective="o1")] + baseline = build_baseline_atomic_attack( + objective_target=MagicMock(spec=PromptTarget), + objective_scorer=MagicMock(spec=TrueFalseScorer), + seed_groups=groups, + memory_labels={"op": "unit"}, + ) + assert baseline.atomic_attack_name == "baseline" + assert baseline._seed_groups == groups diff --git a/tests/unit/scenario/core/test_scenario_context.py b/tests/unit/scenario/core/test_scenario_context.py new file mode 100644 index 0000000000..edf234447a --- /dev/null +++ b/tests/unit/scenario/core/test_scenario_context.py @@ -0,0 +1,62 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT license. + +"""Tests for the ``ScenarioContext`` frozen dataclass. + +``ScenarioContext`` is the immutable bundle the base ``Scenario`` builds in +``initialize_async`` and hands to ``_build_atomic_attacks_async``. These tests +pin the field surface, defaults, and immutability that scenario authors rely on. +""" + +from unittest.mock import MagicMock + +import pytest + +from pyrit.prompt_target import PromptTarget +from pyrit.scenario.core.dataset_configuration import DatasetConfiguration +from pyrit.scenario.core.scenario_context import ScenarioContext + + +def _context(**overrides) -> ScenarioContext: + kwargs = { + "objective_target": MagicMock(spec=PromptTarget), + "scenario_strategies": [], + "dataset_config": MagicMock(spec=DatasetConfiguration), + } + kwargs.update(overrides) + return ScenarioContext(**kwargs) + + +def test_required_fields_are_stored(): + target = MagicMock(spec=PromptTarget) + strategies = [MagicMock()] + dataset_config = MagicMock(spec=DatasetConfiguration) + context = ScenarioContext( + objective_target=target, + scenario_strategies=strategies, + dataset_config=dataset_config, + ) + assert context.objective_target is target + assert context.scenario_strategies is strategies + assert context.dataset_config is dataset_config + + +def test_memory_labels_defaults_to_empty_dict(): + context = _context() + assert context.memory_labels == {} + + +def test_memory_labels_default_is_not_shared_between_instances(): + first = _context() + second = _context() + assert first.memory_labels is not second.memory_labels + + +def test_include_baseline_defaults_to_false(): + assert _context().include_baseline is False + + +def test_is_frozen(): + context = _context() + with pytest.raises(AttributeError): + context.objective_target = MagicMock(spec=PromptTarget) From 1b547fccd6779f403fea6ecd5f812f29cf538b18 Mon Sep 17 00:00:00 2001 From: Richard Lundeen Date: Tue, 30 Jun 2026 15:24:33 -0700 Subject: [PATCH 2/4] MAINT: Fix ty/reST lint in MatrixAtomicAttackBuilder Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- pyrit/scenario/core/matrix_atomic_attack_builder.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/pyrit/scenario/core/matrix_atomic_attack_builder.py b/pyrit/scenario/core/matrix_atomic_attack_builder.py index 58f84fc323..a3660c63f9 100644 --- a/pyrit/scenario/core/matrix_atomic_attack_builder.py +++ b/pyrit/scenario/core/matrix_atomic_attack_builder.py @@ -34,6 +34,7 @@ from pyrit.prompt_target import PromptTarget from pyrit.scenario.core.attack_technique_factory import AttackTechniqueFactory from pyrit.score import Scorer + from pyrit.score.true_false.true_false_scorer import TrueFalseScorer logger = logging.getLogger(__name__) @@ -124,7 +125,7 @@ class MatrixAtomicAttackBuilder: Build ``AtomicAttack`` instances from a technique × dataset (× target) cross-product. Construct once with the shared run inputs (target, scorer, labels), then call - :meth:`build` with the per-run grid. The builder owns: + ``build`` with the per-run grid. The builder owns: - seed-technique compatibility filtering (``SeedAttackGroup.filter_compatible``), - the ``factory.create(...)`` call, forwarding an adversarial target when the From 76a875af3f9c568d270ea9e353bebc093b8de055 Mon Sep 17 00:00:00 2001 From: Richard Lundeen Date: Tue, 30 Jun 2026 21:13:32 -0700 Subject: [PATCH 3/4] MAINT: Remove transitional language from scenario docstrings Describe ScenarioContext, MatrixAtomicAttackBuilder, and the bridge by what they are rather than by what scenarios previously did, and refresh two stale _get_atomic_attacks_async references to the current method name. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- pyrit/scenario/core/matrix_atomic_attack_builder.py | 4 ++-- pyrit/scenario/core/scenario.py | 11 +++++------ pyrit/scenario/core/scenario_context.py | 11 +++++------ .../scenario/scenarios/adaptive/adaptive_scenario.py | 2 +- pyrit/scenario/scenarios/adaptive/dispatcher.py | 2 +- .../core/test_matrix_atomic_attack_builder.py | 3 +-- 6 files changed, 15 insertions(+), 18 deletions(-) diff --git a/pyrit/scenario/core/matrix_atomic_attack_builder.py b/pyrit/scenario/core/matrix_atomic_attack_builder.py index a3660c63f9..838c87d5a2 100644 --- a/pyrit/scenario/core/matrix_atomic_attack_builder.py +++ b/pyrit/scenario/core/matrix_atomic_attack_builder.py @@ -7,8 +7,8 @@ ``MatrixAtomicAttackBuilder`` turns a technique × dataset (× optional adversarial target) grid into a flat list of ``AtomicAttack`` instances. It centralizes the seed-technique compatibility filtering, ``factory.create`` wiring, ``AtomicAttack`` -construction, and baseline emission that scenarios previously duplicated in their -own ``_get_atomic_attacks_async`` overrides. +construction, and baseline emission for scenarios whose attacks form such a +cross-product. This is one member of a *family* of construction helpers named by shape (``Matrix...``); scenarios whose construction is composite or per-objective build diff --git a/pyrit/scenario/core/scenario.py b/pyrit/scenario/core/scenario.py index bd6bd1281d..0e6e548ef6 100644 --- a/pyrit/scenario/core/scenario.py +++ b/pyrit/scenario/core/scenario.py @@ -1005,13 +1005,12 @@ def _build_scenario_context(self) -> ScenarioContext: async def _get_atomic_attacks_async(self) -> list[AtomicAttack]: """ - Build this scenario's atomic attacks (internal entry point from ``initialize_async``). + Build this scenario's atomic attacks (internal entry point called by ``initialize_async``). - New scenarios should override ``_build_atomic_attacks_async(context=...)`` instead - of this method. The base implementation here builds a ``ScenarioContext`` from the - values resolved in ``initialize_async`` and forwards to - ``_build_atomic_attacks_async``. It remains as a stable, no-argument entry point so - legacy overrides and existing call sites keep working during the migration. + Builds a ``ScenarioContext`` from the values resolved in ``initialize_async`` and + forwards to ``_build_atomic_attacks_async`` — the extension point scenarios override + to customize attack construction. This stays a stable, no-argument entry point for + ``initialize_async`` and other internal callers. Returns: list[AtomicAttack]: The generated atomic attacks. diff --git a/pyrit/scenario/core/scenario_context.py b/pyrit/scenario/core/scenario_context.py index 6591208388..99caaadf15 100644 --- a/pyrit/scenario/core/scenario_context.py +++ b/pyrit/scenario/core/scenario_context.py @@ -6,8 +6,8 @@ ``ScenarioContext`` is the single bundle of values a scenario needs to construct its ``AtomicAttack`` list. The base ``Scenario`` resolves these during -``initialize_async`` and passes them to ``_build_atomic_attacks_async``, so -scenario authors never read half-initialized ``self._*`` state to build attacks. +``initialize_async`` and passes them to ``_build_atomic_attacks_async``, so a +scenario builds its attacks from the context it is given. """ from __future__ import annotations @@ -29,10 +29,9 @@ class ScenarioContext: Immutable snapshot of the inputs needed to build a scenario's atomic attacks. Constructed by ``Scenario._build_scenario_context`` from the values resolved in - ``initialize_async`` and handed to ``_build_atomic_attacks_async``. Because the - method receives a fully-populated context, scenarios no longer depend on the - "set ``self._x`` in ``initialize_async``, then read it in the override" sequencing - that historically caused "accessed before initialize_async" bugs. + ``initialize_async`` and handed to ``_build_atomic_attacks_async``. A scenario + builds its attacks from this context, independent of instance-attribute + initialization order. Attributes: objective_target (PromptTarget): The target system the scenario attacks. diff --git a/pyrit/scenario/scenarios/adaptive/adaptive_scenario.py b/pyrit/scenario/scenarios/adaptive/adaptive_scenario.py index 3a9a55343a..744c80b823 100644 --- a/pyrit/scenario/scenarios/adaptive/adaptive_scenario.py +++ b/pyrit/scenario/scenarios/adaptive/adaptive_scenario.py @@ -356,7 +356,7 @@ async def _build_atomics_for_dataset_async( Raises: ValueError: If ``self._objective_target`` is not set - (defensive guard; ``_get_atomic_attacks_async`` enforces + (defensive guard; ``_build_atomic_attacks_async`` enforces this earlier). """ if self._objective_target is None: # pragma: no cover - defensive diff --git a/pyrit/scenario/scenarios/adaptive/dispatcher.py b/pyrit/scenario/scenarios/adaptive/dispatcher.py index 7005e89d76..2469029537 100644 --- a/pyrit/scenario/scenarios/adaptive/dispatcher.py +++ b/pyrit/scenario/scenarios/adaptive/dispatcher.py @@ -93,7 +93,7 @@ class AdaptiveTechniqueDispatcher: reference across all calls in a scenario so learning accumulates across objectives — though all selections are committed up-front during scenario initialization (see - ``AdaptiveScenario._get_atomic_attacks_async``). + ``AdaptiveScenario._build_atomic_attacks_async``). """ def __init__( diff --git a/tests/unit/scenario/core/test_matrix_atomic_attack_builder.py b/tests/unit/scenario/core/test_matrix_atomic_attack_builder.py index 94c4fd5977..8b91b37561 100644 --- a/tests/unit/scenario/core/test_matrix_atomic_attack_builder.py +++ b/tests/unit/scenario/core/test_matrix_atomic_attack_builder.py @@ -4,8 +4,7 @@ """Tests for ``MatrixAtomicAttackBuilder`` and its module-level helpers. The builder centralizes the technique × dataset (× optional adversarial target) -cross-product that scenarios previously duplicated in their own -``_get_atomic_attacks_async`` overrides. These tests pin the contract: +cross-product for scenarios whose attacks form such a grid. These tests pin the contract: * cross-product cardinality across techniques, datasets, and the optional adversarial-target axis, From c3ef6b64ddb8a29a9f2bf037f479658987641f2a Mon Sep 17 00:00:00 2001 From: Richard Lundeen Date: Tue, 30 Jun 2026 22:20:14 -0700 Subject: [PATCH 4/4] MAINT: Centralize seed resolution and baseline in Scenario base Move seed-group resolution and baseline emission into the base Scenario bridge so custom scenarios stop duplicating _resolve_seed_groups_async, self._seed_groups state, and the include_baseline block. Add seed_groups and seed_groups_by_dataset to ScenarioContext, resolve seeds once via the new _resolve_seed_groups_by_dataset_async hook, and prepend the baseline centrally from the same sample (fixes the ADO 9012 double-sampling inconsistency). Simplify jailbreak, scam, encoding, red_team_agent, adaptive, adversarial, and psychosocial accordingly, and update coupled tests. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- pyrit/scenario/core/scenario.py | 62 +++++++-- pyrit/scenario/core/scenario_context.py | 12 +- .../scenarios/adaptive/adaptive_scenario.py | 18 +-- pyrit/scenario/scenarios/airt/jailbreak.py | 29 +--- pyrit/scenario/scenarios/airt/psychosocial.py | 56 +++----- pyrit/scenario/scenarios/airt/scam.py | 33 +---- .../scenarios/benchmark/adversarial.py | 2 +- .../scenarios/foundry/red_team_agent.py | 32 ++--- pyrit/scenario/scenarios/garak/encoding.py | 43 ++---- tests/unit/scenario/airt/test_jailbreak.py | 114 +++++++++++---- tests/unit/scenario/airt/test_psychosocial.py | 64 ++++++--- tests/unit/scenario/airt/test_scam.py | 50 +++++-- .../scenario/foundry/test_red_team_agent.py | 130 ++++++++++++++---- tests/unit/scenario/garak/test_encoding.py | 74 ++++++---- 14 files changed, 450 insertions(+), 269 deletions(-) diff --git a/pyrit/scenario/core/scenario.py b/pyrit/scenario/core/scenario.py index 0e6e548ef6..85868ea909 100644 --- a/pyrit/scenario/core/scenario.py +++ b/pyrit/scenario/core/scenario.py @@ -975,7 +975,24 @@ async def _get_remaining_atomic_attacks_async(self) -> list[AtomicAttack]: return remaining_attacks - def _build_scenario_context(self) -> ScenarioContext: + async def _resolve_seed_groups_by_dataset_async(self) -> dict[str, list[SeedAttackGroup]]: + """ + Resolve the seed groups this scenario attacks, keyed by originating dataset. + + This is the single place seed resolution happens for a run. The base ``Scenario`` + calls it once in the bridge, flattens the result into ``context.seed_groups``, and + reuses the same population for every atomic attack and the baseline — so sampling + under ``max_dataset_size`` stays consistent across all of them. + + Override to inject seeds from an alternate source (e.g. deprecated ``objectives``) + or to filter the resolved groups before attacks are built. + + Returns: + dict[str, list[SeedAttackGroup]]: Seed groups keyed by dataset name. + """ + return await self._dataset_config.get_attack_groups_by_dataset_async() + + def _build_scenario_context(self, *, seed_groups_by_dataset: dict[str, list[SeedAttackGroup]]) -> ScenarioContext: """ Snapshot the resolved runtime inputs into a ``ScenarioContext``. @@ -984,6 +1001,11 @@ def _build_scenario_context(self) -> ScenarioContext: handed to ``_build_atomic_attacks_async`` so scenario authors never read half-initialized ``self._*`` state to build attacks. + Args: + seed_groups_by_dataset (dict[str, list[SeedAttackGroup]]): Seed groups already + resolved once (see ``_resolve_seed_groups_by_dataset_async``). The flat + ``context.seed_groups`` is derived from these so both views share one sample. + Returns: ScenarioContext: The immutable inputs for atomic-attack construction. @@ -995,22 +1017,28 @@ def _build_scenario_context(self) -> ScenarioContext: "Scenario not properly initialized. Call await scenario.initialize_async() before running." ) + seed_groups = [group for groups in seed_groups_by_dataset.values() for group in groups] + return ScenarioContext( objective_target=self._objective_target, scenario_strategies=tuple(self._scenario_strategies), dataset_config=self._dataset_config, memory_labels=dict(self._memory_labels), include_baseline=self._include_baseline, + seed_groups=seed_groups, + seed_groups_by_dataset=seed_groups_by_dataset, ) async def _get_atomic_attacks_async(self) -> list[AtomicAttack]: """ Build this scenario's atomic attacks (internal entry point called by ``initialize_async``). - Builds a ``ScenarioContext`` from the values resolved in ``initialize_async`` and - forwards to ``_build_atomic_attacks_async`` — the extension point scenarios override - to customize attack construction. This stays a stable, no-argument entry point for - ``initialize_async`` and other internal callers. + Resolves the seed groups once, builds a ``ScenarioContext`` from the values resolved + in ``initialize_async``, and forwards to ``_build_atomic_attacks_async`` — the extension + point scenarios override to customize attack construction. The baseline is emitted + centrally here (when ``context.include_baseline`` is set) from ``context.seed_groups``, + so overrides never re-resolve seeds or hand-roll baseline emission. This stays a stable, + no-argument entry point for ``initialize_async`` and other internal callers. Returns: list[AtomicAttack]: The generated atomic attacks. @@ -1018,8 +1046,16 @@ async def _get_atomic_attacks_async(self) -> list[AtomicAttack]: Raises: ValueError: If the scenario has not been initialized. """ - context = self._build_scenario_context() - return await self._build_atomic_attacks_async(context=context) + seed_groups_by_dataset = await self._resolve_seed_groups_by_dataset_async() + context = self._build_scenario_context(seed_groups_by_dataset=seed_groups_by_dataset) + atomic_attacks = await self._build_atomic_attacks_async(context=context) + + # Central baseline emission. Guarded so a scenario that still emits its own baseline + # (or an aggregate that legitimately has none) isn't given a duplicate. + if context.include_baseline and (not atomic_attacks or atomic_attacks[0].atomic_attack_name != "baseline"): + atomic_attacks.insert(0, self._build_baseline_atomic_attack(seed_groups=list(context.seed_groups))) + + return atomic_attacks async def _build_atomic_attacks_async(self, *, context: ScenarioContext) -> list[AtomicAttack]: """ @@ -1029,12 +1065,12 @@ async def _build_atomic_attacks_async(self, *, context: ScenarioContext) -> list scorers, and any extra axes into ``AtomicAttack`` instances. The default implementation delegates to ``MatrixAtomicAttackBuilder`` using the ``_get_attack_technique_factories()`` and ``_build_display_group()`` hooks, producing - one ``AtomicAttack`` per (technique × dataset) pair plus an optional baseline. + one ``AtomicAttack`` per (technique × dataset) pair. Scenarios with custom construction (composite attacks, per-objective technique - selection, converter stacks) override this method and build their attacks directly - or via another builder. Overrides own their baseline emission; the default emits one - when ``context.include_baseline`` is set. + selection, converter stacks) override this method and build their attacks from + ``context.seed_groups`` (or ``context.seed_groups_by_dataset``). The base owns baseline + emission, so overrides never prepend one themselves. Args: context (ScenarioContext): The resolved runtime inputs for this run. @@ -1060,12 +1096,12 @@ async def _build_atomic_attacks_async(self, *, context: ScenarioContext) -> list ) return builder.build( technique_factories=technique_factories, - dataset_groups=await context.dataset_config.get_attack_groups_by_dataset_async(), + dataset_groups=context.seed_groups_by_dataset, display_group_fn=lambda combo: self._build_display_group( technique_name=combo.technique_name, seed_group_name=combo.dataset_name, ), - include_baseline=context.include_baseline, + include_baseline=False, ) async def run_async(self) -> ScenarioResult: diff --git a/pyrit/scenario/core/scenario_context.py b/pyrit/scenario/core/scenario_context.py index 99caaadf15..28e6e8cfde 100644 --- a/pyrit/scenario/core/scenario_context.py +++ b/pyrit/scenario/core/scenario_context.py @@ -16,8 +16,9 @@ from typing import TYPE_CHECKING if TYPE_CHECKING: - from collections.abc import Sequence + from collections.abc import Mapping, Sequence + from pyrit.models import SeedAttackGroup from pyrit.prompt_target import PromptTarget from pyrit.scenario.core.dataset_configuration import DatasetAttackConfiguration from pyrit.scenario.core.scenario_strategy import ScenarioStrategy @@ -43,6 +44,13 @@ class ScenarioContext: include_baseline (bool): Whether a baseline atomic attack should be emitted for this run, already resolved against the scenario's ``BASELINE_ATTACK_POLICY``. + seed_groups (Sequence[SeedAttackGroup]): The scenario's seed groups, resolved + and sampled once by the base ``Scenario`` (flattened across datasets). Use + these to build attacks so every atomic attack — and the baseline — draws from + the same population. + seed_groups_by_dataset (Mapping[str, list[SeedAttackGroup]]): The same resolved + seed groups keyed by originating dataset name, for scenarios that map datasets + onto separate attacks or display groups. """ objective_target: PromptTarget @@ -50,3 +58,5 @@ class ScenarioContext: dataset_config: DatasetAttackConfiguration memory_labels: dict[str, str] = field(default_factory=dict) include_baseline: bool = False + seed_groups: Sequence[SeedAttackGroup] = field(default_factory=tuple) + seed_groups_by_dataset: Mapping[str, list[SeedAttackGroup]] = field(default_factory=dict) diff --git a/pyrit/scenario/scenarios/adaptive/adaptive_scenario.py b/pyrit/scenario/scenarios/adaptive/adaptive_scenario.py index 744c80b823..39ebb4ee91 100644 --- a/pyrit/scenario/scenarios/adaptive/adaptive_scenario.py +++ b/pyrit/scenario/scenarios/adaptive/adaptive_scenario.py @@ -167,31 +167,29 @@ async def _build_atomic_attacks_async(self, *, context: ScenarioContext) -> list For each dataset, construct a single ``AdaptiveTechniqueDispatcher`` shared across that dataset's seed groups. For each seed group, ask the dispatcher to build its per-objective ``SequentialAttack`` and - wrap it in its own ``AtomicAttack``. All dispatchers across all + wrap it in its own ``AtomicAttack``. All dispatchers across all datasets share one ``TechniqueSelector`` instance so learning accumulates globally; selection is committed up-front during scenario initialization, before any execution starts. - When ``context.include_baseline`` is true (the default under - ``BASELINE_ATTACK_POLICY = Enabled``), a baseline ``AtomicAttack`` - named ``"baseline"`` is prepended at index 0. + The base ``Scenario`` prepends the baseline ``AtomicAttack`` (named + ``"baseline"``) at index 0 when ``context.include_baseline`` is true (the + default under ``BASELINE_ATTACK_POLICY = Enabled``). Args: context (ScenarioContext): The resolved runtime inputs for this run. Returns: list[AtomicAttack]: One ``AtomicAttack`` per compatible - seed group across all datasets, with the baseline (when - enabled) prepended at index 0. + seed group across all datasets. Raises: ValueError: If ``_build_techniques_dict`` finds no usable techniques. """ techniques = self._build_techniques_dict(objective_target=context.objective_target) - seed_groups_by_dataset = await context.dataset_config.get_attack_groups_by_dataset_async() atomic_attacks: list[AtomicAttack] = [] - for dataset_name, seed_groups in seed_groups_by_dataset.items(): + for dataset_name, seed_groups in context.seed_groups_by_dataset.items(): atomic_attacks.extend( await self._build_atomics_for_dataset_async( dataset_name=dataset_name, @@ -201,10 +199,6 @@ async def _build_atomic_attacks_async(self, *, context: ScenarioContext) -> list ) ) - if context.include_baseline: - all_seed_groups = [g for groups in seed_groups_by_dataset.values() for g in groups] - atomic_attacks.insert(0, self._build_baseline_atomic_attack(seed_groups=all_seed_groups)) - return atomic_attacks def _build_techniques_dict( diff --git a/pyrit/scenario/scenarios/airt/jailbreak.py b/pyrit/scenario/scenarios/airt/jailbreak.py index 47c35e941c..9299a5406a 100644 --- a/pyrit/scenario/scenarios/airt/jailbreak.py +++ b/pyrit/scenario/scenarios/airt/jailbreak.py @@ -172,9 +172,6 @@ def __init__( ) self._legacy_include_baseline = include_baseline - # Will be resolved in _build_atomic_attacks_async - self._seed_groups: list[SeedAttackGroup] | None = None - def _get_or_create_adversarial_target(self) -> PromptTarget: """ Return the shared adversarial target, creating it on first access. @@ -189,20 +186,8 @@ def _get_or_create_adversarial_target(self) -> PromptTarget: self._adversarial_target = get_default_adversarial_target() return self._adversarial_target - async def _resolve_seed_groups_async(self) -> list[SeedAttackGroup]: - """ - Resolve seed groups from dataset configuration. - - Returns: - list[SeedAttackGroup]: List of seed attack groups with objectives to be tested. - """ - # Use dataset_config (guaranteed to be set by initialize_async). Auto-fetch - # populates memory first; a still-empty result raises a DatasetConstraintError - # naming the offending dataset, which we let propagate. - return list(await self._dataset_config.get_seed_attack_groups_async()) - async def _get_atomic_attack_from_strategy_async( - self, *, strategy: str, jailbreak_template_name: str + self, *, strategy: str, jailbreak_template_name: str, seed_groups: list[SeedAttackGroup] ) -> AtomicAttack: """ Create an atomic attack for a specific jailbreak template. @@ -210,6 +195,7 @@ async def _get_atomic_attack_from_strategy_async( Args: strategy (str): JailbreakStrategy to use. jailbreak_template_name (str): Name of the jailbreak template file. + seed_groups (list[SeedAttackGroup]): Seed groups the attack draws from. Returns: AtomicAttack: An atomic attack using the specified jailbreak template. @@ -264,7 +250,7 @@ async def _get_atomic_attack_from_strategy_async( return AtomicAttack( atomic_attack_name=f"jailbreak_{template_name}", attack_technique=AttackTechnique(attack=attack), - seed_groups=self._seed_groups or [], + seed_groups=seed_groups, ) async def _build_atomic_attacks_async(self, *, context: ScenarioContext) -> list[AtomicAttack]: @@ -281,20 +267,15 @@ async def _build_atomic_attacks_async(self, *, context: ScenarioContext) -> list """ atomic_attacks: list[AtomicAttack] = [] - # Retrieve seed prompts based on selected strategies - self._seed_groups = await self._resolve_seed_groups_async() - + seed_groups = list(context.seed_groups) strategies = {s.value for s in context.scenario_strategies} for strategy in strategies: for template_name in self._jailbreaks: for _ in range(self._num_attempts): atomic_attack = await self._get_atomic_attack_from_strategy_async( - strategy=strategy, jailbreak_template_name=template_name + strategy=strategy, jailbreak_template_name=template_name, seed_groups=seed_groups ) atomic_attacks.append(atomic_attack) - if context.include_baseline: - atomic_attacks.insert(0, self._build_baseline_atomic_attack(seed_groups=self._seed_groups or [])) - return atomic_attacks diff --git a/pyrit/scenario/scenarios/airt/psychosocial.py b/pyrit/scenario/scenarios/airt/psychosocial.py index 48f242b910..678e99bb93 100644 --- a/pyrit/scenario/scenarios/airt/psychosocial.py +++ b/pyrit/scenario/scenarios/airt/psychosocial.py @@ -72,14 +72,6 @@ class SubharmConfig: scoring_rubric_path: str -@dataclass -class ResolvedSeedData: - """Helper dataclass for resolved seed data.""" - - seed_groups: list[SeedAttackGroup] - subharm: str | None - - class PsychosocialStrategy(ScenarioStrategy): """ PsychosocialHarmsStrategy defines a set of strategies for testing model behavior @@ -258,17 +250,21 @@ def __init__( ) self._legacy_include_baseline = include_baseline - # Store deprecated objectives for later resolution in _resolve_seed_groups_async + # Store deprecated objectives for later resolution in _resolve_seed_groups_by_dataset_async self._deprecated_objectives = objectives - # Will be resolved in _build_atomic_attacks_async - self._seed_groups: list[SeedAttackGroup] | None = None - async def _resolve_seed_groups_async(self) -> ResolvedSeedData: + async def _resolve_seed_groups_by_dataset_async(self) -> dict[str, list[SeedAttackGroup]]: """ Resolve seed groups from deprecated objectives or dataset configuration. + Seeds are filtered to the harm category selected by the scenario strategies (e.g. + ``imminent_crisis``); the default ``ALL`` strategy keeps the broad ``psychosocial`` + category. The base ``Scenario`` flattens the result into ``context.seed_groups`` and + reuses it for the strategy attacks and the baseline. + Returns: - ResolvedSeedData: Contains seed groups and optional subharm category. + dict[str, list[SeedAttackGroup]]: Seed groups keyed by dataset (or a synthetic + key for deprecated inline objectives). Raises: ValueError: If both objectives and dataset_config are specified. @@ -282,15 +278,14 @@ async def _resolve_seed_groups_async(self) -> ResolvedSeedData: ) if self._deprecated_objectives is not None: - return ResolvedSeedData( - seed_groups=[SeedAttackGroup(seeds=[SeedObjective(value=obj)]) for obj in self._deprecated_objectives], - subharm=None, - ) + return { + "objectives": [SeedAttackGroup(seeds=[SeedObjective(value=obj)]) for obj in self._deprecated_objectives] + } harm_category_filter = self._extract_harm_category_filter() # Auto-fetch populates memory first; a still-empty result raises a # DatasetConstraintError naming the offending dataset, which we let propagate. - seed_groups = await self._dataset_config.get_seed_attack_groups_async() + seed_groups = list(await self._dataset_config.get_seed_attack_groups_async()) if harm_category_filter: seed_groups = self._filter_by_harm_category( @@ -306,10 +301,7 @@ async def _resolve_seed_groups_async(self) -> ResolvedSeedData: f"No seeds remained after filtering by harm_category '{harm_category_filter}'." ) - return ResolvedSeedData( - seed_groups=list(seed_groups), - subharm=harm_category_filter, - ) + return {harm_category_filter or "psychosocial": seed_groups} def _extract_harm_category_filter(self) -> str | None: """ @@ -414,25 +406,21 @@ async def _build_atomic_attacks_async(self, *, context: ScenarioContext) -> list f"conversations with editable history. Target {type(context.objective_target).__name__} " f"does not satisfy these requirements: {exc}" ) from exc - resolved = await self._resolve_seed_groups_async() - self._seed_groups = resolved.seed_groups - scoring_config = self._create_scoring_config(resolved.subharm) + # Deprecated inline objectives carry no harm category, so they map to no subharm rubric. + subharm = None if self._deprecated_objectives is not None else self._extract_harm_category_filter() + seed_groups = list(context.seed_groups) + scoring_config = self._create_scoring_config(subharm) - atomic_attacks: list[AtomicAttack] = [ - *self._create_single_turn_attacks(scoring_config=scoring_config, seed_groups=self._seed_groups), + return [ + *self._create_single_turn_attacks(scoring_config=scoring_config, seed_groups=seed_groups), self._create_multi_turn_attack( scoring_config=scoring_config, - subharm=resolved.subharm, - seed_groups=self._seed_groups, + subharm=subharm, + seed_groups=seed_groups, ), ] - if context.include_baseline: - atomic_attacks.insert(0, self._build_baseline_atomic_attack(seed_groups=self._seed_groups)) - - return atomic_attacks - def _create_scoring_config(self, subharm: str | None) -> AttackScoringConfig: subharm_config = self._subharm_configs.get(subharm) if subharm else None scorer = self._get_scorer(subharm=subharm) if subharm_config else self._objective_scorer diff --git a/pyrit/scenario/scenarios/airt/scam.py b/pyrit/scenario/scenarios/airt/scam.py index 7c31c7368c..6c8fa05b9d 100644 --- a/pyrit/scenario/scenarios/airt/scam.py +++ b/pyrit/scenario/scenarios/airt/scam.py @@ -168,27 +168,13 @@ def __init__( ) self._legacy_include_baseline = include_baseline - # Will be resolved in _build_atomic_attacks_async - self._seed_groups: list[SeedAttackGroup] | None = None - - async def _resolve_seed_groups_async(self) -> list[SeedAttackGroup]: - """ - Resolve seed groups from dataset configuration. - - Returns: - list[SeedAttackGroup]: List of seed attack groups with objectives to be tested. - """ - # Use dataset_config (guaranteed to be set by initialize_async). Auto-fetch - # populates memory first; a still-empty result raises a DatasetConstraintError - # naming the offending dataset, which we let propagate. - return list(await self._dataset_config.get_seed_attack_groups_async()) - - def _get_atomic_attack_from_strategy(self, strategy: str) -> AtomicAttack: + def _get_atomic_attack_from_strategy(self, *, strategy: str, seed_groups: list[SeedAttackGroup]) -> AtomicAttack: """ Translate the strategies into actual AtomicAttacks. Args: strategy (str): The strategy to create the attack from. + seed_groups (list[SeedAttackGroup]): Seed groups the attack draws from. Returns: AtomicAttack: Configured for the specified strategy. @@ -237,7 +223,7 @@ def _get_atomic_attack_from_strategy(self, strategy: str) -> AtomicAttack: return AtomicAttack( atomic_attack_name=f"scam_{strategy}", attack_technique=AttackTechnique(attack=attack_strategy), - seed_groups=self._seed_groups or [], + seed_groups=seed_groups, memory_labels=self._memory_labels, ) @@ -251,14 +237,9 @@ async def _build_atomic_attacks_async(self, *, context: ScenarioContext) -> list Returns: list[AtomicAttack]: List of atomic attacks to execute. """ - # Resolve seed groups from deprecated objectives or dataset config - self._seed_groups = await self._resolve_seed_groups_async() - + seed_groups = list(context.seed_groups) strategies = {s.value for s in context.scenario_strategies} - atomic_attacks = [self._get_atomic_attack_from_strategy(strategy) for strategy in strategies] - - if context.include_baseline: - atomic_attacks.insert(0, self._build_baseline_atomic_attack(seed_groups=self._seed_groups or [])) - - return atomic_attacks + return [ + self._get_atomic_attack_from_strategy(strategy=strategy, seed_groups=seed_groups) for strategy in strategies + ] diff --git a/pyrit/scenario/scenarios/benchmark/adversarial.py b/pyrit/scenario/scenarios/benchmark/adversarial.py index c82d856fe6..c13d1e5b59 100644 --- a/pyrit/scenario/scenarios/benchmark/adversarial.py +++ b/pyrit/scenario/scenarios/benchmark/adversarial.py @@ -247,7 +247,7 @@ async def _build_atomic_attacks_async(self, *, context: ScenarioContext) -> list # ``{technique}__{target}_{dataset}`` naming preserves the VERSION=2 cache key shape. atomic_attacks = builder.build( technique_factories=technique_factories, - dataset_groups=await context.dataset_config.get_attack_groups_by_dataset_async(), + dataset_groups=context.seed_groups_by_dataset, adversarial_targets=resolved_targets, display_group_fn=lambda combo: combo.target_name or "", include_baseline=False, diff --git a/pyrit/scenario/scenarios/foundry/red_team_agent.py b/pyrit/scenario/scenarios/foundry/red_team_agent.py index a61b65629a..ab9384202b 100644 --- a/pyrit/scenario/scenarios/foundry/red_team_agent.py +++ b/pyrit/scenario/scenarios/foundry/red_team_agent.py @@ -390,15 +390,6 @@ def _strategy_to_composite(strategy: ScenarioStrategy) -> "FoundryComposite": return FoundryComposite(attack=strategy) return FoundryComposite(attack=None, converters=[strategy]) - async def _resolve_seed_groups_async(self) -> list[SeedAttackGroup]: - """ - Resolve seed groups from the dataset configuration. - - Returns: - list[SeedGroup]: The resolved seed groups. - """ - return await self._dataset_config.get_seed_attack_groups_async() - async def _build_atomic_attacks_async(self, *, context: ScenarioContext) -> list[AtomicAttack]: """ Build one ``AtomicAttack`` per resolved FoundryComposite. @@ -409,23 +400,22 @@ async def _build_atomic_attacks_async(self, *, context: ScenarioContext) -> list Returns: list[AtomicAttack]: The list of AtomicAttack instances in this scenario. """ - # Resolve seed groups now that initialize_async has been called - self._seed_groups = await self._resolve_seed_groups_async() - - atomic_attacks = [self._get_attack_from_strategy(composition) for composition in self._scenario_composites] - - if context.include_baseline: - atomic_attacks.insert(0, self._build_baseline_atomic_attack(seed_groups=self._seed_groups)) - - return atomic_attacks - - def _get_attack_from_strategy(self, composite: FoundryComposite) -> AtomicAttack: + seed_groups = list(context.seed_groups) + return [ + self._get_attack_from_strategy(composite=composition, seed_groups=seed_groups) + for composition in self._scenario_composites + ] + + def _get_attack_from_strategy( + self, *, composite: FoundryComposite, seed_groups: list[SeedAttackGroup] + ) -> AtomicAttack: """ Get an atomic attack for the specified FoundryComposite. Args: composite (FoundryComposite): Typed composite with an optional attack strategy and zero or more converter strategies. + seed_groups (list[SeedAttackGroup]): Seed groups the attack draws from. Returns: AtomicAttack: The configured atomic attack. @@ -501,7 +491,7 @@ def _get_attack_from_strategy(self, composite: FoundryComposite) -> AtomicAttack return AtomicAttack( atomic_attack_name=composite.name, attack_technique=AttackTechnique(attack=attack), - seed_groups=self._seed_groups, + seed_groups=seed_groups, adversarial_chat=self._adversarial_chat, objective_scorer=self._attack_scoring_config.objective_scorer, memory_labels=self._memory_labels, diff --git a/pyrit/scenario/scenarios/garak/encoding.py b/pyrit/scenario/scenarios/garak/encoding.py index 179523fa86..2cf255753a 100644 --- a/pyrit/scenario/scenarios/garak/encoding.py +++ b/pyrit/scenario/scenarios/garak/encoding.py @@ -187,23 +187,6 @@ def __init__( ) self._legacy_include_baseline = include_baseline - # Will be resolved in _build_atomic_attacks_async - self._resolved_seed_groups: list[SeedAttackGroup] | None = None - - async def _resolve_seed_groups_async(self) -> list[SeedAttackGroup]: - """ - Resolve seed groups from dataset configuration. - - Returns: - list[SeedAttackGroup]: List of seed attack groups to be encoded and tested. - """ - # Use dataset_config (guaranteed to be set by initialize_async). The configured - # EncodingDatasetConfiguration shapes raw seeds into objective-bearing attack - # groups via its _build_attack_groups override; auto-fetch populates memory first - # when the configured datasets aren't present. A still-empty result raises a - # DatasetConstraintError naming the offending dataset, which we let propagate. - return await self._dataset_config.get_seed_attack_groups_async() - async def _build_atomic_attacks_async(self, *, context: ScenarioContext) -> list[AtomicAttack]: """ Build the encoding atomic attacks for this run. @@ -218,24 +201,19 @@ async def _build_atomic_attacks_async(self, *, context: ScenarioContext) -> list Returns: list[AtomicAttack]: The list of AtomicAttack instances in this scenario. """ - # Resolve seed prompts from deprecated parameter or dataset config - self._resolved_seed_groups = await self._resolve_seed_groups_async() - - atomic_attacks = self._get_converter_attacks() - - if context.include_baseline: - atomic_attacks.insert(0, self._build_baseline_atomic_attack(seed_groups=self._resolved_seed_groups or [])) - - return atomic_attacks + return self._get_converter_attacks(seed_groups=list(context.seed_groups)) # These are the same as Garak encoding attacks - def _get_converter_attacks(self) -> list[AtomicAttack]: + def _get_converter_attacks(self, *, seed_groups: list[SeedAttackGroup]) -> list[AtomicAttack]: """ Get all converter-based atomic attacks. Creates atomic attacks for each encoding scheme specified in the scenario strategies. Each encoding scheme is tested both with and without explicit decoding instructions. + Args: + seed_groups (list[SeedAttackGroup]): Seed groups the attacks draw from. + Returns: list[AtomicAttack]: List of all atomic attacks to execute. """ @@ -272,10 +250,14 @@ def _get_converter_attacks(self) -> list[AtomicAttack]: atomic_attacks = [] for conv, name in converters_with_encodings: - atomic_attacks.extend(self._get_prompt_attacks(converters=conv, encoding_name=name)) + atomic_attacks.extend( + self._get_prompt_attacks(converters=conv, encoding_name=name, seed_groups=seed_groups) + ) return atomic_attacks - def _get_prompt_attacks(self, *, converters: list[PromptConverter], encoding_name: str) -> list[AtomicAttack]: + def _get_prompt_attacks( + self, *, converters: list[PromptConverter], encoding_name: str, seed_groups: list[SeedAttackGroup] + ) -> list[AtomicAttack]: """ Create atomic attacks for a specific encoding scheme. @@ -288,6 +270,7 @@ def _get_prompt_attacks(self, *, converters: list[PromptConverter], encoding_nam Args: converters (list[PromptConverter]): The list of converters to apply to the seed prompts. encoding_name (str): Human-readable name of the encoding scheme (e.g., "Base64", "ROT13"). + seed_groups (list[SeedAttackGroup]): Seed groups the attacks draw from. Returns: list[AtomicAttack]: List of atomic attacks for this encoding scheme. @@ -326,7 +309,7 @@ def _get_prompt_attacks(self, *, converters: list[PromptConverter], encoding_nam AtomicAttack( atomic_attack_name=encoding_name, attack_technique=AttackTechnique(attack=attack), - seed_groups=self._resolved_seed_groups or [], + seed_groups=seed_groups, ) ) diff --git a/tests/unit/scenario/airt/test_jailbreak.py b/tests/unit/scenario/airt/test_jailbreak.py index 888967d411..6376641988 100644 --- a/tests/unit/scenario/airt/test_jailbreak.py +++ b/tests/unit/scenario/airt/test_jailbreak.py @@ -147,7 +147,10 @@ class TestJailbreakInitialization: def test_init_with_scenario_result_id(self, mock_scenario_result_id): """Test initialization with a scenario result ID.""" with patch.object( - Jailbreak, "_resolve_seed_groups_async", new_callable=AsyncMock, return_value=mock_memory_seed_groups + Jailbreak, + "_resolve_seed_groups_by_dataset_async", + new_callable=AsyncMock, + return_value={"memory": mock_memory_seed_groups}, ): scenario = Jailbreak(scenario_result_id=mock_scenario_result_id) assert scenario._scenario_result_id == mock_scenario_result_id @@ -155,7 +158,10 @@ def test_init_with_scenario_result_id(self, mock_scenario_result_id): def test_init_with_default_scorer(self, mock_memory_seed_groups): """Test initialization with default scorer.""" with patch.object( - Jailbreak, "_resolve_seed_groups_async", new_callable=AsyncMock, return_value=mock_memory_seed_groups + Jailbreak, + "_resolve_seed_groups_by_dataset_async", + new_callable=AsyncMock, + return_value={"memory": mock_memory_seed_groups}, ): scenario = Jailbreak() assert scenario._objective_scorer_identifier @@ -163,7 +169,10 @@ def test_init_with_default_scorer(self, mock_memory_seed_groups): def test_init_with_custom_scorer(self, mock_objective_scorer, mock_memory_seed_groups): """Test initialization with custom scorer.""" with patch.object( - Jailbreak, "_resolve_seed_groups_async", new_callable=AsyncMock, return_value=mock_memory_seed_groups + Jailbreak, + "_resolve_seed_groups_by_dataset_async", + new_callable=AsyncMock, + return_value={"memory": mock_memory_seed_groups}, ): scenario = Jailbreak(objective_scorer=mock_objective_scorer) assert scenario._objective_scorer == mock_objective_scorer @@ -171,7 +180,10 @@ def test_init_with_custom_scorer(self, mock_objective_scorer, mock_memory_seed_g def test_init_with_num_templates(self, mock_random_num_templates): """Test initialization with num_templates provided.""" with patch.object( - Jailbreak, "_resolve_seed_groups_async", new_callable=AsyncMock, return_value=mock_memory_seed_groups + Jailbreak, + "_resolve_seed_groups_by_dataset_async", + new_callable=AsyncMock, + return_value={"memory": mock_memory_seed_groups}, ): scenario = Jailbreak(num_templates=mock_random_num_templates) assert scenario._num_templates == mock_random_num_templates @@ -179,7 +191,10 @@ def test_init_with_num_templates(self, mock_random_num_templates): def test_init_with_num_attempts(self, mock_random_num_attempts): """Test initialization with n provided.""" with patch.object( - Jailbreak, "_resolve_seed_groups_async", new_callable=AsyncMock, return_value=mock_memory_seed_groups + Jailbreak, + "_resolve_seed_groups_by_dataset_async", + new_callable=AsyncMock, + return_value={"memory": mock_memory_seed_groups}, ): scenario = Jailbreak(num_attempts=mock_random_num_attempts) assert scenario._num_attempts == mock_random_num_attempts @@ -200,7 +215,10 @@ def test_init_accepts_subdirectory_jailbreak_names(self, mock_objective_scorer, subdir_name = subdir_templates[0] with patch.object( - Jailbreak, "_resolve_seed_groups_async", new_callable=AsyncMock, return_value=mock_memory_seed_groups + Jailbreak, + "_resolve_seed_groups_by_dataset_async", + new_callable=AsyncMock, + return_value={"memory": mock_memory_seed_groups}, ): scenario = Jailbreak(objective_scorer=mock_objective_scorer, jailbreak_names=[subdir_name]) assert scenario._jailbreaks == [subdir_name] @@ -230,7 +248,10 @@ async def test_default_initialize_includes_baseline( ): """initialize_async without include_baseline honors BASELINE_ATTACK_POLICY=Enabled.""" with patch.object( - Jailbreak, "_resolve_seed_groups_async", new_callable=AsyncMock, return_value=mock_memory_seed_groups + Jailbreak, + "_resolve_seed_groups_by_dataset_async", + new_callable=AsyncMock, + return_value={"memory": mock_memory_seed_groups}, ): scenario = Jailbreak(objective_scorer=mock_objective_scorer) await scenario.initialize_async(objective_target=mock_objective_target) @@ -241,7 +262,10 @@ async def test_explicit_include_baseline_false_omits_baseline( ): """Caller can opt out of baseline by passing include_baseline=False.""" with patch.object( - Jailbreak, "_resolve_seed_groups_async", new_callable=AsyncMock, return_value=mock_memory_seed_groups + Jailbreak, + "_resolve_seed_groups_by_dataset_async", + new_callable=AsyncMock, + return_value={"memory": mock_memory_seed_groups}, ): scenario = Jailbreak(objective_scorer=mock_objective_scorer) await scenario.initialize_async( @@ -260,7 +284,10 @@ async def test_attack_generation_for_simple( ): """Test that the simple attack generation works.""" with patch.object( - Jailbreak, "_resolve_seed_groups_async", new_callable=AsyncMock, return_value=mock_memory_seed_groups + Jailbreak, + "_resolve_seed_groups_by_dataset_async", + new_callable=AsyncMock, + return_value={"memory": mock_memory_seed_groups}, ): scenario = Jailbreak(objective_scorer=mock_objective_scorer, num_templates=2) @@ -276,7 +303,10 @@ async def test_attack_generation_for_complex( ): """Test that the complex attack generation works.""" with patch.object( - Jailbreak, "_resolve_seed_groups_async", new_callable=AsyncMock, return_value=mock_memory_seed_groups + Jailbreak, + "_resolve_seed_groups_by_dataset_async", + new_callable=AsyncMock, + return_value={"memory": mock_memory_seed_groups}, ): scenario = Jailbreak(objective_scorer=mock_objective_scorer, num_templates=2) @@ -296,7 +326,10 @@ async def test_attack_generation_for_manyshot( ): """Test that the manyshot attack generation works.""" with patch.object( - Jailbreak, "_resolve_seed_groups_async", new_callable=AsyncMock, return_value=mock_memory_seed_groups + Jailbreak, + "_resolve_seed_groups_by_dataset_async", + new_callable=AsyncMock, + return_value={"memory": mock_memory_seed_groups}, ): scenario = Jailbreak(objective_scorer=mock_objective_scorer, num_templates=2) @@ -314,7 +347,10 @@ async def test_attack_generation_for_promptsending( ): """Test that the prompt sending attack generation works.""" with patch.object( - Jailbreak, "_resolve_seed_groups_async", new_callable=AsyncMock, return_value=mock_memory_seed_groups + Jailbreak, + "_resolve_seed_groups_by_dataset_async", + new_callable=AsyncMock, + return_value={"memory": mock_memory_seed_groups}, ): scenario = Jailbreak(objective_scorer=mock_objective_scorer, num_templates=2) @@ -332,7 +368,10 @@ async def test_attack_generation_for_skeleton( ): """Test that the skelton key attack generation works.""" with patch.object( - Jailbreak, "_resolve_seed_groups_async", new_callable=AsyncMock, return_value=mock_memory_seed_groups + Jailbreak, + "_resolve_seed_groups_by_dataset_async", + new_callable=AsyncMock, + return_value={"memory": mock_memory_seed_groups}, ): scenario = Jailbreak(objective_scorer=mock_objective_scorer, num_templates=2) @@ -350,7 +389,10 @@ async def test_attack_generation_for_roleplay( ): """Test that the roleplaying attack generation works.""" with patch.object( - Jailbreak, "_resolve_seed_groups_async", new_callable=AsyncMock, return_value=mock_memory_seed_groups + Jailbreak, + "_resolve_seed_groups_by_dataset_async", + new_callable=AsyncMock, + return_value={"memory": mock_memory_seed_groups}, ): scenario = Jailbreak(objective_scorer=mock_objective_scorer, num_templates=2) @@ -372,7 +414,10 @@ async def test_attack_runs_include_objectives( Combined coverage previously split across test_get_atomic_attacks_async_returns_attacks. """ with patch.object( - Jailbreak, "_resolve_seed_groups_async", new_callable=AsyncMock, return_value=mock_memory_seed_groups + Jailbreak, + "_resolve_seed_groups_by_dataset_async", + new_callable=AsyncMock, + return_value={"memory": mock_memory_seed_groups}, ): scenario = Jailbreak(objective_scorer=mock_objective_scorer, num_templates=2) @@ -389,7 +434,10 @@ async def test_get_all_jailbreak_templates( ): """Test that all jailbreak templates are found.""" with patch.object( - Jailbreak, "_resolve_seed_groups_async", new_callable=AsyncMock, return_value=mock_memory_seed_groups + Jailbreak, + "_resolve_seed_groups_by_dataset_async", + new_callable=AsyncMock, + return_value={"memory": mock_memory_seed_groups}, ): scenario = Jailbreak( objective_scorer=mock_objective_scorer, @@ -402,7 +450,10 @@ async def test_get_some_jailbreak_templates( ): """Test that random jailbreak template selection works.""" with patch.object( - Jailbreak, "_resolve_seed_groups_async", new_callable=AsyncMock, return_value=mock_memory_seed_groups + Jailbreak, + "_resolve_seed_groups_by_dataset_async", + new_callable=AsyncMock, + return_value={"memory": mock_memory_seed_groups}, ): scenario = Jailbreak(objective_scorer=mock_objective_scorer, num_templates=mock_random_num_templates) await scenario.initialize_async(objective_target=mock_objective_target) @@ -413,7 +464,10 @@ async def test_custom_num_attempts( ): """Test that n successfully tries each jailbreak template n-many times.""" with patch.object( - Jailbreak, "_resolve_seed_groups_async", new_callable=AsyncMock, return_value=mock_memory_seed_groups + Jailbreak, + "_resolve_seed_groups_by_dataset_async", + new_callable=AsyncMock, + return_value={"memory": mock_memory_seed_groups}, ): base_scenario = Jailbreak(objective_scorer=mock_objective_scorer, num_templates=2) await base_scenario.initialize_async(objective_target=mock_objective_target, include_baseline=False) @@ -443,7 +497,10 @@ async def test_initialize_async_with_max_concurrency( ) -> None: """Test initialization with custom max_concurrency.""" with patch.object( - Jailbreak, "_resolve_seed_groups_async", new_callable=AsyncMock, return_value=mock_memory_seed_groups + Jailbreak, + "_resolve_seed_groups_by_dataset_async", + new_callable=AsyncMock, + return_value={"memory": mock_memory_seed_groups}, ): scenario = Jailbreak(objective_scorer=mock_objective_scorer) await scenario.initialize_async(objective_target=mock_objective_target, max_concurrency=20) @@ -460,7 +517,10 @@ async def test_initialize_async_with_memory_labels( memory_labels = {"type": "jailbreak", "category": "scenario"} with patch.object( - Jailbreak, "_resolve_seed_groups_async", new_callable=AsyncMock, return_value=mock_memory_seed_groups + Jailbreak, + "_resolve_seed_groups_by_dataset_async", + new_callable=AsyncMock, + return_value={"memory": mock_memory_seed_groups}, ): scenario = Jailbreak(objective_scorer=mock_objective_scorer) await scenario.initialize_async( @@ -496,7 +556,10 @@ async def test_no_target_duplication_async( ) -> None: """Test that all three targets (adversarial, object, scorer) are distinct.""" with patch.object( - Jailbreak, "_resolve_seed_groups_async", new_callable=AsyncMock, return_value=mock_memory_seed_groups + Jailbreak, + "_resolve_seed_groups_by_dataset_async", + new_callable=AsyncMock, + return_value={"memory": mock_memory_seed_groups}, ): scenario = Jailbreak() await scenario.initialize_async(objective_target=mock_objective_target) @@ -543,7 +606,10 @@ async def test_roleplay_attacks_share_adversarial_target( ) -> None: """Test that multiple role-play attacks share the same adversarial target instance.""" with patch.object( - Jailbreak, "_resolve_seed_groups_async", new_callable=AsyncMock, return_value=mock_memory_seed_groups + Jailbreak, + "_resolve_seed_groups_by_dataset_async", + new_callable=AsyncMock, + return_value={"memory": mock_memory_seed_groups}, ): scenario = Jailbreak(objective_scorer=mock_objective_scorer, num_templates=2) await scenario.initialize_async( @@ -572,8 +638,8 @@ async def test_one_resolution_call_baseline_matches_strategies( seed_groups = [SeedAttackGroup(seeds=[SeedObjective(value=f"obj{i}")]) for i in range(10)] config = DatasetAttackConfiguration(seed_groups=seed_groups, max_dataset_size=3) - first_sample = seed_groups[:3] - second_sample = seed_groups[5:8] + first_sample = [("inline", group) for group in seed_groups[:3]] + second_sample = [("inline", group) for group in seed_groups[5:8]] scenario = Jailbreak(objective_scorer=mock_objective_scorer, num_templates=1) with patch( "pyrit.scenario.core.dataset_configuration.random.sample", diff --git a/tests/unit/scenario/airt/test_psychosocial.py b/tests/unit/scenario/airt/test_psychosocial.py index 45cb6ec517..e62fbeba4b 100644 --- a/tests/unit/scenario/airt/test_psychosocial.py +++ b/tests/unit/scenario/airt/test_psychosocial.py @@ -14,7 +14,7 @@ Psychosocial, PsychosocialStrategy, ) -from pyrit.scenario.scenarios.airt.psychosocial import ResolvedSeedData, SubharmConfig +from pyrit.scenario.scenarios.airt.psychosocial import SubharmConfig from pyrit.score import FloatScaleThresholdScorer SEED_DATASETS_PATH = DATASETS_PATH / "seed_datasets" / "local" / "airt" @@ -28,9 +28,9 @@ def mock_memory_seed_groups() -> list[SeedGroup]: @pytest.fixture -def mock_resolved_seed_data(mock_memory_seed_groups) -> ResolvedSeedData: - """Create mock ResolvedSeedData for patching _resolve_seed_groups.""" - return ResolvedSeedData(seed_groups=mock_memory_seed_groups, subharm=None) +def mock_seed_groups_by_dataset(mock_memory_seed_groups) -> dict[str, list[SeedAttackGroup]]: + """Create mock by-dataset seed groups for patching _resolve_seed_groups_by_dataset_async.""" + return {"psychosocial": mock_memory_seed_groups} @pytest.fixture @@ -183,12 +183,15 @@ async def test_attack_generation_for_all( self, mock_objective_target, mock_objective_scorer, - mock_resolved_seed_data, + mock_seed_groups_by_dataset, mock_dataset_config, ): """Test that _get_atomic_attacks_async returns atomic attacks.""" with patch.object( - Psychosocial, "_resolve_seed_groups_async", new_callable=AsyncMock, return_value=mock_resolved_seed_data + Psychosocial, + "_resolve_seed_groups_by_dataset_async", + new_callable=AsyncMock, + return_value=mock_seed_groups_by_dataset, ): scenario = Psychosocial(objective_scorer=mock_objective_scorer) @@ -203,12 +206,15 @@ async def test_attack_runs_include_objectives_async( *, mock_objective_target: PromptTarget, mock_objective_scorer: FloatScaleThresholdScorer, - mock_resolved_seed_data, + mock_seed_groups_by_dataset, mock_dataset_config, ) -> None: """Test that attack runs include objectives for each seed prompt.""" with patch.object( - Psychosocial, "_resolve_seed_groups_async", new_callable=AsyncMock, return_value=mock_resolved_seed_data + Psychosocial, + "_resolve_seed_groups_by_dataset_async", + new_callable=AsyncMock, + return_value=mock_seed_groups_by_dataset, ): scenario = Psychosocial( objective_scorer=mock_objective_scorer, @@ -225,12 +231,15 @@ async def test_get_atomic_attacks_async_returns_attacks( *, mock_objective_target: PromptTarget, mock_objective_scorer: FloatScaleThresholdScorer, - mock_resolved_seed_data, + mock_seed_groups_by_dataset, mock_dataset_config, ) -> None: """Test that _get_atomic_attacks_async returns atomic attacks.""" with patch.object( - Psychosocial, "_resolve_seed_groups_async", new_callable=AsyncMock, return_value=mock_resolved_seed_data + Psychosocial, + "_resolve_seed_groups_by_dataset_async", + new_callable=AsyncMock, + return_value=mock_seed_groups_by_dataset, ): scenario = Psychosocial( objective_scorer=mock_objective_scorer, @@ -251,12 +260,15 @@ async def test_initialize_async_with_max_concurrency( *, mock_objective_target: PromptTarget, mock_objective_scorer: FloatScaleThresholdScorer, - mock_resolved_seed_data, + mock_seed_groups_by_dataset, mock_dataset_config, ) -> None: """Test initialization with custom max_concurrency.""" with patch.object( - Psychosocial, "_resolve_seed_groups_async", new_callable=AsyncMock, return_value=mock_resolved_seed_data + Psychosocial, + "_resolve_seed_groups_by_dataset_async", + new_callable=AsyncMock, + return_value=mock_seed_groups_by_dataset, ): scenario = Psychosocial(objective_scorer=mock_objective_scorer) await scenario.initialize_async( @@ -269,14 +281,17 @@ async def test_initialize_async_with_memory_labels( *, mock_objective_target: PromptTarget, mock_objective_scorer: FloatScaleThresholdScorer, - mock_resolved_seed_data, + mock_seed_groups_by_dataset, mock_dataset_config, ) -> None: """Test initialization with memory labels.""" memory_labels = {"type": "psychosocial", "category": "crisis"} with patch.object( - Psychosocial, "_resolve_seed_groups_async", new_callable=AsyncMock, return_value=mock_resolved_seed_data + Psychosocial, + "_resolve_seed_groups_by_dataset_async", + new_callable=AsyncMock, + return_value=mock_seed_groups_by_dataset, ): scenario = Psychosocial(objective_scorer=mock_objective_scorer) await scenario.initialize_async( @@ -317,12 +332,15 @@ async def test_no_target_duplication_async( self, *, mock_objective_target: PromptTarget, - mock_resolved_seed_data, + mock_seed_groups_by_dataset, mock_dataset_config, ) -> None: """Test that all three targets (adversarial, objective, scorer) are distinct.""" with patch.object( - Psychosocial, "_resolve_seed_groups_async", new_callable=AsyncMock, return_value=mock_resolved_seed_data + Psychosocial, + "_resolve_seed_groups_by_dataset_async", + new_callable=AsyncMock, + return_value=mock_seed_groups_by_dataset, ): scenario = Psychosocial() await scenario.initialize_async(objective_target=mock_objective_target, dataset_config=mock_dataset_config) @@ -350,12 +368,15 @@ async def test_initialize_async_invokes_target_requirements_validate( self, mock_objective_target, mock_objective_scorer, - mock_resolved_seed_data, + mock_seed_groups_by_dataset, mock_dataset_config, ): """initialize_async must delegate capability validation to TARGET_REQUIREMENTS.validate.""" with patch.object( - Psychosocial, "_resolve_seed_groups_async", new_callable=AsyncMock, return_value=mock_resolved_seed_data + Psychosocial, + "_resolve_seed_groups_by_dataset_async", + new_callable=AsyncMock, + return_value=mock_seed_groups_by_dataset, ): scenario = Psychosocial(objective_scorer=mock_objective_scorer) with patch("pyrit.prompt_target.common.target_requirements.TargetRequirements.validate") as mock_validate: @@ -373,7 +394,7 @@ async def test_initialize_async_invokes_target_requirements_validate( async def test_initialize_async_rejects_target_missing_editable_history( self, mock_objective_scorer, - mock_resolved_seed_data, + mock_seed_groups_by_dataset, mock_dataset_config, ): """A target that does not natively support EDITABLE_HISTORY must be rejected.""" @@ -390,7 +411,10 @@ async def test_initialize_async_rejects_target_missing_editable_history( ) with patch.object( - Psychosocial, "_resolve_seed_groups_async", new_callable=AsyncMock, return_value=mock_resolved_seed_data + Psychosocial, + "_resolve_seed_groups_by_dataset_async", + new_callable=AsyncMock, + return_value=mock_seed_groups_by_dataset, ): scenario = Psychosocial(objective_scorer=mock_objective_scorer) with pytest.raises(ValueError, match="editable_history"): diff --git a/tests/unit/scenario/airt/test_scam.py b/tests/unit/scenario/airt/test_scam.py index 0218ef8973..05b4760e8c 100644 --- a/tests/unit/scenario/airt/test_scam.py +++ b/tests/unit/scenario/airt/test_scam.py @@ -59,6 +59,7 @@ def mock_dataset_config(mock_memory_seed_groups): seed_attack_groups = list(mock_memory_seed_groups) mock_config = MagicMock(spec=DatasetAttackConfiguration) mock_config.get_seed_attack_groups_async = AsyncMock(return_value=seed_attack_groups) + mock_config.get_attack_groups_by_dataset_async = AsyncMock(return_value={"airt_scam": seed_attack_groups}) mock_config.dataset_names = ["airt_scam"] return mock_config @@ -129,7 +130,10 @@ def test_init_with_default_objectives( mock_memory_seed_groups: list[SeedAttackGroup], ) -> None: with patch.object( - Scam, "_resolve_seed_groups_async", new_callable=AsyncMock, return_value=mock_memory_seed_groups + Scam, + "_resolve_seed_groups_by_dataset_async", + new_callable=AsyncMock, + return_value={"memory": mock_memory_seed_groups}, ): scenario = Scam(objective_scorer=mock_objective_scorer) @@ -139,7 +143,10 @@ def test_init_with_default_objectives( def test_init_with_default_scorer(self, mock_memory_seed_groups) -> None: """Test initialization with default scorer.""" with patch.object( - Scam, "_resolve_seed_groups_async", new_callable=AsyncMock, return_value=mock_memory_seed_groups + Scam, + "_resolve_seed_groups_by_dataset_async", + new_callable=AsyncMock, + return_value={"memory": mock_memory_seed_groups}, ): scenario = Scam() assert scenario._objective_scorer_identifier @@ -149,7 +156,10 @@ def test_init_with_custom_scorer(self, *, mock_memory_seed_groups: list[SeedAtta scorer = MagicMock(spec=TrueFalseCompositeScorer) with patch.object( - Scam, "_resolve_seed_groups_async", new_callable=AsyncMock, return_value=mock_memory_seed_groups + Scam, + "_resolve_seed_groups_by_dataset_async", + new_callable=AsyncMock, + return_value={"memory": mock_memory_seed_groups}, ): scenario = Scam(objective_scorer=scorer) assert isinstance(scenario._scorer_config, AttackScoringConfig) @@ -158,7 +168,10 @@ def test_init_default_adversarial_chat( self, *, mock_objective_scorer: TrueFalseCompositeScorer, mock_memory_seed_groups: list[SeedAttackGroup] ) -> None: with patch.object( - Scam, "_resolve_seed_groups_async", new_callable=AsyncMock, return_value=mock_memory_seed_groups + Scam, + "_resolve_seed_groups_by_dataset_async", + new_callable=AsyncMock, + return_value={"memory": mock_memory_seed_groups}, ): scenario = Scam(objective_scorer=mock_objective_scorer) @@ -172,7 +185,10 @@ def test_init_with_adversarial_chat( adversarial_chat.get_identifier.return_value = _mock_target_id("CustomAdversary") with patch.object( - Scam, "_resolve_seed_groups_async", new_callable=AsyncMock, return_value=mock_memory_seed_groups + Scam, + "_resolve_seed_groups_by_dataset_async", + new_callable=AsyncMock, + return_value={"memory": mock_memory_seed_groups}, ): scenario = Scam( adversarial_chat=adversarial_chat, @@ -209,7 +225,10 @@ async def test_attack_generation_for_all( ): """Test that _get_atomic_attacks_async returns atomic attacks.""" with patch.object( - Scam, "_resolve_seed_groups_async", new_callable=AsyncMock, return_value=mock_memory_seed_groups + Scam, + "_resolve_seed_groups_by_dataset_async", + new_callable=AsyncMock, + return_value={"memory": mock_memory_seed_groups}, ): scenario = Scam(objective_scorer=mock_objective_scorer) @@ -363,7 +382,10 @@ async def test_initialize_async_with_max_concurrency( ) -> None: """Test initialization with custom max_concurrency.""" with patch.object( - Scam, "_resolve_seed_groups_async", new_callable=AsyncMock, return_value=mock_memory_seed_groups + Scam, + "_resolve_seed_groups_by_dataset_async", + new_callable=AsyncMock, + return_value={"memory": mock_memory_seed_groups}, ): scenario = Scam(objective_scorer=mock_objective_scorer) await scenario.initialize_async( @@ -383,7 +405,10 @@ async def test_initialize_async_with_memory_labels( memory_labels = {"type": "scam", "category": "scenario"} with patch.object( - Scam, "_resolve_seed_groups_async", new_callable=AsyncMock, return_value=mock_memory_seed_groups + Scam, + "_resolve_seed_groups_by_dataset_async", + new_callable=AsyncMock, + return_value={"memory": mock_memory_seed_groups}, ): scenario = Scam(objective_scorer=mock_objective_scorer) await scenario.initialize_async( @@ -419,7 +444,10 @@ async def test_no_target_duplication_async( ) -> None: """Test that all three targets (adversarial, object, scorer) are distinct.""" with patch.object( - Scam, "_resolve_seed_groups_async", new_callable=AsyncMock, return_value=mock_memory_seed_groups + Scam, + "_resolve_seed_groups_by_dataset_async", + new_callable=AsyncMock, + return_value={"memory": mock_memory_seed_groups}, ): scenario = Scam() await scenario.initialize_async(objective_target=mock_objective_target, dataset_config=mock_dataset_config) @@ -445,8 +473,8 @@ async def test_one_resolution_call_baseline_matches_strategies( seed_groups = [SeedAttackGroup(seeds=[SeedObjective(value=f"obj{i}")]) for i in range(10)] config = DatasetAttackConfiguration(seed_groups=seed_groups, max_dataset_size=3) - first_sample = seed_groups[:3] - second_sample = seed_groups[5:8] + first_sample = [("inline", group) for group in seed_groups[:3]] + second_sample = [("inline", group) for group in seed_groups[5:8]] with patch( "pyrit.scenario.core.dataset_configuration.random.sample", side_effect=[first_sample, second_sample], diff --git a/tests/unit/scenario/foundry/test_red_team_agent.py b/tests/unit/scenario/foundry/test_red_team_agent.py index 1799f6946b..2a97be6db3 100644 --- a/tests/unit/scenario/foundry/test_red_team_agent.py +++ b/tests/unit/scenario/foundry/test_red_team_agent.py @@ -116,7 +116,10 @@ async def test_init_with_single_strategy( ): """Test initialization with a single attack strategy.""" with patch.object( - RedTeamAgent, "_resolve_seed_groups_async", new_callable=AsyncMock, return_value=mock_memory_seed_groups + RedTeamAgent, + "_resolve_seed_groups_by_dataset_async", + new_callable=AsyncMock, + return_value={"memory": mock_memory_seed_groups}, ): scenario = RedTeamAgent( attack_scoring_config=AttackScoringConfig(objective_scorer=mock_objective_scorer), @@ -141,7 +144,10 @@ async def test_init_with_multiple_strategies( ] with patch.object( - RedTeamAgent, "_resolve_seed_groups_async", new_callable=AsyncMock, return_value=mock_memory_seed_groups + RedTeamAgent, + "_resolve_seed_groups_by_dataset_async", + new_callable=AsyncMock, + return_value={"memory": mock_memory_seed_groups}, ): scenario = RedTeamAgent( attack_scoring_config=AttackScoringConfig(objective_scorer=mock_objective_scorer), @@ -180,7 +186,10 @@ async def test_init_with_memory_labels( memory_labels = {"test": "foundry", "category": "attack"} with patch.object( - RedTeamAgent, "_resolve_seed_groups_async", new_callable=AsyncMock, return_value=mock_memory_seed_groups + RedTeamAgent, + "_resolve_seed_groups_by_dataset_async", + new_callable=AsyncMock, + return_value={"memory": mock_memory_seed_groups}, ): scenario = RedTeamAgent( attack_scoring_config=AttackScoringConfig(objective_scorer=mock_objective_scorer), @@ -205,7 +214,10 @@ def test_init_creates_default_scorer_when_not_provided( mock_get_scorer.return_value = mock_scorer_instance with patch.object( - RedTeamAgent, "_resolve_seed_groups_async", new_callable=AsyncMock, return_value=mock_memory_seed_groups + RedTeamAgent, + "_resolve_seed_groups_by_dataset_async", + new_callable=AsyncMock, + return_value={"memory": mock_memory_seed_groups}, ): scenario = RedTeamAgent() @@ -240,7 +252,10 @@ async def test_normalize_easy_strategies( ): """Test that EASY strategy expands to easy attack strategies.""" with patch.object( - RedTeamAgent, "_resolve_seed_groups_async", new_callable=AsyncMock, return_value=mock_memory_seed_groups + RedTeamAgent, + "_resolve_seed_groups_by_dataset_async", + new_callable=AsyncMock, + return_value={"memory": mock_memory_seed_groups}, ): scenario = RedTeamAgent( attack_scoring_config=AttackScoringConfig(objective_scorer=mock_objective_scorer), @@ -259,7 +274,10 @@ async def test_normalize_moderate_strategies( ): """Test that MODERATE strategy expands to moderate attack strategies.""" with patch.object( - RedTeamAgent, "_resolve_seed_groups_async", new_callable=AsyncMock, return_value=mock_memory_seed_groups + RedTeamAgent, + "_resolve_seed_groups_by_dataset_async", + new_callable=AsyncMock, + return_value={"memory": mock_memory_seed_groups}, ): scenario = RedTeamAgent( attack_scoring_config=AttackScoringConfig(objective_scorer=mock_objective_scorer), @@ -278,7 +296,10 @@ async def test_normalize_difficult_strategies( ): """Test that DIFFICULT strategy expands to difficult attack strategies.""" with patch.object( - RedTeamAgent, "_resolve_seed_groups_async", new_callable=AsyncMock, return_value=mock_memory_seed_groups + RedTeamAgent, + "_resolve_seed_groups_by_dataset_async", + new_callable=AsyncMock, + return_value={"memory": mock_memory_seed_groups}, ): # DIFFICULT strategy includes TAP which requires FloatScaleThresholdScorer scenario = RedTeamAgent( @@ -298,7 +319,10 @@ async def test_normalize_mixed_difficulty_levels( ): """Test that multiple difficulty levels expand correctly.""" with patch.object( - RedTeamAgent, "_resolve_seed_groups_async", new_callable=AsyncMock, return_value=mock_memory_seed_groups + RedTeamAgent, + "_resolve_seed_groups_by_dataset_async", + new_callable=AsyncMock, + return_value={"memory": mock_memory_seed_groups}, ): scenario = RedTeamAgent( attack_scoring_config=AttackScoringConfig(objective_scorer=mock_objective_scorer), @@ -317,7 +341,10 @@ async def test_normalize_with_specific_and_difficulty_levels( ): """Test that specific strategies combined with difficulty levels work correctly.""" with patch.object( - RedTeamAgent, "_resolve_seed_groups_async", new_callable=AsyncMock, return_value=mock_memory_seed_groups + RedTeamAgent, + "_resolve_seed_groups_by_dataset_async", + new_callable=AsyncMock, + return_value={"memory": mock_memory_seed_groups}, ): scenario = RedTeamAgent( attack_scoring_config=AttackScoringConfig(objective_scorer=mock_objective_scorer), @@ -344,7 +371,10 @@ async def test_get_attack_from_single_turn_strategy( ): """Test creating an attack from a single-turn strategy.""" with patch.object( - RedTeamAgent, "_resolve_seed_groups_async", new_callable=AsyncMock, return_value=mock_memory_seed_groups + RedTeamAgent, + "_resolve_seed_groups_by_dataset_async", + new_callable=AsyncMock, + return_value={"memory": mock_memory_seed_groups}, ): scenario = RedTeamAgent( attack_scoring_config=AttackScoringConfig(objective_scorer=mock_objective_scorer), @@ -358,7 +388,9 @@ async def test_get_attack_from_single_turn_strategy( # Get the composite strategy that was created during initialization composite_strategy = scenario._scenario_composites[0] - atomic_attack = scenario._get_attack_from_strategy(composite_strategy) + atomic_attack = scenario._get_attack_from_strategy( + composite=composite_strategy, seed_groups=mock_memory_seed_groups + ) assert isinstance(atomic_attack, AtomicAttack) assert atomic_attack.seed_groups == mock_memory_seed_groups @@ -373,7 +405,10 @@ async def test_get_attack_from_multi_turn_strategy( ): """Test creating a multi-turn attack strategy.""" with patch.object( - RedTeamAgent, "_resolve_seed_groups_async", new_callable=AsyncMock, return_value=mock_memory_seed_groups + RedTeamAgent, + "_resolve_seed_groups_by_dataset_async", + new_callable=AsyncMock, + return_value={"memory": mock_memory_seed_groups}, ): scenario = RedTeamAgent( adversarial_chat=mock_adversarial_target, @@ -388,7 +423,9 @@ async def test_get_attack_from_multi_turn_strategy( # Get the composite strategy that was created during initialization composite_strategy = scenario._scenario_composites[0] - atomic_attack = scenario._get_attack_from_strategy(composite_strategy) + atomic_attack = scenario._get_attack_from_strategy( + composite=composite_strategy, seed_groups=mock_memory_seed_groups + ) assert isinstance(atomic_attack, AtomicAttack) assert atomic_attack.seed_groups == mock_memory_seed_groups @@ -403,7 +440,10 @@ async def test_get_attack_single_turn_with_converters( ): """Test creating a single-turn attack with converters.""" with patch.object( - RedTeamAgent, "_resolve_seed_groups_async", new_callable=AsyncMock, return_value=mock_memory_seed_groups + RedTeamAgent, + "_resolve_seed_groups_by_dataset_async", + new_callable=AsyncMock, + return_value={"memory": mock_memory_seed_groups}, ): scenario = RedTeamAgent( attack_scoring_config=AttackScoringConfig(objective_scorer=mock_objective_scorer), @@ -432,7 +472,10 @@ async def test_get_attack_multi_turn_with_adversarial_target( ): """Test creating a multi-turn attack.""" with patch.object( - RedTeamAgent, "_resolve_seed_groups_async", new_callable=AsyncMock, return_value=mock_memory_seed_groups + RedTeamAgent, + "_resolve_seed_groups_by_dataset_async", + new_callable=AsyncMock, + return_value={"memory": mock_memory_seed_groups}, ): scenario = RedTeamAgent( adversarial_chat=mock_adversarial_target, @@ -488,7 +531,10 @@ async def test_all_single_turn_strategies_create_attack_runs( ): """Test that all single-turn strategies can create attack runs.""" with patch.object( - RedTeamAgent, "_resolve_seed_groups_async", new_callable=AsyncMock, return_value=mock_memory_seed_groups + RedTeamAgent, + "_resolve_seed_groups_by_dataset_async", + new_callable=AsyncMock, + return_value={"memory": mock_memory_seed_groups}, ): scenario = RedTeamAgent( attack_scoring_config=AttackScoringConfig(objective_scorer=mock_objective_scorer), @@ -502,7 +548,9 @@ async def test_all_single_turn_strategies_create_attack_runs( # Get the composite strategy that was created during initialization composite_strategy = scenario._scenario_composites[0] - atomic_attack = scenario._get_attack_from_strategy(composite_strategy) + atomic_attack = scenario._get_attack_from_strategy( + composite=composite_strategy, seed_groups=mock_memory_seed_groups + ) assert isinstance(atomic_attack, AtomicAttack) @pytest.mark.parametrize( @@ -523,7 +571,10 @@ async def test_all_multi_turn_strategies_create_attack_runs( ): """Test that all multi-turn strategies can create attack runs.""" with patch.object( - RedTeamAgent, "_resolve_seed_groups_async", new_callable=AsyncMock, return_value=mock_memory_seed_groups + RedTeamAgent, + "_resolve_seed_groups_by_dataset_async", + new_callable=AsyncMock, + return_value={"memory": mock_memory_seed_groups}, ): scenario = RedTeamAgent( adversarial_chat=mock_adversarial_target, @@ -538,7 +589,9 @@ async def test_all_multi_turn_strategies_create_attack_runs( # Get the composite strategy that was created during initialization composite_strategy = scenario._scenario_composites[0] - atomic_attack = scenario._get_attack_from_strategy(composite_strategy) + atomic_attack = scenario._get_attack_from_strategy( + composite=composite_strategy, seed_groups=mock_memory_seed_groups + ) assert isinstance(atomic_attack, AtomicAttack) @@ -553,7 +606,10 @@ async def test_scenario_composites_set_after_initialize( strategies = [FoundryStrategy.Base64, FoundryStrategy.ROT13] with patch.object( - RedTeamAgent, "_resolve_seed_groups_async", new_callable=AsyncMock, return_value=mock_memory_seed_groups + RedTeamAgent, + "_resolve_seed_groups_by_dataset_async", + new_callable=AsyncMock, + return_value={"memory": mock_memory_seed_groups}, ): scenario = RedTeamAgent( attack_scoring_config=AttackScoringConfig(objective_scorer=mock_objective_scorer), @@ -592,7 +648,10 @@ async def test_scenario_atomic_attack_count_matches_strategies( ] with patch.object( - RedTeamAgent, "_resolve_seed_groups_async", new_callable=AsyncMock, return_value=mock_memory_seed_groups + RedTeamAgent, + "_resolve_seed_groups_by_dataset_async", + new_callable=AsyncMock, + return_value={"memory": mock_memory_seed_groups}, ): scenario = RedTeamAgent( attack_scoring_config=AttackScoringConfig(objective_scorer=mock_objective_scorer), @@ -613,7 +672,10 @@ async def test_initialize_with_foundry_composite_directly( composite = FoundryComposite(attack=FoundryStrategy.Crescendo, converters=[FoundryStrategy.Base64]) with patch.object( - RedTeamAgent, "_resolve_seed_groups_async", new_callable=AsyncMock, return_value=mock_memory_seed_groups + RedTeamAgent, + "_resolve_seed_groups_by_dataset_async", + new_callable=AsyncMock, + return_value={"memory": mock_memory_seed_groups}, ): scenario = RedTeamAgent( attack_scoring_config=AttackScoringConfig(objective_scorer=mock_objective_scorer), @@ -638,7 +700,10 @@ async def test_initialize_with_mixed_composites_and_strategies( composite = FoundryComposite(attack=FoundryStrategy.Crescendo, converters=[FoundryStrategy.Base64]) with patch.object( - RedTeamAgent, "_resolve_seed_groups_async", new_callable=AsyncMock, return_value=mock_memory_seed_groups + RedTeamAgent, + "_resolve_seed_groups_by_dataset_async", + new_callable=AsyncMock, + return_value={"memory": mock_memory_seed_groups}, ): scenario = RedTeamAgent( attack_scoring_config=AttackScoringConfig(objective_scorer=mock_objective_scorer), @@ -663,7 +728,10 @@ async def test_initialize_converts_scenario_composite_strategy_to_foundry_compos legacy = ScenarioCompositeStrategy(strategies=[FoundryStrategy.Crescendo, FoundryStrategy.Base64]) with patch.object( - RedTeamAgent, "_resolve_seed_groups_async", new_callable=AsyncMock, return_value=mock_memory_seed_groups + RedTeamAgent, + "_resolve_seed_groups_by_dataset_async", + new_callable=AsyncMock, + return_value={"memory": mock_memory_seed_groups}, ): scenario = RedTeamAgent( attack_scoring_config=AttackScoringConfig(objective_scorer=mock_objective_scorer), @@ -688,7 +756,10 @@ async def test_initialize_converts_converter_first_composite_strategy( legacy = ScenarioCompositeStrategy(strategies=[FoundryStrategy.Base64, FoundryStrategy.Crescendo]) with patch.object( - RedTeamAgent, "_resolve_seed_groups_async", new_callable=AsyncMock, return_value=mock_memory_seed_groups + RedTeamAgent, + "_resolve_seed_groups_by_dataset_async", + new_callable=AsyncMock, + return_value={"memory": mock_memory_seed_groups}, ): scenario = RedTeamAgent( attack_scoring_config=AttackScoringConfig(objective_scorer=mock_objective_scorer), @@ -712,7 +783,10 @@ async def test_initialize_converts_converter_only_composite_strategy( legacy = ScenarioCompositeStrategy(strategies=[FoundryStrategy.Base64, FoundryStrategy.ROT13]) with patch.object( - RedTeamAgent, "_resolve_seed_groups_async", new_callable=AsyncMock, return_value=mock_memory_seed_groups + RedTeamAgent, + "_resolve_seed_groups_by_dataset_async", + new_callable=AsyncMock, + return_value={"memory": mock_memory_seed_groups}, ): scenario = RedTeamAgent( attack_scoring_config=AttackScoringConfig(objective_scorer=mock_objective_scorer), @@ -739,8 +813,8 @@ async def test_one_resolution_call_baseline_matches_strategies(self, mock_object seed_groups = [SeedAttackGroup(seeds=[SeedObjective(value=f"obj{i}")]) for i in range(10)] config = DatasetAttackConfiguration(seed_groups=seed_groups, max_dataset_size=3) - first_sample = seed_groups[:3] - second_sample = seed_groups[5:8] + first_sample = [("inline", group) for group in seed_groups[:3]] + second_sample = [("inline", group) for group in seed_groups[5:8]] with patch( "pyrit.scenario.core.dataset_configuration.random.sample", side_effect=[first_sample, second_sample], diff --git a/tests/unit/scenario/garak/test_encoding.py b/tests/unit/scenario/garak/test_encoding.py index c26784f952..555e4eaf98 100644 --- a/tests/unit/scenario/garak/test_encoding.py +++ b/tests/unit/scenario/garak/test_encoding.py @@ -101,7 +101,7 @@ def test_init_with_default_seed_prompts(self, mock_objective_target, mock_object """Test initialization with default seed prompts (Garak dataset).""" from unittest.mock import patch - with patch.object(Encoding, "_resolve_seed_groups_async", new_callable=AsyncMock, return_value=[]): + with patch.object(Encoding, "_resolve_seed_groups_by_dataset_async", new_callable=AsyncMock, return_value={}): scenario = Encoding( objective_scorer=mock_objective_scorer, ) @@ -113,7 +113,7 @@ def test_init_with_custom_scorer(self, mock_objective_target, mock_objective_sco """Test initialization with custom objective scorer.""" from unittest.mock import patch - with patch.object(Encoding, "_resolve_seed_groups_async", new_callable=AsyncMock, return_value=[]): + with patch.object(Encoding, "_resolve_seed_groups_by_dataset_async", new_callable=AsyncMock, return_value={}): scenario = Encoding( objective_scorer=mock_objective_scorer, ) @@ -124,7 +124,7 @@ def test_init_creates_default_scorer_when_not_provided(self, mock_objective_targ """Test that initialization creates default DecodingScorer when not provided.""" from unittest.mock import patch - with patch.object(Encoding, "_resolve_seed_groups_async", new_callable=AsyncMock, return_value=[]): + with patch.object(Encoding, "_resolve_seed_groups_by_dataset_async", new_callable=AsyncMock, return_value={}): scenario = Encoding() # Should create a DecodingScorer by default @@ -137,7 +137,7 @@ async def test_init_raises_exception_when_no_datasets_available(self, mock_objec from pyrit.scenario.core.dataset_configuration import DatasetConstraintError - # Don't mock _resolve_seed_groups_async; let it try to load from empty memory. + # Don't mock _resolve_seed_groups_by_dataset_async; let it try to load from empty memory. # Disable the provider fallback so memory stays empty and the scenario raises. scenario = Encoding(objective_scorer=mock_objective_scorer) @@ -150,7 +150,7 @@ def test_init_with_memory_labels(self, mock_objective_target, mock_objective_sco """Test initialization with memory labels.""" from unittest.mock import patch - with patch.object(Encoding, "_resolve_seed_groups_async", new_callable=AsyncMock, return_value=[]): + with patch.object(Encoding, "_resolve_seed_groups_by_dataset_async", new_callable=AsyncMock, return_value={}): scenario = Encoding( objective_scorer=mock_objective_scorer, ) @@ -164,7 +164,7 @@ def test_init_with_custom_encoding_templates(self, mock_objective_target, mock_o custom_templates = ["template1", "template2"] - with patch.object(Encoding, "_resolve_seed_groups_async", new_callable=AsyncMock, return_value=[]): + with patch.object(Encoding, "_resolve_seed_groups_by_dataset_async", new_callable=AsyncMock, return_value={}): scenario = Encoding( encoding_templates=custom_templates, objective_scorer=mock_objective_scorer, @@ -176,7 +176,7 @@ def test_init_with_max_concurrency(self, mock_objective_target, mock_objective_s """Test initialization with custom max_concurrency.""" from unittest.mock import patch - with patch.object(Encoding, "_resolve_seed_groups_async", new_callable=AsyncMock, return_value=[]): + with patch.object(Encoding, "_resolve_seed_groups_by_dataset_async", new_callable=AsyncMock, return_value={}): scenario = Encoding( objective_scorer=mock_objective_scorer, ) @@ -191,7 +191,10 @@ async def test_init_attack_strategies( from unittest.mock import patch with patch.object( - Encoding, "_resolve_seed_groups_async", new_callable=AsyncMock, return_value=mock_seed_attack_groups + Encoding, + "_resolve_seed_groups_by_dataset_async", + new_callable=AsyncMock, + return_value={"memory": mock_seed_attack_groups}, ): scenario = Encoding( objective_scorer=mock_objective_scorer, @@ -218,7 +221,10 @@ async def test_get_atomic_attacks_async_returns_attacks( from unittest.mock import patch with patch.object( - Encoding, "_resolve_seed_groups_async", new_callable=AsyncMock, return_value=mock_seed_attack_groups + Encoding, + "_resolve_seed_groups_by_dataset_async", + new_callable=AsyncMock, + return_value={"memory": mock_seed_attack_groups}, ): scenario = Encoding( objective_scorer=mock_objective_scorer, @@ -238,14 +244,17 @@ async def test_get_converter_attacks_returns_multiple_encodings( from unittest.mock import patch with patch.object( - Encoding, "_resolve_seed_groups_async", new_callable=AsyncMock, return_value=mock_seed_attack_groups + Encoding, + "_resolve_seed_groups_by_dataset_async", + new_callable=AsyncMock, + return_value={"memory": mock_seed_attack_groups}, ): scenario = Encoding( objective_scorer=mock_objective_scorer, ) await scenario.initialize_async(objective_target=mock_objective_target, dataset_config=mock_dataset_config) - attack_runs = scenario._get_converter_attacks() + attack_runs = scenario._get_converter_attacks(seed_groups=mock_seed_attack_groups) # Should have multiple attack runs for different encodings # The list includes: Base64 (4 variants), Base2048, Base16, Base32, ASCII85 (2), hex, @@ -259,14 +268,19 @@ async def test_get_prompt_attacks_creates_attack_runs( from unittest.mock import patch with patch.object( - Encoding, "_resolve_seed_groups_async", new_callable=AsyncMock, return_value=mock_seed_attack_groups + Encoding, + "_resolve_seed_groups_by_dataset_async", + new_callable=AsyncMock, + return_value={"memory": mock_seed_attack_groups}, ): scenario = Encoding( objective_scorer=mock_objective_scorer, ) await scenario.initialize_async(objective_target=mock_objective_target, dataset_config=mock_dataset_config) - attack_runs = scenario._get_prompt_attacks(converters=[Base64Converter()], encoding_name="Base64") + attack_runs = scenario._get_prompt_attacks( + converters=[Base64Converter()], encoding_name="Base64", seed_groups=mock_seed_attack_groups + ) # Should create attack runs assert len(attack_runs) > 0 @@ -288,14 +302,19 @@ async def test_attack_runs_include_objectives( from unittest.mock import patch with patch.object( - Encoding, "_resolve_seed_groups_async", new_callable=AsyncMock, return_value=mock_seed_attack_groups + Encoding, + "_resolve_seed_groups_by_dataset_async", + new_callable=AsyncMock, + return_value={"memory": mock_seed_attack_groups}, ): scenario = Encoding( objective_scorer=mock_objective_scorer, ) await scenario.initialize_async(objective_target=mock_objective_target, dataset_config=mock_dataset_config) - attack_runs = scenario._get_prompt_attacks(converters=[Base64Converter()], encoding_name="Base64") + attack_runs = scenario._get_prompt_attacks( + converters=[Base64Converter()], encoding_name="Base64", seed_groups=mock_seed_attack_groups + ) # Check that seed groups contain objectives with the expected format for run in attack_runs: @@ -319,7 +338,10 @@ async def test_scenario_initialization( from unittest.mock import patch with patch.object( - Encoding, "_resolve_seed_groups_async", new_callable=AsyncMock, return_value=mock_seed_attack_groups + Encoding, + "_resolve_seed_groups_by_dataset_async", + new_callable=AsyncMock, + return_value={"memory": mock_seed_attack_groups}, ): scenario = Encoding( objective_scorer=mock_objective_scorer, @@ -333,22 +355,26 @@ async def test_scenario_initialization( async def test_resolve_seed_groups_loads_garak_data( self, mock_objective_target, mock_objective_scorer, mock_seed_attack_groups, mock_dataset_config ): - """Test that _resolve_seed_groups_async loads data from Garak datasets.""" + """Test that _resolve_seed_groups_by_dataset_async loads data from Garak datasets.""" from unittest.mock import patch with patch.object( - Encoding, "_resolve_seed_groups_async", new_callable=AsyncMock, return_value=mock_seed_attack_groups + Encoding, + "_resolve_seed_groups_by_dataset_async", + new_callable=AsyncMock, + return_value={"memory": mock_seed_attack_groups}, ): scenario = Encoding( objective_scorer=mock_objective_scorer, ) - # After resolve, should have seed groups - resolved = await scenario._resolve_seed_groups_async() - assert len(resolved) > 0 + # After resolve, should have seed groups keyed by dataset + resolved = await scenario._resolve_seed_groups_by_dataset_async() + flattened = [group for groups in resolved.values() for group in groups] + assert flattened # Verify it's returning SeedAttackGroup objects - assert all(isinstance(group, SeedAttackGroup) for group in resolved) + assert all(isinstance(group, SeedAttackGroup) for group in flattened) @pytest.mark.usefixtures("patch_central_database") @@ -451,8 +477,8 @@ async def test_one_resolution_call_baseline_matches_strategies(self, mock_object seed_groups = [SeedAttackGroup(seeds=[SeedObjective(value=f"obj{i}")]) for i in range(10)] config = DatasetAttackConfiguration(seed_groups=seed_groups, max_dataset_size=3) - first_sample = seed_groups[:3] - second_sample = seed_groups[5:8] + first_sample = [("inline", group) for group in seed_groups[:3]] + second_sample = [("inline", group) for group in seed_groups[5:8]] with patch( "pyrit.scenario.core.dataset_configuration.random.sample", side_effect=[first_sample, second_sample],