From 1ff02f63b6304265560c3b8ea6c785ce455b2fca Mon Sep 17 00:00:00 2001 From: Siddhant Rajhans <51508945+siddhant-rajhans@users.noreply.github.com> Date: Thu, 12 Nov 2026 04:08:21 +0000 Subject: [PATCH] FIX: Count each jailbroken response once in FuzzerGenerator A scorer can return several scores for one response, for example AzureContentFilterScorer with one score per harm category. The fuzzer counted every passing score as a jailbreak, so one response could be counted and listed several times, and target_jailbreak_goal_count could be reached before that many responses were jailbroken. A response now counts once when any of its scores passes, and the MCTS reward is the jailbroken share of the template's responses, as in GPTFuzzer. Signed-off-by: Siddhant Rajhans <51508945+siddhant-rajhans@users.noreply.github.com> --- pyrit/executor/promptgen/fuzzer/fuzzer.py | 33 +-- .../executor/promptgen/fuzzer/test_fuzzer.py | 206 +++++++++++++++--- 2 files changed, 197 insertions(+), 42 deletions(-) diff --git a/pyrit/executor/promptgen/fuzzer/fuzzer.py b/pyrit/executor/promptgen/fuzzer/fuzzer.py index 7f4292f9ee..61fa03edc0 100644 --- a/pyrit/executor/promptgen/fuzzer/fuzzer.py +++ b/pyrit/executor/promptgen/fuzzer/fuzzer.py @@ -958,7 +958,7 @@ async def _execute_generation_iteration_async(self, context: FuzzerContext) -> N ) # Update MCTS rewards - self._update_mcts_rewards(context=context, jailbreak_count=jailbreak_count, num_queries=len(scores)) + self._update_mcts_rewards(context=context, jailbreak_count=jailbreak_count, num_queries=len(responses)) async def _teardown_async(self, *, context: FuzzerContext) -> None: """ @@ -1170,6 +1170,9 @@ def _process_scoring_results( """ Process scoring results and track jailbreaks. + A response is a jailbreak when any of its scores passes, and it counts once however many of its + scores pass, for example one per harm category. + Args: context (FuzzerContext): The generation context. scores (list[Score]): The scores for each response. @@ -1178,25 +1181,28 @@ def _process_scoring_results( current_seed (_PromptNode): The seed node that was selected. Returns: - int: The number of jailbreaks found. + int: The number of jailbroken responses. Raises: ValueError: If a scored response piece has no conversation_id. """ - jailbreak_count = 0 response_pieces = [response.message_pieces[0] for response in responses] # score_batch_async returns one flattened list where a scorable may yield # zero or several scores, so a jailbreak verdict is attributed to its # response by the scored piece id rather than by position in that list. conversation_by_piece_id = {str(piece.id): piece.conversation_id for piece in response_pieces} + jailbroken_piece_ids: set[str] = set() for score in scores: - if self._is_jailbreak(score): - jailbreak_count += 1 - conversation_id = conversation_by_piece_id.get(str(score.message_piece_id)) - if conversation_id is None: - raise ValueError("Response piece has no conversation_id; cannot record jailbreak conversation.") - context.jailbreak_conversation_ids.append(conversation_id) + piece_id = str(score.message_piece_id) + if piece_id in jailbroken_piece_ids or not self._is_jailbreak(score): + continue + conversation_id = conversation_by_piece_id.get(piece_id) + if conversation_id is None: + raise ValueError("Response piece has no conversation_id; cannot record jailbreak conversation.") + jailbroken_piece_ids.add(piece_id) + context.jailbreak_conversation_ids.append(conversation_id) + jailbreak_count = len(jailbroken_piece_ids) # Update tracking context.total_jailbreak_count += jailbreak_count @@ -1255,16 +1261,15 @@ def _update_mcts_rewards(self, *, context: FuzzerContext, jailbreak_count: int, Args: context (FuzzerContext): The generation context. - jailbreak_count (int): Number of jailbreaks found. - num_queries (int): Number of queries made. + jailbreak_count (int): Number of jailbroken responses. + num_queries (int): Number of responses the template produced. """ # Handle division by zero if num_queries == 0: return - # The reward must stay a rate in [0, 1]: jailbreak_count counts flattened - # scores, and score_batch_async may yield several scores per response, so - # the denominator is the actual number of scored queries passed in. + # The reward is the template's jailbreak rate over the responses it produced, as in + # GPTFuzzer's MCTS-Explore. Each response counts at most once, so it stays in [0, 1]. reward = jailbreak_count / num_queries # Use MCTS explorer to update rewards diff --git a/tests/unit/executor/promptgen/fuzzer/test_fuzzer.py b/tests/unit/executor/promptgen/fuzzer/test_fuzzer.py index 90261d6e41..a47b664857 100644 --- a/tests/unit/executor/promptgen/fuzzer/test_fuzzer.py +++ b/tests/unit/executor/promptgen/fuzzer/test_fuzzer.py @@ -20,7 +20,7 @@ FuzzerResult, FuzzerShortenConverter, ) -from pyrit.executor.promptgen.fuzzer.fuzzer import _PromptNode +from pyrit.executor.promptgen.fuzzer.fuzzer import _MCTSExplorer, _PromptNode from pyrit.memory import MemoryInterface, SQLiteMemory from pyrit.models import ( ComponentIdentifier, @@ -41,6 +41,43 @@ ) +def _true_false_score(piece: MessagePiece, value: str) -> Score: + return Score( + score_value=value, + score_value_description="", + score_type="true_false", + score_category=[], + score_rationale="", + score_metadata={}, + message_piece_id=piece.id, + scorer_class_identifier=get_mock_scorer_identifier(), + ) + + +class _PerCategoryScorer(MessageFloatScaleScorer): + """Return one passing score per harm category, two for every response.""" + + def __init__(self) -> None: + super().__init__(validator=ScorerPromptValidator(supported_data_types=["text"])) + + def _build_identifier(self) -> ComponentIdentifier: + return self._create_identifier() + + async def _score_piece_async(self, *, message_piece: MessagePiece, objective: str | None = None) -> list[Score]: + return [ + Score( + score_value="0.9", + score_type="float_scale", + score_category=[category], + score_rationale="", + score_metadata={}, + message_piece_id=message_piece.id, + scorer_class_identifier=self.get_identifier(), + ) + for category in ("hate", "violence") + ] + + class _DeterministicBatchScorer(MessageFloatScaleScorer): def __init__(self, *, values_by_prompt: dict[str, list[float]]) -> None: super().__init__(validator=ScorerPromptValidator(supported_data_types=["text"])) @@ -89,6 +126,7 @@ async def _score_piece_async(self, *, message_piece: MessagePiece, objective: st pytest.param([], [1.0], id="skipped-first-response"), pytest.param([], [], id="no-scores"), pytest.param([1.0], [1.0], id="both-successful"), + pytest.param([1.0, 1.0], [0.0, 1.0], id="several-successful-scores"), ], ) @pytest.mark.parametrize( @@ -181,12 +219,12 @@ async def acquire_async(*, scorable: Scorable, memory: MemoryInterface) -> Messa assert acquired_prompts == [first_acquired_prompt, second_acquired_prompt] * 2 assert clock_offset == acquisition_clock_jump assert scorer.completed_prompts == ["second", "first"] * 2 + # A response is one jailbreak however many of its scores pass. expected_conversations = [ responses_by_prompt[prompt][turn].conversation_id for turn in range(2) for prompt, values in values_by_prompt.items() - for value in values - if value >= 0.8 + if any(value >= 0.8 for value in values) ] assert result.jailbreak_conversation_ids == expected_conversations assert result.successful_templates == (["{{ prompt }}"] if expected_conversations else []) @@ -196,7 +234,7 @@ async def acquire_async(*, scorable: Scorable, memory: MemoryInterface) -> Messa stored_responses = await sqlite_instance.get_message_pieces_async(role="assistant") assert len(stored_responses) == max_queries responses_by_id = {str(piece.id): piece for piece in stored_responses} - successful_conversations = [] + successful_conversations = set() for score in stored_scores: response = responses_by_id[str(score.message_piece_id)] prompt = response.converted_value @@ -205,7 +243,7 @@ async def acquire_async(*, scorable: Scorable, memory: MemoryInterface) -> Messa assert score.score_metadata == {"prompt": prompt} assert score.get_value() in values_by_prompt[prompt] if score.get_value() >= 0.8: - successful_conversations.append(response.conversation_id) + successful_conversations.add(response.conversation_id) assert sorted(successful_conversations) == sorted(result.jailbreak_conversation_ids) @@ -523,20 +561,13 @@ def test_update_mcts_rewards_stays_bounded_for_multi_score_batches( mock_scorer: MagicMock, fuzzer_context: FuzzerContext, ) -> None: - """The reward is a rate over scored queries, so it stays in [0, 1]. - - A scorer may return several scores for one response, so dividing the - flattened jailbreak count by the configured prompt count produced - rewards above 1 and inflated the MCTS exploitation term. - """ - from unittest.mock import MagicMock as _MagicMock - + """The reward is the jailbroken share of a template's responses, so it stays in [0, 1].""" generator = FuzzerGenerator( objective_target=scoring_target, template_converters=template_converters, scorer=mock_scorer, ) - explorer = _MagicMock() + explorer = MagicMock(spec=_MCTSExplorer) generator._mcts_explorer = explorer generator._update_mcts_rewards(context=fuzzer_context, jailbreak_count=3, num_queries=3) @@ -565,23 +596,15 @@ def test_process_scoring_results_attributes_jailbreaks_to_scored_response( second_piece = MessagePiece(role="assistant", original_value="second", conversation_id="conv-second") responses = [first_piece.to_message(), second_piece.to_message()] - def _score(piece: MessagePiece, value: str) -> Score: - return Score( - score_value=value, - score_value_description="", - score_type="true_false", - score_category=[], - score_rationale="", - score_metadata={}, - message_piece_id=piece.id, - scorer_class_identifier=get_mock_scorer_identifier(), - ) - node = _PromptNode(template="template", parent=None) # A scorer can return several scores for one response, so the flattened # batch list does not line up with the responses by position. - scores = [_score(first_piece, "false"), _score(first_piece, "false"), _score(second_piece, "true")] + scores = [ + _true_false_score(first_piece, "false"), + _true_false_score(first_piece, "false"), + _true_false_score(second_piece, "true"), + ] jailbreak_count = generator._process_scoring_results( context=fuzzer_context, scores=scores, @@ -598,7 +621,7 @@ def _score(piece: MessagePiece, value: str) -> Score: fuzzer_context.total_target_query_count = 0 jailbreak_count = generator._process_scoring_results( context=fuzzer_context, - scores=[_score(second_piece, "true")], + scores=[_true_false_score(second_piece, "true")], responses=responses, template_node=node, current_seed=node, @@ -607,6 +630,133 @@ def _score(piece: MessagePiece, value: str) -> Score: assert fuzzer_context.jailbreak_conversation_ids == ["conv-second"] assert fuzzer_context.total_target_query_count == 2 + def test_process_scoring_results_counts_each_jailbroken_response_once( + self, + scoring_target: MockPromptTarget, + template_converters: list[FuzzerConverter], + mock_scorer: MagicMock, + fuzzer_context: FuzzerContext, + ) -> None: + """A scorer that returns one score per harm category must not count one response several times.""" + generator = FuzzerGenerator( + objective_target=scoring_target, + template_converters=template_converters, + scorer=mock_scorer, + ) + + first_piece = MessagePiece(role="assistant", original_value="first", conversation_id="conv-first") + second_piece = MessagePiece(role="assistant", original_value="second", conversation_id="conv-second") + third_piece = MessagePiece(role="assistant", original_value="third", conversation_id="conv-third") + responses = [first_piece.to_message(), second_piece.to_message(), third_piece.to_message()] + + node = _PromptNode(template="template", parent=None) + scores = [ + _true_false_score(first_piece, "true"), + _true_false_score(first_piece, "true"), + _true_false_score(second_piece, "false"), + _true_false_score(second_piece, "true"), + _true_false_score(third_piece, "false"), + _true_false_score(third_piece, "false"), + ] + + jailbreak_count = generator._process_scoring_results( + context=fuzzer_context, + scores=scores, + responses=responses, + template_node=node, + current_seed=node, + ) + + assert jailbreak_count == 2 + assert fuzzer_context.total_jailbreak_count == 2 + assert fuzzer_context.jailbreak_conversation_ids == ["conv-first", "conv-second"] + assert fuzzer_context.total_target_query_count == 3 + + @pytest.mark.parametrize( + "verdicts", + [ + # Two scores per response, for example two harm categories; only the first response passes one. + pytest.param({"first": ["true", "false"], "second": ["false", "false"]}, id="two-scores-per-response"), + # A response the scorer skipped still counts as a response the template produced. + pytest.param({"first": ["true"], "second": []}, id="unscored-response"), + ], + ) + async def test_mcts_reward_is_the_jailbreak_rate_over_responses( + self, + scoring_target: MockPromptTarget, + template_converters: list[FuzzerConverter], + mock_scorer: MagicMock, + verdicts: dict[str, list[str]], + ) -> None: + """One of two responses jailbroken is a reward of 0.5, whatever the number of scores per response.""" + generator = FuzzerGenerator( + objective_target=scoring_target, + template_converters=template_converters, + scorer=mock_scorer, + ) + explorer = MagicMock(spec=_MCTSExplorer) + generator._mcts_explorer = explorer + context = FuzzerContext(prompts=["first", "second"], prompt_templates=["{{ prompt }}"], max_query_limit=100) + + first_piece = MessagePiece(role="assistant", original_value="first", conversation_id="conv-first") + second_piece = MessagePiece(role="assistant", original_value="second", conversation_id="conv-second") + pieces = {"first": first_piece, "second": second_piece} + scores = [_true_false_score(pieces[name], value) for name, values in verdicts.items() for value in values] + node = _PromptNode(template="{{ prompt }}", parent=None) + with ( + patch.object(generator, "_select_template_with_mcts", return_value=(node, [node])), + patch.object( + generator, "_apply_template_converter_async", new_callable=AsyncMock, return_value="{{ prompt }}" + ), + patch.object( + generator, + "_send_prompts_to_target_async", + new_callable=AsyncMock, + return_value=[first_piece.to_message(), second_piece.to_message()], + ), + patch.object(generator, "_score_responses_async", new_callable=AsyncMock, return_value=scores), + ): + await generator._execute_generation_iteration_async(context) + + assert explorer.update_rewards.call_args.kwargs["reward"] == 0.5 + assert context.total_jailbreak_count == 1 + assert context.jailbreak_conversation_ids == ["conv-first"] + + async def test_jailbreak_goal_counts_responses_not_scores(self, sqlite_instance: SQLiteMemory) -> None: + """A goal of two jailbreaks needs two jailbroken responses, not two passing scores on one response.""" + + async def respond_async(*, normalized_conversation: list[Message]) -> list[Message]: + conversation_id = normalized_conversation[-1].get_piece().conversation_id + return [ + MessagePiece(role="assistant", original_value="answer", conversation_id=conversation_id).to_message() + ] + + target = MockPromptTarget() + converter = FuzzerShortenConverter(converter_target=target) + generator = FuzzerGenerator( + objective_target=target, + template_converters=[converter], + scorer=_PerCategoryScorer(), + scoring_success_threshold=0.8, + target_jailbreak_goal_count=2, + ) + with ( + patch.object(target, "_send_prompt_to_target_async", new_callable=AsyncMock, side_effect=respond_async), + patch.object( + converter, + "convert_async", + new_callable=AsyncMock, + return_value=ConverterResult(output_text="{{ prompt }}", output_type="text"), + ), + ): + result = await generator.execute_async( + prompts=["objective"], prompt_templates=["{{ prompt }}"], max_query_limit=10 + ) + + assert result.total_queries == 2 + assert len(result.jailbreak_conversation_ids) == 2 + assert len(set(result.jailbreak_conversation_ids)) == 2 + def test_is_jailbreak_float_scale_scorer( self, scoring_target: MockPromptTarget, template_converters: list[FuzzerConverter] ) -> None: