FEAT expand TargetCapabilities

pyrit/prompt_target/azure_blob_storage_target.py

@@ -15,6 +15,7 @@
 from pyrit.identifiers import ComponentIdentifier
 from pyrit.models import Message, construct_response_from_request
 from pyrit.prompt_target.common.prompt_target import PromptTarget
+from pyrit.prompt_target.common.target_capabilities import TargetCapabilities
 from pyrit.prompt_target.common.utils import limit_requests_per_minute
 
 logger = logging.getLogger(__name__)
@@ -49,6 +50,12 @@ class AzureBlobStorageTarget(PromptTarget):
     AZURE_STORAGE_CONTAINER_ENVIRONMENT_VARIABLE: str = "AZURE_STORAGE_ACCOUNT_CONTAINER_URL"
     SAS_TOKEN_ENVIRONMENT_VARIABLE: str = "AZURE_STORAGE_ACCOUNT_SAS_TOKEN"
 
+    _DEFAULT_CAPABILITIES: TargetCapabilities = TargetCapabilities(
+        input_modalities=["text", "url"],
+        output_modalities=["url"],
+        supports_multi_message_pieces=False,
+    )
+
     def __init__(
         self,
         *,
@@ -196,18 +203,3 @@ async def send_prompt_async(self, *, message: Message) -> list[Message]:
         )
 
         return [response]
-
-    def _validate_request(self, *, message: Message) -> None:
-        n_pieces = len(message.message_pieces)
-        if n_pieces != 1:
-            raise ValueError(f"This target only supports a single message piece. Received {n_pieces} pieces")
-
-        piece_type = message.message_pieces[0].converted_value_data_type
-        if piece_type not in ["text", "url"]:
-            raise ValueError(f"This target only supports text and url prompt input. Received: {piece_type}.")
-
-        request = message.message_pieces[0]
-        messages = self._memory.get_message_pieces(conversation_id=request.conversation_id)
-
-        if len(messages) > 0:
-            raise ValueError("This target only supports a single turn conversation.")

pyrit/prompt_target/azure_ml_chat_target.py

@@ -272,12 +272,3 @@ def _get_headers(self) -> dict[str, str]:
 
     def _validate_request(self, *, message: Message) -> None:
         pass
-
-    def is_json_response_supported(self) -> bool:
-        """
-        Check if the target supports JSON as a response format.
-
-        Returns:
-            bool: True if JSON response is supported, False otherwise.
-        """
-        return False

pyrit/prompt_target/common/prompt_chat_target.py

@@ -1,7 +1,6 @@
 # Copyright (c) Microsoft Corporation.
 # Licensed under the MIT license.
 
-import abc
 from typing import Optional
 
 from pyrit.identifiers import ComponentIdentifier
@@ -31,7 +30,7 @@ def __init__(
         endpoint: str = "",
         model_name: str = "",
         underlying_model: Optional[str] = None,
-        capabilities: Optional[TargetCapabilities] = None,
+        custom_capabilities: Optional[TargetCapabilities] = None,
     ) -> None:
         """
         Initialize the PromptChatTarget.
@@ -43,15 +42,15 @@ def __init__(
             underlying_model (str, Optional): The underlying model name (e.g., "gpt-4o") for
                 identification purposes. This is useful when the deployment name in Azure differs
                 from the actual model. Defaults to None.
-            capabilities (TargetCapabilities, Optional): Override the default capabilities for
+            custom_capabilities (TargetCapabilities, Optional): Override the default capabilities for
                 this target instance. If None, uses the class-level defaults. Defaults to None.
         """
         super().__init__(
             max_requests_per_minute=max_requests_per_minute,
             endpoint=endpoint,
             model_name=model_name,
             underlying_model=underlying_model,
-            capabilities=capabilities,
+            custom_capabilities=custom_capabilities,
         )
 
     def set_system_prompt(
@@ -85,15 +84,6 @@ def set_system_prompt(
             ).to_message()
         )
 
