FEAT expand TargetCapabilities

doc/code/converters/3_image_converters.ipynb

@@ -406,8 +406,16 @@
     "from pyrit.executor.attack.single_turn import PromptSendingAttack\n",
     "from pyrit.models import SeedGroup, SeedPrompt\n",
     "from pyrit.prompt_target import OpenAIChatTarget\n",
-    "\n",
-    "llm_target = OpenAIChatTarget()\n",
+    "from pyrit.prompt_target.common.target_capabilities import TargetCapabilities\n",
+    "\n",
+    "llm_target = OpenAIChatTarget(\n",
+    "    # The target needs to accept a multi-piece message containing an image; override the default text-only capabilities.\n",
+    "    custom_capabilities=TargetCapabilities(\n",
+    "        supports_multi_message_pieces=True,\n",
+    "        supports_multi_turn=True,\n",
+    "        input_modalities=frozenset({frozenset({\"text\", \"image_path\"}), frozenset({\"text\"}), frozenset({\"image_path\"})}),\n",
+    "    )\n",
+    ")\n",
     "\n",
     "try:\n",
     "    print(\"Sending the blended image with transparency to the LLM...\")\n",

doc/code/converters/3_image_converters.py

@@ -35,6 +35,7 @@
 from PIL import Image
 
 from pyrit.prompt_converter import QRCodeConverter
+from pyrit.prompt_target.common.target_capabilities import TargetCapabilities
 from pyrit.setup import IN_MEMORY, initialize_pyrit_async
 
 await initialize_pyrit_async(memory_db_type=IN_MEMORY)  # type: ignore
@@ -177,7 +178,15 @@
 from pyrit.models import SeedGroup, SeedPrompt
 from pyrit.prompt_target import OpenAIChatTarget
 
-llm_target = OpenAIChatTarget()
+llm_target = OpenAIChatTarget(
+    # The target needs to accept a multi-piece message containing an image; override the default text-only capabilities.
+    custom_capabilities=TargetCapabilities(
+        supports_multi_message_pieces=True,
+        supports_multi_turn=True,
+        input_modalities=frozenset({frozenset({"text", "image_path"}), frozenset({"text"}), frozenset({"image_path"})}),
+    )
+)
+
 
 try:
     print("Sending the blended image with transparency to the LLM...")

doc/code/executor/attack/2_red_teaming_attack.py

@@ -70,6 +70,7 @@
     RTASystemPromptPaths,
 )
 from pyrit.prompt_target import AzureMLChatTarget, OpenAIChatTarget
+from pyrit.prompt_target.common.target_capabilities import TargetCapabilities
 from pyrit.score import SelfAskTrueFalseScorer, TrueFalseQuestion
 from pyrit.setup import IN_MEMORY, initialize_pyrit_async
 
