Skip to content

Commit 654c50c

Browse files
hannahwestra25hannahwestra25
authored
FEAT: Runtime capability discovery for prompt targets (#1699)
1 parent 79ee41a commit 654c50c

11 files changed

Lines changed: 2511 additions & 14 deletions

File tree

‎doc/code/targets/0_prompt_targets.md‎

Lines changed: 15 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -25,7 +25,7 @@ A `PromptTarget` is a generic place to send a prompt. With PyRIT, the idea is th
2525

2626
With some algorithms, you want to send a prompt, set a system prompt, and modify conversation history (including PAIR [@chao2023pair], TAP [@mehrotra2023tap], and flip attack [@li2024flipattack]). These algorithms require a target whose [`TargetCapabilities`](#target-capabilities) declare both `supports_multi_turn=True` and `supports_editable_history=True` — i.e. you can modify a conversation history. Consumers express this requirement via `CHAT_TARGET_REQUIREMENTS` and validate it against `target.configuration` at construction time. See [Target Capabilities](#target-capabilities) below for the full list of capabilities and how they compose into a `TargetConfiguration`.
2727

28-
Note: The previous `PromptChatTarget` class is **deprecated** as of v0.13.0 and will be removed in v0.15.0. Use `PromptTarget` directly with a `TargetConfiguration` declaring `supports_multi_turn=True` and `supports_editable_history=True`. See [Target Capabilities](#target-capabilities) for details.
28+
Note: The previous `PromptChatTarget` class is **deprecated** as of v0.14.0 and will be removed in v0.16.0. Use `PromptTarget` directly with a `TargetConfiguration` declaring `supports_multi_turn=True` and `supports_editable_history=True`. See [Target Capabilities](#target-capabilities) for details.
2929

3030

3131
Here are some examples:
@@ -107,6 +107,20 @@ target = MyHTTPTarget(custom_configuration=config, ...)
107107

108108
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).
109109

110+
### Discovering live target capabilities
111+
112+
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:
113+
114+
```python
115+
from pyrit.prompt_target import discover_target_capabilities_async
116+
117+
# Probe boolean capabilities and input modalities, returning a
118+
# best-effort TargetCapabilities:
119+
queried = await discover_target_capabilities_async(target=target)
120+
```
121+
122+
Each probe sends a minimal request (bounded by `per_probe_timeout_s`, default 30s, with one retry on transient errors) and only marks a capability or modality as supported if the call returns cleanly. `discover_target_capabilities_async` returns a merged view: probed where possible, declared where probing is unavailable or out of scope. "Supported" here means *the request was accepted* — a target that silently ignores a system prompt or `response_format` directive is still reported as supporting it, so validate response content out of band when the distinction matters. This function is not safe to call concurrently with other operations on the same target instance: it temporarily mutates `target._configuration` and writes probe rows to memory (rows are tagged with `prompt_metadata["capability_probe"] == "1"` for filtering). See [Target Capabilities](./6_1_target_capabilities.ipynb) for runnable examples.
123+
110124
## Multi-Modal Targets
111125

112126
Like most of PyRIT, targets can be multi-modal.

‎doc/code/targets/6_1_target_capabilities.ipynb‎

Lines changed: 288 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -53,13 +53,7 @@
5353
"name": "stdout",
5454
"output_type": "stream",
5555
"text": [
56-
"No new upgrade operations detected.\n"
57-
]
58-
},
59-
{
60-
"name": "stdout",
61-
"output_type": "stream",
62-
"text": [
56+
"No new upgrade operations detected.\n",
6357
"supports_multi_turn: True\n",
6458
"supports_editable_history: True\n",
6559
"supports_system_prompt: True\n",
@@ -462,12 +456,297 @@
462456
"try:\n",
463457
" no_editable_history.ensure_can_handle(capability=CapabilityName.EDITABLE_HISTORY)\n",
464458
"except ValueError as exc:\n",
465-
" print(exc)\n",
466-
"# ---"
459+
" print(exc)"
460+
]
461+
},
462+
{
463+
"cell_type": "markdown",
464+
"id": "19",
465+
"metadata": {},
466+
"source": [
467+
"## 7. Discovering live target capabilities\n",
468+
"\n",
469+
"Declared capabilities describe what a target *should* support. For deployments where the actual\n",
470+
"behavior is uncertain — custom OpenAI-compatible endpoints, gateways that strip features, models\n",
471+
"whose support drifts over time — you can probe what the target *actually* accepts at runtime with\n",
472+
"`discover_target_capabilities_async`. It runs both the boolean capability probes and the input\n",
473+
"modality probes and returns a best-effort `TargetCapabilities`.\n",
474+
"\n",
475+
"Internally it walks each capability that has a registered probe (currently\n",
476+
"`SYSTEM_PROMPT`, `MULTI_MESSAGE_PIECES`, `MULTI_TURN`, `JSON_OUTPUT`, `JSON_SCHEMA`), sends a\n",
477+
"minimal request, and includes the capability in the result only if the call succeeds.\n",
478+
"During probing the target's configuration is temporarily replaced with a permissive one so\n",
479+
"`ensure_can_handle` does not short-circuit a probe for a capability the target declares as\n",
480+
"unsupported. The original configuration is restored before the function returns. The same\n",
481+
"treatment is applied to each input modality combination declared in\n",
482+
"`capabilities.input_modalities`, sending a small payload built from optional `test_assets`.\n",
483+
"\n",
484+
"Each probe call is bounded by `per_probe_timeout_s` (default 30s) and is retried once on\n",
485+
"transient errors before being declared failed. The returned `TargetCapabilities` is a merged\n",
486+
"view: probed where possible, declared where probing is unavailable or out of scope.\n",
487+
"\"Supported\" here means *the request was accepted* — a target that silently ignores a system\n",
488+
"prompt or `response_format` directive will still be reported as supporting that capability.\n",
489+
"\n",
490+
"This function is **not safe to call concurrently** with other operations on the same target\n",
491+
"instance: it temporarily mutates `target._configuration` and writes probe rows to\n",
492+
"`target._memory`. Probe-written memory rows are tagged with\n",
493+
"`prompt_metadata[\"capability_probe\"] == \"1\"` so consumers can filter them.\n",
494+
"\n",
495+
"Typical usage against a real endpoint:\n",
496+
"\n",
497+
"```python\n",
498+
"from pyrit.prompt_target import discover_target_capabilities_async\n",
499+
"\n",
500+
"queried = await discover_target_capabilities_async(target=target)\n",
501+
"print(queried)\n",
502+
"```\n",
503+
"\n",
504+
"Below we mock the target's underlying transport (`_send_prompt_to_target_async`) so the notebook\n",
505+
"stays self-contained — the result shape is the same as a live run. We mock the protected method\n",
506+
"rather than `send_prompt_async` so the probe still exercises the real validation and memory\n",
507+
"pipeline."
508+
]
509+
},
510+
{
511+
"cell_type": "code",
512+
"execution_count": null,
513+
"id": "20",
514+
"metadata": {},
515+
"outputs": [
516+
{
517+
"name": "stdout",
518+
"output_type": "stream",
519+
"text": [
520+
"queried capabilities:\n",
521+
" - supports_editable_history\n",
522+
" - supports_json_output\n",
523+
" - supports_json_schema\n",
524+
" - supports_multi_message_pieces\n",
525+
" - supports_multi_turn\n",
526+
" - supports_system_prompt\n"
527+
]
528+
}
529+
],
530+
"source": [
531+
"from unittest.mock import AsyncMock\n",
532+
"\n",
533+
"from pyrit.models import MessagePiece\n",
534+
"from pyrit.prompt_target import discover_target_capabilities_async\n",
535+
"\n",
536+
"\n",
537+
"def _ok_response():\n",
538+
" return [\n",
539+
" Message(\n",
540+
" [\n",
541+
" MessagePiece(\n",
542+
" role=\"assistant\",\n",
543+
" original_value=\"ok\",\n",
544+
" original_value_data_type=\"text\",\n",
545+
" conversation_id=\"probe\",\n",
546+
" response_error=\"none\",\n",
547+
" )\n",
548+
" ]\n",
549+
" )\n",
550+
" ]\n",
551+
"\n",
552+
"\n",
553+
"probe_target = OpenAIChatTarget(model_name=\"gpt-4o\", endpoint=\"https://example.invalid/\", api_key=\"sk-not-a-real-key\")\n",
554+
"probe_target._send_prompt_to_target_async = AsyncMock(return_value=_ok_response()) # type: ignore[method-assign]\n",
555+
"\n",
556+
"queried = await discover_target_capabilities_async(target=probe_target, per_probe_timeout_s=5.0) # type: ignore\n",
557+
"print(\"discover_target_capabilities_async result:\")\n",
558+
"print(f\" supports_multi_turn: {queried.supports_multi_turn}\")\n",
559+
"print(f\" supports_system_prompt: {queried.supports_system_prompt}\")\n",
560+
"print(f\" supports_multi_message_pieces: {queried.supports_multi_message_pieces}\")\n",
561+
"print(f\" supports_json_output: {queried.supports_json_output}\")\n",
562+
"print(f\" supports_json_schema: {queried.supports_json_schema}\")\n",
563+
"print(f\" input_modalities: {sorted(sorted(m) for m in queried.input_modalities)}\")"
564+
]
565+
},
566+
{
567+
"cell_type": "markdown",
568+
"id": "21",
569+
"metadata": {},
570+
"source": [
571+
"To narrow the probe to specific capabilities (faster, fewer calls), pass `capabilities=`:\n",
572+
"\n",
573+
"```python\n",
574+
"from pyrit.prompt_target.common.target_capabilities import CapabilityName\n",
575+
"\n",
576+
"queried = await discover_target_capabilities_async(\n",
577+
" target=target,\n",
578+
" capabilities=[CapabilityName.JSON_SCHEMA, CapabilityName.SYSTEM_PROMPT],\n",
579+
")\n",
580+
"```\n",
581+
"\n",
582+
"Similarly, narrow the modality probe set with `test_modalities=` and override the\n",
583+
"packaged default probe assets with `test_assets=`."
584+
]
585+
},
586+
{
587+
"cell_type": "markdown",
588+
"id": "22",
589+
"metadata": {},
590+
"source": [
591+
"### Discovering undeclared modalities\n",
592+
"\n",
593+
"By default `discover_target_capabilities_async` only probes modality combinations the target already\n",
594+
"**declares** in `capabilities.input_modalities`. For an OpenAI-compatible endpoint that\n",
595+
"claims text-only but might actually accept images, pass `test_modalities=` explicitly to\n",
596+
"probe combinations beyond the declared baseline. Provide `test_assets=` as well if you need\n",
597+
"to override the packaged defaults or probe a modality without one:\n",
598+
"\n",
599+
"```python\n",
600+
"queried = await discover_target_capabilities_async(\n",
601+
" target=target,\n",
602+
" test_modalities={frozenset({\"text\"}), frozenset({\"text\", \"image_path\"})},\n",
603+
" test_assets={\"image_path\": \"/path/to/test_image.png\"},\n",
604+
")\n",
605+
"```\n",
606+
"\n",
607+
"Similarly, when narrowing the probe set with `capabilities=`, capabilities NOT in the\n",
608+
"narrowed set are copied from the target's declared values rather than being reset to\n",
609+
"`False` — narrowing controls *what is re-queried*, not what the returned dataclass\n",
610+
"reports. This makes incremental probing safe:\n",
611+
"\n",
612+
"```python\n",
613+
"# Re-query only JSON support; other declared flags pass through unchanged.\n",
614+
"queried = await discover_target_capabilities_async(\n",
615+
" target=target,\n",
616+
" capabilities={CapabilityName.JSON_OUTPUT, CapabilityName.JSON_SCHEMA},\n",
617+
")\n",
618+
"```"
619+
]
620+
},
621+
{
622+
"cell_type": "markdown",
623+
"id": "23",
624+
"metadata": {},
625+
"source": [
626+
"## 8. Applying probed capabilities back onto the target\n",
627+
"\n",
628+
"`discover_target_capabilities_async` is intentionally pure: it returns a `TargetCapabilities` without\n",
629+
"mutating the target. That lets you inspect (or diff against the declared view, log, gate on\n",
630+
"the result) before committing. Once you're satisfied, call `target.apply_capabilities(...)`\n",
631+
"to install the probed view on the instance. The target's existing\n",
632+
"`CapabilityHandlingPolicy` is preserved — policy expresses user intent (ADAPT vs RAISE),\n",
633+
"which is independent of what the probe found.\n",
634+
"\n",
635+
"Why a two-step pattern rather than auto-apply? Probe results are an upper bound\n",
636+
"(\"the request was accepted\"); a target that silently ignores a feature still passes its\n",
637+
"probe. Keeping discovery separate from application lets callers diff, log, persist, or\n",
638+
"reject the result before it affects subsequent sends.\n",
639+
"\n",
640+
"Below is the end-to-end pattern: construct a target whose declared capabilities are\n",
641+
"pessimistic, discover what the endpoint actually accepts, diff the two views, then apply."
642+
]
643+
},
644+
{
645+
"cell_type": "code",
646+
"execution_count": null,
647+
"id": "24",
648+
"metadata": {},
649+
"outputs": [
650+
{
651+
"name": "stdout",
652+
"output_type": "stream",
653+
"text": [
654+
"declared (before probing):\n",
655+
" supports_multi_turn: False\n",
656+
" supports_system_prompt: False\n",
657+
" supports_json_output: False\n",
658+
"\n",
659+
"probed (returned from discover_target_capabilities_async, target NOT yet updated):\n",
660+
" supports_multi_turn: True\n",
661+
" supports_system_prompt: True\n",
662+
" supports_json_output: True\n",
663+
" target.capabilities.supports_multi_turn (still declared): False\n",
664+
"\n",
665+
"flags probed True that were declared False: ['supports_multi_turn', 'supports_system_prompt', 'supports_multi_message_pieces', 'supports_json_output', 'supports_json_schema']\n",
666+
"\n",
667+
"after apply_capabilities:\n",
668+
" supports_multi_turn: True\n",
669+
" supports_system_prompt: True\n",
670+
" supports_json_output: True\n",
671+
" policy preserved: True\n",
672+
"\n",
673+
"CHAT_TARGET_REQUIREMENTS.validate now passes against the probed target\n"
674+
]
675+
}
676+
],
677+
"source": [
678+
"# Start with an instance that declares fewer capabilities than the endpoint actually has,\n",
679+
"# e.g. a custom gateway whose support we're unsure about.\n",
680+
"pessimistic_config = TargetConfiguration(\n",
681+
" capabilities=TargetCapabilities(\n",
682+
" supports_multi_turn=False,\n",
683+
" supports_system_prompt=False,\n",
684+
" supports_multi_message_pieces=False,\n",
685+
" supports_json_output=False,\n",
686+
" supports_json_schema=False,\n",
687+
" # Editable history has no live probe and falls back to the declared value.\n",
688+
" # Declare it True here so the probed view inherits it.\n",
689+
" supports_editable_history=True,\n",
690+
" ),\n",
691+
")\n",
692+
"endpoint_target = OpenAIChatTarget(\n",
693+
" model_name=\"custom-model\",\n",
694+
" endpoint=\"https://example.invalid/\",\n",
695+
" api_key=\"sk-not-a-real-key\",\n",
696+
" custom_configuration=pessimistic_config,\n",
697+
")\n",
698+
"endpoint_target._send_prompt_to_target_async = AsyncMock(return_value=_ok_response()) # type: ignore[method-assign]\n",
699+
"\n",
700+
"print(\"declared (before probing):\")\n",
701+
"print(f\" supports_multi_turn: {endpoint_target.capabilities.supports_multi_turn}\")\n",
702+
"print(f\" supports_system_prompt: {endpoint_target.capabilities.supports_system_prompt}\")\n",
703+
"print(f\" supports_json_output: {endpoint_target.capabilities.supports_json_output}\")\n",
704+
"\n",
705+
"# Step 1: discover. No mutation yet — `endpoint_target.capabilities` is unchanged.\n",
706+
"probed_caps = await discover_target_capabilities_async(target=endpoint_target, per_probe_timeout_s=5.0) # type: ignore\n",
707+
"\n",
708+
"print(\"\\nprobed (returned from discover_target_capabilities_async, target NOT yet updated):\")\n",
709+
"print(f\" supports_multi_turn: {probed_caps.supports_multi_turn}\")\n",
710+
"print(f\" supports_system_prompt: {probed_caps.supports_system_prompt}\")\n",
711+
"print(f\" supports_json_output: {probed_caps.supports_json_output}\")\n",
712+
"print(f\" target.capabilities.supports_multi_turn (still declared): {endpoint_target.capabilities.supports_multi_turn}\")\n",
713+
"\n",
714+
"# Step 2: diff — see exactly what the probe upgraded.\n",
715+
"declared = pessimistic_config.capabilities\n",
716+
"upgraded = [\n",
717+
" name\n",
718+
" for name in (\n",
719+
" \"supports_multi_turn\",\n",
720+
" \"supports_system_prompt\",\n",
721+
" \"supports_multi_message_pieces\",\n",
722+
" \"supports_json_output\",\n",
723+
" \"supports_json_schema\",\n",
724+
" )\n",
725+
" if getattr(probed_caps, name) and not getattr(declared, name)\n",
726+
"]\n",
727+
"print(f\"\\nflags probed True that were declared False: {upgraded}\")\n",
728+
"\n",
729+
"# Step 3: apply. Policy is preserved; the normalization pipeline is rebuilt.\n",
730+
"original_policy = endpoint_target.configuration.policy\n",
731+
"endpoint_target.apply_capabilities(capabilities=probed_caps)\n",
732+
"\n",
733+
"print(\"\\nafter apply_capabilities:\")\n",
734+
"print(f\" supports_multi_turn: {endpoint_target.capabilities.supports_multi_turn}\")\n",
735+
"print(f\" supports_system_prompt: {endpoint_target.capabilities.supports_system_prompt}\")\n",
736+
"print(f\" supports_json_output: {endpoint_target.capabilities.supports_json_output}\")\n",
737+
"print(f\" policy preserved: {endpoint_target.configuration.policy is original_policy}\")\n",
738+
"\n",
739+
"# Subsequent consumer checks now reflect the probed reality — for example, a chat-style\n",
740+
"# requirement that would have failed against the pessimistic declaration now passes.\n",
741+
"CHAT_TARGET_REQUIREMENTS.validate(target=endpoint_target)\n",
742+
"print(\"\\nCHAT_TARGET_REQUIREMENTS.validate now passes against the probed target\")"
467743
]
468744
}
469745
],
470746
"metadata": {
747+
"jupytext": {
748+
"main_language": "python"
749+
},
471750
"language_info": {
472751
"codemirror_mode": {
473752
"name": "ipython",

0 commit comments

Comments
 (0)