Privatize modality feedback router API

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

Author: Copilot

pyrit/executor/attack/component/__init__.py

@@ -11,9 +11,6 @@
     get_prepended_turn_count,
     mark_messages_as_simulated,
 )
-from pyrit.executor.attack.component.modality_router import (
-    ModalityFeedbackRouter,
-)
 from pyrit.executor.attack.component.prepended_conversation_config import (
     PrependedConversationConfig,
 )
@@ -25,6 +22,5 @@
     "get_adversarial_chat_messages",
     "get_prepended_turn_count",
     "mark_messages_as_simulated",
-    "ModalityFeedbackRouter",
     "PrependedConversationConfig",
 ]

pyrit/executor/attack/component/modality_router.py

@@ -25,7 +25,7 @@
 logger = logging.getLogger(__name__)
 
 
-class ModalityFeedbackRouter:
+class _ModalityFeedbackRouter:
     """
     Capability-aware utility for building ``Message`` objects between an
     adversarial chat and an objective target in multi-turn attacks.

pyrit/executor/attack/multi_turn/crescendo.py

@@ -21,9 +21,9 @@
 )
 from pyrit.executor.attack.component import (
     ConversationManager,
-    ModalityFeedbackRouter,
     PrependedConversationConfig,
 )
+from pyrit.executor.attack.component.modality_router import _ModalityFeedbackRouter
 from pyrit.executor.attack.core import (
     AttackAdversarialConfig,
     AttackConverterConfig,
@@ -226,7 +226,7 @@ def __init__(
         # whether prior media should travel back to the adversarial chat or
         # forward to the objective target, and that fills in adversarial
         # placeholders when ``next_message`` carries seed media.
-        self._modality_router = ModalityFeedbackRouter(
+        self._modality_router = _ModalityFeedbackRouter(
             adversarial_chat=self._adversarial_chat,
             objective_target=objective_target,
         )

pyrit/executor/attack/multi_turn/red_teaming.py

@@ -14,9 +14,9 @@
 from pyrit.exceptions import ComponentRole, execution_context
 from pyrit.executor.attack.component import (
     ConversationManager,
-    ModalityFeedbackRouter,
     get_adversarial_chat_messages,
 )
+from pyrit.executor.attack.component.modality_router import _ModalityFeedbackRouter
 from pyrit.executor.attack.core.attack_config import (
     DEFAULT_ADVERSARIAL_SEED_PROMPT,
     AttackAdversarialConfig,
@@ -161,7 +161,7 @@ def __init__(
         # whether prior media should travel back to the adversarial chat or
         # forward to the objective target, and that fills in adversarial
         # placeholders when ``next_message`` carries seed media.
-        self._modality_router = ModalityFeedbackRouter(
+        self._modality_router = _ModalityFeedbackRouter(
             adversarial_chat=self._adversarial_chat,
             objective_target=objective_target,
         )

pyrit/executor/attack/multi_turn/tree_of_attacks.py

@@ -25,13 +25,13 @@
 )
 from pyrit.executor.attack.component import (
     ConversationManager,
-    ModalityFeedbackRouter,
     PrependedConversationConfig,
     get_prepended_turn_count,
 )
 from pyrit.executor.attack.component.conversation_manager import (
     build_conversation_context_string_async,
 )
+from pyrit.executor.attack.component.modality_router import _ModalityFeedbackRouter
 from pyrit.executor.attack.core.attack_config import (
     AttackAdversarialConfig,
     AttackConverterConfig,
@@ -295,7 +295,7 @@ def __init__(
         auxiliary_scorers: list[Scorer] | None,
         attack_id: ComponentIdentifier,
         attack_strategy_name: str,
-        modality_router: ModalityFeedbackRouter,
+        modality_router: _ModalityFeedbackRouter,
         memory_labels: dict[str, str] | None = None,
         parent_id: str | None = None,
         prompt_normalizer: PromptNormalizer | None = None,
@@ -318,7 +318,7 @@ def __init__(
             auxiliary_scorers (list[Scorer] | None): Additional scorers for the response
             attack_id (ComponentIdentifier): Unique identifier for the attack.
             attack_strategy_name (str): Name of the attack strategy for execution context.
-            modality_router (ModalityFeedbackRouter): Capability-aware router that decides
+            modality_router (_ModalityFeedbackRouter): Capability-aware router that decides
                 whether prior media should travel back to the adversarial chat or forward to
                 the objective target, and fills adversarial-placeholder pieces in seed
                 messages. Typically shared across all nodes of the same attack.
@@ -1430,7 +1430,7 @@ def __init__(
         # forward to the objective target, and that fills in adversarial
         # placeholders when ``next_message`` carries seed media. Shared across
         # all nodes of the tree.
-        self._modality_router = ModalityFeedbackRouter(
+        self._modality_router = _ModalityFeedbackRouter(
             adversarial_chat=self._adversarial_chat,
             objective_target=objective_target,
         )

tests/unit/executor/attack/component/test_modality_router.py

@@ -1,7 +1,7 @@
 # Copyright (c) Microsoft Corporation.
 # Licensed under the MIT license.
 
-"""Unit tests for :class:`ModalityFeedbackRouter`."""
+"""Unit tests for :class:`_ModalityFeedbackRouter`."""
 
 from __future__ import annotations
 
@@ -10,7 +10,7 @@
 
 import pytest
 
-from pyrit.executor.attack.component import ModalityFeedbackRouter
+from pyrit.executor.attack.component.modality_router import _ModalityFeedbackRouter
 from pyrit.models import Message, MessagePiece
 from pyrit.prompt_target.common.target_capabilities import TargetCapabilities
 from pyrit.prompt_target.common.target_configuration import TargetConfiguration
@@ -53,23 +53,23 @@ def _make_response(*, data_type: PromptDataType, value: str = "some-value") -> M
 # ---------------------------------------------------------------------------
 class TestValidateFirstTurnSeed:
     def test_text_only_target_no_seed_required(self):
-        router = ModalityFeedbackRouter(
+        router = _ModalityFeedbackRouter(
             adversarial_chat=_build_target(input_modalities=[{"text"}]),
             objective_target=_build_target(input_modalities=[{"text"}]),
         )
 
         router.validate_first_turn_seed(next_message=None)
 
     def test_multimodal_target_with_text_allowed_no_seed_required(self):
-        router = ModalityFeedbackRouter(
+        router = _ModalityFeedbackRouter(
             adversarial_chat=_build_target(input_modalities=[{"text"}]),
             objective_target=_build_target(input_modalities=[{"text"}, {"text", "image_path"}]),
         )
 
         router.validate_first_turn_seed(next_message=None)
 
     def test_edit_only_target_without_next_message_raises(self):
-        router = ModalityFeedbackRouter(
+        router = _ModalityFeedbackRouter(
             adversarial_chat=_build_target(input_modalities=[{"text"}]),
             objective_target=_build_target(input_modalities=[{"text", "image_path"}]),
         )
@@ -78,7 +78,7 @@ def test_edit_only_target_without_next_message_raises(self):
             router.validate_first_turn_seed(next_message=None)
 
     def test_edit_only_target_with_text_only_next_message_raises(self):
-        router = ModalityFeedbackRouter(
+        router = _ModalityFeedbackRouter(
             adversarial_chat=_build_target(input_modalities=[{"text"}]),
             objective_target=_build_target(input_modalities=[{"text", "image_path"}]),
         )
@@ -88,7 +88,7 @@ def test_edit_only_target_with_text_only_next_message_raises(self):
             router.validate_first_turn_seed(next_message=next_message)
 
     def test_edit_only_target_with_media_seed_passes(self):
-        router = ModalityFeedbackRouter(
+        router = _ModalityFeedbackRouter(
             adversarial_chat=_build_target(input_modalities=[{"text"}]),
             objective_target=_build_target(input_modalities=[{"text", "image_path"}]),
         )
@@ -120,7 +120,7 @@ def test_edit_only_target_with_media_seed_passes(self):
 class TestBuildAdversarialInputMessage:
     @pytest.mark.parametrize("media_type", ["image_path", "audio_path", "video_path"])
     def test_text_only_adversarial_falls_back_to_text(self, media_type):
-        router = ModalityFeedbackRouter(
+        router = _ModalityFeedbackRouter(
             adversarial_chat=_build_target(input_modalities=[{"text"}]),
             objective_target=_build_target(input_modalities=[{"text"}, {"text", media_type}]),
         )
@@ -137,7 +137,7 @@ def test_text_only_adversarial_falls_back_to_text(self, media_type):
 
     @pytest.mark.parametrize("media_type", ["image_path", "audio_path", "video_path"])
     def test_multimodal_adversarial_attaches_media(self, media_type):
-        router = ModalityFeedbackRouter(
+        router = _ModalityFeedbackRouter(
             adversarial_chat=_build_target(input_modalities=[{"text"}, {"text", media_type}]),
             objective_target=_build_target(input_modalities=[{"text"}]),
         )
@@ -157,7 +157,7 @@ def test_multimodal_adversarial_attaches_media(self, media_type):
         assert msg.message_pieces[0].conversation_id == msg.message_pieces[1].conversation_id
 
     def test_no_last_response_returns_text_only(self):
-        router = ModalityFeedbackRouter(
+        router = _ModalityFeedbackRouter(
             adversarial_chat=_build_target(input_modalities=[{"text"}, {"text", "image_path"}]),
             objective_target=_build_target(input_modalities=[{"text"}]),
         )
@@ -172,7 +172,7 @@ def test_no_last_response_returns_text_only(self):
 
     def test_response_media_type_not_supported_falls_back(self):
         # Adversarial accepts {text, image_path} but the response is audio_path.
-        router = ModalityFeedbackRouter(
+        router = _ModalityFeedbackRouter(
             adversarial_chat=_build_target(input_modalities=[{"text"}, {"text", "image_path"}]),
             objective_target=_build_target(input_modalities=[{"text"}, {"text", "audio_path"}]),
         )
@@ -213,7 +213,7 @@ def _make_placeholder_message() -> Message:
         )
 
     def test_replaces_placeholder_only(self):
-        router = ModalityFeedbackRouter(
+        router = _ModalityFeedbackRouter(
             adversarial_chat=_build_target(input_modalities=[{"text"}]),
             objective_target=_build_target(input_modalities=[{"text", "image_path"}]),
         )
@@ -232,7 +232,7 @@ def test_replaces_placeholder_only(self):
         assert filled.message_pieces[1].original_value_data_type == "image_path"
 
     def test_does_not_mutate_input_message(self):
-        router = ModalityFeedbackRouter(
+        router = _ModalityFeedbackRouter(
             adversarial_chat=_build_target(input_modalities=[{"text"}]),
             objective_target=_build_target(input_modalities=[{"text", "image_path"}]),
         )
@@ -248,7 +248,7 @@ def test_does_not_mutate_input_message(self):
         assert message.message_pieces[0].original_value == ""
 
     def test_filled_pieces_share_conversation_id(self):
-        router = ModalityFeedbackRouter(
+        router = _ModalityFeedbackRouter(
             adversarial_chat=_build_target(input_modalities=[{"text"}]),
             objective_target=_build_target(input_modalities=[{"text", "image_path"}]),
         )
@@ -268,7 +268,7 @@ def test_filled_pieces_share_conversation_id(self):
 # ---------------------------------------------------------------------------
 class TestBuildObjectiveInputMessage:
     def test_turn_zero_text_only_target(self):
-        router = ModalityFeedbackRouter(
+        router = _ModalityFeedbackRouter(
             adversarial_chat=_build_target(input_modalities=[{"text"}]),
             objective_target=_build_target(input_modalities=[{"text"}]),
         )
@@ -283,7 +283,7 @@ def test_turn_zero_text_only_target(self):
         assert msg.get_value() == "adversarial prompt"
 
     def test_turn_zero_edit_only_target_raises_defensive(self):
-        router = ModalityFeedbackRouter(
+        router = _ModalityFeedbackRouter(
             adversarial_chat=_build_target(input_modalities=[{"text"}]),
             objective_target=_build_target(input_modalities=[{"text", "image_path"}]),
         )
@@ -296,7 +296,7 @@ def test_turn_zero_edit_only_target_raises_defensive(self):
             )
 
     def test_turn_n_attaches_prev_media_when_supported(self):
-        router = ModalityFeedbackRouter(
+        router = _ModalityFeedbackRouter(
             adversarial_chat=_build_target(input_modalities=[{"text"}]),
             objective_target=_build_target(input_modalities=[{"text"}, {"text", "image_path"}]),
         )
@@ -314,7 +314,7 @@ def test_turn_n_attaches_prev_media_when_supported(self):
         assert msg.message_pieces[1].original_value == "/tmp/prev.png"
 
     def test_turn_n_text_only_target_drops_media(self):
-        router = ModalityFeedbackRouter(
+        router = _ModalityFeedbackRouter(
             adversarial_chat=_build_target(input_modalities=[{"text"}]),
             objective_target=_build_target(input_modalities=[{"text"}]),
         )
@@ -330,7 +330,7 @@ def test_turn_n_text_only_target_drops_media(self):
         assert msg.get_value() == "refine this further"
 
     def test_turn_n_no_prev_response_returns_text_only(self):
-        router = ModalityFeedbackRouter(
+        router = _ModalityFeedbackRouter(
             adversarial_chat=_build_target(input_modalities=[{"text"}]),
             objective_target=_build_target(input_modalities=[{"text"}, {"text", "image_path"}]),
         )
@@ -346,7 +346,7 @@ def test_turn_n_no_prev_response_returns_text_only(self):
 
     @pytest.mark.parametrize("media_type", ["image_path", "audio_path", "video_path"])
     def test_turn_n_generic_media_types(self, media_type):
-        router = ModalityFeedbackRouter(
+        router = _ModalityFeedbackRouter(
             adversarial_chat=_build_target(input_modalities=[{"text"}]),
             objective_target=_build_target(input_modalities=[{"text"}, {"text", media_type}]),
         )
@@ -367,15 +367,15 @@ def test_turn_n_generic_media_types(self, media_type):
 # ---------------------------------------------------------------------------
 class TestProperties:
     def test_objective_target_requires_media_on_first_turn_when_text_excluded(self):
-        router = ModalityFeedbackRouter(
+        router = _ModalityFeedbackRouter(
             adversarial_chat=_build_target(input_modalities=[{"text"}]),
             objective_target=_build_target(input_modalities=[{"text", "image_path"}]),
         )
 
         assert router.objective_target_requires_media_on_first_turn is True
 
     def test_objective_does_not_require_media_when_text_allowed(self):
-        router = ModalityFeedbackRouter(
+        router = _ModalityFeedbackRouter(
             adversarial_chat=_build_target(input_modalities=[{"text"}]),
             objective_target=_build_target(input_modalities=[{"text"}, {"text", "image_path"}]),
         )

tests/unit/executor/attack/multi_turn/test_adversarial_chat_schema_forwarding.py

@@ -14,7 +14,7 @@
 from unit.mocks import MockPromptTarget
 
 from pyrit.executor.attack import AttackAdversarialConfig, AttackParameters
-from pyrit.executor.attack.component import ModalityFeedbackRouter
+from pyrit.executor.attack.component.modality_router import _ModalityFeedbackRouter
 from pyrit.executor.attack.multi_turn.crescendo import (
     ConversationSession,
     CrescendoAttack,
@@ -84,7 +84,7 @@ async def test_tap_forwards_schema_to_adversarial_target(patch_central_database)
         auxiliary_scorers=[],
         attack_id=attack.get_identifier(),
         attack_strategy_name="TreeOfAttacksWithPruningAttack",
-        modality_router=ModalityFeedbackRouter(adversarial_chat=adversarial, objective_target=objective),
+        modality_router=_ModalityFeedbackRouter(adversarial_chat=adversarial, objective_target=objective),
     )
 
     await node._send_to_adversarial_chat_async(prompt_text="hello")

tests/unit/executor/attack/multi_turn/test_crescendo.py

@@ -67,7 +67,7 @@ def create_mock_chat_target(*, name: str = "MockChatTarget") -> MagicMock:
     target.send_prompt_async = AsyncMock()
     target.set_system_prompt = MagicMock()
     target.get_identifier.return_value = _mock_target_id(name)
-    # Sensible default for the ModalityFeedbackRouter: text-only target. Tests
+    # Sensible default for the _ModalityFeedbackRouter: text-only target. Tests
     # that need multimodal behavior override input_modalities on their own copy.
     target.configuration.capabilities.input_modalities = frozenset({frozenset({"text"})})
     target.configuration.capabilities.output_modalities = frozenset({frozenset({"text"})})

tests/unit/executor/attack/multi_turn/test_supports_multi_turn_attacks.py

@@ -373,7 +373,7 @@ class TestTAPNodeDuplicateSystemMessages:
 
     def _make_tap_node(self, *, supports_multi_turn: bool):
         """Create a minimal _TreeOfAttacksNode for testing."""
-        from pyrit.executor.attack.component import ModalityFeedbackRouter
+        from pyrit.executor.attack.component.modality_router import _ModalityFeedbackRouter
         from pyrit.executor.attack.multi_turn.tree_of_attacks import _TreeOfAttacksNode
 
         target = MagicMock()
@@ -408,7 +408,7 @@ def _make_tap_node(self, *, supports_multi_turn: bool):
             auxiliary_scorers=None,
             attack_id=MagicMock(),
             attack_strategy_name="TAP",
-            modality_router=ModalityFeedbackRouter(
+            modality_router=_ModalityFeedbackRouter(
                 adversarial_chat=adversarial_chat,
                 objective_target=target,
             ),
@@ -756,7 +756,7 @@ class TestTAPBranchingPreservesSystemPrompts:
 
     def _make_tap_node(self, *, supports_multi_turn: bool):
         """Create a _TreeOfAttacksNode with real memory."""
-        from pyrit.executor.attack.component import ModalityFeedbackRouter
+        from pyrit.executor.attack.component.modality_router import _ModalityFeedbackRouter
         from pyrit.executor.attack.multi_turn.tree_of_attacks import _TreeOfAttacksNode
 
         target = MagicMock()
@@ -791,7 +791,7 @@ def _make_tap_node(self, *, supports_multi_turn: bool):
             auxiliary_scorers=None,
             attack_id=MagicMock(),
             attack_strategy_name="TAP",
-            modality_router=ModalityFeedbackRouter(
+            modality_router=_ModalityFeedbackRouter(
                 adversarial_chat=adversarial_chat,
                 objective_target=target,
             ),