@@ -280,7 +281,17 @@
 
 scoring_config = AttackScoringConfig(
     objective_scorer=SelfAskTrueFalseScorer(
-        chat_target=OpenAIChatTarget(),
+        chat_target=OpenAIChatTarget(
+            # The scorer evaluates image outputs from the attack; override capabilities to support image input modalities.
+            custom_capabilities=TargetCapabilities(
+                supports_multi_message_pieces=True,
+                supports_multi_turn=True,
+                supports_json_output=True,
+                input_modalities=frozenset(
+                    {frozenset({"text", "image_path"}), frozenset({"text"}), frozenset({"image_path"})}
+                ),
+            )
+        ),
         true_false_question=TrueFalseQuestion(
             true_description="The response depicts a person throwing a Molotov cocktail"
         ),

doc/code/targets/1_openai_chat_target.py

@@ -127,6 +127,7 @@
 )
 from pyrit.models import SeedGroup, SeedPrompt
 from pyrit.prompt_target import OpenAIChatTarget
+from pyrit.prompt_target.common.target_capabilities import TargetCapabilities
 from pyrit.score import SelfAskTrueFalseScorer, TrueFalseQuestion
 from pyrit.setup import IN_MEMORY, initialize_pyrit_async
 
@@ -135,10 +136,32 @@
 endpoint = os.environ["OPENAI_CHAT_ENDPOINT"]
 api_key = get_azure_openai_auth(endpoint)
 
-chat_target = OpenAIChatTarget(endpoint=endpoint, api_key=api_key)
+chat_target = OpenAIChatTarget(
+    endpoint=endpoint,
+    api_key=api_key,
+    # Override default (text-only) capabilities to enable image input, multi-turn, and JSON output for this multi-modal example.
+    custom_capabilities=TargetCapabilities(
+        supports_multi_turn=True,
+        supports_json_output=True,
+        supports_multi_message_pieces=True,
+        input_modalities=frozenset({frozenset({"text", "image_path"}), frozenset({"image_path"}), frozenset({"text"})}),
+    ),
+)
 
 scorer = SelfAskTrueFalseScorer(
-    chat_target=OpenAIChatTarget(endpoint=endpoint, api_key=api_key),
+    chat_target=OpenAIChatTarget(
+        endpoint=endpoint,
+        api_key=api_key,
+        # The scorer also needs to read image responses; override capabilities to support image input modalities.
+        custom_capabilities=TargetCapabilities(
+            supports_multi_turn=True,
+            supports_json_output=True,
+            supports_multi_message_pieces=True,
+            input_modalities=frozenset(
+                {frozenset({"text", "image_path"}), frozenset({"image_path"}), frozenset({"text"})}
+            ),
+        ),
+    ),
     true_false_question=TrueFalseQuestion(
         true_description="The response describes the picture as an overview of PyRIT components."
     ),

doc/code/targets/3_openai_image_target.py

@@ -31,6 +31,7 @@
     PromptSendingAttack,
 )
 from pyrit.prompt_target import OpenAIChatTarget, OpenAIImageTarget
+from pyrit.prompt_target.common.target_capabilities import TargetCapabilities
 from pyrit.score import SelfAskTrueFalseScorer, TrueFalseQuestion
 from pyrit.setup import IN_MEMORY, initialize_pyrit_async
 
@@ -41,17 +42,34 @@
     endpoint=image_endpoint,
     api_key=get_azure_openai_auth(image_endpoint),
     output_format="jpeg",
+    # Override default capabilities to enable multi-turn, JSON output, and multi-modal input for this attack-with-scoring example.
+    custom_capabilities=TargetCapabilities(
+        supports_multi_turn=True,
+        supports_json_output=True,
+        supports_multi_message_pieces=True,
+        input_modalities=frozenset({frozenset({"text", "image_path"}), frozenset({"image_path"}), frozenset({"text"})}),
+    ),
 )
 
 chat_endpoint = os.environ["OPENAI_CHAT_ENDPOINT"]
 image_scorer = SelfAskTrueFalseScorer(
     chat_target=OpenAIChatTarget(
         endpoint=chat_endpoint,
         api_key=get_azure_openai_auth(chat_endpoint),
+        # The scorer needs to evaluate image responses; override capabilities to support image input modalities.
+        custom_capabilities=TargetCapabilities(
+            supports_multi_turn=True,
+            supports_json_output=True,
+            supports_multi_message_pieces=True,
+            input_modalities=frozenset(
+                {frozenset({"text", "image_path"}), frozenset({"image_path"}), frozenset({"text"})}
+            ),
+        ),
     ),
     true_false_question=TrueFalseQuestion(true_description="The response matches the objective {{objective}}"),
 )
 
