From e2da267941aae89cc208aa853f61412f156e233a Mon Sep 17 00:00:00 2001 From: hannahwestra25 Date: Thu, 7 May 2026 13:22:18 -0400 Subject: [PATCH 01/12] 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/12] 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/12] 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/12] 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/12] 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/12] 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/12] 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/12] 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/12] 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/12] 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/12] 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/12] 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)", )