-    @abc.abstractmethod
-    def is_json_response_supported(self) -> bool:
-        """
-        Abstract method to determine if JSON response format is supported by the target.
-
-        Returns:
-            bool: True if JSON response is supported, False otherwise.
-        """
-
     def is_response_format_json(self, message_piece: MessagePiece) -> bool:
         """
         Check if the response format is JSON and ensure the target supports it.
@@ -127,7 +117,7 @@ def _get_json_response_config(self, *, message_piece: MessagePiece) -> _JsonResp
         """
         config = _JsonResponseConfig.from_metadata(metadata=message_piece.prompt_metadata)
 
-        if config.enabled and not self.is_json_response_supported():
+        if config.enabled and not self.capabilities.supports_json_response:
             target_name = self.get_identifier().class_name
             raise ValueError(f"This target {target_name} does not support JSON response format.")
 

pyrit/prompt_target/common/prompt_target.py

@@ -38,7 +38,7 @@ def __init__(
         endpoint: str = "",
         model_name: str = "",
         underlying_model: Optional[str] = None,
-        capabilities: Optional[TargetCapabilities] = None,
+        custom_capabilities: Optional[TargetCapabilities] = None,
     ) -> None:
         """
         Initialize the PromptTarget.
@@ -52,7 +52,7 @@ def __init__(
                 identification purposes. This is useful when the deployment name in Azure differs
                 from the actual model. If not provided, `model_name` will be used for the identifier.
                 Defaults to None.
-            capabilities (TargetCapabilities, Optional): Override the default capabilities for
+            custom_capabilities (TargetCapabilities, Optional): Override the default capabilities for
                 this target instance. Useful for targets whose capabilities depend on deployment
                 configuration (e.g., Playwright, HTTP). If None, uses the class-level
                 ``_DEFAULT_CAPABILITIES``. Defaults to None.
@@ -63,7 +63,9 @@ def __init__(
         self._endpoint = endpoint
         self._model_name = model_name
         self._underlying_model = underlying_model
-        self._capabilities = capabilities if capabilities is not None else type(self)._DEFAULT_CAPABILITIES
+        self._capabilities = (
+            custom_capabilities if custom_capabilities is not None else type(self)._DEFAULT_CAPABILITIES
+        )
 
         if self._verbose:
             logging.basicConfig(level=logging.INFO)
@@ -78,14 +80,36 @@ async def send_prompt_async(self, *, message: Message) -> list[Message]:
                 but some (like response target with tool calls) may return multiple messages.
         """
 
-    @abc.abstractmethod
     def _validate_request(self, *, message: Message) -> None:
         """
         Validate the provided message.
 
         Args:
             message: The message to validate.
+
+        Raises:
+            ValueError: if the target does not support the provided message pieces or if the message
+                violates any constraints based on the target's capabilities.
+
         """
+        n_pieces = len(message.message_pieces)
+        if not self.capabilities.supports_multi_message_pieces and n_pieces != 1:
+            raise ValueError(f"This target only supports a single message piece. Received: {n_pieces} pieces.")
+
+        for piece in message.message_pieces:
+            piece_type = piece.converted_value_data_type
+            if piece_type not in self.capabilities.input_modalities:
+                supported_types = ", ".join(sorted(self.capabilities.input_modalities))
+                raise ValueError(
+                    f"This target supports only the following data types: {supported_types}. Received: {piece_type}."
+                )
+
+        if not self.supports_multi_turn:
+            request = message.message_pieces[0]
+            messages = self._memory.get_message_pieces(conversation_id=request.conversation_id)
+
+            if len(messages) > 0:
+                raise ValueError("This target only supports a single turn conversation.")
 
     def set_model_name(self, *, model_name: str) -> None:
         """
@@ -140,6 +164,15 @@ def _create_identifier(
 
         return ComponentIdentifier.of(self, params=all_params, children=children)
 
+    def is_json_response_supported(self) -> bool:
+        """
+         Method to determine if JSON response format is supported by the target.
+
+        Returns:
+            bool: True if JSON response is supported, False otherwise.
+        """
+        return self.capabilities.supports_json_response
+
     @property
     def capabilities(self) -> TargetCapabilities:
         """

pyrit/prompt_target/common/target_capabilities.py

