MAINT Simplify Foundry technique mapping

pyrit/scenario/scenarios/foundry/red_team_agent.py

@@ -10,9 +10,11 @@
 """
 
 import logging
+from collections.abc import Callable, Mapping
 from dataclasses import dataclass, field
 from inspect import signature
-from typing import TYPE_CHECKING, Any, TypeVar, cast
+from types import MappingProxyType
+from typing import TYPE_CHECKING, Any, ClassVar, TypeVar, cast
 
 from pyrit.common import apply_defaults
 from pyrit.converter import (
@@ -40,7 +42,13 @@
 from pyrit.converter.binary_converter import BinaryConverter
 from pyrit.converter.token_smuggling.ascii_smuggler_converter import AsciiSmugglerConverter
 from pyrit.datasets import TextJailBreak
-from pyrit.executor.attack import CrescendoAttack, PromptSendingAttack, RedTeamingAttack, TreeOfAttacksWithPruningAttack
+from pyrit.executor.attack import (
+    AttackStrategy,
+    CrescendoAttack,
+    PromptSendingAttack,
+    RedTeamingAttack,
+    TreeOfAttacksWithPruningAttack,
+)
 from pyrit.executor.attack.core.attack_config import AttackAdversarialConfig, AttackConverterConfig, AttackScoringConfig
 from pyrit.models import AttackSeedGroup
 from pyrit.prompt_normalizer.converter_configuration import ConverterConfiguration
@@ -57,8 +65,6 @@
 if TYPE_CHECKING:
     from collections.abc import Sequence
 
-    from pyrit.executor.attack.core.attack_strategy import AttackStrategy
-
 AttackStrategyT = TypeVar("AttackStrategyT", bound="AttackStrategy[Any, Any]")
 logger = logging.getLogger(__name__)
 
@@ -193,6 +199,40 @@ def default(cls) -> "FoundryTechnique":
         return cls.EASY
 
 
+@dataclass(frozen=True)
+class _AttackSpecification:
+    """Declarative construction details for a Foundry attack technique."""
+
+    attack_type: type[AttackStrategy[Any, Any]]
+    kwargs: tuple[tuple[str, Any], ...] = ()
+
+
+@dataclass(frozen=True)
+class _ConverterSpecification:
+    """Declarative construction details for a Foundry converter technique."""
+
+    factory: Callable[..., Converter]
+    kwargs: tuple[tuple[str, Any], ...] = ()
+    deferred_kwargs: tuple[tuple[str, Callable[[], Any]], ...] = ()
+    target_kwarg: str | None = None
+
+    def create(self, *, converter_target: PromptTarget) -> Converter:
+        """
+        Create a fresh converter from this specification.
+
+        Args:
+            converter_target (PromptTarget): Target supplied to converters that require one.
+
+        Returns:
+            Converter: The configured converter instance.
+        """
+        kwargs = dict(self.kwargs)
+        kwargs.update((name, factory()) for name, factory in self.deferred_kwargs)
+        if self.target_kwarg:
+            kwargs[self.target_kwarg] = converter_target
+        return self.factory(**kwargs)
+
+
 class RedTeamAgent(Scenario):
     """
     RedTeamAgent is a preconfigured scenario that automatically generates multiple
@@ -208,6 +248,56 @@ class RedTeamAgent(Scenario):
     """
 
     VERSION: int = 1
+    _DEFAULT_ATTACK_SPECIFICATION: ClassVar[_AttackSpecification] = _AttackSpecification(PromptSendingAttack)
+    _ATTACK_SPECIFICATIONS: ClassVar[Mapping[FoundryTechnique, _AttackSpecification]] = MappingProxyType(
+        {
+            FoundryTechnique.Crescendo: _AttackSpecification(CrescendoAttack),
+            FoundryTechnique.MultiTurn: _AttackSpecification(RedTeamingAttack),
+            FoundryTechnique.Pair: _AttackSpecification(
+                TreeOfAttacksWithPruningAttack,
+                kwargs=(("tree_width", 1),),
+            ),
+            FoundryTechnique.Tap: _AttackSpecification(TreeOfAttacksWithPruningAttack),
+        }
+    )
+    _CONVERTER_SPECIFICATIONS: ClassVar[Mapping[FoundryTechnique, _ConverterSpecification]] = MappingProxyType(
+        {
+            FoundryTechnique.AnsiAttack: _ConverterSpecification(AnsiAttackConverter),
+            FoundryTechnique.AsciiArt: _ConverterSpecification(AsciiArtConverter),
+            FoundryTechnique.AsciiSmuggler: _ConverterSpecification(AsciiSmugglerConverter),
+            FoundryTechnique.Atbash: _ConverterSpecification(AtbashConverter),
+            FoundryTechnique.Base64: _ConverterSpecification(Base64Converter),
+            FoundryTechnique.Binary: _ConverterSpecification(BinaryConverter),
+            FoundryTechnique.Caesar: _ConverterSpecification(
+                CaesarConverter,
+                kwargs=(("caesar_offset", 3),),
+            ),
+            FoundryTechnique.CharacterSpace: _ConverterSpecification(CharacterSpaceConverter),
+            FoundryTechnique.CharSwap: _ConverterSpecification(CharSwapConverter),
+            FoundryTechnique.Diacritic: _ConverterSpecification(DiacriticConverter),
+            FoundryTechnique.Flip: _ConverterSpecification(FlipConverter),
+            FoundryTechnique.Leetspeak: _ConverterSpecification(LeetspeakConverter),
+            FoundryTechnique.Morse: _ConverterSpecification(MorseConverter),
+            FoundryTechnique.ROT13: _ConverterSpecification(ROT13Converter),
+            FoundryTechnique.SuffixAppend: _ConverterSpecification(
+                SuffixAppendConverter,
+                kwargs=(("suffix", "!!!"),),
+            ),
+            FoundryTechnique.StringJoin: _ConverterSpecification(StringJoinConverter),
+            FoundryTechnique.Tense: _ConverterSpecification(
+                TenseConverter,
+                kwargs=(("tense", "past"),),
+                target_kwarg="converter_target",
+            ),
+            FoundryTechnique.UnicodeConfusable: _ConverterSpecification(UnicodeConfusableConverter),
+            FoundryTechnique.UnicodeSubstitution: _ConverterSpecification(UnicodeSubstitutionConverter),
+            FoundryTechnique.Url: _ConverterSpecification(UrlConverter),
+            FoundryTechnique.Jailbreak: _ConverterSpecification(
+                TextJailbreakConverter,
+                deferred_kwargs=(("jailbreak_template", lambda: TextJailBreak(random_template=True)),),
+            ),
+        }
+    )
 
     @apply_defaults
     def __init__(
@@ -386,70 +476,22 @@ def _get_attack_from_technique(
         Raises:
             ValueError: If a converter technique in the composite is not recognized.
         """
