FEAT Beam Search for OpenAIResponseTarget

[riedgar-ms]

Description

Use the Lark grammar feature of the OpenAIResponseTarget to create a beam search for PyRIT. This is a single turn attack, where a collection of candidate responses (the beams) are maintained. On each iteration, the model's response is allowed to extend a little for each beam. The beams are scored, with the worst performing ones discarded, and replaced with copies of higher scoring beams.

Tests and Documentation

Have basic unit tests of the classes added, but since this requires features only currently in the OpenAIResponseTarget there didn't seem much point in mocking that. There is a notebook which runs everything E2E.

[riedgar-ms]

This is ready for preliminary review; there aren't any docs or (proper) tests yet. I'd like to make sure that I'm manipulating the database correctly before delving into those.

[riedgar-ms]

I'm not certain I'm handling the context correctly here. I end up making lots of copies of things, which is going to be filling up the database with fragmentary responses. Each time one is extended, it ends up being cloned and a new conversation started.

[riedgar-ms]

See note below. I duplicate the context and the message for each beam on each iteration. I'm not certain that this is the best way to use the database.

[riedgar-ms]

Ready for review, but I will need help running the notebook prior to merge.

[Copilot]

Pull request overview

This pull request implements a beam search attack strategy for PyRIT that leverages the Lark grammar feature of OpenAIResponseTarget. The attack maintains multiple candidate responses (beams) and iteratively extends them character-by-character, scoring and pruning beams at each step to focus on the most promising candidates. This is designed as a single-turn attack where multiple model calls extend the same conversation turn.

Changes:

• Adds BeamSearchAttack class implementing a novel beam search strategy for prompt attacks
• Adds fresh_instance() method to OpenAIResponseTarget to support creating duplicate target instances with modified parameters
• Includes comprehensive unit tests for initialization validation and core components
• Provides E2E documentation with Jupyter notebook examples

Reviewed changes

Copilot reviewed 7 out of 7 changed files in this pull request and generated 25 comments.

Show a summary per file

File	Description
pyrit/executor/attack/single_turn/beam_search.py	Core implementation of beam search attack with Beam, BeamReviewer, TopKBeamReviewer, and BeamSearchAttack classes
pyrit/prompt_target/openai/openai_response_target.py	Adds fresh_instance method to enable creating duplicate targets with same configuration
tests/unit/executor/attack/single_turn/test_beam_search.py	Unit tests for initialization, validation, Beam grammar generation, and TopKBeamReviewer logic
doc/code/executor/attack/beam_search_attack.py	Documentation notebook demonstrating usage of BeamSearchAttack with examples
doc/code/executor/attack/beam_search_attack.ipynb	Jupyter notebook version of the documentation
doc/_toc.yml	Updates table of contents to include new beam search attack documentation
beam_search_test.py	E2E integration test script (located at repository root)

[Copilot]

Pull request overview

Copilot reviewed 7 out of 7 changed files in this pull request and generated 13 comments.

[Copilot]

BeamSearchAttack introduces substantial new control flow (beam propagation/scoring loop), but the unit tests currently only cover init validation and small helpers. Add unit tests that mock PromptNormalizer.send_prompt_async and Scorer.score_response_async to validate iteration behavior (beam cloning/pruning, failure handling, and final result selection) without requiring real OpenAIResponseTarget E2E calls.

[riedgar-ms]

Similar to my comments elsewhere (and in the PR description I believe).... since OpenAIResponseTarget is the only target currently supported, I fear a mock would add a maintenance headache. I am happy to rewrite the sample notebook into an integration test as well, if that would be desired.

[Copilot]

Review details

Comments suppressed due to low confidence (4)

tests/unit/prompt_target/target/test_openai_response_target.py:721

• @pytest.mark.asyncio should not be used in this repo because asyncio_mode = "auto" is configured globally; the decorator is redundant and conflicts with the testing guidelines. Remove the marker and keep the test as a plain async def.

@pytest.mark.asyncio

pyrit/prompt_target/openai/openai_response_target.py:233

