From 818414b4b1c9b2a99b63cb1d05639e1b29e20861 Mon Sep 17 00:00:00 2001 From: Richard Lundeen Date: Tue, 14 Jul 2026 09:29:49 -0700 Subject: [PATCH 1/4] FIX: Fixing json retry JSON retries replayed poisoned conversation history because PromptNormalizer persists the request+response to memory before the caller validates the JSON. Add a memory-aware send_json_with_retry_async helper that rolls memory back to a baseline sequence between attempts, deleting the failed turn and recording a ConversationRetry marker, then rewire the adversarial manager and shared LLM scorer through it. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> Copilot-Session: 244845c9-0655-4fa8-9717-29d75faef9f2 --- .../adversarial_conversation_manager.py | 54 +++++------ ...c3e5d7f9b0_add_retries_to_conversations.py | 35 +++++++ pyrit/memory/memory_interface.py | 80 +++++++++++++++- pyrit/memory/memory_models.py | 13 ++- pyrit/models/__init__.py | 3 + pyrit/models/messages/conversation_retry.py | 34 +++++++ pyrit/models/messages/conversations.py | 18 ++-- pyrit/prompt_normalizer/__init__.py | 2 + pyrit/prompt_normalizer/json_retry.py | 96 +++++++++++++++++++ pyrit/prompt_normalizer/prompt_normalizer.py | 5 + pyrit/score/llm_scoring.py | 85 ++++++++++------ .../test_adversarial_conversation_manager.py | 4 +- .../attack/multi_turn/test_crescendo.py | 2 +- .../attack/multi_turn/test_red_teaming.py | 2 +- tests/unit/score/test_scorer.py | 51 ++++++---- 15 files changed, 397 insertions(+), 87 deletions(-) create mode 100644 pyrit/memory/alembic/versions/a1c3e5d7f9b0_add_retries_to_conversations.py create mode 100644 pyrit/models/messages/conversation_retry.py create mode 100644 pyrit/prompt_normalizer/json_retry.py diff --git a/pyrit/executor/attack/component/adversarial_conversation_manager.py b/pyrit/executor/attack/component/adversarial_conversation_manager.py index ba062fae66..5a617bc470 100644 --- a/pyrit/executor/attack/component/adversarial_conversation_manager.py +++ b/pyrit/executor/attack/component/adversarial_conversation_manager.py @@ -16,7 +16,6 @@ ComponentRole, InvalidJsonException, execution_context, - pyrit_json_retry, remove_markdown_json, ) from pyrit.executor.attack.core.attack_config import ( @@ -34,7 +33,7 @@ SeedPrompt, get_common_json_schema, ) -from pyrit.prompt_normalizer import PromptNormalizer +from pyrit.prompt_normalizer import PromptNormalizer, send_json_with_retry_async if TYPE_CHECKING: from pathlib import Path @@ -440,8 +439,8 @@ def __init__( max_turns: Maximum number of turns; rendered into the adversarial system prompt as ``max_turns``. Defaults to 1. raise_on_invalid_json: When True (default), a reply that fails to match the resolved - schema raises ``InvalidJsonException`` (retried via ``pyrit_json_retry``). When False, - the raw reply text is returned as ``next_message`` instead of raising. + schema raises ``InvalidJsonException`` (retried via ``send_json_with_retry_async``). + When False, the raw reply text is returned as ``next_message`` instead of raising. prompt_normalizer: The prompt normalizer to send through. Defaults to a new one. conversation_id: The adversarial-chat conversation id this manager drives. A fresh id is generated when None. @@ -803,7 +802,6 @@ def _build_objective_message( ) return Message.from_prompt(prompt=reply.next_message, role="user") - @pyrit_json_retry async def _send_and_parse_async( self, *, @@ -814,10 +812,11 @@ async def _send_and_parse_async( """ Send one user turn to the adversarial chat and parse its reply. - This is the single place adversarial-chat JSON retry lives: when the reply fails to match the - resolved schema, ``InvalidJsonException`` propagates and ``pyrit_json_retry`` re-sends the turn - until it parses or the attempt budget is exhausted. When ``raise_on_invalid_json`` is False, an - unparseable reply is returned as raw text instead. + This is the single place adversarial-chat JSON retry lives. It delegates to + ``send_json_with_retry_async`` so each retry sends on a clean conversation history: + the failed turn is rolled back out of memory before the turn is resent, instead of + replaying the model's own malformed reply. When ``raise_on_invalid_json`` is False, an + unparseable reply is returned as raw text (single attempt, no retry). When a ``modality_router`` is configured, the outgoing message is built via ``build_adversarial_input_message`` so first-turn seed media (``seed_message``) and prior @@ -852,6 +851,19 @@ async def _send_and_parse_async( prompt_metadata=prompt_metadata or None, ) + if self._memory_labels: + for piece in message.message_pieces: + piece.labels = self._memory_labels + + schema = self._response_json_schema + + def _parse(response: Message) -> AdversarialReply: + return _parse_adversarial_reply(response.get_value(), schema=schema) + + def _fallback(response: Message) -> AdversarialReply: + raw = response.get_value() + return AdversarialReply(next_message=raw, raw=raw) + with execution_context( component_role=ComponentRole.ADVERSARIAL_CHAT, attack_strategy_name=self._attack_strategy_name, @@ -859,25 +871,11 @@ async def _send_and_parse_async( objective_target_conversation_id=self._objective_target_conversation_id, objective=self._objective, ): - if self._memory_labels: - for piece in message.message_pieces: - piece.labels = self._memory_labels - response = await self._prompt_normalizer.send_prompt_async( + return await send_json_with_retry_async( + normalizer=self._prompt_normalizer, + target=self._adversarial_target, message=message, conversation_id=self._conversation_id, - target=self._adversarial_target, + parse=_parse, + fallback=None if self._raise_on_invalid_json else _fallback, ) - - if not response: - raise ValueError("No response received from adversarial chat") - - raw = response.get_value() - - schema = self._response_json_schema - if not self._raise_on_invalid_json: - try: - return _parse_adversarial_reply(raw, schema=schema) - except InvalidJsonException: - return AdversarialReply(next_message=raw, raw=raw) - - return _parse_adversarial_reply(raw, schema=schema) diff --git a/pyrit/memory/alembic/versions/a1c3e5d7f9b0_add_retries_to_conversations.py b/pyrit/memory/alembic/versions/a1c3e5d7f9b0_add_retries_to_conversations.py new file mode 100644 index 0000000000..5cbf22042d --- /dev/null +++ b/pyrit/memory/alembic/versions/a1c3e5d7f9b0_add_retries_to_conversations.py @@ -0,0 +1,35 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT license. + +""" +add retries to Conversations. + +Revision ID: a1c3e5d7f9b0 +Revises: d4e6f8a0b2c4 +Create Date: 2026-01-01 00:00:00.000000 +""" + +from collections.abc import Sequence + +import sqlalchemy as sa +from alembic import op + +# revision identifiers, used by Alembic. +revision: str = "a1c3e5d7f9b0" +down_revision: str | None = "d4e6f8a0b2c4" +branch_labels: str | Sequence[str] | None = None +depends_on: str | Sequence[str] | None = None + + +def upgrade() -> None: + """Apply this schema upgrade.""" + # ### commands auto generated by Alembic and reviewed by author ### + op.add_column("Conversations", sa.Column("retries", sa.JSON(), nullable=True)) + # ### end Alembic commands ### + + +def downgrade() -> None: + """Revert this schema upgrade.""" + # ### commands auto generated by Alembic and reviewed by author ### + op.drop_column("Conversations", "retries") + # ### end Alembic commands ### diff --git a/pyrit/memory/memory_interface.py b/pyrit/memory/memory_interface.py index 9e2bf380af..8d8dce36f9 100644 --- a/pyrit/memory/memory_interface.py +++ b/pyrit/memory/memory_interface.py @@ -15,7 +15,7 @@ from sqlalchemy import MetaData, and_, not_, or_, select from sqlalchemy.engine.base import Engine from sqlalchemy.exc import SQLAlchemyError -from sqlalchemy.orm.attributes import InstrumentedAttribute +from sqlalchemy.orm.attributes import InstrumentedAttribute, flag_modified if TYPE_CHECKING: from pyrit.memory.memory_embedding import MemoryEmbedding @@ -39,6 +39,8 @@ from pyrit.models import ( AttackResult, Conversation, + ConversationRetry, + ConversationRetryReason, ConversationStats, IdentifierFilter, IdentifierType, @@ -476,6 +478,82 @@ def _insert_conversation(self, *, conversation: Conversation) -> None: logger.exception(f"Error registering conversation {conversation.conversation_id}: {e}") raise + def add_conversation_retry(self, *, conversation_id: str, sequence: int, reason: ConversationRetryReason) -> None: + """ + Append a retry record to the conversation-scoped metadata for ``conversation_id``. + + Records that a turn had to be retried (e.g. because its response failed JSON + validation and was rolled back out of memory). The conversation's ``Conversations`` + row is updated in place; if no row exists yet it is created. This is distinct from + the insert-only ``_insert_conversation``. + + Args: + conversation_id (str): The conversation whose turn was retried. + sequence (int): The sequence the retried turn's request occupies. + reason (ConversationRetryReason): Why the turn was retried. + + Raises: + SQLAlchemyError: If the database update fails; the transaction is rolled back first. + """ + record = ConversationRetry(sequence=sequence, reason=reason).model_dump(mode="json") + with closing(self.get_session()) as session: + try: + entry = session.get(ConversationEntry, str(conversation_id)) + if entry is None: + entry = ConversationEntry(conversation=Conversation(conversation_id=str(conversation_id))) + session.add(entry) + entry.retries = [*(entry.retries or []), record] + flag_modified(entry, "retries") + session.commit() + except SQLAlchemyError as e: + session.rollback() + logger.exception(f"Error recording retry for conversation {conversation_id}: {e}") + raise + + def delete_conversation_pieces_after_sequence(self, *, conversation_id: str, sequence: int) -> int: + """ + Delete all message pieces in a conversation whose sequence is greater than ``sequence``. + + Rolls a conversation back to a baseline so a failed turn can be resent on a clean + history. Dependent ``EmbeddingData`` rows for the deleted pieces are removed first to + avoid orphaned foreign keys. Pieces at or below ``sequence`` (e.g. the system prompt + and any prior good turns) are left intact. + + Args: + conversation_id (str): The conversation to roll back. + sequence (int): The baseline sequence; pieces with a greater sequence are deleted. + + Returns: + int: The number of ``PromptMemoryEntries`` deleted. + + Raises: + SQLAlchemyError: If the deletion fails; the transaction is rolled back first. + """ + with closing(self.get_session()) as session: + try: + pieces = ( + session.query(PromptMemoryEntry) + .filter( + PromptMemoryEntry.conversation_id == str(conversation_id), + PromptMemoryEntry.sequence > sequence, + ) + .all() + ) + if not pieces: + return 0 + piece_ids = [piece.id for piece in pieces] + session.query(EmbeddingDataEntry).filter(EmbeddingDataEntry.id.in_(piece_ids)).delete( + synchronize_session=False + ) + for piece in pieces: + session.delete(piece) + session.commit() + return len(pieces) + except SQLAlchemyError as e: + session.rollback() + logger.exception(f"Error deleting conversation pieces for {conversation_id}: {e}") + raise + @abc.abstractmethod def _add_embeddings_to_memory(self, *, embedding_data: Sequence[EmbeddingDataEntry]) -> None: """ diff --git a/pyrit/memory/memory_models.py b/pyrit/memory/memory_models.py index 581375a895..a5413abc27 100644 --- a/pyrit/memory/memory_models.py +++ b/pyrit/memory/memory_models.py @@ -41,6 +41,7 @@ ComponentIdentifier, Conversation, ConversationReference, + ConversationRetry, ConversationType, EvaluationIdentifier, MessagePiece, @@ -370,6 +371,10 @@ class ConversationEntry(Base): conversation_id = mapped_column(String(36), primary_key=True, nullable=False) target_identifier: Mapped[dict[str, str] | None] = mapped_column(JSON, nullable=True) + # JSON-serialized list of ConversationRetry records (turns that were retried in + # this conversation). Nullable for backwards compatibility with existing databases. + retries: Mapped[list[dict[str, Any]] | None] = mapped_column(JSON, nullable=True) + # Version of PyRIT used when this entry was created. Nullable for backwards # compatibility with existing databases. pyrit_version = mapped_column(String, nullable=True) @@ -383,6 +388,7 @@ def __init__(self, *, conversation: Conversation) -> None: """ self.conversation_id = conversation.conversation_id self.target_identifier = conversation.target_identifier.model_dump() if conversation.target_identifier else None + self.retries = [retry.model_dump(mode="json") for retry in conversation.retries] or None self.pyrit_version = pyrit.__version__ def get_conversation(self) -> Conversation: @@ -394,7 +400,12 @@ def get_conversation(self) -> Conversation: """ stored_version = self.pyrit_version or LEGACY_PYRIT_VERSION target_id = _load_identifier(self.target_identifier, pyrit_version=stored_version) - return Conversation(conversation_id=self.conversation_id, target_identifier=target_id) + retries = [ConversationRetry.model_validate(retry) for retry in self.retries or []] + return Conversation( + conversation_id=self.conversation_id, + target_identifier=target_id, + retries=retries, + ) class EmbeddingDataEntry(Base): diff --git a/pyrit/models/__init__.py b/pyrit/models/__init__.py index 80584e62af..498fc2a0d8 100644 --- a/pyrit/models/__init__.py +++ b/pyrit/models/__init__.py @@ -84,6 +84,7 @@ ToolCall, ) from pyrit.models.messages.conversation_reference import ConversationReference, ConversationType +from pyrit.models.messages.conversation_retry import ConversationRetry, ConversationRetryReason from pyrit.models.parameter import ( ComponentType, Parameter, @@ -138,6 +139,8 @@ "ConverterIdentifier", "Conversation", "ConversationReference", + "ConversationRetry", + "ConversationRetryReason", "ConversationStats", "ConversationType", "construct_response_from_request", diff --git a/pyrit/models/messages/conversation_retry.py b/pyrit/models/messages/conversation_retry.py new file mode 100644 index 0000000000..9d21ce9d48 --- /dev/null +++ b/pyrit/models/messages/conversation_retry.py @@ -0,0 +1,34 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT license. + +from __future__ import annotations + +from enum import Enum + +from pydantic import BaseModel, ConfigDict + + +class ConversationRetryReason(str, Enum): + """Why a turn in a conversation had to be retried.""" + + JSON_PARSING = "json_parsing" + + +class ConversationRetry(BaseModel): + """ + Record of a single retried turn within a conversation. + + A retry happens when a turn's response failed validation (e.g. malformed JSON) + and the failed turn was rolled back out of memory so the turn could be resent on + a clean history. The record is conversation-scoped metadata: it captures which + turn was retried and why, without keeping the discarded turn's pieces around. + """ + + model_config = ConfigDict(frozen=True) + + # The sequence the retried turn's request occupies. Stable across attempts + # because the rollback resets the sequence so the eventual successful request + # lands at the same index. + sequence: int + + reason: ConversationRetryReason diff --git a/pyrit/models/messages/conversations.py b/pyrit/models/messages/conversations.py index d132679bb1..009bcb8989 100644 --- a/pyrit/models/messages/conversations.py +++ b/pyrit/models/messages/conversations.py @@ -7,8 +7,11 @@ from typing import TYPE_CHECKING -from pydantic import BaseModel, ConfigDict +from pydantic import BaseModel, ConfigDict, Field +from pyrit.models.messages.conversation_retry import ( # noqa: TC001 (runtime-required by Pydantic field annotations) + ConversationRetry, +) from pyrit.models.messages.message import Message from pyrit.models.messages.message_piece import MessagePiece from pyrit.models.score import ( # noqa: TC001 (runtime-required by Pydantic field annotations) @@ -25,11 +28,11 @@ class Conversation(BaseModel): """ Conversation-scoped metadata shared by every piece in a conversation. - A ``Conversation`` records identifiers that belong to the conversation as a - whole rather than to any individual ``MessagePiece`` -- most importantly the - target the conversation is held with. Persisting these once per conversation - (instead of stamping them onto every piece/row) is what keeps ``MessagePiece`` - small. + A ``Conversation`` records state that belongs to the conversation as a whole + rather than to any individual ``MessagePiece`` -- most importantly the target + the conversation is held with, plus the record of any turns that were retried. + Persisting the per-conversation identifiers once here (instead of stamping them + onto every piece/row) is what keeps ``MessagePiece`` small. """ model_config = ConfigDict( @@ -41,6 +44,9 @@ class Conversation(BaseModel): conversation_id: str target_identifier: ComponentIdentifierField | None = None + # Turns that were retried (rolled back out of memory and resent) in this conversation. + retries: list[ConversationRetry] = Field(default_factory=list) + def get_all_values(messages: Sequence[Message]) -> list[str]: """ diff --git a/pyrit/prompt_normalizer/__init__.py b/pyrit/prompt_normalizer/__init__.py index dbc09564f7..55351e61ac 100644 --- a/pyrit/prompt_normalizer/__init__.py +++ b/pyrit/prompt_normalizer/__init__.py @@ -9,6 +9,7 @@ """ from pyrit.prompt_normalizer.converter_configuration import ConverterConfiguration +from pyrit.prompt_normalizer.json_retry import send_json_with_retry_async from pyrit.prompt_normalizer.normalizer_request import NormalizerRequest from pyrit.prompt_normalizer.prompt_normalizer import PromptNormalizer @@ -16,4 +17,5 @@ "NormalizerRequest", "ConverterConfiguration", "PromptNormalizer", + "send_json_with_retry_async", ] diff --git a/pyrit/prompt_normalizer/json_retry.py b/pyrit/prompt_normalizer/json_retry.py new file mode 100644 index 0000000000..1e6a42e6b6 --- /dev/null +++ b/pyrit/prompt_normalizer/json_retry.py @@ -0,0 +1,96 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT license. + +from __future__ import annotations + +from typing import TYPE_CHECKING, TypeVar + +from pyrit.exceptions import InvalidJsonException, pyrit_json_retry +from pyrit.models import ConversationRetryReason + +if TYPE_CHECKING: + from collections.abc import Callable + + from pyrit.models import Message + from pyrit.prompt_normalizer.prompt_normalizer import PromptNormalizer + from pyrit.prompt_target import PromptTarget + +T = TypeVar("T") + + +async def send_json_with_retry_async( + *, + normalizer: PromptNormalizer, + target: PromptTarget, + message: Message, + conversation_id: str, + parse: Callable[[Message], T], + fallback: Callable[[Message], T] | None = None, +) -> T: + """ + Send a message expecting a JSON response, retrying on a clean conversation history. + + JSON retries are only useful if each attempt is independent. Because the normalizer + persists the request and response to memory *before* the caller validates the JSON, a + naive retry on a stable ``conversation_id`` replays the failed turn (the target rebuilds + history from memory and re-sees its own malformed reply plus a duplicated user prompt). + + This helper fixes that: it records a baseline sequence for the conversation and, on every + attempt, rolls memory back to that baseline before resending. The first attempt deletes + nothing; each retry deletes the previous failed turn (and records a ``ConversationRetry`` + marker) so the model sees a clean history identical to the first attempt. + + The retry loop keeps the ``@pyrit_json_retry`` (tenacity) decorator so retry logging and + the ``RetryCollector`` attribution driven by ``after=log_exception`` are preserved. + + Args: + normalizer (PromptNormalizer): Normalizer used to send the message. Its memory is the + source of truth that gets rolled back between attempts. + target (PromptTarget): The target to send the message to. + message (Message): The message to send. It is reused across attempts; the normalizer + deep-copies it, so mutating the persisted copy does not affect this object. + conversation_id (str): The conversation the message belongs to. Stays stable across + attempts; only the failed turn's pieces are rolled back. + parse (Callable[[Message], T]): Turns the response into the parsed result. Must raise + ``InvalidJsonException`` on a bad parse to trigger a retry. Other exceptions + (e.g. blocked/empty) propagate without retrying. + fallback (Callable[[Message], T] | None): When provided, the message is sent exactly + once and, on ``InvalidJsonException``, this is used to build the result from the + raw response instead of retrying. Preserves "do not raise / do not retry" + semantics for callers that tolerate invalid JSON. Defaults to None (retry mode). + + Returns: + T: The parsed result. + + Raises: + InvalidJsonException: In retry mode, when parsing still fails after the retry budget. + ValueError: If the target returns no response. + """ + memory = normalizer.memory + existing_pieces = memory.get_message_pieces(conversation_id=conversation_id) + baseline = max((piece.sequence for piece in existing_pieces), default=-1) + + if fallback is not None: + response = await normalizer.send_prompt_async(message=message, conversation_id=conversation_id, target=target) + if not response: + raise ValueError(f"No response received for conversation ID: {conversation_id}") + try: + return parse(response) + except InvalidJsonException: + return fallback(response) + + @pyrit_json_retry + async def _attempt_async() -> T: + deleted = memory.delete_conversation_pieces_after_sequence(conversation_id=conversation_id, sequence=baseline) + if deleted: + memory.add_conversation_retry( + conversation_id=conversation_id, + sequence=baseline + 1, + reason=ConversationRetryReason.JSON_PARSING, + ) + response = await normalizer.send_prompt_async(message=message, conversation_id=conversation_id, target=target) + if not response: + raise ValueError(f"No response received for conversation ID: {conversation_id}") + return parse(response) + + return await _attempt_async() diff --git a/pyrit/prompt_normalizer/prompt_normalizer.py b/pyrit/prompt_normalizer/prompt_normalizer.py index 687f620284..20c7e60eb5 100644 --- a/pyrit/prompt_normalizer/prompt_normalizer.py +++ b/pyrit/prompt_normalizer/prompt_normalizer.py @@ -167,6 +167,11 @@ async def send_prompt_async( # Only apply response converters to the last message (final response) # Intermediate messages are tool calls/outputs that don't need conversion for i, resp in enumerate(responses): + # A response belongs to the conversation it answers. Real targets already stamp this + # (via construct_response_from_request), matching the request pieces stamped above; + # enforcing it here keeps the persisted conversation coherent regardless of target. + for piece in resp.message_pieces: + piece.conversation_id = conversation_id is_last = i == len(responses) - 1 if is_last: await self.convert_values_async( diff --git a/pyrit/score/llm_scoring.py b/pyrit/score/llm_scoring.py index 78183f0a74..b21bf52598 100644 --- a/pyrit/score/llm_scoring.py +++ b/pyrit/score/llm_scoring.py @@ -6,8 +6,13 @@ import uuid from typing import TYPE_CHECKING, Any -from pyrit.exceptions import EmptyResponseException, ScorerLLMResponseBlockedException, pyrit_json_retry +from pyrit.exceptions import ( + EmptyResponseException, + InvalidJsonException, + ScorerLLMResponseBlockedException, +) from pyrit.models import JSON_SCHEMA_METADATA_KEY, Message, MessagePiece +from pyrit.prompt_normalizer import PromptNormalizer, send_json_with_retry_async if TYPE_CHECKING: from collections.abc import Sequence @@ -21,7 +26,6 @@ from pyrit.score.response_handler import ResponseHandler -@pyrit_json_retry async def _run_llm_scoring_async( *, chat_target: PromptTarget, @@ -34,14 +38,19 @@ async def _run_llm_scoring_async( prepended_text: str | None = None, category: Sequence[str] | str | None = None, objective: str | None = None, + normalizer: PromptNormalizer | None = None, ) -> UnvalidatedScore: """ Perform a single scoring round-trip against an LLM target and delegate parsing. This is the shared LLM evaluation mechanism: it sets the system prompt on the target, sends the value to be scored (forwarding ``response_handler.response_schema`` so targets that - support structured output can enforce it), applies the standard JSON retry behavior, and - delegates parsing and validation to ``response_handler``. It is intentionally stateless and + support structured output can enforce it), and delegates parsing and validation to + ``response_handler``. The round-trip is routed through a ``PromptNormalizer`` via + ``send_json_with_retry_async`` so the scorer's question and the target's answer are persisted + to memory (a full audit trail, and a real conversation an attack can link as a SCORE-type + related conversation) and so JSON retries roll memory back to a clean baseline between attempts + instead of replaying the target's own malformed reply. It is intentionally stateless and independent of any particular ``Scorer`` so that scorers can compose it without inheriting LLM machinery. @@ -69,6 +78,9 @@ async def _run_llm_scoring_async( from the response; supplying both is an error. Defaults to None. objective (str | None): The objective associated with the score, used for contextualizing the result. Defaults to None. + normalizer (PromptNormalizer | None): Normalizer used to send the scoring round-trip + and whose memory is rolled back between JSON retries. Injectable for testing; + defaults to a fresh ``PromptNormalizer()`` when not supplied. Returns: UnvalidatedScore: The parsed score, whose ``raw_score_value`` still needs to be @@ -129,39 +141,54 @@ async def _run_llm_scoring_async( ) scorer_llm_request = Message(message_pieces=message_pieces) - try: - response = await chat_target.send_prompt_async(message=scorer_llm_request) - except Exception as ex: - raise Exception(f"Error scoring prompt with original prompt ID: {scored_prompt_id}") from ex # Resolve the text piece that holds the JSON response (score_value + rationale). When it's # absent the scorer produced no parseable output. Guard on the actual failure (no text # piece) rather than "all pieces blocked" so every no-text shape is handled instead of # only the fully-blocked one. A content-filter block surfaces as a dedicated exception # (the calling Scorer owns whether to raise or fall back); any other no-text response is a - # genuine empty/malformed error. Neither is retried by @pyrit_json_retry. - text_piece = next( - (piece for piece in response[0].message_pieces if piece.converted_value_data_type == "text"), None - ) - if text_piece is None: - if any(piece.is_blocked() for piece in response[0].message_pieces): - raise ScorerLLMResponseBlockedException( + # genuine empty/malformed error. Neither is retried; only invalid JSON triggers a retry. + def _parse(response: Message) -> UnvalidatedScore: + text_piece = next( + (piece for piece in response.message_pieces if piece.converted_value_data_type == "text"), None + ) + if text_piece is None: + if any(piece.is_blocked() for piece in response.message_pieces): + raise ScorerLLMResponseBlockedException( + message=( + f"The scorer's LLM response was blocked by content filtering while scoring " + f"prompt ID: {scored_prompt_id}. Consider using a scorer endpoint with " + f"content filtering disabled for red-teaming workflows." + ) + ) + raise EmptyResponseException( message=( - f"The scorer's LLM response was blocked by content filtering while scoring " - f"prompt ID: {scored_prompt_id}. Consider using a scorer endpoint with " - f"content filtering disabled for red-teaming workflows." + f"The scorer's LLM response contained no text to parse while scoring prompt ID: {scored_prompt_id}." ) ) - raise EmptyResponseException( - message=( - f"The scorer's LLM response contained no text to parse while scoring prompt ID: {scored_prompt_id}." - ) + + return response_handler.parse( + response_text=text_piece.converted_value, + scorer_identifier=scorer_identifier, + scored_prompt_id=scored_prompt_id, + category=category, + objective=objective, ) - return response_handler.parse( - response_text=text_piece.converted_value, - scorer_identifier=scorer_identifier, - scored_prompt_id=scored_prompt_id, - category=category, - objective=objective, - ) + # Route the round-trip through the normalizer so the scorer Q&A is persisted and JSON retries + # replay on a clean history. + try: + return await send_json_with_retry_async( + normalizer=normalizer or PromptNormalizer(), + target=chat_target, + message=scorer_llm_request, + conversation_id=conversation_id, + parse=_parse, + fallback=None, + ) + except (ScorerLLMResponseBlockedException, EmptyResponseException, InvalidJsonException): + # Terminal / caller-owned outcomes: propagate unchanged so the calling Scorer can apply + # its own policy (fall back, raise, or -- for invalid JSON -- surface the retry exhaustion). + raise + except Exception as ex: + raise Exception(f"Error scoring prompt with original prompt ID: {scored_prompt_id}") from ex diff --git a/tests/unit/executor/attack/component/test_adversarial_conversation_manager.py b/tests/unit/executor/attack/component/test_adversarial_conversation_manager.py index 5e23990d61..271b248ff6 100644 --- a/tests/unit/executor/attack/component/test_adversarial_conversation_manager.py +++ b/tests/unit/executor/attack/component/test_adversarial_conversation_manager.py @@ -529,7 +529,7 @@ async def test_schema_metadata_forwarded(self): async def test_no_response_raises(self): manager = _manager(prompt_normalizer=_normalizer(None)) - with pytest.raises(ValueError, match="No response received from adversarial chat"): + with pytest.raises(ValueError, match="No response received for conversation ID"): await manager.get_next_message_async(turn_index=1, last_response=_response_message()) async def test_invalid_reply_raises(self): @@ -774,7 +774,7 @@ async def test_schema_metadata_forwarded(self): async def test_no_response_raises(self): manager = _manager(adversarial_system_prompt=_system_prompt(schema=SCHEMA), prompt_normalizer=_normalizer(None)) - with pytest.raises(ValueError, match="No response received from adversarial chat"): + with pytest.raises(ValueError, match="No response received for conversation ID"): await manager.generate_adversarial_reply_async(prompt_text="x") async def test_invalid_json_raises(self): diff --git a/tests/unit/executor/attack/multi_turn/test_crescendo.py b/tests/unit/executor/attack/multi_turn/test_crescendo.py index 9ea9cb9d08..cb74a18fd5 100644 --- a/tests/unit/executor/attack/multi_turn/test_crescendo.py +++ b/tests/unit/executor/attack/multi_turn/test_crescendo.py @@ -872,7 +872,7 @@ async def test_generate_next_prompt_raises_when_adversarial_chat_returns_no_resp # Mock no response mock_prompt_normalizer.send_prompt_async.return_value = None - with pytest.raises(ValueError, match="No response received from adversarial chat"): + with pytest.raises(ValueError, match="No response received for conversation ID"): await attack._generate_next_prompt_async(context=basic_context) async def test_generate_next_prompt_forwards_seed_media_and_schema_to_adversarial_chat( diff --git a/tests/unit/executor/attack/multi_turn/test_red_teaming.py b/tests/unit/executor/attack/multi_turn/test_red_teaming.py index 4dfdaeebba..74721c0fbd 100644 --- a/tests/unit/executor/attack/multi_turn/test_red_teaming.py +++ b/tests/unit/executor/attack/multi_turn/test_red_teaming.py @@ -948,7 +948,7 @@ async def test_generate_next_prompt_raises_on_none_response( basic_context.executed_turns = 1 mock_prompt_normalizer.send_prompt_async.return_value = None - with pytest.raises(ValueError, match="No response received from adversarial chat"): + with pytest.raises(ValueError, match="No response received for conversation ID"): await attack._generate_next_prompt_async(context=basic_context) diff --git a/tests/unit/score/test_scorer.py b/tests/unit/score/test_scorer.py index 0e2435897c..eab82624e5 100644 --- a/tests/unit/score/test_scorer.py +++ b/tests/unit/score/test_scorer.py @@ -168,13 +168,20 @@ def get_scorer_metrics(self): @pytest.mark.parametrize("bad_json", [BAD_JSON, KEY_ERROR_JSON, KEY_ERROR2_JSON]) -async def test_scorer_send_chat_target_async_bad_json_exception_retries(bad_json: str): +async def test_scorer_send_chat_target_async_bad_json_exception_retries(bad_json: str, patch_central_database): chat_target = MagicMock(PromptTarget) chat_target.get_identifier.return_value = get_mock_target_identifier("MockChatTarget") - bad_json_resp = Message( - message_pieces=[MessagePiece(role="assistant", original_value=bad_json, conversation_id="test-convo")] - ) - chat_target.send_prompt_async = AsyncMock(return_value=[bad_json_resp]) + + def _fresh_bad_json_response(*args, **kwargs): + # A real target returns a fresh response (new piece ids) on every call; build one per + # attempt so the retry path doesn't collide on a reused message-piece id in memory. + return [ + Message( + message_pieces=[MessagePiece(role="assistant", original_value=bad_json, conversation_id="test-convo")] + ) + ] + + chat_target.send_prompt_async = AsyncMock(side_effect=_fresh_bad_json_response) scorer = MockScorer() with pytest.raises(InvalidJsonException): await _run_llm_scoring_async( @@ -193,7 +200,7 @@ async def test_scorer_send_chat_target_async_bad_json_exception_retries(bad_json assert chat_target.send_prompt_async.call_count == 2 -async def test_scorer_score_value_with_llm_exception_display_prompt_id(): +async def test_scorer_score_value_with_llm_exception_display_prompt_id(patch_central_database): chat_target = MagicMock(PromptTarget) chat_target.get_identifier.return_value = get_mock_target_identifier("MockChatTarget") chat_target.send_prompt_async = AsyncMock(side_effect=Exception("Test exception")) @@ -214,7 +221,7 @@ async def test_scorer_score_value_with_llm_exception_display_prompt_id(): ) -async def test_scorer_send_chat_target_async_good_response(good_json): +async def test_scorer_send_chat_target_async_good_response(good_json, patch_central_database): chat_target = MagicMock(PromptTarget) chat_target.get_identifier.return_value = get_mock_target_identifier("MockChatTarget") @@ -240,7 +247,7 @@ async def test_scorer_send_chat_target_async_good_response(good_json): assert chat_target.send_prompt_async.call_count == 1 -async def test_scorer_remove_markdown_json_called(good_json): +async def test_scorer_remove_markdown_json_called(good_json, patch_central_database): chat_target = MagicMock(PromptTarget) chat_target.get_identifier.return_value = get_mock_target_identifier("MockChatTarget") good_json_resp = Message( @@ -268,7 +275,9 @@ async def test_scorer_remove_markdown_json_called(good_json): mock_remove_markdown_json.assert_called_once() -async def test_score_value_with_llm_prepended_text_message_piece_creates_multipiece_message(good_json): +async def test_score_value_with_llm_prepended_text_message_piece_creates_multipiece_message( + good_json, patch_central_database, tmp_path +): """Test that prepended_text_message_piece creates a multi-piece message (text context + main content).""" chat_target = MagicMock(PromptTarget) chat_target.get_identifier.return_value = get_mock_target_identifier("MockChatTarget") @@ -279,12 +288,15 @@ async def test_score_value_with_llm_prepended_text_message_piece_creates_multipi scorer = MockScorer() + image_path = tmp_path / "test_image.png" + image_path.write_bytes(b"\x89PNG\r\n\x1a\n") + await _run_llm_scoring_async( chat_target=chat_target, response_handler=JsonSchemaResponseHandler(), scorer_identifier=scorer.get_identifier(), system_prompt="system_prompt", - value="test_image.png", + value=str(image_path), data_type="image_path", scored_prompt_id="123", prepended_text="objective: test\nresponse:", @@ -310,10 +322,10 @@ async def test_score_value_with_llm_prepended_text_message_piece_creates_multipi # Second piece should be the main content (image in this case) main_piece = sent_message.message_pieces[1] assert main_piece.converted_value_data_type == "image_path" - assert main_piece.original_value == "test_image.png" + assert main_piece.original_value == str(image_path) -async def test_score_value_with_llm_no_prepended_text_creates_single_piece_message(good_json): +async def test_score_value_with_llm_no_prepended_text_creates_single_piece_message(good_json, patch_central_database): """Test that without prepended_text_message_piece, only a single piece message is created.""" chat_target = MagicMock(PromptTarget) chat_target.get_identifier.return_value = get_mock_target_identifier("MockChatTarget") @@ -350,7 +362,7 @@ async def test_score_value_with_llm_no_prepended_text_creates_single_piece_messa assert "response: some text" in text_piece.original_value -async def test_score_value_with_llm_prepended_text_works_with_audio(good_json): +async def test_score_value_with_llm_prepended_text_works_with_audio(good_json, patch_central_database, tmp_path): """Test that prepended_text_message_piece works with audio content (type-independent).""" chat_target = MagicMock(PromptTarget) chat_target.get_identifier.return_value = get_mock_target_identifier("MockChatTarget") @@ -361,12 +373,15 @@ async def test_score_value_with_llm_prepended_text_works_with_audio(good_json): scorer = MockScorer() + audio_path = tmp_path / "test_audio.wav" + audio_path.write_bytes(b"RIFF0000WAVE") + await _run_llm_scoring_async( chat_target=chat_target, response_handler=JsonSchemaResponseHandler(), scorer_identifier=scorer.get_identifier(), system_prompt="system_prompt", - value="test_audio.wav", + value=str(audio_path), data_type="audio_path", scored_prompt_id="123", prepended_text="objective: transcribe and evaluate\nresponse:", @@ -388,7 +403,7 @@ async def test_score_value_with_llm_prepended_text_works_with_audio(good_json): # Second piece should be audio audio_piece = sent_message.message_pieces[1] assert audio_piece.converted_value_data_type == "audio_path" - assert audio_piece.original_value == "test_audio.wav" + assert audio_piece.original_value == str(audio_path) def test_scorer_extract_task_from_response(patch_central_database): @@ -1531,7 +1546,7 @@ async def test_text_only_scorer_filters_blocked_via_validator( assert scores[0].get_value() == 0.0 -async def test_score_value_with_llm_skips_reasoning_piece(good_json): +async def test_score_value_with_llm_skips_reasoning_piece(good_json, patch_central_database): """Test that _score_value_with_llm extracts JSON from the text piece, not a reasoning piece.""" chat_target = MagicMock(PromptTarget) chat_target.get_identifier.return_value = get_mock_target_identifier("MockChatTarget") @@ -1571,7 +1586,7 @@ async def test_score_value_with_llm_skips_reasoning_piece(good_json): assert result.score_rationale == "Valid response" -async def test_score_value_with_llm_raises_when_scorer_response_blocked(): +async def test_score_value_with_llm_raises_when_scorer_response_blocked(patch_central_database): """When the scorer's own LLM response is blocked, the transport raises ScorerLLMResponseBlockedException.""" from pyrit.exceptions import ScorerLLMResponseBlockedException @@ -1609,7 +1624,7 @@ async def test_score_value_with_llm_raises_when_scorer_response_blocked(): assert chat_target.send_prompt_async.call_count == 1 -async def test_score_value_with_llm_raises_empty_response_when_no_text_piece(): +async def test_score_value_with_llm_raises_empty_response_when_no_text_piece(patch_central_database): """A no-text response that wasn't content-filtered raises EmptyResponseException, not blocked.""" from pyrit.exceptions import EmptyResponseException From 3cab558c034f683d707d86e1e96b5584ec28c542 Mon Sep 17 00:00:00 2001 From: Richard Lundeen Date: Tue, 14 Jul 2026 09:49:09 -0700 Subject: [PATCH 2/4] MAINT: Remove raise_on_invalid_json escape hatch The raise_on_invalid_json=False path forwarded unparseable adversarial model output as the next attack prompt instead of failing loudly -- a debuggability hazard -- and no shipped attack used it (default True; only a test exercised it). Removing the flag collapses the adversarial send to a single always-raise path and lets send_json_with_retry_async drop its fallback parameter. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> Copilot-Session: 244845c9-0655-4fa8-9717-29d75faef9f2 --- .../adversarial_conversation_manager.py | 17 +++-------------- pyrit/prompt_normalizer/json_retry.py | 18 ++---------------- pyrit/score/llm_scoring.py | 1 - .../test_adversarial_conversation_manager.py | 11 ----------- 4 files changed, 5 insertions(+), 42 deletions(-) diff --git a/pyrit/executor/attack/component/adversarial_conversation_manager.py b/pyrit/executor/attack/component/adversarial_conversation_manager.py index 5a617bc470..662568ab28 100644 --- a/pyrit/executor/attack/component/adversarial_conversation_manager.py +++ b/pyrit/executor/attack/component/adversarial_conversation_manager.py @@ -412,7 +412,6 @@ def __init__( adversarial_first_user_message: SeedPrompt | None = None, adversarial_next_user_message: SeedPrompt | None = None, max_turns: int = 1, - raise_on_invalid_json: bool = True, prompt_normalizer: PromptNormalizer | None = None, conversation_id: str | None = None, objective: str | None = None, @@ -438,9 +437,6 @@ def __init__( text directly. max_turns: Maximum number of turns; rendered into the adversarial system prompt as ``max_turns``. Defaults to 1. - raise_on_invalid_json: When True (default), a reply that fails to match the resolved - schema raises ``InvalidJsonException`` (retried via ``send_json_with_retry_async``). - When False, the raw reply text is returned as ``next_message`` instead of raising. prompt_normalizer: The prompt normalizer to send through. Defaults to a new one. conversation_id: The adversarial-chat conversation id this manager drives. A fresh id is generated when None. @@ -467,7 +463,6 @@ def __init__( self._adversarial_first_user_message = adversarial_first_user_message self._adversarial_next_user_message = adversarial_next_user_message self._max_turns = max_turns - self._raise_on_invalid_json = raise_on_invalid_json self._prompt_normalizer = prompt_normalizer or PromptNormalizer() self._conversation_id = conversation_id or str(uuid4()) self._objective = objective @@ -756,7 +751,7 @@ async def generate_adversarial_reply_async( Raises: ValueError: If no response is received from the adversarial chat. - InvalidJsonException: If ``raise_on_invalid_json`` is True and the reply is invalid. + InvalidJsonException: If the reply is not valid JSON after the retry budget is exhausted. """ return await self._send_and_parse_async( prompt_text=prompt_text, @@ -815,8 +810,7 @@ async def _send_and_parse_async( This is the single place adversarial-chat JSON retry lives. It delegates to ``send_json_with_retry_async`` so each retry sends on a clean conversation history: the failed turn is rolled back out of memory before the turn is resent, instead of - replaying the model's own malformed reply. When ``raise_on_invalid_json`` is False, an - unparseable reply is returned as raw text (single attempt, no retry). + replaying the model's own malformed reply. When a ``modality_router`` is configured, the outgoing message is built via ``build_adversarial_input_message`` so first-turn seed media (``seed_message``) and prior @@ -833,7 +827,7 @@ async def _send_and_parse_async( Raises: ValueError: If no response is received from the adversarial chat. - InvalidJsonException: If ``raise_on_invalid_json`` is True and the reply is invalid. + InvalidJsonException: If the reply is not valid JSON after the retry budget is exhausted. """ prompt_metadata = _build_adversarial_prompt_metadata(response_json_schema=self._response_json_schema) @@ -860,10 +854,6 @@ async def _send_and_parse_async( def _parse(response: Message) -> AdversarialReply: return _parse_adversarial_reply(response.get_value(), schema=schema) - def _fallback(response: Message) -> AdversarialReply: - raw = response.get_value() - return AdversarialReply(next_message=raw, raw=raw) - with execution_context( component_role=ComponentRole.ADVERSARIAL_CHAT, attack_strategy_name=self._attack_strategy_name, @@ -877,5 +867,4 @@ def _fallback(response: Message) -> AdversarialReply: message=message, conversation_id=self._conversation_id, parse=_parse, - fallback=None if self._raise_on_invalid_json else _fallback, ) diff --git a/pyrit/prompt_normalizer/json_retry.py b/pyrit/prompt_normalizer/json_retry.py index 1e6a42e6b6..ee221ff230 100644 --- a/pyrit/prompt_normalizer/json_retry.py +++ b/pyrit/prompt_normalizer/json_retry.py @@ -5,7 +5,7 @@ from typing import TYPE_CHECKING, TypeVar -from pyrit.exceptions import InvalidJsonException, pyrit_json_retry +from pyrit.exceptions import pyrit_json_retry from pyrit.models import ConversationRetryReason if TYPE_CHECKING: @@ -25,7 +25,6 @@ async def send_json_with_retry_async( message: Message, conversation_id: str, parse: Callable[[Message], T], - fallback: Callable[[Message], T] | None = None, ) -> T: """ Send a message expecting a JSON response, retrying on a clean conversation history. @@ -54,31 +53,18 @@ async def send_json_with_retry_async( parse (Callable[[Message], T]): Turns the response into the parsed result. Must raise ``InvalidJsonException`` on a bad parse to trigger a retry. Other exceptions (e.g. blocked/empty) propagate without retrying. - fallback (Callable[[Message], T] | None): When provided, the message is sent exactly - once and, on ``InvalidJsonException``, this is used to build the result from the - raw response instead of retrying. Preserves "do not raise / do not retry" - semantics for callers that tolerate invalid JSON. Defaults to None (retry mode). Returns: T: The parsed result. Raises: - InvalidJsonException: In retry mode, when parsing still fails after the retry budget. + InvalidJsonException: When parsing still fails after the retry budget is exhausted. ValueError: If the target returns no response. """ memory = normalizer.memory existing_pieces = memory.get_message_pieces(conversation_id=conversation_id) baseline = max((piece.sequence for piece in existing_pieces), default=-1) - if fallback is not None: - response = await normalizer.send_prompt_async(message=message, conversation_id=conversation_id, target=target) - if not response: - raise ValueError(f"No response received for conversation ID: {conversation_id}") - try: - return parse(response) - except InvalidJsonException: - return fallback(response) - @pyrit_json_retry async def _attempt_async() -> T: deleted = memory.delete_conversation_pieces_after_sequence(conversation_id=conversation_id, sequence=baseline) diff --git a/pyrit/score/llm_scoring.py b/pyrit/score/llm_scoring.py index b21bf52598..4071942677 100644 --- a/pyrit/score/llm_scoring.py +++ b/pyrit/score/llm_scoring.py @@ -184,7 +184,6 @@ def _parse(response: Message) -> UnvalidatedScore: message=scorer_llm_request, conversation_id=conversation_id, parse=_parse, - fallback=None, ) except (ScorerLLMResponseBlockedException, EmptyResponseException, InvalidJsonException): # Terminal / caller-owned outcomes: propagate unchanged so the calling Scorer can apply diff --git a/tests/unit/executor/attack/component/test_adversarial_conversation_manager.py b/tests/unit/executor/attack/component/test_adversarial_conversation_manager.py index 271b248ff6..e275dbffb1 100644 --- a/tests/unit/executor/attack/component/test_adversarial_conversation_manager.py +++ b/tests/unit/executor/attack/component/test_adversarial_conversation_manager.py @@ -539,17 +539,6 @@ async def test_invalid_reply_raises(self): with pytest.raises(InvalidJsonException): await manager.get_next_message_async(turn_index=1, last_response=_response_message()) - async def test_raise_on_invalid_json_false_returns_raw(self): - normalizer = _normalizer("totally not json") - manager = _manager( - adversarial_system_prompt=_system_prompt(schema=SCHEMA), - raise_on_invalid_json=False, - prompt_normalizer=normalizer, - ) - turn = await manager.get_next_message_async(turn_index=1, last_response=_response_message()) - assert turn.reply is not None and turn.reply.next_message == "totally not json" - assert turn.objective_message.get_value() == "totally not json" - # --- get_next_message_async: bypass path ------------------------------------- From 371c8cc0e49e5fec9a984a7c027bebcf568e8b11 Mon Sep 17 00:00:00 2001 From: Richard Lundeen Date: Tue, 14 Jul 2026 09:55:58 -0700 Subject: [PATCH 3/4] DOC: Clarify normalizer owns memory persistence, targets never write Add a concise note to the framework doc that the prompt_normalizer is the single component that persists requests/responses to memory, and targets never write to memory themselves. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> Copilot-Session: 244845c9-0655-4fa8-9717-29d75faef9f2 --- doc/code/framework.md | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/doc/code/framework.md b/doc/code/framework.md index 6bcb4a6890..4706742da0 100644 --- a/doc/code/framework.md +++ b/doc/code/framework.md @@ -231,7 +231,7 @@ If you are contributing to PyRIT, that work will most likely land in one of the - Targets should use message_normalizer along with TargetConfiguration to transform `Messages` into formats that target supports. - Because targets are so varied, it is reasonable to return multiple tool calls, or none at all. - One attack can have many targets (and in fact, converters and scorers can also use targets to convert/score the prompt). -- **Does not own**: what to send or what to do with the response. A target sends a prepared `Message` and returns a response — it doesn't convert prompts (converters), score (scorers), manage the conversation or decide the next turn (attacks), or apply attack logic. Its retries stay at the target layer (e.g. `RateLimitException`). +- **Does not own**: what to send or what to do with the response. A target sends a prepared `Message` and returns a response — it doesn't convert prompts (converters), score (scorers), manage the conversation or decide the next turn (attacks), apply attack logic, or persist prompts and responses to memory (the `prompt_normalizer` owns that). Its retries stay at the target layer (e.g. `RateLimitException`). **Framework Plans**: @@ -310,7 +310,7 @@ The below talks about responsibilities of most modules in the PyRIT library **Responsibility**: Reshape prompts and conversations so components and targets can interoperate. There are two distinct modules: -- **`prompt_normalizer`** applies converters and dispatches individual prompts to a `PromptTarget` (handling batching and memory persistence). `NormalizerRequest` and `ConverterConfiguration` describe what to send and which converters to apply. +- **`prompt_normalizer`** applies converters and dispatches individual prompts to a `PromptTarget` (handling batching and memory persistence). It is the single component that writes each request and response to memory; targets never persist on their own. `NormalizerRequest` and `ConverterConfiguration` describe what to send and which converters to apply. - **`message_normalizer`** reshapes multi-message conversation payloads into the structure a given model expects — for example, handling system-message behavior (keep / squash / ignore), history squashing, and tokenizer chat templates. ## [Output](./output/0_output) From 4934d02f2464489e319293a092a4f05b519ae294 Mon Sep 17 00:00:00 2001 From: Richard Lundeen Date: Wed, 15 Jul 2026 13:13:55 -0700 Subject: [PATCH 4/4] TEST: Add unit tests for send_json_with_retry_async; trim docstring Address PR review: add dedicated tests covering the happy path, retry-then-success (asserting the poisoned turn is rolled back and a retry marker is recorded), retry exhaustion, no-response ValueError, and non-JSON exception passthrough. Also trim the helper docstring so it describes the retry contract rather than narrating the specific bug it fixed. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> Copilot-Session: 244845c9-0655-4fa8-9717-29d75faef9f2 --- pyrit/prompt_normalizer/json_retry.py | 16 +- .../unit/prompt_normalizer/test_json_retry.py | 158 ++++++++++++++++++ 2 files changed, 164 insertions(+), 10 deletions(-) create mode 100644 tests/unit/prompt_normalizer/test_json_retry.py diff --git a/pyrit/prompt_normalizer/json_retry.py b/pyrit/prompt_normalizer/json_retry.py index ee221ff230..e63fd662e4 100644 --- a/pyrit/prompt_normalizer/json_retry.py +++ b/pyrit/prompt_normalizer/json_retry.py @@ -27,17 +27,13 @@ async def send_json_with_retry_async( parse: Callable[[Message], T], ) -> T: """ - Send a message expecting a JSON response, retrying on a clean conversation history. + Send a message expecting a JSON response, retrying each attempt on a clean conversation history. - JSON retries are only useful if each attempt is independent. Because the normalizer - persists the request and response to memory *before* the caller validates the JSON, a - naive retry on a stable ``conversation_id`` replays the failed turn (the target rebuilds - history from memory and re-sees its own malformed reply plus a duplicated user prompt). - - This helper fixes that: it records a baseline sequence for the conversation and, on every - attempt, rolls memory back to that baseline before resending. The first attempt deletes - nothing; each retry deletes the previous failed turn (and records a ``ConversationRetry`` - marker) so the model sees a clean history identical to the first attempt. + JSON retries are only useful if each attempt is independent. This helper records a baseline + sequence for the conversation and, on every attempt, rolls memory back to that baseline before + resending: the first attempt deletes nothing; each retry deletes the previous failed turn (and + records a ``ConversationRetry`` marker) so the target rebuilds history from memory and sees a + clean conversation identical to the first attempt, instead of replaying its own malformed reply. The retry loop keeps the ``@pyrit_json_retry`` (tenacity) decorator so retry logging and the ``RetryCollector`` attribution driven by ``after=log_exception`` are preserved. diff --git a/tests/unit/prompt_normalizer/test_json_retry.py b/tests/unit/prompt_normalizer/test_json_retry.py new file mode 100644 index 0000000000..46c8559642 --- /dev/null +++ b/tests/unit/prompt_normalizer/test_json_retry.py @@ -0,0 +1,158 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT license. + +import json +from unittest.mock import AsyncMock, MagicMock +from uuid import uuid4 + +import pytest +from unit.mocks import MockPromptTarget + +from pyrit.exceptions import InvalidJsonException +from pyrit.models import ConversationRetryReason, Message, MessagePiece +from pyrit.prompt_normalizer import PromptNormalizer, send_json_with_retry_async + + +class _QueueTarget(MockPromptTarget): + """A MockPromptTarget that replies with a queued sequence of texts, one per send.""" + + def __init__(self, *, responses: list[str]) -> None: + super().__init__() + self._responses = list(responses) + + async def _send_prompt_to_target_async(self, *, normalized_conversation: list[Message]) -> list[Message]: + message = normalized_conversation[-1] + self.prompt_sent.append(message.get_value()) + text = self._responses.pop(0) + return [ + MessagePiece( + role="assistant", + original_value=text, + conversation_id=message.message_pieces[0].conversation_id, + ).to_message() + ] + + +def _parse(response: Message) -> dict: + text = response.get_value() + try: + return json.loads(text) + except json.JSONDecodeError: + raise InvalidJsonException(message=f"invalid json: {text}") from None + + +def _user_message() -> Message: + return Message.from_prompt(prompt="give me json", role="user") + + +@pytest.mark.usefixtures("patch_central_database") +class TestSendJsonWithRetryAsync: + async def test_valid_first_attempt_returns_parsed_result(self): + normalizer = PromptNormalizer() + target = _QueueTarget(responses=['{"answer": 42}']) + conversation_id = str(uuid4()) + + result = await send_json_with_retry_async( + normalizer=normalizer, + target=target, + message=_user_message(), + conversation_id=conversation_id, + parse=_parse, + ) + + assert result == {"answer": 42} + assert len(target.prompt_sent) == 1 + # No retry occurred, so no retry marker is recorded. + conversation = normalizer.memory._get_conversation(conversation_id=conversation_id) + assert conversation is not None + assert conversation.retries == [] + + async def test_retry_then_success_rolls_back_poisoned_turn(self): + normalizer = PromptNormalizer() + target = _QueueTarget(responses=["not json at all", '{"answer": 7}']) + conversation_id = str(uuid4()) + + result = await send_json_with_retry_async( + normalizer=normalizer, + target=target, + message=_user_message(), + conversation_id=conversation_id, + parse=_parse, + ) + + assert result == {"answer": 7} + assert len(target.prompt_sent) == 2 + + # The malformed first turn must not survive in memory: only the clean retried + # request/response pair remains, so a caller replaying history never sees the bad reply. + pieces = normalizer.memory.get_message_pieces(conversation_id=conversation_id) + values = [piece.original_value for piece in pieces] + assert "not json at all" not in values + assert '{"answer": 7}' in values + + # A retry marker was recorded for the rolled-back turn. + conversation = normalizer.memory._get_conversation(conversation_id=conversation_id) + assert conversation is not None + assert len(conversation.retries) == 1 + assert conversation.retries[0].reason == ConversationRetryReason.JSON_PARSING + + async def test_exhausting_retries_raises_invalid_json(self): + normalizer = PromptNormalizer() + # conftest sets RETRY_MAX_NUM_ATTEMPTS=2, so two malformed replies exhaust the budget. + target = _QueueTarget(responses=["nope", "still nope"]) + conversation_id = str(uuid4()) + + with pytest.raises(InvalidJsonException): + await send_json_with_retry_async( + normalizer=normalizer, + target=target, + message=_user_message(), + conversation_id=conversation_id, + parse=_parse, + ) + + # Rollback happens at the start of each attempt, so the first failed turn ("nope") is + # rolled back before the second attempt; a retry marker records it. The final failed turn + # legitimately remains -- there is no further attempt to poison. + pieces = normalizer.memory.get_message_pieces(conversation_id=conversation_id) + assert all(piece.original_value != "nope" for piece in pieces) + conversation = normalizer.memory._get_conversation(conversation_id=conversation_id) + assert conversation is not None + assert len(conversation.retries) == 1 + + async def test_no_response_raises_value_error(self): + normalizer = MagicMock(spec=PromptNormalizer) + normalizer.memory = MagicMock() + normalizer.memory.get_message_pieces.return_value = [] + normalizer.memory.delete_conversation_pieces_after_sequence.return_value = 0 + normalizer.send_prompt_async = AsyncMock(return_value=None) + conversation_id = str(uuid4()) + + with pytest.raises(ValueError, match="No response received for conversation ID"): + await send_json_with_retry_async( + normalizer=normalizer, + target=MagicMock(), + message=_user_message(), + conversation_id=conversation_id, + parse=_parse, + ) + + async def test_non_json_exception_propagates_without_retrying(self): + normalizer = PromptNormalizer() + target = _QueueTarget(responses=['{"answer": 1}']) + conversation_id = str(uuid4()) + + def _parse_raising(response: Message) -> dict: + raise RuntimeError("boom") + + with pytest.raises(RuntimeError, match="boom"): + await send_json_with_retry_async( + normalizer=normalizer, + target=target, + message=_user_message(), + conversation_id=conversation_id, + parse=_parse_raising, + ) + + # A non-JSON error is terminal: the target is hit exactly once (no retry). + assert len(target.prompt_sent) == 1