From ae9bac0e1c1f0fd768b6cba21415f7d2462d818b Mon Sep 17 00:00:00 2001 From: Copilot <223556219+Copilot@users.noreply.github.com> Date: Fri, 7 Aug 2026 07:25:54 -0700 Subject: [PATCH 01/24] Fix request converter scoping Default prepended request conversion to user history and restrict Jailbreak composition to explicitly compatible direct techniques. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> --- doc/code/framework.md | 1 + .../attack/component/conversation_manager.py | 30 +++-- .../prepended_conversation_config.py | 13 +- .../attack/single_turn/prompt_sending.py | 2 + .../scenario/core/attack_technique_factory.py | 34 +++++ pyrit/scenario/scenarios/airt/jailbreak.py | 86 +++++++++--- pyrit/setup/initializers/techniques/airt.py | 1 + pyrit/setup/initializers/techniques/core.py | 1 + .../component/test_conversation_manager.py | 27 +++- .../test_prepended_conversation_config.py | 6 +- .../attack/single_turn/test_prompt_sending.py | 58 +++++++- tests/unit/scenario/airt/test_jailbreak.py | 125 ++++++++++++------ .../core/test_attack_technique_factory.py | 23 ++++ 13 files changed, 322 insertions(+), 85 deletions(-) diff --git a/doc/code/framework.md b/doc/code/framework.md index a3938ef2dc..3f98f0b69a 100644 --- a/doc/code/framework.md +++ b/doc/code/framework.md @@ -182,6 +182,7 @@ If you are contributing to PyRIT, that work will most likely land in one of the - `core` stays deliberately small so a default run doesn't print 200 techniques or take forever; the wider catalog lives in `extra` and is selected on demand. Users pick subsets by passing initializer tags (e.g. `core`, `extra`, `all`) or writing their own initializer, so different runs — including from the CLI — can register different technique sets without changing the catalog. - A technique tied to one scenario is fine; if it's pinned and non-reusable it can stay local to that scenario, but if another scenario could reuse it, promote it to a catalog module and tag it. - Tags describe a technique (behavioral tags like `single_turn`/`multi_turn`, owner tags like `airt`); they don't decide what a scenario runs. There is deliberately **no global `default` tag** — a default is scenario-relative, declared per scenario via `build_technique_class_from_factories` (the `factories` list is the pool, catalog tags become named aggregate presets, and `default_tags` / `default_names` set what runs when nothing is chosen). +- Factories opt into additive request-converter composition with `supports_request_converter_composition=True`. This is a semantic capability, not just constructor-signature detection; the factory validates that opted-in attacks accept `attack_converter_config`. - **Does not own**: the conversation algorithm itself. Branching, turn management, and scoring decisions live in the executor it wraps — a technique only selects and configures existing components, and shouldn't implement new sending, scoring, or branching logic. **Framework Plans**: diff --git a/pyrit/executor/attack/component/conversation_manager.py b/pyrit/executor/attack/component/conversation_manager.py index d92f99bbde..d859e7a12d 100644 --- a/pyrit/executor/attack/component/conversation_manager.py +++ b/pyrit/executor/attack/component/conversation_manager.py @@ -321,10 +321,22 @@ async def initialize_context_async( # single-string fallback path via capability-based routing. is_chat_target = target.configuration.includes(capability=CapabilityName.EDITABLE_HISTORY) if not is_chat_target: + config = prepended_conversation_config or PrependedConversationConfig() + if request_converters: + present_roles = { + piece.api_role for message in prepended_conversation for piece in message.message_pieces + } + excluded_roles = present_roles - set(config.apply_converters_to_roles) + if excluded_roles: + raise ValueError( + "Cannot preserve prepended-conversation converter role scoping for a non-chat target: " + f"the flattened context contains excluded roles {sorted(excluded_roles)}. " + "Use a chat target, remove request converters, or explicitly opt into every prepended role." + ) return await self._handle_non_chat_target_async( context=context, prepended_conversation=prepended_conversation, - config=prepended_conversation_config, + config=config, ) # Process prepended conversation for objective target @@ -453,10 +465,10 @@ async def add_prepended_conversation_to_memory_async( conversation=Conversation(conversation_id=conversation_id, target_identifier=target_identifier) ) - # Get roles that should have converters applied - apply_to_roles = ( - prepended_conversation_config.apply_converters_to_roles if prepended_conversation_config else None - ) + # Assistant history represents simulated target output, so the absent-config + # path must use the same safe role default as an explicit default config. + config = prepended_conversation_config or PrependedConversationConfig() + apply_to_roles = config.apply_converters_to_roles turn_count = 0 @@ -575,7 +587,7 @@ async def _apply_converters_async( *, message: Message, request_converters: list[ConverterConfiguration], - apply_to_roles: list[ChatMessageRole] | None, + apply_to_roles: list[ChatMessageRole], ) -> None: """ Apply converters to message pieces. @@ -583,12 +595,10 @@ async def _apply_converters_async( Args: message: The message containing pieces to convert. request_converters: Converter configurations to apply. - apply_to_roles: If provided, only apply to pieces with these roles. - If None, apply to all roles. + apply_to_roles: Only apply to pieces with these roles. """ for piece in message.message_pieces: - # Filter by role if specified - if apply_to_roles is not None and piece.api_role not in apply_to_roles: + if piece.api_role not in apply_to_roles: continue temp_message = Message(message_pieces=[piece]) diff --git a/pyrit/executor/attack/component/prepended_conversation_config.py b/pyrit/executor/attack/component/prepended_conversation_config.py index a0daedfd6e..7fc85688bb 100644 --- a/pyrit/executor/attack/component/prepended_conversation_config.py +++ b/pyrit/executor/attack/component/prepended_conversation_config.py @@ -4,13 +4,15 @@ from __future__ import annotations from dataclasses import dataclass, field -from typing import get_args +from typing import TYPE_CHECKING from pyrit.message_normalizer import ( ConversationContextNormalizer, MessageStringNormalizer, ) -from pyrit.models import ChatMessageRole + +if TYPE_CHECKING: + from pyrit.models import ChatMessageRole @dataclass @@ -27,10 +29,9 @@ class PrependedConversationConfig: first turn (via ``message_normalizer``; default: ConversationContextNormalizer). """ - # Roles for which request converters should be applied to prepended messages. - # By default, converters are applied to all roles. - # Example: ["user"] to apply converters only to user messages. - apply_converters_to_roles: list[ChatMessageRole] = field(default_factory=lambda: list(get_args(ChatMessageRole))) + # Request converters default to prepended user messages only. Assistant history is + # simulated target output and must be explicitly opted in with ["assistant"]. + apply_converters_to_roles: list[ChatMessageRole] = field(default_factory=lambda: ["user"]) # Optional normalizer to format conversation history into a single text block. # Must implement MessageStringNormalizer (e.g., TokenizerTemplateNormalizer or ConversationContextNormalizer). diff --git a/pyrit/executor/attack/single_turn/prompt_sending.py b/pyrit/executor/attack/single_turn/prompt_sending.py index 508b72d924..ffebec0f2a 100644 --- a/pyrit/executor/attack/single_turn/prompt_sending.py +++ b/pyrit/executor/attack/single_turn/prompt_sending.py @@ -77,6 +77,8 @@ def __init__( prepended_conversation_config (PrependedConversationConfiguration | None): Configuration for how to process prepended conversations. Controls converter application by role, message normalization, and non-chat target behavior. + Request converters apply to prepended user messages by default; include + ``"assistant"`` explicitly to transform simulated assistant history. Raises: ValueError: If the objective scorer is not a true/false scorer. diff --git a/pyrit/scenario/core/attack_technique_factory.py b/pyrit/scenario/core/attack_technique_factory.py index 1d168a0545..c74114c770 100644 --- a/pyrit/scenario/core/attack_technique_factory.py +++ b/pyrit/scenario/core/attack_technique_factory.py @@ -82,6 +82,7 @@ def __init__( adversarial_seed_prompt: SeedPrompt | str | None = None, seed_technique: AttackTechniqueSeedGroup | None = None, uses_adversarial: bool | None = None, + supports_request_converter_composition: bool = False, scorer_override_policy: ScorerOverridePolicy = ScorerOverridePolicy.WARN, ) -> None: """ @@ -119,6 +120,9 @@ def __init__( chat during execution. ``None`` auto-derives from the attack class constructor signature and seed-technique shape. Authors can override the derivation explicitly. + supports_request_converter_composition: Whether callers may safely + append request converters to this technique. This is an explicit + semantic opt-in, not merely constructor-signature detection. scorer_override_policy: What to do when a scenario's scorer is incompatible with the attack's ``attack_scoring_config`` type annotation. Defaults to WARN. @@ -143,11 +147,13 @@ class constructor signature and seed-technique shape. adversarial_system_prompt is not None or adversarial_seed_prompt is not None ) self._seed_technique = seed_technique + self._supports_request_converter_composition = supports_request_converter_composition self._scorer_override_policy = scorer_override_policy self._uses_adversarial = uses_adversarial if uses_adversarial is not None else self._derive_uses_adversarial() self._validate_kwargs() + self._validate_converter_composition() self._validate_adversarial_flags() @classmethod @@ -166,6 +172,7 @@ def with_simulated_conversation( attack_kwargs: dict[str, Any] | None = None, adversarial_chat: PromptTarget | None = None, uses_adversarial: bool | None = None, + supports_request_converter_composition: bool = False, scorer_override_policy: ScorerOverridePolicy = ScorerOverridePolicy.WARN, ) -> AttackTechniqueFactory: """ @@ -217,6 +224,9 @@ def with_simulated_conversation( during execution. ``None`` auto-derives from the attack class constructor signature and seed-technique shape. Forwarded to the factory constructor. + supports_request_converter_composition: Whether callers may safely + append request converters to this technique. Forwarded to the + factory constructor. scorer_override_policy: Policy applied when a scenario's scorer is incompatible with the attack's ``attack_scoring_config`` type annotation. Defaults to ``WARN``. Forwarded to the factory @@ -277,6 +287,7 @@ def with_simulated_conversation( adversarial_chat=adversarial_chat, seed_technique=seed_technique, uses_adversarial=uses_adversarial, + supports_request_converter_composition=supports_request_converter_composition, scorer_override_policy=scorer_override_policy, ) @@ -312,6 +323,23 @@ def _validate_adversarial_flags(self) -> None: f"should not have one wired." ) + def _validate_converter_composition(self) -> None: + """ + Validate that an opt-in factory can receive additive request converters. + + Raises: + ValueError: If composition is enabled but the attack constructor does + not accept ``attack_converter_config``. + """ + if ( + self._supports_request_converter_composition + and "attack_converter_config" not in self._get_accepted_params() + ): + raise ValueError( + f"Factory '{self._name}' declares supports_request_converter_composition=True, " + f"but {self._attack_class.__name__} does not accept 'attack_converter_config'." + ) + def _validate_kwargs(self) -> None: """ Validate that all kwargs are valid parameters for the attack class constructor. @@ -434,6 +462,11 @@ def uses_adversarial(self) -> bool: """Whether this technique drives an adversarial chat during execution.""" return self._uses_adversarial + @property + def supports_request_converter_composition(self) -> bool: + """Whether callers may safely append request converters to this technique.""" + return self._supports_request_converter_composition + @property def scoring_config_type(self) -> type | None: """The required ``attack_scoring_config`` subtype, or ``None`` if any config is accepted.""" @@ -774,6 +807,7 @@ def _build_identifier(self) -> ComponentIdentifier: "attack_class": self._attack_class.__name__, "kwargs": kwargs_for_id, "uses_adversarial": self._uses_adversarial, + "supports_request_converter_composition": self._supports_request_converter_composition, } if self._technique_tags: params["technique_tags"] = list(self._technique_tags) diff --git a/pyrit/scenario/scenarios/airt/jailbreak.py b/pyrit/scenario/scenarios/airt/jailbreak.py index 710c7eba7c..aa744a4864 100644 --- a/pyrit/scenario/scenarios/airt/jailbreak.py +++ b/pyrit/scenario/scenarios/airt/jailbreak.py @@ -72,6 +72,7 @@ def _prompt_sending_factory() -> AttackTechniqueFactory: name=_PROMPT_SENDING, attack_class=PromptSendingAttack, technique_tags=["single_turn"], + supports_request_converter_composition=True, ) @@ -93,6 +94,7 @@ def _jailbreak_system_prompt_factory() -> AttackTechniqueFactory: name=_JAILBREAK_SYSTEM_PROMPT, attack_class=PromptSendingAttack, technique_tags=["single_turn"], + supports_request_converter_composition=True, ) @@ -104,23 +106,34 @@ def _extra_default_factories() -> dict[str, AttackTechniqueFactory]: } +def _is_jailbreak_compatible_factory(factory: AttackTechniqueFactory) -> bool: + """Return whether a factory supports direct jailbreak-converter delivery.""" + has_simulated_conversation = ( + factory.seed_technique is not None and factory.seed_technique.has_simulated_conversation + ) + return ( + "multi_turn" not in factory.technique_tags + and not has_simulated_conversation + and factory.supports_request_converter_composition + ) + + @cache def _build_jailbreak_technique() -> type[ScenarioTechnique]: """ - Build the Jailbreak technique class dynamically from every registered factory plus the - scenario-local defaults. + Build the Jailbreak technique class from compatible direct factories and local deliveries. - The technique axis is the set of *attack techniques* a jailbreak is delivered through: the two - default deliveries (``prompt_sending`` and ``jailbreak_system_prompt``) plus whatever techniques - are registered (``role_play_*``, ``many_shot``, ``tap``, …). Jailbreak templates are a separate - selector (``num_jailbreaks`` / ``jailbreak_names``), so only the two deliveries are on by default - — crossing every template with every registered technique explodes quickly. + Registered multi-turn techniques, simulated-conversation seed techniques, and factories that + have not explicitly opted into additive request-converter composition are excluded. Returns: type[ScenarioTechnique]: The dynamically generated technique enum class. """ registry = AttackTechniqueRegistry.get_registry_singleton() - factories = list(registry.get_factories_or_raise().values()) + list(_extra_default_factories().values()) + registered = [ + factory for factory in registry.get_factories_or_raise().values() if _is_jailbreak_compatible_factory(factory) + ] + factories = registered + list(_extra_default_factories().values()) return AttackTechniqueRegistry.build_technique_class_from_factories( # type: ignore[return-value, ty:invalid-return-type] class_name="JailbreakTechnique", factories=factories, @@ -136,24 +149,25 @@ class Jailbreak(Scenario): selectors: - **dataset** — the harmful objectives (HarmBench). - - **techniques** — the *attack techniques* each jailbreak is delivered through. Two deliveries + - **techniques** — compatible direct deliveries for each jailbreak. Two deliveries are on by default: ``prompt_sending`` (the template rendered inline into the user message) and ``jailbreak_system_prompt`` (the template set as the system prompt with the objective sent as - the user turn). The registry techniques (``role_play_*``, ``many_shot``, ``tap``, …) are - opt-in. + the user turn). Registered direct techniques are opt-in only when their factory explicitly + supports request-converter composition. Multi-turn and simulated-conversation techniques are + not offered. - **jailbreaks** — which jailbreak templates to run (a random ``num_jailbreaks`` sample or an explicit ``jailbreak_names`` set). ``prompt_sending`` applies each template as a ``TextJailbreakConverter`` on the outgoing request, so the objective is rendered inline into the template's ``{{prompt}}`` slot; this keeps that - delivery target-agnostic and lets it compose with every technique. ``jailbreak_system_prompt`` + delivery target-agnostic and lets it compose with compatible direct techniques. ``jailbreak_system_prompt`` instead sets the template as a native system prompt and sends the objective as its own user turn, so it is only built for targets that natively support editable history and system prompts (it is skipped for incapable targets, or raises if it is the only selected technique). Responses are scored to determine whether the jailbreak succeeded (non-refusal). """ - VERSION: int = 3 + VERSION: int = 4 #: Baseline (an un-jailbroken prompt-send over the objectives) is included by default: a model #: that complies with the bare objective is itself interesting signal. Callers opt out per run @@ -232,6 +246,30 @@ def __init__( scenario_result_id=scenario_result_id, ) + def _resolve_scenario_techniques(self, *, scenario_techniques: Any) -> list[ScenarioTechnique]: + """ + Resolve techniques while rejecting stale or incompatible enum members. + + Args: + scenario_techniques (Any): Requested Jailbreak technique members. + + Returns: + list[ScenarioTechnique]: Compatible concrete techniques. + + Raises: + ValueError: If a caller supplies members from an older or different + technique enum. + """ + if scenario_techniques: + incompatible = [item for item in scenario_techniques if not isinstance(item, self._technique_class)] + if incompatible: + values = [getattr(item, "value", repr(item)) for item in incompatible] + raise ValueError( + "Jailbreak received stale or incompatible techniques " + f"{values}. Select a compatible direct-delivery technique from JailbreakTechnique." + ) + return super()._resolve_scenario_techniques(scenario_techniques=scenario_techniques) + def _resolve_templates(self) -> list[str]: """ Resolve the jailbreak templates for this run, replaying the persisted set on resume. @@ -286,10 +324,10 @@ async def _build_atomic_attacks_async(self, *, context: ScenarioContext) -> list """ Build one atomic attack per (technique x jailbreak template x dataset x attempt). - ``prompt_sending`` (and any opt-in registry techniques) deliver each jailbreak template as a + ``prompt_sending`` (and compatible opt-in direct techniques) deliver each jailbreak template as a ``TextJailbreakConverter`` appended to that technique's request converters, so the objective - is rendered inline into the template's ``{{prompt}}`` slot on the wire — target-agnostic and - composable with every technique. ``jailbreak_system_prompt`` instead delivers the template as + is rendered inline into the template's ``{{prompt}}`` slot on the wire. ``jailbreak_system_prompt`` + instead delivers the template as a native system prompt (no converter) with the objective sent as its own user turn, so it is only built when the objective target natively supports editable history and system prompts. Results group by jailbreak template so per-template ASR rolls up naturally. @@ -314,6 +352,21 @@ async def _build_atomic_attacks_async(self, *, context: ScenarioContext) -> list num_attempts = self.params.get("num_jailbreak_attempts", 1) technique_factories = resolve_technique_factories(context=context, extra_factories=_extra_default_factories()) + selected_names = {technique.value for technique in context.scenario_techniques} + missing = selected_names - set(technique_factories) + if missing: + raise ValueError( + "Jailbreak selected techniques that are no longer available: " + f"{sorted(missing)}. Refresh the plan and select compatible direct-delivery techniques." + ) + incompatible = [ + name for name, factory in technique_factories.items() if not _is_jailbreak_compatible_factory(factory) + ] + if incompatible: + raise ValueError( + "Jailbreak cannot compose with multi-turn, simulated-conversation, or " + f"non-composable techniques: {sorted(incompatible)}." + ) # ``jailbreak_system_prompt`` is delivered separately (native system prompt, no converter); # every other technique goes through the inline converter path. @@ -456,6 +509,7 @@ def _build_system_prompt_factory(self, *, template_file_name: str) -> AttackTech attack_class=PromptSendingAttack, technique_tags=["single_turn"], seed_technique=seed_technique, + supports_request_converter_composition=True, ) @staticmethod diff --git a/pyrit/setup/initializers/techniques/airt.py b/pyrit/setup/initializers/techniques/airt.py index 1469ea6e19..5b3c92cec5 100644 --- a/pyrit/setup/initializers/techniques/airt.py +++ b/pyrit/setup/initializers/techniques/airt.py @@ -42,6 +42,7 @@ def get_technique_factories() -> list[AttackTechniqueFactory]: attack_class=PromptSendingAttack, description="Obfuscates the objective by asking for it encoded as the first letter of each word.", technique_tags=["single_turn", "airt", "leakage"], + supports_request_converter_composition=True, attack_kwargs={ "attack_converter_config": AttackConverterConfig( request_converters=ConverterConfiguration.from_converters(converters=[FirstLetterConverter()]) diff --git a/pyrit/setup/initializers/techniques/core.py b/pyrit/setup/initializers/techniques/core.py index b4be579675..e696429af7 100644 --- a/pyrit/setup/initializers/techniques/core.py +++ b/pyrit/setup/initializers/techniques/core.py @@ -159,6 +159,7 @@ def get_technique_factories() -> list[AttackTechniqueFactory]: attack_class=PromptSendingAttack, description="Reverses the objective text so it slips past filters, then asks the target to flip it back.", technique_tags=["single_turn", "light"], + supports_request_converter_composition=True, attack_kwargs={ "attack_converter_config": AttackConverterConfig( request_converters=ConverterConfiguration.from_converters( diff --git a/tests/unit/executor/attack/component/test_conversation_manager.py b/tests/unit/executor/attack/component/test_conversation_manager.py index 550e4b631d..2a25984ece 100644 --- a/tests/unit/executor/attack/component/test_conversation_manager.py +++ b/tests/unit/executor/attack/component/test_conversation_manager.py @@ -23,6 +23,7 @@ import pytest from unit.mocks import get_mock_scorer_identifier +from pyrit.converter import Base64Converter from pyrit.executor.attack import ConversationManager, ConversationState from pyrit.executor.attack.component import PrependedConversationConfig from pyrit.executor.attack.component.conversation_manager import ( @@ -1093,6 +1094,25 @@ async def test_non_chat_target_behavior_normalize_is_default( text_value = context.next_message.get_piece().original_value assert len(text_value) > 0 + async def test_non_chat_target_rejects_converter_scoping_that_excludes_history_roles( + self, + attack_identifier: ComponentIdentifier, + mock_prompt_target: MagicMock, + sample_conversation: list[Message], + ) -> None: + manager = ConversationManager() + context = _TestAttackContext(params=AttackParameters(objective="Test objective")) + context.prepended_conversation = sample_conversation + converter_config = ConverterConfiguration.from_converters(converters=[Base64Converter()]) + + with pytest.raises(ValueError, match="non-chat target.*excluded roles.*assistant"): + await manager.initialize_context_async( + context=context, + target=mock_prompt_target, + conversation_id=str(uuid.uuid4()), + request_converters=converter_config, + ) + async def test_non_chat_target_behavior_normalize_first_turn_creates_next_message( self, attack_identifier: ComponentIdentifier, @@ -1178,13 +1198,13 @@ async def test_non_chat_target_behavior_normalize_returns_empty_state( # apply_converters_to_roles Tests # ------------------------------------------------------------------------- - async def test_apply_converters_to_roles_default_applies_to_all( + async def test_apply_converters_to_roles_default_applies_to_user_only( self, attack_identifier: ComponentIdentifier, mock_chat_target: MagicMock, sample_conversation: list[Message], ) -> None: - """Test that converters are applied to all roles by default.""" + """Test that converters are applied only to user history by default.""" mock_normalizer = MagicMock(spec=PromptNormalizer) mock_normalizer.convert_values_async = AsyncMock() manager = ConversationManager(prompt_normalizer=mock_normalizer) @@ -1201,8 +1221,7 @@ async def test_apply_converters_to_roles_default_applies_to_all( request_converters=converter_config, ) - # convert_values_async should be called for each message (both user and assistant) - assert mock_normalizer.convert_values_async.call_count == 2 + mock_normalizer.convert_values_async.assert_awaited_once() async def test_apply_converters_to_roles_user_only( self, diff --git a/tests/unit/executor/attack/component/test_prepended_conversation_config.py b/tests/unit/executor/attack/component/test_prepended_conversation_config.py index b1c34a6770..546a7f1dbf 100644 --- a/tests/unit/executor/attack/component/test_prepended_conversation_config.py +++ b/tests/unit/executor/attack/component/test_prepended_conversation_config.py @@ -1,17 +1,15 @@ # Copyright (c) Microsoft Corporation. # Licensed under the MIT license. -from typing import get_args from unittest.mock import MagicMock from pyrit.executor.attack.component.prepended_conversation_config import PrependedConversationConfig from pyrit.message_normalizer import ConversationContextNormalizer -from pyrit.models import ChatMessageRole -def test_default_init_apply_converters_to_all_roles(): +def test_default_init_apply_converters_to_user_role(): config = PrependedConversationConfig() - assert config.apply_converters_to_roles == list(get_args(ChatMessageRole)) + assert config.apply_converters_to_roles == ["user"] def test_default_init_message_normalizer_is_none(): diff --git a/tests/unit/executor/attack/single_turn/test_prompt_sending.py b/tests/unit/executor/attack/single_turn/test_prompt_sending.py index 5d4fda2209..7b938b5d35 100644 --- a/tests/unit/executor/attack/single_turn/test_prompt_sending.py +++ b/tests/unit/executor/attack/single_turn/test_prompt_sending.py @@ -5,16 +5,18 @@ from unittest.mock import AsyncMock, MagicMock, patch import pytest -from unit.mocks import get_mock_scorer_identifier, get_mock_target_identifier +from unit.mocks import MockPromptTarget, get_mock_scorer_identifier, get_mock_target_identifier from pyrit.converter import Base64Converter, StringJoinConverter from pyrit.executor.attack import ( AttackConverterConfig, AttackParameters, AttackScoringConfig, + PrependedConversationConfig, PromptSendingAttack, SingleTurnAttackContext, ) +from pyrit.memory import CentralMemory from pyrit.models import ( AttackOutcome, AttackResult, @@ -280,6 +282,60 @@ async def test_setup_updates_conversation_state_with_converters(self, mock_targe memory_labels={}, ) + async def test_default_converter_scoping_preserves_simulated_assistant_history(self): + target = MockPromptTarget() + converter_config = AttackConverterConfig( + request_converters=ConverterConfiguration.from_converters(converters=[Base64Converter()]) + ) + attack = PromptSendingAttack(objective_target=target, attack_converter_config=converter_config) + prepended_user = "prepended user request" + simulated_response = "simulated assistant response" + final_request = "live final request" + + result = await attack.execute_async( + objective="Test objective", + prepended_conversation=[ + Message.from_prompt(prompt=prepended_user, role="user"), + Message.from_prompt(prompt=simulated_response, role="assistant"), + ], + next_message=Message.from_prompt(prompt=final_request, role="user"), + ) + + pieces = CentralMemory.get_memory_instance().get_message_pieces(conversation_id=result.conversation_id) + assistant_piece = next(piece for piece in pieces if piece.original_value == simulated_response) + final_piece = next(piece for piece in pieces if piece.original_value == final_request) + + assert assistant_piece.role == "simulated_assistant" + assert assistant_piece.converted_value == assistant_piece.original_value + assert assistant_piece.converter_identifiers == [] + assert final_piece.converted_value != final_piece.original_value + assert [identifier.class_name for identifier in final_piece.converter_identifiers] == ["Base64Converter"] + + async def test_explicit_assistant_role_opt_in_converts_simulated_history(self): + target = MockPromptTarget() + converter_config = AttackConverterConfig( + request_converters=ConverterConfiguration.from_converters(converters=[Base64Converter()]) + ) + attack = PromptSendingAttack( + objective_target=target, + attack_converter_config=converter_config, + prepended_conversation_config=PrependedConversationConfig(apply_converters_to_roles=["assistant"]), + ) + simulated_response = "assistant history explicitly converted" + + result = await attack.execute_async( + objective="Test objective", + prepended_conversation=[Message.from_prompt(prompt=simulated_response, role="assistant")], + next_message=Message.from_prompt(prompt="live request", role="user"), + ) + + pieces = CentralMemory.get_memory_instance().get_message_pieces(conversation_id=result.conversation_id) + assistant_piece = next(piece for piece in pieces if piece.original_value == simulated_response) + + assert assistant_piece.role == "simulated_assistant" + assert assistant_piece.converted_value != assistant_piece.original_value + assert [identifier.class_name for identifier in assistant_piece.converter_identifiers] == ["Base64Converter"] + @pytest.mark.usefixtures("patch_central_database") class TestPromptPreparation: diff --git a/tests/unit/scenario/airt/test_jailbreak.py b/tests/unit/scenario/airt/test_jailbreak.py index 5cce529015..3a8beb3bfb 100644 --- a/tests/unit/scenario/airt/test_jailbreak.py +++ b/tests/unit/scenario/airt/test_jailbreak.py @@ -16,6 +16,7 @@ from pyrit.prompt_target import PromptTarget from pyrit.registry import TargetRegistry from pyrit.registry.components.attack_technique_registry import AttackTechniqueRegistry +from pyrit.registry.components.scenario_registry import ScenarioRegistry from pyrit.scenario.core import BaselineAttackPolicy from pyrit.scenario.core.attack_technique_factory import AttackTechniqueFactory from pyrit.scenario.scenarios.airt.jailbreak import ( @@ -328,9 +329,7 @@ def _spy_create(self, **kwargs): assert captured, "Expected factory.create to be called" converters = [c for extra in captured if extra for cc in extra for c in cc.converters] - assert any(isinstance(c, TextJailbreakConverter) for c in converters), ( - "Expected a TextJailbreakConverter to be threaded to factory.create" - ) + assert sum(isinstance(c, TextJailbreakConverter) for c in converters) == 1 # The objective seed groups themselves carry no jailbreak framing (converter delivery only). for attack in scenario._atomic_attacks: @@ -376,49 +375,61 @@ def _spy_create(self, **kwargs): "Jailbreak converter must be applied before caller-supplied converters" ) - async def test_simulated_conversation_techniques_produce_attacks_with_jailbreak( + async def test_stale_incompatible_technique_is_rejected( self, mock_objective_target, mock_objective_scorer, mock_memory_seed_groups ): - """Regression: simulated-conversation techniques (``role_play_*``, ``crescendo_*``) must still - produce atomic attacks when crossed with a jailbreak template, and each must receive the - jailbreak converter. - - Converter delivery leaves the objective seed group unframed, so it stays compatible with the - simulated-conversation seed technique. (Delivering the jailbreak as a system-role framing seed - instead collided with that technique's seed range and silently produced zero attacks.) - """ - technique_class = _build_jailbreak_technique() - techniques = [ - technique_class("role_play_movie_script"), - technique_class("crescendo_simulated"), - ] - captured: list[Any] = [] - original_create = AttackTechniqueFactory.create - - def _spy_create(self, **kwargs): - captured.append(kwargs.get("extra_request_converters")) - return original_create(self, **kwargs) + registry_factories = list(AttackTechniqueRegistry.get_registry_singleton().get_factories_or_raise().values()) + legacy_class = AttackTechniqueRegistry.build_technique_class_from_factories( + class_name="LegacyJailbreakTechnique", + factories=registry_factories, + ) + with _patch_seed_groups(mock_memory_seed_groups): + scenario = Jailbreak(objective_scorer=mock_objective_scorer) + scenario.set_params_from_args( + args=_default_args( + mock_objective_target, + scenario_techniques=[legacy_class("tap")], + jailbreak_names=["aim.yaml"], + ) + ) + with pytest.raises(ValueError, match="stale or incompatible"): + await scenario.initialize_async() + async def test_incompatible_runtime_factory_is_rejected( + self, mock_objective_target, mock_objective_scorer, mock_memory_seed_groups + ): + incompatible = AttackTechniqueFactory( + name="legacy_multi_turn", + attack_class=PromptSendingAttack, + technique_tags=["multi_turn"], + supports_request_converter_composition=True, + ) with _patch_seed_groups(mock_memory_seed_groups): - with patch.object(AttackTechniqueFactory, "create", _spy_create): + with patch( + "pyrit.scenario.scenarios.airt.jailbreak.resolve_technique_factories", + return_value={_PROMPT_SENDING: incompatible}, + ): scenario = Jailbreak(objective_scorer=mock_objective_scorer) + technique_class = _build_jailbreak_technique() scenario.set_params_from_args( args=_default_args( - mock_objective_target, scenario_techniques=techniques, jailbreak_names=["aim.yaml"] + mock_objective_target, + scenario_techniques=[technique_class(_PROMPT_SENDING)], + jailbreak_names=["aim.yaml"], ) ) - await scenario.initialize_async() - names = {a.atomic_attack_name for a in scenario._atomic_attacks} - assert "role_play_movie_script_aim_harmbench" in names - assert "crescendo_simulated_aim_harmbench" in names - # Every build of a simulated-conversation technique must still carry the jailbreak - # converter. Assert both techniques captured a non-empty converter stack (so the check - # can't pass vacuously on a dropped/None stack) and each contains the jailbreak converter. - populated = [extra for extra in captured if extra] - assert len(populated) == 2, "Expected both simulated-conversation techniques to receive converters" - assert all( - any(isinstance(c, TextJailbreakConverter) for cc in extra for c in cc.converters) for extra in populated - ), "Each simulated-conversation technique must receive the jailbreak converter" + with pytest.raises(ValueError, match="cannot compose"): + await scenario.initialize_async() + + async def test_missing_runtime_factory_is_rejected( + self, mock_objective_target, mock_objective_scorer, mock_memory_seed_groups + ): + with _patch_seed_groups(mock_memory_seed_groups): + with patch("pyrit.scenario.scenarios.airt.jailbreak.resolve_technique_factories", return_value={}): + scenario = Jailbreak(objective_scorer=mock_objective_scorer) + scenario.set_params_from_args(args=_default_args(mock_objective_target, jailbreak_names=["aim.yaml"])) + with pytest.raises(ValueError, match="no longer available.*prompt_sending"): + await scenario.initialize_async() async def test_all_templates_produce_attacks( self, mock_objective_target, mock_objective_scorer, mock_memory_seed_groups @@ -708,15 +719,41 @@ def test_default_techniques_are_the_two_deliveries(self): assert default_values == set(_DEFAULT_TECHNIQUES) assert default_values == {_PROMPT_SENDING, _JAILBREAK_SYSTEM_PROMPT} - def test_registry_techniques_are_available(self): + def test_only_compatible_direct_registry_techniques_are_available(self): technique_class = _technique_class() available = {t.value for t in technique_class.get_all_techniques()} - assert {_PROMPT_SENDING, _JAILBREAK_SYSTEM_PROMPT}.issubset(available) - # The "normal ones available like from rapid response" are exposed as opt-in techniques. - assert {"role_play_movie_script", "many_shot", "tap"}.issubset(available) - - def test_scenario_version_is_three(self): - assert Jailbreak.VERSION == 3 + incompatible = { + "context_compliance", + "role_play_movie_script", + "role_play_video_game", + "role_play_trivia_game", + "role_play_persuasion", + "role_play_persuasion_written", + "crescendo_simulated", + "crescendo_movie_director", + "crescendo_history_lecture", + "crescendo_journalist_interview", + "red_teaming", + "tap", + "many_shot", + "pair", + } + assert {_PROMPT_SENDING, _JAILBREAK_SYSTEM_PROMPT, "flip"}.issubset(available) + assert incompatible.isdisjoint(available) + + def test_registry_metadata_omits_incompatible_techniques(self): + metadata = ScenarioRegistry()._build_metadata("airt.jailbreak", Jailbreak) + assert { + "context_compliance", + "role_play_movie_script", + "crescendo_simulated", + "red_teaming", + "tap", + "many_shot", + }.isdisjoint(metadata.all_techniques) + + def test_scenario_version_is_four(self): + assert Jailbreak.VERSION == 4 def test_default_dataset_is_harmbench(self): assert Jailbreak.required_datasets() == ["harmbench"] diff --git a/tests/unit/scenario/core/test_attack_technique_factory.py b/tests/unit/scenario/core/test_attack_technique_factory.py index 39e3046685..a52773c301 100644 --- a/tests/unit/scenario/core/test_attack_technique_factory.py +++ b/tests/unit/scenario/core/test_attack_technique_factory.py @@ -178,6 +178,29 @@ def test_validate_kwargs_rejects_invalid_param_on_real_attack_class(self): attack_kwargs={"nonexistent_param": 42}, ) + def test_request_converter_composition_requires_supported_constructor(self): + class _NoConverterAttack: + def __init__(self, *, objective_target, attack_scoring_config=None): + self.objective_target = objective_target + + with pytest.raises(ValueError, match="does not accept 'attack_converter_config'"): + AttackTechniqueFactory( + name="test", + attack_class=_NoConverterAttack, + supports_request_converter_composition=True, + ) + + def test_request_converter_composition_is_explicit_opt_in(self): + default_factory = AttackTechniqueFactory(name="default", attack_class=_StubAttack) + composable_factory = AttackTechniqueFactory( + name="composable", + attack_class=_StubAttack, + supports_request_converter_composition=True, + ) + + assert not default_factory.supports_request_converter_composition + assert composable_factory.supports_request_converter_composition + class TestFactoryCreate: """Tests for AttackTechniqueFactory.create().""" From 5d70865b500034d347487112cedab10a4a218ac8 Mon Sep 17 00:00:00 2001 From: Copilot <223556219+Copilot@users.noreply.github.com> Date: Fri, 7 Aug 2026 21:40:11 -0700 Subject: [PATCH 02/24] Fix non-chat converter role scoping Apply request converters to role-separated prepended history before flattening, while preventing the resulting request from being converted twice. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> Copilot-Session: 6776aa95-0924-4967-a056-e5a228af3faa --- .../attack/component/conversation_manager.py | 223 +++++++++++++----- .../prepended_conversation_config.py | 5 +- pyrit/models/messages/message.py | 17 +- pyrit/prompt_normalizer/prompt_normalizer.py | 9 +- .../component/test_conversation_manager.py | 57 ++++- .../attack/single_turn/test_prompt_sending.py | 29 +++ 6 files changed, 264 insertions(+), 76 deletions(-) diff --git a/pyrit/executor/attack/component/conversation_manager.py b/pyrit/executor/attack/component/conversation_manager.py index d859e7a12d..60d25904ef 100644 --- a/pyrit/executor/attack/component/conversation_manager.py +++ b/pyrit/executor/attack/component/conversation_manager.py @@ -13,7 +13,11 @@ PrependedConversationConfig, ) from pyrit.memory import CentralMemory -from pyrit.message_normalizer import ConversationContextNormalizer, GenericSystemSquashNormalizer +from pyrit.message_normalizer import ( + ConversationContextNormalizer, + GenericSystemSquashNormalizer, + MessageStringNormalizer, +) from pyrit.models import ( ChatMessageRole, ComponentIdentifier, @@ -284,8 +288,9 @@ async def initialize_context_async( - All messages get new UUIDs For non-chat PromptTarget: + - Applies request converters to configured prepended roles before normalization - Normalizes the prepended conversation to a string and prepends it to - ``context.next_message`` (using ``config.message_normalizer`` when provided). + ``context.next_message`` (using ``config.message_normalizer`` when provided) Args: context: The attack context to initialize. @@ -300,8 +305,7 @@ async def initialize_context_async( ConversationState with turn_count and last_assistant_message_scores. Raises: - ValueError: If conversation_id is empty, or if prepended_conversation - requires a chat-capable PromptTarget but target is not one. + ValueError: If conversation_id is empty. """ if not conversation_id: raise ValueError("conversation_id cannot be empty") @@ -322,21 +326,11 @@ async def initialize_context_async( is_chat_target = target.configuration.includes(capability=CapabilityName.EDITABLE_HISTORY) if not is_chat_target: config = prepended_conversation_config or PrependedConversationConfig() - if request_converters: - present_roles = { - piece.api_role for message in prepended_conversation for piece in message.message_pieces - } - excluded_roles = present_roles - set(config.apply_converters_to_roles) - if excluded_roles: - raise ValueError( - "Cannot preserve prepended-conversation converter role scoping for a non-chat target: " - f"the flattened context contains excluded roles {sorted(excluded_roles)}. " - "Use a chat target, remove request converters, or explicitly opt into every prepended role." - ) return await self._handle_non_chat_target_async( context=context, prepended_conversation=prepended_conversation, config=config, + request_converters=request_converters, ) # Process prepended conversation for objective target @@ -355,7 +349,8 @@ async def _handle_non_chat_target_async( *, context: AttackContext[Any], prepended_conversation: list[Message], - config: PrependedConversationConfig | None, + config: PrependedConversationConfig, + request_converters: list[ConverterConfiguration] | None, ) -> ConversationState: """ Handle prepended conversation for non-chat targets. @@ -364,60 +359,164 @@ async def _handle_non_chat_target_async( context: The attack context. prepended_conversation: Messages to prepend. config: Configuration for non-chat target behavior. + request_converters: Converters to apply before flattening. Returns: Empty ConversationState (non-chat targets don't track turns). """ - if config is None: - config = PrependedConversationConfig() - normalizer = config.get_message_normalizer() - messages_to_normalize = prepended_conversation - if isinstance(normalizer, ConversationContextNormalizer): - messages_to_normalize = await GenericSystemSquashNormalizer().normalize_async(prepended_conversation) - - normalized_context = await normalizer.normalize_string_async(messages_to_normalize) - - next_message = context.next_message - if next_message is None: - next_message = Message.from_prompt(prompt=context.objective, role="user") - context.next_message = next_message - - if normalized_context: - # Find an existing text piece to prepend to - text_piece = None - for piece in next_message.message_pieces: - if piece.original_value_data_type == "text": - text_piece = piece - break - - if text_piece: - # Prepend context to the existing text piece - context_prefix = f"{normalized_context}\n\n" - if text_piece.original_value != normalized_context and not text_piece.original_value.startswith( - context_prefix - ): - text_piece.original_value = f"{context_prefix}{text_piece.original_value}" - if text_piece.converted_value != normalized_context and not text_piece.converted_value.startswith( - context_prefix - ): - text_piece.converted_value = f"{context_prefix}{text_piece.converted_value}" - else: - # No text piece found (multimodal message), add a new text piece at the beginning - context_piece = MessagePiece( - id=uuid.uuid4(), - role="user", - original_value=normalized_context, - converted_value=normalized_context, - original_value_data_type="text", - converted_value_data_type="text", - ) - # Create a new message with the context piece prepended - context.next_message = Message(message_pieces=[context_piece] + list(next_message.message_pieces)) + original_context = await self._normalize_non_chat_context_async( + messages=self._build_original_normalizer_view(prepended_conversation), + normalizer=normalizer, + ) + converted_context = original_context + if request_converters: + converted_messages = await self._build_converted_normalizer_view_async( + messages=prepended_conversation, + request_converters=request_converters, + apply_to_roles=config.apply_converters_to_roles, + ) + converted_context = await self._normalize_non_chat_context_async( + messages=converted_messages, + normalizer=normalizer, + ) + + next_message = ( + context.next_message.duplicate() + if context.next_message + else Message.from_prompt( + prompt=context.objective, + role="user", + ) + ) + if request_converters and not next_message.request_converters_applied: + await self._prompt_normalizer.convert_values_async( + converter_configurations=request_converters, + message=next_message, + ) + next_message.mark_request_converters_applied() - logger.debug(f"Normalized prepended conversation for non-chat target: {len(normalized_context)} characters") + self._prepend_non_chat_context( + message=next_message, + original_context=original_context, + converted_context=converted_context, + ) + context.next_message = next_message + + logger.debug(f"Normalized prepended conversation for non-chat target: {len(converted_context)} characters") return ConversationState() + async def _build_converted_normalizer_view_async( + self, + *, + messages: list[Message], + request_converters: list[ConverterConfiguration] | None, + apply_to_roles: list[ChatMessageRole], + ) -> list[Message]: + """ + Build copies containing only the values that should be sent. + + Returns: + list[Message]: Converted message copies ready for string normalization. + """ + converted_messages = [message.duplicate() for message in messages] + if request_converters: + for message in converted_messages: + await self._apply_converters_async( + message=message, + request_converters=request_converters, + apply_to_roles=apply_to_roles, + ) + for message in converted_messages: + for piece in message.message_pieces: + piece.original_value = piece.converted_value + piece.original_value_data_type = piece.converted_value_data_type + return converted_messages + + @staticmethod + def _build_original_normalizer_view(messages: list[Message]) -> list[Message]: + """ + Build copies containing only the original values. + + Returns: + list[Message]: Message copies with converted fields reset to their originals. + """ + original_messages = [message.duplicate() for message in messages] + for message in original_messages: + for piece in message.message_pieces: + piece.converted_value = piece.original_value + piece.converted_value_data_type = piece.original_value_data_type + return original_messages + + @staticmethod + async def _normalize_non_chat_context_async( + *, + messages: list[Message], + normalizer: MessageStringNormalizer, + ) -> str: + """ + Flatten role-separated messages into a context string. + + Returns: + str: The flattened conversation context. + """ + messages_to_normalize = messages + if isinstance(normalizer, ConversationContextNormalizer): + messages_to_normalize = await GenericSystemSquashNormalizer().normalize_async(messages) + return await normalizer.normalize_string_async(messages_to_normalize) + + @staticmethod + def _prepend_non_chat_context( + *, + message: Message, + original_context: str, + converted_context: str, + ) -> None: + """Prepend original and converted context without mixing their values.""" + text_piece = next( + ( + piece + for piece in message.message_pieces + if piece.original_value_data_type == "text" and piece.converted_value_data_type == "text" + ), + None, + ) + if text_piece: + text_piece.original_value = ConversationManager._prepend_context_value( + context=original_context, + value=text_piece.original_value, + ) + text_piece.converted_value = ConversationManager._prepend_context_value( + context=converted_context, + value=text_piece.converted_value, + ) + return + + template_piece = message.get_piece() + context_piece = MessagePiece( + id=uuid.uuid4(), + role=template_piece.role, + original_value=original_context, + converted_value=converted_context, + original_value_data_type="text", + converted_value_data_type="text", + conversation_id=template_piece.conversation_id, + sequence=template_piece.sequence, + ) + message.message_pieces.insert(0, context_piece) + + @staticmethod + def _prepend_context_value(*, context: str, value: str) -> str: + """ + Prepend context once to a message value. + + Returns: + str: The value prefixed with context when it was not already present. + """ + if not context or value == context or value.startswith(f"{context}\n\n"): + return value + return f"{context}\n\n{value}" + async def add_prepended_conversation_to_memory_async( self, *, diff --git a/pyrit/executor/attack/component/prepended_conversation_config.py b/pyrit/executor/attack/component/prepended_conversation_config.py index 7fc85688bb..8930f829af 100644 --- a/pyrit/executor/attack/component/prepended_conversation_config.py +++ b/pyrit/executor/attack/component/prepended_conversation_config.py @@ -25,8 +25,9 @@ class PrependedConversationConfig: - Which message roles should have request converters applied - How to normalize conversation history for non-chat objective targets - Non-chat objective targets always normalize the prepended conversation into the - first turn (via ``message_normalizer``; default: ConversationContextNormalizer). + Non-chat objective targets apply request converters to the configured roles before + normalizing the prepended conversation into the first turn (via ``message_normalizer``; + default: ConversationContextNormalizer). """ # Request converters default to prepended user messages only. Assistant history is diff --git a/pyrit/models/messages/message.py b/pyrit/models/messages/message.py index 7e0a4e301b..b160820292 100644 --- a/pyrit/models/messages/message.py +++ b/pyrit/models/messages/message.py @@ -8,7 +8,7 @@ from datetime import datetime, timezone from typing import TYPE_CHECKING, cast -from pydantic import BaseModel, ConfigDict, model_validator +from pydantic import BaseModel, ConfigDict, PrivateAttr, model_validator from pyrit.models.messages.message_piece import MessagePiece @@ -30,6 +30,7 @@ class Message(BaseModel): ) message_pieces: list[MessagePiece] + _request_converters_applied: bool = PrivateAttr(default=False) # ------------------------------------------------------------------ # # Validators @@ -140,6 +141,15 @@ def get_piece(self, n: int = 0) -> MessagePiece: return self.message_pieces[n] + @property + def request_converters_applied(self) -> bool: + """Whether request converters have already been applied.""" + return self._request_converters_applied + + def mark_request_converters_applied(self) -> None: + """Mark this message as already processed by its request converters.""" + self._request_converters_applied = True + def get_pieces_by_type( self, *, @@ -371,4 +381,7 @@ def duplicate(self) -> Message: piece.id = uuid.uuid4() piece.timestamp = new_timestamp # original_prompt_id intentionally kept the same to track the origin - return Message(message_pieces=new_pieces) + duplicate = Message(message_pieces=new_pieces) + if self.request_converters_applied: + duplicate.mark_request_converters_applied() + return duplicate diff --git a/pyrit/prompt_normalizer/prompt_normalizer.py b/pyrit/prompt_normalizer/prompt_normalizer.py index cf7b458f94..bc5507ad34 100644 --- a/pyrit/prompt_normalizer/prompt_normalizer.py +++ b/pyrit/prompt_normalizer/prompt_normalizer.py @@ -108,8 +108,13 @@ async def send_prompt_async( for piece in request.message_pieces: piece.conversation_id = conversation_id - # Apply request converters - await self.convert_values_async(converter_configurations=request_converter_configurations, message=request) + # A caller may need to apply converters before a lossy normalization step + # such as flattening role-separated history for a non-chat target. + if not request.request_converters_applied: + await self.convert_values_async( + converter_configurations=request_converter_configurations, + message=request, + ) await self._calc_hash_async(request=request) diff --git a/tests/unit/executor/attack/component/test_conversation_manager.py b/tests/unit/executor/attack/component/test_conversation_manager.py index 2a25984ece..4404efcee7 100644 --- a/tests/unit/executor/attack/component/test_conversation_manager.py +++ b/tests/unit/executor/attack/component/test_conversation_manager.py @@ -17,6 +17,7 @@ - get_prepended_turn_count: Counts assistant messages in a conversation """ +import base64 import uuid from unittest.mock import AsyncMock, MagicMock @@ -1094,7 +1095,7 @@ async def test_non_chat_target_behavior_normalize_is_default( text_value = context.next_message.get_piece().original_value assert len(text_value) > 0 - async def test_non_chat_target_rejects_converter_scoping_that_excludes_history_roles( + async def test_non_chat_target_converts_selected_roles_before_flattening( self, attack_identifier: ComponentIdentifier, mock_prompt_target: MagicMock, @@ -1103,15 +1104,55 @@ async def test_non_chat_target_rejects_converter_scoping_that_excludes_history_r manager = ConversationManager() context = _TestAttackContext(params=AttackParameters(objective="Test objective")) context.prepended_conversation = sample_conversation + context.next_message = Message.from_prompt(prompt="live request", role="user") converter_config = ConverterConfiguration.from_converters(converters=[Base64Converter()]) - with pytest.raises(ValueError, match="non-chat target.*excluded roles.*assistant"): - await manager.initialize_context_async( - context=context, - target=mock_prompt_target, - conversation_id=str(uuid.uuid4()), - request_converters=converter_config, - ) + await manager.initialize_context_async( + context=context, + target=mock_prompt_target, + conversation_id=str(uuid.uuid4()), + request_converters=converter_config, + ) + + assert context.next_message is not None + piece = context.next_message.get_piece() + encoded_user = base64.b64encode(b"Hello, how are you?").decode() + encoded_live_request = base64.b64encode(b"live request").decode() + assert piece.original_value == ( + "Turn 1:\nuser: Hello, how are you?\nassistant: I'm doing well, thank you!\n\nlive request" + ) + assert piece.converted_value == ( + f"Turn 1:\nuser: {encoded_user}\nassistant: I'm doing well, thank you!\n\n{encoded_live_request}" + ) + assert context.next_message.request_converters_applied + + async def test_non_chat_target_converts_assistant_history_only_when_opted_in( + self, + attack_identifier: ComponentIdentifier, + mock_prompt_target: MagicMock, + sample_conversation: list[Message], + ) -> None: + manager = ConversationManager() + context = _TestAttackContext(params=AttackParameters(objective="Test objective")) + context.prepended_conversation = sample_conversation + context.next_message = Message.from_prompt(prompt="live request", role="user") + converter_config = ConverterConfiguration.from_converters(converters=[Base64Converter()]) + config = PrependedConversationConfig(apply_converters_to_roles=["assistant"]) + + await manager.initialize_context_async( + context=context, + target=mock_prompt_target, + conversation_id=str(uuid.uuid4()), + request_converters=converter_config, + prepended_conversation_config=config, + ) + + assert context.next_message is not None + encoded_assistant = base64.b64encode(b"I'm doing well, thank you!").decode() + encoded_live_request = base64.b64encode(b"live request").decode() + assert context.next_message.get_piece().converted_value == ( + f"Turn 1:\nuser: Hello, how are you?\nassistant: {encoded_assistant}\n\n{encoded_live_request}" + ) async def test_non_chat_target_behavior_normalize_first_turn_creates_next_message( self, diff --git a/tests/unit/executor/attack/single_turn/test_prompt_sending.py b/tests/unit/executor/attack/single_turn/test_prompt_sending.py index 7b938b5d35..76d0b8cf62 100644 --- a/tests/unit/executor/attack/single_turn/test_prompt_sending.py +++ b/tests/unit/executor/attack/single_turn/test_prompt_sending.py @@ -1,6 +1,7 @@ # Copyright (c) Microsoft Corporation. # Licensed under the MIT license. +import base64 import uuid from unittest.mock import AsyncMock, MagicMock, patch @@ -29,6 +30,8 @@ ) from pyrit.prompt_normalizer import ConverterConfiguration, PromptNormalizer from pyrit.prompt_target import PromptTarget +from pyrit.prompt_target.common.target_capabilities import TargetCapabilities +from pyrit.prompt_target.common.target_configuration import TargetConfiguration from pyrit.score import Scorer, TrueFalseScorer @@ -336,6 +339,32 @@ async def test_explicit_assistant_role_opt_in_converts_simulated_history(self): assert assistant_piece.converted_value != assistant_piece.original_value assert [identifier.class_name for identifier in assistant_piece.converter_identifiers] == ["Base64Converter"] + async def test_non_chat_target_converts_history_by_role_before_flattening(self): + target = MockPromptTarget() + target._configuration = TargetConfiguration(capabilities=TargetCapabilities()) + converter_config = AttackConverterConfig( + request_converters=ConverterConfiguration.from_converters(converters=[Base64Converter()]) + ) + attack = PromptSendingAttack(objective_target=target, attack_converter_config=converter_config) + prepended_user = "prepended user request" + simulated_response = "simulated assistant response" + final_request = "live final request" + + await attack.execute_async( + objective="Test objective", + prepended_conversation=[ + Message.from_prompt(prompt=prepended_user, role="user"), + Message.from_prompt(prompt=simulated_response, role="assistant"), + ], + next_message=Message.from_prompt(prompt=final_request, role="user"), + ) + + encoded_user = base64.b64encode(prepended_user.encode()).decode() + encoded_final_request = base64.b64encode(final_request.encode()).decode() + assert target.prompt_sent == [ + f"Turn 1:\nuser: {encoded_user}\nassistant: {simulated_response}\n\n{encoded_final_request}" + ] + @pytest.mark.usefixtures("patch_central_database") class TestPromptPreparation: From 5c759f7b85fcc1aedeb4554a44f180a3d79ed04f Mon Sep 17 00:00:00 2001 From: Copilot <223556219+Copilot@users.noreply.github.com> Date: Mon, 10 Aug 2026 10:54:31 -0700 Subject: [PATCH 03/24] Clarify converter composition compatibility Explain which factories opt in, what callers append, and why constructor support alone does not guarantee safe converter composition. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> Copilot-Session: 6776aa95-0924-4967-a056-e5a228af3faa --- doc/code/framework.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/doc/code/framework.md b/doc/code/framework.md index 3f98f0b69a..170464baac 100644 --- a/doc/code/framework.md +++ b/doc/code/framework.md @@ -182,7 +182,7 @@ If you are contributing to PyRIT, that work will most likely land in one of the - `core` stays deliberately small so a default run doesn't print 200 techniques or take forever; the wider catalog lives in `extra` and is selected on demand. Users pick subsets by passing initializer tags (e.g. `core`, `extra`, `all`) or writing their own initializer, so different runs — including from the CLI — can register different technique sets without changing the catalog. - A technique tied to one scenario is fine; if it's pinned and non-reusable it can stay local to that scenario, but if another scenario could reuse it, promote it to a catalog module and tag it. - Tags describe a technique (behavioral tags like `single_turn`/`multi_turn`, owner tags like `airt`); they don't decide what a scenario runs. There is deliberately **no global `default` tag** — a default is scenario-relative, declared per scenario via `build_technique_class_from_factories` (the `factories` list is the pool, catalog tags become named aggregate presets, and `default_tags` / `default_names` set what runs when nothing is chosen). -- Factories opt into additive request-converter composition with `supports_request_converter_composition=True`. This is a semantic capability, not just constructor-signature detection; the factory validates that opted-in attacks accept `attack_converter_config`. +- Each `AttackTechniqueFactory` describes how to build one named attack technique. Set `supports_request_converter_composition=True` only when callers can safely append more request converters to that technique's existing converter chain (for example, when the Jailbreak scenario adds a jailbreak-template converter). Accepting an `attack_converter_config` constructor argument is not enough by itself: an attack may accept converters but use them in a way that cannot safely be combined with others. When a technique opts in, the factory also verifies that its attack class accepts `attack_converter_config`. - **Does not own**: the conversation algorithm itself. Branching, turn management, and scoring decisions live in the executor it wraps — a technique only selects and configures existing components, and shouldn't implement new sending, scoring, or branching logic. **Framework Plans**: From 408e066030183fea446b10e25444deec5fbc7c9b Mon Sep 17 00:00:00 2001 From: Copilot <223556219+Copilot@users.noreply.github.com> Date: Mon, 10 Aug 2026 16:12:03 -0700 Subject: [PATCH 04/24] Reject lossy non-chat modality flattening Fail clearly when role-scoped converters produce non-text prepended history that string normalization cannot preserve. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> Copilot-Session: 6776aa95-0924-4967-a056-e5a228af3faa --- .../attack/component/conversation_manager.py | 36 +++++++++++++++++++ .../prepended_conversation_config.py | 3 +- .../component/test_conversation_manager.py | 35 ++++++++++++++++-- 3 files changed, 71 insertions(+), 3 deletions(-) diff --git a/pyrit/executor/attack/component/conversation_manager.py b/pyrit/executor/attack/component/conversation_manager.py index 60d25904ef..7b1a794dbd 100644 --- a/pyrit/executor/attack/component/conversation_manager.py +++ b/pyrit/executor/attack/component/conversation_manager.py @@ -427,12 +427,48 @@ async def _build_converted_normalizer_view_async( request_converters=request_converters, apply_to_roles=apply_to_roles, ) + self._validate_flattenable_converter_output( + source_messages=messages, + converted_messages=converted_messages, + ) for message in converted_messages: for piece in message.message_pieces: piece.original_value = piece.converted_value piece.original_value_data_type = piece.converted_value_data_type return converted_messages + @staticmethod + def _validate_flattenable_converter_output( + *, + source_messages: list[Message], + converted_messages: list[Message], + ) -> None: + """ + Reject converted history that a string normalizer cannot preserve. + + Raises: + ValueError: If an applied converter produced non-text prepended history. + """ + output_types: set[str] = set() + for source_message, converted_message in zip(source_messages, converted_messages, strict=True): + for source_piece, converted_piece in zip( + source_message.message_pieces, + converted_message.message_pieces, + strict=True, + ): + converter_was_applied = len(converted_piece.converter_identifiers) > len( + source_piece.converter_identifiers + ) + if converter_was_applied and converted_piece.converted_value_data_type != "text": + output_types.add(converted_piece.converted_value_data_type) + + if output_types: + raise ValueError( + "Cannot flatten prepended conversation for a non-chat target after request converters " + f"produced non-text output types {sorted(output_types)}. Role-scoped prepended conversion " + "must produce text; use text-output converters or a chat target with editable history." + ) + @staticmethod def _build_original_normalizer_view(messages: list[Message]) -> list[Message]: """ diff --git a/pyrit/executor/attack/component/prepended_conversation_config.py b/pyrit/executor/attack/component/prepended_conversation_config.py index 8930f829af..e03f21adbc 100644 --- a/pyrit/executor/attack/component/prepended_conversation_config.py +++ b/pyrit/executor/attack/component/prepended_conversation_config.py @@ -27,7 +27,8 @@ class PrependedConversationConfig: Non-chat objective targets apply request converters to the configured roles before normalizing the prepended conversation into the first turn (via ``message_normalizer``; - default: ConversationContextNormalizer). + default: ConversationContextNormalizer). Those converters must produce text because + string normalization cannot preserve converted image, audio, or other non-text output. """ # Request converters default to prepended user messages only. Assistant history is diff --git a/tests/unit/executor/attack/component/test_conversation_manager.py b/tests/unit/executor/attack/component/test_conversation_manager.py index 4404efcee7..1f79f0b3ae 100644 --- a/tests/unit/executor/attack/component/test_conversation_manager.py +++ b/tests/unit/executor/attack/component/test_conversation_manager.py @@ -24,7 +24,7 @@ import pytest from unit.mocks import get_mock_scorer_identifier -from pyrit.converter import Base64Converter +from pyrit.converter import Base64Converter, Converter, ConverterResult from pyrit.executor.attack import ConversationManager, ConversationState from pyrit.executor.attack.component import PrependedConversationConfig from pyrit.executor.attack.component.conversation_manager import ( @@ -36,7 +36,7 @@ from pyrit.executor.attack.core import AttackContext from pyrit.executor.attack.core.attack_parameters import AttackParameters from pyrit.message_normalizer import ConversationContextNormalizer -from pyrit.models import ComponentIdentifier, Message, MessagePiece, Score +from pyrit.models import ComponentIdentifier, Message, MessagePiece, PromptDataType, Score from pyrit.prompt_normalizer import ConverterConfiguration, PromptNormalizer from pyrit.prompt_target import PromptTarget @@ -61,6 +61,16 @@ class _TestAttackContext(AttackContext): last_score: Score | None = None +class _ImageOutputConverter(Converter): + """A deterministic text-to-image converter for non-chat flattening tests.""" + + SUPPORTED_INPUT_TYPES: tuple[PromptDataType, ...] = ("text",) + SUPPORTED_OUTPUT_TYPES: tuple[PromptDataType, ...] = ("image_path",) + + async def convert_async(self, *, prompt: str, input_type: PromptDataType = "text") -> ConverterResult: + return ConverterResult(output_text="converted.png", output_type="image_path") + + # ============================================================================= # Fixtures # ============================================================================= @@ -1154,6 +1164,27 @@ async def test_non_chat_target_converts_assistant_history_only_when_opted_in( f"Turn 1:\nuser: Hello, how are you?\nassistant: {encoded_assistant}\n\n{encoded_live_request}" ) + async def test_non_chat_target_rejects_non_text_history_converter_output( + self, + attack_identifier: ComponentIdentifier, + mock_prompt_target: MagicMock, + sample_conversation: list[Message], + ) -> None: + manager = ConversationManager() + context = _TestAttackContext(params=AttackParameters(objective="Test objective")) + context.prepended_conversation = sample_conversation + converter_config = ConverterConfiguration.from_converters(converters=[_ImageOutputConverter()]) + + with pytest.raises(ValueError, match="non-chat target.*non-text output types.*image_path"): + await manager.initialize_context_async( + context=context, + target=mock_prompt_target, + conversation_id=str(uuid.uuid4()), + request_converters=converter_config, + ) + + assert sample_conversation[0].get_piece().converted_value_data_type == "text" + async def test_non_chat_target_behavior_normalize_first_turn_creates_next_message( self, attack_identifier: ComponentIdentifier, From 029a1ede9be11a3a58f1ac76d2f7a46e7e8f7c77 Mon Sep 17 00:00:00 2001 From: Copilot <223556219+Copilot@users.noreply.github.com> Date: Mon, 10 Aug 2026 16:31:04 -0700 Subject: [PATCH 05/24] Fix converter retry and index scoping Reuse prepared non-chat requests across retries and preserve converter piece indexes when applying role-scoped prepended conversion. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> Copilot-Session: 6776aa95-0924-4967-a056-e5a228af3faa --- .../attack/component/conversation_manager.py | 17 ++-- .../component/test_conversation_manager.py | 93 +++++++++++++++++++ 2 files changed, 102 insertions(+), 8 deletions(-) diff --git a/pyrit/executor/attack/component/conversation_manager.py b/pyrit/executor/attack/component/conversation_manager.py index 7b1a794dbd..771cb4dc03 100644 --- a/pyrit/executor/attack/component/conversation_manager.py +++ b/pyrit/executor/attack/component/conversation_manager.py @@ -364,6 +364,9 @@ async def _handle_non_chat_target_async( Returns: Empty ConversationState (non-chat targets don't track turns). """ + if context.next_message and context.next_message.request_converters_applied: + return ConversationState() + normalizer = config.get_message_normalizer() original_context = await self._normalize_non_chat_context_async( messages=self._build_original_normalizer_view(prepended_conversation), @@ -732,12 +735,10 @@ async def _apply_converters_async( request_converters: Converter configurations to apply. apply_to_roles: Only apply to pieces with these roles. """ - for piece in message.message_pieces: - if piece.api_role not in apply_to_roles: - continue + if message.api_role not in apply_to_roles: + return - temp_message = Message(message_pieces=[piece]) - await self._prompt_normalizer.convert_values_async( - message=temp_message, - converter_configurations=request_converters, - ) + await self._prompt_normalizer.convert_values_async( + message=message, + converter_configurations=request_converters, + ) diff --git a/tests/unit/executor/attack/component/test_conversation_manager.py b/tests/unit/executor/attack/component/test_conversation_manager.py index 1f79f0b3ae..1a9d8569aa 100644 --- a/tests/unit/executor/attack/component/test_conversation_manager.py +++ b/tests/unit/executor/attack/component/test_conversation_manager.py @@ -71,6 +71,20 @@ async def convert_async(self, *, prompt: str, input_type: PromptDataType = "text return ConverterResult(output_text="converted.png", output_type="image_path") +class _CountingTextConverter(Converter): + """A text converter that produces a different result on each call.""" + + SUPPORTED_INPUT_TYPES: tuple[PromptDataType, ...] = ("text",) + SUPPORTED_OUTPUT_TYPES: tuple[PromptDataType, ...] = ("text",) + + def __init__(self) -> None: + self.call_count = 0 + + async def convert_async(self, *, prompt: str, input_type: PromptDataType = "text") -> ConverterResult: + self.call_count += 1 + return ConverterResult(output_text=f"conversion-{self.call_count}<{prompt}>", output_type="text") + + # ============================================================================= # Fixtures # ============================================================================= @@ -1185,6 +1199,85 @@ async def test_non_chat_target_rejects_non_text_history_converter_output( assert sample_conversation[0].get_piece().converted_value_data_type == "text" + async def test_non_chat_target_repeated_setup_reuses_prepared_request( + self, + attack_identifier: ComponentIdentifier, + mock_prompt_target: MagicMock, + sample_conversation: list[Message], + ) -> None: + manager = ConversationManager() + context = _TestAttackContext(params=AttackParameters(objective="Test objective")) + context.prepended_conversation = sample_conversation + context.next_message = Message.from_prompt(prompt="live request", role="user") + converter = _CountingTextConverter() + converter_config = ConverterConfiguration.from_converters(converters=[converter]) + + await manager.initialize_context_async( + context=context, + target=mock_prompt_target, + conversation_id=str(uuid.uuid4()), + request_converters=converter_config, + ) + assert context.next_message is not None + first_prepared_value = context.next_message.get_piece().converted_value + assert converter.call_count == 2 + + await manager.initialize_context_async( + context=context, + target=mock_prompt_target, + conversation_id=str(uuid.uuid4()), + request_converters=converter_config, + ) + + assert context.next_message.get_piece().converted_value == first_prepared_value + assert converter.call_count == 2 + + async def test_non_chat_target_preserves_converter_piece_indexes( + self, + attack_identifier: ComponentIdentifier, + mock_prompt_target: MagicMock, + ) -> None: + manager = ConversationManager() + context = _TestAttackContext(params=AttackParameters(objective="Test objective")) + context.prepended_conversation = [ + Message( + message_pieces=[ + MessagePiece( + role="user", + original_value="first piece", + conversation_id="seed", + sequence=0, + ), + MessagePiece( + role="user", + original_value="second piece", + conversation_id="seed", + sequence=0, + ), + ] + ) + ] + context.next_message = Message.from_prompt(prompt="live request", role="user") + converter_config = [ + ConverterConfiguration( + converters=[Base64Converter()], + indexes_to_apply=[0], + ) + ] + + await manager.initialize_context_async( + context=context, + target=mock_prompt_target, + conversation_id=str(uuid.uuid4()), + request_converters=converter_config, + ) + + assert context.next_message is not None + converted_value = context.next_message.get_piece().converted_value + assert base64.b64encode(b"first piece").decode() in converted_value + assert "second piece" in converted_value + assert base64.b64encode(b"second piece").decode() not in converted_value + async def test_non_chat_target_behavior_normalize_first_turn_creates_next_message( self, attack_identifier: ComponentIdentifier, From 91542324bc8c3f2532e8b9b99bce290146da3dfd Mon Sep 17 00:00:00 2001 From: Copilot <223556219+Copilot@users.noreply.github.com> Date: Mon, 10 Aug 2026 16:47:05 -0700 Subject: [PATCH 06/24] Restrict jailbreak converter to prompt sending Limit inline jailbreak-template conversion to the scenario-owned prompt_sending delivery. Keep native system-prompt delivery separate and remove the now-unnecessary cross-technique composition capability. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> Copilot-Session: 6776aa95-0924-4967-a056-e5a228af3faa --- doc/code/framework.md | 1 - .../scenario/core/attack_technique_factory.py | 34 ------- pyrit/scenario/scenarios/airt/jailbreak.py | 95 ++++++------------- pyrit/setup/initializers/techniques/airt.py | 1 - pyrit/setup/initializers/techniques/core.py | 1 - tests/unit/scenario/airt/test_jailbreak.py | 73 ++------------ .../core/test_attack_technique_factory.py | 23 ----- 7 files changed, 36 insertions(+), 192 deletions(-) diff --git a/doc/code/framework.md b/doc/code/framework.md index 170464baac..a3938ef2dc 100644 --- a/doc/code/framework.md +++ b/doc/code/framework.md @@ -182,7 +182,6 @@ If you are contributing to PyRIT, that work will most likely land in one of the - `core` stays deliberately small so a default run doesn't print 200 techniques or take forever; the wider catalog lives in `extra` and is selected on demand. Users pick subsets by passing initializer tags (e.g. `core`, `extra`, `all`) or writing their own initializer, so different runs — including from the CLI — can register different technique sets without changing the catalog. - A technique tied to one scenario is fine; if it's pinned and non-reusable it can stay local to that scenario, but if another scenario could reuse it, promote it to a catalog module and tag it. - Tags describe a technique (behavioral tags like `single_turn`/`multi_turn`, owner tags like `airt`); they don't decide what a scenario runs. There is deliberately **no global `default` tag** — a default is scenario-relative, declared per scenario via `build_technique_class_from_factories` (the `factories` list is the pool, catalog tags become named aggregate presets, and `default_tags` / `default_names` set what runs when nothing is chosen). -- Each `AttackTechniqueFactory` describes how to build one named attack technique. Set `supports_request_converter_composition=True` only when callers can safely append more request converters to that technique's existing converter chain (for example, when the Jailbreak scenario adds a jailbreak-template converter). Accepting an `attack_converter_config` constructor argument is not enough by itself: an attack may accept converters but use them in a way that cannot safely be combined with others. When a technique opts in, the factory also verifies that its attack class accepts `attack_converter_config`. - **Does not own**: the conversation algorithm itself. Branching, turn management, and scoring decisions live in the executor it wraps — a technique only selects and configures existing components, and shouldn't implement new sending, scoring, or branching logic. **Framework Plans**: diff --git a/pyrit/scenario/core/attack_technique_factory.py b/pyrit/scenario/core/attack_technique_factory.py index c74114c770..1d168a0545 100644 --- a/pyrit/scenario/core/attack_technique_factory.py +++ b/pyrit/scenario/core/attack_technique_factory.py @@ -82,7 +82,6 @@ def __init__( adversarial_seed_prompt: SeedPrompt | str | None = None, seed_technique: AttackTechniqueSeedGroup | None = None, uses_adversarial: bool | None = None, - supports_request_converter_composition: bool = False, scorer_override_policy: ScorerOverridePolicy = ScorerOverridePolicy.WARN, ) -> None: """ @@ -120,9 +119,6 @@ def __init__( chat during execution. ``None`` auto-derives from the attack class constructor signature and seed-technique shape. Authors can override the derivation explicitly. - supports_request_converter_composition: Whether callers may safely - append request converters to this technique. This is an explicit - semantic opt-in, not merely constructor-signature detection. scorer_override_policy: What to do when a scenario's scorer is incompatible with the attack's ``attack_scoring_config`` type annotation. Defaults to WARN. @@ -147,13 +143,11 @@ class constructor signature and seed-technique shape. adversarial_system_prompt is not None or adversarial_seed_prompt is not None ) self._seed_technique = seed_technique - self._supports_request_converter_composition = supports_request_converter_composition self._scorer_override_policy = scorer_override_policy self._uses_adversarial = uses_adversarial if uses_adversarial is not None else self._derive_uses_adversarial() self._validate_kwargs() - self._validate_converter_composition() self._validate_adversarial_flags() @classmethod @@ -172,7 +166,6 @@ def with_simulated_conversation( attack_kwargs: dict[str, Any] | None = None, adversarial_chat: PromptTarget | None = None, uses_adversarial: bool | None = None, - supports_request_converter_composition: bool = False, scorer_override_policy: ScorerOverridePolicy = ScorerOverridePolicy.WARN, ) -> AttackTechniqueFactory: """ @@ -224,9 +217,6 @@ def with_simulated_conversation( during execution. ``None`` auto-derives from the attack class constructor signature and seed-technique shape. Forwarded to the factory constructor. - supports_request_converter_composition: Whether callers may safely - append request converters to this technique. Forwarded to the - factory constructor. scorer_override_policy: Policy applied when a scenario's scorer is incompatible with the attack's ``attack_scoring_config`` type annotation. Defaults to ``WARN``. Forwarded to the factory @@ -287,7 +277,6 @@ def with_simulated_conversation( adversarial_chat=adversarial_chat, seed_technique=seed_technique, uses_adversarial=uses_adversarial, - supports_request_converter_composition=supports_request_converter_composition, scorer_override_policy=scorer_override_policy, ) @@ -323,23 +312,6 @@ def _validate_adversarial_flags(self) -> None: f"should not have one wired." ) - def _validate_converter_composition(self) -> None: - """ - Validate that an opt-in factory can receive additive request converters. - - Raises: - ValueError: If composition is enabled but the attack constructor does - not accept ``attack_converter_config``. - """ - if ( - self._supports_request_converter_composition - and "attack_converter_config" not in self._get_accepted_params() - ): - raise ValueError( - f"Factory '{self._name}' declares supports_request_converter_composition=True, " - f"but {self._attack_class.__name__} does not accept 'attack_converter_config'." - ) - def _validate_kwargs(self) -> None: """ Validate that all kwargs are valid parameters for the attack class constructor. @@ -462,11 +434,6 @@ def uses_adversarial(self) -> bool: """Whether this technique drives an adversarial chat during execution.""" return self._uses_adversarial - @property - def supports_request_converter_composition(self) -> bool: - """Whether callers may safely append request converters to this technique.""" - return self._supports_request_converter_composition - @property def scoring_config_type(self) -> type | None: """The required ``attack_scoring_config`` subtype, or ``None`` if any config is accepted.""" @@ -807,7 +774,6 @@ def _build_identifier(self) -> ComponentIdentifier: "attack_class": self._attack_class.__name__, "kwargs": kwargs_for_id, "uses_adversarial": self._uses_adversarial, - "supports_request_converter_composition": self._supports_request_converter_composition, } if self._technique_tags: params["technique_tags"] = list(self._technique_tags) diff --git a/pyrit/scenario/scenarios/airt/jailbreak.py b/pyrit/scenario/scenarios/airt/jailbreak.py index aa744a4864..9bfcc63f19 100644 --- a/pyrit/scenario/scenarios/airt/jailbreak.py +++ b/pyrit/scenario/scenarios/airt/jailbreak.py @@ -72,7 +72,6 @@ def _prompt_sending_factory() -> AttackTechniqueFactory: name=_PROMPT_SENDING, attack_class=PromptSendingAttack, technique_tags=["single_turn"], - supports_request_converter_composition=True, ) @@ -94,7 +93,6 @@ def _jailbreak_system_prompt_factory() -> AttackTechniqueFactory: name=_JAILBREAK_SYSTEM_PROMPT, attack_class=PromptSendingAttack, technique_tags=["single_turn"], - supports_request_converter_composition=True, ) @@ -106,37 +104,17 @@ def _extra_default_factories() -> dict[str, AttackTechniqueFactory]: } -def _is_jailbreak_compatible_factory(factory: AttackTechniqueFactory) -> bool: - """Return whether a factory supports direct jailbreak-converter delivery.""" - has_simulated_conversation = ( - factory.seed_technique is not None and factory.seed_technique.has_simulated_conversation - ) - return ( - "multi_turn" not in factory.technique_tags - and not has_simulated_conversation - and factory.supports_request_converter_composition - ) - - @cache def _build_jailbreak_technique() -> type[ScenarioTechnique]: """ - Build the Jailbreak technique class from compatible direct factories and local deliveries. - - Registered multi-turn techniques, simulated-conversation seed techniques, and factories that - have not explicitly opted into additive request-converter composition are excluded. + Build the Jailbreak technique class from its two scenario-owned delivery methods. Returns: type[ScenarioTechnique]: The dynamically generated technique enum class. """ - registry = AttackTechniqueRegistry.get_registry_singleton() - registered = [ - factory for factory in registry.get_factories_or_raise().values() if _is_jailbreak_compatible_factory(factory) - ] - factories = registered + list(_extra_default_factories().values()) return AttackTechniqueRegistry.build_technique_class_from_factories( # type: ignore[return-value, ty:invalid-return-type] class_name="JailbreakTechnique", - factories=factories, + factories=list(_extra_default_factories().values()), default_names=set(_DEFAULT_TECHNIQUES), ) @@ -149,22 +127,20 @@ class Jailbreak(Scenario): selectors: - **dataset** — the harmful objectives (HarmBench). - - **techniques** — compatible direct deliveries for each jailbreak. Two deliveries - are on by default: ``prompt_sending`` (the template rendered inline into the user message) and + - **techniques** — two delivery methods for each jailbreak: ``prompt_sending`` (the template + rendered inline into the user message) and ``jailbreak_system_prompt`` (the template set as the system prompt with the objective sent as - the user turn). Registered direct techniques are opt-in only when their factory explicitly - supports request-converter composition. Multi-turn and simulated-conversation techniques are - not offered. + the user turn). - **jailbreaks** — which jailbreak templates to run (a random ``num_jailbreaks`` sample or an explicit ``jailbreak_names`` set). ``prompt_sending`` applies each template as a ``TextJailbreakConverter`` on the outgoing request, - so the objective is rendered inline into the template's ``{{prompt}}`` slot; this keeps that - delivery target-agnostic and lets it compose with compatible direct techniques. ``jailbreak_system_prompt`` - instead sets the template as a native system prompt and sends the objective as its own user turn, - so it is only built for targets that natively support editable history and system prompts (it is - skipped for incapable targets, or raises if it is the only selected technique). Responses are - scored to determine whether the jailbreak succeeded (non-refusal). + so the objective is rendered inline into the template's ``{{prompt}}`` slot. + ``jailbreak_system_prompt`` instead sets the template as a native system prompt and sends the + objective as its own user turn, so it is only built for targets that natively support editable + history and system prompts (it is skipped for incapable targets, or raises if it is the only + selected technique). Responses are scored to determine whether the jailbreak succeeded + (non-refusal). """ VERSION: int = 4 @@ -266,7 +242,7 @@ def _resolve_scenario_techniques(self, *, scenario_techniques: Any) -> list[Scen values = [getattr(item, "value", repr(item)) for item in incompatible] raise ValueError( "Jailbreak received stale or incompatible techniques " - f"{values}. Select a compatible direct-delivery technique from JailbreakTechnique." + f"{values}. Select 'prompt_sending' or 'jailbreak_system_prompt'." ) return super()._resolve_scenario_techniques(scenario_techniques=scenario_techniques) @@ -324,13 +300,12 @@ async def _build_atomic_attacks_async(self, *, context: ScenarioContext) -> list """ Build one atomic attack per (technique x jailbreak template x dataset x attempt). - ``prompt_sending`` (and compatible opt-in direct techniques) deliver each jailbreak template as a - ``TextJailbreakConverter`` appended to that technique's request converters, so the objective - is rendered inline into the template's ``{{prompt}}`` slot on the wire. ``jailbreak_system_prompt`` - instead delivers the template as - a native system prompt (no converter) with the objective sent as its own user turn, so it is - only built when the objective target natively supports editable history and system prompts. - Results group by jailbreak template so per-template ASR rolls up naturally. + ``prompt_sending`` delivers each jailbreak template as a ``TextJailbreakConverter`` so the + objective is rendered inline into the template's ``{{prompt}}`` slot on the wire. + ``jailbreak_system_prompt`` instead delivers the template as a native system prompt (no + converter) with the objective sent as its own user turn, so it is only built when the + objective target natively supports editable history and system prompts. Results group by + jailbreak template so per-template ASR rolls up naturally. Args: context (ScenarioContext): The resolved runtime inputs for this run. @@ -357,27 +332,15 @@ async def _build_atomic_attacks_async(self, *, context: ScenarioContext) -> list if missing: raise ValueError( "Jailbreak selected techniques that are no longer available: " - f"{sorted(missing)}. Refresh the plan and select compatible direct-delivery techniques." - ) - incompatible = [ - name for name, factory in technique_factories.items() if not _is_jailbreak_compatible_factory(factory) - ] - if incompatible: - raise ValueError( - "Jailbreak cannot compose with multi-turn, simulated-conversation, or " - f"non-composable techniques: {sorted(incompatible)}." + f"{sorted(missing)}. Refresh the plan and select a supported delivery method." ) - # ``jailbreak_system_prompt`` is delivered separately (native system prompt, no converter); - # every other technique goes through the inline converter path. + prompt_sending_factory = technique_factories.get(_PROMPT_SENDING) system_selected = _JAILBREAK_SYSTEM_PROMPT in technique_factories - converter_factories = { - name: factory for name, factory in technique_factories.items() if name != _JAILBREAK_SYSTEM_PROMPT - } build_system_delivery = system_selected and self._target_supports_system_delivery(self._objective_target) if system_selected and not build_system_delivery: - if not converter_factories: + if prompt_sending_factory is None: raise ValueError( "The 'jailbreak_system_prompt' technique needs a target that natively supports " "editable history and system prompts. Choose a capable target or a different technique." @@ -406,22 +369,21 @@ async def _build_atomic_attacks_async(self, *, context: ScenarioContext) -> list for template_file_name in self._resolved_jailbreaks: template_stem = Path(template_file_name).stem - if converter_factories: + if prompt_sending_factory is not None: jailbreak_converter = TextJailbreakConverter( jailbreak_template=TextJailBreak(template_file_name=template_file_name) ) - # Within the extra-converter stack, apply the jailbreak first (wrap the raw objective - # in the template), then any per-technique converters the caller layered on via - # ``--techniques :converter.*``. (A technique's own built-in converters, if any, - # still run ahead of this extra stack inside the factory.) + # Apply the jailbreak before any caller-supplied prompt_sending converters. technique_converters = { - technique_name: [jailbreak_converter, *self._technique_converters.get(technique_name, [])] - for technique_name in converter_factories + _PROMPT_SENDING: [ + jailbreak_converter, + *self._technique_converters.get(_PROMPT_SENDING, []), + ] } atomic_attacks.extend( self._build_delivery_attacks( builder=builder, - technique_factories=converter_factories, + technique_factories={_PROMPT_SENDING: prompt_sending_factory}, technique_converters=technique_converters, dataset_groups=context.seed_groups_by_dataset, template_stem=template_stem, @@ -509,7 +471,6 @@ def _build_system_prompt_factory(self, *, template_file_name: str) -> AttackTech attack_class=PromptSendingAttack, technique_tags=["single_turn"], seed_technique=seed_technique, - supports_request_converter_composition=True, ) @staticmethod diff --git a/pyrit/setup/initializers/techniques/airt.py b/pyrit/setup/initializers/techniques/airt.py index 5b3c92cec5..1469ea6e19 100644 --- a/pyrit/setup/initializers/techniques/airt.py +++ b/pyrit/setup/initializers/techniques/airt.py @@ -42,7 +42,6 @@ def get_technique_factories() -> list[AttackTechniqueFactory]: attack_class=PromptSendingAttack, description="Obfuscates the objective by asking for it encoded as the first letter of each word.", technique_tags=["single_turn", "airt", "leakage"], - supports_request_converter_composition=True, attack_kwargs={ "attack_converter_config": AttackConverterConfig( request_converters=ConverterConfiguration.from_converters(converters=[FirstLetterConverter()]) diff --git a/pyrit/setup/initializers/techniques/core.py b/pyrit/setup/initializers/techniques/core.py index e696429af7..b4be579675 100644 --- a/pyrit/setup/initializers/techniques/core.py +++ b/pyrit/setup/initializers/techniques/core.py @@ -159,7 +159,6 @@ def get_technique_factories() -> list[AttackTechniqueFactory]: attack_class=PromptSendingAttack, description="Reverses the objective text so it slips past filters, then asks the target to flip it back.", technique_tags=["single_turn", "light"], - supports_request_converter_composition=True, attack_kwargs={ "attack_converter_config": AttackConverterConfig( request_converters=ConverterConfiguration.from_converters( diff --git a/tests/unit/scenario/airt/test_jailbreak.py b/tests/unit/scenario/airt/test_jailbreak.py index 3a8beb3bfb..b055e087aa 100644 --- a/tests/unit/scenario/airt/test_jailbreak.py +++ b/tests/unit/scenario/airt/test_jailbreak.py @@ -42,13 +42,7 @@ def _technique_class(): @pytest.fixture(autouse=True) def reset_technique_registry(): - """Populate the attack-technique registry so the dynamic technique class can be built. - - Mirrors the RapidResponse test setup: reset the registries, register a mock adversarial - target (so factory construction does not fall back to a real target), and register the core - technique factories. The build cache is cleared around each test so the class reflects the - freshly-registered factories. - """ + """Populate the attack-technique registry used by the shared matrix factory resolver.""" AttackTechniqueRegistry.reset_registry_singleton() TargetRegistry.reset_registry_singleton() _build_jailbreak_technique.cache_clear() @@ -310,9 +304,8 @@ async def test_jailbreak_delivered_as_request_converter( """The crux: the jailbreak template reaches the target as a ``TextJailbreakConverter`` on the technique's outgoing requests (not as prepended framing on the seed group). - Delivery via ``factory.create(extra_request_converters=...)`` is what keeps the scenario - target-agnostic and composable with every technique. Also assert the seed groups carry no - prepended jailbreak framing. + Delivery via ``factory.create(extra_request_converters=...)`` keeps the prompt-sending path + target-agnostic. Also assert the seed groups carry no prepended jailbreak framing. """ captured: list[Any] = [] original_create = AttackTechniqueFactory.create @@ -395,32 +388,6 @@ async def test_stale_incompatible_technique_is_rejected( with pytest.raises(ValueError, match="stale or incompatible"): await scenario.initialize_async() - async def test_incompatible_runtime_factory_is_rejected( - self, mock_objective_target, mock_objective_scorer, mock_memory_seed_groups - ): - incompatible = AttackTechniqueFactory( - name="legacy_multi_turn", - attack_class=PromptSendingAttack, - technique_tags=["multi_turn"], - supports_request_converter_composition=True, - ) - with _patch_seed_groups(mock_memory_seed_groups): - with patch( - "pyrit.scenario.scenarios.airt.jailbreak.resolve_technique_factories", - return_value={_PROMPT_SENDING: incompatible}, - ): - scenario = Jailbreak(objective_scorer=mock_objective_scorer) - technique_class = _build_jailbreak_technique() - scenario.set_params_from_args( - args=_default_args( - mock_objective_target, - scenario_techniques=[technique_class(_PROMPT_SENDING)], - jailbreak_names=["aim.yaml"], - ) - ) - with pytest.raises(ValueError, match="cannot compose"): - await scenario.initialize_async() - async def test_missing_runtime_factory_is_rejected( self, mock_objective_target, mock_objective_scorer, mock_memory_seed_groups ): @@ -719,38 +686,14 @@ def test_default_techniques_are_the_two_deliveries(self): assert default_values == set(_DEFAULT_TECHNIQUES) assert default_values == {_PROMPT_SENDING, _JAILBREAK_SYSTEM_PROMPT} - def test_only_compatible_direct_registry_techniques_are_available(self): + def test_only_scenario_delivery_techniques_are_available(self): technique_class = _technique_class() available = {t.value for t in technique_class.get_all_techniques()} - incompatible = { - "context_compliance", - "role_play_movie_script", - "role_play_video_game", - "role_play_trivia_game", - "role_play_persuasion", - "role_play_persuasion_written", - "crescendo_simulated", - "crescendo_movie_director", - "crescendo_history_lecture", - "crescendo_journalist_interview", - "red_teaming", - "tap", - "many_shot", - "pair", - } - assert {_PROMPT_SENDING, _JAILBREAK_SYSTEM_PROMPT, "flip"}.issubset(available) - assert incompatible.isdisjoint(available) - - def test_registry_metadata_omits_incompatible_techniques(self): + assert available == {_PROMPT_SENDING, _JAILBREAK_SYSTEM_PROMPT} + + def test_registry_metadata_lists_only_scenario_deliveries(self): metadata = ScenarioRegistry()._build_metadata("airt.jailbreak", Jailbreak) - assert { - "context_compliance", - "role_play_movie_script", - "crescendo_simulated", - "red_teaming", - "tap", - "many_shot", - }.isdisjoint(metadata.all_techniques) + assert set(metadata.all_techniques) == {_PROMPT_SENDING, _JAILBREAK_SYSTEM_PROMPT} def test_scenario_version_is_four(self): assert Jailbreak.VERSION == 4 diff --git a/tests/unit/scenario/core/test_attack_technique_factory.py b/tests/unit/scenario/core/test_attack_technique_factory.py index a52773c301..39e3046685 100644 --- a/tests/unit/scenario/core/test_attack_technique_factory.py +++ b/tests/unit/scenario/core/test_attack_technique_factory.py @@ -178,29 +178,6 @@ def test_validate_kwargs_rejects_invalid_param_on_real_attack_class(self): attack_kwargs={"nonexistent_param": 42}, ) - def test_request_converter_composition_requires_supported_constructor(self): - class _NoConverterAttack: - def __init__(self, *, objective_target, attack_scoring_config=None): - self.objective_target = objective_target - - with pytest.raises(ValueError, match="does not accept 'attack_converter_config'"): - AttackTechniqueFactory( - name="test", - attack_class=_NoConverterAttack, - supports_request_converter_composition=True, - ) - - def test_request_converter_composition_is_explicit_opt_in(self): - default_factory = AttackTechniqueFactory(name="default", attack_class=_StubAttack) - composable_factory = AttackTechniqueFactory( - name="composable", - attack_class=_StubAttack, - supports_request_converter_composition=True, - ) - - assert not default_factory.supports_request_converter_composition - assert composable_factory.supports_request_converter_composition - class TestFactoryCreate: """Tests for AttackTechniqueFactory.create().""" From 3db85d3da7c23a413e5d19bc9d66055209ec315d Mon Sep 17 00:00:00 2001 From: Copilot <223556219+Copilot@users.noreply.github.com> Date: Mon, 10 Aug 2026 19:00:37 -0700 Subject: [PATCH 07/24] Preserve upstream ShieldGemma exports Restore the exact origin/main blob after the merge's line-ending check normalized this unrelated file. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> Copilot-Session: 6776aa95-0924-4967-a056-e5a228af3faa --- pyrit/score/__init__.py | 550 ++++++++++++++++++++-------------------- 1 file changed, 275 insertions(+), 275 deletions(-) diff --git a/pyrit/score/__init__.py b/pyrit/score/__init__.py index 30e32dfffd..0647026606 100644 --- a/pyrit/score/__init__.py +++ b/pyrit/score/__init__.py @@ -1,275 +1,275 @@ -# Copyright (c) Microsoft Corporation. -# Licensed under the MIT license. - -""" -Scoring functionality for evaluating AI model responses across various dimensions -including harm detection, objective completion, and content classification. -""" - -import importlib -from typing import TYPE_CHECKING - -from pyrit.output.scorer.base import ScorerPrinterBase as ScorerPrinter -from pyrit.score.batch_scorer import BatchScorer -from pyrit.score.conversation_scorer import ConversationScorer, create_conversation_scorer -from pyrit.score.float_scale.azure_content_filter_scorer import AzureContentFilterScorer -from pyrit.score.float_scale.float_scale_score_aggregator import ( - FloatScaleScoreAggregator, - FloatScaleScorerAllCategories, - FloatScaleScorerByCategory, -) -from pyrit.score.float_scale.float_scale_scorer import FloatScaleScorer -from pyrit.score.float_scale.insecure_code_scorer import ( - InsecureCodeScorer, - render_insecure_code_system_prompt, -) -from pyrit.score.float_scale.likert_scale import LikertScale, LikertScaleEntry -from pyrit.score.float_scale.numeric_scale import NumericRange, NumericRubric -from pyrit.score.float_scale.plagiarism_scorer import PlagiarismMetric, PlagiarismScorer -from pyrit.score.float_scale.self_ask_general_float_scale_scorer import SelfAskGeneralFloatScaleScorer -from pyrit.score.float_scale.self_ask_likert_scorer import ( - LikertScaleEvalFiles, - LikertScalePaths, - SelfAskLikertScorer, - render_likert_system_prompt, -) -from pyrit.score.float_scale.self_ask_scale_scorer import ( - SelfAskScaleScorer, - render_scale_system_prompt, -) -from pyrit.score.response_handler import ( - CallableResponseHandler, - JsonSchemaResponseHandler, - ResponseHandler, -) -from pyrit.score.scorer import Scorer -from pyrit.score.scorer_evaluation.metrics_type import MetricsType, RegistryUpdateBehavior -from pyrit.score.scorer_evaluation.scorer_metrics import ( - HarmScorerMetrics, - ObjectiveScorerMetrics, - ScorerMetrics, - ScorerMetricsWithIdentity, -) -from pyrit.score.scorer_evaluation.scorer_metrics_io import ( - find_objective_metrics_by_eval_hash, - get_all_harm_metrics, - get_all_objective_metrics, -) -from pyrit.score.scorer_info import get_scorer_info -from pyrit.score.scorer_prompt_validator import ScorerPromptValidator -from pyrit.score.true_false.decoding_scorer import DecodingScorer -from pyrit.score.true_false.float_scale_threshold_scorer import FloatScaleThresholdScorer -from pyrit.score.true_false.gandalf_scorer import GandalfScorer -from pyrit.score.true_false.llamaguard_parser import LLAMAGUARD_3_CATEGORY_CODES, parse_llamaguard_response -from pyrit.score.true_false.llamaguard_policy import LlamaGuardCategory, LlamaGuardPolicy -from pyrit.score.true_false.llamaguard_scorer import ( - LlamaGuardMessageRole, - LlamaGuardScorer, - render_llamaguard_prompt, -) -from pyrit.score.true_false.prompt_shield_scorer import PromptShieldScorer -from pyrit.score.true_false.question_answer_scorer import QuestionAnswerScorer -from pyrit.score.true_false.regex.anthrax_keyword_scorer import AnthraxKeywordScorer -from pyrit.score.true_false.regex.credential_leak_scorer import CredentialLeakScorer -from pyrit.score.true_false.regex.fentanyl_keyword_scorer import FentanylKeywordScorer -from pyrit.score.true_false.regex.ldap_injection_output_scorer import LDAPInjectionOutputScorer -from pyrit.score.true_false.regex.markdown_injection import MarkdownInjectionScorer -from pyrit.score.true_false.regex.meth_keyword_scorer import MethKeywordScorer -from pyrit.score.true_false.regex.nerve_agent_keyword_scorer import NerveAgentKeywordScorer -from pyrit.score.true_false.regex.open_redirect_output_scorer import OpenRedirectOutputScorer -from pyrit.score.true_false.regex.path_traversal_output_scorer import PathTraversalOutputScorer -from pyrit.score.true_false.regex.regex_scorer import RegexScorer -from pyrit.score.true_false.regex.shell_command_output_scorer import ShellCommandOutputScorer -from pyrit.score.true_false.regex.sql_injection_output_scorer import SQLInjectionOutputScorer -from pyrit.score.true_false.regex.ssrf_output_scorer import SSRFOutputScorer -from pyrit.score.true_false.regex.ssti_output_scorer import SSTIOutputScorer -from pyrit.score.true_false.regex.static_prompt_injection_scorer import StaticPromptInjectionScorer -from pyrit.score.true_false.regex.xss_output_scorer import XSSOutputScorer -from pyrit.score.true_false.regex.xxe_output_scorer import XXEOutputScorer -from pyrit.score.true_false.self_ask_category_scorer import ( - ContentClassifier, - ContentClassifierCategory, - ContentClassifierPaths, - SelfAskCategoryScorer, - render_category_system_prompt, -) -from pyrit.score.true_false.self_ask_general_true_false_scorer import SelfAskGeneralTrueFalseScorer -from pyrit.score.true_false.self_ask_question_answer_scorer import SelfAskQuestionAnswerScorer -from pyrit.score.true_false.self_ask_refusal_scorer import RefusalScorerPaths, SelfAskRefusalScorer -from pyrit.score.true_false.self_ask_true_false_scorer import ( - SelfAskTrueFalseScorer, - TrueFalseQuestion, - TrueFalseQuestionPaths, - render_true_false_system_prompt, -) -from pyrit.score.true_false.shieldgemma_parser import parse_shieldgemma_response -from pyrit.score.true_false.shieldgemma_policy import ( - SHIELDGEMMA_DEFAULT_POLICY_PATH, - ShieldGemmaGuideline, - ShieldGemmaMessageRole, - ShieldGemmaPolicy, -) -from pyrit.score.true_false.shieldgemma_scorer import ( - ShieldGemmaScorer, - render_shieldgemma_prompt, -) -from pyrit.score.true_false.substring_scorer import SubStringScorer -from pyrit.score.true_false.true_false_composite_scorer import TrueFalseCompositeScorer -from pyrit.score.true_false.true_false_inverter_scorer import TrueFalseInverterScorer -from pyrit.score.true_false.true_false_score_aggregator import TrueFalseAggregatorFunc, TrueFalseScoreAggregator -from pyrit.score.true_false.true_false_scorer import TrueFalseScorer - -if TYPE_CHECKING: - from pyrit.score.float_scale.audio_float_scale_scorer import AudioFloatScaleScorer - from pyrit.score.float_scale.video_float_scale_scorer import VideoFloatScaleScorer - from pyrit.score.scorer_evaluation.human_labeled_dataset import ( - HarmHumanLabeledEntry, - HumanLabeledDataset, - HumanLabeledEntry, - ObjectiveHumanLabeledEntry, - ) - from pyrit.score.scorer_evaluation.scorer_evaluator import ( - HarmScorerEvaluator, - ObjectiveScorerEvaluator, - ScorerEvalDatasetFiles, - ScorerEvaluator, - ) - from pyrit.score.true_false.audio_true_false_scorer import AudioTrueFalseScorer - from pyrit.score.true_false.video_true_false_scorer import VideoTrueFalseScorer - -# Lazy imports for modules with heavy third-party dependencies (PEP 562). -# Audio/video scorers import `av` (~1.9s), human_labeled_dataset imports `pandas` (~1.6s), -# scorer_evaluator imports `scipy.stats` (~1s). -_LAZY_IMPORTS: dict[str, str] = { - "AudioFloatScaleScorer": "pyrit.score.float_scale.audio_float_scale_scorer", - "AudioTrueFalseScorer": "pyrit.score.true_false.audio_true_false_scorer", - "VideoFloatScaleScorer": "pyrit.score.float_scale.video_float_scale_scorer", - "VideoTrueFalseScorer": "pyrit.score.true_false.video_true_false_scorer", - "HarmHumanLabeledEntry": "pyrit.score.scorer_evaluation.human_labeled_dataset", - "HumanLabeledDataset": "pyrit.score.scorer_evaluation.human_labeled_dataset", - "HumanLabeledEntry": "pyrit.score.scorer_evaluation.human_labeled_dataset", - "ObjectiveHumanLabeledEntry": "pyrit.score.scorer_evaluation.human_labeled_dataset", - "HarmScorerEvaluator": "pyrit.score.scorer_evaluation.scorer_evaluator", - "ObjectiveScorerEvaluator": "pyrit.score.scorer_evaluation.scorer_evaluator", - "ScorerEvalDatasetFiles": "pyrit.score.scorer_evaluation.scorer_evaluator", - "ScorerEvaluator": "pyrit.score.scorer_evaluation.scorer_evaluator", -} - - -def __getattr__(name: str) -> object: - if name in _LAZY_IMPORTS: - module = importlib.import_module(_LAZY_IMPORTS[name]) - attr = getattr(module, name) - globals()[name] = attr - return attr - raise AttributeError(f"module {__name__!r} has no attribute {name!r}") - - -__all__ = [ - "AnthraxKeywordScorer", - "AudioFloatScaleScorer", - "AudioTrueFalseScorer", - "AzureContentFilterScorer", - "BatchScorer", - "CallableResponseHandler", - "ContentClassifier", - "ContentClassifierCategory", - "ContentClassifierPaths", - "ConversationScorer", - "CredentialLeakScorer", - "DecodingScorer", - "FentanylKeywordScorer", - "create_conversation_scorer", - "FloatScaleScoreAggregator", - "FloatScaleScorerAllCategories", - "FloatScaleScorerByCategory", - "FloatScaleScorer", - "FloatScaleThresholdScorer", - "GandalfScorer", - "HarmHumanLabeledEntry", - "HarmScorerEvaluator", - "HarmScorerMetrics", - "HumanLabeledDataset", - "HumanLabeledEntry", - "InsecureCodeScorer", - "JsonSchemaResponseHandler", - "LDAPInjectionOutputScorer", - "LikertScaleEvalFiles", - "LikertScale", - "LikertScaleEntry", - "LikertScalePaths", - "LLAMAGUARD_3_CATEGORY_CODES", - "LlamaGuardCategory", - "LlamaGuardMessageRole", - "LlamaGuardPolicy", - "LlamaGuardScorer", - "MarkdownInjectionScorer", - "MethKeywordScorer", - "MetricsType", - "NerveAgentKeywordScorer", - "NumericRange", - "NumericRubric", - "ObjectiveHumanLabeledEntry", - "ObjectiveScorerEvaluator", - "ObjectiveScorerMetrics", - "OpenRedirectOutputScorer", - "parse_llamaguard_response", - "parse_shieldgemma_response", - "PathTraversalOutputScorer", - "PlagiarismMetric", - "PlagiarismScorer", - "PromptShieldScorer", - "QuestionAnswerScorer", - "RegexScorer", - "RegistryUpdateBehavior", - "render_category_system_prompt", - "render_insecure_code_system_prompt", - "render_llamaguard_prompt", - "render_likert_system_prompt", - "render_scale_system_prompt", - "render_shieldgemma_prompt", - "render_true_false_system_prompt", - "ResponseHandler", - "Scorer", - "ScorerEvalDatasetFiles", - "ScorerEvaluator", - "ScorerMetrics", - "ScorerMetricsWithIdentity", - "get_all_harm_metrics", - "get_all_objective_metrics", - "get_scorer_info", - "find_objective_metrics_by_eval_hash", - "ScorerPromptValidator", - "SelfAskCategoryScorer", - "SelfAskGeneralFloatScaleScorer", - "SelfAskGeneralTrueFalseScorer", - "SelfAskLikertScorer", - "SelfAskQuestionAnswerScorer", - "RefusalScorerPaths", - "SelfAskRefusalScorer", - "SelfAskScaleScorer", - "SelfAskTrueFalseScorer", - "ScorerPrinter", - "SHIELDGEMMA_DEFAULT_POLICY_PATH", - "ShieldGemmaGuideline", - "ShieldGemmaMessageRole", - "ShieldGemmaPolicy", - "ShieldGemmaScorer", - "ShellCommandOutputScorer", - "SQLInjectionOutputScorer", - "SSRFOutputScorer", - "SSTIOutputScorer", - "StaticPromptInjectionScorer", - "SubStringScorer", - "TrueFalseCompositeScorer", - "TrueFalseInverterScorer", - "TrueFalseQuestion", - "TrueFalseQuestionPaths", - "TrueFalseScoreAggregator", - "TrueFalseAggregatorFunc", - "TrueFalseScorer", - "VideoFloatScaleScorer", - "VideoTrueFalseScorer", - "XSSOutputScorer", - "XXEOutputScorer", -] +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT license. + +""" +Scoring functionality for evaluating AI model responses across various dimensions +including harm detection, objective completion, and content classification. +""" + +import importlib +from typing import TYPE_CHECKING + +from pyrit.output.scorer.base import ScorerPrinterBase as ScorerPrinter +from pyrit.score.batch_scorer import BatchScorer +from pyrit.score.conversation_scorer import ConversationScorer, create_conversation_scorer +from pyrit.score.float_scale.azure_content_filter_scorer import AzureContentFilterScorer +from pyrit.score.float_scale.float_scale_score_aggregator import ( + FloatScaleScoreAggregator, + FloatScaleScorerAllCategories, + FloatScaleScorerByCategory, +) +from pyrit.score.float_scale.float_scale_scorer import FloatScaleScorer +from pyrit.score.float_scale.insecure_code_scorer import ( + InsecureCodeScorer, + render_insecure_code_system_prompt, +) +from pyrit.score.float_scale.likert_scale import LikertScale, LikertScaleEntry +from pyrit.score.float_scale.numeric_scale import NumericRange, NumericRubric +from pyrit.score.float_scale.plagiarism_scorer import PlagiarismMetric, PlagiarismScorer +from pyrit.score.float_scale.self_ask_general_float_scale_scorer import SelfAskGeneralFloatScaleScorer +from pyrit.score.float_scale.self_ask_likert_scorer import ( + LikertScaleEvalFiles, + LikertScalePaths, + SelfAskLikertScorer, + render_likert_system_prompt, +) +from pyrit.score.float_scale.self_ask_scale_scorer import ( + SelfAskScaleScorer, + render_scale_system_prompt, +) +from pyrit.score.response_handler import ( + CallableResponseHandler, + JsonSchemaResponseHandler, + ResponseHandler, +) +from pyrit.score.scorer import Scorer +from pyrit.score.scorer_evaluation.metrics_type import MetricsType, RegistryUpdateBehavior +from pyrit.score.scorer_evaluation.scorer_metrics import ( + HarmScorerMetrics, + ObjectiveScorerMetrics, + ScorerMetrics, + ScorerMetricsWithIdentity, +) +from pyrit.score.scorer_evaluation.scorer_metrics_io import ( + find_objective_metrics_by_eval_hash, + get_all_harm_metrics, + get_all_objective_metrics, +) +from pyrit.score.scorer_info import get_scorer_info +from pyrit.score.scorer_prompt_validator import ScorerPromptValidator +from pyrit.score.true_false.decoding_scorer import DecodingScorer +from pyrit.score.true_false.float_scale_threshold_scorer import FloatScaleThresholdScorer +from pyrit.score.true_false.gandalf_scorer import GandalfScorer +from pyrit.score.true_false.llamaguard_parser import LLAMAGUARD_3_CATEGORY_CODES, parse_llamaguard_response +from pyrit.score.true_false.llamaguard_policy import LlamaGuardCategory, LlamaGuardPolicy +from pyrit.score.true_false.llamaguard_scorer import ( + LlamaGuardMessageRole, + LlamaGuardScorer, + render_llamaguard_prompt, +) +from pyrit.score.true_false.prompt_shield_scorer import PromptShieldScorer +from pyrit.score.true_false.question_answer_scorer import QuestionAnswerScorer +from pyrit.score.true_false.regex.anthrax_keyword_scorer import AnthraxKeywordScorer +from pyrit.score.true_false.regex.credential_leak_scorer import CredentialLeakScorer +from pyrit.score.true_false.regex.fentanyl_keyword_scorer import FentanylKeywordScorer +from pyrit.score.true_false.regex.ldap_injection_output_scorer import LDAPInjectionOutputScorer +from pyrit.score.true_false.regex.markdown_injection import MarkdownInjectionScorer +from pyrit.score.true_false.regex.meth_keyword_scorer import MethKeywordScorer +from pyrit.score.true_false.regex.nerve_agent_keyword_scorer import NerveAgentKeywordScorer +from pyrit.score.true_false.regex.open_redirect_output_scorer import OpenRedirectOutputScorer +from pyrit.score.true_false.regex.path_traversal_output_scorer import PathTraversalOutputScorer +from pyrit.score.true_false.regex.regex_scorer import RegexScorer +from pyrit.score.true_false.regex.shell_command_output_scorer import ShellCommandOutputScorer +from pyrit.score.true_false.regex.sql_injection_output_scorer import SQLInjectionOutputScorer +from pyrit.score.true_false.regex.ssrf_output_scorer import SSRFOutputScorer +from pyrit.score.true_false.regex.ssti_output_scorer import SSTIOutputScorer +from pyrit.score.true_false.regex.static_prompt_injection_scorer import StaticPromptInjectionScorer +from pyrit.score.true_false.regex.xss_output_scorer import XSSOutputScorer +from pyrit.score.true_false.regex.xxe_output_scorer import XXEOutputScorer +from pyrit.score.true_false.self_ask_category_scorer import ( + ContentClassifier, + ContentClassifierCategory, + ContentClassifierPaths, + SelfAskCategoryScorer, + render_category_system_prompt, +) +from pyrit.score.true_false.self_ask_general_true_false_scorer import SelfAskGeneralTrueFalseScorer +from pyrit.score.true_false.self_ask_question_answer_scorer import SelfAskQuestionAnswerScorer +from pyrit.score.true_false.self_ask_refusal_scorer import RefusalScorerPaths, SelfAskRefusalScorer +from pyrit.score.true_false.self_ask_true_false_scorer import ( + SelfAskTrueFalseScorer, + TrueFalseQuestion, + TrueFalseQuestionPaths, + render_true_false_system_prompt, +) +from pyrit.score.true_false.shieldgemma_parser import parse_shieldgemma_response +from pyrit.score.true_false.shieldgemma_policy import ( + SHIELDGEMMA_DEFAULT_POLICY_PATH, + ShieldGemmaGuideline, + ShieldGemmaMessageRole, + ShieldGemmaPolicy, +) +from pyrit.score.true_false.shieldgemma_scorer import ( + ShieldGemmaScorer, + render_shieldgemma_prompt, +) +from pyrit.score.true_false.substring_scorer import SubStringScorer +from pyrit.score.true_false.true_false_composite_scorer import TrueFalseCompositeScorer +from pyrit.score.true_false.true_false_inverter_scorer import TrueFalseInverterScorer +from pyrit.score.true_false.true_false_score_aggregator import TrueFalseAggregatorFunc, TrueFalseScoreAggregator +from pyrit.score.true_false.true_false_scorer import TrueFalseScorer + +if TYPE_CHECKING: + from pyrit.score.float_scale.audio_float_scale_scorer import AudioFloatScaleScorer + from pyrit.score.float_scale.video_float_scale_scorer import VideoFloatScaleScorer + from pyrit.score.scorer_evaluation.human_labeled_dataset import ( + HarmHumanLabeledEntry, + HumanLabeledDataset, + HumanLabeledEntry, + ObjectiveHumanLabeledEntry, + ) + from pyrit.score.scorer_evaluation.scorer_evaluator import ( + HarmScorerEvaluator, + ObjectiveScorerEvaluator, + ScorerEvalDatasetFiles, + ScorerEvaluator, + ) + from pyrit.score.true_false.audio_true_false_scorer import AudioTrueFalseScorer + from pyrit.score.true_false.video_true_false_scorer import VideoTrueFalseScorer + +# Lazy imports for modules with heavy third-party dependencies (PEP 562). +# Audio/video scorers import `av` (~1.9s), human_labeled_dataset imports `pandas` (~1.6s), +# scorer_evaluator imports `scipy.stats` (~1s). +_LAZY_IMPORTS: dict[str, str] = { + "AudioFloatScaleScorer": "pyrit.score.float_scale.audio_float_scale_scorer", + "AudioTrueFalseScorer": "pyrit.score.true_false.audio_true_false_scorer", + "VideoFloatScaleScorer": "pyrit.score.float_scale.video_float_scale_scorer", + "VideoTrueFalseScorer": "pyrit.score.true_false.video_true_false_scorer", + "HarmHumanLabeledEntry": "pyrit.score.scorer_evaluation.human_labeled_dataset", + "HumanLabeledDataset": "pyrit.score.scorer_evaluation.human_labeled_dataset", + "HumanLabeledEntry": "pyrit.score.scorer_evaluation.human_labeled_dataset", + "ObjectiveHumanLabeledEntry": "pyrit.score.scorer_evaluation.human_labeled_dataset", + "HarmScorerEvaluator": "pyrit.score.scorer_evaluation.scorer_evaluator", + "ObjectiveScorerEvaluator": "pyrit.score.scorer_evaluation.scorer_evaluator", + "ScorerEvalDatasetFiles": "pyrit.score.scorer_evaluation.scorer_evaluator", + "ScorerEvaluator": "pyrit.score.scorer_evaluation.scorer_evaluator", +} + + +def __getattr__(name: str) -> object: + if name in _LAZY_IMPORTS: + module = importlib.import_module(_LAZY_IMPORTS[name]) + attr = getattr(module, name) + globals()[name] = attr + return attr + raise AttributeError(f"module {__name__!r} has no attribute {name!r}") + + +__all__ = [ + "AnthraxKeywordScorer", + "AudioFloatScaleScorer", + "AudioTrueFalseScorer", + "AzureContentFilterScorer", + "BatchScorer", + "CallableResponseHandler", + "ContentClassifier", + "ContentClassifierCategory", + "ContentClassifierPaths", + "ConversationScorer", + "CredentialLeakScorer", + "DecodingScorer", + "FentanylKeywordScorer", + "create_conversation_scorer", + "FloatScaleScoreAggregator", + "FloatScaleScorerAllCategories", + "FloatScaleScorerByCategory", + "FloatScaleScorer", + "FloatScaleThresholdScorer", + "GandalfScorer", + "HarmHumanLabeledEntry", + "HarmScorerEvaluator", + "HarmScorerMetrics", + "HumanLabeledDataset", + "HumanLabeledEntry", + "InsecureCodeScorer", + "JsonSchemaResponseHandler", + "LDAPInjectionOutputScorer", + "LikertScaleEvalFiles", + "LikertScale", + "LikertScaleEntry", + "LikertScalePaths", + "LLAMAGUARD_3_CATEGORY_CODES", + "LlamaGuardCategory", + "LlamaGuardMessageRole", + "LlamaGuardPolicy", + "LlamaGuardScorer", + "MarkdownInjectionScorer", + "MethKeywordScorer", + "MetricsType", + "NerveAgentKeywordScorer", + "NumericRange", + "NumericRubric", + "ObjectiveHumanLabeledEntry", + "ObjectiveScorerEvaluator", + "ObjectiveScorerMetrics", + "OpenRedirectOutputScorer", + "parse_llamaguard_response", + "parse_shieldgemma_response", + "PathTraversalOutputScorer", + "PlagiarismMetric", + "PlagiarismScorer", + "PromptShieldScorer", + "QuestionAnswerScorer", + "RegexScorer", + "RegistryUpdateBehavior", + "render_category_system_prompt", + "render_insecure_code_system_prompt", + "render_llamaguard_prompt", + "render_likert_system_prompt", + "render_scale_system_prompt", + "render_shieldgemma_prompt", + "render_true_false_system_prompt", + "ResponseHandler", + "Scorer", + "ScorerEvalDatasetFiles", + "ScorerEvaluator", + "ScorerMetrics", + "ScorerMetricsWithIdentity", + "get_all_harm_metrics", + "get_all_objective_metrics", + "get_scorer_info", + "find_objective_metrics_by_eval_hash", + "ScorerPromptValidator", + "SelfAskCategoryScorer", + "SelfAskGeneralFloatScaleScorer", + "SelfAskGeneralTrueFalseScorer", + "SelfAskLikertScorer", + "SelfAskQuestionAnswerScorer", + "RefusalScorerPaths", + "SelfAskRefusalScorer", + "SelfAskScaleScorer", + "SelfAskTrueFalseScorer", + "ScorerPrinter", + "SHIELDGEMMA_DEFAULT_POLICY_PATH", + "ShieldGemmaGuideline", + "ShieldGemmaMessageRole", + "ShieldGemmaPolicy", + "ShieldGemmaScorer", + "ShellCommandOutputScorer", + "SQLInjectionOutputScorer", + "SSRFOutputScorer", + "SSTIOutputScorer", + "StaticPromptInjectionScorer", + "SubStringScorer", + "TrueFalseCompositeScorer", + "TrueFalseInverterScorer", + "TrueFalseQuestion", + "TrueFalseQuestionPaths", + "TrueFalseScoreAggregator", + "TrueFalseAggregatorFunc", + "TrueFalseScorer", + "VideoFloatScaleScorer", + "VideoTrueFalseScorer", + "XSSOutputScorer", + "XXEOutputScorer", +] From 3530751c6beeec366cc0ce0cd8d94d1015b90623 Mon Sep 17 00:00:00 2001 From: Copilot <223556219+Copilot@users.noreply.github.com> Date: Tue, 11 Aug 2026 15:21:20 -0700 Subject: [PATCH 08/24] Tighten converted view parameter type Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> Copilot-Session: 6776aa95-0924-4967-a056-e5a228af3faa --- pyrit/executor/attack/component/conversation_manager.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pyrit/executor/attack/component/conversation_manager.py b/pyrit/executor/attack/component/conversation_manager.py index 771cb4dc03..ae1e5a4583 100644 --- a/pyrit/executor/attack/component/conversation_manager.py +++ b/pyrit/executor/attack/component/conversation_manager.py @@ -413,7 +413,7 @@ async def _build_converted_normalizer_view_async( self, *, messages: list[Message], - request_converters: list[ConverterConfiguration] | None, + request_converters: list[ConverterConfiguration], apply_to_roles: list[ChatMessageRole], ) -> list[Message]: """ From 18b1df9c15f21bcf2bfbbe061cc9657ce9b9aec9 Mon Sep 17 00:00:00 2001 From: Copilot <223556219+Copilot@users.noreply.github.com> Date: Wed, 12 Aug 2026 05:49:47 -0700 Subject: [PATCH 09/24] Clarify non-chat converter flow Document why non-chat history is converted before flattening, why original and wire views stay separate, and how retry and piece-index safeguards work. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> Copilot-Session: 6776aa95-0924-4967-a056-e5a228af3faa --- .../attack/component/conversation_manager.py | 27 +++++++++++++++++++ 1 file changed, 27 insertions(+) diff --git a/pyrit/executor/attack/component/conversation_manager.py b/pyrit/executor/attack/component/conversation_manager.py index ae1e5a4583..cb6e5b5af0 100644 --- a/pyrit/executor/attack/component/conversation_manager.py +++ b/pyrit/executor/attack/component/conversation_manager.py @@ -364,16 +364,23 @@ async def _handle_non_chat_target_async( Returns: Empty ConversationState (non-chat targets don't track turns). """ + # Context initialization can run again during retry setup. A marked message already contains + # the flattened history and converted live request, so rebuilding it would duplicate the + # prefix and rerun converters that may be stateful or nondeterministic. if context.next_message and context.next_message.request_converters_applied: return ConversationState() normalizer = config.get_message_normalizer() + # Keep separate audit and wire renderings. String normalizers can inspect both value fields, + # so each temporary view collapses them to one representation before flattening. original_context = await self._normalize_non_chat_context_async( messages=self._build_original_normalizer_view(prepended_conversation), normalizer=normalizer, ) converted_context = original_context if request_converters: + # Apply converters while roles are still structural. After flattening, assistant text is + # indistinguishable from user text to a request converter and cannot be safely excluded. converted_messages = await self._build_converted_normalizer_view_async( messages=prepended_conversation, request_converters=request_converters, @@ -384,6 +391,8 @@ async def _handle_non_chat_target_async( normalizer=normalizer, ) + # Build on a copy so a conversion or compatibility failure does not partially mutate the + # attack context. The prepared request is assigned only after every step succeeds. next_message = ( context.next_message.duplicate() if context.next_message @@ -393,6 +402,9 @@ async def _handle_non_chat_target_async( ) ) if request_converters and not next_message.request_converters_applied: + # Convert the live request before attaching history. Letting the normal send path convert + # afterward would apply the converter to the entire flattened string, including roles + # excluded above. The marker tells PromptNormalizer not to run the same chain again. await self._prompt_normalizer.convert_values_async( converter_configurations=request_converters, message=next_message, @@ -422,6 +434,7 @@ async def _build_converted_normalizer_view_async( Returns: list[Message]: Converted message copies ready for string normalization. """ + # Prepended history may also be used by another target path, so conversion must not mutate it. converted_messages = [message.duplicate() for message in messages] if request_converters: for message in converted_messages: @@ -434,6 +447,8 @@ async def _build_converted_normalizer_view_async( source_messages=messages, converted_messages=converted_messages, ) + # ConversationContextNormalizer displays both values when they differ. In this temporary + # wire-only view, align them so flattening emits converted text without audit annotations. for message in converted_messages: for piece in message.message_pieces: piece.original_value = piece.converted_value @@ -459,6 +474,9 @@ def _validate_flattenable_converter_output( converted_message.message_pieces, strict=True, ): + # Existing non-text history already has a defined string representation. Reject only + # modality changes introduced by this conversion pass, which flattening would reduce + # to a placeholder and thereby discard the converter's actual output. converter_was_applied = len(converted_piece.converter_identifiers) > len( source_piece.converter_identifiers ) @@ -501,6 +519,8 @@ async def _normalize_non_chat_context_async( """ messages_to_normalize = messages if isinstance(normalizer, ConversationContextNormalizer): + # ConversationContextNormalizer omits system messages. Squash them into the following + # user message first so non-chat delivery does not silently lose system instructions. messages_to_normalize = await GenericSystemSquashNormalizer().normalize_async(messages) return await normalizer.normalize_string_async(messages_to_normalize) @@ -512,6 +532,8 @@ def _prepend_non_chat_context( converted_context: str, ) -> None: """Prepend original and converted context without mixing their values.""" + # Preserve provenance and wire data independently: memory should retain the unconverted + # conversation while the target receives the role-scoped converted conversation. text_piece = next( ( piece @@ -531,6 +553,8 @@ def _prepend_non_chat_context( ) return + # A multimodal request may have no piece that is text in both views. Add a dedicated text + # piece rather than overwriting or coercing the existing artifact. template_piece = message.get_piece() context_piece = MessagePiece( id=uuid.uuid4(), @@ -738,6 +762,9 @@ async def _apply_converters_async( if message.api_role not in apply_to_roles: return + # Apply to the complete message so ConverterConfiguration.indexes_to_apply remains relative + # to the original piece list. Converting one temporary piece at a time would reset every + # selected piece to index zero. await self._prompt_normalizer.convert_values_async( message=message, converter_configurations=request_converters, From 78f83f18d191b6aef8b977136ad4eb9d0f7ebdb1 Mon Sep 17 00:00:00 2001 From: Copilot <223556219+Copilot@users.noreply.github.com> Date: Wed, 12 Aug 2026 08:27:49 -0700 Subject: [PATCH 10/24] Clarify request converter marker lifetime Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> Copilot-Session: 6776aa95-0924-4967-a056-e5a228af3faa --- pyrit/models/messages/message.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/pyrit/models/messages/message.py b/pyrit/models/messages/message.py index b160820292..779581b1f4 100644 --- a/pyrit/models/messages/message.py +++ b/pyrit/models/messages/message.py @@ -30,6 +30,8 @@ class Message(BaseModel): ) message_pieces: list[MessagePiece] + # Ephemeral guard that keeps context initialization idempotent. + # PrivateAttr excludes it from serialization and DB persistence. _request_converters_applied: bool = PrivateAttr(default=False) # ------------------------------------------------------------------ # From 0b8532e6d814335dd5f6da68e8963fe41e6f2e20 Mon Sep 17 00:00:00 2001 From: Copilot <223556219+Copilot@users.noreply.github.com> Date: Fri, 14 Aug 2026 15:02:14 -0700 Subject: [PATCH 11/24] Move prepended history adaptation to targets Persist prepended conversations structurally for every target, then adapt them only for the first live request when editable history is unavailable. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> Copilot-Session: 6776aa95-0924-4967-a056-e5a228af3faa --- doc/code/framework.md | 4 +- .../attack/component/conversation_manager.py | 331 +++------------ .../prepended_conversation_config.py | 15 +- pyrit/executor/attack/multi_turn/crescendo.py | 2 +- .../executor/attack/multi_turn/red_teaming.py | 2 +- .../attack/multi_turn/tree_of_attacks.py | 3 +- .../attack/single_turn/prompt_sending.py | 2 +- pyrit/message_normalizer/__init__.py | 2 + .../prepended_conversation_normalizer.py | 154 +++++++ pyrit/models/messages/message.py | 19 +- pyrit/prompt_normalizer/prompt_normalizer.py | 40 +- pyrit/prompt_target/common/prompt_target.py | 31 +- .../component/test_conversation_manager.py | 392 +++--------------- .../attack/multi_turn/test_red_teaming.py | 1 + .../test_prompt_normalizer.py | 34 ++ .../test_normalize_async_integration.py | 181 +++++++- 16 files changed, 565 insertions(+), 648 deletions(-) create mode 100644 pyrit/message_normalizer/prepended_conversation_normalizer.py diff --git a/doc/code/framework.md b/doc/code/framework.md index a3938ef2dc..0929a71f77 100644 --- a/doc/code/framework.md +++ b/doc/code/framework.md @@ -314,8 +314,8 @@ 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). 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. +- **`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. Prepended history remains role-structured in memory; when a target cannot edit history, the prompt normalizer passes its formatter to the target for the first live send. +- **`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, prepended-history adaptation, and tokenizer chat templates. These target-specific views are ephemeral and do not replace the logical conversation in memory. ## [Output](./output/0_output) diff --git a/pyrit/executor/attack/component/conversation_manager.py b/pyrit/executor/attack/component/conversation_manager.py index cb6e5b5af0..5c064549aa 100644 --- a/pyrit/executor/attack/component/conversation_manager.py +++ b/pyrit/executor/attack/component/conversation_manager.py @@ -13,11 +13,7 @@ PrependedConversationConfig, ) from pyrit.memory import CentralMemory -from pyrit.message_normalizer import ( - ConversationContextNormalizer, - GenericSystemSquashNormalizer, - MessageStringNormalizer, -) +from pyrit.message_normalizer import ConversationContextNormalizer from pyrit.models import ( ChatMessageRole, ComponentIdentifier, @@ -279,18 +275,12 @@ async def initialize_context_async( This is the primary method for setting up an attack context. It: 1. Merges memory_labels from attack strategy with context labels - 2. Processes prepended_conversation based on target type and config + 2. Persists prepended_conversation structurally with role-scoped converters 3. Updates context.executed_turns for multi-turn attacks - 4. Sets context.next_message if there's an unanswered user message - - For chat-capable PromptTarget: - - Adds prepended messages to memory with simulated_assistant role - - All messages get new UUIDs - For non-chat PromptTarget: - - Applies request converters to configured prepended roles before normalization - - Normalizes the prepended conversation to a string and prepends it to - ``context.next_message`` (using ``config.message_normalizer`` when provided) + For all PromptTarget types, prepended messages are added to memory with + simulated_assistant roles and new UUIDs. Targets without editable history receive + a one-shot formatter that combines this structured history with the first live request. Args: context: The attack context to initialize. @@ -320,21 +310,7 @@ async def initialize_context_async( logger.debug(f"No prepended conversation for context initialization: {conversation_id}") return state - # Targets that don't natively support editable history cannot consume a - # prepended multi-message conversation as-is — route them to the - # single-string fallback path via capability-based routing. - is_chat_target = target.configuration.includes(capability=CapabilityName.EDITABLE_HISTORY) - if not is_chat_target: - config = prepended_conversation_config or PrependedConversationConfig() - return await self._handle_non_chat_target_async( - context=context, - prepended_conversation=prepended_conversation, - config=config, - request_converters=request_converters, - ) - - # Process prepended conversation for objective target - return await self._process_prepended_for_chat_target_async( + return await self._process_prepended_conversation_async( context=context, prepended_conversation=prepended_conversation, conversation_id=conversation_id, @@ -342,244 +318,9 @@ async def initialize_context_async( prepended_conversation_config=prepended_conversation_config, max_turns=max_turns, target_identifier=target.get_identifier(), + target=target, ) - async def _handle_non_chat_target_async( - self, - *, - context: AttackContext[Any], - prepended_conversation: list[Message], - config: PrependedConversationConfig, - request_converters: list[ConverterConfiguration] | None, - ) -> ConversationState: - """ - Handle prepended conversation for non-chat targets. - - Args: - context: The attack context. - prepended_conversation: Messages to prepend. - config: Configuration for non-chat target behavior. - request_converters: Converters to apply before flattening. - - Returns: - Empty ConversationState (non-chat targets don't track turns). - """ - # Context initialization can run again during retry setup. A marked message already contains - # the flattened history and converted live request, so rebuilding it would duplicate the - # prefix and rerun converters that may be stateful or nondeterministic. - if context.next_message and context.next_message.request_converters_applied: - return ConversationState() - - normalizer = config.get_message_normalizer() - # Keep separate audit and wire renderings. String normalizers can inspect both value fields, - # so each temporary view collapses them to one representation before flattening. - original_context = await self._normalize_non_chat_context_async( - messages=self._build_original_normalizer_view(prepended_conversation), - normalizer=normalizer, - ) - converted_context = original_context - if request_converters: - # Apply converters while roles are still structural. After flattening, assistant text is - # indistinguishable from user text to a request converter and cannot be safely excluded. - converted_messages = await self._build_converted_normalizer_view_async( - messages=prepended_conversation, - request_converters=request_converters, - apply_to_roles=config.apply_converters_to_roles, - ) - converted_context = await self._normalize_non_chat_context_async( - messages=converted_messages, - normalizer=normalizer, - ) - - # Build on a copy so a conversion or compatibility failure does not partially mutate the - # attack context. The prepared request is assigned only after every step succeeds. - next_message = ( - context.next_message.duplicate() - if context.next_message - else Message.from_prompt( - prompt=context.objective, - role="user", - ) - ) - if request_converters and not next_message.request_converters_applied: - # Convert the live request before attaching history. Letting the normal send path convert - # afterward would apply the converter to the entire flattened string, including roles - # excluded above. The marker tells PromptNormalizer not to run the same chain again. - await self._prompt_normalizer.convert_values_async( - converter_configurations=request_converters, - message=next_message, - ) - next_message.mark_request_converters_applied() - - self._prepend_non_chat_context( - message=next_message, - original_context=original_context, - converted_context=converted_context, - ) - context.next_message = next_message - - logger.debug(f"Normalized prepended conversation for non-chat target: {len(converted_context)} characters") - return ConversationState() - - async def _build_converted_normalizer_view_async( - self, - *, - messages: list[Message], - request_converters: list[ConverterConfiguration], - apply_to_roles: list[ChatMessageRole], - ) -> list[Message]: - """ - Build copies containing only the values that should be sent. - - Returns: - list[Message]: Converted message copies ready for string normalization. - """ - # Prepended history may also be used by another target path, so conversion must not mutate it. - converted_messages = [message.duplicate() for message in messages] - if request_converters: - for message in converted_messages: - await self._apply_converters_async( - message=message, - request_converters=request_converters, - apply_to_roles=apply_to_roles, - ) - self._validate_flattenable_converter_output( - source_messages=messages, - converted_messages=converted_messages, - ) - # ConversationContextNormalizer displays both values when they differ. In this temporary - # wire-only view, align them so flattening emits converted text without audit annotations. - for message in converted_messages: - for piece in message.message_pieces: - piece.original_value = piece.converted_value - piece.original_value_data_type = piece.converted_value_data_type - return converted_messages - - @staticmethod - def _validate_flattenable_converter_output( - *, - source_messages: list[Message], - converted_messages: list[Message], - ) -> None: - """ - Reject converted history that a string normalizer cannot preserve. - - Raises: - ValueError: If an applied converter produced non-text prepended history. - """ - output_types: set[str] = set() - for source_message, converted_message in zip(source_messages, converted_messages, strict=True): - for source_piece, converted_piece in zip( - source_message.message_pieces, - converted_message.message_pieces, - strict=True, - ): - # Existing non-text history already has a defined string representation. Reject only - # modality changes introduced by this conversion pass, which flattening would reduce - # to a placeholder and thereby discard the converter's actual output. - converter_was_applied = len(converted_piece.converter_identifiers) > len( - source_piece.converter_identifiers - ) - if converter_was_applied and converted_piece.converted_value_data_type != "text": - output_types.add(converted_piece.converted_value_data_type) - - if output_types: - raise ValueError( - "Cannot flatten prepended conversation for a non-chat target after request converters " - f"produced non-text output types {sorted(output_types)}. Role-scoped prepended conversion " - "must produce text; use text-output converters or a chat target with editable history." - ) - - @staticmethod - def _build_original_normalizer_view(messages: list[Message]) -> list[Message]: - """ - Build copies containing only the original values. - - Returns: - list[Message]: Message copies with converted fields reset to their originals. - """ - original_messages = [message.duplicate() for message in messages] - for message in original_messages: - for piece in message.message_pieces: - piece.converted_value = piece.original_value - piece.converted_value_data_type = piece.original_value_data_type - return original_messages - - @staticmethod - async def _normalize_non_chat_context_async( - *, - messages: list[Message], - normalizer: MessageStringNormalizer, - ) -> str: - """ - Flatten role-separated messages into a context string. - - Returns: - str: The flattened conversation context. - """ - messages_to_normalize = messages - if isinstance(normalizer, ConversationContextNormalizer): - # ConversationContextNormalizer omits system messages. Squash them into the following - # user message first so non-chat delivery does not silently lose system instructions. - messages_to_normalize = await GenericSystemSquashNormalizer().normalize_async(messages) - return await normalizer.normalize_string_async(messages_to_normalize) - - @staticmethod - def _prepend_non_chat_context( - *, - message: Message, - original_context: str, - converted_context: str, - ) -> None: - """Prepend original and converted context without mixing their values.""" - # Preserve provenance and wire data independently: memory should retain the unconverted - # conversation while the target receives the role-scoped converted conversation. - text_piece = next( - ( - piece - for piece in message.message_pieces - if piece.original_value_data_type == "text" and piece.converted_value_data_type == "text" - ), - None, - ) - if text_piece: - text_piece.original_value = ConversationManager._prepend_context_value( - context=original_context, - value=text_piece.original_value, - ) - text_piece.converted_value = ConversationManager._prepend_context_value( - context=converted_context, - value=text_piece.converted_value, - ) - return - - # A multimodal request may have no piece that is text in both views. Add a dedicated text - # piece rather than overwriting or coercing the existing artifact. - template_piece = message.get_piece() - context_piece = MessagePiece( - id=uuid.uuid4(), - role=template_piece.role, - original_value=original_context, - converted_value=converted_context, - original_value_data_type="text", - converted_value_data_type="text", - conversation_id=template_piece.conversation_id, - sequence=template_piece.sequence, - ) - message.message_pieces.insert(0, context_piece) - - @staticmethod - def _prepend_context_value(*, context: str, value: str) -> str: - """ - Prepend context once to a message value. - - Returns: - str: The value prefixed with context when it was not already present. - """ - if not context or value == context or value.startswith(f"{context}\n\n"): - return value - return f"{context}\n\n{value}" - async def add_prepended_conversation_to_memory_async( self, *, @@ -589,9 +330,10 @@ async def add_prepended_conversation_to_memory_async( prepended_conversation_config: PrependedConversationConfig | None = None, max_turns: int | None = None, target_identifier: ComponentIdentifier | None = None, + target: PromptTarget | None = None, ) -> int: """ - Add prepended conversation messages to memory for a chat target. + Add prepended conversation messages to memory for a target. This is a lower-level method that handles adding messages to memory without modifying any attack context state. It can be called directly by attacks @@ -611,6 +353,8 @@ async def add_prepended_conversation_to_memory_async( max_turns: If provided, validates that turn count doesn't exceed this limit. target_identifier (ComponentIdentifier | None): The target the conversation is held with, if known. Recorded once per conversation. + target (PromptTarget | None): Target that will receive the first live request. When it + lacks editable history, its target-normalization path receives the configured formatter. Returns: The number of turns (assistant messages) added. @@ -623,6 +367,9 @@ async def add_prepended_conversation_to_memory_async( if not valid_messages: return 0 + if target and target_identifier is None: + target_identifier = target.get_identifier() + self._memory.add_conversation_to_memory( conversation=Conversation(conversation_id=conversation_id, target_identifier=target_identifier) ) @@ -631,6 +378,9 @@ async def add_prepended_conversation_to_memory_async( # path must use the same safe role default as an explicit default config. config = prepended_conversation_config or PrependedConversationConfig() apply_to_roles = config.apply_converters_to_roles + requires_prepended_adaptation = bool( + target and not target.configuration.includes(capability=CapabilityName.EDITABLE_HISTORY) + ) turn_count = 0 @@ -659,14 +409,25 @@ async def add_prepended_conversation_to_memory_async( request_converters=request_converters, apply_to_roles=apply_to_roles, ) + if requires_prepended_adaptation: + self._validate_flattenable_converter_output( + source_message=message, + converted_message=message_copy, + ) # Add to memory self._memory.add_message_to_memory(request=message_copy) logger.debug(f"Added prepended message {i + 1}/{len(valid_messages)} to memory") + if requires_prepended_adaptation: + self._prompt_normalizer.register_prepended_conversation_normalizer( + conversation_id=conversation_id, + message_normalizer=config.get_message_normalizer(), + ) + return turn_count - async def _process_prepended_for_chat_target_async( + async def _process_prepended_conversation_async( self, *, context: AttackContext[Any], @@ -676,9 +437,10 @@ async def _process_prepended_for_chat_target_async( prepended_conversation_config: PrependedConversationConfig | None, max_turns: int | None, target_identifier: ComponentIdentifier | None = None, + target: PromptTarget, ) -> ConversationState: """ - Process prepended conversation for a chat target. + Process prepended conversation for a target. Adds messages to memory with: - New UUIDs for all pieces @@ -694,6 +456,7 @@ async def _process_prepended_for_chat_target_async( max_turns: Maximum turns for validation. target_identifier (ComponentIdentifier | None): The objective target the conversation is held with, if known. + target: The objective target that will receive the conversation. Returns: ConversationState with turn_count and scores. @@ -714,6 +477,7 @@ async def _process_prepended_for_chat_target_async( prepended_conversation_config=prepended_conversation_config, max_turns=max_turns, target_identifier=target_identifier, + target=target, ) # Update context for multi-turn attacks to reflect prepended_conversation @@ -744,6 +508,35 @@ async def _process_prepended_for_chat_target_async( return state + @staticmethod + def _validate_flattenable_converter_output( + *, + source_message: Message, + converted_message: Message, + ) -> None: + """ + Reject non-text output produced by this prepended conversion pass. + + Raises: + ValueError: If an applied converter produced non-text prepended history. + """ + output_types = { + converted_piece.converted_value_data_type + for source_piece, converted_piece in zip( + source_message.message_pieces, + converted_message.message_pieces, + strict=True, + ) + if len(converted_piece.converter_identifiers) > len(source_piece.converter_identifiers) + and converted_piece.converted_value_data_type != "text" + } + if output_types: + raise ValueError( + "Cannot flatten prepended conversation for a target without editable history after " + f"request converters produced non-text output types {sorted(output_types)}. Prepended " + "conversion must produce text." + ) + async def _apply_converters_async( self, *, diff --git a/pyrit/executor/attack/component/prepended_conversation_config.py b/pyrit/executor/attack/component/prepended_conversation_config.py index e03f21adbc..0a511f8489 100644 --- a/pyrit/executor/attack/component/prepended_conversation_config.py +++ b/pyrit/executor/attack/component/prepended_conversation_config.py @@ -23,12 +23,13 @@ class PrependedConversationConfig: This class provides control over: - Which message roles should have request converters applied - - How to normalize conversation history for non-chat objective targets + - How targets without editable history format prepended messages on the first live send - Non-chat objective targets apply request converters to the configured roles before - normalizing the prepended conversation into the first turn (via ``message_normalizer``; - default: ConversationContextNormalizer). Those converters must produce text because - string normalization cannot preserve converted image, audio, or other non-text output. + Prepended messages remain role-structured in memory. Request converters are applied to + configured roles before a target without editable history renders that history into the + first live request (via ``message_normalizer``; default: ConversationContextNormalizer). + Those converters must produce text because string normalization cannot preserve converted + image, audio, or other non-text output. """ # Request converters default to prepended user messages only. Assistant history is @@ -37,8 +38,8 @@ class PrependedConversationConfig: # Optional normalizer to format conversation history into a single text block. # Must implement MessageStringNormalizer (e.g., TokenizerTemplateNormalizer or ConversationContextNormalizer). - # When None and normalization is needed (e.g., for non-chat targets), a default - # ConversationContextNormalizer is used that produces "Turn N: User/Assistant" format. + # When None and adaptation is needed, a default ConversationContextNormalizer is used + # that produces "Turn N: User/Assistant" format. message_normalizer: MessageStringNormalizer | None = None def get_message_normalizer(self) -> MessageStringNormalizer: diff --git a/pyrit/executor/attack/multi_turn/crescendo.py b/pyrit/executor/attack/multi_turn/crescendo.py index 584065aa87..f1feee5d64 100644 --- a/pyrit/executor/attack/multi_turn/crescendo.py +++ b/pyrit/executor/attack/multi_turn/crescendo.py @@ -181,7 +181,7 @@ def __init__( max_turns (int): Maximum number of turns allowed. prepended_conversation_config (PrependedConversationConfiguration | None): Configuration for how to process prepended conversations. Controls converter - application by role, message normalization, and non-chat target behavior. + application by role and first-send formatting for targets without editable history. Raises: ValueError: If objective_target does not natively support editable history. diff --git a/pyrit/executor/attack/multi_turn/red_teaming.py b/pyrit/executor/attack/multi_turn/red_teaming.py index 402eb1303b..c05d2c8994 100644 --- a/pyrit/executor/attack/multi_turn/red_teaming.py +++ b/pyrit/executor/attack/multi_turn/red_teaming.py @@ -171,7 +171,7 @@ def __init__( # Initialize utilities self._prompt_normalizer = prompt_normalizer or PromptNormalizer() - self._conversation_manager = ConversationManager() + self._conversation_manager = ConversationManager(prompt_normalizer=self._prompt_normalizer) # set the maximum number of turns for the attack if max_turns <= 0: diff --git a/pyrit/executor/attack/multi_turn/tree_of_attacks.py b/pyrit/executor/attack/multi_turn/tree_of_attacks.py index 639f0e2427..9cd0550231 100644 --- a/pyrit/executor/attack/multi_turn/tree_of_attacks.py +++ b/pyrit/executor/attack/multi_turn/tree_of_attacks.py @@ -464,6 +464,7 @@ async def initialize_with_prepended_conversation_async( request_converters=self._request_converters, prepended_conversation_config=prepended_conversation_config, target_identifier=self._objective_target.get_identifier(), + target=self._objective_target, ) # Build context string for adversarial chat system prompt (like Crescendo) @@ -1373,7 +1374,7 @@ def __init__( batch_size (int): Number of nodes to process in parallel per batch. Defaults to 10. prepended_conversation_config (PrependedConversationConfig | None): Configuration for how to process prepended conversations. Controls converter - application by role, message normalization, and non-chat target behavior. + application by role and first-send formatting for targets without editable history. Raises: ValueError: If attack_scoring_config uses a non-FloatScaleThresholdScorer objective scorer, diff --git a/pyrit/executor/attack/single_turn/prompt_sending.py b/pyrit/executor/attack/single_turn/prompt_sending.py index ffebec0f2a..1736bfc19d 100644 --- a/pyrit/executor/attack/single_turn/prompt_sending.py +++ b/pyrit/executor/attack/single_turn/prompt_sending.py @@ -76,7 +76,7 @@ def __init__( a params type that rejects certain fields. prepended_conversation_config (PrependedConversationConfiguration | None): Configuration for how to process prepended conversations. Controls converter - application by role, message normalization, and non-chat target behavior. + application by role and first-send formatting for targets without editable history. Request converters apply to prepended user messages by default; include ``"assistant"`` explicitly to transform simulated assistant history. diff --git a/pyrit/message_normalizer/__init__.py b/pyrit/message_normalizer/__init__.py index 79df1cb50d..2ef7e1bc39 100644 --- a/pyrit/message_normalizer/__init__.py +++ b/pyrit/message_normalizer/__init__.py @@ -14,6 +14,7 @@ MessageListNormalizer, MessageStringNormalizer, ) +from pyrit.message_normalizer.prepended_conversation_normalizer import PrependedConversationNormalizer from pyrit.message_normalizer.tokenizer_template_normalizer import TokenizerTemplateNormalizer __all__ = [ @@ -22,6 +23,7 @@ "GenericSystemSquashNormalizer", "HistorySquashNormalizer", "JsonSchemaNormalizer", + "PrependedConversationNormalizer", "TokenizerTemplateNormalizer", "ConversationContextNormalizer", "ChatMessageNormalizer", diff --git a/pyrit/message_normalizer/prepended_conversation_normalizer.py b/pyrit/message_normalizer/prepended_conversation_normalizer.py new file mode 100644 index 0000000000..3640095a4b --- /dev/null +++ b/pyrit/message_normalizer/prepended_conversation_normalizer.py @@ -0,0 +1,154 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT license. + +import copy +import uuid + +from pyrit.message_normalizer.conversation_context_normalizer import ConversationContextNormalizer +from pyrit.message_normalizer.generic_system_squash import GenericSystemSquashNormalizer +from pyrit.message_normalizer.message_normalizer import MessageListNormalizer, MessageStringNormalizer +from pyrit.models import Message, MessagePiece + + +class PrependedConversationNormalizer(MessageListNormalizer[Message]): + """ + Combine prepended history with the first live request for targets without editable history. + + The history remains structured in memory. This normalizer creates an ephemeral target view + that preserves the live request's modalities while prefixing independently rendered original + and converted history. + """ + + def __init__(self, *, message_normalizer: MessageStringNormalizer) -> None: + """ + Initialize the adapter. + + Args: + message_normalizer: Formatter used to render prepended history. + """ + self._message_normalizer = message_normalizer + + async def normalize_async(self, messages: list[Message]) -> list[Message]: + """ + Normalize prepended history into the final request message. + + Args: + messages: Prepended history followed by the first live request. + + Returns: + A single request message containing the rendered history. + """ + if len(messages) < 2: + return copy.deepcopy(messages) + + prepended_messages = messages[:-1] + self._validate_flattenable_converter_output(messages=prepended_messages) + + original_context = await self._normalize_context_async( + messages=self._build_original_view(messages=prepended_messages) + ) + converted_context = original_context + if self._contains_converted_values(messages=prepended_messages): + converted_context = await self._normalize_context_async( + messages=self._build_converted_view(messages=prepended_messages) + ) + + request = copy.deepcopy(messages[-1]) + self._prepend_context( + message=request, + original_context=original_context, + converted_context=converted_context, + ) + return [request] + + async def _normalize_context_async(self, *, messages: list[Message]) -> str: + messages_to_normalize = messages + if isinstance(self._message_normalizer, ConversationContextNormalizer): + messages_to_normalize = await GenericSystemSquashNormalizer().normalize_async(messages) + return await self._message_normalizer.normalize_string_async(messages_to_normalize) + + @staticmethod + def _build_original_view(*, messages: list[Message]) -> list[Message]: + original_messages = copy.deepcopy(messages) + for message in original_messages: + for piece in message.message_pieces: + piece.converted_value = piece.original_value + piece.converted_value_data_type = piece.original_value_data_type + return original_messages + + @staticmethod + def _build_converted_view(*, messages: list[Message]) -> list[Message]: + converted_messages = copy.deepcopy(messages) + for message in converted_messages: + for piece in message.message_pieces: + piece.original_value = piece.converted_value + piece.original_value_data_type = piece.converted_value_data_type + return converted_messages + + @staticmethod + def _contains_converted_values(*, messages: list[Message]) -> bool: + return any( + piece.converter_identifiers + or piece.original_value != piece.converted_value + or piece.original_value_data_type != piece.converted_value_data_type + for message in messages + for piece in message.message_pieces + ) + + @staticmethod + def _validate_flattenable_converter_output(*, messages: list[Message]) -> None: + output_types = { + piece.converted_value_data_type + for message in messages + for piece in message.message_pieces + if piece.converted_value_data_type != "text" + and piece.converted_value_data_type != piece.original_value_data_type + } + if output_types: + raise ValueError( + "Cannot flatten prepended conversation after request converters produced " + f"non-text output types {sorted(output_types)}. Prepended conversion must produce " + "text for a target without editable history." + ) + + @staticmethod + def _prepend_context(*, message: Message, original_context: str, converted_context: str) -> None: + text_piece = next( + ( + piece + for piece in message.message_pieces + if piece.original_value_data_type == "text" and piece.converted_value_data_type == "text" + ), + None, + ) + if text_piece: + text_piece.original_value = PrependedConversationNormalizer._prepend_context_value( + context=original_context, + value=text_piece.original_value, + ) + text_piece.converted_value = PrependedConversationNormalizer._prepend_context_value( + context=converted_context, + value=text_piece.converted_value, + ) + return + + template_piece = message.get_piece() + message.message_pieces.insert( + 0, + MessagePiece( + id=uuid.uuid4(), + role=template_piece.role, + original_value=original_context, + converted_value=converted_context, + original_value_data_type="text", + converted_value_data_type="text", + conversation_id=template_piece.conversation_id, + sequence=template_piece.sequence, + ), + ) + + @staticmethod + def _prepend_context_value(*, context: str, value: str) -> str: + if not context or value == context or value.startswith(f"{context}\n\n"): + return value + return f"{context}\n\n{value}" diff --git a/pyrit/models/messages/message.py b/pyrit/models/messages/message.py index 779581b1f4..7e0a4e301b 100644 --- a/pyrit/models/messages/message.py +++ b/pyrit/models/messages/message.py @@ -8,7 +8,7 @@ from datetime import datetime, timezone from typing import TYPE_CHECKING, cast -from pydantic import BaseModel, ConfigDict, PrivateAttr, model_validator +from pydantic import BaseModel, ConfigDict, model_validator from pyrit.models.messages.message_piece import MessagePiece @@ -30,9 +30,6 @@ class Message(BaseModel): ) message_pieces: list[MessagePiece] - # Ephemeral guard that keeps context initialization idempotent. - # PrivateAttr excludes it from serialization and DB persistence. - _request_converters_applied: bool = PrivateAttr(default=False) # ------------------------------------------------------------------ # # Validators @@ -143,15 +140,6 @@ def get_piece(self, n: int = 0) -> MessagePiece: return self.message_pieces[n] - @property - def request_converters_applied(self) -> bool: - """Whether request converters have already been applied.""" - return self._request_converters_applied - - def mark_request_converters_applied(self) -> None: - """Mark this message as already processed by its request converters.""" - self._request_converters_applied = True - def get_pieces_by_type( self, *, @@ -383,7 +371,4 @@ def duplicate(self) -> Message: piece.id = uuid.uuid4() piece.timestamp = new_timestamp # original_prompt_id intentionally kept the same to track the origin - duplicate = Message(message_pieces=new_pieces) - if self.request_converters_applied: - duplicate.mark_request_converters_applied() - return duplicate + return Message(message_pieces=new_pieces) diff --git a/pyrit/prompt_normalizer/prompt_normalizer.py b/pyrit/prompt_normalizer/prompt_normalizer.py index bc5507ad34..24d995812e 100644 --- a/pyrit/prompt_normalizer/prompt_normalizer.py +++ b/pyrit/prompt_normalizer/prompt_normalizer.py @@ -19,6 +19,7 @@ get_execution_context, ) from pyrit.memory import CentralMemory, MemoryInterface, set_message_piece_sha256_async +from pyrit.message_normalizer import MessageStringNormalizer from pyrit.models import ( ComponentIdentifier, Conversation, @@ -62,6 +63,22 @@ def __init__(self, start_token: str = "⟪", end_token: str = "⟫") -> None: self._start_token = start_token self._end_token = end_token self.id = str(uuid4()) + self._prepended_conversation_normalizers: dict[str, MessageStringNormalizer] = {} + + def register_prepended_conversation_normalizer( + self, + *, + conversation_id: str, + message_normalizer: MessageStringNormalizer, + ) -> None: + """ + Register the formatter used to deliver structured prepended history on the next send. + + Args: + conversation_id: Conversation whose next request should include prepended history. + message_normalizer: Formatter the target should use for that history. + """ + self._prepended_conversation_normalizers[conversation_id] = message_normalizer async def send_prompt_async( self, @@ -108,20 +125,27 @@ async def send_prompt_async( for piece in request.message_pieces: piece.conversation_id = conversation_id - # A caller may need to apply converters before a lossy normalization step - # such as flattening role-separated history for a non-chat target. - if not request.request_converters_applied: - await self.convert_values_async( - converter_configurations=request_converter_configurations, - message=request, - ) + prepended_conversation_normalizer = self._prepended_conversation_normalizers.pop( + request.conversation_id, + None, + ) + await self.convert_values_async( + converter_configurations=request_converter_configurations, + message=request, + ) await self._calc_hash_async(request=request) responses = None try: - responses = await target.send_prompt_async(message=request) + if prepended_conversation_normalizer: + responses = await target.send_prompt_async( + message=request, + prepended_conversation_normalizer=prepended_conversation_normalizer, + ) + else: + responses = await target.send_prompt_async(message=request) self.memory.add_message_to_memory(request=request) except EmptyResponseException: # Empty responses are retried, but we don't want them to stop execution diff --git a/pyrit/prompt_target/common/prompt_target.py b/pyrit/prompt_target/common/prompt_target.py index aaf918f4cf..99b8f5c330 100644 --- a/pyrit/prompt_target/common/prompt_target.py +++ b/pyrit/prompt_target/common/prompt_target.py @@ -6,6 +6,7 @@ from typing import Any, ClassVar, Literal, final from pyrit.memory import CentralMemory, MemoryInterface +from pyrit.message_normalizer import MessageStringNormalizer, PrependedConversationNormalizer from pyrit.models import ( ComponentIdentifier, Conversation, @@ -132,7 +133,12 @@ def __init__( logging.basicConfig(level=logging.INFO) @final - async def send_prompt_async(self, *, message: Message) -> list[Message]: + async def send_prompt_async( + self, + *, + message: Message, + prepended_conversation_normalizer: MessageStringNormalizer | None = None, + ) -> list[Message]: """ Validate, normalize, and send a prompt to the target. @@ -149,6 +155,8 @@ async def send_prompt_async(self, *, message: Message) -> list[Message]: Args: message (Message): The message to send. + prepended_conversation_normalizer (MessageStringNormalizer | None): Optional one-shot + formatter for structured prepended history on a target without editable history. Returns: list[Message]: Response messages from the target. @@ -157,7 +165,10 @@ async def send_prompt_async(self, *, message: Message) -> list[Message]: ValueError: If the message or normalized conversation are empty. """ message.validate() - normalized_conversation = await self._get_normalized_conversation_async(message=message) + normalized_conversation = await self._get_normalized_conversation_async( + message=message, + prepended_conversation_normalizer=prepended_conversation_normalizer, + ) if not normalized_conversation: raise ValueError("Normalization pipeline returned an empty conversation. Cannot send an empty request.") self._validate_request(normalized_conversation=normalized_conversation) @@ -220,7 +231,12 @@ def _validate_request(self, *, normalized_conversation: list[Message]) -> None: if not self.configuration.includes(capability=CapabilityName.MULTI_TURN) and len(normalized_conversation) > 1: raise ValueError(f"This target only supports a single turn conversation. {custom_configuration_message}") - async def _get_normalized_conversation_async(self, *, message: Message) -> list[Message]: + async def _get_normalized_conversation_async( + self, + *, + message: Message, + prepended_conversation_normalizer: MessageStringNormalizer | None = None, + ) -> list[Message]: """ Fetch the conversation from memory, append the current message, and run the normalization pipeline. @@ -235,6 +251,9 @@ async def _get_normalized_conversation_async(self, *, message: Message) -> list[ Args: message (Message): The current message to append. + prepended_conversation_normalizer (MessageStringNormalizer | None): Optional formatter + that combines the existing prepended history with this request before the standard + capability pipeline runs. Returns: list[Message]: The normalized conversation (possibly with system prompt squashed, @@ -245,6 +264,12 @@ async def _get_normalized_conversation_async(self, *, message: Message) -> list[ list(self._memory.get_conversation_messages(conversation_id=conversation_id)) if conversation_id else [] ) conversation.append(message) + if prepended_conversation_normalizer and not self.configuration.includes( + capability=CapabilityName.EDITABLE_HISTORY + ): + conversation = await PrependedConversationNormalizer( + message_normalizer=prepended_conversation_normalizer + ).normalize_async(conversation) normalized = await self.configuration.normalize_async(messages=conversation) if normalized: # Normalizers may create new Message objects (via Message.from_prompt) with diff --git a/tests/unit/executor/attack/component/test_conversation_manager.py b/tests/unit/executor/attack/component/test_conversation_manager.py index 1a9d8569aa..8ddca07ef0 100644 --- a/tests/unit/executor/attack/component/test_conversation_manager.py +++ b/tests/unit/executor/attack/component/test_conversation_manager.py @@ -62,7 +62,7 @@ class _TestAttackContext(AttackContext): class _ImageOutputConverter(Converter): - """A deterministic text-to-image converter for non-chat flattening tests.""" + """A deterministic text-to-image converter for prepended-history adaptation tests.""" SUPPORTED_INPUT_TYPES: tuple[PromptDataType, ...] = ("text",) SUPPORTED_OUTPUT_TYPES: tuple[PromptDataType, ...] = ("image_path",) @@ -71,20 +71,6 @@ async def convert_async(self, *, prompt: str, input_type: PromptDataType = "text return ConverterResult(output_text="converted.png", output_type="image_path") -class _CountingTextConverter(Converter): - """A text converter that produces a different result on each call.""" - - SUPPORTED_INPUT_TYPES: tuple[PromptDataType, ...] = ("text",) - SUPPORTED_OUTPUT_TYPES: tuple[PromptDataType, ...] = ("text",) - - def __init__(self) -> None: - self.call_count = 0 - - async def convert_async(self, *, prompt: str, input_type: PromptDataType = "text") -> ConverterResult: - self.call_count += 1 - return ConverterResult(output_text=f"conversion-{self.call_count}<{prompt}>", output_type="text") - - # ============================================================================= # Fixtures # ============================================================================= @@ -736,105 +722,41 @@ async def test_converts_assistant_to_simulated_assistant( assert stored[0].get_piece().role == "simulated_assistant" assert stored[0].get_piece().api_role == "assistant" - async def test_normalizes_for_non_chat_target_by_default( - self, - attack_identifier: ComponentIdentifier, - mock_prompt_target: MagicMock, - sample_conversation: list[Message], - ) -> None: - """Test that prepended conversation is normalized for non-chat targets by default.""" - manager = ConversationManager() - conversation_id = str(uuid.uuid4()) - context = _TestAttackContext(params=AttackParameters(objective="Test objective")) - context.prepended_conversation = sample_conversation - context.next_message = None - - # By default, should normalize (not raise) - matching PrependedConversationConfig field default - await manager.initialize_context_async( - context=context, - target=mock_prompt_target, - conversation_id=conversation_id, - ) - - # next_message should now contain the normalized prepended context - assert context.next_message is not None - text_value = context.next_message.get_piece().original_value - assert len(text_value) > 0 - - async def test_normalizes_for_non_chat_target_when_configured( + async def test_stores_prepended_conversation_for_non_editable_target( self, attack_identifier: ComponentIdentifier, mock_prompt_target: MagicMock, sample_conversation: list[Message], ) -> None: - """Test that non-chat target normalizes prepended conversation when configured.""" manager = ConversationManager() conversation_id = str(uuid.uuid4()) context = _TestAttackContext(params=AttackParameters(objective="Test objective")) context.prepended_conversation = sample_conversation - context.next_message = Message.from_prompt(prompt="Next message", role="user") - config = PrependedConversationConfig() - - await manager.initialize_context_async( + state = await manager.initialize_context_async( context=context, target=mock_prompt_target, conversation_id=conversation_id, - prepended_conversation_config=config, - ) - - # next_message should now contain the prepended context - assert context.next_message is not None - text_value = context.next_message.get_piece().original_value - assert "Next message" in text_value - assert "Hello" in text_value or "doing well" in text_value - - @pytest.mark.parametrize( - "prepended_conversation_config", - [ - None, - PrependedConversationConfig(message_normalizer=ConversationContextNormalizer()), - ], - ) - async def test_system_prompt_for_non_chat_target_preserves_instruction_and_objective( - self, - attack_identifier: ComponentIdentifier, - mock_prompt_target: MagicMock, - prepended_conversation_config: PrependedConversationConfig | None, - ) -> None: - manager = ConversationManager() - context = _TestAttackContext(params=AttackParameters(objective="Explain saponification")) - context.prepended_conversation = [Message.from_system_prompt("You are a chemistry tutor")] - - await manager.initialize_context_async( - context=context, - target=mock_prompt_target, - conversation_id=str(uuid.uuid4()), - prepended_conversation_config=prepended_conversation_config, ) - assert context.next_message is not None - assert context.next_message.get_value() == "Turn 1:\nuser: You are a chemistry tutor\n\nExplain saponification" - - await manager.initialize_context_async( - context=context, - target=mock_prompt_target, - conversation_id=str(uuid.uuid4()), - prepended_conversation_config=prepended_conversation_config, - ) - - assert context.next_message.get_value() == "Turn 1:\nuser: You are a chemistry tutor\n\nExplain saponification" + stored = manager.get_conversation(conversation_id) + assert len(stored) == 2 + assert [message.api_role for message in stored] == ["user", "assistant"] + assert stored[1].get_piece().role == "simulated_assistant" + assert state.turn_count == 1 + assert context.next_message is None - async def test_system_prompt_for_non_chat_target_preserves_supplied_next_message( + async def test_non_editable_target_does_not_rewrite_supplied_next_message( self, attack_identifier: ComponentIdentifier, mock_prompt_target: MagicMock, ) -> None: manager = ConversationManager() + next_message = Message.from_prompt(prompt="Caller-supplied question", role="user") context = _TestAttackContext( params=AttackParameters( objective="Unused objective", - next_message=Message.from_prompt(prompt="Caller-supplied question", role="user"), + next_message=next_message, ) ) context.prepended_conversation = [Message.from_system_prompt("Follow the policy")] @@ -845,48 +767,35 @@ async def test_system_prompt_for_non_chat_target_preserves_supplied_next_message conversation_id=str(uuid.uuid4()), ) - assert context.next_message is not None - assert context.next_message.get_value() == "Turn 1:\nuser: Follow the policy\n\nCaller-supplied question" + assert context.next_message is next_message + assert context.next_message.get_value() == "Caller-supplied question" - async def test_system_prompt_for_non_chat_target_preserves_multimodal_next_message( + async def test_non_editable_target_registers_custom_first_send_formatter( self, attack_identifier: ComponentIdentifier, + mock_prompt_normalizer: MagicMock, mock_prompt_target: MagicMock, + sample_conversation: list[Message], ) -> None: - manager = ConversationManager() - image_piece = MessagePiece( - role="user", - original_value="diagram.png", - original_value_data_type="image_path", - ) - context = _TestAttackContext( - params=AttackParameters( - objective="Unused objective", - next_message=Message(message_pieces=[image_piece]), - ) - ) - context.prepended_conversation = [Message.from_system_prompt("Describe images precisely")] + manager = ConversationManager(prompt_normalizer=mock_prompt_normalizer) + conversation_id = str(uuid.uuid4()) + context = _TestAttackContext(params=AttackParameters(objective="Test objective")) + context.prepended_conversation = sample_conversation + message_normalizer = MagicMock(spec=ConversationContextNormalizer) + config = PrependedConversationConfig(message_normalizer=message_normalizer) await manager.initialize_context_async( context=context, target=mock_prompt_target, - conversation_id=str(uuid.uuid4()), + conversation_id=conversation_id, + prepended_conversation_config=config, ) - assert context.next_message is not None - assert len(context.next_message.message_pieces) == 2 - assert context.next_message.message_pieces[0].converted_value == "Turn 1:\nuser: Describe images precisely" - assert context.next_message.message_pieces[1].converted_value == "diagram.png" - assert context.next_message.message_pieces[1].original_value_data_type == "image_path" - - await manager.initialize_context_async( - context=context, - target=mock_prompt_target, - conversation_id=str(uuid.uuid4()), + mock_prompt_normalizer.register_prepended_conversation_normalizer.assert_called_once_with( + conversation_id=conversation_id, + message_normalizer=message_normalizer, ) - - assert len(context.next_message.message_pieces) == 2 - assert context.next_message.message_pieces[0].converted_value == "Turn 1:\nuser: Describe images precisely" + message_normalizer.normalize_string_async.assert_not_called() async def test_returns_turn_count_for_multi_turn_attacks( self, @@ -1090,36 +999,7 @@ async def test_prepended_conversation_ignores_true_scores( class TestPrependedConversationConfigSettings: """Tests for PrependedConversationConfig settings in initialize_context_async.""" - # ------------------------------------------------------------------------- - # non_chat_target_behavior Tests - # ------------------------------------------------------------------------- - - async def test_non_chat_target_behavior_normalize_is_default( - self, - attack_identifier: ComponentIdentifier, - mock_prompt_target: MagicMock, - sample_conversation: list[Message], - ) -> None: - """Test that non-chat targets normalize by default (no config), matching dataclass field default.""" - manager = ConversationManager() - conversation_id = str(uuid.uuid4()) - context = _TestAttackContext(params=AttackParameters(objective="Test objective")) - context.prepended_conversation = sample_conversation - context.next_message = None - - # Should normalize by default (matching PrependedConversationConfig field default) - await manager.initialize_context_async( - context=context, - target=mock_prompt_target, - conversation_id=conversation_id, - ) - - # next_message should contain normalized context - assert context.next_message is not None - text_value = context.next_message.get_piece().original_value - assert len(text_value) > 0 - - async def test_non_chat_target_converts_selected_roles_before_flattening( + async def test_non_editable_target_converts_selected_roles_before_storage( self, attack_identifier: ComponentIdentifier, mock_prompt_target: MagicMock, @@ -1131,26 +1011,21 @@ async def test_non_chat_target_converts_selected_roles_before_flattening( context.next_message = Message.from_prompt(prompt="live request", role="user") converter_config = ConverterConfiguration.from_converters(converters=[Base64Converter()]) + conversation_id = str(uuid.uuid4()) await manager.initialize_context_async( context=context, target=mock_prompt_target, - conversation_id=str(uuid.uuid4()), + conversation_id=conversation_id, request_converters=converter_config, ) - assert context.next_message is not None - piece = context.next_message.get_piece() + stored = manager.get_conversation(conversation_id) encoded_user = base64.b64encode(b"Hello, how are you?").decode() - encoded_live_request = base64.b64encode(b"live request").decode() - assert piece.original_value == ( - "Turn 1:\nuser: Hello, how are you?\nassistant: I'm doing well, thank you!\n\nlive request" - ) - assert piece.converted_value == ( - f"Turn 1:\nuser: {encoded_user}\nassistant: I'm doing well, thank you!\n\n{encoded_live_request}" - ) - assert context.next_message.request_converters_applied + assert stored[0].get_piece().converted_value == encoded_user + assert stored[1].get_piece().converted_value == "I'm doing well, thank you!" + assert context.next_message.get_piece().converted_value == "live request" - async def test_non_chat_target_converts_assistant_history_only_when_opted_in( + async def test_non_editable_target_converts_assistant_history_only_when_opted_in( self, attack_identifier: ComponentIdentifier, mock_prompt_target: MagicMock, @@ -1163,22 +1038,21 @@ async def test_non_chat_target_converts_assistant_history_only_when_opted_in( converter_config = ConverterConfiguration.from_converters(converters=[Base64Converter()]) config = PrependedConversationConfig(apply_converters_to_roles=["assistant"]) + conversation_id = str(uuid.uuid4()) await manager.initialize_context_async( context=context, target=mock_prompt_target, - conversation_id=str(uuid.uuid4()), + conversation_id=conversation_id, request_converters=converter_config, prepended_conversation_config=config, ) - assert context.next_message is not None + stored = manager.get_conversation(conversation_id) encoded_assistant = base64.b64encode(b"I'm doing well, thank you!").decode() - encoded_live_request = base64.b64encode(b"live request").decode() - assert context.next_message.get_piece().converted_value == ( - f"Turn 1:\nuser: Hello, how are you?\nassistant: {encoded_assistant}\n\n{encoded_live_request}" - ) + assert stored[0].get_piece().converted_value == "Hello, how are you?" + assert stored[1].get_piece().converted_value == encoded_assistant - async def test_non_chat_target_rejects_non_text_history_converter_output( + async def test_non_editable_target_rejects_non_text_output_from_current_converter( self, attack_identifier: ComponentIdentifier, mock_prompt_target: MagicMock, @@ -1189,7 +1063,7 @@ async def test_non_chat_target_rejects_non_text_history_converter_output( context.prepended_conversation = sample_conversation converter_config = ConverterConfiguration.from_converters(converters=[_ImageOutputConverter()]) - with pytest.raises(ValueError, match="non-chat target.*non-text output types.*image_path"): + with pytest.raises(ValueError, match="non-text output types.*image_path"): await manager.initialize_context_async( context=context, target=mock_prompt_target, @@ -1199,40 +1073,7 @@ async def test_non_chat_target_rejects_non_text_history_converter_output( assert sample_conversation[0].get_piece().converted_value_data_type == "text" - async def test_non_chat_target_repeated_setup_reuses_prepared_request( - self, - attack_identifier: ComponentIdentifier, - mock_prompt_target: MagicMock, - sample_conversation: list[Message], - ) -> None: - manager = ConversationManager() - context = _TestAttackContext(params=AttackParameters(objective="Test objective")) - context.prepended_conversation = sample_conversation - context.next_message = Message.from_prompt(prompt="live request", role="user") - converter = _CountingTextConverter() - converter_config = ConverterConfiguration.from_converters(converters=[converter]) - - await manager.initialize_context_async( - context=context, - target=mock_prompt_target, - conversation_id=str(uuid.uuid4()), - request_converters=converter_config, - ) - assert context.next_message is not None - first_prepared_value = context.next_message.get_piece().converted_value - assert converter.call_count == 2 - - await manager.initialize_context_async( - context=context, - target=mock_prompt_target, - conversation_id=str(uuid.uuid4()), - request_converters=converter_config, - ) - - assert context.next_message.get_piece().converted_value == first_prepared_value - assert converter.call_count == 2 - - async def test_non_chat_target_preserves_converter_piece_indexes( + async def test_non_editable_target_preserves_converter_piece_indexes( self, attack_identifier: ComponentIdentifier, mock_prompt_target: MagicMock, @@ -1265,99 +1106,17 @@ async def test_non_chat_target_preserves_converter_piece_indexes( ) ] - await manager.initialize_context_async( - context=context, - target=mock_prompt_target, - conversation_id=str(uuid.uuid4()), - request_converters=converter_config, - ) - - assert context.next_message is not None - converted_value = context.next_message.get_piece().converted_value - assert base64.b64encode(b"first piece").decode() in converted_value - assert "second piece" in converted_value - assert base64.b64encode(b"second piece").decode() not in converted_value - - async def test_non_chat_target_behavior_normalize_first_turn_creates_next_message( - self, - attack_identifier: ComponentIdentifier, - mock_prompt_target: MagicMock, - sample_conversation: list[Message], - ) -> None: - """Test that normalize_first_turn creates next_message when none exists.""" - manager = ConversationManager() - conversation_id = str(uuid.uuid4()) - context = _TestAttackContext(params=AttackParameters(objective="Test objective")) - context.prepended_conversation = sample_conversation - context.next_message = None - - config = PrependedConversationConfig() - - await manager.initialize_context_async( - context=context, - target=mock_prompt_target, - conversation_id=conversation_id, - prepended_conversation_config=config, - ) - - # Should have created a next_message with the normalized context - assert context.next_message is not None - text_value = context.next_message.get_piece().original_value - assert len(text_value) > 0 - - async def test_non_chat_target_behavior_normalize_first_turn_prepends_to_existing_message( - self, - attack_identifier: ComponentIdentifier, - mock_prompt_target: MagicMock, - sample_conversation: list[Message], - ) -> None: - """Test that normalize_first_turn prepends context to existing next_message.""" - manager = ConversationManager() conversation_id = str(uuid.uuid4()) - context = _TestAttackContext(params=AttackParameters(objective="Test objective")) - context.prepended_conversation = sample_conversation - context.next_message = Message.from_prompt(prompt="My question", role="user") - - config = PrependedConversationConfig() - await manager.initialize_context_async( context=context, target=mock_prompt_target, conversation_id=conversation_id, - prepended_conversation_config=config, - ) - - # Should have prepended context to existing message - text_value = context.next_message.get_piece().original_value - assert "My question" in text_value - # Context should come before the original question - question_index = text_value.find("My question") - assert question_index > 0 # Context should be prepended - - async def test_non_chat_target_behavior_normalize_returns_empty_state( - self, - attack_identifier: ComponentIdentifier, - mock_prompt_target: MagicMock, - sample_conversation: list[Message], - ) -> None: - """Test that normalize_first_turn returns empty ConversationState (no turn tracking).""" - manager = ConversationManager() - conversation_id = str(uuid.uuid4()) - context = _TestAttackContext(params=AttackParameters(objective="Test objective")) - context.prepended_conversation = sample_conversation - - config = PrependedConversationConfig() - - state = await manager.initialize_context_async( - context=context, - target=mock_prompt_target, - conversation_id=conversation_id, - prepended_conversation_config=config, + request_converters=converter_config, ) - # Non-chat targets don't track turns - assert state.turn_count == 0 - assert state.last_assistant_message_scores == [] + stored_pieces = manager.get_conversation(conversation_id)[0].message_pieces + assert stored_pieces[0].converted_value == base64.b64encode(b"first piece").decode() + assert stored_pieces[1].converted_value == "second piece" # ------------------------------------------------------------------------- # apply_converters_to_roles Tests @@ -1479,66 +1238,25 @@ async def test_apply_converters_to_roles_empty_list_skips_all( async def test_message_normalizer_default_uses_conversation_context_normalizer( self, attack_identifier: ComponentIdentifier, + mock_prompt_normalizer: MagicMock, mock_prompt_target: MagicMock, sample_conversation: list[Message], ) -> None: - """Test that default normalizer produces Turn N format.""" - manager = ConversationManager() - conversation_id = str(uuid.uuid4()) - context = _TestAttackContext(params=AttackParameters(objective="Test objective")) - context.prepended_conversation = sample_conversation - context.next_message = None - - config = PrependedConversationConfig() - - await manager.initialize_context_async( - context=context, - target=mock_prompt_target, - conversation_id=conversation_id, - prepended_conversation_config=config, - ) - - # Default ConversationContextNormalizer produces "Turn N:" format - assert context.next_message is not None - text_value = context.next_message.get_piece().original_value - assert "Turn 1" in text_value or "turn 1" in text_value.lower() - - async def test_message_normalizer_custom_normalizer_is_used( - self, - attack_identifier: ComponentIdentifier, - mock_prompt_target: MagicMock, - sample_conversation: list[Message], - ) -> None: - """Test that custom message_normalizer is used when provided.""" - from pyrit.message_normalizer import MessageStringNormalizer - - # Create a mock normalizer that returns a specific format - mock_normalizer = MagicMock(spec=MessageStringNormalizer) - mock_normalizer.normalize_string_async = AsyncMock(return_value="CUSTOM_FORMAT: test content") - - manager = ConversationManager() + """Test that the default formatter is registered for target-side adaptation.""" + manager = ConversationManager(prompt_normalizer=mock_prompt_normalizer) conversation_id = str(uuid.uuid4()) context = _TestAttackContext(params=AttackParameters(objective="Test objective")) context.prepended_conversation = sample_conversation - context.next_message = None - - config = PrependedConversationConfig( - message_normalizer=mock_normalizer, - ) await manager.initialize_context_async( context=context, target=mock_prompt_target, conversation_id=conversation_id, - prepended_conversation_config=config, ) - # Verify custom normalizer was called - mock_normalizer.normalize_string_async.assert_called_once() - # Verify the custom format is in the message - assert context.next_message is not None - text_value = context.next_message.get_piece().original_value - assert "CUSTOM_FORMAT: test content" in text_value + registered = mock_prompt_normalizer.register_prepended_conversation_normalizer.call_args.kwargs + assert registered["conversation_id"] == conversation_id + assert isinstance(registered["message_normalizer"], ConversationContextNormalizer) # ------------------------------------------------------------------------- # Chat Target Behavior (Config has no effect) 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 4603152590..51b0f326b7 100644 --- a/tests/unit/executor/attack/multi_turn/test_red_teaming.py +++ b/tests/unit/executor/attack/multi_turn/test_red_teaming.py @@ -298,6 +298,7 @@ def test_init_with_all_custom_configurations( assert attack._request_converters == converter_config.request_converters assert attack._response_converters == converter_config.response_converters assert attack._prompt_normalizer == mock_prompt_normalizer + assert attack._conversation_manager._prompt_normalizer is mock_prompt_normalizer assert attack._max_turns == 20 def test_init_without_objective_scorer_raises_error( diff --git a/tests/unit/prompt_normalizer/test_prompt_normalizer.py b/tests/unit/prompt_normalizer/test_prompt_normalizer.py index 025d6d6c77..01ad9bf435 100644 --- a/tests/unit/prompt_normalizer/test_prompt_normalizer.py +++ b/tests/unit/prompt_normalizer/test_prompt_normalizer.py @@ -25,6 +25,7 @@ get_execution_context, ) from pyrit.memory import CentralMemory +from pyrit.message_normalizer import ConversationContextNormalizer from pyrit.models import ( Message, MessagePiece, @@ -115,6 +116,39 @@ async def test_send_prompt_async_multiple_converters(mock_memory_instance, seed_ assert prompt_target.prompt_sent == ["S_G_V_s_b_G_8_="] +async def test_send_prompt_async_passes_registered_prepended_formatter_once(mock_memory_instance): + prompt_target = MagicMock(spec=PromptTarget) + prompt_target.get_identifier.return_value = get_mock_target_identifier("MockTarget") + prompt_target.send_prompt_async = AsyncMock( + side_effect=[ + [MessagePiece(role="assistant", original_value="first").to_message()], + [MessagePiece(role="assistant", original_value="second").to_message()], + ] + ) + normalizer = PromptNormalizer() + formatter = ConversationContextNormalizer() + conversation_id = "prepended-conversation" + normalizer.register_prepended_conversation_normalizer( + conversation_id=conversation_id, + message_normalizer=formatter, + ) + + await normalizer.send_prompt_async( + message=Message.from_prompt(prompt="first request", role="user"), + target=prompt_target, + conversation_id=conversation_id, + ) + await normalizer.send_prompt_async( + message=Message.from_prompt(prompt="second request", role="user"), + target=prompt_target, + conversation_id=conversation_id, + ) + + first_call, second_call = prompt_target.send_prompt_async.await_args_list + assert first_call.kwargs["prepended_conversation_normalizer"] is formatter + assert "prepended_conversation_normalizer" not in second_call.kwargs + + async def test_send_prompt_async_no_response_adds_memory(mock_memory_instance, seed_group): prompt_target = MagicMock() prompt_target.send_prompt_async = AsyncMock(return_value=None) diff --git a/tests/unit/prompt_target/target/test_normalize_async_integration.py b/tests/unit/prompt_target/target/test_normalize_async_integration.py index 407f1222a2..438d0ab8f8 100644 --- a/tests/unit/prompt_target/target/test_normalize_async_integration.py +++ b/tests/unit/prompt_target/target/test_normalize_async_integration.py @@ -13,9 +13,11 @@ import pytest from openai.types.chat import ChatCompletion from openai.types.responses import ResponseOutputMessage, ResponseOutputText +from unit.mocks import MockPromptTarget from pyrit.memory.memory_interface import MemoryInterface -from pyrit.models import Message, MessagePiece +from pyrit.message_normalizer import ConversationContextNormalizer, MessageStringNormalizer +from pyrit.models import ComponentIdentifier, Message, MessagePiece from pyrit.prompt_target import AzureMLChatTarget, OpenAIChatTarget from pyrit.prompt_target.common.target_capabilities import ( CapabilityHandlingPolicy, @@ -489,3 +491,180 @@ async def test_get_normalized_conversation_passthrough_when_no_adaptation_needed assert result[0].get_value() == "be nice" assert result[1].get_piece().api_role == "user" assert result[1].get_value() == "hello" + + +# --------------------------------------------------------------------------- +# Prepended history adaptation +# --------------------------------------------------------------------------- + + +@pytest.mark.usefixtures("patch_central_database") +async def test_non_editable_target_adapts_prepended_history_without_mutating_memory(): + target = MockPromptTarget() + target._configuration = TargetConfiguration(capabilities=TargetCapabilities()) + + prepended_user = _make_message(role="user", content="original history") + prepended_user.get_piece().converted_value = "converted history" + prepended_user.get_piece().converter_identifiers = [ + ComponentIdentifier(class_name="TestConverter", class_module="tests") + ] + prepended_assistant = _make_message(role="simulated_assistant", content="assistant history") + live_request = _make_message(role="user", content="original live") + live_request.get_piece().converted_value = "converted live" + memory_messages: MutableSequence[Message] = [prepended_user, prepended_assistant] + + mock_memory = MagicMock(spec=MemoryInterface) + mock_memory.get_conversation_messages.return_value = memory_messages + target._memory = mock_memory + + result = await target._get_normalized_conversation_async( + message=live_request, + prepended_conversation_normalizer=ConversationContextNormalizer(), + ) + + assert len(result) == 1 + assert result[0].get_piece().original_value == ( + "Turn 1:\nuser: original history\nassistant: assistant history\n\noriginal live" + ) + assert result[0].get_piece().converted_value == ( + "Turn 1:\nuser: converted history\nassistant: assistant history\n\nconverted live" + ) + assert len(memory_messages) == 2 + assert memory_messages[0].get_piece().original_value == "original history" + + +@pytest.mark.usefixtures("patch_central_database") +async def test_non_editable_target_preserves_system_history_and_multimodal_live_request(): + target = MockPromptTarget() + target._configuration = TargetConfiguration( + capabilities=TargetCapabilities( + supports_multi_message_pieces=True, + input_modalities=frozenset({frozenset({"text", "image_path"})}), + ) + ) + system_message = _make_message(role="system", content="Describe images precisely") + live_request = Message( + message_pieces=[ + MessagePiece( + role="user", + original_value="diagram.png", + converted_value="diagram.png", + original_value_data_type="image_path", + converted_value_data_type="image_path", + conversation_id="conv1", + ) + ] + ) + mock_memory = MagicMock(spec=MemoryInterface) + mock_memory.get_conversation_messages.return_value = [system_message] + target._memory = mock_memory + + result = await target._get_normalized_conversation_async( + message=live_request, + prepended_conversation_normalizer=ConversationContextNormalizer(), + ) + + assert len(result) == 1 + assert len(result[0].message_pieces) == 2 + assert result[0].message_pieces[0].converted_value == "Turn 1:\nuser: Describe images precisely" + assert result[0].message_pieces[1].converted_value == "diagram.png" + assert result[0].message_pieces[1].converted_value_data_type == "image_path" + + +@pytest.mark.usefixtures("patch_central_database") +async def test_prepended_history_adapter_is_used_only_when_explicitly_passed(): + target = MockPromptTarget() + target._configuration = TargetConfiguration( + capabilities=TargetCapabilities( + supports_multi_turn=True, + supports_system_prompt=True, + ) + ) + prepended = _make_message(role="user", content="prepended") + prior_live = _make_message(role="user", content="first live") + prior_response = _make_message(role="assistant", content="first response") + second_live = _make_message(role="user", content="second live") + + mock_memory = MagicMock(spec=MemoryInterface) + mock_memory.get_conversation_messages.side_effect = [ + [prepended], + [prepended, prior_live, prior_response], + ] + target._memory = mock_memory + + await target.send_prompt_async( + message=prior_live, + prepended_conversation_normalizer=ConversationContextNormalizer(), + ) + await target.send_prompt_async(message=second_live) + + assert target.prompt_sent == ["Turn 1:\nuser: prepended\n\nfirst live", "second live"] + + +@pytest.mark.usefixtures("patch_central_database") +async def test_non_editable_target_uses_custom_prepended_formatter(): + target = MockPromptTarget() + target._configuration = TargetConfiguration(capabilities=TargetCapabilities()) + mock_memory = MagicMock(spec=MemoryInterface) + mock_memory.get_conversation_messages.return_value = [_make_message(role="user", content="prepended")] + target._memory = mock_memory + formatter = MagicMock(spec=MessageStringNormalizer) + formatter.normalize_string_async = AsyncMock(return_value="CUSTOM HISTORY") + + result = await target._get_normalized_conversation_async( + message=_make_message(role="user", content="live"), + prepended_conversation_normalizer=formatter, + ) + + assert result[0].get_value() == "CUSTOM HISTORY\n\nlive" + formatter.normalize_string_async.assert_awaited_once() + + +@pytest.mark.usefixtures("patch_central_database") +async def test_non_editable_target_rejects_non_text_converted_prepended_history(): + target = MockPromptTarget() + target._configuration = TargetConfiguration(capabilities=TargetCapabilities()) + prepended = _make_message(role="user", content="original") + prepended.get_piece().converted_value = "converted.png" + prepended.get_piece().converted_value_data_type = "image_path" + prepended.get_piece().converter_identifiers = [ + ComponentIdentifier(class_name="ImageConverter", class_module="tests") + ] + mock_memory = MagicMock(spec=MemoryInterface) + mock_memory.get_conversation_messages.return_value = [prepended] + target._memory = mock_memory + + with pytest.raises(ValueError, match="non-text output types.*image_path"): + await target._get_normalized_conversation_async( + message=_make_message(role="user", content="live"), + prepended_conversation_normalizer=ConversationContextNormalizer(), + ) + + +@pytest.mark.usefixtures("patch_central_database") +async def test_non_editable_target_allows_preexisting_non_text_history_with_converter_provenance(): + target = MockPromptTarget() + target._configuration = TargetConfiguration(capabilities=TargetCapabilities()) + prepended = Message( + message_pieces=[ + MessagePiece( + role="user", + original_value="existing.png", + converted_value="existing.png", + original_value_data_type="image_path", + converted_value_data_type="image_path", + conversation_id="conv1", + converter_identifiers=[ComponentIdentifier(class_name="PriorConverter", class_module="tests")], + ) + ] + ) + mock_memory = MagicMock(spec=MemoryInterface) + mock_memory.get_conversation_messages.return_value = [prepended] + target._memory = mock_memory + + result = await target._get_normalized_conversation_async( + message=_make_message(role="user", content="live"), + prepended_conversation_normalizer=ConversationContextNormalizer(), + ) + + assert result[0].get_value() == "Turn 1:\nuser: [Image_path]\n\nlive" From aff10068187c46ce1b3cf6d1cfdaf0fce600ef44 Mon Sep 17 00:00:00 2001 From: Copilot <223556219+Copilot@users.noreply.github.com> Date: Wed, 19 Aug 2026 07:06:24 -0700 Subject: [PATCH 12/24] Refine prepended history target normalization Replace shared formatter registration with an explicit one-shot target normalization context. Keep prepended history structured in memory, stage role-scoped conversion before persistence, and adapt history with the first live request immediately before provider invocation. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> Copilot-Session: 6776aa95-0924-4967-a056-e5a228af3faa --- doc/code/framework.md | 4 +- .../attack/component/conversation_manager.py | 96 +++-- .../prepended_conversation_config.py | 6 +- pyrit/executor/attack/core/attack_strategy.py | 8 + .../attack/multi_turn/chunked_request.py | 1 + pyrit/executor/attack/multi_turn/crescendo.py | 1 + .../attack/multi_turn/multi_prompt_sending.py | 1 + .../multi_turn/multi_turn_attack_strategy.py | 9 +- .../executor/attack/multi_turn/red_teaming.py | 1 + .../attack/multi_turn/tree_of_attacks.py | 25 +- .../attack/single_turn/prompt_sending.py | 1 + pyrit/message_normalizer/__init__.py | 4 +- .../first_turn_history_normalizer.py | 183 ++++++++++ .../prepended_conversation_normalizer.py | 154 -------- pyrit/prompt_normalizer/prompt_normalizer.py | 40 +-- pyrit/prompt_target/__init__.py | 6 + pyrit/prompt_target/common/prompt_target.py | 85 +++-- .../common/target_normalization_context.py | 138 ++++++++ .../component/test_conversation_manager.py | 160 ++++++++- .../test_supports_multi_turn_attacks.py | 22 ++ .../attack/multi_turn/test_tree_of_attacks.py | 2 + .../attack/single_turn/test_prompt_sending.py | 24 +- .../test_prompt_normalizer.py | 120 ++++++- .../test_normalize_async_integration.py | 335 +++++++++++++++++- .../test_target_normalization_context.py | 60 ++++ 25 files changed, 1201 insertions(+), 285 deletions(-) create mode 100644 pyrit/message_normalizer/first_turn_history_normalizer.py delete mode 100644 pyrit/message_normalizer/prepended_conversation_normalizer.py create mode 100644 pyrit/prompt_target/common/target_normalization_context.py create mode 100644 tests/unit/prompt_target/target/test_target_normalization_context.py diff --git a/doc/code/framework.md b/doc/code/framework.md index 0929a71f77..11fdb02af8 100644 --- a/doc/code/framework.md +++ b/doc/code/framework.md @@ -314,9 +314,11 @@ 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). 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. Prepended history remains role-structured in memory; when a target cannot edit history, the prompt normalizer passes its formatter to the target for the first live send. +- **`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. Prepended history remains role-structured in memory. Attacks pass an ephemeral `TargetNormalizationContext` with the first live send when a target cannot edit history; the target consumes that context immediately before provider invocation. - **`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, prepended-history adaptation, and tokenizer chat templates. These target-specific views are ephemeral and do not replace the logical conversation in memory. +For prepended history on a target without editable history, processing order is: convert and persist the structured prepended messages, convert the live request, apply first-turn history normalization, run the target's ordinary capability normalizers, then serialize and invoke the provider. + ## [Output](./output/0_output) **Responsibility**: The Output module is responsible for writing different components in different formats to different places. diff --git a/pyrit/executor/attack/component/conversation_manager.py b/pyrit/executor/attack/component/conversation_manager.py index 5c064549aa..89470202f3 100644 --- a/pyrit/executor/attack/component/conversation_manager.py +++ b/pyrit/executor/attack/component/conversation_manager.py @@ -13,7 +13,7 @@ PrependedConversationConfig, ) from pyrit.memory import CentralMemory -from pyrit.message_normalizer import ConversationContextNormalizer +from pyrit.message_normalizer import ConversationContextNormalizer, FirstTurnHistoryNormalizer from pyrit.models import ( ChatMessageRole, ComponentIdentifier, @@ -23,7 +23,7 @@ Score, ) from pyrit.prompt_normalizer.prompt_normalizer import PromptNormalizer -from pyrit.prompt_target import CapabilityName, PromptTarget +from pyrit.prompt_target import CapabilityName, PromptTarget, TargetNormalizationContext if TYPE_CHECKING: from collections.abc import Sequence @@ -280,7 +280,7 @@ async def initialize_context_async( For all PromptTarget types, prepended messages are added to memory with simulated_assistant roles and new UUIDs. Targets without editable history receive - a one-shot formatter that combines this structured history with the first live request. + explicit one-shot normalization state on the attack context. Args: context: The attack context to initialize. @@ -302,6 +302,7 @@ async def initialize_context_async( # Merge memory labels: attack strategy labels + context labels context.memory_labels = combine_dict(existing_dict=memory_labels, new_dict=context.memory_labels) + context.target_normalization_context = None state = ConversationState() prepended_conversation = context.prepended_conversation @@ -353,8 +354,7 @@ async def add_prepended_conversation_to_memory_async( max_turns: If provided, validates that turn count doesn't exceed this limit. target_identifier (ComponentIdentifier | None): The target the conversation is held with, if known. Recorded once per conversation. - target (PromptTarget | None): Target that will receive the first live request. When it - lacks editable history, its target-normalization path receives the configured formatter. + target (PromptTarget | None): Target that will receive the first live request. Returns: The number of turns (assistant messages) added. @@ -362,18 +362,13 @@ async def add_prepended_conversation_to_memory_async( Raises: ValueError: If max_turns is exceeded by the prepended conversation. """ - # Filter valid messages - valid_messages = [msg for msg in prepended_conversation if msg and msg.message_pieces] + valid_messages = self.get_persistable_prepended_messages(prepended_conversation=prepended_conversation) if not valid_messages: return 0 if target and target_identifier is None: target_identifier = target.get_identifier() - self._memory.add_conversation_to_memory( - conversation=Conversation(conversation_id=conversation_id, target_identifier=target_identifier) - ) - # Assistant history represents simulated target output, so the absent-config # path must use the same safe role default as an explicit default config. config = prepended_conversation_config or PrependedConversationConfig() @@ -383,8 +378,9 @@ async def add_prepended_conversation_to_memory_async( ) turn_count = 0 + prepared_messages: list[Message] = [] - for i, message in enumerate(valid_messages): + for message in valid_messages: message_copy = message.duplicate() message_copy.set_simulated_role() @@ -415,18 +411,66 @@ async def add_prepended_conversation_to_memory_async( converted_message=message_copy, ) - # Add to memory - self._memory.add_message_to_memory(request=message_copy) - logger.debug(f"Added prepended message {i + 1}/{len(valid_messages)} to memory") + prepared_messages.append(message_copy) - if requires_prepended_adaptation: - self._prompt_normalizer.register_prepended_conversation_normalizer( - conversation_id=conversation_id, - message_normalizer=config.get_message_normalizer(), - ) + self._memory.add_conversation_to_memory( + conversation=Conversation(conversation_id=conversation_id, target_identifier=target_identifier) + ) + for i, message in enumerate(prepared_messages): + self._memory.add_message_to_memory(request=message) + logger.debug(f"Added prepended message {i + 1}/{len(prepared_messages)} to memory") return turn_count + @staticmethod + def get_persistable_prepended_messages( + *, + prepended_conversation: list[Message], + ) -> list[Message]: + """ + Return prepended messages that can be recovered from memory at send time. + + Args: + prepended_conversation: Candidate prepended messages. + + Returns: + list[Message]: Non-empty messages containing at least one persistable piece. + """ + return [ + message + for message in prepended_conversation + if message and message.message_pieces and any(not piece.not_in_memory for piece in message.message_pieces) + ] + + @staticmethod + def create_target_normalization_context( + *, + target: PromptTarget, + conversation_id: str, + prepended_message_count: int, + prepended_conversation_config: PrependedConversationConfig | None = None, + ) -> TargetNormalizationContext | None: + """ + Build first-send normalization state for a target without editable history. + + Returns: + TargetNormalizationContext | None: First-send state, or ``None`` for + editable-history targets or empty prepended history. + """ + if prepended_message_count < 1 or target.configuration.includes(capability=CapabilityName.EDITABLE_HISTORY): + return None + + config = prepended_conversation_config or PrependedConversationConfig() + return TargetNormalizationContext( + conversation_id=conversation_id, + normalizers=( + FirstTurnHistoryNormalizer( + message_normalizer=config.get_message_normalizer(), + prepended_message_count=prepended_message_count, + ), + ), + ) + async def _process_prepended_conversation_async( self, *, @@ -464,11 +508,17 @@ async def _process_prepended_conversation_async( state = ConversationState() is_multi_turn = max_turns is not None - # Filter valid messages - valid_messages = [msg for msg in prepended_conversation if msg and msg.message_pieces] + valid_messages = self.get_persistable_prepended_messages(prepended_conversation=prepended_conversation) if not valid_messages: return state + target_normalization_context = self.create_target_normalization_context( + target=target, + conversation_id=conversation_id, + prepended_message_count=len(valid_messages), + prepended_conversation_config=prepended_conversation_config, + ) + # Use the lower-level method to add messages to memory state.turn_count = await self.add_prepended_conversation_to_memory_async( prepended_conversation=prepended_conversation, @@ -479,6 +529,7 @@ async def _process_prepended_conversation_async( target_identifier=target_identifier, target=target, ) + context.target_normalization_context = target_normalization_context # Update context for multi-turn attacks to reflect prepended_conversation @@ -528,6 +579,7 @@ def _validate_flattenable_converter_output( strict=True, ) if len(converted_piece.converter_identifiers) > len(source_piece.converter_identifiers) + and not converted_piece.not_in_memory and converted_piece.converted_value_data_type != "text" } if output_types: diff --git a/pyrit/executor/attack/component/prepended_conversation_config.py b/pyrit/executor/attack/component/prepended_conversation_config.py index 0a511f8489..e6c2639a18 100644 --- a/pyrit/executor/attack/component/prepended_conversation_config.py +++ b/pyrit/executor/attack/component/prepended_conversation_config.py @@ -26,8 +26,8 @@ class PrependedConversationConfig: - How targets without editable history format prepended messages on the first live send Prepended messages remain role-structured in memory. Request converters are applied to - configured roles before a target without editable history renders that history into the - first live request (via ``message_normalizer``; default: ConversationContextNormalizer). + configured roles before a target without editable history renders that history and the + first live request together (via ``message_normalizer``; default: ConversationContextNormalizer). Those converters must produce text because string normalization cannot preserve converted image, audio, or other non-text output. """ @@ -36,7 +36,7 @@ class PrependedConversationConfig: # simulated target output and must be explicitly opted in with ["assistant"]. apply_converters_to_roles: list[ChatMessageRole] = field(default_factory=lambda: ["user"]) - # Optional normalizer to format conversation history into a single text block. + # Optional normalizer to format prepended history and the first live request as one text block. # Must implement MessageStringNormalizer (e.g., TokenizerTemplateNormalizer or ConversationContextNormalizer). # When None and adaptation is needed, a default ConversationContextNormalizer is used # that produces "Turn N: User/Assistant" format. diff --git a/pyrit/executor/attack/core/attack_strategy.py b/pyrit/executor/attack/core/attack_strategy.py index 11c58b60ad..f1dc2353bd 100644 --- a/pyrit/executor/attack/core/attack_strategy.py +++ b/pyrit/executor/attack/core/attack_strategy.py @@ -48,6 +48,7 @@ ) from pyrit.executor.attack.core.attack_result_attribution import AttackResultAttribution from pyrit.prompt_target import PromptTarget + from pyrit.prompt_target.common.target_normalization_context import TargetNormalizationContext AttackStrategyContextT = TypeVar("AttackStrategyContextT", bound="AttackContext[Any]") AttackStrategyResultT = TypeVar("AttackStrategyResultT", bound="AttackResult") @@ -88,6 +89,13 @@ class AttackContext(StrategyContext, ABC, Generic[AttackParamsT]): _prepended_conversation_override: list[Message] | None = None _memory_labels_override: dict[str, str] | None = None + # Ephemeral first-send target normalization state. This is never persisted. + target_normalization_context: TargetNormalizationContext | None = field( + default=None, + repr=False, + compare=False, + ) + # Optional attribution from an upstream orchestrator (e.g. Scenario). When # set, the persistence path stamps attribution_parent_id + attribution_data # onto the resulting AttackResult so it can be located later for hydration diff --git a/pyrit/executor/attack/multi_turn/chunked_request.py b/pyrit/executor/attack/multi_turn/chunked_request.py index 6db7537722..a5f3f064c7 100644 --- a/pyrit/executor/attack/multi_turn/chunked_request.py +++ b/pyrit/executor/attack/multi_turn/chunked_request.py @@ -288,6 +288,7 @@ async def _perform_async(self, *, context: ChunkedRequestAttackContext) -> Attac conversation_id=context.session.conversation_id, request_converter_configurations=self._request_converters, response_converter_configurations=self._response_converters, + target_normalization_context=context.target_normalization_context, ) # Store the response diff --git a/pyrit/executor/attack/multi_turn/crescendo.py b/pyrit/executor/attack/multi_turn/crescendo.py index f1feee5d64..80796a9ca8 100644 --- a/pyrit/executor/attack/multi_turn/crescendo.py +++ b/pyrit/executor/attack/multi_turn/crescendo.py @@ -648,6 +648,7 @@ async def _send_prompt_to_objective_target_async( conversation_id=context.session.conversation_id, request_converter_configurations=self._request_converters, response_converter_configurations=self._response_converters, + target_normalization_context=context.target_normalization_context, ) if not response: diff --git a/pyrit/executor/attack/multi_turn/multi_prompt_sending.py b/pyrit/executor/attack/multi_turn/multi_prompt_sending.py index d431ec20db..24fd562284 100644 --- a/pyrit/executor/attack/multi_turn/multi_prompt_sending.py +++ b/pyrit/executor/attack/multi_turn/multi_prompt_sending.py @@ -366,6 +366,7 @@ async def _send_prompt_to_objective_target_async( conversation_id=context.session.conversation_id, request_converter_configurations=self._request_converters, response_converter_configurations=self._response_converters, + target_normalization_context=context.target_normalization_context, ) async def _evaluate_response_async(self, *, response: Message, objective: str) -> Score | None: diff --git a/pyrit/executor/attack/multi_turn/multi_turn_attack_strategy.py b/pyrit/executor/attack/multi_turn/multi_turn_attack_strategy.py index db7a44755b..657a6f5e79 100644 --- a/pyrit/executor/attack/multi_turn/multi_turn_attack_strategy.py +++ b/pyrit/executor/attack/multi_turn/multi_turn_attack_strategy.py @@ -112,8 +112,9 @@ def _rotate_conversation_for_single_turn_target( For multi-turn targets this method is a no-op. - This should be called before each turn (except the first) when sending prompts to the - objective target. + This should be called before each turn when sending prompts to the objective target. + A pending first-turn normalization context suppresses rotation even when prepended + assistant messages have already increased ``executed_turns``. Args: context: The current attack context. @@ -121,6 +122,9 @@ def _rotate_conversation_for_single_turn_target( if self._objective_target.configuration.includes(capability=CapabilityName.MULTI_TURN): return + if context.target_normalization_context and not context.target_normalization_context.is_consumed: + return + if context.executed_turns == 0: return @@ -150,6 +154,7 @@ def _rotate_conversation_for_single_turn_target( context.session.conversation_id = new_conversation_id else: context.session.conversation_id = str(uuid.uuid4()) + context.target_normalization_context = None self._logger.debug( f"Rotated conversation_id for single-turn target: " diff --git a/pyrit/executor/attack/multi_turn/red_teaming.py b/pyrit/executor/attack/multi_turn/red_teaming.py index c05d2c8994..6dfaefaa9e 100644 --- a/pyrit/executor/attack/multi_turn/red_teaming.py +++ b/pyrit/executor/attack/multi_turn/red_teaming.py @@ -487,6 +487,7 @@ async def _send_prompt_to_objective_target_async( request_converter_configurations=self._request_converters, response_converter_configurations=self._response_converters, target=self._objective_target, + target_normalization_context=context.target_normalization_context, ) if response is None: diff --git a/pyrit/executor/attack/multi_turn/tree_of_attacks.py b/pyrit/executor/attack/multi_turn/tree_of_attacks.py index 9cd0550231..9cc9bc6d67 100644 --- a/pyrit/executor/attack/multi_turn/tree_of_attacks.py +++ b/pyrit/executor/attack/multi_turn/tree_of_attacks.py @@ -53,7 +53,7 @@ SeedPrompt, ) from pyrit.prompt_normalizer import ConverterConfiguration, PromptNormalizer -from pyrit.prompt_target import CapabilityName, PromptTarget +from pyrit.prompt_target import CapabilityName, PromptTarget, TargetNormalizationContext from pyrit.prompt_target.common.target_requirements import TargetRequirements from pyrit.score import ( FloatScaleThresholdScorer, @@ -409,6 +409,7 @@ def __init__( self.last_prompt_sent: str | None = None self.last_response: Message | None = None self.error_message: str | None = None + self._target_normalization_context: TargetNormalizationContext | None = None # Context from prepended conversation (for adversarial chat system prompt) self._conversation_context: str | None = None @@ -458,6 +459,15 @@ async def initialize_with_prepended_conversation_async( prompt_normalizer=self._prompt_normalizer, ) + valid_message_count = len( + conversation_manager.get_persistable_prepended_messages(prepended_conversation=prepended_conversation) + ) + target_normalization_context = conversation_manager.create_target_normalization_context( + target=self._objective_target, + conversation_id=self.objective_target_conversation_id, + prepended_message_count=valid_message_count, + prepended_conversation_config=prepended_conversation_config, + ) await conversation_manager.add_prepended_conversation_to_memory_async( prepended_conversation=prepended_conversation, conversation_id=self.objective_target_conversation_id, @@ -466,6 +476,7 @@ async def initialize_with_prepended_conversation_async( target_identifier=self._objective_target.get_identifier(), target=self._objective_target, ) + self._target_normalization_context = target_normalization_context # Build context string for adversarial chat system prompt (like Crescendo) # The adversarial chat uses this in its system prompt rather than in conversation history @@ -600,8 +611,11 @@ async def _send_prompt_to_target_async(self, prompt: str) -> Message: """ # For single-turn targets, generate a fresh conversation ID before each send # to ensure the target always receives a clean conversation without prior history. - if not self._objective_target.configuration.includes(capability=CapabilityName.MULTI_TURN): + if not self._objective_target.configuration.includes(capability=CapabilityName.MULTI_TURN) and not ( + self._target_normalization_context and self._target_normalization_context.is_pending + ): self.objective_target_conversation_id = str(uuid.uuid4()) + self._target_normalization_context = None # Build the request message via the modality router so prior media (if any) # is included when the objective target accepts it. @@ -625,6 +639,7 @@ async def _send_prompt_to_target_async(self, prompt: str) -> Message: response_converter_configurations=self._response_converters, conversation_id=self.objective_target_conversation_id, target=self._objective_target, + target_normalization_context=self._target_normalization_context, ) # Store the full response so subsequent turns can forward media when supported. @@ -661,8 +676,11 @@ async def _send_initial_prompt_to_target_async(self) -> Message: raise ValueError("_initial_prompt must be set before calling this method") # For single-turn targets, generate a fresh conversation ID - if not self._objective_target.configuration.includes(capability=CapabilityName.MULTI_TURN): + if not self._objective_target.configuration.includes(capability=CapabilityName.MULTI_TURN) and not ( + self._target_normalization_context and self._target_normalization_context.is_pending + ): self.objective_target_conversation_id = str(uuid.uuid4()) + self._target_normalization_context = None assert self._objective is not None initial_prompt = self._initial_prompt @@ -701,6 +719,7 @@ async def _send_initial_prompt_to_target_async(self) -> Message: response_converter_configurations=self._response_converters, conversation_id=self.objective_target_conversation_id, target=self._objective_target, + target_normalization_context=self._target_normalization_context, ) # Store the full response so subsequent turns can forward media when supported. diff --git a/pyrit/executor/attack/single_turn/prompt_sending.py b/pyrit/executor/attack/single_turn/prompt_sending.py index 1736bfc19d..0a9fcebc1c 100644 --- a/pyrit/executor/attack/single_turn/prompt_sending.py +++ b/pyrit/executor/attack/single_turn/prompt_sending.py @@ -325,6 +325,7 @@ async def _send_prompt_to_objective_target_async( conversation_id=context.conversation_id, request_converter_configurations=self._request_converters, response_converter_configurations=self._response_converters, + target_normalization_context=context.target_normalization_context, ) async def _evaluate_response_async( diff --git a/pyrit/message_normalizer/__init__.py b/pyrit/message_normalizer/__init__.py index 2ef7e1bc39..0580bd48b1 100644 --- a/pyrit/message_normalizer/__init__.py +++ b/pyrit/message_normalizer/__init__.py @@ -7,6 +7,7 @@ from pyrit.message_normalizer.chat_message_normalizer import ChatMessageNormalizer from pyrit.message_normalizer.conversation_context_normalizer import ConversationContextNormalizer +from pyrit.message_normalizer.first_turn_history_normalizer import FirstTurnHistoryNormalizer from pyrit.message_normalizer.generic_system_squash import GenericSystemSquashNormalizer from pyrit.message_normalizer.history_squash_normalizer import HistorySquashNormalizer from pyrit.message_normalizer.json_schema_normalizer import JsonSchemaNormalizer @@ -14,16 +15,15 @@ MessageListNormalizer, MessageStringNormalizer, ) -from pyrit.message_normalizer.prepended_conversation_normalizer import PrependedConversationNormalizer from pyrit.message_normalizer.tokenizer_template_normalizer import TokenizerTemplateNormalizer __all__ = [ "MessageListNormalizer", "MessageStringNormalizer", + "FirstTurnHistoryNormalizer", "GenericSystemSquashNormalizer", "HistorySquashNormalizer", "JsonSchemaNormalizer", - "PrependedConversationNormalizer", "TokenizerTemplateNormalizer", "ConversationContextNormalizer", "ChatMessageNormalizer", diff --git a/pyrit/message_normalizer/first_turn_history_normalizer.py b/pyrit/message_normalizer/first_turn_history_normalizer.py new file mode 100644 index 0000000000..8f0541ace7 --- /dev/null +++ b/pyrit/message_normalizer/first_turn_history_normalizer.py @@ -0,0 +1,183 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT license. + +import copy +import uuid + +from pyrit.message_normalizer.conversation_context_normalizer import ConversationContextNormalizer +from pyrit.message_normalizer.generic_system_squash import GenericSystemSquashNormalizer +from pyrit.message_normalizer.message_normalizer import MessageListNormalizer, MessageStringNormalizer +from pyrit.models import Message, MessagePiece + + +class FirstTurnHistoryNormalizer(MessageListNormalizer[Message]): + """ + Combine structured prepended history with the first live request. + + The configured string normalizer receives both the prepended history and + the live request, so tokenizer templates place generation markers after the + live user content. Non-text live pieces remain separate in the target view. + """ + + def __init__( + self, + *, + message_normalizer: MessageStringNormalizer, + prepended_message_count: int, + ) -> None: + """ + Initialize the normalizer. + + Args: + message_normalizer: Formatter for the target-facing text. + prepended_message_count: Number of leading messages that belong to + the prepended conversation. + + Raises: + ValueError: If prepended_message_count is less than one. + """ + if prepended_message_count < 1: + raise ValueError("prepended_message_count must be at least 1") + self._message_normalizer = message_normalizer + self._prepended_message_count = prepended_message_count + + async def normalize_async(self, messages: list[Message]) -> list[Message]: + """ + Return one target-facing request containing the prepended history. + + Args: + messages: Prepended history followed by exactly one live request. + + Returns: + A single request message with formatted text and preserved live + non-text pieces. + + Raises: + ValueError: If the input does not contain the configured prepended + message count followed by one live request, or if converted + prepended history cannot be represented as text. + """ + expected_count = self._prepended_message_count + 1 + if len(messages) != expected_count: + raise ValueError( + "First-turn history normalization expected " + f"{self._prepended_message_count} prepended messages and one live request, " + f"but received {len(messages)} messages." + ) + + prepended_messages = messages[: self._prepended_message_count] + live_request = messages[-1] + self._validate_flattenable_converter_output(messages=prepended_messages) + + original_view = self._build_original_view(messages=messages) + converted_view = self._build_converted_view(messages=messages) + original_text = await self._normalize_context_async(messages=original_view) + converted_text = original_text + if self._contains_converted_values(messages=messages): + converted_text = await self._normalize_context_async(messages=converted_view) + + return [ + self._build_target_request( + live_request=live_request, + original_text=original_text, + converted_text=converted_text, + ) + ] + + async def _normalize_context_async(self, *, messages: list[Message]) -> str: + messages_to_normalize = self._filter_live_non_text_pieces(messages=messages) + if isinstance(self._message_normalizer, ConversationContextNormalizer): + messages_to_normalize = await GenericSystemSquashNormalizer().normalize_async(messages_to_normalize) + return await self._message_normalizer.normalize_string_async(messages_to_normalize) + + @staticmethod + def _filter_live_non_text_pieces(*, messages: list[Message]) -> list[Message]: + filtered = copy.deepcopy(messages) + live_message = filtered[-1] + text_pieces = [piece for piece in live_message.message_pieces if piece.converted_value_data_type == "text"] + if not text_pieces: + filtered.pop() + else: + live_message.message_pieces = text_pieces + return filtered + + @staticmethod + def _build_original_view(*, messages: list[Message]) -> list[Message]: + original_messages = copy.deepcopy(messages) + for message in original_messages: + for piece in message.message_pieces: + piece.converted_value = piece.original_value + piece.converted_value_data_type = piece.original_value_data_type + return original_messages + + @staticmethod + def _build_converted_view(*, messages: list[Message]) -> list[Message]: + converted_messages = copy.deepcopy(messages) + for message in converted_messages: + for piece in message.message_pieces: + piece.original_value = piece.converted_value + piece.original_value_data_type = piece.converted_value_data_type + return converted_messages + + @staticmethod + def _contains_converted_values(*, messages: list[Message]) -> bool: + return any( + piece.converter_identifiers + or piece.original_value != piece.converted_value + or piece.original_value_data_type != piece.converted_value_data_type + for message in messages + for piece in message.message_pieces + ) + + @staticmethod + def _validate_flattenable_converter_output(*, messages: list[Message]) -> None: + output_types = { + piece.converted_value_data_type + for message in messages + for piece in message.message_pieces + if piece.converted_value_data_type != "text" + and ( + piece.original_value != piece.converted_value + or piece.original_value_data_type != piece.converted_value_data_type + ) + } + if output_types: + raise ValueError( + "Cannot flatten prepended conversation after request converters produced " + f"non-text output types {sorted(output_types)}. Prepended conversion must produce " + "text for a target without editable history." + ) + + @staticmethod + def _build_target_request( + *, + live_request: Message, + original_text: str, + converted_text: str, + ) -> Message: + request = copy.deepcopy(live_request) + template_piece = request.get_piece() + text_piece = MessagePiece( + id=uuid.uuid4(), + role=template_piece.role, + original_value=original_text, + converted_value=converted_text, + original_value_data_type="text", + converted_value_data_type="text", + conversation_id=template_piece.conversation_id, + sequence=template_piece.sequence, + prompt_metadata=dict(template_piece.prompt_metadata), + ) + target_pieces: list[MessagePiece] = [] + text_inserted = False + for piece in request.message_pieces: + if piece.converted_value_data_type == "text": + if not text_inserted: + target_pieces.append(text_piece) + text_inserted = True + continue + target_pieces.append(piece) + if not text_inserted: + target_pieces.insert(0, text_piece) + request.message_pieces = target_pieces + return request diff --git a/pyrit/message_normalizer/prepended_conversation_normalizer.py b/pyrit/message_normalizer/prepended_conversation_normalizer.py deleted file mode 100644 index 3640095a4b..0000000000 --- a/pyrit/message_normalizer/prepended_conversation_normalizer.py +++ /dev/null @@ -1,154 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# Licensed under the MIT license. - -import copy -import uuid - -from pyrit.message_normalizer.conversation_context_normalizer import ConversationContextNormalizer -from pyrit.message_normalizer.generic_system_squash import GenericSystemSquashNormalizer -from pyrit.message_normalizer.message_normalizer import MessageListNormalizer, MessageStringNormalizer -from pyrit.models import Message, MessagePiece - - -class PrependedConversationNormalizer(MessageListNormalizer[Message]): - """ - Combine prepended history with the first live request for targets without editable history. - - The history remains structured in memory. This normalizer creates an ephemeral target view - that preserves the live request's modalities while prefixing independently rendered original - and converted history. - """ - - def __init__(self, *, message_normalizer: MessageStringNormalizer) -> None: - """ - Initialize the adapter. - - Args: - message_normalizer: Formatter used to render prepended history. - """ - self._message_normalizer = message_normalizer - - async def normalize_async(self, messages: list[Message]) -> list[Message]: - """ - Normalize prepended history into the final request message. - - Args: - messages: Prepended history followed by the first live request. - - Returns: - A single request message containing the rendered history. - """ - if len(messages) < 2: - return copy.deepcopy(messages) - - prepended_messages = messages[:-1] - self._validate_flattenable_converter_output(messages=prepended_messages) - - original_context = await self._normalize_context_async( - messages=self._build_original_view(messages=prepended_messages) - ) - converted_context = original_context - if self._contains_converted_values(messages=prepended_messages): - converted_context = await self._normalize_context_async( - messages=self._build_converted_view(messages=prepended_messages) - ) - - request = copy.deepcopy(messages[-1]) - self._prepend_context( - message=request, - original_context=original_context, - converted_context=converted_context, - ) - return [request] - - async def _normalize_context_async(self, *, messages: list[Message]) -> str: - messages_to_normalize = messages - if isinstance(self._message_normalizer, ConversationContextNormalizer): - messages_to_normalize = await GenericSystemSquashNormalizer().normalize_async(messages) - return await self._message_normalizer.normalize_string_async(messages_to_normalize) - - @staticmethod - def _build_original_view(*, messages: list[Message]) -> list[Message]: - original_messages = copy.deepcopy(messages) - for message in original_messages: - for piece in message.message_pieces: - piece.converted_value = piece.original_value - piece.converted_value_data_type = piece.original_value_data_type - return original_messages - - @staticmethod - def _build_converted_view(*, messages: list[Message]) -> list[Message]: - converted_messages = copy.deepcopy(messages) - for message in converted_messages: - for piece in message.message_pieces: - piece.original_value = piece.converted_value - piece.original_value_data_type = piece.converted_value_data_type - return converted_messages - - @staticmethod - def _contains_converted_values(*, messages: list[Message]) -> bool: - return any( - piece.converter_identifiers - or piece.original_value != piece.converted_value - or piece.original_value_data_type != piece.converted_value_data_type - for message in messages - for piece in message.message_pieces - ) - - @staticmethod - def _validate_flattenable_converter_output(*, messages: list[Message]) -> None: - output_types = { - piece.converted_value_data_type - for message in messages - for piece in message.message_pieces - if piece.converted_value_data_type != "text" - and piece.converted_value_data_type != piece.original_value_data_type - } - if output_types: - raise ValueError( - "Cannot flatten prepended conversation after request converters produced " - f"non-text output types {sorted(output_types)}. Prepended conversion must produce " - "text for a target without editable history." - ) - - @staticmethod - def _prepend_context(*, message: Message, original_context: str, converted_context: str) -> None: - text_piece = next( - ( - piece - for piece in message.message_pieces - if piece.original_value_data_type == "text" and piece.converted_value_data_type == "text" - ), - None, - ) - if text_piece: - text_piece.original_value = PrependedConversationNormalizer._prepend_context_value( - context=original_context, - value=text_piece.original_value, - ) - text_piece.converted_value = PrependedConversationNormalizer._prepend_context_value( - context=converted_context, - value=text_piece.converted_value, - ) - return - - template_piece = message.get_piece() - message.message_pieces.insert( - 0, - MessagePiece( - id=uuid.uuid4(), - role=template_piece.role, - original_value=original_context, - converted_value=converted_context, - original_value_data_type="text", - converted_value_data_type="text", - conversation_id=template_piece.conversation_id, - sequence=template_piece.sequence, - ), - ) - - @staticmethod - def _prepend_context_value(*, context: str, value: str) -> str: - if not context or value == context or value.startswith(f"{context}\n\n"): - return value - return f"{context}\n\n{value}" diff --git a/pyrit/prompt_normalizer/prompt_normalizer.py b/pyrit/prompt_normalizer/prompt_normalizer.py index 24d995812e..e7022ec0e0 100644 --- a/pyrit/prompt_normalizer/prompt_normalizer.py +++ b/pyrit/prompt_normalizer/prompt_normalizer.py @@ -19,7 +19,6 @@ get_execution_context, ) from pyrit.memory import CentralMemory, MemoryInterface, set_message_piece_sha256_async -from pyrit.message_normalizer import MessageStringNormalizer from pyrit.models import ( ComponentIdentifier, Conversation, @@ -30,6 +29,7 @@ from pyrit.prompt_normalizer import ConverterConfiguration, NormalizerRequest from pyrit.prompt_target import PromptTarget from pyrit.prompt_target.batch_helper import batch_task_async +from pyrit.prompt_target.common.target_normalization_context import TargetNormalizationContext logger = logging.getLogger(__name__) @@ -63,22 +63,6 @@ def __init__(self, start_token: str = "⟪", end_token: str = "⟫") -> None: self._start_token = start_token self._end_token = end_token self.id = str(uuid4()) - self._prepended_conversation_normalizers: dict[str, MessageStringNormalizer] = {} - - def register_prepended_conversation_normalizer( - self, - *, - conversation_id: str, - message_normalizer: MessageStringNormalizer, - ) -> None: - """ - Register the formatter used to deliver structured prepended history on the next send. - - Args: - conversation_id: Conversation whose next request should include prepended history. - message_normalizer: Formatter the target should use for that history. - """ - self._prepended_conversation_normalizers[conversation_id] = message_normalizer async def send_prompt_async( self, @@ -88,6 +72,7 @@ async def send_prompt_async( conversation_id: str | None = None, request_converter_configurations: list[ConverterConfiguration] | None = None, response_converter_configurations: list[ConverterConfiguration] | None = None, + target_normalization_context: TargetNormalizationContext | None = None, ) -> Message: """ Send a single request to a target. @@ -100,6 +85,7 @@ async def send_prompt_async( converting the request. Defaults to an empty list. response_converter_configurations (list[ConverterConfiguration], optional): Configurations for converting the response. Defaults to an empty list. + target_normalization_context: Optional per-conversation target normalization state. Returns: Message: The response received from the target. @@ -125,10 +111,6 @@ async def send_prompt_async( for piece in request.message_pieces: piece.conversation_id = conversation_id - prepended_conversation_normalizer = self._prepended_conversation_normalizers.pop( - request.conversation_id, - None, - ) await self.convert_values_async( converter_configurations=request_converter_configurations, message=request, @@ -139,15 +121,19 @@ async def send_prompt_async( responses = None try: - if prepended_conversation_normalizer: + if target_normalization_context: responses = await target.send_prompt_async( message=request, - prepended_conversation_normalizer=prepended_conversation_normalizer, + target_normalization_context=target_normalization_context, ) else: responses = await target.send_prompt_async(message=request) self.memory.add_message_to_memory(request=request) - except EmptyResponseException: + except EmptyResponseException as ex: + if target_normalization_context and not target_normalization_context.is_consumed: + cid = request.message_pieces[0].conversation_id if request.message_pieces else None + raise Exception(f"Error sending prompt with conversation ID: {cid}") from ex + # Empty responses are retried, but we don't want them to stop execution self.memory.add_message_to_memory(request=request) @@ -161,6 +147,12 @@ async def send_prompt_async( ] except Exception as ex: + if target_normalization_context and not target_normalization_context.is_consumed: + # The provider was never invoked, so leave memory unchanged and allow + # the same first-turn context to be retried. + cid = request.message_pieces[0].conversation_id if request.message_pieces else None + raise Exception(f"Error sending prompt with conversation ID: {cid}") from ex + # Ensure request to memory before processing exception self.memory.add_message_to_memory(request=request) diff --git a/pyrit/prompt_target/__init__.py b/pyrit/prompt_target/__init__.py index 9f45cb4ee5..69518ca613 100644 --- a/pyrit/prompt_target/__init__.py +++ b/pyrit/prompt_target/__init__.py @@ -27,6 +27,10 @@ get_known_capabilities, ) from pyrit.prompt_target.common.target_configuration import TargetConfiguration +from pyrit.prompt_target.common.target_normalization_context import ( + TargetNormalizationContext, + TargetNormalizationContextState, +) from pyrit.prompt_target.common.target_requirements import CHAT_TARGET_REQUIREMENTS, TargetRequirements from pyrit.prompt_target.common.utils import limit_requests_per_minute from pyrit.prompt_target.gandalf_target import GandalfLevel, GandalfTarget @@ -106,6 +110,8 @@ def __getattr__(name: str) -> object: "RoundRobinTarget", "TargetCapabilities", "TargetConfiguration", + "TargetNormalizationContext", + "TargetNormalizationContextState", "TargetRequirements", "UnsupportedCapabilityBehavior", "TextTarget", diff --git a/pyrit/prompt_target/common/prompt_target.py b/pyrit/prompt_target/common/prompt_target.py index 99b8f5c330..bde5edcda1 100644 --- a/pyrit/prompt_target/common/prompt_target.py +++ b/pyrit/prompt_target/common/prompt_target.py @@ -6,7 +6,6 @@ from typing import Any, ClassVar, Literal, final from pyrit.memory import CentralMemory, MemoryInterface -from pyrit.message_normalizer import MessageStringNormalizer, PrependedConversationNormalizer from pyrit.models import ( ComponentIdentifier, Conversation, @@ -22,6 +21,7 @@ get_known_capabilities, ) from pyrit.prompt_target.common.target_configuration import TargetConfiguration +from pyrit.prompt_target.common.target_normalization_context import TargetNormalizationContext logger = logging.getLogger(__name__) @@ -137,26 +137,27 @@ async def send_prompt_async( self, *, message: Message, - prepended_conversation_normalizer: MessageStringNormalizer | None = None, + target_normalization_context: TargetNormalizationContext | None = None, ) -> list[Message]: """ Validate, normalize, and send a prompt to the target. This is the public entry point called by the prompt normalizer. It: - 1. Validates the message, fetches the conversation from memory, appends ``message``, and runs - the normalization pipeline (system‑squash, history‑squash, etc.). - 2. Validates the normalized conversation against the target's capabilities. - 3. Delegates to ``_send_prompt_to_target_async`` with the normalized - conversation. + 1. Validates the message and acquires an optional one-shot target context. + 2. Loads memory history unless the context was already consumed, applies + acquired context normalizers, then runs the target's ordinary pipeline. + 3. Validates the normalized conversation against the target's capabilities. + 4. Marks the context consumed and delegates to + ``_send_prompt_to_target_async`` with the normalized conversation. Subclasses MUST NOT override this method. Override ``_send_prompt_to_target_async`` instead. Args: message (Message): The message to send. - prepended_conversation_normalizer (MessageStringNormalizer | None): Optional one-shot - formatter for structured prepended history on a target without editable history. + target_normalization_context: Optional per-conversation normalizers + and one-shot lifecycle state. Returns: list[Message]: Response messages from the target. @@ -165,13 +166,31 @@ async def send_prompt_async( ValueError: If the message or normalized conversation are empty. """ message.validate() - normalized_conversation = await self._get_normalized_conversation_async( - message=message, - prepended_conversation_normalizer=prepended_conversation_normalizer, - ) - if not normalized_conversation: - raise ValueError("Normalization pipeline returned an empty conversation. Cannot send an empty request.") - self._validate_request(normalized_conversation=normalized_conversation) + conversation_id = message.message_pieces[0].conversation_id + should_apply_context = False + if target_normalization_context: + should_apply_context = target_normalization_context.begin_normalization( + conversation_id=conversation_id or "" + ) + + try: + normalized_conversation = await self._get_normalized_conversation_async( + message=message, + target_normalization_context=target_normalization_context, + should_apply_context=should_apply_context, + ) + if not normalized_conversation: + raise ValueError("Normalization pipeline returned an empty conversation. Cannot send an empty request.") + self._validate_request(normalized_conversation=normalized_conversation) + except BaseException: + if target_normalization_context and should_apply_context: + target_normalization_context.restore_pending() + raise + + if target_normalization_context and should_apply_context: + # Target-level retry decorators reuse this normalized payload. Once + # provider invocation starts, attack-level retries must not replay history. + target_normalization_context.mark_consumed() return await self._send_prompt_to_target_async(normalized_conversation=normalized_conversation) @abc.abstractmethod @@ -235,11 +254,15 @@ async def _get_normalized_conversation_async( self, *, message: Message, - prepended_conversation_normalizer: MessageStringNormalizer | None = None, + target_normalization_context: TargetNormalizationContext | None = None, + should_apply_context: bool = False, ) -> list[Message]: """ - Fetch the conversation from memory, append the current message, and run the - normalization pipeline. + Build the target-facing conversation and run the normalization pipeline. + + A consumed target context supplies only the current message so retained + target history is not replayed. Otherwise, memory history is loaded and the + current message is appended before any acquired context normalizers run. The original conversation in memory is never mutated. The returned list is an ephemeral copy intended only for building the API request body. @@ -251,25 +274,23 @@ async def _get_normalized_conversation_async( Args: message (Message): The current message to append. - prepended_conversation_normalizer (MessageStringNormalizer | None): Optional formatter - that combines the existing prepended history with this request before the standard - capability pipeline runs. + target_normalization_context: Optional per-conversation normalization state. + should_apply_context: Whether this send acquired the one-shot context. Returns: list[Message]: The normalized conversation (possibly with system prompt squashed, history squashed, etc.). """ conversation_id = message.message_pieces[0].conversation_id - conversation = ( - list(self._memory.get_conversation_messages(conversation_id=conversation_id)) if conversation_id else [] - ) - conversation.append(message) - if prepended_conversation_normalizer and not self.configuration.includes( - capability=CapabilityName.EDITABLE_HISTORY - ): - conversation = await PrependedConversationNormalizer( - message_normalizer=prepended_conversation_normalizer - ).normalize_async(conversation) + if target_normalization_context and not should_apply_context: + conversation = [message] + else: + conversation = ( + list(self._memory.get_conversation_messages(conversation_id=conversation_id)) if conversation_id else [] + ) + conversation.append(message) + if target_normalization_context: + conversation = await target_normalization_context.normalize_async(messages=conversation) normalized = await self.configuration.normalize_async(messages=conversation) if normalized: # Normalizers may create new Message objects (via Message.from_prompt) with diff --git a/pyrit/prompt_target/common/target_normalization_context.py b/pyrit/prompt_target/common/target_normalization_context.py new file mode 100644 index 0000000000..7cc88323d4 --- /dev/null +++ b/pyrit/prompt_target/common/target_normalization_context.py @@ -0,0 +1,138 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT license. + +from __future__ import annotations + +from dataclasses import dataclass, field +from enum import Enum +from typing import TYPE_CHECKING + +if TYPE_CHECKING: + from pyrit.message_normalizer import MessageListNormalizer + from pyrit.models import Message + + +class TargetNormalizationContextState(str, Enum): + """Lifecycle state for per-send target normalization.""" + + PENDING = "pending" + PREPARING = "preparing" + CONSUMED = "consumed" + + +@dataclass +class TargetNormalizationContext: + """ + Ephemeral normalization state for one target conversation. + + The context is owned by an attack execution and passed explicitly with each + send. It is never persisted to memory or stored on a shared target. + """ + + conversation_id: str + normalizers: tuple[MessageListNormalizer[Message], ...] + _state: TargetNormalizationContextState = field( + default=TargetNormalizationContextState.PENDING, + init=False, + repr=False, + ) + + def __post_init__(self) -> None: + """ + Validate required context data. + + Raises: + ValueError: If the conversation ID is empty or no normalizers are configured. + """ + if not self.conversation_id: + raise ValueError("conversation_id cannot be empty") + if not self.normalizers: + raise ValueError("At least one target normalizer is required") + + @property + def state(self) -> TargetNormalizationContextState: + """The current lifecycle state.""" + return self._state + + @property + def is_pending(self) -> bool: + """Whether normalization can be attempted.""" + return self._state == TargetNormalizationContextState.PENDING + + @property + def is_consumed(self) -> bool: + """Whether provider invocation has started.""" + return self._state == TargetNormalizationContextState.CONSUMED + + def begin_normalization(self, *, conversation_id: str) -> bool: + """ + Acquire the context for normalization. + + Args: + conversation_id: Conversation ID on the outgoing request. + + Returns: + bool: ``True`` when the caller acquired the context, or ``False`` + after it has already been consumed. + + Raises: + ValueError: If the request belongs to another conversation. + RuntimeError: If another send is already preparing the first request. + """ + if conversation_id != self.conversation_id: + raise ValueError( + "Target normalization context belongs to conversation " + f"'{self.conversation_id}', not '{conversation_id}'." + ) + if self._state == TargetNormalizationContextState.CONSUMED: + return False + if self._state == TargetNormalizationContextState.PREPARING: + # Reject rather than queue so two callers cannot both believe they own + # the first target-facing request. + raise RuntimeError("Target normalization is already in progress for this conversation.") + + self._state = TargetNormalizationContextState.PREPARING + return True + + async def normalize_async(self, *, messages: list[Message]) -> list[Message]: + """ + Run the per-send normalizers while the context is acquired. + + Args: + messages: Target conversation to normalize. + + Returns: + list[Message]: The normalized target conversation. + + Raises: + RuntimeError: If the context has not been acquired. + """ + if self._state != TargetNormalizationContextState.PREPARING: + raise RuntimeError("Target normalization context must be acquired before use.") + + normalized = list(messages) + for normalizer in self.normalizers: + normalized = await normalizer.normalize_async(normalized) + return normalized + + def restore_pending(self) -> None: + """ + Allow another attempt after pre-provider normalization fails. + + Raises: + RuntimeError: If the context is not currently preparing. + """ + if self._state != TargetNormalizationContextState.PREPARING: + raise RuntimeError("Only a preparing target normalization context can be restored.") + self._state = TargetNormalizationContextState.PENDING + + def mark_consumed(self) -> None: + """ + Consume the context immediately before provider invocation. + + Raises: + RuntimeError: If the context is not currently preparing. + """ + if self._state != TargetNormalizationContextState.PREPARING: + raise RuntimeError("Only a preparing target normalization context can be consumed.") + self._state = TargetNormalizationContextState.CONSUMED diff --git a/tests/unit/executor/attack/component/test_conversation_manager.py b/tests/unit/executor/attack/component/test_conversation_manager.py index 8ddca07ef0..40146b3d58 100644 --- a/tests/unit/executor/attack/component/test_conversation_manager.py +++ b/tests/unit/executor/attack/component/test_conversation_manager.py @@ -35,7 +35,7 @@ ) from pyrit.executor.attack.core import AttackContext from pyrit.executor.attack.core.attack_parameters import AttackParameters -from pyrit.message_normalizer import ConversationContextNormalizer +from pyrit.message_normalizer import ConversationContextNormalizer, FirstTurnHistoryNormalizer from pyrit.models import ComponentIdentifier, Message, MessagePiece, PromptDataType, Score from pyrit.prompt_normalizer import ConverterConfiguration, PromptNormalizer from pyrit.prompt_target import PromptTarget @@ -71,6 +71,16 @@ async def convert_async(self, *, prompt: str, input_type: PromptDataType = "text return ConverterResult(output_text="converted.png", output_type="image_path") +class _ImageToImageConverter(Converter): + """A deterministic image-to-image converter for lossy adaptation tests.""" + + SUPPORTED_INPUT_TYPES: tuple[PromptDataType, ...] = ("image_path",) + SUPPORTED_OUTPUT_TYPES: tuple[PromptDataType, ...] = ("image_path",) + + async def convert_async(self, *, prompt: str, input_type: PromptDataType = "image_path") -> ConverterResult: + return ConverterResult(output_text="converted.png", output_type="image_path") + + # ============================================================================= # Fixtures # ============================================================================= @@ -770,7 +780,7 @@ async def test_non_editable_target_does_not_rewrite_supplied_next_message( assert context.next_message is next_message assert context.next_message.get_value() == "Caller-supplied question" - async def test_non_editable_target_registers_custom_first_send_formatter( + async def test_non_editable_target_sets_custom_first_send_context( self, attack_identifier: ComponentIdentifier, mock_prompt_normalizer: MagicMock, @@ -791,10 +801,11 @@ async def test_non_editable_target_registers_custom_first_send_formatter( prepended_conversation_config=config, ) - mock_prompt_normalizer.register_prepended_conversation_normalizer.assert_called_once_with( - conversation_id=conversation_id, - message_normalizer=message_normalizer, - ) + assert context.target_normalization_context is not None + assert context.target_normalization_context.conversation_id == conversation_id + normalizer = context.target_normalization_context.normalizers[0] + assert isinstance(normalizer, FirstTurnHistoryNormalizer) + assert normalizer._message_normalizer is message_normalizer message_normalizer.normalize_string_async.assert_not_called() async def test_returns_turn_count_for_multi_turn_attacks( @@ -1073,6 +1084,134 @@ async def test_non_editable_target_rejects_non_text_output_from_current_converte assert sample_conversation[0].get_piece().converted_value_data_type == "text" + async def test_non_editable_target_rejects_same_modality_non_text_output( + self, + attack_identifier: ComponentIdentifier, + mock_prompt_target: MagicMock, + ) -> None: + manager = ConversationManager() + context = _TestAttackContext(params=AttackParameters(objective="Test objective")) + context.prepended_conversation = [ + Message( + message_pieces=[ + MessagePiece( + role="user", + original_value="original.png", + converted_value="original.png", + original_value_data_type="image_path", + converted_value_data_type="image_path", + conversation_id="seed", + ) + ] + ) + ] + converter_config = ConverterConfiguration.from_converters(converters=[_ImageToImageConverter()]) + + with pytest.raises(ValueError, match="non-text output types.*image_path"): + await manager.initialize_context_async( + context=context, + target=mock_prompt_target, + conversation_id=str(uuid.uuid4()), + request_converters=converter_config, + ) + + async def test_prepended_conversion_failure_does_not_partially_write_history( + self, + attack_identifier: ComponentIdentifier, + mock_prompt_normalizer: MagicMock, + mock_prompt_target: MagicMock, + sample_conversation: list[Message], + ) -> None: + manager = ConversationManager(prompt_normalizer=mock_prompt_normalizer) + context = _TestAttackContext(params=AttackParameters(objective="Test objective")) + context.prepended_conversation = sample_conversation + mock_prompt_normalizer.convert_values_async.side_effect = [None, ValueError("second conversion failed")] + config = PrependedConversationConfig(apply_converters_to_roles=["user", "assistant"]) + conversation_id = str(uuid.uuid4()) + + with pytest.raises(ValueError, match="second conversion failed"): + await manager.initialize_context_async( + context=context, + target=mock_prompt_target, + conversation_id=conversation_id, + request_converters=[ConverterConfiguration(converters=[])], + prepended_conversation_config=config, + ) + + assert manager.get_conversation(conversation_id) == [] + assert context.target_normalization_context is None + + async def test_non_persisted_prepended_message_is_not_counted_in_context( + self, + attack_identifier: ComponentIdentifier, + mock_prompt_target: MagicMock, + ) -> None: + manager = ConversationManager() + context = _TestAttackContext(params=AttackParameters(objective="Test objective")) + piece = MessagePiece( + role="user", + original_value="ephemeral", + conversation_id="seed", + ) + piece.not_in_memory = True + context.prepended_conversation = [Message(message_pieces=[piece])] + conversation_id = str(uuid.uuid4()) + + await manager.initialize_context_async( + context=context, + target=mock_prompt_target, + conversation_id=conversation_id, + ) + + assert manager.get_conversation(conversation_id) == [] + assert context.target_normalization_context is None + + async def test_non_persisted_piece_does_not_constrain_flattening( + self, + attack_identifier: ComponentIdentifier, + mock_prompt_target: MagicMock, + ) -> None: + manager = ConversationManager() + context = _TestAttackContext(params=AttackParameters(objective="Test objective")) + ephemeral_piece = MessagePiece( + role="user", + original_value="ephemeral", + conversation_id="seed", + sequence=0, + ) + ephemeral_piece.not_in_memory = True + context.prepended_conversation = [ + Message( + message_pieces=[ + MessagePiece( + role="user", + original_value="persisted", + conversation_id="seed", + sequence=0, + ), + ephemeral_piece, + ] + ) + ] + conversation_id = str(uuid.uuid4()) + + await manager.initialize_context_async( + context=context, + target=mock_prompt_target, + conversation_id=conversation_id, + request_converters=[ + ConverterConfiguration( + converters=[_ImageOutputConverter()], + indexes_to_apply=[1], + ) + ], + ) + + stored = manager.get_conversation(conversation_id) + assert len(stored) == 1 + assert [piece.converted_value for piece in stored[0].message_pieces] == ["persisted"] + assert context.target_normalization_context is not None + async def test_non_editable_target_preserves_converter_piece_indexes( self, attack_identifier: ComponentIdentifier, @@ -1242,7 +1381,7 @@ async def test_message_normalizer_default_uses_conversation_context_normalizer( mock_prompt_target: MagicMock, sample_conversation: list[Message], ) -> None: - """Test that the default formatter is registered for target-side adaptation.""" + """Test that the default formatter is carried in explicit target-side state.""" manager = ConversationManager(prompt_normalizer=mock_prompt_normalizer) conversation_id = str(uuid.uuid4()) context = _TestAttackContext(params=AttackParameters(objective="Test objective")) @@ -1254,9 +1393,10 @@ async def test_message_normalizer_default_uses_conversation_context_normalizer( conversation_id=conversation_id, ) - registered = mock_prompt_normalizer.register_prepended_conversation_normalizer.call_args.kwargs - assert registered["conversation_id"] == conversation_id - assert isinstance(registered["message_normalizer"], ConversationContextNormalizer) + assert context.target_normalization_context is not None + normalizer = context.target_normalization_context.normalizers[0] + assert isinstance(normalizer, FirstTurnHistoryNormalizer) + assert isinstance(normalizer._message_normalizer, ConversationContextNormalizer) # ------------------------------------------------------------------------- # Chat Target Behavior (Config has no effect) diff --git a/tests/unit/executor/attack/multi_turn/test_supports_multi_turn_attacks.py b/tests/unit/executor/attack/multi_turn/test_supports_multi_turn_attacks.py index ec81a8ff6e..348529fe12 100644 --- a/tests/unit/executor/attack/multi_turn/test_supports_multi_turn_attacks.py +++ b/tests/unit/executor/attack/multi_turn/test_supports_multi_turn_attacks.py @@ -11,7 +11,9 @@ MultiTurnAttackContext, ) from pyrit.memory import CentralMemory +from pyrit.message_normalizer import ConversationContextNormalizer, FirstTurnHistoryNormalizer from pyrit.models import ConversationType, MessagePiece +from pyrit.prompt_target import TargetNormalizationContext def _make_context() -> MultiTurnAttackContext: @@ -88,6 +90,26 @@ def test_noop_on_first_turn(self): assert context.session.conversation_id == original_id assert len(context.related_conversations) == 0 + def test_pending_first_turn_context_suppresses_rotation_after_prepended_turns(self): + strategy = _make_strategy(supports_multi_turn=False) + context = _make_context() + context.executed_turns = 2 + original_id = context.session.conversation_id + context.target_normalization_context = TargetNormalizationContext( + conversation_id=original_id, + normalizers=( + FirstTurnHistoryNormalizer( + message_normalizer=ConversationContextNormalizer(), + prepended_message_count=4, + ), + ), + ) + + strategy._rotate_conversation_for_single_turn_target(context=context) + + assert context.session.conversation_id == original_id + assert len(context.related_conversations) == 0 + def test_rotates_on_second_turn_for_single_turn_target(self): strategy = _make_strategy(supports_multi_turn=False) context = _make_context() diff --git a/tests/unit/executor/attack/multi_turn/test_tree_of_attacks.py b/tests/unit/executor/attack/multi_turn/test_tree_of_attacks.py index 239f7c444c..cbcf203413 100644 --- a/tests/unit/executor/attack/multi_turn/test_tree_of_attacks.py +++ b/tests/unit/executor/attack/multi_turn/test_tree_of_attacks.py @@ -1585,6 +1585,7 @@ def test_node_duplicate_creates_child(self, node_components): """Test that duplicate() creates a proper child node.""" parent_node = _TreeOfAttacksNode(**node_components) parent_node.node_id = "parent_node_id" + parent_node._target_normalization_context = MagicMock() # Mock memory duplicate conversation with patch.object(parent_node._memory, "duplicate_conversation", return_value="new_conv_id"): @@ -1593,6 +1594,7 @@ def test_node_duplicate_creates_child(self, node_components): assert child_node.node_id != parent_node.node_id assert child_node.parent_id == parent_node.node_id assert child_node.completed is False + assert child_node._target_normalization_context is None def _node_with_schema(self, node_components, schema): """Build a real node whose adversarial system prompt advertises ``schema``. diff --git a/tests/unit/executor/attack/single_turn/test_prompt_sending.py b/tests/unit/executor/attack/single_turn/test_prompt_sending.py index 76d0b8cf62..cd2017e349 100644 --- a/tests/unit/executor/attack/single_turn/test_prompt_sending.py +++ b/tests/unit/executor/attack/single_turn/test_prompt_sending.py @@ -362,9 +362,31 @@ async def test_non_chat_target_converts_history_by_role_before_flattening(self): encoded_user = base64.b64encode(prepended_user.encode()).decode() encoded_final_request = base64.b64encode(final_request.encode()).decode() assert target.prompt_sent == [ - f"Turn 1:\nuser: {encoded_user}\nassistant: {simulated_response}\n\n{encoded_final_request}" + (f"Turn 1:\nuser: {encoded_user}\nassistant: {simulated_response}\nTurn 2:\nuser: {encoded_final_request}") ] + async def test_retry_setup_creates_fresh_first_turn_context(self, basic_context): + target = MockPromptTarget() + target._configuration = TargetConfiguration(capabilities=TargetCapabilities()) + attack = PromptSendingAttack(objective_target=target) + basic_context.prepended_conversation = [ + Message.from_prompt(prompt="prepended", role="user"), + ] + + await attack._setup_async(context=basic_context) + first_context = basic_context.target_normalization_context + first_conversation_id = basic_context.conversation_id + assert first_context is not None + assert first_context.begin_normalization(conversation_id=first_conversation_id) + first_context.mark_consumed() + + await attack._setup_async(context=basic_context) + + assert basic_context.conversation_id != first_conversation_id + assert basic_context.target_normalization_context is not None + assert basic_context.target_normalization_context is not first_context + assert basic_context.target_normalization_context.is_pending + @pytest.mark.usefixtures("patch_central_database") class TestPromptPreparation: diff --git a/tests/unit/prompt_normalizer/test_prompt_normalizer.py b/tests/unit/prompt_normalizer/test_prompt_normalizer.py index 01ad9bf435..b314568a18 100644 --- a/tests/unit/prompt_normalizer/test_prompt_normalizer.py +++ b/tests/unit/prompt_normalizer/test_prompt_normalizer.py @@ -25,7 +25,7 @@ get_execution_context, ) from pyrit.memory import CentralMemory -from pyrit.message_normalizer import ConversationContextNormalizer +from pyrit.message_normalizer import ConversationContextNormalizer, FirstTurnHistoryNormalizer from pyrit.models import ( Message, MessagePiece, @@ -37,7 +37,7 @@ from pyrit.prompt_normalizer.converter_configuration import ( ConverterConfiguration, ) -from pyrit.prompt_target import PromptTarget +from pyrit.prompt_target import PromptTarget, TargetNormalizationContext @pytest.fixture @@ -94,6 +94,14 @@ def output_supported(self, output_type: PromptDataType) -> bool: return output_type == "text" +class ContextFailingConverter(Converter): + SUPPORTED_INPUT_TYPES: tuple[PromptDataType, ...] = ("text",) + SUPPORTED_OUTPUT_TYPES: tuple[PromptDataType, ...] = ("text",) + + async def convert_async(self, *, prompt: str, input_type: PromptDataType = "text") -> ConverterResult: + raise ValueError("conversion failed") + + def assert_message_piece_hashes_set(request: Message): assert request assert request.message_pieces @@ -116,37 +124,117 @@ async def test_send_prompt_async_multiple_converters(mock_memory_instance, seed_ assert prompt_target.prompt_sent == ["S_G_V_s_b_G_8_="] -async def test_send_prompt_async_passes_registered_prepended_formatter_once(mock_memory_instance): +async def test_send_prompt_async_forwards_target_normalization_context(mock_memory_instance): prompt_target = MagicMock(spec=PromptTarget) prompt_target.get_identifier.return_value = get_mock_target_identifier("MockTarget") prompt_target.send_prompt_async = AsyncMock( - side_effect=[ - [MessagePiece(role="assistant", original_value="first").to_message()], - [MessagePiece(role="assistant", original_value="second").to_message()], - ] + return_value=[MessagePiece(role="assistant", original_value="first").to_message()] ) normalizer = PromptNormalizer() - formatter = ConversationContextNormalizer() conversation_id = "prepended-conversation" - normalizer.register_prepended_conversation_normalizer( + target_context = TargetNormalizationContext( conversation_id=conversation_id, - message_normalizer=formatter, + normalizers=( + FirstTurnHistoryNormalizer( + message_normalizer=ConversationContextNormalizer(), + prepended_message_count=1, + ), + ), ) await normalizer.send_prompt_async( message=Message.from_prompt(prompt="first request", role="user"), target=prompt_target, conversation_id=conversation_id, + target_normalization_context=target_context, ) - await normalizer.send_prompt_async( - message=Message.from_prompt(prompt="second request", role="user"), - target=prompt_target, + + call = prompt_target.send_prompt_async.await_args + assert call.kwargs["target_normalization_context"] is target_context + + +async def test_send_prompt_async_conversion_failure_leaves_context_pending(mock_memory_instance): + prompt_target = MagicMock(spec=PromptTarget) + prompt_target.get_identifier.return_value = get_mock_target_identifier("MockTarget") + prompt_target.send_prompt_async = AsyncMock() + conversation_id = "prepended-conversation" + target_context = TargetNormalizationContext( conversation_id=conversation_id, + normalizers=( + FirstTurnHistoryNormalizer( + message_normalizer=ConversationContextNormalizer(), + prepended_message_count=1, + ), + ), ) + converter_config = ConverterConfiguration.from_converters(converters=[ContextFailingConverter()]) + + with pytest.raises(ValueError, match="conversion failed"): + await PromptNormalizer().send_prompt_async( + message=Message.from_prompt(prompt="request", role="user"), + target=prompt_target, + conversation_id=conversation_id, + request_converter_configurations=converter_config, + target_normalization_context=target_context, + ) + + assert target_context.is_pending + prompt_target.send_prompt_async.assert_not_awaited() + mock_memory_instance.add_message_to_memory.assert_not_called() + + +async def test_send_prompt_async_pre_provider_failure_is_not_persisted(mock_memory_instance): + prompt_target = MagicMock(spec=PromptTarget) + prompt_target.get_identifier.return_value = get_mock_target_identifier("MockTarget") + prompt_target.send_prompt_async = AsyncMock(side_effect=ValueError("normalization failed")) + conversation_id = "prepended-conversation" + target_context = TargetNormalizationContext( + conversation_id=conversation_id, + normalizers=( + FirstTurnHistoryNormalizer( + message_normalizer=ConversationContextNormalizer(), + prepended_message_count=1, + ), + ), + ) + + with pytest.raises(Exception, match="Error sending prompt with conversation ID"): + await PromptNormalizer().send_prompt_async( + message=Message.from_prompt(prompt="request", role="user"), + target=prompt_target, + conversation_id=conversation_id, + target_normalization_context=target_context, + ) + + assert target_context.is_pending + mock_memory_instance.add_message_to_memory.assert_not_called() + + +async def test_send_prompt_async_pre_provider_empty_response_is_not_persisted(mock_memory_instance): + prompt_target = MagicMock(spec=PromptTarget) + prompt_target.get_identifier.return_value = get_mock_target_identifier("MockTarget") + prompt_target.send_prompt_async = AsyncMock(side_effect=EmptyResponseException(message="normalization failed")) + conversation_id = "prepended-conversation" + target_context = TargetNormalizationContext( + conversation_id=conversation_id, + normalizers=( + FirstTurnHistoryNormalizer( + message_normalizer=ConversationContextNormalizer(), + prepended_message_count=1, + ), + ), + ) + + with pytest.raises(Exception, match="Error sending prompt with conversation ID"): + await PromptNormalizer().send_prompt_async( + message=Message.from_prompt(prompt="request", role="user"), + target=prompt_target, + conversation_id=conversation_id, + target_normalization_context=target_context, + ) - first_call, second_call = prompt_target.send_prompt_async.await_args_list - assert first_call.kwargs["prepended_conversation_normalizer"] is formatter - assert "prepended_conversation_normalizer" not in second_call.kwargs + assert target_context.is_pending + mock_memory_instance.add_message_to_memory.assert_not_called() async def test_send_prompt_async_no_response_adds_memory(mock_memory_instance, seed_group): diff --git a/tests/unit/prompt_target/target/test_normalize_async_integration.py b/tests/unit/prompt_target/target/test_normalize_async_integration.py index 438d0ab8f8..d8e4feab9d 100644 --- a/tests/unit/prompt_target/target/test_normalize_async_integration.py +++ b/tests/unit/prompt_target/target/test_normalize_async_integration.py @@ -3,6 +3,7 @@ from __future__ import annotations +import asyncio import json from typing import TYPE_CHECKING from unittest.mock import AsyncMock, MagicMock, patch @@ -16,9 +17,14 @@ from unit.mocks import MockPromptTarget from pyrit.memory.memory_interface import MemoryInterface -from pyrit.message_normalizer import ConversationContextNormalizer, MessageStringNormalizer +from pyrit.message_normalizer import ( + ConversationContextNormalizer, + FirstTurnHistoryNormalizer, + MessageStringNormalizer, + TokenizerTemplateNormalizer, +) from pyrit.models import ComponentIdentifier, Message, MessagePiece -from pyrit.prompt_target import AzureMLChatTarget, OpenAIChatTarget +from pyrit.prompt_target import AzureMLChatTarget, OpenAIChatTarget, TargetNormalizationContext from pyrit.prompt_target.common.target_capabilities import ( CapabilityHandlingPolicy, CapabilityName, @@ -44,6 +50,23 @@ def _make_message(*, role: str, content: str, conversation_id: str = "conv1") -> return Message(message_pieces=[_make_message_piece(role=role, content=content, conversation_id=conversation_id)]) +def _make_target_normalization_context( + *, + prepended_message_count: int, + formatter: MessageStringNormalizer | None = None, + conversation_id: str = "conv1", +) -> TargetNormalizationContext: + return TargetNormalizationContext( + conversation_id=conversation_id, + normalizers=( + FirstTurnHistoryNormalizer( + message_normalizer=formatter or ConversationContextNormalizer(), + prepended_message_count=prepended_message_count, + ), + ), + ) + + def _create_mock_chat_completion(content: str = "hi") -> MagicMock: mock = MagicMock(spec=ChatCompletion) mock.choices = [MagicMock()] @@ -516,18 +539,21 @@ async def test_non_editable_target_adapts_prepended_history_without_mutating_mem mock_memory = MagicMock(spec=MemoryInterface) mock_memory.get_conversation_messages.return_value = memory_messages target._memory = mock_memory + target_context = _make_target_normalization_context(prepended_message_count=2) + assert target_context.begin_normalization(conversation_id="conv1") result = await target._get_normalized_conversation_async( message=live_request, - prepended_conversation_normalizer=ConversationContextNormalizer(), + target_normalization_context=target_context, + should_apply_context=True, ) assert len(result) == 1 assert result[0].get_piece().original_value == ( - "Turn 1:\nuser: original history\nassistant: assistant history\n\noriginal live" + "Turn 1:\nuser: original history\nassistant: assistant history\nTurn 2:\nuser: original live" ) assert result[0].get_piece().converted_value == ( - "Turn 1:\nuser: converted history\nassistant: assistant history\n\nconverted live" + "Turn 1:\nuser: converted history\nassistant: assistant history\nTurn 2:\nuser: converted live" ) assert len(memory_messages) == 2 assert memory_messages[0].get_piece().original_value == "original history" @@ -558,10 +584,13 @@ async def test_non_editable_target_preserves_system_history_and_multimodal_live_ mock_memory = MagicMock(spec=MemoryInterface) mock_memory.get_conversation_messages.return_value = [system_message] target._memory = mock_memory + target_context = _make_target_normalization_context(prepended_message_count=1) + assert target_context.begin_normalization(conversation_id="conv1") result = await target._get_normalized_conversation_async( message=live_request, - prepended_conversation_normalizer=ConversationContextNormalizer(), + target_normalization_context=target_context, + should_apply_context=True, ) assert len(result) == 1 @@ -571,6 +600,54 @@ async def test_non_editable_target_preserves_system_history_and_multimodal_live_ assert result[0].message_pieces[1].converted_value_data_type == "image_path" +@pytest.mark.usefixtures("patch_central_database") +async def test_first_turn_normalization_preserves_live_multimodal_piece_order(): + target = MockPromptTarget() + target._configuration = TargetConfiguration( + capabilities=TargetCapabilities( + supports_multi_message_pieces=True, + input_modalities=frozenset({frozenset({"text", "image_path"})}), + ) + ) + live_request = Message( + message_pieces=[ + MessagePiece( + role="user", + original_value="diagram.png", + converted_value="diagram.png", + original_value_data_type="image_path", + converted_value_data_type="image_path", + conversation_id="conv1", + sequence=0, + ), + MessagePiece( + role="user", + original_value="What does this show?", + converted_value="What does this show?", + original_value_data_type="text", + converted_value_data_type="text", + conversation_id="conv1", + sequence=0, + ), + ] + ) + mock_memory = MagicMock(spec=MemoryInterface) + mock_memory.get_conversation_messages.return_value = [_make_message(role="user", content="prepended")] + target._memory = mock_memory + target_context = _make_target_normalization_context(prepended_message_count=1) + assert target_context.begin_normalization(conversation_id="conv1") + + result = await target._get_normalized_conversation_async( + message=live_request, + target_normalization_context=target_context, + should_apply_context=True, + ) + + assert [piece.converted_value_data_type for piece in result[0].message_pieces] == ["image_path", "text"] + assert result[0].message_pieces[0].converted_value == "diagram.png" + assert "What does this show?" in result[0].message_pieces[1].converted_value + + @pytest.mark.usefixtures("patch_central_database") async def test_prepended_history_adapter_is_used_only_when_explicitly_passed(): target = MockPromptTarget() @@ -591,14 +668,58 @@ async def test_prepended_history_adapter_is_used_only_when_explicitly_passed(): [prepended, prior_live, prior_response], ] target._memory = mock_memory + target_context = _make_target_normalization_context(prepended_message_count=1) await target.send_prompt_async( message=prior_live, - prepended_conversation_normalizer=ConversationContextNormalizer(), + target_normalization_context=target_context, + ) + await target.send_prompt_async( + message=second_live, + target_normalization_context=target_context, ) - await target.send_prompt_async(message=second_live) - assert target.prompt_sent == ["Turn 1:\nuser: prepended\n\nfirst live", "second live"] + assert target.prompt_sent == ["Turn 1:\nuser: prepended\nTurn 2:\nuser: first live", "second live"] + assert target_context.is_consumed + + +@pytest.mark.usefixtures("patch_central_database") +async def test_consumed_context_sends_only_the_current_target_facing_request(): + target = MockPromptTarget() + target._configuration = TargetConfiguration( + capabilities=TargetCapabilities( + supports_multi_turn=True, + supports_system_prompt=True, + ) + ) + prepended = _make_message(role="user", content="prepended") + first_live = _make_message(role="user", content="first live") + second_live = _make_message(role="user", content="second live") + prior_response = _make_message(role="assistant", content="first response") + mock_memory = MagicMock(spec=MemoryInterface) + mock_memory.get_conversation_messages.side_effect = [ + [prepended], + [prepended, first_live, prior_response], + ] + target._memory = mock_memory + target._send_prompt_to_target_async = AsyncMock( # type: ignore[method-assign] + return_value=[_make_message(role="assistant", content="response")] + ) + target_context = _make_target_normalization_context(prepended_message_count=1) + + await target.send_prompt_async( + message=first_live, + target_normalization_context=target_context, + ) + await target.send_prompt_async( + message=second_live, + target_normalization_context=target_context, + ) + + first_payload, second_payload = target._send_prompt_to_target_async.await_args_list + assert len(first_payload.kwargs["normalized_conversation"]) == 1 + assert len(second_payload.kwargs["normalized_conversation"]) == 1 + assert second_payload.kwargs["normalized_conversation"][0].get_value() == "second live" @pytest.mark.usefixtures("patch_central_database") @@ -610,13 +731,19 @@ async def test_non_editable_target_uses_custom_prepended_formatter(): target._memory = mock_memory formatter = MagicMock(spec=MessageStringNormalizer) formatter.normalize_string_async = AsyncMock(return_value="CUSTOM HISTORY") + target_context = _make_target_normalization_context( + prepended_message_count=1, + formatter=formatter, + ) + assert target_context.begin_normalization(conversation_id="conv1") result = await target._get_normalized_conversation_async( message=_make_message(role="user", content="live"), - prepended_conversation_normalizer=formatter, + target_normalization_context=target_context, + should_apply_context=True, ) - assert result[0].get_value() == "CUSTOM HISTORY\n\nlive" + assert result[0].get_value() == "CUSTOM HISTORY" formatter.normalize_string_async.assert_awaited_once() @@ -633,12 +760,44 @@ async def test_non_editable_target_rejects_non_text_converted_prepended_history( mock_memory = MagicMock(spec=MemoryInterface) mock_memory.get_conversation_messages.return_value = [prepended] target._memory = mock_memory + target_context = _make_target_normalization_context(prepended_message_count=1) with pytest.raises(ValueError, match="non-text output types.*image_path"): - await target._get_normalized_conversation_async( + await target.send_prompt_async( message=_make_message(role="user", content="live"), - prepended_conversation_normalizer=ConversationContextNormalizer(), + target_normalization_context=target_context, ) + assert target_context.is_pending + + +@pytest.mark.usefixtures("patch_central_database") +async def test_non_editable_target_rejects_same_modality_non_text_conversion(): + target = MockPromptTarget() + target._configuration = TargetConfiguration(capabilities=TargetCapabilities()) + prepended = Message( + message_pieces=[ + MessagePiece( + role="user", + original_value="original.png", + converted_value="converted.png", + original_value_data_type="image_path", + converted_value_data_type="image_path", + conversation_id="conv1", + ) + ] + ) + mock_memory = MagicMock(spec=MemoryInterface) + mock_memory.get_conversation_messages.return_value = [prepended] + target._memory = mock_memory + target_context = _make_target_normalization_context(prepended_message_count=1) + + with pytest.raises(ValueError, match="non-text output types.*image_path"): + await target.send_prompt_async( + message=_make_message(role="user", content="live"), + target_normalization_context=target_context, + ) + + assert target_context.is_pending @pytest.mark.usefixtures("patch_central_database") @@ -661,10 +820,156 @@ async def test_non_editable_target_allows_preexisting_non_text_history_with_conv mock_memory = MagicMock(spec=MemoryInterface) mock_memory.get_conversation_messages.return_value = [prepended] target._memory = mock_memory + target_context = _make_target_normalization_context(prepended_message_count=1) + assert target_context.begin_normalization(conversation_id="conv1") + + result = await target._get_normalized_conversation_async( + message=_make_message(role="user", content="live"), + target_normalization_context=target_context, + should_apply_context=True, + ) + + assert result[0].get_value() == "Turn 1:\nuser: [Image_path]\nTurn 2:\nuser: live" + + +@pytest.mark.usefixtures("patch_central_database") +async def test_target_normalization_failure_can_be_retried(): + target = MockPromptTarget() + target._configuration = TargetConfiguration(capabilities=TargetCapabilities()) + mock_memory = MagicMock(spec=MemoryInterface) + mock_memory.get_conversation_messages.return_value = [_make_message(role="user", content="prepended")] + target._memory = mock_memory + formatter = MagicMock(spec=MessageStringNormalizer) + formatter.normalize_string_async = AsyncMock(side_effect=[ValueError("format failed"), "formatted request"]) + target_context = _make_target_normalization_context( + prepended_message_count=1, + formatter=formatter, + ) + live_request = _make_message(role="user", content="live") + + with pytest.raises(ValueError, match="format failed"): + await target.send_prompt_async( + message=live_request, + target_normalization_context=target_context, + ) + + assert target_context.is_pending + await target.send_prompt_async( + message=live_request, + target_normalization_context=target_context, + ) + assert target_context.is_consumed + assert target.prompt_sent == ["formatted request"] + + +@pytest.mark.usefixtures("patch_central_database") +async def test_target_normalization_cancellation_restores_pending_state(): + target = MockPromptTarget() + target._configuration = TargetConfiguration(capabilities=TargetCapabilities()) + mock_memory = MagicMock(spec=MemoryInterface) + mock_memory.get_conversation_messages.return_value = [_make_message(role="user", content="prepended")] + target._memory = mock_memory + formatter = MagicMock(spec=MessageStringNormalizer) + formatter.normalize_string_async = AsyncMock(side_effect=asyncio.CancelledError()) + target_context = _make_target_normalization_context( + prepended_message_count=1, + formatter=formatter, + ) + + with pytest.raises(asyncio.CancelledError): + await target.send_prompt_async( + message=_make_message(role="user", content="live"), + target_normalization_context=target_context, + ) + + assert target_context.is_pending + + +@pytest.mark.usefixtures("patch_central_database") +async def test_provider_failure_leaves_target_normalization_context_consumed(): + target = MockPromptTarget() + target._configuration = TargetConfiguration(capabilities=TargetCapabilities()) + mock_memory = MagicMock(spec=MemoryInterface) + mock_memory.get_conversation_messages.return_value = [_make_message(role="user", content="prepended")] + target._memory = mock_memory + target._send_prompt_to_target_async = AsyncMock(side_effect=RuntimeError("provider failed")) # type: ignore[method-assign] + target_context = _make_target_normalization_context(prepended_message_count=1) + + with pytest.raises(RuntimeError, match="provider failed"): + await target.send_prompt_async( + message=_make_message(role="user", content="live"), + target_normalization_context=target_context, + ) + + assert target_context.is_consumed + + +@pytest.mark.usefixtures("patch_central_database") +async def test_concurrent_first_sends_do_not_both_apply_prepended_history(): + target = MockPromptTarget() + target._configuration = TargetConfiguration(capabilities=TargetCapabilities()) + mock_memory = MagicMock(spec=MemoryInterface) + mock_memory.get_conversation_messages.return_value = [_make_message(role="user", content="prepended")] + target._memory = mock_memory + started = asyncio.Event() + release = asyncio.Event() + + async def wait_to_format(messages: list[Message]) -> str: + started.set() + await release.wait() + return "formatted request" + + formatter = MagicMock(spec=MessageStringNormalizer) + formatter.normalize_string_async = AsyncMock(side_effect=wait_to_format) + target_context = _make_target_normalization_context( + prepended_message_count=1, + formatter=formatter, + ) + first_send = asyncio.create_task( + target.send_prompt_async( + message=_make_message(role="user", content="first"), + target_normalization_context=target_context, + ) + ) + await started.wait() + + try: + with pytest.raises(RuntimeError, match="already in progress"): + await target.send_prompt_async( + message=_make_message(role="user", content="second"), + target_normalization_context=target_context, + ) + finally: + release.set() + + await first_send + assert target.prompt_sent == ["formatted request"] + assert target_context.is_consumed + + +@pytest.mark.usefixtures("patch_central_database") +async def test_tokenizer_formatter_receives_live_request_before_generation_prompt(): + target = MockPromptTarget() + target._configuration = TargetConfiguration(capabilities=TargetCapabilities()) + mock_memory = MagicMock(spec=MemoryInterface) + mock_memory.get_conversation_messages.return_value = [_make_message(role="user", content="prepended")] + target._memory = mock_memory + tokenizer = MagicMock() + tokenizer.apply_chat_template.return_value = "TOKENIZED REQUEST" + formatter = TokenizerTemplateNormalizer(tokenizer=tokenizer) + target_context = _make_target_normalization_context( + prepended_message_count=1, + formatter=formatter, + ) + assert target_context.begin_normalization(conversation_id="conv1") result = await target._get_normalized_conversation_async( message=_make_message(role="user", content="live"), - prepended_conversation_normalizer=ConversationContextNormalizer(), + target_normalization_context=target_context, + should_apply_context=True, ) - assert result[0].get_value() == "Turn 1:\nuser: [Image_path]\n\nlive" + tokenizer_messages = tokenizer.apply_chat_template.call_args.args[0] + assert tokenizer_messages[-1] == {"role": "user", "content": "live"} + assert tokenizer.apply_chat_template.call_args.kwargs["add_generation_prompt"] is True + assert result[0].get_value() == "TOKENIZED REQUEST" diff --git a/tests/unit/prompt_target/target/test_target_normalization_context.py b/tests/unit/prompt_target/target/test_target_normalization_context.py new file mode 100644 index 0000000000..304a82d709 --- /dev/null +++ b/tests/unit/prompt_target/target/test_target_normalization_context.py @@ -0,0 +1,60 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT license. + +from unittest.mock import AsyncMock, MagicMock + +import pytest + +from pyrit.message_normalizer import MessageListNormalizer +from pyrit.models import Message +from pyrit.prompt_target import TargetNormalizationContext, TargetNormalizationContextState + + +def _make_context() -> tuple[TargetNormalizationContext, MagicMock]: + normalizer = MagicMock(spec=MessageListNormalizer) + normalizer.normalize_async = AsyncMock(side_effect=lambda messages: messages) + context = TargetNormalizationContext( + conversation_id="conversation", + normalizers=(normalizer,), + ) + return context, normalizer + + +async def test_context_normalizes_and_is_consumed_once(): + context, normalizer = _make_context() + messages = [Message.from_prompt(prompt="request", role="user")] + + assert context.begin_normalization(conversation_id="conversation") + assert context.state == TargetNormalizationContextState.PREPARING + assert await context.normalize_async(messages=messages) == messages + context.mark_consumed() + + assert context.is_consumed + assert context.begin_normalization(conversation_id="conversation") is False + normalizer.normalize_async.assert_awaited_once_with(messages) + + +def test_context_can_restore_pending_after_preparation_failure(): + context, _ = _make_context() + + assert context.begin_normalization(conversation_id="conversation") + context.restore_pending() + + assert context.is_pending + assert context.begin_normalization(conversation_id="conversation") + + +def test_context_rejects_concurrent_preparation(): + context, _ = _make_context() + + assert context.begin_normalization(conversation_id="conversation") + + with pytest.raises(RuntimeError, match="already in progress"): + context.begin_normalization(conversation_id="conversation") + + +def test_context_rejects_another_conversation(): + context, _ = _make_context() + + with pytest.raises(ValueError, match="belongs to conversation"): + context.begin_normalization(conversation_id="other") From 14ce81ebaa0016fdc17053dbfdf3f69b7fd452dc Mon Sep 17 00:00:00 2001 From: Copilot <223556219+Copilot@users.noreply.github.com> Date: Thu, 20 Aug 2026 06:52:56 -0700 Subject: [PATCH 13/24] Preserve system context across target rotation Create a fresh target normalization context when single-turn target rotation carries system messages into a new conversation. This keeps the next request single-message while retaining its system framing. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> Copilot-Session: 6776aa95-0924-4967-a056-e5a228af3faa --- .../multi_turn/multi_turn_attack_strategy.py | 8 ++- .../test_supports_multi_turn_attacks.py | 68 ++++++++++++++++++- 2 files changed, 73 insertions(+), 3 deletions(-) diff --git a/pyrit/executor/attack/multi_turn/multi_turn_attack_strategy.py b/pyrit/executor/attack/multi_turn/multi_turn_attack_strategy.py index 657a6f5e79..59d7b2df08 100644 --- a/pyrit/executor/attack/multi_turn/multi_turn_attack_strategy.py +++ b/pyrit/executor/attack/multi_turn/multi_turn_attack_strategy.py @@ -10,6 +10,7 @@ from typing import TYPE_CHECKING, Any, TypeVar from pyrit.common.logger import logger +from pyrit.executor.attack.component.conversation_manager import ConversationManager from pyrit.executor.attack.core.attack_parameters import AttackParameters, AttackParamsT from pyrit.executor.attack.core.attack_strategy import ( AttackContext, @@ -152,9 +153,14 @@ def _rotate_conversation_for_single_turn_target( ) memory.add_message_pieces_to_memory(message_pieces=pieces) context.session.conversation_id = new_conversation_id + context.target_normalization_context = ConversationManager.create_target_normalization_context( + target=self._objective_target, + conversation_id=new_conversation_id, + prepended_message_count=len(system_messages), + ) else: context.session.conversation_id = str(uuid.uuid4()) - context.target_normalization_context = None + context.target_normalization_context = None self._logger.debug( f"Rotated conversation_id for single-turn target: " diff --git a/tests/unit/executor/attack/multi_turn/test_supports_multi_turn_attacks.py b/tests/unit/executor/attack/multi_turn/test_supports_multi_turn_attacks.py index 348529fe12..3a2d86facb 100644 --- a/tests/unit/executor/attack/multi_turn/test_supports_multi_turn_attacks.py +++ b/tests/unit/executor/attack/multi_turn/test_supports_multi_turn_attacks.py @@ -12,8 +12,29 @@ ) from pyrit.memory import CentralMemory from pyrit.message_normalizer import ConversationContextNormalizer, FirstTurnHistoryNormalizer -from pyrit.models import ConversationType, MessagePiece -from pyrit.prompt_target import TargetNormalizationContext +from pyrit.models import ConversationType, Message, MessagePiece +from pyrit.prompt_target import PromptTarget, TargetNormalizationContext +from pyrit.prompt_target.common.target_capabilities import TargetCapabilities +from pyrit.prompt_target.common.target_configuration import TargetConfiguration + + +class _SingleTurnPromptTarget(PromptTarget): + _DEFAULT_CONFIGURATION = TargetConfiguration(capabilities=TargetCapabilities()) + + def __init__(self) -> None: + super().__init__() + self.normalized_conversations: list[list[Message]] = [] + + async def _send_prompt_to_target_async(self, *, normalized_conversation: list[Message]) -> list[Message]: + self.normalized_conversations.append(normalized_conversation) + request_piece = normalized_conversation[-1].get_piece() + return [ + MessagePiece( + role="assistant", + original_value="response", + conversation_id=request_piece.conversation_id, + ).to_message() + ] def _make_context() -> MultiTurnAttackContext: @@ -198,6 +219,49 @@ def test_system_prompt_preserved_across_multiple_rotations(self): ) memory.add_message_pieces_to_memory(message_pieces=[user_piece]) + async def test_rotated_system_prompt_is_normalized_with_next_request(self): + target = _SingleTurnPromptTarget() + strategy = _make_strategy(supports_multi_turn=False) + strategy._objective_target = target + context = _make_context() + old_id = context.session.conversation_id + _seed_conversation( + conversation_id=old_id, + system_prompt="You are a helpful assistant.", + user_text="First request", + ) + previous_context = TargetNormalizationContext( + conversation_id=old_id, + normalizers=( + FirstTurnHistoryNormalizer( + message_normalizer=ConversationContextNormalizer(), + prepended_message_count=1, + ), + ), + ) + previous_context.begin_normalization(conversation_id=old_id) + previous_context.mark_consumed() + context.target_normalization_context = previous_context + context.executed_turns = 1 + + strategy._rotate_conversation_for_single_turn_target(context=context) + + next_request = Message.from_prompt(prompt="Second request", role="user") + next_request.get_piece().conversation_id = context.session.conversation_id + await target.send_prompt_async( + message=next_request, + target_normalization_context=context.target_normalization_context, + ) + + assert context.target_normalization_context is not None + assert context.target_normalization_context.is_consumed + assert len(target.normalized_conversations) == 1 + normalized_conversation = target.normalized_conversations[0] + assert len(normalized_conversation) == 1 + assert normalized_conversation[0].get_value() == ( + "Turn 1:\nuser: ### Instructions ###\n\nYou are a helpful assistant.\n\n######\n\nSecond request" + ) + def test_no_system_prompt_yields_fresh_conversation_id(self): """When there is no system prompt, rotation still generates a new conversation_id.""" strategy = _make_strategy(supports_multi_turn=False) From 290a5b8f684683160bae26cd15828222c0494591 Mon Sep 17 00:00:00 2001 From: Copilot <223556219+Copilot@users.noreply.github.com> Date: Thu, 20 Aug 2026 16:29:39 -0700 Subject: [PATCH 14/24] Unify target history normalization Use HistorySquashNormalizer for both one-shot prepended-history adaptation and ordinary single-turn capability adaptation, with shared multimodal handling and documented lifecycle scopes. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> Copilot-Session: 6776aa95-0924-4967-a056-e5a228af3faa --- doc/code/framework.md | 12 +- doc/code/targets/11_message_normalizer.ipynb | 64 +++-- doc/code/targets/11_message_normalizer.py | 27 +- .../attack/component/conversation_manager.py | 6 +- pyrit/message_normalizer/__init__.py | 2 - pyrit/message_normalizer/_helpers.py | 45 +++- .../conversation_context_normalizer.py | 19 +- .../first_turn_history_normalizer.py | 183 ------------- .../history_squash_normalizer.py | 251 ++++++++++++++++-- pyrit/prompt_target/common/prompt_target.py | 8 +- .../component/test_conversation_manager.py | 7 +- .../test_supports_multi_turn_attacks.py | 10 +- .../test_history_squash_normalizer.py | 156 ++++++++++- .../test_prompt_normalizer.py | 18 +- .../test_normalize_async_integration.py | 6 +- 15 files changed, 525 insertions(+), 289 deletions(-) delete mode 100644 pyrit/message_normalizer/first_turn_history_normalizer.py diff --git a/doc/code/framework.md b/doc/code/framework.md index 11fdb02af8..f954e751d8 100644 --- a/doc/code/framework.md +++ b/doc/code/framework.md @@ -317,7 +317,17 @@ The below talks about responsibilities of most modules in the PyRIT library - **`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. Prepended history remains role-structured in memory. Attacks pass an ephemeral `TargetNormalizationContext` with the first live send when a target cannot edit history; the target consumes that context immediately before provider invocation. - **`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, prepended-history adaptation, and tokenizer chat templates. These target-specific views are ephemeral and do not replace the logical conversation in memory. -For prepended history on a target without editable history, processing order is: convert and persist the structured prepended messages, convert the live request, apply first-turn history normalization, run the target's ordinary capability normalizers, then serialize and invoke the provider. +`supports_multi_turn` and `supports_editable_history` answer different questions. Multi-turn support means the target can continue a conversation across sends. Editable-history support means PyRIT can supply or rewrite earlier turns before sending the current message. The corresponding normalizers therefore have different scopes: + +| Target capabilities | Prepended-history behavior | +| --- | --- | +| Multi-turn with editable history | The target receives the structured history directly; neither normalizer is needed. | +| Multi-turn without editable history | A context-scoped `HistorySquashNormalizer` encodes the prepended history and first live message into one initial request. Later turns use the target's own conversation state. No history squasher is added to the ordinary capability pipeline because the target supports multiple turns. | +| Single-turn without editable history | The context-scoped `HistorySquashNormalizer` encodes the prepended history and first live message. The ordinary capability pipeline contains another use of the same normalizer, but it receives the already combined message and is therefore a no-op for that first request. | + +The first-turn distinction is therefore a matter of lifecycle and configuration, not a separate flattening implementation. `TargetNormalizationContext` applies its configured `HistorySquashNormalizer` once to bootstrap prepended history for any target that cannot accept caller-supplied prior turns. A target's ordinary capability pipeline independently applies `HistorySquashNormalizer` whenever the target cannot receive a multi-message conversation. Both uses preserve non-text pieces from the current request while rendering historical content as text. + +For prepended history on a target without editable history, processing order is: convert and persist the structured prepended messages, convert the live request, apply context-scoped history squashing, run the target's ordinary capability normalizers, then serialize and invoke the provider. ## [Output](./output/0_output) diff --git a/doc/code/targets/11_message_normalizer.ipynb b/doc/code/targets/11_message_normalizer.ipynb index 1fd52230fa..0261daa565 100644 --- a/doc/code/targets/11_message_normalizer.ipynb +++ b/doc/code/targets/11_message_normalizer.ipynb @@ -25,10 +25,40 @@ "Some normalizers implement both interfaces." ] }, + { + "cell_type": "markdown", + "id": "1", + "metadata": {}, + "source": [ + "## Two History-Squashing Scopes\n", + "\n", + "PyRIT uses `HistorySquashNormalizer` in two places that adapt different target capabilities:\n", + "\n", + "- **Multi-turn support** answers whether the target can continue a conversation across sends.\n", + "- **Editable-history support** answers whether PyRIT can supply or rewrite earlier turns before\n", + " sending the current message.\n", + "\n", + "| Target capabilities | Prepended-history behavior |\n", + "|---|---|\n", + "| Multi-turn with editable history | Send the structured history directly; no history squashing is needed. |\n", + "| Multi-turn without editable history | A context-scoped `HistorySquashNormalizer` encodes the prepended history and first live message into the initial request. Later turns use the target's own conversation state. |\n", + "| Single-turn without editable history | The context-scoped `HistorySquashNormalizer` first produces one message. The target's ordinary use of the same normalizer then sees one message and does nothing. |\n", + "\n", + "The first-turn distinction comes from `TargetNormalizationContext`, not from a separate normalizer\n", + "implementation. The context applies its configured `HistorySquashNormalizer` once to bootstrap\n", + "prepended history for a target that cannot accept caller-supplied prior turns. Its key use case is a\n", + "multi-turn, server-managed target without editable history.\n", + "\n", + "The target's ordinary capability pipeline independently uses `HistorySquashNormalizer` for a target\n", + "without multi-turn support. Both scopes preserve original-versus-converted text views and keep\n", + "non-text pieces from the current request separate. The context-scoped use can supply a custom\n", + "formatter; the ordinary use defaults to `[Conversation History]` and `[Current Message]` sections." + ] + }, { "cell_type": "code", "execution_count": null, - "id": "1", + "id": "2", "metadata": {}, "outputs": [ { @@ -61,7 +91,7 @@ }, { "cell_type": "markdown", - "id": "2", + "id": "3", "metadata": {}, "source": [ "## ChatMessageNormalizer\n", @@ -77,7 +107,7 @@ { "cell_type": "code", "execution_count": null, - "id": "3", + "id": "4", "metadata": {}, "outputs": [ { @@ -107,7 +137,7 @@ { "cell_type": "code", "execution_count": null, - "id": "4", + "id": "5", "metadata": {}, "outputs": [ { @@ -135,7 +165,7 @@ { "cell_type": "code", "execution_count": null, - "id": "5", + "id": "6", "metadata": {}, "outputs": [ { @@ -173,7 +203,7 @@ }, { "cell_type": "markdown", - "id": "6", + "id": "7", "metadata": {}, "source": [ "## GenericSystemSquashNormalizer\n", @@ -201,7 +231,7 @@ { "cell_type": "code", "execution_count": null, - "id": "7", + "id": "8", "metadata": {}, "outputs": [ { @@ -236,7 +266,7 @@ }, { "cell_type": "markdown", - "id": "8", + "id": "9", "metadata": {}, "source": [ "## ConversationContextNormalizer\n", @@ -261,7 +291,7 @@ { "cell_type": "code", "execution_count": null, - "id": "9", + "id": "10", "metadata": {}, "outputs": [ { @@ -289,7 +319,7 @@ }, { "cell_type": "markdown", - "id": "10", + "id": "11", "metadata": {}, "source": [ "## TokenizerTemplateNormalizer\n", @@ -316,7 +346,7 @@ { "cell_type": "code", "execution_count": null, - "id": "11", + "id": "12", "metadata": {}, "outputs": [ { @@ -357,7 +387,7 @@ }, { "cell_type": "markdown", - "id": "12", + "id": "13", "metadata": {}, "source": [ "### System Message Behavior\n", @@ -373,7 +403,7 @@ { "cell_type": "code", "execution_count": null, - "id": "13", + "id": "14", "metadata": {}, "outputs": [ { @@ -416,7 +446,7 @@ }, { "cell_type": "markdown", - "id": "14", + "id": "15", "metadata": {}, "source": [ "### Using Custom Models\n", @@ -427,7 +457,7 @@ { "cell_type": "code", "execution_count": null, - "id": "15", + "id": "16", "metadata": {}, "outputs": [ { @@ -467,7 +497,7 @@ }, { "cell_type": "markdown", - "id": "16", + "id": "17", "metadata": {}, "source": [ "## Creating Custom Normalizers\n", @@ -478,7 +508,7 @@ { "cell_type": "code", "execution_count": null, - "id": "17", + "id": "18", "metadata": {}, "outputs": [ { diff --git a/doc/code/targets/11_message_normalizer.py b/doc/code/targets/11_message_normalizer.py index 87e8cb4441..52e1ef59b8 100644 --- a/doc/code/targets/11_message_normalizer.py +++ b/doc/code/targets/11_message_normalizer.py @@ -5,7 +5,7 @@ # extension: .py # format_name: percent # format_version: '1.3' -# jupytext_version: 1.19.4 +# jupytext_version: 1.19.5 # --- # %% [markdown] @@ -28,6 +28,31 @@ # # Some normalizers implement both interfaces. +# %% [markdown] +# ## Two History-Squashing Scopes +# +# PyRIT uses `HistorySquashNormalizer` in two places that adapt different target capabilities: +# +# - **Multi-turn support** answers whether the target can continue a conversation across sends. +# - **Editable-history support** answers whether PyRIT can supply or rewrite earlier turns before +# sending the current message. +# +# | Target capabilities | Prepended-history behavior | +# |---|---| +# | Multi-turn with editable history | Send the structured history directly; no history squashing is needed. | +# | Multi-turn without editable history | A context-scoped `HistorySquashNormalizer` encodes the prepended history and first live message into the initial request. Later turns use the target's own conversation state. | +# | Single-turn without editable history | The context-scoped `HistorySquashNormalizer` first produces one message. The target's ordinary use of the same normalizer then sees one message and does nothing. | +# +# The first-turn distinction comes from `TargetNormalizationContext`, not from a separate normalizer +# implementation. The context applies its configured `HistorySquashNormalizer` once to bootstrap +# prepended history for a target that cannot accept caller-supplied prior turns. Its key use case is a +# multi-turn, server-managed target without editable history. +# +# The target's ordinary capability pipeline independently uses `HistorySquashNormalizer` for a target +# without multi-turn support. Both scopes preserve original-versus-converted text views and keep +# non-text pieces from the current request separate. The context-scoped use can supply a custom +# formatter; the ordinary use defaults to `[Conversation History]` and `[Current Message]` sections. + # %% from pyrit.models import Message diff --git a/pyrit/executor/attack/component/conversation_manager.py b/pyrit/executor/attack/component/conversation_manager.py index 89470202f3..3320b6efd1 100644 --- a/pyrit/executor/attack/component/conversation_manager.py +++ b/pyrit/executor/attack/component/conversation_manager.py @@ -13,7 +13,7 @@ PrependedConversationConfig, ) from pyrit.memory import CentralMemory -from pyrit.message_normalizer import ConversationContextNormalizer, FirstTurnHistoryNormalizer +from pyrit.message_normalizer import ConversationContextNormalizer, HistorySquashNormalizer from pyrit.models import ( ChatMessageRole, ComponentIdentifier, @@ -464,9 +464,9 @@ def create_target_normalization_context( return TargetNormalizationContext( conversation_id=conversation_id, normalizers=( - FirstTurnHistoryNormalizer( + HistorySquashNormalizer( message_normalizer=config.get_message_normalizer(), - prepended_message_count=prepended_message_count, + expected_history_message_count=prepended_message_count, ), ), ) diff --git a/pyrit/message_normalizer/__init__.py b/pyrit/message_normalizer/__init__.py index 0580bd48b1..79df1cb50d 100644 --- a/pyrit/message_normalizer/__init__.py +++ b/pyrit/message_normalizer/__init__.py @@ -7,7 +7,6 @@ from pyrit.message_normalizer.chat_message_normalizer import ChatMessageNormalizer from pyrit.message_normalizer.conversation_context_normalizer import ConversationContextNormalizer -from pyrit.message_normalizer.first_turn_history_normalizer import FirstTurnHistoryNormalizer from pyrit.message_normalizer.generic_system_squash import GenericSystemSquashNormalizer from pyrit.message_normalizer.history_squash_normalizer import HistorySquashNormalizer from pyrit.message_normalizer.json_schema_normalizer import JsonSchemaNormalizer @@ -20,7 +19,6 @@ __all__ = [ "MessageListNormalizer", "MessageStringNormalizer", - "FirstTurnHistoryNormalizer", "GenericSystemSquashNormalizer", "HistorySquashNormalizer", "JsonSchemaNormalizer", diff --git a/pyrit/message_normalizer/_helpers.py b/pyrit/message_normalizer/_helpers.py index 081ff5b2fe..cbc73e16fd 100644 --- a/pyrit/message_normalizer/_helpers.py +++ b/pyrit/message_normalizer/_helpers.py @@ -2,18 +2,43 @@ # Licensed under the MIT license. """ -Internal helpers shared by squash-style message normalizers. - -Squash normalizers (e.g. ``HistorySquashNormalizer``, -``GenericSystemSquashNormalizer``) collapse several input messages into one -fresh user-role ``Message`` built via ``Message.from_prompt``. Because that -factory creates a brand-new piece with empty ``prompt_metadata``, callers must -explicitly carry request-level metadata (such as the JSON schema key) forward -so downstream normalizers in the pipeline still see it. ``build_squashed_user_message`` -centralizes that propagation rule. +Internal helpers for squash-style message normalizers. + +Some squash normalizers, such as ``GenericSystemSquashNormalizer``, collapse +several input messages into one fresh user-role ``Message`` built via +``Message.from_prompt``. Because that factory creates a brand-new piece with +empty ``prompt_metadata``, callers must explicitly carry request-level metadata +forward so downstream normalizers still see it. +``build_squashed_user_message`` centralizes that propagation rule. """ -from pyrit.models import Message +from pyrit.models import Message, MessagePiece + + +def format_message_piece_for_context(*, piece: MessagePiece) -> str: + """ + Format one message piece for inclusion in textual conversation history. + + Non-text pieces use their context description when available and otherwise + use a modality placeholder, so local asset paths are not exposed as prompt + text. + + Args: + piece (MessagePiece): The piece to represent as context. + + Returns: + str: The textual representation of the piece. + """ + data_type = piece.converted_value_data_type or piece.original_value_data_type + if data_type != "text": + description = piece.prompt_metadata.get("context_description") + if description: + return f"[{data_type.capitalize()} - {description}]" + return f"[{data_type.capitalize()}]" + + if piece.original_value != piece.converted_value: + return f"{piece.converted_value} (original: {piece.original_value})" + return piece.converted_value def build_squashed_user_message(*, new_message_content: str, source_messages: list[Message]) -> Message: diff --git a/pyrit/message_normalizer/conversation_context_normalizer.py b/pyrit/message_normalizer/conversation_context_normalizer.py index a7b7d60a17..b3a7bb78e7 100644 --- a/pyrit/message_normalizer/conversation_context_normalizer.py +++ b/pyrit/message_normalizer/conversation_context_normalizer.py @@ -1,7 +1,7 @@ # Copyright (c) Microsoft Corporation. # Licensed under the MIT license. - +from pyrit.message_normalizer._helpers import format_message_piece_for_context from pyrit.message_normalizer.message_normalizer import MessageStringNormalizer from pyrit.models import Message, MessagePiece @@ -72,19 +72,4 @@ def _format_piece_content(self, piece: MessagePiece) -> str: Returns: The formatted content string. """ - data_type = piece.converted_value_data_type or piece.original_value_data_type - - # For non-text pieces, use metadata description or placeholder - if data_type != "text": - if piece.prompt_metadata and "context_description" in piece.prompt_metadata: - description = piece.prompt_metadata["context_description"] - return f"[{data_type.capitalize()} - {description}]" - return f"[{data_type.capitalize()}]" - - # For text pieces, include both original and converted if different - original = piece.original_value - converted = piece.converted_value - - if original != converted: - return f"{converted} (original: {original})" - return converted + return format_message_piece_for_context(piece=piece) diff --git a/pyrit/message_normalizer/first_turn_history_normalizer.py b/pyrit/message_normalizer/first_turn_history_normalizer.py deleted file mode 100644 index 8f0541ace7..0000000000 --- a/pyrit/message_normalizer/first_turn_history_normalizer.py +++ /dev/null @@ -1,183 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# Licensed under the MIT license. - -import copy -import uuid - -from pyrit.message_normalizer.conversation_context_normalizer import ConversationContextNormalizer -from pyrit.message_normalizer.generic_system_squash import GenericSystemSquashNormalizer -from pyrit.message_normalizer.message_normalizer import MessageListNormalizer, MessageStringNormalizer -from pyrit.models import Message, MessagePiece - - -class FirstTurnHistoryNormalizer(MessageListNormalizer[Message]): - """ - Combine structured prepended history with the first live request. - - The configured string normalizer receives both the prepended history and - the live request, so tokenizer templates place generation markers after the - live user content. Non-text live pieces remain separate in the target view. - """ - - def __init__( - self, - *, - message_normalizer: MessageStringNormalizer, - prepended_message_count: int, - ) -> None: - """ - Initialize the normalizer. - - Args: - message_normalizer: Formatter for the target-facing text. - prepended_message_count: Number of leading messages that belong to - the prepended conversation. - - Raises: - ValueError: If prepended_message_count is less than one. - """ - if prepended_message_count < 1: - raise ValueError("prepended_message_count must be at least 1") - self._message_normalizer = message_normalizer - self._prepended_message_count = prepended_message_count - - async def normalize_async(self, messages: list[Message]) -> list[Message]: - """ - Return one target-facing request containing the prepended history. - - Args: - messages: Prepended history followed by exactly one live request. - - Returns: - A single request message with formatted text and preserved live - non-text pieces. - - Raises: - ValueError: If the input does not contain the configured prepended - message count followed by one live request, or if converted - prepended history cannot be represented as text. - """ - expected_count = self._prepended_message_count + 1 - if len(messages) != expected_count: - raise ValueError( - "First-turn history normalization expected " - f"{self._prepended_message_count} prepended messages and one live request, " - f"but received {len(messages)} messages." - ) - - prepended_messages = messages[: self._prepended_message_count] - live_request = messages[-1] - self._validate_flattenable_converter_output(messages=prepended_messages) - - original_view = self._build_original_view(messages=messages) - converted_view = self._build_converted_view(messages=messages) - original_text = await self._normalize_context_async(messages=original_view) - converted_text = original_text - if self._contains_converted_values(messages=messages): - converted_text = await self._normalize_context_async(messages=converted_view) - - return [ - self._build_target_request( - live_request=live_request, - original_text=original_text, - converted_text=converted_text, - ) - ] - - async def _normalize_context_async(self, *, messages: list[Message]) -> str: - messages_to_normalize = self._filter_live_non_text_pieces(messages=messages) - if isinstance(self._message_normalizer, ConversationContextNormalizer): - messages_to_normalize = await GenericSystemSquashNormalizer().normalize_async(messages_to_normalize) - return await self._message_normalizer.normalize_string_async(messages_to_normalize) - - @staticmethod - def _filter_live_non_text_pieces(*, messages: list[Message]) -> list[Message]: - filtered = copy.deepcopy(messages) - live_message = filtered[-1] - text_pieces = [piece for piece in live_message.message_pieces if piece.converted_value_data_type == "text"] - if not text_pieces: - filtered.pop() - else: - live_message.message_pieces = text_pieces - return filtered - - @staticmethod - def _build_original_view(*, messages: list[Message]) -> list[Message]: - original_messages = copy.deepcopy(messages) - for message in original_messages: - for piece in message.message_pieces: - piece.converted_value = piece.original_value - piece.converted_value_data_type = piece.original_value_data_type - return original_messages - - @staticmethod - def _build_converted_view(*, messages: list[Message]) -> list[Message]: - converted_messages = copy.deepcopy(messages) - for message in converted_messages: - for piece in message.message_pieces: - piece.original_value = piece.converted_value - piece.original_value_data_type = piece.converted_value_data_type - return converted_messages - - @staticmethod - def _contains_converted_values(*, messages: list[Message]) -> bool: - return any( - piece.converter_identifiers - or piece.original_value != piece.converted_value - or piece.original_value_data_type != piece.converted_value_data_type - for message in messages - for piece in message.message_pieces - ) - - @staticmethod - def _validate_flattenable_converter_output(*, messages: list[Message]) -> None: - output_types = { - piece.converted_value_data_type - for message in messages - for piece in message.message_pieces - if piece.converted_value_data_type != "text" - and ( - piece.original_value != piece.converted_value - or piece.original_value_data_type != piece.converted_value_data_type - ) - } - if output_types: - raise ValueError( - "Cannot flatten prepended conversation after request converters produced " - f"non-text output types {sorted(output_types)}. Prepended conversion must produce " - "text for a target without editable history." - ) - - @staticmethod - def _build_target_request( - *, - live_request: Message, - original_text: str, - converted_text: str, - ) -> Message: - request = copy.deepcopy(live_request) - template_piece = request.get_piece() - text_piece = MessagePiece( - id=uuid.uuid4(), - role=template_piece.role, - original_value=original_text, - converted_value=converted_text, - original_value_data_type="text", - converted_value_data_type="text", - conversation_id=template_piece.conversation_id, - sequence=template_piece.sequence, - prompt_metadata=dict(template_piece.prompt_metadata), - ) - target_pieces: list[MessagePiece] = [] - text_inserted = False - for piece in request.message_pieces: - if piece.converted_value_data_type == "text": - if not text_inserted: - target_pieces.append(text_piece) - text_inserted = True - continue - target_pieces.append(piece) - if not text_inserted: - target_pieces.insert(0, text_piece) - request.message_pieces = target_pieces - return request diff --git a/pyrit/message_normalizer/history_squash_normalizer.py b/pyrit/message_normalizer/history_squash_normalizer.py index f0369af682..0a08ca0778 100644 --- a/pyrit/message_normalizer/history_squash_normalizer.py +++ b/pyrit/message_normalizer/history_squash_normalizer.py @@ -1,64 +1,257 @@ # Copyright (c) Microsoft Corporation. # Licensed under the MIT license. -from pyrit.message_normalizer._helpers import build_squashed_user_message -from pyrit.message_normalizer.message_normalizer import MessageListNormalizer -from pyrit.models import Message +import copy +import uuid + +from pyrit.message_normalizer._helpers import format_message_piece_for_context +from pyrit.message_normalizer.conversation_context_normalizer import ConversationContextNormalizer +from pyrit.message_normalizer.generic_system_squash import GenericSystemSquashNormalizer +from pyrit.message_normalizer.message_normalizer import MessageListNormalizer, MessageStringNormalizer +from pyrit.models import Message, MessagePiece class HistorySquashNormalizer(MessageListNormalizer[Message]): """ - Squashes a multi-turn conversation into a single user message. + Combine conversation history and the current request into one message. + + The same implementation serves two normalization scopes. A + ``TargetNormalizationContext`` configures it with a string formatter and an + expected history count to bootstrap prepended history exactly once for a target + without editable history. The ordinary target capability pipeline uses the + default formatter whenever a target does not support multiple turns. - Previous turns are formatted as labeled context and prepended to the - latest message. Used by the normalization pipeline to adapt prompts - for targets that do not support multi-turn conversations. + The surrounding pipeline controls when the normalizer runs; this class does not + track turn state. In both scopes, historical content becomes text while non-text + pieces from the current request remain separate target-facing pieces. """ + def __init__( + self, + *, + message_normalizer: MessageStringNormalizer | None = None, + expected_history_message_count: int | None = None, + ) -> None: + """ + Initialize the normalizer. + + Args: + message_normalizer (MessageStringNormalizer | None): Optional formatter + for the combined text. When + omitted, use the labeled conversation-history format. + expected_history_message_count (int | None): Optional exact number of + messages expected before the current request. The one-shot + prepended-history scope sets this value; the ordinary capability + pipeline does not. + + Raises: + ValueError: If expected_history_message_count is less than one. + """ + if expected_history_message_count is not None and expected_history_message_count < 1: + raise ValueError("expected_history_message_count must be at least 1") + self._message_normalizer = message_normalizer + self._expected_history_message_count = expected_history_message_count + async def normalize_async(self, messages: list[Message]) -> list[Message]: """ - Combine all messages into a single user message. + Combine history and the current request into one target-facing message. - When there is only one message it is returned unchanged. Otherwise - all prior turns are formatted as ``Role: content`` lines under a - ``[Conversation History]`` header and the last message's content - appears under a ``[Current Message]`` header. + When there is only one message it is returned unchanged. Otherwise, the + configured formatter receives both history and current text. The combined + text replaces the current request's text pieces, while current non-text + pieces retain their original positions. Args: - messages: The conversation messages to squash. + messages (list[Message]): The conversation messages to squash. Returns: - list[Message]: A single-element list containing the squashed message. + list[Message]: A single-element list containing the target-facing message. Raises: - ValueError: If the messages list is empty. + ValueError: If messages is empty, the expected history count does not + match, or converted history cannot be represented as text. """ if not messages: raise ValueError("Messages list cannot be empty") + self._validate_expected_message_count(messages=messages) if len(messages) == 1: return list(messages) - history_lines = self._format_history(messages=messages[:-1]) - current_parts = [piece.converted_value for piece in messages[-1].message_pieces] + history = messages[:-1] + live_request = messages[-1] + self._validate_flattenable_converter_output(messages=history) - combined = ( - "[Conversation History]\n" + "\n".join(history_lines) + "\n\n[Current Message]\n" + "\n".join(current_parts) - ) + original_view = self._build_original_view(messages=messages) + converted_view = self._build_converted_view(messages=messages) + original_text = await self._normalize_context_async(messages=original_view) + converted_text = original_text + if self._contains_converted_values(messages=messages): + converted_text = await self._normalize_context_async(messages=converted_view) + + return [ + self._build_target_request( + live_request=live_request, + original_text=original_text, + converted_text=converted_text, + ) + ] + + def _validate_expected_message_count(self, *, messages: list[Message]) -> None: + """ + Validate the optional history-boundary contract. + + Args: + messages (list[Message]): History followed by the current request. + + Raises: + ValueError: If the configured history count does not match the input. + """ + if self._expected_history_message_count is None: + return - return [build_squashed_user_message(new_message_content=combined, source_messages=messages)] + expected_count = self._expected_history_message_count + 1 + if len(messages) != expected_count: + raise ValueError( + "History squash expected " + f"{self._expected_history_message_count} history messages and one current request, " + f"but received {len(messages)} messages." + ) - def _format_history(self, *, messages: list[Message]) -> list[str]: + async def _normalize_context_async(self, *, messages: list[Message]) -> str: """ - Format prior messages as ``Role: content`` lines. + Format the text portion of history and the current request. Args: - messages: The history messages to format. + messages (list[Message]): Original-view or converted-view messages to + format. Returns: - list[str]: One line per message piece. + str: The combined target-facing text. """ - lines: list[str] = [] - for msg in messages: - lines.extend(f"{piece.api_role.capitalize()}: {piece.converted_value}" for piece in msg.message_pieces) - return lines + if self._message_normalizer is None: + return self._format_default_context(messages=messages) + + messages_to_normalize = self._filter_live_non_text_pieces(messages=messages) + if isinstance(self._message_normalizer, ConversationContextNormalizer): + messages_to_normalize = await GenericSystemSquashNormalizer().normalize_async(messages_to_normalize) + return await self._message_normalizer.normalize_string_async(messages_to_normalize) + + @staticmethod + def _format_default_context(*, messages: list[Message]) -> str: + """ + Format history with labeled roles and current text in a separate section. + + Args: + messages (list[Message]): History followed by the current request. + + Returns: + str: The default labeled history representation. + """ + history_lines = [ + f"{piece.api_role.capitalize()}: {format_message_piece_for_context(piece=piece)}" + for message in messages[:-1] + for piece in message.message_pieces + ] + current_parts = [ + piece.converted_value for piece in messages[-1].message_pieces if piece.converted_value_data_type == "text" + ] + + sections = ["[Conversation History]\n" + "\n".join(history_lines)] + if current_parts: + sections.append("[Current Message]\n" + "\n".join(current_parts)) + return "\n\n".join(sections) + + @staticmethod + def _filter_live_non_text_pieces(*, messages: list[Message]) -> list[Message]: + filtered = copy.deepcopy(messages) + live_message = filtered[-1] + text_pieces = [piece for piece in live_message.message_pieces if piece.converted_value_data_type == "text"] + if not text_pieces: + filtered.pop() + else: + live_message.message_pieces = text_pieces + return filtered + + @staticmethod + def _build_original_view(*, messages: list[Message]) -> list[Message]: + original_messages = copy.deepcopy(messages) + for message in original_messages: + for piece in message.message_pieces: + piece.converted_value = piece.original_value + piece.converted_value_data_type = piece.original_value_data_type + return original_messages + + @staticmethod + def _build_converted_view(*, messages: list[Message]) -> list[Message]: + converted_messages = copy.deepcopy(messages) + for message in converted_messages: + for piece in message.message_pieces: + piece.original_value = piece.converted_value + piece.original_value_data_type = piece.converted_value_data_type + return converted_messages + + @staticmethod + def _contains_converted_values(*, messages: list[Message]) -> bool: + return any( + piece.converter_identifiers + or piece.original_value != piece.converted_value + or piece.original_value_data_type != piece.converted_value_data_type + for message in messages + for piece in message.message_pieces + ) + + @staticmethod + def _validate_flattenable_converter_output(*, messages: list[Message]) -> None: + output_types = { + piece.converted_value_data_type + for message in messages + for piece in message.message_pieces + if piece.converted_value_data_type != "text" + and ( + piece.original_value != piece.converted_value + or piece.original_value_data_type != piece.converted_value_data_type + ) + } + if output_types: + raise ValueError( + "Cannot flatten conversation history after request converters produced " + f"non-text output types {sorted(output_types)}. Historical conversion must produce text." + ) + + @staticmethod + def _build_target_request( + *, + live_request: Message, + original_text: str, + converted_text: str, + ) -> Message: + request = copy.deepcopy(live_request) + template_piece = next( + (piece for piece in request.message_pieces if piece.converted_value_data_type == "text"), + request.get_piece(), + ) + text_piece = MessagePiece( + id=uuid.uuid4(), + role=template_piece.role, + original_value=original_text, + converted_value=converted_text, + original_value_data_type="text", + converted_value_data_type="text", + conversation_id=template_piece.conversation_id, + sequence=template_piece.sequence, + prompt_metadata=dict(template_piece.prompt_metadata), + ) + target_pieces: list[MessagePiece] = [] + text_inserted = False + for piece in request.message_pieces: + if piece.converted_value_data_type == "text": + if not text_inserted: + target_pieces.append(text_piece) + text_inserted = True + continue + target_pieces.append(piece) + if not text_inserted: + target_pieces.insert(0, text_piece) + request.message_pieces = target_pieces + return request diff --git a/pyrit/prompt_target/common/prompt_target.py b/pyrit/prompt_target/common/prompt_target.py index bde5edcda1..84e7ca12b2 100644 --- a/pyrit/prompt_target/common/prompt_target.py +++ b/pyrit/prompt_target/common/prompt_target.py @@ -318,10 +318,10 @@ def _propagate_lineage(*, source: Message, target_message: Message) -> None: """ Copy request-lineage metadata from ``source`` onto every piece in ``target_message``. - Normalizers may create brand-new ``Message`` objects (e.g. ``HistorySquashNormalizer`` - uses ``Message.from_prompt``) that carry fresh random ``conversation_id`` values and - lack request lineage. This method restores the original metadata so that the response - built from the normalized message stays part of the correct conversation and retains + Normalizers may create brand-new messages or pieces, such as the combined + text piece from ``HistorySquashNormalizer``, that lack request lineage. + This method restores the original metadata so that the response built from + the normalized message stays part of the correct conversation and retains traceability. ``prompt_metadata`` is handled by provenance so that metadata-editing normalizers diff --git a/tests/unit/executor/attack/component/test_conversation_manager.py b/tests/unit/executor/attack/component/test_conversation_manager.py index 40146b3d58..0462301403 100644 --- a/tests/unit/executor/attack/component/test_conversation_manager.py +++ b/tests/unit/executor/attack/component/test_conversation_manager.py @@ -35,7 +35,7 @@ ) from pyrit.executor.attack.core import AttackContext from pyrit.executor.attack.core.attack_parameters import AttackParameters -from pyrit.message_normalizer import ConversationContextNormalizer, FirstTurnHistoryNormalizer +from pyrit.message_normalizer import ConversationContextNormalizer, HistorySquashNormalizer from pyrit.models import ComponentIdentifier, Message, MessagePiece, PromptDataType, Score from pyrit.prompt_normalizer import ConverterConfiguration, PromptNormalizer from pyrit.prompt_target import PromptTarget @@ -804,8 +804,9 @@ async def test_non_editable_target_sets_custom_first_send_context( assert context.target_normalization_context is not None assert context.target_normalization_context.conversation_id == conversation_id normalizer = context.target_normalization_context.normalizers[0] - assert isinstance(normalizer, FirstTurnHistoryNormalizer) + assert isinstance(normalizer, HistorySquashNormalizer) assert normalizer._message_normalizer is message_normalizer + assert normalizer._expected_history_message_count == len(sample_conversation) message_normalizer.normalize_string_async.assert_not_called() async def test_returns_turn_count_for_multi_turn_attacks( @@ -1395,7 +1396,7 @@ async def test_message_normalizer_default_uses_conversation_context_normalizer( assert context.target_normalization_context is not None normalizer = context.target_normalization_context.normalizers[0] - assert isinstance(normalizer, FirstTurnHistoryNormalizer) + assert isinstance(normalizer, HistorySquashNormalizer) assert isinstance(normalizer._message_normalizer, ConversationContextNormalizer) # ------------------------------------------------------------------------- diff --git a/tests/unit/executor/attack/multi_turn/test_supports_multi_turn_attacks.py b/tests/unit/executor/attack/multi_turn/test_supports_multi_turn_attacks.py index 3a2d86facb..40aa36be24 100644 --- a/tests/unit/executor/attack/multi_turn/test_supports_multi_turn_attacks.py +++ b/tests/unit/executor/attack/multi_turn/test_supports_multi_turn_attacks.py @@ -11,7 +11,7 @@ MultiTurnAttackContext, ) from pyrit.memory import CentralMemory -from pyrit.message_normalizer import ConversationContextNormalizer, FirstTurnHistoryNormalizer +from pyrit.message_normalizer import ConversationContextNormalizer, HistorySquashNormalizer from pyrit.models import ConversationType, Message, MessagePiece from pyrit.prompt_target import PromptTarget, TargetNormalizationContext from pyrit.prompt_target.common.target_capabilities import TargetCapabilities @@ -119,9 +119,9 @@ def test_pending_first_turn_context_suppresses_rotation_after_prepended_turns(se context.target_normalization_context = TargetNormalizationContext( conversation_id=original_id, normalizers=( - FirstTurnHistoryNormalizer( + HistorySquashNormalizer( message_normalizer=ConversationContextNormalizer(), - prepended_message_count=4, + expected_history_message_count=4, ), ), ) @@ -233,9 +233,9 @@ async def test_rotated_system_prompt_is_normalized_with_next_request(self): previous_context = TargetNormalizationContext( conversation_id=old_id, normalizers=( - FirstTurnHistoryNormalizer( + HistorySquashNormalizer( message_normalizer=ConversationContextNormalizer(), - prepended_message_count=1, + expected_history_message_count=1, ), ), ) diff --git a/tests/unit/message_normalizer/test_history_squash_normalizer.py b/tests/unit/message_normalizer/test_history_squash_normalizer.py index 7ac3648271..c9dad2c2fc 100644 --- a/tests/unit/message_normalizer/test_history_squash_normalizer.py +++ b/tests/unit/message_normalizer/test_history_squash_normalizer.py @@ -1,10 +1,12 @@ # Copyright (c) Microsoft Corporation. # Licensed under the MIT license. +from unittest.mock import AsyncMock, MagicMock + import pytest -from pyrit.message_normalizer import HistorySquashNormalizer -from pyrit.models import Message, MessagePiece +from pyrit.message_normalizer import HistorySquashNormalizer, MessageStringNormalizer +from pyrit.models import JSON_SCHEMA_METADATA_KEY, Message, MessagePiece from pyrit.models.literals import ChatMessageRole @@ -25,6 +27,21 @@ async def test_history_squash_single_message_returns_unchanged(): assert result[0].api_role == "user" +def test_history_squash_rejects_invalid_expected_history_count(): + with pytest.raises(ValueError, match="expected_history_message_count must be at least 1"): + HistorySquashNormalizer(expected_history_message_count=0) + + +async def test_history_squash_rejects_unexpected_history_count(): + messages = [ + _make_message("user", "history"), + _make_message("user", "current"), + ] + + with pytest.raises(ValueError, match="expected 2 history messages.*received 2 messages"): + await HistorySquashNormalizer(expected_history_message_count=2).normalize_async(messages) + + async def test_history_squash_two_turns(): messages = [ _make_message("user", "hello"), @@ -44,6 +61,23 @@ async def test_history_squash_two_turns(): assert "how are you?" in text +async def test_history_squash_uses_configured_formatter(): + formatter = MagicMock(spec=MessageStringNormalizer) + formatter.normalize_string_async = AsyncMock(return_value="custom format") + messages = [ + _make_message("user", "history"), + _make_message("user", "current"), + ] + + result = await HistorySquashNormalizer( + message_normalizer=formatter, + expected_history_message_count=1, + ).normalize_async(messages) + + assert result[0].get_value() == "custom format" + formatter.normalize_string_async.assert_awaited_once() + + async def test_history_squash_includes_system_in_history(): messages = [ _make_message("system", "You are helpful"), @@ -80,6 +114,124 @@ async def test_history_squash_multi_piece_message(): assert "part2" in text +async def test_history_squash_preserves_original_and_converted_views(): + history = _make_message("user", "original history") + history.get_piece().converted_value = "converted history" + current = _make_message("user", "original current") + current.get_piece().converted_value = "converted current" + + result = await HistorySquashNormalizer().normalize_async([history, current]) + + piece = result[0].get_piece() + assert "User: original history" in piece.original_value + assert "original current" in piece.original_value + assert "User: converted history" in piece.converted_value + assert "converted current" in piece.converted_value + + +async def test_history_squash_preserves_live_multimodal_piece_order(): + conversation_id = "test-conv-id" + current = Message( + message_pieces=[ + MessagePiece( + role="user", + original_value="diagram.png", + original_value_data_type="image_path", + conversation_id=conversation_id, + ), + MessagePiece( + role="user", + original_value="What does this show?", + conversation_id=conversation_id, + ), + ] + ) + + result = await HistorySquashNormalizer().normalize_async([_make_message("assistant", "Earlier response"), current]) + + pieces = result[0].message_pieces + assert [piece.converted_value_data_type for piece in pieces] == ["image_path", "text"] + assert pieces[0].converted_value == "diagram.png" + assert "What does this show?" in pieces[1].converted_value + assert "diagram.png" not in pieces[1].converted_value + + +async def test_history_squash_preserves_entirely_non_text_live_request(): + conversation_id = "test-conv-id" + current = Message( + message_pieces=[ + MessagePiece( + role="user", + original_value="diagram.png", + original_value_data_type="image_path", + conversation_id=conversation_id, + ) + ] + ) + + result = await HistorySquashNormalizer().normalize_async([_make_message("assistant", "Earlier response"), current]) + + pieces = result[0].message_pieces + assert [piece.converted_value_data_type for piece in pieces] == ["text", "image_path"] + assert pieces[0].converted_value == "[Conversation History]\nAssistant: Earlier response" + assert pieces[1].converted_value == "diagram.png" + + +async def test_history_squash_describes_non_text_history_without_exposing_path(): + history = Message( + message_pieces=[ + MessagePiece( + role="user", + original_value="C:\\private\\diagram.png", + original_value_data_type="image_path", + prompt_metadata={"context_description": "architecture diagram"}, + ) + ] + ) + + result = await HistorySquashNormalizer().normalize_async([history, _make_message("user", "What does it show?")]) + + text = result[0].get_value() + assert "User: [Image_path - architecture diagram]" in text + assert "C:\\private\\diagram.png" not in text + + +async def test_history_squash_uses_live_text_metadata_for_combined_piece(): + schema = {"type": "object"} + conversation_id = "test-conv-id" + current = Message( + message_pieces=[ + MessagePiece( + role="user", + original_value="diagram.png", + original_value_data_type="image_path", + conversation_id=conversation_id, + ), + MessagePiece( + role="user", + original_value="Describe this image", + conversation_id=conversation_id, + prompt_metadata={JSON_SCHEMA_METADATA_KEY: schema}, + ), + ] + ) + + result = await HistorySquashNormalizer().normalize_async([_make_message("assistant", "Earlier response"), current]) + + text_piece = result[0].message_pieces[1] + assert text_piece.converted_value_data_type == "text" + assert text_piece.prompt_metadata == {JSON_SCHEMA_METADATA_KEY: schema} + + +async def test_history_squash_rejects_converted_non_text_history(): + history = _make_message("user", "original") + history.get_piece().converted_value = "converted.png" + history.get_piece().converted_value_data_type = "image_path" + + with pytest.raises(ValueError, match="non-text output types.*image_path"): + await HistorySquashNormalizer().normalize_async([history, _make_message("user", "current")]) + + async def test_history_squash_preserves_original_list(): """Normalize should not mutate the input list.""" messages = [ diff --git a/tests/unit/prompt_normalizer/test_prompt_normalizer.py b/tests/unit/prompt_normalizer/test_prompt_normalizer.py index b314568a18..e27860b228 100644 --- a/tests/unit/prompt_normalizer/test_prompt_normalizer.py +++ b/tests/unit/prompt_normalizer/test_prompt_normalizer.py @@ -25,7 +25,7 @@ get_execution_context, ) from pyrit.memory import CentralMemory -from pyrit.message_normalizer import ConversationContextNormalizer, FirstTurnHistoryNormalizer +from pyrit.message_normalizer import ConversationContextNormalizer, HistorySquashNormalizer from pyrit.models import ( Message, MessagePiece, @@ -135,9 +135,9 @@ async def test_send_prompt_async_forwards_target_normalization_context(mock_memo target_context = TargetNormalizationContext( conversation_id=conversation_id, normalizers=( - FirstTurnHistoryNormalizer( + HistorySquashNormalizer( message_normalizer=ConversationContextNormalizer(), - prepended_message_count=1, + expected_history_message_count=1, ), ), ) @@ -161,9 +161,9 @@ async def test_send_prompt_async_conversion_failure_leaves_context_pending(mock_ target_context = TargetNormalizationContext( conversation_id=conversation_id, normalizers=( - FirstTurnHistoryNormalizer( + HistorySquashNormalizer( message_normalizer=ConversationContextNormalizer(), - prepended_message_count=1, + expected_history_message_count=1, ), ), ) @@ -191,9 +191,9 @@ async def test_send_prompt_async_pre_provider_failure_is_not_persisted(mock_memo target_context = TargetNormalizationContext( conversation_id=conversation_id, normalizers=( - FirstTurnHistoryNormalizer( + HistorySquashNormalizer( message_normalizer=ConversationContextNormalizer(), - prepended_message_count=1, + expected_history_message_count=1, ), ), ) @@ -218,9 +218,9 @@ async def test_send_prompt_async_pre_provider_empty_response_is_not_persisted(mo target_context = TargetNormalizationContext( conversation_id=conversation_id, normalizers=( - FirstTurnHistoryNormalizer( + HistorySquashNormalizer( message_normalizer=ConversationContextNormalizer(), - prepended_message_count=1, + expected_history_message_count=1, ), ), ) diff --git a/tests/unit/prompt_target/target/test_normalize_async_integration.py b/tests/unit/prompt_target/target/test_normalize_async_integration.py index d8e4feab9d..efdeb561df 100644 --- a/tests/unit/prompt_target/target/test_normalize_async_integration.py +++ b/tests/unit/prompt_target/target/test_normalize_async_integration.py @@ -19,7 +19,7 @@ from pyrit.memory.memory_interface import MemoryInterface from pyrit.message_normalizer import ( ConversationContextNormalizer, - FirstTurnHistoryNormalizer, + HistorySquashNormalizer, MessageStringNormalizer, TokenizerTemplateNormalizer, ) @@ -59,9 +59,9 @@ def _make_target_normalization_context( return TargetNormalizationContext( conversation_id=conversation_id, normalizers=( - FirstTurnHistoryNormalizer( + HistorySquashNormalizer( message_normalizer=formatter or ConversationContextNormalizer(), - prepended_message_count=prepended_message_count, + expected_history_message_count=prepended_message_count, ), ), ) From a7318a2604d349e55917a3d647ecf5c193d1c0d4 Mon Sep 17 00:00:00 2001 From: Copilot <223556219+Copilot@users.noreply.github.com> Date: Thu, 20 Aug 2026 16:45:01 -0700 Subject: [PATCH 15/24] Update converter scoping documentation Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> Copilot-Session: 6776aa95-0924-4967-a056-e5a228af3faa --- .../executor/3_attack_configuration.ipynb | 24 ++++++++- doc/code/executor/3_attack_configuration.py | 22 +++++++- doc/code/scenarios/0_scenarios.ipynb | 54 +++++++++++++------ doc/code/scenarios/0_scenarios.py | 9 ++-- doc/code/targets/0_prompt_targets.md | 28 ++++++++-- doc/scanner/1_pyrit_scan.ipynb | 49 +++++++---------- doc/scanner/airt.ipynb | 16 +++--- doc/scanner/airt.py | 16 +++--- 8 files changed, 147 insertions(+), 71 deletions(-) diff --git a/doc/code/executor/3_attack_configuration.ipynb b/doc/code/executor/3_attack_configuration.ipynb index b5d1df1477..0602296fdf 100644 --- a/doc/code/executor/3_attack_configuration.ipynb +++ b/doc/code/executor/3_attack_configuration.ipynb @@ -462,7 +462,7 @@ "Beyond the call arguments, attacks are tuned at construction time with three configuration objects:\n", "\n", "- **`AttackConverterConfig`** — request/response [converters](../converters/0_converters.ipynb)\n", - " applied to every prompt and response.\n", + " applied to live attack prompts and responses, plus selected roles in prepended history.\n", "- **`AttackScoringConfig`** — the objective scorer plus any auxiliary\n", " [scorers](../scoring/0_scoring.ipynb).\n", "- **`AttackAdversarialConfig`** — the adversarial target (a model PyRIT controls) that multi-turn\n", @@ -470,7 +470,27 @@ "\n", "Converter and scoring configs apply to single- and multi-turn attacks alike; the adversarial config\n", "only applies to attacks that drive a conversation. Below builds a converter config — it's just a\n", - "plain object you hand to the attack constructor." + "plain object you hand to the attack constructor.\n", + "\n", + "Request converters apply to prepended `user` messages by default. Prepended `assistant` messages\n", + "represent simulated target output, so PyRIT leaves them unchanged unless the attack explicitly\n", + "opts in. For example:\n", + "\n", + "```python\n", + "from pyrit.executor.attack import PrependedConversationConfig\n", + "\n", + "attack = PromptSendingAttack(\n", + " objective_target=target,\n", + " attack_converter_config=converter_config,\n", + " prepended_conversation_config=PrependedConversationConfig(\n", + " apply_converters_to_roles=[\"user\", \"assistant\"],\n", + " ),\n", + ")\n", + "```\n", + "\n", + "PyRIT applies these role-specific conversions while the prepended messages are still structured.\n", + "If the target cannot accept editable history, target normalization then formats the converted and\n", + "unconverted history with the first live request without broadening the selected converter scope." ] }, { diff --git a/doc/code/executor/3_attack_configuration.py b/doc/code/executor/3_attack_configuration.py index d7597a2e27..74e2cd6bdd 100644 --- a/doc/code/executor/3_attack_configuration.py +++ b/doc/code/executor/3_attack_configuration.py @@ -166,7 +166,7 @@ # Beyond the call arguments, attacks are tuned at construction time with three configuration objects: # # - **`AttackConverterConfig`** — request/response [converters](../converters/0_converters.ipynb) -# applied to every prompt and response. +# applied to live attack prompts and responses, plus selected roles in prepended history. # - **`AttackScoringConfig`** — the objective scorer plus any auxiliary # [scorers](../scoring/0_scoring.ipynb). # - **`AttackAdversarialConfig`** — the adversarial target (a model PyRIT controls) that multi-turn @@ -175,6 +175,26 @@ # Converter and scoring configs apply to single- and multi-turn attacks alike; the adversarial config # only applies to attacks that drive a conversation. Below builds a converter config — it's just a # plain object you hand to the attack constructor. +# +# Request converters apply to prepended `user` messages by default. Prepended `assistant` messages +# represent simulated target output, so PyRIT leaves them unchanged unless the attack explicitly +# opts in. For example: +# +# ```python +# from pyrit.executor.attack import PrependedConversationConfig +# +# attack = PromptSendingAttack( +# objective_target=target, +# attack_converter_config=converter_config, +# prepended_conversation_config=PrependedConversationConfig( +# apply_converters_to_roles=["user", "assistant"], +# ), +# ) +# ``` +# +# PyRIT applies these role-specific conversions while the prepended messages are still structured. +# If the target cannot accept editable history, target normalization then formats the converted and +# unconverted history with the first live request without broadening the selected converter scope. # %% from pyrit.converter import Base64Converter diff --git a/doc/code/scenarios/0_scenarios.ipynb b/doc/code/scenarios/0_scenarios.ipynb index c96b9b1abd..9ed5f78114 100644 --- a/doc/code/scenarios/0_scenarios.ipynb +++ b/doc/code/scenarios/0_scenarios.ipynb @@ -83,7 +83,7 @@ " - `max_retries`: Number of retry attempts on failure (default: 0)\n", " - `memory_labels`: Optional labels for tracking (optional)\n", " - `include_baseline`: Whether to prepend a baseline attack (defaults to the scenario type's\n", - " `BASELINE_ATTACK_POLICY`; most scenarios default it on, `Jailbreak` defaults it off)\n", + " `BASELINE_ATTACK_POLICY`; most scenarios, including `Jailbreak`, default it on)\n", "\n", "### Example Structure\n", "\n", @@ -241,17 +241,42 @@ " airt.jailbreak\u001b[0m\n", " Class: Jailbreak\n", " Description:\n", - " Jailbreak scenario implementation for PyRIT. This scenario tests how\n", - " vulnerable models are to jailbreak attacks by applying various\n", - " single-turn jailbreak templates to a set of test prompts. The responses\n", - " are scored to determine if the jailbreak was successful.\n", + " Jailbreak scenario implementation for PyRIT. Tests how vulnerable a\n", + " model is to jailbreak templates. A run is the cross-product of three\n", + " selectors: - **dataset** — the harmful objectives (HarmBench). -\n", + " **techniques** — two delivery methods for each jailbreak:\n", + " ``prompt_sending`` (the template rendered inline into the user message)\n", + " and ``jailbreak_system_prompt`` (the template set as the system prompt\n", + " with the objective sent as the user turn). - **jailbreaks** — which\n", + " jailbreak templates to run (a random ``num_jailbreaks`` sample or an\n", + " explicit ``jailbreak_names`` set). ``prompt_sending`` applies each\n", + " template as a ``TextJailbreakConverter`` on the outgoing request, so the\n", + " objective is rendered inline into the template's ``{{prompt}}`` slot.\n", + " ``jailbreak_system_prompt`` instead sets the template as a native system\n", + " prompt and sends the objective as its own user turn, so it is only built\n", + " for targets that natively support editable history and system prompts\n", + " (it is skipped for incapable targets, or raises if it is the only\n", + " selected technique). Responses are scored to determine whether the\n", + " jailbreak succeeded (non-refusal).\n", " Aggregate Techniques:\n", - " - all, simple, complex\n", - " Available Techniques (4):\n", - " prompt_sending, many_shot, skeleton, role_play\n", - " Default Technique: simple\n", - " Default Datasets (1, max 4 per dataset):\n", - " airt_harms\n", + " - all, default, single_turn\n", + " Available Techniques (2):\n", + " prompt_sending, jailbreak_system_prompt\n", + " Default Technique: default\n", + " Default Datasets (1):\n", + " harmbench\n", + " Supported Parameters:\n", + " - objective_target (any): Target system under attack: a registered target name or a PromptTarget instance.\n", + " - scenario_techniques (any): Techniques to execute; defaults to the scenario's default aggregate when omitted.\n", + " - technique_converters (any): Mapping of concrete technique name to extra request converters to append.\n", + " - dataset_config (any): Dataset source configuration; defaults to the scenario's default when omitted.\n", + " - memory_labels (any): Additional labels applied to every attack run in the scenario.\n", + " - max_concurrency (int) [default: 4]: Maximum number of concurrent units of work for the scenario.\n", + " - max_retries (int) [default: 0]: Maximum number of automatic retries if the scenario raises an exception.\n", + " - include_baseline (bool): Whether to prepend a baseline atomic attack; None defers to BASELINE_ATTACK_POLICY.\n", + " - num_jailbreaks (int): Draw this many random jailbreak templates for the run. Mutually exclusive with jailbreak_names.\n", + " - num_jailbreak_attempts (int) [default: 1]: Number of times to try each (technique x jailbreak template x objective).\n", + " - jailbreak_names (list[str]): Explicit jailbreak template file names to run (e.g. aim.yaml dan_11.yaml). When omitted, a random sample is drawn. Mutually exclusive with num_jailbreaks.\n", "\u001b[1m\u001b[36m\n", " airt.leakage\u001b[0m\n", " Class: Leakage\n", @@ -439,8 +464,8 @@ "each objective directly to the target without any converters or multi-turn techniques. This is\n", "controlled by the `include_baseline` scenario parameter, supplied through the CLI, config, or\n", "`set_params_from_args` before `initialize_async`; when omitted, each scenario falls back to its\n", - "own `BASELINE_ATTACK_POLICY` class attribute (most scenarios default it on; `Jailbreak` defaults\n", - "it off). See\n", + "own `BASELINE_ATTACK_POLICY` class attribute (most scenarios, including `Jailbreak`, default it\n", + "on). See\n", "[Common Scenario Parameters](./1_common_scenario_parameters.ipynb) for a worked example.\n", "\n", "Custom scenarios should choose their `BASELINE_ATTACK_POLICY` based on whether an unmodified\n", @@ -449,8 +474,7 @@ "- **`Enabled`** — the baseline is prepended by default and the caller can opt out. Use when an\n", " unmodified-prompt run is a meaningful comparison point (most scenarios).\n", "- **`Disabled`** — the baseline is supported but omitted by default; the caller must opt in. Use\n", - " when the scenario is already dominated by a large set of templates/techniques that already\n", - " exercise the unmodified surface (e.g., `Jailbreak`).\n", + " when an unmodified-prompt comparison is valid but not useful enough to run by default.\n", "- **`Forbidden`** — the baseline is unavailable and passing `include_baseline=True` raises. Use\n", " when the scenario's semantics make a single-shot unmodified prompt meaningless as a comparator\n", " (e.g., benchmarks comparing across adversarial models, or multi-turn-only scenarios)." diff --git a/doc/code/scenarios/0_scenarios.py b/doc/code/scenarios/0_scenarios.py index 459fb6754f..aa01af4930 100644 --- a/doc/code/scenarios/0_scenarios.py +++ b/doc/code/scenarios/0_scenarios.py @@ -85,7 +85,7 @@ # - `max_retries`: Number of retry attempts on failure (default: 0) # - `memory_labels`: Optional labels for tracking (optional) # - `include_baseline`: Whether to prepend a baseline attack (defaults to the scenario type's -# `BASELINE_ATTACK_POLICY`; most scenarios default it on, `Jailbreak` defaults it off) +# `BASELINE_ATTACK_POLICY`; most scenarios, including `Jailbreak`, default it on) # # ### Example Structure # @@ -180,8 +180,8 @@ async def _build_atomic_attacks_async(self, *, context): # each objective directly to the target without any converters or multi-turn techniques. This is # controlled by the `include_baseline` scenario parameter, supplied through the CLI, config, or # `set_params_from_args` before `initialize_async`; when omitted, each scenario falls back to its -# own `BASELINE_ATTACK_POLICY` class attribute (most scenarios default it on; `Jailbreak` defaults -# it off). See +# own `BASELINE_ATTACK_POLICY` class attribute (most scenarios, including `Jailbreak`, default it +# on). See # [Common Scenario Parameters](./1_common_scenario_parameters.ipynb) for a worked example. # # Custom scenarios should choose their `BASELINE_ATTACK_POLICY` based on whether an unmodified @@ -190,8 +190,7 @@ async def _build_atomic_attacks_async(self, *, context): # - **`Enabled`** — the baseline is prepended by default and the caller can opt out. Use when an # unmodified-prompt run is a meaningful comparison point (most scenarios). # - **`Disabled`** — the baseline is supported but omitted by default; the caller must opt in. Use -# when the scenario is already dominated by a large set of templates/techniques that already -# exercise the unmodified surface (e.g., `Jailbreak`). +# when an unmodified-prompt comparison is valid but not useful enough to run by default. # - **`Forbidden`** — the baseline is unavailable and passing `include_baseline=True` raises. Use # when the scenario's semantics make a single-shot unmodified prompt meaningless as a comparator # (e.g., benchmarks comparing across adversarial models, or multi-turn-only scenarios). diff --git a/doc/code/targets/0_prompt_targets.md b/doc/code/targets/0_prompt_targets.md index 3a7943e45d..98cd8c0a9c 100644 --- a/doc/code/targets/0_prompt_targets.md +++ b/doc/code/targets/0_prompt_targets.md @@ -9,15 +9,33 @@ Prompt Targets are endpoints for where to send prompts. For example, a target co Prompt targets are found [here](https://github.com/microsoft/PyRIT/tree/main/pyrit/prompt_target/) in code. -## Send_Prompt_Async +## `send_prompt_async` -The main entry method follow the following signature: +The main entry method has the following signature: +```python +async def send_prompt_async( + self, + *, + message: Message, + target_normalization_context: TargetNormalizationContext | None = None, +) -> list[Message]: ``` -async def send_prompt_async(self, *, message: Message) -> Message: -``` -A `Message` object is a normalized object with all the information a target will need to send a prompt, including a way to get a history for that prompt (in the cases that also needs to be sent). This is discussed in more depth [here](../memory/3_memory_data_types.md). +A `Message` object contains the current request and the identifiers needed to load its conversation +history. This is discussed in more depth [here](../memory/3_memory_data_types.md). + +`target_normalization_context` is an internal, per-conversation handoff used by attacks that prepend +structured history to a target without editable history. Ordinary callers do not construct it. +Before the first provider send, `PromptTarget` loads memory history, applies the context's one-shot +normalizers, and then runs the target's ordinary capability-normalization pipeline. Request +converters have already run by this point, so role-specific converter choices remain intact even +when the target must receive one flattened request. The context is ephemeral and does not replace +the structured messages stored in memory. + +`send_prompt_async` is the final public orchestration method. Custom target subclasses implement +`_send_prompt_to_target_async(*, normalized_conversation: list[Message]) -> list[Message]` instead of +overriding `send_prompt_async`. ## Chat-style targets vs general targets diff --git a/doc/scanner/1_pyrit_scan.ipynb b/doc/scanner/1_pyrit_scan.ipynb index 5edd3c5c51..711975eb20 100644 --- a/doc/scanner/1_pyrit_scan.ipynb +++ b/doc/scanner/1_pyrit_scan.ipynb @@ -220,33 +220,24 @@ " Jailbreak scenario implementation for PyRIT. Tests how vulnerable a\n", " model is to jailbreak templates. A run is the cross-product of three\n", " selectors: - **dataset** — the harmful objectives (HarmBench). -\n", - " **techniques** — the *attack techniques* each jailbreak is delivered\n", - " through. Two deliveries are on by default: ``prompt_sending`` (the\n", - " template rendered inline into the user message) and\n", - " ``jailbreak_system_prompt`` (the template set as the system prompt with\n", - " the objective sent as the user turn). The registry techniques\n", - " (``role_play_*``, ``many_shot``, ``tap``, …) are opt-in. -\n", - " **jailbreaks** — which jailbreak templates to run (a random\n", - " ``num_jailbreaks`` sample or an explicit ``jailbreak_names`` set).\n", - " ``prompt_sending`` applies each template as a ``TextJailbreakConverter``\n", - " on the outgoing request, so the objective is rendered inline into the\n", - " template's ``{{prompt}}`` slot; this keeps that delivery target-agnostic\n", - " and lets it compose with every technique. ``jailbreak_system_prompt``\n", - " instead sets the template as a native system prompt and sends the\n", - " objective as its own user turn, so it is only built for targets that\n", - " natively support editable history and system prompts (it is skipped for\n", - " incapable targets, or raises if it is the only selected technique).\n", - " Responses are scored to determine whether the jailbreak succeeded\n", - " (non-refusal).\n", + " **techniques** — two delivery methods for each jailbreak:\n", + " ``prompt_sending`` (the template rendered inline into the user message)\n", + " and ``jailbreak_system_prompt`` (the template set as the system prompt\n", + " with the objective sent as the user turn). - **jailbreaks** — which\n", + " jailbreak templates to run (a random ``num_jailbreaks`` sample or an\n", + " explicit ``jailbreak_names`` set). ``prompt_sending`` applies each\n", + " template as a ``TextJailbreakConverter`` on the outgoing request, so the\n", + " objective is rendered inline into the template's ``{{prompt}}`` slot.\n", + " ``jailbreak_system_prompt`` instead sets the template as a native system\n", + " prompt and sends the objective as its own user turn, so it is only built\n", + " for targets that natively support editable history and system prompts\n", + " (it is skipped for incapable targets, or raises if it is the only\n", + " selected technique). Responses are scored to determine whether the\n", + " jailbreak succeeded (non-refusal).\n", " Aggregate Techniques:\n", - " - all, default, core, light, multi_turn, single_turn\n", - " Available Techniques (16):\n", - " context_compliance, crescendo_history_lecture,\n", - " crescendo_journalist_interview, crescendo_movie_director,\n", - " crescendo_simulated, flip, many_shot, red_teaming,\n", - " role_play_movie_script, role_play_persuasion,\n", - " role_play_persuasion_written, role_play_trivia_game,\n", - " role_play_video_game, tap, prompt_sending, jailbreak_system_prompt\n", + " - all, default, single_turn\n", + " Available Techniques (2):\n", + " prompt_sending, jailbreak_system_prompt\n", " Default Technique: default\n", " Default Datasets (1):\n", " harmbench\n", @@ -256,11 +247,11 @@ " - technique_converters (any): Mapping of concrete technique name to extra request converters to append.\n", " - dataset_config (any): Dataset source configuration; defaults to the scenario's default when omitted.\n", " - memory_labels (any): Additional labels applied to every attack run in the scenario.\n", - " - max_concurrency (int) [default: '4']: Maximum number of concurrent units of work for the scenario.\n", - " - max_retries (int) [default: '0']: Maximum number of automatic retries if the scenario raises an exception.\n", + " - max_concurrency (int) [default: 4]: Maximum number of concurrent units of work for the scenario.\n", + " - max_retries (int) [default: 0]: Maximum number of automatic retries if the scenario raises an exception.\n", " - include_baseline (bool): Whether to prepend a baseline atomic attack; None defers to BASELINE_ATTACK_POLICY.\n", " - num_jailbreaks (int): Draw this many random jailbreak templates for the run. Mutually exclusive with jailbreak_names.\n", - " - num_jailbreak_attempts (int) [default: '1']: Number of times to try each (technique x jailbreak template x objective).\n", + " - num_jailbreak_attempts (int) [default: 1]: Number of times to try each (technique x jailbreak template x objective).\n", " - jailbreak_names (list[str]): Explicit jailbreak template file names to run (e.g. aim.yaml dan_11.yaml). When omitted, a random sample is drawn. Mutually exclusive with num_jailbreaks.\n", "\u001b[1m\u001b[36m\n", " airt.leakage\u001b[0m\n", diff --git a/doc/scanner/airt.ipynb b/doc/scanner/airt.ipynb index 6c5a75a604..74c871973a 100644 --- a/doc/scanner/airt.ipynb +++ b/doc/scanner/airt.ipynb @@ -486,9 +486,11 @@ "objective inline into the template as a request converter (target-agnostic), and\n", "`jailbreak_system_prompt` sets the template as a native system prompt with the objective sent as\n", "the user turn (only for targets that natively support editable history + system prompts — it is\n", - "skipped for incapable targets). Registry techniques like `role_play_*`, `many_shot`, and `tap` are\n", - "opt-in. Results are grouped by jailbreak template, and a baseline (the un-jailbroken objective) is\n", - "included by default so complying with the bare objective is itself visible.\n", + "skipped for incapable targets). These are the only delivery techniques exposed by Jailbreak.\n", + "Generic simulated, multi-turn, or non-composable registry techniques are intentionally excluded\n", + "because they cannot preserve Jailbreak's per-template delivery semantics. Results are grouped by\n", + "jailbreak template, and a baseline (the un-jailbroken objective) is included by default so\n", + "complying with the bare objective is itself visible.\n", "\n", "```bash\n", "pyrit_scan run airt.jailbreak \\\n", @@ -498,10 +500,10 @@ " --max-dataset-size 1\n", "```\n", "\n", - "**Available techniques:** ALL, DEFAULT (`prompt_sending` + `jailbreak_system_prompt`), plus registry\n", - "techniques (`role_play_*`, `many_shot`, `tap`, …). By default a small random sample of jailbreak\n", - "templates runs; pass `num_jailbreaks` (random count) or `jailbreak_names` (explicit) to widen or\n", - "pin the selection." + "**Available technique selectors:** ALL, DEFAULT, and SINGLE_TURN currently select both\n", + "`prompt_sending` and `jailbreak_system_prompt`; either concrete technique can also be selected\n", + "directly. By default a small random sample of jailbreak templates runs; pass `num_jailbreaks`\n", + "(random count) or `jailbreak_names` (explicit) to widen or pin the selection." ] }, { diff --git a/doc/scanner/airt.py b/doc/scanner/airt.py index cc583894d9..1e9366d0c4 100644 --- a/doc/scanner/airt.py +++ b/doc/scanner/airt.py @@ -177,9 +177,11 @@ # objective inline into the template as a request converter (target-agnostic), and # `jailbreak_system_prompt` sets the template as a native system prompt with the objective sent as # the user turn (only for targets that natively support editable history + system prompts — it is -# skipped for incapable targets). Registry techniques like `role_play_*`, `many_shot`, and `tap` are -# opt-in. Results are grouped by jailbreak template, and a baseline (the un-jailbroken objective) is -# included by default so complying with the bare objective is itself visible. +# skipped for incapable targets). These are the only delivery techniques exposed by Jailbreak. +# Generic simulated, multi-turn, or non-composable registry techniques are intentionally excluded +# because they cannot preserve Jailbreak's per-template delivery semantics. Results are grouped by +# jailbreak template, and a baseline (the un-jailbroken objective) is included by default so +# complying with the bare objective is itself visible. # # ```bash # pyrit_scan run airt.jailbreak \ @@ -189,10 +191,10 @@ # --max-dataset-size 1 # ``` # -# **Available techniques:** ALL, DEFAULT (`prompt_sending` + `jailbreak_system_prompt`), plus registry -# techniques (`role_play_*`, `many_shot`, `tap`, …). By default a small random sample of jailbreak -# templates runs; pass `num_jailbreaks` (random count) or `jailbreak_names` (explicit) to widen or -# pin the selection. +# **Available technique selectors:** ALL, DEFAULT, and SINGLE_TURN currently select both +# `prompt_sending` and `jailbreak_system_prompt`; either concrete technique can also be selected +# directly. By default a small random sample of jailbreak templates runs; pass `num_jailbreaks` +# (random count) or `jailbreak_names` (explicit) to widen or pin the selection. # %% from pyrit.scenario.airt import Jailbreak, JailbreakTechnique From 45a388d5e0856e0dda29e70275f8e8984b9e2220 Mon Sep 17 00:00:00 2001 From: Richard Lundeen Date: Thu, 20 Aug 2026 18:22:22 -0700 Subject: [PATCH 16/24] FIX Preserve target formatting across conversation rotation Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> Copilot-Session: 420eef57-7a1c-4dde-8aad-a93539e8da64 --- .../multi_turn/multi_turn_attack_strategy.py | 5 + .../executor/attack/multi_turn/red_teaming.py | 11 +- .../attack/multi_turn/tree_of_attacks.py | 14 +++ .../attack/multi_turn/test_red_teaming.py | 86 ++++++++++++++ .../test_supports_multi_turn_attacks.py | 106 ++++++++++++++++-- 5 files changed, 213 insertions(+), 9 deletions(-) diff --git a/pyrit/executor/attack/multi_turn/multi_turn_attack_strategy.py b/pyrit/executor/attack/multi_turn/multi_turn_attack_strategy.py index 59d7b2df08..b61d1ec340 100644 --- a/pyrit/executor/attack/multi_turn/multi_turn_attack_strategy.py +++ b/pyrit/executor/attack/multi_turn/multi_turn_attack_strategy.py @@ -22,6 +22,7 @@ from pyrit.prompt_target import CapabilityName if TYPE_CHECKING: + from pyrit.executor.attack.component import PrependedConversationConfig from pyrit.models import ( Message, Score, @@ -100,6 +101,7 @@ def _rotate_conversation_for_single_turn_target( self, *, context: MultiTurnAttackContext[Any], + prepended_conversation_config: PrependedConversationConfig | None = None, ) -> None: """ Create a fresh conversation_id for the objective target if it is a single-turn target. @@ -119,6 +121,8 @@ def _rotate_conversation_for_single_turn_target( Args: context: The current attack context. + prepended_conversation_config: Configuration that controls target-facing + formatting for the carried system messages. """ if self._objective_target.configuration.includes(capability=CapabilityName.MULTI_TURN): return @@ -157,6 +161,7 @@ def _rotate_conversation_for_single_turn_target( target=self._objective_target, conversation_id=new_conversation_id, prepended_message_count=len(system_messages), + prepended_conversation_config=prepended_conversation_config, ) else: context.session.conversation_id = str(uuid.uuid4()) diff --git a/pyrit/executor/attack/multi_turn/red_teaming.py b/pyrit/executor/attack/multi_turn/red_teaming.py index 6dfaefaa9e..b0ec823ee2 100644 --- a/pyrit/executor/attack/multi_turn/red_teaming.py +++ b/pyrit/executor/attack/multi_turn/red_teaming.py @@ -14,6 +14,7 @@ from pyrit.exceptions import ComponentRole, execution_context from pyrit.executor.attack.component import ( ConversationManager, + PrependedConversationConfig, _AdversarialConversationManager, get_adversarial_chat_messages, ) @@ -95,6 +96,7 @@ def __init__( attack_converter_config: AttackConverterConfig | None = None, attack_scoring_config: AttackScoringConfig | None = None, prompt_normalizer: PromptNormalizer | None = None, + prepended_conversation_config: PrependedConversationConfig | None = None, max_turns: int = 10, score_last_turn_only: bool = False, ) -> None: @@ -107,6 +109,8 @@ def __init__( attack_converter_config: Configuration for attack converters. Defaults to None. attack_scoring_config: Configuration for attack scoring. Defaults to None. prompt_normalizer: The prompt normalizer to use for sending prompts. Defaults to None. + prepended_conversation_config: Configuration for prepended-conversation + converter roles and target-facing formatting. Defaults to None. max_turns (int): Maximum number of turns for the attack. Defaults to 10. score_last_turn_only (bool): If True, only score the final turn instead of every turn. This reduces LLM calls when intermediate scores are not needed (e.g., for @@ -170,6 +174,7 @@ def __init__( # Initialize utilities self._prompt_normalizer = prompt_normalizer or PromptNormalizer() + self._prepended_conversation_config = prepended_conversation_config self._conversation_manager = ConversationManager(prompt_normalizer=self._prompt_normalizer) @@ -270,6 +275,7 @@ async def _setup_async(self, *, context: MultiTurnAttackContext[Any]) -> None: target=self._objective_target, conversation_id=context.session.conversation_id, request_converters=self._request_converters, + prepended_conversation_config=self._prepended_conversation_config, max_turns=self._max_turns, memory_labels=self._memory_labels, ) @@ -471,7 +477,10 @@ async def _send_prompt_to_objective_target_async( logger.info(f"Sending prompt to target: {message.get_value()[:50]}...") # For single-turn targets, rotate conversation_id so each turn starts fresh - self._rotate_conversation_for_single_turn_target(context=context) + self._rotate_conversation_for_single_turn_target( + context=context, + prepended_conversation_config=self._prepended_conversation_config, + ) with execution_context( component_role=ComponentRole.OBJECTIVE_TARGET, diff --git a/pyrit/executor/attack/multi_turn/tree_of_attacks.py b/pyrit/executor/attack/multi_turn/tree_of_attacks.py index 9cc9bc6d67..9d5d90f916 100644 --- a/pyrit/executor/attack/multi_turn/tree_of_attacks.py +++ b/pyrit/executor/attack/multi_turn/tree_of_attacks.py @@ -341,6 +341,7 @@ def __init__( parent_id: str | None = None, prompt_normalizer: PromptNormalizer | None = None, initial_prompt: Message | None = None, + prepended_conversation_config: PrependedConversationConfig | None = None, ) -> None: """ Initialize a tree node. @@ -368,6 +369,9 @@ def __init__( prompt_normalizer (PromptNormalizer | None): Normalizer for handling prompts and responses. initial_prompt (Message | None): Initial message to send for the first turn, bypassing adversarial chat generation. Supports multimodal messages. + prepended_conversation_config (PrependedConversationConfig | None): + Configuration for prepended-conversation converter roles and + target-facing formatting. """ # Store configuration self._objective_target = objective_target @@ -385,6 +389,7 @@ def __init__( self._attack_strategy_name = attack_strategy_name self._memory_labels = memory_labels or {} self._modality_router = modality_router + self._prepended_conversation_config = prepended_conversation_config # Initialize utilities self._memory = CentralMemory.get_memory_instance() @@ -453,6 +458,7 @@ async def initialize_with_prepended_conversation_async( """ if not prepended_conversation: return + self._prepended_conversation_config = prepended_conversation_config # Use ConversationManager to add messages to memory conversation_manager = ConversationManager( @@ -899,6 +905,7 @@ def duplicate(self) -> _TreeOfAttacksNode: desired_response_prefix=self._desired_response_prefix, parent_id=self.node_id, prompt_normalizer=self._prompt_normalizer, + prepended_conversation_config=self._prepended_conversation_config, ) # Duplicate the conversations to preserve history @@ -921,6 +928,12 @@ def duplicate(self) -> _TreeOfAttacksNode: ) self._memory.add_message_pieces_to_memory(message_pieces=pieces) duplicate_node.objective_target_conversation_id = new_id + duplicate_node._target_normalization_context = ConversationManager.create_target_normalization_context( + target=self._objective_target, + conversation_id=new_id, + prepended_message_count=len(system_messages), + prepended_conversation_config=self._prepended_conversation_config, + ) else: duplicate_node.objective_target_conversation_id = str(uuid.uuid4()) @@ -2059,6 +2072,7 @@ def _create_attack_node( parent_id=parent_id, prompt_normalizer=self._prompt_normalizer, initial_prompt=initial_prompt, + prepended_conversation_config=self._prepended_conversation_config, ) # Add the adversarial chat conversation ID to the context's tracking (ensuring uniqueness) 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 51b0f326b7..edc62087eb 100644 --- a/tests/unit/executor/attack/multi_turn/test_red_teaming.py +++ b/tests/unit/executor/attack/multi_turn/test_red_teaming.py @@ -19,7 +19,10 @@ RedTeamingAttack, RTASystemPromptPaths, ) +from pyrit.executor.attack.component import PrependedConversationConfig from pyrit.executor.attack.core.attack_config import DEFAULT_ADVERSARIAL_FIRST_MESSAGE +from pyrit.memory import CentralMemory +from pyrit.message_normalizer import MessageStringNormalizer from pyrit.models import ( AttackOutcome, AttackResult, @@ -33,7 +36,10 @@ ) from pyrit.prompt_normalizer import PromptNormalizer from pyrit.prompt_target import PromptTarget +from pyrit.prompt_target.common.target_capabilities import TargetCapabilities +from pyrit.prompt_target.common.target_configuration import TargetConfiguration from pyrit.score import Scorer, TrueFalseScorer +from tests.unit.mocks import MockPromptTarget def _adversarial_reply_message(next_message: str = "Adversarial next message") -> Message: @@ -645,6 +651,32 @@ async def test_setup_initializes_conversation_session( assert basic_context.session is not None assert isinstance(basic_context.session, ConversationSession) + async def test_setup_forwards_prepended_conversation_config( + self, + mock_objective_target: MagicMock, + mock_objective_scorer: MagicMock, + mock_adversarial_chat: MagicMock, + basic_context: MultiTurnAttackContext, + ): + """Setup must use the configured prepended-conversation formatter.""" + prepended_conversation_config = PrependedConversationConfig() + attack = RedTeamingAttack( + objective_target=mock_objective_target, + attack_adversarial_config=AttackAdversarialConfig(target=mock_adversarial_chat), + attack_scoring_config=AttackScoringConfig(objective_scorer=mock_objective_scorer), + prepended_conversation_config=prepended_conversation_config, + ) + + with patch.object( + attack._conversation_manager, + "initialize_context_async", + new_callable=AsyncMock, + return_value=ConversationState(turn_count=0), + ) as mock_initialize: + await attack._setup_async(context=basic_context) + + assert mock_initialize.call_args.kwargs["prepended_conversation_config"] is prepended_conversation_config + async def test_setup_updates_turn_count_from_prepended_conversation( self, mock_objective_target: MagicMock, @@ -917,6 +949,60 @@ async def test_generate_next_prompt_raises_on_none_response( await attack._generate_next_prompt_async(context=basic_context) +@pytest.mark.usefixtures("patch_central_database") +class TestObjectiveTargetSending: + """Tests for sending prompts to the objective target.""" + + async def test_second_turn_rotation_uses_configured_message_normalizer( + self, + mock_objective_scorer: MagicMock, + mock_adversarial_chat: MagicMock, + basic_context: MultiTurnAttackContext, + ) -> None: + """A rotated request must keep the configured prepended-history formatter.""" + objective_target = MockPromptTarget() + objective_target._configuration = TargetConfiguration(capabilities=TargetCapabilities()) + message_normalizer = MagicMock(spec=MessageStringNormalizer) + message_normalizer.normalize_string_async = AsyncMock(return_value="custom formatted request") + attack = RedTeamingAttack( + objective_target=objective_target, + attack_adversarial_config=AttackAdversarialConfig(target=mock_adversarial_chat), + attack_scoring_config=AttackScoringConfig(objective_scorer=mock_objective_scorer), + prepended_conversation_config=PrependedConversationConfig(message_normalizer=message_normalizer), + ) + basic_context.session = ConversationSession() + old_conversation_id = basic_context.session.conversation_id + memory = CentralMemory.get_memory_instance() + memory.add_message_pieces_to_memory( + message_pieces=[ + MessagePiece( + original_value="You are a helpful assistant.", + role="system", + conversation_id=old_conversation_id, + sequence=0, + ), + MessagePiece( + original_value="First request", + role="user", + conversation_id=old_conversation_id, + sequence=1, + ), + ] + ) + basic_context.executed_turns = 1 + + await attack._send_prompt_to_objective_target_async( + context=basic_context, + message=Message.from_prompt(prompt="Second request", role="user"), + ) + + assert basic_context.session.conversation_id != old_conversation_id + assert basic_context.target_normalization_context is not None + assert basic_context.target_normalization_context.is_consumed + assert objective_target.prompt_sent == ["custom formatted request"] + message_normalizer.normalize_string_async.assert_awaited_once() + + @pytest.mark.usefixtures("patch_central_database") class TestResponseScoring: """Tests for response scoring logic.""" diff --git a/tests/unit/executor/attack/multi_turn/test_supports_multi_turn_attacks.py b/tests/unit/executor/attack/multi_turn/test_supports_multi_turn_attacks.py index 3a2d86facb..396c0896c5 100644 --- a/tests/unit/executor/attack/multi_turn/test_supports_multi_turn_attacks.py +++ b/tests/unit/executor/attack/multi_turn/test_supports_multi_turn_attacks.py @@ -1,17 +1,23 @@ # Copyright (c) Microsoft Corporation. # Licensed under the MIT license. +from typing import Any from unittest.mock import AsyncMock, MagicMock, patch import pytest +from pyrit.executor.attack.component import PrependedConversationConfig from pyrit.executor.attack.core.attack_parameters import AttackParameters from pyrit.executor.attack.multi_turn.multi_turn_attack_strategy import ( ConversationSession, MultiTurnAttackContext, ) from pyrit.memory import CentralMemory -from pyrit.message_normalizer import ConversationContextNormalizer, FirstTurnHistoryNormalizer +from pyrit.message_normalizer import ( + ConversationContextNormalizer, + FirstTurnHistoryNormalizer, + MessageStringNormalizer, +) from pyrit.models import ConversationType, Message, MessagePiece from pyrit.prompt_target import PromptTarget, TargetNormalizationContext from pyrit.prompt_target.common.target_capabilities import TargetCapabilities @@ -262,6 +268,49 @@ async def test_rotated_system_prompt_is_normalized_with_next_request(self): "Turn 1:\nuser: ### Instructions ###\n\nYou are a helpful assistant.\n\n######\n\nSecond request" ) + async def test_rotation_preserves_configured_message_normalizer(self): + """Rotation must retain the formatter used for carried system messages.""" + target = _SingleTurnPromptTarget() + strategy = _make_strategy(supports_multi_turn=False) + strategy._objective_target = target + context = _make_context() + old_id = context.session.conversation_id + _seed_conversation( + conversation_id=old_id, + system_prompt="You are a helpful assistant.", + user_text="First request", + ) + previous_context = TargetNormalizationContext( + conversation_id=old_id, + normalizers=( + FirstTurnHistoryNormalizer( + message_normalizer=ConversationContextNormalizer(), + prepended_message_count=1, + ), + ), + ) + previous_context.begin_normalization(conversation_id=old_id) + previous_context.mark_consumed() + context.target_normalization_context = previous_context + context.executed_turns = 1 + message_normalizer = MagicMock(spec=MessageStringNormalizer) + message_normalizer.normalize_string_async = AsyncMock(return_value="custom formatted request") + + strategy._rotate_conversation_for_single_turn_target( + context=context, + prepended_conversation_config=PrependedConversationConfig(message_normalizer=message_normalizer), + ) + + next_request = Message.from_prompt(prompt="Second request", role="user") + next_request.get_piece().conversation_id = context.session.conversation_id + await target.send_prompt_async( + message=next_request, + target_normalization_context=context.target_normalization_context, + ) + + assert target.normalized_conversations[0][0].get_value() == "custom formatted request" + message_normalizer.normalize_string_async.assert_awaited_once() + def test_no_system_prompt_yields_fresh_conversation_id(self): """When there is no system prompt, rotation still generates a new conversation_id.""" strategy = _make_strategy(supports_multi_turn=False) @@ -840,17 +889,24 @@ async def test_chunked_request_raises_for_single_turn_target(self): class TestTAPBranchingPreservesSystemPrompts: """Integration test: TAP branching with real memory verifies system prompt carryover.""" - def _make_tap_node(self, *, supports_multi_turn: bool): + def _make_tap_node( + self, + *, + supports_multi_turn: bool, + prepended_conversation_config: PrependedConversationConfig | None = None, + objective_target: PromptTarget | None = None, + ) -> Any: """Create a _TreeOfAttacksNode with real memory.""" from pyrit.executor.attack.component.modality_router import _ModalityFeedbackRouter from pyrit.executor.attack.multi_turn.tree_of_attacks import _TreeOfAttacksNode - target = MagicMock() - target.capabilities.supports_multi_turn = supports_multi_turn - target.configuration.includes.return_value = supports_multi_turn - target.configuration.capabilities.input_modalities = frozenset({frozenset({"text"})}) - target.configuration.capabilities.output_modalities = frozenset({frozenset({"text"})}) - target.get_identifier.return_value = {"__type__": "MockTarget", "__module__": "test", "id": "mock-id"} + target = objective_target or MagicMock() + if objective_target is None: + target.capabilities.supports_multi_turn = supports_multi_turn + target.configuration.includes.return_value = supports_multi_turn + target.configuration.capabilities.input_modalities = frozenset({frozenset({"text"})}) + target.configuration.capabilities.output_modalities = frozenset({frozenset({"text"})}) + target.get_identifier.return_value = {"__type__": "MockTarget", "__module__": "test", "id": "mock-id"} adversarial_chat = MagicMock() adversarial_chat.get_identifier.return_value = {"__type__": "MockTarget", "__module__": "test", "id": "mock-id"} @@ -881,6 +937,7 @@ def _make_tap_node(self, *, supports_multi_turn: bool): adversarial_chat=adversarial_chat, objective_target=target, ), + prepended_conversation_config=prepended_conversation_config, ) def test_branching_single_turn_target_preserves_system_across_depths(self): @@ -951,6 +1008,39 @@ def test_branching_single_turn_target_preserves_system_across_depths(self): assert branch2_msgs[0].api_role == "system" assert branch2_msgs[0].get_value() == "You are a red team assistant." + async def test_branching_single_turn_target_retains_pending_normalization_context(self): + """A TAP child must send copied system context with the next request.""" + message_normalizer = MagicMock(spec=MessageStringNormalizer) + message_normalizer.normalize_string_async = AsyncMock(return_value="custom formatted request") + target = _SingleTurnPromptTarget() + node = self._make_tap_node( + supports_multi_turn=False, + prepended_conversation_config=PrependedConversationConfig(message_normalizer=message_normalizer), + objective_target=target, + ) + memory = CentralMemory.get_memory_instance() + memory.add_message_pieces_to_memory( + message_pieces=[ + MessagePiece( + original_value="You are a red team assistant.", + role="system", + conversation_id=node.objective_target_conversation_id, + sequence=0, + ) + ] + ) + + branch = node.duplicate() + branch_conversation_id = branch.objective_target_conversation_id + + await branch._send_prompt_to_target_async("next request") + + assert branch.objective_target_conversation_id == branch_conversation_id + assert branch._target_normalization_context is not None + assert branch._target_normalization_context.is_consumed + assert target.normalized_conversations[0][0].get_value() == "custom formatted request" + message_normalizer.normalize_string_async.assert_awaited_once() + def test_branching_multi_turn_target_preserves_full_history(self): """For multi-turn targets, branching should preserve the full conversation.""" memory = CentralMemory.get_memory_instance() From 4f02e909945c4c767e907a38f484c00ea868fba3 Mon Sep 17 00:00:00 2001 From: Richard Lundeen Date: Thu, 20 Aug 2026 19:41:45 -0700 Subject: [PATCH 17/24] FIX Make prepended history capability-driven Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> Copilot-Session: 420eef57-7a1c-4dde-8aad-a93539e8da64 --- doc/code/framework.md | 4 +- .../attack/component/conversation_manager.py | 47 +- .../prepended_conversation_config.py | 36 +- pyrit/executor/attack/core/attack_strategy.py | 8 - .../attack/multi_turn/chunked_request.py | 11 +- pyrit/executor/attack/multi_turn/crescendo.py | 10 +- .../attack/multi_turn/multi_prompt_sending.py | 11 +- .../multi_turn/multi_turn_attack_strategy.py | 80 -- .../executor/attack/multi_turn/red_teaming.py | 12 +- .../attack/multi_turn/tree_of_attacks.py | 81 +- .../attack/single_turn/prompt_sending.py | 8 +- pyrit/executor/attack/streaming/barge_in.py | 6 + .../first_turn_history_normalizer.py | 84 +- pyrit/prompt_normalizer/normalizer_request.py | 7 + pyrit/prompt_normalizer/prompt_normalizer.py | 28 +- pyrit/prompt_target/__init__.py | 6 - .../conversation_normalization_pipeline.py | 25 +- pyrit/prompt_target/common/prompt_target.py | 77 +- .../common/target_capabilities.py | 5 +- .../common/target_configuration.py | 17 +- .../common/target_normalization_context.py | 138 -- .../common/target_requirements.py | 6 +- .../component/test_conversation_manager.py | 40 +- .../attack/multi_turn/test_crescendo.py | 10 +- .../attack/multi_turn/test_red_teaming.py | 8 +- .../test_supports_multi_turn_attacks.py | 1174 +++-------------- .../attack/multi_turn/test_tree_of_attacks.py | 9 +- .../attack/single_turn/test_prompt_sending.py | 11 +- .../test_prompt_normalizer.py | 73 +- .../prompt_target/target/test_image_target.py | 12 +- .../test_normalize_async_integration.py | 176 +-- .../test_prompt_target_azure_blob_storage.py | 13 +- .../target/test_target_capabilities.py | 11 +- .../target/test_target_configuration.py | 23 +- .../test_target_normalization_context.py | 60 - .../target/test_target_requirements.py | 6 +- .../prompt_target/target/test_tts_target.py | 14 +- 37 files changed, 557 insertions(+), 1790 deletions(-) delete mode 100644 pyrit/prompt_target/common/target_normalization_context.py delete mode 100644 tests/unit/prompt_target/target/test_target_normalization_context.py diff --git a/doc/code/framework.md b/doc/code/framework.md index 11fdb02af8..be09144b53 100644 --- a/doc/code/framework.md +++ b/doc/code/framework.md @@ -314,10 +314,10 @@ 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). 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. Prepended history remains role-structured in memory. Attacks pass an ephemeral `TargetNormalizationContext` with the first live send when a target cannot edit history; the target consumes that context immediately before provider invocation. +- **`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. Prepended history remains role-structured in memory. Attacks can pass per-send normalizer overrides from `PrependedConversationConfig`; the target's capability pipeline uses the `EDITABLE_HISTORY` override only when adaptation is required. - **`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, prepended-history adaptation, and tokenizer chat templates. These target-specific views are ephemeral and do not replace the logical conversation in memory. -For prepended history on a target without editable history, processing order is: convert and persist the structured prepended messages, convert the live request, apply first-turn history normalization, run the target's ordinary capability normalizers, then serialize and invoke the provider. +For prepended history on a target without editable history, processing order is: convert and persist the structured prepended messages, convert the live request, apply the capability-driven history normalizer, run the remaining target capability normalizers, then serialize and invoke the provider. The history normalizer derives whether a real target reply exists from memory roles; attacks do not store target-normalization lifecycle state. ## [Output](./output/0_output) diff --git a/pyrit/executor/attack/component/conversation_manager.py b/pyrit/executor/attack/component/conversation_manager.py index 89470202f3..657ae43ac3 100644 --- a/pyrit/executor/attack/component/conversation_manager.py +++ b/pyrit/executor/attack/component/conversation_manager.py @@ -13,7 +13,7 @@ PrependedConversationConfig, ) from pyrit.memory import CentralMemory -from pyrit.message_normalizer import ConversationContextNormalizer, FirstTurnHistoryNormalizer +from pyrit.message_normalizer import ConversationContextNormalizer from pyrit.models import ( ChatMessageRole, ComponentIdentifier, @@ -23,7 +23,7 @@ Score, ) from pyrit.prompt_normalizer.prompt_normalizer import PromptNormalizer -from pyrit.prompt_target import CapabilityName, PromptTarget, TargetNormalizationContext +from pyrit.prompt_target import CapabilityName, PromptTarget if TYPE_CHECKING: from collections.abc import Sequence @@ -279,8 +279,7 @@ async def initialize_context_async( 3. Updates context.executed_turns for multi-turn attacks For all PromptTarget types, prepended messages are added to memory with - simulated_assistant roles and new UUIDs. Targets without editable history receive - explicit one-shot normalization state on the attack context. + simulated_assistant roles and new UUIDs. Args: context: The attack context to initialize. @@ -302,8 +301,6 @@ async def initialize_context_async( # Merge memory labels: attack strategy labels + context labels context.memory_labels = combine_dict(existing_dict=memory_labels, new_dict=context.memory_labels) - context.target_normalization_context = None - state = ConversationState() prepended_conversation = context.prepended_conversation @@ -442,35 +439,6 @@ def get_persistable_prepended_messages( if message and message.message_pieces and any(not piece.not_in_memory for piece in message.message_pieces) ] - @staticmethod - def create_target_normalization_context( - *, - target: PromptTarget, - conversation_id: str, - prepended_message_count: int, - prepended_conversation_config: PrependedConversationConfig | None = None, - ) -> TargetNormalizationContext | None: - """ - Build first-send normalization state for a target without editable history. - - Returns: - TargetNormalizationContext | None: First-send state, or ``None`` for - editable-history targets or empty prepended history. - """ - if prepended_message_count < 1 or target.configuration.includes(capability=CapabilityName.EDITABLE_HISTORY): - return None - - config = prepended_conversation_config or PrependedConversationConfig() - return TargetNormalizationContext( - conversation_id=conversation_id, - normalizers=( - FirstTurnHistoryNormalizer( - message_normalizer=config.get_message_normalizer(), - prepended_message_count=prepended_message_count, - ), - ), - ) - async def _process_prepended_conversation_async( self, *, @@ -512,13 +480,6 @@ async def _process_prepended_conversation_async( if not valid_messages: return state - target_normalization_context = self.create_target_normalization_context( - target=target, - conversation_id=conversation_id, - prepended_message_count=len(valid_messages), - prepended_conversation_config=prepended_conversation_config, - ) - # Use the lower-level method to add messages to memory state.turn_count = await self.add_prepended_conversation_to_memory_async( prepended_conversation=prepended_conversation, @@ -529,8 +490,6 @@ async def _process_prepended_conversation_async( target_identifier=target_identifier, target=target, ) - context.target_normalization_context = target_normalization_context - # Update context for multi-turn attacks to reflect prepended_conversation final_prepended_message = valid_messages[-1] diff --git a/pyrit/executor/attack/component/prepended_conversation_config.py b/pyrit/executor/attack/component/prepended_conversation_config.py index e6c2639a18..ce9398eecb 100644 --- a/pyrit/executor/attack/component/prepended_conversation_config.py +++ b/pyrit/executor/attack/component/prepended_conversation_config.py @@ -8,11 +8,15 @@ from pyrit.message_normalizer import ( ConversationContextNormalizer, + FirstTurnHistoryNormalizer, + MessageListNormalizer, MessageStringNormalizer, ) +from pyrit.prompt_target.common.target_capabilities import CapabilityName if TYPE_CHECKING: - from pyrit.models import ChatMessageRole + from pyrit.models import ChatMessageRole, Message + from pyrit.prompt_target.common.prompt_target import PromptTarget @dataclass @@ -23,11 +27,11 @@ class PrependedConversationConfig: This class provides control over: - Which message roles should have request converters applied - - How targets without editable history format prepended messages on the first live send + - How targets without editable history format prepended messages with live requests Prepended messages remain role-structured in memory. Request converters are applied to configured roles before a target without editable history renders that history and the - first live request together (via ``message_normalizer``; default: ConversationContextNormalizer). + applicable live request together (via ``message_normalizer``; default: ConversationContextNormalizer). Those converters must produce text because string normalization cannot preserve converted image, audio, or other non-text output. """ @@ -36,7 +40,7 @@ class PrependedConversationConfig: # simulated target output and must be explicitly opted in with ["assistant"]. apply_converters_to_roles: list[ChatMessageRole] = field(default_factory=lambda: ["user"]) - # Optional normalizer to format prepended history and the first live request as one text block. + # Optional normalizer to format prepended history and a live request as one text block. # Must implement MessageStringNormalizer (e.g., TokenizerTemplateNormalizer or ConversationContextNormalizer). # When None and adaptation is needed, a default ConversationContextNormalizer is used # that produces "Turn N: User/Assistant" format. @@ -51,3 +55,27 @@ def get_message_normalizer(self) -> MessageStringNormalizer: ConversationContextNormalizer if none was configured. """ return self.message_normalizer or ConversationContextNormalizer() + + def get_normalizer_overrides( + self, + *, + target: PromptTarget, + ) -> dict[CapabilityName, MessageListNormalizer[Message]]: + """ + Build per-send target normalizer overrides for prepended history. + + Args: + target: Target that receives the live request. + + Returns: + Overrides keyed by the capability they adapt. + """ + if target.configuration.includes(capability=CapabilityName.EDITABLE_HISTORY): + return {} + + return { + CapabilityName.EDITABLE_HISTORY: FirstTurnHistoryNormalizer( + message_normalizer=self.get_message_normalizer(), + target_supports_multi_turn=target.configuration.includes(capability=CapabilityName.MULTI_TURN), + ) + } diff --git a/pyrit/executor/attack/core/attack_strategy.py b/pyrit/executor/attack/core/attack_strategy.py index f1dc2353bd..11c58b60ad 100644 --- a/pyrit/executor/attack/core/attack_strategy.py +++ b/pyrit/executor/attack/core/attack_strategy.py @@ -48,7 +48,6 @@ ) from pyrit.executor.attack.core.attack_result_attribution import AttackResultAttribution from pyrit.prompt_target import PromptTarget - from pyrit.prompt_target.common.target_normalization_context import TargetNormalizationContext AttackStrategyContextT = TypeVar("AttackStrategyContextT", bound="AttackContext[Any]") AttackStrategyResultT = TypeVar("AttackStrategyResultT", bound="AttackResult") @@ -89,13 +88,6 @@ class AttackContext(StrategyContext, ABC, Generic[AttackParamsT]): _prepended_conversation_override: list[Message] | None = None _memory_labels_override: dict[str, str] | None = None - # Ephemeral first-send target normalization state. This is never persisted. - target_normalization_context: TargetNormalizationContext | None = field( - default=None, - repr=False, - compare=False, - ) - # Optional attribution from an upstream orchestrator (e.g. Scenario). When # set, the persistence path stamps attribution_parent_id + attribution_data # onto the resulting AttackResult so it can be located later for hydration diff --git a/pyrit/executor/attack/multi_turn/chunked_request.py b/pyrit/executor/attack/multi_turn/chunked_request.py index a5f3f064c7..772eb3df41 100644 --- a/pyrit/executor/attack/multi_turn/chunked_request.py +++ b/pyrit/executor/attack/multi_turn/chunked_request.py @@ -9,7 +9,7 @@ from pyrit.common.apply_defaults import REQUIRED_VALUE, apply_defaults from pyrit.exceptions import ComponentRole, execution_context -from pyrit.executor.attack.component import ConversationManager +from pyrit.executor.attack.component import ConversationManager, PrependedConversationConfig from pyrit.executor.attack.core.attack_config import ( AttackConverterConfig, AttackScoringConfig, @@ -105,6 +105,7 @@ def __init__( attack_converter_config: AttackConverterConfig | None = None, attack_scoring_config: AttackScoringConfig | None = None, prompt_normalizer: PromptNormalizer | None = None, + prepended_conversation_config: PrependedConversationConfig | None = None, ) -> None: """ Initialize the chunked request attack strategy. @@ -119,6 +120,8 @@ def __init__( attack_converter_config (AttackConverterConfig | None): Configuration for converters. attack_scoring_config (AttackScoringConfig | None): Configuration for scoring components. prompt_normalizer (PromptNormalizer | None): Normalizer for handling prompts. + prepended_conversation_config: Configuration for prepended-conversation + conversion and target-facing formatting. Raises: ValueError: If chunk_size or total_length are invalid. @@ -171,6 +174,7 @@ def __init__( # Initialize prompt normalizer and conversation manager self._prompt_normalizer = prompt_normalizer or PromptNormalizer() + self._prepended_conversation_config = prepended_conversation_config or PrependedConversationConfig() self._conversation_manager = ConversationManager( prompt_normalizer=self._prompt_normalizer, ) @@ -246,6 +250,7 @@ async def _setup_async(self, *, context: ChunkedRequestAttackContext) -> None: target=self._objective_target, conversation_id=context.session.conversation_id, request_converters=self._request_converters, + prepended_conversation_config=self._prepended_conversation_config, memory_labels=self._memory_labels, ) @@ -288,7 +293,9 @@ async def _perform_async(self, *, context: ChunkedRequestAttackContext) -> Attac conversation_id=context.session.conversation_id, request_converter_configurations=self._request_converters, response_converter_configurations=self._response_converters, - target_normalization_context=context.target_normalization_context, + normalizer_overrides=self._prepended_conversation_config.get_normalizer_overrides( + target=self._objective_target + ), ) # Store the response diff --git a/pyrit/executor/attack/multi_turn/crescendo.py b/pyrit/executor/attack/multi_turn/crescendo.py index 80796a9ca8..ef90fec474 100644 --- a/pyrit/executor/attack/multi_turn/crescendo.py +++ b/pyrit/executor/attack/multi_turn/crescendo.py @@ -181,10 +181,10 @@ def __init__( max_turns (int): Maximum number of turns allowed. prepended_conversation_config (PrependedConversationConfiguration | None): Configuration for how to process prepended conversations. Controls converter - application by role and first-send formatting for targets without editable history. + application by role and request formatting for targets without editable history. Raises: - ValueError: If objective_target does not natively support editable history. + ValueError: If the objective target does not natively support multi-turn conversations. """ # Initialize base class super().__init__(objective_target=objective_target, logger=logger, context_type=CrescendoAttackContext) @@ -272,7 +272,7 @@ def __init__( self._max_turns = max_turns # Store the prepended conversation configuration - self._prepended_conversation_config = prepended_conversation_config + self._prepended_conversation_config = prepended_conversation_config or PrependedConversationConfig() def get_attack_scoring_config(self) -> AttackScoringConfig | None: """ @@ -648,7 +648,9 @@ async def _send_prompt_to_objective_target_async( conversation_id=context.session.conversation_id, request_converter_configurations=self._request_converters, response_converter_configurations=self._response_converters, - target_normalization_context=context.target_normalization_context, + normalizer_overrides=self._prepended_conversation_config.get_normalizer_overrides( + target=self._objective_target + ), ) if not response: diff --git a/pyrit/executor/attack/multi_turn/multi_prompt_sending.py b/pyrit/executor/attack/multi_turn/multi_prompt_sending.py index 24fd562284..4547173e38 100644 --- a/pyrit/executor/attack/multi_turn/multi_prompt_sending.py +++ b/pyrit/executor/attack/multi_turn/multi_prompt_sending.py @@ -10,7 +10,7 @@ from pyrit.common.apply_defaults import REQUIRED_VALUE, apply_defaults from pyrit.common.utils import get_kwarg_param from pyrit.exceptions import ComponentRole, execution_context -from pyrit.executor.attack.component import ConversationManager +from pyrit.executor.attack.component import ConversationManager, PrependedConversationConfig from pyrit.executor.attack.core.attack_config import ( AttackConverterConfig, AttackScoringConfig, @@ -142,6 +142,7 @@ def __init__( attack_converter_config: AttackConverterConfig | None = None, attack_scoring_config: AttackScoringConfig | None = None, prompt_normalizer: PromptNormalizer | None = None, + prepended_conversation_config: PrependedConversationConfig | None = None, ) -> None: """ Initialize the multi-prompt sending attack strategy. @@ -151,6 +152,8 @@ def __init__( attack_converter_config (AttackConverterConfig | None): Configuration for converters. attack_scoring_config (AttackScoringConfig | None): Configuration for scoring components. prompt_normalizer (PromptNormalizer | None): Normalizer for handling prompts. + prepended_conversation_config: Configuration for prepended-conversation + conversion and target-facing formatting. Raises: ValueError: If the objective scorer is not a true/false scorer. @@ -176,6 +179,7 @@ def __init__( # Initialize prompt normalizer and conversation manager self._prompt_normalizer = prompt_normalizer or PromptNormalizer() + self._prepended_conversation_config = prepended_conversation_config or PrependedConversationConfig() self._conversation_manager = ConversationManager( prompt_normalizer=self._prompt_normalizer, ) @@ -224,6 +228,7 @@ async def _setup_async(self, *, context: MultiTurnAttackContext[Any]) -> None: target=self._objective_target, conversation_id=context.session.conversation_id, request_converters=self._request_converters, + prepended_conversation_config=self._prepended_conversation_config, memory_labels=self._memory_labels, ) @@ -366,7 +371,9 @@ async def _send_prompt_to_objective_target_async( conversation_id=context.session.conversation_id, request_converter_configurations=self._request_converters, response_converter_configurations=self._response_converters, - target_normalization_context=context.target_normalization_context, + normalizer_overrides=self._prepended_conversation_config.get_normalizer_overrides( + target=self._objective_target + ), ) async def _evaluate_response_async(self, *, response: Message, objective: str) -> Score | None: diff --git a/pyrit/executor/attack/multi_turn/multi_turn_attack_strategy.py b/pyrit/executor/attack/multi_turn/multi_turn_attack_strategy.py index b61d1ec340..102c702668 100644 --- a/pyrit/executor/attack/multi_turn/multi_turn_attack_strategy.py +++ b/pyrit/executor/attack/multi_turn/multi_turn_attack_strategy.py @@ -10,19 +10,14 @@ from typing import TYPE_CHECKING, Any, TypeVar from pyrit.common.logger import logger -from pyrit.executor.attack.component.conversation_manager import ConversationManager from pyrit.executor.attack.core.attack_parameters import AttackParameters, AttackParamsT from pyrit.executor.attack.core.attack_strategy import ( AttackContext, AttackStrategy, AttackStrategyResultT, ) -from pyrit.memory import CentralMemory -from pyrit.models import Conversation, ConversationReference, ConversationType -from pyrit.prompt_target import CapabilityName if TYPE_CHECKING: - from pyrit.executor.attack.component import PrependedConversationConfig from pyrit.models import ( Message, Score, @@ -96,78 +91,3 @@ def __init__( params_type=params_type, logger=logger, ) - - def _rotate_conversation_for_single_turn_target( - self, - *, - context: MultiTurnAttackContext[Any], - prepended_conversation_config: PrependedConversationConfig | None = None, - ) -> None: - """ - Create a fresh conversation_id for the objective target if it is a single-turn target. - - For single-turn targets, each turn must use a separate conversation_id because the target - rejects conversations with prior messages. The prior turn's conversation_id is recorded - as a PRUNED related conversation on the attack context. - - System messages (e.g., from prepended conversation) are duplicated into the new - conversation so that the target retains its system prompt context. - - For multi-turn targets this method is a no-op. - - This should be called before each turn when sending prompts to the objective target. - A pending first-turn normalization context suppresses rotation even when prepended - assistant messages have already increased ``executed_turns``. - - Args: - context: The current attack context. - prepended_conversation_config: Configuration that controls target-facing - formatting for the carried system messages. - """ - if self._objective_target.configuration.includes(capability=CapabilityName.MULTI_TURN): - return - - if context.target_normalization_context and not context.target_normalization_context.is_consumed: - return - - if context.executed_turns == 0: - return - - old_conversation_id = context.session.conversation_id - context.related_conversations.add( - ConversationReference( - conversation_id=old_conversation_id, - conversation_type=ConversationType.PRUNED, - description=f"single-turn target prior turn {context.executed_turns}", - ) - ) - - # Duplicate system messages (e.g., system prompt from prepended conversation) - # into the new conversation so the target retains its configuration. - memory = CentralMemory.get_memory_instance() - messages = memory.get_conversation_messages(conversation_id=old_conversation_id) - system_messages = [m for m in messages if m.api_role == "system"] - - if system_messages: - new_conversation_id, pieces = memory.duplicate_messages(messages=system_messages) - memory.add_conversation_to_memory( - conversation=Conversation( - conversation_id=new_conversation_id, target_identifier=self._objective_target.get_identifier() - ) - ) - memory.add_message_pieces_to_memory(message_pieces=pieces) - context.session.conversation_id = new_conversation_id - context.target_normalization_context = ConversationManager.create_target_normalization_context( - target=self._objective_target, - conversation_id=new_conversation_id, - prepended_message_count=len(system_messages), - prepended_conversation_config=prepended_conversation_config, - ) - else: - context.session.conversation_id = str(uuid.uuid4()) - context.target_normalization_context = None - - self._logger.debug( - f"Rotated conversation_id for single-turn target: " - f"{old_conversation_id} -> {context.session.conversation_id}" - ) diff --git a/pyrit/executor/attack/multi_turn/red_teaming.py b/pyrit/executor/attack/multi_turn/red_teaming.py index b0ec823ee2..23342ebade 100644 --- a/pyrit/executor/attack/multi_turn/red_teaming.py +++ b/pyrit/executor/attack/multi_turn/red_teaming.py @@ -174,7 +174,7 @@ def __init__( # Initialize utilities self._prompt_normalizer = prompt_normalizer or PromptNormalizer() - self._prepended_conversation_config = prepended_conversation_config + self._prepended_conversation_config = prepended_conversation_config or PrependedConversationConfig() self._conversation_manager = ConversationManager(prompt_normalizer=self._prompt_normalizer) @@ -476,12 +476,6 @@ async def _send_prompt_to_objective_target_async( """ logger.info(f"Sending prompt to target: {message.get_value()[:50]}...") - # For single-turn targets, rotate conversation_id so each turn starts fresh - self._rotate_conversation_for_single_turn_target( - context=context, - prepended_conversation_config=self._prepended_conversation_config, - ) - with execution_context( component_role=ComponentRole.OBJECTIVE_TARGET, attack_strategy_name=self.__class__.__name__, @@ -496,7 +490,9 @@ async def _send_prompt_to_objective_target_async( request_converter_configurations=self._request_converters, response_converter_configurations=self._response_converters, target=self._objective_target, - target_normalization_context=context.target_normalization_context, + normalizer_overrides=self._prepended_conversation_config.get_normalizer_overrides( + target=self._objective_target + ), ) if response is None: diff --git a/pyrit/executor/attack/multi_turn/tree_of_attacks.py b/pyrit/executor/attack/multi_turn/tree_of_attacks.py index 9d5d90f916..941c0c3aab 100644 --- a/pyrit/executor/attack/multi_turn/tree_of_attacks.py +++ b/pyrit/executor/attack/multi_turn/tree_of_attacks.py @@ -44,7 +44,6 @@ AttackOutcome, AttackResult, ComponentIdentifier, - Conversation, ConversationReference, ConversationType, Message, @@ -53,7 +52,7 @@ SeedPrompt, ) from pyrit.prompt_normalizer import ConverterConfiguration, PromptNormalizer -from pyrit.prompt_target import CapabilityName, PromptTarget, TargetNormalizationContext +from pyrit.prompt_target import CapabilityName, PromptTarget from pyrit.prompt_target.common.target_requirements import TargetRequirements from pyrit.score import ( FloatScaleThresholdScorer, @@ -389,7 +388,7 @@ def __init__( self._attack_strategy_name = attack_strategy_name self._memory_labels = memory_labels or {} self._modality_router = modality_router - self._prepended_conversation_config = prepended_conversation_config + self._prepended_conversation_config = prepended_conversation_config or PrependedConversationConfig() # Initialize utilities self._memory = CentralMemory.get_memory_instance() @@ -414,8 +413,6 @@ def __init__( self.last_prompt_sent: str | None = None self.last_response: Message | None = None self.error_message: str | None = None - self._target_normalization_context: TargetNormalizationContext | None = None - # Context from prepended conversation (for adversarial chat system prompt) self._conversation_context: str | None = None @@ -458,22 +455,14 @@ async def initialize_with_prepended_conversation_async( """ if not prepended_conversation: return - self._prepended_conversation_config = prepended_conversation_config + if prepended_conversation_config: + self._prepended_conversation_config = prepended_conversation_config # Use ConversationManager to add messages to memory conversation_manager = ConversationManager( prompt_normalizer=self._prompt_normalizer, ) - valid_message_count = len( - conversation_manager.get_persistable_prepended_messages(prepended_conversation=prepended_conversation) - ) - target_normalization_context = conversation_manager.create_target_normalization_context( - target=self._objective_target, - conversation_id=self.objective_target_conversation_id, - prepended_message_count=valid_message_count, - prepended_conversation_config=prepended_conversation_config, - ) await conversation_manager.add_prepended_conversation_to_memory_async( prepended_conversation=prepended_conversation, conversation_id=self.objective_target_conversation_id, @@ -482,8 +471,6 @@ async def initialize_with_prepended_conversation_async( target_identifier=self._objective_target.get_identifier(), target=self._objective_target, ) - self._target_normalization_context = target_normalization_context - # Build context string for adversarial chat system prompt (like Crescendo) # The adversarial chat uses this in its system prompt rather than in conversation history self._conversation_context = await build_conversation_context_string_async(prepended_conversation) @@ -615,14 +602,6 @@ async def _send_prompt_to_target_async(self, prompt: str) -> Message: Side Effects: - Sets self.last_response to the target's response text """ - # For single-turn targets, generate a fresh conversation ID before each send - # to ensure the target always receives a clean conversation without prior history. - if not self._objective_target.configuration.includes(capability=CapabilityName.MULTI_TURN) and not ( - self._target_normalization_context and self._target_normalization_context.is_pending - ): - self.objective_target_conversation_id = str(uuid.uuid4()) - self._target_normalization_context = None - # Build the request message via the modality router so prior media (if any) # is included when the objective target accepts it. message = self._modality_router.build_objective_input_message( @@ -645,7 +624,9 @@ async def _send_prompt_to_target_async(self, prompt: str) -> Message: response_converter_configurations=self._response_converters, conversation_id=self.objective_target_conversation_id, target=self._objective_target, - target_normalization_context=self._target_normalization_context, + normalizer_overrides=self._prepended_conversation_config.get_normalizer_overrides( + target=self._objective_target + ), ) # Store the full response so subsequent turns can forward media when supported. @@ -681,13 +662,6 @@ async def _send_initial_prompt_to_target_async(self) -> Message: if self._initial_prompt is None: raise ValueError("_initial_prompt must be set before calling this method") - # For single-turn targets, generate a fresh conversation ID - if not self._objective_target.configuration.includes(capability=CapabilityName.MULTI_TURN) and not ( - self._target_normalization_context and self._target_normalization_context.is_pending - ): - self.objective_target_conversation_id = str(uuid.uuid4()) - self._target_normalization_context = None - assert self._objective is not None initial_prompt = self._initial_prompt self._initial_prompt = None # Clear for future turns @@ -725,7 +699,9 @@ async def _send_initial_prompt_to_target_async(self) -> Message: response_converter_configurations=self._response_converters, conversation_id=self.objective_target_conversation_id, target=self._objective_target, - target_normalization_context=self._target_normalization_context, + normalizer_overrides=self._prepended_conversation_config.get_normalizer_overrides( + target=self._objective_target + ), ) # Store the full response so subsequent turns can forward media when supported. @@ -908,34 +884,11 @@ def duplicate(self) -> _TreeOfAttacksNode: prepended_conversation_config=self._prepended_conversation_config, ) - # Duplicate the conversations to preserve history - # For single-turn targets, duplicate only the system messages (e.g., system prompt - # from prepended conversation) so the target retains its configuration without - # carrying over attack turn history that would cause validation errors. - if self._objective_target.configuration.includes(capability=CapabilityName.MULTI_TURN): - duplicate_node.objective_target_conversation_id = self._memory.duplicate_conversation( - conversation_id=self.objective_target_conversation_id - ) - else: - messages = self._memory.get_conversation_messages(conversation_id=self.objective_target_conversation_id) - system_messages = [m for m in messages if m.api_role == "system"] - if system_messages: - new_id, pieces = self._memory.duplicate_messages(messages=system_messages) - self._memory.add_conversation_to_memory( - conversation=Conversation( - conversation_id=new_id, target_identifier=self._objective_target.get_identifier() - ) - ) - self._memory.add_message_pieces_to_memory(message_pieces=pieces) - duplicate_node.objective_target_conversation_id = new_id - duplicate_node._target_normalization_context = ConversationManager.create_target_normalization_context( - target=self._objective_target, - conversation_id=new_id, - prepended_message_count=len(system_messages), - prepended_conversation_config=self._prepended_conversation_config, - ) - else: - duplicate_node.objective_target_conversation_id = str(uuid.uuid4()) + # Duplicate the complete conversation. The capability pipeline adapts the + # history when the objective target cannot edit it. + duplicate_node.objective_target_conversation_id = self._memory.duplicate_conversation( + conversation_id=self.objective_target_conversation_id + ) duplicate_node.adversarial_chat_conversation_id = self._memory.duplicate_conversation( conversation_id=self.adversarial_chat_conversation_id @@ -1406,7 +1359,7 @@ def __init__( batch_size (int): Number of nodes to process in parallel per batch. Defaults to 10. prepended_conversation_config (PrependedConversationConfig | None): Configuration for how to process prepended conversations. Controls converter - application by role and first-send formatting for targets without editable history. + application by role and request formatting for targets without editable history. Raises: ValueError: If attack_scoring_config uses a non-FloatScaleThresholdScorer objective scorer, @@ -1541,7 +1494,7 @@ def __init__( self._prompt_normalizer = prompt_normalizer or PromptNormalizer() # Store the prepended conversation configuration - self._prepended_conversation_config = prepended_conversation_config + self._prepended_conversation_config = prepended_conversation_config or PrependedConversationConfig() def _load_adversarial_prompts(self) -> None: """Load the adversarial chat prompt template and seed prompt from the default paths.""" diff --git a/pyrit/executor/attack/single_turn/prompt_sending.py b/pyrit/executor/attack/single_turn/prompt_sending.py index 0a9fcebc1c..63a45f86a0 100644 --- a/pyrit/executor/attack/single_turn/prompt_sending.py +++ b/pyrit/executor/attack/single_turn/prompt_sending.py @@ -76,7 +76,7 @@ def __init__( a params type that rejects certain fields. prepended_conversation_config (PrependedConversationConfiguration | None): Configuration for how to process prepended conversations. Controls converter - application by role and first-send formatting for targets without editable history. + application by role and request formatting for targets without editable history. Request converters apply to prepended user messages by default; include ``"assistant"`` explicitly to transform simulated assistant history. @@ -118,7 +118,7 @@ def __init__( self._max_attempts_on_failure = max_attempts_on_failure # Store the prepended conversation configuration - self._prepended_conversation_config = prepended_conversation_config + self._prepended_conversation_config = prepended_conversation_config or PrependedConversationConfig() def get_attack_scoring_config(self) -> AttackScoringConfig | None: """ @@ -325,7 +325,9 @@ async def _send_prompt_to_objective_target_async( conversation_id=context.conversation_id, request_converter_configurations=self._request_converters, response_converter_configurations=self._response_converters, - target_normalization_context=context.target_normalization_context, + normalizer_overrides=self._prepended_conversation_config.get_normalizer_overrides( + target=self._objective_target + ), ) async def _evaluate_response_async( diff --git a/pyrit/executor/attack/streaming/barge_in.py b/pyrit/executor/attack/streaming/barge_in.py index 717e269d71..cd6a3560e4 100644 --- a/pyrit/executor/attack/streaming/barge_in.py +++ b/pyrit/executor/attack/streaming/barge_in.py @@ -11,6 +11,7 @@ from typing import TYPE_CHECKING, Any, ClassVar, cast from pyrit.common.apply_defaults import REQUIRED_VALUE, apply_defaults +from pyrit.executor.attack.component import PrependedConversationConfig from pyrit.executor.attack.component.conversation_manager import ConversationManager from pyrit.executor.attack.core.attack_config import AttackConverterConfig from pyrit.executor.attack.core.attack_parameters import AttackParameters, AttackParamsT @@ -71,6 +72,7 @@ def __init__( attack_converter_config: AttackConverterConfig | None = None, prompt_normalizer: PromptNormalizer | None = None, params_type: type[AttackParamsT] = AttackParameters, # type: ignore[ty:invalid-parameter-default] + prepended_conversation_config: PrependedConversationConfig | None = None, ) -> None: """ Initialize the streaming barge-in attack. @@ -82,6 +84,8 @@ def __init__( prompt_normalizer: Normalizer used to apply converters and persist messages. Defaults to a fresh ``PromptNormalizer``. params_type: Attack parameter dataclass type. + prepended_conversation_config: Configuration for prepended-conversation + conversion and formatting. Raises: ValueError: If ``objective_target`` does not declare the ``STREAMING_AUDIO`` @@ -98,6 +102,7 @@ def __init__( self._request_converters = attack_converter_config.request_converters self._response_converters = attack_converter_config.response_converters self._prompt_normalizer = prompt_normalizer or PromptNormalizer() + self._prepended_conversation_config = prepended_conversation_config or PrependedConversationConfig() self._conversation_manager = ConversationManager( prompt_normalizer=self._prompt_normalizer, ) @@ -135,6 +140,7 @@ async def _setup_async(self, *, context: BargeInAttackContext[Any]) -> None: target=self._objective_target, conversation_id=context.conversation_id, request_converters=self._request_converters, + prepended_conversation_config=self._prepended_conversation_config, ) async def _teardown_async(self, *, context: BargeInAttackContext[Any]) -> None: diff --git a/pyrit/message_normalizer/first_turn_history_normalizer.py b/pyrit/message_normalizer/first_turn_history_normalizer.py index 8f0541ace7..617d213531 100644 --- a/pyrit/message_normalizer/first_turn_history_normalizer.py +++ b/pyrit/message_normalizer/first_turn_history_normalizer.py @@ -23,57 +23,71 @@ def __init__( self, *, message_normalizer: MessageStringNormalizer, - prepended_message_count: int, + target_supports_multi_turn: bool, ) -> None: """ Initialize the normalizer. Args: message_normalizer: Formatter for the target-facing text. - prepended_message_count: Number of leading messages that belong to - the prepended conversation. - - Raises: - ValueError: If prepended_message_count is less than one. + target_supports_multi_turn: Whether the target retains state across + live requests. """ - if prepended_message_count < 1: - raise ValueError("prepended_message_count must be at least 1") self._message_normalizer = message_normalizer - self._prepended_message_count = prepended_message_count + self._target_supports_multi_turn = target_supports_multi_turn async def normalize_async(self, messages: list[Message]) -> list[Message]: """ - Return one target-facing request containing the prepended history. + Adapt memory-backed history for a target whose history is not editable. + + Before the target has replied, persisted prepended history is formatted + with the first live request. Failed live request/error pairs are excluded. + After a real target reply, a stateful target receives only the current + request. A stateless target receives the original prepended prefix plus + the current request on every send. Args: - messages: Prepended history followed by exactly one live request. + messages: Persisted conversation history followed by the current request. Returns: A single request message with formatted text and preserved live non-text pieces. Raises: - ValueError: If the input does not contain the configured prepended - message count followed by one live request, or if converted - prepended history cannot be represented as text. + ValueError: If messages is empty or converted prepended history + cannot be represented as text. """ - expected_count = self._prepended_message_count + 1 - if len(messages) != expected_count: - raise ValueError( - "First-turn history normalization expected " - f"{self._prepended_message_count} prepended messages and one live request, " - f"but received {len(messages)} messages." - ) - - prepended_messages = messages[: self._prepended_message_count] - live_request = messages[-1] + if not messages: + raise ValueError("Messages list cannot be empty") + if len(messages) == 1: + return list(messages) + + history = self._remove_failed_live_requests(messages=messages[:-1]) + messages = [*history, messages[-1]] + if len(messages) == 1: + return messages + + first_response_index = self._find_first_target_response_index(messages=history) + if first_response_index is None: + messages_to_format = messages + elif self._target_supports_multi_turn: + return [messages[-1]] + else: + prepended_end = max(first_response_index - 1, 0) + prepended_messages = messages[:prepended_end] + if not prepended_messages: + return [messages[-1]] + messages_to_format = [*prepended_messages, messages[-1]] + + prepended_messages = messages_to_format[:-1] + live_request = messages_to_format[-1] self._validate_flattenable_converter_output(messages=prepended_messages) - original_view = self._build_original_view(messages=messages) - converted_view = self._build_converted_view(messages=messages) + original_view = self._build_original_view(messages=messages_to_format) + converted_view = self._build_converted_view(messages=messages_to_format) original_text = await self._normalize_context_async(messages=original_view) converted_text = original_text - if self._contains_converted_values(messages=messages): + if self._contains_converted_values(messages=messages_to_format): converted_text = await self._normalize_context_async(messages=converted_view) return [ @@ -84,6 +98,22 @@ async def normalize_async(self, messages: list[Message]) -> list[Message]: ) ] + @staticmethod + def _find_first_target_response_index(*, messages: list[Message]) -> int | None: + for index, message in enumerate(messages): + if any(piece.role == "assistant" and piece.response_error == "none" for piece in message.message_pieces): + return index + return None + + @staticmethod + def _remove_failed_live_requests(*, messages: list[Message]) -> list[Message]: + history = list(messages) + while len(history) >= 2 and any( + piece.role == "assistant" and piece.response_error != "none" for piece in history[-1].message_pieces + ): + history = history[:-2] + return history + async def _normalize_context_async(self, *, messages: list[Message]) -> str: messages_to_normalize = self._filter_live_non_text_pieces(messages=messages) if isinstance(self._message_normalizer, ConversationContextNormalizer): diff --git a/pyrit/prompt_normalizer/normalizer_request.py b/pyrit/prompt_normalizer/normalizer_request.py index cf62e8ebe0..4d42937467 100644 --- a/pyrit/prompt_normalizer/normalizer_request.py +++ b/pyrit/prompt_normalizer/normalizer_request.py @@ -1,12 +1,15 @@ # Copyright (c) Microsoft Corporation. # Licensed under the MIT license. +from collections.abc import Mapping from dataclasses import dataclass +from pyrit.message_normalizer import MessageListNormalizer from pyrit.models import Message from pyrit.prompt_normalizer.converter_configuration import ( ConverterConfiguration, ) +from pyrit.prompt_target.common.target_capabilities import CapabilityName @dataclass @@ -19,6 +22,7 @@ class NormalizerRequest: request_converter_configurations: list[ConverterConfiguration] response_converter_configurations: list[ConverterConfiguration] conversation_id: str | None + normalizer_overrides: dict[CapabilityName, MessageListNormalizer[Message]] def __init__( self, @@ -27,6 +31,7 @@ def __init__( request_converter_configurations: list[ConverterConfiguration] | None = None, response_converter_configurations: list[ConverterConfiguration] | None = None, conversation_id: str | None = None, + normalizer_overrides: Mapping[CapabilityName, MessageListNormalizer[Message]] | None = None, ) -> None: """ Initialize a normalizer request. @@ -38,6 +43,7 @@ def __init__( response_converter_configurations (list[ConverterConfiguration]): Configurations for converting the response. Defaults to an empty list. conversation_id (str | None): The ID of the conversation. Defaults to None. + normalizer_overrides: Optional per-send target normalizer overrides. """ if response_converter_configurations is None: response_converter_configurations = [] @@ -47,3 +53,4 @@ def __init__( self.request_converter_configurations = request_converter_configurations self.response_converter_configurations = response_converter_configurations self.conversation_id = conversation_id + self.normalizer_overrides = dict(normalizer_overrides or {}) diff --git a/pyrit/prompt_normalizer/prompt_normalizer.py b/pyrit/prompt_normalizer/prompt_normalizer.py index e7022ec0e0..9d3a4b9c1b 100644 --- a/pyrit/prompt_normalizer/prompt_normalizer.py +++ b/pyrit/prompt_normalizer/prompt_normalizer.py @@ -8,6 +8,7 @@ import tempfile import traceback import wave +from collections.abc import Mapping from pathlib import Path from typing import Any from uuid import uuid4 @@ -19,6 +20,7 @@ get_execution_context, ) from pyrit.memory import CentralMemory, MemoryInterface, set_message_piece_sha256_async +from pyrit.message_normalizer import MessageListNormalizer from pyrit.models import ( ComponentIdentifier, Conversation, @@ -27,9 +29,8 @@ construct_response_from_request, ) from pyrit.prompt_normalizer import ConverterConfiguration, NormalizerRequest -from pyrit.prompt_target import PromptTarget +from pyrit.prompt_target import CapabilityName, PromptTarget from pyrit.prompt_target.batch_helper import batch_task_async -from pyrit.prompt_target.common.target_normalization_context import TargetNormalizationContext logger = logging.getLogger(__name__) @@ -72,7 +73,7 @@ async def send_prompt_async( conversation_id: str | None = None, request_converter_configurations: list[ConverterConfiguration] | None = None, response_converter_configurations: list[ConverterConfiguration] | None = None, - target_normalization_context: TargetNormalizationContext | None = None, + normalizer_overrides: Mapping[CapabilityName, MessageListNormalizer[Message]] | None = None, ) -> Message: """ Send a single request to a target. @@ -85,7 +86,7 @@ async def send_prompt_async( converting the request. Defaults to an empty list. response_converter_configurations (list[ConverterConfiguration], optional): Configurations for converting the response. Defaults to an empty list. - target_normalization_context: Optional per-conversation target normalization state. + normalizer_overrides: Optional per-send target normalizer overrides. Returns: Message: The response received from the target. @@ -121,19 +122,15 @@ async def send_prompt_async( responses = None try: - if target_normalization_context: + if normalizer_overrides: responses = await target.send_prompt_async( message=request, - target_normalization_context=target_normalization_context, + normalizer_overrides=normalizer_overrides, ) else: responses = await target.send_prompt_async(message=request) self.memory.add_message_to_memory(request=request) except EmptyResponseException as ex: - if target_normalization_context and not target_normalization_context.is_consumed: - cid = request.message_pieces[0].conversation_id if request.message_pieces else None - raise Exception(f"Error sending prompt with conversation ID: {cid}") from ex - # Empty responses are retried, but we don't want them to stop execution self.memory.add_message_to_memory(request=request) @@ -147,12 +144,6 @@ async def send_prompt_async( ] except Exception as ex: - if target_normalization_context and not target_normalization_context.is_consumed: - # The provider was never invoked, so leave memory unchanged and allow - # the same first-turn context to be retried. - cid = request.message_pieces[0].conversation_id if request.message_pieces else None - raise Exception(f"Error sending prompt with conversation ID: {cid}") from ex - # Ensure request to memory before processing exception self.memory.add_message_to_memory(request=request) @@ -228,6 +219,7 @@ async def send_prompt_batch_to_target_async( [request.request_converter_configurations for request in requests], [request.response_converter_configurations for request in requests], [request.conversation_id for request in requests], + [request.normalizer_overrides for request in requests], ] batch_item_keys = [ @@ -235,9 +227,10 @@ async def send_prompt_batch_to_target_async( "request_converter_configurations", "response_converter_configurations", "conversation_id", + "normalizer_overrides", ] - return await batch_task_async( + responses: list[Message] = await batch_task_async( prompt_target=target, batch_size=batch_size, items_to_batch=batch_items, @@ -245,6 +238,7 @@ async def send_prompt_batch_to_target_async( task_arguments=batch_item_keys, target=target, ) + return responses async def convert_values_async( self, diff --git a/pyrit/prompt_target/__init__.py b/pyrit/prompt_target/__init__.py index 69518ca613..9f45cb4ee5 100644 --- a/pyrit/prompt_target/__init__.py +++ b/pyrit/prompt_target/__init__.py @@ -27,10 +27,6 @@ get_known_capabilities, ) from pyrit.prompt_target.common.target_configuration import TargetConfiguration -from pyrit.prompt_target.common.target_normalization_context import ( - TargetNormalizationContext, - TargetNormalizationContextState, -) from pyrit.prompt_target.common.target_requirements import CHAT_TARGET_REQUIREMENTS, TargetRequirements from pyrit.prompt_target.common.utils import limit_requests_per_minute from pyrit.prompt_target.gandalf_target import GandalfLevel, GandalfTarget @@ -110,8 +106,6 @@ def __getattr__(name: str) -> object: "RoundRobinTarget", "TargetCapabilities", "TargetConfiguration", - "TargetNormalizationContext", - "TargetNormalizationContextState", "TargetRequirements", "UnsupportedCapabilityBehavior", "TextTarget", diff --git a/pyrit/prompt_target/common/conversation_normalization_pipeline.py b/pyrit/prompt_target/common/conversation_normalization_pipeline.py index 37d8b024bb..9b349fb734 100644 --- a/pyrit/prompt_target/common/conversation_normalization_pipeline.py +++ b/pyrit/prompt_target/common/conversation_normalization_pipeline.py @@ -2,10 +2,12 @@ # Licensed under the MIT license. import logging -from collections.abc import Mapping +from collections.abc import Callable, Mapping from typing import Any from pyrit.message_normalizer import ( + ConversationContextNormalizer, + FirstTurnHistoryNormalizer, GenericSystemSquashNormalizer, HistorySquashNormalizer, JsonSchemaNormalizer, @@ -26,10 +28,19 @@ # Single registry: add new normalizable capabilities here and nowhere else. # Order in the list determines pipeline execution order. # --------------------------------------------------------------------------- -_NORMALIZER_REGISTRY: list[tuple[CapabilityName, MessageListNormalizer[Message]]] = [ - (CapabilityName.SYSTEM_PROMPT, GenericSystemSquashNormalizer()), - (CapabilityName.MULTI_TURN, HistorySquashNormalizer()), - (CapabilityName.JSON_SCHEMA, JsonSchemaNormalizer()), +NormalizerFactory = Callable[[TargetCapabilities], MessageListNormalizer[Message]] + +_NORMALIZER_REGISTRY: list[tuple[CapabilityName, NormalizerFactory]] = [ + (CapabilityName.SYSTEM_PROMPT, lambda _: GenericSystemSquashNormalizer()), + ( + CapabilityName.EDITABLE_HISTORY, + lambda capabilities: FirstTurnHistoryNormalizer( + message_normalizer=ConversationContextNormalizer(), + target_supports_multi_turn=capabilities.supports_multi_turn, + ), + ), + (CapabilityName.MULTI_TURN, lambda _: HistorySquashNormalizer()), + (CapabilityName.JSON_SCHEMA, lambda _: JsonSchemaNormalizer()), ] # Derived constant — no manual maintenance required. @@ -95,7 +106,7 @@ def from_capabilities( overrides = normalizer_overrides or {} normalizers: list[MessageListNormalizer[Message]] = [] - for capability, default_normalizer in _NORMALIZER_REGISTRY: + for capability, default_normalizer_factory in _NORMALIZER_REGISTRY: if capabilities.includes(capability=capability): continue @@ -112,7 +123,7 @@ def from_capabilities( # which should be called in the request flow once the full end-to-end # workflow is implemented. if behavior == UnsupportedCapabilityBehavior.ADAPT: - normalizer = overrides.get(capability, default_normalizer) + normalizer = overrides.get(capability) or default_normalizer_factory(capabilities) normalizers.append(normalizer) return cls(normalizers=tuple(normalizers)) diff --git a/pyrit/prompt_target/common/prompt_target.py b/pyrit/prompt_target/common/prompt_target.py index bde5edcda1..7d89168088 100644 --- a/pyrit/prompt_target/common/prompt_target.py +++ b/pyrit/prompt_target/common/prompt_target.py @@ -3,9 +3,11 @@ import abc import logging +from collections.abc import Mapping from typing import Any, ClassVar, Literal, final from pyrit.memory import CentralMemory, MemoryInterface +from pyrit.message_normalizer import MessageListNormalizer from pyrit.models import ( ComponentIdentifier, Conversation, @@ -21,7 +23,6 @@ get_known_capabilities, ) from pyrit.prompt_target.common.target_configuration import TargetConfiguration -from pyrit.prompt_target.common.target_normalization_context import TargetNormalizationContext logger = logging.getLogger(__name__) @@ -137,27 +138,24 @@ async def send_prompt_async( self, *, message: Message, - target_normalization_context: TargetNormalizationContext | None = None, + normalizer_overrides: Mapping[CapabilityName, MessageListNormalizer[Message]] | None = None, ) -> list[Message]: """ Validate, normalize, and send a prompt to the target. This is the public entry point called by the prompt normalizer. It: - 1. Validates the message and acquires an optional one-shot target context. - 2. Loads memory history unless the context was already consumed, applies - acquired context normalizers, then runs the target's ordinary pipeline. + 1. Validates the message. + 2. Loads memory history and runs the target's normalization pipeline. 3. Validates the normalized conversation against the target's capabilities. - 4. Marks the context consumed and delegates to - ``_send_prompt_to_target_async`` with the normalized conversation. + 4. Delegates to ``_send_prompt_to_target_async`` with the normalized conversation. Subclasses MUST NOT override this method. Override ``_send_prompt_to_target_async`` instead. Args: message (Message): The message to send. - target_normalization_context: Optional per-conversation normalizers - and one-shot lifecycle state. + normalizer_overrides: Optional per-send target normalizer overrides. Returns: list[Message]: Response messages from the target. @@ -166,31 +164,13 @@ async def send_prompt_async( ValueError: If the message or normalized conversation are empty. """ message.validate() - conversation_id = message.message_pieces[0].conversation_id - should_apply_context = False - if target_normalization_context: - should_apply_context = target_normalization_context.begin_normalization( - conversation_id=conversation_id or "" - ) - - try: - normalized_conversation = await self._get_normalized_conversation_async( - message=message, - target_normalization_context=target_normalization_context, - should_apply_context=should_apply_context, - ) - if not normalized_conversation: - raise ValueError("Normalization pipeline returned an empty conversation. Cannot send an empty request.") - self._validate_request(normalized_conversation=normalized_conversation) - except BaseException: - if target_normalization_context and should_apply_context: - target_normalization_context.restore_pending() - raise - - if target_normalization_context and should_apply_context: - # Target-level retry decorators reuse this normalized payload. Once - # provider invocation starts, attack-level retries must not replay history. - target_normalization_context.mark_consumed() + normalized_conversation = await self._get_normalized_conversation_async( + message=message, + normalizer_overrides=normalizer_overrides, + ) + if not normalized_conversation: + raise ValueError("Normalization pipeline returned an empty conversation. Cannot send an empty request.") + self._validate_request(normalized_conversation=normalized_conversation) return await self._send_prompt_to_target_async(normalized_conversation=normalized_conversation) @abc.abstractmethod @@ -254,15 +234,13 @@ async def _get_normalized_conversation_async( self, *, message: Message, - target_normalization_context: TargetNormalizationContext | None = None, - should_apply_context: bool = False, + normalizer_overrides: Mapping[CapabilityName, MessageListNormalizer[Message]] | None = None, ) -> list[Message]: """ Build the target-facing conversation and run the normalization pipeline. - A consumed target context supplies only the current message so retained - target history is not replayed. Otherwise, memory history is loaded and the - current message is appended before any acquired context normalizers run. + Memory history is loaded and the current message is appended before the + target normalization pipeline runs. The original conversation in memory is never mutated. The returned list is an ephemeral copy intended only for building the API request body. @@ -274,24 +252,21 @@ async def _get_normalized_conversation_async( Args: message (Message): The current message to append. - target_normalization_context: Optional per-conversation normalization state. - should_apply_context: Whether this send acquired the one-shot context. + normalizer_overrides: Optional per-send target normalizer overrides. Returns: list[Message]: The normalized conversation (possibly with system prompt squashed, history squashed, etc.). """ conversation_id = message.message_pieces[0].conversation_id - if target_normalization_context and not should_apply_context: - conversation = [message] - else: - conversation = ( - list(self._memory.get_conversation_messages(conversation_id=conversation_id)) if conversation_id else [] - ) - conversation.append(message) - if target_normalization_context: - conversation = await target_normalization_context.normalize_async(messages=conversation) - normalized = await self.configuration.normalize_async(messages=conversation) + conversation = ( + list(self._memory.get_conversation_messages(conversation_id=conversation_id)) if conversation_id else [] + ) + conversation.append(message) + normalized = await self.configuration.normalize_async( + messages=conversation, + normalizer_overrides=normalizer_overrides, + ) if normalized: # Normalizers may create new Message objects (via Message.from_prompt) with # random conversation_ids. Stamp the correct conversation_id on every diff --git a/pyrit/prompt_target/common/target_capabilities.py b/pyrit/prompt_target/common/target_capabilities.py index 382feed1c3..ac16b7b418 100644 --- a/pyrit/prompt_target/common/target_capabilities.py +++ b/pyrit/prompt_target/common/target_capabilities.py @@ -39,7 +39,7 @@ class CapabilityHandlingPolicy: Design invariants ----------------- * The policy is never consulted if the capability is already supported. - * Non-adaptable capabilities (e.g. ``supports_editable_history``) are not + * Non-adaptable capabilities (e.g. ``supports_multi_message_pieces``) are not represented here; requesting them on a target that lacks them always raises immediately. """ @@ -47,6 +47,7 @@ class CapabilityHandlingPolicy: behaviors: Mapping[CapabilityName, UnsupportedCapabilityBehavior] = field( default_factory=lambda: { CapabilityName.MULTI_TURN: UnsupportedCapabilityBehavior.RAISE, + CapabilityName.EDITABLE_HISTORY: UnsupportedCapabilityBehavior.ADAPT, CapabilityName.SYSTEM_PROMPT: UnsupportedCapabilityBehavior.RAISE, CapabilityName.JSON_SCHEMA: UnsupportedCapabilityBehavior.ADAPT, } @@ -64,7 +65,7 @@ def get_behavior(self, *, capability: CapabilityName) -> UnsupportedCapabilityBe Raises: KeyError: If no behavior exists for the capability. This occurs for - non-adaptable capabilities (e.g., supports_editable_history). + non-adaptable capabilities (e.g., supports_multi_message_pieces). """ try: return self.behaviors[capability] diff --git a/pyrit/prompt_target/common/target_configuration.py b/pyrit/prompt_target/common/target_configuration.py index 6613b9dbb0..ffbbd2f532 100644 --- a/pyrit/prompt_target/common/target_configuration.py +++ b/pyrit/prompt_target/common/target_configuration.py @@ -119,17 +119,30 @@ def ensure_can_handle(self, *, capability: CapabilityName) -> None: if behavior == UnsupportedCapabilityBehavior.RAISE: raise ValueError(f"Target does not support '{capability.value}' and the handling policy is RAISE.") - async def normalize_async(self, *, messages: list[Message]) -> list[Message]: + async def normalize_async( + self, + *, + messages: list[Message], + normalizer_overrides: Mapping[CapabilityName, MessageListNormalizer[Any]] | None = None, + ) -> list[Message]: """ Run the normalization pipeline over the given messages. Args: messages (list[Message]): The full conversation to normalize. + normalizer_overrides: Per-send replacements for capability normalizers. Returns: list[Message]: The (possibly adapted) message list. """ - return await self._pipeline.normalize_async(messages=messages) + pipeline = self._pipeline + if normalizer_overrides: + pipeline = ConversationNormalizationPipeline.from_capabilities( + capabilities=self._capabilities, + policy=self._policy, + normalizer_overrides=normalizer_overrides, + ) + return await pipeline.normalize_async(messages=messages) def as_identifier_params(self) -> dict[str, Any]: """ diff --git a/pyrit/prompt_target/common/target_normalization_context.py b/pyrit/prompt_target/common/target_normalization_context.py deleted file mode 100644 index 7cc88323d4..0000000000 --- a/pyrit/prompt_target/common/target_normalization_context.py +++ /dev/null @@ -1,138 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# Licensed under the MIT license. - -from __future__ import annotations - -from dataclasses import dataclass, field -from enum import Enum -from typing import TYPE_CHECKING - -if TYPE_CHECKING: - from pyrit.message_normalizer import MessageListNormalizer - from pyrit.models import Message - - -class TargetNormalizationContextState(str, Enum): - """Lifecycle state for per-send target normalization.""" - - PENDING = "pending" - PREPARING = "preparing" - CONSUMED = "consumed" - - -@dataclass -class TargetNormalizationContext: - """ - Ephemeral normalization state for one target conversation. - - The context is owned by an attack execution and passed explicitly with each - send. It is never persisted to memory or stored on a shared target. - """ - - conversation_id: str - normalizers: tuple[MessageListNormalizer[Message], ...] - _state: TargetNormalizationContextState = field( - default=TargetNormalizationContextState.PENDING, - init=False, - repr=False, - ) - - def __post_init__(self) -> None: - """ - Validate required context data. - - Raises: - ValueError: If the conversation ID is empty or no normalizers are configured. - """ - if not self.conversation_id: - raise ValueError("conversation_id cannot be empty") - if not self.normalizers: - raise ValueError("At least one target normalizer is required") - - @property - def state(self) -> TargetNormalizationContextState: - """The current lifecycle state.""" - return self._state - - @property - def is_pending(self) -> bool: - """Whether normalization can be attempted.""" - return self._state == TargetNormalizationContextState.PENDING - - @property - def is_consumed(self) -> bool: - """Whether provider invocation has started.""" - return self._state == TargetNormalizationContextState.CONSUMED - - def begin_normalization(self, *, conversation_id: str) -> bool: - """ - Acquire the context for normalization. - - Args: - conversation_id: Conversation ID on the outgoing request. - - Returns: - bool: ``True`` when the caller acquired the context, or ``False`` - after it has already been consumed. - - Raises: - ValueError: If the request belongs to another conversation. - RuntimeError: If another send is already preparing the first request. - """ - if conversation_id != self.conversation_id: - raise ValueError( - "Target normalization context belongs to conversation " - f"'{self.conversation_id}', not '{conversation_id}'." - ) - if self._state == TargetNormalizationContextState.CONSUMED: - return False - if self._state == TargetNormalizationContextState.PREPARING: - # Reject rather than queue so two callers cannot both believe they own - # the first target-facing request. - raise RuntimeError("Target normalization is already in progress for this conversation.") - - self._state = TargetNormalizationContextState.PREPARING - return True - - async def normalize_async(self, *, messages: list[Message]) -> list[Message]: - """ - Run the per-send normalizers while the context is acquired. - - Args: - messages: Target conversation to normalize. - - Returns: - list[Message]: The normalized target conversation. - - Raises: - RuntimeError: If the context has not been acquired. - """ - if self._state != TargetNormalizationContextState.PREPARING: - raise RuntimeError("Target normalization context must be acquired before use.") - - normalized = list(messages) - for normalizer in self.normalizers: - normalized = await normalizer.normalize_async(normalized) - return normalized - - def restore_pending(self) -> None: - """ - Allow another attempt after pre-provider normalization fails. - - Raises: - RuntimeError: If the context is not currently preparing. - """ - if self._state != TargetNormalizationContextState.PREPARING: - raise RuntimeError("Only a preparing target normalization context can be restored.") - self._state = TargetNormalizationContextState.PENDING - - def mark_consumed(self) -> None: - """ - Consume the context immediately before provider invocation. - - Raises: - RuntimeError: If the context is not currently preparing. - """ - if self._state != TargetNormalizationContextState.PREPARING: - raise RuntimeError("Only a preparing target normalization context can be consumed.") - self._state = TargetNormalizationContextState.CONSUMED diff --git a/pyrit/prompt_target/common/target_requirements.py b/pyrit/prompt_target/common/target_requirements.py index dc7ea6527a..5225b9165b 100644 --- a/pyrit/prompt_target/common/target_requirements.py +++ b/pyrit/prompt_target/common/target_requirements.py @@ -123,12 +123,12 @@ def _build_chat_target_requirements() -> TargetRequirements: Returns: TargetRequirements: The requirements for a chat-style target. """ - return TargetRequirements(required=frozenset({CapabilityName.MULTI_TURN, CapabilityName.EDITABLE_HISTORY})) + return TargetRequirements(native_required=frozenset({CapabilityName.MULTI_TURN, CapabilityName.EDITABLE_HISTORY})) CHAT_TARGET_REQUIREMENTS: TargetRequirements = _build_chat_target_requirements() """ -Standard requirements for a chat-style target: must support multi-turn conversations -with an editable history. Consumers validate their target against +Standard requirements for a chat-style target: must natively support multi-turn +conversations with editable history. Consumers validate their target against these requirements at construction time. """ diff --git a/tests/unit/executor/attack/component/test_conversation_manager.py b/tests/unit/executor/attack/component/test_conversation_manager.py index 40146b3d58..2ab295c79d 100644 --- a/tests/unit/executor/attack/component/test_conversation_manager.py +++ b/tests/unit/executor/attack/component/test_conversation_manager.py @@ -38,7 +38,7 @@ from pyrit.message_normalizer import ConversationContextNormalizer, FirstTurnHistoryNormalizer from pyrit.models import ComponentIdentifier, Message, MessagePiece, PromptDataType, Score from pyrit.prompt_normalizer import ConverterConfiguration, PromptNormalizer -from pyrit.prompt_target import PromptTarget +from pyrit.prompt_target import CapabilityName, PromptTarget def _mock_target_id(name: str = "MockTarget") -> ComponentIdentifier: @@ -780,7 +780,7 @@ async def test_non_editable_target_does_not_rewrite_supplied_next_message( assert context.next_message is next_message assert context.next_message.get_value() == "Caller-supplied question" - async def test_non_editable_target_sets_custom_first_send_context( + async def test_non_editable_target_persists_history_without_using_formatter( self, attack_identifier: ComponentIdentifier, mock_prompt_normalizer: MagicMock, @@ -801,11 +801,7 @@ async def test_non_editable_target_sets_custom_first_send_context( prepended_conversation_config=config, ) - assert context.target_normalization_context is not None - assert context.target_normalization_context.conversation_id == conversation_id - normalizer = context.target_normalization_context.normalizers[0] - assert isinstance(normalizer, FirstTurnHistoryNormalizer) - assert normalizer._message_normalizer is message_normalizer + assert len(manager.get_conversation(conversation_id)) == len(sample_conversation) message_normalizer.normalize_string_async.assert_not_called() async def test_returns_turn_count_for_multi_turn_attacks( @@ -1139,7 +1135,6 @@ async def test_prepended_conversion_failure_does_not_partially_write_history( ) assert manager.get_conversation(conversation_id) == [] - assert context.target_normalization_context is None async def test_non_persisted_prepended_message_is_not_counted_in_context( self, @@ -1164,7 +1159,6 @@ async def test_non_persisted_prepended_message_is_not_counted_in_context( ) assert manager.get_conversation(conversation_id) == [] - assert context.target_normalization_context is None async def test_non_persisted_piece_does_not_constrain_flattening( self, @@ -1210,7 +1204,6 @@ async def test_non_persisted_piece_does_not_constrain_flattening( stored = manager.get_conversation(conversation_id) assert len(stored) == 1 assert [piece.converted_value for piece in stored[0].message_pieces] == ["persisted"] - assert context.target_normalization_context is not None async def test_non_editable_target_preserves_converter_piece_indexes( self, @@ -1381,23 +1374,22 @@ async def test_message_normalizer_default_uses_conversation_context_normalizer( mock_prompt_target: MagicMock, sample_conversation: list[Message], ) -> None: - """Test that the default formatter is carried in explicit target-side state.""" - manager = ConversationManager(prompt_normalizer=mock_prompt_normalizer) - conversation_id = str(uuid.uuid4()) - context = _TestAttackContext(params=AttackParameters(objective="Test objective")) - context.prepended_conversation = sample_conversation - - await manager.initialize_context_async( - context=context, - target=mock_prompt_target, - conversation_id=conversation_id, - ) - - assert context.target_normalization_context is not None - normalizer = context.target_normalization_context.normalizers[0] + """Test that the default configuration creates the history normalizer.""" + config = PrependedConversationConfig() + normalizer = config.get_normalizer_overrides(target=mock_prompt_target)[CapabilityName.EDITABLE_HISTORY] assert isinstance(normalizer, FirstTurnHistoryNormalizer) assert isinstance(normalizer._message_normalizer, ConversationContextNormalizer) + def test_message_normalizer_is_not_overridden_for_editable_target( + self, + mock_chat_target: MagicMock, + ) -> None: + mock_chat_target.configuration.includes.return_value = True + + overrides = PrependedConversationConfig().get_normalizer_overrides(target=mock_chat_target) + + assert overrides == {} + # ------------------------------------------------------------------------- # Chat Target Behavior (Config has no effect) # ------------------------------------------------------------------------- diff --git a/tests/unit/executor/attack/multi_turn/test_crescendo.py b/tests/unit/executor/attack/multi_turn/test_crescendo.py index e962dddf41..0e11d8150e 100644 --- a/tests/unit/executor/attack/multi_turn/test_crescendo.py +++ b/tests/unit/executor/attack/multi_turn/test_crescendo.py @@ -458,12 +458,14 @@ def test_init_rejects_adversarial_chat_missing_native_capability( """Adversarial chat must natively support MULTI_TURN and SYSTEM_PROMPT.""" from pyrit.prompt_target.common.target_capabilities import CapabilityName - mock_adversarial_chat.configuration.includes.side_effect = lambda *, capability: ( - capability != CapabilityName(missing_capability) - ) + missing = { + "multi_turn": CapabilityName.MULTI_TURN, + "system_prompt": CapabilityName.SYSTEM_PROMPT, + }[missing_capability] + mock_adversarial_chat.configuration.includes.side_effect = lambda *, capability: capability != missing adversarial_config = AttackAdversarialConfig(target=mock_adversarial_chat) - with pytest.raises(ValueError, match=f"CrescendoAttack .*{missing_capability}"): + with pytest.raises(ValueError, match=f"supports_{missing_capability}"): CrescendoAttack( objective_target=mock_objective_target, attack_adversarial_config=adversarial_config, 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 edc62087eb..85229dc291 100644 --- a/tests/unit/executor/attack/multi_turn/test_red_teaming.py +++ b/tests/unit/executor/attack/multi_turn/test_red_teaming.py @@ -953,13 +953,13 @@ async def test_generate_next_prompt_raises_on_none_response( class TestObjectiveTargetSending: """Tests for sending prompts to the objective target.""" - async def test_second_turn_rotation_uses_configured_message_normalizer( + async def test_second_turn_uses_configured_message_normalizer_without_rotation( self, mock_objective_scorer: MagicMock, mock_adversarial_chat: MagicMock, basic_context: MultiTurnAttackContext, ) -> None: - """A rotated request must keep the configured prepended-history formatter.""" + """A stateless target must reuse formatting without attack-specific rotation.""" objective_target = MockPromptTarget() objective_target._configuration = TargetConfiguration(capabilities=TargetCapabilities()) message_normalizer = MagicMock(spec=MessageStringNormalizer) @@ -996,9 +996,7 @@ async def test_second_turn_rotation_uses_configured_message_normalizer( message=Message.from_prompt(prompt="Second request", role="user"), ) - assert basic_context.session.conversation_id != old_conversation_id - assert basic_context.target_normalization_context is not None - assert basic_context.target_normalization_context.is_consumed + assert basic_context.session.conversation_id == old_conversation_id assert objective_target.prompt_sent == ["custom formatted request"] message_normalizer.normalize_string_async.assert_awaited_once() diff --git a/tests/unit/executor/attack/multi_turn/test_supports_multi_turn_attacks.py b/tests/unit/executor/attack/multi_turn/test_supports_multi_turn_attacks.py index 396c0896c5..2c777105b8 100644 --- a/tests/unit/executor/attack/multi_turn/test_supports_multi_turn_attacks.py +++ b/tests/unit/executor/attack/multi_turn/test_supports_multi_turn_attacks.py @@ -1,27 +1,20 @@ # Copyright (c) Microsoft Corporation. # Licensed under the MIT license. -from typing import Any -from unittest.mock import AsyncMock, MagicMock, patch +from unittest.mock import AsyncMock, MagicMock import pytest -from pyrit.executor.attack.component import PrependedConversationConfig -from pyrit.executor.attack.core.attack_parameters import AttackParameters -from pyrit.executor.attack.multi_turn.multi_turn_attack_strategy import ( - ConversationSession, - MultiTurnAttackContext, -) +from pyrit.executor.attack.component import ConversationManager, PrependedConversationConfig +from pyrit.executor.attack.component.modality_router import _ModalityFeedbackRouter +from pyrit.executor.attack.core import AttackAdversarialConfig, AttackScoringConfig +from pyrit.executor.attack.multi_turn.tree_of_attacks import _TreeOfAttacksNode from pyrit.memory import CentralMemory -from pyrit.message_normalizer import ( - ConversationContextNormalizer, - FirstTurnHistoryNormalizer, - MessageStringNormalizer, -) -from pyrit.models import ConversationType, Message, MessagePiece -from pyrit.prompt_target import PromptTarget, TargetNormalizationContext -from pyrit.prompt_target.common.target_capabilities import TargetCapabilities -from pyrit.prompt_target.common.target_configuration import TargetConfiguration +from pyrit.message_normalizer import MessageStringNormalizer +from pyrit.models import Message, MessagePiece +from pyrit.prompt_normalizer import PromptNormalizer +from pyrit.prompt_target import PromptTarget, TargetCapabilities, TargetConfiguration +from pyrit.score import TrueFalseScorer class _SingleTurnPromptTarget(PromptTarget): @@ -29,1059 +22,188 @@ class _SingleTurnPromptTarget(PromptTarget): def __init__(self) -> None: super().__init__() - self.normalized_conversations: list[list[Message]] = [] + self.prompt_sent: list[str] = [] async def _send_prompt_to_target_async(self, *, normalized_conversation: list[Message]) -> list[Message]: - self.normalized_conversations.append(normalized_conversation) - request_piece = normalized_conversation[-1].get_piece() + request = normalized_conversation[-1] + self.prompt_sent.append(request.get_value()) return [ MessagePiece( role="assistant", original_value="response", - conversation_id=request_piece.conversation_id, + conversation_id=request.get_piece().conversation_id, ).to_message() ] -def _make_context() -> MultiTurnAttackContext: - return MultiTurnAttackContext( - params=AttackParameters(objective="Test objective"), - session=ConversationSession(), +@pytest.mark.usefixtures("patch_central_database") +async def test_single_turn_target_reuses_prepended_history_without_conversation_rotation(): + target = _SingleTurnPromptTarget() + prompt_normalizer = PromptNormalizer() + conversation_manager = ConversationManager(prompt_normalizer=prompt_normalizer) + formatter = MagicMock(spec=MessageStringNormalizer) + formatter.normalize_string_async = AsyncMock(side_effect=lambda messages: f"formatted: {messages[-1].get_value()}") + config = PrependedConversationConfig(message_normalizer=formatter) + conversation_id = "conversation" + + await conversation_manager.add_prepended_conversation_to_memory_async( + prepended_conversation=[Message.from_system_prompt("system")], + conversation_id=conversation_id, + prepended_conversation_config=config, + target=target, ) + overrides = config.get_normalizer_overrides(target=target) - -def _make_strategy(*, supports_multi_turn: bool): - """Create a minimal MultiTurnAttackStrategy for testing.""" - from pyrit.executor.attack.multi_turn.multi_turn_attack_strategy import MultiTurnAttackStrategy - - target = MagicMock() - target.capabilities.supports_multi_turn = supports_multi_turn - target.configuration.includes.return_value = supports_multi_turn - target.get_identifier.return_value = {"__type__": "MockTarget", "__module__": "test", "id": "mock-id"} - - with patch.multiple( - MultiTurnAttackStrategy, - __abstractmethods__=frozenset(), - _perform_async=AsyncMock(), - _setup_async=AsyncMock(), - ): - strategy = MultiTurnAttackStrategy( - objective_target=target, - context_type=MultiTurnAttackContext, - ) - return strategy # noqa: RET504 - - -def _seed_conversation(*, conversation_id: str, system_prompt: str, user_text: str = "Hello") -> None: - """Add a system message and a user message to memory under the given conversation_id.""" - memory = CentralMemory.get_memory_instance() - - sys_piece = MessagePiece( - original_value=system_prompt, - role="system", + await prompt_normalizer.send_prompt_async( + message=Message.from_prompt(prompt="first", role="user"), + target=target, conversation_id=conversation_id, - sequence=0, + normalizer_overrides=overrides, ) - user_piece = MessagePiece( - original_value=user_text, - role="user", + await prompt_normalizer.send_prompt_async( + message=Message.from_prompt(prompt="second", role="user"), + target=target, conversation_id=conversation_id, - sequence=1, + normalizer_overrides=overrides, ) - memory.add_message_pieces_to_memory(message_pieces=[sys_piece, user_piece]) - - -@pytest.mark.usefixtures("patch_central_database") -class TestRotateConversationForSingleTurnTarget: - """Test the _rotate_conversation_for_single_turn_target helper.""" - - def test_noop_for_multi_turn_target(self): - strategy = _make_strategy(supports_multi_turn=True) - context = _make_context() - context.executed_turns = 1 - original_id = context.session.conversation_id - - strategy._rotate_conversation_for_single_turn_target(context=context) - - assert context.session.conversation_id == original_id - assert len(context.related_conversations) == 0 - - def test_noop_on_first_turn(self): - strategy = _make_strategy(supports_multi_turn=False) - context = _make_context() - context.executed_turns = 0 - original_id = context.session.conversation_id - - strategy._rotate_conversation_for_single_turn_target(context=context) - - assert context.session.conversation_id == original_id - assert len(context.related_conversations) == 0 - - def test_pending_first_turn_context_suppresses_rotation_after_prepended_turns(self): - strategy = _make_strategy(supports_multi_turn=False) - context = _make_context() - context.executed_turns = 2 - original_id = context.session.conversation_id - context.target_normalization_context = TargetNormalizationContext( - conversation_id=original_id, - normalizers=( - FirstTurnHistoryNormalizer( - message_normalizer=ConversationContextNormalizer(), - prepended_message_count=4, - ), - ), - ) - - strategy._rotate_conversation_for_single_turn_target(context=context) - - assert context.session.conversation_id == original_id - assert len(context.related_conversations) == 0 - - def test_rotates_on_second_turn_for_single_turn_target(self): - strategy = _make_strategy(supports_multi_turn=False) - context = _make_context() - context.executed_turns = 1 - original_id = context.session.conversation_id - - strategy._rotate_conversation_for_single_turn_target(context=context) - - assert context.session.conversation_id != original_id - assert len(context.related_conversations) == 1 - ref = next(iter(context.related_conversations)) - assert ref.conversation_id == original_id - assert ref.conversation_type == ConversationType.PRUNED - def test_rotates_multiple_turns(self): - strategy = _make_strategy(supports_multi_turn=False) - context = _make_context() - seen_ids = {context.session.conversation_id} + assert target.prompt_sent == ["formatted: first", "formatted: second"] + assert formatter.normalize_string_async.await_count == 2 + second_messages = formatter.normalize_string_async.await_args_list[1].args[0] + assert [message.get_value() for message in second_messages] == ["system", "second"] - for turn in range(1, 4): - context.executed_turns = turn - strategy._rotate_conversation_for_single_turn_target(context=context) - assert context.session.conversation_id not in seen_ids - seen_ids.add(context.session.conversation_id) - assert len(context.related_conversations) == 3 +def _make_tap_node(*, supports_multi_turn: bool) -> _TreeOfAttacksNode: + target = MagicMock() + target.configuration.includes.return_value = supports_multi_turn + target.configuration.capabilities.input_modalities = frozenset({frozenset({"text"})}) + target.configuration.capabilities.output_modalities = frozenset({frozenset({"text"})}) + target.get_identifier.return_value = {"__type__": "MockTarget", "__module__": "test", "id": "target"} + + adversarial_chat = MagicMock() + adversarial_chat.get_identifier.return_value = { + "__type__": "MockTarget", + "__module__": "test", + "id": "adversarial", + } + adversarial_chat.configuration.capabilities.input_modalities = frozenset({frozenset({"text"})}) + adversarial_chat.configuration.capabilities.output_modalities = frozenset({frozenset({"text"})}) + + scorer = MagicMock() + scorer.get_identifier.return_value = {"__type__": "MockScorer", "__module__": "test", "id": "scorer"} + seed = MagicMock() + seed.render_template_value.return_value = "template" + + return _TreeOfAttacksNode( + objective_target=target, + adversarial_chat=adversarial_chat, + adversarial_chat_seed_prompt=seed, + adversarial_chat_prompt_template=seed, + adversarial_chat_system_seed_prompt=seed, + desired_response_prefix="Sure,", + objective_scorer=scorer, + on_topic_scorer=None, + request_converters=[], + response_converters=[], + auxiliary_scorers=None, + attack_id=MagicMock(), + attack_strategy_name="TAP", + modality_router=_ModalityFeedbackRouter( + adversarial_chat=adversarial_chat, + objective_target=target, + ), + ) @pytest.mark.usefixtures("patch_central_database") -class TestSystemPromptCarryoverOnRotation: - """Test that system prompts are duplicated into the new conversation on rotation.""" - - def test_system_prompt_duplicated_into_new_conversation(self): - """When rotating, system messages must be copied to the new conversation_id.""" - strategy = _make_strategy(supports_multi_turn=False) - context = _make_context() - old_id = context.session.conversation_id - - _seed_conversation( - conversation_id=old_id, - system_prompt="You are a helpful assistant.", - user_text="Turn 1 user message", - ) - - context.executed_turns = 1 - strategy._rotate_conversation_for_single_turn_target(context=context) - - new_id = context.session.conversation_id - assert new_id != old_id - - memory = CentralMemory.get_memory_instance() - new_messages = memory.get_conversation_messages(conversation_id=new_id) - - # Only the system message should be in the new conversation (not the user message) - assert len(new_messages) == 1 - assert new_messages[0].api_role == "system" - assert new_messages[0].get_value() == "You are a helpful assistant." - - def test_system_prompt_preserved_across_multiple_rotations(self): - """System prompt must carry over through successive rotations.""" - strategy = _make_strategy(supports_multi_turn=False) - context = _make_context() - memory = CentralMemory.get_memory_instance() - - # Seed initial conversation with system prompt + user message - _seed_conversation( - conversation_id=context.session.conversation_id, - system_prompt="You are an expert.", - ) - - for turn in range(1, 4): - context.executed_turns = turn - strategy._rotate_conversation_for_single_turn_target(context=context) - - messages = memory.get_conversation_messages(conversation_id=context.session.conversation_id) - system_msgs = [m for m in messages if m.api_role == "system"] - assert len(system_msgs) == 1, f"Turn {turn}: expected 1 system message, got {len(system_msgs)}" - assert system_msgs[0].get_value() == "You are an expert." - - # Simulate a user message for the next turn's conversation - user_piece = MessagePiece( - original_value=f"User message turn {turn}", +@pytest.mark.parametrize("supports_multi_turn", [False, True]) +def test_tap_branch_duplicates_full_history(supports_multi_turn: bool): + node = _make_tap_node(supports_multi_turn=supports_multi_turn) + memory = CentralMemory.get_memory_instance() + memory.add_message_pieces_to_memory( + message_pieces=[ + MessagePiece( + original_value="system", + role="system", + conversation_id=node.objective_target_conversation_id, + sequence=0, + ), + MessagePiece( + original_value="request", role="user", - conversation_id=context.session.conversation_id, + conversation_id=node.objective_target_conversation_id, sequence=1, - ) - memory.add_message_pieces_to_memory(message_pieces=[user_piece]) - - async def test_rotated_system_prompt_is_normalized_with_next_request(self): - target = _SingleTurnPromptTarget() - strategy = _make_strategy(supports_multi_turn=False) - strategy._objective_target = target - context = _make_context() - old_id = context.session.conversation_id - _seed_conversation( - conversation_id=old_id, - system_prompt="You are a helpful assistant.", - user_text="First request", - ) - previous_context = TargetNormalizationContext( - conversation_id=old_id, - normalizers=( - FirstTurnHistoryNormalizer( - message_normalizer=ConversationContextNormalizer(), - prepended_message_count=1, - ), ), - ) - previous_context.begin_normalization(conversation_id=old_id) - previous_context.mark_consumed() - context.target_normalization_context = previous_context - context.executed_turns = 1 - - strategy._rotate_conversation_for_single_turn_target(context=context) - - next_request = Message.from_prompt(prompt="Second request", role="user") - next_request.get_piece().conversation_id = context.session.conversation_id - await target.send_prompt_async( - message=next_request, - target_normalization_context=context.target_normalization_context, - ) - - assert context.target_normalization_context is not None - assert context.target_normalization_context.is_consumed - assert len(target.normalized_conversations) == 1 - normalized_conversation = target.normalized_conversations[0] - assert len(normalized_conversation) == 1 - assert normalized_conversation[0].get_value() == ( - "Turn 1:\nuser: ### Instructions ###\n\nYou are a helpful assistant.\n\n######\n\nSecond request" - ) - - async def test_rotation_preserves_configured_message_normalizer(self): - """Rotation must retain the formatter used for carried system messages.""" - target = _SingleTurnPromptTarget() - strategy = _make_strategy(supports_multi_turn=False) - strategy._objective_target = target - context = _make_context() - old_id = context.session.conversation_id - _seed_conversation( - conversation_id=old_id, - system_prompt="You are a helpful assistant.", - user_text="First request", - ) - previous_context = TargetNormalizationContext( - conversation_id=old_id, - normalizers=( - FirstTurnHistoryNormalizer( - message_normalizer=ConversationContextNormalizer(), - prepended_message_count=1, - ), + MessagePiece( + original_value="response", + role="assistant", + conversation_id=node.objective_target_conversation_id, + sequence=2, ), - ) - previous_context.begin_normalization(conversation_id=old_id) - previous_context.mark_consumed() - context.target_normalization_context = previous_context - context.executed_turns = 1 - message_normalizer = MagicMock(spec=MessageStringNormalizer) - message_normalizer.normalize_string_async = AsyncMock(return_value="custom formatted request") - - strategy._rotate_conversation_for_single_turn_target( - context=context, - prepended_conversation_config=PrependedConversationConfig(message_normalizer=message_normalizer), - ) - - next_request = Message.from_prompt(prompt="Second request", role="user") - next_request.get_piece().conversation_id = context.session.conversation_id - await target.send_prompt_async( - message=next_request, - target_normalization_context=context.target_normalization_context, - ) - - assert target.normalized_conversations[0][0].get_value() == "custom formatted request" - message_normalizer.normalize_string_async.assert_awaited_once() - - def test_no_system_prompt_yields_fresh_conversation_id(self): - """When there is no system prompt, rotation still generates a new conversation_id.""" - strategy = _make_strategy(supports_multi_turn=False) - context = _make_context() - memory = CentralMemory.get_memory_instance() - - # Seed conversation with only a user message (no system prompt) - user_piece = MessagePiece( - original_value="Just a user message", - role="user", - conversation_id=context.session.conversation_id, - sequence=0, - ) - memory.add_message_pieces_to_memory(message_pieces=[user_piece]) - - old_id = context.session.conversation_id - context.executed_turns = 1 - strategy._rotate_conversation_for_single_turn_target(context=context) - - assert context.session.conversation_id != old_id - new_messages = memory.get_conversation_messages(conversation_id=context.session.conversation_id) - assert len(new_messages) == 0 - - def test_user_messages_not_carried_over(self): - """Only system messages should be carried to the new conversation, not user/assistant.""" - strategy = _make_strategy(supports_multi_turn=False) - context = _make_context() - memory = CentralMemory.get_memory_instance() - - # Seed with system + user + assistant - sys_piece = MessagePiece( - original_value="System prompt", - role="system", - conversation_id=context.session.conversation_id, - sequence=0, - ) - user_piece = MessagePiece( - original_value="User message", - role="user", - conversation_id=context.session.conversation_id, - sequence=1, - ) - asst_piece = MessagePiece( - original_value="Assistant response", - role="assistant", - conversation_id=context.session.conversation_id, - sequence=2, - ) - memory.add_message_pieces_to_memory(message_pieces=[sys_piece, user_piece, asst_piece]) - - context.executed_turns = 1 - strategy._rotate_conversation_for_single_turn_target(context=context) - - new_messages = memory.get_conversation_messages(conversation_id=context.session.conversation_id) - roles = [m.api_role for m in new_messages] - assert roles == ["system"], f"Expected only system, got {roles}" - - def test_multiple_system_messages_all_carried_over(self): - """When multiple system messages exist (different sequences), all are duplicated.""" - strategy = _make_strategy(supports_multi_turn=False) - context = _make_context() - memory = CentralMemory.get_memory_instance() - - # Two system messages at different sequences (e.g., system prompt + safety instructions - # injected at different points in a prepended conversation) - sys1 = MessagePiece( - original_value="System prompt 1", - role="system", - conversation_id=context.session.conversation_id, - sequence=0, - ) - user_piece = MessagePiece( - original_value="Hello", - role="user", - conversation_id=context.session.conversation_id, - sequence=1, - ) - sys2 = MessagePiece( - original_value="Safety instructions", - role="system", - conversation_id=context.session.conversation_id, - sequence=2, - ) - memory.add_message_pieces_to_memory(message_pieces=[sys1, user_piece, sys2]) - - context.executed_turns = 1 - strategy._rotate_conversation_for_single_turn_target(context=context) - - new_messages = memory.get_conversation_messages(conversation_id=context.session.conversation_id) - system_values = sorted(m.get_value() for m in new_messages if m.api_role == "system") - assert system_values == ["Safety instructions", "System prompt 1"] - assert all(m.api_role == "system" for m in new_messages) - - def test_empty_conversation_yields_fresh_id(self): - """When the conversation has zero messages, rotation still produces a new ID.""" - strategy = _make_strategy(supports_multi_turn=False) - context = _make_context() - old_id = context.session.conversation_id - - # Don't seed any messages — conversation is completely empty - context.executed_turns = 1 - strategy._rotate_conversation_for_single_turn_target(context=context) - - assert context.session.conversation_id != old_id - memory = CentralMemory.get_memory_instance() - new_messages = memory.get_conversation_messages(conversation_id=context.session.conversation_id) - assert len(new_messages) == 0 - - def test_only_system_messages_all_carried_over(self): - """When the conversation contains only system messages (no user/assistant), all are carried.""" - strategy = _make_strategy(supports_multi_turn=False) - context = _make_context() - memory = CentralMemory.get_memory_instance() - - sys_piece = MessagePiece( - original_value="Only a system message", - role="system", - conversation_id=context.session.conversation_id, - sequence=0, - ) - memory.add_message_pieces_to_memory(message_pieces=[sys_piece]) - - context.executed_turns = 1 - strategy._rotate_conversation_for_single_turn_target(context=context) - - new_messages = memory.get_conversation_messages(conversation_id=context.session.conversation_id) - assert len(new_messages) == 1 - assert new_messages[0].api_role == "system" - assert new_messages[0].get_value() == "Only a system message" + ] + ) - def test_multipiece_system_message_fully_duplicated(self): - """A system Message with multiple pieces (same sequence) is fully duplicated.""" - strategy = _make_strategy(supports_multi_turn=False) - context = _make_context() - memory = CentralMemory.get_memory_instance() + duplicate = node.duplicate() - # Two system pieces at the same sequence form one multi-piece Message - sys_text = MessagePiece( - original_value="System text instruction", - role="system", - conversation_id=context.session.conversation_id, - sequence=0, - ) - sys_image = MessagePiece( - original_value="image_placeholder", - role="system", - conversation_id=context.session.conversation_id, - sequence=0, - original_value_data_type="image_path", - ) - user_piece = MessagePiece( - original_value="Hello", - role="user", - conversation_id=context.session.conversation_id, - sequence=1, - ) - memory.add_message_pieces_to_memory(message_pieces=[sys_text, sys_image, user_piece]) + messages = memory.get_conversation_messages(conversation_id=duplicate.objective_target_conversation_id) + assert [message.api_role for message in messages] == ["system", "user", "assistant"] - context.executed_turns = 1 - strategy._rotate_conversation_for_single_turn_target(context=context) - new_messages = memory.get_conversation_messages(conversation_id=context.session.conversation_id) - assert len(new_messages) == 1 - assert new_messages[0].api_role == "system" - # Both pieces should be present in the duplicated message - assert len(new_messages[0].message_pieces) == 2 - values = {p.converted_value for p in new_messages[0].message_pieces} - assert values == {"System text instruction", "image_placeholder"} +@pytest.fixture +def single_turn_target() -> MagicMock: + target = MagicMock() + target.configuration = TargetConfiguration( + capabilities=TargetCapabilities(supports_multi_turn=False, supports_system_prompt=True) + ) + target.get_identifier.return_value = {"__type__": "MockTarget", "__module__": "test", "id": "target"} + return target - def test_old_conversation_untouched_after_rotation(self): - """Rotation must not alter messages in the old conversation.""" - strategy = _make_strategy(supports_multi_turn=False) - context = _make_context() - memory = CentralMemory.get_memory_instance() - old_id = context.session.conversation_id - _seed_conversation( - conversation_id=old_id, - system_prompt="Original system prompt", - user_text="Original user message", - ) +@pytest.fixture +def adversarial_config() -> AttackAdversarialConfig: + adversarial_chat = MagicMock() + adversarial_chat.get_identifier.return_value = { + "__type__": "MockTarget", + "__module__": "test", + "id": "adversarial", + } + return AttackAdversarialConfig(target=adversarial_chat) - context.executed_turns = 1 - strategy._rotate_conversation_for_single_turn_target(context=context) - # Old conversation should still have both messages intact - old_messages = memory.get_conversation_messages(conversation_id=old_id) - old_roles = [m.api_role for m in old_messages] - assert old_roles == ["system", "user"] +@pytest.fixture +def scoring_config() -> AttackScoringConfig: + scorer = MagicMock(spec=TrueFalseScorer) + scorer.get_identifier.return_value = {"__type__": "MockScorer", "__module__": "test", "id": "scorer"} + return AttackScoringConfig(objective_scorer=scorer) @pytest.mark.usefixtures("patch_central_database") -class TestTAPNodeDuplicateSystemMessages: - """Test that TAP's duplicate correctly handles system messages.""" - - def _make_tap_node(self, *, supports_multi_turn: bool): - """Create a minimal _TreeOfAttacksNode for testing.""" - from pyrit.executor.attack.component.modality_router import _ModalityFeedbackRouter - from pyrit.executor.attack.multi_turn.tree_of_attacks import _TreeOfAttacksNode - - target = MagicMock() - target.capabilities.supports_multi_turn = supports_multi_turn - target.configuration.includes.return_value = supports_multi_turn - target.configuration.capabilities.input_modalities = frozenset({frozenset({"text"})}) - target.configuration.capabilities.output_modalities = frozenset({frozenset({"text"})}) - target.get_identifier.return_value = {"__type__": "MockTarget", "__module__": "test", "id": "mock-id"} - - adversarial_chat = MagicMock() - adversarial_chat.get_identifier.return_value = {"__type__": "MockTarget", "__module__": "test", "id": "mock-id"} - adversarial_chat.configuration.capabilities.input_modalities = frozenset({frozenset({"text"})}) - adversarial_chat.configuration.capabilities.output_modalities = frozenset({frozenset({"text"})}) - - scorer = MagicMock() - scorer.get_identifier.return_value = {"__type__": "MockTarget", "__module__": "test", "id": "mock-id"} +def test_crescendo_requires_native_multi_turn( + single_turn_target: MagicMock, + adversarial_config: AttackAdversarialConfig, + scoring_config: AttackScoringConfig, +): + from pyrit.executor.attack.multi_turn.crescendo import CrescendoAttack - seed = MagicMock() - seed.render_template_value.return_value = "template" - - return _TreeOfAttacksNode( - objective_target=target, - adversarial_chat=adversarial_chat, - adversarial_chat_seed_prompt=seed, - adversarial_chat_prompt_template=seed, - adversarial_chat_system_seed_prompt=seed, - desired_response_prefix="Sure,", - objective_scorer=scorer, - on_topic_scorer=None, - request_converters=[], - response_converters=[], - auxiliary_scorers=None, - attack_id=MagicMock(), - attack_strategy_name="TAP", - modality_router=_ModalityFeedbackRouter( - adversarial_chat=adversarial_chat, - objective_target=target, - ), + with pytest.raises(ValueError, match="supports_multi_turn"): + CrescendoAttack( + objective_target=single_turn_target, + attack_adversarial_config=adversarial_config, + attack_scoring_config=scoring_config, ) - def test_single_turn_target_duplicates_only_system_messages(self): - """For single-turn targets, only system messages are copied to the duplicate node.""" - node = self._make_tap_node(supports_multi_turn=False) - memory = CentralMemory.get_memory_instance() - - # Seed the node's conversation with system + user + assistant messages - sys_piece = MessagePiece( - original_value="TAP system prompt", - role="system", - conversation_id=node.objective_target_conversation_id, - sequence=0, - ) - user_piece = MessagePiece( - original_value="Attack prompt turn 1", - role="user", - conversation_id=node.objective_target_conversation_id, - sequence=1, - ) - asst_piece = MessagePiece( - original_value="Target response turn 1", - role="assistant", - conversation_id=node.objective_target_conversation_id, - sequence=2, - ) - memory.add_message_pieces_to_memory(message_pieces=[sys_piece, user_piece, asst_piece]) - - duplicate = node.duplicate() - - # The duplicate should have a different conversation_id - assert duplicate.objective_target_conversation_id != node.objective_target_conversation_id - - # The duplicate's conversation should contain only the system message - dup_messages = memory.get_conversation_messages(conversation_id=duplicate.objective_target_conversation_id) - assert len(dup_messages) == 1 - assert dup_messages[0].api_role == "system" - assert dup_messages[0].get_value() == "TAP system prompt" - - def test_multi_turn_target_duplicates_full_conversation(self): - """For multi-turn targets, the full conversation is duplicated.""" - node = self._make_tap_node(supports_multi_turn=True) - memory = CentralMemory.get_memory_instance() - - # Seed the node's conversation with system + user + assistant - sys_piece = MessagePiece( - original_value="TAP system prompt", - role="system", - conversation_id=node.objective_target_conversation_id, - sequence=0, - ) - user_piece = MessagePiece( - original_value="Attack prompt", - role="user", - conversation_id=node.objective_target_conversation_id, - sequence=1, - ) - asst_piece = MessagePiece( - original_value="Target response", - role="assistant", - conversation_id=node.objective_target_conversation_id, - sequence=2, - ) - memory.add_message_pieces_to_memory(message_pieces=[sys_piece, user_piece, asst_piece]) - - duplicate = node.duplicate() - - assert duplicate.objective_target_conversation_id != node.objective_target_conversation_id - - dup_messages = memory.get_conversation_messages(conversation_id=duplicate.objective_target_conversation_id) - roles = [m.api_role for m in dup_messages] - assert roles == ["system", "user", "assistant"] - - def test_single_turn_no_system_messages_yields_fresh_id(self): - """For single-turn targets with no system messages, a fresh empty conversation is created.""" - node = self._make_tap_node(supports_multi_turn=False) - memory = CentralMemory.get_memory_instance() - - # Seed with only user/assistant (no system prompt) - user_piece = MessagePiece( - original_value="Attack prompt", - role="user", - conversation_id=node.objective_target_conversation_id, - sequence=0, - ) - memory.add_message_pieces_to_memory(message_pieces=[user_piece]) - - duplicate = node.duplicate() - - assert duplicate.objective_target_conversation_id != node.objective_target_conversation_id - dup_messages = memory.get_conversation_messages(conversation_id=duplicate.objective_target_conversation_id) - assert len(dup_messages) == 0 - - def test_adversarial_chat_always_fully_duplicated(self): - """The adversarial chat conversation should always be fully duplicated regardless of target type.""" - node = self._make_tap_node(supports_multi_turn=False) - memory = CentralMemory.get_memory_instance() - - # Seed adversarial chat conversation - sys_piece = MessagePiece( - original_value="Adversarial system prompt", - role="system", - conversation_id=node.adversarial_chat_conversation_id, - sequence=0, - ) - user_piece = MessagePiece( - original_value="Adversarial user message", - role="user", - conversation_id=node.adversarial_chat_conversation_id, - sequence=1, - ) - memory.add_message_pieces_to_memory(message_pieces=[sys_piece, user_piece]) - - # Also seed objective target conversation so it doesn't error - target_piece = MessagePiece( - original_value="Target user", - role="user", - conversation_id=node.objective_target_conversation_id, - sequence=0, - ) - memory.add_message_pieces_to_memory(message_pieces=[target_piece]) - - duplicate = node.duplicate() - - dup_adv_messages = memory.get_conversation_messages(conversation_id=duplicate.adversarial_chat_conversation_id) - roles = [m.api_role for m in dup_adv_messages] - assert roles == ["system", "user"] - - def test_single_turn_multiple_system_messages_all_duplicated(self): - """For single-turn targets with multiple system messages, all are duplicated.""" - node = self._make_tap_node(supports_multi_turn=False) - memory = CentralMemory.get_memory_instance() - - sys1 = MessagePiece( - original_value="System prompt A", - role="system", - conversation_id=node.objective_target_conversation_id, - sequence=0, - ) - user_piece = MessagePiece( - original_value="Attack prompt", - role="user", - conversation_id=node.objective_target_conversation_id, - sequence=1, - ) - sys2 = MessagePiece( - original_value="System prompt B", - role="system", - conversation_id=node.objective_target_conversation_id, - sequence=2, - ) - memory.add_message_pieces_to_memory(message_pieces=[sys1, user_piece, sys2]) - - duplicate = node.duplicate() - - dup_messages = memory.get_conversation_messages(conversation_id=duplicate.objective_target_conversation_id) - assert all(m.api_role == "system" for m in dup_messages) - dup_values = sorted(m.get_value() for m in dup_messages) - assert dup_values == ["System prompt A", "System prompt B"] - - def test_single_turn_empty_conversation_yields_fresh_id(self): - """For single-turn targets with empty conversation, a fresh ID is produced.""" - node = self._make_tap_node(supports_multi_turn=False) - memory = CentralMemory.get_memory_instance() - - # Don't seed any messages - duplicate = node.duplicate() - - assert duplicate.objective_target_conversation_id != node.objective_target_conversation_id - dup_messages = memory.get_conversation_messages(conversation_id=duplicate.objective_target_conversation_id) - assert len(dup_messages) == 0 - - def test_duplicate_node_has_correct_parent_id(self): - """The duplicate node's parent_id should be the original node's node_id.""" - node = self._make_tap_node(supports_multi_turn=False) - - duplicate = node.duplicate() - - assert duplicate.parent_id == node.node_id - assert duplicate.node_id != node.node_id - - def test_duplicate_node_copies_conversation_context(self): - """The duplicate node should inherit the _conversation_context from the original.""" - node = self._make_tap_node(supports_multi_turn=False) - node._conversation_context = "Some prior conversation context" - - duplicate = node.duplicate() - - assert duplicate._conversation_context == "Some prior conversation context" - - def test_system_message_content_preserved_exactly(self): - """The duplicated system message text must match the original exactly.""" - node = self._make_tap_node(supports_multi_turn=False) - memory = CentralMemory.get_memory_instance() - - long_prompt = "You are a helpful assistant.\n\nRules:\n1. Be concise\n2. Be accurate\n\n Special chars: àéîöü" - sys_piece = MessagePiece( - original_value=long_prompt, - role="system", - conversation_id=node.objective_target_conversation_id, - sequence=0, - ) - user_piece = MessagePiece( - original_value="Hello", - role="user", - conversation_id=node.objective_target_conversation_id, - sequence=1, - ) - memory.add_message_pieces_to_memory(message_pieces=[sys_piece, user_piece]) - - duplicate = node.duplicate() - - dup_messages = memory.get_conversation_messages(conversation_id=duplicate.objective_target_conversation_id) - assert len(dup_messages) == 1 - assert dup_messages[0].get_value() == long_prompt - - def test_original_conversation_untouched_after_duplicate(self): - """Duplicating must not alter the original node's conversation.""" - node = self._make_tap_node(supports_multi_turn=False) - memory = CentralMemory.get_memory_instance() - - sys_piece = MessagePiece( - original_value="System prompt", - role="system", - conversation_id=node.objective_target_conversation_id, - sequence=0, - ) - user_piece = MessagePiece( - original_value="User attack", - role="user", - conversation_id=node.objective_target_conversation_id, - sequence=1, - ) - memory.add_message_pieces_to_memory(message_pieces=[sys_piece, user_piece]) - - node.duplicate() - - # Original conversation should still have both messages - orig_messages = memory.get_conversation_messages(conversation_id=node.objective_target_conversation_id) - orig_roles = [m.api_role for m in orig_messages] - assert orig_roles == ["system", "user"] - - def test_single_turn_multipiece_system_message_duplicated(self): - """A multi-piece system Message (same sequence) is fully duplicated in TAP.""" - node = self._make_tap_node(supports_multi_turn=False) - memory = CentralMemory.get_memory_instance() - - sys_text = MessagePiece( - original_value="System text", - role="system", - conversation_id=node.objective_target_conversation_id, - sequence=0, - ) - sys_image = MessagePiece( - original_value="image_data", - role="system", - conversation_id=node.objective_target_conversation_id, - sequence=0, - original_value_data_type="image_path", - ) - user_piece = MessagePiece( - original_value="Attack", - role="user", - conversation_id=node.objective_target_conversation_id, - sequence=1, - ) - memory.add_message_pieces_to_memory(message_pieces=[sys_text, sys_image, user_piece]) - - duplicate = node.duplicate() - - dup_messages = memory.get_conversation_messages(conversation_id=duplicate.objective_target_conversation_id) - assert len(dup_messages) == 1 - assert dup_messages[0].api_role == "system" - assert len(dup_messages[0].message_pieces) == 2 - dup_values = {p.converted_value for p in dup_messages[0].message_pieces} - assert dup_values == {"System text", "image_data"} - @pytest.mark.usefixtures("patch_central_database") -class TestValueErrorGuards: - """Test that incompatible attacks raise ValueError for single-turn targets.""" - - def _make_single_turn_target(self): - from pyrit.prompt_target.common.target_capabilities import TargetCapabilities - from pyrit.prompt_target.common.target_configuration import TargetConfiguration - - target = MagicMock() - target.capabilities.supports_multi_turn = False - target.configuration = TargetConfiguration( - capabilities=TargetCapabilities(supports_multi_turn=False, supports_system_prompt=True), - ) - target.get_identifier.return_value = {"__type__": "MockTarget", "__module__": "test", "id": "mock-id"} - return target - - def _make_adversarial_config(self): - from pyrit.executor.attack.core.attack_config import AttackAdversarialConfig - - adversarial_chat = MagicMock() - adversarial_chat.get_identifier.return_value = {"__type__": "MockTarget", "__module__": "test", "id": "mock-id"} - return AttackAdversarialConfig(target=adversarial_chat) - - def _make_scoring_config(self): - from pyrit.executor.attack.core.attack_config import AttackScoringConfig - from pyrit.score import TrueFalseScorer - - scorer = MagicMock(spec=TrueFalseScorer) - scorer.get_identifier.return_value = {"__type__": "MockTarget", "__module__": "test", "id": "mock-id"} - return AttackScoringConfig(objective_scorer=scorer) - - async def test_crescendo_raises_for_single_turn_target(self): - from pyrit.executor.attack.multi_turn.crescendo import CrescendoAttack - - target = self._make_single_turn_target() - - with pytest.raises(ValueError, match="supports_multi_turn"): - CrescendoAttack( - objective_target=target, - attack_adversarial_config=self._make_adversarial_config(), - attack_scoring_config=self._make_scoring_config(), - ) - - async def test_multi_prompt_sending_raises_for_single_turn_target(self): - from pyrit.executor.attack.multi_turn.multi_prompt_sending import MultiPromptSendingAttack - - target = self._make_single_turn_target() +def test_multi_prompt_sending_requires_native_multi_turn(single_turn_target: MagicMock): + from pyrit.executor.attack.multi_turn.multi_prompt_sending import MultiPromptSendingAttack - with pytest.raises(ValueError, match="supports_multi_turn"): - MultiPromptSendingAttack(objective_target=target) - - async def test_chunked_request_raises_for_single_turn_target(self): - from pyrit.executor.attack.multi_turn.chunked_request import ChunkedRequestAttack - - target = self._make_single_turn_target() - - with pytest.raises(ValueError, match="supports_multi_turn"): - ChunkedRequestAttack(objective_target=target) + with pytest.raises(ValueError, match="supports_multi_turn"): + MultiPromptSendingAttack(objective_target=single_turn_target) @pytest.mark.usefixtures("patch_central_database") -class TestTAPBranchingPreservesSystemPrompts: - """Integration test: TAP branching with real memory verifies system prompt carryover.""" - - def _make_tap_node( - self, - *, - supports_multi_turn: bool, - prepended_conversation_config: PrependedConversationConfig | None = None, - objective_target: PromptTarget | None = None, - ) -> Any: - """Create a _TreeOfAttacksNode with real memory.""" - from pyrit.executor.attack.component.modality_router import _ModalityFeedbackRouter - from pyrit.executor.attack.multi_turn.tree_of_attacks import _TreeOfAttacksNode - - target = objective_target or MagicMock() - if objective_target is None: - target.capabilities.supports_multi_turn = supports_multi_turn - target.configuration.includes.return_value = supports_multi_turn - target.configuration.capabilities.input_modalities = frozenset({frozenset({"text"})}) - target.configuration.capabilities.output_modalities = frozenset({frozenset({"text"})}) - target.get_identifier.return_value = {"__type__": "MockTarget", "__module__": "test", "id": "mock-id"} - - adversarial_chat = MagicMock() - adversarial_chat.get_identifier.return_value = {"__type__": "MockTarget", "__module__": "test", "id": "mock-id"} - adversarial_chat.configuration.capabilities.input_modalities = frozenset({frozenset({"text"})}) - adversarial_chat.configuration.capabilities.output_modalities = frozenset({frozenset({"text"})}) - - scorer = MagicMock() - scorer.get_identifier.return_value = {"__type__": "MockTarget", "__module__": "test", "id": "mock-id"} - - seed = MagicMock() - seed.render_template_value.return_value = "template" - - return _TreeOfAttacksNode( - objective_target=target, - adversarial_chat=adversarial_chat, - adversarial_chat_seed_prompt=seed, - adversarial_chat_prompt_template=seed, - adversarial_chat_system_seed_prompt=seed, - desired_response_prefix="Sure,", - objective_scorer=scorer, - on_topic_scorer=None, - request_converters=[], - response_converters=[], - auxiliary_scorers=None, - attack_id=MagicMock(), - attack_strategy_name="TAP", - modality_router=_ModalityFeedbackRouter( - adversarial_chat=adversarial_chat, - objective_target=target, - ), - prepended_conversation_config=prepended_conversation_config, - ) - - def test_branching_single_turn_target_preserves_system_across_depths(self): - """Simulate TAP branching across 2 depths and verify system prompts survive. - - Depth 1: Create a node, seed system + user + assistant messages. - Depth 2: duplicate() the node (simulating branching). For single-turn targets, - only the system message should be in the duplicate's conversation. - Then simulate another turn on the duplicate (add user + assistant). - Depth 3: duplicate() again. System message should still be there. - """ - memory = CentralMemory.get_memory_instance() - node = self._make_tap_node(supports_multi_turn=False) - - # Depth 1: seed the conversation with system prompt + a completed turn - sys_piece = MessagePiece( - original_value="You are a red team assistant.", - role="system", - conversation_id=node.objective_target_conversation_id, - sequence=0, - ) - user_piece = MessagePiece( - original_value="Tell me about X", - role="user", - conversation_id=node.objective_target_conversation_id, - sequence=1, - ) - asst_piece = MessagePiece( - original_value="Here is info about X", - role="assistant", - conversation_id=node.objective_target_conversation_id, - sequence=2, - ) - memory.add_message_pieces_to_memory(message_pieces=[sys_piece, user_piece, asst_piece]) - - # Depth 2: branch (duplicate) — single-turn means only system msg is copied - branch1 = node.duplicate() - - branch1_msgs = memory.get_conversation_messages(conversation_id=branch1.objective_target_conversation_id) - assert len(branch1_msgs) == 1 - assert branch1_msgs[0].api_role == "system" - assert branch1_msgs[0].get_value() == "You are a red team assistant." - - # Simulate depth-2 turn on branch1: add user + assistant on branch1's conversation - user2 = MessagePiece( - original_value="Now tell me about Y", - role="user", - conversation_id=branch1.objective_target_conversation_id, - sequence=1, - ) - asst2 = MessagePiece( - original_value="Here is info about Y", - role="assistant", - conversation_id=branch1.objective_target_conversation_id, - sequence=2, - ) - memory.add_message_pieces_to_memory(message_pieces=[user2, asst2]) - - # Verify branch1 now has system + user + assistant - branch1_full = memory.get_conversation_messages(conversation_id=branch1.objective_target_conversation_id) - assert [m.api_role for m in branch1_full] == ["system", "user", "assistant"] - - # Depth 3: branch again from branch1 - branch2 = branch1.duplicate() - - branch2_msgs = memory.get_conversation_messages(conversation_id=branch2.objective_target_conversation_id) - assert len(branch2_msgs) == 1 - assert branch2_msgs[0].api_role == "system" - assert branch2_msgs[0].get_value() == "You are a red team assistant." - - async def test_branching_single_turn_target_retains_pending_normalization_context(self): - """A TAP child must send copied system context with the next request.""" - message_normalizer = MagicMock(spec=MessageStringNormalizer) - message_normalizer.normalize_string_async = AsyncMock(return_value="custom formatted request") - target = _SingleTurnPromptTarget() - node = self._make_tap_node( - supports_multi_turn=False, - prepended_conversation_config=PrependedConversationConfig(message_normalizer=message_normalizer), - objective_target=target, - ) - memory = CentralMemory.get_memory_instance() - memory.add_message_pieces_to_memory( - message_pieces=[ - MessagePiece( - original_value="You are a red team assistant.", - role="system", - conversation_id=node.objective_target_conversation_id, - sequence=0, - ) - ] - ) - - branch = node.duplicate() - branch_conversation_id = branch.objective_target_conversation_id - - await branch._send_prompt_to_target_async("next request") - - assert branch.objective_target_conversation_id == branch_conversation_id - assert branch._target_normalization_context is not None - assert branch._target_normalization_context.is_consumed - assert target.normalized_conversations[0][0].get_value() == "custom formatted request" - message_normalizer.normalize_string_async.assert_awaited_once() - - def test_branching_multi_turn_target_preserves_full_history(self): - """For multi-turn targets, branching should preserve the full conversation.""" - memory = CentralMemory.get_memory_instance() - node = self._make_tap_node(supports_multi_turn=True) - - # Seed system + user + assistant - sys_piece = MessagePiece( - original_value="System prompt", - role="system", - conversation_id=node.objective_target_conversation_id, - sequence=0, - ) - user_piece = MessagePiece( - original_value="User message", - role="user", - conversation_id=node.objective_target_conversation_id, - sequence=1, - ) - asst_piece = MessagePiece( - original_value="Assistant response", - role="assistant", - conversation_id=node.objective_target_conversation_id, - sequence=2, - ) - memory.add_message_pieces_to_memory(message_pieces=[sys_piece, user_piece, asst_piece]) - - branch = node.duplicate() - - branch_msgs = memory.get_conversation_messages(conversation_id=branch.objective_target_conversation_id) - assert [m.api_role for m in branch_msgs] == ["system", "user", "assistant"] - - # Add another turn on the branch - user2 = MessagePiece( - original_value="Follow-up", - role="user", - conversation_id=branch.objective_target_conversation_id, - sequence=3, - ) - memory.add_message_pieces_to_memory(message_pieces=[user2]) +def test_chunked_request_requires_native_multi_turn(single_turn_target: MagicMock): + from pyrit.executor.attack.multi_turn.chunked_request import ChunkedRequestAttack - # Branch again — should have all 4 messages - branch2 = branch.duplicate() - branch2_msgs = memory.get_conversation_messages(conversation_id=branch2.objective_target_conversation_id) - assert [m.api_role for m in branch2_msgs] == ["system", "user", "assistant", "user"] + with pytest.raises(ValueError, match="supports_multi_turn"): + ChunkedRequestAttack(objective_target=single_turn_target) diff --git a/tests/unit/executor/attack/multi_turn/test_tree_of_attacks.py b/tests/unit/executor/attack/multi_turn/test_tree_of_attacks.py index cbcf203413..bdc12f3599 100644 --- a/tests/unit/executor/attack/multi_turn/test_tree_of_attacks.py +++ b/tests/unit/executor/attack/multi_turn/test_tree_of_attacks.py @@ -1585,7 +1585,6 @@ def test_node_duplicate_creates_child(self, node_components): """Test that duplicate() creates a proper child node.""" parent_node = _TreeOfAttacksNode(**node_components) parent_node.node_id = "parent_node_id" - parent_node._target_normalization_context = MagicMock() # Mock memory duplicate conversation with patch.object(parent_node._memory, "duplicate_conversation", return_value="new_conv_id"): @@ -1594,7 +1593,6 @@ def test_node_duplicate_creates_child(self, node_components): assert child_node.node_id != parent_node.node_id assert child_node.parent_id == parent_node.node_id assert child_node.completed is False - assert child_node._target_normalization_context is None def _node_with_schema(self, node_components, schema): """Build a real node whose adversarial system prompt advertises ``schema``. @@ -1902,8 +1900,8 @@ def mock_score_response(*args, **kwargs): assert node.auxiliary_scores["AuxScorer2"].get_value() == 0.6 @pytest.mark.asyncio - async def test_node_single_turn_target_generates_new_conv_id(self, node_components): - """Test that single-turn targets get a fresh conversation_id before each send.""" + async def test_node_single_turn_target_keeps_conversation_id(self, node_components): + """Test that target normalization removes the need for conversation rotation.""" node_components["objective_target"].capabilities.supports_multi_turn = False node_components["objective_target"].configuration.includes.side_effect = lambda capability: False node = _TreeOfAttacksNode(**node_components) @@ -1926,8 +1924,7 @@ async def test_node_single_turn_target_generates_new_conv_id(self, node_componen with patch.object(node, "_score_response_async", new_callable=AsyncMock): await node._send_prompt_to_target_async("test prompt") - # Conversation ID should have changed for single-turn target - assert node.objective_target_conversation_id != original_conv_id + assert node.objective_target_conversation_id == original_conv_id @pytest.mark.asyncio async def test_node_multi_turn_target_keeps_conv_id(self, node_components): diff --git a/tests/unit/executor/attack/single_turn/test_prompt_sending.py b/tests/unit/executor/attack/single_turn/test_prompt_sending.py index cd2017e349..3a1da94b59 100644 --- a/tests/unit/executor/attack/single_turn/test_prompt_sending.py +++ b/tests/unit/executor/attack/single_turn/test_prompt_sending.py @@ -281,7 +281,7 @@ async def test_setup_updates_conversation_state_with_converters(self, mock_targe target=mock_target, conversation_id=basic_context.conversation_id, request_converters=converter_config, - prepended_conversation_config=None, + prepended_conversation_config=PrependedConversationConfig(), memory_labels={}, ) @@ -365,7 +365,7 @@ async def test_non_chat_target_converts_history_by_role_before_flattening(self): (f"Turn 1:\nuser: {encoded_user}\nassistant: {simulated_response}\nTurn 2:\nuser: {encoded_final_request}") ] - async def test_retry_setup_creates_fresh_first_turn_context(self, basic_context): + async def test_retry_setup_creates_fresh_conversation(self, basic_context): target = MockPromptTarget() target._configuration = TargetConfiguration(capabilities=TargetCapabilities()) attack = PromptSendingAttack(objective_target=target) @@ -374,18 +374,11 @@ async def test_retry_setup_creates_fresh_first_turn_context(self, basic_context) ] await attack._setup_async(context=basic_context) - first_context = basic_context.target_normalization_context first_conversation_id = basic_context.conversation_id - assert first_context is not None - assert first_context.begin_normalization(conversation_id=first_conversation_id) - first_context.mark_consumed() await attack._setup_async(context=basic_context) assert basic_context.conversation_id != first_conversation_id - assert basic_context.target_normalization_context is not None - assert basic_context.target_normalization_context is not first_context - assert basic_context.target_normalization_context.is_pending @pytest.mark.usefixtures("patch_central_database") diff --git a/tests/unit/prompt_normalizer/test_prompt_normalizer.py b/tests/unit/prompt_normalizer/test_prompt_normalizer.py index b314568a18..ab4952054e 100644 --- a/tests/unit/prompt_normalizer/test_prompt_normalizer.py +++ b/tests/unit/prompt_normalizer/test_prompt_normalizer.py @@ -25,7 +25,7 @@ get_execution_context, ) from pyrit.memory import CentralMemory -from pyrit.message_normalizer import ConversationContextNormalizer, FirstTurnHistoryNormalizer +from pyrit.message_normalizer import MessageListNormalizer from pyrit.models import ( Message, MessagePiece, @@ -37,7 +37,7 @@ from pyrit.prompt_normalizer.converter_configuration import ( ConverterConfiguration, ) -from pyrit.prompt_target import PromptTarget, TargetNormalizationContext +from pyrit.prompt_target import CapabilityName, PromptTarget @pytest.fixture @@ -124,7 +124,7 @@ async def test_send_prompt_async_multiple_converters(mock_memory_instance, seed_ assert prompt_target.prompt_sent == ["S_G_V_s_b_G_8_="] -async def test_send_prompt_async_forwards_target_normalization_context(mock_memory_instance): +async def test_send_prompt_async_forwards_normalizer_overrides(mock_memory_instance): prompt_target = MagicMock(spec=PromptTarget) prompt_target.get_identifier.return_value = get_mock_target_identifier("MockTarget") prompt_target.send_prompt_async = AsyncMock( @@ -132,41 +132,25 @@ async def test_send_prompt_async_forwards_target_normalization_context(mock_memo ) normalizer = PromptNormalizer() conversation_id = "prepended-conversation" - target_context = TargetNormalizationContext( - conversation_id=conversation_id, - normalizers=( - FirstTurnHistoryNormalizer( - message_normalizer=ConversationContextNormalizer(), - prepended_message_count=1, - ), - ), - ) + message_normalizer = MagicMock(spec=MessageListNormalizer) + normalizer_overrides = {CapabilityName.EDITABLE_HISTORY: message_normalizer} await normalizer.send_prompt_async( message=Message.from_prompt(prompt="first request", role="user"), target=prompt_target, conversation_id=conversation_id, - target_normalization_context=target_context, + normalizer_overrides=normalizer_overrides, ) call = prompt_target.send_prompt_async.await_args - assert call.kwargs["target_normalization_context"] is target_context + assert call.kwargs["normalizer_overrides"] == normalizer_overrides -async def test_send_prompt_async_conversion_failure_leaves_context_pending(mock_memory_instance): +async def test_send_prompt_async_conversion_failure_does_not_call_target(mock_memory_instance): prompt_target = MagicMock(spec=PromptTarget) prompt_target.get_identifier.return_value = get_mock_target_identifier("MockTarget") prompt_target.send_prompt_async = AsyncMock() conversation_id = "prepended-conversation" - target_context = TargetNormalizationContext( - conversation_id=conversation_id, - normalizers=( - FirstTurnHistoryNormalizer( - message_normalizer=ConversationContextNormalizer(), - prepended_message_count=1, - ), - ), - ) converter_config = ConverterConfiguration.from_converters(converters=[ContextFailingConverter()]) with pytest.raises(ValueError, match="conversion failed"): @@ -175,66 +159,41 @@ async def test_send_prompt_async_conversion_failure_leaves_context_pending(mock_ target=prompt_target, conversation_id=conversation_id, request_converter_configurations=converter_config, - target_normalization_context=target_context, ) - assert target_context.is_pending prompt_target.send_prompt_async.assert_not_awaited() mock_memory_instance.add_message_to_memory.assert_not_called() -async def test_send_prompt_async_pre_provider_failure_is_not_persisted(mock_memory_instance): +async def test_send_prompt_async_target_failure_is_persisted(mock_memory_instance): prompt_target = MagicMock(spec=PromptTarget) prompt_target.get_identifier.return_value = get_mock_target_identifier("MockTarget") prompt_target.send_prompt_async = AsyncMock(side_effect=ValueError("normalization failed")) conversation_id = "prepended-conversation" - target_context = TargetNormalizationContext( - conversation_id=conversation_id, - normalizers=( - FirstTurnHistoryNormalizer( - message_normalizer=ConversationContextNormalizer(), - prepended_message_count=1, - ), - ), - ) with pytest.raises(Exception, match="Error sending prompt with conversation ID"): await PromptNormalizer().send_prompt_async( message=Message.from_prompt(prompt="request", role="user"), target=prompt_target, conversation_id=conversation_id, - target_normalization_context=target_context, ) - assert target_context.is_pending - mock_memory_instance.add_message_to_memory.assert_not_called() + assert mock_memory_instance.add_message_to_memory.call_count == 2 -async def test_send_prompt_async_pre_provider_empty_response_is_not_persisted(mock_memory_instance): +async def test_send_prompt_async_empty_response_exception_is_persisted(mock_memory_instance): prompt_target = MagicMock(spec=PromptTarget) prompt_target.get_identifier.return_value = get_mock_target_identifier("MockTarget") prompt_target.send_prompt_async = AsyncMock(side_effect=EmptyResponseException(message="normalization failed")) conversation_id = "prepended-conversation" - target_context = TargetNormalizationContext( + response = await PromptNormalizer().send_prompt_async( + message=Message.from_prompt(prompt="request", role="user"), + target=prompt_target, conversation_id=conversation_id, - normalizers=( - FirstTurnHistoryNormalizer( - message_normalizer=ConversationContextNormalizer(), - prepended_message_count=1, - ), - ), ) - with pytest.raises(Exception, match="Error sending prompt with conversation ID"): - await PromptNormalizer().send_prompt_async( - message=Message.from_prompt(prompt="request", role="user"), - target=prompt_target, - conversation_id=conversation_id, - target_normalization_context=target_context, - ) - - assert target_context.is_pending - mock_memory_instance.add_message_to_memory.assert_not_called() + assert response.get_piece().response_error == "empty" + assert mock_memory_instance.add_message_to_memory.call_count == 2 async def test_send_prompt_async_no_response_adds_memory(mock_memory_instance, seed_group): diff --git a/tests/unit/prompt_target/target/test_image_target.py b/tests/unit/prompt_target/target/test_image_target.py index f214439694..b9d42f5e81 100644 --- a/tests/unit/prompt_target/target/test_image_target.py +++ b/tests/unit/prompt_target/target/test_image_target.py @@ -486,7 +486,7 @@ async def test_validate_piece_type(image_target: OpenAIImageTarget): os.remove(audio_piece.original_value) -async def test_validate_previous_conversations( +async def test_adapt_previous_conversations( image_target: OpenAIImageTarget, sample_conversations: MutableSequence[MessagePiece] ): message_piece = sample_conversations[0] @@ -501,12 +501,10 @@ async def test_validate_previous_conversations( request = Message(message_pieces=[message_piece]) - with pytest.raises( - ValueError, - match="This target only supports a single turn conversation.*If your target does support this, set the" - " custom_configuration parameter accordingly", - ): - await image_target.send_prompt_async(message=request) + normalized = await image_target._get_normalized_conversation_async(message=request) + + assert len(normalized) == 1 + assert message_piece.converted_value in normalized[0].get_value() def test_background_param_stored(patch_central_database): diff --git a/tests/unit/prompt_target/target/test_normalize_async_integration.py b/tests/unit/prompt_target/target/test_normalize_async_integration.py index d8e4feab9d..a21d83b3c4 100644 --- a/tests/unit/prompt_target/target/test_normalize_async_integration.py +++ b/tests/unit/prompt_target/target/test_normalize_async_integration.py @@ -24,7 +24,7 @@ TokenizerTemplateNormalizer, ) from pyrit.models import ComponentIdentifier, Message, MessagePiece -from pyrit.prompt_target import AzureMLChatTarget, OpenAIChatTarget, TargetNormalizationContext +from pyrit.prompt_target import AzureMLChatTarget, OpenAIChatTarget from pyrit.prompt_target.common.target_capabilities import ( CapabilityHandlingPolicy, CapabilityName, @@ -50,21 +50,17 @@ def _make_message(*, role: str, content: str, conversation_id: str = "conv1") -> return Message(message_pieces=[_make_message_piece(role=role, content=content, conversation_id=conversation_id)]) -def _make_target_normalization_context( +def _make_normalizer_overrides( *, - prepended_message_count: int, formatter: MessageStringNormalizer | None = None, - conversation_id: str = "conv1", -) -> TargetNormalizationContext: - return TargetNormalizationContext( - conversation_id=conversation_id, - normalizers=( - FirstTurnHistoryNormalizer( - message_normalizer=formatter or ConversationContextNormalizer(), - prepended_message_count=prepended_message_count, - ), - ), - ) + target_supports_multi_turn: bool = False, +) -> dict[CapabilityName, FirstTurnHistoryNormalizer]: + return { + CapabilityName.EDITABLE_HISTORY: FirstTurnHistoryNormalizer( + message_normalizer=formatter or ConversationContextNormalizer(), + target_supports_multi_turn=target_supports_multi_turn, + ) + } def _create_mock_chat_completion(content: str = "hi") -> MagicMock: @@ -539,13 +535,11 @@ async def test_non_editable_target_adapts_prepended_history_without_mutating_mem mock_memory = MagicMock(spec=MemoryInterface) mock_memory.get_conversation_messages.return_value = memory_messages target._memory = mock_memory - target_context = _make_target_normalization_context(prepended_message_count=2) - assert target_context.begin_normalization(conversation_id="conv1") + normalizer_overrides = _make_normalizer_overrides() result = await target._get_normalized_conversation_async( message=live_request, - target_normalization_context=target_context, - should_apply_context=True, + normalizer_overrides=normalizer_overrides, ) assert len(result) == 1 @@ -584,13 +578,11 @@ async def test_non_editable_target_preserves_system_history_and_multimodal_live_ mock_memory = MagicMock(spec=MemoryInterface) mock_memory.get_conversation_messages.return_value = [system_message] target._memory = mock_memory - target_context = _make_target_normalization_context(prepended_message_count=1) - assert target_context.begin_normalization(conversation_id="conv1") + normalizer_overrides = _make_normalizer_overrides() result = await target._get_normalized_conversation_async( message=live_request, - target_normalization_context=target_context, - should_apply_context=True, + normalizer_overrides=normalizer_overrides, ) assert len(result) == 1 @@ -634,13 +626,11 @@ async def test_first_turn_normalization_preserves_live_multimodal_piece_order(): mock_memory = MagicMock(spec=MemoryInterface) mock_memory.get_conversation_messages.return_value = [_make_message(role="user", content="prepended")] target._memory = mock_memory - target_context = _make_target_normalization_context(prepended_message_count=1) - assert target_context.begin_normalization(conversation_id="conv1") + normalizer_overrides = _make_normalizer_overrides() result = await target._get_normalized_conversation_async( message=live_request, - target_normalization_context=target_context, - should_apply_context=True, + normalizer_overrides=normalizer_overrides, ) assert [piece.converted_value_data_type for piece in result[0].message_pieces] == ["image_path", "text"] @@ -668,23 +658,22 @@ async def test_prepended_history_adapter_is_used_only_when_explicitly_passed(): [prepended, prior_live, prior_response], ] target._memory = mock_memory - target_context = _make_target_normalization_context(prepended_message_count=1) + normalizer_overrides = _make_normalizer_overrides(target_supports_multi_turn=True) await target.send_prompt_async( message=prior_live, - target_normalization_context=target_context, + normalizer_overrides=normalizer_overrides, ) await target.send_prompt_async( message=second_live, - target_normalization_context=target_context, + normalizer_overrides=normalizer_overrides, ) assert target.prompt_sent == ["Turn 1:\nuser: prepended\nTurn 2:\nuser: first live", "second live"] - assert target_context.is_consumed @pytest.mark.usefixtures("patch_central_database") -async def test_consumed_context_sends_only_the_current_target_facing_request(): +async def test_non_editable_multi_turn_target_sends_only_current_request_after_response(): target = MockPromptTarget() target._configuration = TargetConfiguration( capabilities=TargetCapabilities( @@ -705,15 +694,15 @@ async def test_consumed_context_sends_only_the_current_target_facing_request(): target._send_prompt_to_target_async = AsyncMock( # type: ignore[method-assign] return_value=[_make_message(role="assistant", content="response")] ) - target_context = _make_target_normalization_context(prepended_message_count=1) + normalizer_overrides = _make_normalizer_overrides(target_supports_multi_turn=True) await target.send_prompt_async( message=first_live, - target_normalization_context=target_context, + normalizer_overrides=normalizer_overrides, ) await target.send_prompt_async( message=second_live, - target_normalization_context=target_context, + normalizer_overrides=normalizer_overrides, ) first_payload, second_payload = target._send_prompt_to_target_async.await_args_list @@ -731,16 +720,13 @@ async def test_non_editable_target_uses_custom_prepended_formatter(): target._memory = mock_memory formatter = MagicMock(spec=MessageStringNormalizer) formatter.normalize_string_async = AsyncMock(return_value="CUSTOM HISTORY") - target_context = _make_target_normalization_context( - prepended_message_count=1, + normalizer_overrides = _make_normalizer_overrides( formatter=formatter, ) - assert target_context.begin_normalization(conversation_id="conv1") result = await target._get_normalized_conversation_async( message=_make_message(role="user", content="live"), - target_normalization_context=target_context, - should_apply_context=True, + normalizer_overrides=normalizer_overrides, ) assert result[0].get_value() == "CUSTOM HISTORY" @@ -760,14 +746,13 @@ async def test_non_editable_target_rejects_non_text_converted_prepended_history( mock_memory = MagicMock(spec=MemoryInterface) mock_memory.get_conversation_messages.return_value = [prepended] target._memory = mock_memory - target_context = _make_target_normalization_context(prepended_message_count=1) + normalizer_overrides = _make_normalizer_overrides() with pytest.raises(ValueError, match="non-text output types.*image_path"): await target.send_prompt_async( message=_make_message(role="user", content="live"), - target_normalization_context=target_context, + normalizer_overrides=normalizer_overrides, ) - assert target_context.is_pending @pytest.mark.usefixtures("patch_central_database") @@ -789,16 +774,14 @@ async def test_non_editable_target_rejects_same_modality_non_text_conversion(): mock_memory = MagicMock(spec=MemoryInterface) mock_memory.get_conversation_messages.return_value = [prepended] target._memory = mock_memory - target_context = _make_target_normalization_context(prepended_message_count=1) + normalizer_overrides = _make_normalizer_overrides() with pytest.raises(ValueError, match="non-text output types.*image_path"): await target.send_prompt_async( message=_make_message(role="user", content="live"), - target_normalization_context=target_context, + normalizer_overrides=normalizer_overrides, ) - assert target_context.is_pending - @pytest.mark.usefixtures("patch_central_database") async def test_non_editable_target_allows_preexisting_non_text_history_with_converter_provenance(): @@ -820,13 +803,11 @@ async def test_non_editable_target_allows_preexisting_non_text_history_with_conv mock_memory = MagicMock(spec=MemoryInterface) mock_memory.get_conversation_messages.return_value = [prepended] target._memory = mock_memory - target_context = _make_target_normalization_context(prepended_message_count=1) - assert target_context.begin_normalization(conversation_id="conv1") + normalizer_overrides = _make_normalizer_overrides() result = await target._get_normalized_conversation_async( message=_make_message(role="user", content="live"), - target_normalization_context=target_context, - should_apply_context=True, + normalizer_overrides=normalizer_overrides, ) assert result[0].get_value() == "Turn 1:\nuser: [Image_path]\nTurn 2:\nuser: live" @@ -841,8 +822,7 @@ async def test_target_normalization_failure_can_be_retried(): target._memory = mock_memory formatter = MagicMock(spec=MessageStringNormalizer) formatter.normalize_string_async = AsyncMock(side_effect=[ValueError("format failed"), "formatted request"]) - target_context = _make_target_normalization_context( - prepended_message_count=1, + normalizer_overrides = _make_normalizer_overrides( formatter=formatter, ) live_request = _make_message(role="user", content="live") @@ -850,20 +830,51 @@ async def test_target_normalization_failure_can_be_retried(): with pytest.raises(ValueError, match="format failed"): await target.send_prompt_async( message=live_request, - target_normalization_context=target_context, + normalizer_overrides=normalizer_overrides, ) - assert target_context.is_pending await target.send_prompt_async( message=live_request, - target_normalization_context=target_context, + normalizer_overrides=normalizer_overrides, ) - assert target_context.is_consumed assert target.prompt_sent == ["formatted request"] @pytest.mark.usefixtures("patch_central_database") -async def test_target_normalization_cancellation_restores_pending_state(): +async def test_retry_ignores_persisted_failed_request_and_error_response(): + target = MockPromptTarget() + target._configuration = TargetConfiguration( + capabilities=TargetCapabilities( + supports_multi_turn=True, + supports_system_prompt=True, + ) + ) + prepended = _make_message(role="user", content="prepended") + failed_request = _make_message(role="user", content="failed request") + error_response = _make_message(role="assistant", content="processing error") + error_response.get_piece().response_error = "processing" + mock_memory = MagicMock(spec=MemoryInterface) + mock_memory.get_conversation_messages.return_value = [prepended, failed_request, error_response] + target._memory = mock_memory + formatter = MagicMock(spec=MessageStringNormalizer) + formatter.normalize_string_async = AsyncMock(return_value="formatted retry") + normalizer_overrides = _make_normalizer_overrides( + formatter=formatter, + target_supports_multi_turn=True, + ) + + result = await target._get_normalized_conversation_async( + message=_make_message(role="user", content="retry"), + normalizer_overrides=normalizer_overrides, + ) + + assert result[0].get_value() == "formatted retry" + formatted_messages = formatter.normalize_string_async.await_args.args[0] + assert [message.get_value() for message in formatted_messages] == ["prepended", "retry"] + + +@pytest.mark.usefixtures("patch_central_database") +async def test_target_normalization_cancellation_propagates(): target = MockPromptTarget() target._configuration = TargetConfiguration(capabilities=TargetCapabilities()) mock_memory = MagicMock(spec=MemoryInterface) @@ -871,41 +882,36 @@ async def test_target_normalization_cancellation_restores_pending_state(): target._memory = mock_memory formatter = MagicMock(spec=MessageStringNormalizer) formatter.normalize_string_async = AsyncMock(side_effect=asyncio.CancelledError()) - target_context = _make_target_normalization_context( - prepended_message_count=1, + normalizer_overrides = _make_normalizer_overrides( formatter=formatter, ) with pytest.raises(asyncio.CancelledError): await target.send_prompt_async( message=_make_message(role="user", content="live"), - target_normalization_context=target_context, + normalizer_overrides=normalizer_overrides, ) - assert target_context.is_pending - @pytest.mark.usefixtures("patch_central_database") -async def test_provider_failure_leaves_target_normalization_context_consumed(): +async def test_provider_failure_propagates_after_normalization(): target = MockPromptTarget() target._configuration = TargetConfiguration(capabilities=TargetCapabilities()) mock_memory = MagicMock(spec=MemoryInterface) mock_memory.get_conversation_messages.return_value = [_make_message(role="user", content="prepended")] target._memory = mock_memory target._send_prompt_to_target_async = AsyncMock(side_effect=RuntimeError("provider failed")) # type: ignore[method-assign] - target_context = _make_target_normalization_context(prepended_message_count=1) + normalizer_overrides = _make_normalizer_overrides() with pytest.raises(RuntimeError, match="provider failed"): await target.send_prompt_async( message=_make_message(role="user", content="live"), - target_normalization_context=target_context, + normalizer_overrides=normalizer_overrides, ) - assert target_context.is_consumed - @pytest.mark.usefixtures("patch_central_database") -async def test_concurrent_first_sends_do_not_both_apply_prepended_history(): +async def test_concurrent_sends_apply_stateless_normalizer_independently(): target = MockPromptTarget() target._configuration = TargetConfiguration(capabilities=TargetCapabilities()) mock_memory = MagicMock(spec=MemoryInterface) @@ -921,30 +927,27 @@ async def wait_to_format(messages: list[Message]) -> str: formatter = MagicMock(spec=MessageStringNormalizer) formatter.normalize_string_async = AsyncMock(side_effect=wait_to_format) - target_context = _make_target_normalization_context( - prepended_message_count=1, + normalizer_overrides = _make_normalizer_overrides( formatter=formatter, ) first_send = asyncio.create_task( target.send_prompt_async( message=_make_message(role="user", content="first"), - target_normalization_context=target_context, + normalizer_overrides=normalizer_overrides, ) ) await started.wait() + second_send = asyncio.create_task( + target.send_prompt_async( + message=_make_message(role="user", content="second"), + normalizer_overrides=normalizer_overrides, + ) + ) + release.set() - try: - with pytest.raises(RuntimeError, match="already in progress"): - await target.send_prompt_async( - message=_make_message(role="user", content="second"), - target_normalization_context=target_context, - ) - finally: - release.set() - - await first_send - assert target.prompt_sent == ["formatted request"] - assert target_context.is_consumed + await asyncio.gather(first_send, second_send) + assert target.prompt_sent == ["formatted request", "formatted request"] + assert formatter.normalize_string_async.await_count == 2 @pytest.mark.usefixtures("patch_central_database") @@ -957,16 +960,13 @@ async def test_tokenizer_formatter_receives_live_request_before_generation_promp tokenizer = MagicMock() tokenizer.apply_chat_template.return_value = "TOKENIZED REQUEST" formatter = TokenizerTemplateNormalizer(tokenizer=tokenizer) - target_context = _make_target_normalization_context( - prepended_message_count=1, + normalizer_overrides = _make_normalizer_overrides( formatter=formatter, ) - assert target_context.begin_normalization(conversation_id="conv1") result = await target._get_normalized_conversation_async( message=_make_message(role="user", content="live"), - target_normalization_context=target_context, - should_apply_context=True, + normalizer_overrides=normalizer_overrides, ) tokenizer_messages = tokenizer.apply_chat_template.call_args.args[0] diff --git a/tests/unit/prompt_target/target/test_prompt_target_azure_blob_storage.py b/tests/unit/prompt_target/target/test_prompt_target_azure_blob_storage.py index 4feb350bb4..86ca212624 100644 --- a/tests/unit/prompt_target/target/test_prompt_target_azure_blob_storage.py +++ b/tests/unit/prompt_target/target/test_prompt_target_azure_blob_storage.py @@ -96,23 +96,18 @@ async def test_azure_blob_storage_validate_prompt_type( await azure_blob_storage_target.send_prompt_async(message=request) -@patch("azure.storage.blob.aio.ContainerClient.upload_blob") async def test_azure_blob_storage_validate_prev_convs( - mock_upload_async, azure_blob_storage_target: AzureBlobStorageTarget, sample_entries: MutableSequence[MessagePiece], ): - mock_upload_async.return_value = None message_piece = sample_entries[0] azure_blob_storage_target._memory.add_message_to_memory(request=Message(message_pieces=[message_piece])) request = Message(message_pieces=[message_piece]) - with pytest.raises( - ValueError, - match="This target only supports a single turn conversation.*If your target does support this, set the" - " custom_configuration parameter accordingly", - ): - await azure_blob_storage_target.send_prompt_async(message=request) + normalized = await azure_blob_storage_target._get_normalized_conversation_async(message=request) + + assert len(normalized) == 1 + assert message_piece.converted_value in normalized[0].get_value() @patch.object(AzureBlobStorageTarget, "_create_container_client_async", new_callable=AsyncMock) diff --git a/tests/unit/prompt_target/target/test_target_capabilities.py b/tests/unit/prompt_target/target/test_target_capabilities.py index 30c47ce1ca..9aa40322d9 100644 --- a/tests/unit/prompt_target/target/test_target_capabilities.py +++ b/tests/unit/prompt_target/target/test_target_capabilities.py @@ -35,6 +35,7 @@ def test_capability_handling_policy_defaults(self): policy = CapabilityHandlingPolicy() assert policy.behaviors == { CapabilityName.MULTI_TURN: UnsupportedCapabilityBehavior.RAISE, + CapabilityName.EDITABLE_HISTORY: UnsupportedCapabilityBehavior.ADAPT, CapabilityName.SYSTEM_PROMPT: UnsupportedCapabilityBehavior.RAISE, CapabilityName.JSON_SCHEMA: UnsupportedCapabilityBehavior.ADAPT, } @@ -62,6 +63,7 @@ def test_capability_handling_policy_get_behavior_for_all_default_keys(self): policy = CapabilityHandlingPolicy() expected = { CapabilityName.MULTI_TURN: UnsupportedCapabilityBehavior.RAISE, + CapabilityName.EDITABLE_HISTORY: UnsupportedCapabilityBehavior.ADAPT, CapabilityName.SYSTEM_PROMPT: UnsupportedCapabilityBehavior.RAISE, CapabilityName.JSON_SCHEMA: UnsupportedCapabilityBehavior.ADAPT, } @@ -71,11 +73,11 @@ def test_capability_handling_policy_get_behavior_for_all_default_keys(self): def test_capability_handling_policy_rejects_capability_without_policy(self): policy = CapabilityHandlingPolicy() - with pytest.raises(KeyError, match="No policy for capability 'supports_editable_history'"): - policy.get_behavior(capability=CapabilityName.EDITABLE_HISTORY) + with pytest.raises(KeyError, match="No policy for capability 'supports_multi_message_pieces'"): + policy.get_behavior(capability=CapabilityName.MULTI_MESSAGE_PIECES) - with pytest.raises(AttributeError, match="supports_editable_history"): - _ = policy.supports_editable_history + with pytest.raises(AttributeError, match="supports_multi_message_pieces"): + _ = policy.supports_multi_message_pieces def test_capability_handling_policy_rejects_unknown_attribute(self): policy = CapabilityHandlingPolicy() @@ -88,6 +90,7 @@ def test_normalizable_capabilities(self): frozenset( { CapabilityName.MULTI_TURN, + CapabilityName.EDITABLE_HISTORY, CapabilityName.SYSTEM_PROMPT, CapabilityName.JSON_SCHEMA, } diff --git a/tests/unit/prompt_target/target/test_target_configuration.py b/tests/unit/prompt_target/target/test_target_configuration.py index 36f0c62fc5..882590f5ca 100644 --- a/tests/unit/prompt_target/target/test_target_configuration.py +++ b/tests/unit/prompt_target/target/test_target_configuration.py @@ -47,14 +47,22 @@ def _make(role: ChatMessageRole, content: str) -> Message: def test_init_with_defaults_uses_raise_policy(): - caps = TargetCapabilities(supports_multi_turn=True, supports_system_prompt=True) + caps = TargetCapabilities( + supports_multi_turn=True, + supports_editable_history=True, + supports_system_prompt=True, + ) config = TargetConfiguration(capabilities=caps) # Default policy is RAISE for all adaptable capabilities assert config.policy.get_behavior(capability=CapabilityName.MULTI_TURN) == UnsupportedCapabilityBehavior.RAISE def test_init_with_explicit_policy(adapt_all_policy): - caps = TargetCapabilities(supports_multi_turn=True, supports_system_prompt=True) + caps = TargetCapabilities( + supports_multi_turn=True, + supports_editable_history=True, + supports_system_prompt=True, + ) config = TargetConfiguration(capabilities=caps, policy=adapt_all_policy) assert config.policy is adapt_all_policy @@ -92,7 +100,11 @@ def test_init_missing_capability_raise_policy_skips_normalizer(): def test_init_missing_json_schema_default_policy_adds_normalizer(): # Default policy adapts JSON_SCHEMA; a target lacking native support gets the JSON-schema normalizer. - caps = TargetCapabilities(supports_multi_turn=True, supports_system_prompt=True) + caps = TargetCapabilities( + supports_multi_turn=True, + supports_editable_history=True, + supports_system_prompt=True, + ) config = TargetConfiguration(capabilities=caps) assert len(config.pipeline.normalizers) == 1 assert isinstance(config.pipeline.normalizers[0], JsonSchemaNormalizer) @@ -101,6 +113,7 @@ def test_init_missing_json_schema_default_policy_adds_normalizer(): def test_init_supports_json_schema_no_normalizer(): caps = TargetCapabilities( supports_multi_turn=True, + supports_editable_history=True, supports_system_prompt=True, supports_json_schema=True, ) @@ -233,10 +246,10 @@ def test_ensure_can_handle_raises_when_capability_missing_from_policy(): def test_ensure_can_handle_raises_valueerror_for_non_normalizable_capability(): - caps = TargetCapabilities(supports_multi_turn=True, supports_system_prompt=True, supports_editable_history=False) + caps = TargetCapabilities(supports_multi_turn=True, supports_system_prompt=True) config = TargetConfiguration(capabilities=caps) with pytest.raises(ValueError, match="no handling policy"): - config.ensure_can_handle(capability=CapabilityName.EDITABLE_HISTORY) + config.ensure_can_handle(capability=CapabilityName.MULTI_MESSAGE_PIECES) # --------------------------------------------------------------------------- diff --git a/tests/unit/prompt_target/target/test_target_normalization_context.py b/tests/unit/prompt_target/target/test_target_normalization_context.py deleted file mode 100644 index 304a82d709..0000000000 --- a/tests/unit/prompt_target/target/test_target_normalization_context.py +++ /dev/null @@ -1,60 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# Licensed under the MIT license. - -from unittest.mock import AsyncMock, MagicMock - -import pytest - -from pyrit.message_normalizer import MessageListNormalizer -from pyrit.models import Message -from pyrit.prompt_target import TargetNormalizationContext, TargetNormalizationContextState - - -def _make_context() -> tuple[TargetNormalizationContext, MagicMock]: - normalizer = MagicMock(spec=MessageListNormalizer) - normalizer.normalize_async = AsyncMock(side_effect=lambda messages: messages) - context = TargetNormalizationContext( - conversation_id="conversation", - normalizers=(normalizer,), - ) - return context, normalizer - - -async def test_context_normalizes_and_is_consumed_once(): - context, normalizer = _make_context() - messages = [Message.from_prompt(prompt="request", role="user")] - - assert context.begin_normalization(conversation_id="conversation") - assert context.state == TargetNormalizationContextState.PREPARING - assert await context.normalize_async(messages=messages) == messages - context.mark_consumed() - - assert context.is_consumed - assert context.begin_normalization(conversation_id="conversation") is False - normalizer.normalize_async.assert_awaited_once_with(messages) - - -def test_context_can_restore_pending_after_preparation_failure(): - context, _ = _make_context() - - assert context.begin_normalization(conversation_id="conversation") - context.restore_pending() - - assert context.is_pending - assert context.begin_normalization(conversation_id="conversation") - - -def test_context_rejects_concurrent_preparation(): - context, _ = _make_context() - - assert context.begin_normalization(conversation_id="conversation") - - with pytest.raises(RuntimeError, match="already in progress"): - context.begin_normalization(conversation_id="conversation") - - -def test_context_rejects_another_conversation(): - context, _ = _make_context() - - with pytest.raises(ValueError, match="belongs to conversation"): - context.begin_normalization(conversation_id="other") diff --git a/tests/unit/prompt_target/target/test_target_requirements.py b/tests/unit/prompt_target/target/test_target_requirements.py index 9c34ce326e..adb896ef6a 100644 --- a/tests/unit/prompt_target/target/test_target_requirements.py +++ b/tests/unit/prompt_target/target/test_target_requirements.py @@ -36,7 +36,8 @@ def test_construction_from_frozenset(): def test_chat_target_requirements_shape(): - assert CHAT_TARGET_REQUIREMENTS.required == { + assert CHAT_TARGET_REQUIREMENTS.required == set() + assert CHAT_TARGET_REQUIREMENTS.native_required == { CapabilityName.EDITABLE_HISTORY, CapabilityName.MULTI_TURN, } @@ -62,8 +63,7 @@ def test_validate_passes_on_native_support(): def test_validate_passes_when_policy_is_adapt(): - # Note: EDITABLE_HISTORY is not adaptable, so this test uses a custom - # requirement over capabilities that the policy can adapt. + # Use a custom requirement over capabilities that the policy can adapt. reqs = TargetRequirements(required=frozenset({CapabilityName.MULTI_TURN, CapabilityName.SYSTEM_PROMPT})) target = _make_target( configuration=TargetConfiguration( diff --git a/tests/unit/prompt_target/target/test_tts_target.py b/tests/unit/prompt_target/target/test_tts_target.py index cfea0175ec..12350356f0 100644 --- a/tests/unit/prompt_target/target/test_tts_target.py +++ b/tests/unit/prompt_target/target/test_tts_target.py @@ -82,7 +82,7 @@ async def test_tts_validate_prompt_type(tts_target: OpenAITTSTarget): await tts_target.send_prompt_async(message=request) -async def test_tts_validate_previous_conversations( +async def test_tts_adapts_previous_conversations( tts_target: OpenAITTSTarget, sample_conversations: MutableSequence[MessagePiece] ): message_piece = sample_conversations[0] @@ -97,14 +97,10 @@ async def test_tts_validate_previous_conversations( request = Message(message_pieces=[message_piece]) - with patch("pyrit.common.net_utility.make_request_and_raise_if_error_async") as mock_request: - mock_request.return_value = MagicMock(content=b"audio data") - with pytest.raises( - ValueError, - match="This target only supports a single turn conversation.*If your target does support this, set the" - " custom_configuration parameter accordingly", - ): - await tts_target.send_prompt_async(message=request) + normalized = await tts_target._get_normalized_conversation_async(message=request) + + assert len(normalized) == 1 + assert message_piece.converted_value in normalized[0].get_value() @pytest.mark.parametrize("response_format", ["mp3", "ogg"]) From 913197b0fbe936597d58a694c9c52ade78171b12 Mon Sep 17 00:00:00 2001 From: Copilot <223556219+Copilot@users.noreply.github.com> Date: Fri, 21 Aug 2026 10:07:30 -0700 Subject: [PATCH 18/24] Fix TAP seed boundary replay Keep cloned and retained stateless TAP branches pinned to their original prepended seed, correlate failed exchanges using guarded persisted adjacency, and make normalization ownership explicit. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> Copilot-Session: cb50fd66-f0e0-4435-a71d-1d42fef7d307 --- doc/code/framework.md | 4 +- .../attack/multi_turn/tree_of_attacks.py | 39 +-- pyrit/executor/attack/streaming/barge_in.py | 8 +- .../message_normalizer/message_normalizer.py | 6 + pyrit/prompt_target/common/prompt_target.py | 45 +--- .../common/target_normalization_context.py | 111 +++++++-- .../test_prepended_history_normalization.py | 235 +++++++++++++----- .../attack/streaming/test_barge_in.py | 13 +- .../test_normalize_async_integration.py | 179 ++++++++----- .../target/test_prompt_target.py | 66 +++-- .../test_target_normalization_context.py | 91 ++++++- 11 files changed, 534 insertions(+), 263 deletions(-) diff --git a/doc/code/framework.md b/doc/code/framework.md index 798af25e42..55d568b55d 100644 --- a/doc/code/framework.md +++ b/doc/code/framework.md @@ -329,10 +329,12 @@ The below talks about responsibilities of most modules in the PyRIT library The stateful/stateless distinction is therefore a matter of lifecycle and configuration, not a separate flattening implementation. The prepended-conversation flow records the persisted seed message IDs in a `TargetNormalizationContext` and supplies a per-send `HistorySquashNormalizer` override while that explicit history should be included. A target's ordinary capability pipeline independently applies `HistorySquashNormalizer` only when its `MULTI_TURN` behavior is explicitly configured for adaptation. Both uses preserve non-text pieces from the current request while rendering historical content as text. -For prepended history on a target without editable history, processing order is: convert and persist the structured prepended messages, convert the live request, apply the per-send editable-history override, run the target's ordinary capability normalizers, then serialize and invoke the provider. Stateful targets consume the seed boundary when provider invocation begins and send only the current request afterward. Stateless targets replay the original seed with each current request. Pre-provider failures leave the seed reusable, while provider-attempt failures do not replay it for stateful targets. Failed request/error exchanges are excluded from later target-facing history. A concurrent send sharing an active normalization context is rejected so stateful bootstrap context cannot be duplicated. +For prepended history on a target without editable history, processing order is: convert and persist the structured prepended messages, convert the live request, apply the per-send editable-history override, run the target's ordinary capability normalizers, then serialize and invoke the provider. Stateful targets consume the seed boundary when provider invocation begins and send only the current request afterward. Stateless targets replay the original seed with each current request. Unseeded stateless targets send only the current request. Branching attacks remap only the original seed boundary when duplicating memory conversations; copied live turns never become seed history. Pre-provider failures leave the seed reusable, while provider-attempt failures do not replay it for stateful targets. Failed `processing` and `unknown` responses are excluded from later target-facing history together with an immediately preceding user request only when it has the same conversation and adjacent sequence. `blocked` and `empty` responses remain replayable provider turns. A concurrent send sharing an active normalization context is rejected so stateful bootstrap context cannot be duplicated. The prepended-history formatter belongs to the attack configuration rather than the target instance: it affects only sends owned by that attack. Its behavioral identifier is therefore recorded as a child of the attack identifier, including formatter-specific configuration such as tokenizer templates and system-message handling. +Streaming attacks that bypass `PromptTarget.send_prompt_async` do not use target-normalization overrides. `BargeInAttack` persists prepended messages and honors converter-role selection, but its direct realtime session uses only the leading system prompt and live audio; the prepended-history formatter and target-normalization context are intentionally inactive. + ## [Output](./output/0_output) **Responsibility**: The Output module is responsible for writing different components in different formats to different places. diff --git a/pyrit/executor/attack/multi_turn/tree_of_attacks.py b/pyrit/executor/attack/multi_turn/tree_of_attacks.py index feb5d662bf..983862c16f 100644 --- a/pyrit/executor/attack/multi_turn/tree_of_attacks.py +++ b/pyrit/executor/attack/multi_turn/tree_of_attacks.py @@ -749,19 +749,6 @@ def _rotate_unseeded_single_turn_conversation(self) -> None: ): self.objective_target_conversation_id = str(uuid.uuid4()) - def _refresh_stateless_branch_boundary(self) -> None: - """Use the complete current branch as the next stateless target payload boundary.""" - if self._objective_target.configuration.includes(capability=CapabilityName.MULTI_TURN): - return - - messages = list(self._memory.get_conversation_messages(conversation_id=self.objective_target_conversation_id)) - replayable_messages = ConversationManager.get_persistable_prepended_messages(prepended_conversation=messages) - self._target_normalization_context = ConversationManager.create_target_normalization_context( - target=self._objective_target, - conversation_id=self.objective_target_conversation_id, - prepended_messages=replayable_messages, - ) - async def _score_response_async(self, *, response: Message, objective: str) -> None: """ Score the response from the objective target using the configured scorers. @@ -937,22 +924,23 @@ def duplicate(self) -> _TreeOfAttacksNode: prepended_conversation_config=self._prepended_conversation_config, ) - # Duplicate the complete logical branch. A non-editable target receives - # this explicit branch boundary once (stateful) or on every send - # (stateless), without inferring it from assistant/error rows. + source_messages = list( + self._memory.get_conversation_messages(conversation_id=self.objective_target_conversation_id) + ) duplicate_node.objective_target_conversation_id = self._memory.duplicate_conversation( conversation_id=self.objective_target_conversation_id ) duplicated_messages = list( self._memory.get_conversation_messages(conversation_id=duplicate_node.objective_target_conversation_id) ) - replayable_messages = ConversationManager.get_persistable_prepended_messages( - prepended_conversation=duplicated_messages - ) - duplicate_node._target_normalization_context = ConversationManager.create_target_normalization_context( - target=self._objective_target, - conversation_id=duplicate_node.objective_target_conversation_id, - prepended_messages=replayable_messages, + duplicate_node._target_normalization_context = ( + self._target_normalization_context.remap_for_duplicate_conversation( + conversation_id=duplicate_node.objective_target_conversation_id, + source_messages=source_messages, + duplicated_messages=duplicated_messages, + ) + if self._target_normalization_context + else None ) duplicate_node.adversarial_chat_conversation_id = self._memory.duplicate_conversation( @@ -1971,11 +1959,6 @@ def _branch_existing_nodes(self, context: TAPAttackContext) -> None: cloned_nodes = [] for node in context.nodes: - # Stateless targets have no provider-side conversation to retain. - # Refresh the original node as well as its clones so every outgoing - # branch sees the same complete logical history. This also covers a - # branching factor of one, where no clone is created. - node._refresh_stateless_branch_boundary() for _ in range(self._configuration.branching_factor - 1): cloned_node = node.duplicate() # Add the adversarial chat conversation ID of the duplicated node to the context's tracking diff --git a/pyrit/executor/attack/streaming/barge_in.py b/pyrit/executor/attack/streaming/barge_in.py index cd6a3560e4..5ebfbe690b 100644 --- a/pyrit/executor/attack/streaming/barge_in.py +++ b/pyrit/executor/attack/streaming/barge_in.py @@ -85,7 +85,8 @@ def __init__( Defaults to a fresh ``PromptNormalizer``. params_type: Attack parameter dataclass type. prepended_conversation_config: Configuration for prepended-conversation - conversion and formatting. + converter-role selection. Its message formatter is not used by the + direct realtime streaming path. Raises: ValueError: If ``objective_target`` does not declare the ``STREAMING_AUDIO`` @@ -131,7 +132,9 @@ async def _setup_async(self, *, context: BargeInAttackContext[Any]) -> None: Prepended messages are recorded in memory but are NOT pushed into the live realtime session beyond the system prompt — the model only conditions on the system message - and live audio chunks. + and live audio chunks. The direct streaming path does not call + ``PromptTarget.send_prompt_async``, so it cannot consume a target-normalization + context or a prepended-history formatter. """ if not context.conversation_id: context.conversation_id = str(uuid.uuid4()) @@ -142,6 +145,7 @@ async def _setup_async(self, *, context: BargeInAttackContext[Any]) -> None: request_converters=self._request_converters, prepended_conversation_config=self._prepended_conversation_config, ) + context.target_normalization_context = None async def _teardown_async(self, *, context: BargeInAttackContext[Any]) -> None: """No-op teardown — connection / dispatcher are closed inside the session's ``run_async``.""" diff --git a/pyrit/message_normalizer/message_normalizer.py b/pyrit/message_normalizer/message_normalizer.py index d2fe3fd20a..dfd941dd6f 100644 --- a/pyrit/message_normalizer/message_normalizer.py +++ b/pyrit/message_normalizer/message_normalizer.py @@ -38,6 +38,12 @@ async def normalize_async(self, messages: list[Message]) -> list[T]: Returns: A list of normalized items of type T. + + Note: + Output metadata is authoritative. A normalizer that creates replacement + ``Message`` or ``MessagePiece`` objects must preserve any + ``prompt_metadata`` required by downstream consumers. The target stamps + the active conversation ID but does not merge removed metadata back in. """ async def normalize_to_dicts_async(self, messages: list[Message]) -> list[dict[str, Any]]: diff --git a/pyrit/prompt_target/common/prompt_target.py b/pyrit/prompt_target/common/prompt_target.py index ceaaffa6ee..1a7e552c96 100644 --- a/pyrit/prompt_target/common/prompt_target.py +++ b/pyrit/prompt_target/common/prompt_target.py @@ -268,10 +268,9 @@ async def _get_normalized_conversation_async( The original conversation in memory is never mutated. The returned list is an ephemeral copy intended only for building the API request body. - After normalization, the metadata from the original ``message`` is copied - onto the last normalized message so that downstream code (e.g. - ``construct_response_from_request``) propagates the correct - ``conversation_id`` and request lineage to the response. + After normalization, every output piece is stamped with the current + conversation ID. Normalizers own all other output metadata; removed + ``prompt_metadata`` keys are not restored. Args: message (Message): The current message to append. @@ -298,49 +297,11 @@ async def _get_normalized_conversation_async( normalizer_overrides=normalizer_overrides, ) if normalized: - # Normalizers may create new Message objects (via Message.from_prompt) with - # random conversation_ids. Stamp the correct conversation_id on every - # message (idempotent for originals, fixes new ones). Full lineage is only - # propagated to the last message — it's the one targets use to build the - # response, and earlier messages carry their own legitimate metadata. for msg in normalized: for piece in msg.message_pieces: piece.conversation_id = conversation_id - self._propagate_lineage(source=message, target_message=normalized[-1]) - if len(normalized) > len(conversation): - logger.warning( - "Normalization produced more messages than the input conversation " - "(%d → %d). Only the last normalized message has full lineage " - "metadata. Additional new messages have conversation_id set but " - "require manual lineage updates if needed.", - len(conversation), - len(normalized), - ) return normalized - @staticmethod - def _propagate_lineage(*, source: Message, target_message: Message) -> None: - """ - Copy request-lineage metadata from ``source`` onto every piece in ``target_message``. - - Normalizers may create brand-new messages or pieces, such as the combined - text piece from ``HistorySquashNormalizer``, that lack request lineage. - This method restores the original metadata so that the response built from - the normalized message stays part of the correct conversation and retains - traceability. - - Normalizers own ``prompt_metadata`` on their output. The final metadata is - authoritative so a later normalizer can intentionally remove keys such as - ``json_schema`` without lineage propagation restoring them. - - Args: - source: The original (pre-normalization) message whose metadata is authoritative. - target_message: The normalized message whose pieces will be updated in place. - """ - conversation_id = source.get_piece().conversation_id - for piece in target_message.message_pieces: - piece.conversation_id = conversation_id - def set_model_name(self, *, model_name: str) -> None: """ Set the model name for this target. diff --git a/pyrit/prompt_target/common/target_normalization_context.py b/pyrit/prompt_target/common/target_normalization_context.py index d3e444819c..8d574c33fc 100644 --- a/pyrit/prompt_target/common/target_normalization_context.py +++ b/pyrit/prompt_target/common/target_normalization_context.py @@ -29,26 +29,31 @@ def filter_non_replayable_messages(*, messages: list[Message]) -> list[Message]: Messages that are safe to include in a later target-facing payload. """ non_replayable_errors = {"processing", "unknown"} - failed_request_piece_ids = { - piece.original_prompt_id - for message in messages - for piece in message.message_pieces - if piece.response_error in non_replayable_errors and piece.original_prompt_id is not None - } - failed_request_sequences = { - piece.sequence - 1 - for message in messages - for piece in message.message_pieces - if piece.response_error in non_replayable_errors and piece.sequence > 0 - } - - return [ - message - for message in messages - if not any(piece.response_error in non_replayable_errors for piece in message.message_pieces) - and not any(piece.id in failed_request_piece_ids for piece in message.message_pieces) - and not any(piece.sequence in failed_request_sequences for piece in message.message_pieces) - ] + excluded_indexes: set[int] = set() + + for index, error_response in enumerate(messages): + if not any(piece.response_error in non_replayable_errors for piece in error_response.message_pieces): + continue + + excluded_indexes.add(index) + if index == 0: + continue + + request = messages[index - 1] + request_piece = request.get_piece() + error_piece = error_response.get_piece() + is_adjacent_request = ( + request.api_role == "user" + and error_response.api_role == "assistant" + and bool(request_piece.conversation_id) + and request_piece.conversation_id == error_piece.conversation_id + and request_piece.sequence >= 0 + and error_piece.sequence == request_piece.sequence + 1 + ) + if is_adjacent_request: + excluded_indexes.add(index - 1) + + return [message for index, message in enumerate(messages) if index not in excluded_indexes] @dataclass(slots=True) @@ -180,3 +185,69 @@ def select_history(self, *, messages: list[Message]) -> list[Message]: f"Missing {len(missing_ids)} message(s)." ) return [messages_by_id[message_id] for message_id in self.history_message_ids] + + def remap_for_duplicate_conversation( + self, + *, + conversation_id: str, + source_messages: list[Message], + duplicated_messages: list[Message], + ) -> TargetNormalizationContext: + """ + Remap this explicit boundary to a duplicated conversation. + + Only message pieces already identified as history are remapped. Live turns + copied into the new memory conversation never become part of the boundary. + + Args: + conversation_id: Conversation ID assigned to the duplicated messages. + source_messages: Messages from the source conversation in persisted order. + duplicated_messages: Their duplicates in the same persisted order. + + Returns: + A new context with boundary IDs from the duplicated conversation. + + Raises: + ValueError: If the duplicated messages do not match the source structure + or an explicit history piece cannot be remapped. + """ + if len(source_messages) != len(duplicated_messages): + raise ValueError("Duplicated conversation does not match the source message count.") + + duplicated_ids_by_source_id: dict[uuid.UUID, uuid.UUID] = {} + for source_message, duplicated_message in zip(source_messages, duplicated_messages, strict=True): + if ( + source_message.api_role != duplicated_message.api_role + or source_message.sequence != duplicated_message.sequence + or len(source_message.message_pieces) != len(duplicated_message.message_pieces) + ): + raise ValueError("Duplicated conversation does not preserve the source message structure.") + + for source_piece, duplicated_piece in zip( + source_message.message_pieces, + duplicated_message.message_pieces, + strict=True, + ): + if ( + source_piece.api_role != duplicated_piece.api_role + or source_piece.sequence != duplicated_piece.sequence + or duplicated_piece.conversation_id != conversation_id + ): + raise ValueError("Duplicated conversation does not preserve the source piece structure.") + duplicated_ids_by_source_id[source_piece.id] = duplicated_piece.id + + missing_ids = [ + message_id for message_id in self.history_message_ids if message_id not in duplicated_ids_by_source_id + ] + if missing_ids: + raise ValueError(f"Could not remap {len(missing_ids)} explicit history message(s).") + + duplicated_context = TargetNormalizationContext( + conversation_id=conversation_id, + history_message_ids=tuple( + duplicated_ids_by_source_id[message_id] for message_id in self.history_message_ids + ), + replay_history_each_send=self.replay_history_each_send, + ) + duplicated_context._history_consumed = self._history_consumed + return duplicated_context diff --git a/tests/unit/executor/attack/multi_turn/test_prepended_history_normalization.py b/tests/unit/executor/attack/multi_turn/test_prepended_history_normalization.py index 2df3efd802..ef8a04a994 100644 --- a/tests/unit/executor/attack/multi_turn/test_prepended_history_normalization.py +++ b/tests/unit/executor/attack/multi_turn/test_prepended_history_normalization.py @@ -3,10 +3,12 @@ """Focused regression tests for prepended-history target normalization.""" +from pathlib import Path from unittest.mock import AsyncMock, MagicMock import pytest +from pyrit.converter import Converter, ConverterResult from pyrit.executor.attack.component import ConversationManager, PrependedConversationConfig from pyrit.executor.attack.component.modality_router import _ModalityFeedbackRouter from pyrit.executor.attack.core import AttackAdversarialConfig, AttackScoringConfig @@ -28,8 +30,9 @@ ConversationType, Message, MessagePiece, + PromptDataType, ) -from pyrit.prompt_normalizer import PromptNormalizer +from pyrit.prompt_normalizer import ConverterConfiguration, PromptNormalizer from pyrit.prompt_target import ( CapabilityName, PromptTarget, @@ -50,10 +53,12 @@ def __init__(self, *, supports_multi_turn: bool = False, supports_editable_histo ) ) self.prompt_sent: list[str] = [] + self.normalized_requests: list[Message] = [] async def _send_prompt_to_target_async(self, *, normalized_conversation: list[Message]) -> list[Message]: request = normalized_conversation[-1] self.prompt_sent.append(request.get_value()) + self.normalized_requests.append(request) return [ MessagePiece( role="assistant", @@ -63,6 +68,18 @@ async def _send_prompt_to_target_async(self, *, normalized_conversation: list[Me ] +class _ImageOutputConverter(Converter): + SUPPORTED_INPUT_TYPES: tuple[PromptDataType, ...] = ("text",) + SUPPORTED_OUTPUT_TYPES: tuple[PromptDataType, ...] = ("image_path",) + + def __init__(self, *, output_path: str) -> None: + super().__init__() + self._output_path = output_path + + async def convert_async(self, *, prompt: str, input_type: PromptDataType = "text") -> ConverterResult: + return ConverterResult(output_text=self._output_path, output_type="image_path") + + def _make_context() -> MultiTurnAttackContext[AttackParameters]: return MultiTurnAttackContext(params=AttackParameters(objective="test objective")) @@ -261,30 +278,91 @@ def _make_tap_node(*, target: PromptTarget) -> _TreeOfAttacksNode: ) -@pytest.mark.parametrize("supports_multi_turn", [False, True]) -@pytest.mark.usefixtures("patch_central_database") -def test_tap_branch_duplicates_full_history_with_explicit_boundary(supports_multi_turn: bool): - target = _RecordingTarget(supports_multi_turn=supports_multi_turn) - node = _make_tap_node(target=target) +def _set_tap_seed_boundary( + *, + node: _TreeOfAttacksNode, + target: PromptTarget, + seed_messages: list[Message], +) -> None: _seed_conversation( conversation_id=node.objective_target_conversation_id, target=target, - messages=[ - Message.from_system_prompt("system"), - Message.from_prompt(prompt="request", role="user"), - Message.from_prompt(prompt="response", role="assistant"), - ], + messages=seed_messages, ) + node._target_normalization_context = ConversationManager.create_target_normalization_context( + target=target, + conversation_id=node.objective_target_conversation_id, + prepended_messages=seed_messages, + ) + assert node._target_normalization_context is not None - duplicate = node.duplicate() - messages = CentralMemory.get_memory_instance().get_conversation_messages( - conversation_id=duplicate.objective_target_conversation_id +def _branch_tap_node( + *, + node: _TreeOfAttacksNode, + branching_factor: int, +) -> list[_TreeOfAttacksNode]: + context = MagicMock() + context.nodes = [node] + context.related_conversations = set() + attack = MagicMock() + attack._configuration.branching_factor = branching_factor + TreeOfAttacksWithPruningAttack._branch_existing_nodes(attack, context) + return context.nodes + + +@pytest.mark.usefixtures("patch_central_database") +async def test_tap_seeded_stateless_retained_and_cloned_branches_replay_only_original_seed(): + target = _RecordingTarget() + formatter = MagicMock(spec=MessageStringNormalizer) + + async def format_messages(messages: list[Message]) -> str: + return "|".join(message.get_value() for message in messages) + + formatter.normalize_string_async = AsyncMock(side_effect=format_messages) + node = _make_tap_node(target=target) + node._prepended_conversation_config = PrependedConversationConfig(message_normalizer=formatter) + seed = Message.from_prompt(prompt="original seed", role="user") + _set_tap_seed_boundary(node=node, target=target, seed_messages=[seed]) + node._objective = "objective" + await node._send_prompt_to_target_async("depth one") + original_context = node._target_normalization_context + assert original_context is not None + + retained, cloned = _branch_tap_node(node=node, branching_factor=2) + assert retained is node + assert retained._target_normalization_context is original_context + assert cloned._target_normalization_context is not None + assert cloned._target_normalization_context.history_message_count == 1 + cloned_messages = CentralMemory.get_memory_instance().get_conversation_messages( + conversation_id=cloned.objective_target_conversation_id ) - assert [message.api_role for message in messages] == ["system", "user", "assistant"] - assert duplicate._target_normalization_context is not None - assert duplicate._target_normalization_context.history_message_count == 3 - assert duplicate._target_normalization_context.replay_history_each_send is not supports_multi_turn + assert cloned._target_normalization_context.history_message_ids == (cloned_messages[0].get_piece().id,) + assert cloned._target_normalization_context.history_message_ids != original_context.history_message_ids + + for branch, prompt in [(retained, "depth two retained"), (cloned, "depth two cloned")]: + branch._objective = "objective" + await branch._send_prompt_to_target_async(prompt) + + deep_clone = cloned.duplicate() + deep_clone._objective = "objective" + await deep_clone._send_prompt_to_target_async("depth three cloned") + + assert target.prompt_sent == [ + "original seed|depth one", + "original seed|depth two retained", + "original seed|depth two cloned", + "original seed|depth three cloned", + ] + formatted_values = [ + [message.get_value() for message in call.args[0]] for call in formatter.normalize_string_async.await_args_list + ] + assert formatted_values == [ + ["original seed", "depth one"], + ["original seed", "depth two retained"], + ["original seed", "depth two cloned"], + ["original seed", "depth three cloned"], + ] @pytest.mark.usefixtures("patch_central_database") @@ -331,7 +409,33 @@ async def test_tap_unseeded_stateless_send_retains_current_only_payload(): @pytest.mark.usefixtures("patch_central_database") -async def test_tap_stateless_branch_sends_original_boundary_plus_each_current_request(): +async def test_tap_unseeded_stateless_retained_and_cloned_branches_send_current_only(): + target = _RecordingTarget() + node = _make_tap_node(target=target) + node._objective = "objective" + await node._send_prompt_to_target_async("depth one") + + retained, cloned = _branch_tap_node(node=node, branching_factor=2) + assert retained._target_normalization_context is None + assert cloned._target_normalization_context is None + for branch, prompt in [(retained, "depth two retained"), (cloned, "depth two cloned")]: + branch._objective = "objective" + await branch._send_prompt_to_target_async(prompt) + + deep_clone = cloned.duplicate() + deep_clone._objective = "objective" + await deep_clone._send_prompt_to_target_async("depth three cloned") + + assert target.prompt_sent == [ + "depth one", + "depth two retained", + "depth two cloned", + "depth three cloned", + ] + + +@pytest.mark.usefixtures("patch_central_database") +async def test_tap_branching_factor_one_preserves_retained_seed_boundary(): target = _RecordingTarget() formatter = MagicMock(spec=MessageStringNormalizer) @@ -341,37 +445,46 @@ async def format_messages(messages: list[Message]) -> str: formatter.normalize_string_async = AsyncMock(side_effect=format_messages) node = _make_tap_node(target=target) node._prepended_conversation_config = PrependedConversationConfig(message_normalizer=formatter) - _seed_conversation( - conversation_id=node.objective_target_conversation_id, + _set_tap_seed_boundary( + node=node, target=target, - messages=[ - Message.from_prompt(prompt="branch request", role="user"), - Message.from_prompt(prompt="branch response", role="assistant"), - ], + seed_messages=[Message.from_prompt(prompt="original seed", role="user")], ) - duplicate = node.duplicate() - duplicate._objective = "objective" + node._objective = "objective" + await node._send_prompt_to_target_async("depth one") + original_context = node._target_normalization_context + original_boundary = original_context.history_message_ids if original_context else () - await duplicate._send_prompt_to_target_async("first current") - await duplicate._send_prompt_to_target_async("second current") + branches = _branch_tap_node(node=node, branching_factor=1) + assert branches == [node] + assert node._target_normalization_context is original_context + assert node._target_normalization_context is not None + assert node._target_normalization_context.history_message_ids == original_boundary + await node._send_prompt_to_target_async("depth two retained") assert target.prompt_sent == [ - "branch request|branch response|first current", - "branch request|branch response|second current", - ] - formatted_values = [ - [message.get_value() for message in call.args[0]] for call in formatter.normalize_string_async.await_args_list - ] - assert formatted_values == [ - ["branch request", "branch response", "first current"], - ["branch request", "branch response", "second current"], + "original seed|depth one", + "original seed|depth two retained", ] -@pytest.mark.parametrize("branching_factor", [1, 2]) @pytest.mark.usefixtures("patch_central_database") -async def test_tap_retained_and_cloned_stateless_branches_replay_same_full_history(branching_factor: int): +async def test_tap_clone_does_not_replay_non_text_live_converter_output(tmp_path: Path): + image_path = tmp_path / "converted.png" + image_path.write_bytes(b"test image") target = _RecordingTarget() + target._configuration = TargetConfiguration( + capabilities=TargetCapabilities( + supports_multi_message_pieces=True, + input_modalities=frozenset( + { + frozenset({"text"}), + frozenset({"image_path"}), + frozenset({"text", "image_path"}), + } + ), + ) + ) formatter = MagicMock(spec=MessageStringNormalizer) async def format_messages(messages: list[Message]) -> str: @@ -379,28 +492,36 @@ async def format_messages(messages: list[Message]) -> str: formatter.normalize_string_async = AsyncMock(side_effect=format_messages) node = _make_tap_node(target=target) + node._request_converters = ConverterConfiguration.from_converters( + converters=[_ImageOutputConverter(output_path=str(image_path))] + ) node._prepended_conversation_config = PrependedConversationConfig(message_normalizer=formatter) - _seed_conversation( - conversation_id=node.objective_target_conversation_id, + _set_tap_seed_boundary( + node=node, target=target, - messages=[ - Message.from_prompt(prompt="branch request", role="user"), - Message.from_prompt(prompt="branch response", role="assistant"), - ], + seed_messages=[Message.from_prompt(prompt="original seed", role="user")], ) - context = MagicMock() - context.nodes = [node] - context.related_conversations = set() - attack = MagicMock() - attack._configuration.branching_factor = branching_factor + node._objective = "objective" - TreeOfAttacksWithPruningAttack._branch_existing_nodes(attack, context) + await node._send_prompt_to_target_async("depth one") + cloned = node.duplicate() + cloned._objective = "objective" + await cloned._send_prompt_to_target_async("depth two") - assert len(context.nodes) == branching_factor - for branch in context.nodes: - branch._objective = "objective" - await branch._send_prompt_to_target_async("current request") - assert target.prompt_sent == ["branch request|branch response|current request"] * branching_factor + assert [piece.converted_value_data_type for piece in target.normalized_requests[-1].message_pieces] == [ + "text", + "image_path", + ] + assert target.normalized_requests[-1].get_values() == ["original seed", str(image_path)] + formatted_values = [ + [message.get_value() for message in call.args[0]] for call in formatter.normalize_string_async.await_args_list + ] + assert formatted_values == [ + ["original seed", "depth one"], + ["original seed"], + ["original seed", "depth two"], + ["original seed"], + ] @pytest.fixture diff --git a/tests/unit/executor/attack/streaming/test_barge_in.py b/tests/unit/executor/attack/streaming/test_barge_in.py index 7a4bc289fe..029b1e6677 100644 --- a/tests/unit/executor/attack/streaming/test_barge_in.py +++ b/tests/unit/executor/attack/streaming/test_barge_in.py @@ -13,7 +13,7 @@ from pyrit.executor.attack import BargeInAttack, BargeInAttackContext from pyrit.executor.attack.core import AttackParameters from pyrit.models import AttackOutcome, Message, MessagePiece -from pyrit.prompt_target import RealtimeTarget +from pyrit.prompt_target import RealtimeTarget, TargetNormalizationContext if TYPE_CHECKING: from collections.abc import AsyncIterator @@ -156,15 +156,21 @@ async def test_setup_async_persists_prepended_conversation_to_memory(vad_target) # All three messages share the context's conversation_id post-setup. for m in add_calls: assert m.message_pieces[0].conversation_id == ctx.conversation_id + assert ctx.target_normalization_context is None -async def test_setup_async_no_op_when_prepended_conversation_empty(vad_target): - """Empty prepended_conversation: no memory writes, no crash.""" +async def test_setup_async_clears_unused_normalization_context_when_prepended_empty(vad_target): + """The direct streaming path does not retain target normalization state.""" attack = BargeInAttack(objective_target=vad_target) ctx = BargeInAttackContext( params=AttackParameters(objective="o"), # no prepended_conversation audio_chunks=_aiter([b"\x00" * 96]), ) + ctx.target_normalization_context = TargetNormalizationContext( + conversation_id=ctx.conversation_id, + history_message_ids=(Message.from_prompt(prompt="unused", role="user").get_piece().id,), + replay_history_each_send=True, + ) add_calls: list[Any] = [] with patch.object(attack._conversation_manager._memory, "add_message_to_memory") as mock_add: @@ -172,6 +178,7 @@ async def test_setup_async_no_op_when_prepended_conversation_empty(vad_target): await attack._setup_async(context=ctx) assert add_calls == [] + assert ctx.target_normalization_context is None # ---- _perform_async: session factory passthrough ---------------------------------------------- diff --git a/tests/unit/prompt_target/target/test_normalize_async_integration.py b/tests/unit/prompt_target/target/test_normalize_async_integration.py index 4a37d34b11..c480f6ffd1 100644 --- a/tests/unit/prompt_target/target/test_normalize_async_integration.py +++ b/tests/unit/prompt_target/target/test_normalize_async_integration.py @@ -16,14 +16,17 @@ from openai.types.responses import ResponseOutputMessage, ResponseOutputText from unit.mocks import MockPromptTarget +from pyrit.memory import CentralMemory from pyrit.memory.memory_interface import MemoryInterface from pyrit.message_normalizer import ( ConversationContextNormalizer, HistorySquashNormalizer, + MessageListNormalizer, MessageStringNormalizer, TokenizerTemplateNormalizer, ) from pyrit.models import ComponentIdentifier, Message, MessagePiece, PromptResponseError +from pyrit.prompt_normalizer import PromptNormalizer from pyrit.prompt_target import AzureMLChatTarget, OpenAIChatTarget, TargetNormalizationContext from pyrit.prompt_target.common.target_capabilities import ( CapabilityHandlingPolicy, @@ -752,6 +755,34 @@ async def test_history_squash_does_not_restore_adapted_json_schema_metadata(): assert '"description"' in text_piece.converted_value +@pytest.mark.usefixtures("patch_central_database") +async def test_custom_normalizer_output_metadata_is_authoritative(): + target = MockPromptTarget() + target._configuration = TargetConfiguration( + capabilities=TargetCapabilities( + supports_multi_turn=True, + supports_multi_message_pieces=True, + supports_system_prompt=True, + ) + ) + source = _make_message(role="user", content="source") + source.get_piece().prompt_metadata = {"source-only": "must not be restored"} + replacement = Message.from_prompt(prompt="replacement", role="user") + normalizer = MagicMock(spec=MessageListNormalizer) + normalizer.normalize_async = AsyncMock(return_value=[replacement]) + mock_memory = MagicMock(spec=MemoryInterface) + mock_memory.get_conversation_messages.return_value = [] + target._memory = mock_memory + + result = await target._get_normalized_conversation_async( + message=source, + normalizer_overrides={CapabilityName.EDITABLE_HISTORY: normalizer}, + ) + + assert result[0].get_piece().conversation_id == "conv1" + assert result[0].get_piece().prompt_metadata == {} + + @pytest.mark.usefixtures("patch_central_database") async def test_prepended_history_adapter_is_used_only_when_explicitly_passed(): target = MockPromptTarget() @@ -1053,81 +1084,101 @@ async def test_target_normalization_failure_can_be_retried(): @pytest.mark.usefixtures("patch_central_database") -async def test_retry_ignores_persisted_failed_request_and_error_response(): +async def test_prompt_normalizer_retry_excludes_persisted_processing_exchange(): target = MockPromptTarget() - target._configuration = TargetConfiguration( - capabilities=TargetCapabilities( - supports_multi_turn=True, - supports_system_prompt=True, + prompt_normalizer = PromptNormalizer() + conversation_id = "processing-retry" + successful_response = _make_message(role="assistant", content="successful response") + retry_response = _make_message(role="assistant", content="retry response") + + with patch.object(target, "_send_prompt_to_target_async", new_callable=AsyncMock) as send: + send.side_effect = [[successful_response], RuntimeError("private provider failure"), [retry_response]] + await prompt_normalizer.send_prompt_async( + message=_make_message(role="user", content="successful request"), + target=target, + conversation_id=conversation_id, + ) + with pytest.raises(Exception, match="Error sending prompt"): + await prompt_normalizer.send_prompt_async( + message=_make_message(role="user", content="failed request"), + target=target, + conversation_id=conversation_id, + ) + await prompt_normalizer.send_prompt_async( + message=_make_message(role="user", content="retry"), + target=target, + conversation_id=conversation_id, ) - ) - prepended = _make_message(role="user", content="prepended") - failed_request = _make_message(role="user", content="failed request") - error_response = _make_message(role="assistant", content="processing error") - error_response.get_piece().response_error = "processing" - mock_memory = MagicMock(spec=MemoryInterface) - mock_memory.get_conversation_messages.return_value = [prepended, failed_request, error_response] - target._memory = mock_memory - formatter = MagicMock(spec=MessageStringNormalizer) - formatter.normalize_string_async = AsyncMock(return_value="formatted retry") - target_context = _make_target_normalization_context( - prepended_messages=[prepended], - target_supports_multi_turn=True, - ) - normalizer_overrides = _make_normalizer_overrides( - target_normalization_context=target_context, - formatter=formatter, - ) - result = await target._get_normalized_conversation_async( - message=_make_message(role="user", content="retry"), - normalizer_overrides=normalizer_overrides, - target_normalization_context=target_context, + retry_payload = send.await_args_list[2].kwargs["normalized_conversation"] + assert [message.get_value() for message in retry_payload] == [ + "successful request", + "successful response", + "retry", + ] + persisted = list(CentralMemory.get_memory_instance().get_conversation_messages(conversation_id=conversation_id)) + processing_index = next( + index for index, message in enumerate(persisted) if message.get_piece().response_error == "processing" ) - - assert result[0].get_value() == "formatted retry" - formatted_messages = formatter.normalize_string_async.await_args.args[0] - assert [message.get_value() for message in formatted_messages] == ["prepended", "retry"] + failed_request = persisted[processing_index - 1].get_piece() + processing_error = persisted[processing_index].get_piece() + assert failed_request.api_role == "user" + assert processing_error.original_prompt_id != failed_request.id + assert processing_error.sequence == failed_request.sequence + 1 @pytest.mark.usefixtures("patch_central_database") -async def test_ordinary_history_excludes_processing_exchange_and_traceback(): +async def test_prompt_normalizer_retry_excludes_persisted_unknown_exchange(): target = MockPromptTarget() - target._configuration = TargetConfiguration( - capabilities=TargetCapabilities( - supports_multi_turn=True, - supports_editable_history=True, + prompt_normalizer = PromptNormalizer() + unknown_response = _make_message(role="assistant", content="unknown provider failure") + unknown_response.get_piece().response_error = "unknown" + retry_response = _make_message(role="assistant", content="retry response") + + with patch.object(target, "_send_prompt_to_target_async", new_callable=AsyncMock) as send: + send.side_effect = [[unknown_response], [retry_response]] + await prompt_normalizer.send_prompt_async( + message=_make_message(role="user", content="failed request"), + target=target, + conversation_id="unknown-retry", + ) + await prompt_normalizer.send_prompt_async( + message=_make_message(role="user", content="retry"), + target=target, + conversation_id="unknown-retry", ) - ) - successful_request = _make_message(role="user", content="successful request") - successful_request.get_piece().sequence = 0 - successful_response = _make_message(role="assistant", content="successful response") - successful_response.get_piece().sequence = 1 - failed_request = _make_message(role="user", content="failed request") - failed_request.get_piece().sequence = 2 - processing_response = _make_message( - role="assistant", - content="RuntimeError\nTraceback (most recent call last): secret stack", - ) - processing_response.get_piece().sequence = 3 - processing_response.get_piece().response_error = "processing" - mock_memory = MagicMock(spec=MemoryInterface) - mock_memory.get_conversation_messages.return_value = [ - successful_request, - successful_response, - failed_request, - processing_response, - ] - target._memory = mock_memory - result = await target._get_normalized_conversation_async( - message=_make_message(role="user", content="retry"), - ) + retry_payload = send.await_args_list[1].kwargs["normalized_conversation"] + assert [message.get_value() for message in retry_payload] == ["retry"] - assert [message.get_value() for message in result] == [ - "successful request", - "successful response", - "retry", + +@pytest.mark.parametrize("response_error", ["blocked", "empty"]) +@pytest.mark.usefixtures("patch_central_database") +async def test_prompt_normalizer_retains_provider_round_trip(response_error: PromptResponseError): + target = MockPromptTarget() + prompt_normalizer = PromptNormalizer() + provider_response = _make_message(role="assistant", content=f"{response_error} response") + provider_response.get_piece().response_error = response_error + next_response = _make_message(role="assistant", content="next response") + + with patch.object(target, "_send_prompt_to_target_async", new_callable=AsyncMock) as send: + send.side_effect = [[provider_response], [next_response]] + await prompt_normalizer.send_prompt_async( + message=_make_message(role="user", content="first request"), + target=target, + conversation_id=f"{response_error}-round-trip", + ) + await prompt_normalizer.send_prompt_async( + message=_make_message(role="user", content="second request"), + target=target, + conversation_id=f"{response_error}-round-trip", + ) + + second_payload = send.await_args_list[1].kwargs["normalized_conversation"] + assert [message.get_value() for message in second_payload] == [ + "first request", + f"{response_error} response", + "second request", ] diff --git a/tests/unit/prompt_target/target/test_prompt_target.py b/tests/unit/prompt_target/target/test_prompt_target.py index 3de86263e7..453870e382 100644 --- a/tests/unit/prompt_target/target/test_prompt_target.py +++ b/tests/unit/prompt_target/target/test_prompt_target.py @@ -153,7 +153,7 @@ async def test_send_prompt_async_with_delay( # --------------------------------------------------------------------------- -# _propagate_lineage — metadata preservation after normalization +# Normalizer metadata and conversation ownership # --------------------------------------------------------------------------- _LINEAGE_CONVERSATION_ID = "original-conv-id-12345" @@ -192,8 +192,8 @@ def _make_mock_chat_completion(content: str = "response") -> MagicMock: @pytest.mark.usefixtures("patch_central_database") async def test_history_squash_preserves_metadata_on_normalized_message(): """ - After history squash, _propagate_lineage should restore the original request's - conversation ID and prompt metadata onto the squashed message. + History squash preserves the current request's metadata, and the target stamps + the active conversation ID on its output. """ target = OpenAIChatTarget( model_name="gpt-4o", @@ -236,8 +236,7 @@ async def test_history_squash_preserves_metadata_on_normalized_message(): async def test_response_preserves_metadata_after_history_squash(): """ End-to-end: after history squash the response must carry the original - request's conversation ID and prompt metadata, not the random values - created by the normalizer. + request's conversation ID and prompt metadata. """ target = OpenAIChatTarget( model_name="gpt-4o", @@ -282,8 +281,8 @@ async def test_response_preserves_metadata_after_history_squash(): @pytest.mark.usefixtures("patch_central_database") async def test_system_squash_preserves_metadata(): """ - GenericSystemSquashNormalizer also creates messages via Message.from_prompt. - _propagate_lineage should restore the original metadata after system squash too. + GenericSystemSquashNormalizer preserves the current request's metadata when + it builds the replacement user message. """ target = OpenAIChatTarget( model_name="gpt-4o", @@ -324,10 +323,10 @@ async def test_system_squash_preserves_metadata(): @pytest.mark.usefixtures("patch_central_database") -async def test_history_squash_propagates_lineage_to_all_pieces(): +async def test_history_squash_preserves_metadata_on_all_output_pieces(): """ - When the squashed message contains multiple pieces, _propagate_lineage - must stamp every piece — not just the first one. + Every piece produced by history squash keeps the current request's metadata + and receives the active conversation ID. """ target = OpenAIChatTarget( model_name="gpt-4o", @@ -372,12 +371,10 @@ async def test_history_squash_propagates_lineage_to_all_pieces(): @pytest.mark.usefixtures("patch_central_database") -async def test_conversation_id_stamped_on_all_but_full_lineage_only_on_last(): +async def test_conversation_id_stamped_without_merging_normalizer_output_metadata(): """ - conversation_id is stamped on every normalized message (including new ones - created by the normalizer). Full lineage is only propagated to the last - message. Earlier messages keep their own metadata. A warning is logged when - the normalizer increases message count. + The target stamps conversation_id on every normalized output while leaving + each normalizer-produced message's metadata authoritative. """ target = OpenAIChatTarget( model_name="gpt-4o", @@ -405,14 +402,19 @@ async def test_conversation_id_stamped_on_all_but_full_lineage_only_on_last(): converted_value_data_type="text", ) new_msg = Message(message_pieces=[new_piece]) + replacement_piece = MessagePiece( + role="user", + conversation_id="another-normalizer-uuid", + original_value="replacement", + converted_value="replacement", + original_value_data_type="text", + converted_value_data_type="text", + ) + replacement_msg = Message(message_pieces=[replacement_piece]) with patch.object(target.configuration, "normalize_async", new_callable=AsyncMock) as mock_normalize: - mock_normalize.return_value = [history_msg, new_msg, user_msg] - - import logging - - with patch.object(logging.getLogger("pyrit.prompt_target.common.prompt_target"), "warning") as mock_warn: - normalized = await target._get_normalized_conversation_async(message=user_msg) + mock_normalize.return_value = [history_msg, new_msg, replacement_msg] + normalized = await target._get_normalized_conversation_async(message=user_msg) # All messages should carry the correct conversation_id. for msg in normalized: @@ -422,24 +424,17 @@ async def test_conversation_id_stamped_on_all_but_full_lineage_only_on_last(): # History message's other metadata should be untouched. assert normalized[0].message_pieces[0].prompt_metadata == {"original": "history_meta"} - # New middle message should NOT have full lineage overwritten. + # New messages keep exactly the metadata produced by the normalizer. assert normalized[1].message_pieces[0].prompt_metadata == {} - - # Last message should carry full lineage. - last_piece = normalized[-1].message_pieces[0] - assert last_piece.prompt_metadata == _LINEAGE_PROMPT_METADATA - - # Warning should fire because message count increased (2 → 3). - mock_warn.assert_called_once() + assert normalized[-1].message_pieces[0].prompt_metadata == {} @pytest.mark.usefixtures("patch_central_database") -async def test_json_schema_stripped_for_non_schema_target_survives_lineage(): +async def test_json_schema_stripped_for_non_schema_target_remains_authoritative(): """ Regression: for a non-schema target (default ADAPT) the embedded json_schema is - removed by JsonSchemaNormalizer and must NOT be re-introduced by - _propagate_lineage copying the original (unstripped) request metadata back onto - the normalized message. + removed by JsonSchemaNormalizer and must not be reintroduced from the source + request metadata. """ target = OpenAIChatTarget( model_name="gpt-4o", @@ -471,11 +466,10 @@ async def test_json_schema_stripped_for_non_schema_target_survives_lineage(): @pytest.mark.usefixtures("patch_central_database") -async def test_json_schema_only_metadata_fully_stripped_survives_lineage(): +async def test_json_schema_only_metadata_fully_stripped_remains_authoritative(): """ Regression: even when json_schema is the ONLY metadata key, the strip leaves empty - metadata and _propagate_lineage must not restore the original json_schema (the piece - is the same logical piece, identified by id, so its stripped metadata is authoritative). + metadata and the target must not restore the original json_schema. """ target = OpenAIChatTarget( model_name="gpt-4o", diff --git a/tests/unit/prompt_target/target/test_target_normalization_context.py b/tests/unit/prompt_target/target/test_target_normalization_context.py index a0b235b0c6..359840faac 100644 --- a/tests/unit/prompt_target/target/test_target_normalization_context.py +++ b/tests/unit/prompt_target/target/test_target_normalization_context.py @@ -12,9 +12,17 @@ ) -def _message(*, role: ChatMessageRole, value: str, sequence: int) -> Message: +def _message( + *, + role: ChatMessageRole, + value: str, + sequence: int, + conversation_id: str = "conversation", +) -> Message: message = Message.from_prompt(prompt=value, role=role) - message.get_piece().sequence = sequence + piece = message.get_piece() + piece.sequence = sequence + piece.conversation_id = conversation_id return message @@ -122,27 +130,42 @@ def test_context_rejects_missing_persisted_boundary_message(): @pytest.mark.parametrize("response_error", ["processing", "unknown"]) -def test_filter_removes_failed_exchange_by_request_id(response_error: PromptResponseError): +def test_filter_removes_adjacent_failed_exchange(response_error: PromptResponseError): successful = _message(role="user", value="successful", sequence=0) failed_request = _message(role="user", value="failed request", sequence=1) error_response = _message(role="assistant", value="private stack trace", sequence=2) error_response.get_piece().response_error = response_error - error_response.get_piece().original_prompt_id = failed_request.get_piece().id filtered = filter_non_replayable_messages(messages=[successful, failed_request, error_response]) assert filtered == [successful] -def test_filter_removes_duplicated_failed_exchange_by_sequence(): - failed_request = _message(role="user", value="failed request", sequence=4) +@pytest.mark.parametrize( + ("preceding_role", "preceding_conversation_id", "preceding_sequence"), + [ + ("assistant", "conversation", 4), + ("user", "other-conversation", 4), + ("user", "conversation", 3), + ], +) +def test_filter_does_not_remove_unrelated_preceding_message( + preceding_role: ChatMessageRole, + preceding_conversation_id: str, + preceding_sequence: int, +): + preceding = _message( + role=preceding_role, + value="unrelated", + sequence=preceding_sequence, + conversation_id=preceding_conversation_id, + ) error_response = _message(role="assistant", value="private stack trace", sequence=5) error_response.get_piece().response_error = "processing" - error_response.get_piece().original_prompt_id = uuid.uuid4() - filtered = filter_non_replayable_messages(messages=[failed_request, error_response]) + filtered = filter_non_replayable_messages(messages=[preceding, error_response]) - assert filtered == [] + assert filtered == [preceding] @pytest.mark.parametrize("response_error", ["blocked", "empty"]) @@ -150,8 +173,56 @@ def test_filter_retains_provider_round_trip(response_error: PromptResponseError) request = _message(role="user", value="request", sequence=0) response = _message(role="assistant", value="provider response", sequence=1) response.get_piece().response_error = response_error - response.get_piece().original_prompt_id = request.get_piece().id filtered = filter_non_replayable_messages(messages=[request, response]) assert filtered == [request, response] + + +def test_context_remaps_only_explicit_history_for_duplicate_conversation(): + seed = _message(role="system", value="seed", sequence=0) + live_request = _message(role="user", value="live", sequence=1) + source_messages = [seed, live_request] + duplicated_messages = [message.duplicate() for message in source_messages] + for message in duplicated_messages: + for piece in message.message_pieces: + piece.conversation_id = "duplicate" + + context = TargetNormalizationContext( + conversation_id="conversation", + history_message_ids=(seed.get_piece().id,), + replay_history_each_send=True, + ) + + duplicate = context.remap_for_duplicate_conversation( + conversation_id="duplicate", + source_messages=source_messages, + duplicated_messages=duplicated_messages, + ) + + assert duplicate.history_message_ids == (duplicated_messages[0].get_piece().id,) + assert duplicated_messages[1].get_piece().id not in duplicate.history_message_ids + assert duplicate.select_history(messages=duplicated_messages) == [duplicated_messages[0]] + + +def test_context_remap_preserves_consumed_state(): + seed = _message(role="user", value="seed", sequence=0) + duplicated_seed = seed.duplicate() + duplicated_seed.get_piece().conversation_id = "duplicate" + context = TargetNormalizationContext( + conversation_id="conversation", + history_message_ids=(seed.get_piece().id,), + replay_history_each_send=False, + ) + context.begin_send() + context.mark_provider_attempted() + context.finish_send() + + duplicate = context.remap_for_duplicate_conversation( + conversation_id="duplicate", + source_messages=[seed], + duplicated_messages=[duplicated_seed], + ) + + assert duplicate.is_consumed + assert duplicate.select_history(messages=[duplicated_seed]) == [] From 8acc5199351720a864453f5118cd119acad8ee29 Mon Sep 17 00:00:00 2001 From: Copilot <223556219+Copilot@users.noreply.github.com> Date: Fri, 21 Aug 2026 11:02:41 -0700 Subject: [PATCH 19/24] Fix cloned target seed consumption Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> Copilot-Session: 6776aa95-0924-4967-a056-e5a228af3faa --- .../common/target_normalization_context.py | 5 +- .../test_prepended_history_normalization.py | 53 +++++++++++++++++++ .../test_supports_multi_turn_attacks.py | 7 ++- .../test_target_normalization_context.py | 6 +-- 4 files changed, 63 insertions(+), 8 deletions(-) diff --git a/pyrit/prompt_target/common/target_normalization_context.py b/pyrit/prompt_target/common/target_normalization_context.py index 8d574c33fc..bc50138632 100644 --- a/pyrit/prompt_target/common/target_normalization_context.py +++ b/pyrit/prompt_target/common/target_normalization_context.py @@ -206,6 +206,7 @@ def remap_for_duplicate_conversation( Returns: A new context with boundary IDs from the duplicated conversation. + A new logical conversation starts with an unconsumed boundary. Raises: ValueError: If the duplicated messages do not match the source structure @@ -249,5 +250,7 @@ def remap_for_duplicate_conversation( ), replay_history_each_send=self.replay_history_each_send, ) - duplicated_context._history_consumed = self._history_consumed + # Provider bootstrap consumption belongs to the logical conversation, not copied memory. + if conversation_id == self.conversation_id: + duplicated_context._history_consumed = self._history_consumed return duplicated_context diff --git a/tests/unit/executor/attack/multi_turn/test_prepended_history_normalization.py b/tests/unit/executor/attack/multi_turn/test_prepended_history_normalization.py index ef8a04a994..796207af56 100644 --- a/tests/unit/executor/attack/multi_turn/test_prepended_history_normalization.py +++ b/tests/unit/executor/attack/multi_turn/test_prepended_history_normalization.py @@ -68,6 +68,19 @@ async def _send_prompt_to_target_async(self, *, normalized_conversation: list[Me ] +class _ConversationKeyedRecordingTarget(_RecordingTarget): + def __init__(self) -> None: + super().__init__(supports_multi_turn=True) + self.prompts_by_conversation: dict[str, list[str]] = {} + + async def _send_prompt_to_target_async(self, *, normalized_conversation: list[Message]) -> list[Message]: + request = normalized_conversation[-1] + conversation_id = request.get_piece().conversation_id + assert conversation_id + self.prompts_by_conversation.setdefault(conversation_id, []).append(request.get_value()) + return await super()._send_prompt_to_target_async(normalized_conversation=normalized_conversation) + + class _ImageOutputConverter(Converter): SUPPORTED_INPUT_TYPES: tuple[PromptDataType, ...] = ("text",) SUPPORTED_OUTPUT_TYPES: tuple[PromptDataType, ...] = ("image_path",) @@ -365,6 +378,46 @@ async def format_messages(messages: list[Message]) -> str: ] +@pytest.mark.usefixtures("patch_central_database") +async def test_tap_stateful_clone_replays_seed_once_for_new_conversation(): + target = _ConversationKeyedRecordingTarget() + formatter = MagicMock(spec=MessageStringNormalizer) + + async def format_messages(messages: list[Message]) -> str: + return "|".join(message.get_value() for message in messages) + + formatter.normalize_string_async = AsyncMock(side_effect=format_messages) + node = _make_tap_node(target=target) + node._prepended_conversation_config = PrependedConversationConfig(message_normalizer=formatter) + _set_tap_seed_boundary( + node=node, + target=target, + seed_messages=[Message.from_prompt(prompt="original seed", role="user")], + ) + node._objective = "objective" + parent_conversation_id = node.objective_target_conversation_id + + await node._send_prompt_to_target_async("parent first") + assert node._target_normalization_context + assert node._target_normalization_context.is_consumed + + cloned = node.duplicate() + cloned._objective = "objective" + cloned_conversation_id = cloned.objective_target_conversation_id + assert cloned_conversation_id != parent_conversation_id + assert cloned._target_normalization_context + assert not cloned._target_normalization_context.is_consumed + + await node._send_prompt_to_target_async("parent second") + await cloned._send_prompt_to_target_async("clone first") + await cloned._send_prompt_to_target_async("clone second") + + assert target.prompts_by_conversation == { + parent_conversation_id: ["original seed|parent first", "parent second"], + cloned_conversation_id: ["original seed|clone first", "clone second"], + } + + @pytest.mark.usefixtures("patch_central_database") def test_tap_branch_preserves_multimodal_last_response(): target = _RecordingTarget() diff --git a/tests/unit/executor/attack/multi_turn/test_supports_multi_turn_attacks.py b/tests/unit/executor/attack/multi_turn/test_supports_multi_turn_attacks.py index 69c155ea4c..bc699eeabb 100644 --- a/tests/unit/executor/attack/multi_turn/test_supports_multi_turn_attacks.py +++ b/tests/unit/executor/attack/multi_turn/test_supports_multi_turn_attacks.py @@ -490,8 +490,8 @@ def _make_tap_node(self, *, supports_multi_turn: bool): ), ) - def test_single_turn_target_duplicates_full_conversation(self): - """Single-turn branches retain their logical history for payload adaptation.""" + def test_single_turn_target_duplicates_logical_history_without_seed_boundary(self): + """Single-turn branches retain memory without replaying live turns as seed history.""" node = self._make_tap_node(supports_multi_turn=False) memory = CentralMemory.get_memory_instance() @@ -523,8 +523,7 @@ def test_single_turn_target_duplicates_full_conversation(self): dup_messages = memory.get_conversation_messages(conversation_id=duplicate.objective_target_conversation_id) assert [message.api_role for message in dup_messages] == ["system", "user", "assistant"] - assert duplicate._target_normalization_context is not None - assert duplicate._target_normalization_context.history_message_count == 3 + assert duplicate._target_normalization_context is None def test_multi_turn_target_duplicates_full_conversation(self): """For multi-turn targets, the full conversation is duplicated.""" diff --git a/tests/unit/prompt_target/target/test_target_normalization_context.py b/tests/unit/prompt_target/target/test_target_normalization_context.py index 359840faac..521b74547e 100644 --- a/tests/unit/prompt_target/target/test_target_normalization_context.py +++ b/tests/unit/prompt_target/target/test_target_normalization_context.py @@ -205,7 +205,7 @@ def test_context_remaps_only_explicit_history_for_duplicate_conversation(): assert duplicate.select_history(messages=duplicated_messages) == [duplicated_messages[0]] -def test_context_remap_preserves_consumed_state(): +def test_context_remap_resets_consumed_state_for_new_conversation(): seed = _message(role="user", value="seed", sequence=0) duplicated_seed = seed.duplicate() duplicated_seed.get_piece().conversation_id = "duplicate" @@ -224,5 +224,5 @@ def test_context_remap_preserves_consumed_state(): duplicated_messages=[duplicated_seed], ) - assert duplicate.is_consumed - assert duplicate.select_history(messages=[duplicated_seed]) == [] + assert not duplicate.is_consumed + assert duplicate.select_history(messages=[duplicated_seed]) == [duplicated_seed] From 3338f94e6db6b09103da4d54204899edafe128f1 Mon Sep 17 00:00:00 2001 From: Richard Lundeen Date: Fri, 21 Aug 2026 10:40:47 -0700 Subject: [PATCH 20/24] FIX Centralize prepended conversation policy Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> Copilot-Session: 420eef57-7a1c-4dde-8aad-a93539e8da64 (cherry picked from commit 3d0ba4ec697487ab9148f0822e1e85ddec9a3e41) --- doc/code/framework.md | 23 +--- doc/code/targets/11_message_normalizer.py | 27 +++++ .../attack/component/conversation_manager.py | 6 +- pyrit/executor/attack/core/attack_strategy.py | 43 ++++++- .../attack/multi_turn/chunked_request.py | 5 +- pyrit/executor/attack/multi_turn/crescendo.py | 13 ++- .../attack/multi_turn/multi_prompt_sending.py | 5 +- .../multi_turn/multi_turn_attack_strategy.py | 7 ++ .../executor/attack/multi_turn/red_teaming.py | 11 +- .../attack/multi_turn/tree_of_attacks.py | 10 +- .../attack/single_turn/prompt_sending.py | 7 +- .../single_turn_attack_strategy.py | 7 ++ .../attack/single_turn/skeleton_key.py | 7 ++ pyrit/executor/attack/streaming/barge_in.py | 4 +- .../history_squash_normalizer.py | 23 ++++ .../common/target_configuration.py | 8 +- .../component/test_conversation_manager.py | 56 ++++++++++ .../test_attack_strategy_prepended_policy.py | 105 ++++++++++++++++++ .../test_supports_multi_turn_attacks.py | 2 + .../test_normalize_async_integration.py | 33 ++++++ 20 files changed, 347 insertions(+), 55 deletions(-) create mode 100644 tests/unit/executor/attack/core/test_attack_strategy_prepended_policy.py diff --git a/doc/code/framework.md b/doc/code/framework.md index 55d568b55d..3a47fdba69 100644 --- a/doc/code/framework.md +++ b/doc/code/framework.md @@ -314,26 +314,13 @@ The below talks about responsibilities of most modules in the PyRIT library ## [Normalizers](./targets/11_message_normalizer) -**Responsibility**: Reshape prompts and conversations so components and targets can interoperate. There are two distinct modules: +**Responsibility**: Reshape prompts and conversations so components and targets can interoperate. -- **`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. Prepended history remains role-structured in memory. Attacks can pass per-send normalizer overrides from `PrependedConversationConfig`; the target's capability pipeline uses the `EDITABLE_HISTORY` override only when adaptation is required. -- **`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, prepended-history adaptation, and tokenizer chat templates. These target-specific views are ephemeral and do not replace the logical conversation in memory. +- **`prompt_normalizer`** applies converters, persists requests and responses, and dispatches prompts to a `PromptTarget`. Targets do not persist messages. +- **`message_normalizer`** reshapes conversations into target-compatible payloads. It owns target-facing representation, not attack policy or conversation state. +- **Does not own**: the conversation of record. Memory is canonical; a normalized payload is an ephemeral target-facing view that is never written back. -`supports_multi_turn` and `supports_editable_history` answer different questions. Multi-turn support means the target can continue a conversation across sends. Editable-history support means PyRIT can supply or rewrite earlier turns before sending the current message. The corresponding normalizers therefore have different scopes: - -| Target capabilities | Prepended-history behavior | -| --- | --- | -| Multi-turn with editable history | The target receives the structured history directly; neither normalizer is needed. | -| Multi-turn without editable history | A per-send `HistorySquashNormalizer` override encodes the prepended history and first live message into one initial request. Later turns use the target's own conversation state. | -| Single-turn without editable history | A per-send `HistorySquashNormalizer` override encodes the original prepended history with each current request. Prior live requests and responses are not replayed. The target's ordinary `MULTI_TURN` policy remains unchanged because the override produces one target-facing message. | - -The stateful/stateless distinction is therefore a matter of lifecycle and configuration, not a separate flattening implementation. The prepended-conversation flow records the persisted seed message IDs in a `TargetNormalizationContext` and supplies a per-send `HistorySquashNormalizer` override while that explicit history should be included. A target's ordinary capability pipeline independently applies `HistorySquashNormalizer` only when its `MULTI_TURN` behavior is explicitly configured for adaptation. Both uses preserve non-text pieces from the current request while rendering historical content as text. - -For prepended history on a target without editable history, processing order is: convert and persist the structured prepended messages, convert the live request, apply the per-send editable-history override, run the target's ordinary capability normalizers, then serialize and invoke the provider. Stateful targets consume the seed boundary when provider invocation begins and send only the current request afterward. Stateless targets replay the original seed with each current request. Unseeded stateless targets send only the current request. Branching attacks remap only the original seed boundary when duplicating memory conversations; copied live turns never become seed history. Pre-provider failures leave the seed reusable, while provider-attempt failures do not replay it for stateful targets. Failed `processing` and `unknown` responses are excluded from later target-facing history together with an immediately preceding user request only when it has the same conversation and adjacent sequence. `blocked` and `empty` responses remain replayable provider turns. A concurrent send sharing an active normalization context is rejected so stateful bootstrap context cannot be duplicated. - -The prepended-history formatter belongs to the attack configuration rather than the target instance: it affects only sends owned by that attack. Its behavioral identifier is therefore recorded as a child of the attack identifier, including formatter-specific configuration such as tokenizer templates and system-message handling. - -Streaming attacks that bypass `PromptTarget.send_prompt_async` do not use target-normalization overrides. `BargeInAttack` persists prepended messages and honors converter-role selection, but its direct realtime session uses only the leading system prompt and live audio; the prepended-history formatter and target-normalization context are intentionally inactive. +See [message normalizers](./targets/11_message_normalizer) for capability behavior, processing order, and prepended-history lifecycle details. ## [Output](./output/0_output) diff --git a/doc/code/targets/11_message_normalizer.py b/doc/code/targets/11_message_normalizer.py index 52e1ef59b8..aebe69d3a1 100644 --- a/doc/code/targets/11_message_normalizer.py +++ b/doc/code/targets/11_message_normalizer.py @@ -20,6 +20,33 @@ # # The `MessageNormalizer` classes handle these conversions, making it easy to work with any target regardless of its expected input format. # +# ## Memory is canonical, the normalized payload is not +# +# A normalizer builds a **target-facing view at send time**. It is never written back to memory. +# +# This matters most for a target that cannot accept editable history. If you prepend eight +# structured turns to a conversation, memory keeps eight structured turns, and the UI, scorers, +# resume, and exports all see eight turns. `HistorySquashNormalizer` flattens those turns into +# a single prompt only for the wire, then discards the flattened copy. The two representations +# are expected to differ. +# +# Flattening the conversation in memory instead would break scoring, resume, and evaluation for +# the sake of one target's wire format. If a prepended piece is not text (an image, for example), +# the flattened view holds a text placeholder and the normalizer logs a warning, because the +# target receives a description instead of the media. +# +# For prepended history on a target without editable history, PyRIT: +# +# 1. Converts and persists the structured prepended messages. +# 2. Converts the live request. +# 3. Applies the per-send `EDITABLE_HISTORY` normalizer selected by +# `PrependedConversationConfig`. +# 4. Applies the target's remaining capability normalizers. +# 5. Serializes the normalized view and invokes the provider. +# +# `TargetNormalizationContext` records the persisted prepended-message boundary. Stateful targets +# consume it after the first provider attempt; stateless targets reuse it for each current request. +# # ## Base Classes # # There are two base normalizer types: diff --git a/pyrit/executor/attack/component/conversation_manager.py b/pyrit/executor/attack/component/conversation_manager.py index 9161b15af1..bed543bd4f 100644 --- a/pyrit/executor/attack/component/conversation_manager.py +++ b/pyrit/executor/attack/component/conversation_manager.py @@ -539,7 +539,11 @@ async def _process_prepended_conversation_async( # true_false scores with score_value=False so attacks can use the rationale for # feedback without re-scoring. memory_pieces = self._memory.get_message_pieces(conversation_id=conversation_id) - assistant_piece_ids = [str(piece.id) for piece in memory_pieces if piece.api_role == "assistant"] + assistant_pieces = [piece for piece in memory_pieces if piece.api_role == "assistant"] + last_assistant_sequence = max((piece.sequence for piece in assistant_pieces), default=None) + assistant_piece_ids = [ + str(piece.id) for piece in assistant_pieces if piece.sequence == last_assistant_sequence + ] existing_scores = ( self._memory.get_prompt_scores(prompt_ids=assistant_piece_ids) if assistant_piece_ids else [] ) diff --git a/pyrit/executor/attack/core/attack_strategy.py b/pyrit/executor/attack/core/attack_strategy.py index 3147d3204f..9c390d6f4c 100644 --- a/pyrit/executor/attack/core/attack_strategy.py +++ b/pyrit/executor/attack/core/attack_strategy.py @@ -42,12 +42,17 @@ from pyrit.prompt_target.common.target_requirements import TargetRequirements if TYPE_CHECKING: + from pyrit.executor.attack.component.prepended_conversation_config import ( + PrependedConversationConfig, + ) from pyrit.executor.attack.core.attack_config import ( AttackAdversarialConfig, AttackScoringConfig, ) from pyrit.executor.attack.core.attack_result_attribution import AttackResultAttribution + from pyrit.message_normalizer import MessageListNormalizer from pyrit.prompt_target import PromptTarget + from pyrit.prompt_target.common.target_capabilities import CapabilityName from pyrit.prompt_target.common.target_normalization_context import TargetNormalizationContext AttackStrategyContextT = TypeVar("AttackStrategyContextT", bound="AttackContext[Any]") @@ -430,6 +435,7 @@ def __init__( objective_target: PromptTarget, context_type: type[AttackStrategyContextT], params_type: type[AttackParamsT] = AttackParameters, # type: ignore[ty:invalid-parameter-default] + prepended_conversation_config: PrependedConversationConfig | None = None, logger: logging.Logger = logger, ) -> None: """ @@ -441,6 +447,9 @@ def __init__( params_type (type[AttackParamsT]): The type of parameters this strategy accepts. Defaults to AttackParameters. Use AttackParameters.excluding() to create a params type that rejects certain fields. + prepended_conversation_config (PrependedConversationConfig | None): Policy for + prepended conversations. Controls converter role scope and target-facing + history formatting. logger (logging.Logger): Logger instance for logging events. """ super().__init__( @@ -450,15 +459,40 @@ def __init__( ), logger=logger, ) + # Local import avoids the component package's import cycle through attack config. + from pyrit.executor.attack.component.prepended_conversation_config import ( + PrependedConversationConfig, + ) + type(self).TARGET_REQUIREMENTS.validate(target=objective_target) self._objective_target = objective_target self._params_type = params_type + self._prepended_conversation_config = prepended_conversation_config or PrependedConversationConfig() # Guard so subclasses that set converters before calling super() aren't clobbered if not hasattr(self, "_request_converters"): self._request_converters: list[Any] = [] if not hasattr(self, "_response_converters"): self._response_converters: list[Any] = [] + def _get_prepended_normalizer_overrides( + self, + *, + target_normalization_context: TargetNormalizationContext | None, + ) -> dict[CapabilityName, MessageListNormalizer[Message]]: + """ + Resolve prepended-history overrides for one target send. + + Args: + target_normalization_context: Persisted seed boundary for this execution. + + Returns: + Overrides keyed by the capability they adapt. + """ + return self._prepended_conversation_config.get_normalizer_overrides( + target=self._objective_target, + target_normalization_context=target_normalization_context, + ) + def _create_identifier( self, *, @@ -486,12 +520,9 @@ def _create_identifier( objective_target = TargetIdentifier.from_component_identifier(self.get_objective_target().get_identifier()) - prepended_config = getattr(self, "_prepended_conversation_config", None) - if prepended_config is not None: - merged_params["prepended_conversation_converter_roles"] = list(prepended_config.apply_converters_to_roles) - all_children["prepended_conversation_formatter"] = ( - prepended_config.get_message_normalizer().get_identifier() - ) + prepended_config = self._prepended_conversation_config + merged_params["prepended_conversation_converter_roles"] = list(prepended_config.apply_converters_to_roles) + all_children["prepended_conversation_formatter"] = prepended_config.get_message_normalizer().get_identifier() # Add scorer if present objective_scorer: ScorerIdentifier | None = None diff --git a/pyrit/executor/attack/multi_turn/chunked_request.py b/pyrit/executor/attack/multi_turn/chunked_request.py index eef334c588..252665c846 100644 --- a/pyrit/executor/attack/multi_turn/chunked_request.py +++ b/pyrit/executor/attack/multi_turn/chunked_request.py @@ -153,6 +153,7 @@ def __init__( logger=logger, context_type=ChunkedRequestAttackContext, params_type=ChunkedRequestAttackParameters, + prepended_conversation_config=prepended_conversation_config, ) # Store chunk configuration @@ -174,7 +175,6 @@ def __init__( # Initialize prompt normalizer and conversation manager self._prompt_normalizer = prompt_normalizer or PromptNormalizer() - self._prepended_conversation_config = prepended_conversation_config or PrependedConversationConfig() self._conversation_manager = ConversationManager( prompt_normalizer=self._prompt_normalizer, ) @@ -293,8 +293,7 @@ async def _perform_async(self, *, context: ChunkedRequestAttackContext) -> Attac conversation_id=context.session.conversation_id, request_converter_configurations=self._request_converters, response_converter_configurations=self._response_converters, - normalizer_overrides=self._prepended_conversation_config.get_normalizer_overrides( - target=self._objective_target, + normalizer_overrides=self._get_prepended_normalizer_overrides( target_normalization_context=context.target_normalization_context, ), target_normalization_context=context.target_normalization_context, diff --git a/pyrit/executor/attack/multi_turn/crescendo.py b/pyrit/executor/attack/multi_turn/crescendo.py index 061f28cba6..dc96a5eefd 100644 --- a/pyrit/executor/attack/multi_turn/crescendo.py +++ b/pyrit/executor/attack/multi_turn/crescendo.py @@ -177,7 +177,12 @@ def __init__( and editable history. """ # Initialize base class - super().__init__(objective_target=objective_target, logger=logger, context_type=CrescendoAttackContext) + super().__init__( + objective_target=objective_target, + logger=logger, + context_type=CrescendoAttackContext, + prepended_conversation_config=prepended_conversation_config, + ) self._memory = CentralMemory.get_memory_instance() @@ -261,9 +266,6 @@ def __init__( self._max_backtracks = max_backtracks self._max_turns = max_turns - # Store the prepended conversation configuration - self._prepended_conversation_config = prepended_conversation_config or PrependedConversationConfig() - def get_attack_scoring_config(self) -> AttackScoringConfig | None: """ Get the attack scoring configuration used by this strategy. @@ -639,8 +641,7 @@ async def _send_prompt_to_objective_target_async( conversation_id=context.session.conversation_id, request_converter_configurations=self._request_converters, response_converter_configurations=self._response_converters, - normalizer_overrides=self._prepended_conversation_config.get_normalizer_overrides( - target=self._objective_target, + normalizer_overrides=self._get_prepended_normalizer_overrides( target_normalization_context=context.target_normalization_context, ), target_normalization_context=context.target_normalization_context, diff --git a/pyrit/executor/attack/multi_turn/multi_prompt_sending.py b/pyrit/executor/attack/multi_turn/multi_prompt_sending.py index 7b6d318661..8613a4fa4e 100644 --- a/pyrit/executor/attack/multi_turn/multi_prompt_sending.py +++ b/pyrit/executor/attack/multi_turn/multi_prompt_sending.py @@ -164,6 +164,7 @@ def __init__( logger=logger, context_type=MultiTurnAttackContext, params_type=MultiPromptSendingAttackParameters, + prepended_conversation_config=prepended_conversation_config, ) # Initialize the converter configuration @@ -179,7 +180,6 @@ def __init__( # Initialize prompt normalizer and conversation manager self._prompt_normalizer = prompt_normalizer or PromptNormalizer() - self._prepended_conversation_config = prepended_conversation_config or PrependedConversationConfig() self._conversation_manager = ConversationManager( prompt_normalizer=self._prompt_normalizer, ) @@ -371,8 +371,7 @@ async def _send_prompt_to_objective_target_async( conversation_id=context.session.conversation_id, request_converter_configurations=self._request_converters, response_converter_configurations=self._response_converters, - normalizer_overrides=self._prepended_conversation_config.get_normalizer_overrides( - target=self._objective_target, + normalizer_overrides=self._get_prepended_normalizer_overrides( target_normalization_context=context.target_normalization_context, ), target_normalization_context=context.target_normalization_context, diff --git a/pyrit/executor/attack/multi_turn/multi_turn_attack_strategy.py b/pyrit/executor/attack/multi_turn/multi_turn_attack_strategy.py index b3584587e8..c035ec971c 100644 --- a/pyrit/executor/attack/multi_turn/multi_turn_attack_strategy.py +++ b/pyrit/executor/attack/multi_turn/multi_turn_attack_strategy.py @@ -22,6 +22,9 @@ from pyrit.prompt_target import CapabilityName if TYPE_CHECKING: + from pyrit.executor.attack.component.prepended_conversation_config import ( + PrependedConversationConfig, + ) from pyrit.models import ( Message, Score, @@ -78,6 +81,7 @@ def __init__( objective_target: PromptTarget, context_type: type[MultiTurnAttackStrategyContextT], params_type: type[AttackParamsT] = AttackParameters, # type: ignore[ty:invalid-parameter-default] + prepended_conversation_config: PrependedConversationConfig | None = None, logger: logging.Logger = logger, ) -> None: """ @@ -87,12 +91,15 @@ def __init__( objective_target (PromptTarget): The target system to attack. context_type (type[MultiTurnAttackContext]): The type of context this strategy will use. params_type (type[AttackParamsT]): The type of parameters this strategy accepts. + prepended_conversation_config (PrependedConversationConfig | None): Policy for + prepended conversations. See ``AttackStrategy``. logger (logging.Logger): Logger instance for logging events and messages. """ super().__init__( objective_target=objective_target, context_type=context_type, params_type=params_type, + prepended_conversation_config=prepended_conversation_config, logger=logger, ) diff --git a/pyrit/executor/attack/multi_turn/red_teaming.py b/pyrit/executor/attack/multi_turn/red_teaming.py index 4431b50a4d..a66e45e157 100644 --- a/pyrit/executor/attack/multi_turn/red_teaming.py +++ b/pyrit/executor/attack/multi_turn/red_teaming.py @@ -119,7 +119,12 @@ def __init__( ValueError: If objective_scorer is not provided in attack_scoring_config. """ # Initialize base class - super().__init__(objective_target=objective_target, logger=logger, context_type=MultiTurnAttackContext) + super().__init__( + objective_target=objective_target, + logger=logger, + context_type=MultiTurnAttackContext, + prepended_conversation_config=prepended_conversation_config, + ) self._memory = CentralMemory.get_memory_instance() # Initialize converter configuration @@ -172,7 +177,6 @@ def __init__( # Initialize utilities self._prompt_normalizer = prompt_normalizer or PromptNormalizer() - self._prepended_conversation_config = prepended_conversation_config or PrependedConversationConfig() self._conversation_manager = ConversationManager(prompt_normalizer=self._prompt_normalizer) @@ -490,8 +494,7 @@ async def _send_prompt_to_objective_target_async( request_converter_configurations=self._request_converters, response_converter_configurations=self._response_converters, target=self._objective_target, - normalizer_overrides=self._prepended_conversation_config.get_normalizer_overrides( - target=self._objective_target, + normalizer_overrides=self._get_prepended_normalizer_overrides( target_normalization_context=context.target_normalization_context, ), target_normalization_context=context.target_normalization_context, diff --git a/pyrit/executor/attack/multi_turn/tree_of_attacks.py b/pyrit/executor/attack/multi_turn/tree_of_attacks.py index 983862c16f..e68c7b4e31 100644 --- a/pyrit/executor/attack/multi_turn/tree_of_attacks.py +++ b/pyrit/executor/attack/multi_turn/tree_of_attacks.py @@ -1507,7 +1507,12 @@ def __init__( ) # Initialize base class - super().__init__(objective_target=objective_target, logger=logger, context_type=TAPAttackContext) + super().__init__( + objective_target=objective_target, + logger=logger, + context_type=TAPAttackContext, + prepended_conversation_config=prepended_conversation_config, + ) self._memory = CentralMemory.get_memory_instance() self._node_executor = _TreeOfAttacksNodeExecutor( @@ -1619,9 +1624,6 @@ def __init__( self._prompt_normalizer = prompt_normalizer or PromptNormalizer() - # Store the prepended conversation configuration - self._prepended_conversation_config = prepended_conversation_config or PrependedConversationConfig() - def _load_adversarial_prompts(self) -> None: """Load the adversarial chat prompt template and seed prompt from the default paths.""" # Load prompt template diff --git a/pyrit/executor/attack/single_turn/prompt_sending.py b/pyrit/executor/attack/single_turn/prompt_sending.py index 62edb3b5f4..de94336eb4 100644 --- a/pyrit/executor/attack/single_turn/prompt_sending.py +++ b/pyrit/executor/attack/single_turn/prompt_sending.py @@ -89,6 +89,7 @@ def __init__( logger=logger, context_type=SingleTurnAttackContext, params_type=params_type, + prepended_conversation_config=prepended_conversation_config, ) # Initialize the converter configuration @@ -117,9 +118,6 @@ def __init__( self._max_attempts_on_failure = max_attempts_on_failure - # Store the prepended conversation configuration - self._prepended_conversation_config = prepended_conversation_config or PrependedConversationConfig() - def get_attack_scoring_config(self) -> AttackScoringConfig | None: """ Get the attack scoring configuration used by this strategy. @@ -325,8 +323,7 @@ async def _send_prompt_to_objective_target_async( conversation_id=context.conversation_id, request_converter_configurations=self._request_converters, response_converter_configurations=self._response_converters, - normalizer_overrides=self._prepended_conversation_config.get_normalizer_overrides( - target=self._objective_target, + normalizer_overrides=self._get_prepended_normalizer_overrides( target_normalization_context=context.target_normalization_context, ), target_normalization_context=context.target_normalization_context, diff --git a/pyrit/executor/attack/single_turn/single_turn_attack_strategy.py b/pyrit/executor/attack/single_turn/single_turn_attack_strategy.py index 1f699e3654..7dd7c32d17 100644 --- a/pyrit/executor/attack/single_turn/single_turn_attack_strategy.py +++ b/pyrit/executor/attack/single_turn/single_turn_attack_strategy.py @@ -15,6 +15,9 @@ from pyrit.models import AttackResult if TYPE_CHECKING: + from pyrit.executor.attack.component.prepended_conversation_config import ( + PrependedConversationConfig, + ) from pyrit.prompt_target import PromptTarget @@ -48,6 +51,7 @@ def __init__( objective_target: PromptTarget, context_type: type[SingleTurnAttackContext[Any]] = SingleTurnAttackContext, params_type: type[AttackParamsT] = AttackParameters, # type: ignore[ty:invalid-parameter-default] + prepended_conversation_config: PrependedConversationConfig | None = None, logger: logging.Logger = logger, ) -> None: """ @@ -57,11 +61,14 @@ def __init__( objective_target (PromptTarget): The target system to attack. context_type (type[SingleTurnAttackContext]): The type of context this strategy will use. params_type (type[AttackParamsT]): The type of parameters this strategy accepts. + prepended_conversation_config (PrependedConversationConfig | None): Policy for + prepended conversations. See ``AttackStrategy``. logger (logging.Logger): Logger instance for logging events and messages. """ super().__init__( objective_target=objective_target, context_type=context_type, params_type=params_type, + prepended_conversation_config=prepended_conversation_config, logger=logger, ) diff --git a/pyrit/executor/attack/single_turn/skeleton_key.py b/pyrit/executor/attack/single_turn/skeleton_key.py index 5e9d33fbf3..e283153587 100644 --- a/pyrit/executor/attack/single_turn/skeleton_key.py +++ b/pyrit/executor/attack/single_turn/skeleton_key.py @@ -7,6 +7,7 @@ from pyrit.common.apply_defaults import REQUIRED_VALUE, apply_defaults from pyrit.common.path import EXECUTOR_SEED_PROMPT_PATH +from pyrit.executor.attack.component import PrependedConversationConfig from pyrit.executor.attack.core.attack_config import AttackConverterConfig, AttackScoringConfig from pyrit.executor.attack.core.attack_parameters import AttackParameters from pyrit.executor.attack.single_turn.prompt_sending import PromptSendingAttack @@ -59,6 +60,7 @@ def __init__( skeleton_key_prompt: str | None = None, skeleton_key_acceptance: str | None = None, max_attempts_on_failure: int = 0, + prepended_conversation_config: PrependedConversationConfig | None = None, ) -> None: """ Initialize the skeleton key attack strategy. @@ -73,6 +75,10 @@ def __init__( skeleton_key_acceptance (str | None): The simulated assistant acceptance response to prepend. If not provided, uses the default acceptance response. max_attempts_on_failure (int): Maximum number of attempts to retry on failure. + prepended_conversation_config (PrependedConversationConfig | None): Policy for the + skeleton key exchange this attack prepends. Controls which roles receive request + converters and how a target without editable history renders that exchange with + the live request. """ super().__init__( objective_target=objective_target, @@ -81,6 +87,7 @@ def __init__( prompt_normalizer=prompt_normalizer, max_attempts_on_failure=max_attempts_on_failure, params_type=SkeletonKeyAttackParameters, + prepended_conversation_config=prepended_conversation_config, ) self._skeleton_key_prompt = ( diff --git a/pyrit/executor/attack/streaming/barge_in.py b/pyrit/executor/attack/streaming/barge_in.py index 5ebfbe690b..894f6e48d8 100644 --- a/pyrit/executor/attack/streaming/barge_in.py +++ b/pyrit/executor/attack/streaming/barge_in.py @@ -11,7 +11,6 @@ from typing import TYPE_CHECKING, Any, ClassVar, cast from pyrit.common.apply_defaults import REQUIRED_VALUE, apply_defaults -from pyrit.executor.attack.component import PrependedConversationConfig from pyrit.executor.attack.component.conversation_manager import ConversationManager from pyrit.executor.attack.core.attack_config import AttackConverterConfig from pyrit.executor.attack.core.attack_parameters import AttackParameters, AttackParamsT @@ -29,6 +28,7 @@ if TYPE_CHECKING: from collections.abc import AsyncIterator + from pyrit.executor.attack.component import PrependedConversationConfig from pyrit.prompt_target import PromptTarget, RealtimeTarget logger = logging.getLogger(__name__) @@ -96,6 +96,7 @@ def __init__( objective_target=objective_target, context_type=BargeInAttackContext, params_type=params_type, + prepended_conversation_config=prepended_conversation_config, logger=logger, ) self._realtime_target = cast("RealtimeTarget", objective_target) @@ -103,7 +104,6 @@ def __init__( self._request_converters = attack_converter_config.request_converters self._response_converters = attack_converter_config.response_converters self._prompt_normalizer = prompt_normalizer or PromptNormalizer() - self._prepended_conversation_config = prepended_conversation_config or PrependedConversationConfig() self._conversation_manager = ConversationManager( prompt_normalizer=self._prompt_normalizer, ) diff --git a/pyrit/message_normalizer/history_squash_normalizer.py b/pyrit/message_normalizer/history_squash_normalizer.py index 8ff5b7f307..b437c02f0c 100644 --- a/pyrit/message_normalizer/history_squash_normalizer.py +++ b/pyrit/message_normalizer/history_squash_normalizer.py @@ -2,6 +2,7 @@ # Licensed under the MIT license. import copy +import logging import uuid from pyrit.message_normalizer._helpers import format_message_piece_for_context @@ -10,6 +11,8 @@ from pyrit.message_normalizer.message_normalizer import MessageListNormalizer, MessageStringNormalizer from pyrit.models import Message, MessagePiece +logger = logging.getLogger(__name__) + class HistorySquashNormalizer(MessageListNormalizer[Message]): """ @@ -81,6 +84,7 @@ async def normalize_async(self, messages: list[Message]) -> list[Message]: history = messages[:-1] live_request = messages[-1] self._validate_flattenable_converter_output(messages=history) + self._warn_on_non_text_history(messages=history) original_view = self._build_original_view(messages=messages) converted_view = self._build_converted_view(messages=messages) @@ -219,6 +223,25 @@ def _validate_flattenable_converter_output(*, messages: list[Message]) -> None: f"non-text output types {sorted(output_types)}. Historical conversion must produce text." ) + @staticmethod + def _warn_on_non_text_history(*, messages: list[Message]) -> None: + """Warn when native non-text history becomes target-facing text.""" + flattened_types = sorted( + { + piece.converted_value_data_type + for message in messages + for piece in message.message_pieces + if piece.converted_value_data_type != "text" + } + ) + if flattened_types: + logger.warning( + "Conversation history contains non-text pieces %s. History squashing " + "represents them as text placeholders for the target; memory keeps the " + "original pieces.", + flattened_types, + ) + @staticmethod def _build_target_request( *, diff --git a/pyrit/prompt_target/common/target_configuration.py b/pyrit/prompt_target/common/target_configuration.py index 7aea861185..c5b19fa259 100644 --- a/pyrit/prompt_target/common/target_configuration.py +++ b/pyrit/prompt_target/common/target_configuration.py @@ -188,9 +188,11 @@ def as_identifier_params(self) -> dict[str, Any]: for capability, behavior in self._policy.behaviors.items() if not caps.includes(capability=capability) }, - # Stable, ordered representation of the resolved normalization - # pipeline. Captures the effect of ``normalizer_overrides`` since - # the pipeline is built from defaults + overrides. + # Stable, ordered representation of the pipeline this configuration was + # built with. Per-send ``normalizer_overrides`` are NOT reflected here: + # ``normalize_async`` builds a throwaway pipeline and never mutates + # ``self._pipeline``. Attack-owned overrides are represented by the + # attack identifier instead. "normalization_pipeline": [ f"{type(normalizer).__module__}.{type(normalizer).__qualname__}" for normalizer in self._pipeline.normalizers diff --git a/tests/unit/executor/attack/component/test_conversation_manager.py b/tests/unit/executor/attack/component/test_conversation_manager.py index 2d61cc79fd..b6faf11a1c 100644 --- a/tests/unit/executor/attack/component/test_conversation_manager.py +++ b/tests/unit/executor/attack/component/test_conversation_manager.py @@ -914,6 +914,62 @@ async def test_multipart_message_extracts_scores_from_all_pieces( assert score1.id in returned_ids assert score2.id in returned_ids + async def test_scores_come_only_from_the_last_assistant_turn( + self, + attack_identifier: ComponentIdentifier, + mock_chat_target: MagicMock, + ) -> None: + """Only the final assistant turn's scores are surfaced, not every assistant turn.""" + manager = ConversationManager() + conversation_id = str(uuid.uuid4()) + context = _TestAttackContext(params=AttackParameters(objective="Test objective")) + + early_piece = MessagePiece( + role="assistant", + original_value="early reply", + conversation_id=str(uuid.uuid4()), + ) + final_piece = MessagePiece( + role="assistant", + original_value="final reply", + conversation_id=str(uuid.uuid4()), + ) + manager._memory.add_message_pieces_to_memory(message_pieces=[early_piece, final_piece]) + + def _false_score(piece: MessagePiece, rationale: str) -> Score: + return Score( + score_type="true_false", + score_value="false", + score_category=["test"], + score_value_description=rationale, + score_rationale=rationale, + score_metadata={}, + message_piece_id=str(piece.id), + scorer_class_identifier=get_mock_scorer_identifier(), + ) + + early_score = _false_score(early_piece, "early") + final_score = _false_score(final_piece, "final") + manager._memory.add_scores_to_memory(scores=[early_score, final_score]) + + context.prepended_conversation = [ + Message.from_prompt(prompt="first ask", role="user"), + Message(message_pieces=[early_piece]), + Message.from_prompt(prompt="second ask", role="user"), + Message(message_pieces=[final_piece]), + ] + + state = await manager.initialize_context_async( + context=context, + target=mock_chat_target, + conversation_id=conversation_id, + max_turns=10, + ) + + assert [score.id for score in state.last_assistant_message_scores] == [final_score.id] + assert context.last_score is not None + assert context.last_score.id == final_score.id + async def test_prepended_conversation_ignores_true_scores( self, attack_identifier: ComponentIdentifier, diff --git a/tests/unit/executor/attack/core/test_attack_strategy_prepended_policy.py b/tests/unit/executor/attack/core/test_attack_strategy_prepended_policy.py new file mode 100644 index 0000000000..939c2b22c2 --- /dev/null +++ b/tests/unit/executor/attack/core/test_attack_strategy_prepended_policy.py @@ -0,0 +1,105 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT license. + +"""Tests for prepended-conversation policy owned by ``AttackStrategy``.""" + +import inspect +import uuid + +import pytest + +from pyrit.executor.attack.component import PrependedConversationConfig +from pyrit.executor.attack.core.attack_config import AttackScoringConfig +from pyrit.executor.attack.multi_turn.chunked_request import ChunkedRequestAttack +from pyrit.executor.attack.multi_turn.crescendo import CrescendoAttack +from pyrit.executor.attack.multi_turn.multi_prompt_sending import MultiPromptSendingAttack +from pyrit.executor.attack.multi_turn.pair import PAIRAttack +from pyrit.executor.attack.multi_turn.red_teaming import RedTeamingAttack +from pyrit.executor.attack.multi_turn.tree_of_attacks import TreeOfAttacksWithPruningAttack +from pyrit.executor.attack.single_turn.prompt_sending import PromptSendingAttack +from pyrit.executor.attack.single_turn.skeleton_key import SkeletonKeyAttack +from pyrit.message_normalizer import HistorySquashNormalizer +from pyrit.models import Message, MessagePiece +from pyrit.prompt_target import CapabilityName, PromptTarget, TargetCapabilities, TargetConfiguration +from pyrit.prompt_target.common.target_normalization_context import TargetNormalizationContext +from pyrit.scenario.core.attack_technique_factory import AttackTechniqueFactory + + +class _NonEditableHistoryTarget(PromptTarget): + _DEFAULT_CONFIGURATION = TargetConfiguration(capabilities=TargetCapabilities(supports_editable_history=False)) + + async def _send_prompt_to_target_async(self, *, normalized_conversation: list[Message]) -> list[Message]: + request = normalized_conversation[-1] + return [ + MessagePiece( + role="assistant", + original_value="response", + conversation_id=request.get_piece().conversation_id, + ).to_message() + ] + + +@pytest.mark.usefixtures("patch_central_database") +def test_attack_strategy_owns_prepended_policy_and_identifier(): + config = PrependedConversationConfig(apply_converters_to_roles=["user", "assistant"]) + attack = PromptSendingAttack( + objective_target=_NonEditableHistoryTarget(), + prepended_conversation_config=config, + ) + + identifier = attack.get_identifier() + assert attack._prepended_conversation_config is config + assert identifier.params["prepended_conversation_converter_roles"] == ["user", "assistant"] + assert "prepended_conversation_formatter" in identifier.children + + +@pytest.mark.usefixtures("patch_central_database") +def test_attack_strategy_resolves_per_send_history_override(): + attack = PromptSendingAttack(objective_target=_NonEditableHistoryTarget()) + context = TargetNormalizationContext( + conversation_id="conversation", + history_message_ids=(uuid.uuid4(),), + replay_history_each_send=False, + ) + + overrides = attack._get_prepended_normalizer_overrides(target_normalization_context=context) + + assert isinstance(overrides[CapabilityName.EDITABLE_HISTORY], HistorySquashNormalizer) + + +@pytest.mark.parametrize( + "attack_class", + [ + ChunkedRequestAttack, + CrescendoAttack, + MultiPromptSendingAttack, + PAIRAttack, + PromptSendingAttack, + RedTeamingAttack, + SkeletonKeyAttack, + TreeOfAttacksWithPruningAttack, + ], +) +def test_techniques_can_specify_prepended_policy(attack_class): + """Each attack that creates or accepts prepended history exposes the policy.""" + parameter = inspect.signature(attack_class.__init__).parameters.get("prepended_conversation_config") + + assert parameter is not None + assert parameter.kind is inspect.Parameter.KEYWORD_ONLY + + +@pytest.mark.usefixtures("patch_central_database") +def test_technique_factory_forwards_prepended_policy(): + config = PrependedConversationConfig(apply_converters_to_roles=["user", "assistant"]) + factory = AttackTechniqueFactory( + name="policy_test", + attack_class=PromptSendingAttack, + attack_kwargs={"prepended_conversation_config": config}, + ) + + attack = factory.create( + objective_target=_NonEditableHistoryTarget(), + attack_scoring_config=AttackScoringConfig(), + ).attack + + assert attack._prepended_conversation_config is config diff --git a/tests/unit/executor/attack/multi_turn/test_supports_multi_turn_attacks.py b/tests/unit/executor/attack/multi_turn/test_supports_multi_turn_attacks.py index bc699eeabb..9809ce0a47 100644 --- a/tests/unit/executor/attack/multi_turn/test_supports_multi_turn_attacks.py +++ b/tests/unit/executor/attack/multi_turn/test_supports_multi_turn_attacks.py @@ -523,6 +523,8 @@ def test_single_turn_target_duplicates_logical_history_without_seed_boundary(sel dup_messages = memory.get_conversation_messages(conversation_id=duplicate.objective_target_conversation_id) assert [message.api_role for message in dup_messages] == ["system", "user", "assistant"] + # No explicit prepended seed was initialized, so copied live turns do not + # become a new replay boundary. assert duplicate._target_normalization_context is None def test_multi_turn_target_duplicates_full_conversation(self): diff --git a/tests/unit/prompt_target/target/test_normalize_async_integration.py b/tests/unit/prompt_target/target/test_normalize_async_integration.py index c480f6ffd1..acb13d453a 100644 --- a/tests/unit/prompt_target/target/test_normalize_async_integration.py +++ b/tests/unit/prompt_target/target/test_normalize_async_integration.py @@ -5,6 +5,7 @@ import asyncio import json +import logging from typing import TYPE_CHECKING from unittest.mock import AsyncMock, MagicMock, patch @@ -1051,6 +1052,38 @@ async def test_non_editable_target_allows_preexisting_non_text_history_with_conv assert result[0].get_value() == "Turn 1:\nuser: [Image_path]\nTurn 2:\nuser: live" +@pytest.mark.usefixtures("patch_central_database") +async def test_non_editable_target_warns_when_non_text_history_becomes_a_placeholder(caplog): + target = MockPromptTarget() + target._configuration = TargetConfiguration(capabilities=TargetCapabilities()) + prepended = Message( + message_pieces=[ + MessagePiece( + role="user", + original_value="existing.png", + converted_value="existing.png", + original_value_data_type="image_path", + converted_value_data_type="image_path", + conversation_id="conv1", + ) + ] + ) + mock_memory = MagicMock(spec=MemoryInterface) + mock_memory.get_conversation_messages.return_value = [prepended] + target._memory = mock_memory + target_context = _make_target_normalization_context(prepended_messages=[prepended]) + + with caplog.at_level(logging.WARNING, logger="pyrit.message_normalizer.history_squash_normalizer"): + await target._get_normalized_conversation_async( + message=_make_message(role="user", content="live"), + normalizer_overrides=_make_normalizer_overrides(target_normalization_context=target_context), + target_normalization_context=target_context, + ) + + assert "image_path" in caplog.text + assert "text placeholders" in caplog.text + + @pytest.mark.usefixtures("patch_central_database") async def test_target_normalization_failure_can_be_retried(): target = MockPromptTarget() From 2ce05b7626129a96d293e299c7cf59d104c45fdc Mon Sep 17 00:00:00 2001 From: Copilot <223556219+Copilot@users.noreply.github.com> Date: Fri, 21 Aug 2026 11:33:57 -0700 Subject: [PATCH 21/24] Synchronize centralized prepended policy docs Add complete annotations to the new regression tests and synchronize the paired message-normalizer notebook without dropping outputs. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> Copilot-Session: 6776aa95-0924-4967-a056-e5a228af3faa --- doc/code/targets/11_message_normalizer.ipynb | 27 +++++++++++++++++++ .../test_attack_strategy_prepended_policy.py | 10 ++++--- .../test_normalize_async_integration.py | 4 ++- 3 files changed, 36 insertions(+), 5 deletions(-) diff --git a/doc/code/targets/11_message_normalizer.ipynb b/doc/code/targets/11_message_normalizer.ipynb index 0261daa565..200c842d61 100644 --- a/doc/code/targets/11_message_normalizer.ipynb +++ b/doc/code/targets/11_message_normalizer.ipynb @@ -16,6 +16,33 @@ "\n", "The `MessageNormalizer` classes handle these conversions, making it easy to work with any target regardless of its expected input format.\n", "\n", + "## Memory is canonical, the normalized payload is not\n", + "\n", + "A normalizer builds a **target-facing view at send time**. It is never written back to memory.\n", + "\n", + "This matters most for a target that cannot accept editable history. If you prepend eight\n", + "structured turns to a conversation, memory keeps eight structured turns, and the UI, scorers,\n", + "resume, and exports all see eight turns. `HistorySquashNormalizer` flattens those turns into\n", + "a single prompt only for the wire, then discards the flattened copy. The two representations\n", + "are expected to differ.\n", + "\n", + "Flattening the conversation in memory instead would break scoring, resume, and evaluation for\n", + "the sake of one target's wire format. If a prepended piece is not text (an image, for example),\n", + "the flattened view holds a text placeholder and the normalizer logs a warning, because the\n", + "target receives a description instead of the media.\n", + "\n", + "For prepended history on a target without editable history, PyRIT:\n", + "\n", + "1. Converts and persists the structured prepended messages.\n", + "2. Converts the live request.\n", + "3. Applies the per-send `EDITABLE_HISTORY` normalizer selected by\n", + " `PrependedConversationConfig`.\n", + "4. Applies the target's remaining capability normalizers.\n", + "5. Serializes the normalized view and invokes the provider.\n", + "\n", + "`TargetNormalizationContext` records the persisted prepended-message boundary. Stateful targets\n", + "consume it after the first provider attempt; stateless targets reuse it for each current request.\n", + "\n", "## Base Classes\n", "\n", "There are two base normalizer types:\n", diff --git a/tests/unit/executor/attack/core/test_attack_strategy_prepended_policy.py b/tests/unit/executor/attack/core/test_attack_strategy_prepended_policy.py index 939c2b22c2..df01cfa0ae 100644 --- a/tests/unit/executor/attack/core/test_attack_strategy_prepended_policy.py +++ b/tests/unit/executor/attack/core/test_attack_strategy_prepended_policy.py @@ -5,10 +5,12 @@ import inspect import uuid +from typing import Any import pytest from pyrit.executor.attack.component import PrependedConversationConfig +from pyrit.executor.attack.core import AttackStrategy from pyrit.executor.attack.core.attack_config import AttackScoringConfig from pyrit.executor.attack.multi_turn.chunked_request import ChunkedRequestAttack from pyrit.executor.attack.multi_turn.crescendo import CrescendoAttack @@ -40,7 +42,7 @@ async def _send_prompt_to_target_async(self, *, normalized_conversation: list[Me @pytest.mark.usefixtures("patch_central_database") -def test_attack_strategy_owns_prepended_policy_and_identifier(): +def test_attack_strategy_owns_prepended_policy_and_identifier() -> None: config = PrependedConversationConfig(apply_converters_to_roles=["user", "assistant"]) attack = PromptSendingAttack( objective_target=_NonEditableHistoryTarget(), @@ -54,7 +56,7 @@ def test_attack_strategy_owns_prepended_policy_and_identifier(): @pytest.mark.usefixtures("patch_central_database") -def test_attack_strategy_resolves_per_send_history_override(): +def test_attack_strategy_resolves_per_send_history_override() -> None: attack = PromptSendingAttack(objective_target=_NonEditableHistoryTarget()) context = TargetNormalizationContext( conversation_id="conversation", @@ -80,7 +82,7 @@ def test_attack_strategy_resolves_per_send_history_override(): TreeOfAttacksWithPruningAttack, ], ) -def test_techniques_can_specify_prepended_policy(attack_class): +def test_techniques_can_specify_prepended_policy(attack_class: type[AttackStrategy[Any, Any]]) -> None: """Each attack that creates or accepts prepended history exposes the policy.""" parameter = inspect.signature(attack_class.__init__).parameters.get("prepended_conversation_config") @@ -89,7 +91,7 @@ def test_techniques_can_specify_prepended_policy(attack_class): @pytest.mark.usefixtures("patch_central_database") -def test_technique_factory_forwards_prepended_policy(): +def test_technique_factory_forwards_prepended_policy() -> None: config = PrependedConversationConfig(apply_converters_to_roles=["user", "assistant"]) factory = AttackTechniqueFactory( name="policy_test", diff --git a/tests/unit/prompt_target/target/test_normalize_async_integration.py b/tests/unit/prompt_target/target/test_normalize_async_integration.py index acb13d453a..065fa111fe 100644 --- a/tests/unit/prompt_target/target/test_normalize_async_integration.py +++ b/tests/unit/prompt_target/target/test_normalize_async_integration.py @@ -1053,7 +1053,9 @@ async def test_non_editable_target_allows_preexisting_non_text_history_with_conv @pytest.mark.usefixtures("patch_central_database") -async def test_non_editable_target_warns_when_non_text_history_becomes_a_placeholder(caplog): +async def test_non_editable_target_warns_when_non_text_history_becomes_a_placeholder( + caplog: pytest.LogCaptureFixture, +) -> None: target = MockPromptTarget() target._configuration = TargetConfiguration(capabilities=TargetCapabilities()) prepended = Message( From 8a1dbd5fcbfcae33d5b94c837b4434402f2d7450 Mon Sep 17 00:00:00 2001 From: Copilot <223556219+Copilot@users.noreply.github.com> Date: Fri, 21 Aug 2026 14:33:57 -0700 Subject: [PATCH 22/24] Move prepended history state to attacks Keep the provider-attempt state machine intact while hiding it behind an internal target send protocol and removing the concrete context from the prompt-target API. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> Copilot-Session: 6776aa95-0924-4967-a056-e5a228af3faa --- doc/code/framework.md | 2 + doc/code/targets/0_prompt_targets.md | 17 +- doc/code/targets/11_message_normalizer.ipynb | 14 +- doc/code/targets/11_message_normalizer.py | 14 +- .../attack/component/conversation_manager.py | 22 +- .../prepended_conversation_config.py | 14 +- .../prepended_history_send_context.py} | 142 ++++------- pyrit/executor/attack/core/attack_strategy.py | 14 +- .../attack/multi_turn/chunked_request.py | 4 +- pyrit/executor/attack/multi_turn/crescendo.py | 4 +- .../attack/multi_turn/multi_prompt_sending.py | 4 +- .../multi_turn/multi_turn_attack_strategy.py | 4 +- .../executor/attack/multi_turn/red_teaming.py | 4 +- .../attack/multi_turn/tree_of_attacks.py | 24 +- .../attack/single_turn/prompt_sending.py | 4 +- pyrit/executor/attack/streaming/barge_in.py | 2 +- .../history_squash_normalizer.py | 2 +- pyrit/prompt_normalizer/normalizer_request.py | 10 +- pyrit/prompt_normalizer/prompt_normalizer.py | 18 +- pyrit/prompt_target/__init__.py | 2 - pyrit/prompt_target/common/prompt_target.py | 43 ++-- pyrit/prompt_target/common/target_history.py | 46 ++++ .../common/target_send_context.py | 36 +++ .../component/test_conversation_manager.py | 14 +- .../test_prepended_history_send_context.py | 177 ++++++++++++++ .../test_attack_strategy_prepended_policy.py | 12 +- .../attack/multi_turn/test_chunked_request.py | 14 +- .../multi_turn/test_multi_prompt_sending.py | 14 +- .../test_prepended_history_normalization.py | 52 ++-- .../attack/multi_turn/test_red_teaming.py | 2 +- .../test_supports_multi_turn_attacks.py | 21 +- .../attack/multi_turn/test_tree_of_attacks.py | 2 +- .../attack/streaming/test_barge_in.py | 15 +- .../test_prompt_normalizer.py | 41 ++-- .../test_normalize_async_integration.py | 171 ++++++------- .../target/test_target_history.py | 71 ++++++ .../test_target_normalization_context.py | 228 ------------------ 37 files changed, 680 insertions(+), 600 deletions(-) rename pyrit/{prompt_target/common/target_normalization_context.py => executor/attack/component/prepended_history_send_context.py} (54%) create mode 100644 pyrit/prompt_target/common/target_history.py create mode 100644 pyrit/prompt_target/common/target_send_context.py create mode 100644 tests/unit/executor/attack/component/test_prepended_history_send_context.py create mode 100644 tests/unit/prompt_target/target/test_target_history.py delete mode 100644 tests/unit/prompt_target/target/test_target_normalization_context.py diff --git a/doc/code/framework.md b/doc/code/framework.md index 3a47fdba69..cab56e3dc5 100644 --- a/doc/code/framework.md +++ b/doc/code/framework.md @@ -233,6 +233,7 @@ If you are contributing to PyRIT, that work will most likely land in one of the - This is often an LLM, but it doesn't have to be. For Cross-Domain Prompt Injection Attacks, the target might be a storage account that a later target has a reference to. Message and conversation should be generic enough to handle this extra data. - Target capabilities should be used to see if a target is compatible with the capabilities that the other components want to use. - Targets should use message_normalizer along with TargetConfiguration to transform `Messages` into formats that target supports. +- A target may observe an internal, caller-owned send context at the provider-invocation boundary, but the caller owns any seed-history identity, replay, or branching state. - 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), 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`). @@ -318,6 +319,7 @@ The below talks about responsibilities of most modules in the PyRIT library - **`prompt_normalizer`** applies converters, persists requests and responses, and dispatches prompts to a `PromptTarget`. Targets do not persist messages. - **`message_normalizer`** reshapes conversations into target-compatible payloads. It owns target-facing representation, not attack policy or conversation state. +- Prepended-history identity and delivery state belong to the attack execution context. Targets and normalizers consume only the narrow send-time view needed for provider adaptation. - **Does not own**: the conversation of record. Memory is canonical; a normalized payload is an ephemeral target-facing view that is never written back. See [message normalizers](./targets/11_message_normalizer) for capability behavior, processing order, and prepended-history lifecycle details. diff --git a/doc/code/targets/0_prompt_targets.md b/doc/code/targets/0_prompt_targets.md index 98cd8c0a9c..963ee8fca2 100644 --- a/doc/code/targets/0_prompt_targets.md +++ b/doc/code/targets/0_prompt_targets.md @@ -18,20 +18,21 @@ async def send_prompt_async( self, *, message: Message, - target_normalization_context: TargetNormalizationContext | None = None, + send_context: TargetSendContext | None = None, ) -> list[Message]: ``` A `Message` object contains the current request and the identifiers needed to load its conversation history. This is discussed in more depth [here](../memory/3_memory_data_types.md). -`target_normalization_context` is an internal, per-conversation handoff used by attacks that prepend -structured history to a target without editable history. Ordinary callers do not construct it. -Before the first provider send, `PromptTarget` loads memory history, applies the context's one-shot -normalizers, and then runs the target's ordinary capability-normalization pipeline. Request -converters have already run by this point, so role-specific converter choices remain intact even -when the target must receive one flattened request. The context is ephemeral and does not replace -the structured messages stored in memory. +`send_context` is an internal protocol that lets caller-owned execution state select persisted history +and observe the provider-attempt boundary. Attacks with prepended history own the concrete +`PrependedHistorySendContext`; targets do not construct it, clone it, or decide whether its seed should +be replayed. Before a provider send, `PromptTarget` loads memory history, asks the protocol for the +caller-approved target view, and then runs the target's capability-normalization pipeline. Request +converters have already run by this point, so role-specific converter choices remain intact even when +the target must receive one flattened request. The context is ephemeral and does not replace the +structured messages stored in memory. `send_prompt_async` is the final public orchestration method. Custom target subclasses implement `_send_prompt_to_target_async(*, normalized_conversation: list[Message]) -> list[Message]` instead of diff --git a/doc/code/targets/11_message_normalizer.ipynb b/doc/code/targets/11_message_normalizer.ipynb index 200c842d61..43cb028149 100644 --- a/doc/code/targets/11_message_normalizer.ipynb +++ b/doc/code/targets/11_message_normalizer.ipynb @@ -40,8 +40,10 @@ "4. Applies the target's remaining capability normalizers.\n", "5. Serializes the normalized view and invokes the provider.\n", "\n", - "`TargetNormalizationContext` records the persisted prepended-message boundary. Stateful targets\n", - "consume it after the first provider attempt; stateless targets reuse it for each current request.\n", + "The attack-owned `PrependedHistorySendContext` records the persisted prepended-message boundary.\n", + "Stateful targets consume it after the first provider attempt; stateless targets reuse it for each\n", + "current request. Targets interact with that state only through an internal `TargetSendContext`\n", + "protocol at the send boundary.\n", "\n", "## Base Classes\n", "\n", @@ -71,10 +73,10 @@ "| Multi-turn without editable history | A context-scoped `HistorySquashNormalizer` encodes the prepended history and first live message into the initial request. Later turns use the target's own conversation state. |\n", "| Single-turn without editable history | The context-scoped `HistorySquashNormalizer` first produces one message. The target's ordinary use of the same normalizer then sees one message and does nothing. |\n", "\n", - "The first-turn distinction comes from `TargetNormalizationContext`, not from a separate normalizer\n", - "implementation. The context applies its configured `HistorySquashNormalizer` once to bootstrap\n", - "prepended history for a target that cannot accept caller-supplied prior turns. Its key use case is a\n", - "multi-turn, server-managed target without editable history.\n", + "The first-turn distinction comes from the attack-owned prepended-history send context, not from a\n", + "separate normalizer implementation. The context applies its configured `HistorySquashNormalizer`\n", + "once to bootstrap prepended history for a target that cannot accept caller-supplied prior turns.\n", + "Its key use case is a multi-turn, server-managed target without editable history.\n", "\n", "The target's ordinary capability pipeline independently uses `HistorySquashNormalizer` for a target\n", "without multi-turn support. Both scopes preserve original-versus-converted text views and keep\n", diff --git a/doc/code/targets/11_message_normalizer.py b/doc/code/targets/11_message_normalizer.py index aebe69d3a1..e6fbf13177 100644 --- a/doc/code/targets/11_message_normalizer.py +++ b/doc/code/targets/11_message_normalizer.py @@ -44,8 +44,10 @@ # 4. Applies the target's remaining capability normalizers. # 5. Serializes the normalized view and invokes the provider. # -# `TargetNormalizationContext` records the persisted prepended-message boundary. Stateful targets -# consume it after the first provider attempt; stateless targets reuse it for each current request. +# The attack-owned `PrependedHistorySendContext` records the persisted prepended-message boundary. +# Stateful targets consume it after the first provider attempt; stateless targets reuse it for each +# current request. Targets interact with that state only through an internal `TargetSendContext` +# protocol at the send boundary. # # ## Base Classes # @@ -70,10 +72,10 @@ # | Multi-turn without editable history | A context-scoped `HistorySquashNormalizer` encodes the prepended history and first live message into the initial request. Later turns use the target's own conversation state. | # | Single-turn without editable history | The context-scoped `HistorySquashNormalizer` first produces one message. The target's ordinary use of the same normalizer then sees one message and does nothing. | # -# The first-turn distinction comes from `TargetNormalizationContext`, not from a separate normalizer -# implementation. The context applies its configured `HistorySquashNormalizer` once to bootstrap -# prepended history for a target that cannot accept caller-supplied prior turns. Its key use case is a -# multi-turn, server-managed target without editable history. +# The first-turn distinction comes from the attack-owned prepended-history send context, not from a +# separate normalizer implementation. The context applies its configured `HistorySquashNormalizer` +# once to bootstrap prepended history for a target that cannot accept caller-supplied prior turns. +# Its key use case is a multi-turn, server-managed target without editable history. # # The target's ordinary capability pipeline independently uses `HistorySquashNormalizer` for a target # without multi-turn support. Both scopes preserve original-versus-converted text views and keep diff --git a/pyrit/executor/attack/component/conversation_manager.py b/pyrit/executor/attack/component/conversation_manager.py index bed543bd4f..731dc9cfc0 100644 --- a/pyrit/executor/attack/component/conversation_manager.py +++ b/pyrit/executor/attack/component/conversation_manager.py @@ -12,6 +12,9 @@ from pyrit.executor.attack.component.prepended_conversation_config import ( PrependedConversationConfig, ) +from pyrit.executor.attack.component.prepended_history_send_context import ( + PrependedHistorySendContext, +) from pyrit.memory import CentralMemory from pyrit.message_normalizer import ConversationContextNormalizer from pyrit.models import ( @@ -24,10 +27,7 @@ ) from pyrit.prompt_normalizer.prompt_normalizer import PromptNormalizer from pyrit.prompt_target import CapabilityName, PromptTarget -from pyrit.prompt_target.common.target_normalization_context import ( - TargetNormalizationContext, - filter_non_replayable_messages, -) +from pyrit.prompt_target.common.target_history import filter_non_replayable_messages if TYPE_CHECKING: from collections.abc import Sequence @@ -445,26 +445,26 @@ def get_persistable_prepended_messages( return filter_non_replayable_messages(messages=persistable_messages) @staticmethod - def create_target_normalization_context( + def create_prepended_history_send_context( *, target: PromptTarget, conversation_id: str, prepended_messages: list[Message], - ) -> TargetNormalizationContext | None: + ) -> PrependedHistorySendContext | None: """ Build persisted-prefix state for a target without editable history. Returns: - TargetNormalizationContext | None: Per-send state, or ``None`` for + PrependedHistorySendContext | None: Per-send state, or ``None`` for editable-history targets or empty prepended history. """ if not prepended_messages or target.configuration.includes(capability=CapabilityName.EDITABLE_HISTORY): return None - return TargetNormalizationContext( + return PrependedHistorySendContext( conversation_id=conversation_id, - history_message_ids=tuple(message.get_piece().id for message in prepended_messages), - replay_history_each_send=not target.configuration.includes(capability=CapabilityName.MULTI_TURN), + seed_message_ids=tuple(message.get_piece().id for message in prepended_messages), + replay_seed_each_send=not target.configuration.includes(capability=CapabilityName.MULTI_TURN), ) async def _process_prepended_conversation_async( @@ -519,7 +519,7 @@ async def _process_prepended_conversation_async( target=target, ) persisted_messages = self.get_conversation(conversation_id) - context.target_normalization_context = self.create_target_normalization_context( + context.prepended_history_send_context = self.create_prepended_history_send_context( target=target, conversation_id=conversation_id, prepended_messages=persisted_messages, diff --git a/pyrit/executor/attack/component/prepended_conversation_config.py b/pyrit/executor/attack/component/prepended_conversation_config.py index 948ac4ebe7..d277921f13 100644 --- a/pyrit/executor/attack/component/prepended_conversation_config.py +++ b/pyrit/executor/attack/component/prepended_conversation_config.py @@ -15,9 +15,11 @@ from pyrit.prompt_target.common.target_capabilities import CapabilityName if TYPE_CHECKING: + from pyrit.executor.attack.component.prepended_history_send_context import ( + PrependedHistorySendContext, + ) from pyrit.models import ChatMessageRole, Message from pyrit.prompt_target.common.prompt_target import PromptTarget - from pyrit.prompt_target.common.target_normalization_context import TargetNormalizationContext @dataclass @@ -61,14 +63,14 @@ def get_normalizer_overrides( self, *, target: PromptTarget, - target_normalization_context: TargetNormalizationContext | None, + prepended_history_send_context: PrependedHistorySendContext | None, ) -> dict[CapabilityName, MessageListNormalizer[Message]]: """ Build per-send target normalizer overrides for prepended history. Args: target: Target that receives the live request. - target_normalization_context: Explicit persisted history boundary for + prepended_history_send_context: Explicit persisted seed boundary for this attack execution. Returns: @@ -76,14 +78,14 @@ def get_normalizer_overrides( """ if ( target.configuration.includes(capability=CapabilityName.EDITABLE_HISTORY) - or target_normalization_context is None - or not target_normalization_context.should_include_history + or prepended_history_send_context is None + or not prepended_history_send_context.should_include_seed ): return {} return { CapabilityName.EDITABLE_HISTORY: HistorySquashNormalizer( - expected_history_message_count=target_normalization_context.history_message_count, + expected_history_message_count=prepended_history_send_context.seed_message_count, message_normalizer=self.get_message_normalizer(), ) } diff --git a/pyrit/prompt_target/common/target_normalization_context.py b/pyrit/executor/attack/component/prepended_history_send_context.py similarity index 54% rename from pyrit/prompt_target/common/target_normalization_context.py rename to pyrit/executor/attack/component/prepended_history_send_context.py index bc50138632..e4b55350ef 100644 --- a/pyrit/prompt_target/common/target_normalization_context.py +++ b/pyrit/executor/attack/component/prepended_history_send_context.py @@ -14,95 +14,53 @@ from pyrit.models import Message -def filter_non_replayable_messages(*, messages: list[Message]) -> list[Message]: - """ - Remove failed request/error-response pairs from target-facing history. - - ``processing`` and ``unknown`` responses represent failed exchanges, not - provider-authored conversation turns. ``blocked`` and ``empty`` responses - are retained because they are real provider round trips. - - Args: - messages: Persisted messages that may contain failed exchanges. - - Returns: - Messages that are safe to include in a later target-facing payload. - """ - non_replayable_errors = {"processing", "unknown"} - excluded_indexes: set[int] = set() - - for index, error_response in enumerate(messages): - if not any(piece.response_error in non_replayable_errors for piece in error_response.message_pieces): - continue - - excluded_indexes.add(index) - if index == 0: - continue - - request = messages[index - 1] - request_piece = request.get_piece() - error_piece = error_response.get_piece() - is_adjacent_request = ( - request.api_role == "user" - and error_response.api_role == "assistant" - and bool(request_piece.conversation_id) - and request_piece.conversation_id == error_piece.conversation_id - and request_piece.sequence >= 0 - and error_piece.sequence == request_piece.sequence + 1 - ) - if is_adjacent_request: - excluded_indexes.add(index - 1) - - return [message for index, message in enumerate(messages) if index not in excluded_indexes] - - @dataclass(slots=True) -class TargetNormalizationContext: +class PrependedHistorySendContext: """ - Per-execution state for adapting an explicit persisted history prefix. + Per-execution state for delivering an explicit persisted seed prefix. - The prefix is identified by persisted message-piece IDs instead of inferred - from response roles or error codes. A context may be used by only one send at - a time. Stateful targets consume the prefix when provider invocation begins; + The seed is identified by persisted message-piece IDs instead of inferred + from response roles or error codes. A context may be used by only one send + at a time. Stateful targets consume the seed when provider invocation begins; stateless targets replay it for every send. """ conversation_id: str - history_message_ids: tuple[uuid.UUID, ...] - replay_history_each_send: bool - _history_consumed: bool = field(default=False, init=False, repr=False) + seed_message_ids: tuple[uuid.UUID, ...] + replay_seed_each_send: bool + _seed_consumed: bool = field(default=False, init=False, repr=False) _send_in_progress: bool = field(default=False, init=False, repr=False) _provider_attempted_task: asyncio.Task[Any] | None = field(default=None, init=False, repr=False) _provider_attempt_count: int = field(default=0, init=False, repr=False) def __post_init__(self) -> None: """ - Validate the persisted prefix identity. + Validate the persisted seed identity. Raises: - ValueError: If the conversation ID or history IDs are invalid. + ValueError: If the conversation ID or seed IDs are invalid. """ if not self.conversation_id: raise ValueError("conversation_id must not be empty") - if not self.history_message_ids: - raise ValueError("history_message_ids must not be empty") - if len(set(self.history_message_ids)) != len(self.history_message_ids): - raise ValueError("history_message_ids must be unique") + if not self.seed_message_ids: + raise ValueError("seed_message_ids must not be empty") + if len(set(self.seed_message_ids)) != len(self.seed_message_ids): + raise ValueError("seed_message_ids must be unique") @property - def history_message_count(self) -> int: - """Number of messages in the explicit history prefix.""" - return len(self.history_message_ids) + def seed_message_count(self) -> int: + """Number of messages in the explicit seed prefix.""" + return len(self.seed_message_ids) @property - def should_include_history(self) -> bool: - """Whether the prefix should be included in the next send.""" - return self.replay_history_each_send or not self._history_consumed + def should_include_seed(self) -> bool: + """Whether the seed prefix should be included in the next send.""" + return self.replay_seed_each_send or not self._seed_consumed @property - def is_consumed(self) -> bool: - """Whether a stateful target has consumed the prefix.""" - return self._history_consumed + def is_seed_consumed(self) -> bool: + """Whether a stateful target has consumed the seed prefix.""" + return self._seed_consumed @property def provider_attempt_count(self) -> int: @@ -127,7 +85,7 @@ def begin_send(self) -> None: """ if self._send_in_progress: raise RuntimeError( - "Concurrent sends for the same target normalization context are not supported. " + "Concurrent sends for the same prepended history send context are not supported. " "Wait for the active send to finish before sending another request." ) self._send_in_progress = True @@ -137,8 +95,8 @@ def mark_provider_attempted(self) -> None: """ Record that provider invocation has begun for the active send. - Stateful targets consume their bootstrap history at this boundary even - if the provider later fails or the task is cancelled. + Stateful targets consume their seed history at this boundary even if the + provider later fails or the task is cancelled. Raises: RuntimeError: If no send currently owns the context. @@ -150,8 +108,8 @@ def mark_provider_attempted(self) -> None: except RuntimeError: self._provider_attempted_task = None self._provider_attempt_count += 1 - if not self.replay_history_each_send: - self._history_consumed = True + if not self.replay_seed_each_send: + self._seed_consumed = True def finish_send(self) -> None: """Release this context after the active send completes or is cancelled.""" @@ -161,30 +119,30 @@ def finish_send(self) -> None: def select_history(self, *, messages: list[Message]) -> list[Message]: """ - Select the explicit persisted prefix from conversation history. + Select the explicit persisted seed prefix from conversation history. Args: messages: Persisted conversation messages after failed exchanges have been removed. Returns: - The prefix messages in their original persisted order, or an empty - list after a stateful target has consumed the prefix. + The seed messages in their original persisted order, or an empty + list after a stateful target has consumed the seed. Raises: - ValueError: If an expected persisted prefix message is missing. + ValueError: If an expected persisted seed message is missing. """ - if not self.should_include_history: + if not self.should_include_seed: return [] messages_by_id = {message.get_piece().id: message for message in messages} - missing_ids = [message_id for message_id in self.history_message_ids if message_id not in messages_by_id] + missing_ids = [message_id for message_id in self.seed_message_ids if message_id not in messages_by_id] if missing_ids: raise ValueError( - "The persisted prepended history no longer matches the target normalization context. " + "The persisted prepended history no longer matches the prepended history send context. " f"Missing {len(missing_ids)} message(s)." ) - return [messages_by_id[message_id] for message_id in self.history_message_ids] + return [messages_by_id[message_id] for message_id in self.seed_message_ids] def remap_for_duplicate_conversation( self, @@ -192,12 +150,12 @@ def remap_for_duplicate_conversation( conversation_id: str, source_messages: list[Message], duplicated_messages: list[Message], - ) -> TargetNormalizationContext: + ) -> PrependedHistorySendContext: """ - Remap this explicit boundary to a duplicated conversation. + Remap this explicit seed boundary to a duplicated conversation. - Only message pieces already identified as history are remapped. Live turns - copied into the new memory conversation never become part of the boundary. + Only message pieces already identified as seed history are remapped. Live + turns copied into the new memory conversation never become part of the boundary. Args: conversation_id: Conversation ID assigned to the duplicated messages. @@ -205,12 +163,12 @@ def remap_for_duplicate_conversation( duplicated_messages: Their duplicates in the same persisted order. Returns: - A new context with boundary IDs from the duplicated conversation. - A new logical conversation starts with an unconsumed boundary. + A new context with seed IDs from the duplicated conversation. + A new logical conversation starts with an unconsumed seed boundary. Raises: ValueError: If the duplicated messages do not match the source structure - or an explicit history piece cannot be remapped. + or an explicit seed piece cannot be remapped. """ if len(source_messages) != len(duplicated_messages): raise ValueError("Duplicated conversation does not match the source message count.") @@ -238,19 +196,17 @@ def remap_for_duplicate_conversation( duplicated_ids_by_source_id[source_piece.id] = duplicated_piece.id missing_ids = [ - message_id for message_id in self.history_message_ids if message_id not in duplicated_ids_by_source_id + message_id for message_id in self.seed_message_ids if message_id not in duplicated_ids_by_source_id ] if missing_ids: - raise ValueError(f"Could not remap {len(missing_ids)} explicit history message(s).") + raise ValueError(f"Could not remap {len(missing_ids)} explicit seed message(s).") - duplicated_context = TargetNormalizationContext( + duplicated_context = PrependedHistorySendContext( conversation_id=conversation_id, - history_message_ids=tuple( - duplicated_ids_by_source_id[message_id] for message_id in self.history_message_ids - ), - replay_history_each_send=self.replay_history_each_send, + seed_message_ids=tuple(duplicated_ids_by_source_id[message_id] for message_id in self.seed_message_ids), + replay_seed_each_send=self.replay_seed_each_send, ) # Provider bootstrap consumption belongs to the logical conversation, not copied memory. if conversation_id == self.conversation_id: - duplicated_context._history_consumed = self._history_consumed + duplicated_context._seed_consumed = self._seed_consumed return duplicated_context diff --git a/pyrit/executor/attack/core/attack_strategy.py b/pyrit/executor/attack/core/attack_strategy.py index 9c390d6f4c..a6fee0f286 100644 --- a/pyrit/executor/attack/core/attack_strategy.py +++ b/pyrit/executor/attack/core/attack_strategy.py @@ -45,6 +45,9 @@ from pyrit.executor.attack.component.prepended_conversation_config import ( PrependedConversationConfig, ) + from pyrit.executor.attack.component.prepended_history_send_context import ( + PrependedHistorySendContext, + ) from pyrit.executor.attack.core.attack_config import ( AttackAdversarialConfig, AttackScoringConfig, @@ -53,7 +56,6 @@ from pyrit.message_normalizer import MessageListNormalizer from pyrit.prompt_target import PromptTarget from pyrit.prompt_target.common.target_capabilities import CapabilityName - from pyrit.prompt_target.common.target_normalization_context import TargetNormalizationContext AttackStrategyContextT = TypeVar("AttackStrategyContextT", bound="AttackContext[Any]") AttackStrategyResultT = TypeVar("AttackStrategyResultT", bound="AttackResult") @@ -94,8 +96,8 @@ class AttackContext(StrategyContext, ABC, Generic[AttackParamsT]): _prepended_conversation_override: list[Message] | None = None _memory_labels_override: dict[str, str] | None = None - # Per-execution target-facing boundary and send lifecycle. Never persisted. - target_normalization_context: TargetNormalizationContext | None = field( + # Per-execution prepended-history boundary and send lifecycle. Never persisted. + prepended_history_send_context: PrependedHistorySendContext | None = field( default=None, repr=False, compare=False, @@ -477,20 +479,20 @@ def __init__( def _get_prepended_normalizer_overrides( self, *, - target_normalization_context: TargetNormalizationContext | None, + prepended_history_send_context: PrependedHistorySendContext | None, ) -> dict[CapabilityName, MessageListNormalizer[Message]]: """ Resolve prepended-history overrides for one target send. Args: - target_normalization_context: Persisted seed boundary for this execution. + prepended_history_send_context: Persisted seed boundary for this execution. Returns: Overrides keyed by the capability they adapt. """ return self._prepended_conversation_config.get_normalizer_overrides( target=self._objective_target, - target_normalization_context=target_normalization_context, + prepended_history_send_context=prepended_history_send_context, ) def _create_identifier( diff --git a/pyrit/executor/attack/multi_turn/chunked_request.py b/pyrit/executor/attack/multi_turn/chunked_request.py index 252665c846..b90eca4d09 100644 --- a/pyrit/executor/attack/multi_turn/chunked_request.py +++ b/pyrit/executor/attack/multi_turn/chunked_request.py @@ -294,9 +294,9 @@ async def _perform_async(self, *, context: ChunkedRequestAttackContext) -> Attac request_converter_configurations=self._request_converters, response_converter_configurations=self._response_converters, normalizer_overrides=self._get_prepended_normalizer_overrides( - target_normalization_context=context.target_normalization_context, + prepended_history_send_context=context.prepended_history_send_context, ), - target_normalization_context=context.target_normalization_context, + send_context=context.prepended_history_send_context, ) # Store the response diff --git a/pyrit/executor/attack/multi_turn/crescendo.py b/pyrit/executor/attack/multi_turn/crescendo.py index dc96a5eefd..b9fb98a677 100644 --- a/pyrit/executor/attack/multi_turn/crescendo.py +++ b/pyrit/executor/attack/multi_turn/crescendo.py @@ -642,9 +642,9 @@ async def _send_prompt_to_objective_target_async( request_converter_configurations=self._request_converters, response_converter_configurations=self._response_converters, normalizer_overrides=self._get_prepended_normalizer_overrides( - target_normalization_context=context.target_normalization_context, + prepended_history_send_context=context.prepended_history_send_context, ), - target_normalization_context=context.target_normalization_context, + send_context=context.prepended_history_send_context, ) if not response: diff --git a/pyrit/executor/attack/multi_turn/multi_prompt_sending.py b/pyrit/executor/attack/multi_turn/multi_prompt_sending.py index 8613a4fa4e..eb065fb6d0 100644 --- a/pyrit/executor/attack/multi_turn/multi_prompt_sending.py +++ b/pyrit/executor/attack/multi_turn/multi_prompt_sending.py @@ -372,9 +372,9 @@ async def _send_prompt_to_objective_target_async( request_converter_configurations=self._request_converters, response_converter_configurations=self._response_converters, normalizer_overrides=self._get_prepended_normalizer_overrides( - target_normalization_context=context.target_normalization_context, + prepended_history_send_context=context.prepended_history_send_context, ), - target_normalization_context=context.target_normalization_context, + send_context=context.prepended_history_send_context, ) async def _evaluate_response_async(self, *, response: Message, objective: str) -> Score | None: diff --git a/pyrit/executor/attack/multi_turn/multi_turn_attack_strategy.py b/pyrit/executor/attack/multi_turn/multi_turn_attack_strategy.py index c035ec971c..0aa621b7f5 100644 --- a/pyrit/executor/attack/multi_turn/multi_turn_attack_strategy.py +++ b/pyrit/executor/attack/multi_turn/multi_turn_attack_strategy.py @@ -117,7 +117,7 @@ def _rotate_conversation_for_single_turn_target( """ if self._objective_target.configuration.includes(capability=CapabilityName.MULTI_TURN): return - if context.target_normalization_context: + if context.prepended_history_send_context: return if context.executed_turns == 0: return @@ -146,7 +146,7 @@ def _rotate_conversation_for_single_turn_target( memory.add_message_pieces_to_memory(message_pieces=pieces) context.session.conversation_id = new_conversation_id persisted_messages = list(memory.get_conversation_messages(conversation_id=new_conversation_id)) - context.target_normalization_context = ConversationManager.create_target_normalization_context( + context.prepended_history_send_context = ConversationManager.create_prepended_history_send_context( target=self._objective_target, conversation_id=new_conversation_id, prepended_messages=persisted_messages, diff --git a/pyrit/executor/attack/multi_turn/red_teaming.py b/pyrit/executor/attack/multi_turn/red_teaming.py index a66e45e157..44893bb755 100644 --- a/pyrit/executor/attack/multi_turn/red_teaming.py +++ b/pyrit/executor/attack/multi_turn/red_teaming.py @@ -495,9 +495,9 @@ async def _send_prompt_to_objective_target_async( response_converter_configurations=self._response_converters, target=self._objective_target, normalizer_overrides=self._get_prepended_normalizer_overrides( - target_normalization_context=context.target_normalization_context, + prepended_history_send_context=context.prepended_history_send_context, ), - target_normalization_context=context.target_normalization_context, + send_context=context.prepended_history_send_context, ) if response is None: diff --git a/pyrit/executor/attack/multi_turn/tree_of_attacks.py b/pyrit/executor/attack/multi_turn/tree_of_attacks.py index e68c7b4e31..ea833770a2 100644 --- a/pyrit/executor/attack/multi_turn/tree_of_attacks.py +++ b/pyrit/executor/attack/multi_turn/tree_of_attacks.py @@ -72,8 +72,10 @@ from collections.abc import AsyncIterator from pathlib import Path + from pyrit.executor.attack.component.prepended_history_send_context import ( + PrependedHistorySendContext, + ) from pyrit.models.literals import PromptDataType - from pyrit.prompt_target.common.target_normalization_context import TargetNormalizationContext logger = logging.getLogger(__name__) @@ -427,7 +429,7 @@ def __init__( self.last_prompt_sent: str | None = None self.last_response: Message | None = None self.error_message: str | None = None - self._target_normalization_context: TargetNormalizationContext | None = None + self._prepended_history_send_context: PrependedHistorySendContext | None = None # Context from prepended conversation (for adversarial chat system prompt) self._conversation_context: str | None = None @@ -489,7 +491,7 @@ async def initialize_with_prepended_conversation_async( persisted_messages = list( self._memory.get_conversation_messages(conversation_id=self.objective_target_conversation_id) ) - self._target_normalization_context = conversation_manager.create_target_normalization_context( + self._prepended_history_send_context = conversation_manager.create_prepended_history_send_context( target=self._objective_target, conversation_id=self.objective_target_conversation_id, prepended_messages=persisted_messages, @@ -651,9 +653,9 @@ async def _send_prompt_to_target_async(self, prompt: str) -> Message: target=self._objective_target, normalizer_overrides=self._prepended_conversation_config.get_normalizer_overrides( target=self._objective_target, - target_normalization_context=self._target_normalization_context, + prepended_history_send_context=self._prepended_history_send_context, ), - target_normalization_context=self._target_normalization_context, + send_context=self._prepended_history_send_context, ) # Store the full response so subsequent turns can forward media when supported. @@ -730,9 +732,9 @@ async def _send_initial_prompt_to_target_async(self) -> Message: target=self._objective_target, normalizer_overrides=self._prepended_conversation_config.get_normalizer_overrides( target=self._objective_target, - target_normalization_context=self._target_normalization_context, + prepended_history_send_context=self._prepended_history_send_context, ), - target_normalization_context=self._target_normalization_context, + send_context=self._prepended_history_send_context, ) # Store the full response so subsequent turns can forward media when supported. @@ -745,7 +747,7 @@ def _rotate_unseeded_single_turn_conversation(self) -> None: """Isolate unseeded single-turn sends without discarding an explicit branch boundary.""" if ( not self._objective_target.configuration.includes(capability=CapabilityName.MULTI_TURN) - and self._target_normalization_context is None + and self._prepended_history_send_context is None ): self.objective_target_conversation_id = str(uuid.uuid4()) @@ -933,13 +935,13 @@ def duplicate(self) -> _TreeOfAttacksNode: duplicated_messages = list( self._memory.get_conversation_messages(conversation_id=duplicate_node.objective_target_conversation_id) ) - duplicate_node._target_normalization_context = ( - self._target_normalization_context.remap_for_duplicate_conversation( + duplicate_node._prepended_history_send_context = ( + self._prepended_history_send_context.remap_for_duplicate_conversation( conversation_id=duplicate_node.objective_target_conversation_id, source_messages=source_messages, duplicated_messages=duplicated_messages, ) - if self._target_normalization_context + if self._prepended_history_send_context else None ) diff --git a/pyrit/executor/attack/single_turn/prompt_sending.py b/pyrit/executor/attack/single_turn/prompt_sending.py index de94336eb4..c37889d7e9 100644 --- a/pyrit/executor/attack/single_turn/prompt_sending.py +++ b/pyrit/executor/attack/single_turn/prompt_sending.py @@ -324,9 +324,9 @@ async def _send_prompt_to_objective_target_async( request_converter_configurations=self._request_converters, response_converter_configurations=self._response_converters, normalizer_overrides=self._get_prepended_normalizer_overrides( - target_normalization_context=context.target_normalization_context, + prepended_history_send_context=context.prepended_history_send_context, ), - target_normalization_context=context.target_normalization_context, + send_context=context.prepended_history_send_context, ) async def _evaluate_response_async( diff --git a/pyrit/executor/attack/streaming/barge_in.py b/pyrit/executor/attack/streaming/barge_in.py index 894f6e48d8..f28be0a211 100644 --- a/pyrit/executor/attack/streaming/barge_in.py +++ b/pyrit/executor/attack/streaming/barge_in.py @@ -145,7 +145,7 @@ async def _setup_async(self, *, context: BargeInAttackContext[Any]) -> None: request_converters=self._request_converters, prepended_conversation_config=self._prepended_conversation_config, ) - context.target_normalization_context = None + context.prepended_history_send_context = None async def _teardown_async(self, *, context: BargeInAttackContext[Any]) -> None: """No-op teardown — connection / dispatcher are closed inside the session's ``run_async``.""" diff --git a/pyrit/message_normalizer/history_squash_normalizer.py b/pyrit/message_normalizer/history_squash_normalizer.py index b437c02f0c..8fd428c2f8 100644 --- a/pyrit/message_normalizer/history_squash_normalizer.py +++ b/pyrit/message_normalizer/history_squash_normalizer.py @@ -20,7 +20,7 @@ class HistorySquashNormalizer(MessageListNormalizer[Message]): The same implementation serves two normalization scopes. Prepended-conversation flows create per-send overrides with a string formatter and the explicit history - count from ``TargetNormalizationContext``. The ordinary target capability + count from the caller-owned prepended-history send context. The ordinary target capability pipeline uses the default formatter whenever a target does not support multiple turns. diff --git a/pyrit/prompt_normalizer/normalizer_request.py b/pyrit/prompt_normalizer/normalizer_request.py index fdde9ceddb..ccd2bf22d3 100644 --- a/pyrit/prompt_normalizer/normalizer_request.py +++ b/pyrit/prompt_normalizer/normalizer_request.py @@ -10,7 +10,7 @@ ConverterConfiguration, ) from pyrit.prompt_target.common.target_capabilities import CapabilityName -from pyrit.prompt_target.common.target_normalization_context import TargetNormalizationContext +from pyrit.prompt_target.common.target_send_context import TargetSendContext @dataclass @@ -24,7 +24,7 @@ class NormalizerRequest: response_converter_configurations: list[ConverterConfiguration] conversation_id: str | None normalizer_overrides: dict[CapabilityName, MessageListNormalizer[Message]] - target_normalization_context: TargetNormalizationContext | None + send_context: TargetSendContext | None def __init__( self, @@ -34,7 +34,7 @@ def __init__( response_converter_configurations: list[ConverterConfiguration] | None = None, conversation_id: str | None = None, normalizer_overrides: Mapping[CapabilityName, MessageListNormalizer[Message]] | None = None, - target_normalization_context: TargetNormalizationContext | None = None, + send_context: TargetSendContext | None = None, ) -> None: """ Initialize a normalizer request. @@ -47,7 +47,7 @@ def __init__( the response. Defaults to an empty list. conversation_id (str | None): The ID of the conversation. Defaults to None. normalizer_overrides: Optional per-send target normalizer overrides. - target_normalization_context: Optional explicit persisted-history boundary. + send_context: Optional internal target-send coordination contract. """ if response_converter_configurations is None: response_converter_configurations = [] @@ -58,4 +58,4 @@ def __init__( self.response_converter_configurations = response_converter_configurations self.conversation_id = conversation_id self.normalizer_overrides = dict(normalizer_overrides or {}) - self.target_normalization_context = target_normalization_context + self.send_context = send_context diff --git a/pyrit/prompt_normalizer/prompt_normalizer.py b/pyrit/prompt_normalizer/prompt_normalizer.py index 17d3567210..8c9e6cad54 100644 --- a/pyrit/prompt_normalizer/prompt_normalizer.py +++ b/pyrit/prompt_normalizer/prompt_normalizer.py @@ -31,7 +31,7 @@ from pyrit.prompt_normalizer import ConverterConfiguration, NormalizerRequest from pyrit.prompt_target import CapabilityName, PromptTarget from pyrit.prompt_target.batch_helper import batch_task_async -from pyrit.prompt_target.common.target_normalization_context import TargetNormalizationContext +from pyrit.prompt_target.common.target_send_context import TargetSendContext logger = logging.getLogger(__name__) @@ -75,7 +75,7 @@ async def send_prompt_async( request_converter_configurations: list[ConverterConfiguration] | None = None, response_converter_configurations: list[ConverterConfiguration] | None = None, normalizer_overrides: Mapping[CapabilityName, MessageListNormalizer[Message]] | None = None, - target_normalization_context: TargetNormalizationContext | None = None, + send_context: TargetSendContext | None = None, ) -> Message: """ Send a single request to a target. @@ -89,8 +89,8 @@ async def send_prompt_async( response_converter_configurations (list[ConverterConfiguration], optional): Configurations for converting the response. Defaults to an empty list. normalizer_overrides: Optional per-send target normalizer overrides. - target_normalization_context: Optional explicit persisted-history - boundary and send lifecycle state. + send_context: Optional internal coordination contract for caller-owned + history selection and send lifecycle state. Returns: Message: The response received from the target. @@ -128,11 +128,11 @@ async def send_prompt_async( responses = await target.send_prompt_async( message=request, normalizer_overrides=normalizer_overrides, - target_normalization_context=target_normalization_context, + send_context=send_context, ) self.memory.add_message_to_memory(request=request) except EmptyResponseException as ex: - if target_normalization_context and not target_normalization_context.provider_attempted_by_current_task: + if send_context and not send_context.provider_attempted_by_current_task: cid = request.message_pieces[0].conversation_id if request.message_pieces else None raise Exception(f"Error normalizing prompt with conversation ID: {cid}") from ex @@ -149,7 +149,7 @@ async def send_prompt_async( ] except Exception as ex: - if target_normalization_context and not target_normalization_context.provider_attempted_by_current_task: + if send_context and not send_context.provider_attempted_by_current_task: cid = request.message_pieces[0].conversation_id if request.message_pieces else None raise Exception(f"Error normalizing prompt with conversation ID: {cid}") from ex @@ -229,7 +229,7 @@ async def send_prompt_batch_to_target_async( [request.response_converter_configurations for request in requests], [request.conversation_id for request in requests], [request.normalizer_overrides for request in requests], - [request.target_normalization_context for request in requests], + [request.send_context for request in requests], ] batch_item_keys = [ @@ -238,7 +238,7 @@ async def send_prompt_batch_to_target_async( "response_converter_configurations", "conversation_id", "normalizer_overrides", - "target_normalization_context", + "send_context", ] responses: list[Message] = await batch_task_async( diff --git a/pyrit/prompt_target/__init__.py b/pyrit/prompt_target/__init__.py index feef27b85b..1d5608d691 100644 --- a/pyrit/prompt_target/__init__.py +++ b/pyrit/prompt_target/__init__.py @@ -27,7 +27,6 @@ get_known_capabilities, ) from pyrit.prompt_target.common.target_configuration import TargetConfiguration -from pyrit.prompt_target.common.target_normalization_context import TargetNormalizationContext from pyrit.prompt_target.common.target_requirements import CHAT_TARGET_REQUIREMENTS, TargetRequirements from pyrit.prompt_target.common.utils import limit_requests_per_minute from pyrit.prompt_target.gandalf_target import GandalfLevel, GandalfTarget @@ -108,7 +107,6 @@ def __getattr__(name: str) -> object: "RoundRobinTarget", "TargetCapabilities", "TargetConfiguration", - "TargetNormalizationContext", "TargetRequirements", "UnsupportedCapabilityBehavior", "TextTarget", diff --git a/pyrit/prompt_target/common/prompt_target.py b/pyrit/prompt_target/common/prompt_target.py index 1a7e552c96..ca97638e83 100644 --- a/pyrit/prompt_target/common/prompt_target.py +++ b/pyrit/prompt_target/common/prompt_target.py @@ -23,10 +23,8 @@ get_known_capabilities, ) from pyrit.prompt_target.common.target_configuration import TargetConfiguration -from pyrit.prompt_target.common.target_normalization_context import ( - TargetNormalizationContext, - filter_non_replayable_messages, -) +from pyrit.prompt_target.common.target_history import filter_non_replayable_messages +from pyrit.prompt_target.common.target_send_context import TargetSendContext logger = logging.getLogger(__name__) @@ -143,7 +141,7 @@ async def send_prompt_async( *, message: Message, normalizer_overrides: Mapping[CapabilityName, MessageListNormalizer[Message]] | None = None, - target_normalization_context: TargetNormalizationContext | None = None, + send_context: TargetSendContext | None = None, ) -> list[Message]: """ Validate, normalize, and send a prompt to the target. @@ -161,8 +159,8 @@ async def send_prompt_async( Args: message (Message): The message to send. normalizer_overrides: Optional per-send target normalizer overrides. - target_normalization_context: Optional explicit persisted-history - boundary and send lifecycle state. + send_context: Optional internal coordination contract for caller-owned + history selection and send lifecycle state. Returns: list[Message]: Response messages from the target. @@ -172,28 +170,26 @@ async def send_prompt_async( """ message.validate() conversation_id = message.get_piece().conversation_id or "" - if target_normalization_context and target_normalization_context.conversation_id != conversation_id: - raise ValueError( - "Target normalization context conversation_id does not match the current request conversation_id." - ) - if target_normalization_context: - target_normalization_context.begin_send() + if send_context and send_context.conversation_id != conversation_id: + raise ValueError("Target send context conversation_id does not match the current request conversation_id.") + if send_context: + send_context.begin_send() try: normalized_conversation = await self._get_normalized_conversation_async( message=message, normalizer_overrides=normalizer_overrides, - target_normalization_context=target_normalization_context, + send_context=send_context, ) if not normalized_conversation: raise ValueError("Normalization pipeline returned an empty conversation. Cannot send an empty request.") self._validate_request(normalized_conversation=normalized_conversation) - if target_normalization_context: - target_normalization_context.mark_provider_attempted() + if send_context: + send_context.mark_provider_attempted() return await self._send_prompt_to_target_async(normalized_conversation=normalized_conversation) finally: - if target_normalization_context: - target_normalization_context.finish_send() + if send_context: + send_context.finish_send() @abc.abstractmethod async def _send_prompt_to_target_async(self, *, normalized_conversation: list[Message]) -> list[Message]: @@ -257,7 +253,7 @@ async def _get_normalized_conversation_async( *, message: Message, normalizer_overrides: Mapping[CapabilityName, MessageListNormalizer[Message]] | None = None, - target_normalization_context: TargetNormalizationContext | None = None, + send_context: TargetSendContext | None = None, ) -> list[Message]: """ Build the target-facing conversation and run the normalization pipeline. @@ -275,7 +271,8 @@ async def _get_normalized_conversation_async( Args: message (Message): The current message to append. normalizer_overrides: Optional per-send target normalizer overrides. - target_normalization_context: Optional explicit persisted-history boundary. + send_context: Optional internal coordination contract for caller-approved + persisted history. Returns: list[Message]: The normalized conversation (possibly with system prompt squashed, @@ -286,11 +283,7 @@ async def _get_normalized_conversation_async( list(self._memory.get_conversation_messages(conversation_id=conversation_id)) if conversation_id else [] ) persisted_messages = filter_non_replayable_messages(messages=persisted_messages) - conversation = ( - target_normalization_context.select_history(messages=persisted_messages) - if target_normalization_context - else persisted_messages - ) + conversation = send_context.select_history(messages=persisted_messages) if send_context else persisted_messages conversation.append(message) normalized = await self.configuration.normalize_async( messages=conversation, diff --git a/pyrit/prompt_target/common/target_history.py b/pyrit/prompt_target/common/target_history.py new file mode 100644 index 0000000000..17b1faa0da --- /dev/null +++ b/pyrit/prompt_target/common/target_history.py @@ -0,0 +1,46 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT license. + +from pyrit.models import Message + + +def filter_non_replayable_messages(*, messages: list[Message]) -> list[Message]: + """ + Remove failed request/error-response pairs from target-facing history. + + ``processing`` and ``unknown`` responses represent failed exchanges, not + provider-authored conversation turns. ``blocked`` and ``empty`` responses + are retained because they are real provider round trips. + + Args: + messages: Persisted messages that may contain failed exchanges. + + Returns: + Messages that are safe to include in a later target-facing payload. + """ + non_replayable_errors = {"processing", "unknown"} + excluded_indexes: set[int] = set() + + for index, error_response in enumerate(messages): + if not any(piece.response_error in non_replayable_errors for piece in error_response.message_pieces): + continue + + excluded_indexes.add(index) + if index == 0: + continue + + request = messages[index - 1] + request_piece = request.get_piece() + error_piece = error_response.get_piece() + is_adjacent_request = ( + request.api_role == "user" + and error_response.api_role == "assistant" + and bool(request_piece.conversation_id) + and request_piece.conversation_id == error_piece.conversation_id + and request_piece.sequence >= 0 + and error_piece.sequence == request_piece.sequence + 1 + ) + if is_adjacent_request: + excluded_indexes.add(index - 1) + + return [message for index, message in enumerate(messages) if index not in excluded_indexes] diff --git a/pyrit/prompt_target/common/target_send_context.py b/pyrit/prompt_target/common/target_send_context.py new file mode 100644 index 0000000000..02f66551e1 --- /dev/null +++ b/pyrit/prompt_target/common/target_send_context.py @@ -0,0 +1,36 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT license. + +from __future__ import annotations + +from typing import TYPE_CHECKING, Protocol + +if TYPE_CHECKING: + from pyrit.models import Message + + +class TargetSendContext(Protocol): + """Internal contract coordinating one target send with caller-owned state.""" + + conversation_id: str + + @property + def provider_attempted_by_current_task(self) -> bool: + """Whether the current task reached provider invocation.""" + ... + + def begin_send(self) -> None: + """Acquire caller-owned state for one complete send.""" + ... + + def select_history(self, *, messages: list[Message]) -> list[Message]: + """Select the caller-approved persisted history for this send.""" + ... + + def mark_provider_attempted(self) -> None: + """Record that provider invocation has begun.""" + ... + + def finish_send(self) -> None: + """Release caller-owned state after the send.""" + ... diff --git a/tests/unit/executor/attack/component/test_conversation_manager.py b/tests/unit/executor/attack/component/test_conversation_manager.py index b6faf11a1c..72fe89e431 100644 --- a/tests/unit/executor/attack/component/test_conversation_manager.py +++ b/tests/unit/executor/attack/component/test_conversation_manager.py @@ -801,16 +801,16 @@ async def test_non_editable_target_persists_history_without_using_formatter( prepended_conversation_config=config, ) - assert context.target_normalization_context is not None - assert context.target_normalization_context.conversation_id == conversation_id + assert context.prepended_history_send_context is not None + assert context.prepended_history_send_context.conversation_id == conversation_id stored = manager.get_conversation(conversation_id) assert len(stored) == len(sample_conversation) - assert context.target_normalization_context.history_message_ids == tuple( + assert context.prepended_history_send_context.seed_message_ids == tuple( message.get_piece().id for message in stored ) normalizer = config.get_normalizer_overrides( target=mock_prompt_target, - target_normalization_context=context.target_normalization_context, + prepended_history_send_context=context.prepended_history_send_context, )[CapabilityName.EDITABLE_HISTORY] assert isinstance(normalizer, HistorySquashNormalizer) assert normalizer._message_normalizer is message_normalizer @@ -1455,10 +1455,10 @@ async def test_message_normalizer_default_uses_conversation_context_normalizer( conversation_id=conversation_id, ) - assert context.target_normalization_context is not None + assert context.prepended_history_send_context is not None normalizer = PrependedConversationConfig().get_normalizer_overrides( target=mock_prompt_target, - target_normalization_context=context.target_normalization_context, + prepended_history_send_context=context.prepended_history_send_context, )[CapabilityName.EDITABLE_HISTORY] assert isinstance(normalizer, HistorySquashNormalizer) assert isinstance(normalizer._message_normalizer, ConversationContextNormalizer) @@ -1471,7 +1471,7 @@ def test_message_normalizer_is_not_overridden_for_editable_target( overrides = PrependedConversationConfig().get_normalizer_overrides( target=mock_chat_target, - target_normalization_context=None, + prepended_history_send_context=None, ) assert overrides == {} diff --git a/tests/unit/executor/attack/component/test_prepended_history_send_context.py b/tests/unit/executor/attack/component/test_prepended_history_send_context.py new file mode 100644 index 0000000000..5138645828 --- /dev/null +++ b/tests/unit/executor/attack/component/test_prepended_history_send_context.py @@ -0,0 +1,177 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT license. + +import uuid + +import pytest + +from pyrit.executor.attack.component.prepended_history_send_context import ( + PrependedHistorySendContext, +) +from pyrit.models import ChatMessageRole, Message + + +def _message( + *, + role: ChatMessageRole, + value: str, + sequence: int, + conversation_id: str = "conversation", +) -> Message: + message = Message.from_prompt(prompt=value, role=role) + piece = message.get_piece() + piece.sequence = sequence + piece.conversation_id = conversation_id + return message + + +def test_context_requires_unique_persisted_seed_ids() -> None: + message_id = uuid.uuid4() + + with pytest.raises(ValueError, match="conversation_id"): + PrependedHistorySendContext( + conversation_id="", + seed_message_ids=(message_id,), + replay_seed_each_send=False, + ) + with pytest.raises(ValueError, match="seed_message_ids"): + PrependedHistorySendContext( + conversation_id="conversation", + seed_message_ids=(), + replay_seed_each_send=False, + ) + with pytest.raises(ValueError, match="unique"): + PrependedHistorySendContext( + conversation_id="conversation", + seed_message_ids=(message_id, message_id), + replay_seed_each_send=False, + ) + + +def test_stateful_context_consumes_explicit_boundary_at_provider_attempt() -> None: + first = _message(role="system", value="system", sequence=0) + second = _message(role="user", value="seed", sequence=1) + unrelated = _message(role="assistant", value="later response", sequence=2) + context = PrependedHistorySendContext( + conversation_id="conversation", + seed_message_ids=(first.get_piece().id, second.get_piece().id), + replay_seed_each_send=False, + ) + + context.begin_send() + selected = context.select_history(messages=[unrelated, second, first]) + context.mark_provider_attempted() + context.finish_send() + + assert selected == [first, second] + assert context.is_seed_consumed + assert context.provider_attempt_count == 1 + + context.begin_send() + assert context.select_history(messages=[first, second, unrelated]) == [] + context.finish_send() + + +def test_stateless_context_replays_explicit_boundary_after_provider_attempt() -> None: + seed = _message(role="user", value="seed", sequence=0) + context = PrependedHistorySendContext( + conversation_id="conversation", + seed_message_ids=(seed.get_piece().id,), + replay_seed_each_send=True, + ) + + for _ in range(2): + context.begin_send() + assert context.select_history(messages=[seed]) == [seed] + context.mark_provider_attempted() + context.finish_send() + + assert not context.is_seed_consumed + assert context.provider_attempt_count == 2 + + +def test_context_rejects_concurrent_send_until_active_send_finishes() -> None: + context = PrependedHistorySendContext( + conversation_id="conversation", + seed_message_ids=(uuid.uuid4(),), + replay_seed_each_send=False, + ) + + context.begin_send() + with pytest.raises(RuntimeError, match="Concurrent sends"): + context.begin_send() + + context.finish_send() + context.begin_send() + context.finish_send() + + +def test_context_rejects_provider_attempt_without_active_send() -> None: + context = PrependedHistorySendContext( + conversation_id="conversation", + seed_message_ids=(uuid.uuid4(),), + replay_seed_each_send=False, + ) + + with pytest.raises(RuntimeError, match="without an active send"): + context.mark_provider_attempted() + + +def test_context_rejects_missing_persisted_boundary_message() -> None: + context = PrependedHistorySendContext( + conversation_id="conversation", + seed_message_ids=(uuid.uuid4(),), + replay_seed_each_send=False, + ) + + with pytest.raises(ValueError, match="Missing 1 message"): + context.select_history(messages=[]) + + +def test_context_remaps_only_explicit_seed_for_duplicate_conversation() -> None: + seed = _message(role="system", value="seed", sequence=0) + live_request = _message(role="user", value="live", sequence=1) + source_messages = [seed, live_request] + duplicated_messages = [message.duplicate() for message in source_messages] + for message in duplicated_messages: + for piece in message.message_pieces: + piece.conversation_id = "duplicate" + + context = PrependedHistorySendContext( + conversation_id="conversation", + seed_message_ids=(seed.get_piece().id,), + replay_seed_each_send=True, + ) + + duplicate = context.remap_for_duplicate_conversation( + conversation_id="duplicate", + source_messages=source_messages, + duplicated_messages=duplicated_messages, + ) + + assert duplicate.seed_message_ids == (duplicated_messages[0].get_piece().id,) + assert duplicated_messages[1].get_piece().id not in duplicate.seed_message_ids + assert duplicate.select_history(messages=duplicated_messages) == [duplicated_messages[0]] + + +def test_context_remap_resets_consumed_state_for_new_conversation() -> None: + seed = _message(role="user", value="seed", sequence=0) + duplicated_seed = seed.duplicate() + duplicated_seed.get_piece().conversation_id = "duplicate" + context = PrependedHistorySendContext( + conversation_id="conversation", + seed_message_ids=(seed.get_piece().id,), + replay_seed_each_send=False, + ) + context.begin_send() + context.mark_provider_attempted() + context.finish_send() + + duplicate = context.remap_for_duplicate_conversation( + conversation_id="duplicate", + source_messages=[seed], + duplicated_messages=[duplicated_seed], + ) + + assert not duplicate.is_seed_consumed + assert duplicate.select_history(messages=[duplicated_seed]) == [duplicated_seed] diff --git a/tests/unit/executor/attack/core/test_attack_strategy_prepended_policy.py b/tests/unit/executor/attack/core/test_attack_strategy_prepended_policy.py index df01cfa0ae..e812d4cf25 100644 --- a/tests/unit/executor/attack/core/test_attack_strategy_prepended_policy.py +++ b/tests/unit/executor/attack/core/test_attack_strategy_prepended_policy.py @@ -10,6 +10,9 @@ import pytest from pyrit.executor.attack.component import PrependedConversationConfig +from pyrit.executor.attack.component.prepended_history_send_context import ( + PrependedHistorySendContext, +) from pyrit.executor.attack.core import AttackStrategy from pyrit.executor.attack.core.attack_config import AttackScoringConfig from pyrit.executor.attack.multi_turn.chunked_request import ChunkedRequestAttack @@ -23,7 +26,6 @@ from pyrit.message_normalizer import HistorySquashNormalizer from pyrit.models import Message, MessagePiece from pyrit.prompt_target import CapabilityName, PromptTarget, TargetCapabilities, TargetConfiguration -from pyrit.prompt_target.common.target_normalization_context import TargetNormalizationContext from pyrit.scenario.core.attack_technique_factory import AttackTechniqueFactory @@ -58,13 +60,13 @@ def test_attack_strategy_owns_prepended_policy_and_identifier() -> None: @pytest.mark.usefixtures("patch_central_database") def test_attack_strategy_resolves_per_send_history_override() -> None: attack = PromptSendingAttack(objective_target=_NonEditableHistoryTarget()) - context = TargetNormalizationContext( + context = PrependedHistorySendContext( conversation_id="conversation", - history_message_ids=(uuid.uuid4(),), - replay_history_each_send=False, + seed_message_ids=(uuid.uuid4(),), + replay_seed_each_send=False, ) - overrides = attack._get_prepended_normalizer_overrides(target_normalization_context=context) + overrides = attack._get_prepended_normalizer_overrides(prepended_history_send_context=context) assert isinstance(overrides[CapabilityName.EDITABLE_HISTORY], HistorySquashNormalizer) diff --git a/tests/unit/executor/attack/multi_turn/test_chunked_request.py b/tests/unit/executor/attack/multi_turn/test_chunked_request.py index c21514e6a3..106499dbbe 100644 --- a/tests/unit/executor/attack/multi_turn/test_chunked_request.py +++ b/tests/unit/executor/attack/multi_turn/test_chunked_request.py @@ -11,6 +11,9 @@ import pytest from pyrit.executor.attack.component import PrependedConversationConfig +from pyrit.executor.attack.component.prepended_history_send_context import ( + PrependedHistorySendContext, +) from pyrit.executor.attack.core.attack_parameters import AttackParameters from pyrit.executor.attack.multi_turn import ( ChunkedRequestAttack, @@ -24,7 +27,6 @@ PromptTarget, TargetCapabilities, TargetConfiguration, - TargetNormalizationContext, ) @@ -283,12 +285,12 @@ async def test_perform_async_forwards_prepended_formatter_override(self): total_length=100, ) context = ChunkedRequestAttackContext(params=AttackParameters(objective="Extract the secret")) - target_context = TargetNormalizationContext( + target_context = PrependedHistorySendContext( conversation_id=context.session.conversation_id, - history_message_ids=(uuid.uuid4(),), - replay_history_each_send=False, + seed_message_ids=(uuid.uuid4(),), + replay_seed_each_send=False, ) - context.target_normalization_context = target_context + context.prepended_history_send_context = target_context await attack._perform_async(context=context) @@ -296,7 +298,7 @@ async def test_perform_async_forwards_prepended_formatter_override(self): override = send_kwargs["normalizer_overrides"][CapabilityName.EDITABLE_HISTORY] assert isinstance(override, HistorySquashNormalizer) assert override._message_normalizer is formatter - assert send_kwargs["target_normalization_context"] is target_context + assert send_kwargs["send_context"] is target_context async def test_perform_async_sets_atomic_attack_identifier(self): """Test that _perform_async sets atomic_attack_identifier in the correct AtomicAttack format.""" diff --git a/tests/unit/executor/attack/multi_turn/test_multi_prompt_sending.py b/tests/unit/executor/attack/multi_turn/test_multi_prompt_sending.py index 397d91cbe2..40dee2a434 100644 --- a/tests/unit/executor/attack/multi_turn/test_multi_prompt_sending.py +++ b/tests/unit/executor/attack/multi_turn/test_multi_prompt_sending.py @@ -17,6 +17,9 @@ MultiTurnAttackContext, ) from pyrit.executor.attack.component import PrependedConversationConfig +from pyrit.executor.attack.component.prepended_history_send_context import ( + PrependedHistorySendContext, +) from pyrit.message_normalizer import HistorySquashNormalizer, MessageStringNormalizer from pyrit.models import ( AttackOutcome, @@ -32,7 +35,6 @@ PromptTarget, TargetCapabilities, TargetConfiguration, - TargetNormalizationContext, ) from pyrit.score import Scorer, TrueFalseScorer @@ -329,12 +331,12 @@ async def test_send_prompt_forwards_prepended_formatter_override( prompt_normalizer=mock_prompt_normalizer, prepended_conversation_config=PrependedConversationConfig(message_normalizer=formatter), ) - target_context = TargetNormalizationContext( + target_context = PrependedHistorySendContext( conversation_id=basic_context.session.conversation_id, - history_message_ids=(uuid.uuid4(),), - replay_history_each_send=False, + seed_message_ids=(uuid.uuid4(),), + replay_seed_each_send=False, ) - basic_context.target_normalization_context = target_context + basic_context.prepended_history_send_context = target_context mock_prompt_normalizer.send_prompt_async.return_value = sample_response await attack._send_prompt_to_objective_target_async( @@ -346,7 +348,7 @@ async def test_send_prompt_forwards_prepended_formatter_override( override = send_kwargs["normalizer_overrides"][CapabilityName.EDITABLE_HISTORY] assert isinstance(override, HistorySquashNormalizer) assert override._message_normalizer is formatter - assert send_kwargs["target_normalization_context"] is target_context + assert send_kwargs["send_context"] is target_context async def test_send_prompt_to_target_with_all_configurations( self, mock_target, mock_prompt_normalizer, basic_context, sample_response diff --git a/tests/unit/executor/attack/multi_turn/test_prepended_history_normalization.py b/tests/unit/executor/attack/multi_turn/test_prepended_history_normalization.py index 796207af56..727e735e30 100644 --- a/tests/unit/executor/attack/multi_turn/test_prepended_history_normalization.py +++ b/tests/unit/executor/attack/multi_turn/test_prepended_history_normalization.py @@ -144,7 +144,7 @@ async def test_single_turn_target_replays_seed_without_rotation(): target=target, ) persisted = manager.get_conversation(conversation_id) - target_context = manager.create_target_normalization_context( + target_context = manager.create_prepended_history_send_context( target=target, conversation_id=conversation_id, prepended_messages=persisted, @@ -158,9 +158,9 @@ async def test_single_turn_target_replays_seed_without_rotation(): conversation_id=conversation_id, normalizer_overrides=config.get_normalizer_overrides( target=target, - target_normalization_context=target_context, + prepended_history_send_context=target_context, ), - target_normalization_context=target_context, + send_context=target_context, ) assert target.prompt_sent == [ @@ -189,7 +189,7 @@ def test_rotation_is_noop_for_seeded_single_turn_target(): context.executed_turns = 2 original_id = context.session.conversation_id seed = Message.from_prompt(prompt="seed", role="user") - context.target_normalization_context = ConversationManager.create_target_normalization_context( + context.prepended_history_send_context = ConversationManager.create_prepended_history_send_context( target=target, conversation_id=original_id, prepended_messages=[seed], @@ -238,7 +238,7 @@ async def test_rotation_preserves_system_payload_for_later_single_turn_send(): _rotate(target=target, context=context) - assert context.target_normalization_context is not None + assert context.prepended_history_send_context is not None config = PrependedConversationConfig() await PromptNormalizer().send_prompt_async( message=Message.from_prompt(prompt="current request", role="user"), @@ -246,9 +246,9 @@ async def test_rotation_preserves_system_payload_for_later_single_turn_send(): conversation_id=context.session.conversation_id, normalizer_overrides=config.get_normalizer_overrides( target=target, - target_normalization_context=context.target_normalization_context, + prepended_history_send_context=context.prepended_history_send_context, ), - target_normalization_context=context.target_normalization_context, + send_context=context.prepended_history_send_context, ) assert target.prompt_sent == ["Turn 1:\nuser: ### Instructions ###\n\nsystem\n\n######\n\ncurrent request"] @@ -302,12 +302,12 @@ def _set_tap_seed_boundary( target=target, messages=seed_messages, ) - node._target_normalization_context = ConversationManager.create_target_normalization_context( + node._prepended_history_send_context = ConversationManager.create_prepended_history_send_context( target=target, conversation_id=node.objective_target_conversation_id, prepended_messages=seed_messages, ) - assert node._target_normalization_context is not None + assert node._prepended_history_send_context is not None def _branch_tap_node( @@ -339,19 +339,19 @@ async def format_messages(messages: list[Message]) -> str: _set_tap_seed_boundary(node=node, target=target, seed_messages=[seed]) node._objective = "objective" await node._send_prompt_to_target_async("depth one") - original_context = node._target_normalization_context + original_context = node._prepended_history_send_context assert original_context is not None retained, cloned = _branch_tap_node(node=node, branching_factor=2) assert retained is node - assert retained._target_normalization_context is original_context - assert cloned._target_normalization_context is not None - assert cloned._target_normalization_context.history_message_count == 1 + assert retained._prepended_history_send_context is original_context + assert cloned._prepended_history_send_context is not None + assert cloned._prepended_history_send_context.seed_message_count == 1 cloned_messages = CentralMemory.get_memory_instance().get_conversation_messages( conversation_id=cloned.objective_target_conversation_id ) - assert cloned._target_normalization_context.history_message_ids == (cloned_messages[0].get_piece().id,) - assert cloned._target_normalization_context.history_message_ids != original_context.history_message_ids + assert cloned._prepended_history_send_context.seed_message_ids == (cloned_messages[0].get_piece().id,) + assert cloned._prepended_history_send_context.seed_message_ids != original_context.seed_message_ids for branch, prompt in [(retained, "depth two retained"), (cloned, "depth two cloned")]: branch._objective = "objective" @@ -398,15 +398,15 @@ async def format_messages(messages: list[Message]) -> str: parent_conversation_id = node.objective_target_conversation_id await node._send_prompt_to_target_async("parent first") - assert node._target_normalization_context - assert node._target_normalization_context.is_consumed + assert node._prepended_history_send_context + assert node._prepended_history_send_context.is_seed_consumed cloned = node.duplicate() cloned._objective = "objective" cloned_conversation_id = cloned.objective_target_conversation_id assert cloned_conversation_id != parent_conversation_id - assert cloned._target_normalization_context - assert not cloned._target_normalization_context.is_consumed + assert cloned._prepended_history_send_context + assert not cloned._prepended_history_send_context.is_seed_consumed await node._send_prompt_to_target_async("parent second") await cloned._send_prompt_to_target_async("clone first") @@ -469,8 +469,8 @@ async def test_tap_unseeded_stateless_retained_and_cloned_branches_send_current_ await node._send_prompt_to_target_async("depth one") retained, cloned = _branch_tap_node(node=node, branching_factor=2) - assert retained._target_normalization_context is None - assert cloned._target_normalization_context is None + assert retained._prepended_history_send_context is None + assert cloned._prepended_history_send_context is None for branch, prompt in [(retained, "depth two retained"), (cloned, "depth two cloned")]: branch._objective = "objective" await branch._send_prompt_to_target_async(prompt) @@ -505,15 +505,15 @@ async def format_messages(messages: list[Message]) -> str: ) node._objective = "objective" await node._send_prompt_to_target_async("depth one") - original_context = node._target_normalization_context - original_boundary = original_context.history_message_ids if original_context else () + original_context = node._prepended_history_send_context + original_boundary = original_context.seed_message_ids if original_context else () branches = _branch_tap_node(node=node, branching_factor=1) assert branches == [node] - assert node._target_normalization_context is original_context - assert node._target_normalization_context is not None - assert node._target_normalization_context.history_message_ids == original_boundary + assert node._prepended_history_send_context is original_context + assert node._prepended_history_send_context is not None + assert node._prepended_history_send_context.seed_message_ids == original_boundary await node._send_prompt_to_target_async("depth two retained") assert target.prompt_sent == [ "original seed|depth one", 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 03768c2c11..8ba062cc05 100644 --- a/tests/unit/executor/attack/multi_turn/test_red_teaming.py +++ b/tests/unit/executor/attack/multi_turn/test_red_teaming.py @@ -990,7 +990,7 @@ async def test_second_turn_uses_configured_message_normalizer_without_rotation( ), ] ) - basic_context.target_normalization_context = ConversationManager.create_target_normalization_context( + basic_context.prepended_history_send_context = ConversationManager.create_prepended_history_send_context( target=objective_target, conversation_id=old_conversation_id, prepended_messages=[system_piece.to_message()], diff --git a/tests/unit/executor/attack/multi_turn/test_supports_multi_turn_attacks.py b/tests/unit/executor/attack/multi_turn/test_supports_multi_turn_attacks.py index 9809ce0a47..7cc2ce0625 100644 --- a/tests/unit/executor/attack/multi_turn/test_supports_multi_turn_attacks.py +++ b/tests/unit/executor/attack/multi_turn/test_supports_multi_turn_attacks.py @@ -7,6 +7,9 @@ import pytest from pyrit.executor.attack.component import PrependedConversationConfig +from pyrit.executor.attack.component.prepended_history_send_context import ( + PrependedHistorySendContext, +) from pyrit.executor.attack.core.attack_parameters import AttackParameters from pyrit.executor.attack.multi_turn.multi_turn_attack_strategy import ( ConversationSession, @@ -14,7 +17,7 @@ ) from pyrit.memory import CentralMemory from pyrit.models import ConversationType, Message, MessagePiece -from pyrit.prompt_target import PromptTarget, TargetNormalizationContext +from pyrit.prompt_target import PromptTarget from pyrit.prompt_target.common.target_capabilities import TargetCapabilities from pyrit.prompt_target.common.target_configuration import TargetConfiguration @@ -117,10 +120,10 @@ def test_pending_first_turn_context_suppresses_rotation_after_prepended_turns(se context = _make_context() context.executed_turns = 2 original_id = context.session.conversation_id - context.target_normalization_context = TargetNormalizationContext( + context.prepended_history_send_context = PrependedHistorySendContext( conversation_id=original_id, - history_message_ids=(uuid.uuid4(),), - replay_history_each_send=True, + seed_message_ids=(uuid.uuid4(),), + replay_seed_each_send=True, ) strategy._rotate_conversation_for_single_turn_target(context=context) @@ -238,13 +241,13 @@ async def test_rotated_system_prompt_is_normalized_with_next_request(self): message=next_request, normalizer_overrides=config.get_normalizer_overrides( target=target, - target_normalization_context=context.target_normalization_context, + prepended_history_send_context=context.prepended_history_send_context, ), - target_normalization_context=context.target_normalization_context, + send_context=context.prepended_history_send_context, ) - assert context.target_normalization_context is not None - assert not context.target_normalization_context.is_consumed + assert context.prepended_history_send_context is not None + assert not context.prepended_history_send_context.is_seed_consumed assert len(target.normalized_conversations) == 1 normalized_conversation = target.normalized_conversations[0] assert len(normalized_conversation) == 1 @@ -525,7 +528,7 @@ def test_single_turn_target_duplicates_logical_history_without_seed_boundary(sel assert [message.api_role for message in dup_messages] == ["system", "user", "assistant"] # No explicit prepended seed was initialized, so copied live turns do not # become a new replay boundary. - assert duplicate._target_normalization_context is None + assert duplicate._prepended_history_send_context is None def test_multi_turn_target_duplicates_full_conversation(self): """For multi-turn targets, the full conversation is duplicated.""" diff --git a/tests/unit/executor/attack/multi_turn/test_tree_of_attacks.py b/tests/unit/executor/attack/multi_turn/test_tree_of_attacks.py index e0e7beb588..671cd2c1c4 100644 --- a/tests/unit/executor/attack/multi_turn/test_tree_of_attacks.py +++ b/tests/unit/executor/attack/multi_turn/test_tree_of_attacks.py @@ -1951,7 +1951,7 @@ async def test_node_unseeded_single_turn_target_rotates_conversation_id(self, no assert node.objective_target_conversation_id != original_conv_id send_kwargs = node._prompt_normalizer.send_prompt_async.await_args.kwargs assert send_kwargs["conversation_id"] == node.objective_target_conversation_id - assert send_kwargs["target_normalization_context"] is None + assert send_kwargs["send_context"] is None @pytest.mark.asyncio async def test_node_multi_turn_target_keeps_conv_id(self, node_components): diff --git a/tests/unit/executor/attack/streaming/test_barge_in.py b/tests/unit/executor/attack/streaming/test_barge_in.py index 029b1e6677..ba86ee5671 100644 --- a/tests/unit/executor/attack/streaming/test_barge_in.py +++ b/tests/unit/executor/attack/streaming/test_barge_in.py @@ -11,9 +11,12 @@ import pytest from pyrit.executor.attack import BargeInAttack, BargeInAttackContext +from pyrit.executor.attack.component.prepended_history_send_context import ( + PrependedHistorySendContext, +) from pyrit.executor.attack.core import AttackParameters from pyrit.models import AttackOutcome, Message, MessagePiece -from pyrit.prompt_target import RealtimeTarget, TargetNormalizationContext +from pyrit.prompt_target import RealtimeTarget if TYPE_CHECKING: from collections.abc import AsyncIterator @@ -156,7 +159,7 @@ async def test_setup_async_persists_prepended_conversation_to_memory(vad_target) # All three messages share the context's conversation_id post-setup. for m in add_calls: assert m.message_pieces[0].conversation_id == ctx.conversation_id - assert ctx.target_normalization_context is None + assert ctx.prepended_history_send_context is None async def test_setup_async_clears_unused_normalization_context_when_prepended_empty(vad_target): @@ -166,10 +169,10 @@ async def test_setup_async_clears_unused_normalization_context_when_prepended_em params=AttackParameters(objective="o"), # no prepended_conversation audio_chunks=_aiter([b"\x00" * 96]), ) - ctx.target_normalization_context = TargetNormalizationContext( + ctx.prepended_history_send_context = PrependedHistorySendContext( conversation_id=ctx.conversation_id, - history_message_ids=(Message.from_prompt(prompt="unused", role="user").get_piece().id,), - replay_history_each_send=True, + seed_message_ids=(Message.from_prompt(prompt="unused", role="user").get_piece().id,), + replay_seed_each_send=True, ) add_calls: list[Any] = [] @@ -178,7 +181,7 @@ async def test_setup_async_clears_unused_normalization_context_when_prepended_em await attack._setup_async(context=ctx) assert add_calls == [] - assert ctx.target_normalization_context is None + assert ctx.prepended_history_send_context is None # ---- _perform_async: session factory passthrough ---------------------------------------------- diff --git a/tests/unit/prompt_normalizer/test_prompt_normalizer.py b/tests/unit/prompt_normalizer/test_prompt_normalizer.py index 58a2f4995b..bdc0f00ae8 100644 --- a/tests/unit/prompt_normalizer/test_prompt_normalizer.py +++ b/tests/unit/prompt_normalizer/test_prompt_normalizer.py @@ -25,6 +25,9 @@ execution_context, get_execution_context, ) +from pyrit.executor.attack.component.prepended_history_send_context import ( + PrependedHistorySendContext, +) from pyrit.memory import CentralMemory from pyrit.message_normalizer import MessageListNormalizer from pyrit.models import ( @@ -38,7 +41,7 @@ from pyrit.prompt_normalizer.converter_configuration import ( ConverterConfiguration, ) -from pyrit.prompt_target import CapabilityName, PromptTarget, TargetNormalizationContext +from pyrit.prompt_target import CapabilityName, PromptTarget from pyrit.prompt_target.common.target_capabilities import TargetCapabilities from pyrit.prompt_target.common.target_configuration import TargetConfiguration @@ -137,10 +140,10 @@ async def test_send_prompt_async_forwards_normalizer_overrides_and_context(mock_ conversation_id = "prepended-conversation" message_normalizer = MagicMock(spec=MessageListNormalizer) normalizer_overrides = {CapabilityName.EDITABLE_HISTORY: message_normalizer} - target_context = TargetNormalizationContext( + target_context = PrependedHistorySendContext( conversation_id=conversation_id, - history_message_ids=(uuid4(),), - replay_history_each_send=False, + seed_message_ids=(uuid4(),), + replay_seed_each_send=False, ) await normalizer.send_prompt_async( @@ -148,12 +151,12 @@ async def test_send_prompt_async_forwards_normalizer_overrides_and_context(mock_ target=prompt_target, conversation_id=conversation_id, normalizer_overrides=normalizer_overrides, - target_normalization_context=target_context, + send_context=target_context, ) call = prompt_target.send_prompt_async.await_args assert call.kwargs["normalizer_overrides"] == normalizer_overrides - assert call.kwargs["target_normalization_context"] is target_context + assert call.kwargs["send_context"] is target_context async def test_send_prompt_async_conversion_failure_does_not_call_target(mock_memory_instance): @@ -161,10 +164,10 @@ async def test_send_prompt_async_conversion_failure_does_not_call_target(mock_me prompt_target.get_identifier.return_value = get_mock_target_identifier("MockTarget") prompt_target.send_prompt_async = AsyncMock() conversation_id = "prepended-conversation" - target_context = TargetNormalizationContext( + target_context = PrependedHistorySendContext( conversation_id=conversation_id, - history_message_ids=(uuid4(),), - replay_history_each_send=False, + seed_message_ids=(uuid4(),), + replay_seed_each_send=False, ) converter_config = ConverterConfiguration.from_converters(converters=[ContextFailingConverter()]) @@ -174,7 +177,7 @@ async def test_send_prompt_async_conversion_failure_does_not_call_target(mock_me target=prompt_target, conversation_id=conversation_id, request_converter_configurations=converter_config, - target_normalization_context=target_context, + send_context=target_context, ) prompt_target.send_prompt_async.assert_not_awaited() @@ -186,10 +189,10 @@ async def test_send_prompt_async_target_failure_is_persisted(mock_memory_instanc prompt_target.get_identifier.return_value = get_mock_target_identifier("MockTarget") prompt_target.send_prompt_async = AsyncMock(side_effect=ValueError("normalization failed")) conversation_id = "prepended-conversation" - target_context = TargetNormalizationContext( + target_context = PrependedHistorySendContext( conversation_id=conversation_id, - history_message_ids=(uuid4(),), - replay_history_each_send=False, + seed_message_ids=(uuid4(),), + replay_seed_each_send=False, ) with pytest.raises(Exception, match="Error normalizing prompt with conversation ID"): @@ -197,7 +200,7 @@ async def test_send_prompt_async_target_failure_is_persisted(mock_memory_instanc message=Message.from_prompt(prompt="request", role="user"), target=prompt_target, conversation_id=conversation_id, - target_normalization_context=target_context, + send_context=target_context, ) assert target_context.provider_attempt_count == 0 @@ -220,10 +223,10 @@ async def wait_in_provider(*, normalized_conversation: list[Message]) -> list[Me raise RuntimeError("provider failed") target._send_prompt_to_target_async = AsyncMock(side_effect=wait_in_provider) # type: ignore[method-assign] - target_context = TargetNormalizationContext( + target_context = PrependedHistorySendContext( conversation_id=conversation_id, - history_message_ids=(seed.get_piece().id,), - replay_history_each_send=False, + seed_message_ids=(seed.get_piece().id,), + replay_seed_each_send=False, ) normalizer = PromptNormalizer() first_send = asyncio.create_task( @@ -231,7 +234,7 @@ async def wait_in_provider(*, normalized_conversation: list[Message]) -> list[Me message=Message.from_prompt(prompt="first request", role="user"), target=target, conversation_id=conversation_id, - target_normalization_context=target_context, + send_context=target_context, ) ) await provider_started.wait() @@ -241,7 +244,7 @@ async def wait_in_provider(*, normalized_conversation: list[Message]) -> list[Me message=Message.from_prompt(prompt="concurrent request", role="user"), target=target, conversation_id=conversation_id, - target_normalization_context=target_context, + send_context=target_context, ) mock_memory_instance.add_message_to_memory.assert_not_called() diff --git a/tests/unit/prompt_target/target/test_normalize_async_integration.py b/tests/unit/prompt_target/target/test_normalize_async_integration.py index 065fa111fe..29b14b7137 100644 --- a/tests/unit/prompt_target/target/test_normalize_async_integration.py +++ b/tests/unit/prompt_target/target/test_normalize_async_integration.py @@ -17,6 +17,9 @@ from openai.types.responses import ResponseOutputMessage, ResponseOutputText from unit.mocks import MockPromptTarget +from pyrit.executor.attack.component.prepended_history_send_context import ( + PrependedHistorySendContext, +) from pyrit.memory import CentralMemory from pyrit.memory.memory_interface import MemoryInterface from pyrit.message_normalizer import ( @@ -28,7 +31,7 @@ ) from pyrit.models import ComponentIdentifier, Message, MessagePiece, PromptResponseError from pyrit.prompt_normalizer import PromptNormalizer -from pyrit.prompt_target import AzureMLChatTarget, OpenAIChatTarget, TargetNormalizationContext +from pyrit.prompt_target import AzureMLChatTarget, OpenAIChatTarget from pyrit.prompt_target.common.target_capabilities import ( CapabilityHandlingPolicy, CapabilityName, @@ -56,29 +59,29 @@ def _make_message(*, role: str, content: str, conversation_id: str = "conv1") -> def _make_normalizer_overrides( *, - target_normalization_context: TargetNormalizationContext, + send_context: PrependedHistorySendContext, formatter: MessageStringNormalizer | None = None, ) -> dict[CapabilityName, HistorySquashNormalizer]: - if not target_normalization_context.should_include_history: + if not send_context.should_include_seed: return {} return { CapabilityName.EDITABLE_HISTORY: HistorySquashNormalizer( message_normalizer=formatter or ConversationContextNormalizer(), - expected_history_message_count=target_normalization_context.history_message_count, + expected_history_message_count=send_context.seed_message_count, ) } -def _make_target_normalization_context( +def _make_prepended_history_send_context( *, prepended_messages: list[Message], target_supports_multi_turn: bool = False, conversation_id: str = "conv1", -) -> TargetNormalizationContext: - return TargetNormalizationContext( +) -> PrependedHistorySendContext: + return PrependedHistorySendContext( conversation_id=conversation_id, - history_message_ids=tuple(message.get_piece().id for message in prepended_messages), - replay_history_each_send=not target_supports_multi_turn, + seed_message_ids=tuple(message.get_piece().id for message in prepended_messages), + replay_seed_each_send=not target_supports_multi_turn, ) @@ -554,13 +557,13 @@ async def test_non_editable_target_adapts_prepended_history_without_mutating_mem mock_memory = MagicMock(spec=MemoryInterface) mock_memory.get_conversation_messages.return_value = memory_messages target._memory = mock_memory - target_context = _make_target_normalization_context(prepended_messages=list(memory_messages)) - normalizer_overrides = _make_normalizer_overrides(target_normalization_context=target_context) + target_context = _make_prepended_history_send_context(prepended_messages=list(memory_messages)) + normalizer_overrides = _make_normalizer_overrides(send_context=target_context) result = await target._get_normalized_conversation_async( message=live_request, normalizer_overrides=normalizer_overrides, - target_normalization_context=target_context, + send_context=target_context, ) assert len(result) == 1 @@ -599,13 +602,13 @@ async def test_non_editable_target_preserves_system_history_and_multimodal_live_ mock_memory = MagicMock(spec=MemoryInterface) mock_memory.get_conversation_messages.return_value = [system_message] target._memory = mock_memory - target_context = _make_target_normalization_context(prepended_messages=[system_message]) - normalizer_overrides = _make_normalizer_overrides(target_normalization_context=target_context) + target_context = _make_prepended_history_send_context(prepended_messages=[system_message]) + normalizer_overrides = _make_normalizer_overrides(send_context=target_context) result = await target._get_normalized_conversation_async( message=live_request, normalizer_overrides=normalizer_overrides, - target_normalization_context=target_context, + send_context=target_context, ) assert len(result) == 1 @@ -636,15 +639,15 @@ async def test_editable_history_override_runs_before_system_prompt_adaptation(): mock_memory = MagicMock(spec=MemoryInterface) mock_memory.get_conversation_messages.return_value = [prepended] target._memory = mock_memory - target_context = _make_target_normalization_context( + target_context = _make_prepended_history_send_context( prepended_messages=[prepended], target_supports_multi_turn=True, ) result = await target._get_normalized_conversation_async( message=live_request, - normalizer_overrides=_make_normalizer_overrides(target_normalization_context=target_context), - target_normalization_context=target_context, + normalizer_overrides=_make_normalizer_overrides(send_context=target_context), + send_context=target_context, ) assert [message.get_value() for message in result] == [ @@ -689,13 +692,13 @@ async def test_first_turn_normalization_preserves_live_multimodal_piece_order(): prepended = _make_message(role="user", content="prepended") mock_memory.get_conversation_messages.return_value = [prepended] target._memory = mock_memory - target_context = _make_target_normalization_context(prepended_messages=[prepended]) - normalizer_overrides = _make_normalizer_overrides(target_normalization_context=target_context) + target_context = _make_prepended_history_send_context(prepended_messages=[prepended]) + normalizer_overrides = _make_normalizer_overrides(send_context=target_context) result = await target._get_normalized_conversation_async( message=live_request, normalizer_overrides=normalizer_overrides, - target_normalization_context=target_context, + send_context=target_context, ) assert [piece.converted_value_data_type for piece in result[0].message_pieces] == ["image_path", "text"] @@ -742,12 +745,12 @@ async def test_history_squash_does_not_restore_adapted_json_schema_metadata(): mock_memory = MagicMock(spec=MemoryInterface) mock_memory.get_conversation_messages.return_value = [prepended] target._memory = mock_memory - target_context = _make_target_normalization_context(prepended_messages=[prepended]) + target_context = _make_prepended_history_send_context(prepended_messages=[prepended]) result = await target._get_normalized_conversation_async( message=live_request, - normalizer_overrides=_make_normalizer_overrides(target_normalization_context=target_context), - target_normalization_context=target_context, + normalizer_overrides=_make_normalizer_overrides(send_context=target_context), + send_context=target_context, ) image_piece, text_piece = result[0].message_pieces @@ -804,20 +807,20 @@ async def test_prepended_history_adapter_is_used_only_when_explicitly_passed(): [prepended, prior_live, prior_response], ] target._memory = mock_memory - target_context = _make_target_normalization_context( + target_context = _make_prepended_history_send_context( prepended_messages=[prepended], target_supports_multi_turn=True, ) await target.send_prompt_async( message=prior_live, - normalizer_overrides=_make_normalizer_overrides(target_normalization_context=target_context), - target_normalization_context=target_context, + normalizer_overrides=_make_normalizer_overrides(send_context=target_context), + send_context=target_context, ) await target.send_prompt_async( message=second_live, - normalizer_overrides=_make_normalizer_overrides(target_normalization_context=target_context), - target_normalization_context=target_context, + normalizer_overrides=_make_normalizer_overrides(send_context=target_context), + send_context=target_context, ) assert target.prompt_sent == ["Turn 1:\nuser: prepended\nTurn 2:\nuser: first live", "second live"] @@ -845,20 +848,20 @@ async def test_non_editable_multi_turn_target_sends_only_current_request_after_r target._send_prompt_to_target_async = AsyncMock( # type: ignore[method-assign] return_value=[_make_message(role="assistant", content="response")] ) - target_context = _make_target_normalization_context( + target_context = _make_prepended_history_send_context( prepended_messages=[prepended], target_supports_multi_turn=True, ) await target.send_prompt_async( message=first_live, - normalizer_overrides=_make_normalizer_overrides(target_normalization_context=target_context), - target_normalization_context=target_context, + normalizer_overrides=_make_normalizer_overrides(send_context=target_context), + send_context=target_context, ) await target.send_prompt_async( message=second_live, - normalizer_overrides=_make_normalizer_overrides(target_normalization_context=target_context), - target_normalization_context=target_context, + normalizer_overrides=_make_normalizer_overrides(send_context=target_context), + send_context=target_context, ) first_payload, second_payload = target._send_prompt_to_target_async.await_args_list @@ -881,17 +884,17 @@ async def test_stateless_target_replays_only_seed_and_current_request(): [prepended, first_live, first_response], ] target._memory = mock_memory - target_context = _make_target_normalization_context(prepended_messages=[prepended]) + target_context = _make_prepended_history_send_context(prepended_messages=[prepended]) await target.send_prompt_async( message=first_live, - normalizer_overrides=_make_normalizer_overrides(target_normalization_context=target_context), - target_normalization_context=target_context, + normalizer_overrides=_make_normalizer_overrides(send_context=target_context), + send_context=target_context, ) await target.send_prompt_async( message=second_live, - normalizer_overrides=_make_normalizer_overrides(target_normalization_context=target_context), - target_normalization_context=target_context, + normalizer_overrides=_make_normalizer_overrides(send_context=target_context), + send_context=target_context, ) assert target.prompt_sent == [ @@ -919,20 +922,20 @@ async def test_stateful_target_consumes_seed_after_provider_outcome(response_err target._send_prompt_to_target_async = AsyncMock( # type: ignore[method-assign] side_effect=[[provider_response], [_make_message(role="assistant", content="second response")]] ) - target_context = _make_target_normalization_context( + target_context = _make_prepended_history_send_context( prepended_messages=[prepended], target_supports_multi_turn=True, ) await target.send_prompt_async( message=first_live, - normalizer_overrides=_make_normalizer_overrides(target_normalization_context=target_context), - target_normalization_context=target_context, + normalizer_overrides=_make_normalizer_overrides(send_context=target_context), + send_context=target_context, ) await target.send_prompt_async( message=second_live, - normalizer_overrides=_make_normalizer_overrides(target_normalization_context=target_context), - target_normalization_context=target_context, + normalizer_overrides=_make_normalizer_overrides(send_context=target_context), + send_context=target_context, ) first_payload, second_payload = target._send_prompt_to_target_async.await_args_list @@ -950,16 +953,16 @@ async def test_non_editable_target_uses_custom_prepended_formatter(): target._memory = mock_memory formatter = MagicMock(spec=MessageStringNormalizer) formatter.normalize_string_async = AsyncMock(return_value="CUSTOM HISTORY") - target_context = _make_target_normalization_context(prepended_messages=[prepended]) + target_context = _make_prepended_history_send_context(prepended_messages=[prepended]) normalizer_overrides = _make_normalizer_overrides( - target_normalization_context=target_context, + send_context=target_context, formatter=formatter, ) result = await target._get_normalized_conversation_async( message=_make_message(role="user", content="live"), normalizer_overrides=normalizer_overrides, - target_normalization_context=target_context, + send_context=target_context, ) assert result[0].get_value() == "CUSTOM HISTORY" @@ -979,14 +982,14 @@ async def test_non_editable_target_rejects_non_text_converted_prepended_history( mock_memory = MagicMock(spec=MemoryInterface) mock_memory.get_conversation_messages.return_value = [prepended] target._memory = mock_memory - target_context = _make_target_normalization_context(prepended_messages=[prepended]) - normalizer_overrides = _make_normalizer_overrides(target_normalization_context=target_context) + target_context = _make_prepended_history_send_context(prepended_messages=[prepended]) + normalizer_overrides = _make_normalizer_overrides(send_context=target_context) with pytest.raises(ValueError, match="non-text output types.*image_path"): await target.send_prompt_async( message=_make_message(role="user", content="live"), normalizer_overrides=normalizer_overrides, - target_normalization_context=target_context, + send_context=target_context, ) @@ -1009,14 +1012,14 @@ async def test_non_editable_target_rejects_same_modality_non_text_conversion(): mock_memory = MagicMock(spec=MemoryInterface) mock_memory.get_conversation_messages.return_value = [prepended] target._memory = mock_memory - target_context = _make_target_normalization_context(prepended_messages=[prepended]) - normalizer_overrides = _make_normalizer_overrides(target_normalization_context=target_context) + target_context = _make_prepended_history_send_context(prepended_messages=[prepended]) + normalizer_overrides = _make_normalizer_overrides(send_context=target_context) with pytest.raises(ValueError, match="non-text output types.*image_path"): await target.send_prompt_async( message=_make_message(role="user", content="live"), normalizer_overrides=normalizer_overrides, - target_normalization_context=target_context, + send_context=target_context, ) @@ -1040,13 +1043,13 @@ async def test_non_editable_target_allows_preexisting_non_text_history_with_conv mock_memory = MagicMock(spec=MemoryInterface) mock_memory.get_conversation_messages.return_value = [prepended] target._memory = mock_memory - target_context = _make_target_normalization_context(prepended_messages=[prepended]) - normalizer_overrides = _make_normalizer_overrides(target_normalization_context=target_context) + target_context = _make_prepended_history_send_context(prepended_messages=[prepended]) + normalizer_overrides = _make_normalizer_overrides(send_context=target_context) result = await target._get_normalized_conversation_async( message=_make_message(role="user", content="live"), normalizer_overrides=normalizer_overrides, - target_normalization_context=target_context, + send_context=target_context, ) assert result[0].get_value() == "Turn 1:\nuser: [Image_path]\nTurn 2:\nuser: live" @@ -1073,13 +1076,13 @@ async def test_non_editable_target_warns_when_non_text_history_becomes_a_placeho mock_memory = MagicMock(spec=MemoryInterface) mock_memory.get_conversation_messages.return_value = [prepended] target._memory = mock_memory - target_context = _make_target_normalization_context(prepended_messages=[prepended]) + target_context = _make_prepended_history_send_context(prepended_messages=[prepended]) with caplog.at_level(logging.WARNING, logger="pyrit.message_normalizer.history_squash_normalizer"): await target._get_normalized_conversation_async( message=_make_message(role="user", content="live"), - normalizer_overrides=_make_normalizer_overrides(target_normalization_context=target_context), - target_normalization_context=target_context, + normalizer_overrides=_make_normalizer_overrides(send_context=target_context), + send_context=target_context, ) assert "image_path" in caplog.text @@ -1096,9 +1099,9 @@ async def test_target_normalization_failure_can_be_retried(): target._memory = mock_memory formatter = MagicMock(spec=MessageStringNormalizer) formatter.normalize_string_async = AsyncMock(side_effect=[ValueError("format failed"), "formatted request"]) - target_context = _make_target_normalization_context(prepended_messages=[prepended]) + target_context = _make_prepended_history_send_context(prepended_messages=[prepended]) normalizer_overrides = _make_normalizer_overrides( - target_normalization_context=target_context, + send_context=target_context, formatter=formatter, ) live_request = _make_message(role="user", content="live") @@ -1107,13 +1110,13 @@ async def test_target_normalization_failure_can_be_retried(): await target.send_prompt_async( message=live_request, normalizer_overrides=normalizer_overrides, - target_normalization_context=target_context, + send_context=target_context, ) await target.send_prompt_async( message=live_request, normalizer_overrides=normalizer_overrides, - target_normalization_context=target_context, + send_context=target_context, ) assert target.prompt_sent == ["formatted request"] @@ -1227,9 +1230,9 @@ async def test_target_normalization_cancellation_propagates(): target._memory = mock_memory formatter = MagicMock(spec=MessageStringNormalizer) formatter.normalize_string_async = AsyncMock(side_effect=asyncio.CancelledError()) - target_context = _make_target_normalization_context(prepended_messages=[prepended]) + target_context = _make_prepended_history_send_context(prepended_messages=[prepended]) normalizer_overrides = _make_normalizer_overrides( - target_normalization_context=target_context, + send_context=target_context, formatter=formatter, ) @@ -1237,10 +1240,10 @@ async def test_target_normalization_cancellation_propagates(): await target.send_prompt_async( message=_make_message(role="user", content="live"), normalizer_overrides=normalizer_overrides, - target_normalization_context=target_context, + send_context=target_context, ) assert target_context.provider_attempt_count == 0 - assert not target_context.is_consumed + assert not target_context.is_seed_consumed @pytest.mark.usefixtures("patch_central_database") @@ -1252,20 +1255,20 @@ async def test_provider_failure_propagates_after_normalization(): mock_memory.get_conversation_messages.return_value = [prepended] target._memory = mock_memory target._send_prompt_to_target_async = AsyncMock(side_effect=RuntimeError("provider failed")) # type: ignore[method-assign] - target_context = _make_target_normalization_context( + target_context = _make_prepended_history_send_context( prepended_messages=[prepended], target_supports_multi_turn=True, ) - normalizer_overrides = _make_normalizer_overrides(target_normalization_context=target_context) + normalizer_overrides = _make_normalizer_overrides(send_context=target_context) with pytest.raises(RuntimeError, match="provider failed"): await target.send_prompt_async( message=_make_message(role="user", content="live"), normalizer_overrides=normalizer_overrides, - target_normalization_context=target_context, + send_context=target_context, ) assert target_context.provider_attempt_count == 1 - assert target_context.is_consumed + assert target_context.is_seed_consumed @pytest.mark.usefixtures("patch_central_database") @@ -1290,15 +1293,15 @@ async def wait_in_provider(*, normalized_conversation: list[Message]) -> list[Me return [_make_message(role="assistant", content="response")] target._send_prompt_to_target_async = AsyncMock(side_effect=wait_in_provider) # type: ignore[method-assign] - target_context = _make_target_normalization_context( + target_context = _make_prepended_history_send_context( prepended_messages=[prepended], target_supports_multi_turn=True, ) first_send = asyncio.create_task( target.send_prompt_async( message=first_live, - normalizer_overrides=_make_normalizer_overrides(target_normalization_context=target_context), - target_normalization_context=target_context, + normalizer_overrides=_make_normalizer_overrides(send_context=target_context), + send_context=target_context, ) ) await provider_started.wait() @@ -1306,7 +1309,7 @@ async def wait_in_provider(*, normalized_conversation: list[Message]) -> list[Me with pytest.raises(asyncio.CancelledError): await first_send - assert target_context.is_consumed + assert target_context.is_seed_consumed assert target_context.provider_attempt_count == 1 target._send_prompt_to_target_async = AsyncMock( # type: ignore[method-assign] @@ -1314,8 +1317,8 @@ async def wait_in_provider(*, normalized_conversation: list[Message]) -> list[Me ) await target.send_prompt_async( message=second_live, - normalizer_overrides=_make_normalizer_overrides(target_normalization_context=target_context), - target_normalization_context=target_context, + normalizer_overrides=_make_normalizer_overrides(send_context=target_context), + send_context=target_context, ) payload = target._send_prompt_to_target_async.await_args.kwargs["normalized_conversation"] assert [message.get_value() for message in payload] == ["second live"] @@ -1339,16 +1342,16 @@ async def wait_to_format(messages: list[Message]) -> str: formatter = MagicMock(spec=MessageStringNormalizer) formatter.normalize_string_async = AsyncMock(side_effect=wait_to_format) - target_context = _make_target_normalization_context(prepended_messages=[prepended]) + target_context = _make_prepended_history_send_context(prepended_messages=[prepended]) normalizer_overrides = _make_normalizer_overrides( - target_normalization_context=target_context, + send_context=target_context, formatter=formatter, ) first_send = asyncio.create_task( target.send_prompt_async( message=_make_message(role="user", content="first"), normalizer_overrides=normalizer_overrides, - target_normalization_context=target_context, + send_context=target_context, ) ) await started.wait() @@ -1357,7 +1360,7 @@ async def wait_to_format(messages: list[Message]) -> str: await target.send_prompt_async( message=_make_message(role="user", content="second"), normalizer_overrides=normalizer_overrides, - target_normalization_context=target_context, + send_context=target_context, ) finally: release.set() @@ -1378,16 +1381,16 @@ async def test_tokenizer_formatter_receives_live_request_before_generation_promp tokenizer = MagicMock() tokenizer.apply_chat_template.return_value = "TOKENIZED REQUEST" formatter = TokenizerTemplateNormalizer(tokenizer=tokenizer) - target_context = _make_target_normalization_context(prepended_messages=[prepended]) + target_context = _make_prepended_history_send_context(prepended_messages=[prepended]) normalizer_overrides = _make_normalizer_overrides( - target_normalization_context=target_context, + send_context=target_context, formatter=formatter, ) result = await target._get_normalized_conversation_async( message=_make_message(role="user", content="live"), normalizer_overrides=normalizer_overrides, - target_normalization_context=target_context, + send_context=target_context, ) tokenizer_messages = tokenizer.apply_chat_template.call_args.args[0] diff --git a/tests/unit/prompt_target/target/test_target_history.py b/tests/unit/prompt_target/target/test_target_history.py new file mode 100644 index 0000000000..dda5313525 --- /dev/null +++ b/tests/unit/prompt_target/target/test_target_history.py @@ -0,0 +1,71 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT license. + +import pytest + +from pyrit.models import ChatMessageRole, Message, PromptResponseError +from pyrit.prompt_target.common.target_history import filter_non_replayable_messages + + +def _message( + *, + role: ChatMessageRole, + value: str, + sequence: int, + conversation_id: str = "conversation", +) -> Message: + message = Message.from_prompt(prompt=value, role=role) + piece = message.get_piece() + piece.sequence = sequence + piece.conversation_id = conversation_id + return message + + +@pytest.mark.parametrize("response_error", ["processing", "unknown"]) +def test_filter_removes_adjacent_failed_exchange(response_error: PromptResponseError) -> None: + successful = _message(role="user", value="successful", sequence=0) + failed_request = _message(role="user", value="failed request", sequence=1) + error_response = _message(role="assistant", value="private stack trace", sequence=2) + error_response.get_piece().response_error = response_error + + filtered = filter_non_replayable_messages(messages=[successful, failed_request, error_response]) + + assert filtered == [successful] + + +@pytest.mark.parametrize( + ("preceding_role", "preceding_conversation_id", "preceding_sequence"), + [ + ("assistant", "conversation", 4), + ("user", "other-conversation", 4), + ("user", "conversation", 3), + ], +) +def test_filter_does_not_remove_unrelated_preceding_message( + preceding_role: ChatMessageRole, + preceding_conversation_id: str, + preceding_sequence: int, +) -> None: + preceding = _message( + role=preceding_role, + value="unrelated", + sequence=preceding_sequence, + conversation_id=preceding_conversation_id, + ) + error_response = _message(role="assistant", value="private stack trace", sequence=5) + error_response.get_piece().response_error = "processing" + + filtered = filter_non_replayable_messages(messages=[preceding, error_response]) + + assert filtered == [preceding] + + +@pytest.mark.parametrize("response_error", ["blocked", "empty"]) +def test_filter_retains_provider_round_trip(response_error: PromptResponseError) -> None: + request = _message(role="user", value="request", sequence=0) + response = _message(role="assistant", value="provider response", sequence=1) + response.get_piece().response_error = response_error + + filtered = filter_non_replayable_messages(messages=[request, response]) + + assert filtered == [request, response] diff --git a/tests/unit/prompt_target/target/test_target_normalization_context.py b/tests/unit/prompt_target/target/test_target_normalization_context.py deleted file mode 100644 index 521b74547e..0000000000 --- a/tests/unit/prompt_target/target/test_target_normalization_context.py +++ /dev/null @@ -1,228 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# Licensed under the MIT license. - -import uuid - -import pytest - -from pyrit.models import ChatMessageRole, Message, PromptResponseError -from pyrit.prompt_target.common.target_normalization_context import ( - TargetNormalizationContext, - filter_non_replayable_messages, -) - - -def _message( - *, - role: ChatMessageRole, - value: str, - sequence: int, - conversation_id: str = "conversation", -) -> Message: - message = Message.from_prompt(prompt=value, role=role) - piece = message.get_piece() - piece.sequence = sequence - piece.conversation_id = conversation_id - return message - - -def test_context_requires_unique_persisted_history_ids(): - message_id = uuid.uuid4() - - with pytest.raises(ValueError, match="conversation_id"): - TargetNormalizationContext( - conversation_id="", - history_message_ids=(message_id,), - replay_history_each_send=False, - ) - with pytest.raises(ValueError, match="history_message_ids"): - TargetNormalizationContext( - conversation_id="conversation", - history_message_ids=(), - replay_history_each_send=False, - ) - with pytest.raises(ValueError, match="unique"): - TargetNormalizationContext( - conversation_id="conversation", - history_message_ids=(message_id, message_id), - replay_history_each_send=False, - ) - - -def test_stateful_context_consumes_explicit_boundary_at_provider_attempt(): - first = _message(role="system", value="system", sequence=0) - second = _message(role="user", value="seed", sequence=1) - unrelated = _message(role="assistant", value="later response", sequence=2) - context = TargetNormalizationContext( - conversation_id="conversation", - history_message_ids=(first.get_piece().id, second.get_piece().id), - replay_history_each_send=False, - ) - - context.begin_send() - selected = context.select_history(messages=[unrelated, second, first]) - context.mark_provider_attempted() - context.finish_send() - - assert selected == [first, second] - assert context.is_consumed - assert context.provider_attempt_count == 1 - - context.begin_send() - assert context.select_history(messages=[first, second, unrelated]) == [] - context.finish_send() - - -def test_stateless_context_replays_explicit_boundary_after_provider_attempt(): - seed = _message(role="user", value="seed", sequence=0) - context = TargetNormalizationContext( - conversation_id="conversation", - history_message_ids=(seed.get_piece().id,), - replay_history_each_send=True, - ) - - for _ in range(2): - context.begin_send() - assert context.select_history(messages=[seed]) == [seed] - context.mark_provider_attempted() - context.finish_send() - - assert not context.is_consumed - assert context.provider_attempt_count == 2 - - -def test_context_rejects_concurrent_send_until_active_send_finishes(): - context = TargetNormalizationContext( - conversation_id="conversation", - history_message_ids=(uuid.uuid4(),), - replay_history_each_send=False, - ) - - context.begin_send() - with pytest.raises(RuntimeError, match="Concurrent sends"): - context.begin_send() - - context.finish_send() - context.begin_send() - context.finish_send() - - -def test_context_rejects_provider_attempt_without_active_send(): - context = TargetNormalizationContext( - conversation_id="conversation", - history_message_ids=(uuid.uuid4(),), - replay_history_each_send=False, - ) - - with pytest.raises(RuntimeError, match="without an active send"): - context.mark_provider_attempted() - - -def test_context_rejects_missing_persisted_boundary_message(): - context = TargetNormalizationContext( - conversation_id="conversation", - history_message_ids=(uuid.uuid4(),), - replay_history_each_send=False, - ) - - with pytest.raises(ValueError, match="Missing 1 message"): - context.select_history(messages=[]) - - -@pytest.mark.parametrize("response_error", ["processing", "unknown"]) -def test_filter_removes_adjacent_failed_exchange(response_error: PromptResponseError): - successful = _message(role="user", value="successful", sequence=0) - failed_request = _message(role="user", value="failed request", sequence=1) - error_response = _message(role="assistant", value="private stack trace", sequence=2) - error_response.get_piece().response_error = response_error - - filtered = filter_non_replayable_messages(messages=[successful, failed_request, error_response]) - - assert filtered == [successful] - - -@pytest.mark.parametrize( - ("preceding_role", "preceding_conversation_id", "preceding_sequence"), - [ - ("assistant", "conversation", 4), - ("user", "other-conversation", 4), - ("user", "conversation", 3), - ], -) -def test_filter_does_not_remove_unrelated_preceding_message( - preceding_role: ChatMessageRole, - preceding_conversation_id: str, - preceding_sequence: int, -): - preceding = _message( - role=preceding_role, - value="unrelated", - sequence=preceding_sequence, - conversation_id=preceding_conversation_id, - ) - error_response = _message(role="assistant", value="private stack trace", sequence=5) - error_response.get_piece().response_error = "processing" - - filtered = filter_non_replayable_messages(messages=[preceding, error_response]) - - assert filtered == [preceding] - - -@pytest.mark.parametrize("response_error", ["blocked", "empty"]) -def test_filter_retains_provider_round_trip(response_error: PromptResponseError): - request = _message(role="user", value="request", sequence=0) - response = _message(role="assistant", value="provider response", sequence=1) - response.get_piece().response_error = response_error - - filtered = filter_non_replayable_messages(messages=[request, response]) - - assert filtered == [request, response] - - -def test_context_remaps_only_explicit_history_for_duplicate_conversation(): - seed = _message(role="system", value="seed", sequence=0) - live_request = _message(role="user", value="live", sequence=1) - source_messages = [seed, live_request] - duplicated_messages = [message.duplicate() for message in source_messages] - for message in duplicated_messages: - for piece in message.message_pieces: - piece.conversation_id = "duplicate" - - context = TargetNormalizationContext( - conversation_id="conversation", - history_message_ids=(seed.get_piece().id,), - replay_history_each_send=True, - ) - - duplicate = context.remap_for_duplicate_conversation( - conversation_id="duplicate", - source_messages=source_messages, - duplicated_messages=duplicated_messages, - ) - - assert duplicate.history_message_ids == (duplicated_messages[0].get_piece().id,) - assert duplicated_messages[1].get_piece().id not in duplicate.history_message_ids - assert duplicate.select_history(messages=duplicated_messages) == [duplicated_messages[0]] - - -def test_context_remap_resets_consumed_state_for_new_conversation(): - seed = _message(role="user", value="seed", sequence=0) - duplicated_seed = seed.duplicate() - duplicated_seed.get_piece().conversation_id = "duplicate" - context = TargetNormalizationContext( - conversation_id="conversation", - history_message_ids=(seed.get_piece().id,), - replay_history_each_send=False, - ) - context.begin_send() - context.mark_provider_attempted() - context.finish_send() - - duplicate = context.remap_for_duplicate_conversation( - conversation_id="duplicate", - source_messages=[seed], - duplicated_messages=[duplicated_seed], - ) - - assert not duplicate.is_consumed - assert duplicate.select_history(messages=[duplicated_seed]) == [duplicated_seed] From 73336d9a26a87633f01a7b202f6c4dc6bcbe3fb9 Mon Sep 17 00:00:00 2001 From: Copilot <223556219+Copilot@users.noreply.github.com> Date: Fri, 21 Aug 2026 21:25:23 -0700 Subject: [PATCH 23/24] Address prepended history lifecycle feedback Restore stateful branch bootstrap and reconnect history, align BargeIn delivery with persisted conversion, and consume prepended state only when provider interaction begins. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> Copilot-Session: 6776aa95-0924-4967-a056-e5a228af3faa --- doc/code/framework.md | 4 +- doc/code/targets/0_prompt_targets.md | 15 +++ .../attack/component/conversation_manager.py | 10 +- .../prepended_conversation_config.py | 2 +- .../prepended_history_send_context.py | 80 +++++++---- .../attack/multi_turn/tree_of_attacks.py | 83 +++++++++--- pyrit/executor/attack/streaming/barge_in.py | 8 ++ pyrit/prompt_target/common/prompt_target.py | 32 ++++- .../common/target_send_context.py | 29 ++++ pyrit/prompt_target/common/utils.py | 13 ++ .../http_target/httpx_api_target.py | 3 + .../playwright_copilot_target.py | 12 ++ pyrit/prompt_target/playwright_target.py | 2 + .../prompt_target/websocket_copilot_target.py | 6 +- pyrit/prompt_target/websocket_target.py | 2 + .../test_prepended_history_send_context.py | 47 ++++++- .../test_prepended_history_normalization.py | 107 ++++++++++++++- .../attack/streaming/test_barge_in.py | 58 +++++++- .../test_normalize_async_integration.py | 50 ++++++- .../target/test_playwright_copilot_target.py | 37 ++++- .../target/test_websocket_target.py | 127 ++++++++++++++++++ 21 files changed, 669 insertions(+), 58 deletions(-) diff --git a/doc/code/framework.md b/doc/code/framework.md index cab56e3dc5..b75ecf231b 100644 --- a/doc/code/framework.md +++ b/doc/code/framework.md @@ -233,7 +233,9 @@ If you are contributing to PyRIT, that work will most likely land in one of the - This is often an LLM, but it doesn't have to be. For Cross-Domain Prompt Injection Attacks, the target might be a storage account that a later target has a reference to. Message and conversation should be generic enough to handle this extra data. - Target capabilities should be used to see if a target is compatible with the capabilities that the other components want to use. - Targets should use message_normalizer along with TargetConfiguration to transform `Messages` into formats that target supports. -- A target may observe an internal, caller-owned send context at the provider-invocation boundary, but the caller owns any seed-history identity, replay, or branching state. +- A target may observe an internal, caller-owned send context at the provider-invocation boundary, + after target-side waits and immediately before irreversible provider I/O, but the caller owns any + bootstrap-history identity, replay, or branching state. - 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), 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`). diff --git a/doc/code/targets/0_prompt_targets.md b/doc/code/targets/0_prompt_targets.md index 963ee8fca2..1da66a6bfd 100644 --- a/doc/code/targets/0_prompt_targets.md +++ b/doc/code/targets/0_prompt_targets.md @@ -34,6 +34,21 @@ converters have already run by this point, so role-specific converter choices re the target must receive one flattened request. The context is ephemeral and does not replace the structured messages stored in memory. +For a stateful target without editable history, the initial bootstrap is flattened once. Later sends +retain replayable memory history in the normalized view so a target such as `WebsocketTarget` can +restore a replaced provider session, while the existing provider session still receives only the +current request. A stateful TAP clone bootstraps its new provider session with the complete replayable +duplicated branch, even when the attack started without an explicit prepended seed. Stateless targets +continue to receive only the explicit prepended seed plus each current request, never prior live branch +turns. A stateful TAP target without editable history can flatten only text converter output when +branching; non-text converter output requires an editable-history target, a stateless target, or +`branching_factor=1` so copied media is never replayed under a different role. + +The provider-attempt signal is emitted after shared target-side rate limiting. Targets that need more +precise setup, such as WebSocket, Playwright, or conversation-keyed HTTP targets, emit it immediately +before the irreversible provider operation. Cancellation before that point leaves a one-time +bootstrap available for retry. + `send_prompt_async` is the final public orchestration method. Custom target subclasses implement `_send_prompt_to_target_async(*, normalized_conversation: list[Message]) -> list[Message]` instead of overriding `send_prompt_async`. diff --git a/pyrit/executor/attack/component/conversation_manager.py b/pyrit/executor/attack/component/conversation_manager.py index 731dc9cfc0..893da6c752 100644 --- a/pyrit/executor/attack/component/conversation_manager.py +++ b/pyrit/executor/attack/component/conversation_manager.py @@ -508,6 +508,10 @@ async def _process_prepended_conversation_async( if not valid_messages: return state + existing_message_ids = { + message.get_piece().id for message in self.get_conversation(conversation_id=conversation_id) + } + # Use the lower-level method to add messages to memory state.turn_count = await self.add_prepended_conversation_to_memory_async( prepended_conversation=prepended_conversation, @@ -518,7 +522,11 @@ async def _process_prepended_conversation_async( target_identifier=target_identifier, target=target, ) - persisted_messages = self.get_conversation(conversation_id) + persisted_messages = [ + message + for message in self.get_conversation(conversation_id) + if message.get_piece().id not in existing_message_ids + ] context.prepended_history_send_context = self.create_prepended_history_send_context( target=target, conversation_id=conversation_id, diff --git a/pyrit/executor/attack/component/prepended_conversation_config.py b/pyrit/executor/attack/component/prepended_conversation_config.py index d277921f13..3ec3e35b61 100644 --- a/pyrit/executor/attack/component/prepended_conversation_config.py +++ b/pyrit/executor/attack/component/prepended_conversation_config.py @@ -85,7 +85,7 @@ def get_normalizer_overrides( return { CapabilityName.EDITABLE_HISTORY: HistorySquashNormalizer( - expected_history_message_count=prepended_history_send_context.seed_message_count, + expected_history_message_count=prepended_history_send_context.bootstrap_message_count, message_normalizer=self.get_message_normalizer(), ) } diff --git a/pyrit/executor/attack/component/prepended_history_send_context.py b/pyrit/executor/attack/component/prepended_history_send_context.py index e4b55350ef..a11df417fc 100644 --- a/pyrit/executor/attack/component/prepended_history_send_context.py +++ b/pyrit/executor/attack/component/prepended_history_send_context.py @@ -17,19 +17,23 @@ @dataclass(slots=True) class PrependedHistorySendContext: """ - Per-execution state for delivering an explicit persisted seed prefix. - - The seed is identified by persisted message-piece IDs instead of inferred - from response roles or error codes. A context may be used by only one send - at a time. Stateful targets consume the seed when provider invocation begins; - stateless targets replay it for every send. + Per-execution state for delivering persisted bootstrap history. + + An explicit seed is identified by persisted message-piece IDs instead of + inferred from response roles or error codes. A stateful cloned conversation + may instead have an empty seed plus copied branch messages to bootstrap its + new provider session. A context may be used by only one send at a time. + Stateful targets consume the bootstrap when provider invocation begins; + stateless targets replay their explicit seed for every send. """ conversation_id: str seed_message_ids: tuple[uuid.UUID, ...] replay_seed_each_send: bool + bootstrap_message_ids: tuple[uuid.UUID, ...] | None = None _seed_consumed: bool = field(default=False, init=False, repr=False) _send_in_progress: bool = field(default=False, init=False, repr=False) + _provider_attempt_marked: bool = field(default=False, init=False, repr=False) _provider_attempted_task: asyncio.Task[Any] | None = field(default=None, init=False, repr=False) _provider_attempt_count: int = field(default=0, init=False, repr=False) @@ -42,16 +46,28 @@ def __post_init__(self) -> None: """ if not self.conversation_id: raise ValueError("conversation_id must not be empty") - if not self.seed_message_ids: - raise ValueError("seed_message_ids must not be empty") + if not self.seed_message_ids and not self.bootstrap_message_ids: + raise ValueError("seed_message_ids must not be empty unless bootstrap_message_ids are provided") if len(set(self.seed_message_ids)) != len(self.seed_message_ids): raise ValueError("seed_message_ids must be unique") + if self.bootstrap_message_ids is not None: + if not self.bootstrap_message_ids: + raise ValueError("bootstrap_message_ids must not be empty") + if len(set(self.bootstrap_message_ids)) != len(self.bootstrap_message_ids): + raise ValueError("bootstrap_message_ids must be unique") + if not set(self.seed_message_ids).issubset(self.bootstrap_message_ids): + raise ValueError("bootstrap_message_ids must include every seed message") @property def seed_message_count(self) -> int: """Number of messages in the explicit seed prefix.""" return len(self.seed_message_ids) + @property + def bootstrap_message_count(self) -> int: + """Number of historical messages needed to bootstrap the next provider session.""" + return len(self.bootstrap_message_ids or self.seed_message_ids) + @property def should_include_seed(self) -> bool: """Whether the seed prefix should be included in the next send.""" @@ -89,6 +105,7 @@ def begin_send(self) -> None: "Wait for the active send to finish before sending another request." ) self._send_in_progress = True + self._provider_attempt_marked = False self._provider_attempted_task = None def mark_provider_attempted(self) -> None: @@ -103,10 +120,14 @@ def mark_provider_attempted(self) -> None: """ if not self._send_in_progress: raise RuntimeError("Cannot mark a provider attempt without an active send.") + if self._provider_attempt_marked: + return try: - self._provider_attempted_task = asyncio.current_task() + current_task = asyncio.current_task() except RuntimeError: - self._provider_attempted_task = None + current_task = None + self._provider_attempt_marked = True + self._provider_attempted_task = current_task self._provider_attempt_count += 1 if not self.replay_seed_each_send: self._seed_consumed = True @@ -119,30 +140,32 @@ def finish_send(self) -> None: def select_history(self, *, messages: list[Message]) -> list[Message]: """ - Select the explicit persisted seed prefix from conversation history. + Select history needed for delivery or provider-session restoration. Args: messages: Persisted conversation messages after failed exchanges have been removed. Returns: - The seed messages in their original persisted order, or an empty - list after a stateful target has consumed the seed. + The pending bootstrap messages before a stateful target consumes + them, the explicit seed for a stateless target, or all replayable + persisted history after a stateful target has been bootstrapped. Raises: ValueError: If an expected persisted seed message is missing. """ if not self.should_include_seed: - return [] + return list(messages) messages_by_id = {message.get_piece().id: message for message in messages} - missing_ids = [message_id for message_id in self.seed_message_ids if message_id not in messages_by_id] + selected_message_ids = self.bootstrap_message_ids or self.seed_message_ids + missing_ids = [message_id for message_id in selected_message_ids if message_id not in messages_by_id] if missing_ids: raise ValueError( "The persisted prepended history no longer matches the prepended history send context. " f"Missing {len(missing_ids)} message(s)." ) - return [messages_by_id[message_id] for message_id in self.seed_message_ids] + return [messages_by_id[message_id] for message_id in selected_message_ids] def remap_for_duplicate_conversation( self, @@ -154,8 +177,10 @@ def remap_for_duplicate_conversation( """ Remap this explicit seed boundary to a duplicated conversation. - Only message pieces already identified as seed history are remapped. Live - turns copied into the new memory conversation never become part of the boundary. + The explicit prepended seed identity is always remapped. A stateless clone + continues to replay only that seed. A stateful clone opens a new provider + session, so every replayable duplicated branch message becomes its one-time + bootstrap boundary. Args: conversation_id: Conversation ID assigned to the duplicated messages. @@ -195,16 +220,25 @@ def remap_for_duplicate_conversation( raise ValueError("Duplicated conversation does not preserve the source piece structure.") duplicated_ids_by_source_id[source_piece.id] = duplicated_piece.id - missing_ids = [ - message_id for message_id in self.seed_message_ids if message_id not in duplicated_ids_by_source_id - ] + current_bootstrap_ids = self.bootstrap_message_ids or self.seed_message_ids + ids_to_remap = set(self.seed_message_ids) | set(current_bootstrap_ids) + missing_ids = [message_id for message_id in ids_to_remap if message_id not in duplicated_ids_by_source_id] if missing_ids: - raise ValueError(f"Could not remap {len(missing_ids)} explicit seed message(s).") + raise ValueError(f"Could not remap {len(missing_ids)} bootstrap message(s).") + + remapped_seed_ids = tuple(duplicated_ids_by_source_id[message_id] for message_id in self.seed_message_ids) + if conversation_id != self.conversation_id and not self.replay_seed_each_send: + remapped_bootstrap_ids = tuple(message.get_piece().id for message in duplicated_messages) + else: + remapped_bootstrap_ids = tuple( + duplicated_ids_by_source_id[message_id] for message_id in current_bootstrap_ids + ) duplicated_context = PrependedHistorySendContext( conversation_id=conversation_id, - seed_message_ids=tuple(duplicated_ids_by_source_id[message_id] for message_id in self.seed_message_ids), + seed_message_ids=remapped_seed_ids, replay_seed_each_send=self.replay_seed_each_send, + bootstrap_message_ids=remapped_bootstrap_ids, ) # Provider bootstrap consumption belongs to the logical conversation, not copied memory. if conversation_id == self.conversation_id: diff --git a/pyrit/executor/attack/multi_turn/tree_of_attacks.py b/pyrit/executor/attack/multi_turn/tree_of_attacks.py index ea833770a2..8ad0bed536 100644 --- a/pyrit/executor/attack/multi_turn/tree_of_attacks.py +++ b/pyrit/executor/attack/multi_turn/tree_of_attacks.py @@ -32,6 +32,9 @@ build_conversation_context_string_async, ) from pyrit.executor.attack.component.modality_router import _ModalityFeedbackRouter +from pyrit.executor.attack.component.prepended_history_send_context import ( + PrependedHistorySendContext, +) from pyrit.executor.attack.core.attack_config import ( AttackAdversarialConfig, AttackConverterConfig, @@ -54,6 +57,7 @@ ) from pyrit.prompt_normalizer import ConverterConfiguration, PromptNormalizer from pyrit.prompt_target import CapabilityName, PromptTarget +from pyrit.prompt_target.common.target_history import filter_non_replayable_messages from pyrit.prompt_target.common.target_requirements import TargetRequirements from pyrit.score import ( FloatScaleThresholdScorer, @@ -72,9 +76,6 @@ from collections.abc import AsyncIterator from pathlib import Path - from pyrit.executor.attack.component.prepended_history_send_context import ( - PrependedHistorySendContext, - ) from pyrit.models.literals import PromptDataType logger = logging.getLogger(__name__) @@ -95,6 +96,42 @@ class TAPSystemPromptPaths(enum.Enum): ) +def _validate_stateful_clone_history_compatibility( + *, + objective_target: PromptTarget, + messages: list[Message], +) -> None: + """ + Reject converted history that cannot be replayed into a cloned provider session. + + Raises: + ValueError: If a stateful target cannot preserve converted media while cloning. + """ + if not objective_target.configuration.includes( + capability=CapabilityName.MULTI_TURN + ) or objective_target.configuration.includes(capability=CapabilityName.EDITABLE_HISTORY): + return + + non_text_output_types = { + piece.converted_value_data_type + for message in messages + for piece in message.message_pieces + if piece.converted_value_data_type != "text" + and ( + piece.original_value != piece.converted_value + or piece.original_value_data_type != piece.converted_value_data_type + ) + } + if non_text_output_types: + raise ValueError( + "Tree of Attacks cannot clone a stateful objective-target conversation without editable history " + "when persisted request or response history contains converted non-text output " + f"{sorted(non_text_output_types)}. Copied media cannot be flattened without changing converter " + "role scoping. Use an editable-history target, text-output converters, a stateless target, " + "or branching_factor=1." + ) + + class TAPAttackScoringConfig(AttackScoringConfig): """ Scoring configuration specifically for Tree of Attacks with Pruning (TAP). @@ -904,6 +941,13 @@ def duplicate(self) -> _TreeOfAttacksNode: of multiple attack variations from promising nodes. The tree expands by duplicating successful nodes and pruning unsuccessful ones. """ + source_messages = filter_non_replayable_messages( + messages=list(self._memory.get_conversation_messages(conversation_id=self.objective_target_conversation_id)) + ) + _validate_stateful_clone_history_compatibility( + objective_target=self._objective_target, + messages=source_messages, + ) duplicate_node = _TreeOfAttacksNode( objective_target=self._objective_target, adversarial_chat=self._adversarial_chat, @@ -926,24 +970,33 @@ def duplicate(self) -> _TreeOfAttacksNode: prepended_conversation_config=self._prepended_conversation_config, ) - source_messages = list( - self._memory.get_conversation_messages(conversation_id=self.objective_target_conversation_id) - ) duplicate_node.objective_target_conversation_id = self._memory.duplicate_conversation( conversation_id=self.objective_target_conversation_id ) - duplicated_messages = list( - self._memory.get_conversation_messages(conversation_id=duplicate_node.objective_target_conversation_id) + duplicated_messages = filter_non_replayable_messages( + messages=list( + self._memory.get_conversation_messages(conversation_id=duplicate_node.objective_target_conversation_id) + ) ) - duplicate_node._prepended_history_send_context = ( - self._prepended_history_send_context.remap_for_duplicate_conversation( + if self._prepended_history_send_context: + duplicate_node._prepended_history_send_context = ( + self._prepended_history_send_context.remap_for_duplicate_conversation( + conversation_id=duplicate_node.objective_target_conversation_id, + source_messages=source_messages, + duplicated_messages=duplicated_messages, + ) + ) + elif ( + duplicated_messages + and self._objective_target.configuration.includes(capability=CapabilityName.MULTI_TURN) + and not self._objective_target.configuration.includes(capability=CapabilityName.EDITABLE_HISTORY) + ): + duplicate_node._prepended_history_send_context = PrependedHistorySendContext( conversation_id=duplicate_node.objective_target_conversation_id, - source_messages=source_messages, - duplicated_messages=duplicated_messages, + seed_message_ids=(), + replay_seed_each_send=False, + bootstrap_message_ids=tuple(message.get_piece().id for message in duplicated_messages), ) - if self._prepended_history_send_context - else None - ) duplicate_node.adversarial_chat_conversation_id = self._memory.duplicate_conversation( conversation_id=self.adversarial_chat_conversation_id diff --git a/pyrit/executor/attack/streaming/barge_in.py b/pyrit/executor/attack/streaming/barge_in.py index f28be0a211..4e892bea41 100644 --- a/pyrit/executor/attack/streaming/barge_in.py +++ b/pyrit/executor/attack/streaming/barge_in.py @@ -138,6 +138,10 @@ async def _setup_async(self, *, context: BargeInAttackContext[Any]) -> None: """ if not context.conversation_id: context.conversation_id = str(uuid.uuid4()) + existing_message_ids = { + message.get_piece().id for message in self._conversation_manager.get_conversation(context.conversation_id) + } + context.prepended_history_send_context = None await self._conversation_manager.initialize_context_async( context=context, target=self._objective_target, @@ -145,6 +149,10 @@ async def _setup_async(self, *, context: BargeInAttackContext[Any]) -> None: request_converters=self._request_converters, prepended_conversation_config=self._prepended_conversation_config, ) + persisted_messages = self._conversation_manager.get_conversation(context.conversation_id) + context.prepended_conversation = [ + message for message in persisted_messages if message.get_piece().id not in existing_message_ids + ] context.prepended_history_send_context = None async def _teardown_async(self, *, context: BargeInAttackContext[Any]) -> None: diff --git a/pyrit/prompt_target/common/prompt_target.py b/pyrit/prompt_target/common/prompt_target.py index ca97638e83..f3abfd211b 100644 --- a/pyrit/prompt_target/common/prompt_target.py +++ b/pyrit/prompt_target/common/prompt_target.py @@ -24,7 +24,13 @@ ) from pyrit.prompt_target.common.target_configuration import TargetConfiguration from pyrit.prompt_target.common.target_history import filter_non_replayable_messages -from pyrit.prompt_target.common.target_send_context import TargetSendContext +from pyrit.prompt_target.common.target_send_context import ( + TargetSendContext, + _activate_target_send_context, + _mark_active_provider_attempted, + _reset_target_send_context, +) +from pyrit.prompt_target.common.utils import _marks_provider_attempt logger = logging.getLogger(__name__) @@ -62,6 +68,7 @@ class PromptTarget(Identifiable): # Per-instance overrides are also possible via the ``custom_configuration`` # constructor parameter, which takes precedence over the class-level value. _DEFAULT_CONFIGURATION: TargetConfiguration = TargetConfiguration(capabilities=TargetCapabilities()) + _MANAGES_PROVIDER_ATTEMPT_BOUNDARY: ClassVar[bool] = False # Declarative auth facts consumed by the create-target service and catalog. # Kept off ``TargetCapabilities`` (auth is a construction/credential axis, not @@ -88,6 +95,8 @@ def __init_subclass__(cls, **kwargs: object) -> None: after ``self``. """ super().__init_subclass__(**kwargs) + if "_send_prompt_to_target_async" in cls.__dict__ and "_MANAGES_PROVIDER_ATTEMPT_BOUNDARY" not in cls.__dict__: + cls._MANAGES_PROVIDER_ATTEMPT_BOUNDARY = False # Local import to avoid a circular dependency at package init time. from pyrit.common.brick_contract import enforce_keyword_only_init @@ -184,13 +193,28 @@ async def send_prompt_async( if not normalized_conversation: raise ValueError("Normalization pipeline returned an empty conversation. Cannot send an empty request.") self._validate_request(normalized_conversation=normalized_conversation) - if send_context: - send_context.mark_provider_attempted() - return await self._send_prompt_to_target_async(normalized_conversation=normalized_conversation) + active_context_token = ( + _activate_target_send_context(send_context=send_context) if send_context is not None else None + ) + try: + target_send = self._send_prompt_to_target_async + target_marks_provider_attempt = self._MANAGES_PROVIDER_ATTEMPT_BOUNDARY or _marks_provider_attempt( + target_send + ) + if send_context and not target_marks_provider_attempt: + send_context.mark_provider_attempted() + return await target_send(normalized_conversation=normalized_conversation) + finally: + if active_context_token is not None: + _reset_target_send_context(token=active_context_token) finally: if send_context: send_context.finish_send() + def _mark_provider_attempted(self) -> None: + """Notify caller-owned state immediately before irreversible provider I/O.""" + _mark_active_provider_attempted() + @abc.abstractmethod async def _send_prompt_to_target_async(self, *, normalized_conversation: list[Message]) -> list[Message]: """ diff --git a/pyrit/prompt_target/common/target_send_context.py b/pyrit/prompt_target/common/target_send_context.py index 02f66551e1..de589c3374 100644 --- a/pyrit/prompt_target/common/target_send_context.py +++ b/pyrit/prompt_target/common/target_send_context.py @@ -3,6 +3,7 @@ from __future__ import annotations +from contextvars import ContextVar, Token from typing import TYPE_CHECKING, Protocol if TYPE_CHECKING: @@ -34,3 +35,31 @@ def mark_provider_attempted(self) -> None: def finish_send(self) -> None: """Release caller-owned state after the send.""" ... + + +_ACTIVE_TARGET_SEND_CONTEXT: ContextVar[TargetSendContext | None] = ContextVar( + "_ACTIVE_TARGET_SEND_CONTEXT", + default=None, +) + + +def _activate_target_send_context(*, send_context: TargetSendContext) -> Token[TargetSendContext | None]: + """ + Expose one caller-owned context to target-side provider-boundary helpers. + + Returns: + A token that restores the previous task-local context. + """ + return _ACTIVE_TARGET_SEND_CONTEXT.set(send_context) + + +def _reset_target_send_context(*, token: Token[TargetSendContext | None]) -> None: + """Restore the task-local target-send context after target invocation.""" + _ACTIVE_TARGET_SEND_CONTEXT.reset(token) + + +def _mark_active_provider_attempted() -> None: + """Mark provider invocation for the active target send, if one exists.""" + send_context = _ACTIVE_TARGET_SEND_CONTEXT.get() + if send_context: + send_context.mark_provider_attempted() diff --git a/pyrit/prompt_target/common/utils.py b/pyrit/prompt_target/common/utils.py index a28b164252..25bf593903 100644 --- a/pyrit/prompt_target/common/utils.py +++ b/pyrit/prompt_target/common/utils.py @@ -4,6 +4,7 @@ import asyncio import logging from collections.abc import Callable +from functools import wraps from typing import Any from pyrit.exceptions import PyritException @@ -14,9 +15,17 @@ TokenUsage, construct_response_from_request, ) +from pyrit.prompt_target.common.target_send_context import _mark_active_provider_attempted logger = logging.getLogger(__name__) +_PROVIDER_ATTEMPT_MARKING_WRAPPERS: set[Callable[..., Any]] = set() + + +def _marks_provider_attempt(func: Callable[..., Any]) -> bool: + unbound_func = getattr(func, "__func__", func) + return unbound_func in _PROVIDER_ATTEMPT_MARKING_WRAPPERS + def validate_temperature(temperature: float | None) -> None: """ @@ -58,14 +67,18 @@ def limit_requests_per_minute(func: Callable[..., Any]) -> Callable[..., Any]: Callable: The decorated function with a sleep introduced. """ + @wraps(func) async def set_max_rpm_async(*args: Any, **kwargs: Any) -> Any: self = args[0] rpm = getattr(self, "_max_requests_per_minute", None) if rpm and rpm > 0: await asyncio.sleep(60 / rpm) + if not getattr(self, "_MANAGES_PROVIDER_ATTEMPT_BOUNDARY", False): + _mark_active_provider_attempted() return await func(*args, **kwargs) + _PROVIDER_ATTEMPT_MARKING_WRAPPERS.add(set_max_rpm_async) return set_max_rpm_async diff --git a/pyrit/prompt_target/http_target/httpx_api_target.py b/pyrit/prompt_target/http_target/httpx_api_target.py index 465ae0fa1f..02ba17a0f0 100644 --- a/pyrit/prompt_target/http_target/httpx_api_target.py +++ b/pyrit/prompt_target/http_target/httpx_api_target.py @@ -50,6 +50,7 @@ class HTTPXAPITarget(HTTPTarget): ), ) ) + _MANAGES_PROVIDER_ATTEMPT_BOUNDARY = True def __init__( self, @@ -166,6 +167,7 @@ async def _send_prompt_to_target_async(self, *, normalized_conversation: list[Me logger.info(f"HTTPXApiTarget: uploading file={filename} via {self.method} to {self.http_url}") + self._mark_provider_attempted() response = await client.request( method=self.method, url=self.http_url, @@ -177,6 +179,7 @@ async def _send_prompt_to_target_async(self, *, normalized_conversation: list[Me else: # No file upload, handle based on HTTP method logger.info(f"HTTPXApiTarget: sending {self.method} to {self.http_url} with possible JSON/form.") + self._mark_provider_attempted() response = await client.request( method=self.method, url=self.http_url, diff --git a/pyrit/prompt_target/playwright_copilot_target.py b/pyrit/prompt_target/playwright_copilot_target.py index a662e7f1a1..ce937868ed 100644 --- a/pyrit/prompt_target/playwright_copilot_target.py +++ b/pyrit/prompt_target/playwright_copilot_target.py @@ -94,6 +94,7 @@ class PlaywrightCopilotTarget(PromptTarget): ), ) ) + _MANAGES_PROVIDER_ATTEMPT_BOUNDARY = True # Placeholder text constants PLACEHOLDER_GENERATING_RESPONSE: str = "generating response" @@ -258,6 +259,8 @@ async def _interact_with_copilot_async(self, message: Message) -> str | list[tup either as a single text string or a list of (data, data_type) tuples. """ selectors = self._get_selectors() + if any(piece.converted_value_data_type == "text" for piece in message.message_pieces): + await self._clear_text_input_async(input_selector=selectors.input_selector) # Handle multimodal input - process all pieces in the request for piece in message.message_pieces: @@ -290,6 +293,7 @@ async def _wait_for_response_async(self, selectors: CopilotSelectors) -> str | l initial_group_count = len(initial_ai_message_groups) logger.debug(f"Initial message group count before sending: {initial_group_count}") + self._mark_provider_attempted() await self._page.click(selectors.send_button_selector) # Wait for the next AI message to appear @@ -797,6 +801,13 @@ async def _send_text_async(self, *, text: str, input_selector: str) -> None: await self._page.locator(input_selector).click() # Focus first await self._page.locator(input_selector).type(text) + async def _clear_text_input_async(self, *, input_selector: str) -> None: + """Clear locally staged text so a cancelled send can be retried safely.""" + input_locator = self._page.locator(input_selector) + await input_locator.click() + await input_locator.press("ControlOrMeta+A") + await input_locator.press("Backspace") + async def _upload_image_async(self, image_path: str) -> None: """ Handle image upload through Copilot's dropdown interface. @@ -817,6 +828,7 @@ async def _upload_image_async(self, image_path: str) -> None: async with self._page.expect_file_chooser() as fc_info: await add_files_button.click() file_chooser = await fc_info.value + self._mark_provider_attempted() await file_chooser.set_files(image_path) # Check for login requirement in Consumer Copilot diff --git a/pyrit/prompt_target/playwright_target.py b/pyrit/prompt_target/playwright_target.py index 6dcd5d378e..9d29316843 100644 --- a/pyrit/prompt_target/playwright_target.py +++ b/pyrit/prompt_target/playwright_target.py @@ -65,6 +65,7 @@ class PlaywrightTarget(PromptTarget): ), ) ) + _MANAGES_PROVIDER_ATTEMPT_BOUNDARY = True def __init__( self, @@ -118,6 +119,7 @@ async def _send_prompt_to_target_async(self, *, normalized_conversation: list[Me ) try: + self._mark_provider_attempted() text = await self._interaction_func(self._page, message) except Exception as e: raise RuntimeError(f"An error occurred during interaction: {str(e)}") from e diff --git a/pyrit/prompt_target/websocket_copilot_target.py b/pyrit/prompt_target/websocket_copilot_target.py index 1652e176d7..26ccb78da8 100644 --- a/pyrit/prompt_target/websocket_copilot_target.py +++ b/pyrit/prompt_target/websocket_copilot_target.py @@ -91,6 +91,7 @@ class WebSocketCopilotTarget(PromptTarget): ), ) ) + _MANAGES_PROVIDER_ATTEMPT_BOUNDARY = True def __init__( self, @@ -526,9 +527,10 @@ async def _connect_and_send_async( ) as websocket: for input_msg in inputs: payload = self._dict_to_websocket(input_msg) - await websocket.send(payload) - is_user_input = input_msg.get("type") == CopilotMessageType.USER_PROMPT + if is_user_input: + self._mark_provider_attempted() + await websocket.send(payload) max_message_iterations = 1000 iteration_count = 0 diff --git a/pyrit/prompt_target/websocket_target.py b/pyrit/prompt_target/websocket_target.py index b2ab2be695..ba4cc6cd97 100644 --- a/pyrit/prompt_target/websocket_target.py +++ b/pyrit/prompt_target/websocket_target.py @@ -38,6 +38,7 @@ class WebsocketTarget(PromptTarget): _DEFAULT_CONFIGURATION: TargetConfiguration = TargetConfiguration( capabilities=TargetCapabilities(supports_multi_turn=True) ) + _MANAGES_PROVIDER_ATTEMPT_BOUNDARY = True def __init__( self, @@ -247,6 +248,7 @@ async def _send_prompt_to_target_async(self, *, normalized_conversation: list[Me conversation_id=conversation_id, conversation_history=normalized_conversation[:-1], ) + self._mark_provider_attempted() result = await self._send_text_async( text=request.converted_value, conversation_id=conversation_id, diff --git a/tests/unit/executor/attack/component/test_prepended_history_send_context.py b/tests/unit/executor/attack/component/test_prepended_history_send_context.py index 5138645828..43561ce4e5 100644 --- a/tests/unit/executor/attack/component/test_prepended_history_send_context.py +++ b/tests/unit/executor/attack/component/test_prepended_history_send_context.py @@ -68,7 +68,7 @@ def test_stateful_context_consumes_explicit_boundary_at_provider_attempt() -> No assert context.provider_attempt_count == 1 context.begin_send() - assert context.select_history(messages=[first, second, unrelated]) == [] + assert context.select_history(messages=[first, second, unrelated]) == [first, second, unrelated] context.finish_send() @@ -117,6 +117,21 @@ def test_context_rejects_provider_attempt_without_active_send() -> None: context.mark_provider_attempted() +def test_context_counts_one_provider_attempt_per_send_task() -> None: + context = PrependedHistorySendContext( + conversation_id="conversation", + seed_message_ids=(uuid.uuid4(),), + replay_seed_each_send=False, + ) + + context.begin_send() + context.mark_provider_attempted() + context.mark_provider_attempted() + context.finish_send() + + assert context.provider_attempt_count == 1 + + def test_context_rejects_missing_persisted_boundary_message() -> None: context = PrependedHistorySendContext( conversation_id="conversation", @@ -150,10 +165,40 @@ def test_context_remaps_only_explicit_seed_for_duplicate_conversation() -> None: ) assert duplicate.seed_message_ids == (duplicated_messages[0].get_piece().id,) + assert duplicate.bootstrap_message_ids == (duplicated_messages[0].get_piece().id,) assert duplicated_messages[1].get_piece().id not in duplicate.seed_message_ids assert duplicate.select_history(messages=duplicated_messages) == [duplicated_messages[0]] +def test_stateful_context_bootstraps_full_duplicated_branch() -> None: + seed = _message(role="system", value="seed", sequence=0) + live_request = _message(role="user", value="live", sequence=1) + response = _message(role="assistant", value="response", sequence=2) + source_messages = [seed, live_request, response] + duplicated_messages = [ + _message(role=message.api_role, value=message.get_value(), sequence=message.sequence) + for message in source_messages + ] + for message in duplicated_messages: + message.get_piece().conversation_id = "duplicate" + + context = PrependedHistorySendContext( + conversation_id="source", + seed_message_ids=(seed.get_piece().id,), + replay_seed_each_send=False, + ) + + duplicate = context.remap_for_duplicate_conversation( + conversation_id="duplicate", + source_messages=source_messages, + duplicated_messages=duplicated_messages, + ) + + assert duplicate.seed_message_ids == (duplicated_messages[0].get_piece().id,) + assert duplicate.bootstrap_message_ids == tuple(message.get_piece().id for message in duplicated_messages) + assert duplicate.select_history(messages=duplicated_messages) == duplicated_messages + + def test_context_remap_resets_consumed_state_for_new_conversation() -> None: seed = _message(role="user", value="seed", sequence=0) duplicated_seed = seed.duplicate() diff --git a/tests/unit/executor/attack/multi_turn/test_prepended_history_normalization.py b/tests/unit/executor/attack/multi_turn/test_prepended_history_normalization.py index 727e735e30..067d3b4b6b 100644 --- a/tests/unit/executor/attack/multi_turn/test_prepended_history_normalization.py +++ b/tests/unit/executor/attack/multi_turn/test_prepended_history_normalization.py @@ -93,6 +93,14 @@ async def convert_async(self, *, prompt: str, input_type: PromptDataType = "text return ConverterResult(output_text=self._output_path, output_type="image_path") +class _TextOutputConverter(Converter): + SUPPORTED_INPUT_TYPES: tuple[PromptDataType, ...] = ("image_path",) + SUPPORTED_OUTPUT_TYPES: tuple[PromptDataType, ...] = ("text",) + + async def convert_async(self, *, prompt: str, input_type: PromptDataType = "text") -> ConverterResult: + return ConverterResult(output_text="converted text", output_type="text") + + def _make_context() -> MultiTurnAttackContext[AttackParameters]: return MultiTurnAttackContext(params=AttackParameters(objective="test objective")) @@ -379,7 +387,7 @@ async def format_messages(messages: list[Message]) -> str: @pytest.mark.usefixtures("patch_central_database") -async def test_tap_stateful_clone_replays_seed_once_for_new_conversation(): +async def test_tap_stateful_clone_bootstraps_duplicated_branch_once(): target = _ConversationKeyedRecordingTarget() formatter = MagicMock(spec=MessageStringNormalizer) @@ -414,7 +422,41 @@ async def format_messages(messages: list[Message]) -> str: assert target.prompts_by_conversation == { parent_conversation_id: ["original seed|parent first", "parent second"], - cloned_conversation_id: ["original seed|clone first", "clone second"], + cloned_conversation_id: ["original seed|parent first|response|clone first", "clone second"], + } + + +@pytest.mark.usefixtures("patch_central_database") +async def test_tap_unseeded_stateful_clone_bootstraps_duplicated_branch_once(): + target = _ConversationKeyedRecordingTarget() + formatter = MagicMock(spec=MessageStringNormalizer) + + async def format_messages(messages: list[Message]) -> str: + return "|".join(message.get_value() for message in messages) + + formatter.normalize_string_async = AsyncMock(side_effect=format_messages) + node = _make_tap_node(target=target) + node._prepended_conversation_config = PrependedConversationConfig(message_normalizer=formatter) + node._objective = "objective" + parent_conversation_id = node.objective_target_conversation_id + + await node._send_prompt_to_target_async("parent first") + assert node._prepended_history_send_context is None + + cloned = node.duplicate() + cloned._objective = "objective" + cloned_conversation_id = cloned.objective_target_conversation_id + assert cloned._prepended_history_send_context + assert cloned._prepended_history_send_context.seed_message_count == 0 + assert cloned._prepended_history_send_context.bootstrap_message_count == 2 + + await node._send_prompt_to_target_async("parent second") + await cloned._send_prompt_to_target_async("clone first") + await cloned._send_prompt_to_target_async("clone second") + + assert target.prompts_by_conversation == { + parent_conversation_id: ["parent first", "parent second"], + cloned_conversation_id: ["parent first|response|clone first", "clone second"], } @@ -577,6 +619,67 @@ async def format_messages(messages: list[Message]) -> str: ] +@pytest.mark.usefixtures("patch_central_database") +async def test_tap_stateful_clone_rejects_non_text_converter_history(tmp_path: Path): + image_path = tmp_path / "converted.png" + image_path.write_bytes(b"test image") + target = _RecordingTarget() + target._configuration = TargetConfiguration( + capabilities=TargetCapabilities( + supports_multi_turn=True, + supports_multi_message_pieces=True, + input_modalities=frozenset( + { + frozenset({"text"}), + frozenset({"image_path"}), + frozenset({"text", "image_path"}), + } + ), + ) + ) + node = _make_tap_node(target=target) + node._request_converters = ConverterConfiguration.from_converters( + converters=[_ImageOutputConverter(output_path=str(image_path))] + ) + node._objective = "objective" + + await node._send_prompt_to_target_async("depth one") + + with pytest.raises(ValueError, match="cannot clone.*non-text output.*image_path"): + node.duplicate() + + +@pytest.mark.usefixtures("patch_central_database") +async def test_tap_stateful_clone_accepts_converter_pipeline_with_final_text_output(tmp_path: Path): + image_path = tmp_path / "converted.png" + image_path.write_bytes(b"test image") + target = _ConversationKeyedRecordingTarget() + formatter = MagicMock(spec=MessageStringNormalizer) + + async def format_messages(messages: list[Message]) -> str: + return "|".join(message.get_value() for message in messages) + + formatter.normalize_string_async = AsyncMock(side_effect=format_messages) + node = _make_tap_node(target=target) + node._request_converters = ConverterConfiguration.from_converters( + converters=[ + _ImageOutputConverter(output_path=str(image_path)), + _TextOutputConverter(), + ] + ) + node._prepended_conversation_config = PrependedConversationConfig(message_normalizer=formatter) + node._objective = "objective" + + await node._send_prompt_to_target_async("depth one") + cloned = node.duplicate() + cloned._objective = "objective" + await cloned._send_prompt_to_target_async("depth two") + + assert cloned._prepended_history_send_context + assert cloned._prepended_history_send_context.is_seed_consumed + assert target.normalized_requests[-1].get_piece().converted_value == "converted text|response|converted text" + + @pytest.fixture def adversarial_config() -> AttackAdversarialConfig: target = MagicMock(spec=PromptTarget) diff --git a/tests/unit/executor/attack/streaming/test_barge_in.py b/tests/unit/executor/attack/streaming/test_barge_in.py index ba86ee5671..60f3a49e95 100644 --- a/tests/unit/executor/attack/streaming/test_barge_in.py +++ b/tests/unit/executor/attack/streaming/test_barge_in.py @@ -10,12 +10,15 @@ import pytest +from pyrit.converter import Base64Converter from pyrit.executor.attack import BargeInAttack, BargeInAttackContext +from pyrit.executor.attack.component import PrependedConversationConfig from pyrit.executor.attack.component.prepended_history_send_context import ( PrependedHistorySendContext, ) -from pyrit.executor.attack.core import AttackParameters +from pyrit.executor.attack.core import AttackConverterConfig, AttackParameters from pyrit.models import AttackOutcome, Message, MessagePiece +from pyrit.prompt_normalizer import ConverterConfiguration from pyrit.prompt_target import RealtimeTarget if TYPE_CHECKING: @@ -162,6 +165,59 @@ async def test_setup_async_persists_prepended_conversation_to_memory(vad_target) assert ctx.prepended_history_send_context is None +async def test_converted_system_prompt_is_passed_to_streaming_session(vad_target): + attack = BargeInAttack( + objective_target=vad_target, + attack_converter_config=AttackConverterConfig( + request_converters=ConverterConfiguration.from_converters(converters=[Base64Converter()]) + ), + prepended_conversation_config=PrependedConversationConfig(apply_converters_to_roles=["system"]), + ) + ctx = BargeInAttackContext( + params=AttackParameters( + objective="o", + prepended_conversation=[Message.from_prompt(prompt="You are strict.", role="system")], + ), + audio_chunks=_aiter([b"\x00" * 96]), + ) + await attack._setup_async(context=ctx) + fake_session = _fake_session() + + with patch.object(RealtimeTarget, "open_streaming_session", return_value=fake_session) as factory: + await attack._perform_async(context=ctx) + + persisted = attack._conversation_manager.get_conversation(ctx.conversation_id) + passed_to_session = factory.call_args.kwargs["prepended_conversation"] + assert passed_to_session == persisted + assert passed_to_session[0].get_piece().converted_value == "WW91IGFyZSBzdHJpY3Qu" + + +async def test_setup_reusing_conversation_passes_only_new_prepended_messages(vad_target): + attack = BargeInAttack(objective_target=vad_target) + conversation_id = "existing-conversation" + await attack._conversation_manager.add_prepended_conversation_to_memory_async( + prepended_conversation=[Message.from_prompt(prompt="old system", role="system")], + conversation_id=conversation_id, + target=vad_target, + ) + ctx = BargeInAttackContext( + params=AttackParameters( + objective="o", + prepended_conversation=[Message.from_prompt(prompt="new system", role="system")], + ), + audio_chunks=_aiter([b"\x00" * 96]), + conversation_id=conversation_id, + ) + + await attack._setup_async(context=ctx) + + assert [message.get_value() for message in ctx.prepended_conversation] == ["new system"] + assert [message.get_value() for message in attack._conversation_manager.get_conversation(conversation_id)] == [ + "old system", + "new system", + ] + + async def test_setup_async_clears_unused_normalization_context_when_prepended_empty(vad_target): """The direct streaming path does not retain target normalization state.""" attack = BargeInAttack(objective_target=vad_target) diff --git a/tests/unit/prompt_target/target/test_normalize_async_integration.py b/tests/unit/prompt_target/target/test_normalize_async_integration.py index 29b14b7137..f6a61e1882 100644 --- a/tests/unit/prompt_target/target/test_normalize_async_integration.py +++ b/tests/unit/prompt_target/target/test_normalize_async_integration.py @@ -827,7 +827,7 @@ async def test_prepended_history_adapter_is_used_only_when_explicitly_passed(): @pytest.mark.usefixtures("patch_central_database") -async def test_non_editable_multi_turn_target_sends_only_current_request_after_response(): +async def test_non_editable_multi_turn_target_retains_history_after_response(): target = MockPromptTarget() target._configuration = TargetConfiguration( capabilities=TargetCapabilities( @@ -866,8 +866,12 @@ async def test_non_editable_multi_turn_target_sends_only_current_request_after_r first_payload, second_payload = target._send_prompt_to_target_async.await_args_list assert len(first_payload.kwargs["normalized_conversation"]) == 1 - assert len(second_payload.kwargs["normalized_conversation"]) == 1 - assert second_payload.kwargs["normalized_conversation"][0].get_value() == "second live" + assert [message.get_value() for message in second_payload.kwargs["normalized_conversation"]] == [ + "prepended", + "first live", + "first response", + "second live", + ] @pytest.mark.usefixtures("patch_central_database") @@ -940,7 +944,7 @@ async def test_stateful_target_consumes_seed_after_provider_outcome(response_err first_payload, second_payload = target._send_prompt_to_target_async.await_args_list assert "prepended" in first_payload.kwargs["normalized_conversation"][0].get_value() - assert [message.get_value() for message in second_payload.kwargs["normalized_conversation"]] == ["second live"] + assert second_payload.kwargs["normalized_conversation"][-1].get_value() == "second live" @pytest.mark.usefixtures("patch_central_database") @@ -1246,6 +1250,42 @@ async def test_target_normalization_cancellation_propagates(): assert not target_context.is_seed_consumed +@pytest.mark.usefixtures("patch_central_database") +async def test_rate_limit_cancellation_does_not_consume_stateful_seed(): + target = MockPromptTarget(rpm=1) + target._configuration = TargetConfiguration(capabilities=TargetCapabilities(supports_multi_turn=True)) + mock_memory = MagicMock(spec=MemoryInterface) + prepended = _make_message(role="user", content="prepended") + mock_memory.get_conversation_messages.return_value = [prepended] + target._memory = mock_memory + target_context = _make_prepended_history_send_context( + prepended_messages=[prepended], + target_supports_multi_turn=True, + ) + sleep_started = asyncio.Event() + sleep_release = asyncio.Event() + + async def wait_for_rate_limit(delay: float) -> None: + sleep_started.set() + await sleep_release.wait() + + with patch("pyrit.prompt_target.common.utils.asyncio.sleep", side_effect=wait_for_rate_limit): + send_task = asyncio.create_task( + target.send_prompt_async( + message=_make_message(role="user", content="live"), + normalizer_overrides=_make_normalizer_overrides(send_context=target_context), + send_context=target_context, + ) + ) + await sleep_started.wait() + send_task.cancel() + with pytest.raises(asyncio.CancelledError): + await send_task + + assert target_context.provider_attempt_count == 0 + assert not target_context.is_seed_consumed + + @pytest.mark.usefixtures("patch_central_database") async def test_provider_failure_propagates_after_normalization(): target = MockPromptTarget() @@ -1321,7 +1361,7 @@ async def wait_in_provider(*, normalized_conversation: list[Message]) -> list[Me send_context=target_context, ) payload = target._send_prompt_to_target_async.await_args.kwargs["normalized_conversation"] - assert [message.get_value() for message in payload] == ["second live"] + assert [message.get_value() for message in payload] == ["prepended", "first live", "second live"] @pytest.mark.usefixtures("patch_central_database") diff --git a/tests/unit/prompt_target/target/test_playwright_copilot_target.py b/tests/unit/prompt_target/target/test_playwright_copilot_target.py index f933f4cd13..5e22f125ad 100644 --- a/tests/unit/prompt_target/target/test_playwright_copilot_target.py +++ b/tests/unit/prompt_target/target/test_playwright_copilot_target.py @@ -1,6 +1,7 @@ # Copyright (c) Microsoft Corporation. # Licensed under the MIT license. +import asyncio from unittest.mock import AsyncMock, MagicMock, call, patch import pytest @@ -296,11 +297,13 @@ async def __aexit__(self, *args): mock_page.expect_file_chooser = MagicMock(return_value=MockFileChooserContextManager(mock_file_chooser)) - await target._upload_image_async("/path/to/image.jpg") + with patch.object(target, "_mark_provider_attempted") as mark_provider_attempted: + await target._upload_image_async("/path/to/image.jpg") dropdown_locator.click.assert_awaited_once() file_picker_locator.wait_for.assert_awaited_once_with(state="visible", timeout=5000) file_picker_locator.click.assert_awaited_once() + mark_provider_attempted.assert_called_once_with() mock_file_chooser.set_files.assert_awaited_once_with("/path/to/image.jpg") async def test_wait_for_response_async_success(self, mock_page): @@ -314,10 +317,14 @@ async def test_wait_for_response_async_success(self, mock_page): mock_page.query_selector_all.return_value = [AsyncMock()] # Mock response extraction - with patch.object(target, "_extract_multimodal_content_async", return_value="Response text") as mock_extract: + with ( + patch.object(target, "_extract_multimodal_content_async", return_value="Response text") as mock_extract, + patch.object(target, "_mark_provider_attempted") as mark_provider_attempted, + ): result = await target._wait_for_response_async(selectors) assert result == "Response text" + mark_provider_attempted.assert_called_once_with() mock_page.click.assert_awaited_once_with(selectors.send_button_selector) mock_extract.assert_awaited_once() @@ -372,6 +379,7 @@ async def test_interact_with_copilot_async_multimodal(self, mock_page, multimoda # Mock the helper methods with ( + patch.object(target, "_clear_text_input_async") as mock_clear_text, patch.object(target, "_send_text_async") as mock_send_text, patch.object(target, "_upload_image_async") as mock_upload_image, patch.object(target, "_wait_for_response_async", return_value="AI response") as mock_wait, @@ -379,11 +387,36 @@ async def test_interact_with_copilot_async_multimodal(self, mock_page, multimoda result = await target._interact_with_copilot_async(multimodal_request) # Verify text and image handling + mock_clear_text.assert_awaited_once() mock_send_text.assert_awaited_once() mock_upload_image.assert_awaited_once_with("/path/to/image.jpg") mock_wait.assert_awaited_once() assert result == "AI response" + async def test_interact_cancellation_while_staging_text_does_not_mark_provider_attempt( + self, mock_page, text_request_piece + ): + target = PlaywrightCopilotTarget(page=mock_page) + request = Message(message_pieces=[text_request_piece]) + staging_started = asyncio.Event() + + async def stage_text(*, text: str, input_selector: str) -> None: + staging_started.set() + await asyncio.Event().wait() + + with ( + patch.object(target, "_clear_text_input_async"), + patch.object(target, "_send_text_async", side_effect=stage_text), + patch.object(target, "_mark_provider_attempted") as mark_provider_attempted, + ): + task = asyncio.create_task(target._interact_with_copilot_async(request)) + await staging_started.wait() + task.cancel() + with pytest.raises(asyncio.CancelledError): + await task + + mark_provider_attempted.assert_not_called() + def test_constants(self, mock_page): """Test that class constants are defined correctly.""" target = PlaywrightCopilotTarget(page=mock_page) diff --git a/tests/unit/prompt_target/target/test_websocket_target.py b/tests/unit/prompt_target/target/test_websocket_target.py index 281dc5d5cd..5eac8e842a 100644 --- a/tests/unit/prompt_target/target/test_websocket_target.py +++ b/tests/unit/prompt_target/target/test_websocket_target.py @@ -13,11 +13,26 @@ from websockets.protocol import State from pyrit.exceptions import EmptyResponseException +from pyrit.executor.attack.component.prepended_history_send_context import ( + PrependedHistorySendContext, +) from pyrit.memory import SQLiteMemory from pyrit.models import Message, MessagePiece from pyrit.prompt_target import WebsocketTarget +class _OverriddenWebsocketTarget(WebsocketTarget): + async def _send_prompt_to_target_async(self, *, normalized_conversation: list[Message]) -> list[Message]: + request = normalized_conversation[-1].get_piece() + return [ + MessagePiece( + role="assistant", + original_value="response", + conversation_id=request.conversation_id, + ).to_message() + ] + + @pytest.fixture def response_parser() -> Callable[[str | bytes], str | None]: def parse_response(message: str | bytes) -> str | None: @@ -78,6 +93,10 @@ def test_init_invalid_endpoint_raises( ) +def test_overridden_managed_target_uses_compatibility_boundary() -> None: + assert not _OverriddenWebsocketTarget._MANAGES_PROVIDER_ATTEMPT_BOUNDARY + + def test_init_empty_protocol_identifier_raises( response_parser: Callable[[str | bytes], str | None], message_builder: Callable[[str], str | bytes], @@ -249,6 +268,45 @@ async def test_send_prompt_async_failure_discards_connection(websocket_target: W assert "conversation" not in websocket_target._existing_conversation +async def test_cancellation_before_websocket_send_does_not_consume_seed( + websocket_target: WebsocketTarget, + sqlite_instance: SQLiteMemory, +) -> None: + seed = create_message(value="Seed") + sqlite_instance.add_message_to_memory(request=seed) + send_context = PrependedHistorySendContext( + conversation_id="conversation", + seed_message_ids=(seed.get_piece().id,), + replay_seed_each_send=False, + ) + connection_started = asyncio.Event() + connection_release = asyncio.Event() + + async def wait_for_connection( + *, + conversation_id: str, + conversation_history: list[Message], + ) -> ClientConnection: + connection_started.set() + await connection_release.wait() + return AsyncMock(spec=ClientConnection) + + with patch.object(websocket_target, "_get_or_create_connection_async", side_effect=wait_for_connection): + send_task = asyncio.create_task( + websocket_target.send_prompt_async( + message=create_message(value="Current"), + send_context=send_context, + ) + ) + await connection_started.wait() + send_task.cancel() + with pytest.raises(asyncio.CancelledError): + await send_task + + assert send_context.provider_attempt_count == 0 + assert not send_context.is_seed_consumed + + async def test_get_or_create_connection_async_restores_history_on_reconnect( response_parser: Callable[[str | bytes], str | None], message_builder: Callable[[str], str | bytes], @@ -289,6 +347,75 @@ async def test_get_or_create_connection_async_restores_history_on_reconnect( assert target._existing_conversation == {"conversation": replacement_connection} +async def test_consumed_context_retains_history_for_reconnect( + response_parser: Callable[[str | bytes], str | None], + message_builder: Callable[[str], str | bytes], + sqlite_instance: SQLiteMemory, +) -> None: + stale_connection = AsyncMock(spec=ClientConnection) + stale_connection.state = State.CLOSED + replacement_connection = AsyncMock(spec=ClientConnection) + replacement_connection.state = State.OPEN + restore_callback = AsyncMock() + target = WebsocketTarget( + endpoint="wss://example.com", + protocol_identifier="test-protocol", + initialization_strings=[], + response_parser=response_parser, + message_builder=message_builder, + conversation_restore_callback=restore_callback, + discard_initial_messages=0, + existing_convo={"conversation": stale_connection}, + ) + seed = create_message(value="Seed") + first_request = create_message(value="First request") + first_response = MessagePiece( + role="assistant", + original_value="First response", + converted_value="First response", + conversation_id="conversation", + ).to_message() + for sequence, message in enumerate([seed, first_request, first_response]): + for piece in message.message_pieces: + piece.sequence = sequence + sqlite_instance.add_message_to_memory(request=message) + + send_context = PrependedHistorySendContext( + conversation_id="conversation", + seed_message_ids=(seed.get_piece().id,), + replay_seed_each_send=False, + ) + send_context.begin_send() + send_context.mark_provider_attempted() + send_context.finish_send() + + with ( + patch.object( + target, + "_connect_async", + new_callable=AsyncMock, + return_value=replacement_connection, + ), + patch.object( + target, + "_send_text_async", + new_callable=AsyncMock, + return_value="Second response", + ), + ): + await target.send_prompt_async( + message=create_message(value="Second request"), + send_context=send_context, + ) + + restored_history = restore_callback.await_args.args[1] + assert [message.get_value() for message in restored_history] == [ + "Seed", + "First request", + "First response", + ] + + async def test_get_or_create_connection_async_fails_when_history_cannot_be_restored( websocket_target: WebsocketTarget, ) -> None: From 217a1f59e5655fb5727a56d237aaea55a03e0285 Mon Sep 17 00:00:00 2001 From: Roman Lutz Date: Sat, 22 Aug 2026 06:55:34 -0700 Subject: [PATCH 24/24] Harden prepended conversation role handling Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> Copilot-Session: 6776aa95-0924-4967-a056-e5a228af3faa --- doc/code/executor/3_attack_configuration.ipynb | 7 ++++--- doc/code/executor/3_attack_configuration.py | 7 ++++--- doc/code/targets/0_prompt_targets.md | 3 +++ .../executor/attack/component/conversation_manager.py | 1 + .../attack/component/prepended_conversation_config.py | 10 ++++++++-- pyrit/executor/attack/streaming/barge_in.py | 1 - .../attack/component/test_conversation_manager.py | 9 +++++++++ .../component/test_prepended_conversation_config.py | 5 +++++ 8 files changed, 34 insertions(+), 9 deletions(-) diff --git a/doc/code/executor/3_attack_configuration.ipynb b/doc/code/executor/3_attack_configuration.ipynb index 0602296fdf..4366b4f6d7 100644 --- a/doc/code/executor/3_attack_configuration.ipynb +++ b/doc/code/executor/3_attack_configuration.ipynb @@ -472,9 +472,10 @@ "only applies to attacks that drive a conversation. Below builds a converter config — it's just a\n", "plain object you hand to the attack constructor.\n", "\n", - "Request converters apply to prepended `user` messages by default. Prepended `assistant` messages\n", - "represent simulated target output, so PyRIT leaves them unchanged unless the attack explicitly\n", - "opts in. For example:\n", + "Request converters apply only to prepended `user` messages by default. PyRIT leaves every other\n", + "role, including `system`, `developer`, `tool`, and `assistant` / `simulated_assistant`, unchanged\n", + "unless the attack explicitly opts in. `simulated_assistant` is accepted as an alias for\n", + "`assistant`. For example:\n", "\n", "```python\n", "from pyrit.executor.attack import PrependedConversationConfig\n", diff --git a/doc/code/executor/3_attack_configuration.py b/doc/code/executor/3_attack_configuration.py index 74e2cd6bdd..6c47ed2f29 100644 --- a/doc/code/executor/3_attack_configuration.py +++ b/doc/code/executor/3_attack_configuration.py @@ -176,9 +176,10 @@ # only applies to attacks that drive a conversation. Below builds a converter config — it's just a # plain object you hand to the attack constructor. # -# Request converters apply to prepended `user` messages by default. Prepended `assistant` messages -# represent simulated target output, so PyRIT leaves them unchanged unless the attack explicitly -# opts in. For example: +# Request converters apply only to prepended `user` messages by default. PyRIT leaves every other +# role, including `system`, `developer`, `tool`, and `assistant` / `simulated_assistant`, unchanged +# unless the attack explicitly opts in. `simulated_assistant` is accepted as an alias for +# `assistant`. For example: # # ```python # from pyrit.executor.attack import PrependedConversationConfig diff --git a/doc/code/targets/0_prompt_targets.md b/doc/code/targets/0_prompt_targets.md index 1da66a6bfd..4a4d3f17b2 100644 --- a/doc/code/targets/0_prompt_targets.md +++ b/doc/code/targets/0_prompt_targets.md @@ -34,6 +34,9 @@ converters have already run by this point, so role-specific converter choices re the target must receive one flattened request. The context is ephemeral and does not replace the structured messages stored in memory. +Prepended request converters apply only to `user` messages by default. Every other role, including +`system`, `developer`, `tool`, and `assistant` / `simulated_assistant`, requires explicit opt-in. + For a stateful target without editable history, the initial bootstrap is flattened once. Later sends retain replayable memory history in the normalized view so a target such as `WebsocketTarget` can restore a replaced provider session, while the existing provider session still receives only the diff --git a/pyrit/executor/attack/component/conversation_manager.py b/pyrit/executor/attack/component/conversation_manager.py index 893da6c752..40ac2fe3e3 100644 --- a/pyrit/executor/attack/component/conversation_manager.py +++ b/pyrit/executor/attack/component/conversation_manager.py @@ -307,6 +307,7 @@ async def initialize_context_async( context.memory_labels = combine_dict(existing_dict=memory_labels, new_dict=context.memory_labels) state = ConversationState() prepended_conversation = context.prepended_conversation + context.prepended_history_send_context = None if not prepended_conversation: logger.debug(f"No prepended conversation for context initialization: {conversation_id}") diff --git a/pyrit/executor/attack/component/prepended_conversation_config.py b/pyrit/executor/attack/component/prepended_conversation_config.py index 3ec3e35b61..30c40d1a9e 100644 --- a/pyrit/executor/attack/component/prepended_conversation_config.py +++ b/pyrit/executor/attack/component/prepended_conversation_config.py @@ -39,8 +39,8 @@ class PrependedConversationConfig: image, audio, or other non-text output. """ - # Request converters default to prepended user messages only. Assistant history is - # simulated target output and must be explicitly opted in with ["assistant"]. + # Request converters default to prepended user messages only. Every non-user role, + # including system and simulated assistant history, requires explicit opt-in. apply_converters_to_roles: list[ChatMessageRole] = field(default_factory=lambda: ["user"]) # Optional normalizer to format prepended history and a live request as one text block. @@ -49,6 +49,12 @@ class PrependedConversationConfig: # that produces "Turn N: User/Assistant" format. message_normalizer: MessageStringNormalizer | None = None + def __post_init__(self) -> None: + """Normalize simulated assistant opt-in to its API-compatible role.""" + self.apply_converters_to_roles = [ + "assistant" if role == "simulated_assistant" else role for role in self.apply_converters_to_roles + ] + def get_message_normalizer(self) -> MessageStringNormalizer: """ Get the normalizer for objective target context, with a default fallback. diff --git a/pyrit/executor/attack/streaming/barge_in.py b/pyrit/executor/attack/streaming/barge_in.py index 4e892bea41..947b2c3221 100644 --- a/pyrit/executor/attack/streaming/barge_in.py +++ b/pyrit/executor/attack/streaming/barge_in.py @@ -141,7 +141,6 @@ async def _setup_async(self, *, context: BargeInAttackContext[Any]) -> None: existing_message_ids = { message.get_piece().id for message in self._conversation_manager.get_conversation(context.conversation_id) } - context.prepended_history_send_context = None await self._conversation_manager.initialize_context_async( context=context, target=self._objective_target, diff --git a/tests/unit/executor/attack/component/test_conversation_manager.py b/tests/unit/executor/attack/component/test_conversation_manager.py index 72fe89e431..61d5e7dd40 100644 --- a/tests/unit/executor/attack/component/test_conversation_manager.py +++ b/tests/unit/executor/attack/component/test_conversation_manager.py @@ -33,6 +33,9 @@ get_prepended_turn_count, mark_messages_as_simulated, ) +from pyrit.executor.attack.component.prepended_history_send_context import ( + PrependedHistorySendContext, +) from pyrit.executor.attack.core import AttackContext from pyrit.executor.attack.core.attack_parameters import AttackParameters from pyrit.message_normalizer import ConversationContextNormalizer, HistorySquashNormalizer @@ -653,6 +656,11 @@ async def test_returns_default_state_for_no_prepended_conversation( """Test that no prepended conversation returns default state.""" manager = ConversationManager() conversation_id = str(uuid.uuid4()) + mock_attack_context.prepended_history_send_context = PrependedHistorySendContext( + conversation_id="stale-conversation", + seed_message_ids=(uuid.uuid4(),), + replay_seed_each_send=False, + ) state = await manager.initialize_context_async( context=mock_attack_context, @@ -663,6 +671,7 @@ async def test_returns_default_state_for_no_prepended_conversation( assert isinstance(state, ConversationState) assert state.turn_count == 0 assert state.last_assistant_message_scores == [] + assert mock_attack_context.prepended_history_send_context is None async def test_merges_memory_labels( self, diff --git a/tests/unit/executor/attack/component/test_prepended_conversation_config.py b/tests/unit/executor/attack/component/test_prepended_conversation_config.py index 546a7f1dbf..5968102b35 100644 --- a/tests/unit/executor/attack/component/test_prepended_conversation_config.py +++ b/tests/unit/executor/attack/component/test_prepended_conversation_config.py @@ -12,6 +12,11 @@ def test_default_init_apply_converters_to_user_role(): assert config.apply_converters_to_roles == ["user"] +def test_simulated_assistant_converter_role_normalizes_to_assistant(): + config = PrependedConversationConfig(apply_converters_to_roles=["simulated_assistant"]) + assert config.apply_converters_to_roles == ["assistant"] + + def test_default_init_message_normalizer_is_none(): config = PrependedConversationConfig() assert config.message_normalizer is None