• The fresh_instance docstring uses legacy Optional[...] type syntax, but this repo’s style guide requires modern X | None types in docstrings/comments as well. Update these arg type annotations to match the signature.

        Args:
            extra_body_parameters (Optional[dict[str, Any]]): Optional overrides for the
                extra body parameters of the new instance.
            grammar_name (Optional[str]): Optional override for the grammar name of the
                new instance.

pyrit/executor/attack/single_turn/beam_search.py:374

• _propagate_beam_async updates beam.id before the model call but does not clear beam.response_message / beam.score on failure. If a later propagation fails, the beam can retain stale response/score from a previous iteration while pointing at a new conversation_id, which breaks traceability and can cause the sorter to pick a beam with no valid response. Reset per-beam state before attempting propagation.

        message = self._get_message(current_context)
        beam.id = current_context.conversation_id

pyrit/executor/attack/single_turn/init.py:6

• This import line is likely to exceed the project’s line-length limit and will trip format/lint checks. Split it into a parenthesized multi-line import to keep the module ruff/isort compliant.

from pyrit.executor.attack.single_turn.beam_search import Beam, BeamReviewer, BeamSearchAttack, TopKBeamReviewer

• Files reviewed: 9/9 changed files
• Comments generated: 2
• Review effort level: Low

[riedgar-ms]

Addressed in f927b6de. Beam replication now uses the number of retained beams rather than k, preventing an IndexError when fewer than k beams are available. A boundary regression test was added.

[Copilot]

Review details

Comments suppressed due to low confidence (6)

pyrit/prompt_target/openai/openai_response_target.py:233

• The docstring uses legacy Optional[...] type syntax, but the project style guide requires modern union syntax in docstrings too (e.g., dict[str, Any] | None).

        Args:
            extra_body_parameters (Optional[dict[str, Any]]): Optional overrides for the
                extra body parameters of the new instance.
            grammar_name (Optional[str]): Optional override for the grammar name of the
                new instance.

pyrit/executor/attack/single_turn/beam_search.py:142

• TopKBeamReviewer.review can raise IndexError when k is larger than the number of beams (e.g., new_beams is shorter than self.k, but the code indexes new_beams[i % self.k]). This should degrade gracefully by clamping k to the available beam count.

        _extra_beam_count = self.desired_beam_count or len(beams)

        for i in range(_extra_beam_count - len(new_beams)):
            nxt = copy.deepcopy(new_beams[i % self.k])
            if self.drop_chars > 0 and len(nxt.text) > self.drop_chars:

pyrit/executor/attack/single_turn/beam_search.py:373

• _propagate_beam_async calls _setup_async, which mutates shared attack state (self._start_context) and is invoked concurrently via asyncio.gather. This creates a race condition and can cause beams to stomp each other's start context. Instead, initialize a per-beam context without modifying self._start_context.

        current_context = copy.deepcopy(self._start_context)
        await self._setup_async(context=current_context)

        message = self._get_message(current_context)
        beam.id = current_context.conversation_id

pyrit/executor/attack/single_turn/beam_search.py:104

• This docstring uses Optional[...] type syntax; the project style guide requires modern union syntax in docstrings (e.g., int | None).

            desired_beam_count (Optional[int]): The desired total number of beams after review.
                If None, it will be set to the supplied number of beams.

pyrit/executor/attack/single_turn/beam_search.py:182

• This docstring uses Optional[...] type syntax; the project style guide requires modern union syntax in docstrings (e.g., AttackConverterConfig | None).

            attack_converter_config (Optional[AttackConverterConfig]): Configuration for prompt converters.
            attack_scoring_config (Optional[AttackScoringConfig]): Configuration for scoring components.
            prompt_normalizer (Optional[PromptNormalizer]): The prompt normalizer to use.

pyrit/executor/attack/single_turn/beam_search.py:185

• This docstring uses Optional[...] type syntax; the project style guide requires modern union syntax in docstrings (e.g., PrependedConversationConfig | None).

            prepended_conversation_config (Optional[PrependedConversationConfig]): Configuration for prepended
                conversation.

• Files reviewed: 9/9 changed files
• Comments generated: 2
• Review effort level: Low