-        attack: AttackStrategy[Any, Any]
-
-        attack_type: type[AttackStrategy[Any, Any]] = PromptSendingAttack
-        attack_kwargs: dict[str, Any] = {}
-        if composite.attack is not None:
-            if composite.attack == FoundryTechnique.Crescendo:
-                attack_type = CrescendoAttack
-            elif composite.attack == FoundryTechnique.MultiTurn:
-                attack_type = RedTeamingAttack
-            elif composite.attack == FoundryTechnique.Pair:
-                attack_type = TreeOfAttacksWithPruningAttack
-                attack_kwargs = {"tree_width": 1}
-            elif composite.attack == FoundryTechnique.Tap:
-                attack_type = TreeOfAttacksWithPruningAttack
-
+        attack_specification = self._ATTACK_SPECIFICATIONS.get(
+            composite.attack,
+            self._DEFAULT_ATTACK_SPECIFICATION,
+        )
         converters: list[Converter] = []
         for technique in composite.converters:
-            if technique == FoundryTechnique.AnsiAttack:
-                converters.append(AnsiAttackConverter())
-            elif technique == FoundryTechnique.AsciiArt:
-                converters.append(AsciiArtConverter())
-            elif technique == FoundryTechnique.AsciiSmuggler:
-                converters.append(AsciiSmugglerConverter())
-            elif technique == FoundryTechnique.Atbash:
-                converters.append(AtbashConverter())
-            elif technique == FoundryTechnique.Base64:
-                converters.append(Base64Converter())
-            elif technique == FoundryTechnique.Binary:
-                converters.append(BinaryConverter())
-            elif technique == FoundryTechnique.Caesar:
-                converters.append(CaesarConverter(caesar_offset=3))
-            elif technique == FoundryTechnique.CharacterSpace:
-                converters.append(CharacterSpaceConverter())
-            elif technique == FoundryTechnique.CharSwap:
-                converters.append(CharSwapConverter())
-            elif technique == FoundryTechnique.Diacritic:
-                converters.append(DiacriticConverter())
-            elif technique == FoundryTechnique.Flip:
-                converters.append(FlipConverter())
-            elif technique == FoundryTechnique.Leetspeak:
-                converters.append(LeetspeakConverter())
-            elif technique == FoundryTechnique.Morse:
-                converters.append(MorseConverter())
-            elif technique == FoundryTechnique.ROT13:
-                converters.append(ROT13Converter())
-            elif technique == FoundryTechnique.SuffixAppend:
-                converters.append(SuffixAppendConverter(suffix="!!!"))
-            elif technique == FoundryTechnique.StringJoin:
-                converters.append(StringJoinConverter())
-            elif technique == FoundryTechnique.Tense:
-                converters.append(TenseConverter(tense="past", converter_target=self._adversarial_chat))
-            elif technique == FoundryTechnique.UnicodeConfusable:
-                converters.append(UnicodeConfusableConverter())
-            elif technique == FoundryTechnique.UnicodeSubstitution:
-                converters.append(UnicodeSubstitutionConverter())
-            elif technique == FoundryTechnique.Url:
-                converters.append(UrlConverter())
-            elif technique == FoundryTechnique.Jailbreak:
-                jailbreak_template = TextJailBreak(random_template=True)
-                converters.append(TextJailbreakConverter(jailbreak_template=jailbreak_template))
-            else:
+            converter_specification = self._CONVERTER_SPECIFICATIONS.get(technique)
+            if converter_specification is None:
                 raise ValueError(f"Unknown technique: {technique}")
+            converters.append(converter_specification.create(converter_target=self._adversarial_chat))
 
-        attack = self._get_attack(attack_type=attack_type, converters=converters, attack_kwargs=attack_kwargs)
+        attack = self._get_attack(
+            attack_type=attack_specification.attack_type,
+            converters=converters,
+            attack_kwargs=dict(attack_specification.kwargs),
+        )
 
         return AtomicAttack(
             atomic_attack_name=composite.name,