FEAT Beam Search for OpenAIResponseTarget#1346
Conversation
riedgar-ms
left a comment
riedgar-ms
left a comment
There was a problem hiding this comment.
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.
| **deepcopy(kwargs), | ||
| } | ||
|
|
||
| def fresh_instance(self) -> "OpenAIResponseTarget": |
There was a problem hiding this comment.
These changes are required because the OpenAI API takes a grammar as a tool and PyRIT makes the tool list part of the object, not the send_prompt_async() API. Since we have multiple beams being managed asynchronously, each task needs its own copy of the OpenAIResponseTarget
| return new_beams | ||
|
|
||
|
|
||
| class BeamSearchAttack(SingleTurnAttackStrategy): |
There was a problem hiding this comment.
This is largely copied from the PromptSendingAttack
| target = self._get_target_for_beam(beam) | ||
|
|
||
| current_context = copy.deepcopy(self._start_context) | ||
| await self._setup_async(context=current_context) |
There was a problem hiding this comment.
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.
| objective=context.objective, | ||
| ) | ||
|
|
||
| aux_scores = scoring_results["auxiliary_scores"] |
There was a problem hiding this comment.
So the auxilliary scorer is required; it is used to assess the beams as they develop
|
|
||
| new_beams = list(reversed(sorted_beams[: self.k])) | ||
| for i in range(len(beams) - len(new_beams)): | ||
| nxt = copy.deepcopy(new_beams[i % self.k]) |
There was a problem hiding this comment.
We don't just duplicate the highest scoring beam, in the hope of maintaining some variety
| Args: | ||
| context (SingleTurnAttackContext): The attack context containing attack parameters. | ||
| """ | ||
| self._start_context = copy.deepcopy(context) |
There was a problem hiding this comment.
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.
| @@ -0,0 +1,110 @@ | |||
| # --- | |||
There was a problem hiding this comment.
This is a conversion of the (to be deleted) file beam_search_test.py in the repo root. I can't run this, since our endpoints forbid key-based auth.
riedgar-ms
changed the title
|
Ready for review, but I will need help running the notebook prior to merge. |
riedgar-ms
changed the title
romanlutz
left a comment
romanlutz
left a comment
There was a problem hiding this comment.
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.
| raise ValueError("Start context must be set before propagating beams") | ||
| target = self._get_target_for_beam(beam) | ||
|
|
||
| current_context = copy.deepcopy(self._start_context) |
There was a problem hiding this comment.
Copilot found a potential concurrency bug that I am also worried about. See below...
AttackExecutor can run this strategy instance concurrently for multiple objectives, but _setup_async stores each execution's context in self._start_context. A beam can therefore copy whichever objective wrote this shared field last and persist the wrong response under another objective. Please keep the base context on the per-execution context or pass it into _propagate_beam_async; execution state should not live on the reusable strategy instance.
| "tool_choice": "required", | ||
| } | ||
|
|
||
| return self._objective_target.fresh_instance(extra_body_parameters=ebp, grammar_name=str(grammar_tool["name"])) |
There was a problem hiding this comment.
Passing extra_body_parameters here replaces all settings already configured on the target. That silently drops metadata, routing controls, and store=False; for the OpenAI Responses API, omitting store=False restores the default 30-day application-state retention. Please preserve the caller's body parameters and override only the beam-specific reasoning, tools, and tool_choice fields.
| self._logger.debug(f"Beam {i} score: {beam.score}") | ||
|
|
||
| # Sort the list of beams | ||
| beams = sorted(beams, key=lambda b: b.score, reverse=True) |
There was a problem hiding this comment.
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.
| piece.converted_value for piece in assistant_pieces if isinstance(piece.converted_value, str) | ||
| ) | ||
| beam.response_message = model_response | ||
| except Exception as e: |
There was a problem hiding this comment.
If propagation fails, this leaves text, response_message, score, and objective_score from the prior iteration intact. The failed beam remains scoreable and can win or report success using a stale response. Please clear per-iteration state before the call and remove failed beams; unexpected target or configuration errors should propagate instead of being converted into an ordinary failed attempt.
| raise ValueError("BeamSearchAttack requires all auxiliary scorers to be instances of FloatScaleScorer") | ||
|
|
||
| self._auxiliary_scorers = attack_scoring_config.auxiliary_scorers | ||
| self._objective_scorer = attack_scoring_config.objective_scorer |
There was a problem hiding this comment.
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.
| for piece in model_response.message_pieces | ||
| if isinstance(getattr(piece, "converted_value", None), str) | ||
| ] | ||
| beam.text = "".join( |
There was a problem hiding this comment.
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.
| objective=context.objective, | ||
| ) | ||
|
|
||
| aux_scores = scoring_results["auxiliary_scores"] |
There was a problem hiding this comment.
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.
Description
Use the Lark grammar feature of the
OpenAIResponseTargetto 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
OpenAIResponseTargetthere didn't seem much point in mocking that. There is a notebook which runs everything E2E.