Skip to content

Commit 73d90a7

Browse files
FIX TAP and PAIR run their first turn after a prepended conversation (#2822)
Co-authored-by: hannahwestra25 <hannahwestra@microsoft.com>
1 parent 7d80d95 commit 73d90a7

3 files changed

Lines changed: 166 additions & 9 deletions

File tree

‎pyrit/executor/attack/multi_turn/tree_of_attacks.py‎

Lines changed: 10 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -470,6 +470,9 @@ def __init__(
470470
self._prepended_history_send_context: PrependedHistorySendContext | None = None
471471
# Context from prepended conversation (for adversarial chat system prompt)
472472
self._conversation_context: str | None = None
473+
# Number of prepended messages persisted in the objective target conversation. They are
474+
# history, not attack turns, so they do not count when deciding whether this is the first turn.
475+
self._prepended_message_count = 0
473476

474477
# Initial prompt for first turn (bypasses adversarial chat generation)
475478
# This supports multimodal messages
@@ -529,6 +532,7 @@ async def initialize_with_prepended_conversation_async(
529532
persisted_messages = list(
530533
self._memory.get_conversation_messages(conversation_id=self.objective_target_conversation_id)
531534
)
535+
self._prepended_message_count = len(persisted_messages)
532536
self._prepended_history_send_context = conversation_manager.create_prepended_history_send_context(
533537
target=self._objective_target,
534538
conversation_id=self.objective_target_conversation_id,
@@ -1021,6 +1025,7 @@ def duplicate(self) -> _TreeOfAttacksNode:
10211025

10221026
# Copy conversation context for adversarial chat system prompt
10231027
duplicate_node._conversation_context = self._conversation_context
1028+
duplicate_node._prepended_message_count = self._prepended_message_count
10241029
duplicate_node.last_response = copy.deepcopy(self.last_response)
10251030

10261031
# Copy visualization position so the clone starts from the same tree position
@@ -1210,14 +1215,15 @@ def _is_first_turn(self) -> bool:
12101215
Check if this is the first turn of the conversation.
12111216
12121217
This method determines whether the node is executing its initial attack turn by
1213-
examining the objective target conversation history.
1218+
examining the objective target conversation history. Prepended messages are history
1219+
rather than attack turns, so they are not counted.
12141220
12151221
Returns:
1216-
bool: True if no messages exist in the objective target conversation (first turn),
1217-
False if the conversation already contains messages (subsequent turns).
1222+
bool: True if the objective target conversation contains no messages beyond the
1223+
prepended conversation (first turn), False otherwise (subsequent turns).
12181224
"""
12191225
target_messages = self._memory.get_conversation_messages(conversation_id=self.objective_target_conversation_id)
1220-
return not target_messages
1226+
return len(target_messages) <= self._prepended_message_count
12211227

12221228
async def _generate_first_turn_prompt_async(self, objective: str) -> str:
12231229
"""
Lines changed: 151 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,151 @@
1+
# Copyright (c) Microsoft Corporation.
2+
# Licensed under the MIT license.
3+
4+
"""
5+
End-to-end coverage that TAP and PAIR still run their first attack turn when a prepended
6+
conversation is supplied. Prepended messages are persisted in the objective target
7+
conversation, but they are history rather than attack turns: the first live turn must still
8+
set the adversarial system prompt and honor ``next_message``.
9+
"""
10+
11+
import json
12+
13+
import pytest
14+
from unit.mocks import MockPromptTarget
15+
16+
from pyrit.executor.attack import (
17+
AttackAdversarialConfig,
18+
AttackScoringConfig,
19+
PAIRAttack,
20+
TreeOfAttacksWithPruningAttack,
21+
)
22+
from pyrit.models import AttackOutcome, Message, MessagePiece, Score
23+
from pyrit.score import FloatScaleThresholdScorer
24+
from pyrit.score.float_scale.float_scale_scorer import MessageFloatScaleScorer
25+
from pyrit.score.scorer_prompt_validator import ScorerPromptValidator
26+
27+
28+
class _AdversarialTarget(MockPromptTarget):
29+
"""Adversarial chat that records its system prompts and replies with schema-valid JSON."""
30+
31+
def __init__(self) -> None:
32+
super().__init__()
33+
self.system_prompts: list[str] = []
34+
35+
def set_system_prompt(self, *, system_prompt: str, conversation_id: str, **kwargs) -> None:
36+
self.system_prompts.append(system_prompt)
37+
super().set_system_prompt(system_prompt=system_prompt, conversation_id=conversation_id, **kwargs)
38+
39+
async def _send_prompt_to_target_async(self, *, normalized_conversation: list[Message]) -> list[Message]:
40+
message = normalized_conversation[-1]
41+
self.prompt_sent.append(message.get_value())
42+
reply = json.dumps(
43+
{
44+
"next_message": f"attack prompt {len(self.prompt_sent)}",
45+
"rationale": "r",
46+
"last_response_summary": "s",
47+
}
48+
)
49+
return [
50+
MessagePiece(
51+
role="assistant", original_value=reply, conversation_id=message.message_pieces[0].conversation_id
52+
).to_message()
53+
]
54+
55+
56+
class _FixedFloatScorer(MessageFloatScaleScorer):
57+
def __init__(self, *, value: str) -> None:
58+
super().__init__(validator=ScorerPromptValidator(supported_data_types=["text"]))
59+
self._value = value
60+
61+
def _build_identifier(self):
62+
return self._create_identifier()
63+
64+
async def _score_piece_async(self, message_piece: MessagePiece, *, objective: str | None = None) -> list[Score]:
65+
return [
66+
Score(
67+
score_value=self._value,
68+
score_type="float_scale",
69+
score_category=["test"],
70+
score_rationale="rationale",
71+
score_metadata=None,
72+
message_piece_id=message_piece.id,
73+
score_value_description="description",
74+
scorer_class_identifier=self.get_identifier(),
75+
)
76+
]
77+
78+
79+
def _build_attack(attack_cls, *, objective_target, adversarial_chat, score_value: str):
80+
scorer = FloatScaleThresholdScorer(scorer=_FixedFloatScorer(value=score_value), threshold=0.7)
81+
kwargs = (
82+
{"branching_factor": 2, "on_topic_checking_enabled": False}
83+
if attack_cls is TreeOfAttacksWithPruningAttack
84+
else {}
85+
)
86+
return attack_cls(
87+
objective_target=objective_target,
88+
attack_adversarial_config=AttackAdversarialConfig(target=adversarial_chat),
89+
attack_scoring_config=AttackScoringConfig(objective_scorer=scorer),
90+
tree_width=2,
91+
tree_depth=3,
92+
**kwargs,
93+
)
94+
95+
96+
@pytest.mark.parametrize("attack_cls", [TreeOfAttacksWithPruningAttack, PAIRAttack])
97+
async def test_system_only_prepended_conversation_still_attacks(attack_cls, sqlite_instance):
98+
objective_target = MockPromptTarget()
99+
adversarial_chat = _AdversarialTarget()
100+
attack = _build_attack(
101+
attack_cls, objective_target=objective_target, adversarial_chat=adversarial_chat, score_value="0.9"
102+
)
103+
104+
result = await attack.execute_async(
105+
objective="do X", prepended_conversation=[Message.from_system_prompt("You are a helpful assistant.")]
106+
)
107+
108+
assert objective_target.prompt_sent, "no prompt reached the objective target"
109+
assert adversarial_chat.system_prompts, "the adversarial system prompt was never set"
110+
assert result.outcome == AttackOutcome.SUCCESS
111+
112+
113+
@pytest.mark.parametrize("attack_cls", [TreeOfAttacksWithPruningAttack, PAIRAttack])
114+
async def test_prepended_conversation_sets_adversarial_system_prompt_with_context(attack_cls, sqlite_instance):
115+
objective_target = MockPromptTarget()
116+
adversarial_chat = _AdversarialTarget()
117+
attack = _build_attack(
118+
attack_cls, objective_target=objective_target, adversarial_chat=adversarial_chat, score_value="0.1"
119+
)
120+
prepended = [
121+
Message.from_prompt(prompt="PRIOR-USER-TURN", role="user"),
122+
Message.from_prompt(prompt="PRIOR-ASSISTANT-TURN", role="assistant"),
123+
]
124+
125+
await attack.execute_async(objective="do X", prepended_conversation=prepended)
126+
127+
assert adversarial_chat.system_prompts
128+
assert all("PRIOR-ASSISTANT-TURN" in system_prompt for system_prompt in adversarial_chat.system_prompts)
129+
130+
131+
async def test_prepended_conversation_sends_next_message_first(sqlite_instance):
132+
objective_target = MockPromptTarget()
133+
adversarial_chat = _AdversarialTarget()
134+
attack = _build_attack(
135+
TreeOfAttacksWithPruningAttack,
136+
objective_target=objective_target,
137+
adversarial_chat=adversarial_chat,
138+
score_value="0.1",
139+
)
140+
prepended = [
141+
Message.from_prompt(prompt="PRIOR-USER-TURN", role="user"),
142+
Message.from_prompt(prompt="PRIOR-ASSISTANT-TURN", role="assistant"),
143+
]
144+
145+
await attack.execute_async(
146+
objective="do X",
147+
prepended_conversation=prepended,
148+
next_message=Message.from_prompt(prompt="MY-CUSTOM-FIRST-PROMPT", role="user"),
149+
)
150+
151+
assert objective_target.prompt_sent[0] == "MY-CUSTOM-FIRST-PROMPT"

‎tests/unit/executor/attack/test_attack_parameter_consistency.py‎

Lines changed: 5 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -946,13 +946,13 @@ async def test_tap_attack_adds_prepended_to_memory(
946946
next_message=multimodal_text_message, # Required when prepended_conversation is provided
947947
)
948948

949-
# TAP prunes all branches with these mocks, so result.conversation_id is empty. The prepended
950-
# messages were duplicated into the single node conversation; resolve that id from memory.
951-
assert not result.conversation_id
949+
# next_message is sent on the first live turn and scored as a success, so the result points at
950+
# the single node conversation that holds the prepended messages.
951+
assert result.conversation_id
952952
memory = CentralMemory.get_memory_instance()
953953
node_conversation_ids = {piece.conversation_id for piece in memory.get_message_pieces()}
954-
assert len(node_conversation_ids) == 1, f"Expected one conversation in memory, got {node_conversation_ids}"
955-
conversation = list(memory.get_conversation_messages(conversation_id=node_conversation_ids.pop()))
954+
assert node_conversation_ids == {result.conversation_id}
955+
conversation = list(memory.get_conversation_messages(conversation_id=result.conversation_id))
956956

957957
# Should have exactly the prepended messages in memory (mock normalizer doesn't add responses)
958958
assert len(conversation) == 2, f"Expected exactly 2 prepended messages, got {len(conversation)}"

0 commit comments

Comments
 (0)