MAINT: Minor Scenario class improvements (microsoft#1190)

Co-authored-by: Victor Valbuena <50061128+ValbuenaVC@users.noreply.github.com>

Author: rlundeen2

pyrit/scenarios/scenario.py

@@ -79,7 +79,6 @@ def __init__(
         name: str,
         version: int,
         strategy_class: Type[ScenarioStrategy],
-        default_aggregate: ScenarioStrategy,
         objective_scorer_identifier: Optional[Dict[str, str]] = None,
         include_default_baseline: bool = True,
         scenario_result_id: Optional[Union[uuid.UUID, str]] = None,
@@ -91,8 +90,6 @@ def __init__(
             name (str): Descriptive name for the scenario.
             version (int): Version number of the scenario.
             strategy_class (Type[ScenarioStrategy]): The strategy enum class for this scenario.
-            default_aggregate (ScenarioStrategy): The default aggregate strategy to use when
-                scenario_strategies is None in initialize_async.
             objective_scorer_identifier (Optional[Dict[str, str]]): Identifier for the objective scorer.
             include_default_baseline (bool): Whether to include a baseline atomic attack that sends all objectives
                 from the first atomic attack without modifications. Most scenarios should have some kind of
@@ -119,7 +116,6 @@ def __init__(
 
         # Store strategy configuration for use in initialize_async
         self._strategy_class = strategy_class
-        self._default_aggregate = default_aggregate
 
         # These will be set in initialize_async
         self._objective_target: Optional[PromptTarget] = None
@@ -284,7 +280,7 @@ async def initialize_async(
 
         # Prepare scenario strategies using the stored configuration
         self._scenario_composites = self._strategy_class.prepare_scenario_strategies(
-            scenario_strategies, default_aggregate=self._default_aggregate
+            scenario_strategies, default_aggregate=self.get_default_strategy()
         )
 
         self._atomic_attacks = await self._get_atomic_attacks_async()

pyrit/scenarios/scenario_strategy.py

@@ -393,6 +393,47 @@ def is_single_strategy(self) -> bool:
         """Check if this composition contains only a single strategy."""
         return len(self._strategies) == 1
 
+    @staticmethod
+    def extract_single_strategy_values(
+        composites: Sequence["ScenarioCompositeStrategy"], *, strategy_type: type[T]
+    ) -> Set[str]:
+        """
+        Extract strategy values from single-strategy composites.
+
+        This is a helper method for scenarios that don't support composition and need
+        to filter or map strategies by their values. It flattens the composites into
+        a simple set of strategy values.
+
+        This method enforces that all composites contain only a single strategy. If any
+        composite contains multiple strategies, a ValueError is raised.
+
+        Args:
+            composites (Sequence[ScenarioCompositeStrategy]): List of composite strategies.
+                Each composite must contain only a single strategy.
+            strategy_type (type[T]): The strategy enum type to filter by.
+
+        Returns:
+            Set[str]: Set of strategy values (e.g., {"base64", "rot13", "morse_code"}).
+
+        Raises:
+            ValueError: If any composite contains multiple strategies.
+        """
+        # Check that all composites are single-strategy
+        multi_strategy_composites = [comp for comp in composites if not comp.is_single_strategy]
+        if multi_strategy_composites:
+            composite_names = [comp.name for comp in multi_strategy_composites]
+            raise ValueError(
+                f"extract_single_strategy_values() requires all composites to contain a single strategy. "
+                f"Found composites with multiple strategies: {composite_names}"
+            )
+
+        return {
+            strategy.value
+            for composite in composites
+            for strategy in composite.strategies
+            if isinstance(strategy, strategy_type)
+        }
+
     @staticmethod
     def get_composite_name(strategies: Sequence[ScenarioStrategy]) -> str:
         """

pyrit/scenarios/scenarios/encoding_scenario.py

@@ -36,6 +36,7 @@
 from pyrit.scenarios.atomic_attack import AtomicAttack
 from pyrit.scenarios.scenario import Scenario
 from pyrit.scenarios.scenario_strategy import (
+    ScenarioCompositeStrategy,
     ScenarioStrategy,
 )
 from pyrit.score import TrueFalseScorer
@@ -153,7 +154,6 @@ def __init__(
             name="Encoding Scenario",
             version=self.version,
             strategy_class=EncodingStrategy,
-            default_aggregate=EncodingStrategy.ALL,
             objective_scorer_identifier=objective_scorer.get_identifier(),
             include_default_baseline=include_baseline,
             scenario_result_id=scenario_result_id,
@@ -224,8 +224,9 @@ def _get_converter_attacks(self) -> list[AtomicAttack]:
         ]
 
         # Filter to only include selected strategies
-        # Extract strategy names from composites (each has exactly one strategy since composition not supported)
-        selected_encoding_names = {comp.strategies[0].value for comp in self._scenario_composites if comp.strategies}
+        selected_encoding_names = ScenarioCompositeStrategy.extract_single_strategy_values(
+            self._scenario_composites, strategy_type=EncodingStrategy
+        )
         converters_with_encodings = [
             (conv, name) for conv, name in all_converters_with_encodings if name in selected_encoding_names
         ]

pyrit/scenarios/scenarios/foundry_scenario.py

@@ -275,7 +275,6 @@ def __init__(
             name="Foundry Scenario",
             version=self.version,
             strategy_class=FoundryStrategy,
-            default_aggregate=FoundryStrategy.EASY,
             objective_scorer_identifier=self._objective_scorer.get_identifier(),
             include_default_baseline=include_baseline,
             scenario_result_id=scenario_result_id,

tests/unit/scenarios/test_scenario.py

@@ -94,18 +94,17 @@ def __init__(self, atomic_attacks_to_return=None, **kwargs):
         # Default include_default_baseline=False for tests unless explicitly specified
         kwargs.setdefault("include_default_baseline", False)
 
-        # Add required strategy_class and default_aggregate if not provided
+        # Add required strategy_class if not provided
 
         class TestStrategy(ScenarioStrategy):
-            TEST = ("test", set())
+            TEST = ("test", {"concrete"})  # Tagged as concrete, not aggregate
             ALL = ("all", {"all"})
 
             @classmethod
             def get_aggregate_tags(cls) -> set[str]:
                 return {"all"}
 
         kwargs.setdefault("strategy_class", TestStrategy)
-        kwargs.setdefault("default_aggregate", TestStrategy.ALL)
 
         super().__init__(**kwargs)
         self._atomic_attacks_to_return = atomic_attacks_to_return or []
@@ -118,7 +117,7 @@ def get_strategy_class(cls):
 
         # Return a simple mock strategy class for testing
         class TestStrategy(ScenarioStrategy):
-            TEST = ("test", set())
+            TEST = ("test", {"concrete"})  # Tagged as concrete, not aggregate
             ALL = ("all", {"all"})
 
             @classmethod
@@ -130,7 +129,7 @@ def get_aggregate_tags(cls) -> set[str]:
     @classmethod
     def get_default_strategy(cls):
         """Return the default strategy for testing."""
-        return cls.get_strategy_class().TEST
+        return cls.get_strategy_class().ALL
 
     async def _get_atomic_attacks_async(self):
         return self._atomic_attacks_to_return

tests/unit/scenarios/test_scenario_partial_results.py

@@ -51,9 +51,8 @@ def __init__(self, *, atomic_attacks_to_return=None, **kwargs):
 
         # Get strategy_class from kwargs or use default
         strategy_class = kwargs.pop("strategy_class", None) or self.get_strategy_class()
-        default_aggregate = kwargs.pop("default_aggregate", None) or self.get_default_strategy()
 
-        super().__init__(strategy_class=strategy_class, default_aggregate=default_aggregate, **kwargs)
+        super().__init__(strategy_class=strategy_class, **kwargs)
         self._test_atomic_attacks = atomic_attacks_to_return or []
 
     async def _get_atomic_attacks_async(self):
@@ -63,7 +62,7 @@ async def _get_atomic_attacks_async(self):
     def get_strategy_class(cls):
 
         class TestStrategy(ScenarioStrategy):
-            CONCRETE = ("concrete", set())
+            CONCRETE = ("concrete", {"concrete"})
             ALL = ("all", {"all"})
 
             @classmethod

tests/unit/scenarios/test_scenario_retry.py

@@ -122,9 +122,8 @@ def __init__(self, atomic_attacks_to_return=None, **kwargs):
 
         # Get strategy_class from kwargs or use default
         strategy_class = kwargs.pop("strategy_class", None) or self.get_strategy_class()
-        default_aggregate = kwargs.pop("default_aggregate", None) or self.get_default_strategy()
 
-        super().__init__(strategy_class=strategy_class, default_aggregate=default_aggregate, **kwargs)
+        super().__init__(strategy_class=strategy_class, **kwargs)
         self._atomic_attacks_to_return = atomic_attacks_to_return or []
 
     @classmethod
@@ -133,7 +132,7 @@ def get_strategy_class(cls):
 
         # Return a simple mock strategy class for testing
         class TestStrategy(ScenarioStrategy):
-            CONCRETE = ("concrete", set())
+            CONCRETE = ("concrete", {"concrete"})
             ALL = ("all", {"all"})
 
             @classmethod

tests/unit/scenarios/test_strategy_validation.py

@@ -5,7 +5,7 @@
 
 import pytest
 
-from pyrit.scenarios import EncodingStrategy, FoundryStrategy
+from pyrit.scenarios import EncodingStrategy, FoundryStrategy, ScenarioCompositeStrategy
 
 
 class TestStrategyValidation:
@@ -49,3 +49,53 @@ def test_foundry_validation_rejects_attacks_with_converters_and_another_attack(s
             FoundryStrategy.validate_composition(
                 [FoundryStrategy.Base64, FoundryStrategy.Crescendo, FoundryStrategy.MultiTurn]
             )
+
+
+class TestScenarioCompositeStrategyExtraction:
+    """Test extraction of strategy values from composite strategies."""
+
+    def test_extract_single_strategy_values_with_single_strategies(self):
+        """Test extracting values from single-strategy composites."""
+        composites = [
+            ScenarioCompositeStrategy(strategies=[EncodingStrategy.Base64]),
+            ScenarioCompositeStrategy(strategies=[EncodingStrategy.ROT13]),
+            ScenarioCompositeStrategy(strategies=[EncodingStrategy.Atbash]),
+        ]
+
+        values = ScenarioCompositeStrategy.extract_single_strategy_values(composites, strategy_type=EncodingStrategy)
+
+        assert values == {"base64", "rot13", "atbash"}
+
+    def test_extract_single_strategy_values_filters_by_type(self):
+        """Test that extraction filters by strategy type."""
+        composites = [
+            ScenarioCompositeStrategy(strategies=[EncodingStrategy.Base64]),
+            ScenarioCompositeStrategy(strategies=[FoundryStrategy.ROT13]),
+        ]
+
+        # Extract only EncodingStrategy values
+        encoding_values = ScenarioCompositeStrategy.extract_single_strategy_values(
+            composites, strategy_type=EncodingStrategy
+        )
+        assert encoding_values == {"base64"}
+
+        # Extract only FoundryStrategy values
+        foundry_values = ScenarioCompositeStrategy.extract_single_strategy_values(
+            composites, strategy_type=FoundryStrategy
+        )
+        assert foundry_values == {"rot13"}
+
+    def test_extract_single_strategy_values_rejects_multi_strategy_composites(self):
+        """Test that extraction raises error if any composite has multiple strategies."""
+        composites = [
+            ScenarioCompositeStrategy(strategies=[FoundryStrategy.Base64]),
+            ScenarioCompositeStrategy(strategies=[FoundryStrategy.ROT13, FoundryStrategy.Atbash]),  # Multi-strategy!
+        ]
+
+        with pytest.raises(ValueError, match="extract_single_strategy_values.*requires all composites"):
+            ScenarioCompositeStrategy.extract_single_strategy_values(composites, strategy_type=FoundryStrategy)
+
+    def test_extract_single_strategy_values_with_empty_list(self):
+        """Test that extraction handles empty composite list."""
+        values = ScenarioCompositeStrategy.extract_single_strategy_values([], strategy_type=EncodingStrategy)
+        assert values == set()