Add input/output modality validation to TargetRequirements.validate() (#1778)

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

Author: hannahwestra25

doc/code/targets/0_prompt_targets.md

@@ -78,6 +78,19 @@ CHAT_TARGET_REQUIREMENTS.validate(target=target)
 
 `TargetRequirements.validate` collects every missing capability and raises a single `ValueError`. For one-off checks against a single capability you can also call `target.configuration.ensure_can_handle(capability=...)` directly.
 
+`TargetRequirements` can also enforce **modality** constraints via `required_input_modalities` and `required_output_modalities`. Each entry is a set of `PromptDataType` values the consumer needs the target to accept (or produce). At least one of the target's modality combos must be a superset of each required combo:
+
+```python
+from pyrit.prompt_target import TargetRequirements
+
+# A consumer that requires image input and text output
+VISION_REQUIREMENTS = TargetRequirements(
+    required_input_modalities=frozenset({frozenset({"image_path"})}),
+    required_output_modalities=frozenset({frozenset({"text"})}),
+)
+VISION_REQUIREMENTS.validate(target=target)
+```
+
 ### Adapting vs raising
 
 Some capability gaps can be papered over by PyRIT itself. For example, a single-turn target can be made to *appear* multi-turn by flattening the conversation history into a single prompt before sending. The `CapabilityHandlingPolicy` controls this on a per-capability basis:

doc/code/targets/6_1_target_capabilities.ipynb

@@ -154,7 +154,12 @@
     "multi-turn + editable history — the replacement for the deprecated `PromptChatTarget` type check.\n",
     "\n",
     "`TargetRequirements.validate` collects every missing capability and raises a single `ValueError` so\n",
-    "callers see all violations at once."
+    "callers see all violations at once.\n",
+    "\n",
+    "`TargetRequirements` can also enforce **modality** constraints via `required_input_modalities` and\n",
+    "`required_output_modalities`. Each entry is a set of `PromptDataType` values the consumer needs\n",
+    "the target to accept (or produce). At least one of the target's modality combos must be a superset\n",
+    "of each required combo."
    ]
   },
   {

doc/code/targets/6_1_target_capabilities.py

@@ -92,6 +92,11 @@
 #
 # `TargetRequirements.validate` collects every missing capability and raises a single `ValueError` so
 # callers see all violations at once.
+#
+# `TargetRequirements` can also enforce **modality** constraints via `required_input_modalities` and
+# `required_output_modalities`. Each entry is a set of `PromptDataType` values the consumer needs
+# the target to accept (or produce). At least one of the target's modality combos must be a superset
+# of each required combo.
 
 # %%
 from pyrit.prompt_target import CHAT_TARGET_REQUIREMENTS

pyrit/prompt_target/common/target_requirements.py

@@ -9,6 +9,7 @@
 from pyrit.prompt_target.common.target_capabilities import CapabilityName
 
 if TYPE_CHECKING:
+    from pyrit.models import PromptDataType
     from pyrit.prompt_target.common.prompt_target import PromptTarget
 
 
@@ -33,10 +34,20 @@ class TargetRequirements:
       consumer's semantics (e.g. an attack that depends on the target
       remembering prior turns, where history-squash normalization would
       collapse the conversation into a single prompt).
+
+    Modality requirements are also supported:
+
+    * ``required_input_modalities`` — each entry is a frozenset of
+      :class:`PromptDataType` values the consumer needs the target to
+      accept. At least one of the target's input modality combos must be
+      a superset of each required combo.
+    * ``required_output_modalities`` — same semantics for outputs.
     """
 
     required: frozenset[CapabilityName] = field(default_factory=frozenset)
     native_required: frozenset[CapabilityName] = field(default_factory=frozenset)
+    required_input_modalities: frozenset[frozenset[PromptDataType]] = field(default_factory=frozenset)
+    required_output_modalities: frozenset[frozenset[PromptDataType]] = field(default_factory=frozenset)
 
     def validate(self, *, target: PromptTarget) -> None:
         """
@@ -52,7 +63,9 @@ def validate(self, *, target: PromptTarget) -> None:
         Raises:
             ValueError: If any ``native_required`` capability is not natively
                 supported, or if any ``required`` capability is not supported
