FIX: Count each jailbroken response once in FuzzerGenerator

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

tests/unit/executor/promptgen/fuzzer/test_fuzzer.py

@@ -89,6 +89,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 +182,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 +197,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 +206,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)
 
 
@@ -607,6 +608,114 @@ 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()]
+
+        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)
+        scores = [
+            _score(first_piece, "true"),
+            _score(first_piece, "true"),
+            _score(second_piece, "false"),
+            _score(second_piece, "true"),
+            _score(third_piece, "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
+
+    async def test_mcts_reward_is_the_jailbreak_rate_over_responses(
+        self,
+        scoring_target: MockPromptTarget,
+        template_converters: list[FuzzerConverter],
+        mock_scorer: MagicMock,
+    ) -> 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()
+        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")
+
+        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(),
+            )
+
+        # Two scores per response, for example two harm categories; only the first response passes one.
+        scores = [_score(first_piece, "true"), _score(first_piece, "false")]
+        scores += [_score(second_piece, "false"), _score(second_piece, "false")]
+        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"]
+
     def test_is_jailbreak_float_scale_scorer(
         self, scoring_target: MockPromptTarget, template_converters: list[FuzzerConverter]
     ) -> None: