|
| 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" |
0 commit comments