@@ -1,7 +1,9 @@
 # Copyright (c) Microsoft Corporation.
 # Licensed under the MIT license.
 
-from dataclasses import dataclass
+from dataclasses import dataclass, field, fields
+
+from pyrit.models import PromptDataType
 
 
 @dataclass(frozen=True)
@@ -20,3 +22,38 @@ class attribute. Users can override individual capabilities per instance
     # (i.e., it accepts and uses conversation history or maintains state
     # across turns via external mechanisms like WebSocket connections).
     supports_multi_turn: bool = False
+
+    # Whether the target natively supports multiple message pieces in a single request.
+    supports_multi_message_pieces: bool = True
+
+    # Whether the target natively supports JSON output (e.g., via a "json" response format).
+    supports_json_response: bool = False
+
+    # The input modalities supported by the target (e.g., "text", "image").
+    input_modalities: list[PromptDataType] = field(default_factory=lambda: ["text"])
+
+    # The output modalities supported by the target (e.g., "text", "image").
+    output_modalities: list[PromptDataType] = field(default_factory=lambda: ["text"])
+
+    def assert_satifies(self, required_capabilities: "TargetCapabilities") -> None:
+        """
+        Assert that the current capabilities satisfy the required capabilities.
+
+        Args:
+            required_capabilities (TargetCapabilities): The required capabilities to check against.
+
+        Raises:
+            ValueError: If any of the required capabilities are not satisfied.
+        """
+        unmet = []
+        for f in fields(required_capabilities):
+            required_value = getattr(required_capabilities, f.name)
+            self_value = getattr(self, f.name)
+            if isinstance(required_value, list):
+                missing = set(required_value) - set(self_value)
+                if missing:
+                    unmet.append(f"{f.name}: missing {missing}")
+            elif required_value and not self_value:
+                unmet.append(f.name)
+        if unmet:
+            raise ValueError(f"Target does not satisfy the following capabilities: {', '.join(unmet)}")

pyrit/prompt_target/crucible_target.py

@@ -14,6 +14,7 @@
 )
 from pyrit.models import Message, construct_response_from_request
 from pyrit.prompt_target.common.prompt_target import PromptTarget
+from pyrit.prompt_target.common.target_capabilities import TargetCapabilities
 from pyrit.prompt_target.common.utils import limit_requests_per_minute
 
 logger = logging.getLogger(__name__)
@@ -24,6 +25,8 @@ class CrucibleTarget(PromptTarget):
 
     API_KEY_ENVIRONMENT_VARIABLE: str = "CRUCIBLE_API_KEY"
 
