diff --git a/doc/code/targets/0_prompt_targets.md b/doc/code/targets/0_prompt_targets.md index e00983f76..9b1b5dab3 100644 --- a/doc/code/targets/0_prompt_targets.md +++ b/doc/code/targets/0_prompt_targets.md @@ -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). +### Querying 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 ( + query_target_capabilities_async, + verify_target_async, + verify_target_modalities_async, +) + +# Probe a single dimension: +verified_caps = await query_target_capabilities_async(target=target) +verified_modalities = await verify_target_modalities_async(target=target) + +# Or do both at once and get a populated TargetCapabilities back: +verified = await verify_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. "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. diff --git a/doc/code/targets/6_1_target_capabilities.ipynb b/doc/code/targets/6_1_target_capabilities.ipynb index 18c190206..11393fc1a 100644 --- a/doc/code/targets/6_1_target_capabilities.ipynb +++ b/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", @@ -462,8 +456,199 @@ "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": "19", + "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`, `verify_target_modalities_async`, or the convenience wrapper\n", + "`verify_target_async` that runs both and returns a populated `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", + "`verify_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. \"Supported\" here means *the request was\n", + "accepted* — a target that silently ignores a system prompt or `response_format` directive will\n", + "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 verify_target_async\n", + "\n", + "verified = await verify_target_async(target=target)\n", + "print(verified)\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": "20", + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "verified 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", + " query_target_capabilities_async,\n", + " verify_target_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", + "verified = await query_target_capabilities_async(target=probe_target, per_probe_timeout_s=5.0) # type: ignore\n", + "print(\"verified capabilities:\")\n", + "for capability in sorted(verified, key=lambda c: c.value):\n", + " print(f\" - {capability.value}\")" + ] + }, + { + "cell_type": "markdown", + "id": "21", + "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", + "verified = await query_target_capabilities_async(\n", + " target=target,\n", + " capabilities=[CapabilityName.JSON_SCHEMA, CapabilityName.SYSTEM_PROMPT],\n", + ")\n", + "```\n", + "\n", + "`verify_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": "22", + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "verify_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", + "verified_caps = await verify_target_async(target=probe_target, per_probe_timeout_s=5.0) # type: ignore\n", + "print(\"verify_target_async result:\")\n", + "print(f\" supports_multi_turn: {verified_caps.supports_multi_turn}\")\n", + "print(f\" supports_system_prompt: {verified_caps.supports_system_prompt}\")\n", + "print(f\" supports_multi_message_pieces: {verified_caps.supports_multi_message_pieces}\")\n", + "print(f\" supports_json_output: {verified_caps.supports_json_output}\")\n", + "print(f\" supports_json_schema: {verified_caps.supports_json_schema}\")\n", + "print(f\" input_modalities: {sorted(sorted(m) for m in verified_caps.input_modalities)}\")" + ] + }, + { + "cell_type": "markdown", + "id": "23", + "metadata": {}, + "source": [ + "### Discovering undeclared modalities\n", + "\n", + "By default `verify_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=` (and matching\n", + "`test_assets=`) explicitly to probe combinations beyond the declared baseline:\n", + "\n", + "```python\n", + "verified = await verify_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-verified*, not what the returned dataclass\n", + "reports. This makes incremental probing safe:\n", + "\n", + "```python\n", + "# Re-verify only JSON support; other declared flags pass through unchanged.\n", + "verified = await verify_target_async(\n", + " target=target,\n", + " capabilities={CapabilityName.JSON_OUTPUT, CapabilityName.JSON_SCHEMA},\n", + ")\n", + "```" ] } ], diff --git a/doc/code/targets/6_1_target_capabilities.py b/doc/code/targets/6_1_target_capabilities.py index 985374357..485bef347 100644 --- a/doc/code/targets/6_1_target_capabilities.py +++ b/doc/code/targets/6_1_target_capabilities.py @@ -245,4 +245,138 @@ no_editable_history.ensure_can_handle(capability=CapabilityName.EDITABLE_HISTORY) except ValueError as exc: print(exc) -# --- + +# %% [markdown] +# ## 7. Querying live target capabilities +# +# Declared capabilities describe what a target *should* support. For deployments where the actual +# behavior is uncertain — custom OpenAI-compatible endpoints, gateways that strip features, models +# whose support drifts over time — you can probe what the target *actually* accepts at runtime with +# `query_target_capabilities_async`, `verify_target_modalities_async`, or the convenience wrapper +# `verify_target_async` that runs both and returns a populated `TargetCapabilities`. +# +# `query_target_capabilities_async` walks each capability that has a registered probe (currently +# `SYSTEM_PROMPT`, `MULTI_MESSAGE_PIECES`, `MULTI_TURN`, `JSON_OUTPUT`, `JSON_SCHEMA`), sends a +# minimal request, and includes the capability in the returned set only if the call succeeds. +# During probing the target's configuration is temporarily replaced with a permissive one so +# `ensure_can_handle` does not short-circuit a probe for a capability the target declares as +# unsupported. The original configuration is restored before the function returns. +# +# `verify_target_modalities_async` does the same for input modality combinations declared in +# `capabilities.input_modalities`, sending a small payload built from optional `test_assets`. +# +# Each probe call is bounded by `per_probe_timeout_s` (default 30s) and is retried once on +# transient errors before being declared failed. "Supported" here means *the request was +# accepted* — a target that silently ignores a system prompt or `response_format` directive will +# still be reported as supporting that capability. +# +# 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 +# `target._memory`. Probe-written memory rows are tagged with +# `prompt_metadata["capability_probe"] == "1"` so consumers can filter them. +# +# Typical usage against a real endpoint: +# +# ```python +# from pyrit.prompt_target import verify_target_async +# +# verified = await verify_target_async(target=target) +# print(verified) +# ``` +# +# Below we mock the target's underlying transport (`_send_prompt_to_target_async`) so the notebook +# stays self-contained — the result shape is the same as a live run. We mock the protected method +# rather than `send_prompt_async` so the probe still exercises the real validation and memory +# pipeline. + +# %% +from unittest.mock import AsyncMock + +from pyrit.models import MessagePiece +from pyrit.prompt_target import ( + query_target_capabilities_async, + verify_target_async, +) + + +def _ok_response(): + return [ + Message( + [ + MessagePiece( + role="assistant", + original_value="ok", + original_value_data_type="text", + conversation_id="probe", + response_error="none", + ) + ] + ) + ] + + +probe_target = OpenAIChatTarget(model_name="gpt-4o", endpoint="https://example.invalid/", api_key="sk-not-a-real-key") +probe_target._send_prompt_to_target_async = AsyncMock(return_value=_ok_response()) # type: ignore[method-assign] + +verified = await query_target_capabilities_async(target=probe_target, per_probe_timeout_s=5.0) # type: ignore +print("verified capabilities:") +for capability in sorted(verified, key=lambda c: c.value): + print(f" - {capability.value}") + +# %% [markdown] +# To narrow the probe to specific capabilities (faster, fewer calls), pass `capabilities=`: +# +# ```python +# from pyrit.prompt_target.common.target_capabilities import CapabilityName +# +# verified = await query_target_capabilities_async( +# target=target, +# capabilities=[CapabilityName.JSON_SCHEMA, CapabilityName.SYSTEM_PROMPT], +# ) +# ``` +# +# `verify_target_async` is the most common entry point: it runs both the capability and modality +# probes and assembles a `TargetCapabilities` you can drop straight into a `TargetConfiguration`, +# so the rest of PyRIT (attacks, scorers, the normalization pipeline) operates on capabilities +# that have been observed to work end-to-end. + +# %% +probe_target._send_prompt_to_target_async = AsyncMock(return_value=_ok_response()) # type: ignore[method-assign] + +verified_caps = await verify_target_async(target=probe_target, per_probe_timeout_s=5.0) # type: ignore +print("verify_target_async result:") +print(f" supports_multi_turn: {verified_caps.supports_multi_turn}") +print(f" supports_system_prompt: {verified_caps.supports_system_prompt}") +print(f" supports_multi_message_pieces: {verified_caps.supports_multi_message_pieces}") +print(f" supports_json_output: {verified_caps.supports_json_output}") +print(f" supports_json_schema: {verified_caps.supports_json_schema}") +print(f" input_modalities: {sorted(sorted(m) for m in verified_caps.input_modalities)}") + +# %% [markdown] +# ### Discovering undeclared modalities +# +# By default `verify_target_async` only probes modality combinations the target already +# **declares** in `capabilities.input_modalities`. For an OpenAI-compatible endpoint that +# claims text-only but might actually accept images, pass `test_modalities=` (and matching +# `test_assets=`) explicitly to probe combinations beyond the declared baseline: +# +# ```python +# verified = await verify_target_async( +# target=target, +# test_modalities={frozenset({"text"}), frozenset({"text", "image_path"})}, +# test_assets={"image_path": "/path/to/test_image.png"}, +# ) +# ``` +# +# Similarly, when narrowing the probe set with `capabilities=`, capabilities NOT in the +# narrowed set are copied from the target's declared values rather than being reset to +# `False` — narrowing controls *what is re-verified*, not what the returned dataclass +# reports. This makes incremental probing safe: +# +# ```python +# # Re-verify only JSON support; other declared flags pass through unchanged. +# verified = await verify_target_async( +# target=target, +# capabilities={CapabilityName.JSON_OUTPUT, CapabilityName.JSON_SCHEMA}, +# ) +# ``` diff --git a/pyrit/prompt_target/__init__.py b/pyrit/prompt_target/__init__.py index 489fe3490..ef682b2a4 100644 --- a/pyrit/prompt_target/__init__.py +++ b/pyrit/prompt_target/__init__.py @@ -16,6 +16,11 @@ from pyrit.prompt_target.common.conversation_normalization_pipeline import ConversationNormalizationPipeline from pyrit.prompt_target.common.prompt_chat_target import PromptChatTarget from pyrit.prompt_target.common.prompt_target import PromptTarget +from pyrit.prompt_target.common.query_target_capabilities import ( + query_target_capabilities_async, + verify_target_async, + verify_target_modalities_async, +) from pyrit.prompt_target.common.target_capabilities import ( CapabilityHandlingPolicy, CapabilityName, @@ -97,11 +102,14 @@ def __getattr__(name: str) -> object: "PromptChatTarget", "PromptShieldTarget", "PromptTarget", + "query_target_capabilities_async", "RealtimeTarget", "TargetCapabilities", "TargetConfiguration", "TargetRequirements", "UnsupportedCapabilityBehavior", "TextTarget", + "verify_target_async", + "verify_target_modalities_async", "WebSocketCopilotTarget", ] diff --git a/pyrit/prompt_target/common/query_target_capabilities.py b/pyrit/prompt_target/common/query_target_capabilities.py new file mode 100644 index 000000000..9a8d18fa6 --- /dev/null +++ b/pyrit/prompt_target/common/query_target_capabilities.py @@ -0,0 +1,775 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT license. + +""" +Runtime capability and modality discovery for prompt targets. + +This module exposes two complementary probes: + +* :func:`query_target_capabilities_async` probes the boolean capability flags + defined on :class:`TargetCapabilities` (e.g. ``supports_system_prompt``, + ``supports_multi_message_pieces``). For each capability that has a probe + defined, a minimal request is sent to the target. If the request succeeds, + the capability is included in the returned set. Capabilities without a + registered probe fall back to whatever the target declares via its + :class:`TargetConfiguration`. +* :func:`verify_target_modalities_async` probes which input modality + combinations a target actually supports by sending a minimal test request + for each combination declared in ``TargetCapabilities.input_modalities``. + +.. note:: + Output modality probing is intentionally not provided. Unlike inputs, + output modality is largely a property of the endpoint type (chat models + return text, image models return images, TTS endpoints return audio) + rather than something the caller controls per request, and there is no + PyRIT-level ``response_format=image`` style hint to assert against. + Eliciting non-text output reliably depends on prompt phrasing, costs + real compute per probe, and is prone to false negatives from safety + filters. Trust ``target.capabilities.output_modalities`` as declared. + +.. warning:: + These probes only verify that a request was *accepted* (the call returned + without raising and the response had no error). They cannot detect a + target that silently ignores a feature. For example, an endpoint that + accepts a ``system`` role but discards it, or that accepts a + ``response_format="json"`` hint but returns prose, will be reported as + supporting those capabilities. Treat the returned sets as an upper bound + on actual support and validate response content out of band when the + distinction matters (e.g. parse JSON responses, assert that the model + honored the system prompt). +""" + +import asyncio +import json +import logging +import os +import uuid +from collections.abc import Awaitable, Callable, Iterable, Iterator +from contextlib import contextmanager +from dataclasses import replace + +from pyrit.models import Message, MessagePiece, PromptDataType +from pyrit.prompt_target.common.prompt_target import PromptTarget +from pyrit.prompt_target.common.target_capabilities import ( + CapabilityHandlingPolicy, + CapabilityName, + TargetCapabilities, + UnsupportedCapabilityBehavior, +) +from pyrit.prompt_target.common.target_configuration import TargetConfiguration + +logger = logging.getLogger(__name__) + +# Per-call timeout (seconds) applied to every probe request. Override per-call via +# the ``per_probe_timeout_s`` parameter on the public functions. +DEFAULT_PROBE_TIMEOUT_SECONDS: float = 30.0 + +# Marker stamped onto every MessagePiece this module writes to memory. Consumers +# that aggregate or display memory rows can filter probe-written rows by checking +# ``piece.prompt_metadata.get("capability_probe") == "1"``. Memory does not yet +# expose a delete-by-conversation-id API, so tagging is the cleanup mechanism. +PROBE_METADATA_KEY: str = "capability_probe" +PROBE_METADATA_VALUE: str = "1" + +_CapabilityProbe = Callable[[PromptTarget, float, int], Awaitable[bool]] + + +_PROBE_POLICY = CapabilityHandlingPolicy( + behaviors={ + CapabilityName.MULTI_TURN: UnsupportedCapabilityBehavior.RAISE, + CapabilityName.SYSTEM_PROMPT: UnsupportedCapabilityBehavior.RAISE, + } +) + + +# Every text probe sends a text-only payload. Permissive overrides therefore +# always include this combination so that ``_validate_request``'s per-piece +# data-type check does not reject text probes against text-less targets. +_TEXT_MODALITY: frozenset[frozenset[PromptDataType]] = frozenset({frozenset({"text"})}) + + +@contextmanager +def _permissive_configuration( + *, + target: PromptTarget, + extra_input_modalities: Iterable[frozenset[PromptDataType]] | None = None, +) -> Iterator[None]: + """ + Temporarily replace ``target``'s configuration with one that declares every + boolean capability as natively supported. + + This bypasses :meth:`PromptTarget._validate_request`, which would otherwise + short-circuit probes for capabilities the target declares as unsupported + before any API call is made. The original configuration is restored on exit. + + Args: + target (PromptTarget): The target whose configuration is temporarily replaced. + extra_input_modalities (Iterable[frozenset[PromptDataType]] | None): + Additional modality combinations to include in ``input_modalities`` + during the override. Used by modality probes so that + ``_validate_request``'s per-piece data-type check does not reject + combinations the caller asked us to test but the target does not + yet declare. Defaults to None. + + Yields: + None: Control returns to the ``with`` block while the permissive + configuration is in effect. + """ + original = target.configuration + merged_modalities = original.capabilities.input_modalities | _TEXT_MODALITY + if extra_input_modalities is not None: + merged_modalities = frozenset(merged_modalities | frozenset(extra_input_modalities)) + permissive_caps = replace( + original.capabilities, + supports_multi_turn=True, + supports_multi_message_pieces=True, + supports_json_schema=True, + supports_json_output=True, + supports_editable_history=True, + supports_system_prompt=True, + input_modalities=merged_modalities, + ) + target._configuration = TargetConfiguration( + capabilities=permissive_caps, + policy=_PROBE_POLICY, + ) + try: + yield + finally: + target._configuration = original + + +def _new_conversation_id() -> str: + """ + Generate a unique conversation id for a single capability probe. + + Returns: + str: A conversation id of the form ``"capability-probe-"``. + """ + return f"capability-probe-{uuid.uuid4()}" + + +def _probe_metadata(extra: dict[str, str | int] | None = None) -> dict[str, str | int]: + """Return a fresh ``prompt_metadata`` dict tagged as a capability probe.""" + metadata: dict[str, str | int] = {PROBE_METADATA_KEY: PROBE_METADATA_VALUE} + if extra: + metadata.update(extra) + return metadata + + +def _user_text_piece(*, value: str, conversation_id: str) -> MessagePiece: + """ + Build a single user-role text :class:`MessagePiece` for use in a probe. + + The piece's ``prompt_metadata`` is tagged with :data:`PROBE_METADATA_KEY` + so that consumers aggregating memory can filter out probe-written rows. + + Args: + value (str): The text payload to send. + conversation_id (str): The conversation id to attach to the piece. + + Returns: + MessagePiece: A user-role text piece bound to ``conversation_id``. + """ + return MessagePiece( + role="user", + original_value=value, + original_value_data_type="text", + conversation_id=conversation_id, + prompt_metadata=_probe_metadata(), + ) + + +async def _send_and_check_async( + *, + target: PromptTarget, + message: Message, + timeout_s: float, + retries: int = 1, + label: str = "Capability probe", +) -> bool: + """ + Send ``message`` and report whether the call succeeded cleanly. + + Each attempt is bounded by ``timeout_s``. Exceptions (network errors, + timeouts, validation failures) trigger up to ``retries`` retries before + the probe is declared failed; an explicit error response from the target + is treated as deterministic and never retried. + + Args: + target (PromptTarget): The target to send the probe message to. + message (Message): The probe message to send. + timeout_s (float): Per-attempt timeout in seconds. + retries (int): Number of additional attempts after the first failure. + Only exceptions are retried; a non-error response is final. + Defaults to 1. + label (str): Short label used in log messages. Defaults to + ``"Capability probe"``. + + Returns: + bool: ``True`` iff the call returned without raising and every response + piece reported ``response_error == "none"``; ``False`` otherwise. + Any other ``response_error`` value (``"blocked"``, ``"processing"``, + ``"empty"``, ``"unknown"``) is treated as failure. An empty response + list (or responses with no message pieces) is also treated as a failure. + """ + attempts = max(1, retries + 1) + last_exc: Exception | None = None + for attempt in range(attempts): + try: + responses = await asyncio.wait_for(target.send_prompt_async(message=message), timeout=timeout_s) + except asyncio.TimeoutError: + last_exc = TimeoutError(f"timed out after {timeout_s}s") + logger.debug("%s timed out (attempt %d/%d)", label, attempt + 1, attempts) + continue + except Exception as exc: + last_exc = exc + logger.debug("%s failed (attempt %d/%d): %s", label, attempt + 1, attempts, exc) + continue + + if not responses or not any(r.message_pieces for r in responses): + logger.debug("%s returned an empty response; treating as failure", label) + return False + for response in responses: + for piece in response.message_pieces: + if piece.response_error != "none": + logger.debug("%s returned error response: %s", label, piece.converted_value) + return False + return True + + logger.info("%s exhausted %d attempt(s); last error: %s", label, attempts, last_exc) + return False + + +async def _probe_system_prompt_async(target: PromptTarget, timeout_s: float, retries: int = 1) -> bool: + """ + Probe whether ``target`` accepts a system prompt followed by a user message. + + Writes a system-role :class:`MessagePiece` directly to ``target._memory`` + rather than calling :meth:`PromptTarget.set_system_prompt`. ``set_system_prompt`` + can be overridden by subclasses (e.g. mocks) to do nothing or to perform + extra work, which would mask whether the underlying API actually accepts a + system message. A direct memory write guarantees the probe sees the same + multi-piece, system-then-user payload the target's wire layer would see + via the standard pipeline. + + Args: + target (PromptTarget): The target to probe. + timeout_s (float): Per-attempt timeout in seconds. + + Returns: + bool: ``True`` if the system + user request succeeded; ``False`` otherwise. + """ + conversation_id = _new_conversation_id() + system_piece = MessagePiece( + role="system", + original_value="You are a helpful assistant.", + original_value_data_type="text", + conversation_id=conversation_id, + prompt_metadata=_probe_metadata(), + ) + try: + target._memory.add_message_to_memory(request=Message([system_piece])) + except Exception as exc: + logger.debug("System-prompt probe could not seed system message: %s", exc) + return False + user_piece = _user_text_piece(value="hi", conversation_id=conversation_id) + return await _send_and_check_async( + target=target, + message=Message([user_piece]), + timeout_s=timeout_s, + retries=retries, + label="System-prompt probe", + ) + + +async def _probe_multi_message_pieces_async(target: PromptTarget, timeout_s: float, retries: int = 1) -> bool: + """ + Probe whether ``target`` accepts a single message containing multiple pieces. + + Args: + target (PromptTarget): The target to probe. + timeout_s (float): Per-attempt timeout in seconds. + + Returns: + bool: ``True`` if the multi-piece request succeeded; ``False`` otherwise. + """ + conversation_id = _new_conversation_id() + pieces = [ + _user_text_piece(value="part one", conversation_id=conversation_id), + _user_text_piece(value="part two", conversation_id=conversation_id), + ] + return await _send_and_check_async( + target=target, + message=Message(pieces), + timeout_s=timeout_s, + retries=retries, + label="Multi-message-pieces probe", + ) + + +async def _probe_multi_turn_async(target: PromptTarget, timeout_s: float, retries: int = 1) -> bool: + """ + Probe whether ``target`` accepts a request that includes prior conversation history. + + ``PromptTarget.send_prompt_async`` reads conversation history from memory but + does not write to it (persistence normally happens in the orchestrator + layer). To exercise true multi-turn behavior, this probe: + + 1. Sends an initial user message. + 2. Persists that user message and a synthetic assistant reply directly to + the target's memory under the same ``conversation_id``. + 3. Sends a second user message; ``send_prompt_async`` then fetches the + 2-message history and the target receives a real 3-message + multi-turn payload. + + The synthetic assistant reply's content is irrelevant — we are testing + whether the target's API accepts a multi-turn payload, not whether the + model recalls anything. + + Args: + target (PromptTarget): The target to probe. + timeout_s (float): Per-attempt timeout in seconds. + + Returns: + bool: ``True`` if both turns succeeded; ``False`` if either turn failed. + """ + conversation_id = _new_conversation_id() + first = _user_text_piece(value="My favorite color is blue.", conversation_id=conversation_id) + if not await _send_and_check_async( + target=target, message=Message([first]), timeout_s=timeout_s, retries=retries, label="Multi-turn probe (turn 1)" + ): + return False + + # Seed memory so the second send sees real prior history. + try: + target._memory.add_message_to_memory(request=Message([first])) + assistant_reply = MessagePiece( + role="assistant", + original_value="Got it.", + original_value_data_type="text", + conversation_id=conversation_id, + prompt_metadata=_probe_metadata(), + ).to_message() + target._memory.add_message_to_memory(request=assistant_reply) + except Exception as exc: + logger.debug("Multi-turn probe could not seed conversation history: %s", exc) + return False + + second = _user_text_piece(value="What did I just tell you?", conversation_id=conversation_id) + return await _send_and_check_async( + target=target, + message=Message([second]), + timeout_s=timeout_s, + retries=retries, + label="Multi-turn probe (turn 2)", + ) + + +async def _probe_json_output_async(target: PromptTarget, timeout_s: float, retries: int = 1) -> bool: + """ + Probe whether ``target`` accepts a request asking for JSON-mode output. + + Args: + target (PromptTarget): The target to probe. + timeout_s (float): Per-attempt timeout in seconds. + + Returns: + bool: ``True`` if the JSON-mode request succeeded; ``False`` otherwise. + """ + conversation_id = _new_conversation_id() + piece = MessagePiece( + role="user", + original_value='Respond with a JSON object: {"ok": true}.', + original_value_data_type="text", + conversation_id=conversation_id, + prompt_metadata=_probe_metadata({"response_format": "json"}), + ) + return await _send_and_check_async( + target=target, message=Message([piece]), timeout_s=timeout_s, retries=retries, label="JSON-output probe" + ) + + +async def _probe_json_schema_async(target: PromptTarget, timeout_s: float, retries: int = 1) -> bool: + """ + Probe whether ``target`` accepts a request constrained by a JSON schema. + + Args: + target (PromptTarget): The target to probe. + timeout_s (float): Per-attempt timeout in seconds. + + Returns: + bool: ``True`` if the schema-constrained request succeeded; ``False`` otherwise. + """ + schema = { + "type": "object", + "properties": {"ok": {"type": "boolean"}}, + "required": ["ok"], + "additionalProperties": False, + } + conversation_id = _new_conversation_id() + piece = MessagePiece( + role="user", + original_value='Respond with a JSON object matching the schema: {"ok": true}.', + original_value_data_type="text", + conversation_id=conversation_id, + prompt_metadata=_probe_metadata( + { + "response_format": "json", + "json_schema": json.dumps(schema), + } + ), + ) + return await _send_and_check_async( + target=target, message=Message([piece]), timeout_s=timeout_s, retries=retries, label="JSON-schema probe" + ) + + +# Registry of capabilities that can be verified via a live API call. +# Capabilities not present here fall back to the target's declared support. +_CAPABILITY_PROBES: dict[CapabilityName, _CapabilityProbe] = { + CapabilityName.SYSTEM_PROMPT: _probe_system_prompt_async, + CapabilityName.MULTI_MESSAGE_PIECES: _probe_multi_message_pieces_async, + CapabilityName.MULTI_TURN: _probe_multi_turn_async, + CapabilityName.JSON_OUTPUT: _probe_json_output_async, + CapabilityName.JSON_SCHEMA: _probe_json_schema_async, +} + + +async def query_target_capabilities_async( + *, + target: PromptTarget, + capabilities: Iterable[CapabilityName] | None = None, + per_probe_timeout_s: float = DEFAULT_PROBE_TIMEOUT_SECONDS, + retries: int = 1, +) -> set[CapabilityName]: + """ + Probe ``target`` to determine which capabilities it actually supports. + + For each requested capability that has a registered probe, a minimal + request is sent to the target. The capability is treated as supported + only if the call returns successfully with no error response. For + capabilities without a registered probe, the target's declared + **native** support (``target.capabilities.includes(...)``) is used as + a fallback. We deliberately do *not* consult + ``target.configuration.includes(...)`` here, because that would also + return ``True`` for capabilities the target lacks but PyRIT + ``ADAPT``s via the :class:`CapabilityHandlingPolicy` — and adaptation + is an emulation by PyRIT, not evidence that the target itself supports + the capability. + + .. warning:: + "Supported" here means "the request was accepted", not "the feature + was actually applied". A target that silently ignores a system + prompt, ``response_format``, or schema directive will still be + reported as supporting that capability. Validate response content + out of band when correctness matters. + + .. warning:: + This function is **not safe to call concurrently** with other + operations on the same ``target`` instance. It temporarily mutates + ``target._configuration`` and writes probe rows to ``target._memory``; + concurrent callers may observe the permissive configuration or + interleaved memory rows. Probe-written memory rows are tagged with + ``prompt_metadata["capability_probe"] == "1"`` so consumers can + filter them; memory does not currently expose a delete-by-conversation + API, so probe rows persist for the lifetime of the memory backend. + + During probing, the target's configuration is temporarily replaced with + one that declares every boolean capability as supported, so that + :meth:`PromptTarget._validate_request` does not short-circuit probes for + capabilities the target declares as unsupported. The original + configuration is restored before this function returns. + + Args: + target (PromptTarget): The target to probe. + capabilities (Iterable[CapabilityName] | None): Capabilities to check. + Defaults to every member of :class:`CapabilityName`. + per_probe_timeout_s (float): Per-attempt timeout (seconds) applied to + each probe request. Defaults to + :data:`DEFAULT_PROBE_TIMEOUT_SECONDS`. + retries (int): Number of additional attempts after the first failure + for each probe. Only exceptions/timeouts are retried; an explicit + error response is final. Set to ``0`` to disable retries. + Defaults to 1. + + Returns: + set[CapabilityName]: The capabilities verified to work against the target. + """ + capabilities_to_check: list[CapabilityName] = ( + list(capabilities) if capabilities is not None else list(CapabilityName) + ) + + verified: set[CapabilityName] = set() + with _permissive_configuration(target=target): + for capability in capabilities_to_check: + probe = _CAPABILITY_PROBES.get(capability) + if probe is None: + # No live probe; fall back to whatever the (original) configuration declared. + # We're inside the permissive override, so consult the saved configuration directly. + continue + + try: + if await probe(target, per_probe_timeout_s, retries): + verified.add(capability) + except Exception as exc: + logger.debug("Probe for %s raised: %s", capability.value, exc) + + # Add capabilities without a probe based on the original (now-restored) NATIVE + # support. Using target.capabilities.includes (native flags) rather than + # target.configuration.includes (which also returns True for ADAPT'd capabilities) + # keeps this function's contract honest: we report only what the target itself + # supports, never what PyRIT emulates on top of it. + for capability in capabilities_to_check: + if capability not in _CAPABILITY_PROBES and target.capabilities.includes(capability=capability): + verified.add(capability) + + return verified + + +# --------------------------------------------------------------------------- +# Modality verification +# --------------------------------------------------------------------------- + + +# Default mapping of non-text modalities to test asset paths. Callers can +# override via the ``test_assets`` parameter of +# :func:`verify_target_modalities_async`. Modalities whose assets do not +# exist on disk are skipped (logged and excluded from the result). +DEFAULT_TEST_ASSETS: dict[PromptDataType, str] = {} + + +async def verify_target_modalities_async( + *, + target: PromptTarget, + test_modalities: set[frozenset[PromptDataType]] | None = None, + test_assets: dict[PromptDataType, str] | None = None, + per_probe_timeout_s: float = DEFAULT_PROBE_TIMEOUT_SECONDS, + retries: int = 1, +) -> set[frozenset[PromptDataType]]: + """ + Probe ``target`` to determine which input modality combinations it supports. + + Each combination is exercised with a minimal request built by + :func:`_create_test_message`. A combination is considered supported only + if the request returns successfully with no error response. + + During probing the target's configuration is temporarily replaced with + one that declares every boolean capability as natively supported and + that includes every probed modality combination in ``input_modalities``, + so :meth:`PromptTarget._validate_request` does not short-circuit a probe + before any API call is made. The original configuration is restored + before this function returns. + + .. warning:: + "Supported" here means the target accepted the request. A target + that accepts e.g. an ``image_path`` piece but ignores its content + will still be reported as supporting that modality. + + .. warning:: + This function is **not safe to call concurrently** with other + operations on the same ``target`` instance. It temporarily mutates + ``target._configuration``. + + Args: + target (PromptTarget): The target to probe. + test_modalities (set[frozenset[PromptDataType]] | None): Specific + modality combinations to test. Defaults to the combinations + declared in ``target.capabilities.input_modalities``. + test_assets (dict[PromptDataType, str] | None): Mapping from + non-text modality to a file path used as the probe payload. + Defaults to :data:`DEFAULT_TEST_ASSETS`. Combinations whose + non-text assets are missing on disk are skipped. + per_probe_timeout_s (float): Per-attempt timeout (seconds) applied to + each probe request. Defaults to + :data:`DEFAULT_PROBE_TIMEOUT_SECONDS`. + retries (int): Number of additional attempts after the first failure + for each probe. Only exceptions/timeouts are retried; an explicit + error response is final. Set to ``0`` to disable retries. + Defaults to 1. + + Returns: + set[frozenset[PromptDataType]]: The modality combinations verified + to work against the target. + """ + if test_modalities is None: + declared = target.capabilities.input_modalities + test_modalities = set(declared) + + assets = test_assets if test_assets is not None else DEFAULT_TEST_ASSETS + + verified: set[frozenset[PromptDataType]] = set() + with _permissive_configuration(target=target, extra_input_modalities=test_modalities): + for combination in test_modalities: + try: + message = _create_test_message(modalities=combination, test_assets=assets) + except FileNotFoundError as exc: + logger.info("Skipping modality %s: %s", combination, exc) + continue + except ValueError as exc: + logger.info("Skipping modality %s: %s", combination, exc) + continue + + if await _send_and_check_async( + target=target, + message=message, + timeout_s=per_probe_timeout_s, + retries=retries, + label=f"Modality probe {sorted(combination)}", + ): + verified.add(combination) + + return verified + + +async def verify_target_async( + *, + target: PromptTarget, + per_probe_timeout_s: float = DEFAULT_PROBE_TIMEOUT_SECONDS, + test_modalities: set[frozenset[PromptDataType]] | None = None, + test_assets: dict[PromptDataType, str] | None = None, + capabilities: Iterable[CapabilityName] | None = None, + retries: int = 1, +) -> TargetCapabilities: + """ + Probe both capabilities and modalities and return a combined result. + + Calls :func:`query_target_capabilities_async` and + :func:`verify_target_modalities_async` and returns a + :class:`TargetCapabilities` populated from the verified results, so + callers don't need to assemble the dataclass themselves. + + Boolean capability flags not covered by + :data:`_CAPABILITY_PROBES` (e.g. ``supports_editable_history``) are + copied from ``target.capabilities`` (the target's declared native flags). + When ``capabilities`` narrows the probe set, capabilities not in the + narrowed set are also copied from declared values rather than reset to + ``False`` — narrowing controls *what is re-verified*, not what the + returned dataclass reports. + + .. warning:: + By default ``test_modalities`` is sourced from + ``target.capabilities.input_modalities`` (the target's *declared* + modalities). This means the modality probe cannot discover modalities + the target does not already declare. Pass ``test_modalities=`` (and + matching ``test_assets=``) explicitly to probe combinations beyond + the declared baseline. + + Args: + target (PromptTarget): The target to probe. + per_probe_timeout_s (float): Per-attempt timeout (seconds) applied to + each probe request. + test_modalities (set[frozenset[PromptDataType]] | None): Specific + modality combinations to probe. See + :func:`verify_target_modalities_async`. Defaults to the + target's declared ``input_modalities``. + test_assets (dict[PromptDataType, str] | None): Mapping from non-text + modality to a file path. See :func:`verify_target_modalities_async`. + capabilities (Iterable[CapabilityName] | None): Capabilities to probe. + See :func:`query_target_capabilities_async`. Defaults to every + member of :class:`CapabilityName`. + retries (int): Number of additional attempts after the first failure + for each probe. Only exceptions/timeouts are retried; an explicit + error response is final. Set to ``0`` to disable retries. + Defaults to 1. + + Returns: + TargetCapabilities: A dataclass reflecting verified capabilities and + modalities. ``output_modalities`` is copied from + ``target.capabilities.output_modalities`` because outputs cannot be + verified by sending a request. + """ + verified_caps = await query_target_capabilities_async( + target=target, + capabilities=capabilities, + per_probe_timeout_s=per_probe_timeout_s, + retries=retries, + ) + verified_modalities = await verify_target_modalities_async( + target=target, + test_modalities=test_modalities, + test_assets=test_assets, + per_probe_timeout_s=per_probe_timeout_s, + retries=retries, + ) + + declared = target.capabilities + # When ``capabilities`` narrows the probe set, capabilities NOT in the + # narrowed set were never probed and must fall back to declared values + # rather than being silently reset to False. + probed: set[CapabilityName] = set(capabilities) if capabilities is not None else set(CapabilityName) + + def _resolve(name: CapabilityName) -> bool: + if name in probed: + return name in verified_caps + return bool(getattr(declared, name.value)) + + return TargetCapabilities( + supports_multi_turn=_resolve(CapabilityName.MULTI_TURN), + supports_multi_message_pieces=_resolve(CapabilityName.MULTI_MESSAGE_PIECES), + supports_json_schema=_resolve(CapabilityName.JSON_SCHEMA), + supports_json_output=_resolve(CapabilityName.JSON_OUTPUT), + supports_editable_history=declared.supports_editable_history, + supports_system_prompt=_resolve(CapabilityName.SYSTEM_PROMPT), + input_modalities=frozenset(verified_modalities), + output_modalities=declared.output_modalities, + ) + + +def _create_test_message( + *, + modalities: frozenset[PromptDataType], + test_assets: dict[PromptDataType, str], +) -> Message: + """ + Build a minimal :class:`Message` that exercises ``modalities``. + + Args: + modalities (frozenset[PromptDataType]): The modalities to include. + test_assets (dict[PromptDataType, str]): Mapping from non-text + modality to a file path used for the probe. + + Returns: + Message: A message containing one piece per modality. + + Raises: + FileNotFoundError: If a configured asset path does not exist. + ValueError: If a non-text modality has no configured asset, or if + no pieces could be constructed. + """ + conversation_id = f"modality-probe-{uuid.uuid4()}" + pieces: list[MessagePiece] = [] + + for modality in modalities: + if modality == "text": + pieces.append( + MessagePiece( + role="user", + original_value="test", + original_value_data_type="text", + conversation_id=conversation_id, + prompt_metadata=_probe_metadata(), + ) + ) + continue + + asset_path = test_assets.get(modality) + if asset_path is None: + raise ValueError(f"No test asset configured for modality '{modality}'.") + if not os.path.isfile(asset_path): + raise FileNotFoundError(f"Test asset for modality '{modality}' not found at: {asset_path}") + + pieces.append( + MessagePiece( + role="user", + original_value=asset_path, + original_value_data_type=modality, + conversation_id=conversation_id, + prompt_metadata=_probe_metadata(), + ) + ) + + if not pieces: + raise ValueError(f"Could not create test message for modalities: {modalities}") + + return Message(pieces) diff --git a/tests/unit/prompt_target/test_query_target_capabilities.py b/tests/unit/prompt_target/test_query_target_capabilities.py new file mode 100644 index 000000000..df12f4763 --- /dev/null +++ b/tests/unit/prompt_target/test_query_target_capabilities.py @@ -0,0 +1,763 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT license. + +import asyncio +import json +from pathlib import Path +from unittest.mock import AsyncMock, MagicMock + +import pytest + +from pyrit.models import Message, MessagePiece, PromptDataType +from pyrit.prompt_target.common.prompt_target import PromptTarget +from pyrit.prompt_target.common.query_target_capabilities import ( + _CAPABILITY_PROBES, + _create_test_message, + _permissive_configuration, + query_target_capabilities_async, + verify_target_async, + verify_target_modalities_async, +) +from pyrit.prompt_target.common.target_capabilities import ( + CapabilityName, + TargetCapabilities, +) +from pyrit.prompt_target.common.target_configuration import TargetConfiguration +from tests.unit.mocks import MockPromptTarget + + +class _RealValidationTarget(PromptTarget): + """ + Bare ``PromptTarget`` subclass that does NOT override ``_validate_request``. + + Tests that need to verify ``_permissive_configuration`` actually bypasses + the validation guard use this instead of ``MockPromptTarget`` (which + no-ops ``_validate_request``). + """ + + _DEFAULT_CONFIGURATION: TargetConfiguration = TargetConfiguration( + capabilities=TargetCapabilities(), + ) + + async def _send_prompt_to_target_async(self, *, normalized_conversation: list[Message]) -> list[Message]: + return _ok_response() + + +def _ok_response(*, conversation_id: str = "probe", text: str = "ok") -> list[Message]: + return [ + Message( + [ + MessagePiece( + role="assistant", + original_value=text, + original_value_data_type="text", + conversation_id=conversation_id, + response_error="none", + ) + ] + ) + ] + + +def _error_response(*, conversation_id: str = "probe") -> list[Message]: + return [ + Message( + [ + MessagePiece( + role="assistant", + original_value="blocked", + original_value_data_type="text", + conversation_id=conversation_id, + response_error="blocked", + ) + ] + ) + ] + + +@pytest.mark.usefixtures("patch_central_database") +class TestPermissiveConfiguration: + def test_replaces_and_restores_configuration(self) -> None: + target = MockPromptTarget() + original = target.configuration + + with _permissive_configuration(target=target): + permissive = target.configuration + assert permissive is not original + for capability in CapabilityName: + assert permissive.includes(capability=capability) + + assert target.configuration is original + + def test_restores_on_exception(self) -> None: + target = MockPromptTarget() + original = target.configuration + + with pytest.raises(RuntimeError): + with _permissive_configuration(target=target): + raise RuntimeError("boom") + + assert target.configuration is original + + +@pytest.mark.usefixtures("patch_central_database") +class TestQueryTargetCapabilitiesAsync: + async def test_returns_only_supported_when_all_probes_succeed(self) -> None: + target = MockPromptTarget() + target._send_prompt_to_target_async = AsyncMock(return_value=_ok_response()) # type: ignore[method-assign] + + result = await query_target_capabilities_async(target=target) + + # Every capability with a probe should be in the result. + for capability in _CAPABILITY_PROBES: + assert capability in result + + async def test_excludes_capabilities_when_probe_fails(self) -> None: + target = MockPromptTarget() + target._send_prompt_to_target_async = AsyncMock(side_effect=Exception("nope")) # type: ignore[method-assign] + + result = await query_target_capabilities_async(target=target) + + for capability in _CAPABILITY_PROBES: + assert capability not in result + + async def test_excludes_capabilities_when_response_has_error(self) -> None: + target = MockPromptTarget() + target._send_prompt_to_target_async = AsyncMock(return_value=_error_response()) # type: ignore[method-assign] + + result = await query_target_capabilities_async(target=target) + + for capability in _CAPABILITY_PROBES: + assert capability not in result + + async def test_filters_by_requested_capabilities(self) -> None: + target = MockPromptTarget() + target._send_prompt_to_target_async = AsyncMock(return_value=_ok_response()) # type: ignore[method-assign] + + requested = {CapabilityName.SYSTEM_PROMPT, CapabilityName.MULTI_TURN} + result = await query_target_capabilities_async(target=target, capabilities=requested) + + assert result == requested + + async def test_capability_without_probe_falls_back_to_declared_support(self) -> None: + target = MockPromptTarget() + # Override the configuration so editable_history is declared as supported. + target._configuration = TargetConfiguration( + capabilities=TargetCapabilities(supports_editable_history=True), + ) + target._send_prompt_to_target_async = AsyncMock(return_value=_ok_response()) # type: ignore[method-assign] + + result = await query_target_capabilities_async( + target=target, + capabilities={CapabilityName.EDITABLE_HISTORY}, + ) + + assert result == {CapabilityName.EDITABLE_HISTORY} + + async def test_capability_without_probe_excluded_when_not_declared(self) -> None: + target = MockPromptTarget() + # Override to a configuration that does NOT declare editable_history. + target._configuration = TargetConfiguration(capabilities=TargetCapabilities()) + target._send_prompt_to_target_async = AsyncMock(return_value=_ok_response()) # type: ignore[method-assign] + + result = await query_target_capabilities_async( + target=target, + capabilities={CapabilityName.EDITABLE_HISTORY}, + ) + + assert result == set() + + async def test_capability_without_probe_excluded_when_only_adapted(self, monkeypatch: pytest.MonkeyPatch) -> None: + """ + ADAPT in the policy must NOT count as native support for the fallback. + + Today every adaptable capability also has a probe, so this scenario only + arises if a future capability is declared adaptable without a probe. + We simulate that by removing SYSTEM_PROMPT from the registry and + configuring the target with ``ADAPT`` for it but no native support. + """ + from pyrit.prompt_target.common import query_target_capabilities as qtc + from pyrit.prompt_target.common.target_capabilities import ( + CapabilityHandlingPolicy, + UnsupportedCapabilityBehavior, + ) + + patched_probes = {k: v for k, v in qtc._CAPABILITY_PROBES.items() if k is not CapabilityName.SYSTEM_PROMPT} + monkeypatch.setattr(qtc, "_CAPABILITY_PROBES", patched_probes) + + target = MockPromptTarget() + target._configuration = TargetConfiguration( + capabilities=TargetCapabilities(), # no native SYSTEM_PROMPT + policy=CapabilityHandlingPolicy( + behaviors={ + CapabilityName.SYSTEM_PROMPT: UnsupportedCapabilityBehavior.ADAPT, + CapabilityName.MULTI_TURN: UnsupportedCapabilityBehavior.RAISE, + } + ), + ) + + result = await query_target_capabilities_async( + target=target, + capabilities={CapabilityName.SYSTEM_PROMPT}, + ) + + assert result == set() + + async def test_accepts_single_pass_iterable(self) -> None: + """Passing a generator must not silently drop fallback (non-probed) capabilities.""" + target = MockPromptTarget() + target._configuration = TargetConfiguration( + capabilities=TargetCapabilities(supports_editable_history=True), + ) + target._send_prompt_to_target_async = AsyncMock(return_value=_ok_response()) # type: ignore[method-assign] + + gen = (c for c in [CapabilityName.SYSTEM_PROMPT, CapabilityName.EDITABLE_HISTORY]) + result = await query_target_capabilities_async(target=target, capabilities=gen) + + assert CapabilityName.SYSTEM_PROMPT in result + assert CapabilityName.EDITABLE_HISTORY in result + + async def test_retries_zero_disables_retry(self) -> None: + target = MockPromptTarget() + target._send_prompt_to_target_async = AsyncMock(side_effect=Exception("boom")) # type: ignore[method-assign] + + result = await query_target_capabilities_async( + target=target, + capabilities={CapabilityName.JSON_OUTPUT}, + retries=0, + ) + + assert result == set() + assert target._send_prompt_to_target_async.await_count == 1 + + async def test_restores_configuration_after_probing(self) -> None: + target = MockPromptTarget() + original = target.configuration + target._send_prompt_to_target_async = AsyncMock(return_value=_ok_response()) # type: ignore[method-assign] + + await query_target_capabilities_async(target=target) + + assert target.configuration is original + + async def test_multi_turn_probe_sends_history_on_second_call(self) -> None: + target = MockPromptTarget() + target._send_prompt_to_target_async = AsyncMock(return_value=_ok_response()) # type: ignore[method-assign] + + await query_target_capabilities_async( + target=target, + capabilities={CapabilityName.MULTI_TURN}, + ) + + # Multi-turn probe sends two requests on the same conversation_id, and + # seeds memory between them so the second call carries real history. + calls = target._send_prompt_to_target_async.await_args_list + assert len(calls) == 2 + + first_conv = calls[0].kwargs["normalized_conversation"] + second_conv = calls[1].kwargs["normalized_conversation"] + + first_conv_id = first_conv[-1].message_pieces[0].conversation_id + second_conv_id = second_conv[-1].message_pieces[0].conversation_id + assert first_conv_id == second_conv_id + + # First call is a single-turn user message; the second call must include + # the seeded user + assistant history followed by the new user turn. + assert len(first_conv) == 1 + assert len(second_conv) >= 3 + roles = [msg.message_pieces[0]._role for msg in second_conv] + assert roles[-3:] == ["user", "assistant", "user"] + + async def test_multi_turn_probe_short_circuits_on_first_failure(self) -> None: + target = MockPromptTarget() + target._send_prompt_to_target_async = AsyncMock(side_effect=Exception("first call fails")) # type: ignore[method-assign] + + result = await query_target_capabilities_async( + target=target, + capabilities={CapabilityName.MULTI_TURN}, + ) + + assert result == set() + # _send_and_check_async retries once on exception, so the failing + # first turn is attempted twice; the second turn is never reached. + assert target._send_prompt_to_target_async.await_count == 2 + + async def test_json_schema_probe_sends_schema_in_metadata(self) -> None: + target = MockPromptTarget() + target._send_prompt_to_target_async = AsyncMock(return_value=_ok_response()) # type: ignore[method-assign] + + await query_target_capabilities_async( + target=target, + capabilities={CapabilityName.JSON_SCHEMA}, + ) + + normalized: list[Message] = target._send_prompt_to_target_async.await_args.kwargs["normalized_conversation"] + metadata = normalized[-1].message_pieces[0].prompt_metadata + assert metadata is not None + assert metadata["response_format"] == "json" + # Schema is JSON-encoded into a string for prompt_metadata's value type. + schema = json.loads(metadata["json_schema"]) + assert schema["type"] == "object" + + async def test_system_prompt_probe_installs_system_message_and_sends_user(self) -> None: + target = MockPromptTarget() + target._send_prompt_to_target_async = AsyncMock(return_value=_ok_response()) # type: ignore[method-assign] + + await query_target_capabilities_async( + target=target, + capabilities={CapabilityName.SYSTEM_PROMPT}, + ) + + # The probe writes a system message directly to memory (bypassing + # PromptTarget.set_system_prompt, which subclasses can override) and + # then sends a user-role message. Message.validate forbids mixed + # roles in a single Message, so the system and user turns are + # separate. Verify the system message is in memory and the wire + # payload contains the system + user history. + normalized: list[Message] = target._send_prompt_to_target_async.await_args.kwargs["normalized_conversation"] + roles_sent = [piece._role for msg in normalized for piece in msg.message_pieces] + assert "system" in roles_sent + assert roles_sent[-1] == "user" + # The last sent Message itself should be user-only. + assert [piece._role for piece in normalized[-1].message_pieces] == ["user"] + + async def test_multi_message_pieces_probe_sends_two_pieces(self) -> None: + target = MockPromptTarget() + target._send_prompt_to_target_async = AsyncMock(return_value=_ok_response()) # type: ignore[method-assign] + + await query_target_capabilities_async( + target=target, + capabilities={CapabilityName.MULTI_MESSAGE_PIECES}, + ) + + normalized: list[Message] = target._send_prompt_to_target_async.await_args.kwargs["normalized_conversation"] + assert len(normalized[-1].message_pieces) == 2 + + async def test_probes_run_under_permissive_configuration(self) -> None: + """ + Even when the target declares no boolean capabilities, the probe should + still execute because the configuration is temporarily permissive. + + Uses ``_RealValidationTarget`` so that ``_validate_request`` actually + runs and would reject the multi-piece probe were the override absent. + """ + target = _RealValidationTarget() + send_mock = AsyncMock(return_value=_ok_response()) + target._send_prompt_to_target_async = send_mock # type: ignore[method-assign] + + result = await query_target_capabilities_async( + target=target, + capabilities={CapabilityName.MULTI_MESSAGE_PIECES}, + ) + + # Probe was actually invoked through the full send_prompt_async pipeline, + # which means _validate_request ran and was satisfied by the permissive + # override (the bare target declares no capabilities natively). + assert send_mock.await_count >= 1 + assert CapabilityName.MULTI_MESSAGE_PIECES in result + + +@pytest.mark.usefixtures("patch_central_database") +class TestQueryTargetCapabilitiesIsolatedTarget: + """Tests using a bare PromptTarget subclass (no PromptChatTarget extras).""" + + async def test_with_minimal_target_subclass(self) -> None: + class _MinimalTarget(PromptTarget): + async def _send_prompt_to_target_async(self, *, normalized_conversation: list[Message]) -> list[Message]: + return _ok_response() + + target = _MinimalTarget() + target._send_prompt_to_target_async = AsyncMock(return_value=_ok_response()) # type: ignore[method-assign] + + result = await query_target_capabilities_async(target=target) + + for capability in _CAPABILITY_PROBES: + assert capability in result + + +# --------------------------------------------------------------------------- +# Modality verification tests +# --------------------------------------------------------------------------- + + +def _set_input_modalities( + *, + target: MockPromptTarget, + modalities: set[frozenset[PromptDataType]], +) -> None: + target._configuration = TargetConfiguration( + capabilities=TargetCapabilities( + input_modalities=frozenset(modalities), + ), + ) + + +@pytest.fixture +def image_asset(tmp_path: Path) -> str: + """Create a tiny placeholder file usable as an image_path asset.""" + asset = tmp_path / "test_image.png" + asset.write_bytes(b"\x89PNG\r\n\x1a\n") + return str(asset) + + +@pytest.mark.usefixtures("patch_central_database") +class TestCreateTestMessage: + def test_text_only(self) -> None: + msg = _create_test_message(modalities=frozenset({"text"}), test_assets={}) + assert len(msg.message_pieces) == 1 + assert msg.message_pieces[0].original_value_data_type == "text" + + def test_multimodal_uses_assets(self, image_asset: str) -> None: + msg = _create_test_message( + modalities=frozenset({"text", "image_path"}), + test_assets={"image_path": image_asset}, + ) + types = {piece.original_value_data_type for piece in msg.message_pieces} + assert types == {"text", "image_path"} + + # All pieces share the same conversation_id (Message.validate requires it). + conv_ids = {piece.conversation_id for piece in msg.message_pieces} + assert len(conv_ids) == 1 + + def test_missing_asset_file_raises_filenotfound(self, tmp_path: Path) -> None: + missing_path = str(tmp_path / "does_not_exist.png") + with pytest.raises(FileNotFoundError): + _create_test_message( + modalities=frozenset({"image_path"}), + test_assets={"image_path": missing_path}, + ) + + def test_unconfigured_modality_raises_valueerror(self) -> None: + with pytest.raises(ValueError, match="No test asset configured"): + _create_test_message( + modalities=frozenset({"image_path"}), + test_assets={}, + ) + + +@pytest.mark.usefixtures("patch_central_database") +class TestVerifyTargetModalitiesAsync: + async def test_all_combinations_supported(self) -> None: + target = MockPromptTarget() + _set_input_modalities(target=target, modalities={frozenset({"text"})}) + target._send_prompt_to_target_async = AsyncMock(return_value=_ok_response()) # type: ignore[method-assign] + + result = await verify_target_modalities_async(target=target) + + assert frozenset({"text"}) in result + + async def test_exception_excludes_combination(self) -> None: + target = MockPromptTarget() + _set_input_modalities(target=target, modalities={frozenset({"text"})}) + target._send_prompt_to_target_async = AsyncMock(side_effect=Exception("nope")) # type: ignore[method-assign] + + result = await verify_target_modalities_async(target=target) + + assert result == set() + + async def test_error_response_excludes_combination(self) -> None: + target = MockPromptTarget() + _set_input_modalities(target=target, modalities={frozenset({"text"})}) + target._send_prompt_to_target_async = AsyncMock(return_value=_error_response()) # type: ignore[method-assign] + + result = await verify_target_modalities_async(target=target) + + assert result == set() + + async def test_partial_support_via_selective_failure(self, image_asset: str) -> None: + target = MockPromptTarget() + _set_input_modalities( + target=target, + modalities={frozenset({"text"}), frozenset({"text", "image_path"})}, + ) + + async def selective_send(*, normalized_conversation: list[Message]) -> list[Message]: + message = normalized_conversation[-1] + types = {p.original_value_data_type for p in message.message_pieces} + if "image_path" in types: + raise Exception("image not supported") + return _ok_response() + + target._send_prompt_to_target_async = selective_send # type: ignore[method-assign] + + result = await verify_target_modalities_async( + target=target, + test_assets={"image_path": image_asset}, + ) + + assert frozenset({"text"}) in result + assert frozenset({"text", "image_path"}) not in result + + async def test_explicit_test_modalities_overrides_declared(self, image_asset: str) -> None: + target = MockPromptTarget() + # Declared as text-only, but caller asks us to probe text+image too. + _set_input_modalities(target=target, modalities={frozenset({"text"})}) + target._send_prompt_to_target_async = AsyncMock(return_value=_ok_response()) # type: ignore[method-assign] + + result = await verify_target_modalities_async( + target=target, + test_modalities={frozenset({"text"}), frozenset({"text", "image_path"})}, + test_assets={"image_path": image_asset}, + ) + + assert frozenset({"text"}) in result + assert frozenset({"text", "image_path"}) in result + + async def test_combination_skipped_when_asset_missing(self, tmp_path: Path) -> None: + target = MockPromptTarget() + _set_input_modalities(target=target, modalities={frozenset({"text", "image_path"})}) + target._send_prompt_to_target_async = AsyncMock(return_value=_ok_response()) # type: ignore[method-assign] + + # No assets provided — image_path combinations are skipped, not probed. + result = await verify_target_modalities_async(target=target) + + assert result == set() + assert target._send_prompt_to_target_async.await_count == 0 + + async def test_explicit_test_modalities_runs_under_permissive_configuration(self, image_asset: str) -> None: + """ + Probing a modality combination the target does NOT declare must still + succeed. Uses ``_RealValidationTarget`` so ``_validate_request`` runs + and would reject the multi-piece, non-text payload were the + permissive override absent. + """ + target = _RealValidationTarget() + send_mock = AsyncMock(return_value=_ok_response()) + target._send_prompt_to_target_async = send_mock # type: ignore[method-assign] + + result = await verify_target_modalities_async( + target=target, + test_modalities={frozenset({"text", "image_path"})}, + test_assets={"image_path": image_asset}, + ) + + assert send_mock.await_count == 1 + assert frozenset({"text", "image_path"}) in result + + +@pytest.mark.usefixtures("patch_central_database") +class TestSendAndCheckTimeout: + async def test_timeout_returns_false_after_retries(self) -> None: + """ + When ``send_prompt_async`` exceeds ``per_probe_timeout_s``, the probe + is treated as failed. ``_send_and_check_async`` retries once on + timeout, so the underlying mock is awaited twice and the capability + is excluded from the verified set. + """ + target = MockPromptTarget() + + async def _hang(**_kwargs: object) -> list[Message]: + await asyncio.sleep(10) + return _ok_response() + + target._send_prompt_to_target_async = AsyncMock(side_effect=_hang) # type: ignore[method-assign] + + result = await query_target_capabilities_async( + target=target, + capabilities={CapabilityName.JSON_OUTPUT}, + per_probe_timeout_s=0.01, + ) + + assert result == set() + # One initial attempt plus one retry. + assert target._send_prompt_to_target_async.await_count == 2 + + +@pytest.mark.usefixtures("patch_central_database") +class TestSystemPromptProbeMemoryFailure: + async def test_returns_false_when_memory_seed_raises(self) -> None: + """ + If seeding the system message into memory raises (e.g. backend + offline), the system-prompt probe returns False without attempting + the user send. + """ + target = MockPromptTarget() + target._memory.add_message_to_memory = MagicMock(side_effect=RuntimeError("memory offline")) # type: ignore[method-assign] + send_mock = AsyncMock(return_value=_ok_response()) + target._send_prompt_to_target_async = send_mock # type: ignore[method-assign] + + result = await query_target_capabilities_async( + target=target, + capabilities={CapabilityName.SYSTEM_PROMPT}, + ) + + assert result == set() + # The user send is never attempted because seeding failed. + send_mock.assert_not_awaited() + + +@pytest.mark.usefixtures("patch_central_database") +class TestVerifyTargetAsync: + async def test_returns_target_capabilities_assembled_from_probes(self) -> None: + """ + ``verify_target_async`` runs both the capability and modality probes + and assembles a :class:`TargetCapabilities` populated from the + verified results, copying ``supports_editable_history`` and + ``output_modalities`` from the target's declared capabilities. + """ + declared = TargetCapabilities( + supports_editable_history=True, + input_modalities=frozenset({frozenset({"text"})}), + output_modalities=frozenset({frozenset({"text"})}), + ) + target = MockPromptTarget() + target._configuration = TargetConfiguration(capabilities=declared) + target._send_prompt_to_target_async = AsyncMock(return_value=_ok_response()) # type: ignore[method-assign] + + result = await verify_target_async(target=target, per_probe_timeout_s=5.0) + + assert isinstance(result, TargetCapabilities) + # Single-piece probes that don't touch memory always succeed when + # the underlying send returns a clean response. + assert result.supports_multi_message_pieces is True + assert result.supports_json_schema is True + assert result.supports_json_output is True + # Non-probed flags are copied from target.capabilities. + assert result.supports_editable_history is True + # Modalities returned from the modality probe (text combination). + assert frozenset({"text"}) in result.input_modalities + # Output modalities copied through (not probed). + assert result.output_modalities == declared.output_modalities + + async def test_excludes_capabilities_when_probe_send_fails(self) -> None: + """ + When the underlying send raises, no capability or modality is + verified, but ``supports_editable_history`` and ``output_modalities`` + are still copied from the declared capabilities. + """ + declared = TargetCapabilities( + supports_editable_history=True, + output_modalities=frozenset({frozenset({"text"})}), + ) + target = MockPromptTarget() + target._configuration = TargetConfiguration(capabilities=declared) + target._send_prompt_to_target_async = AsyncMock(side_effect=RuntimeError("boom")) # type: ignore[method-assign] + + result = await verify_target_async(target=target, per_probe_timeout_s=0.5) + + assert result.supports_multi_turn is False + assert result.supports_system_prompt is False + assert result.supports_json_output is False + assert result.supports_json_schema is False + assert result.supports_multi_message_pieces is False + # Non-probed flag preserved. + assert result.supports_editable_history is True + # No modalities verified because send always fails. + assert result.input_modalities == frozenset() + # Output modalities still copied. + assert result.output_modalities == declared.output_modalities + + async def test_empty_response_treated_as_failure(self) -> None: + """A target returning an empty response list must NOT be reported as supporting probes.""" + target = MockPromptTarget() + target._send_prompt_to_target_async = AsyncMock(return_value=[]) # type: ignore[method-assign] + + result = await query_target_capabilities_async( + target=target, + capabilities={CapabilityName.JSON_OUTPUT, CapabilityName.MULTI_MESSAGE_PIECES}, + ) + + assert result == set() + + async def test_response_with_no_pieces_treated_as_failure(self) -> None: + """Responses whose Messages have no pieces must also be rejected.""" + target = MockPromptTarget() + target._send_prompt_to_target_async = AsyncMock( # type: ignore[method-assign] + return_value=[Message.__new__(Message)] + ) + # Bypass __init__ to construct a Message with no pieces (Message.__init__ rejects empty). + empty_msg = target._send_prompt_to_target_async.return_value[0] + empty_msg.message_pieces = [] + + result = await query_target_capabilities_async( + target=target, + capabilities={CapabilityName.JSON_OUTPUT}, + ) + + assert result == set() + + async def test_verify_target_async_forwards_test_modalities(self, image_asset: str) -> None: + declared = TargetCapabilities(input_modalities=frozenset({frozenset({"text"})})) + target = MockPromptTarget() + target._configuration = TargetConfiguration(capabilities=declared) + target._send_prompt_to_target_async = AsyncMock(return_value=_ok_response()) + + extra_combo = frozenset({"text", "image_path"}) + result = await verify_target_async( + target=target, + test_modalities={extra_combo}, + test_assets={"image_path": image_asset}, + per_probe_timeout_s=2.0, + ) + + # The undeclared combination is in the result only if test_modalities was forwarded. + assert extra_combo in result.input_modalities + + async def test_verify_target_async_forwards_capabilities(self) -> None: + """``verify_target_async`` must forward ``capabilities`` to narrow the probe set.""" + target = MockPromptTarget() + target._send_prompt_to_target_async = AsyncMock(return_value=_ok_response()) # type: ignore[method-assign] + + await verify_target_async( + target=target, + capabilities={CapabilityName.JSON_OUTPUT}, + per_probe_timeout_s=2.0, + ) + + # Only the JSON_OUTPUT probe (1 send) and the modality probe(s) should run; + # if `capabilities` were ignored, all 5 capability probes would fire (>= 6 sends + # because multi-turn issues 2 sends). + assert target._send_prompt_to_target_async.await_count <= 3 + + async def test_verify_target_async_preserves_declared_when_capabilities_narrowed(self) -> None: + """ + When ``capabilities`` narrows the probe set, capabilities NOT in the + narrowed set must fall back to the target's declared values rather + than being silently reset to False. + """ + declared = TargetCapabilities( + supports_multi_turn=True, + supports_system_prompt=True, + supports_json_schema=True, + supports_editable_history=True, + ) + target = MockPromptTarget() + target._configuration = TargetConfiguration(capabilities=declared) + target._send_prompt_to_target_async = AsyncMock(return_value=_ok_response()) # type: ignore[method-assign] + + result = await verify_target_async( + target=target, + capabilities={CapabilityName.JSON_OUTPUT}, + per_probe_timeout_s=2.0, + ) + + # The probed capability reflects the verified result. + assert result.supports_json_output is True + # Non-probed capabilities fall back to declared values. + assert result.supports_multi_turn is True + assert result.supports_system_prompt is True + assert result.supports_json_schema is True + assert result.supports_editable_history is True + + +@pytest.mark.usefixtures("patch_central_database") +class TestMultiTurnProbeMemoryFailure: + async def test_returns_false_when_history_seed_raises(self) -> None: + """ + If seeding conversation history into memory raises, the multi-turn + probe returns False rather than proceeding with a half-seeded + conversation that would produce a false positive. + """ + target = MockPromptTarget() + send_mock = AsyncMock(return_value=_ok_response()) + target._send_prompt_to_target_async = send_mock # type: ignore[method-assign] + target._memory.add_message_to_memory = MagicMock(side_effect=RuntimeError("memory offline")) # type: ignore[method-assign] + + result = await query_target_capabilities_async( + target=target, + capabilities={CapabilityName.MULTI_TURN}, + ) + + assert result == set() + # The first turn ran (1 send); the second turn must NOT run because + # seeding failed, otherwise the probe would falsely succeed. + assert send_mock.await_count == 1