diff --git a/pyrit/models/seeds/attack_seed_group.py b/pyrit/models/seeds/attack_seed_group.py index b96cdccf0c..99d325dd48 100644 --- a/pyrit/models/seeds/attack_seed_group.py +++ b/pyrit/models/seeds/attack_seed_group.py @@ -21,6 +21,7 @@ from collections.abc import Sequence from pyrit.models.seeds.attack_technique_seed_group import AttackTechniqueSeedGroup + from pyrit.models.seeds.seed import Seed class AttackSeedGroup(SeedGroup): @@ -144,6 +145,8 @@ def with_technique(self, *, technique: AttackTechniqueSeedGroup) -> AttackSeedGr Raises: ValueError: If the technique contains a SeedSimulatedConversation whose sequence range overlaps with existing prompt sequences. + ValueError: If preserving prompt placement combines conflicting roles at + the same sequence. """ # Pre-merge compatibility check with a clear error message if not self.is_compatible_with_technique(technique=technique): @@ -157,18 +160,12 @@ def with_technique(self, *, technique: AttackTechniqueSeedGroup) -> AttackSeedGr f"overlapping the simulated conversation range are incompatible." ) - base = list(self.seeds) + base_seeds = [copy.deepcopy(seed) for seed in self.seeds] + technique_seeds = [copy.deepcopy(seed) for seed in technique.seeds] idx = technique.insertion_index - technique_seeds = list(technique.seeds) - merged_seeds = base + technique_seeds if idx is None else base[:idx] + technique_seeds + base[idx:] - - # ``self`` and ``technique`` may be shared across multiple ``with_technique`` - # calls (e.g. the dispatcher reuses one ``bundle.seed_technique`` instance - # across every objective). Deepcopy first so the per-seed mutation below - # and the fresh group_id assigned by ``AttackSeedGroup.__init__`` only - # touch the returned group, leaving the originals untouched as the - # docstring promises. - merged_seeds = [copy.deepcopy(seed) for seed in merged_seeds] + merged_seeds = ( + base_seeds + technique_seeds if idx is None else base_seeds[:idx] + technique_seeds + base_seeds[idx:] + ) # Clear group IDs so the new group assigns a fresh one. # ``_enforce_consistent_group_id`` in the constructor will overwrite @@ -176,19 +173,33 @@ def with_technique(self, *, technique: AttackTechniqueSeedGroup) -> AttackSeedGr for seed in merged_seeds: seed.prompt_group_id = None - # Normalize prompt sequences to dense, 0-based order preserving relative - # ordering. A technique whose seed leads the conversation (e.g. a system - # prompt built at ``sequence=-1`` by ``from_system_prompt``) is thereby - # prepended cleanly: it lands at sequence 0 and the existing turns shift - # up (user 0 -> 1, assistant 1 -> 2, ...), rather than leaving a negative - # or sparse sequence. This keeps the merge robust no matter how the base - # group was numbered. Skipped when a simulated conversation is present, - # since its ``sequence_range`` is absolute and self-consistent. - has_simulated = any(isinstance(seed, SeedSimulatedConversation) for seed in merged_seeds) - if not has_simulated: - prompt_seeds = [seed for seed in merged_seeds if isinstance(seed, SeedPrompt)] - rank_by_sequence = {seq: rank for rank, seq in enumerate(sorted({p.sequence for p in prompt_seeds}))} - for seed in prompt_seeds: - seed.sequence = rank_by_sequence[seed.sequence] + self._normalize_prompt_sequences( + base_seeds=base_seeds, + technique_seeds=technique_seeds, + prepend_technique=technique.prompt_placement == "prepend", + ) return AttackSeedGroup(seeds=merged_seeds) + + @staticmethod + def _normalize_prompt_sequences( + *, + base_seeds: Sequence[Seed], + technique_seeds: Sequence[Seed], + prepend_technique: bool, + ) -> None: + """Normalize merged prompt sequences while preserving source-relative order.""" + all_seeds = [*base_seeds, *technique_seeds] + # Simulated conversations reserve an absolute sequence range; renumbering only prompts + # could invalidate that range or create an overlap. + if any(isinstance(seed, SeedSimulatedConversation) for seed in all_seeds): + return + + seed_groups = (technique_seeds, base_seeds) if prepend_technique else (all_seeds,) + next_sequence = 0 + for seeds in seed_groups: + prompts = [seed for seed in seeds if isinstance(seed, SeedPrompt)] + rank_by_sequence = {value: rank for rank, value in enumerate(sorted({p.sequence for p in prompts}))} + for prompt in prompts: + prompt.sequence = next_sequence + rank_by_sequence[prompt.sequence] + next_sequence += len(rank_by_sequence) diff --git a/pyrit/models/seeds/attack_technique_seed_group.py b/pyrit/models/seeds/attack_technique_seed_group.py index 0e8341450c..d9f0316cf3 100644 --- a/pyrit/models/seeds/attack_technique_seed_group.py +++ b/pyrit/models/seeds/attack_technique_seed_group.py @@ -11,6 +11,10 @@ from __future__ import annotations +from typing import Literal + +from pydantic import Field + from pyrit.models.seeds.seed_group import SeedGroup from pyrit.models.seeds.seed_objective import SeedObjective from pyrit.models.seeds.seed_prompt import SeedPrompt @@ -31,6 +35,15 @@ class AttackTechniqueSeedGroup(SeedGroup): # ``None`` (default) appends at the end; an integer inserts before that position. insertion_index: int | None = None + prompt_placement: Literal["preserve", "prepend"] = Field( + default="preserve", + description=( + '"preserve" combines existing sequence relationships. During AttackSeedGroup construction, ' + "prompts at the same sequence are grouped when roles are the same and rejected when roles conflict. " + '"prepend" places technique prompts before base prompts.' + ), + ) + @classmethod def from_system_prompt(cls, system_prompt: str, *, insertion_index: int | None = None) -> AttackTechniqueSeedGroup: """ @@ -41,13 +54,9 @@ def from_system_prompt(cls, system_prompt: str, *, insertion_index: int | None = value is wrapped verbatim (``is_jinja_template=False``), so any literal ``{{ ... }}`` in ``system_prompt`` is preserved rather than re-rendered. - The seed is built at ``sequence=-1`` as an internal "lead" marker so it orders ahead of - any user turn. When merged via ``AttackSeedGroup.with_technique`` the merged sequences are - normalized to dense 0-based order, so the system framing lands at sequence 0 and the - objective's turns shift up (user 0 -> 1, assistant 1 -> 2, ...). Without leading it, merging - onto a seed group that carries a user prompt at the default ``sequence=0`` would raise - ``Inconsistent roles found for sequence 0`` (one ``sequence`` maps to one ``Message``, which - requires a single role). + The group declares ``prompt_placement="prepend"`` so ``AttackSeedGroup.with_technique`` + places the system framing before the base prompts without relying on a reserved sequence + value. Args: system_prompt (str): The system-role instruction text. @@ -58,10 +67,9 @@ def from_system_prompt(cls, system_prompt: str, *, insertion_index: int | None = AttackTechniqueSeedGroup: A group with a single general-technique system seed. """ return cls( - seeds=[ - SeedPrompt(value=system_prompt, data_type="text", role="system", is_general_technique=True, sequence=-1) - ], + seeds=[SeedPrompt(value=system_prompt, data_type="text", role="system", is_general_technique=True)], insertion_index=insertion_index, + prompt_placement="prepend", ) def _check_invariants(self) -> None: diff --git a/pyrit/scenario/core/attack_technique.py b/pyrit/scenario/core/attack_technique.py index 7a4f135e65..cea0bbf1ea 100644 --- a/pyrit/scenario/core/attack_technique.py +++ b/pyrit/scenario/core/attack_technique.py @@ -53,20 +53,23 @@ def _build_identifier(self) -> ComponentIdentifier: Build the behavioral identity for this attack technique. The identifier always contains the attack technique as ``children["attack"]``. - When a seed technique is present, its seeds are added as - ``children["technique_seeds"]``. + When a seed technique is present, its seeds and prompt placement are included. Returns: ComponentIdentifier: The frozen identity snapshot. """ technique_seeds: list[SeedIdentifier] | None = None + identifier_params: dict[str, Any] | None = None if self._seed_technique is not None: technique_seed_ids = [SeedIdentifier.from_seed(seed) for seed in self._seed_technique.seeds] if technique_seed_ids: technique_seeds = list(technique_seed_ids) + if self._seed_technique.prompt_placement != "preserve": + identifier_params = {"prompt_placement": self._seed_technique.prompt_placement} return AttackTechniqueIdentifier.of( self, + params=identifier_params, attack=self._attack.get_identifier(), technique_seeds=technique_seeds, ) diff --git a/pyrit/scenario/scenarios/airt/jailbreak.py b/pyrit/scenario/scenarios/airt/jailbreak.py index 4bde50110e..710c7eba7c 100644 --- a/pyrit/scenario/scenarios/airt/jailbreak.py +++ b/pyrit/scenario/scenarios/airt/jailbreak.py @@ -12,7 +12,7 @@ from pyrit.converter import TextJailbreakConverter from pyrit.datasets import TextJailBreak from pyrit.executor.attack.single_turn.prompt_sending import PromptSendingAttack -from pyrit.models import AttackTechniqueSeedGroup, Parameter, SeedPrompt +from pyrit.models import AttackTechniqueSeedGroup, Parameter from pyrit.prompt_target import CapabilityName from pyrit.registry.components.attack_technique_registry import AttackTechniqueRegistry from pyrit.scenario.core.attack_technique_factory import AttackTechniqueFactory @@ -450,12 +450,7 @@ def _build_system_prompt_factory(self, *, template_file_name: str) -> AttackTech the rendered jailbreak framing. """ framing = TextJailBreak(template_file_name=template_file_name).get_jailbreak_system_prompt() - # sequence=-1 orders the system framing ahead of any user turn, so a caller-supplied seed - # group carrying a user prompt at the default sequence 0 does not raise a same-sequence - # role collision when this technique is merged in. - seed_technique = AttackTechniqueSeedGroup( - seeds=[SeedPrompt(value=framing, data_type="text", role="system", is_general_technique=True, sequence=-1)] - ) + seed_technique = AttackTechniqueSeedGroup.from_system_prompt(framing) return AttackTechniqueFactory( name=_JAILBREAK_SYSTEM_PROMPT, attack_class=PromptSendingAttack, diff --git a/tests/unit/models/test_attack_technique_seed_group.py b/tests/unit/models/test_attack_technique_seed_group.py index 8bb6b41ace..3358a60fdd 100644 --- a/tests/unit/models/test_attack_technique_seed_group.py +++ b/tests/unit/models/test_attack_technique_seed_group.py @@ -204,6 +204,7 @@ def test_default_insertion_index_is_none(self): seeds=[SeedPrompt(value="s", data_type="text", is_general_technique=True)], ) assert group.insertion_index is None + assert group.prompt_placement == "preserve" def test_insertion_index_set_to_int(self): """Test that insertion_index can be set to an integer.""" @@ -234,10 +235,11 @@ def test_builds_single_system_general_technique_seed(self): assert isinstance(seed, SeedPrompt) assert seed.value == "Follow these rules." assert seed.role == "system" - assert seed.sequence == -1 + assert seed.sequence == 0 assert seed.data_type == "text" assert seed.is_general_technique is True assert group.insertion_index is None + assert group.prompt_placement == "prepend" def test_preserves_literal_braces_without_rendering(self): """Test that literal Jinja braces are preserved (is_jinja_template stays False).""" @@ -249,6 +251,14 @@ def test_respects_insertion_index(self): group = AttackTechniqueSeedGroup.from_system_prompt("s", insertion_index=0) assert group.insertion_index == 0 + def test_prompt_placement_survives_serialization_round_trip(self): + """Test that prepend intent is preserved when the group is serialized.""" + group = AttackTechniqueSeedGroup.from_system_prompt("Follow these rules.") + + restored = AttackTechniqueSeedGroup.model_validate_json(group.model_dump_json()) + + assert restored.prompt_placement == "prepend" + class TestAttackTechniqueSeedGroupRepr: """Tests for AttackTechniqueSeedGroup.__repr__ method.""" diff --git a/tests/unit/models/test_seed_group.py b/tests/unit/models/test_seed_group.py index e207fd168f..bcb0c754fb 100644 --- a/tests/unit/models/test_seed_group.py +++ b/tests/unit/models/test_seed_group.py @@ -647,6 +647,25 @@ def test_system_prompt_technique_merges_onto_user_turn_at_sequence_zero(self): ("user", 3), ] + def test_system_prompt_technique_prepends_when_base_uses_negative_sequence(self): + """Explicit prepend placement must not reserve a sequence value in the base group.""" + base = AttackSeedGroup( + seeds=[ + SeedObjective(value="objective"), + SeedPrompt(value="opening user turn", data_type="text", role="user", sequence=-1), + SeedPrompt(value="assistant reply", data_type="text", role="assistant", sequence=4), + ] + ) + technique = AttackTechniqueSeedGroup.from_system_prompt("Follow these rules.") + + merged = base.with_technique(technique=technique) + + assert [(p.role, p.sequence) for p in merged.prompts] == [ + ("system", 0), + ("user", 1), + ("assistant", 2), + ] + def test_raises_when_technique_has_simulated_conversation_and_prompts_overlap(self): """Merging a technique with SeedSimulatedConversation into a group with overlapping prompts raises.""" base = AttackSeedGroup( diff --git a/tests/unit/scenario/airt/test_jailbreak.py b/tests/unit/scenario/airt/test_jailbreak.py index 080d4a91c5..5cce529015 100644 --- a/tests/unit/scenario/airt/test_jailbreak.py +++ b/tests/unit/scenario/airt/test_jailbreak.py @@ -525,6 +525,7 @@ async def test_system_delivery_attaches_system_role_framing_seed( seed_technique = system_attack.attack_technique.seed_technique assert seed_technique is not None assert [s.role for s in seed_technique.seeds] == ["system"] + assert seed_technique.prompt_placement == "prepend" assert seed_technique.seeds[0].value async def test_system_delivery_uses_no_jailbreak_converter( @@ -629,8 +630,8 @@ async def test_system_delivery_coexists_with_custom_user_prompt_seed_group(self) """A caller-supplied seed group carrying a user prompt at the default sequence 0 must not collide with the system framing seed when native system delivery merges in. - The framing seed is ordered at ``sequence=-1`` precisely so this merge succeeds; without it - the group would raise ``Inconsistent roles found for sequence 0`` at runtime. + The framing technique declares prepend placement so this merge does not depend on the + caller's sequence values. """ from pyrit.memory import CentralMemory from pyrit.score import SubStringScorer diff --git a/tests/unit/scenario/core/test_attack_technique.py b/tests/unit/scenario/core/test_attack_technique.py index ab33b34cc0..2117be1277 100644 --- a/tests/unit/scenario/core/test_attack_technique.py +++ b/tests/unit/scenario/core/test_attack_technique.py @@ -116,6 +116,38 @@ def test_technique_seeds_present_when_provided(self): assert "technique_seeds" in result.children assert len(result.children["technique_seeds"]) == 2 + def test_prompt_placement_is_part_of_identifier(self): + mock_attack = MagicMock(spec=AttackStrategy) + mock_attack.get_identifier.return_value = ComponentIdentifier( + class_name="PromptSendingAttack", class_module="pyrit.executor.attack" + ) + seed_technique = AttackTechniqueSeedGroup.from_system_prompt("Follow these rules.") + technique = AttackTechnique(attack=mock_attack, seed_technique=seed_technique) + + result = technique.get_identifier() + + assert result.params["prompt_placement"] == "prepend" + + def test_prompt_placement_changes_identifier_hash(self): + attack_id = ComponentIdentifier(class_name="PromptSendingAttack", class_module="pyrit.executor.attack") + preserve_attack = MagicMock(spec=AttackStrategy) + preserve_attack.get_identifier.return_value = attack_id + prepend_attack = MagicMock(spec=AttackStrategy) + prepend_attack.get_identifier.return_value = attack_id + seeds = [SeedPrompt(value="technique", data_type="text", is_general_technique=True)] + preserve = AttackTechnique( + attack=preserve_attack, + seed_technique=AttackTechniqueSeedGroup(seeds=seeds, prompt_placement="preserve"), + ) + prepend = AttackTechnique( + attack=prepend_attack, + seed_technique=AttackTechniqueSeedGroup(seeds=seeds, prompt_placement="prepend"), + ) + + assert "prompt_placement" not in preserve.get_identifier().params + assert prepend.get_identifier().params["prompt_placement"] == "prepend" + assert preserve.get_identifier().hash != prepend.get_identifier().hash + def test_identifier_is_cached(self): mock_attack = MagicMock(spec=AttackStrategy) mock_attack.get_identifier.return_value = ComponentIdentifier( diff --git a/tests/unit/setup/techniques/test_core_techniques.py b/tests/unit/setup/techniques/test_core_techniques.py index 20ad4bc30c..5c0ff650a8 100644 --- a/tests/unit/setup/techniques/test_core_techniques.py +++ b/tests/unit/setup/techniques/test_core_techniques.py @@ -42,8 +42,9 @@ def test_factory_shape(self): assert factory.seed_technique is not None seed = factory.seed_technique.seeds[0] assert seed.role == "system" - assert seed.sequence == -1 + assert seed.sequence == 0 assert seed.is_general_technique is True + assert factory.seed_technique.prompt_placement == "prepend" assert "flipping each word" in seed.value def test_merges_onto_group_with_user_turn_at_sequence_zero(self):