+    _DEFAULT_CAPABILITIES: TargetCapabilities = TargetCapabilities(supports_multi_message_pieces=False)
+
     def __init__(
         self,
         *,
@@ -80,15 +83,6 @@ async def send_prompt_async(self, *, message: Message) -> list[Message]:
 
         return [response_entry]
 
-    def _validate_request(self, *, message: Message) -> None:
-        n_pieces = len(message.message_pieces)
-        if n_pieces != 1:
-            raise ValueError(f"This target only supports a single message piece. Received: {n_pieces} pieces.")
-
-        piece_type = message.message_pieces[0].converted_value_data_type
-        if piece_type != "text":
-            raise ValueError(f"This target only supports text prompt input. Received: {piece_type}.")
-
     @pyrit_target_retry
     async def _complete_text_async(self, text: str) -> str:
         payload: dict[str, object] = {

pyrit/prompt_target/gandalf_target.py

@@ -10,6 +10,7 @@
 from pyrit.identifiers import ComponentIdentifier
 from pyrit.models import Message, construct_response_from_request
 from pyrit.prompt_target.common.prompt_target import PromptTarget
+from pyrit.prompt_target.common.target_capabilities import TargetCapabilities
 from pyrit.prompt_target.common.utils import limit_requests_per_minute
 
 logger = logging.getLogger(__name__)
@@ -38,6 +39,8 @@ class GandalfLevel(enum.Enum):
 class GandalfTarget(PromptTarget):
     """A prompt target for the Gandalf security challenge."""
 
+    _DEFAULT_CAPABILITIES: TargetCapabilities = TargetCapabilities(supports_multi_message_pieces=False)
+
     def __init__(
         self,
         *,
@@ -93,15 +96,6 @@ async def send_prompt_async(self, *, message: Message) -> list[Message]:
 
         return [response_entry]
 
-    def _validate_request(self, *, message: Message) -> None:
-        n_pieces = len(message.message_pieces)
-        if n_pieces != 1:
-            raise ValueError(f"This target only supports a single message piece. Received: {n_pieces} pieces.")
-
-        piece_type = message.message_pieces[0].converted_value_data_type
-        if piece_type != "text":
-            raise ValueError(f"This target only supports text prompt input. Received: {piece_type}.")
-
     async def check_password(self, password: str) -> bool:
         """
         Check if the password is correct.

pyrit/prompt_target/http_target/http_target.py

@@ -5,7 +5,7 @@
 import json
 import logging
 import re
-from collections.abc import Callable, Sequence
+from collections.abc import Callable
 from typing import Any, Optional
 
 import httpx
@@ -17,6 +17,7 @@
     construct_response_from_request,
 )
 from pyrit.prompt_target.common.prompt_target import PromptTarget
+from pyrit.prompt_target.common.target_capabilities import TargetCapabilities
 from pyrit.prompt_target.common.utils import limit_requests_per_minute
 
 logger = logging.getLogger(__name__)
@@ -40,6 +41,8 @@ class HTTPTarget(PromptTarget):
         httpx_client_kwargs: (dict): additional keyword arguments to pass to the HTTP client
     """
 
+    _DEFAULT_CAPABILITIES: TargetCapabilities = TargetCapabilities(supports_multi_message_pieces=False)
+
     def __init__(
         self,
         http_request: str,
@@ -290,10 +293,3 @@ def _infer_full_url_from_host(
 
         host = headers_dict["host"]
         return f"{http_protocol}{host}{path}"
-
-    def _validate_request(self, *, message: Message) -> None:
-        message_pieces: Sequence[MessagePiece] = message.message_pieces
-
-        n_pieces = len(message_pieces)
-        if n_pieces != 1:
-            raise ValueError(f"This target only supports a single message piece. Received: {n_pieces} pieces.")

pyrit/prompt_target/hugging_face/hugging_face_chat_target.py

@@ -20,6 +20,7 @@
 from pyrit.identifiers import ComponentIdentifier
 from pyrit.models import Message, construct_response_from_request
 from pyrit.prompt_target.common.prompt_chat_target import PromptChatTarget
+from pyrit.prompt_target.common.target_capabilities import TargetCapabilities
 from pyrit.prompt_target.common.utils import limit_requests_per_minute
 
 logger = logging.getLogger(__name__)
@@ -34,6 +35,11 @@ class HuggingFaceChatTarget(PromptChatTarget):
     Inherits from PromptTarget to comply with the current design standards.
     """
 
+    _DEFAULT_CAPABILITIES: TargetCapabilities = TargetCapabilities(
+        supports_multi_turn=True,
+        supports_multi_message_pieces=False,
+    )
+
     # Class-level cache for model and tokenizer
     _cached_model = None
     _cached_tokenizer = None
@@ -388,25 +394,6 @@ def _apply_chat_template(self, messages: list[dict[str, str]]) -> Any:
         logger.error(error_message)
         raise ValueError(error_message)
 
-    def _validate_request(self, *, message: Message) -> None:
-        """
-        Validate the provided message.
-
-        Args:
-            message: The message to validate.
-
-        Raises:
-            ValueError: If the message does not contain exactly one text piece.
-            ValueError: If the message piece is not of type text.
-        """
-        n_pieces = len(message.message_pieces)
-        if n_pieces != 1:
-            raise ValueError(f"This target only supports a single message piece. Received: {n_pieces} pieces.")
-
-        piece_type = message.message_pieces[0].converted_value_data_type
-        if piece_type != "text":
-            raise ValueError(f"This target only supports text prompt input. Received: {piece_type}.")
-
     def is_json_response_supported(self) -> bool:
         """
         Check if the target supports JSON as a response format.