Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
61 changes: 36 additions & 25 deletions pyrit/models/seeds/attack_seed_group.py
Original file line number Diff line number Diff line change
Expand Up @@ -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):
Expand Down Expand Up @@ -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):
Expand All @@ -157,38 +160,46 @@ 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
# all of them with a single new UUID.
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)
28 changes: 18 additions & 10 deletions pyrit/models/seeds/attack_technique_seed_group.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -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:
"""
Expand All @@ -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.
Expand All @@ -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:
Expand Down
7 changes: 5 additions & 2 deletions pyrit/scenario/core/attack_technique.py
Original file line number Diff line number Diff line change
Expand Up @@ -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,
)
9 changes: 2 additions & 7 deletions pyrit/scenario/scenarios/airt/jailbreak.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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,
Expand Down
12 changes: 11 additions & 1 deletion tests/unit/models/test_attack_technique_seed_group.py
Original file line number Diff line number Diff line change
Expand Up @@ -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."""
Expand Down Expand Up @@ -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)."""
Expand All @@ -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."""
Expand Down
19 changes: 19 additions & 0 deletions tests/unit/models/test_seed_group.py
Original file line number Diff line number Diff line change
Expand Up @@ -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(
Expand Down
5 changes: 3 additions & 2 deletions tests/unit/scenario/airt/test_jailbreak.py
Original file line number Diff line number Diff line change
Expand Up @@ -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(
Expand Down Expand Up @@ -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
Expand Down
32 changes: 32 additions & 0 deletions tests/unit/scenario/core/test_attack_technique.py
Original file line number Diff line number Diff line change
Expand Up @@ -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(
Expand Down
3 changes: 2 additions & 1 deletion tests/unit/setup/techniques/test_core_techniques.py
Original file line number Diff line number Diff line change
Expand Up @@ -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):
Expand Down