FEAT: Runtime capability discovery for prompt targets

doc/code/targets/0_prompt_targets.md

@@ -25,7 +25,7 @@ A `PromptTarget` is a generic place to send a prompt. With PyRIT, the idea is th
 
 With some algorithms, you want to send a prompt, set a system prompt, and modify conversation history (including PAIR [@chao2023pair], TAP [@mehrotra2023tap], and flip attack [@li2024flipattack]). These algorithms require a target whose [`TargetCapabilities`](#target-capabilities) declare both `supports_multi_turn=True` and `supports_editable_history=True` — i.e. you can modify a conversation history. Consumers express this requirement via `CHAT_TARGET_REQUIREMENTS` and validate it against `target.configuration` at construction time. See [Target Capabilities](#target-capabilities) below for the full list of capabilities and how they compose into a `TargetConfiguration`.
 
-Note: The previous `PromptChatTarget` class is **deprecated** as of v0.13.0 and will be removed in v0.15.0. Use `PromptTarget` directly with a `TargetConfiguration` declaring `supports_multi_turn=True` and `supports_editable_history=True`. See [Target Capabilities](#target-capabilities) for details.
+Note: The previous `PromptChatTarget` class is **deprecated** as of v0.14.0 and will be removed in v0.16.0. Use `PromptTarget` directly with a `TargetConfiguration` declaring `supports_multi_turn=True` and `supports_editable_history=True`. See [Target Capabilities](#target-capabilities) for details.
 
 
 Here are some examples:
@@ -107,6 +107,27 @@ target = MyHTTPTarget(custom_configuration=config, ...)
 
 The full implementation lives in [`pyrit/prompt_target/common/target_capabilities.py`](https://github.com/microsoft/PyRIT/blob/main/pyrit/prompt_target/common/target_capabilities.py) and [`pyrit/prompt_target/common/target_configuration.py`](https://github.com/microsoft/PyRIT/blob/main/pyrit/prompt_target/common/target_configuration.py). For runnable examples — inspecting capabilities on a real target, comparing known model profiles, and `ADAPT` vs `RAISE` in action — see [Target Capabilities](./6_1_target_capabilities.ipynb).
 
+### Discovering live target capabilities
+
+Declared capabilities describe what a target *should* support. For deployments where actual behavior is uncertain — custom OpenAI-compatible endpoints, gateways that strip features, models whose support drifts — you can probe what the target *actually* accepts at runtime:
+
+```python
+from pyrit.prompt_target import (
+    discover_target_capabilities_async,
+    discover_target_async,
+    discover_target_modalities_async,
+)
+
+# Probe a single dimension:
+queried_caps = await discover_target_capabilities_async(target=target)
+queried_modalities = await discover_target_modalities_async(target=target)
+
+# Or do both at once and get a best-effort TargetCapabilities back:
+queried = await discover_target_async(target=target)
+```
+
+Each probe sends a minimal request (bounded by `per_probe_timeout_s`, default 30s, with one retry on transient errors) and only marks a capability or modality as supported if the call returns cleanly. `discover_target_async` returns a merged view: probed where possible, declared where probing is unavailable or out of scope. "Supported" here means *the request was accepted* — a target that silently ignores a system prompt or `response_format` directive is still reported as supporting it, so validate response content out of band when the distinction matters. These functions are not safe to call concurrently with other operations on the same target instance: they temporarily mutate `target._configuration` and write probe rows to memory (rows are tagged with `prompt_metadata["capability_probe"] == "1"` for filtering). See [Target Capabilities](./6_1_target_capabilities.ipynb) for runnable examples.
+
 ## Multi-Modal Targets
 
 Like most of PyRIT, targets can be multi-modal.

doc/code/targets/6_1_target_capabilities.ipynb

@@ -53,13 +53,7 @@
      "name": "stdout",
      "output_type": "stream",
      "text": [
-      "No new upgrade operations detected.\n"
-     ]
-    },
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
+      "No new upgrade operations detected.\n",
       "supports_multi_turn:         True\n",
       "supports_editable_history:   True\n",
       "supports_system_prompt:      True\n",
@@ -382,35 +376,94 @@
     }
    ],
    "source": [
-    "from pyrit.models import Message\n",
-    "\n",
-    "conversation = [\n",
-    "    Message.from_prompt(prompt=\"What is the capital of France?\", role=\"user\"),\n",
-    "    Message.from_prompt(prompt=\"Paris.\", role=\"assistant\"),\n",
-    "    Message.from_prompt(prompt=\"And of Germany?\", role=\"user\"),\n",
-    "]\n",
-    "\n",
-    "normalized = await adapt_target.configuration.normalize_async(messages=conversation)  # type: ignore\n",
-    "print(f\"original turns:   {len(conversation)}\")\n",
-    "print(f\"normalized turns: {len(normalized)}\")\n",
-    "print(\"flattened text:\")\n",
-    "print(normalized[-1].message_pieces[0].original_value)"
+    "from unittest.mock import AsyncMock\n",
+    "\n",
+    "from pyrit.models import MessagePiece\n",
+    "from pyrit.prompt_target import (\n",
+    "    discover_target_async,\n",
+    "    discover_target_capabilities_async,\n",
+    "    discover_target_modalities_async,\n",
+    ")\n",
+    "\n",
+    "\n",
+    "def _ok_response():\n",
+    "    return [\n",
+    "        Message(\n",
+    "            [\n",
+    "                MessagePiece(\n",
+    "                    role=\"assistant\",\n",
+    "                    original_value=\"ok\",\n",
+    "                    original_value_data_type=\"text\",\n",
+    "                    conversation_id=\"probe\",\n",
+    "                    response_error=\"none\",\n",
+    "                )\n",
+    "            ]\n",
+    "        )\n",
+    "    ]\n",
+    "\n",
+    "\n",
+    "probe_target = OpenAIChatTarget(model_name=\"gpt-4o\", endpoint=\"https://example.invalid/\", api_key=\"sk-not-a-real-key\")\n",
+    "probe_target._send_prompt_to_target_async = AsyncMock(return_value=_ok_response())  # type: ignore[method-assign]\n",
+    "\n",
+    "queried = await discover_target_capabilities_async(target=probe_target, per_probe_timeout_s=5.0)  # type: ignore\n",
+    "print(\"queried capabilities:\")\n",
+    "for capability in sorted(queried, key=lambda c: c.value):\n",
+    "    print(f\"  - {capability.value}\")"
    ]
   },
   {
    "cell_type": "markdown",
    "id": "15",
    "metadata": {},
    "source": [
-    "By contrast, the `RAISE` configuration validates eagerly: any consumer requiring `MULTI_TURN` will\n",
-    "get a `ValueError` before a single prompt is sent."
+    "To narrow the probe to specific capabilities (faster, fewer calls), pass `capabilities=`:\n",
+    "\n",
+    "```python\n",
+    "from pyrit.prompt_target.common.target_capabilities import CapabilityName\n",
+    "\n",
+    "queried = await query_target_capabilities_async(\n",
+    "    target=target,\n",
+    "    capabilities=[CapabilityName.JSON_SCHEMA, CapabilityName.SYSTEM_PROMPT],\n",
+    ")\n",
+    "```\n",
+    "\n",
+    "If you only care about accepted input combinations, call\n",
+    "`query_target_modalities_async` directly. The example below uses the\n",
+    "packaged default probe assets for the non-text modalities PyRIT ships.\n",
+    "Pass `test_assets=` only when you want to override those defaults or probe\n",
+    "a modality without a packaged asset.\n",
+    "\n",
+    "`query_target_async` is the most common entry point: it runs both the capability and modality\n",
+    "probes and assembles a best-effort `TargetCapabilities` you can drop into a\n",
+    "`TargetConfiguration`, so the rest of PyRIT operates on probed values where available and\n",
+    "declared values otherwise."
    ]
   },
   {
    "cell_type": "code",
    "execution_count": null,
    "id": "16",
    "metadata": {},
+   "outputs": [],
+   "source": [
+    "probe_target._send_prompt_to_target_async = AsyncMock(return_value=_ok_response())  # type: ignore[method-assign]\n",
+    "\n",
+    "queried_modalities = await discover_target_modalities_async(\n",
+    "    target=probe_target,\n",
+    "    test_modalities={frozenset({\"text\"}), frozenset({\"text\", \"image_path\"})},\n",
+    "    per_probe_timeout_s=5.0,\n",
+    ")  # type: ignore\n",
+    "\n",
+    "print(\"query_target_modalities_async result:\")\n",
+    "for combination in sorted(sorted(m) for m in queried_modalities):\n",
+    "    print(f\"  - {combination}\")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "17",
+   "metadata": {},
    "outputs": [
     {
      "name": "stdout",
@@ -429,7 +482,7 @@
   },
   {
    "cell_type": "markdown",
-   "id": "17",
+   "id": "18",
    "metadata": {},
    "source": [
     "## 6. Non-adaptable capabilities\n",
@@ -443,7 +496,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "18",
+   "id": "19",
    "metadata": {},
    "outputs": [
     {
@@ -462,12 +515,208 @@
     "try:\n",
     "    no_editable_history.ensure_can_handle(capability=CapabilityName.EDITABLE_HISTORY)\n",
     "except ValueError as exc:\n",
-    "    print(exc)\n",
-    "# ---"
+    "    print(exc)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "20",
+   "metadata": {},
+   "source": [
+    "## 7. Querying live target capabilities\n",
+    "\n",
+    "Declared capabilities describe what a target *should* support. For deployments where the actual\n",
+    "behavior is uncertain — custom OpenAI-compatible endpoints, gateways that strip features, models\n",
+    "whose support drifts over time — you can probe what the target *actually* accepts at runtime with\n",
+    "`query_target_capabilities_async`, `query_target_modalities_async`, or the convenience wrapper\n",
+    "`query_target_async` that runs both and returns a best-effort `TargetCapabilities`.\n",
+    "\n",
+    "`query_target_capabilities_async` walks each capability that has a registered probe (currently\n",
+    "`SYSTEM_PROMPT`, `MULTI_MESSAGE_PIECES`, `MULTI_TURN`, `JSON_OUTPUT`, `JSON_SCHEMA`), sends a\n",
+    "minimal request, and includes the capability in the returned set only if the call succeeds.\n",
+    "During probing the target's configuration is temporarily replaced with a permissive one so\n",
+    "`ensure_can_handle` does not short-circuit a probe for a capability the target declares as\n",
+    "unsupported. The original configuration is restored before the function returns.\n",
+    "\n",
+    "`query_target_modalities_async` does the same for input modality combinations declared in\n",
+    "`capabilities.input_modalities`, sending a small payload built from optional `test_assets`.\n",
+    "\n",
+    "Each probe call is bounded by `per_probe_timeout_s` (default 30s) and is retried once on\n",
+    "transient errors before being declared failed. `query_target_async` returns a merged view:\n",
+    "probed where possible, declared where probing is unavailable or out of scope. \"Supported\" here\n",
+    "means *the request was accepted* — a target that silently ignores a system prompt or\n",
+    "`response_format` directive will still be reported as supporting that capability.\n",
+    "\n",
+    "These functions are **not safe to call concurrently** with other operations on the same target\n",
+    "instance: they temporarily mutate `target._configuration` and write probe rows to\n",
+    "`target._memory`. Probe-written memory rows are tagged with\n",
+    "`prompt_metadata[\"capability_probe\"] == \"1\"` so consumers can filter them.\n",
+    "\n",
+    "Typical usage against a real endpoint:\n",
+    "\n",
+    "```python\n",
+    "from pyrit.prompt_target import query_target_async\n",
+    "\n",
+    "queried = await query_target_async(target=target)\n",
+    "print(queried)\n",
+    "```\n",
+    "\n",
+    "Below we mock the target's underlying transport (`_send_prompt_to_target_async`) so the notebook\n",
+    "stays self-contained — the result shape is the same as a live run. We mock the protected method\n",
+    "rather than `send_prompt_async` so the probe still exercises the real validation and memory\n",
+    "pipeline."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "21",
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "queried capabilities:\n",
+      "  - supports_editable_history\n",
+      "  - supports_json_output\n",
+      "  - supports_json_schema\n",
+      "  - supports_multi_message_pieces\n",
+      "  - supports_multi_turn\n",
+      "  - supports_system_prompt\n"
+     ]
+    }
+   ],
+   "source": [
+    "from unittest.mock import AsyncMock\n",
+    "\n",
+    "from pyrit.models import MessagePiece\n",
+    "from pyrit.prompt_target import (\n",
+    "    discover_target_async,\n",
+    "    discover_target_capabilities_async,\n",
+    ")\n",
+    "\n",
+    "\n",
+    "def _ok_response():\n",
+    "    return [\n",
+    "        Message(\n",
+    "            [\n",
+    "                MessagePiece(\n",
+    "                    role=\"assistant\",\n",
+    "                    original_value=\"ok\",\n",
+    "                    original_value_data_type=\"text\",\n",
+    "                    conversation_id=\"probe\",\n",
+    "                    response_error=\"none\",\n",
+    "                )\n",
+    "            ]\n",
+    "        )\n",
+    "    ]\n",
+    "\n",
+    "\n",
+    "probe_target = OpenAIChatTarget(model_name=\"gpt-4o\", endpoint=\"https://example.invalid/\", api_key=\"sk-not-a-real-key\")\n",
+    "probe_target._send_prompt_to_target_async = AsyncMock(return_value=_ok_response())  # type: ignore[method-assign]\n",
+    "\n",
+    "queried = await discover_target_capabilities_async(target=probe_target, per_probe_timeout_s=5.0)  # type: ignore\n",
+    "print(\"queried capabilities:\")\n",
+    "for capability in sorted(queried, key=lambda c: c.value):\n",
+    "    print(f\"  - {capability.value}\")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "22",
+   "metadata": {},
+   "source": [
+    "To narrow the probe to specific capabilities (faster, fewer calls), pass `capabilities=`:\n",
+    "\n",
+    "```python\n",
+    "from pyrit.prompt_target.common.target_capabilities import CapabilityName\n",
+    "\n",
+    "queried = await query_target_capabilities_async(\n",
+    "    target=target,\n",
+    "    capabilities=[CapabilityName.JSON_SCHEMA, CapabilityName.SYSTEM_PROMPT],\n",
+    ")\n",
+    "```\n",
+    "\n",
+    "`query_target_async` is the most common entry point: it runs both the capability and modality\n",
+    "probes and assembles a `TargetCapabilities` you can drop straight into a `TargetConfiguration`,\n",
+    "so the rest of PyRIT (attacks, scorers, the normalization pipeline) operates on capabilities\n",
+    "that have been observed to work end-to-end."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "23",
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "query_target_async result:\n",
+      "  supports_multi_turn:           True\n",
+      "  supports_system_prompt:        True\n",
+      "  supports_multi_message_pieces: True\n",
+      "  supports_json_output:          True\n",
+      "  supports_json_schema:          True\n",
+      "  input_modalities:              [['text']]\n"
+     ]
+    }
+   ],
+   "source": [
+    "probe_target._send_prompt_to_target_async = AsyncMock(return_value=_ok_response())  # type: ignore[method-assign]\n",
+    "\n",
+    "queried_caps = await discover_target_async(target=probe_target, per_probe_timeout_s=5.0)  # type: ignore\n",
+    "print(\"query_target_async result:\")\n",
+    "print(f\"  supports_multi_turn:           {queried_caps.supports_multi_turn}\")\n",
+    "print(f\"  supports_system_prompt:        {queried_caps.supports_system_prompt}\")\n",
+    "print(f\"  supports_multi_message_pieces: {queried_caps.supports_multi_message_pieces}\")\n",
+    "print(f\"  supports_json_output:          {queried_caps.supports_json_output}\")\n",
+    "print(f\"  supports_json_schema:          {queried_caps.supports_json_schema}\")\n",
+    "print(f\"  input_modalities:              {sorted(sorted(m) for m in queried_caps.input_modalities)}\")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "24",
+   "metadata": {},
+   "source": [
+    "### Discovering undeclared modalities\n",
+    "\n",
+    "By default `query_target_async` only probes modality combinations the target already\n",
+    "**declares** in `capabilities.input_modalities`. For an OpenAI-compatible endpoint that\n",
+    "claims text-only but might actually accept images, pass `test_modalities=` explicitly to\n",
+    "probe combinations beyond the declared baseline. Provide `test_assets=` as well if you need\n",
+    "to override the packaged defaults or probe a modality without one:\n",
+    "\n",
+    "```python\n",
+    "queried = await query_target_async(\n",
+    "    target=target,\n",
+    "    test_modalities={frozenset({\"text\"}), frozenset({\"text\", \"image_path\"})},\n",
+    "    test_assets={\"image_path\": \"/path/to/test_image.png\"},\n",
+    ")\n",
+    "```\n",
+    "\n",
+    "Similarly, when narrowing the probe set with `capabilities=`, capabilities NOT in the\n",
+    "narrowed set are copied from the target's declared values rather than being reset to\n",
+    "`False` — narrowing controls *what is re-queried*, not what the returned dataclass\n",
+    "reports. This makes incremental probing safe:\n",
+    "\n",
+    "```python\n",
+    "# Re-query only JSON support; other declared flags pass through unchanged.\n",
+    "queried = await query_target_async(\n",
+    "    target=target,\n",
+    "    capabilities={CapabilityName.JSON_OUTPUT, CapabilityName.JSON_SCHEMA},\n",
+    ")\n",
+    "```"
    ]
   }
  ],
  "metadata": {
+  "jupytext": {
+   "main_language": "python"
+  },
   "language_info": {
    "codemirror_mode": {
     "name": "ipython",