-                natively and has no ``ADAPT`` entry in the target's policy.
+                natively and has no ``ADAPT`` entry in the target's policy,
+                or if the target's modalities do not satisfy
+                ``required_input_modalities`` / ``required_output_modalities``.
         """
         errors: list[str] = [
             f"Target must natively support '{capability.value}'; adaptation is not acceptable for this consumer."
@@ -66,12 +79,42 @@ def validate(self, *, target: PromptTarget) -> None:
             except ValueError as exc:
                 errors.append(str(exc))
 
+        errors.extend(
+            self._check_modalities(
+                required=self.required_input_modalities,
+                supported=target.configuration.capabilities.input_modalities,
+                direction="input",
+            )
+        )
+        errors.extend(
+            self._check_modalities(
+                required=self.required_output_modalities,
+                supported=target.configuration.capabilities.output_modalities,
+                direction="output",
+            )
+        )
+
         if errors:
             raise ValueError(
                 f"Target does not satisfy {len(errors)} required capability(ies):\n"
                 + "\n".join(f"  - {e}" for e in errors)
             )
 
+    @staticmethod
+    def _check_modalities(
+        *,
+        required: frozenset[frozenset[PromptDataType]],
+        supported: frozenset[frozenset[PromptDataType]],
+        direction: str,
+    ) -> list[str]:
+        """Return error strings for each required modality combo not covered by *supported*."""
+        return [
+            f"Target must support {direction} modality {{{', '.join(sorted(combo))}}}; "
+            f"supported: {[sorted(s) for s in sorted(supported, key=lambda s: sorted(s))]}."
+            for combo in sorted(required, key=lambda c: sorted(c))
+            if not any(combo <= sup for sup in supported)
+        ]
+
 
 def _build_chat_target_requirements() -> TargetRequirements:
     """

tests/unit/prompt_target/target/test_target_requirements.py

@@ -262,3 +262,104 @@ def test_validate_aggregates_all_violations():
     assert "2 required capability" in message
     assert CapabilityName.MULTI_TURN.value in message
     assert CapabilityName.JSON_OUTPUT.value in message
+
+
+# ---------------------------------------------------------------------------
+# Modality validation
+# ---------------------------------------------------------------------------
+
+
+def test_validate_passes_when_no_modality_requirements():
+    """Backward compat: empty modality requirements should always pass."""
+    target = _make_target(
+        configuration=TargetConfiguration(
+            capabilities=TargetCapabilities(),
+        ),
+    )
+    TargetRequirements().validate(target=target)
+
+
+def test_validate_passes_when_input_modality_matches():
+    reqs = TargetRequirements(
+        required_input_modalities=frozenset({frozenset({"text"})}),
+    )
+    target = _make_target(
+        configuration=TargetConfiguration(
+            capabilities=TargetCapabilities(
+                input_modalities=frozenset({frozenset({"text"})}),
+            ),
+        ),
+    )
+    reqs.validate(target=target)
+
+
+def test_validate_passes_when_target_modality_is_superset():
+    reqs = TargetRequirements(
+        required_input_modalities=frozenset({frozenset({"text"})}),
+    )
+    target = _make_target(
+        configuration=TargetConfiguration(
+            capabilities=TargetCapabilities(
+                input_modalities=frozenset({frozenset({"text", "image_path"})}),
+            ),
+        ),
+    )
+    reqs.validate(target=target)
+
+
+def test_validate_fails_on_missing_input_modality():
+    reqs = TargetRequirements(
+        required_input_modalities=frozenset({frozenset({"image_path"})}),
+    )
+    target = _make_target(
+        configuration=TargetConfiguration(
+            capabilities=TargetCapabilities(
+                input_modalities=frozenset({frozenset({"text"})}),
+            ),
+        ),
+    )
+    with pytest.raises(ValueError, match="input modality"):
+        reqs.validate(target=target)
+
+
+def test_validate_fails_on_missing_output_modality():
+    reqs = TargetRequirements(
+        required_output_modalities=frozenset({frozenset({"audio_path"})}),
+    )
+    target = _make_target(
+        configuration=TargetConfiguration(
+            capabilities=TargetCapabilities(
+                output_modalities=frozenset({frozenset({"text"})}),
+            ),
+        ),
+    )
+    with pytest.raises(ValueError, match="output modality"):
+        reqs.validate(target=target)
+
+
+def test_validate_aggregates_modality_and_capability_errors():
+    reqs = TargetRequirements(
+        native_required=frozenset({CapabilityName.MULTI_TURN}),
+        required_input_modalities=frozenset({frozenset({"image_path"})}),
+    )
+    target = _make_target(
+        configuration=TargetConfiguration(
+            capabilities=TargetCapabilities(
+                supports_multi_turn=False,
+                input_modalities=frozenset({frozenset({"text"})}),
+            ),
+        ),
+    )
+    with pytest.raises(ValueError) as exc_info:
+        reqs.validate(target=target)
+
+    message = str(exc_info.value)
+    assert "2 required capability" in message
+    assert CapabilityName.MULTI_TURN.value in message
+    assert "input modality" in message
+
+
+def test_default_modality_requirements_are_empty():
+    reqs = TargetRequirements()
+    assert reqs.required_input_modalities == frozenset()
+    assert reqs.required_output_modalities == frozenset()