From e2da267941aae89cc208aa853f61412f156e233a Mon Sep 17 00:00:00 2001 From: hannahwestra25 Date: Thu, 7 May 2026 13:22:18 -0400 Subject: [PATCH 01/27] add ability to query target capabilities --- .../targets/6_1_target_capabilities.ipynb | 282 +++++------ doc/code/targets/6_1_target_capabilities.py | 80 ++- pyrit/prompt_target/__init__.py | 6 + .../common/query_target_capabilities.py | 462 ++++++++++++++++++ .../test_query_target_capabilities.py | 404 +++++++++++++++ 5 files changed, 1078 insertions(+), 156 deletions(-) create mode 100644 pyrit/prompt_target/common/query_target_capabilities.py create mode 100644 tests/unit/prompt_target/test_query_target_capabilities.py diff --git a/doc/code/targets/6_1_target_capabilities.ipynb b/doc/code/targets/6_1_target_capabilities.ipynb index 18c1902062..d76c68ef5b 100644 --- a/doc/code/targets/6_1_target_capabilities.ipynb +++ b/doc/code/targets/6_1_target_capabilities.ipynb @@ -2,7 +2,7 @@ "cells": [ { "cell_type": "markdown", - "id": "0", + "id": "670645e2", "metadata": {}, "source": [ "# 6.1 Target Capabilities\n", @@ -26,7 +26,7 @@ }, { "cell_type": "markdown", - "id": "1", + "id": "53b22857", "metadata": {}, "source": [ "## 1. Inspect a real target's configuration\n", @@ -38,38 +38,9 @@ { "cell_type": "code", "execution_count": null, - "id": "2", + "id": "8593f07b", "metadata": {}, - "outputs": [ - { - "name": "stdout", - "output_type": "stream", - "text": [ - "Found default environment files: ['./.pyrit/.env']\n", - "Loaded environment file: ./.pyrit/.env\n" - ] - }, - { - "name": "stdout", - "output_type": "stream", - "text": [ - "No new upgrade operations detected.\n" - ] - }, - { - "name": "stdout", - "output_type": "stream", - "text": [ - "supports_multi_turn: True\n", - "supports_editable_history: True\n", - "supports_system_prompt: True\n", - "supports_json_output: True\n", - "supports_json_schema: False\n", - "input_modalities: [['image_path'], ['image_path', 'text'], ['text']]\n", - "output_modalities: [['text']]\n" - ] - } - ], + "outputs": [], "source": [ "from pyrit.prompt_target import OpenAIChatTarget\n", "from pyrit.setup import IN_MEMORY, initialize_pyrit_async\n", @@ -90,7 +61,7 @@ }, { "cell_type": "markdown", - "id": "3", + "id": "8643c933", "metadata": {}, "source": [ "## 2. Default configurations and known model profiles\n", @@ -105,23 +76,9 @@ { "cell_type": "code", "execution_count": null, - "id": "4", + "id": "396b2cc6", "metadata": {}, - "outputs": [ - { - "name": "stdout", - "output_type": "stream", - "text": [ - "capability class default gpt-4o gpt-5 unknown \n", - "--------------------------------------------------------------------------------\n", - "supports_multi_turn True True True True \n", - "supports_editable_history True True True True \n", - "supports_system_prompt True True True True \n", - "supports_json_output True True True True \n", - "supports_json_schema False False True False \n" - ] - } - ], + "outputs": [], "source": [ "class_default = OpenAIChatTarget._DEFAULT_CONFIGURATION.capabilities\n", "gpt_4o = OpenAIChatTarget.get_default_configuration(underlying_model=\"gpt-4o\").capabilities\n", @@ -149,7 +106,7 @@ }, { "cell_type": "markdown", - "id": "5", + "id": "178f940c", "metadata": {}, "source": [ "## 3. Declare and validate consumer requirements\n", @@ -165,17 +122,9 @@ { "cell_type": "code", "execution_count": null, - "id": "6", + "id": "8b76ea2b", "metadata": {}, - "outputs": [ - { - "name": "stdout", - "output_type": "stream", - "text": [ - "OpenAIChatTarget satisfies CHAT_TARGET_REQUIREMENTS\n" - ] - } - ], + "outputs": [], "source": [ "from pyrit.prompt_target import CHAT_TARGET_REQUIREMENTS\n", "\n", @@ -185,7 +134,7 @@ }, { "cell_type": "markdown", - "id": "7", + "id": "0d5f8336", "metadata": {}, "source": [ "To check a single capability, call `target.configuration.ensure_can_handle(capability=...)` directly." @@ -194,17 +143,9 @@ { "cell_type": "code", "execution_count": null, - "id": "8", + "id": "ea6bca27", "metadata": {}, - "outputs": [ - { - "name": "stdout", - "output_type": "stream", - "text": [ - "Multi-turn check passed\n" - ] - } - ], + "outputs": [], "source": [ "from pyrit.prompt_target.common.target_capabilities import CapabilityName\n", "\n", @@ -214,7 +155,7 @@ }, { "cell_type": "markdown", - "id": "9", + "id": "ac56ae78", "metadata": {}, "source": [ "## 4. Override the configuration per instance\n", @@ -228,23 +169,9 @@ { "cell_type": "code", "execution_count": null, - "id": "10", + "id": "3b8b5936", "metadata": {}, - "outputs": [ - { - "name": "stdout", - "output_type": "stream", - "text": [ - "class default supports_multi_turn: True\n", - "instance supports_multi_turn: False\n", - "\n", - "Validation failed as expected:\n", - "Target does not satisfy 2 required capability(ies):\n", - " - Target does not support 'supports_editable_history' and no handling policy exists for it.\n", - " - Target does not support 'supports_multi_turn' and the handling policy is RAISE.\n" - ] - } - ], + "outputs": [], "source": [ "from pyrit.prompt_target.common.target_capabilities import TargetCapabilities\n", "from pyrit.prompt_target.common.target_configuration import TargetConfiguration\n", @@ -275,7 +202,7 @@ }, { "cell_type": "markdown", - "id": "11", + "id": "c80c4fdb", "metadata": {}, "source": [ "## 5. ADAPT vs RAISE\n", @@ -294,18 +221,9 @@ { "cell_type": "code", "execution_count": null, - "id": "12", + "id": "50e54468", "metadata": {}, - "outputs": [ - { - "name": "stdout", - "output_type": "stream", - "text": [ - "RAISE pipeline normalizers: []\n", - "ADAPT pipeline normalizers: ['GenericSystemSquashNormalizer', 'HistorySquashNormalizer']\n" - ] - } - ], + "outputs": [], "source": [ "from pyrit.prompt_target.common.target_capabilities import (\n", " CapabilityHandlingPolicy,\n", @@ -352,7 +270,7 @@ }, { "cell_type": "markdown", - "id": "13", + "id": "df1b9a25", "metadata": {}, "source": [ "With `ADAPT`, running a multi-turn conversation through `normalize_async` collapses it into a single\n", @@ -362,25 +280,9 @@ { "cell_type": "code", "execution_count": null, - "id": "14", + "id": "94e60fa5", "metadata": {}, - "outputs": [ - { - "name": "stdout", - "output_type": "stream", - "text": [ - "original turns: 3\n", - "normalized turns: 1\n", - "flattened text:\n", - "[Conversation History]\n", - "User: What is the capital of France?\n", - "Assistant: Paris.\n", - "\n", - "[Current Message]\n", - "And of Germany?\n" - ] - } - ], + "outputs": [], "source": [ "from pyrit.models import Message\n", "\n", @@ -399,7 +301,7 @@ }, { "cell_type": "markdown", - "id": "15", + "id": "d5fe7b82", "metadata": {}, "source": [ "By contrast, the `RAISE` configuration validates eagerly: any consumer requiring `MULTI_TURN` will\n", @@ -409,17 +311,9 @@ { "cell_type": "code", "execution_count": null, - "id": "16", + "id": "9dff84f4", "metadata": {}, - "outputs": [ - { - "name": "stdout", - "output_type": "stream", - "text": [ - "Target does not support 'supports_multi_turn' and the handling policy is RAISE.\n" - ] - } - ], + "outputs": [], "source": [ "try:\n", " raise_target.configuration.ensure_can_handle(capability=CapabilityName.MULTI_TURN)\n", @@ -429,7 +323,7 @@ }, { "cell_type": "markdown", - "id": "17", + "id": "fa5e1ef8", "metadata": {}, "source": [ "## 6. Non-adaptable capabilities\n", @@ -443,17 +337,9 @@ { "cell_type": "code", "execution_count": null, - "id": "18", + "id": "911183b3", "metadata": {}, - "outputs": [ - { - "name": "stdout", - "output_type": "stream", - "text": [ - "Target does not support 'supports_editable_history' and no handling policy exists for it.\n" - ] - } - ], + "outputs": [], "source": [ "no_editable_history = TargetConfiguration(\n", " capabilities=TargetCapabilities(supports_multi_turn=True, supports_editable_history=False),\n", @@ -462,23 +348,109 @@ "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": "85ab10c1", + "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` and `verify_target_modalities_async`.\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", + "Typical usage against a real endpoint:\n", + "\n", + "```python\n", + "from pyrit.prompt_target import query_target_capabilities_async\n", + "\n", + "verified = await query_target_capabilities_async(target=target)\n", + "print(verified)\n", + "```\n", + "\n", + "Below we mock `send_prompt_async` so the notebook stays self-contained — the result shape is\n", + "the same as a live run." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "c6f1b47d", + "metadata": {}, + "outputs": [], + "source": [ + "from unittest.mock import AsyncMock\n", + "\n", + "from pyrit.models import MessagePiece\n", + "from pyrit.prompt_target import query_target_capabilities_async\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(\n", + " model_name=\"gpt-4o\", endpoint=\"https://example.invalid/\", api_key=\"sk-not-a-real-key\"\n", + ")\n", + "probe_target.send_prompt_async = AsyncMock(return_value=_ok_response()) # type: ignore[method-assign]\n", + "\n", + "verified = await query_target_capabilities_async(target=probe_target) # 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": "5e86d09a", + "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", + "A common workflow is to probe the live target and then construct a `TargetConfiguration` from the\n", + "verified set, so the rest of PyRIT (attacks, scorers, the normalization pipeline) operates on\n", + "capabilities that have been observed to work end-to-end." ] } ], "metadata": { - "language_info": { - "codemirror_mode": { - "name": "ipython", - "version": 3 - }, - "file_extension": ".py", - "mimetype": "text/x-python", - "name": "python", - "nbconvert_exporter": "python", - "pygments_lexer": "ipython3", - "version": "3.11.15" + "jupytext": { + "main_language": "python" } }, "nbformat": 4, diff --git a/doc/code/targets/6_1_target_capabilities.py b/doc/code/targets/6_1_target_capabilities.py index 985374357b..9cea0f7e34 100644 --- a/doc/code/targets/6_1_target_capabilities.py +++ b/doc/code/targets/6_1_target_capabilities.py @@ -245,4 +245,82 @@ 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` and `verify_target_modalities_async`. +# +# `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`. +# +# Typical usage against a real endpoint: +# +# ```python +# from pyrit.prompt_target import query_target_capabilities_async +# +# verified = await query_target_capabilities_async(target=target) +# print(verified) +# ``` +# +# Below we mock `send_prompt_async` so the notebook stays self-contained — the result shape is +# the same as a live run. + +# %% +from unittest.mock import AsyncMock + +from pyrit.models import MessagePiece +from pyrit.prompt_target import query_target_capabilities_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_async = AsyncMock(return_value=_ok_response()) # type: ignore[method-assign] + +verified = await query_target_capabilities_async(target=probe_target) # 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], +# ) +# ``` +# +# A common workflow is to probe the live target and then construct a `TargetConfiguration` from the +# verified set, so the rest of PyRIT (attacks, scorers, the normalization pipeline) operates on +# capabilities that have been observed to work end-to-end. diff --git a/pyrit/prompt_target/__init__.py b/pyrit/prompt_target/__init__.py index 489fe34900..0867e3c34b 100644 --- a/pyrit/prompt_target/__init__.py +++ b/pyrit/prompt_target/__init__.py @@ -22,6 +22,10 @@ TargetCapabilities, UnsupportedCapabilityBehavior, ) +from pyrit.prompt_target.common.query_target_capabilities import ( + query_target_capabilities_async, + verify_target_modalities_async, +) from pyrit.prompt_target.common.target_configuration import TargetConfiguration from pyrit.prompt_target.common.target_requirements import CHAT_TARGET_REQUIREMENTS, TargetRequirements from pyrit.prompt_target.common.utils import limit_requests_per_minute @@ -97,11 +101,13 @@ def __getattr__(name: str) -> object: "PromptChatTarget", "PromptShieldTarget", "PromptTarget", + "query_target_capabilities_async", "RealtimeTarget", "TargetCapabilities", "TargetConfiguration", "TargetRequirements", "UnsupportedCapabilityBehavior", "TextTarget", + "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 0000000000..e10dc67ecc --- /dev/null +++ b/pyrit/prompt_target/common/query_target_capabilities.py @@ -0,0 +1,462 @@ +# 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``. +""" + +import json +import logging +import os +import uuid +from collections.abc import Awaitable, Iterable +from contextlib import contextmanager +from dataclasses import replace +from typing import Callable, Iterator, cast + +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__) + + +_CapabilityProbe = Callable[[PromptTarget], Awaitable[bool]] + + +_PERMISSIVE_POLICY = CapabilityHandlingPolicy( + behaviors={ + CapabilityName.MULTI_TURN: UnsupportedCapabilityBehavior.RAISE, + CapabilityName.SYSTEM_PROMPT: UnsupportedCapabilityBehavior.RAISE, + } +) + + +@contextmanager +def _permissive_configuration(*, target: PromptTarget) -> 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. + + Yields: + None: Control returns to the ``with`` block while the permissive + configuration is in effect. + """ + original = target.configuration + 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, + ) + target._configuration = TargetConfiguration( + capabilities=permissive_caps, + policy=_PERMISSIVE_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 _user_text_piece(*, value: str, conversation_id: str) -> MessagePiece: + """ + Build a single user-role text :class:`MessagePiece` for use in a probe. + + 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, + ) + + +async def _send_and_check_async(*, target: PromptTarget, message: Message) -> bool: + """ + Send ``message`` and report whether the call succeeded cleanly. + + Args: + target (PromptTarget): The target to send the probe message to. + message (Message): The probe message to send. + + Returns: + bool: ``True`` iff the call returned without raising and every response + piece reported ``response_error == "none"``; ``False`` otherwise. + """ + try: + responses = await target.send_prompt_async(message=message) + except Exception as exc: + logger.info("Capability probe failed: %s", exc) + return False + + for response in responses: + for piece in response.message_pieces: + if piece.response_error != "none": + logger.info("Capability probe returned error response: %s", piece.converted_value) + return False + return True + + +async def _probe_system_prompt_async(target: PromptTarget) -> bool: + """ + Probe whether ``target`` accepts a system message alongside a user message. + + Args: + target (PromptTarget): The target to probe. + + 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, + ) + user_piece = _user_text_piece(value="hi", conversation_id=conversation_id) + return await _send_and_check_async(target=target, message=Message([system_piece, user_piece])) + + +async def _probe_multi_message_pieces_async(target: PromptTarget) -> bool: + """ + Probe whether ``target`` accepts a single message containing multiple pieces. + + Args: + target (PromptTarget): The target to probe. + + 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)) + + +async def _probe_multi_turn_async(target: PromptTarget) -> bool: + """ + Probe whether ``target`` accepts two sequential messages on the same conversation. + + Args: + target (PromptTarget): The target to probe. + + Returns: + bool: ``True`` if both turns succeeded; ``False`` if either turn failed. + """ + conversation_id = _new_conversation_id() + first = _user_text_piece(value="hello", conversation_id=conversation_id) + if not await _send_and_check_async(target=target, message=Message([first])): + return False + second = _user_text_piece(value="and again", conversation_id=conversation_id) + return await _send_and_check_async(target=target, message=Message([second])) + + +async def _probe_json_output_async(target: PromptTarget) -> bool: + """ + Probe whether ``target`` accepts a request asking for JSON-mode output. + + Args: + target (PromptTarget): The target to probe. + + 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={"response_format": "json"}, + ) + return await _send_and_check_async(target=target, message=Message([piece])) + + +async def _probe_json_schema_async(target: PromptTarget) -> bool: + """ + Probe whether ``target`` accepts a request constrained by a JSON schema. + + Args: + target (PromptTarget): The target to probe. + + 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={ + "response_format": "json", + "json_schema": json.dumps(schema), + }, + ) + return await _send_and_check_async(target=target, message=Message([piece])) + + +# 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, +) -> 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 support + (``target.configuration.includes(...)``) is used as a fallback. + + 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`. + + Returns: + set[CapabilityName]: The capabilities verified to work against the target. + """ + capabilities_to_check: Iterable[CapabilityName] = capabilities if capabilities is not None else 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): + verified.add(capability) + except Exception as exc: + logger.info("Probe for %s raised: %s", capability.value, exc) + + # Add capabilities without a probe based on the original (now-restored) declared support. + for capability in capabilities_to_check: + if capability not in _CAPABILITY_PROBES and target.configuration.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, +) -> 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. + + 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. + + 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 = cast("set[frozenset[PromptDataType]]", set(declared)) + + assets = test_assets if test_assets is not None else DEFAULT_TEST_ASSETS + + verified: set[frozenset[PromptDataType]] = set() + 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 _test_modality_combination_async(target=target, message=message): + verified.add(combination) + + return verified + + +async def _test_modality_combination_async(*, target: PromptTarget, message: Message) -> bool: + """ + Send a modality probe ``message`` and report whether the call succeeded cleanly. + + Args: + target (PromptTarget): The target to send the probe message to. + message (Message): The probe message exercising a specific modality combination. + + Returns: + bool: ``True`` iff the call returned without raising and every response + piece reported ``response_error == "none"``; ``False`` otherwise. + """ + try: + responses = await target.send_prompt_async(message=message) + except Exception as exc: + logger.info("Modality probe failed: %s", exc) + return False + + for response in responses: + for piece in response.message_pieces: + if piece.response_error != "none": + logger.info("Modality probe returned error response: %s", piece.converted_value) + return False + return True + + +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, + ) + ) + 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, + ) + ) + + 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 0000000000..3af059f753 --- /dev/null +++ b/tests/unit/prompt_target/test_query_target_capabilities.py @@ -0,0 +1,404 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT license. + +import json +from pathlib import Path +from unittest.mock import AsyncMock + +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_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 + + +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_async = AsyncMock(return_value=_ok_response()) + + 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_async = AsyncMock(side_effect=Exception("nope")) + + 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_async = AsyncMock(return_value=_error_response()) + + 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_async = AsyncMock(return_value=_ok_response()) + + 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_async = AsyncMock(return_value=_ok_response()) + + 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() + target.send_prompt_async = AsyncMock(return_value=_ok_response()) + + result = await query_target_capabilities_async( + target=target, + capabilities={CapabilityName.EDITABLE_HISTORY}, + ) + + assert result == set() + + async def test_restores_configuration_after_probing(self) -> None: + target = MockPromptTarget() + original = target.configuration + target.send_prompt_async = AsyncMock(return_value=_ok_response()) + + await query_target_capabilities_async(target=target) + + assert target.configuration is original + + async def test_multi_turn_probe_makes_two_calls_with_same_conversation_id(self) -> None: + target = MockPromptTarget() + target.send_prompt_async = AsyncMock(return_value=_ok_response()) + + await query_target_capabilities_async( + target=target, + capabilities={CapabilityName.MULTI_TURN}, + ) + + # Multi-turn probe sends two messages on the same conversation_id. + calls = target.send_prompt_async.await_args_list + assert len(calls) == 2 + first_conv_id = calls[0].kwargs["message"].message_pieces[0].conversation_id + second_conv_id = calls[1].kwargs["message"].message_pieces[0].conversation_id + assert first_conv_id == second_conv_id + + async def test_multi_turn_probe_short_circuits_on_first_failure(self) -> None: + target = MockPromptTarget() + target.send_prompt_async = AsyncMock(side_effect=Exception("first call fails")) + + result = await query_target_capabilities_async( + target=target, + capabilities={CapabilityName.MULTI_TURN}, + ) + + assert result == set() + assert target.send_prompt_async.await_count == 1 + + async def test_json_schema_probe_sends_schema_in_metadata(self) -> None: + target = MockPromptTarget() + target.send_prompt_async = AsyncMock(return_value=_ok_response()) + + await query_target_capabilities_async( + target=target, + capabilities={CapabilityName.JSON_SCHEMA}, + ) + + message: Message = target.send_prompt_async.await_args.kwargs["message"] + metadata = message.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_sends_system_role(self) -> None: + target = MockPromptTarget() + target.send_prompt_async = AsyncMock(return_value=_ok_response()) + + await query_target_capabilities_async( + target=target, + capabilities={CapabilityName.SYSTEM_PROMPT}, + ) + + message: Message = target.send_prompt_async.await_args.kwargs["message"] + roles = [piece.role for piece in message.message_pieces] + assert "system" in roles + + async def test_multi_message_pieces_probe_sends_two_pieces(self) -> None: + target = MockPromptTarget() + target.send_prompt_async = AsyncMock(return_value=_ok_response()) + + await query_target_capabilities_async( + target=target, + capabilities={CapabilityName.MULTI_MESSAGE_PIECES}, + ) + + message: Message = target.send_prompt_async.await_args.kwargs["message"] + assert len(message.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. + """ + target = MockPromptTarget() + target._configuration = TargetConfiguration(capabilities=TargetCapabilities()) + target.send_prompt_async = AsyncMock(return_value=_ok_response()) + + result = await query_target_capabilities_async( + target=target, + capabilities={CapabilityName.MULTI_MESSAGE_PIECES}, + ) + + # Probe was actually invoked. + assert target.send_prompt_async.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_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_async = AsyncMock(return_value=_ok_response()) + + 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_async = AsyncMock(side_effect=Exception("nope")) + + 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_async = AsyncMock(return_value=_error_response()) + + 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(*, message: Message) -> list[Message]: + 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_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_async = AsyncMock(return_value=_ok_response()) + + 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_async = AsyncMock(return_value=_ok_response()) + + # 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_async.await_count == 0 From 7aeef2f9d079eb345ab495f6147745aa54732dba Mon Sep 17 00:00:00 2001 From: hannahwestra25 Date: Fri, 8 May 2026 10:39:33 -0400 Subject: [PATCH 02/27] FEAT deprecate prompt chat (#1678) --- .../targets/6_1_target_capabilities.ipynb | 183 +++++++++++++++--- 1 file changed, 153 insertions(+), 30 deletions(-) diff --git a/doc/code/targets/6_1_target_capabilities.ipynb b/doc/code/targets/6_1_target_capabilities.ipynb index d76c68ef5b..391f5fd52f 100644 --- a/doc/code/targets/6_1_target_capabilities.ipynb +++ b/doc/code/targets/6_1_target_capabilities.ipynb @@ -2,7 +2,7 @@ "cells": [ { "cell_type": "markdown", - "id": "670645e2", + "id": "0", "metadata": {}, "source": [ "# 6.1 Target Capabilities\n", @@ -26,7 +26,7 @@ }, { "cell_type": "markdown", - "id": "53b22857", + "id": "1", "metadata": {}, "source": [ "## 1. Inspect a real target's configuration\n", @@ -38,9 +38,38 @@ { "cell_type": "code", "execution_count": null, - "id": "8593f07b", + "id": "2", "metadata": {}, - "outputs": [], + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "Found default environment files: ['./.pyrit/.env']\n", + "Loaded environment file: ./.pyrit/.env\n" + ] + }, + { + "name": "stdout", + "output_type": "stream", + "text": [ + "No new upgrade operations detected.\n" + ] + }, + { + "name": "stdout", + "output_type": "stream", + "text": [ + "supports_multi_turn: True\n", + "supports_editable_history: True\n", + "supports_system_prompt: True\n", + "supports_json_output: True\n", + "supports_json_schema: False\n", + "input_modalities: [['image_path'], ['image_path', 'text'], ['text']]\n", + "output_modalities: [['text']]\n" + ] + } + ], "source": [ "from pyrit.prompt_target import OpenAIChatTarget\n", "from pyrit.setup import IN_MEMORY, initialize_pyrit_async\n", @@ -61,7 +90,7 @@ }, { "cell_type": "markdown", - "id": "8643c933", + "id": "3", "metadata": {}, "source": [ "## 2. Default configurations and known model profiles\n", @@ -76,9 +105,23 @@ { "cell_type": "code", "execution_count": null, - "id": "396b2cc6", + "id": "4", "metadata": {}, - "outputs": [], + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "capability class default gpt-4o gpt-5 unknown \n", + "--------------------------------------------------------------------------------\n", + "supports_multi_turn True True True True \n", + "supports_editable_history True True True True \n", + "supports_system_prompt True True True True \n", + "supports_json_output True True True True \n", + "supports_json_schema False False True False \n" + ] + } + ], "source": [ "class_default = OpenAIChatTarget._DEFAULT_CONFIGURATION.capabilities\n", "gpt_4o = OpenAIChatTarget.get_default_configuration(underlying_model=\"gpt-4o\").capabilities\n", @@ -106,7 +149,7 @@ }, { "cell_type": "markdown", - "id": "178f940c", + "id": "5", "metadata": {}, "source": [ "## 3. Declare and validate consumer requirements\n", @@ -122,9 +165,17 @@ { "cell_type": "code", "execution_count": null, - "id": "8b76ea2b", + "id": "6", "metadata": {}, - "outputs": [], + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "OpenAIChatTarget satisfies CHAT_TARGET_REQUIREMENTS\n" + ] + } + ], "source": [ "from pyrit.prompt_target import CHAT_TARGET_REQUIREMENTS\n", "\n", @@ -134,7 +185,7 @@ }, { "cell_type": "markdown", - "id": "0d5f8336", + "id": "7", "metadata": {}, "source": [ "To check a single capability, call `target.configuration.ensure_can_handle(capability=...)` directly." @@ -143,9 +194,17 @@ { "cell_type": "code", "execution_count": null, - "id": "ea6bca27", + "id": "8", "metadata": {}, - "outputs": [], + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "Multi-turn check passed\n" + ] + } + ], "source": [ "from pyrit.prompt_target.common.target_capabilities import CapabilityName\n", "\n", @@ -155,7 +214,7 @@ }, { "cell_type": "markdown", - "id": "ac56ae78", + "id": "9", "metadata": {}, "source": [ "## 4. Override the configuration per instance\n", @@ -169,9 +228,23 @@ { "cell_type": "code", "execution_count": null, - "id": "3b8b5936", + "id": "10", "metadata": {}, - "outputs": [], + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "class default supports_multi_turn: True\n", + "instance supports_multi_turn: False\n", + "\n", + "Validation failed as expected:\n", + "Target does not satisfy 2 required capability(ies):\n", + " - Target does not support 'supports_editable_history' and no handling policy exists for it.\n", + " - Target does not support 'supports_multi_turn' and the handling policy is RAISE.\n" + ] + } + ], "source": [ "from pyrit.prompt_target.common.target_capabilities import TargetCapabilities\n", "from pyrit.prompt_target.common.target_configuration import TargetConfiguration\n", @@ -202,7 +275,7 @@ }, { "cell_type": "markdown", - "id": "c80c4fdb", + "id": "11", "metadata": {}, "source": [ "## 5. ADAPT vs RAISE\n", @@ -221,9 +294,18 @@ { "cell_type": "code", "execution_count": null, - "id": "50e54468", + "id": "12", "metadata": {}, - "outputs": [], + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "RAISE pipeline normalizers: []\n", + "ADAPT pipeline normalizers: ['GenericSystemSquashNormalizer', 'HistorySquashNormalizer']\n" + ] + } + ], "source": [ "from pyrit.prompt_target.common.target_capabilities import (\n", " CapabilityHandlingPolicy,\n", @@ -270,7 +352,7 @@ }, { "cell_type": "markdown", - "id": "df1b9a25", + "id": "13", "metadata": {}, "source": [ "With `ADAPT`, running a multi-turn conversation through `normalize_async` collapses it into a single\n", @@ -280,9 +362,25 @@ { "cell_type": "code", "execution_count": null, - "id": "94e60fa5", + "id": "14", "metadata": {}, - "outputs": [], + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "original turns: 3\n", + "normalized turns: 1\n", + "flattened text:\n", + "[Conversation History]\n", + "User: What is the capital of France?\n", + "Assistant: Paris.\n", + "\n", + "[Current Message]\n", + "And of Germany?\n" + ] + } + ], "source": [ "from pyrit.models import Message\n", "\n", @@ -301,7 +399,7 @@ }, { "cell_type": "markdown", - "id": "d5fe7b82", + "id": "15", "metadata": {}, "source": [ "By contrast, the `RAISE` configuration validates eagerly: any consumer requiring `MULTI_TURN` will\n", @@ -311,9 +409,17 @@ { "cell_type": "code", "execution_count": null, - "id": "9dff84f4", + "id": "16", "metadata": {}, - "outputs": [], + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "Target does not support 'supports_multi_turn' and the handling policy is RAISE.\n" + ] + } + ], "source": [ "try:\n", " raise_target.configuration.ensure_can_handle(capability=CapabilityName.MULTI_TURN)\n", @@ -323,7 +429,7 @@ }, { "cell_type": "markdown", - "id": "fa5e1ef8", + "id": "17", "metadata": {}, "source": [ "## 6. Non-adaptable capabilities\n", @@ -337,9 +443,17 @@ { "cell_type": "code", "execution_count": null, - "id": "911183b3", + "id": "18", "metadata": {}, - "outputs": [], + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "Target does not support 'supports_editable_history' and no handling policy exists for it.\n" + ] + } + ], "source": [ "no_editable_history = TargetConfiguration(\n", " capabilities=TargetCapabilities(supports_multi_turn=True, supports_editable_history=False),\n", @@ -449,8 +563,17 @@ } ], "metadata": { - "jupytext": { - "main_language": "python" + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.11.15" } }, "nbformat": 4, From a3873b53398836d1afe7eeb73322bb541123e774 Mon Sep 17 00:00:00 2001 From: hannahwestra25 Date: Fri, 8 May 2026 12:24:10 -0400 Subject: [PATCH 03/27] merge --- .../targets/6_1_target_capabilities.ipynb | 29 ++++++++++++++----- doc/code/targets/6_1_target_capabilities.py | 4 +-- pyrit/prompt_target/__init__.py | 8 ++--- .../common/query_target_capabilities.py | 10 ++----- .../test_query_target_capabilities.py | 5 +--- 5 files changed, 31 insertions(+), 25 deletions(-) diff --git a/doc/code/targets/6_1_target_capabilities.ipynb b/doc/code/targets/6_1_target_capabilities.ipynb index 391f5fd52f..a06ecc8367 100644 --- a/doc/code/targets/6_1_target_capabilities.ipynb +++ b/doc/code/targets/6_1_target_capabilities.ipynb @@ -53,6 +53,7 @@ "name": "stdout", "output_type": "stream", "text": [ +<<<<<<< HEAD "No new upgrade operations detected.\n" ] }, @@ -60,6 +61,9 @@ "name": "stdout", "output_type": "stream", "text": [ +======= + "No new upgrade operations detected.\n", +>>>>>>> c4f012e4f (merge) "supports_multi_turn: True\n", "supports_editable_history: True\n", "supports_system_prompt: True\n", @@ -467,7 +471,7 @@ }, { "cell_type": "markdown", - "id": "85ab10c1", + "id": "19", "metadata": {}, "source": [ "## 7. Querying live target capabilities\n", @@ -503,9 +507,22 @@ { "cell_type": "code", "execution_count": null, - "id": "c6f1b47d", + "id": "20", "metadata": {}, - "outputs": [], + "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" + ] + } + ], "source": [ "from unittest.mock import AsyncMock\n", "\n", @@ -529,9 +546,7 @@ " ]\n", "\n", "\n", - "probe_target = OpenAIChatTarget(\n", - " model_name=\"gpt-4o\", endpoint=\"https://example.invalid/\", api_key=\"sk-not-a-real-key\"\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_async = AsyncMock(return_value=_ok_response()) # type: ignore[method-assign]\n", "\n", "verified = await query_target_capabilities_async(target=probe_target) # type: ignore\n", @@ -542,7 +557,7 @@ }, { "cell_type": "markdown", - "id": "5e86d09a", + "id": "21", "metadata": {}, "source": [ "To narrow the probe to specific capabilities (faster, fewer calls), pass `capabilities=`:\n", diff --git a/doc/code/targets/6_1_target_capabilities.py b/doc/code/targets/6_1_target_capabilities.py index 9cea0f7e34..c5c900de96 100644 --- a/doc/code/targets/6_1_target_capabilities.py +++ b/doc/code/targets/6_1_target_capabilities.py @@ -299,9 +299,7 @@ def _ok_response(): ] -probe_target = OpenAIChatTarget( - model_name="gpt-4o", endpoint="https://example.invalid/", api_key="sk-not-a-real-key" -) +probe_target = OpenAIChatTarget(model_name="gpt-4o", endpoint="https://example.invalid/", api_key="sk-not-a-real-key") probe_target.send_prompt_async = AsyncMock(return_value=_ok_response()) # type: ignore[method-assign] verified = await query_target_capabilities_async(target=probe_target) # type: ignore diff --git a/pyrit/prompt_target/__init__.py b/pyrit/prompt_target/__init__.py index 0867e3c34b..0163c77ee5 100644 --- a/pyrit/prompt_target/__init__.py +++ b/pyrit/prompt_target/__init__.py @@ -16,16 +16,16 @@ 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_modalities_async, +) from pyrit.prompt_target.common.target_capabilities import ( CapabilityHandlingPolicy, CapabilityName, TargetCapabilities, UnsupportedCapabilityBehavior, ) -from pyrit.prompt_target.common.query_target_capabilities import ( - query_target_capabilities_async, - verify_target_modalities_async, -) from pyrit.prompt_target.common.target_configuration import TargetConfiguration from pyrit.prompt_target.common.target_requirements import CHAT_TARGET_REQUIREMENTS, TargetRequirements from pyrit.prompt_target.common.utils import limit_requests_per_minute diff --git a/pyrit/prompt_target/common/query_target_capabilities.py b/pyrit/prompt_target/common/query_target_capabilities.py index e10dc67ecc..ec976d25a1 100644 --- a/pyrit/prompt_target/common/query_target_capabilities.py +++ b/pyrit/prompt_target/common/query_target_capabilities.py @@ -22,17 +22,15 @@ import logging import os import uuid -from collections.abc import Awaitable, Iterable +from collections.abc import Awaitable, Callable, Iterable, Iterator from contextlib import contextmanager from dataclasses import replace -from typing import Callable, Iterator, cast 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 @@ -356,7 +354,7 @@ async def verify_target_modalities_async( """ if test_modalities is None: declared = target.capabilities.input_modalities - test_modalities = cast("set[frozenset[PromptDataType]]", set(declared)) + test_modalities = set(declared) assets = test_assets if test_assets is not None else DEFAULT_TEST_ASSETS @@ -443,9 +441,7 @@ def _create_test_message( 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}" - ) + raise FileNotFoundError(f"Test asset for modality '{modality}' not found at: {asset_path}") pieces.append( MessagePiece( diff --git a/tests/unit/prompt_target/test_query_target_capabilities.py b/tests/unit/prompt_target/test_query_target_capabilities.py index 3af059f753..e684824311 100644 --- a/tests/unit/prompt_target/test_query_target_capabilities.py +++ b/tests/unit/prompt_target/test_query_target_capabilities.py @@ -21,7 +21,6 @@ TargetCapabilities, ) from pyrit.prompt_target.common.target_configuration import TargetConfiguration - from tests.unit.mocks import MockPromptTarget @@ -251,9 +250,7 @@ class TestQueryTargetCapabilitiesIsolatedTarget: 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]: + async def _send_prompt_to_target_async(self, *, normalized_conversation: list[Message]) -> list[Message]: return _ok_response() target = _MinimalTarget() From c912285f6bccc91567ccdeb87cf5f93be84e3301 Mon Sep 17 00:00:00 2001 From: hannahwestra25 Date: Fri, 8 May 2026 12:36:46 -0400 Subject: [PATCH 04/27] resolve merge --- doc/code/targets/6_1_target_capabilities.ipynb | 10 ---------- 1 file changed, 10 deletions(-) diff --git a/doc/code/targets/6_1_target_capabilities.ipynb b/doc/code/targets/6_1_target_capabilities.ipynb index a06ecc8367..e6663a8870 100644 --- a/doc/code/targets/6_1_target_capabilities.ipynb +++ b/doc/code/targets/6_1_target_capabilities.ipynb @@ -53,17 +53,7 @@ "name": "stdout", "output_type": "stream", "text": [ -<<<<<<< HEAD - "No new upgrade operations detected.\n" - ] - }, - { - "name": "stdout", - "output_type": "stream", - "text": [ -======= "No new upgrade operations detected.\n", ->>>>>>> c4f012e4f (merge) "supports_multi_turn: True\n", "supports_editable_history: True\n", "supports_system_prompt: True\n", From 026aa54d8cfaedbcf1772130d39d238a037f4cd1 Mon Sep 17 00:00:00 2001 From: hannahwestra25 Date: Fri, 8 May 2026 15:32:14 -0400 Subject: [PATCH 05/27] fix / clarify misc issues --- .../targets/6_1_target_capabilities.ipynb | 245 ++++++++--- doc/code/targets/6_1_target_capabilities.py | 51 ++- pyrit/prompt_target/__init__.py | 2 + .../common/query_target_capabilities.py | 379 ++++++++++++++---- .../test_query_target_capabilities.py | 184 +++++++-- 5 files changed, 685 insertions(+), 176 deletions(-) diff --git a/doc/code/targets/6_1_target_capabilities.ipynb b/doc/code/targets/6_1_target_capabilities.ipynb index e6663a8870..5560914dc4 100644 --- a/doc/code/targets/6_1_target_capabilities.ipynb +++ b/doc/code/targets/6_1_target_capabilities.ipynb @@ -2,7 +2,7 @@ "cells": [ { "cell_type": "markdown", - "id": "0", + "id": "47584e7f", "metadata": {}, "source": [ "# 6.1 Target Capabilities\n", @@ -26,7 +26,7 @@ }, { "cell_type": "markdown", - "id": "1", + "id": "02d9f5ba", "metadata": {}, "source": [ "## 1. Inspect a real target's configuration\n", @@ -37,16 +37,23 @@ }, { "cell_type": "code", - "execution_count": null, - "id": "2", - "metadata": {}, + "execution_count": 1, + "id": "d3eb107b", + "metadata": { + "execution": { + "iopub.execute_input": "2026-05-08T19:20:44.364073Z", + "iopub.status.busy": "2026-05-08T19:20:44.363623Z", + "iopub.status.idle": "2026-05-08T19:20:52.981077Z", + "shell.execute_reply": "2026-05-08T19:20:52.979594Z" + } + }, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ - "Found default environment files: ['./.pyrit/.env']\n", - "Loaded environment file: ./.pyrit/.env\n" + "Found default environment files: ['/home/vscode/.pyrit/.env']\n", + "Loaded environment file: /home/vscode/.pyrit/.env\n" ] }, { @@ -84,7 +91,7 @@ }, { "cell_type": "markdown", - "id": "3", + "id": "4287a821", "metadata": {}, "source": [ "## 2. Default configurations and known model profiles\n", @@ -98,9 +105,16 @@ }, { "cell_type": "code", - "execution_count": null, - "id": "4", - "metadata": {}, + "execution_count": 2, + "id": "bf8b20f1", + "metadata": { + "execution": { + "iopub.execute_input": "2026-05-08T19:20:52.983872Z", + "iopub.status.busy": "2026-05-08T19:20:52.983461Z", + "iopub.status.idle": "2026-05-08T19:20:52.991617Z", + "shell.execute_reply": "2026-05-08T19:20:52.990162Z" + } + }, "outputs": [ { "name": "stdout", @@ -143,7 +157,7 @@ }, { "cell_type": "markdown", - "id": "5", + "id": "d19340c0", "metadata": {}, "source": [ "## 3. Declare and validate consumer requirements\n", @@ -158,9 +172,16 @@ }, { "cell_type": "code", - "execution_count": null, - "id": "6", - "metadata": {}, + "execution_count": 3, + "id": "6a1e09ef", + "metadata": { + "execution": { + "iopub.execute_input": "2026-05-08T19:20:52.994167Z", + "iopub.status.busy": "2026-05-08T19:20:52.993923Z", + "iopub.status.idle": "2026-05-08T19:20:53.002172Z", + "shell.execute_reply": "2026-05-08T19:20:53.000425Z" + } + }, "outputs": [ { "name": "stdout", @@ -179,7 +200,7 @@ }, { "cell_type": "markdown", - "id": "7", + "id": "39c9f98e", "metadata": {}, "source": [ "To check a single capability, call `target.configuration.ensure_can_handle(capability=...)` directly." @@ -187,9 +208,16 @@ }, { "cell_type": "code", - "execution_count": null, - "id": "8", - "metadata": {}, + "execution_count": 4, + "id": "0f21674f", + "metadata": { + "execution": { + "iopub.execute_input": "2026-05-08T19:20:53.004666Z", + "iopub.status.busy": "2026-05-08T19:20:53.004435Z", + "iopub.status.idle": "2026-05-08T19:20:53.010111Z", + "shell.execute_reply": "2026-05-08T19:20:53.008857Z" + } + }, "outputs": [ { "name": "stdout", @@ -208,7 +236,7 @@ }, { "cell_type": "markdown", - "id": "9", + "id": "1fe8b880", "metadata": {}, "source": [ "## 4. Override the configuration per instance\n", @@ -221,9 +249,16 @@ }, { "cell_type": "code", - "execution_count": null, - "id": "10", - "metadata": {}, + "execution_count": 5, + "id": "ff9dea78", + "metadata": { + "execution": { + "iopub.execute_input": "2026-05-08T19:20:53.012407Z", + "iopub.status.busy": "2026-05-08T19:20:53.012206Z", + "iopub.status.idle": "2026-05-08T19:20:53.041744Z", + "shell.execute_reply": "2026-05-08T19:20:53.040420Z" + } + }, "outputs": [ { "name": "stdout", @@ -269,7 +304,7 @@ }, { "cell_type": "markdown", - "id": "11", + "id": "78340d48", "metadata": {}, "source": [ "## 5. ADAPT vs RAISE\n", @@ -287,9 +322,16 @@ }, { "cell_type": "code", - "execution_count": null, - "id": "12", - "metadata": {}, + "execution_count": 6, + "id": "0ea85378", + "metadata": { + "execution": { + "iopub.execute_input": "2026-05-08T19:20:53.044041Z", + "iopub.status.busy": "2026-05-08T19:20:53.043856Z", + "iopub.status.idle": "2026-05-08T19:20:53.099310Z", + "shell.execute_reply": "2026-05-08T19:20:53.097936Z" + } + }, "outputs": [ { "name": "stdout", @@ -346,7 +388,7 @@ }, { "cell_type": "markdown", - "id": "13", + "id": "1dd59c3f", "metadata": {}, "source": [ "With `ADAPT`, running a multi-turn conversation through `normalize_async` collapses it into a single\n", @@ -355,9 +397,16 @@ }, { "cell_type": "code", - "execution_count": null, - "id": "14", - "metadata": {}, + "execution_count": 7, + "id": "350866c5", + "metadata": { + "execution": { + "iopub.execute_input": "2026-05-08T19:20:53.101739Z", + "iopub.status.busy": "2026-05-08T19:20:53.101541Z", + "iopub.status.idle": "2026-05-08T19:20:53.107941Z", + "shell.execute_reply": "2026-05-08T19:20:53.106664Z" + } + }, "outputs": [ { "name": "stdout", @@ -393,7 +442,7 @@ }, { "cell_type": "markdown", - "id": "15", + "id": "8c1a0ca8", "metadata": {}, "source": [ "By contrast, the `RAISE` configuration validates eagerly: any consumer requiring `MULTI_TURN` will\n", @@ -402,9 +451,16 @@ }, { "cell_type": "code", - "execution_count": null, - "id": "16", - "metadata": {}, + "execution_count": 8, + "id": "3ceb9e9b", + "metadata": { + "execution": { + "iopub.execute_input": "2026-05-08T19:20:53.110181Z", + "iopub.status.busy": "2026-05-08T19:20:53.109984Z", + "iopub.status.idle": "2026-05-08T19:20:53.116259Z", + "shell.execute_reply": "2026-05-08T19:20:53.115087Z" + } + }, "outputs": [ { "name": "stdout", @@ -423,7 +479,7 @@ }, { "cell_type": "markdown", - "id": "17", + "id": "b42a46dd", "metadata": {}, "source": [ "## 6. Non-adaptable capabilities\n", @@ -436,9 +492,16 @@ }, { "cell_type": "code", - "execution_count": null, - "id": "18", - "metadata": {}, + "execution_count": 9, + "id": "e7bcb64f", + "metadata": { + "execution": { + "iopub.execute_input": "2026-05-08T19:20:53.118380Z", + "iopub.status.busy": "2026-05-08T19:20:53.118197Z", + "iopub.status.idle": "2026-05-08T19:20:53.126291Z", + "shell.execute_reply": "2026-05-08T19:20:53.124843Z" + } + }, "outputs": [ { "name": "stdout", @@ -461,7 +524,7 @@ }, { "cell_type": "markdown", - "id": "19", + "id": "f87bdac3", "metadata": {}, "source": [ "## 7. Querying live target capabilities\n", @@ -469,7 +532,8 @@ "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` and `verify_target_modalities_async`.\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", @@ -481,24 +545,43 @@ "`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 query_target_capabilities_async\n", + "from pyrit.prompt_target import verify_target_async\n", "\n", - "verified = await query_target_capabilities_async(target=target)\n", + "verified = await verify_target_async(target=target)\n", "print(verified)\n", "```\n", "\n", - "Below we mock `send_prompt_async` so the notebook stays self-contained — the result shape is\n", - "the same as a live run." + "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": {}, + "execution_count": 10, + "id": "2d74b445", + "metadata": { + "execution": { + "iopub.execute_input": "2026-05-08T19:20:53.128398Z", + "iopub.status.busy": "2026-05-08T19:20:53.128214Z", + "iopub.status.idle": "2026-05-08T19:20:53.219597Z", + "shell.execute_reply": "2026-05-08T19:20:53.218424Z" + } + }, "outputs": [ { "name": "stdout", @@ -509,7 +592,8 @@ " - supports_json_output\n", " - supports_json_schema\n", " - supports_multi_message_pieces\n", - " - supports_multi_turn\n" + " - supports_multi_turn\n", + " - supports_system_prompt\n" ] } ], @@ -517,7 +601,10 @@ "from unittest.mock import AsyncMock\n", "\n", "from pyrit.models import MessagePiece\n", - "from pyrit.prompt_target import query_target_capabilities_async\n", + "from pyrit.prompt_target import (\n", + " query_target_capabilities_async,\n", + " verify_target_async,\n", + ")\n", "\n", "\n", "def _ok_response():\n", @@ -537,9 +624,9 @@ "\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_async = AsyncMock(return_value=_ok_response()) # type: ignore[method-assign]\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) # type: ignore\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}\")" @@ -547,7 +634,7 @@ }, { "cell_type": "markdown", - "id": "21", + "id": "108ec658", "metadata": {}, "source": [ "To narrow the probe to specific capabilities (faster, fewer calls), pass `capabilities=`:\n", @@ -561,13 +648,59 @@ ")\n", "```\n", "\n", - "A common workflow is to probe the live target and then construct a `TargetConfiguration` from the\n", - "verified set, so the rest of PyRIT (attacks, scorers, the normalization pipeline) operates on\n", - "capabilities that have been observed to work end-to-end." + "`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": 11, + "id": "7d92dbfc", + "metadata": { + "execution": { + "iopub.execute_input": "2026-05-08T19:20:53.222310Z", + "iopub.status.busy": "2026-05-08T19:20:53.222097Z", + "iopub.status.idle": "2026-05-08T19:20:53.239327Z", + "shell.execute_reply": "2026-05-08T19:20:53.238113Z" + } + }, + "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)}\")" ] } ], "metadata": { + "kernelspec": { + "display_name": "Python (pyrit-dev)", + "language": "python", + "name": "pyrit-dev" + }, "language_info": { "codemirror_mode": { "name": "ipython", diff --git a/doc/code/targets/6_1_target_capabilities.py b/doc/code/targets/6_1_target_capabilities.py index c5c900de96..185312bbb9 100644 --- a/doc/code/targets/6_1_target_capabilities.py +++ b/doc/code/targets/6_1_target_capabilities.py @@ -252,7 +252,8 @@ # 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` and `verify_target_modalities_async`. +# `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 @@ -264,23 +265,38 @@ # `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 query_target_capabilities_async +# from pyrit.prompt_target import verify_target_async # -# verified = await query_target_capabilities_async(target=target) +# verified = await verify_target_async(target=target) # print(verified) # ``` # -# Below we mock `send_prompt_async` so the notebook stays self-contained — the result shape is -# the same as a live run. +# 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 +from pyrit.prompt_target import ( + query_target_capabilities_async, + verify_target_async, +) def _ok_response(): @@ -300,9 +316,9 @@ def _ok_response(): probe_target = OpenAIChatTarget(model_name="gpt-4o", endpoint="https://example.invalid/", api_key="sk-not-a-real-key") -probe_target.send_prompt_async = AsyncMock(return_value=_ok_response()) # type: ignore[method-assign] +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) # type: ignore +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}") @@ -319,6 +335,19 @@ def _ok_response(): # ) # ``` # -# A common workflow is to probe the live target and then construct a `TargetConfiguration` from the -# verified set, so the rest of PyRIT (attacks, scorers, the normalization pipeline) operates on -# capabilities that have been observed to work end-to-end. +# `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)}") diff --git a/pyrit/prompt_target/__init__.py b/pyrit/prompt_target/__init__.py index 0163c77ee5..ef682b2a44 100644 --- a/pyrit/prompt_target/__init__.py +++ b/pyrit/prompt_target/__init__.py @@ -18,6 +18,7 @@ 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 ( @@ -108,6 +109,7 @@ def __getattr__(name: str) -> object: "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 index ec976d25a1..24222a6405 100644 --- a/pyrit/prompt_target/common/query_target_capabilities.py +++ b/pyrit/prompt_target/common/query_target_capabilities.py @@ -16,8 +16,30 @@ * :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 @@ -31,17 +53,28 @@ 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], Awaitable[bool]] +_CapabilityProbe = Callable[[PromptTarget, float], Awaitable[bool]] -_PERMISSIVE_POLICY = CapabilityHandlingPolicy( +_PROBE_POLICY = CapabilityHandlingPolicy( behaviors={ CapabilityName.MULTI_TURN: UnsupportedCapabilityBehavior.RAISE, CapabilityName.SYSTEM_PROMPT: UnsupportedCapabilityBehavior.RAISE, @@ -49,8 +82,18 @@ ) +# 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) -> Iterator[None]: +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. @@ -61,12 +104,21 @@ def _permissive_configuration(*, target: PromptTarget) -> Iterator[None]: 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, @@ -75,10 +127,11 @@ def _permissive_configuration(*, target: PromptTarget) -> Iterator[None]: supports_json_output=True, supports_editable_history=True, supports_system_prompt=True, + input_modalities=merged_modalities, ) target._configuration = TargetConfiguration( capabilities=permissive_caps, - policy=_PERMISSIVE_POLICY, + policy=_PROBE_POLICY, ) try: yield @@ -96,10 +149,21 @@ def _new_conversation_id() -> str: 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. @@ -112,41 +176,80 @@ def _user_text_piece(*, value: str, conversation_id: str) -> MessagePiece: 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) -> bool: +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. """ - try: - responses = await target.send_prompt_async(message=message) - except Exception as exc: - logger.info("Capability probe failed: %s", exc) - return False + 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.info("%s timed out (attempt %d/%d)", label, attempt + 1, attempts) + continue + except Exception as exc: + last_exc = exc + logger.info("%s failed (attempt %d/%d): %s", label, attempt + 1, attempts, exc) + continue + + for response in responses: + for piece in response.message_pieces: + if piece.response_error != "none": + logger.info("%s returned error response: %s", label, piece.converted_value) + return False + return True - for response in responses: - for piece in response.message_pieces: - if piece.response_error != "none": - logger.info("Capability probe returned error response: %s", 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) -> bool: +async def _probe_system_prompt_async(target: PromptTarget, timeout_s: float) -> bool: """ - Probe whether ``target`` accepts a system message alongside a user message. + 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. @@ -157,17 +260,29 @@ async def _probe_system_prompt_async(target: PromptTarget) -> bool: 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.info("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([system_piece, user_piece])) + return await _send_and_check_async( + target=target, + message=Message([user_piece]), + timeout_s=timeout_s, + label="System-prompt probe", + ) -async def _probe_multi_message_pieces_async(target: PromptTarget) -> bool: +async def _probe_multi_message_pieces_async(target: PromptTarget, timeout_s: float) -> 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. @@ -177,33 +292,71 @@ async def _probe_multi_message_pieces_async(target: PromptTarget) -> bool: _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)) + return await _send_and_check_async( + target=target, + message=Message(pieces), + timeout_s=timeout_s, + label="Multi-message-pieces probe", + ) -async def _probe_multi_turn_async(target: PromptTarget) -> bool: +async def _probe_multi_turn_async(target: PromptTarget, timeout_s: float) -> bool: """ - Probe whether ``target`` accepts two sequential messages on the same conversation. + 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="hello", conversation_id=conversation_id) - if not await _send_and_check_async(target=target, message=Message([first])): + 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, label="Multi-turn probe (turn 1)" + ): return False - second = _user_text_piece(value="and again", conversation_id=conversation_id) - return await _send_and_check_async(target=target, message=Message([second])) + + # Seed memory so the second send sees real prior history. + 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) + + 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, label="Multi-turn probe (turn 2)" + ) -async def _probe_json_output_async(target: PromptTarget) -> bool: +async def _probe_json_output_async(target: PromptTarget, timeout_s: float) -> 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. @@ -214,17 +367,20 @@ async def _probe_json_output_async(target: PromptTarget) -> bool: original_value='Respond with a JSON object: {"ok": true}.', original_value_data_type="text", conversation_id=conversation_id, - prompt_metadata={"response_format": "json"}, + prompt_metadata=_probe_metadata({"response_format": "json"}), + ) + return await _send_and_check_async( + target=target, message=Message([piece]), timeout_s=timeout_s, label="JSON-output probe" ) - return await _send_and_check_async(target=target, message=Message([piece])) -async def _probe_json_schema_async(target: PromptTarget) -> bool: +async def _probe_json_schema_async(target: PromptTarget, timeout_s: float) -> 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. @@ -241,12 +397,16 @@ async def _probe_json_schema_async(target: PromptTarget) -> bool: original_value='Respond with a JSON object matching the schema: {"ok": true}.', original_value_data_type="text", conversation_id=conversation_id, - prompt_metadata={ - "response_format": "json", - "json_schema": json.dumps(schema), - }, + 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, label="JSON-schema probe" ) - return await _send_and_check_async(target=target, message=Message([piece])) # Registry of capabilities that can be verified via a live API call. @@ -264,6 +424,7 @@ async def query_target_capabilities_async( *, target: PromptTarget, capabilities: Iterable[CapabilityName] | None = None, + per_probe_timeout_s: float = DEFAULT_PROBE_TIMEOUT_SECONDS, ) -> set[CapabilityName]: """ Probe ``target`` to determine which capabilities it actually supports. @@ -271,8 +432,31 @@ async def query_target_capabilities_async( 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 support - (``target.configuration.includes(...)``) is used as a fallback. + 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 @@ -284,6 +468,9 @@ async def query_target_capabilities_async( 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`. Returns: set[CapabilityName]: The capabilities verified to work against the target. @@ -300,14 +487,18 @@ async def query_target_capabilities_async( continue try: - if await probe(target): + if await probe(target, per_probe_timeout_s): verified.add(capability) except Exception as exc: logger.info("Probe for %s raised: %s", capability.value, exc) - # Add capabilities without a probe based on the original (now-restored) declared support. + # 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.configuration.includes(capability=capability): + if capability not in _CAPABILITY_PROBES and target.capabilities.includes(capability=capability): verified.add(capability) return verified @@ -330,6 +521,7 @@ 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, ) -> set[frozenset[PromptDataType]]: """ Probe ``target`` to determine which input modality combinations it supports. @@ -338,6 +530,23 @@ async def verify_target_modalities_async( :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 @@ -347,6 +556,9 @@ async def verify_target_modalities_async( 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`. Returns: set[frozenset[PromptDataType]]: The modality combinations verified @@ -359,46 +571,75 @@ async def verify_target_modalities_async( assets = test_assets if test_assets is not None else DEFAULT_TEST_ASSETS verified: set[frozenset[PromptDataType]] = set() - 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 + 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 _test_modality_combination_async(target=target, message=message): - verified.add(combination) + if await _send_and_check_async( + target=target, + message=message, + timeout_s=per_probe_timeout_s, + label=f"Modality probe {sorted(combination)}", + ): + verified.add(combination) return verified -async def _test_modality_combination_async(*, target: PromptTarget, message: Message) -> bool: +async def verify_target_async( + *, + target: PromptTarget, + per_probe_timeout_s: float = DEFAULT_PROBE_TIMEOUT_SECONDS, + test_assets: dict[PromptDataType, str] | None = None, +) -> TargetCapabilities: """ - Send a modality probe ``message`` and report whether the call succeeded cleanly. + 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). Args: - target (PromptTarget): The target to send the probe message to. - message (Message): The probe message exercising a specific modality combination. + target (PromptTarget): The target to probe. + per_probe_timeout_s (float): Per-attempt timeout (seconds) applied to + each probe request. + test_assets (dict[PromptDataType, str] | None): Mapping from non-text + modality to a file path. See :func:`verify_target_modalities_async`. Returns: - bool: ``True`` iff the call returned without raising and every response - piece reported ``response_error == "none"``; ``False`` otherwise. + 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. """ - try: - responses = await target.send_prompt_async(message=message) - except Exception as exc: - logger.info("Modality probe failed: %s", exc) - return False + verified_caps = await query_target_capabilities_async(target=target, per_probe_timeout_s=per_probe_timeout_s) + verified_modalities = await verify_target_modalities_async( + target=target, test_assets=test_assets, per_probe_timeout_s=per_probe_timeout_s + ) - for response in responses: - for piece in response.message_pieces: - if piece.response_error != "none": - logger.info("Modality probe returned error response: %s", piece.converted_value) - return False - return True + declared = target.capabilities + return TargetCapabilities( + supports_multi_turn=CapabilityName.MULTI_TURN in verified_caps, + supports_multi_message_pieces=CapabilityName.MULTI_MESSAGE_PIECES in verified_caps, + supports_json_schema=CapabilityName.JSON_SCHEMA in verified_caps, + supports_json_output=CapabilityName.JSON_OUTPUT in verified_caps, + supports_editable_history=declared.supports_editable_history, + supports_system_prompt=CapabilityName.SYSTEM_PROMPT in verified_caps, + input_modalities=frozenset(verified_modalities), + output_modalities=declared.output_modalities, + ) def _create_test_message( diff --git a/tests/unit/prompt_target/test_query_target_capabilities.py b/tests/unit/prompt_target/test_query_target_capabilities.py index e684824311..7213c070aa 100644 --- a/tests/unit/prompt_target/test_query_target_capabilities.py +++ b/tests/unit/prompt_target/test_query_target_capabilities.py @@ -24,6 +24,23 @@ 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( @@ -85,7 +102,7 @@ def test_restores_on_exception(self) -> None: class TestQueryTargetCapabilitiesAsync: async def test_returns_only_supported_when_all_probes_succeed(self) -> None: target = MockPromptTarget() - target.send_prompt_async = AsyncMock(return_value=_ok_response()) + target._send_prompt_to_target_async = AsyncMock(return_value=_ok_response()) # type: ignore[method-assign] result = await query_target_capabilities_async(target=target) @@ -95,7 +112,7 @@ async def test_returns_only_supported_when_all_probes_succeed(self) -> None: async def test_excludes_capabilities_when_probe_fails(self) -> None: target = MockPromptTarget() - target.send_prompt_async = AsyncMock(side_effect=Exception("nope")) + target._send_prompt_to_target_async = AsyncMock(side_effect=Exception("nope")) # type: ignore[method-assign] result = await query_target_capabilities_async(target=target) @@ -104,7 +121,7 @@ async def test_excludes_capabilities_when_probe_fails(self) -> None: async def test_excludes_capabilities_when_response_has_error(self) -> None: target = MockPromptTarget() - target.send_prompt_async = AsyncMock(return_value=_error_response()) + target._send_prompt_to_target_async = AsyncMock(return_value=_error_response()) # type: ignore[method-assign] result = await query_target_capabilities_async(target=target) @@ -113,7 +130,7 @@ async def test_excludes_capabilities_when_response_has_error(self) -> None: async def test_filters_by_requested_capabilities(self) -> None: target = MockPromptTarget() - target.send_prompt_async = AsyncMock(return_value=_ok_response()) + 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) @@ -126,7 +143,7 @@ async def test_capability_without_probe_falls_back_to_declared_support(self) -> target._configuration = TargetConfiguration( capabilities=TargetCapabilities(supports_editable_history=True), ) - target.send_prompt_async = AsyncMock(return_value=_ok_response()) + target._send_prompt_to_target_async = AsyncMock(return_value=_ok_response()) # type: ignore[method-assign] result = await query_target_capabilities_async( target=target, @@ -137,7 +154,9 @@ async def test_capability_without_probe_falls_back_to_declared_support(self) -> async def test_capability_without_probe_excluded_when_not_declared(self) -> None: target = MockPromptTarget() - target.send_prompt_async = AsyncMock(return_value=_ok_response()) + # 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, @@ -146,34 +165,82 @@ async def test_capability_without_probe_excluded_when_not_declared(self) -> None 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_restores_configuration_after_probing(self) -> None: target = MockPromptTarget() original = target.configuration - target.send_prompt_async = AsyncMock(return_value=_ok_response()) + 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_makes_two_calls_with_same_conversation_id(self) -> None: + async def test_multi_turn_probe_sends_history_on_second_call(self) -> None: target = MockPromptTarget() - target.send_prompt_async = AsyncMock(return_value=_ok_response()) + 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 messages on the same conversation_id. - calls = target.send_prompt_async.await_args_list + # 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_id = calls[0].kwargs["message"].message_pieces[0].conversation_id - second_conv_id = calls[1].kwargs["message"].message_pieces[0].conversation_id + + 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_async = AsyncMock(side_effect=Exception("first call fails")) + 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, @@ -181,66 +248,82 @@ async def test_multi_turn_probe_short_circuits_on_first_failure(self) -> None: ) assert result == set() - assert target.send_prompt_async.await_count == 1 + # _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_async = AsyncMock(return_value=_ok_response()) + 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}, ) - message: Message = target.send_prompt_async.await_args.kwargs["message"] - metadata = message.message_pieces[0].prompt_metadata + 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_sends_system_role(self) -> None: + async def test_system_prompt_probe_installs_system_message_and_sends_user(self) -> None: target = MockPromptTarget() - target.send_prompt_async = AsyncMock(return_value=_ok_response()) + 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}, ) - message: Message = target.send_prompt_async.await_args.kwargs["message"] - roles = [piece.role for piece in message.message_pieces] - assert "system" in roles + # 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_async = AsyncMock(return_value=_ok_response()) + 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}, ) - message: Message = target.send_prompt_async.await_args.kwargs["message"] - assert len(message.message_pieces) == 2 + 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 = MockPromptTarget() - target._configuration = TargetConfiguration(capabilities=TargetCapabilities()) - target.send_prompt_async = AsyncMock(return_value=_ok_response()) + 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. - assert target.send_prompt_async.await_count >= 1 + # 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 @@ -254,7 +337,7 @@ async def _send_prompt_to_target_async(self, *, normalized_conversation: list[Me return _ok_response() target = _MinimalTarget() - target.send_prompt_async = AsyncMock(return_value=_ok_response()) # type: ignore[method-assign] + target._send_prompt_to_target_async = AsyncMock(return_value=_ok_response()) # type: ignore[method-assign] result = await query_target_capabilities_async(target=target) @@ -327,7 +410,7 @@ class TestVerifyTargetModalitiesAsync: async def test_all_combinations_supported(self) -> None: target = MockPromptTarget() _set_input_modalities(target=target, modalities={frozenset({"text"})}) - target.send_prompt_async = AsyncMock(return_value=_ok_response()) + target._send_prompt_to_target_async = AsyncMock(return_value=_ok_response()) # type: ignore[method-assign] result = await verify_target_modalities_async(target=target) @@ -336,7 +419,7 @@ async def test_all_combinations_supported(self) -> None: async def test_exception_excludes_combination(self) -> None: target = MockPromptTarget() _set_input_modalities(target=target, modalities={frozenset({"text"})}) - target.send_prompt_async = AsyncMock(side_effect=Exception("nope")) + target._send_prompt_to_target_async = AsyncMock(side_effect=Exception("nope")) # type: ignore[method-assign] result = await verify_target_modalities_async(target=target) @@ -345,7 +428,7 @@ async def test_exception_excludes_combination(self) -> None: async def test_error_response_excludes_combination(self) -> None: target = MockPromptTarget() _set_input_modalities(target=target, modalities={frozenset({"text"})}) - target.send_prompt_async = AsyncMock(return_value=_error_response()) + target._send_prompt_to_target_async = AsyncMock(return_value=_error_response()) # type: ignore[method-assign] result = await verify_target_modalities_async(target=target) @@ -358,13 +441,14 @@ async def test_partial_support_via_selective_failure(self, image_asset: str) -> modalities={frozenset({"text"}), frozenset({"text", "image_path"})}, ) - async def selective_send(*, message: Message) -> list[Message]: + 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_async = selective_send # type: ignore[method-assign] + target._send_prompt_to_target_async = selective_send # type: ignore[method-assign] result = await verify_target_modalities_async( target=target, @@ -378,7 +462,7 @@ async def test_explicit_test_modalities_overrides_declared(self, image_asset: st 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_async = AsyncMock(return_value=_ok_response()) + target._send_prompt_to_target_async = AsyncMock(return_value=_ok_response()) # type: ignore[method-assign] result = await verify_target_modalities_async( target=target, @@ -392,10 +476,30 @@ async def test_explicit_test_modalities_overrides_declared(self, image_asset: st 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_async = AsyncMock(return_value=_ok_response()) + 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_async.await_count == 0 + 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 From d5d902a7b8eccd3b59399c8c0bd3c4ad5da2bf8b Mon Sep 17 00:00:00 2001 From: hannahwestra25 Date: Fri, 8 May 2026 15:48:55 -0400 Subject: [PATCH 06/27] add more tests --- .../test_query_target_capabilities.py | 117 +++++++++++++++++- 1 file changed, 116 insertions(+), 1 deletion(-) diff --git a/tests/unit/prompt_target/test_query_target_capabilities.py b/tests/unit/prompt_target/test_query_target_capabilities.py index 7213c070aa..ac7c18c338 100644 --- a/tests/unit/prompt_target/test_query_target_capabilities.py +++ b/tests/unit/prompt_target/test_query_target_capabilities.py @@ -1,9 +1,10 @@ # Copyright (c) Microsoft Corporation. # Licensed under the MIT license. +import asyncio import json from pathlib import Path -from unittest.mock import AsyncMock +from unittest.mock import AsyncMock, MagicMock import pytest @@ -14,6 +15,7 @@ _create_test_message, _permissive_configuration, query_target_capabilities_async, + verify_target_async, verify_target_modalities_async, ) from pyrit.prompt_target.common.target_capabilities import ( @@ -503,3 +505,116 @@ async def test_explicit_test_modalities_runs_under_permissive_configuration(self 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 From 632a11080ec2db8af1dd7e922f084864d58a8adc Mon Sep 17 00:00:00 2001 From: hannahwestra25 Date: Fri, 8 May 2026 15:57:29 -0400 Subject: [PATCH 07/27] pre-commit --- .../targets/6_1_target_capabilities.ipynb | 176 +++++------------- 1 file changed, 47 insertions(+), 129 deletions(-) diff --git a/doc/code/targets/6_1_target_capabilities.ipynb b/doc/code/targets/6_1_target_capabilities.ipynb index 5560914dc4..2adfc84fd4 100644 --- a/doc/code/targets/6_1_target_capabilities.ipynb +++ b/doc/code/targets/6_1_target_capabilities.ipynb @@ -2,7 +2,7 @@ "cells": [ { "cell_type": "markdown", - "id": "47584e7f", + "id": "0", "metadata": {}, "source": [ "# 6.1 Target Capabilities\n", @@ -26,7 +26,7 @@ }, { "cell_type": "markdown", - "id": "02d9f5ba", + "id": "1", "metadata": {}, "source": [ "## 1. Inspect a real target's configuration\n", @@ -37,23 +37,16 @@ }, { "cell_type": "code", - "execution_count": 1, - "id": "d3eb107b", - "metadata": { - "execution": { - "iopub.execute_input": "2026-05-08T19:20:44.364073Z", - "iopub.status.busy": "2026-05-08T19:20:44.363623Z", - "iopub.status.idle": "2026-05-08T19:20:52.981077Z", - "shell.execute_reply": "2026-05-08T19:20:52.979594Z" - } - }, + "execution_count": null, + "id": "2", + "metadata": {}, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ - "Found default environment files: ['/home/vscode/.pyrit/.env']\n", - "Loaded environment file: /home/vscode/.pyrit/.env\n" + "Found default environment files: ['./.pyrit/.env']\n", + "Loaded environment file: ./.pyrit/.env\n" ] }, { @@ -91,7 +84,7 @@ }, { "cell_type": "markdown", - "id": "4287a821", + "id": "3", "metadata": {}, "source": [ "## 2. Default configurations and known model profiles\n", @@ -105,16 +98,9 @@ }, { "cell_type": "code", - "execution_count": 2, - "id": "bf8b20f1", - "metadata": { - "execution": { - "iopub.execute_input": "2026-05-08T19:20:52.983872Z", - "iopub.status.busy": "2026-05-08T19:20:52.983461Z", - "iopub.status.idle": "2026-05-08T19:20:52.991617Z", - "shell.execute_reply": "2026-05-08T19:20:52.990162Z" - } - }, + "execution_count": null, + "id": "4", + "metadata": {}, "outputs": [ { "name": "stdout", @@ -157,7 +143,7 @@ }, { "cell_type": "markdown", - "id": "d19340c0", + "id": "5", "metadata": {}, "source": [ "## 3. Declare and validate consumer requirements\n", @@ -172,16 +158,9 @@ }, { "cell_type": "code", - "execution_count": 3, - "id": "6a1e09ef", - "metadata": { - "execution": { - "iopub.execute_input": "2026-05-08T19:20:52.994167Z", - "iopub.status.busy": "2026-05-08T19:20:52.993923Z", - "iopub.status.idle": "2026-05-08T19:20:53.002172Z", - "shell.execute_reply": "2026-05-08T19:20:53.000425Z" - } - }, + "execution_count": null, + "id": "6", + "metadata": {}, "outputs": [ { "name": "stdout", @@ -200,7 +179,7 @@ }, { "cell_type": "markdown", - "id": "39c9f98e", + "id": "7", "metadata": {}, "source": [ "To check a single capability, call `target.configuration.ensure_can_handle(capability=...)` directly." @@ -208,16 +187,9 @@ }, { "cell_type": "code", - "execution_count": 4, - "id": "0f21674f", - "metadata": { - "execution": { - "iopub.execute_input": "2026-05-08T19:20:53.004666Z", - "iopub.status.busy": "2026-05-08T19:20:53.004435Z", - "iopub.status.idle": "2026-05-08T19:20:53.010111Z", - "shell.execute_reply": "2026-05-08T19:20:53.008857Z" - } - }, + "execution_count": null, + "id": "8", + "metadata": {}, "outputs": [ { "name": "stdout", @@ -236,7 +208,7 @@ }, { "cell_type": "markdown", - "id": "1fe8b880", + "id": "9", "metadata": {}, "source": [ "## 4. Override the configuration per instance\n", @@ -249,16 +221,9 @@ }, { "cell_type": "code", - "execution_count": 5, - "id": "ff9dea78", - "metadata": { - "execution": { - "iopub.execute_input": "2026-05-08T19:20:53.012407Z", - "iopub.status.busy": "2026-05-08T19:20:53.012206Z", - "iopub.status.idle": "2026-05-08T19:20:53.041744Z", - "shell.execute_reply": "2026-05-08T19:20:53.040420Z" - } - }, + "execution_count": null, + "id": "10", + "metadata": {}, "outputs": [ { "name": "stdout", @@ -304,7 +269,7 @@ }, { "cell_type": "markdown", - "id": "78340d48", + "id": "11", "metadata": {}, "source": [ "## 5. ADAPT vs RAISE\n", @@ -322,16 +287,9 @@ }, { "cell_type": "code", - "execution_count": 6, - "id": "0ea85378", - "metadata": { - "execution": { - "iopub.execute_input": "2026-05-08T19:20:53.044041Z", - "iopub.status.busy": "2026-05-08T19:20:53.043856Z", - "iopub.status.idle": "2026-05-08T19:20:53.099310Z", - "shell.execute_reply": "2026-05-08T19:20:53.097936Z" - } - }, + "execution_count": null, + "id": "12", + "metadata": {}, "outputs": [ { "name": "stdout", @@ -388,7 +346,7 @@ }, { "cell_type": "markdown", - "id": "1dd59c3f", + "id": "13", "metadata": {}, "source": [ "With `ADAPT`, running a multi-turn conversation through `normalize_async` collapses it into a single\n", @@ -397,16 +355,9 @@ }, { "cell_type": "code", - "execution_count": 7, - "id": "350866c5", - "metadata": { - "execution": { - "iopub.execute_input": "2026-05-08T19:20:53.101739Z", - "iopub.status.busy": "2026-05-08T19:20:53.101541Z", - "iopub.status.idle": "2026-05-08T19:20:53.107941Z", - "shell.execute_reply": "2026-05-08T19:20:53.106664Z" - } - }, + "execution_count": null, + "id": "14", + "metadata": {}, "outputs": [ { "name": "stdout", @@ -442,7 +393,7 @@ }, { "cell_type": "markdown", - "id": "8c1a0ca8", + "id": "15", "metadata": {}, "source": [ "By contrast, the `RAISE` configuration validates eagerly: any consumer requiring `MULTI_TURN` will\n", @@ -451,16 +402,9 @@ }, { "cell_type": "code", - "execution_count": 8, - "id": "3ceb9e9b", - "metadata": { - "execution": { - "iopub.execute_input": "2026-05-08T19:20:53.110181Z", - "iopub.status.busy": "2026-05-08T19:20:53.109984Z", - "iopub.status.idle": "2026-05-08T19:20:53.116259Z", - "shell.execute_reply": "2026-05-08T19:20:53.115087Z" - } - }, + "execution_count": null, + "id": "16", + "metadata": {}, "outputs": [ { "name": "stdout", @@ -479,7 +423,7 @@ }, { "cell_type": "markdown", - "id": "b42a46dd", + "id": "17", "metadata": {}, "source": [ "## 6. Non-adaptable capabilities\n", @@ -492,16 +436,9 @@ }, { "cell_type": "code", - "execution_count": 9, - "id": "e7bcb64f", - "metadata": { - "execution": { - "iopub.execute_input": "2026-05-08T19:20:53.118380Z", - "iopub.status.busy": "2026-05-08T19:20:53.118197Z", - "iopub.status.idle": "2026-05-08T19:20:53.126291Z", - "shell.execute_reply": "2026-05-08T19:20:53.124843Z" - } - }, + "execution_count": null, + "id": "18", + "metadata": {}, "outputs": [ { "name": "stdout", @@ -524,7 +461,7 @@ }, { "cell_type": "markdown", - "id": "f87bdac3", + "id": "19", "metadata": {}, "source": [ "## 7. Querying live target capabilities\n", @@ -572,16 +509,9 @@ }, { "cell_type": "code", - "execution_count": 10, - "id": "2d74b445", - "metadata": { - "execution": { - "iopub.execute_input": "2026-05-08T19:20:53.128398Z", - "iopub.status.busy": "2026-05-08T19:20:53.128214Z", - "iopub.status.idle": "2026-05-08T19:20:53.219597Z", - "shell.execute_reply": "2026-05-08T19:20:53.218424Z" - } - }, + "execution_count": null, + "id": "20", + "metadata": {}, "outputs": [ { "name": "stdout", @@ -634,7 +564,7 @@ }, { "cell_type": "markdown", - "id": "108ec658", + "id": "21", "metadata": {}, "source": [ "To narrow the probe to specific capabilities (faster, fewer calls), pass `capabilities=`:\n", @@ -656,16 +586,9 @@ }, { "cell_type": "code", - "execution_count": 11, - "id": "7d92dbfc", - "metadata": { - "execution": { - "iopub.execute_input": "2026-05-08T19:20:53.222310Z", - "iopub.status.busy": "2026-05-08T19:20:53.222097Z", - "iopub.status.idle": "2026-05-08T19:20:53.239327Z", - "shell.execute_reply": "2026-05-08T19:20:53.238113Z" - } - }, + "execution_count": null, + "id": "22", + "metadata": {}, "outputs": [ { "name": "stdout", @@ -696,11 +619,6 @@ } ], "metadata": { - "kernelspec": { - "display_name": "Python (pyrit-dev)", - "language": "python", - "name": "pyrit-dev" - }, "language_info": { "codemirror_mode": { "name": "ipython", From bfe4fbf4cc288ccf509f474c17429a17f24c6271 Mon Sep 17 00:00:00 2001 From: hannahwestra25 Date: Fri, 8 May 2026 15:58:16 -0400 Subject: [PATCH 08/27] add documentation --- doc/code/targets/0_prompt_targets.md | 21 +++++++++++++++++++++ 1 file changed, 21 insertions(+) diff --git a/doc/code/targets/0_prompt_targets.md b/doc/code/targets/0_prompt_targets.md index e00983f769..9b1b5dab39 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. From 84a37361aa599f0a28c39b74374b04131faf8d11 Mon Sep 17 00:00:00 2001 From: hannahwestra25 Date: Fri, 8 May 2026 16:28:49 -0400 Subject: [PATCH 09/27] add retries --- .../common/query_target_capabilities.py | 50 +++++++++++++------ .../test_query_target_capabilities.py | 27 ++++++++++ 2 files changed, 63 insertions(+), 14 deletions(-) diff --git a/pyrit/prompt_target/common/query_target_capabilities.py b/pyrit/prompt_target/common/query_target_capabilities.py index 24222a6405..95f4937599 100644 --- a/pyrit/prompt_target/common/query_target_capabilities.py +++ b/pyrit/prompt_target/common/query_target_capabilities.py @@ -71,7 +71,7 @@ PROBE_METADATA_KEY: str = "capability_probe" PROBE_METADATA_VALUE: str = "1" -_CapabilityProbe = Callable[[PromptTarget, float], Awaitable[bool]] +_CapabilityProbe = Callable[[PromptTarget, float, int], Awaitable[bool]] _PROBE_POLICY = CapabilityHandlingPolicy( @@ -235,7 +235,7 @@ async def _send_and_check_async( return False -async def _probe_system_prompt_async(target: PromptTarget, timeout_s: float) -> bool: +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. @@ -272,11 +272,12 @@ async def _probe_system_prompt_async(target: PromptTarget, timeout_s: float) -> 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) -> bool: +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. @@ -296,11 +297,12 @@ async def _probe_multi_message_pieces_async(target: PromptTarget, timeout_s: flo 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) -> bool: +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. @@ -329,7 +331,7 @@ async def _probe_multi_turn_async(target: PromptTarget, timeout_s: float) -> boo 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, label="Multi-turn probe (turn 1)" + target=target, message=Message([first]), timeout_s=timeout_s, retries=retries, label="Multi-turn probe (turn 1)" ): return False @@ -346,11 +348,11 @@ async def _probe_multi_turn_async(target: PromptTarget, timeout_s: float) -> boo 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, label="Multi-turn probe (turn 2)" + 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) -> bool: +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. @@ -370,11 +372,11 @@ async def _probe_json_output_async(target: PromptTarget, timeout_s: float) -> bo prompt_metadata=_probe_metadata({"response_format": "json"}), ) return await _send_and_check_async( - target=target, message=Message([piece]), timeout_s=timeout_s, label="JSON-output probe" + 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) -> bool: +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. @@ -405,7 +407,7 @@ async def _probe_json_schema_async(target: PromptTarget, timeout_s: float) -> bo ), ) return await _send_and_check_async( - target=target, message=Message([piece]), timeout_s=timeout_s, label="JSON-schema probe" + target=target, message=Message([piece]), timeout_s=timeout_s, retries=retries, label="JSON-schema probe" ) @@ -425,6 +427,7 @@ 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. @@ -471,11 +474,17 @@ async def query_target_capabilities_async( 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: Iterable[CapabilityName] = capabilities if capabilities is not None else CapabilityName + capabilities_to_check: list[CapabilityName] = ( + list(capabilities) if capabilities is not None else list(CapabilityName) + ) verified: set[CapabilityName] = set() with _permissive_configuration(target=target): @@ -487,7 +496,7 @@ async def query_target_capabilities_async( continue try: - if await probe(target, per_probe_timeout_s): + if await probe(target, per_probe_timeout_s, retries): verified.add(capability) except Exception as exc: logger.info("Probe for %s raised: %s", capability.value, exc) @@ -522,6 +531,7 @@ async def verify_target_modalities_async( 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. @@ -559,6 +569,10 @@ async def verify_target_modalities_async( 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 @@ -586,6 +600,7 @@ async def verify_target_modalities_async( target=target, message=message, timeout_s=per_probe_timeout_s, + retries=retries, label=f"Modality probe {sorted(combination)}", ): verified.add(combination) @@ -598,6 +613,7 @@ async def verify_target_async( target: PromptTarget, per_probe_timeout_s: float = DEFAULT_PROBE_TIMEOUT_SECONDS, test_assets: dict[PromptDataType, str] | None = None, + retries: int = 1, ) -> TargetCapabilities: """ Probe both capabilities and modalities and return a combined result. @@ -617,6 +633,10 @@ async def verify_target_async( each probe request. test_assets (dict[PromptDataType, str] | None): Mapping from non-text modality to a file path. See :func:`verify_target_modalities_async`. + 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 @@ -624,9 +644,11 @@ async def verify_target_async( ``target.capabilities.output_modalities`` because outputs cannot be verified by sending a request. """ - verified_caps = await query_target_capabilities_async(target=target, per_probe_timeout_s=per_probe_timeout_s) + verified_caps = await query_target_capabilities_async( + target=target, per_probe_timeout_s=per_probe_timeout_s, retries=retries + ) verified_modalities = await verify_target_modalities_async( - target=target, test_assets=test_assets, per_probe_timeout_s=per_probe_timeout_s + target=target, test_assets=test_assets, per_probe_timeout_s=per_probe_timeout_s, retries=retries ) declared = target.capabilities diff --git a/tests/unit/prompt_target/test_query_target_capabilities.py b/tests/unit/prompt_target/test_query_target_capabilities.py index ac7c18c338..b8db7f2bb1 100644 --- a/tests/unit/prompt_target/test_query_target_capabilities.py +++ b/tests/unit/prompt_target/test_query_target_capabilities.py @@ -203,6 +203,33 @@ async def test_capability_without_probe_excluded_when_only_adapted(self, monkeyp 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 From 6e061e3f0862013a3499ff2abacc5667ed3c013a Mon Sep 17 00:00:00 2001 From: hannahwestra25 Date: Fri, 8 May 2026 16:49:55 -0400 Subject: [PATCH 10/27] treat empty response as failure and expose test_modalities and capabilities in verify_target_async --- .../common/query_target_capabilities.py | 35 ++++++++++- .../test_query_target_capabilities.py | 62 +++++++++++++++++++ 2 files changed, 95 insertions(+), 2 deletions(-) diff --git a/pyrit/prompt_target/common/query_target_capabilities.py b/pyrit/prompt_target/common/query_target_capabilities.py index 95f4937599..cb80741c2d 100644 --- a/pyrit/prompt_target/common/query_target_capabilities.py +++ b/pyrit/prompt_target/common/query_target_capabilities.py @@ -209,6 +209,8 @@ async def _send_and_check_async( Returns: bool: ``True`` iff the call returned without raising and every response piece reported ``response_error == "none"``; ``False`` otherwise. + An empty response list (or responses with no message pieces) is treated + as a failure rather than a success. """ attempts = max(1, retries + 1) last_exc: Exception | None = None @@ -224,6 +226,9 @@ async def _send_and_check_async( logger.info("%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.info("%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": @@ -612,7 +617,9 @@ 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: """ @@ -627,12 +634,27 @@ async def verify_target_async( :data:`_CAPABILITY_PROBES` (e.g. ``supports_editable_history``) are copied from ``target.capabilities`` (the target's declared native flags). + .. 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. @@ -645,10 +667,17 @@ async def verify_target_async( verified by sending a request. """ verified_caps = await query_target_capabilities_async( - target=target, per_probe_timeout_s=per_probe_timeout_s, retries=retries + target=target, + capabilities=capabilities, + per_probe_timeout_s=per_probe_timeout_s, + retries=retries, ) verified_modalities = await verify_target_modalities_async( - target=target, test_assets=test_assets, per_probe_timeout_s=per_probe_timeout_s, retries=retries + target=target, + test_modalities=test_modalities, + test_assets=test_assets, + per_probe_timeout_s=per_probe_timeout_s, + retries=retries, ) declared = target.capabilities @@ -696,6 +725,7 @@ def _create_test_message( original_value="test", original_value_data_type="text", conversation_id=conversation_id, + prompt_metadata=_probe_metadata(), ) ) continue @@ -712,6 +742,7 @@ def _create_test_message( original_value=asset_path, original_value_data_type=modality, conversation_id=conversation_id, + prompt_metadata=_probe_metadata(), ) ) diff --git a/tests/unit/prompt_target/test_query_target_capabilities.py b/tests/unit/prompt_target/test_query_target_capabilities.py index b8db7f2bb1..8ff7e150a8 100644 --- a/tests/unit/prompt_target/test_query_target_capabilities.py +++ b/tests/unit/prompt_target/test_query_target_capabilities.py @@ -645,3 +645,65 @@ async def test_excludes_capabilities_when_probe_send_fails(self) -> None: 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 From 9c3fad1fd1425f6ad74ee30cfb5490f1e80ec57d Mon Sep 17 00:00:00 2001 From: hannahwestra25 Date: Fri, 8 May 2026 17:07:08 -0400 Subject: [PATCH 11/27] preserve capabilities and debug flag --- .../targets/6_1_target_capabilities.ipynb | 38 ++++++++++- doc/code/targets/6_1_target_capabilities.py | 29 +++++++++ .../common/query_target_capabilities.py | 63 ++++++++++++------- .../test_query_target_capabilities.py | 54 ++++++++++++++++ 4 files changed, 160 insertions(+), 24 deletions(-) diff --git a/doc/code/targets/6_1_target_capabilities.ipynb b/doc/code/targets/6_1_target_capabilities.ipynb index 2adfc84fd4..ee25c333a7 100644 --- a/doc/code/targets/6_1_target_capabilities.ipynb +++ b/doc/code/targets/6_1_target_capabilities.ipynb @@ -45,8 +45,8 @@ "name": "stdout", "output_type": "stream", "text": [ - "Found default environment files: ['./.pyrit/.env']\n", - "Loaded environment file: ./.pyrit/.env\n" + "Found default environment files: ['/home/vscode/.pyrit/.env']\n", + "Loaded environment file: /home/vscode/.pyrit/.env\n" ] }, { @@ -616,6 +616,40 @@ "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", + "```" + ] } ], "metadata": { diff --git a/doc/code/targets/6_1_target_capabilities.py b/doc/code/targets/6_1_target_capabilities.py index 185312bbb9..485bef3477 100644 --- a/doc/code/targets/6_1_target_capabilities.py +++ b/doc/code/targets/6_1_target_capabilities.py @@ -351,3 +351,32 @@ def _ok_response(): 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/common/query_target_capabilities.py b/pyrit/prompt_target/common/query_target_capabilities.py index cb80741c2d..8e7170c4ed 100644 --- a/pyrit/prompt_target/common/query_target_capabilities.py +++ b/pyrit/prompt_target/common/query_target_capabilities.py @@ -209,8 +209,9 @@ async def _send_and_check_async( Returns: bool: ``True`` iff the call returned without raising and every response piece reported ``response_error == "none"``; ``False`` otherwise. - An empty response list (or responses with no message pieces) is treated - as a failure rather than a success. + 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 @@ -219,20 +220,20 @@ async def _send_and_check_async( 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.info("%s timed out (attempt %d/%d)", label, attempt + 1, attempts) + logger.debug("%s timed out (attempt %d/%d)", label, attempt + 1, attempts) continue except Exception as exc: last_exc = exc - logger.info("%s failed (attempt %d/%d): %s", label, attempt + 1, attempts, 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.info("%s returned an empty response; treating as failure", label) + 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.info("%s returned error response: %s", label, piece.converted_value) + logger.debug("%s returned error response: %s", label, piece.converted_value) return False return True @@ -270,7 +271,7 @@ async def _probe_system_prompt_async(target: PromptTarget, timeout_s: float, ret try: target._memory.add_message_to_memory(request=Message([system_piece])) except Exception as exc: - logger.info("System-prompt probe could not seed system message: %s", 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( @@ -341,15 +342,19 @@ async def _probe_multi_turn_async(target: PromptTarget, timeout_s: float, retrie return False # Seed memory so the second send sees real prior history. - 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) + 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( @@ -504,7 +509,7 @@ async def query_target_capabilities_async( if await probe(target, per_probe_timeout_s, retries): verified.add(capability) except Exception as exc: - logger.info("Probe for %s raised: %s", capability.value, 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 @@ -633,6 +638,10 @@ async def verify_target_async( 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 @@ -681,13 +690,23 @@ async def verify_target_async( ) 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=CapabilityName.MULTI_TURN in verified_caps, - supports_multi_message_pieces=CapabilityName.MULTI_MESSAGE_PIECES in verified_caps, - supports_json_schema=CapabilityName.JSON_SCHEMA in verified_caps, - supports_json_output=CapabilityName.JSON_OUTPUT in verified_caps, + 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=CapabilityName.SYSTEM_PROMPT in verified_caps, + supports_system_prompt=_resolve(CapabilityName.SYSTEM_PROMPT), input_modalities=frozenset(verified_modalities), output_modalities=declared.output_modalities, ) diff --git a/tests/unit/prompt_target/test_query_target_capabilities.py b/tests/unit/prompt_target/test_query_target_capabilities.py index 8ff7e150a8..df12f4763a 100644 --- a/tests/unit/prompt_target/test_query_target_capabilities.py +++ b/tests/unit/prompt_target/test_query_target_capabilities.py @@ -707,3 +707,57 @@ async def test_verify_target_async_forwards_capabilities(self) -> None: # 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 From cb8fc68b17a05e824771ce2be823ad0d71a82f35 Mon Sep 17 00:00:00 2001 From: hannahwestra25 Date: Fri, 8 May 2026 17:16:32 -0400 Subject: [PATCH 12/27] pre-commit --- doc/code/targets/6_1_target_capabilities.ipynb | 4 ++-- pyrit/prompt_target/common/query_target_capabilities.py | 6 +++++- 2 files changed, 7 insertions(+), 3 deletions(-) diff --git a/doc/code/targets/6_1_target_capabilities.ipynb b/doc/code/targets/6_1_target_capabilities.ipynb index ee25c333a7..11393fc1a8 100644 --- a/doc/code/targets/6_1_target_capabilities.ipynb +++ b/doc/code/targets/6_1_target_capabilities.ipynb @@ -45,8 +45,8 @@ "name": "stdout", "output_type": "stream", "text": [ - "Found default environment files: ['/home/vscode/.pyrit/.env']\n", - "Loaded environment file: /home/vscode/.pyrit/.env\n" + "Found default environment files: ['./.pyrit/.env']\n", + "Loaded environment file: ./.pyrit/.env\n" ] }, { diff --git a/pyrit/prompt_target/common/query_target_capabilities.py b/pyrit/prompt_target/common/query_target_capabilities.py index 8e7170c4ed..9a8d18fa68 100644 --- a/pyrit/prompt_target/common/query_target_capabilities.py +++ b/pyrit/prompt_target/common/query_target_capabilities.py @@ -358,7 +358,11 @@ async def _probe_multi_turn_async(target: PromptTarget, timeout_s: float, retrie 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)" + target=target, + message=Message([second]), + timeout_s=timeout_s, + retries=retries, + label="Multi-turn probe (turn 2)", ) From d085ee42cd25c1cf0dff79668687f781947eb18f Mon Sep 17 00:00:00 2001 From: hannahwestra25 Date: Mon, 11 May 2026 10:36:20 -0400 Subject: [PATCH 13/27] address comments --- doc/code/targets/0_prompt_targets.md | 10 +-- .../targets/6_1_target_capabilities.ipynb | 57 +++++++-------- doc/code/targets/6_1_target_capabilities.py | 50 +++++++------- pyrit/prompt_target/__init__.py | 8 +-- .../common/query_target_capabilities.py | 69 +++++++++++-------- .../test_query_target_capabilities.py | 50 +++++++------- 6 files changed, 131 insertions(+), 113 deletions(-) diff --git a/doc/code/targets/0_prompt_targets.md b/doc/code/targets/0_prompt_targets.md index 9b1b5dab39..4d866a0a8b 100644 --- a/doc/code/targets/0_prompt_targets.md +++ b/doc/code/targets/0_prompt_targets.md @@ -114,16 +114,16 @@ Declared capabilities describe what a target *should* support. For deployments w ```python from pyrit.prompt_target import ( query_target_capabilities_async, - verify_target_async, - verify_target_modalities_async, + query_target_async, + query_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) +queried_caps = await query_target_capabilities_async(target=target) +queried_modalities = await query_target_modalities_async(target=target) # Or do both at once and get a populated TargetCapabilities back: -verified = await verify_target_async(target=target) +queried = await query_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. diff --git a/doc/code/targets/6_1_target_capabilities.ipynb b/doc/code/targets/6_1_target_capabilities.ipynb index 11393fc1a8..27e3e25305 100644 --- a/doc/code/targets/6_1_target_capabilities.ipynb +++ b/doc/code/targets/6_1_target_capabilities.ipynb @@ -469,8 +469,8 @@ "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", + "`query_target_capabilities_async`, `query_target_modalities_async`, or the convenience wrapper\n", + "`query_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", @@ -479,7 +479,7 @@ "`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", + "`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", @@ -495,10 +495,10 @@ "Typical usage against a real endpoint:\n", "\n", "```python\n", - "from pyrit.prompt_target import verify_target_async\n", + "from pyrit.prompt_target import query_target_async\n", "\n", - "verified = await verify_target_async(target=target)\n", - "print(verified)\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", @@ -517,7 +517,7 @@ "name": "stdout", "output_type": "stream", "text": [ - "verified capabilities:\n", + "queried capabilities:\n", " - supports_editable_history\n", " - supports_json_output\n", " - supports_json_schema\n", @@ -532,8 +532,8 @@ "\n", "from pyrit.models import MessagePiece\n", "from pyrit.prompt_target import (\n", + " query_target_async,\n", " query_target_capabilities_async,\n", - " verify_target_async,\n", ")\n", "\n", "\n", @@ -556,9 +556,9 @@ "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", + "queried = await query_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}\")" ] }, @@ -572,13 +572,13 @@ "```python\n", "from pyrit.prompt_target.common.target_capabilities import CapabilityName\n", "\n", - "verified = await query_target_capabilities_async(\n", + "queried = 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", + "`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." @@ -594,7 +594,7 @@ "name": "stdout", "output_type": "stream", "text": [ - "verify_target_async result:\n", + "query_target_async result:\n", " supports_multi_turn: True\n", " supports_system_prompt: True\n", " supports_multi_message_pieces: True\n", @@ -607,14 +607,14 @@ "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)}\")" + "queried_caps = await query_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)}\")" ] }, { @@ -624,13 +624,13 @@ "source": [ "### Discovering undeclared modalities\n", "\n", - "By default `verify_target_async` only probes modality combinations the target already\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=` (and matching\n", "`test_assets=`) explicitly to probe combinations beyond the declared baseline:\n", "\n", "```python\n", - "verified = await verify_target_async(\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", @@ -639,12 +639,12 @@ "\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", + "`False` — narrowing controls *what is re-queried*, 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", + "# 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", @@ -653,6 +653,9 @@ } ], "metadata": { + "jupytext": { + "main_language": "python" + }, "language_info": { "codemirror_mode": { "name": "ipython", diff --git a/doc/code/targets/6_1_target_capabilities.py b/doc/code/targets/6_1_target_capabilities.py index 485bef3477..91405a7433 100644 --- a/doc/code/targets/6_1_target_capabilities.py +++ b/doc/code/targets/6_1_target_capabilities.py @@ -252,8 +252,8 @@ # 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`, `query_target_modalities_async`, or the convenience wrapper +# `query_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 @@ -262,7 +262,7 @@ # `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 +# `query_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 @@ -278,10 +278,10 @@ # Typical usage against a real endpoint: # # ```python -# from pyrit.prompt_target import verify_target_async +# from pyrit.prompt_target import query_target_async # -# verified = await verify_target_async(target=target) -# print(verified) +# queried = await query_target_async(target=target) +# print(queried) # ``` # # Below we mock the target's underlying transport (`_send_prompt_to_target_async`) so the notebook @@ -294,8 +294,8 @@ from pyrit.models import MessagePiece from pyrit.prompt_target import ( + query_target_async, query_target_capabilities_async, - verify_target_async, ) @@ -318,9 +318,9 @@ def _ok_response(): 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): +queried = await query_target_capabilities_async(target=probe_target, per_probe_timeout_s=5.0) # type: ignore +print("queried capabilities:") +for capability in sorted(queried, key=lambda c: c.value): print(f" - {capability.value}") # %% [markdown] @@ -329,13 +329,13 @@ def _ok_response(): # ```python # from pyrit.prompt_target.common.target_capabilities import CapabilityName # -# verified = await query_target_capabilities_async( +# queried = 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 +# `query_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. @@ -343,25 +343,25 @@ def _ok_response(): # %% 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)}") +queried_caps = await query_target_async(target=probe_target, per_probe_timeout_s=5.0) # type: ignore +print("query_target_async result:") +print(f" supports_multi_turn: {queried_caps.supports_multi_turn}") +print(f" supports_system_prompt: {queried_caps.supports_system_prompt}") +print(f" supports_multi_message_pieces: {queried_caps.supports_multi_message_pieces}") +print(f" supports_json_output: {queried_caps.supports_json_output}") +print(f" supports_json_schema: {queried_caps.supports_json_schema}") +print(f" input_modalities: {sorted(sorted(m) for m in queried_caps.input_modalities)}") # %% [markdown] # ### Discovering undeclared modalities # -# By default `verify_target_async` only probes modality combinations the target already +# By default `query_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( +# queried = await query_target_async( # target=target, # test_modalities={frozenset({"text"}), frozenset({"text", "image_path"})}, # test_assets={"image_path": "/path/to/test_image.png"}, @@ -370,12 +370,12 @@ def _ok_response(): # # 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 +# `False` — narrowing controls *what is re-queried*, 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( +# # Re-query only JSON support; other declared flags pass through unchanged. +# queried = await query_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 ef682b2a44..cb46c577c4 100644 --- a/pyrit/prompt_target/__init__.py +++ b/pyrit/prompt_target/__init__.py @@ -17,9 +17,9 @@ 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_async, query_target_capabilities_async, - verify_target_async, - verify_target_modalities_async, + query_target_modalities_async, ) from pyrit.prompt_target.common.target_capabilities import ( CapabilityHandlingPolicy, @@ -109,7 +109,7 @@ def __getattr__(name: str) -> object: "TargetRequirements", "UnsupportedCapabilityBehavior", "TextTarget", - "verify_target_async", - "verify_target_modalities_async", + "query_target_async", + "query_target_modalities_async", "WebSocketCopilotTarget", ] diff --git a/pyrit/prompt_target/common/query_target_capabilities.py b/pyrit/prompt_target/common/query_target_capabilities.py index 9a8d18fa68..09d29a4a28 100644 --- a/pyrit/prompt_target/common/query_target_capabilities.py +++ b/pyrit/prompt_target/common/query_target_capabilities.py @@ -13,7 +13,7 @@ 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 +* :func:`query_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``. @@ -425,7 +425,7 @@ async def _probe_json_schema_async(target: PromptTarget, timeout_s: float, retri ) -# Registry of capabilities that can be verified via a live API call. +# Registry of capabilities that can be queried 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, @@ -494,13 +494,13 @@ async def query_target_capabilities_async( Defaults to 1. Returns: - set[CapabilityName]: The capabilities verified to work against the target. + set[CapabilityName]: The capabilities confirmed to work against the target. """ capabilities_to_check: list[CapabilityName] = ( list(capabilities) if capabilities is not None else list(CapabilityName) ) - verified: set[CapabilityName] = set() + queried: set[CapabilityName] = set() with _permissive_configuration(target=target): for capability in capabilities_to_check: probe = _CAPABILITY_PROBES.get(capability) @@ -511,7 +511,7 @@ async def query_target_capabilities_async( try: if await probe(target, per_probe_timeout_s, retries): - verified.add(capability) + queried.add(capability) except Exception as exc: logger.debug("Probe for %s raised: %s", capability.value, exc) @@ -522,24 +522,24 @@ async def query_target_capabilities_async( # 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) + queried.add(capability) - return verified + return queried # --------------------------------------------------------------------------- -# Modality verification +# Modality query # --------------------------------------------------------------------------- # 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 +# :func:`query_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( +async def query_target_modalities_async( *, target: PromptTarget, test_modalities: set[frozenset[PromptDataType]] | None = None, @@ -569,7 +569,13 @@ async def verify_target_modalities_async( .. warning:: This function is **not safe to call concurrently** with other operations on the same ``target`` instance. It temporarily mutates - ``target._configuration``. + ``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. Args: target (PromptTarget): The target to probe. @@ -589,7 +595,7 @@ async def verify_target_modalities_async( Defaults to 1. Returns: - set[frozenset[PromptDataType]]: The modality combinations verified + set[frozenset[PromptDataType]]: The modality combinations confirmed to work against the target. """ if test_modalities is None: @@ -598,7 +604,7 @@ async def verify_target_modalities_async( assets = test_assets if test_assets is not None else DEFAULT_TEST_ASSETS - verified: set[frozenset[PromptDataType]] = set() + queried: set[frozenset[PromptDataType]] = set() with _permissive_configuration(target=target, extra_input_modalities=test_modalities): for combination in test_modalities: try: @@ -617,12 +623,12 @@ async def verify_target_modalities_async( retries=retries, label=f"Modality probe {sorted(combination)}", ): - verified.add(combination) + queried.add(combination) - return verified + return queried -async def verify_target_async( +async def query_target_async( *, target: PromptTarget, per_probe_timeout_s: float = DEFAULT_PROBE_TIMEOUT_SECONDS, @@ -635,8 +641,8 @@ async def verify_target_async( 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 + :func:`query_target_modalities_async` and returns a + :class:`TargetCapabilities` populated from the queried results, so callers don't need to assemble the dataclass themselves. Boolean capability flags not covered by @@ -644,7 +650,7 @@ async def verify_target_async( 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 + ``False`` — narrowing controls *what is re-queried*, not what the returned dataclass reports. .. warning:: @@ -655,16 +661,25 @@ async def verify_target_async( matching ``test_assets=``) explicitly to probe combinations beyond the declared baseline. + .. 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``. 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. + 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 + :func:`query_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`. + modality to a file path. See :func:`query_target_modalities_async`. capabilities (Iterable[CapabilityName] | None): Capabilities to probe. See :func:`query_target_capabilities_async`. Defaults to every member of :class:`CapabilityName`. @@ -674,18 +689,18 @@ async def verify_target_async( Defaults to 1. Returns: - TargetCapabilities: A dataclass reflecting verified capabilities and + TargetCapabilities: A dataclass reflecting queried capabilities and modalities. ``output_modalities`` is copied from ``target.capabilities.output_modalities`` because outputs cannot be - verified by sending a request. + queried by sending a request. """ - verified_caps = await query_target_capabilities_async( + queried_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( + queried_modalities = await query_target_modalities_async( target=target, test_modalities=test_modalities, test_assets=test_assets, @@ -701,7 +716,7 @@ async def verify_target_async( def _resolve(name: CapabilityName) -> bool: if name in probed: - return name in verified_caps + return name in queried_caps return bool(getattr(declared, name.value)) return TargetCapabilities( @@ -711,7 +726,7 @@ def _resolve(name: CapabilityName) -> bool: 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), + input_modalities=frozenset(queried_modalities), output_modalities=declared.output_modalities, ) diff --git a/tests/unit/prompt_target/test_query_target_capabilities.py b/tests/unit/prompt_target/test_query_target_capabilities.py index df12f4763a..eb0ae47002 100644 --- a/tests/unit/prompt_target/test_query_target_capabilities.py +++ b/tests/unit/prompt_target/test_query_target_capabilities.py @@ -14,9 +14,9 @@ _CAPABILITY_PROBES, _create_test_message, _permissive_configuration, + query_target_async, query_target_capabilities_async, - verify_target_async, - verify_target_modalities_async, + query_target_modalities_async, ) from pyrit.prompt_target.common.target_capabilities import ( CapabilityName, @@ -375,7 +375,7 @@ async def _send_prompt_to_target_async(self, *, normalized_conversation: list[Me # --------------------------------------------------------------------------- -# Modality verification tests +# Modality query tests # --------------------------------------------------------------------------- @@ -441,7 +441,7 @@ async def test_all_combinations_supported(self) -> None: _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) + result = await query_target_modalities_async(target=target) assert frozenset({"text"}) in result @@ -450,7 +450,7 @@ async def test_exception_excludes_combination(self) -> None: _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) + result = await query_target_modalities_async(target=target) assert result == set() @@ -459,7 +459,7 @@ async def test_error_response_excludes_combination(self) -> None: _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) + result = await query_target_modalities_async(target=target) assert result == set() @@ -479,7 +479,7 @@ async def selective_send(*, normalized_conversation: list[Message]) -> list[Mess target._send_prompt_to_target_async = selective_send # type: ignore[method-assign] - result = await verify_target_modalities_async( + result = await query_target_modalities_async( target=target, test_assets={"image_path": image_asset}, ) @@ -493,7 +493,7 @@ async def test_explicit_test_modalities_overrides_declared(self, image_asset: st _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( + result = await query_target_modalities_async( target=target, test_modalities={frozenset({"text"}), frozenset({"text", "image_path"})}, test_assets={"image_path": image_asset}, @@ -508,7 +508,7 @@ async def test_combination_skipped_when_asset_missing(self, tmp_path: Path) -> N 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) + result = await query_target_modalities_async(target=target) assert result == set() assert target._send_prompt_to_target_async.await_count == 0 @@ -524,7 +524,7 @@ async def test_explicit_test_modalities_runs_under_permissive_configuration(self 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( + result = await query_target_modalities_async( target=target, test_modalities={frozenset({"text", "image_path"})}, test_assets={"image_path": image_asset}, @@ -541,7 +541,7 @@ 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. + is excluded from the queried set. """ target = MockPromptTarget() @@ -589,9 +589,9 @@ async def test_returns_false_when_memory_seed_raises(self) -> None: class TestVerifyTargetAsync: async def test_returns_target_capabilities_assembled_from_probes(self) -> None: """ - ``verify_target_async`` runs both the capability and modality probes + ``query_target_async`` runs both the capability and modality probes and assembles a :class:`TargetCapabilities` populated from the - verified results, copying ``supports_editable_history`` and + queried results, copying ``supports_editable_history`` and ``output_modalities`` from the target's declared capabilities. """ declared = TargetCapabilities( @@ -603,7 +603,7 @@ async def test_returns_target_capabilities_assembled_from_probes(self) -> None: 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) + result = await query_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 @@ -621,7 +621,7 @@ async def test_returns_target_capabilities_assembled_from_probes(self) -> None: 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`` + queried, but ``supports_editable_history`` and ``output_modalities`` are still copied from the declared capabilities. """ declared = TargetCapabilities( @@ -632,7 +632,7 @@ async def test_excludes_capabilities_when_probe_send_fails(self) -> None: 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) + result = await query_target_async(target=target, per_probe_timeout_s=0.5) assert result.supports_multi_turn is False assert result.supports_system_prompt is False @@ -641,7 +641,7 @@ async def test_excludes_capabilities_when_probe_send_fails(self) -> None: 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. + # No modalities queried because send always fails. assert result.input_modalities == frozenset() # Output modalities still copied. assert result.output_modalities == declared.output_modalities @@ -675,14 +675,14 @@ async def test_response_with_no_pieces_treated_as_failure(self) -> None: assert result == set() - async def test_verify_target_async_forwards_test_modalities(self, image_asset: str) -> None: + async def test_query_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( + result = await query_target_async( target=target, test_modalities={extra_combo}, test_assets={"image_path": image_asset}, @@ -692,12 +692,12 @@ async def test_verify_target_async_forwards_test_modalities(self, image_asset: s # 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.""" + async def test_query_target_async_forwards_capabilities(self) -> None: + """``query_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( + await query_target_async( target=target, capabilities={CapabilityName.JSON_OUTPUT}, per_probe_timeout_s=2.0, @@ -708,7 +708,7 @@ async def test_verify_target_async_forwards_capabilities(self) -> None: # 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: + async def test_query_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 @@ -724,13 +724,13 @@ async def test_verify_target_async_preserves_declared_when_capabilities_narrowed 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( + result = await query_target_async( target=target, capabilities={CapabilityName.JSON_OUTPUT}, per_probe_timeout_s=2.0, ) - # The probed capability reflects the verified result. + # The probed capability reflects the queried result. assert result.supports_json_output is True # Non-probed capabilities fall back to declared values. assert result.supports_multi_turn is True From d892f06ddcad3eac5e11fa1f4e612bc5a6870ebf Mon Sep 17 00:00:00 2001 From: hannahwestra25 Date: Mon, 11 May 2026 11:30:44 -0400 Subject: [PATCH 14/27] fix docstring --- .../common/query_target_capabilities.py | 15 +++++++++++++++ 1 file changed, 15 insertions(+) diff --git a/pyrit/prompt_target/common/query_target_capabilities.py b/pyrit/prompt_target/common/query_target_capabilities.py index 09d29a4a28..ca01da68a8 100644 --- a/pyrit/prompt_target/common/query_target_capabilities.py +++ b/pyrit/prompt_target/common/query_target_capabilities.py @@ -256,6 +256,9 @@ async def _probe_system_prompt_async(target: PromptTarget, timeout_s: float, ret Args: target (PromptTarget): The target to probe. timeout_s (float): Per-attempt timeout in seconds. + retries (int): Number of additional attempts after the first failure. + Only exceptions/timeouts are retried; an explicit error response + is final. Defaults to 1. Returns: bool: ``True`` if the system + user request succeeded; ``False`` otherwise. @@ -290,6 +293,9 @@ async def _probe_multi_message_pieces_async(target: PromptTarget, timeout_s: flo Args: target (PromptTarget): The target to probe. timeout_s (float): Per-attempt timeout in seconds. + retries (int): Number of additional attempts after the first failure. + Only exceptions/timeouts are retried; an explicit error response + is final. Defaults to 1. Returns: bool: ``True`` if the multi-piece request succeeded; ``False`` otherwise. @@ -330,6 +336,9 @@ async def _probe_multi_turn_async(target: PromptTarget, timeout_s: float, retrie Args: target (PromptTarget): The target to probe. timeout_s (float): Per-attempt timeout in seconds. + retries (int): Number of additional attempts after the first failure. + Only exceptions/timeouts are retried; an explicit error response + is final. Defaults to 1. Returns: bool: ``True`` if both turns succeeded; ``False`` if either turn failed. @@ -373,6 +382,9 @@ async def _probe_json_output_async(target: PromptTarget, timeout_s: float, retri Args: target (PromptTarget): The target to probe. timeout_s (float): Per-attempt timeout in seconds. + retries (int): Number of additional attempts after the first failure. + Only exceptions/timeouts are retried; an explicit error response + is final. Defaults to 1. Returns: bool: ``True`` if the JSON-mode request succeeded; ``False`` otherwise. @@ -397,6 +409,9 @@ async def _probe_json_schema_async(target: PromptTarget, timeout_s: float, retri Args: target (PromptTarget): The target to probe. timeout_s (float): Per-attempt timeout in seconds. + retries (int): Number of additional attempts after the first failure. + Only exceptions/timeouts are retried; an explicit error response + is final. Defaults to 1. Returns: bool: ``True`` if the schema-constrained request succeeded; ``False`` otherwise. From 0f5be0a5b77ade20e5a14cf82b3c4c911eab83d2 Mon Sep 17 00:00:00 2001 From: hannahwestra25 Date: Mon, 11 May 2026 15:26:45 -0400 Subject: [PATCH 15/27] clean up comments --- doc/code/targets/0_prompt_targets.md | 6 +- .../targets/6_1_target_capabilities.ipynb | 26 ++- doc/code/targets/6_1_target_capabilities.py | 15 +- .../common/query_target_capabilities.py | 182 +++++------------- .../test_query_target_capabilities.py | 87 ++++++++- 5 files changed, 163 insertions(+), 153 deletions(-) diff --git a/doc/code/targets/0_prompt_targets.md b/doc/code/targets/0_prompt_targets.md index 4d866a0a8b..3922c18927 100644 --- a/doc/code/targets/0_prompt_targets.md +++ b/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: @@ -122,11 +122,11 @@ from pyrit.prompt_target import ( queried_caps = await query_target_capabilities_async(target=target) queried_modalities = await query_target_modalities_async(target=target) -# Or do both at once and get a populated TargetCapabilities back: +# Or do both at once and get a best-effort TargetCapabilities back: queried = await query_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. +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. `query_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 diff --git a/doc/code/targets/6_1_target_capabilities.ipynb b/doc/code/targets/6_1_target_capabilities.ipynb index 27e3e25305..cbe1d02ddb 100644 --- a/doc/code/targets/6_1_target_capabilities.ipynb +++ b/doc/code/targets/6_1_target_capabilities.ipynb @@ -396,8 +396,21 @@ "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", + "`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." ] }, { @@ -470,7 +483,7 @@ "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 populated `TargetCapabilities`.\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", @@ -483,9 +496,10 @@ "`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", + "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", diff --git a/doc/code/targets/6_1_target_capabilities.py b/doc/code/targets/6_1_target_capabilities.py index 91405a7433..d1cf9aa9f2 100644 --- a/doc/code/targets/6_1_target_capabilities.py +++ b/doc/code/targets/6_1_target_capabilities.py @@ -253,7 +253,7 @@ # 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`, `query_target_modalities_async`, or the convenience wrapper -# `query_target_async` that runs both and returns a populated `TargetCapabilities`. +# `query_target_async` that runs both and returns a best-effort `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 @@ -266,9 +266,10 @@ # `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. +# transient errors before being declared failed. `query_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 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 @@ -336,9 +337,9 @@ def _ok_response(): # ``` # # `query_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. +# probes and assembles a best-effort `TargetCapabilities` you can drop into a +# `TargetConfiguration`, so the rest of PyRIT operates on probed values where available and +# declared values otherwise. # %% probe_target._send_prompt_to_target_async = AsyncMock(return_value=_ok_response()) # type: ignore[method-assign] diff --git a/pyrit/prompt_target/common/query_target_capabilities.py b/pyrit/prompt_target/common/query_target_capabilities.py index ca01da68a8..9e26c61a20 100644 --- a/pyrit/prompt_target/common/query_target_capabilities.py +++ b/pyrit/prompt_target/common/query_target_capabilities.py @@ -10,9 +10,9 @@ 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`. + the capability is included in the returned set. Capabilities without a + registered probe fall back to the target's declared native support from + ``target.capabilities``. * :func:`query_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``. @@ -51,10 +51,8 @@ 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 @@ -74,14 +72,6 @@ _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. @@ -129,10 +119,13 @@ def _permissive_configuration( supports_system_prompt=True, input_modalities=merged_modalities, ) - target._configuration = TargetConfiguration( - capabilities=permissive_caps, - policy=_PROBE_POLICY, - ) + probe_configuration = object.__new__(TargetConfiguration) + probe_configuration._capabilities = permissive_caps + probe_configuration._policy = original.policy + # Keep the original normalization pipeline intact so probing exercises the + # target's real request shaping, including custom normalizer overrides. + probe_configuration._pipeline = original.pipeline + target._configuration = probe_configuration try: yield finally: @@ -459,42 +452,10 @@ async def query_target_capabilities_async( 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. + Probe which capabilities ``target`` accepts. + + Registered capabilities are checked with live requests. Capabilities + without a live probe fall back to declared native support. Args: target (PromptTarget): The target to probe. @@ -520,21 +481,21 @@ async def query_target_capabilities_async( 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. + # Capabilities without a probe are handled after the permissive + # override is removed so we can read the target's native flags. continue try: + # "Supported" means the request was accepted. A target can + # still ignore the feature semantics after accepting the call. if await probe(target, per_probe_timeout_s, retries): queried.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. + # Read unprobed capabilities from target.capabilities, not + # target.configuration, so ADAPTed behavior is not reported as native + # support. for capability in capabilities_to_check: if capability not in _CAPABILITY_PROBES and target.capabilities.includes(capability=capability): queried.add(capability) @@ -563,34 +524,10 @@ async def query_target_modalities_async( 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`` 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. + Probe which input modality combinations ``target`` accepts. + + Each modality combination is checked with a minimal request built from the + supplied test assets. Args: target (PromptTarget): The target to probe. @@ -625,12 +562,15 @@ async def query_target_modalities_async( try: message = _create_test_message(modalities=combination, test_assets=assets) except FileNotFoundError as exc: + # Skip combinations we cannot construct a valid probe payload for. logger.info("Skipping modality %s: %s", combination, exc) continue except ValueError as exc: logger.info("Skipping modality %s: %s", combination, exc) continue + # "Supported" means the request was accepted. A target may still + # ignore the non-text payload after accepting it. if await _send_and_check_async( target=target, message=message, @@ -653,37 +593,11 @@ async def query_target_async( retries: int = 1, ) -> TargetCapabilities: """ - Probe both capabilities and modalities and return a combined result. - - Calls :func:`query_target_capabilities_async` and - :func:`query_target_modalities_async` and returns a - :class:`TargetCapabilities` populated from the queried 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-queried*, 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. - - .. 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``. 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. + Probe capabilities and modalities and return a merged result. + + This wraps :func:`query_target_capabilities_async` and + :func:`query_target_modalities_async` and returns a best-effort + :class:`TargetCapabilities`. Args: target (PromptTarget): The target to probe. @@ -704,14 +618,14 @@ async def query_target_async( Defaults to 1. Returns: - TargetCapabilities: A dataclass reflecting queried capabilities and - modalities. ``output_modalities`` is copied from - ``target.capabilities.output_modalities`` because outputs cannot be - queried by sending a request. + TargetCapabilities: A merged capability view: probed where possible, + declared where probing is unavailable or out of scope. """ + capabilities_to_probe = list(capabilities) if capabilities is not None else None + queried_caps = await query_target_capabilities_async( target=target, - capabilities=capabilities, + capabilities=capabilities_to_probe, per_probe_timeout_s=per_probe_timeout_s, retries=retries, ) @@ -724,24 +638,32 @@ async def query_target_async( ) 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) + # If the caller narrows the capability set, leave the rest at their + # declared values instead of silently forcing them to False. + probed: set[CapabilityName] = ( + set(capabilities_to_probe) if capabilities_to_probe is not None else set(CapabilityName) + ) def _resolve(name: CapabilityName) -> bool: if name in probed: return name in queried_caps return bool(getattr(declared, name.value)) + resolved_multi_turn = _resolve(CapabilityName.MULTI_TURN) + # Editable history is only meaningful if multi-turn probing/declaration + # also resolved to True. + resolved_editable_history = declared.supports_editable_history and resolved_multi_turn + return TargetCapabilities( - supports_multi_turn=_resolve(CapabilityName.MULTI_TURN), + supports_multi_turn=resolved_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_editable_history=resolved_editable_history, supports_system_prompt=_resolve(CapabilityName.SYSTEM_PROMPT), input_modalities=frozenset(queried_modalities), + # Output modalities are still declarative because probing them would + # require target-specific response inspection. output_modalities=declared.output_modalities, ) diff --git a/tests/unit/prompt_target/test_query_target_capabilities.py b/tests/unit/prompt_target/test_query_target_capabilities.py index eb0ae47002..09ae588076 100644 --- a/tests/unit/prompt_target/test_query_target_capabilities.py +++ b/tests/unit/prompt_target/test_query_target_capabilities.py @@ -19,8 +19,10 @@ query_target_modalities_async, ) from pyrit.prompt_target.common.target_capabilities import ( + CapabilityHandlingPolicy, CapabilityName, TargetCapabilities, + UnsupportedCapabilityBehavior, ) from pyrit.prompt_target.common.target_configuration import TargetConfiguration from tests.unit.mocks import MockPromptTarget @@ -355,6 +357,33 @@ async def test_probes_run_under_permissive_configuration(self) -> None: assert send_mock.await_count >= 1 assert CapabilityName.MULTI_MESSAGE_PIECES in result + async def test_probed_capability_excluded_when_only_adapted(self) -> None: + target = MockPromptTarget() + target._configuration = TargetConfiguration( + capabilities=TargetCapabilities(supports_system_prompt=False), + policy=CapabilityHandlingPolicy( + behaviors={ + CapabilityName.SYSTEM_PROMPT: UnsupportedCapabilityBehavior.ADAPT, + CapabilityName.MULTI_TURN: UnsupportedCapabilityBehavior.RAISE, + } + ), + ) + + async def reject_system_roles(*, normalized_conversation: list[Message]) -> list[Message]: + roles = [piece.role for message in normalized_conversation for piece in message.message_pieces] + if "system" in roles: + raise RuntimeError("system messages are not natively supported") + return _ok_response() + + target._send_prompt_to_target_async = AsyncMock(side_effect=reject_system_roles) # type: ignore[method-assign] + + result = await query_target_capabilities_async( + target=target, + capabilities={CapabilityName.SYSTEM_PROMPT}, + ) + + assert result == set() + @pytest.mark.usefixtures("patch_central_database") class TestQueryTargetCapabilitiesIsolatedTarget: @@ -591,8 +620,8 @@ async def test_returns_target_capabilities_assembled_from_probes(self) -> None: """ ``query_target_async`` runs both the capability and modality probes and assembles a :class:`TargetCapabilities` populated from the - queried results, copying ``supports_editable_history`` and - ``output_modalities`` from the target's declared capabilities. + queried results, copying ``output_modalities`` from the target's + declared capabilities and deriving editable history conservatively. """ declared = TargetCapabilities( supports_editable_history=True, @@ -611,8 +640,9 @@ async def test_returns_target_capabilities_assembled_from_probes(self) -> None: 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 + # Editable history is conservative and therefore cannot remain true + # when multi-turn support was not confirmed by probing. + assert result.supports_editable_history is False # Modalities returned from the modality probe (text combination). assert frozenset({"text"}) in result.input_modalities # Output modalities copied through (not probed). @@ -622,7 +652,7 @@ async def test_excludes_capabilities_when_probe_send_fails(self) -> None: """ When the underlying send raises, no capability or modality is queried, but ``supports_editable_history`` and ``output_modalities`` - are still copied from the declared capabilities. + are still copied conservatively from the declared capabilities. """ declared = TargetCapabilities( supports_editable_history=True, @@ -639,8 +669,9 @@ async def test_excludes_capabilities_when_probe_send_fails(self) -> None: 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 + # Editable history is derived conservatively and must fall when + # multi-turn probing disproves the prerequisite capability. + assert result.supports_editable_history is False # No modalities queried because send always fails. assert result.input_modalities == frozenset() # Output modalities still copied. @@ -738,6 +769,48 @@ async def test_query_target_async_preserves_declared_when_capabilities_narrowed( assert result.supports_json_schema is True assert result.supports_editable_history is True + async def test_query_target_async_drops_editable_history_when_multi_turn_probe_fails(self) -> None: + """Editable history must not remain true when probing disproves multi-turn support.""" + declared = TargetCapabilities( + supports_multi_turn=True, + supports_editable_history=True, + output_modalities=frozenset({frozenset({"text"})}), + ) + target = MockPromptTarget() + target._configuration = TargetConfiguration(capabilities=declared) + + async def selective_send(*, normalized_conversation: list[Message]) -> list[Message]: + latest_text = normalized_conversation[-1].message_pieces[0].original_value + if latest_text == "My favorite color is blue." or latest_text == "What did I just tell you?": + raise RuntimeError("multi-turn unsupported") + return _ok_response() + + target._send_prompt_to_target_async = AsyncMock(side_effect=selective_send) # type: ignore[method-assign] + + result = await query_target_async(target=target, per_probe_timeout_s=2.0) + + assert result.supports_multi_turn is False + assert result.supports_editable_history is False + + async def test_query_target_async_accepts_single_pass_iterable(self) -> None: + declared = TargetCapabilities( + supports_multi_turn=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] + + gen = (c for c in [CapabilityName.JSON_OUTPUT, CapabilityName.EDITABLE_HISTORY]) + result = await query_target_async( + target=target, + capabilities=gen, + per_probe_timeout_s=2.0, + ) + + assert result.supports_json_output is True + assert result.supports_editable_history is True + @pytest.mark.usefixtures("patch_central_database") class TestMultiTurnProbeMemoryFailure: From a347b1c5f8a1b24215f849e1dde0edb06094454e Mon Sep 17 00:00:00 2001 From: hannahwestra25 Date: Mon, 11 May 2026 16:25:09 -0400 Subject: [PATCH 16/27] add assets and comments --- .../targets/6_1_target_capabilities.ipynb | 91 +++++++++++++----- doc/code/targets/6_1_target_capabilities.py | 25 ++++- .../target_capabilities/probe_audio.wav | Bin 0 -> 44 bytes .../target_capabilities/probe_image.png | Bin 0 -> 68 bytes .../common/query_target_capabilities.py | 33 ++++--- .../test_query_target_capabilities.py | 15 ++- 6 files changed, 125 insertions(+), 39 deletions(-) create mode 100644 pyrit/datasets/prompt_target/target_capabilities/probe_audio.wav create mode 100644 pyrit/datasets/prompt_target/target_capabilities/probe_image.png diff --git a/doc/code/targets/6_1_target_capabilities.ipynb b/doc/code/targets/6_1_target_capabilities.ipynb index cbe1d02ddb..503d2b2d02 100644 --- a/doc/code/targets/6_1_target_capabilities.ipynb +++ b/doc/code/targets/6_1_target_capabilities.ipynb @@ -376,19 +376,39 @@ } ], "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", + " query_target_async,\n", + " query_target_capabilities_async,\n", + " query_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 query_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}\")" ] }, { @@ -407,6 +427,12 @@ ")\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", @@ -418,6 +444,26 @@ "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 query_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", @@ -436,7 +482,7 @@ }, { "cell_type": "markdown", - "id": "17", + "id": "18", "metadata": {}, "source": [ "## 6. Non-adaptable capabilities\n", @@ -450,7 +496,7 @@ { "cell_type": "code", "execution_count": null, - "id": "18", + "id": "19", "metadata": {}, "outputs": [ { @@ -474,7 +520,7 @@ }, { "cell_type": "markdown", - "id": "19", + "id": "20", "metadata": {}, "source": [ "## 7. Querying live target capabilities\n", @@ -524,7 +570,7 @@ { "cell_type": "code", "execution_count": null, - "id": "20", + "id": "21", "metadata": {}, "outputs": [ { @@ -578,7 +624,7 @@ }, { "cell_type": "markdown", - "id": "21", + "id": "22", "metadata": {}, "source": [ "To narrow the probe to specific capabilities (faster, fewer calls), pass `capabilities=`:\n", @@ -601,7 +647,7 @@ { "cell_type": "code", "execution_count": null, - "id": "22", + "id": "23", "metadata": {}, "outputs": [ { @@ -633,15 +679,16 @@ }, { "cell_type": "markdown", - "id": "23", + "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=` (and matching\n", - "`test_assets=`) explicitly to probe combinations beyond the declared baseline:\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", diff --git a/doc/code/targets/6_1_target_capabilities.py b/doc/code/targets/6_1_target_capabilities.py index d1cf9aa9f2..68983be2bf 100644 --- a/doc/code/targets/6_1_target_capabilities.py +++ b/doc/code/targets/6_1_target_capabilities.py @@ -297,6 +297,7 @@ from pyrit.prompt_target import ( query_target_async, query_target_capabilities_async, + query_target_modalities_async, ) @@ -336,6 +337,12 @@ def _ok_response(): # ) # ``` # +# If you only care about accepted input combinations, call +# `query_target_modalities_async` directly. The example below uses the +# packaged default probe assets for the non-text modalities PyRIT ships. +# Pass `test_assets=` only when you want to override those defaults or probe +# a modality without a packaged asset. +# # `query_target_async` is the most common entry point: it runs both the capability and modality # probes and assembles a best-effort `TargetCapabilities` you can drop into a # `TargetConfiguration`, so the rest of PyRIT operates on probed values where available and @@ -344,6 +351,19 @@ def _ok_response(): # %% probe_target._send_prompt_to_target_async = AsyncMock(return_value=_ok_response()) # type: ignore[method-assign] +queried_modalities = await query_target_modalities_async( + target=probe_target, + test_modalities={frozenset({"text"}), frozenset({"text", "image_path"})}, + per_probe_timeout_s=5.0, +) # type: ignore + +print("query_target_modalities_async result:") +for combination in sorted(sorted(m) for m in queried_modalities): + print(f" - {combination}") + +# %% +probe_target._send_prompt_to_target_async = AsyncMock(return_value=_ok_response()) # type: ignore[method-assign] + queried_caps = await query_target_async(target=probe_target, per_probe_timeout_s=5.0) # type: ignore print("query_target_async result:") print(f" supports_multi_turn: {queried_caps.supports_multi_turn}") @@ -358,8 +378,9 @@ def _ok_response(): # # By default `query_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: +# claims text-only but might actually accept images, pass `test_modalities=` explicitly to +# probe combinations beyond the declared baseline. Provide `test_assets=` as well if you need +# to override the packaged defaults or probe a modality without one: # # ```python # queried = await query_target_async( diff --git a/pyrit/datasets/prompt_target/target_capabilities/probe_audio.wav b/pyrit/datasets/prompt_target/target_capabilities/probe_audio.wav new file mode 100644 index 0000000000000000000000000000000000000000..8dbde9545c9f3bc0f0edf6804a28471e5c5cc00f GIT binary patch literal 44 vcmWIYbaPW-U|Zci7-kcwN$fBwreFf%hTyr1>< QK2Vs!)78&qol`;+0Fh4*761SM literal 0 HcmV?d00001 diff --git a/pyrit/prompt_target/common/query_target_capabilities.py b/pyrit/prompt_target/common/query_target_capabilities.py index 9e26c61a20..dc961c612d 100644 --- a/pyrit/prompt_target/common/query_target_capabilities.py +++ b/pyrit/prompt_target/common/query_target_capabilities.py @@ -28,15 +28,11 @@ 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). + These probes only verify that a request was *accepted*. They do not prove + that the endpoint enforced the feature, and the JSON probes are only + meaningful for targets that translate ``prompt_metadata`` JSON hints into + provider request fields. Treat the results as an upper bound on support and + validate response content separately when that distinction matters. """ import asyncio @@ -48,6 +44,7 @@ from contextlib import contextmanager from dataclasses import replace +from pyrit.common.path import DATASETS_PATH from pyrit.models import Message, MessagePiece, PromptDataType from pyrit.prompt_target.common.prompt_target import PromptTarget from pyrit.prompt_target.common.target_capabilities import ( @@ -77,6 +74,9 @@ # data-type check does not reject text probes against text-less targets. _TEXT_MODALITY: frozenset[frozenset[PromptDataType]] = frozenset({frozenset({"text"})}) +# Packaged fallback assets for non-text modality probes. +_TARGET_CAPABILITIES_DATASET_PATH = DATASETS_PATH / "prompt_target" / "target_capabilities" + @contextmanager def _permissive_configuration( @@ -388,6 +388,8 @@ async def _probe_json_output_async(target: PromptTarget, timeout_s: float, retri original_value='Respond with a JSON object: {"ok": true}.', original_value_data_type="text", conversation_id=conversation_id, + # This only becomes a real JSON-mode request on targets that honor + # PyRIT's JSON metadata contract when building the provider payload. prompt_metadata=_probe_metadata({"response_format": "json"}), ) return await _send_and_check_async( @@ -421,6 +423,8 @@ async def _probe_json_schema_async(target: PromptTarget, timeout_s: float, retri original_value='Respond with a JSON object matching the schema: {"ok": true}.', original_value_data_type="text", conversation_id=conversation_id, + # As above, this probe is only strong for targets that map these + # metadata keys to native JSON-schema request parameters. prompt_metadata=_probe_metadata( { "response_format": "json", @@ -508,11 +512,14 @@ async def query_target_capabilities_async( # --------------------------------------------------------------------------- -# Default mapping of non-text modalities to test asset paths. Callers can +# Default mapping of non-text modalities to packaged probe assets. Callers can # override via the ``test_assets`` parameter of -# :func:`query_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] = {} +# :func:`query_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] = { + "audio_path": str(_TARGET_CAPABILITIES_DATASET_PATH / "probe_audio.wav"), + "image_path": str(_TARGET_CAPABILITIES_DATASET_PATH / "probe_image.png"), +} async def query_target_modalities_async( diff --git a/tests/unit/prompt_target/test_query_target_capabilities.py b/tests/unit/prompt_target/test_query_target_capabilities.py index 09ae588076..891cedd9cc 100644 --- a/tests/unit/prompt_target/test_query_target_capabilities.py +++ b/tests/unit/prompt_target/test_query_target_capabilities.py @@ -12,6 +12,7 @@ from pyrit.prompt_target.common.prompt_target import PromptTarget from pyrit.prompt_target.common.query_target_capabilities import ( _CAPABILITY_PROBES, + DEFAULT_TEST_ASSETS, _create_test_message, _permissive_configuration, query_target_async, @@ -430,6 +431,15 @@ def image_asset(tmp_path: Path) -> str: @pytest.mark.usefixtures("patch_central_database") class TestCreateTestMessage: + def test_default_assets_exist_for_packaged_modalities(self) -> None: + msg = _create_test_message( + modalities=frozenset({"audio_path", "image_path"}), + test_assets=DEFAULT_TEST_ASSETS, + ) + + types = {piece.original_value_data_type for piece in msg.message_pieces} + assert types == {"audio_path", "image_path"} + def test_text_only(self) -> None: msg = _create_test_message(modalities=frozenset({"text"}), test_assets={}) assert len(msg.message_pieces) == 1 @@ -536,8 +546,9 @@ async def test_combination_skipped_when_asset_missing(self, tmp_path: Path) -> N _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 query_target_modalities_async(target=target) + # An explicit empty mapping disables the packaged defaults, so + # image_path combinations are skipped instead of probed. + result = await query_target_modalities_async(target=target, test_assets={}) assert result == set() assert target._send_prompt_to_target_async.await_count == 0 From 2cbae682b7771962788e6a52f0dee019faaeb272 Mon Sep 17 00:00:00 2001 From: hannahwestra25 Date: Mon, 11 May 2026 16:55:47 -0400 Subject: [PATCH 17/27] fix bad test --- tests/unit/common/test_common_net_utility.py | 10 ++++++-- .../test_query_target_capabilities.py | 24 +++++++++---------- 2 files changed, 20 insertions(+), 14 deletions(-) diff --git a/tests/unit/common/test_common_net_utility.py b/tests/unit/common/test_common_net_utility.py index 58fff4b222..5088a166de 100644 --- a/tests/unit/common/test_common_net_utility.py +++ b/tests/unit/common/test_common_net_utility.py @@ -77,8 +77,14 @@ def response_callback(request): async def test_debug_is_false_by_default(): with patch("pyrit.common.net_utility.get_httpx_client") as mock_get_httpx_client: - mock_client_instance = MagicMock() - mock_get_httpx_client.return_value = mock_client_instance + mock_client_context = MagicMock() + mock_client = MagicMock() + mock_client.request = AsyncMock( + return_value=httpx.Response(status_code=200, request=httpx.Request("GET", "http://example.com")) + ) + mock_client_context.__aenter__ = AsyncMock(return_value=mock_client) + mock_client_context.__aexit__ = AsyncMock(return_value=None) + mock_get_httpx_client.return_value = mock_client_context await make_request_and_raise_if_error_async(endpoint_uri="http://example.com", method="GET") diff --git a/tests/unit/prompt_target/test_query_target_capabilities.py b/tests/unit/prompt_target/test_query_target_capabilities.py index 891cedd9cc..2980597306 100644 --- a/tests/unit/prompt_target/test_query_target_capabilities.py +++ b/tests/unit/prompt_target/test_query_target_capabilities.py @@ -4,7 +4,7 @@ import asyncio import json from pathlib import Path -from unittest.mock import AsyncMock, MagicMock +from unittest.mock import AsyncMock, MagicMock, patch import pytest @@ -12,9 +12,9 @@ from pyrit.prompt_target.common.prompt_target import PromptTarget from pyrit.prompt_target.common.query_target_capabilities import ( _CAPABILITY_PROBES, - DEFAULT_TEST_ASSETS, _create_test_message, _permissive_configuration, + DEFAULT_TEST_ASSETS, query_target_async, query_target_capabilities_async, query_target_modalities_async, @@ -611,14 +611,14 @@ async def test_returns_false_when_memory_seed_raises(self) -> None: 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}, - ) + with patch.object(target._memory, "add_message_to_memory", side_effect=RuntimeError("memory offline")): + result = await query_target_capabilities_async( + target=target, + capabilities={CapabilityName.SYSTEM_PROMPT}, + ) assert result == set() # The user send is never attempted because seeding failed. @@ -834,12 +834,12 @@ async def test_returns_false_when_history_seed_raises(self) -> None: 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}, - ) + with patch.object(target._memory, "add_message_to_memory", side_effect=RuntimeError("memory offline")): + 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 From 972777c7e3c7ce1d5cd04ac096f97928e19a5108 Mon Sep 17 00:00:00 2001 From: hannahwestra25 Date: Mon, 11 May 2026 17:08:01 -0400 Subject: [PATCH 18/27] pre-commit --- tests/unit/prompt_target/test_query_target_capabilities.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/tests/unit/prompt_target/test_query_target_capabilities.py b/tests/unit/prompt_target/test_query_target_capabilities.py index 2980597306..d1ac3dbd43 100644 --- a/tests/unit/prompt_target/test_query_target_capabilities.py +++ b/tests/unit/prompt_target/test_query_target_capabilities.py @@ -4,7 +4,7 @@ import asyncio import json from pathlib import Path -from unittest.mock import AsyncMock, MagicMock, patch +from unittest.mock import AsyncMock, patch import pytest @@ -12,9 +12,9 @@ from pyrit.prompt_target.common.prompt_target import PromptTarget from pyrit.prompt_target.common.query_target_capabilities import ( _CAPABILITY_PROBES, + DEFAULT_TEST_ASSETS, _create_test_message, _permissive_configuration, - DEFAULT_TEST_ASSETS, query_target_async, query_target_capabilities_async, query_target_modalities_async, From 79500e0a3976c300c6ecd73626581584f52f66c0 Mon Sep 17 00:00:00 2001 From: hannahwestra25 Date: Mon, 11 May 2026 17:25:41 -0400 Subject: [PATCH 19/27] fix test --- tests/unit/prompt_target/test_query_target_capabilities.py | 1 - 1 file changed, 1 deletion(-) diff --git a/tests/unit/prompt_target/test_query_target_capabilities.py b/tests/unit/prompt_target/test_query_target_capabilities.py index d1ac3dbd43..2d8f3a1975 100644 --- a/tests/unit/prompt_target/test_query_target_capabilities.py +++ b/tests/unit/prompt_target/test_query_target_capabilities.py @@ -635,7 +635,6 @@ async def test_returns_target_capabilities_assembled_from_probes(self) -> None: declared capabilities and deriving editable history conservatively. """ declared = TargetCapabilities( - supports_editable_history=True, input_modalities=frozenset({frozenset({"text"})}), output_modalities=frozenset({frozenset({"text"})}), ) From 3934ed5a920d88a49baf5a0a6c102704f1937dff Mon Sep 17 00:00:00 2001 From: hannahwestra25 Date: Mon, 11 May 2026 18:22:59 -0400 Subject: [PATCH 20/27] rename to discover --- doc/code/targets/0_prompt_targets.md | 16 ++-- .../targets/6_1_target_capabilities.ipynb | 18 ++-- doc/code/targets/6_1_target_capabilities.py | 44 ++++----- pyrit/prompt_target/__init__.py | 12 +-- .../common/query_target_capabilities.py | 30 +++--- .../test_query_target_capabilities.py | 94 +++++++++---------- 6 files changed, 107 insertions(+), 107 deletions(-) diff --git a/doc/code/targets/0_prompt_targets.md b/doc/code/targets/0_prompt_targets.md index 3922c18927..ef46d29148 100644 --- a/doc/code/targets/0_prompt_targets.md +++ b/doc/code/targets/0_prompt_targets.md @@ -107,26 +107,26 @@ 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 +### 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 ( - query_target_capabilities_async, - query_target_async, - query_target_modalities_async, + discover_target_capabilities_async, + discover_target_async, + discover_target_modalities_async, ) # Probe a single dimension: -queried_caps = await query_target_capabilities_async(target=target) -queried_modalities = await query_target_modalities_async(target=target) +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 query_target_async(target=target) +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. `query_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. +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 diff --git a/doc/code/targets/6_1_target_capabilities.ipynb b/doc/code/targets/6_1_target_capabilities.ipynb index 503d2b2d02..a296b85bde 100644 --- a/doc/code/targets/6_1_target_capabilities.ipynb +++ b/doc/code/targets/6_1_target_capabilities.ipynb @@ -380,9 +380,9 @@ "\n", "from pyrit.models import MessagePiece\n", "from pyrit.prompt_target import (\n", - " query_target_async,\n", - " query_target_capabilities_async,\n", - " query_target_modalities_async,\n", + " discover_target_async,\n", + " discover_target_capabilities_async,\n", + " discover_target_modalities_async,\n", ")\n", "\n", "\n", @@ -405,7 +405,7 @@ "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 query_target_capabilities_async(target=probe_target, per_probe_timeout_s=5.0) # type: ignore\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}\")" @@ -448,7 +448,7 @@ "source": [ "probe_target._send_prompt_to_target_async = AsyncMock(return_value=_ok_response()) # type: ignore[method-assign]\n", "\n", - "queried_modalities = await query_target_modalities_async(\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", @@ -592,8 +592,8 @@ "\n", "from pyrit.models import MessagePiece\n", "from pyrit.prompt_target import (\n", - " query_target_async,\n", - " query_target_capabilities_async,\n", + " discover_target_async,\n", + " discover_target_capabilities_async,\n", ")\n", "\n", "\n", @@ -616,7 +616,7 @@ "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 query_target_capabilities_async(target=probe_target, per_probe_timeout_s=5.0) # type: ignore\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}\")" @@ -667,7 +667,7 @@ "source": [ "probe_target._send_prompt_to_target_async = AsyncMock(return_value=_ok_response()) # type: ignore[method-assign]\n", "\n", - "queried_caps = await query_target_async(target=probe_target, per_probe_timeout_s=5.0) # type: ignore\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", diff --git a/doc/code/targets/6_1_target_capabilities.py b/doc/code/targets/6_1_target_capabilities.py index 68983be2bf..c3104a54df 100644 --- a/doc/code/targets/6_1_target_capabilities.py +++ b/doc/code/targets/6_1_target_capabilities.py @@ -247,26 +247,26 @@ print(exc) # %% [markdown] -# ## 7. Querying live target capabilities +# ## 7. Discovering 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`, `query_target_modalities_async`, or the convenience wrapper -# `query_target_async` that runs both and returns a best-effort `TargetCapabilities`. +# `discover_target_capabilities_async`, `discover_target_modalities_async`, or the convenience wrapper +# `discover_target_async` that runs both and returns a best-effort `TargetCapabilities`. # -# `query_target_capabilities_async` walks each capability that has a registered probe (currently +# `discover_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. # -# `query_target_modalities_async` does the same for input modality combinations declared in +# `discover_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. `query_target_async` returns a merged view: +# transient errors before being declared failed. `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 will still be reported as supporting that capability. @@ -279,9 +279,9 @@ # Typical usage against a real endpoint: # # ```python -# from pyrit.prompt_target import query_target_async +# from pyrit.prompt_target import discover_target_async # -# queried = await query_target_async(target=target) +# queried = await discover_target_async(target=target) # print(queried) # ``` # @@ -295,9 +295,9 @@ from pyrit.models import MessagePiece from pyrit.prompt_target import ( - query_target_async, - query_target_capabilities_async, - query_target_modalities_async, + discover_target_async, + discover_target_capabilities_async, + discover_target_modalities_async, ) @@ -320,7 +320,7 @@ def _ok_response(): 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] -queried = await query_target_capabilities_async(target=probe_target, per_probe_timeout_s=5.0) # type: ignore +queried = await discover_target_capabilities_async(target=probe_target, per_probe_timeout_s=5.0) # type: ignore print("queried capabilities:") for capability in sorted(queried, key=lambda c: c.value): print(f" - {capability.value}") @@ -331,19 +331,19 @@ def _ok_response(): # ```python # from pyrit.prompt_target.common.target_capabilities import CapabilityName # -# queried = await query_target_capabilities_async( +# queried = await discover_target_capabilities_async( # target=target, # capabilities=[CapabilityName.JSON_SCHEMA, CapabilityName.SYSTEM_PROMPT], # ) # ``` # # If you only care about accepted input combinations, call -# `query_target_modalities_async` directly. The example below uses the +# `discover_target_modalities_async` directly. The example below uses the # packaged default probe assets for the non-text modalities PyRIT ships. # Pass `test_assets=` only when you want to override those defaults or probe # a modality without a packaged asset. # -# `query_target_async` is the most common entry point: it runs both the capability and modality +# `discover_target_async` is the most common entry point: it runs both the capability and modality # probes and assembles a best-effort `TargetCapabilities` you can drop into a # `TargetConfiguration`, so the rest of PyRIT operates on probed values where available and # declared values otherwise. @@ -351,21 +351,21 @@ def _ok_response(): # %% probe_target._send_prompt_to_target_async = AsyncMock(return_value=_ok_response()) # type: ignore[method-assign] -queried_modalities = await query_target_modalities_async( +queried_modalities = await discover_target_modalities_async( target=probe_target, test_modalities={frozenset({"text"}), frozenset({"text", "image_path"})}, per_probe_timeout_s=5.0, ) # type: ignore -print("query_target_modalities_async result:") +print("discover_target_modalities_async result:") for combination in sorted(sorted(m) for m in queried_modalities): print(f" - {combination}") # %% probe_target._send_prompt_to_target_async = AsyncMock(return_value=_ok_response()) # type: ignore[method-assign] -queried_caps = await query_target_async(target=probe_target, per_probe_timeout_s=5.0) # type: ignore -print("query_target_async result:") +queried_caps = await discover_target_async(target=probe_target, per_probe_timeout_s=5.0) # type: ignore +print("discover_target_async result:") print(f" supports_multi_turn: {queried_caps.supports_multi_turn}") print(f" supports_system_prompt: {queried_caps.supports_system_prompt}") print(f" supports_multi_message_pieces: {queried_caps.supports_multi_message_pieces}") @@ -376,14 +376,14 @@ def _ok_response(): # %% [markdown] # ### Discovering undeclared modalities # -# By default `query_target_async` only probes modality combinations the target already +# By default `discover_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=` explicitly to # probe combinations beyond the declared baseline. Provide `test_assets=` as well if you need # to override the packaged defaults or probe a modality without one: # # ```python -# queried = await query_target_async( +# queried = await discover_target_async( # target=target, # test_modalities={frozenset({"text"}), frozenset({"text", "image_path"})}, # test_assets={"image_path": "/path/to/test_image.png"}, @@ -397,7 +397,7 @@ def _ok_response(): # # ```python # # Re-query only JSON support; other declared flags pass through unchanged. -# queried = await query_target_async( +# queried = await discover_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 cb46c577c4..7cf7020a76 100644 --- a/pyrit/prompt_target/__init__.py +++ b/pyrit/prompt_target/__init__.py @@ -17,9 +17,9 @@ 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_async, - query_target_capabilities_async, - query_target_modalities_async, + discover_target_async, + discover_target_capabilities_async, + discover_target_modalities_async, ) from pyrit.prompt_target.common.target_capabilities import ( CapabilityHandlingPolicy, @@ -102,14 +102,14 @@ def __getattr__(name: str) -> object: "PromptChatTarget", "PromptShieldTarget", "PromptTarget", - "query_target_capabilities_async", + "discover_target_capabilities_async", "RealtimeTarget", "TargetCapabilities", "TargetConfiguration", "TargetRequirements", "UnsupportedCapabilityBehavior", "TextTarget", - "query_target_async", - "query_target_modalities_async", + "discover_target_async", + "discover_target_modalities_async", "WebSocketCopilotTarget", ] diff --git a/pyrit/prompt_target/common/query_target_capabilities.py b/pyrit/prompt_target/common/query_target_capabilities.py index dc961c612d..f0f41d1373 100644 --- a/pyrit/prompt_target/common/query_target_capabilities.py +++ b/pyrit/prompt_target/common/query_target_capabilities.py @@ -6,14 +6,14 @@ This module exposes two complementary probes: -* :func:`query_target_capabilities_async` probes the boolean capability flags +* :func:`discover_target_capabilities_async` discovers 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 the target's declared native support from ``target.capabilities``. -* :func:`query_target_modalities_async` probes which input modality +* :func:`discover_target_modalities_async` discovers which input modality combinations a target actually supports by sending a minimal test request for each combination declared in ``TargetCapabilities.input_modalities``. @@ -55,7 +55,7 @@ logger = logging.getLogger(__name__) -# Per-call timeout (seconds) applied to every probe request. Override per-call via +# Per-call timeout (seconds) applied to every discovery request. Override per-call via # the ``per_probe_timeout_s`` parameter on the public functions. DEFAULT_PROBE_TIMEOUT_SECONDS: float = 30.0 @@ -74,7 +74,7 @@ # data-type check does not reject text probes against text-less targets. _TEXT_MODALITY: frozenset[frozenset[PromptDataType]] = frozenset({frozenset({"text"})}) -# Packaged fallback assets for non-text modality probes. +# Packaged fallback assets for non-text modality discovery. _TARGET_CAPABILITIES_DATASET_PATH = DATASETS_PATH / "prompt_target" / "target_capabilities" @@ -448,7 +448,7 @@ async def _probe_json_schema_async(target: PromptTarget, timeout_s: float, retri } -async def query_target_capabilities_async( +async def discover_target_capabilities_async( *, target: PromptTarget, capabilities: Iterable[CapabilityName] | None = None, @@ -514,7 +514,7 @@ async def query_target_capabilities_async( # Default mapping of non-text modalities to packaged probe assets. Callers can # override via the ``test_assets`` parameter of -# :func:`query_target_modalities_async`. Modalities whose assets do not exist +# :func:`discover_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] = { "audio_path": str(_TARGET_CAPABILITIES_DATASET_PATH / "probe_audio.wav"), @@ -522,7 +522,7 @@ async def query_target_capabilities_async( } -async def query_target_modalities_async( +async def discover_target_modalities_async( *, target: PromptTarget, test_modalities: set[frozenset[PromptDataType]] | None = None, @@ -590,7 +590,7 @@ async def query_target_modalities_async( return queried -async def query_target_async( +async def discover_target_async( *, target: PromptTarget, per_probe_timeout_s: float = DEFAULT_PROBE_TIMEOUT_SECONDS, @@ -602,8 +602,8 @@ async def query_target_async( """ Probe capabilities and modalities and return a merged result. - This wraps :func:`query_target_capabilities_async` and - :func:`query_target_modalities_async` and returns a best-effort + This wraps :func:`discover_target_capabilities_async` and + :func:`discover_target_modalities_async` and returns a best-effort :class:`TargetCapabilities`. Args: @@ -612,12 +612,12 @@ async def query_target_async( each probe request. test_modalities (set[frozenset[PromptDataType]] | None): Specific modality combinations to probe. See - :func:`query_target_modalities_async`. Defaults to the + :func:`discover_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:`query_target_modalities_async`. + modality to a file path. See :func:`discover_target_modalities_async`. capabilities (Iterable[CapabilityName] | None): Capabilities to probe. - See :func:`query_target_capabilities_async`. Defaults to every + See :func:`discover_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 @@ -630,13 +630,13 @@ async def query_target_async( """ capabilities_to_probe = list(capabilities) if capabilities is not None else None - queried_caps = await query_target_capabilities_async( + queried_caps = await discover_target_capabilities_async( target=target, capabilities=capabilities_to_probe, per_probe_timeout_s=per_probe_timeout_s, retries=retries, ) - queried_modalities = await query_target_modalities_async( + queried_modalities = await discover_target_modalities_async( target=target, test_modalities=test_modalities, test_assets=test_assets, diff --git a/tests/unit/prompt_target/test_query_target_capabilities.py b/tests/unit/prompt_target/test_query_target_capabilities.py index 2d8f3a1975..81859e522e 100644 --- a/tests/unit/prompt_target/test_query_target_capabilities.py +++ b/tests/unit/prompt_target/test_query_target_capabilities.py @@ -15,9 +15,9 @@ DEFAULT_TEST_ASSETS, _create_test_message, _permissive_configuration, - query_target_async, - query_target_capabilities_async, - query_target_modalities_async, + discover_target_async, + discover_target_capabilities_async, + discover_target_modalities_async, ) from pyrit.prompt_target.common.target_capabilities import ( CapabilityHandlingPolicy, @@ -109,7 +109,7 @@ 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) + result = await discover_target_capabilities_async(target=target) # Every capability with a probe should be in the result. for capability in _CAPABILITY_PROBES: @@ -119,7 +119,7 @@ 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) + result = await discover_target_capabilities_async(target=target) for capability in _CAPABILITY_PROBES: assert capability not in result @@ -128,7 +128,7 @@ 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) + result = await discover_target_capabilities_async(target=target) for capability in _CAPABILITY_PROBES: assert capability not in result @@ -138,7 +138,7 @@ async def test_filters_by_requested_capabilities(self) -> None: 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) + result = await discover_target_capabilities_async(target=target, capabilities=requested) assert result == requested @@ -150,7 +150,7 @@ async def test_capability_without_probe_falls_back_to_declared_support(self) -> ) target._send_prompt_to_target_async = AsyncMock(return_value=_ok_response()) # type: ignore[method-assign] - result = await query_target_capabilities_async( + result = await discover_target_capabilities_async( target=target, capabilities={CapabilityName.EDITABLE_HISTORY}, ) @@ -163,7 +163,7 @@ async def test_capability_without_probe_excluded_when_not_declared(self) -> None 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( + result = await discover_target_capabilities_async( target=target, capabilities={CapabilityName.EDITABLE_HISTORY}, ) @@ -199,7 +199,7 @@ async def test_capability_without_probe_excluded_when_only_adapted(self, monkeyp ), ) - result = await query_target_capabilities_async( + result = await discover_target_capabilities_async( target=target, capabilities={CapabilityName.SYSTEM_PROMPT}, ) @@ -215,7 +215,7 @@ async def test_accepts_single_pass_iterable(self) -> None: 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) + result = await discover_target_capabilities_async(target=target, capabilities=gen) assert CapabilityName.SYSTEM_PROMPT in result assert CapabilityName.EDITABLE_HISTORY in result @@ -224,7 +224,7 @@ 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( + result = await discover_target_capabilities_async( target=target, capabilities={CapabilityName.JSON_OUTPUT}, retries=0, @@ -238,7 +238,7 @@ async def test_restores_configuration_after_probing(self) -> None: 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) + await discover_target_capabilities_async(target=target) assert target.configuration is original @@ -246,7 +246,7 @@ 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( + await discover_target_capabilities_async( target=target, capabilities={CapabilityName.MULTI_TURN}, ) @@ -274,7 +274,7 @@ 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( + result = await discover_target_capabilities_async( target=target, capabilities={CapabilityName.MULTI_TURN}, ) @@ -288,7 +288,7 @@ 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( + await discover_target_capabilities_async( target=target, capabilities={CapabilityName.JSON_SCHEMA}, ) @@ -305,7 +305,7 @@ async def test_system_prompt_probe_installs_system_message_and_sends_user(self) target = MockPromptTarget() target._send_prompt_to_target_async = AsyncMock(return_value=_ok_response()) # type: ignore[method-assign] - await query_target_capabilities_async( + await discover_target_capabilities_async( target=target, capabilities={CapabilityName.SYSTEM_PROMPT}, ) @@ -327,7 +327,7 @@ 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( + await discover_target_capabilities_async( target=target, capabilities={CapabilityName.MULTI_MESSAGE_PIECES}, ) @@ -347,7 +347,7 @@ async def test_probes_run_under_permissive_configuration(self) -> None: 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( + result = await discover_target_capabilities_async( target=target, capabilities={CapabilityName.MULTI_MESSAGE_PIECES}, ) @@ -378,7 +378,7 @@ async def reject_system_roles(*, normalized_conversation: list[Message]) -> list target._send_prompt_to_target_async = AsyncMock(side_effect=reject_system_roles) # type: ignore[method-assign] - result = await query_target_capabilities_async( + result = await discover_target_capabilities_async( target=target, capabilities={CapabilityName.SYSTEM_PROMPT}, ) @@ -398,7 +398,7 @@ async def _send_prompt_to_target_async(self, *, normalized_conversation: list[Me 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) + result = await discover_target_capabilities_async(target=target) for capability in _CAPABILITY_PROBES: assert capability in result @@ -480,7 +480,7 @@ async def test_all_combinations_supported(self) -> None: _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 query_target_modalities_async(target=target) + result = await discover_target_modalities_async(target=target) assert frozenset({"text"}) in result @@ -489,7 +489,7 @@ async def test_exception_excludes_combination(self) -> None: _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 query_target_modalities_async(target=target) + result = await discover_target_modalities_async(target=target) assert result == set() @@ -498,7 +498,7 @@ async def test_error_response_excludes_combination(self) -> None: _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 query_target_modalities_async(target=target) + result = await discover_target_modalities_async(target=target) assert result == set() @@ -518,7 +518,7 @@ async def selective_send(*, normalized_conversation: list[Message]) -> list[Mess target._send_prompt_to_target_async = selective_send # type: ignore[method-assign] - result = await query_target_modalities_async( + result = await discover_target_modalities_async( target=target, test_assets={"image_path": image_asset}, ) @@ -532,7 +532,7 @@ async def test_explicit_test_modalities_overrides_declared(self, image_asset: st _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 query_target_modalities_async( + result = await discover_target_modalities_async( target=target, test_modalities={frozenset({"text"}), frozenset({"text", "image_path"})}, test_assets={"image_path": image_asset}, @@ -548,7 +548,7 @@ async def test_combination_skipped_when_asset_missing(self, tmp_path: Path) -> N # An explicit empty mapping disables the packaged defaults, so # image_path combinations are skipped instead of probed. - result = await query_target_modalities_async(target=target, test_assets={}) + result = await discover_target_modalities_async(target=target, test_assets={}) assert result == set() assert target._send_prompt_to_target_async.await_count == 0 @@ -564,7 +564,7 @@ async def test_explicit_test_modalities_runs_under_permissive_configuration(self send_mock = AsyncMock(return_value=_ok_response()) target._send_prompt_to_target_async = send_mock # type: ignore[method-assign] - result = await query_target_modalities_async( + result = await discover_target_modalities_async( target=target, test_modalities={frozenset({"text", "image_path"})}, test_assets={"image_path": image_asset}, @@ -591,7 +591,7 @@ async def _hang(**_kwargs: object) -> list[Message]: target._send_prompt_to_target_async = AsyncMock(side_effect=_hang) # type: ignore[method-assign] - result = await query_target_capabilities_async( + result = await discover_target_capabilities_async( target=target, capabilities={CapabilityName.JSON_OUTPUT}, per_probe_timeout_s=0.01, @@ -615,7 +615,7 @@ async def test_returns_false_when_memory_seed_raises(self) -> None: target._send_prompt_to_target_async = send_mock # type: ignore[method-assign] with patch.object(target._memory, "add_message_to_memory", side_effect=RuntimeError("memory offline")): - result = await query_target_capabilities_async( + result = await discover_target_capabilities_async( target=target, capabilities={CapabilityName.SYSTEM_PROMPT}, ) @@ -629,7 +629,7 @@ async def test_returns_false_when_memory_seed_raises(self) -> None: class TestVerifyTargetAsync: async def test_returns_target_capabilities_assembled_from_probes(self) -> None: """ - ``query_target_async`` runs both the capability and modality probes + ``discover_target_async`` runs both the capability and modality probes and assembles a :class:`TargetCapabilities` populated from the queried results, copying ``output_modalities`` from the target's declared capabilities and deriving editable history conservatively. @@ -642,7 +642,7 @@ async def test_returns_target_capabilities_assembled_from_probes(self) -> None: target._configuration = TargetConfiguration(capabilities=declared) target._send_prompt_to_target_async = AsyncMock(return_value=_ok_response()) # type: ignore[method-assign] - result = await query_target_async(target=target, per_probe_timeout_s=5.0) + result = await discover_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 @@ -672,7 +672,7 @@ async def test_excludes_capabilities_when_probe_send_fails(self) -> None: target._configuration = TargetConfiguration(capabilities=declared) target._send_prompt_to_target_async = AsyncMock(side_effect=RuntimeError("boom")) # type: ignore[method-assign] - result = await query_target_async(target=target, per_probe_timeout_s=0.5) + result = await discover_target_async(target=target, per_probe_timeout_s=0.5) assert result.supports_multi_turn is False assert result.supports_system_prompt is False @@ -692,7 +692,7 @@ async def test_empty_response_treated_as_failure(self) -> None: target = MockPromptTarget() target._send_prompt_to_target_async = AsyncMock(return_value=[]) # type: ignore[method-assign] - result = await query_target_capabilities_async( + result = await discover_target_capabilities_async( target=target, capabilities={CapabilityName.JSON_OUTPUT, CapabilityName.MULTI_MESSAGE_PIECES}, ) @@ -709,21 +709,21 @@ async def test_response_with_no_pieces_treated_as_failure(self) -> None: empty_msg = target._send_prompt_to_target_async.return_value[0] empty_msg.message_pieces = [] - result = await query_target_capabilities_async( + result = await discover_target_capabilities_async( target=target, capabilities={CapabilityName.JSON_OUTPUT}, ) assert result == set() - async def test_query_target_async_forwards_test_modalities(self, image_asset: str) -> None: + async def test_discover_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 query_target_async( + result = await discover_target_async( target=target, test_modalities={extra_combo}, test_assets={"image_path": image_asset}, @@ -733,12 +733,12 @@ async def test_query_target_async_forwards_test_modalities(self, image_asset: st # The undeclared combination is in the result only if test_modalities was forwarded. assert extra_combo in result.input_modalities - async def test_query_target_async_forwards_capabilities(self) -> None: - """``query_target_async`` must forward ``capabilities`` to narrow the probe set.""" + async def test_discover_target_async_forwards_capabilities(self) -> None: + """``discover_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 query_target_async( + await discover_target_async( target=target, capabilities={CapabilityName.JSON_OUTPUT}, per_probe_timeout_s=2.0, @@ -749,7 +749,7 @@ async def test_query_target_async_forwards_capabilities(self) -> None: # because multi-turn issues 2 sends). assert target._send_prompt_to_target_async.await_count <= 3 - async def test_query_target_async_preserves_declared_when_capabilities_narrowed(self) -> None: + async def test_discover_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 @@ -765,7 +765,7 @@ async def test_query_target_async_preserves_declared_when_capabilities_narrowed( target._configuration = TargetConfiguration(capabilities=declared) target._send_prompt_to_target_async = AsyncMock(return_value=_ok_response()) # type: ignore[method-assign] - result = await query_target_async( + result = await discover_target_async( target=target, capabilities={CapabilityName.JSON_OUTPUT}, per_probe_timeout_s=2.0, @@ -779,7 +779,7 @@ async def test_query_target_async_preserves_declared_when_capabilities_narrowed( assert result.supports_json_schema is True assert result.supports_editable_history is True - async def test_query_target_async_drops_editable_history_when_multi_turn_probe_fails(self) -> None: + async def test_discover_target_async_drops_editable_history_when_multi_turn_probe_fails(self) -> None: """Editable history must not remain true when probing disproves multi-turn support.""" declared = TargetCapabilities( supports_multi_turn=True, @@ -797,12 +797,12 @@ async def selective_send(*, normalized_conversation: list[Message]) -> list[Mess target._send_prompt_to_target_async = AsyncMock(side_effect=selective_send) # type: ignore[method-assign] - result = await query_target_async(target=target, per_probe_timeout_s=2.0) + result = await discover_target_async(target=target, per_probe_timeout_s=2.0) assert result.supports_multi_turn is False assert result.supports_editable_history is False - async def test_query_target_async_accepts_single_pass_iterable(self) -> None: + async def test_discover_target_async_accepts_single_pass_iterable(self) -> None: declared = TargetCapabilities( supports_multi_turn=True, supports_editable_history=True, @@ -812,7 +812,7 @@ async def test_query_target_async_accepts_single_pass_iterable(self) -> None: target._send_prompt_to_target_async = AsyncMock(return_value=_ok_response()) # type: ignore[method-assign] gen = (c for c in [CapabilityName.JSON_OUTPUT, CapabilityName.EDITABLE_HISTORY]) - result = await query_target_async( + result = await discover_target_async( target=target, capabilities=gen, per_probe_timeout_s=2.0, @@ -835,7 +835,7 @@ async def test_returns_false_when_history_seed_raises(self) -> None: target._send_prompt_to_target_async = send_mock # type: ignore[method-assign] with patch.object(target._memory, "add_message_to_memory", side_effect=RuntimeError("memory offline")): - result = await query_target_capabilities_async( + result = await discover_target_capabilities_async( target=target, capabilities={CapabilityName.MULTI_TURN}, ) From 8363beb14459e8652483224a6b29ba9c9b5fcdc7 Mon Sep 17 00:00:00 2001 From: hannahwestra25 Date: Tue, 12 May 2026 11:39:00 -0400 Subject: [PATCH 21/27] add backoff and use fresh config --- .../common/query_target_capabilities.py | 48 +++++++++++--- .../test_query_target_capabilities.py | 65 ++++++++++++++++++- 2 files changed, 102 insertions(+), 11 deletions(-) diff --git a/pyrit/prompt_target/common/query_target_capabilities.py b/pyrit/prompt_target/common/query_target_capabilities.py index f0f41d1373..de024611e2 100644 --- a/pyrit/prompt_target/common/query_target_capabilities.py +++ b/pyrit/prompt_target/common/query_target_capabilities.py @@ -58,6 +58,8 @@ # Per-call timeout (seconds) applied to every discovery request. Override per-call via # the ``per_probe_timeout_s`` parameter on the public functions. DEFAULT_PROBE_TIMEOUT_SECONDS: float = 30.0 +DEFAULT_PROBE_RETRY_BACKOFF_SECONDS: float = 0.1 +MAX_PROBE_RETRY_BACKOFF_SECONDS: float = 1.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 @@ -119,12 +121,10 @@ def _permissive_configuration( supports_system_prompt=True, input_modalities=merged_modalities, ) - probe_configuration = object.__new__(TargetConfiguration) - probe_configuration._capabilities = permissive_caps - probe_configuration._policy = original.policy - # Keep the original normalization pipeline intact so probing exercises the - # target's real request shaping, including custom normalizer overrides. - probe_configuration._pipeline = original.pipeline + # Rebuild a fresh configuration from the instance's native capabilities so + # probes bypass preflight validation without inheriting ADAPT policy or + # custom normalizer overrides from the target's runtime configuration. + probe_configuration = TargetConfiguration(capabilities=permissive_caps) target._configuration = probe_configuration try: yield @@ -186,8 +186,9 @@ async def _send_and_check_async( 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. + the probe is declared failed, with a short exponential backoff between + retry attempts; 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. @@ -195,7 +196,8 @@ async def _send_and_check_async( 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. + Retry attempts use exponential backoff starting at + :data:`DEFAULT_PROBE_RETRY_BACKOFF_SECONDS`. Defaults to 1. label (str): Short label used in log messages. Defaults to ``"Capability probe"``. @@ -214,10 +216,14 @@ async def _send_and_check_async( 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) + if attempt + 1 < attempts: + await _sleep_before_retry_async(attempt=attempt) continue except Exception as exc: last_exc = exc logger.debug("%s failed (attempt %d/%d): %s", label, attempt + 1, attempts, exc) + if attempt + 1 < attempts: + await _sleep_before_retry_async(attempt=attempt) continue if not responses or not any(r.message_pieces for r in responses): @@ -234,6 +240,16 @@ async def _send_and_check_async( return False +def _retry_backoff_seconds(*, attempt: int) -> float: + """Return the exponential backoff delay for a retry attempt.""" + return min(DEFAULT_PROBE_RETRY_BACKOFF_SECONDS * (2**attempt), MAX_PROBE_RETRY_BACKOFF_SECONDS) + + +async def _sleep_before_retry_async(*, attempt: int) -> None: + """Sleep for the retry backoff associated with ``attempt``.""" + await asyncio.sleep(_retry_backoff_seconds(attempt=attempt)) + + 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. @@ -372,6 +388,9 @@ async def _probe_json_output_async(target: PromptTarget, timeout_s: float, retri """ Probe whether ``target`` accepts a request asking for JSON-mode output. + This probe is only meaningful for targets that translate PyRIT's JSON + metadata hints into native provider request fields. + Args: target (PromptTarget): The target to probe. timeout_s (float): Per-attempt timeout in seconds. @@ -401,6 +420,9 @@ async def _probe_json_schema_async(target: PromptTarget, timeout_s: float, retri """ Probe whether ``target`` accepts a request constrained by a JSON schema. + This probe is only meaningful for targets that translate PyRIT's JSON + metadata hints into native provider request fields. + Args: target (PromptTarget): The target to probe. timeout_s (float): Per-attempt timeout in seconds. @@ -660,6 +682,12 @@ def _resolve(name: CapabilityName) -> bool: # Editable history is only meaningful if multi-turn probing/declaration # also resolved to True. resolved_editable_history = declared.supports_editable_history and resolved_multi_turn + if test_modalities is None: + resolved_input_modalities = frozenset(queried_modalities) + else: + resolved_input_modalities = frozenset( + queried_modalities | (declared.input_modalities - frozenset(test_modalities)) + ) return TargetCapabilities( supports_multi_turn=resolved_multi_turn, @@ -668,7 +696,7 @@ def _resolve(name: CapabilityName) -> bool: supports_json_output=_resolve(CapabilityName.JSON_OUTPUT), supports_editable_history=resolved_editable_history, supports_system_prompt=_resolve(CapabilityName.SYSTEM_PROMPT), - input_modalities=frozenset(queried_modalities), + input_modalities=resolved_input_modalities, # Output modalities are still declarative because probing them would # require target-specific response inspection. output_modalities=declared.output_modalities, diff --git a/tests/unit/prompt_target/test_query_target_capabilities.py b/tests/unit/prompt_target/test_query_target_capabilities.py index 81859e522e..ec53895355 100644 --- a/tests/unit/prompt_target/test_query_target_capabilities.py +++ b/tests/unit/prompt_target/test_query_target_capabilities.py @@ -233,6 +233,23 @@ async def test_retries_zero_disables_retry(self) -> None: assert result == set() assert target._send_prompt_to_target_async.await_count == 1 + async def test_retries_use_exponential_backoff(self) -> None: + target = MockPromptTarget() + target._send_prompt_to_target_async = AsyncMock(side_effect=Exception("boom")) # type: ignore[method-assign] + + with patch( + "pyrit.prompt_target.common.query_target_capabilities.asyncio.sleep", new_callable=AsyncMock + ) as sleep_mock: + result = await discover_target_capabilities_async( + target=target, + capabilities={CapabilityName.JSON_OUTPUT}, + retries=2, + ) + + assert result == set() + assert sleep_mock.await_args_list[0].args == (0.1,) + assert sleep_mock.await_args_list[1].args == (0.2,) + async def test_restores_configuration_after_probing(self) -> None: target = MockPromptTarget() original = target.configuration @@ -371,7 +388,7 @@ async def test_probed_capability_excluded_when_only_adapted(self) -> None: ) async def reject_system_roles(*, normalized_conversation: list[Message]) -> list[Message]: - roles = [piece.role for message in normalized_conversation for piece in message.message_pieces] + roles = [piece._role for message in normalized_conversation for piece in message.message_pieces] if "system" in roles: raise RuntimeError("system messages are not natively supported") return _ok_response() @@ -385,6 +402,33 @@ async def reject_system_roles(*, normalized_conversation: list[Message]) -> list assert result == set() + async def test_probe_configuration_does_not_reuse_adapted_pipeline(self) -> None: + target = MockPromptTarget() + target._configuration = TargetConfiguration( + capabilities=TargetCapabilities(supports_system_prompt=False), + policy=CapabilityHandlingPolicy( + behaviors={ + CapabilityName.SYSTEM_PROMPT: UnsupportedCapabilityBehavior.ADAPT, + CapabilityName.MULTI_TURN: UnsupportedCapabilityBehavior.RAISE, + } + ), + ) + + async def require_native_system_role(*, normalized_conversation: list[Message]) -> list[Message]: + roles = [piece._role for message in normalized_conversation for piece in message.message_pieces] + if "system" not in roles: + raise RuntimeError("probe used adapted system-prompt shaping") + return _ok_response() + + target._send_prompt_to_target_async = AsyncMock(side_effect=require_native_system_role) # type: ignore[method-assign] + + result = await discover_target_capabilities_async( + target=target, + capabilities={CapabilityName.SYSTEM_PROMPT}, + ) + + assert result == {CapabilityName.SYSTEM_PROMPT} + @pytest.mark.usefixtures("patch_central_database") class TestQueryTargetCapabilitiesIsolatedTarget: @@ -733,6 +777,25 @@ async def test_discover_target_async_forwards_test_modalities(self, image_asset: # The undeclared combination is in the result only if test_modalities was forwarded. assert extra_combo in result.input_modalities + async def test_discover_target_async_preserves_declared_modalities_when_test_modalities_narrowed( + self, image_asset: str + ) -> None: + declared_combo = frozenset({"text"}) + probed_combo = frozenset({"text", "image_path"}) + declared = TargetCapabilities(input_modalities=frozenset({declared_combo, probed_combo})) + target = MockPromptTarget() + target._configuration = TargetConfiguration(capabilities=declared) + target._send_prompt_to_target_async = AsyncMock(return_value=_ok_response()) + + result = await discover_target_async( + target=target, + test_modalities={probed_combo}, + test_assets={"image_path": image_asset}, + per_probe_timeout_s=2.0, + ) + + assert result.input_modalities == frozenset({declared_combo, probed_combo}) + async def test_discover_target_async_forwards_capabilities(self) -> None: """``discover_target_async`` must forward ``capabilities`` to narrow the probe set.""" target = MockPromptTarget() From 52c56e6e03bf3a680bb0fd0cf67dec10b87022a0 Mon Sep 17 00:00:00 2001 From: hannahwestra25 Date: Tue, 12 May 2026 17:59:42 -0400 Subject: [PATCH 22/27] add logging for classes that don't enforce json capabilities --- .../common/query_target_capabilities.py | 35 ++++++++ .../test_query_target_capabilities.py | 89 +++++++++++++++++++ 2 files changed, 124 insertions(+) diff --git a/pyrit/prompt_target/common/query_target_capabilities.py b/pyrit/prompt_target/common/query_target_capabilities.py index de024611e2..48064542f6 100644 --- a/pyrit/prompt_target/common/query_target_capabilities.py +++ b/pyrit/prompt_target/common/query_target_capabilities.py @@ -71,6 +71,24 @@ _CapabilityProbe = Callable[[PromptTarget, float, int], Awaitable[bool]] +def _json_enforcing_target_types() -> tuple[type[PromptTarget], ...]: + """ + Return the tuple of target classes that translate ``prompt_metadata`` JSON + hints (``response_format``, ``json_schema``) into native provider request + fields. Used to suppress the "JSON probe is upper-bound only" debug log for + these targets and their subclasses. + + Imports are lazy to avoid a circular dependency at module load time and to + keep this concern entirely within the discovery module. Class objects (not + strings) are returned so renames are caught by import errors rather than + silently flipping the log behavior. + """ + from pyrit.prompt_target.openai.openai_chat_target import OpenAIChatTarget + from pyrit.prompt_target.openai.openai_response_target import OpenAIResponseTarget + + return (OpenAIChatTarget, OpenAIResponseTarget) + + # 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. @@ -503,6 +521,8 @@ async def discover_target_capabilities_async( ) queried: set[CapabilityName] = set() + json_capabilities = {CapabilityName.JSON_OUTPUT, CapabilityName.JSON_SCHEMA} + queried_json_capabilities: set[CapabilityName] = set() with _permissive_configuration(target=target): for capability in capabilities_to_check: probe = _CAPABILITY_PROBES.get(capability) @@ -516,9 +536,24 @@ async def discover_target_capabilities_async( # still ignore the feature semantics after accepting the call. if await probe(target, per_probe_timeout_s, retries): queried.add(capability) + if capability in json_capabilities: + queried_json_capabilities.add(capability) except Exception as exc: logger.debug("Probe for %s raised: %s", capability.value, exc) + # JSON probes only verify the target accepted the request, not that the + # target translated the JSON metadata into provider request fields. Emit + # a single summary line when probes succeeded against a target that does + # not enforce JSON hints, so the result is treated as an upper bound. + # ``isinstance`` covers user-defined subclasses of enforcing targets. + if queried_json_capabilities and not isinstance(target, _json_enforcing_target_types()): + logger.debug( + "JSON capability probes %s succeeded for %s, but this target does not translate " + "prompt_metadata JSON hints into provider request fields; treat the result as upper-bound support only.", + sorted(c.value for c in queried_json_capabilities), + type(target).__name__, + ) + # Read unprobed capabilities from target.capabilities, not # target.configuration, so ADAPTed behavior is not reported as native # support. diff --git a/tests/unit/prompt_target/test_query_target_capabilities.py b/tests/unit/prompt_target/test_query_target_capabilities.py index ec53895355..ddf38e2fb4 100644 --- a/tests/unit/prompt_target/test_query_target_capabilities.py +++ b/tests/unit/prompt_target/test_query_target_capabilities.py @@ -3,6 +3,7 @@ import asyncio import json +import logging from pathlib import Path from unittest.mock import AsyncMock, patch @@ -318,6 +319,94 @@ async def test_json_schema_probe_sends_schema_in_metadata(self) -> None: schema = json.loads(metadata["json_schema"]) assert schema["type"] == "object" + @pytest.mark.parametrize("capability", [CapabilityName.JSON_OUTPUT, CapabilityName.JSON_SCHEMA]) + async def test_logs_debug_for_unenforced_json_probe( + self, capability: CapabilityName, caplog: pytest.LogCaptureFixture + ) -> None: + target = MockPromptTarget() + target._send_prompt_to_target_async = AsyncMock(return_value=_ok_response()) # type: ignore[method-assign] + + with caplog.at_level(logging.DEBUG): + result = await discover_target_capabilities_async(target=target, capabilities={capability}) + + assert result == {capability} + matching = [r for r in caplog.records if r.message.startswith("JSON capability probes")] + assert len(matching) == 1 + assert capability.value in matching[0].message + + async def test_logs_unenforced_json_probe_summary_once_for_both_capabilities( + self, caplog: pytest.LogCaptureFixture + ) -> None: + target = MockPromptTarget() + target._send_prompt_to_target_async = AsyncMock(return_value=_ok_response()) # type: ignore[method-assign] + + with caplog.at_level(logging.DEBUG): + await discover_target_capabilities_async( + target=target, + capabilities={CapabilityName.JSON_OUTPUT, CapabilityName.JSON_SCHEMA}, + ) + + # A single summary line covers both probed JSON capabilities. + matching = [r for r in caplog.records if r.message.startswith("JSON capability probes")] + assert len(matching) == 1 + assert CapabilityName.JSON_OUTPUT.value in matching[0].message + assert CapabilityName.JSON_SCHEMA.value in matching[0].message + + async def test_does_not_log_unenforced_json_probe_when_probe_fails(self, caplog: pytest.LogCaptureFixture) -> None: + target = MockPromptTarget() + target._send_prompt_to_target_async = AsyncMock(side_effect=Exception("boom")) # type: ignore[method-assign] + + with caplog.at_level(logging.DEBUG): + result = await discover_target_capabilities_async( + target=target, + capabilities={CapabilityName.JSON_OUTPUT, CapabilityName.JSON_SCHEMA}, + retries=0, + ) + + assert result == set() + assert not any(r.message.startswith("JSON capability probes") for r in caplog.records) + + async def test_does_not_log_debug_for_enforced_json_probe(self, caplog: pytest.LogCaptureFixture) -> None: + target_type = type("FakeEnforcingTarget", (MockPromptTarget,), {}) + target = target_type() + target._send_prompt_to_target_async = AsyncMock(return_value=_ok_response()) # type: ignore[method-assign] + + with ( + patch( + "pyrit.prompt_target.common.query_target_capabilities._json_enforcing_target_types", + return_value=(target_type,), + ), + caplog.at_level(logging.DEBUG), + ): + result = await discover_target_capabilities_async( + target=target, + capabilities={CapabilityName.JSON_OUTPUT}, + ) + + assert result == {CapabilityName.JSON_OUTPUT} + assert not any(r.message.startswith("JSON capability probes") for r in caplog.records) + + async def test_subclass_of_enforced_target_does_not_log(self, caplog: pytest.LogCaptureFixture) -> None: + # ``isinstance`` covers user-defined subclasses of enforcing targets. + base = type("EnforcingBase", (MockPromptTarget,), {}) + sub = type("UserSubclass", (base,), {}) + target = sub() + target._send_prompt_to_target_async = AsyncMock(return_value=_ok_response()) # type: ignore[method-assign] + + with ( + patch( + "pyrit.prompt_target.common.query_target_capabilities._json_enforcing_target_types", + return_value=(base,), + ), + caplog.at_level(logging.DEBUG), + ): + await discover_target_capabilities_async( + target=target, + capabilities={CapabilityName.JSON_OUTPUT}, + ) + + assert not any(r.message.startswith("JSON capability probes") for r in caplog.records) + 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] From b8478a25a04a457256d757cddab1a73f0b798bef Mon Sep 17 00:00:00 2001 From: hannahwestra25 Date: Tue, 12 May 2026 18:25:17 -0400 Subject: [PATCH 23/27] ghcp pr review --- .../common/query_target_capabilities.py | 59 ++++++++++++------ .../test_query_target_capabilities.py | 61 +++++++++++++++++-- 2 files changed, 97 insertions(+), 23 deletions(-) diff --git a/pyrit/prompt_target/common/query_target_capabilities.py b/pyrit/prompt_target/common/query_target_capabilities.py index 48064542f6..76dfff614b 100644 --- a/pyrit/prompt_target/common/query_target_capabilities.py +++ b/pyrit/prompt_target/common/query_target_capabilities.py @@ -61,6 +61,15 @@ DEFAULT_PROBE_RETRY_BACKOFF_SECONDS: float = 0.1 MAX_PROBE_RETRY_BACKOFF_SECONDS: float = 1.0 +# Exceptions that are deterministic on the probe payload and will not become +# valid on a retry (malformed Message, type errors, missing attributes, etc.). +# These fail the probe immediately rather than wasting backoff time. +_NON_RETRYABLE_PROBE_EXCEPTIONS: tuple[type[BaseException], ...] = ( + ValueError, + TypeError, + AttributeError, +) + # 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 @@ -202,20 +211,23 @@ async def _send_and_check_async( """ 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, with a short exponential backoff between - retry attempts; an explicit error response from the target is treated as - deterministic and never retried. + Each attempt is bounded by ``timeout_s``. Transient errors (timeouts, + connection/OS errors) trigger up to ``retries`` retries with a short + exponential backoff. Deterministic errors that will not become valid on + a retry (``ValueError``, ``TypeError``, ``AttributeError`` — typically + from message validation or programmer error in a probe payload) fail + the probe immediately. 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. - Retry attempts use exponential backoff starting at - :data:`DEFAULT_PROBE_RETRY_BACKOFF_SECONDS`. Defaults to 1. + Only transient errors are retried; non-retryable errors and + non-error responses are final. Retry attempts use exponential + backoff starting at :data:`DEFAULT_PROBE_RETRY_BACKOFF_SECONDS`. + Defaults to 1. label (str): Short label used in log messages. Defaults to ``"Capability probe"``. @@ -237,6 +249,10 @@ async def _send_and_check_async( if attempt + 1 < attempts: await _sleep_before_retry_async(attempt=attempt) continue + except _NON_RETRYABLE_PROBE_EXCEPTIONS as exc: + # Deterministic on the probe payload — retrying will not help. + logger.debug("%s failed with non-retryable error: %s", label, exc) + return False except Exception as exc: last_exc = exc logger.debug("%s failed (attempt %d/%d): %s", label, attempt + 1, attempts, exc) @@ -244,7 +260,7 @@ async def _send_and_check_async( await _sleep_before_retry_async(attempt=attempt) continue - if not responses or not any(r.message_pieces for r in responses): + if not responses or any(not r.message_pieces for r in responses): logger.debug("%s returned an empty response; treating as failure", label) return False for response in responses: @@ -273,10 +289,13 @@ async def _probe_system_prompt_async(target: PromptTarget, timeout_s: float, ret 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 + rather than calling :meth:`pyrit.prompt_target.PromptChatTarget.set_system_prompt` + (which is only defined on ``PromptChatTarget`` subclasses anyway). + ``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 also works + uniformly for plain ``PromptTarget`` subclasses that have no + ``set_system_prompt`` method, and guarantees the probe sees the same multi-piece, system-then-user payload the target's wire layer would see via the standard pipeline. @@ -617,6 +636,9 @@ async def discover_target_modalities_async( if test_modalities is None: declared = target.capabilities.input_modalities test_modalities = set(declared) + elif not test_modalities: + logger.info("discover_target_modalities_async called with an empty test_modalities set; nothing to probe.") + return set() assets = test_assets if test_assets is not None else DEFAULT_TEST_ASSETS @@ -718,7 +740,10 @@ def _resolve(name: CapabilityName) -> bool: # also resolved to True. resolved_editable_history = declared.supports_editable_history and resolved_multi_turn if test_modalities is None: - resolved_input_modalities = frozenset(queried_modalities) + # Mirror the boolean fallback: combinations the probe could not confirm + # fall back to the target's declared support rather than being silently + # dropped (e.g. on transient network failure). + resolved_input_modalities = frozenset(queried_modalities | declared.input_modalities) else: resolved_input_modalities = frozenset( queried_modalities | (declared.input_modalities - frozenset(test_modalities)) @@ -756,8 +781,7 @@ def _create_test_message( 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. + ValueError: If a non-text modality has no configured asset. """ conversation_id = f"modality-probe-{uuid.uuid4()}" pieces: list[MessagePiece] = [] @@ -791,7 +815,4 @@ def _create_test_message( ) ) - 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 index ddf38e2fb4..01a0cec86c 100644 --- a/tests/unit/prompt_target/test_query_target_capabilities.py +++ b/tests/unit/prompt_target/test_query_target_capabilities.py @@ -251,6 +251,32 @@ async def test_retries_use_exponential_backoff(self) -> None: assert sleep_mock.await_args_list[0].args == (0.1,) assert sleep_mock.await_args_list[1].args == (0.2,) + async def test_non_retryable_validation_errors_fail_fast(self) -> None: + """ + Deterministic errors (ValueError/TypeError/AttributeError) come from + malformed payloads or programmer error and will not become valid on + a retry. They must fail the probe immediately without consuming the + retry budget or sleeping for backoff. + """ + target = MockPromptTarget() + target._send_prompt_to_target_async = AsyncMock( # type: ignore[method-assign] + side_effect=ValueError("malformed payload") + ) + + with patch( + "pyrit.prompt_target.common.query_target_capabilities.asyncio.sleep", new_callable=AsyncMock + ) as sleep_mock: + result = await discover_target_capabilities_async( + target=target, + capabilities={CapabilityName.JSON_OUTPUT}, + retries=3, + ) + + assert result == set() + # No retries consumed and no backoff sleeps issued. + assert target._send_prompt_to_target_async.await_count == 1 + sleep_mock.assert_not_awaited() + async def test_restores_configuration_after_probing(self) -> None: target = MockPromptTarget() original = target.configuration @@ -686,6 +712,16 @@ async def test_combination_skipped_when_asset_missing(self, tmp_path: Path) -> N assert result == set() assert target._send_prompt_to_target_async.await_count == 0 + async def test_empty_test_modalities_returns_empty_without_probing(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 discover_target_modalities_async(target=target, test_modalities=set()) + + 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 @@ -794,8 +830,9 @@ async def test_returns_target_capabilities_assembled_from_probes(self) -> None: async def test_excludes_capabilities_when_probe_send_fails(self) -> None: """ When the underlying send raises, no capability or modality is - queried, but ``supports_editable_history`` and ``output_modalities`` - are still copied conservatively from the declared capabilities. + queried, but ``supports_editable_history``, ``output_modalities``, + and declared ``input_modalities`` are still preserved conservatively + from the declared capabilities. """ declared = TargetCapabilities( supports_editable_history=True, @@ -815,8 +852,9 @@ async def test_excludes_capabilities_when_probe_send_fails(self) -> None: # Editable history is derived conservatively and must fall when # multi-turn probing disproves the prerequisite capability. assert result.supports_editable_history is False - # No modalities queried because send always fails. - assert result.input_modalities == frozenset() + # When probing cannot confirm modalities, declared modalities are + # preserved (mirroring the boolean fallback semantics). + assert result.input_modalities == declared.input_modalities # Output modalities still copied. assert result.output_modalities == declared.output_modalities @@ -849,6 +887,21 @@ async def test_response_with_no_pieces_treated_as_failure(self) -> None: assert result == set() + async def test_mixed_empty_message_in_response_treated_as_failure(self) -> None: + """Any empty Message in a multi-message response must cause the probe to fail.""" + target = MockPromptTarget() + ok = _ok_response()[0] + empty = Message.__new__(Message) + empty.message_pieces = [] + target._send_prompt_to_target_async = AsyncMock(return_value=[ok, empty]) # type: ignore[method-assign] + + result = await discover_target_capabilities_async( + target=target, + capabilities={CapabilityName.JSON_OUTPUT}, + ) + + assert result == set() + async def test_discover_target_async_forwards_test_modalities(self, image_asset: str) -> None: declared = TargetCapabilities(input_modalities=frozenset({frozenset({"text"})})) target = MockPromptTarget() From 18259d203ad076aca2dff68aa11c79c2e33fe81f Mon Sep 17 00:00:00 2001 From: hannahwestra25 Date: Wed, 13 May 2026 11:44:14 -0400 Subject: [PATCH 24/27] add apply function and update docs --- .../targets/6_1_target_capabilities.ipynb | 303 ++++++++++++------ doc/code/targets/6_1_target_capabilities.py | 85 +++++ pyrit/prompt_target/common/prompt_target.py | 23 ++ .../target/test_prompt_target.py | 47 +++ 4 files changed, 357 insertions(+), 101 deletions(-) diff --git a/doc/code/targets/6_1_target_capabilities.ipynb b/doc/code/targets/6_1_target_capabilities.ipynb index a296b85bde..31ea0af03b 100644 --- a/doc/code/targets/6_1_target_capabilities.ipynb +++ b/doc/code/targets/6_1_target_capabilities.ipynb @@ -376,39 +376,19 @@ } ], "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", - " 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}\")" + "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)" ] }, { @@ -416,27 +396,8 @@ "id": "15", "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", - "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." + "By contrast, the `RAISE` configuration validates eagerly: any consumer requiring `MULTI_TURN` will\n", + "get a `ValueError` before a single prompt is sent." ] }, { @@ -444,26 +405,6 @@ "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", @@ -482,7 +423,7 @@ }, { "cell_type": "markdown", - "id": "18", + "id": "17", "metadata": {}, "source": [ "## 6. Non-adaptable capabilities\n", @@ -496,7 +437,7 @@ { "cell_type": "code", "execution_count": null, - "id": "19", + "id": "18", "metadata": {}, "outputs": [ { @@ -520,29 +461,29 @@ }, { "cell_type": "markdown", - "id": "20", + "id": "19", "metadata": {}, "source": [ - "## 7. Querying live target capabilities\n", + "## 7. Discovering 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", + "`discover_target_capabilities_async`, `discover_target_modalities_async`, or the convenience wrapper\n", + "`discover_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", + "`discover_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", + "`discover_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", + "transient errors before being declared failed. `discover_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", @@ -555,9 +496,9 @@ "Typical usage against a real endpoint:\n", "\n", "```python\n", - "from pyrit.prompt_target import query_target_async\n", + "from pyrit.prompt_target import discover_target_async\n", "\n", - "queried = await query_target_async(target=target)\n", + "queried = await discover_target_async(target=target)\n", "print(queried)\n", "```\n", "\n", @@ -570,7 +511,7 @@ { "cell_type": "code", "execution_count": null, - "id": "21", + "id": "20", "metadata": {}, "outputs": [ { @@ -594,6 +535,7 @@ "from pyrit.prompt_target import (\n", " discover_target_async,\n", " discover_target_capabilities_async,\n", + " discover_target_modalities_async,\n", ")\n", "\n", "\n", @@ -624,7 +566,7 @@ }, { "cell_type": "markdown", - "id": "22", + "id": "21", "metadata": {}, "source": [ "To narrow the probe to specific capabilities (faster, fewer calls), pass `capabilities=`:\n", @@ -632,16 +574,52 @@ "```python\n", "from pyrit.prompt_target.common.target_capabilities import CapabilityName\n", "\n", - "queried = await query_target_capabilities_async(\n", + "queried = await discover_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." + "If you only care about accepted input combinations, call\n", + "`discover_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", + "`discover_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": "22", + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "discover_target_modalities_async result:\n", + " - ['image_path', 'text']\n", + " - ['text']\n" + ] + } + ], + "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(\"discover_target_modalities_async result:\")\n", + "for combination in sorted(sorted(m) for m in queried_modalities):\n", + " print(f\" - {combination}\")" ] }, { @@ -654,13 +632,13 @@ "name": "stdout", "output_type": "stream", "text": [ - "query_target_async result:\n", + "discover_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" + " input_modalities: [['image_path'], ['image_path', 'text'], ['text']]\n" ] } ], @@ -668,7 +646,7 @@ "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(\"discover_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", @@ -684,14 +662,14 @@ "source": [ "### Discovering undeclared modalities\n", "\n", - "By default `query_target_async` only probes modality combinations the target already\n", + "By default `discover_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", + "queried = await discover_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", @@ -705,18 +683,141 @@ "\n", "```python\n", "# Re-query only JSON support; other declared flags pass through unchanged.\n", - "queried = await query_target_async(\n", + "queried = await discover_target_async(\n", " target=target,\n", " capabilities={CapabilityName.JSON_OUTPUT, CapabilityName.JSON_SCHEMA},\n", ")\n", "```" ] + }, + { + "cell_type": "markdown", + "id": "25", + "metadata": {}, + "source": [ + "## 8. Applying probed capabilities back onto the target\n", + "\n", + "`discover_target_async` is intentionally pure: it returns a `TargetCapabilities` without\n", + "mutating the target. That lets you inspect (or diff against the declared view, log, gate on\n", + "the result) before committing. Once you're satisfied, call `target.apply_capabilities(...)`\n", + "to install the probed view on the instance. The target's existing\n", + "`CapabilityHandlingPolicy` is preserved — policy expresses user intent (ADAPT vs RAISE),\n", + "which is independent of what the probe found.\n", + "\n", + "Why a two-step pattern rather than auto-apply? Probe results are an upper bound\n", + "(\"the request was accepted\"); a target that silently ignores a feature still passes its\n", + "probe. Keeping discovery separate from application lets callers diff, log, persist, or\n", + "reject the result before it affects subsequent sends.\n", + "\n", + "Below is the end-to-end pattern: construct a target whose declared capabilities are\n", + "pessimistic, discover what the endpoint actually accepts, diff the two views, then apply." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "26", + "metadata": { + "lines_to_next_cell": 2 + }, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "declared (before probing):\n", + " supports_multi_turn: False\n", + " supports_system_prompt: False\n", + " supports_json_output: False\n", + "\n", + "probed (returned from discover_target_async, target NOT yet updated):\n", + " supports_multi_turn: True\n", + " supports_system_prompt: True\n", + " supports_json_output: True\n", + " target.capabilities.supports_multi_turn (still declared): False\n", + "\n", + "flags probed True that were declared False: ['supports_multi_turn', 'supports_system_prompt', 'supports_multi_message_pieces', 'supports_json_output', 'supports_json_schema']\n", + "\n", + "after apply_capabilities:\n", + " supports_multi_turn: True\n", + " supports_system_prompt: True\n", + " supports_json_output: True\n", + " policy preserved: True\n", + "\n", + "CHAT_TARGET_REQUIREMENTS.validate now passes against the probed target\n" + ] + } + ], + "source": [ + "# Start with an instance that declares fewer capabilities than the endpoint actually has,\n", + "# e.g. a custom gateway whose support we're unsure about.\n", + "pessimistic_config = TargetConfiguration(\n", + " capabilities=TargetCapabilities(\n", + " supports_multi_turn=False,\n", + " supports_system_prompt=False,\n", + " supports_multi_message_pieces=False,\n", + " supports_json_output=False,\n", + " supports_json_schema=False,\n", + " # Editable history has no live probe and falls back to the declared value.\n", + " # Declare it True here so the probed view inherits it.\n", + " supports_editable_history=True,\n", + " ),\n", + ")\n", + "endpoint_target = OpenAIChatTarget(\n", + " model_name=\"custom-model\",\n", + " endpoint=\"https://example.invalid/\",\n", + " api_key=\"sk-not-a-real-key\",\n", + " custom_configuration=pessimistic_config,\n", + ")\n", + "endpoint_target._send_prompt_to_target_async = AsyncMock(return_value=_ok_response()) # type: ignore[method-assign]\n", + "\n", + "print(\"declared (before probing):\")\n", + "print(f\" supports_multi_turn: {endpoint_target.capabilities.supports_multi_turn}\")\n", + "print(f\" supports_system_prompt: {endpoint_target.capabilities.supports_system_prompt}\")\n", + "print(f\" supports_json_output: {endpoint_target.capabilities.supports_json_output}\")\n", + "\n", + "# Step 1: discover. No mutation yet — `endpoint_target.capabilities` is unchanged.\n", + "probed_caps = await discover_target_async(target=endpoint_target, per_probe_timeout_s=5.0) # type: ignore\n", + "\n", + "print(\"\\nprobed (returned from discover_target_async, target NOT yet updated):\")\n", + "print(f\" supports_multi_turn: {probed_caps.supports_multi_turn}\")\n", + "print(f\" supports_system_prompt: {probed_caps.supports_system_prompt}\")\n", + "print(f\" supports_json_output: {probed_caps.supports_json_output}\")\n", + "print(f\" target.capabilities.supports_multi_turn (still declared): {endpoint_target.capabilities.supports_multi_turn}\")\n", + "\n", + "# Step 2: diff — see exactly what the probe upgraded.\n", + "declared = pessimistic_config.capabilities\n", + "upgraded = [\n", + " name\n", + " for name in (\n", + " \"supports_multi_turn\",\n", + " \"supports_system_prompt\",\n", + " \"supports_multi_message_pieces\",\n", + " \"supports_json_output\",\n", + " \"supports_json_schema\",\n", + " )\n", + " if getattr(probed_caps, name) and not getattr(declared, name)\n", + "]\n", + "print(f\"\\nflags probed True that were declared False: {upgraded}\")\n", + "\n", + "# Step 3: apply. Policy is preserved; the normalization pipeline is rebuilt.\n", + "original_policy = endpoint_target.configuration.policy\n", + "endpoint_target.apply_capabilities(capabilities=probed_caps)\n", + "\n", + "print(\"\\nafter apply_capabilities:\")\n", + "print(f\" supports_multi_turn: {endpoint_target.capabilities.supports_multi_turn}\")\n", + "print(f\" supports_system_prompt: {endpoint_target.capabilities.supports_system_prompt}\")\n", + "print(f\" supports_json_output: {endpoint_target.capabilities.supports_json_output}\")\n", + "print(f\" policy preserved: {endpoint_target.configuration.policy is original_policy}\")\n", + "\n", + "# Subsequent consumer checks now reflect the probed reality — for example, a chat-style\n", + "# requirement that would have failed against the pessimistic declaration now passes.\n", + "CHAT_TARGET_REQUIREMENTS.validate(target=endpoint_target)\n", + "print(\"\\nCHAT_TARGET_REQUIREMENTS.validate now passes against the probed target\")" + ] } ], "metadata": { - "jupytext": { - "main_language": "python" - }, "language_info": { "codemirror_mode": { "name": "ipython", diff --git a/doc/code/targets/6_1_target_capabilities.py b/doc/code/targets/6_1_target_capabilities.py index c3104a54df..8952410e8f 100644 --- a/doc/code/targets/6_1_target_capabilities.py +++ b/doc/code/targets/6_1_target_capabilities.py @@ -402,3 +402,88 @@ def _ok_response(): # capabilities={CapabilityName.JSON_OUTPUT, CapabilityName.JSON_SCHEMA}, # ) # ``` + +# %% [markdown] +# ## 8. Applying probed capabilities back onto the target +# +# `discover_target_async` is intentionally pure: it returns a `TargetCapabilities` without +# mutating the target. That lets you inspect (or diff against the declared view, log, gate on +# the result) before committing. Once you're satisfied, call `target.apply_capabilities(...)` +# to install the probed view on the instance. The target's existing +# `CapabilityHandlingPolicy` is preserved — policy expresses user intent (ADAPT vs RAISE), +# which is independent of what the probe found. +# +# Why a two-step pattern rather than auto-apply? Probe results are an upper bound +# ("the request was accepted"); a target that silently ignores a feature still passes its +# probe. Keeping discovery separate from application lets callers diff, log, persist, or +# reject the result before it affects subsequent sends. +# +# Below is the end-to-end pattern: construct a target whose declared capabilities are +# pessimistic, discover what the endpoint actually accepts, diff the two views, then apply. + +# %% +# Start with an instance that declares fewer capabilities than the endpoint actually has, +# e.g. a custom gateway whose support we're unsure about. +pessimistic_config = TargetConfiguration( + capabilities=TargetCapabilities( + supports_multi_turn=False, + supports_system_prompt=False, + supports_multi_message_pieces=False, + supports_json_output=False, + supports_json_schema=False, + # Editable history has no live probe and falls back to the declared value. + # Declare it True here so the probed view inherits it. + supports_editable_history=True, + ), +) +endpoint_target = OpenAIChatTarget( + model_name="custom-model", + endpoint="https://example.invalid/", + api_key="sk-not-a-real-key", + custom_configuration=pessimistic_config, +) +endpoint_target._send_prompt_to_target_async = AsyncMock(return_value=_ok_response()) # type: ignore[method-assign] + +print("declared (before probing):") +print(f" supports_multi_turn: {endpoint_target.capabilities.supports_multi_turn}") +print(f" supports_system_prompt: {endpoint_target.capabilities.supports_system_prompt}") +print(f" supports_json_output: {endpoint_target.capabilities.supports_json_output}") + +# Step 1: discover. No mutation yet — `endpoint_target.capabilities` is unchanged. +probed_caps = await discover_target_async(target=endpoint_target, per_probe_timeout_s=5.0) # type: ignore + +print("\nprobed (returned from discover_target_async, target NOT yet updated):") +print(f" supports_multi_turn: {probed_caps.supports_multi_turn}") +print(f" supports_system_prompt: {probed_caps.supports_system_prompt}") +print(f" supports_json_output: {probed_caps.supports_json_output}") +print(f" target.capabilities.supports_multi_turn (still declared): {endpoint_target.capabilities.supports_multi_turn}") + +# Step 2: diff — see exactly what the probe upgraded. +declared = pessimistic_config.capabilities +upgraded = [ + name + for name in ( + "supports_multi_turn", + "supports_system_prompt", + "supports_multi_message_pieces", + "supports_json_output", + "supports_json_schema", + ) + if getattr(probed_caps, name) and not getattr(declared, name) +] +print(f"\nflags probed True that were declared False: {upgraded}") + +# Step 3: apply. Policy is preserved; the normalization pipeline is rebuilt. +original_policy = endpoint_target.configuration.policy +endpoint_target.apply_capabilities(capabilities=probed_caps) + +print("\nafter apply_capabilities:") +print(f" supports_multi_turn: {endpoint_target.capabilities.supports_multi_turn}") +print(f" supports_system_prompt: {endpoint_target.capabilities.supports_system_prompt}") +print(f" supports_json_output: {endpoint_target.capabilities.supports_json_output}") +print(f" policy preserved: {endpoint_target.configuration.policy is original_policy}") + +# Subsequent consumer checks now reflect the probed reality — for example, a chat-style +# requirement that would have failed against the pessimistic declaration now passes. +CHAT_TARGET_REQUIREMENTS.validate(target=endpoint_target) +print("\nCHAT_TARGET_REQUIREMENTS.validate now passes against the probed target") diff --git a/pyrit/prompt_target/common/prompt_target.py b/pyrit/prompt_target/common/prompt_target.py index a4cf2a8a96..42f2688b01 100644 --- a/pyrit/prompt_target/common/prompt_target.py +++ b/pyrit/prompt_target/common/prompt_target.py @@ -406,6 +406,29 @@ def capabilities(self) -> TargetCapabilities: """ return self._configuration.capabilities + def apply_capabilities(self, *, capabilities: TargetCapabilities) -> None: + """ + Replace this target's capabilities, preserving the existing handling policy. + The normalization pipeline is rebuilt from the input capabilities and the + current policy. + + Policy is preserved because it expresses user intent (ADAPT vs RAISE), + independent of what the probe found. To change policy or normalizer + overrides, build a new :class:`TargetConfiguration` and pass it via + ``custom_configuration`` at construction time instead. + + Note: + This mutates the target's identifier (derived from the configuration). + + Args: + capabilities (TargetCapabilities): The capabilities to install on + this instance. + """ + self._configuration = TargetConfiguration( + capabilities=capabilities, + policy=self._configuration.policy, + ) + @classmethod def get_default_configuration(cls, underlying_model: str | None = None) -> TargetConfiguration: """ diff --git a/tests/unit/prompt_target/target/test_prompt_target.py b/tests/unit/prompt_target/target/test_prompt_target.py index 93bcc0bf19..f3174c2649 100644 --- a/tests/unit/prompt_target/target/test_prompt_target.py +++ b/tests/unit/prompt_target/target/test_prompt_target.py @@ -634,3 +634,50 @@ async def normalize_async(self, messages): # pragma: no cover - not exercised ) assert a.get_identifier().hash != b.get_identifier().hash + + +def test_apply_capabilities_replaces_capabilities_and_preserves_policy(patch_central_database): + initial_policy = CapabilityHandlingPolicy( + behaviors={ + CapabilityName.MULTI_TURN: UnsupportedCapabilityBehavior.ADAPT, + CapabilityName.SYSTEM_PROMPT: UnsupportedCapabilityBehavior.RAISE, + } + ) + target = OpenAIChatTarget( + model_name="gpt-4o", + endpoint="https://mock.azure.com/", + api_key="mock-api-key", + custom_configuration=TargetConfiguration( + capabilities=TargetCapabilities(supports_multi_turn=False, supports_system_prompt=False), + policy=initial_policy, + ), + ) + + new_caps = TargetCapabilities(supports_multi_turn=True, supports_system_prompt=True) + target.apply_capabilities(capabilities=new_caps) + + assert target.capabilities == new_caps + # Policy is preserved by identity, not just by value. + assert target.configuration.policy is initial_policy + + +def test_apply_capabilities_rebuilds_pipeline(patch_central_database): + adapt_policy = CapabilityHandlingPolicy( + behaviors={ + CapabilityName.MULTI_TURN: UnsupportedCapabilityBehavior.ADAPT, + CapabilityName.SYSTEM_PROMPT: UnsupportedCapabilityBehavior.RAISE, + } + ) + target = OpenAIChatTarget( + model_name="gpt-4o", + endpoint="https://mock.azure.com/", + api_key="mock-api-key", + custom_configuration=TargetConfiguration( + capabilities=TargetCapabilities(supports_multi_turn=False, supports_system_prompt=True), + policy=adapt_policy, + ), + ) + assert target.configuration.pipeline._normalizers, "Expected ADAPT pipeline to be non-empty" + + target.apply_capabilities(capabilities=TargetCapabilities(supports_multi_turn=True, supports_system_prompt=True)) + assert not target.configuration.pipeline._normalizers, "Expected pipeline to be rebuilt as empty" From 5fa257294613a69e2434f531894f2613849d65f1 Mon Sep 17 00:00:00 2001 From: hannahwestra25 Date: Wed, 13 May 2026 14:02:40 -0400 Subject: [PATCH 25/27] rename and add apply flag --- .../targets/6_1_target_capabilities.ipynb | 5 + pyrit/prompt_target/__init__.py | 6 +- .../common/discover_target_capabilities.py | 829 ++++++++++++++++++ .../common/query_target_capabilities.py | 14 +- ...y => test_discover_target_capabilities.py} | 45 +- 5 files changed, 886 insertions(+), 13 deletions(-) create mode 100644 pyrit/prompt_target/common/discover_target_capabilities.py rename tests/unit/prompt_target/{test_query_target_capabilities.py => test_discover_target_capabilities.py} (95%) diff --git a/doc/code/targets/6_1_target_capabilities.ipynb b/doc/code/targets/6_1_target_capabilities.ipynb index 31ea0af03b..a54290c893 100644 --- a/doc/code/targets/6_1_target_capabilities.ipynb +++ b/doc/code/targets/6_1_target_capabilities.ipynb @@ -818,6 +818,11 @@ } ], "metadata": { + "kernelspec": { + "display_name": "pyrit", + "language": "python", + "name": "python3" + }, "language_info": { "codemirror_mode": { "name": "ipython", diff --git a/pyrit/prompt_target/__init__.py b/pyrit/prompt_target/__init__.py index 7cf7020a76..3c74da5b17 100644 --- a/pyrit/prompt_target/__init__.py +++ b/pyrit/prompt_target/__init__.py @@ -14,13 +14,13 @@ from pyrit.prompt_target.azure_blob_storage_target import AzureBlobStorageTarget from pyrit.prompt_target.azure_ml_chat_target import AzureMLChatTarget 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 ( +from pyrit.prompt_target.common.discover_target_capabilities import ( discover_target_async, discover_target_capabilities_async, discover_target_modalities_async, ) +from pyrit.prompt_target.common.prompt_chat_target import PromptChatTarget +from pyrit.prompt_target.common.prompt_target import PromptTarget from pyrit.prompt_target.common.target_capabilities import ( CapabilityHandlingPolicy, CapabilityName, diff --git a/pyrit/prompt_target/common/discover_target_capabilities.py b/pyrit/prompt_target/common/discover_target_capabilities.py new file mode 100644 index 0000000000..9982713515 --- /dev/null +++ b/pyrit/prompt_target/common/discover_target_capabilities.py @@ -0,0 +1,829 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT license. + +""" +Runtime capability and modality discovery for prompt targets. + +This module exposes two complementary probes: + +* :func:`discover_target_capabilities_async` discovers 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 the target's declared native support from + ``target.capabilities``. +* :func:`discover_target_modalities_async` discovers 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*. They do not prove + that the endpoint enforced the feature, and the JSON probes are only + meaningful for targets that translate ``prompt_metadata`` JSON hints into + provider request fields. Treat the results as an upper bound on support and + validate response content separately when that distinction matters. +""" + +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.common.path import DATASETS_PATH +from pyrit.models import Message, MessagePiece, PromptDataType +from pyrit.prompt_target.common.prompt_target import PromptTarget +from pyrit.prompt_target.common.target_capabilities import ( + CapabilityName, + TargetCapabilities, +) +from pyrit.prompt_target.common.target_configuration import TargetConfiguration + +logger = logging.getLogger(__name__) + +# Per-call timeout (seconds) applied to every discovery request. Override per-call via +# the ``per_probe_timeout_s`` parameter on the public functions. +DEFAULT_PROBE_TIMEOUT_SECONDS: float = 30.0 +DEFAULT_PROBE_RETRY_BACKOFF_SECONDS: float = 0.1 +MAX_PROBE_RETRY_BACKOFF_SECONDS: float = 1.0 + +# Exceptions that are deterministic on the probe payload and will not become +# valid on a retry (malformed Message, type errors, missing attributes, etc.). +# These fail the probe immediately rather than wasting backoff time. +_NON_RETRYABLE_PROBE_EXCEPTIONS: tuple[type[BaseException], ...] = ( + ValueError, + TypeError, + AttributeError, +) + +# 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]] + + +def _json_enforcing_target_types() -> tuple[type[PromptTarget], ...]: + """ + Return the tuple of target classes that translate ``prompt_metadata`` JSON + hints (``response_format``, ``json_schema``) into native provider request + fields. Used to suppress the "JSON probe is upper-bound only" debug log for + these targets and their subclasses. + + Imports are lazy to avoid a circular dependency at module load time and to + keep this concern entirely within the discovery module. Class objects (not + strings) are returned so renames are caught by import errors rather than + silently flipping the log behavior. + """ + from pyrit.prompt_target.openai.openai_chat_target import OpenAIChatTarget + from pyrit.prompt_target.openai.openai_response_target import OpenAIResponseTarget + + return (OpenAIChatTarget, OpenAIResponseTarget) + + +# 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"})}) + +# Packaged fallback assets for non-text modality discovery. +_TARGET_CAPABILITIES_DATASET_PATH = DATASETS_PATH / "prompt_target" / "target_capabilities" + + +@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, + ) + # Rebuild a fresh configuration from the instance's native capabilities so + # probes bypass preflight validation without inheriting ADAPT policy or + # custom normalizer overrides from the target's runtime configuration. + probe_configuration = TargetConfiguration(capabilities=permissive_caps) + target._configuration = probe_configuration + 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``. Transient errors (timeouts, + connection/OS errors) trigger up to ``retries`` retries with a short + exponential backoff. Deterministic errors that will not become valid on + a retry (``ValueError``, ``TypeError``, ``AttributeError`` — typically + from message validation or programmer error in a probe payload) fail + the probe immediately. 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 transient errors are retried; non-retryable errors and + non-error responses are final. Retry attempts use exponential + backoff starting at :data:`DEFAULT_PROBE_RETRY_BACKOFF_SECONDS`. + 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) + if attempt + 1 < attempts: + await _sleep_before_retry_async(attempt=attempt) + continue + except _NON_RETRYABLE_PROBE_EXCEPTIONS as exc: + # Deterministic on the probe payload — retrying will not help. + logger.debug("%s failed with non-retryable error: %s", label, exc) + return False + except Exception as exc: + last_exc = exc + logger.debug("%s failed (attempt %d/%d): %s", label, attempt + 1, attempts, exc) + if attempt + 1 < attempts: + await _sleep_before_retry_async(attempt=attempt) + continue + + if not responses or any(not 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 + + +def _retry_backoff_seconds(*, attempt: int) -> float: + """Return the exponential backoff delay for a retry attempt.""" + return min(DEFAULT_PROBE_RETRY_BACKOFF_SECONDS * (2**attempt), MAX_PROBE_RETRY_BACKOFF_SECONDS) + + +async def _sleep_before_retry_async(*, attempt: int) -> None: + """Sleep for the retry backoff associated with ``attempt``.""" + await asyncio.sleep(_retry_backoff_seconds(attempt=attempt)) + + +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:`pyrit.prompt_target.PromptChatTarget.set_system_prompt` + (which is only defined on ``PromptChatTarget`` subclasses anyway). + ``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 also works + uniformly for plain ``PromptTarget`` subclasses that have no + ``set_system_prompt`` method, and 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. + retries (int): Number of additional attempts after the first failure. + Only exceptions/timeouts are retried; an explicit error response + is final. Defaults to 1. + + 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. + retries (int): Number of additional attempts after the first failure. + Only exceptions/timeouts are retried; an explicit error response + is final. Defaults to 1. + + 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. + retries (int): Number of additional attempts after the first failure. + Only exceptions/timeouts are retried; an explicit error response + is final. Defaults to 1. + + 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. + + This probe is only meaningful for targets that translate PyRIT's JSON + metadata hints into native provider request fields. + + Args: + target (PromptTarget): The target to probe. + timeout_s (float): Per-attempt timeout in seconds. + retries (int): Number of additional attempts after the first failure. + Only exceptions/timeouts are retried; an explicit error response + is final. Defaults to 1. + + 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, + # This only becomes a real JSON-mode request on targets that honor + # PyRIT's JSON metadata contract when building the provider payload. + 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. + + This probe is only meaningful for targets that translate PyRIT's JSON + metadata hints into native provider request fields. + + Args: + target (PromptTarget): The target to probe. + timeout_s (float): Per-attempt timeout in seconds. + retries (int): Number of additional attempts after the first failure. + Only exceptions/timeouts are retried; an explicit error response + is final. Defaults to 1. + + 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, + # As above, this probe is only strong for targets that map these + # metadata keys to native JSON-schema request parameters. + 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 queried 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 discover_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 which capabilities ``target`` accepts. + + Registered capabilities are checked with live requests. Capabilities + without a live probe fall back to declared native support. + + 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 confirmed to work against the target. + """ + capabilities_to_check: list[CapabilityName] = ( + list(capabilities) if capabilities is not None else list(CapabilityName) + ) + + queried: set[CapabilityName] = set() + json_capabilities = {CapabilityName.JSON_OUTPUT, CapabilityName.JSON_SCHEMA} + queried_json_capabilities: set[CapabilityName] = set() + with _permissive_configuration(target=target): + for capability in capabilities_to_check: + probe = _CAPABILITY_PROBES.get(capability) + if probe is None: + # Capabilities without a probe are handled after the permissive + # override is removed so we can read the target's native flags. + continue + + try: + # "Supported" means the request was accepted. A target can + # still ignore the feature semantics after accepting the call. + if await probe(target, per_probe_timeout_s, retries): + queried.add(capability) + if capability in json_capabilities: + queried_json_capabilities.add(capability) + except Exception as exc: + logger.debug("Probe for %s raised: %s", capability.value, exc) + + # JSON probes only verify the target accepted the request, not that the + # target translated the JSON metadata into provider request fields. Emit + # a single summary line when probes succeeded against a target that does + # not enforce JSON hints, so the result is treated as an upper bound. + # ``isinstance`` covers user-defined subclasses of enforcing targets. + if queried_json_capabilities and not isinstance(target, _json_enforcing_target_types()): + logger.debug( + "JSON capability probes %s succeeded for %s, but this target does not translate " + "prompt_metadata JSON hints into provider request fields; treat the result as upper-bound support only.", + sorted(c.value for c in queried_json_capabilities), + type(target).__name__, + ) + + # Read unprobed capabilities from target.capabilities, not + # target.configuration, so ADAPTed behavior is not reported as native + # support. + for capability in capabilities_to_check: + if capability not in _CAPABILITY_PROBES and target.capabilities.includes(capability=capability): + queried.add(capability) + + return queried + + +# --------------------------------------------------------------------------- +# Modality query +# --------------------------------------------------------------------------- + + +# Default mapping of non-text modalities to packaged probe assets. Callers can +# override via the ``test_assets`` parameter of +# :func:`discover_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] = { + "audio_path": str(_TARGET_CAPABILITIES_DATASET_PATH / "probe_audio.wav"), + "image_path": str(_TARGET_CAPABILITIES_DATASET_PATH / "probe_image.png"), +} + + +async def discover_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 which input modality combinations ``target`` accepts. + + Each modality combination is checked with a minimal request built from the + supplied test assets. + + 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 confirmed + to work against the target. + """ + if test_modalities is None: + declared = target.capabilities.input_modalities + test_modalities = set(declared) + elif not test_modalities: + logger.info("discover_target_modalities_async called with an empty test_modalities set; nothing to probe.") + return set() + + assets = test_assets if test_assets is not None else DEFAULT_TEST_ASSETS + + queried: 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: + # Skip combinations we cannot construct a valid probe payload for. + logger.info("Skipping modality %s: %s", combination, exc) + continue + except ValueError as exc: + logger.info("Skipping modality %s: %s", combination, exc) + continue + + # "Supported" means the request was accepted. A target may still + # ignore the non-text payload after accepting it. + if await _send_and_check_async( + target=target, + message=message, + timeout_s=per_probe_timeout_s, + retries=retries, + label=f"Modality probe {sorted(combination)}", + ): + queried.add(combination) + + return queried + + +async def discover_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, + apply: bool = False, +) -> TargetCapabilities: + """ + Probe capabilities and modalities and return a merged result. + + This wraps :func:`discover_target_capabilities_async` and + :func:`discover_target_modalities_async` and returns a best-effort + :class:`TargetCapabilities`. + + 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:`discover_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:`discover_target_modalities_async`. + capabilities (Iterable[CapabilityName] | None): Capabilities to probe. + See :func:`discover_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. + apply (bool): If True, install the discovered capabilities on ``target`` + via :meth:`PromptTarget.apply_capabilities` before returning. + Probe results are an upper bound (the request was accepted, not + necessarily honored), so leave this False when you want to inspect + or diff the result before committing to it. Defaults to False. + + Returns: + TargetCapabilities: A merged capability view: probed where possible, + declared where probing is unavailable or out of scope. + """ + capabilities_to_probe = list(capabilities) if capabilities is not None else None + + queried_caps = await discover_target_capabilities_async( + target=target, + capabilities=capabilities_to_probe, + per_probe_timeout_s=per_probe_timeout_s, + retries=retries, + ) + queried_modalities = await discover_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 + # If the caller narrows the capability set, leave the rest at their + # declared values instead of silently forcing them to False. + probed: set[CapabilityName] = ( + set(capabilities_to_probe) if capabilities_to_probe is not None else set(CapabilityName) + ) + + def _resolve(name: CapabilityName) -> bool: + if name in probed: + return name in queried_caps + return bool(getattr(declared, name.value)) + + resolved_multi_turn = _resolve(CapabilityName.MULTI_TURN) + # Editable history is only meaningful if multi-turn probing/declaration + # also resolved to True. + resolved_editable_history = declared.supports_editable_history and resolved_multi_turn + if test_modalities is None: + # Mirror the boolean fallback: combinations the probe could not confirm + # fall back to the target's declared support rather than being silently + # dropped (e.g. on transient network failure). + resolved_input_modalities = frozenset(queried_modalities | declared.input_modalities) + else: + resolved_input_modalities = frozenset( + queried_modalities | (declared.input_modalities - frozenset(test_modalities)) + ) + + resolved = TargetCapabilities( + supports_multi_turn=resolved_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=resolved_editable_history, + supports_system_prompt=_resolve(CapabilityName.SYSTEM_PROMPT), + input_modalities=resolved_input_modalities, + # Output modalities are still declarative because probing them would + # require target-specific response inspection. + output_modalities=declared.output_modalities, + ) + + if apply: + target.apply_capabilities(capabilities=resolved) + + return resolved + + +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. + """ + 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(), + ) + ) + + return Message(pieces) diff --git a/pyrit/prompt_target/common/query_target_capabilities.py b/pyrit/prompt_target/common/query_target_capabilities.py index 76dfff614b..eac20e12ff 100644 --- a/pyrit/prompt_target/common/query_target_capabilities.py +++ b/pyrit/prompt_target/common/query_target_capabilities.py @@ -677,6 +677,7 @@ async def discover_target_async( test_assets: dict[PromptDataType, str] | None = None, capabilities: Iterable[CapabilityName] | None = None, retries: int = 1, + apply: bool = False, ) -> TargetCapabilities: """ Probe capabilities and modalities and return a merged result. @@ -702,6 +703,12 @@ async def discover_target_async( for each probe. Only exceptions/timeouts are retried; an explicit error response is final. Set to ``0`` to disable retries. Defaults to 1. + apply (bool): If True, install the discovered capabilities on ``target`` + via :meth:`PromptTarget.apply_capabilities` before returning. The + target's existing handling policy is preserved. Probe results are + an upper bound (the request was accepted, not necessarily honored), + so leave this False when you want to inspect or diff the result + before committing to it. Defaults to False. Returns: TargetCapabilities: A merged capability view: probed where possible, @@ -749,7 +756,7 @@ def _resolve(name: CapabilityName) -> bool: queried_modalities | (declared.input_modalities - frozenset(test_modalities)) ) - return TargetCapabilities( + resolved = TargetCapabilities( supports_multi_turn=resolved_multi_turn, supports_multi_message_pieces=_resolve(CapabilityName.MULTI_MESSAGE_PIECES), supports_json_schema=_resolve(CapabilityName.JSON_SCHEMA), @@ -762,6 +769,11 @@ def _resolve(name: CapabilityName) -> bool: output_modalities=declared.output_modalities, ) + if apply: + target.apply_capabilities(capabilities=resolved) + + return resolved + def _create_test_message( *, diff --git a/tests/unit/prompt_target/test_query_target_capabilities.py b/tests/unit/prompt_target/test_discover_target_capabilities.py similarity index 95% rename from tests/unit/prompt_target/test_query_target_capabilities.py rename to tests/unit/prompt_target/test_discover_target_capabilities.py index 01a0cec86c..44e4a50d04 100644 --- a/tests/unit/prompt_target/test_query_target_capabilities.py +++ b/tests/unit/prompt_target/test_discover_target_capabilities.py @@ -10,8 +10,7 @@ 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 ( +from pyrit.prompt_target.common.discover_target_capabilities import ( _CAPABILITY_PROBES, DEFAULT_TEST_ASSETS, _create_test_message, @@ -20,6 +19,7 @@ discover_target_capabilities_async, discover_target_modalities_async, ) +from pyrit.prompt_target.common.prompt_target import PromptTarget from pyrit.prompt_target.common.target_capabilities import ( CapabilityHandlingPolicy, CapabilityName, @@ -105,7 +105,7 @@ def test_restores_on_exception(self) -> None: @pytest.mark.usefixtures("patch_central_database") -class TestQueryTargetCapabilitiesAsync: +class TestDiscoverTargetCapabilitiesAsync: 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] @@ -180,7 +180,7 @@ async def test_capability_without_probe_excluded_when_only_adapted(self, monkeyp 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 import discover_target_capabilities as qtc from pyrit.prompt_target.common.target_capabilities import ( CapabilityHandlingPolicy, UnsupportedCapabilityBehavior, @@ -239,7 +239,7 @@ async def test_retries_use_exponential_backoff(self) -> None: target._send_prompt_to_target_async = AsyncMock(side_effect=Exception("boom")) # type: ignore[method-assign] with patch( - "pyrit.prompt_target.common.query_target_capabilities.asyncio.sleep", new_callable=AsyncMock + "pyrit.prompt_target.common.discover_target_capabilities.asyncio.sleep", new_callable=AsyncMock ) as sleep_mock: result = await discover_target_capabilities_async( target=target, @@ -264,7 +264,7 @@ async def test_non_retryable_validation_errors_fail_fast(self) -> None: ) with patch( - "pyrit.prompt_target.common.query_target_capabilities.asyncio.sleep", new_callable=AsyncMock + "pyrit.prompt_target.common.discover_target_capabilities.asyncio.sleep", new_callable=AsyncMock ) as sleep_mock: result = await discover_target_capabilities_async( target=target, @@ -399,7 +399,7 @@ async def test_does_not_log_debug_for_enforced_json_probe(self, caplog: pytest.L with ( patch( - "pyrit.prompt_target.common.query_target_capabilities._json_enforcing_target_types", + "pyrit.prompt_target.common.discover_target_capabilities._json_enforcing_target_types", return_value=(target_type,), ), caplog.at_level(logging.DEBUG), @@ -421,7 +421,7 @@ async def test_subclass_of_enforced_target_does_not_log(self, caplog: pytest.Log with ( patch( - "pyrit.prompt_target.common.query_target_capabilities._json_enforcing_target_types", + "pyrit.prompt_target.common.discover_target_capabilities._json_enforcing_target_types", return_value=(base,), ), caplog.at_level(logging.DEBUG), @@ -546,7 +546,7 @@ async def require_native_system_role(*, normalized_conversation: list[Message]) @pytest.mark.usefixtures("patch_central_database") -class TestQueryTargetCapabilitiesIsolatedTarget: +class TestDiscoverTargetCapabilitiesIsolatedTarget: """Tests using a bare PromptTarget subclass (no PromptChatTarget extras).""" async def test_with_minimal_target_subclass(self) -> None: @@ -1026,6 +1026,33 @@ async def test_discover_target_async_accepts_single_pass_iterable(self) -> None: assert result.supports_json_output is True assert result.supports_editable_history is True + async def test_discover_target_async_apply_installs_capabilities_on_target(self) -> None: + """When ``apply=True``, the discovered capabilities are installed on the target.""" + declared = TargetCapabilities(supports_multi_turn=False, supports_system_prompt=False) + target = MockPromptTarget() + target._configuration = TargetConfiguration(capabilities=declared) + target._send_prompt_to_target_async = AsyncMock(return_value=_ok_response()) # type: ignore[method-assign] + + assert target.capabilities.supports_multi_turn is False + + result = await discover_target_async(target=target, per_probe_timeout_s=2.0, apply=True) + + assert target.capabilities == result + assert target.capabilities.supports_multi_turn is True + + async def test_discover_target_async_apply_defaults_to_false(self) -> None: + """By default, ``discover_target_async`` must not mutate the target.""" + declared = TargetCapabilities(supports_multi_turn=False, supports_system_prompt=False) + target = MockPromptTarget() + target._configuration = TargetConfiguration(capabilities=declared) + target._send_prompt_to_target_async = AsyncMock(return_value=_ok_response()) # type: ignore[method-assign] + + result = await discover_target_async(target=target, per_probe_timeout_s=2.0) + + # Result reflects probe; target capabilities remain at declared values. + assert result.supports_multi_turn is True + assert target.capabilities.supports_multi_turn is False + @pytest.mark.usefixtures("patch_central_database") class TestMultiTurnProbeMemoryFailure: From 96cffc84df232cb2db8405eb24658eeb508f3843 Mon Sep 17 00:00:00 2001 From: hannahwestra25 Date: Wed, 13 May 2026 14:57:18 -0400 Subject: [PATCH 26/27] move to private functions --- doc/code/targets/0_prompt_targets.md | 17 +- .../targets/6_1_target_capabilities.ipynb | 150 +++++------------- doc/code/targets/6_1_target_capabilities.py | 96 ++++------- pyrit/prompt_target/__init__.py | 6 +- .../common/discover_target_capabilities.py | 48 +++--- .../test_discover_target_capabilities.py | 126 +++++++-------- 6 files changed, 166 insertions(+), 277 deletions(-) diff --git a/doc/code/targets/0_prompt_targets.md b/doc/code/targets/0_prompt_targets.md index ef46d29148..25fc2b476c 100644 --- a/doc/code/targets/0_prompt_targets.md +++ b/doc/code/targets/0_prompt_targets.md @@ -112,21 +112,14 @@ The full implementation lives in [`pyrit/prompt_target/common/target_capabilitie 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) +from pyrit.prompt_target import discover_target_capabilities_async -# Or do both at once and get a best-effort TargetCapabilities back: -queried = await discover_target_async(target=target) +# Probe boolean capabilities and input modalities, returning a +# best-effort TargetCapabilities: +queried = await discover_target_capabilities_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. +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_capabilities_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. 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 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 diff --git a/doc/code/targets/6_1_target_capabilities.ipynb b/doc/code/targets/6_1_target_capabilities.ipynb index a54290c893..7630a560f6 100644 --- a/doc/code/targets/6_1_target_capabilities.ipynb +++ b/doc/code/targets/6_1_target_capabilities.ipynb @@ -469,36 +469,35 @@ "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", - "`discover_target_capabilities_async`, `discover_target_modalities_async`, or the convenience wrapper\n", - "`discover_target_async` that runs both and returns a best-effort `TargetCapabilities`.\n", + "`discover_target_capabilities_async`. It runs both the boolean capability probes and the input\n", + "modality probes and returns a best-effort `TargetCapabilities`.\n", "\n", - "`discover_target_capabilities_async` walks each capability that has a registered probe (currently\n", + "Internally it 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", + "minimal request, and includes the capability in the result 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", - "`discover_target_modalities_async` does the same for input modality combinations declared in\n", + "unsupported. The original configuration is restored before the function returns. The same\n", + "treatment is applied to each input modality combination 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. `discover_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", + "transient errors before being declared failed. The returned `TargetCapabilities` is a merged\n", + "view: probed where possible, declared where probing is unavailable or out of scope.\n", + "\"Supported\" here means *the request was accepted* — a target that silently ignores a system\n", + "prompt or `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", + "This function is **not safe to call concurrently** with other operations on the same target\n", + "instance: it temporarily mutates `target._configuration` and writes 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 discover_target_async\n", + "from pyrit.prompt_target import discover_target_capabilities_async\n", "\n", - "queried = await discover_target_async(target=target)\n", + "queried = await discover_target_capabilities_async(target=target)\n", "print(queried)\n", "```\n", "\n", @@ -532,11 +531,7 @@ "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", + "from pyrit.prompt_target import discover_target_capabilities_async\n", "\n", "\n", "def _ok_response():\n", @@ -559,9 +554,13 @@ "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}\")" + "print(\"discover_target_capabilities_async result:\")\n", + "print(f\" supports_multi_turn: {queried.supports_multi_turn}\")\n", + "print(f\" supports_system_prompt: {queried.supports_system_prompt}\")\n", + "print(f\" supports_multi_message_pieces: {queried.supports_multi_message_pieces}\")\n", + "print(f\" supports_json_output: {queried.supports_json_output}\")\n", + "print(f\" supports_json_schema: {queried.supports_json_schema}\")\n", + "print(f\" input_modalities: {sorted(sorted(m) for m in queried.input_modalities)}\")" ] }, { @@ -580,96 +579,25 @@ ")\n", "```\n", "\n", - "If you only care about accepted input combinations, call\n", - "`discover_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", - "`discover_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": "22", - "metadata": {}, - "outputs": [ - { - "name": "stdout", - "output_type": "stream", - "text": [ - "discover_target_modalities_async result:\n", - " - ['image_path', 'text']\n", - " - ['text']\n" - ] - } - ], - "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(\"discover_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": "23", - "metadata": {}, - "outputs": [ - { - "name": "stdout", - "output_type": "stream", - "text": [ - "discover_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: [['image_path'], ['image_path', 'text'], ['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(\"discover_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)}\")" + "Similarly, narrow the modality probe set with `test_modalities=` and override the\n", + "packaged default probe assets with `test_assets=`." ] }, { "cell_type": "markdown", - "id": "24", + "id": "22", "metadata": {}, "source": [ "### Discovering undeclared modalities\n", "\n", - "By default `discover_target_async` only probes modality combinations the target already\n", + "By default `discover_target_capabilities_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 discover_target_async(\n", + "queried = await discover_target_capabilities_async(\n", " target=target,\n", " test_modalities={frozenset({\"text\"}), frozenset({\"text\", \"image_path\"})},\n", " test_assets={\"image_path\": \"/path/to/test_image.png\"},\n", @@ -683,7 +611,7 @@ "\n", "```python\n", "# Re-query only JSON support; other declared flags pass through unchanged.\n", - "queried = await discover_target_async(\n", + "queried = await discover_target_capabilities_async(\n", " target=target,\n", " capabilities={CapabilityName.JSON_OUTPUT, CapabilityName.JSON_SCHEMA},\n", ")\n", @@ -692,12 +620,12 @@ }, { "cell_type": "markdown", - "id": "25", + "id": "23", "metadata": {}, "source": [ "## 8. Applying probed capabilities back onto the target\n", "\n", - "`discover_target_async` is intentionally pure: it returns a `TargetCapabilities` without\n", + "`discover_target_capabilities_async` is intentionally pure: it returns a `TargetCapabilities` without\n", "mutating the target. That lets you inspect (or diff against the declared view, log, gate on\n", "the result) before committing. Once you're satisfied, call `target.apply_capabilities(...)`\n", "to install the probed view on the instance. The target's existing\n", @@ -716,10 +644,8 @@ { "cell_type": "code", "execution_count": null, - "id": "26", - "metadata": { - "lines_to_next_cell": 2 - }, + "id": "24", + "metadata": {}, "outputs": [ { "name": "stdout", @@ -730,7 +656,7 @@ " supports_system_prompt: False\n", " supports_json_output: False\n", "\n", - "probed (returned from discover_target_async, target NOT yet updated):\n", + "probed (returned from discover_target_capabilities_async, target NOT yet updated):\n", " supports_multi_turn: True\n", " supports_system_prompt: True\n", " supports_json_output: True\n", @@ -777,9 +703,9 @@ "print(f\" supports_json_output: {endpoint_target.capabilities.supports_json_output}\")\n", "\n", "# Step 1: discover. No mutation yet — `endpoint_target.capabilities` is unchanged.\n", - "probed_caps = await discover_target_async(target=endpoint_target, per_probe_timeout_s=5.0) # type: ignore\n", + "probed_caps = await discover_target_capabilities_async(target=endpoint_target, per_probe_timeout_s=5.0) # type: ignore\n", "\n", - "print(\"\\nprobed (returned from discover_target_async, target NOT yet updated):\")\n", + "print(\"\\nprobed (returned from discover_target_capabilities_async, target NOT yet updated):\")\n", "print(f\" supports_multi_turn: {probed_caps.supports_multi_turn}\")\n", "print(f\" supports_system_prompt: {probed_caps.supports_system_prompt}\")\n", "print(f\" supports_json_output: {probed_caps.supports_json_output}\")\n", @@ -818,10 +744,8 @@ } ], "metadata": { - "kernelspec": { - "display_name": "pyrit", - "language": "python", - "name": "python3" + "jupytext": { + "main_language": "python" }, "language_info": { "codemirror_mode": { diff --git a/doc/code/targets/6_1_target_capabilities.py b/doc/code/targets/6_1_target_capabilities.py index 8952410e8f..2793083d12 100644 --- a/doc/code/targets/6_1_target_capabilities.py +++ b/doc/code/targets/6_1_target_capabilities.py @@ -5,7 +5,7 @@ # extension: .py # format_name: percent # format_version: '1.3' -# jupytext_version: 1.19.0 +# jupytext_version: 1.19.1 # --- # %% [markdown] @@ -252,36 +252,35 @@ # 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 -# `discover_target_capabilities_async`, `discover_target_modalities_async`, or the convenience wrapper -# `discover_target_async` that runs both and returns a best-effort `TargetCapabilities`. +# `discover_target_capabilities_async`. It runs both the boolean capability probes and the input +# modality probes and returns a best-effort `TargetCapabilities`. # -# `discover_target_capabilities_async` walks each capability that has a registered probe (currently +# Internally it 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. +# minimal request, and includes the capability in the result 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. -# -# `discover_target_modalities_async` does the same for input modality combinations declared in +# unsupported. The original configuration is restored before the function returns. The same +# treatment is applied to each input modality combination 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. `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 will still be reported as supporting that capability. +# transient errors before being declared failed. The returned `TargetCapabilities` is 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 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 +# 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`. 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 discover_target_async +# from pyrit.prompt_target import discover_target_capabilities_async # -# queried = await discover_target_async(target=target) +# queried = await discover_target_capabilities_async(target=target) # print(queried) # ``` # @@ -294,11 +293,7 @@ from unittest.mock import AsyncMock from pyrit.models import MessagePiece -from pyrit.prompt_target import ( - discover_target_async, - discover_target_capabilities_async, - discover_target_modalities_async, -) +from pyrit.prompt_target import discover_target_capabilities_async def _ok_response(): @@ -321,9 +316,13 @@ def _ok_response(): probe_target._send_prompt_to_target_async = AsyncMock(return_value=_ok_response()) # type: ignore[method-assign] queried = await discover_target_capabilities_async(target=probe_target, per_probe_timeout_s=5.0) # type: ignore -print("queried capabilities:") -for capability in sorted(queried, key=lambda c: c.value): - print(f" - {capability.value}") +print("discover_target_capabilities_async result:") +print(f" supports_multi_turn: {queried.supports_multi_turn}") +print(f" supports_system_prompt: {queried.supports_system_prompt}") +print(f" supports_multi_message_pieces: {queried.supports_multi_message_pieces}") +print(f" supports_json_output: {queried.supports_json_output}") +print(f" supports_json_schema: {queried.supports_json_schema}") +print(f" input_modalities: {sorted(sorted(m) for m in queried.input_modalities)}") # %% [markdown] # To narrow the probe to specific capabilities (faster, fewer calls), pass `capabilities=`: @@ -337,53 +336,20 @@ def _ok_response(): # ) # ``` # -# If you only care about accepted input combinations, call -# `discover_target_modalities_async` directly. The example below uses the -# packaged default probe assets for the non-text modalities PyRIT ships. -# Pass `test_assets=` only when you want to override those defaults or probe -# a modality without a packaged asset. -# -# `discover_target_async` is the most common entry point: it runs both the capability and modality -# probes and assembles a best-effort `TargetCapabilities` you can drop into a -# `TargetConfiguration`, so the rest of PyRIT operates on probed values where available and -# declared values otherwise. - -# %% -probe_target._send_prompt_to_target_async = AsyncMock(return_value=_ok_response()) # type: ignore[method-assign] - -queried_modalities = await discover_target_modalities_async( - target=probe_target, - test_modalities={frozenset({"text"}), frozenset({"text", "image_path"})}, - per_probe_timeout_s=5.0, -) # type: ignore - -print("discover_target_modalities_async result:") -for combination in sorted(sorted(m) for m in queried_modalities): - print(f" - {combination}") - -# %% -probe_target._send_prompt_to_target_async = AsyncMock(return_value=_ok_response()) # type: ignore[method-assign] - -queried_caps = await discover_target_async(target=probe_target, per_probe_timeout_s=5.0) # type: ignore -print("discover_target_async result:") -print(f" supports_multi_turn: {queried_caps.supports_multi_turn}") -print(f" supports_system_prompt: {queried_caps.supports_system_prompt}") -print(f" supports_multi_message_pieces: {queried_caps.supports_multi_message_pieces}") -print(f" supports_json_output: {queried_caps.supports_json_output}") -print(f" supports_json_schema: {queried_caps.supports_json_schema}") -print(f" input_modalities: {sorted(sorted(m) for m in queried_caps.input_modalities)}") +# Similarly, narrow the modality probe set with `test_modalities=` and override the +# packaged default probe assets with `test_assets=`. # %% [markdown] # ### Discovering undeclared modalities # -# By default `discover_target_async` only probes modality combinations the target already +# By default `discover_target_capabilities_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=` explicitly to # probe combinations beyond the declared baseline. Provide `test_assets=` as well if you need # to override the packaged defaults or probe a modality without one: # # ```python -# queried = await discover_target_async( +# queried = await discover_target_capabilities_async( # target=target, # test_modalities={frozenset({"text"}), frozenset({"text", "image_path"})}, # test_assets={"image_path": "/path/to/test_image.png"}, @@ -397,7 +363,7 @@ def _ok_response(): # # ```python # # Re-query only JSON support; other declared flags pass through unchanged. -# queried = await discover_target_async( +# queried = await discover_target_capabilities_async( # target=target, # capabilities={CapabilityName.JSON_OUTPUT, CapabilityName.JSON_SCHEMA}, # ) @@ -406,7 +372,7 @@ def _ok_response(): # %% [markdown] # ## 8. Applying probed capabilities back onto the target # -# `discover_target_async` is intentionally pure: it returns a `TargetCapabilities` without +# `discover_target_capabilities_async` is intentionally pure: it returns a `TargetCapabilities` without # mutating the target. That lets you inspect (or diff against the declared view, log, gate on # the result) before committing. Once you're satisfied, call `target.apply_capabilities(...)` # to install the probed view on the instance. The target's existing @@ -450,9 +416,9 @@ def _ok_response(): print(f" supports_json_output: {endpoint_target.capabilities.supports_json_output}") # Step 1: discover. No mutation yet — `endpoint_target.capabilities` is unchanged. -probed_caps = await discover_target_async(target=endpoint_target, per_probe_timeout_s=5.0) # type: ignore +probed_caps = await discover_target_capabilities_async(target=endpoint_target, per_probe_timeout_s=5.0) # type: ignore -print("\nprobed (returned from discover_target_async, target NOT yet updated):") +print("\nprobed (returned from discover_target_capabilities_async, target NOT yet updated):") print(f" supports_multi_turn: {probed_caps.supports_multi_turn}") print(f" supports_system_prompt: {probed_caps.supports_system_prompt}") print(f" supports_json_output: {probed_caps.supports_json_output}") diff --git a/pyrit/prompt_target/__init__.py b/pyrit/prompt_target/__init__.py index 3c74da5b17..82f897c156 100644 --- a/pyrit/prompt_target/__init__.py +++ b/pyrit/prompt_target/__init__.py @@ -15,9 +15,7 @@ from pyrit.prompt_target.azure_ml_chat_target import AzureMLChatTarget from pyrit.prompt_target.common.conversation_normalization_pipeline import ConversationNormalizationPipeline from pyrit.prompt_target.common.discover_target_capabilities import ( - discover_target_async, discover_target_capabilities_async, - discover_target_modalities_async, ) from pyrit.prompt_target.common.prompt_chat_target import PromptChatTarget from pyrit.prompt_target.common.prompt_target import PromptTarget @@ -102,14 +100,12 @@ def __getattr__(name: str) -> object: "PromptChatTarget", "PromptShieldTarget", "PromptTarget", - "discover_target_capabilities_async", "RealtimeTarget", "TargetCapabilities", "TargetConfiguration", "TargetRequirements", "UnsupportedCapabilityBehavior", "TextTarget", - "discover_target_async", - "discover_target_modalities_async", + "discover_target_capabilities_async", "WebSocketCopilotTarget", ] diff --git a/pyrit/prompt_target/common/discover_target_capabilities.py b/pyrit/prompt_target/common/discover_target_capabilities.py index 9982713515..859d07d428 100644 --- a/pyrit/prompt_target/common/discover_target_capabilities.py +++ b/pyrit/prompt_target/common/discover_target_capabilities.py @@ -6,14 +6,14 @@ This module exposes two complementary probes: -* :func:`discover_target_capabilities_async` discovers the boolean capability flags +* :func:`_discover_capability_flags_async` discovers 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 the target's declared native support from ``target.capabilities``. -* :func:`discover_target_modalities_async` discovers which input modality +* :func:`_discover_input_modalities_async` discovers which input modality combinations a target actually supports by sending a minimal test request for each combination declared in ``TargetCapabilities.input_modalities``. @@ -91,6 +91,9 @@ def _json_enforcing_target_types() -> tuple[type[PromptTarget], ...]: keep this concern entirely within the discovery module. Class objects (not strings) are returned so renames are caught by import errors rather than silently flipping the log behavior. + + Returns: + tuple[type[PromptTarget], ...]: The target classes that enforce JSON hints. """ from pyrit.prompt_target.openai.openai_chat_target import OpenAIChatTarget from pyrit.prompt_target.openai.openai_response_target import OpenAIResponseTarget @@ -507,7 +510,7 @@ async def _probe_json_schema_async(target: PromptTarget, timeout_s: float, retri } -async def discover_target_capabilities_async( +async def _discover_capability_flags_async( *, target: PromptTarget, capabilities: Iterable[CapabilityName] | None = None, @@ -590,7 +593,7 @@ async def discover_target_capabilities_async( # Default mapping of non-text modalities to packaged probe assets. Callers can # override via the ``test_assets`` parameter of -# :func:`discover_target_modalities_async`. Modalities whose assets do not exist +# :func:`_discover_input_modalities_async`. Modalities whose assets do not exist # on disk are skipped (logged and excluded from the result). DEFAULT_TEST_ASSETS: dict[PromptDataType, str] = { "audio_path": str(_TARGET_CAPABILITIES_DATASET_PATH / "probe_audio.wav"), @@ -598,7 +601,7 @@ async def discover_target_capabilities_async( } -async def discover_target_modalities_async( +async def _discover_input_modalities_async( *, target: PromptTarget, test_modalities: set[frozenset[PromptDataType]] | None = None, @@ -637,7 +640,7 @@ async def discover_target_modalities_async( declared = target.capabilities.input_modalities test_modalities = set(declared) elif not test_modalities: - logger.info("discover_target_modalities_async called with an empty test_modalities set; nothing to probe.") + logger.info("_discover_input_modalities_async called with an empty test_modalities set; nothing to probe.") return set() assets = test_assets if test_assets is not None else DEFAULT_TEST_ASSETS @@ -669,7 +672,7 @@ async def discover_target_modalities_async( return queried -async def discover_target_async( +async def discover_target_capabilities_async( *, target: PromptTarget, per_probe_timeout_s: float = DEFAULT_PROBE_TIMEOUT_SECONDS, @@ -680,25 +683,32 @@ async def discover_target_async( apply: bool = False, ) -> TargetCapabilities: """ - Probe capabilities and modalities and return a merged result. - - This wraps :func:`discover_target_capabilities_async` and - :func:`discover_target_modalities_async` and returns a best-effort + Probe both the boolean capability flags and the input modality combinations + that ``target`` accepts, and return a merged best-effort :class:`TargetCapabilities`. + Boolean capabilities with a registered probe are checked with live + requests; capabilities without a probe fall back to the target's + declared native support. Each input modality combination is checked + with a minimal request built from the supplied test assets. + "Supported" means the request was accepted — a target that silently + ignores a feature is still reported as supporting it. + 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:`discover_target_modalities_async`. Defaults to the - target's declared ``input_modalities``. + modality combinations to probe. Defaults to the target's declared + ``input_modalities``. Combinations not listed here fall back to + the target's declared support. test_assets (dict[PromptDataType, str] | None): Mapping from non-text - modality to a file path. See :func:`discover_target_modalities_async`. + 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. capabilities (Iterable[CapabilityName] | None): Capabilities to probe. - See :func:`discover_target_capabilities_async`. Defaults to every - member of :class:`CapabilityName`. + Defaults to every member of :class:`CapabilityName`. Capabilities + not listed here fall back to the target's declared support. 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. @@ -715,13 +725,13 @@ async def discover_target_async( """ capabilities_to_probe = list(capabilities) if capabilities is not None else None - queried_caps = await discover_target_capabilities_async( + queried_caps = await _discover_capability_flags_async( target=target, capabilities=capabilities_to_probe, per_probe_timeout_s=per_probe_timeout_s, retries=retries, ) - queried_modalities = await discover_target_modalities_async( + queried_modalities = await _discover_input_modalities_async( target=target, test_modalities=test_modalities, test_assets=test_assets, diff --git a/tests/unit/prompt_target/test_discover_target_capabilities.py b/tests/unit/prompt_target/test_discover_target_capabilities.py index 44e4a50d04..e21f24a6f2 100644 --- a/tests/unit/prompt_target/test_discover_target_capabilities.py +++ b/tests/unit/prompt_target/test_discover_target_capabilities.py @@ -14,10 +14,10 @@ _CAPABILITY_PROBES, DEFAULT_TEST_ASSETS, _create_test_message, + _discover_capability_flags_async, + _discover_input_modalities_async, _permissive_configuration, - discover_target_async, discover_target_capabilities_async, - discover_target_modalities_async, ) from pyrit.prompt_target.common.prompt_target import PromptTarget from pyrit.prompt_target.common.target_capabilities import ( @@ -110,7 +110,7 @@ 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 discover_target_capabilities_async(target=target) + result = await _discover_capability_flags_async(target=target) # Every capability with a probe should be in the result. for capability in _CAPABILITY_PROBES: @@ -120,7 +120,7 @@ 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 discover_target_capabilities_async(target=target) + result = await _discover_capability_flags_async(target=target) for capability in _CAPABILITY_PROBES: assert capability not in result @@ -129,7 +129,7 @@ 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 discover_target_capabilities_async(target=target) + result = await _discover_capability_flags_async(target=target) for capability in _CAPABILITY_PROBES: assert capability not in result @@ -139,7 +139,7 @@ async def test_filters_by_requested_capabilities(self) -> None: target._send_prompt_to_target_async = AsyncMock(return_value=_ok_response()) # type: ignore[method-assign] requested = {CapabilityName.SYSTEM_PROMPT, CapabilityName.MULTI_TURN} - result = await discover_target_capabilities_async(target=target, capabilities=requested) + result = await _discover_capability_flags_async(target=target, capabilities=requested) assert result == requested @@ -151,7 +151,7 @@ async def test_capability_without_probe_falls_back_to_declared_support(self) -> ) target._send_prompt_to_target_async = AsyncMock(return_value=_ok_response()) # type: ignore[method-assign] - result = await discover_target_capabilities_async( + result = await _discover_capability_flags_async( target=target, capabilities={CapabilityName.EDITABLE_HISTORY}, ) @@ -164,7 +164,7 @@ async def test_capability_without_probe_excluded_when_not_declared(self) -> None target._configuration = TargetConfiguration(capabilities=TargetCapabilities()) target._send_prompt_to_target_async = AsyncMock(return_value=_ok_response()) # type: ignore[method-assign] - result = await discover_target_capabilities_async( + result = await _discover_capability_flags_async( target=target, capabilities={CapabilityName.EDITABLE_HISTORY}, ) @@ -200,7 +200,7 @@ async def test_capability_without_probe_excluded_when_only_adapted(self, monkeyp ), ) - result = await discover_target_capabilities_async( + result = await _discover_capability_flags_async( target=target, capabilities={CapabilityName.SYSTEM_PROMPT}, ) @@ -216,7 +216,7 @@ async def test_accepts_single_pass_iterable(self) -> None: 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 discover_target_capabilities_async(target=target, capabilities=gen) + result = await _discover_capability_flags_async(target=target, capabilities=gen) assert CapabilityName.SYSTEM_PROMPT in result assert CapabilityName.EDITABLE_HISTORY in result @@ -225,7 +225,7 @@ 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 discover_target_capabilities_async( + result = await _discover_capability_flags_async( target=target, capabilities={CapabilityName.JSON_OUTPUT}, retries=0, @@ -241,7 +241,7 @@ async def test_retries_use_exponential_backoff(self) -> None: with patch( "pyrit.prompt_target.common.discover_target_capabilities.asyncio.sleep", new_callable=AsyncMock ) as sleep_mock: - result = await discover_target_capabilities_async( + result = await _discover_capability_flags_async( target=target, capabilities={CapabilityName.JSON_OUTPUT}, retries=2, @@ -266,7 +266,7 @@ async def test_non_retryable_validation_errors_fail_fast(self) -> None: with patch( "pyrit.prompt_target.common.discover_target_capabilities.asyncio.sleep", new_callable=AsyncMock ) as sleep_mock: - result = await discover_target_capabilities_async( + result = await _discover_capability_flags_async( target=target, capabilities={CapabilityName.JSON_OUTPUT}, retries=3, @@ -282,7 +282,7 @@ async def test_restores_configuration_after_probing(self) -> None: original = target.configuration target._send_prompt_to_target_async = AsyncMock(return_value=_ok_response()) # type: ignore[method-assign] - await discover_target_capabilities_async(target=target) + await _discover_capability_flags_async(target=target) assert target.configuration is original @@ -290,7 +290,7 @@ 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 discover_target_capabilities_async( + await _discover_capability_flags_async( target=target, capabilities={CapabilityName.MULTI_TURN}, ) @@ -318,7 +318,7 @@ 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 discover_target_capabilities_async( + result = await _discover_capability_flags_async( target=target, capabilities={CapabilityName.MULTI_TURN}, ) @@ -332,7 +332,7 @@ 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 discover_target_capabilities_async( + await _discover_capability_flags_async( target=target, capabilities={CapabilityName.JSON_SCHEMA}, ) @@ -353,7 +353,7 @@ async def test_logs_debug_for_unenforced_json_probe( target._send_prompt_to_target_async = AsyncMock(return_value=_ok_response()) # type: ignore[method-assign] with caplog.at_level(logging.DEBUG): - result = await discover_target_capabilities_async(target=target, capabilities={capability}) + result = await _discover_capability_flags_async(target=target, capabilities={capability}) assert result == {capability} matching = [r for r in caplog.records if r.message.startswith("JSON capability probes")] @@ -367,7 +367,7 @@ async def test_logs_unenforced_json_probe_summary_once_for_both_capabilities( target._send_prompt_to_target_async = AsyncMock(return_value=_ok_response()) # type: ignore[method-assign] with caplog.at_level(logging.DEBUG): - await discover_target_capabilities_async( + await _discover_capability_flags_async( target=target, capabilities={CapabilityName.JSON_OUTPUT, CapabilityName.JSON_SCHEMA}, ) @@ -383,7 +383,7 @@ async def test_does_not_log_unenforced_json_probe_when_probe_fails(self, caplog: target._send_prompt_to_target_async = AsyncMock(side_effect=Exception("boom")) # type: ignore[method-assign] with caplog.at_level(logging.DEBUG): - result = await discover_target_capabilities_async( + result = await _discover_capability_flags_async( target=target, capabilities={CapabilityName.JSON_OUTPUT, CapabilityName.JSON_SCHEMA}, retries=0, @@ -404,7 +404,7 @@ async def test_does_not_log_debug_for_enforced_json_probe(self, caplog: pytest.L ), caplog.at_level(logging.DEBUG), ): - result = await discover_target_capabilities_async( + result = await _discover_capability_flags_async( target=target, capabilities={CapabilityName.JSON_OUTPUT}, ) @@ -426,7 +426,7 @@ async def test_subclass_of_enforced_target_does_not_log(self, caplog: pytest.Log ), caplog.at_level(logging.DEBUG), ): - await discover_target_capabilities_async( + await _discover_capability_flags_async( target=target, capabilities={CapabilityName.JSON_OUTPUT}, ) @@ -437,7 +437,7 @@ async def test_system_prompt_probe_installs_system_message_and_sends_user(self) target = MockPromptTarget() target._send_prompt_to_target_async = AsyncMock(return_value=_ok_response()) # type: ignore[method-assign] - await discover_target_capabilities_async( + await _discover_capability_flags_async( target=target, capabilities={CapabilityName.SYSTEM_PROMPT}, ) @@ -459,7 +459,7 @@ 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 discover_target_capabilities_async( + await _discover_capability_flags_async( target=target, capabilities={CapabilityName.MULTI_MESSAGE_PIECES}, ) @@ -479,7 +479,7 @@ async def test_probes_run_under_permissive_configuration(self) -> None: send_mock = AsyncMock(return_value=_ok_response()) target._send_prompt_to_target_async = send_mock # type: ignore[method-assign] - result = await discover_target_capabilities_async( + result = await _discover_capability_flags_async( target=target, capabilities={CapabilityName.MULTI_MESSAGE_PIECES}, ) @@ -510,7 +510,7 @@ async def reject_system_roles(*, normalized_conversation: list[Message]) -> list target._send_prompt_to_target_async = AsyncMock(side_effect=reject_system_roles) # type: ignore[method-assign] - result = await discover_target_capabilities_async( + result = await _discover_capability_flags_async( target=target, capabilities={CapabilityName.SYSTEM_PROMPT}, ) @@ -537,7 +537,7 @@ async def require_native_system_role(*, normalized_conversation: list[Message]) target._send_prompt_to_target_async = AsyncMock(side_effect=require_native_system_role) # type: ignore[method-assign] - result = await discover_target_capabilities_async( + result = await _discover_capability_flags_async( target=target, capabilities={CapabilityName.SYSTEM_PROMPT}, ) @@ -557,7 +557,7 @@ async def _send_prompt_to_target_async(self, *, normalized_conversation: list[Me target = _MinimalTarget() target._send_prompt_to_target_async = AsyncMock(return_value=_ok_response()) # type: ignore[method-assign] - result = await discover_target_capabilities_async(target=target) + result = await _discover_capability_flags_async(target=target) for capability in _CAPABILITY_PROBES: assert capability in result @@ -639,7 +639,7 @@ async def test_all_combinations_supported(self) -> None: _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 discover_target_modalities_async(target=target) + result = await _discover_input_modalities_async(target=target) assert frozenset({"text"}) in result @@ -648,7 +648,7 @@ async def test_exception_excludes_combination(self) -> None: _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 discover_target_modalities_async(target=target) + result = await _discover_input_modalities_async(target=target) assert result == set() @@ -657,7 +657,7 @@ async def test_error_response_excludes_combination(self) -> None: _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 discover_target_modalities_async(target=target) + result = await _discover_input_modalities_async(target=target) assert result == set() @@ -677,7 +677,7 @@ async def selective_send(*, normalized_conversation: list[Message]) -> list[Mess target._send_prompt_to_target_async = selective_send # type: ignore[method-assign] - result = await discover_target_modalities_async( + result = await _discover_input_modalities_async( target=target, test_assets={"image_path": image_asset}, ) @@ -691,7 +691,7 @@ async def test_explicit_test_modalities_overrides_declared(self, image_asset: st _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 discover_target_modalities_async( + result = await _discover_input_modalities_async( target=target, test_modalities={frozenset({"text"}), frozenset({"text", "image_path"})}, test_assets={"image_path": image_asset}, @@ -707,7 +707,7 @@ async def test_combination_skipped_when_asset_missing(self, tmp_path: Path) -> N # An explicit empty mapping disables the packaged defaults, so # image_path combinations are skipped instead of probed. - result = await discover_target_modalities_async(target=target, test_assets={}) + result = await _discover_input_modalities_async(target=target, test_assets={}) assert result == set() assert target._send_prompt_to_target_async.await_count == 0 @@ -717,7 +717,7 @@ async def test_empty_test_modalities_returns_empty_without_probing(self) -> None _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 discover_target_modalities_async(target=target, test_modalities=set()) + result = await _discover_input_modalities_async(target=target, test_modalities=set()) assert result == set() assert target._send_prompt_to_target_async.await_count == 0 @@ -733,7 +733,7 @@ async def test_explicit_test_modalities_runs_under_permissive_configuration(self send_mock = AsyncMock(return_value=_ok_response()) target._send_prompt_to_target_async = send_mock # type: ignore[method-assign] - result = await discover_target_modalities_async( + result = await _discover_input_modalities_async( target=target, test_modalities={frozenset({"text", "image_path"})}, test_assets={"image_path": image_asset}, @@ -760,7 +760,7 @@ async def _hang(**_kwargs: object) -> list[Message]: target._send_prompt_to_target_async = AsyncMock(side_effect=_hang) # type: ignore[method-assign] - result = await discover_target_capabilities_async( + result = await _discover_capability_flags_async( target=target, capabilities={CapabilityName.JSON_OUTPUT}, per_probe_timeout_s=0.01, @@ -784,7 +784,7 @@ async def test_returns_false_when_memory_seed_raises(self) -> None: target._send_prompt_to_target_async = send_mock # type: ignore[method-assign] with patch.object(target._memory, "add_message_to_memory", side_effect=RuntimeError("memory offline")): - result = await discover_target_capabilities_async( + result = await _discover_capability_flags_async( target=target, capabilities={CapabilityName.SYSTEM_PROMPT}, ) @@ -798,7 +798,7 @@ async def test_returns_false_when_memory_seed_raises(self) -> None: class TestVerifyTargetAsync: async def test_returns_target_capabilities_assembled_from_probes(self) -> None: """ - ``discover_target_async`` runs both the capability and modality probes + ``discover_target_capabilities_async`` runs both the capability and modality probes and assembles a :class:`TargetCapabilities` populated from the queried results, copying ``output_modalities`` from the target's declared capabilities and deriving editable history conservatively. @@ -811,7 +811,7 @@ async def test_returns_target_capabilities_assembled_from_probes(self) -> None: target._configuration = TargetConfiguration(capabilities=declared) target._send_prompt_to_target_async = AsyncMock(return_value=_ok_response()) # type: ignore[method-assign] - result = await discover_target_async(target=target, per_probe_timeout_s=5.0) + result = await discover_target_capabilities_async(target=target, per_probe_timeout_s=5.0) assert isinstance(result, TargetCapabilities) # Single-piece probes that don't touch memory always succeed when @@ -842,7 +842,7 @@ async def test_excludes_capabilities_when_probe_send_fails(self) -> None: target._configuration = TargetConfiguration(capabilities=declared) target._send_prompt_to_target_async = AsyncMock(side_effect=RuntimeError("boom")) # type: ignore[method-assign] - result = await discover_target_async(target=target, per_probe_timeout_s=0.5) + result = await discover_target_capabilities_async(target=target, per_probe_timeout_s=0.5) assert result.supports_multi_turn is False assert result.supports_system_prompt is False @@ -863,7 +863,7 @@ async def test_empty_response_treated_as_failure(self) -> None: target = MockPromptTarget() target._send_prompt_to_target_async = AsyncMock(return_value=[]) # type: ignore[method-assign] - result = await discover_target_capabilities_async( + result = await _discover_capability_flags_async( target=target, capabilities={CapabilityName.JSON_OUTPUT, CapabilityName.MULTI_MESSAGE_PIECES}, ) @@ -880,7 +880,7 @@ async def test_response_with_no_pieces_treated_as_failure(self) -> None: empty_msg = target._send_prompt_to_target_async.return_value[0] empty_msg.message_pieces = [] - result = await discover_target_capabilities_async( + result = await _discover_capability_flags_async( target=target, capabilities={CapabilityName.JSON_OUTPUT}, ) @@ -895,21 +895,21 @@ async def test_mixed_empty_message_in_response_treated_as_failure(self) -> None: empty.message_pieces = [] target._send_prompt_to_target_async = AsyncMock(return_value=[ok, empty]) # type: ignore[method-assign] - result = await discover_target_capabilities_async( + result = await _discover_capability_flags_async( target=target, capabilities={CapabilityName.JSON_OUTPUT}, ) assert result == set() - async def test_discover_target_async_forwards_test_modalities(self, image_asset: str) -> None: + async def test_discover_target_capabilities_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 discover_target_async( + result = await discover_target_capabilities_async( target=target, test_modalities={extra_combo}, test_assets={"image_path": image_asset}, @@ -919,7 +919,7 @@ async def test_discover_target_async_forwards_test_modalities(self, image_asset: # The undeclared combination is in the result only if test_modalities was forwarded. assert extra_combo in result.input_modalities - async def test_discover_target_async_preserves_declared_modalities_when_test_modalities_narrowed( + async def test_discover_target_capabilities_async_preserves_declared_modalities_when_test_modalities_narrowed( self, image_asset: str ) -> None: declared_combo = frozenset({"text"}) @@ -929,7 +929,7 @@ async def test_discover_target_async_preserves_declared_modalities_when_test_mod target._configuration = TargetConfiguration(capabilities=declared) target._send_prompt_to_target_async = AsyncMock(return_value=_ok_response()) - result = await discover_target_async( + result = await discover_target_capabilities_async( target=target, test_modalities={probed_combo}, test_assets={"image_path": image_asset}, @@ -938,12 +938,12 @@ async def test_discover_target_async_preserves_declared_modalities_when_test_mod assert result.input_modalities == frozenset({declared_combo, probed_combo}) - async def test_discover_target_async_forwards_capabilities(self) -> None: - """``discover_target_async`` must forward ``capabilities`` to narrow the probe set.""" + async def test_discover_target_capabilities_async_forwards_capabilities(self) -> None: + """``discover_target_capabilities_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 discover_target_async( + await discover_target_capabilities_async( target=target, capabilities={CapabilityName.JSON_OUTPUT}, per_probe_timeout_s=2.0, @@ -954,7 +954,7 @@ async def test_discover_target_async_forwards_capabilities(self) -> None: # because multi-turn issues 2 sends). assert target._send_prompt_to_target_async.await_count <= 3 - async def test_discover_target_async_preserves_declared_when_capabilities_narrowed(self) -> None: + async def test_discover_target_capabilities_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 @@ -970,7 +970,7 @@ async def test_discover_target_async_preserves_declared_when_capabilities_narrow target._configuration = TargetConfiguration(capabilities=declared) target._send_prompt_to_target_async = AsyncMock(return_value=_ok_response()) # type: ignore[method-assign] - result = await discover_target_async( + result = await discover_target_capabilities_async( target=target, capabilities={CapabilityName.JSON_OUTPUT}, per_probe_timeout_s=2.0, @@ -984,7 +984,7 @@ async def test_discover_target_async_preserves_declared_when_capabilities_narrow assert result.supports_json_schema is True assert result.supports_editable_history is True - async def test_discover_target_async_drops_editable_history_when_multi_turn_probe_fails(self) -> None: + async def test_discover_target_capabilities_async_drops_editable_history_when_multi_turn_probe_fails(self) -> None: """Editable history must not remain true when probing disproves multi-turn support.""" declared = TargetCapabilities( supports_multi_turn=True, @@ -1002,12 +1002,12 @@ async def selective_send(*, normalized_conversation: list[Message]) -> list[Mess target._send_prompt_to_target_async = AsyncMock(side_effect=selective_send) # type: ignore[method-assign] - result = await discover_target_async(target=target, per_probe_timeout_s=2.0) + result = await discover_target_capabilities_async(target=target, per_probe_timeout_s=2.0) assert result.supports_multi_turn is False assert result.supports_editable_history is False - async def test_discover_target_async_accepts_single_pass_iterable(self) -> None: + async def test_discover_target_capabilities_async_accepts_single_pass_iterable(self) -> None: declared = TargetCapabilities( supports_multi_turn=True, supports_editable_history=True, @@ -1017,7 +1017,7 @@ async def test_discover_target_async_accepts_single_pass_iterable(self) -> None: target._send_prompt_to_target_async = AsyncMock(return_value=_ok_response()) # type: ignore[method-assign] gen = (c for c in [CapabilityName.JSON_OUTPUT, CapabilityName.EDITABLE_HISTORY]) - result = await discover_target_async( + result = await discover_target_capabilities_async( target=target, capabilities=gen, per_probe_timeout_s=2.0, @@ -1026,7 +1026,7 @@ async def test_discover_target_async_accepts_single_pass_iterable(self) -> None: assert result.supports_json_output is True assert result.supports_editable_history is True - async def test_discover_target_async_apply_installs_capabilities_on_target(self) -> None: + async def test_discover_target_capabilities_async_apply_installs_capabilities_on_target(self) -> None: """When ``apply=True``, the discovered capabilities are installed on the target.""" declared = TargetCapabilities(supports_multi_turn=False, supports_system_prompt=False) target = MockPromptTarget() @@ -1035,19 +1035,19 @@ async def test_discover_target_async_apply_installs_capabilities_on_target(self) assert target.capabilities.supports_multi_turn is False - result = await discover_target_async(target=target, per_probe_timeout_s=2.0, apply=True) + result = await discover_target_capabilities_async(target=target, per_probe_timeout_s=2.0, apply=True) assert target.capabilities == result assert target.capabilities.supports_multi_turn is True - async def test_discover_target_async_apply_defaults_to_false(self) -> None: - """By default, ``discover_target_async`` must not mutate the target.""" + async def test_discover_target_capabilities_async_apply_defaults_to_false(self) -> None: + """By default, ``discover_target_capabilities_async`` must not mutate the target.""" declared = TargetCapabilities(supports_multi_turn=False, supports_system_prompt=False) target = MockPromptTarget() target._configuration = TargetConfiguration(capabilities=declared) target._send_prompt_to_target_async = AsyncMock(return_value=_ok_response()) # type: ignore[method-assign] - result = await discover_target_async(target=target, per_probe_timeout_s=2.0) + result = await discover_target_capabilities_async(target=target, per_probe_timeout_s=2.0) # Result reflects probe; target capabilities remain at declared values. assert result.supports_multi_turn is True @@ -1067,7 +1067,7 @@ async def test_returns_false_when_history_seed_raises(self) -> None: target._send_prompt_to_target_async = send_mock # type: ignore[method-assign] with patch.object(target._memory, "add_message_to_memory", side_effect=RuntimeError("memory offline")): - result = await discover_target_capabilities_async( + result = await _discover_capability_flags_async( target=target, capabilities={CapabilityName.MULTI_TURN}, ) From 25d0a9f4c7c2d5f70dd962665e47e9859db1fbe6 Mon Sep 17 00:00:00 2001 From: hannahwestra25 Date: Wed, 13 May 2026 15:24:42 -0400 Subject: [PATCH 27/27] remove query_target_capabilities --- .../common/query_target_capabilities.py | 830 ------------------ 1 file changed, 830 deletions(-) delete mode 100644 pyrit/prompt_target/common/query_target_capabilities.py diff --git a/pyrit/prompt_target/common/query_target_capabilities.py b/pyrit/prompt_target/common/query_target_capabilities.py deleted file mode 100644 index eac20e12ff..0000000000 --- a/pyrit/prompt_target/common/query_target_capabilities.py +++ /dev/null @@ -1,830 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# Licensed under the MIT license. - -""" -Runtime capability and modality discovery for prompt targets. - -This module exposes two complementary probes: - -* :func:`discover_target_capabilities_async` discovers 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 the target's declared native support from - ``target.capabilities``. -* :func:`discover_target_modalities_async` discovers 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*. They do not prove - that the endpoint enforced the feature, and the JSON probes are only - meaningful for targets that translate ``prompt_metadata`` JSON hints into - provider request fields. Treat the results as an upper bound on support and - validate response content separately when that distinction matters. -""" - -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.common.path import DATASETS_PATH -from pyrit.models import Message, MessagePiece, PromptDataType -from pyrit.prompt_target.common.prompt_target import PromptTarget -from pyrit.prompt_target.common.target_capabilities import ( - CapabilityName, - TargetCapabilities, -) -from pyrit.prompt_target.common.target_configuration import TargetConfiguration - -logger = logging.getLogger(__name__) - -# Per-call timeout (seconds) applied to every discovery request. Override per-call via -# the ``per_probe_timeout_s`` parameter on the public functions. -DEFAULT_PROBE_TIMEOUT_SECONDS: float = 30.0 -DEFAULT_PROBE_RETRY_BACKOFF_SECONDS: float = 0.1 -MAX_PROBE_RETRY_BACKOFF_SECONDS: float = 1.0 - -# Exceptions that are deterministic on the probe payload and will not become -# valid on a retry (malformed Message, type errors, missing attributes, etc.). -# These fail the probe immediately rather than wasting backoff time. -_NON_RETRYABLE_PROBE_EXCEPTIONS: tuple[type[BaseException], ...] = ( - ValueError, - TypeError, - AttributeError, -) - -# 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]] - - -def _json_enforcing_target_types() -> tuple[type[PromptTarget], ...]: - """ - Return the tuple of target classes that translate ``prompt_metadata`` JSON - hints (``response_format``, ``json_schema``) into native provider request - fields. Used to suppress the "JSON probe is upper-bound only" debug log for - these targets and their subclasses. - - Imports are lazy to avoid a circular dependency at module load time and to - keep this concern entirely within the discovery module. Class objects (not - strings) are returned so renames are caught by import errors rather than - silently flipping the log behavior. - """ - from pyrit.prompt_target.openai.openai_chat_target import OpenAIChatTarget - from pyrit.prompt_target.openai.openai_response_target import OpenAIResponseTarget - - return (OpenAIChatTarget, OpenAIResponseTarget) - - -# 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"})}) - -# Packaged fallback assets for non-text modality discovery. -_TARGET_CAPABILITIES_DATASET_PATH = DATASETS_PATH / "prompt_target" / "target_capabilities" - - -@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, - ) - # Rebuild a fresh configuration from the instance's native capabilities so - # probes bypass preflight validation without inheriting ADAPT policy or - # custom normalizer overrides from the target's runtime configuration. - probe_configuration = TargetConfiguration(capabilities=permissive_caps) - target._configuration = probe_configuration - 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``. Transient errors (timeouts, - connection/OS errors) trigger up to ``retries`` retries with a short - exponential backoff. Deterministic errors that will not become valid on - a retry (``ValueError``, ``TypeError``, ``AttributeError`` — typically - from message validation or programmer error in a probe payload) fail - the probe immediately. 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 transient errors are retried; non-retryable errors and - non-error responses are final. Retry attempts use exponential - backoff starting at :data:`DEFAULT_PROBE_RETRY_BACKOFF_SECONDS`. - 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) - if attempt + 1 < attempts: - await _sleep_before_retry_async(attempt=attempt) - continue - except _NON_RETRYABLE_PROBE_EXCEPTIONS as exc: - # Deterministic on the probe payload — retrying will not help. - logger.debug("%s failed with non-retryable error: %s", label, exc) - return False - except Exception as exc: - last_exc = exc - logger.debug("%s failed (attempt %d/%d): %s", label, attempt + 1, attempts, exc) - if attempt + 1 < attempts: - await _sleep_before_retry_async(attempt=attempt) - continue - - if not responses or any(not 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 - - -def _retry_backoff_seconds(*, attempt: int) -> float: - """Return the exponential backoff delay for a retry attempt.""" - return min(DEFAULT_PROBE_RETRY_BACKOFF_SECONDS * (2**attempt), MAX_PROBE_RETRY_BACKOFF_SECONDS) - - -async def _sleep_before_retry_async(*, attempt: int) -> None: - """Sleep for the retry backoff associated with ``attempt``.""" - await asyncio.sleep(_retry_backoff_seconds(attempt=attempt)) - - -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:`pyrit.prompt_target.PromptChatTarget.set_system_prompt` - (which is only defined on ``PromptChatTarget`` subclasses anyway). - ``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 also works - uniformly for plain ``PromptTarget`` subclasses that have no - ``set_system_prompt`` method, and 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. - retries (int): Number of additional attempts after the first failure. - Only exceptions/timeouts are retried; an explicit error response - is final. Defaults to 1. - - 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. - retries (int): Number of additional attempts after the first failure. - Only exceptions/timeouts are retried; an explicit error response - is final. Defaults to 1. - - 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. - retries (int): Number of additional attempts after the first failure. - Only exceptions/timeouts are retried; an explicit error response - is final. Defaults to 1. - - 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. - - This probe is only meaningful for targets that translate PyRIT's JSON - metadata hints into native provider request fields. - - Args: - target (PromptTarget): The target to probe. - timeout_s (float): Per-attempt timeout in seconds. - retries (int): Number of additional attempts after the first failure. - Only exceptions/timeouts are retried; an explicit error response - is final. Defaults to 1. - - 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, - # This only becomes a real JSON-mode request on targets that honor - # PyRIT's JSON metadata contract when building the provider payload. - 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. - - This probe is only meaningful for targets that translate PyRIT's JSON - metadata hints into native provider request fields. - - Args: - target (PromptTarget): The target to probe. - timeout_s (float): Per-attempt timeout in seconds. - retries (int): Number of additional attempts after the first failure. - Only exceptions/timeouts are retried; an explicit error response - is final. Defaults to 1. - - 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, - # As above, this probe is only strong for targets that map these - # metadata keys to native JSON-schema request parameters. - 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 queried 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 discover_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 which capabilities ``target`` accepts. - - Registered capabilities are checked with live requests. Capabilities - without a live probe fall back to declared native support. - - 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 confirmed to work against the target. - """ - capabilities_to_check: list[CapabilityName] = ( - list(capabilities) if capabilities is not None else list(CapabilityName) - ) - - queried: set[CapabilityName] = set() - json_capabilities = {CapabilityName.JSON_OUTPUT, CapabilityName.JSON_SCHEMA} - queried_json_capabilities: set[CapabilityName] = set() - with _permissive_configuration(target=target): - for capability in capabilities_to_check: - probe = _CAPABILITY_PROBES.get(capability) - if probe is None: - # Capabilities without a probe are handled after the permissive - # override is removed so we can read the target's native flags. - continue - - try: - # "Supported" means the request was accepted. A target can - # still ignore the feature semantics after accepting the call. - if await probe(target, per_probe_timeout_s, retries): - queried.add(capability) - if capability in json_capabilities: - queried_json_capabilities.add(capability) - except Exception as exc: - logger.debug("Probe for %s raised: %s", capability.value, exc) - - # JSON probes only verify the target accepted the request, not that the - # target translated the JSON metadata into provider request fields. Emit - # a single summary line when probes succeeded against a target that does - # not enforce JSON hints, so the result is treated as an upper bound. - # ``isinstance`` covers user-defined subclasses of enforcing targets. - if queried_json_capabilities and not isinstance(target, _json_enforcing_target_types()): - logger.debug( - "JSON capability probes %s succeeded for %s, but this target does not translate " - "prompt_metadata JSON hints into provider request fields; treat the result as upper-bound support only.", - sorted(c.value for c in queried_json_capabilities), - type(target).__name__, - ) - - # Read unprobed capabilities from target.capabilities, not - # target.configuration, so ADAPTed behavior is not reported as native - # support. - for capability in capabilities_to_check: - if capability not in _CAPABILITY_PROBES and target.capabilities.includes(capability=capability): - queried.add(capability) - - return queried - - -# --------------------------------------------------------------------------- -# Modality query -# --------------------------------------------------------------------------- - - -# Default mapping of non-text modalities to packaged probe assets. Callers can -# override via the ``test_assets`` parameter of -# :func:`discover_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] = { - "audio_path": str(_TARGET_CAPABILITIES_DATASET_PATH / "probe_audio.wav"), - "image_path": str(_TARGET_CAPABILITIES_DATASET_PATH / "probe_image.png"), -} - - -async def discover_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 which input modality combinations ``target`` accepts. - - Each modality combination is checked with a minimal request built from the - supplied test assets. - - 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 confirmed - to work against the target. - """ - if test_modalities is None: - declared = target.capabilities.input_modalities - test_modalities = set(declared) - elif not test_modalities: - logger.info("discover_target_modalities_async called with an empty test_modalities set; nothing to probe.") - return set() - - assets = test_assets if test_assets is not None else DEFAULT_TEST_ASSETS - - queried: 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: - # Skip combinations we cannot construct a valid probe payload for. - logger.info("Skipping modality %s: %s", combination, exc) - continue - except ValueError as exc: - logger.info("Skipping modality %s: %s", combination, exc) - continue - - # "Supported" means the request was accepted. A target may still - # ignore the non-text payload after accepting it. - if await _send_and_check_async( - target=target, - message=message, - timeout_s=per_probe_timeout_s, - retries=retries, - label=f"Modality probe {sorted(combination)}", - ): - queried.add(combination) - - return queried - - -async def discover_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, - apply: bool = False, -) -> TargetCapabilities: - """ - Probe capabilities and modalities and return a merged result. - - This wraps :func:`discover_target_capabilities_async` and - :func:`discover_target_modalities_async` and returns a best-effort - :class:`TargetCapabilities`. - - 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:`discover_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:`discover_target_modalities_async`. - capabilities (Iterable[CapabilityName] | None): Capabilities to probe. - See :func:`discover_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. - apply (bool): If True, install the discovered capabilities on ``target`` - via :meth:`PromptTarget.apply_capabilities` before returning. The - target's existing handling policy is preserved. Probe results are - an upper bound (the request was accepted, not necessarily honored), - so leave this False when you want to inspect or diff the result - before committing to it. Defaults to False. - - Returns: - TargetCapabilities: A merged capability view: probed where possible, - declared where probing is unavailable or out of scope. - """ - capabilities_to_probe = list(capabilities) if capabilities is not None else None - - queried_caps = await discover_target_capabilities_async( - target=target, - capabilities=capabilities_to_probe, - per_probe_timeout_s=per_probe_timeout_s, - retries=retries, - ) - queried_modalities = await discover_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 - # If the caller narrows the capability set, leave the rest at their - # declared values instead of silently forcing them to False. - probed: set[CapabilityName] = ( - set(capabilities_to_probe) if capabilities_to_probe is not None else set(CapabilityName) - ) - - def _resolve(name: CapabilityName) -> bool: - if name in probed: - return name in queried_caps - return bool(getattr(declared, name.value)) - - resolved_multi_turn = _resolve(CapabilityName.MULTI_TURN) - # Editable history is only meaningful if multi-turn probing/declaration - # also resolved to True. - resolved_editable_history = declared.supports_editable_history and resolved_multi_turn - if test_modalities is None: - # Mirror the boolean fallback: combinations the probe could not confirm - # fall back to the target's declared support rather than being silently - # dropped (e.g. on transient network failure). - resolved_input_modalities = frozenset(queried_modalities | declared.input_modalities) - else: - resolved_input_modalities = frozenset( - queried_modalities | (declared.input_modalities - frozenset(test_modalities)) - ) - - resolved = TargetCapabilities( - supports_multi_turn=resolved_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=resolved_editable_history, - supports_system_prompt=_resolve(CapabilityName.SYSTEM_PROMPT), - input_modalities=resolved_input_modalities, - # Output modalities are still declarative because probing them would - # require target-specific response inspection. - output_modalities=declared.output_modalities, - ) - - if apply: - target.apply_capabilities(capabilities=resolved) - - return resolved - - -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. - """ - 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(), - ) - ) - - return Message(pieces)