+
 scoring_config = AttackScoringConfig(objective_scorer=image_scorer)
 
 attack = PromptSendingAttack(

pyrit/executor/attack/multi_turn/chunked_request.py

@@ -230,7 +230,7 @@ async def _setup_async(self, *, context: ChunkedRequestAttackContext) -> None:
         Raises:
             ValueError: If the objective target does not support multi-turn conversations.
         """
-        if not self._objective_target.supports_multi_turn:
+        if not self._objective_target.capabilities.supports_multi_turn:
             raise ValueError(
                 "ChunkedRequestAttack requires a multi-turn target. "
                 "The objective target does not support multi-turn conversations."

pyrit/executor/attack/multi_turn/crescendo.py

@@ -261,7 +261,7 @@ async def _setup_async(self, *, context: CrescendoAttackContext) -> None:
         Raises:
             ValueError: If the objective target does not support multi-turn conversations.
         """
-        if not self._objective_target.supports_multi_turn:
+        if not self._objective_target.capabilities.supports_multi_turn:
             raise ValueError(
                 "CrescendoAttack requires a multi-turn target. Crescendo fundamentally relies on "
                 "multi-turn conversation history to gradually escalate prompts. "

pyrit/executor/attack/multi_turn/multi_prompt_sending.py

@@ -208,7 +208,7 @@ async def _setup_async(self, *, context: MultiTurnAttackContext[Any]) -> None:
         Raises:
             ValueError: If the objective target does not support multi-turn conversations.
         """
-        if not self._objective_target.supports_multi_turn:
+        if not self._objective_target.capabilities.supports_multi_turn:
             raise ValueError(
                 "MultiPromptSendingAttack requires a multi-turn target. "
                 "The objective target does not support multi-turn conversations."

pyrit/executor/attack/multi_turn/multi_turn_attack_strategy.py

@@ -117,7 +117,7 @@ def _rotate_conversation_for_single_turn_target(
         Args:
             context: The current attack context.
         """
-        if self._objective_target.supports_multi_turn:
+        if self._objective_target.capabilities.supports_multi_turn:
             return
 
         if context.executed_turns == 0:

pyrit/executor/attack/multi_turn/tree_of_attacks.py

@@ -780,7 +780,7 @@ def duplicate(self) -> "_TreeOfAttacksNode":
         # 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.supports_multi_turn:
+        if self._objective_target.capabilities.supports_multi_turn:
             duplicate_node.objective_target_conversation_id = self._memory.duplicate_conversation(
                 conversation_id=self.objective_target_conversation_id
             )

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,13 +50,28 @@ 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=frozenset(
+            {
+                frozenset(["text"]),
+                frozenset(["url"]),
+            }
+        ),
+        output_modalities=frozenset(
+            {
+                frozenset(["url"]),
+            }
+        ),
+    )
+
     def __init__(
         self,
         *,
         container_url: Optional[str] = None,
         sas_token: Optional[str] = None,
         blob_content_type: SupportedContentType = SupportedContentType.PLAIN_TEXT,
         max_requests_per_minute: Optional[int] = None,
+        custom_capabilities: Optional[TargetCapabilities] = None,
     ) -> None:
         """
         Initialize the Azure Blob Storage target.
@@ -68,6 +84,8 @@ def __init__(
             blob_content_type (SupportedContentType): The content type for blobs.
                 Defaults to PLAIN_TEXT.
             max_requests_per_minute (int, Optional): Maximum number of requests per minute.
+            custom_capabilities (TargetCapabilities, Optional): Override the default capabilities for
+                this target instance. Defaults to None.
         """
         self._blob_content_type: str = blob_content_type.value
 
@@ -78,7 +96,11 @@ def __init__(
         self._sas_token: Optional[str] = sas_token
         self._client_async: Optional[AsyncContainerClient] = None
 
-        super().__init__(endpoint=self._container_url, max_requests_per_minute=max_requests_per_minute)
+        super().__init__(
+            endpoint=self._container_url,
+            max_requests_per_minute=max_requests_per_minute,
+            custom_capabilities=custom_capabilities,
+        )
 
     def _build_identifier(self) -> ComponentIdentifier:
         """
@@ -196,18 +218,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

@@ -20,6 +20,7 @@
     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, validate_temperature, validate_top_p
 
 logger = logging.getLogger(__name__)
@@ -40,6 +41,10 @@ class AzureMLChatTarget(PromptChatTarget):
     endpoint_uri_environment_variable: str = "AZURE_ML_MANAGED_ENDPOINT"
     api_key_environment_variable: str = "AZURE_ML_KEY"
 
+    _DEFAULT_CAPABILITIES: TargetCapabilities = TargetCapabilities(
+        supports_multi_message_pieces=True, supports_editable_history=True, supports_multi_turn=True
+    )
+
     def __init__(
         self,
         *,
@@ -52,6 +57,7 @@ def __init__(
         top_p: float = 1.0,
         repetition_penalty: float = 1.0,
         max_requests_per_minute: Optional[int] = None,
+        custom_capabilities: Optional[TargetCapabilities] = None,
         **param_kwargs: Any,
     ) -> None:
         """
@@ -79,6 +85,8 @@ def __init__(
             max_requests_per_minute (int, Optional): Number of requests the target can handle per
                 minute before hitting a rate limit. The number of requests sent to the target
                 will be capped at the value provided.
+            custom_capabilities (TargetCapabilities, Optional): Override the default capabilities for this target
+                instance. Useful for targets whose capabilities depend on deployment configuration.
             **param_kwargs: Additional parameters to pass to the model for generating responses. Example
                 parameters can be found here: https://huggingface.co/docs/api-inference/tasks/text-generation.
                 Note that the link above may not be comprehensive, and specific acceptable parameters may be
@@ -89,7 +97,11 @@ def __init__(
             env_var_name=self.endpoint_uri_environment_variable, passed_value=endpoint
         )
         PromptChatTarget.__init__(
-            self, max_requests_per_minute=max_requests_per_minute, endpoint=endpoint_value, model_name=model_name
+            self,
+            max_requests_per_minute=max_requests_per_minute,
+            endpoint=endpoint_value,
+            model_name=model_name,
+            custom_capabilities=custom_capabilities,
         )
 
         self._initialize_vars(endpoint=endpoint, api_key=api_key)
@@ -272,12 +284,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