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.

[romanlutz]

Thanks for your patience with this one! I'm busy with the v1 release until EOM but hoping to get this in as soon as v1 is out (~mid-week). Some of these comments are questions/concerns that came up from reviewing with GHCP. Nothing is set in stone, of course, and I'm happy to be convinced otherwise.

[romanlutz]

The final candidate is selected only by auxiliary score, and then only that beam's objective score determines the outcome. A lower-ranked beam can have achieved the objective while this returns FAILURE and marks that successful conversation as pruned. Please prefer an objective-success beam (and ideally stop once one succeeds), using auxiliary score only to rank ties or unsuccessful candidates.

[riedgar-ms]

The problem I see is that the objective_scorer is marked as a TrueFalseScorer. To do the hill climbing implicit in beam search, you need a floating point output, so you can do things more gradually, hence my use of the auxiliary scorers.

[romanlutz]

This attack uses scoring but does not override get_attack_scoring_config, and its identifier omits the beam count, iteration count, continuation size, and reviewer policy. Substantially different beam-search configurations therefore produce the same identifier. Please include the scoring configuration and behavior-defining beam parameters so persisted results remain attributable and reproducible.

[riedgar-ms]

Addressed in f927b6de. BeamSearchAttack now exposes its scoring configuration and includes the beam count, iteration count, continuation size, and an identifiable reviewer policy in its identifier. Tests verify that changing the search or reviewer configuration changes the identifier.

[romanlutz]

PromptNormalizer has already applied response converters here. Feeding converted_value back as the next raw grammar prefix means Base64, ROT13, translation, and similar converters constrain the model to continue the transformed representation rather than its original output. Please build the continuation prefix from original_value, or reject response converters for this attack.

[riedgar-ms]

Addressed in f927b6de. Beam continuation prefixes are now built from each response piece's original_value, while the converted response remains available for scoring and persistence. I also added a regression test covering differing original and converted values.

[romanlutz]

Auxiliary scorers are defined as additional metrics/custom evaluations; adding one should not change attack control flow. Here their values are summed and used as the beam-search policy, so an observability scorer changes which candidates survive and which result is returned. Arbitrary float scorers can also have incompatible scales and semantics, making the unweighted sum ill-defined. Is it possible to follow the TAP pattern (?): introduce a beam-search scoring config whose objective scorer is a FloatScaleThresholdScorer, use its underlying float value to rank partial beams and its thresholded result to determine success, and keep auxiliary scores telemetry-only.

[riedgar-ms]

That might be the solution to the issue I pointed out above.

[Copilot]

Review details

Comments suppressed due to low confidence (2)

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

• _setup_async is invoked once by the strategy lifecycle, but it is also called from _propagate_beam_async for each beam. As written, it overwrites self._start_context every time, which can be clobbered by concurrent beam propagation and change the starting context mid-attack.

Only capture _start_context on the initial setup.

        self._start_context = copy.deepcopy(context)

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

• beam.id is updated before the model call succeeds. If the call fails, the exception handler leaves beam.response_message/beam.text from a prior iteration but with a new beam.id, which breaks traceability (conversation_id no longer matches the stored response) and can cause stale responses to be treated as current.

Update beam.id only after a successful response is obtained.

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

        try:
            with execution_context(

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

[riedgar-ms]

Addressed in f927b6de. Removed @pytest.mark.asyncio and now rely on the repository's automatic asyncio discovery.

[riedgar-ms]

Addressed in f927b6de. Updated the fresh_instance docstring to use dict[str, Any] | None and str | None.

[Copilot]

Review details

Comments suppressed due to low confidence (4)

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

• TopKBeamReviewer uses deepcopy() to create replacement beams, which preserves prior iteration state (score, objective_score, response_message). If propagation fails in a later iteration, these copied beams can be (re)scored using stale response_message/objective_score that no longer matches the (possibly truncated) text, producing incorrect beam selection.

        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:
                nxt.text = nxt.text[: -self.drop_chars]
            new_beams.append(nxt)

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

• _propagate_beam_async calls _setup_async for each beam while beams are propagated concurrently (asyncio.gather). _setup_async mutates shared instance state (self._start_context) and the passed context.conversation_id, which can cause race conditions and cross-talk between beams. Each beam should initialize its own context without re-running the strategy-level setup.

        if self._start_context is None:
            raise ValueError("Start context must be set before propagating beams")
        target = self._get_target_for_beam(beam)

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

pyrit/prompt_target/openai/openai_response_target.py:233

• fresh_instance docstring uses legacy Optional[...] type syntax in the Args section. The style guide requires modern X | None / built-in generics in docstrings and comments as well.

        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.

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

• Project test instructions explicitly discourage using @pytest.mark.asyncio because asyncio_mode = "auto" is configured repo-wide. This decorator should be removed to match the established test convention.

@pytest.mark.asyncio
async def test_build_input_for_multi_modal_async_filters_reasoning(target: OpenAIResponseTarget):

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

[riedgar-ms]

Addressed in f927b6de. A grammar_name override now updates the grammar tool itself before constructing the fresh target, keeping the tool name and _grammar_name synchronized. It also raises an explicit error unless exactly one grammar tool is present.

[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