diff --git a/pyrit/executor/attack/multi_turn/tree_of_attacks.py b/pyrit/executor/attack/multi_turn/tree_of_attacks.py index ec1ab0f8e3..639f0e2427 100644 --- a/pyrit/executor/attack/multi_turn/tree_of_attacks.py +++ b/pyrit/executor/attack/multi_turn/tree_of_attacks.py @@ -157,6 +157,35 @@ def threshold(self) -> float: return self.objective_scorer.threshold # type: ignore[ty:unresolved-attribute] +@dataclass(frozen=True, slots=True) +class _TAPAttackConfiguration: + """Immutable configuration for the TAP search and node execution.""" + + tree_width: int + tree_depth: int + branching_factor: int + on_topic_checking_enabled: bool + desired_response_prefix: str + batch_size: int + + def __post_init__(self) -> None: + """ + Validate the TAP search limits. + + Raises: + ValueError: If a search limit is less than one. + """ + validations = ( + (self.tree_depth, "The tree depth must be at least 1."), + (self.tree_width, "The tree width must be at least 1."), + (self.branching_factor, "The branching factor must be at least 1."), + (self.batch_size, "The batch size must be at least 1."), + ) + for value, message in validations: + if value < 1: + raise ValueError(message) + + @dataclass class TAPAttackContext(MultiTurnAttackContext[Any]): """ @@ -179,6 +208,9 @@ class TAPAttackContext(MultiTurnAttackContext[Any]): best_objective_score: Score | None = None best_adversarial_conversation_id: str | None = None + # Visualization parent for first-level nodes in this execution + visualization_root_id: str = "root" + @property def conversation_id(self) -> str | None: """The best objective-target conversation, or the first active branch.""" @@ -1357,32 +1389,20 @@ def __init__( ``score_blocked_content=True`` on the objective scorer (requires ``prompt_metadata["partial_content"]`` on the blocked piece). """ - # Validate tree parameters - if tree_depth < 1: - raise ValueError("The tree depth must be at least 1.") - if tree_width < 1: - raise ValueError("The tree width must be at least 1.") - if branching_factor < 1: - raise ValueError("The branching factor must be at least 1.") - if batch_size < 1: - raise ValueError("The batch size must be at least 1.") + self._configuration = _TAPAttackConfiguration( + tree_width=tree_width, + tree_depth=tree_depth, + branching_factor=branching_factor, + on_topic_checking_enabled=on_topic_checking_enabled, + desired_response_prefix=desired_response_prefix, + batch_size=batch_size, + ) # Initialize base class super().__init__(objective_target=objective_target, logger=logger, context_type=TAPAttackContext) self._memory = CentralMemory.get_memory_instance() - # Store tree configuration - self._tree_width = tree_width - self._tree_depth = tree_depth - self._branching_factor = branching_factor - self._vis_root_id = "root" - - # Store execution configuration - self._on_topic_checking_enabled = on_topic_checking_enabled - self._desired_response_prefix = desired_response_prefix - self._batch_size = batch_size - # Initialize adversarial configuration self._adversarial_chat = attack_adversarial_config.target @@ -1482,7 +1502,7 @@ def __init__( # Use the adversarial chat target for scoring, as in CrescendoAttack self._scoring_target = self._adversarial_chat - if self._on_topic_checking_enabled and not self._scoring_target: + if self._configuration.on_topic_checking_enabled and not self._scoring_target: raise ValueError("On-topic checking is enabled but no scoring target is available.") self._prompt_normalizer = prompt_normalizer or PromptNormalizer() @@ -1573,6 +1593,7 @@ async def _setup_async(self, *, context: TAPAttackContext) -> None: context.tree_visualization = Tree() context.tree_visualization.create_node("Root", "root") + context.visualization_root_id = "root" context.nodes = [] context.best_conversation_id = None @@ -1586,10 +1607,10 @@ async def _setup_async(self, *, context: TAPAttackContext) -> None: context.executed_turns = get_prepended_turn_count(context.prepended_conversation) # Validate that prepended conversation doesn't exceed tree_depth - if context.executed_turns >= self._tree_depth: + if context.executed_turns >= self._configuration.tree_depth: raise ValueError( f"Prepended conversation has {context.executed_turns} turns, " - f"which equals or exceeds tree_depth={self._tree_depth}. " + f"which equals or exceeds tree_depth={self._configuration.tree_depth}. " f"Reduce prepended turns or increase tree_depth." ) @@ -1601,7 +1622,7 @@ async def _setup_async(self, *, context: TAPAttackContext) -> None: node_id = f"prepended_{turn}" context.tree_visualization.create_node(f"{turn}: (prepended)", node_id, parent=vis_parent) vis_parent = node_id - self._vis_root_id = vis_parent + context.visualization_root_id = vis_parent async def _perform_async(self, *, context: TAPAttackContext) -> TAPAttackResult: """ @@ -1629,11 +1650,13 @@ async def _perform_async(self, *, context: TAPAttackContext) -> TAPAttackResult: """ self._logger.info(f"Starting TAP attack with objective: {context.objective}") self._logger.info( - f"Tree dimensions - Width: {self._tree_width}, Depth: {self._tree_depth}, " - f"Branching factor: {self._branching_factor}" + f"Tree dimensions - Width: {self._configuration.tree_width}, " + f"Depth: {self._configuration.tree_depth}, " + f"Branching factor: {self._configuration.branching_factor}" ) self._logger.info( - f"Execution settings - Batch size: {self._batch_size}, On-topic checking: {self._on_topic_checking_enabled}" + f"Execution settings - Batch size: {self._configuration.batch_size}, " + f"On-topic checking: {self._configuration.on_topic_checking_enabled}" ) # TAP Attack Execution Algorithm: @@ -1656,9 +1679,9 @@ async def _perform_async(self, *, context: TAPAttackContext) -> TAPAttackResult: # Execute tree exploration iterations # Note: executed_turns is initialized in _setup_async with prepended conversation count # Start from executed_turns + 1 so prepended turns count toward tree_depth - for turn in range(context.executed_turns + 1, self._tree_depth + 1): + for turn in range(context.executed_turns + 1, self._configuration.tree_depth + 1): context.executed_turns = turn - self._logger.info(f"Starting TAP turn {turn}/{self._tree_depth}") + self._logger.info(f"Starting TAP turn {turn}/{self._configuration.tree_depth}") # Prepare nodes for current iteration await self._prepare_nodes_for_iteration_async(context) @@ -1787,7 +1810,7 @@ async def _initialize_first_level_nodes_async(self, context: TAPAttackContext) - context.next_message is not None and self._modality_router.objective_target_requires_media_on_first_turn ) - for i in range(self._tree_width): + for i in range(self._configuration.tree_width): # Historically only node 0 consumed next_message so sibling roots could # explore alternative starts. For edit-only objectives (no {text} path), # every root must receive seed media to build a valid first-turn request. @@ -1804,7 +1827,7 @@ async def _initialize_first_level_nodes_async(self, context: TAPAttackContext) - ) context.nodes.append(node) - node._vis_node_id = self._vis_root_id + node._vis_node_id = context.visualization_root_id # Clear next_message after initialization (it's been used by the first node) context.next_message = None @@ -1824,7 +1847,7 @@ def _branch_existing_nodes(self, context: TAPAttackContext) -> None: cloned_nodes = [] for node in context.nodes: - for _ in range(self._branching_factor - 1): + 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 context.related_conversations.add( @@ -1863,12 +1886,12 @@ async def _send_prompts_to_all_nodes_async(self, context: TAPAttackContext) -> N node._vis_node_id = vis_id # Process nodes in batches - for batch_start in range(0, len(context.nodes), self._batch_size): - batch_end = min(batch_start + self._batch_size, len(context.nodes)) + for batch_start in range(0, len(context.nodes), self._configuration.batch_size): + batch_end = min(batch_start + self._configuration.batch_size, len(context.nodes)) batch_nodes = context.nodes[batch_start:batch_end] self._logger.debug( - f"Processing batch {batch_start // self._batch_size + 1} " + f"Processing batch {batch_start // self._configuration.batch_size + 1} " f"(nodes {batch_start + 1}-{batch_end} of {len(context.nodes)})" ) @@ -1918,8 +1941,8 @@ def _prune_nodes_to_maintain_width(self, context: TAPAttackContext) -> None: completed_nodes = self._get_completed_nodes_sorted_by_score(context.nodes) # Keep nodes up to width limit - nodes_to_keep = completed_nodes[: self._tree_width] - nodes_to_prune = completed_nodes[self._tree_width :] + nodes_to_keep = completed_nodes[: self._configuration.tree_width] + nodes_to_prune = completed_nodes[self._configuration.tree_width :] # Mark pruned nodes in visualization and track their conversation IDs for node in nodes_to_prune: @@ -2012,7 +2035,7 @@ def _create_attack_node( attack_strategy_name=self.__class__.__name__, modality_router=self._modality_router, memory_labels=context.memory_labels, - desired_response_prefix=self._desired_response_prefix, + desired_response_prefix=self._configuration.desired_response_prefix, parent_id=parent_id, prompt_normalizer=self._prompt_normalizer, initial_prompt=initial_prompt, @@ -2109,7 +2132,7 @@ def _create_on_topic_scorer(self, objective: str) -> Scorer | None: - `None` if `on_topic_checking_enabled` is `False` or no scoring_target is available """ - if not self._on_topic_checking_enabled: + if not self._configuration.on_topic_checking_enabled: return None return TrueFalseInverterScorer( diff --git a/tests/unit/executor/attack/multi_turn/test_pair.py b/tests/unit/executor/attack/multi_turn/test_pair.py index 343dfd5800..662f216bac 100644 --- a/tests/unit/executor/attack/multi_turn/test_pair.py +++ b/tests/unit/executor/attack/multi_turn/test_pair.py @@ -101,10 +101,10 @@ def test_init_applies_pair_structural_defaults(self, objective_target, adversari attack_adversarial_config=adversarial_config, ) - assert attack._tree_width == 3 - assert attack._tree_depth == 5 - assert attack._branching_factor == 1 - assert attack._on_topic_checking_enabled is False + assert attack._configuration.tree_width == 3 + assert attack._configuration.tree_depth == 5 + assert attack._configuration.branching_factor == 1 + assert attack._configuration.on_topic_checking_enabled is False def test_branching_factor_is_not_exposed_in_signature(self): """branching_factor is definitional for PAIR (always 1) and must not be a public init kwarg.""" @@ -126,8 +126,8 @@ def test_tree_width_override(self, objective_target, adversarial_config): attack_adversarial_config=adversarial_config, tree_width=7, ) - assert attack._tree_width == 7 - assert attack._branching_factor == 1 + assert attack._configuration.tree_width == 7 + assert attack._configuration.branching_factor == 1 def test_tree_depth_override(self, objective_target, adversarial_config): attack = PAIRAttack( @@ -135,8 +135,8 @@ def test_tree_depth_override(self, objective_target, adversarial_config): attack_adversarial_config=adversarial_config, tree_depth=12, ) - assert attack._tree_depth == 12 - assert attack._on_topic_checking_enabled is False + assert attack._configuration.tree_depth == 12 + assert attack._configuration.on_topic_checking_enabled is False def test_is_subclass_of_tap(self): assert issubclass(PAIRAttack, TreeOfAttacksWithPruningAttack) 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 f4f962c224..239f7c444c 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 @@ -2,10 +2,11 @@ # Licensed under the MIT license. import asyncio +import inspect import json import logging import uuid -from dataclasses import dataclass, field +from dataclasses import FrozenInstanceError, dataclass, field, replace from typing import Any, cast from unittest.mock import AsyncMock, MagicMock, patch @@ -25,6 +26,7 @@ from pyrit.executor.attack.multi_turn.tree_of_attacks import ( AttackScoringConfig, TAPAttackScoringConfig, + _TAPAttackConfiguration, _TreeOfAttacksNode, ) from pyrit.models import ( @@ -485,42 +487,91 @@ def test_init_with_minimal_required_parameters(self, attack_builder): """Test that attack initializes correctly with only required parameters.""" attack = attack_builder.with_default_mocks().build() - assert attack._tree_width == 3 - assert attack._tree_depth == 5 - assert attack._branching_factor == 2 - assert attack._on_topic_checking_enabled is True - assert attack._batch_size == 10 + assert attack._configuration == _TAPAttackConfiguration( + tree_width=3, + tree_depth=5, + branching_factor=2, + on_topic_checking_enabled=True, + desired_response_prefix="Sure, here is", + batch_size=10, + ) def test_init_with_custom_tree_parameters(self, attack_builder): """Test initialization with custom tree parameters.""" attack = ( attack_builder.with_default_mocks() - .with_tree_params(tree_width=5, tree_depth=10, branching_factor=3, batch_size=20) + .with_tree_params( + tree_width=5, + tree_depth=10, + branching_factor=3, + on_topic_checking_enabled=False, + desired_response_prefix="Absolutely", + batch_size=20, + ) .build() ) - assert attack._tree_width == 5 - assert attack._tree_depth == 10 - assert attack._branching_factor == 3 - assert attack._batch_size == 20 + assert attack._configuration == _TAPAttackConfiguration( + tree_width=5, + tree_depth=10, + branching_factor=3, + on_topic_checking_enabled=False, + desired_response_prefix="Absolutely", + batch_size=20, + ) + + def test_configuration_is_immutable(self, basic_attack): + with pytest.raises(FrozenInstanceError): + basic_attack._configuration.tree_width = 4 # type: ignore[misc] + + def test_constructor_preserves_legacy_keyword_contract(self): + signature = inspect.signature(TreeOfAttacksWithPruningAttack.__init__) + + assert list(signature.parameters) == [ + "self", + "objective_target", + "attack_adversarial_config", + "attack_converter_config", + "attack_scoring_config", + "prompt_normalizer", + "tree_width", + "tree_depth", + "branching_factor", + "on_topic_checking_enabled", + "desired_response_prefix", + "batch_size", + "prepended_conversation_config", + ] + assert all( + parameter.kind is inspect.Parameter.KEYWORD_ONLY + for name, parameter in signature.parameters.items() + if name != "self" + ) + assert signature.parameters["tree_width"].default == 3 + assert signature.parameters["tree_depth"].default == 5 + assert signature.parameters["branching_factor"].default == 2 + assert signature.parameters["on_topic_checking_enabled"].default is True + assert signature.parameters["desired_response_prefix"].default == "Sure, here is" + assert signature.parameters["batch_size"].default == 10 @pytest.mark.parametrize( "tree_params,expected_error", [ - ({"tree_width": 0}, "tree width must be at least 1"), - ({"tree_depth": 0}, "tree depth must be at least 1"), - ({"branching_factor": 0}, "branching factor must be at least 1"), - ({"batch_size": 0}, "batch size must be at least 1"), - ({"tree_width": -1}, "tree width must be at least 1"), - ({"tree_depth": -1}, "tree depth must be at least 1"), - ({"branching_factor": -1}, "branching factor must be at least 1"), - ({"batch_size": -1}, "batch size must be at least 1"), + ({"tree_width": 0}, "The tree width must be at least 1."), + ({"tree_depth": 0}, "The tree depth must be at least 1."), + ({"branching_factor": 0}, "The branching factor must be at least 1."), + ({"batch_size": 0}, "The batch size must be at least 1."), + ({"tree_width": -1}, "The tree width must be at least 1."), + ({"tree_depth": -1}, "The tree depth must be at least 1."), + ({"branching_factor": -1}, "The branching factor must be at least 1."), + ({"batch_size": -1}, "The batch size must be at least 1."), ], ) def test_init_with_invalid_tree_parameters(self, attack_builder, tree_params, expected_error): """Test that invalid tree parameters raise ValueError.""" - with pytest.raises(ValueError, match=expected_error): + with pytest.raises(ValueError) as exc_info: attack_builder.with_default_mocks().with_tree_params(**tree_params).build() + assert str(exc_info.value) == expected_error def test_init_with_auxiliary_scorers(self, attack_builder): """Test initialization with auxiliary scorers.""" @@ -602,6 +653,43 @@ async def test_tree_depth_validation_with_prepended_conversation(self, attack_bu next_message=next_message, ) + async def test_interleaved_contexts_keep_independent_visualization_roots(self, basic_attack, helpers): + prepended_context = helpers.create_basic_context() + prepended_context.prepended_conversation = [ + Message.from_prompt(prompt="Hello", role="user"), + Message.from_prompt(prompt="Hi", role="assistant"), + ] + plain_context = helpers.create_basic_context() + first_setup_complete = asyncio.Event() + second_setup_complete = asyncio.Event() + + def create_node(**_: Any) -> MagicMock: + node = MagicMock(spec=_TreeOfAttacksNode) + node.adversarial_chat_conversation_id = str(uuid.uuid4()) + node.initialize_with_prepended_conversation_async = AsyncMock() + return node + + async def initialize_context_async(*, context: TAPAttackContext, first: bool) -> None: + await basic_attack._setup_async(context=context) + if first: + first_setup_complete.set() + await second_setup_complete.wait() + else: + await first_setup_complete.wait() + second_setup_complete.set() + await basic_attack._initialize_first_level_nodes_async(context) + + with patch.object(basic_attack, "_create_attack_node", side_effect=create_node): + await asyncio.gather( + initialize_context_async(context=prepended_context, first=True), + initialize_context_async(context=plain_context, first=False), + ) + + assert prepended_context.visualization_root_id == "prepended_1" + assert {node._vis_node_id for node in prepended_context.nodes} == {"prepended_1"} + assert plain_context.visualization_root_id == "root" + assert {node._vis_node_id for node in plain_context.nodes} == {"root"} + def test_default_scorer_detects_text_output_modalities(self): """Test that default scorer detects text output modalities from target capabilities.""" builder = AttackBuilder() @@ -643,7 +731,7 @@ def test_prune_nodes_to_maintain_width_removes_lowest_scoring_nodes(self, basic_ nodes = node_factory.create_nodes_with_scores([0.9, 0.7, 0.5, 0.3, 0.1]) context.nodes = nodes helpers.add_nodes_to_tree(context, nodes) - basic_attack._tree_width = 3 + basic_attack._configuration = replace(basic_attack._configuration, tree_width=3) # Execute pruning basic_attack._prune_nodes_to_maintain_width(context=context) @@ -921,7 +1009,7 @@ def test_prune_blocked_nodes_with_score_zero(self, attack_builder, node_factory, def test_no_pruning_when_below_width(self, basic_attack, node_factory, helpers): """Test that blocked nodes are not pruned when completed list is below tree_width.""" - basic_attack._tree_width = 5 + basic_attack._configuration = replace(basic_attack._configuration, tree_width=5) context = helpers.create_basic_context() @@ -1105,7 +1193,7 @@ class TestBranchingLogic: def test_branch_existing_nodes(self, basic_attack, node_factory, helpers): """Test that nodes are branched correctly.""" context = helpers.create_basic_context() - basic_attack._branching_factor = 3 + basic_attack._configuration = replace(basic_attack._configuration, branching_factor=3) # Create initial nodes initial_nodes = node_factory.create_nodes_with_scores([0.8, 0.7]) @@ -2157,7 +2245,7 @@ def test_branch_existing_nodes_tracks_adversarial_chat_conversation_ids(self, ba helpers.add_nodes_to_tree(context, nodes) # Set up branching factor to create additional nodes - basic_attack._branching_factor = 3 + basic_attack._configuration = replace(basic_attack._configuration, branching_factor=3) # Branch the nodes basic_attack._branch_existing_nodes(context) @@ -2191,7 +2279,7 @@ def test_initialize_first_level_nodes_tracks_adversarial_chat_conversation_ids(s context = helpers.create_basic_context() # Set tree width to create multiple nodes - basic_attack._tree_width = 3 + basic_attack._configuration = replace(basic_attack._configuration, tree_width=3) # Initialize first level nodes asyncio.run(basic_attack._initialize_first_level_nodes_async(context)) @@ -2237,7 +2325,7 @@ def test_initialize_first_level_nodes_edit_only_objective_seeds_all_roots(self, def test_initialize_first_level_nodes_text_objective_keeps_seed_on_first_root_only(self, basic_attack, helpers): """Text-capable objectives keep historical behavior: node 0 consumes next_message.""" context = helpers.create_basic_context() - basic_attack._tree_width = 3 + basic_attack._configuration = replace(basic_attack._configuration, tree_width=3) context.next_message = Message( message_pieces=[ MessagePiece.adversarial_placeholder(), diff --git a/tests/unit/scenario/foundry/test_red_team_agent.py b/tests/unit/scenario/foundry/test_red_team_agent.py index 1992b5c462..824420bf26 100644 --- a/tests/unit/scenario/foundry/test_red_team_agent.py +++ b/tests/unit/scenario/foundry/test_red_team_agent.py @@ -673,7 +673,7 @@ def test_attack_technique_maps_to_expected_type_and_configuration( attack = atomic_attack.attack_technique.attack assert type(attack) is expected_attack_type if expected_tree_width is not None: - assert attack._tree_width == expected_tree_width + assert attack._configuration.tree_width == expected_tree_width assert attack._objective_scorer is mock_float_threshold_scorer def test_mapping_cases_cover_all_concrete_foundry_techniques(self) -> None: