diff --git a/doc/code/scenarios/3_adaptive_scenarios.ipynb b/doc/code/scenarios/3_adaptive_scenarios.ipynb new file mode 100644 index 0000000000..026c4e083f --- /dev/null +++ b/doc/code/scenarios/3_adaptive_scenarios.ipynb @@ -0,0 +1,722 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "id": "0", + "metadata": {}, + "source": [ + "# Adaptive Scenarios\n", + "\n", + "An **adaptive scenario** doesn't run every attack technique against every objective.\n", + "Instead, it picks which technique to try next per-objective, learns from what worked,\n", + "and stops as soon as one technique succeeds. This concentrates spend on techniques\n", + "that actually work on your target.\n", + "\n", + "## How it works (high level)\n", + "\n", + "For each objective, the scenario tries up to `max_attempts_per_objective` techniques:\n", + "\n", + "- With probability `epsilon`, it **explores** — picks a random technique.\n", + "- Otherwise it **exploits** — picks the technique with the highest observed success\n", + " rate so far.\n", + "- It records the outcome and stops early on success.\n", + "\n", + "Unseen techniques are tried first, so the first few objectives effectively round-robin\n", + "through every technique before the scenario settles on the best performers.\n", + "\n", + "## Adaptive vs. static scenarios\n", + "\n", + "| Feature | Static scenarios | Adaptive scenarios |\n", + "|---------------------|-----------------------------------|------------------------------------|\n", + "| Technique selection | Run every selected technique | Pick per-objective from outcomes |\n", + "| Early stopping | No | Yes — stops on first success |\n", + "| Cost | O(techniques × objectives) | O(max_attempts × objectives) |\n", + "\n", + "`AdaptiveScenario` is the modality-agnostic base class.\n", + "[`TextAdaptive`](../../../pyrit/scenario/scenarios/adaptive/text_adaptive.py) is the\n", + "text subclass used in the examples below." + ] + }, + { + "cell_type": "markdown", + "id": "1", + "metadata": {}, + "source": [ + "## Setup" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "2", + "metadata": {}, + "outputs": [ + { + "name": "stderr", + "output_type": "stream", + "text": [ + "/tmp/ipykernel_49587/458917033.py:5: DeprecationWarning: pyrit.scenario.printer.console_printer.ConsoleScenarioResultPrinter is deprecated and will be removed in 0.16.0. Use pyrit.output.scenario_result.pretty.PrettyScenarioResultMemoryPrinter instead.\n", + " from pyrit.scenario.printer.console_printer import ConsoleScenarioResultPrinter\n" + ] + }, + { + "name": "stdout", + "output_type": "stream", + "text": [ + "Found default environment files: ['./.pyrit/.env']\n", + "Loaded environment file: ./.pyrit/.env\n", + "No new upgrade operations detected.\n" + ] + }, + { + "name": "stderr", + "output_type": "stream", + "text": [ + "Skipping target 'platform_openai_chat': PLATFORM_OPENAI_CHAT_GPT4O_MODEL is not set. All declared env vars (endpoint, key, model) must be present for this target to register.\n", + "Skipping target 'azure_foundry_phi4': AZURE_FOUNDRY_PHI4_MODEL is not set. All declared env vars (endpoint, key, model) must be present for this target to register.\n", + "Skipping scorer main: required target not found in TargetRegistry\n", + "TextAdaptive: _EXCLUDED_TECHNIQUES entries ['prompt_sending'] are not in the current scenario-techniques catalog ['context_compliance', 'crescendo_history_lecture', 'crescendo_journalist_interview', 'crescendo_movie_director', 'crescendo_simulated', 'many_shot', 'pair', 'red_teaming', 'role_play', 'tap']; the exclusion is a no-op for those entries. Remove stale entries or update the catalog.\n", + "TargetRegistry entry 'adversarial_chat' not found. Falling back to default OpenAIChatTarget.\n", + "TargetRegistry entry 'objective_scorer_chat' not found. Falling back to default OpenAIChatTarget.\n", + "TargetRegistry entry 'adversarial_chat' not found. Falling back to default OpenAIChatTarget.\n", + "TargetRegistry entry 'adversarial_chat' not found. Falling back to default OpenAIChatTarget.\n" + ] + } + ], + "source": [ + "from pathlib import Path\n", + "\n", + "from pyrit.registry import TargetRegistry\n", + "from pyrit.scenario import DatasetConfiguration\n", + "from pyrit.scenario.printer.console_printer import ConsoleScenarioResultPrinter\n", + "from pyrit.scenario.scenarios.adaptive import TextAdaptive\n", + "from pyrit.setup import initialize_from_config_async\n", + "\n", + "await initialize_from_config_async(config_path=Path(\"../../scanner/pyrit_conf.yaml\")) # type: ignore\n", + "\n", + "objective_target = TargetRegistry.get_registry_singleton().get_instance_by_name(\"openai_chat\")\n", + "printer = ConsoleScenarioResultPrinter()" + ] + }, + { + "cell_type": "markdown", + "id": "3", + "metadata": {}, + "source": [ + "## Basic usage\n", + "\n", + "Defaults: `max_attempts_per_objective=3`, epsilon-greedy selector with `epsilon=0.2`,\n", + "the subclass's default datasets." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "4", + "metadata": {}, + "outputs": [ + { + "name": "stderr", + "output_type": "stream", + "text": [ + "TargetRegistry entry 'adversarial_chat' not found. Falling back to default OpenAIChatTarget.\n", + "TargetRegistry entry 'adversarial_chat' not found. Falling back to default OpenAIChatTarget.\n", + "/tmp/ipykernel_49587/3812235746.py:3: DeprecationWarning: Implicit baseline injection for TextAdaptive._get_atomic_attacks_async() is deprecated and will be removed in 0.16.0. Use explicit emission via self._build_baseline_atomic_attack(seed_groups=...) in the override instead.\n", + " await scenario.initialize_async( # type: ignore\n" + ] + }, + { + "data": { + "application/vnd.jupyter.widget-view+json": { + "model_id": "58a5ce9c44f24d3ab737073e6ecd2430", + "version_major": 2, + "version_minor": 0 + }, + "text/plain": [ + "Executing TextAdaptive: 0%| | 0/22 [00:00 str:\n", + " \"\"\"Display name for the attack strategy that produced ``result``.\"\"\"\n", + " attack_id = result.get_attack_strategy_identifier()\n", + " return attack_id.class_name if attack_id else \"\"\n", + "\n", + "\n", + "total_picks: Counter[str] = Counter()\n", + "total_wins: Counter[str] = Counter()\n", + "\n", + "for group_name, results in display_groups.items():\n", + " print(f\"\\n=== Group: {group_name} ===\")\n", + "\n", + " # Collect every child id referenced by any envelope in this group so we\n", + " # can skip the per-attempt child rows when printing per-objective lines.\n", + " # Baseline rows have no envelope and pass through untouched.\n", + " child_ids: set[str] = set()\n", + " for r in results:\n", + " child_ids.update(r.metadata.get(\"child_attack_result_ids\", []) or [])\n", + "\n", + " for r in results:\n", + " if r.attack_result_id in child_ids:\n", + " continue\n", + " child_id_list = r.metadata.get(\"child_attack_result_ids\", []) or []\n", + " trail_parts: list[str] = []\n", + " for child_id in child_id_list:\n", + " child = results_by_id.get(child_id)\n", + " if child is None:\n", + " continue\n", + " trail_parts.append(f\"{_technique_label(child)}({child.outcome.value})\")\n", + " trail = \" → \".join(trail_parts)\n", + " print(f\" [{r.outcome.value:7s}] {r.objective!r}: {trail}\")\n", + "\n", + " picks: Counter[str] = Counter()\n", + " wins: Counter[str] = Counter()\n", + " for r in results:\n", + " if r.attack_result_id not in child_ids:\n", + " continue\n", + " technique = _technique_label(r)\n", + " picks[technique] += 1\n", + " total_picks[technique] += 1\n", + " if r.outcome.value == \"success\":\n", + " wins[technique] += 1\n", + " total_wins[technique] += 1\n", + "\n", + " print(\"\\n Technique wins / picks rate\")\n", + " for technique, n in picks.most_common():\n", + " print(f\" {technique:40s} {wins[technique]:>4} / {n:<4} {wins[technique] / n:.0%}\")\n", + "\n", + "print(\"\\n=== Overall ===\")\n", + "print(\"Technique wins / picks rate\")\n", + "for technique, n in total_picks.most_common():\n", + " print(f\"{technique:40s} {total_wins[technique]:>4} / {n:<4} {total_wins[technique] / n:.0%}\")" + ] + }, + { + "cell_type": "markdown", + "id": "11", + "metadata": {}, + "source": [ + "## Running from the scanner CLI\n", + "\n", + "You can run `TextAdaptive` directly from the `pyrit_scan` CLI without writing Python:\n", + "\n", + "```bash\n", + "# Basic run with defaults\n", + "pyrit_scan --scenario TextAdaptive --target openai_chat\n", + "\n", + "# Tune max attempts and restrict strategies\n", + "pyrit_scan --scenario TextAdaptive --target openai_chat \\\n", + " --params max_attempts_per_objective=5 \\\n", + " --strategies single_turn\n", + "\n", + "# Use specific datasets and limit size\n", + "pyrit_scan --scenario TextAdaptive --target openai_chat \\\n", + " --datasets airt_hate airt_violence \\\n", + " --max-dataset-size 10\n", + "```" + ] + } + ], + "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, + "nbformat_minor": 5 +} diff --git a/doc/code/scenarios/3_adaptive_scenarios.py b/doc/code/scenarios/3_adaptive_scenarios.py new file mode 100644 index 0000000000..6239a24077 --- /dev/null +++ b/doc/code/scenarios/3_adaptive_scenarios.py @@ -0,0 +1,250 @@ +# --- +# jupyter: +# jupytext: +# text_representation: +# extension: .py +# format_name: percent +# format_version: '1.3' +# jupytext_version: 1.18.1 +# --- + +# %% [markdown] +# # Adaptive Scenarios +# +# An **adaptive scenario** doesn't run every attack technique against every objective. +# Instead, it picks which technique to try next per-objective, learns from what worked, +# and stops as soon as one technique succeeds. This concentrates spend on techniques +# that actually work on your target. +# +# ## How it works (high level) +# +# For each objective, the scenario tries up to `max_attempts_per_objective` techniques: +# +# - With probability `epsilon`, it **explores** — picks a random technique. +# - Otherwise it **exploits** — picks the technique with the highest observed success +# rate so far. +# - It records the outcome and stops early on success. +# +# Unseen techniques are tried first, so the first few objectives effectively round-robin +# through every technique before the scenario settles on the best performers. +# +# ## Adaptive vs. static scenarios +# +# | Feature | Static scenarios | Adaptive scenarios | +# |---------------------|-----------------------------------|------------------------------------| +# | Technique selection | Run every selected technique | Pick per-objective from outcomes | +# | Early stopping | No | Yes — stops on first success | +# | Cost | O(techniques × objectives) | O(max_attempts × objectives) | +# +# `AdaptiveScenario` is the modality-agnostic base class. +# [`TextAdaptive`](../../../pyrit/scenario/scenarios/adaptive/text_adaptive.py) is the +# text subclass used in the examples below. + +# %% [markdown] +# ## Setup + +# %% +from pathlib import Path + +from pyrit.registry import TargetRegistry +from pyrit.scenario import DatasetConfiguration +from pyrit.scenario.printer.console_printer import ConsoleScenarioResultPrinter +from pyrit.scenario.scenarios.adaptive import TextAdaptive +from pyrit.setup import initialize_from_config_async + +await initialize_from_config_async(config_path=Path("../../scanner/pyrit_conf.yaml")) # type: ignore + +objective_target = TargetRegistry.get_registry_singleton().get_instance_by_name("openai_chat") +printer = ConsoleScenarioResultPrinter() + +# %% [markdown] +# ## Basic usage +# +# Defaults: `max_attempts_per_objective=3`, epsilon-greedy selector with `epsilon=0.2`, +# the subclass's default datasets. + +# %% +scenario = TextAdaptive() + +await scenario.initialize_async( # type: ignore + objective_target=objective_target, +) +result = await scenario.run_async() # type: ignore +await printer.write_async(result) # type: ignore + +# %% [markdown] +# ## Configuring a run +# +# - **`max_attempts_per_objective`** — caps techniques tried per objective. Higher means +# more chances to succeed and more API calls. Set via `set_params_from_args`. +# - **`selector`** — a pre-built `TechniqueSelector` instance. Pass an +# `EpsilonGreedyTechniqueSelector(epsilon=..., random_seed=...)` +# to tune the selection algorithm. Defaults to an epsilon-greedy selector with +# `epsilon=0.2`. +# - **`scenario_strategies`** (on `initialize_async`) — restricts which techniques the +# selector can pick from. Use `TextAdaptive.get_strategy_class()` to access the enum. +# +# The cell below exercises all of them at once. + +# %% +from pyrit.scenario.scenarios.adaptive import EpsilonGreedyTechniqueSelector + +strategy_class = TextAdaptive.get_strategy_class() + +configured_scenario = TextAdaptive( + selector=EpsilonGreedyTechniqueSelector( + epsilon=0.3, + random_seed=42, + ), +) +configured_scenario.set_params_from_args(args={"max_attempts_per_objective": 5}) + +await configured_scenario.initialize_async( # type: ignore + objective_target=objective_target, + scenario_strategies=[strategy_class("single_turn")], + dataset_config=DatasetConfiguration( + dataset_names=["airt_hate", "airt_violence"], + max_dataset_size=4, + ), +) +configured_result = await configured_scenario.run_async() # type: ignore +await printer.write_async(configured_result) # type: ignore + +# %% [markdown] +# ## Resuming a run +# +# Adaptive scenarios are resumable — pass `scenario_result_id=...` to the `TextAdaptive` +# constructor and the run picks up where it left off. Resume must use the same +# configuration as the original run. + +# %% +resumed_scenario = TextAdaptive( + selector=EpsilonGreedyTechniqueSelector( + epsilon=0.3, + random_seed=42, + ), + scenario_result_id=str(configured_result.id), +) +resumed_scenario.set_params_from_args(args={"max_attempts_per_objective": 5}) + +await resumed_scenario.initialize_async( # type: ignore + objective_target=objective_target, + scenario_strategies=[strategy_class("single_turn")], + dataset_config=DatasetConfiguration( + dataset_names=["airt_hate", "airt_violence"], + max_dataset_size=4, + ), +) +resumed_result = await resumed_scenario.run_async() # type: ignore +await printer.write_async(resumed_result) # type: ignore + +# %% [markdown] +# ## Inspecting which techniques were tried +# +# Every adaptive run persists both the per-objective envelope (a +# `SequentialAttackResult`) AND its per-attempt child rows. Each child row +# carries its own `atomic_attack_identifier`, so the persisted data alone is +# enough to reconstruct the per-attempt trail — no envelope-side metadata, no +# scenario-side lookup tables needed. +# +# Walk the children via the envelope's `child_attack_result_ids` (joined +# against the flat results list), then read each child's attack strategy +# identifier with `child.get_attack_strategy_identifier()`. The returned +# `ComponentIdentifier` exposes `class_name` (e.g. `"CrescendoAttack"`) for a +# human-readable label, and `unique_name` (e.g. `"CrescendoAttack::a1b2c3d4"`) +# when you need to distinguish two factories that wrap the same attack class +# with different configurations. +# +# Use `result.get_display_groups()` to aggregate `attack_results` by the +# per-dataset display label set by the scenario. +# +# If the trail of attacks attempted is shorter than `max_attempts_per_objective`, +# the compatible-technique pool for that seed group was smaller than the cap — +# the run exhausted the pool. + +# %% +from collections import Counter + +# Per-group: one line per objective (the envelope) showing the per-attempt +# trail, plus a per-technique success-rate table within the group. The child +# rows that compose each envelope are filtered out of the per-objective list so +# it stays one line per objective. Aggregate across groups for a grand-total. +display_groups = resumed_result.get_display_groups() + +# Flatten every persisted row across every group so we can look up a child +# AttackResult by its attack_result_id when reconstructing per-envelope trails. +results_by_id = {r.attack_result_id: r for results in display_groups.values() for r in results} + + +def _technique_label(result) -> str: + """Display name for the attack strategy that produced ``result``.""" + attack_id = result.get_attack_strategy_identifier() + return attack_id.class_name if attack_id else "" + + +total_picks: Counter[str] = Counter() +total_wins: Counter[str] = Counter() + +for group_name, results in display_groups.items(): + print(f"\n=== Group: {group_name} ===") + + # Collect every child id referenced by any envelope in this group so we + # can skip the per-attempt child rows when printing per-objective lines. + # Baseline rows have no envelope and pass through untouched. + child_ids: set[str] = set() + for r in results: + child_ids.update(r.metadata.get("child_attack_result_ids", []) or []) + + for r in results: + if r.attack_result_id in child_ids: + continue + child_id_list = r.metadata.get("child_attack_result_ids", []) or [] + trail_parts: list[str] = [] + for child_id in child_id_list: + child = results_by_id.get(child_id) + if child is None: + continue + trail_parts.append(f"{_technique_label(child)}({child.outcome.value})") + trail = " → ".join(trail_parts) + print(f" [{r.outcome.value:7s}] {r.objective!r}: {trail}") + + picks: Counter[str] = Counter() + wins: Counter[str] = Counter() + for r in results: + if r.attack_result_id not in child_ids: + continue + technique = _technique_label(r) + picks[technique] += 1 + total_picks[technique] += 1 + if r.outcome.value == "success": + wins[technique] += 1 + total_wins[technique] += 1 + + print("\n Technique wins / picks rate") + for technique, n in picks.most_common(): + print(f" {technique:40s} {wins[technique]:>4} / {n:<4} {wins[technique] / n:.0%}") + +print("\n=== Overall ===") +print("Technique wins / picks rate") +for technique, n in total_picks.most_common(): + print(f"{technique:40s} {total_wins[technique]:>4} / {n:<4} {total_wins[technique] / n:.0%}") + +# %% [markdown] +# ## Running from the scanner CLI +# +# You can run `TextAdaptive` directly from the `pyrit_scan` CLI without writing Python: +# +# ```bash +# # Basic run with defaults +# pyrit_scan --scenario TextAdaptive --target openai_chat +# +# # Tune max attempts and restrict strategies +# pyrit_scan --scenario TextAdaptive --target openai_chat \ +# --params max_attempts_per_objective=5 \ +# --strategies single_turn +# +# # Use specific datasets and limit size +# pyrit_scan --scenario TextAdaptive --target openai_chat \ +# --datasets airt_hate airt_violence \ +# --max-dataset-size 10 +# ``` diff --git a/doc/myst.yml b/doc/myst.yml index a8d5780298..eb879f134f 100644 --- a/doc/myst.yml +++ b/doc/myst.yml @@ -199,6 +199,7 @@ project: children: - file: code/scenarios/1_common_scenario_parameters.ipynb - file: code/scenarios/2_custom_scenario_parameters.ipynb + - file: code/scenarios/3_adaptive_scenarios.ipynb - file: code/registry/0_registry.md children: - file: code/registry/1_class_registry.ipynb diff --git a/pyrit/analytics/technique_analysis.py b/pyrit/analytics/technique_analysis.py new file mode 100644 index 0000000000..a091d29821 --- /dev/null +++ b/pyrit/analytics/technique_analysis.py @@ -0,0 +1,84 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT license. + +"""Scenario-level analytics: technique success rates and related helpers.""" + +from __future__ import annotations + +from typing import TYPE_CHECKING + +from pyrit.analytics.result_analysis import AttackStats, _compute_stats +from pyrit.memory import CentralMemory +from pyrit.models import AttackOutcome + +if TYPE_CHECKING: + from collections.abc import Sequence + + from pyrit.memory.memory_interface import MemoryInterface + + +def compute_technique_stats( + *, + technique_eval_hashes: Sequence[str], + scenario_result_id: str | None = None, + targeted_harm_categories: Sequence[str] | None = None, + memory: MemoryInterface | None = None, +) -> dict[str, AttackStats]: + """ + Compute per-technique outcome statistics from persisted attack results. + + Queries memory for all ``AttackResult`` rows whose + ``atomic_attack_identifier.eval_hash`` matches one of + ``technique_eval_hashes``, then aggregates outcomes into per-technique + ``AttackStats``. The eval hash is auto-stamped on every persisted result + by ``AtomicAttackEvaluationIdentifier`` and is the canonical primitive + for behavioral-equivalence aggregation (seeds excluded, scorer excluded, + only behavior-relevant target params included). + + Args: + technique_eval_hashes (Sequence[str]): Eval hashes to aggregate. + Returned dict is keyed by these. + scenario_result_id (str | None): Restrict to a single scenario run. + Defaults to ``None`` (aggregate across all runs). + targeted_harm_categories (Sequence[str] | None): Restrict to results + whose prompts targeted these harm categories. Defaults to ``None``. + memory (MemoryInterface | None): Memory backend to query. Defaults to + ``CentralMemory.get_memory_instance()``. + + Returns: + dict[str, AttackStats]: Stats per technique eval hash. Hashes with no + historical results are omitted from the result. + """ + if not technique_eval_hashes: + return {} + + if memory is None: + memory = CentralMemory.get_memory_instance() + results = memory.get_attack_results( + atomic_attack_eval_hashes=list(technique_eval_hashes), + scenario_result_id=scenario_result_id, + targeted_harm_categories=targeted_harm_categories, + ) + + requested = set(technique_eval_hashes) + counts: dict[str, tuple[int, int, int, int]] = {} + for result in results: + identifier = result.atomic_attack_identifier + eval_hash = identifier.eval_hash if identifier is not None else None + if eval_hash is None or eval_hash not in requested: + continue + + s, f, u, e = counts.get(eval_hash, (0, 0, 0, 0)) + if result.outcome == AttackOutcome.SUCCESS: + counts[eval_hash] = (s + 1, f, u, e) + elif result.outcome == AttackOutcome.FAILURE: + counts[eval_hash] = (s, f + 1, u, e) + elif result.outcome == AttackOutcome.ERROR: + counts[eval_hash] = (s, f, u, e + 1) + else: + counts[eval_hash] = (s, f, u + 1, e) + + return { + eval_hash: _compute_stats(successes=s, failures=f, undetermined=u, errors=e) + for eval_hash, (s, f, u, e) in counts.items() + } diff --git a/pyrit/cli/api_client.py b/pyrit/cli/api_client.py index 103653ca46..bfd75ca420 100644 --- a/pyrit/cli/api_client.py +++ b/pyrit/cli/api_client.py @@ -33,14 +33,20 @@ class PyRITApiClient: scenarios = await client.list_scenarios_async() """ - def __init__(self, *, base_url: str) -> None: + def __init__(self, *, base_url: str, request_timeout: float | None = None) -> None: """ Initialize the API client. Args: base_url (str): Base URL of the PyRIT backend (e.g., ``"http://localhost:8000"``). + request_timeout (float | None): Read timeout in seconds applied to every + non-polling request (catalog, results, cancel, start, etc.). Polling + the live scenario-run endpoint always uses ``read=None`` regardless + of this value, because the server may legitimately take many seconds + to respond while a scenario is executing. Defaults to ``60.0``. """ self._base_url = base_url.rstrip("/") + self._request_timeout = request_timeout if request_timeout is not None else 60.0 self._client: Any = None # httpx.AsyncClient (typed Any to avoid top-level import) async def __aenter__(self) -> PyRITApiClient: @@ -52,7 +58,7 @@ async def __aenter__(self) -> PyRITApiClient: """ import httpx - self._client = httpx.AsyncClient(base_url=self._base_url, timeout=60.0) + self._client = httpx.AsyncClient(base_url=self._base_url, timeout=self._request_timeout) return self async def __aexit__(self, exc_type: Any, exc_val: Any, exc_tb: Any) -> None: @@ -149,7 +155,7 @@ async def register_initializer_async(self, *, name: str, script_content: str) -> if resp.status_code == 403: detail = resp.json().get("detail", "Custom initializer operations are disabled on the server.") raise ServerNotAvailableError(detail) - resp.raise_for_status() + self._raise_for_status(resp) return resp.json() # ------------------------------------------------------------------ @@ -181,17 +187,41 @@ async def start_scenario_run_async(self, *, request: dict[str, Any]) -> dict[str """ client = self._get_client() resp = await client.post("/api/scenarios/runs", json=request) - resp.raise_for_status() + self._raise_for_status(resp) return resp.json() async def get_scenario_run_async(self, *, scenario_result_id: str) -> dict[str, Any]: """ Get the current status of a scenario run. + This is the endpoint the CLI polls while waiting for a run to finish. + It uses ``read=None`` (wait indefinitely for a response) so a server + busy executing a long-running scenario doesn't trip the client's + default read timeout. The other endpoints keep the configured timeout. + Returns: dict: ``ScenarioRunSummary`` payload. + + Raises: + ServerNotAvailableError: If the server cannot be reached. """ - return await self._get_json_async(path=f"/api/scenarios/runs/{scenario_result_id}") + import httpx + + client = self._get_client() + try: + resp = await client.get( + f"/api/scenarios/runs/{scenario_result_id}", + params=None, + timeout=httpx.Timeout(connect=10.0, read=None, write=30.0, pool=10.0), + ) + except httpx.ConnectError as exc: + raise ServerNotAvailableError( + f"Cannot connect to PyRIT server at {self._base_url}.\n" + "Hint: Use '--start-server' to launch a local backend, " + "or pass '--server-url '." + ) from exc + self._raise_for_status(resp) + return resp.json() async def get_scenario_run_results_async(self, *, scenario_result_id: str) -> dict[str, Any]: """ @@ -211,7 +241,7 @@ async def cancel_scenario_run_async(self, *, scenario_result_id: str) -> dict[st """ client = self._get_client() resp = await client.post(f"/api/scenarios/runs/{scenario_result_id}/cancel") - resp.raise_for_status() + self._raise_for_status(resp) return resp.json() async def list_scenario_runs_async(self, *, limit: int = 100) -> dict[str, Any]: @@ -275,5 +305,45 @@ async def _get_json_async(self, *, path: str, params: dict[str, Any] | None = No "Hint: Use '--start-server' to launch a local backend, " "or pass '--server-url '." ) from exc - resp.raise_for_status() + self._raise_for_status(resp) return resp.json() + + @staticmethod + def _raise_for_status(resp: Any) -> None: + """ + Raise an HTTP error with the response body appended to the message. + + Behaves like ``httpx.Response.raise_for_status`` but includes the + ``detail`` field from the response body (falling back to raw text) so + CLI users can see the actual server-side reason instead of just the + HTTP status line. The exception type is preserved so existing callers + / tests continue to work. + + Raises: + httpx.HTTPStatusError: When the response carries a 4xx or 5xx status. + """ + import httpx + + try: + resp.raise_for_status() + except httpx.HTTPStatusError as exc: + detail: str | None = None + try: + payload = resp.json() + except Exception: + payload = None + if isinstance(payload, dict): + detail_value = payload.get("detail") + if isinstance(detail_value, str) and detail_value.strip(): + detail = detail_value + elif detail_value is not None: + detail = str(detail_value) + if detail is None: + text = getattr(resp, "text", "") or "" + text = text.strip() + if text: + detail = text + if detail is None: + raise + message = f"{exc}: {detail}" + raise httpx.HTTPStatusError(message, request=exc.request, response=exc.response) from exc diff --git a/pyrit/cli/pyrit_scan.py b/pyrit/cli/pyrit_scan.py index e2652b3be6..57ee0a3326 100644 --- a/pyrit/cli/pyrit_scan.py +++ b/pyrit/cli/pyrit_scan.py @@ -30,6 +30,44 @@ _TERMINAL_STATUSES = {"COMPLETED", "FAILED", "CANCELLED"} +def _print_cli_exception(*, exc: BaseException) -> None: + """ + Print a user-facing error line for an exception that bubbled out of the CLI. + + Surfaces the exception class (so callers can tell ``ReadTimeout`` apart from + ``HTTPStatusError``) and dumps the traceback when log-level is ``DEBUG``. + Adds a specific hint for ``httpx.ReadTimeout`` since that case usually means + the server is taking longer than ``--request-timeout`` to respond and the + default bare ``str(exc)`` is empty. + + Args: + exc (BaseException): The exception caught by the CLI. + """ + import traceback + + try: + import httpx + + is_read_timeout = isinstance(exc, httpx.ReadTimeout) + except Exception: + is_read_timeout = False + + cls_name = type(exc).__name__ + detail = str(exc) or repr(exc) + + if is_read_timeout: + print( + "\nError (ReadTimeout): server did not respond in time. " + "Pass '--request-timeout ' to wait longer, or check the " + "server logs for a blocked event loop." + ) + else: + print(f"\nError ({cls_name}): {detail}") + + if logging.getLogger().isEnabledFor(logging.DEBUG): + traceback.print_exception(type(exc), exc, exc.__traceback__) + + _DESCRIPTION = """PyRIT Scanner - Run AI security scenarios from the command line. Requires a running PyRIT backend server. Use --start-server to launch one, @@ -113,6 +151,16 @@ def _build_base_parser(*, add_help: bool = True) -> ArgumentParser: default=logging.WARNING, help="Logging level (DEBUG, INFO, WARNING, ERROR, CRITICAL) (default: WARNING)", ) + server_group.add_argument( + "--request-timeout", + type=float, + default=None, + help=( + "HTTP read timeout in seconds for non-polling server requests " + "(catalog/results/cancel/etc). Defaults to 60. Polling a live " + "scenario run always waits indefinitely regardless of this value." + ), + ) # -- Discovery -- discovery_group = parser.add_argument_group("discovery") @@ -696,7 +744,10 @@ async def _run_async(*, parsed_args: Namespace) -> int: return 0 try: - async with PyRITApiClient(base_url=base_url_result) as client: + async with PyRITApiClient( + base_url=base_url_result, + request_timeout=getattr(parsed_args, "request_timeout", None), + ) as client: return await _dispatch_with_client_async(client=client, parsed_args=parsed_args) except ServerNotAvailableError as exc: _output.print_error_with_hint( @@ -705,7 +756,7 @@ async def _run_async(*, parsed_args: Namespace) -> int: ) return 1 except Exception as exc: - print(f"\nError: {exc}") + _print_cli_exception(exc=exc) return 1 diff --git a/pyrit/executor/attack/multi_turn/red_teaming.py b/pyrit/executor/attack/multi_turn/red_teaming.py index 0525b185a4..84cb085fca 100644 --- a/pyrit/executor/attack/multi_turn/red_teaming.py +++ b/pyrit/executor/attack/multi_turn/red_teaming.py @@ -247,20 +247,6 @@ async def _setup_async(self, *, context: MultiTurnAttackContext[Any]) -> None: memory_labels=self._memory_labels, ) - # Set up adversarial chat with prepended conversation - if context.prepended_conversation: - # Get adversarial messages with swapped roles - adversarial_messages = get_adversarial_chat_messages( - prepended_conversation=context.prepended_conversation, - adversarial_chat_conversation_id=context.session.adversarial_chat_conversation_id, - attack_identifier=self.get_identifier(), - adversarial_chat_target_identifier=self._adversarial_chat.get_identifier(), - labels=context.memory_labels, - ) - - for msg in adversarial_messages: - self._memory.add_message_to_memory(request=msg) - adversarial_system_prompt = self._adversarial_chat_system_prompt_template.render_template_value( objective=context.objective, max_turns=self._max_turns, @@ -268,6 +254,9 @@ async def _setup_async(self, *, context: MultiTurnAttackContext[Any]) -> None: if not adversarial_system_prompt: raise ValueError("Adversarial chat system prompt must be defined") + # ``set_system_prompt`` rejects any conversation that already has messages, + # so it must run before we hydrate the adversarial chat with the swapped + # prepended turns below. self._adversarial_chat.set_system_prompt( system_prompt=adversarial_system_prompt, conversation_id=context.session.adversarial_chat_conversation_id, @@ -275,6 +264,20 @@ async def _setup_async(self, *, context: MultiTurnAttackContext[Any]) -> None: labels=context.memory_labels, # deprecated ) + # Set up adversarial chat with prepended conversation + if context.prepended_conversation: + # Get adversarial messages with swapped roles + adversarial_messages = get_adversarial_chat_messages( + prepended_conversation=context.prepended_conversation, + adversarial_chat_conversation_id=context.session.adversarial_chat_conversation_id, + attack_identifier=self.get_identifier(), + adversarial_chat_target_identifier=self._adversarial_chat.get_identifier(), + labels=context.memory_labels, + ) + + for msg in adversarial_messages: + self._memory.add_message_to_memory(request=msg) + async def _perform_async(self, *, context: MultiTurnAttackContext[Any]) -> AttackResult: """ Execute the red teaming attack by iteratively generating prompts, diff --git a/pyrit/memory/memory_interface.py b/pyrit/memory/memory_interface.py index c49cca44a8..26448f5b6c 100644 --- a/pyrit/memory/memory_interface.py +++ b/pyrit/memory/memory_interface.py @@ -1722,6 +1722,7 @@ def get_attack_results( outcome: Optional[str] = None, attack_class: Optional[str] = None, attack_classes: Optional[Sequence[str]] = None, + atomic_attack_eval_hashes: Optional[Sequence[str]] = None, converter_classes: Optional[Sequence[str]] = None, converter_classes_match: Literal["all", "any"] = "all", has_converters: Optional[bool] = None, @@ -1747,6 +1748,11 @@ def get_attack_results( attack_classes (Optional[Sequence[str]], optional): Filter by exact attack class_name in attack_identifier. Returns attacks matching ANY of the listed class names (OR logic, case-sensitive). An empty sequence applies no filter. Defaults to None. + atomic_attack_eval_hashes (Optional[Sequence[str]], optional): Filter by behavioral + equivalence hash on ``atomic_attack_identifier.eval_hash`` (auto-stamped on persistence + by ``AtomicAttackEvaluationIdentifier``). Returns results matching ANY of the listed + hashes (OR logic, case-sensitive). Designed for ASR aggregation by technique + configuration. An empty sequence applies no filter. Defaults to None. converter_classes (Optional[Sequence[str]], optional): Filter by converter class names. Combination semantics for multiple entries are controlled by ``converter_classes_match``. An empty sequence filters to attacks that used no converters; ``None`` applies no @@ -1835,6 +1841,24 @@ def get_attack_results( ) ) + if atomic_attack_eval_hashes: + # Single JSON path query on the auto-stamped eval_hash. OR-combined across + # supplied hashes so callers can fetch history for multiple technique + # configurations in one round trip. + conditions.append( + or_( + *[ + self._get_condition_json_property_match( + json_column=AttackResultEntry.atomic_attack_identifier, + property_path="$.eval_hash", + value=h, + case_sensitive=True, + ) + for h in atomic_attack_eval_hashes + ] + ) + ) + if converter_classes is not None: # Non-empty sequence: filter to attacks that used ALL (or ANY, depending on # converter_classes_match) of the listed converters. diff --git a/pyrit/models/identifiers/__init__.py b/pyrit/models/identifiers/__init__.py index 0d0995fe7b..4e260fdf8c 100644 --- a/pyrit/models/identifiers/__init__.py +++ b/pyrit/models/identifiers/__init__.py @@ -26,6 +26,7 @@ ObjectiveTargetEvaluationIdentifier, ScorerEvaluationIdentifier, compute_eval_hash, + compute_inner_attack_eval_hash, ) from pyrit.models.identifiers.identifier_filters import IdentifierFilter, IdentifierType @@ -43,6 +44,7 @@ "class_name_to_snake_case", "ComponentIdentifier", "compute_eval_hash", + "compute_inner_attack_eval_hash", "EvaluationIdentifier", "Identifiable", "ObjectiveTargetEvaluationIdentifier", diff --git a/pyrit/models/identifiers/evaluation_identifier.py b/pyrit/models/identifiers/evaluation_identifier.py index 56ce818277..b2fc1b996d 100644 --- a/pyrit/models/identifiers/evaluation_identifier.py +++ b/pyrit/models/identifiers/evaluation_identifier.py @@ -23,12 +23,15 @@ from __future__ import annotations from abc import ABC -from typing import Any, ClassVar, Optional +from typing import TYPE_CHECKING, Any, ClassVar, Optional from pydantic import BaseModel, ConfigDict, Field from pyrit.models.identifiers.component_identifier import ComponentIdentifier, config_hash +if TYPE_CHECKING: + from pyrit.executor.attack.core.attack_strategy import AttackStrategy + # Behavioral params that define model output quality for scoring. TARGET_EVAL_PARAMS: frozenset[str] = frozenset({"underlying_model_name", "temperature", "top_p"}) TARGET_EVAL_PARAM_FALLBACKS: dict[str, str] = {"underlying_model_name": "model_name"} @@ -358,3 +361,25 @@ class ObjectiveTargetEvaluationIdentifier(EvaluationIdentifier): included_params=TARGET_EVAL_PARAMS, param_fallbacks=TARGET_EVAL_PARAM_FALLBACKS, ) + + +def compute_inner_attack_eval_hash(*, attack: AttackStrategy) -> str: + """ + Predict the eval hash the executor will stamp on persisted child rows + for this attack. + + Mirrors the inner-attack write path so callers can look up historical + results matching the same behavioral configuration *before* any row is + written. Use this rather than reconstructing the recipe inline. + + Args: + attack (AttackStrategy): Inner attack strategy. + + Returns: + str: The eval hash that will appear on persisted child rows. + """ + # Local import avoids a circular dependency inside the identifiers package. + from pyrit.models.identifiers.atomic_attack_identifier import build_atomic_attack_identifier + + composite = build_atomic_attack_identifier(attack_identifier=attack.get_identifier()) + return AtomicAttackEvaluationIdentifier(composite).eval_hash diff --git a/pyrit/models/scenario_result.py b/pyrit/models/scenario_result.py index 3680d83b7d..3bddaf776f 100644 --- a/pyrit/models/scenario_result.py +++ b/pyrit/models/scenario_result.py @@ -9,7 +9,7 @@ from typing import TYPE_CHECKING, Any, Literal import pyrit -from pyrit.models import AttackOutcome, AttackResult +from pyrit.models.attack_result import AttackOutcome, AttackResult if TYPE_CHECKING: from pyrit.models.identifiers.component_identifier import ComponentIdentifier diff --git a/pyrit/models/seeds/seed_attack_group.py b/pyrit/models/seeds/seed_attack_group.py index f52f9dac7f..cb077755db 100644 --- a/pyrit/models/seeds/seed_attack_group.py +++ b/pyrit/models/seeds/seed_attack_group.py @@ -9,6 +9,7 @@ from __future__ import annotations +import copy from typing import TYPE_CHECKING from pyrit.models.seeds.seed_group import SeedGroup @@ -159,9 +160,17 @@ def with_technique(self, *, technique: SeedAttackTechniqueGroup) -> SeedAttackGr technique_seeds = list(technique.seeds) merged_seeds = base + technique_seeds if idx is None else base[:idx] + technique_seeds + base[idx:] + # ``self`` and ``technique`` may be shared across multiple ``with_technique`` + # calls (e.g. the dispatcher reuses one ``bundle.seed_technique`` instance + # across every objective). Deepcopy first so the per-seed mutation below + # and the fresh group_id assigned by ``SeedAttackGroup.__init__`` only + # touch the returned group, leaving the originals untouched as the + # docstring promises. + merged_seeds = [copy.deepcopy(seed) for seed in merged_seeds] + # Clear group IDs so the new group assigns a fresh one. - # This mutates the seed objects, but _enforce_consistent_group_id - # in the constructor will immediately overwrite with a new UUID. + # ``_enforce_consistent_group_id`` in the constructor will overwrite + # all of them with a single new UUID. for seed in merged_seeds: seed.prompt_group_id = None diff --git a/pyrit/models/seeds/seed_prompt.py b/pyrit/models/seeds/seed_prompt.py index 027b6935e9..9656211be8 100644 --- a/pyrit/models/seeds/seed_prompt.py +++ b/pyrit/models/seeds/seed_prompt.py @@ -15,7 +15,7 @@ from tinytag import TinyTag from pyrit.common.path import PATHS_DICT -from pyrit.models import DataTypeSerializer +from pyrit.models.data_type_serializer import DataTypeSerializer from pyrit.models.literals import ( # noqa: TC001 (runtime-required by Pydantic field annotations) ChatMessageRole, PromptDataType, diff --git a/pyrit/scenario/__init__.py b/pyrit/scenario/__init__.py index 02d725c583..1be2caea53 100644 --- a/pyrit/scenario/__init__.py +++ b/pyrit/scenario/__init__.py @@ -31,17 +31,20 @@ # Import scenario submodules directly and register them as virtual subpackages # This allows: from pyrit.scenario.airt import ContentHarms # without needing separate pyrit/scenario/airt/ directories +from pyrit.scenario.scenarios import adaptive as _adaptive_module from pyrit.scenario.scenarios import airt as _airt_module from pyrit.scenario.scenarios import benchmark as _benchmark_module from pyrit.scenario.scenarios import foundry as _foundry_module from pyrit.scenario.scenarios import garak as _garak_module +sys.modules["pyrit.scenario.adaptive"] = _adaptive_module sys.modules["pyrit.scenario.airt"] = _airt_module sys.modules["pyrit.scenario.benchmark"] = _benchmark_module sys.modules["pyrit.scenario.garak"] = _garak_module sys.modules["pyrit.scenario.foundry"] = _foundry_module # Also expose as attributes for IDE support +adaptive = _adaptive_module airt = _airt_module benchmark = _benchmark_module garak = _garak_module @@ -59,6 +62,7 @@ "ScenarioStrategy", "ScenarioIdentifier", "ScenarioResult", + "adaptive", "airt", "benchmark", "garak", diff --git a/pyrit/scenario/core/attack_technique_factory.py b/pyrit/scenario/core/attack_technique_factory.py index 1c4c1c4d32..3a1b3c6315 100644 --- a/pyrit/scenario/core/attack_technique_factory.py +++ b/pyrit/scenario/core/attack_technique_factory.py @@ -332,6 +332,11 @@ def uses_adversarial(self) -> bool: """Whether this technique drives an adversarial chat during execution.""" return self._uses_adversarial + @property + def scoring_config_type(self) -> type | None: + """The required ``attack_scoring_config`` subtype, or ``None`` if any config is accepted.""" + return self._get_scoring_config_type() + def create( self, *, diff --git a/pyrit/scenario/scenarios/adaptive/__init__.py b/pyrit/scenario/scenarios/adaptive/__init__.py new file mode 100644 index 0000000000..98bdf444dc --- /dev/null +++ b/pyrit/scenario/scenarios/adaptive/__init__.py @@ -0,0 +1,28 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT license. + +"""Adaptive scenario classes.""" + +from pyrit.scenario.scenarios.adaptive.adaptive_scenario import AdaptiveScenario +from pyrit.scenario.scenarios.adaptive.dispatcher import ( + ADAPTIVE_ATTEMPT_LABEL, + AdaptiveTechniqueDispatcher, + TechniqueBundle, +) +from pyrit.scenario.scenarios.adaptive.selectors import ( + EpsilonGreedyTechniqueSelector, + SelectorScope, + TechniqueSelector, +) +from pyrit.scenario.scenarios.adaptive.text_adaptive import TextAdaptive + +__all__ = [ + "ADAPTIVE_ATTEMPT_LABEL", + "AdaptiveScenario", + "AdaptiveTechniqueDispatcher", + "EpsilonGreedyTechniqueSelector", + "SelectorScope", + "TechniqueBundle", + "TechniqueSelector", + "TextAdaptive", +] diff --git a/pyrit/scenario/scenarios/adaptive/adaptive_scenario.py b/pyrit/scenario/scenarios/adaptive/adaptive_scenario.py new file mode 100644 index 0000000000..081bfca989 --- /dev/null +++ b/pyrit/scenario/scenarios/adaptive/adaptive_scenario.py @@ -0,0 +1,400 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT license. + +""" +``AdaptiveScenario`` — modality-agnostic base for scenarios that pick attack +techniques per-objective using a ``TechniqueSelector``. + +Owns selector wiring, dispatcher construction, and per-dataset atomic-attack +emission. Concrete subclasses (``TextAdaptive``, future ``ImageAdaptive`` / +``AudioAdaptive``) only declare strategy class, default datasets, version, +and atomic-attack prefix. + +Baseline policy is ``Enabled``: prompt_sending runs as a separate baseline +comparison and is excluded from the adaptive technique pool. +""" + +from __future__ import annotations + +import logging +from abc import abstractmethod +from typing import TYPE_CHECKING, ClassVar + +from pyrit.common.utils import to_sha256 +from pyrit.executor.attack import AttackScoringConfig +from pyrit.models.identifiers import compute_inner_attack_eval_hash +from pyrit.scenario.core.atomic_attack import AtomicAttack +from pyrit.scenario.core.attack_technique import AttackTechnique +from pyrit.scenario.core.scenario import Scenario +from pyrit.scenario.core.scenario_target_defaults import get_default_adversarial_target +from pyrit.scenario.scenarios.adaptive.dispatcher import ( + AdaptiveTechniqueDispatcher, + TechniqueBundle, +) +from pyrit.scenario.scenarios.adaptive.selectors import ( + EpsilonGreedyTechniqueSelector, + TechniqueSelector, +) + +if TYPE_CHECKING: + from pyrit.models import SeedAttackGroup + from pyrit.prompt_target import PromptTarget + from pyrit.scenario.core.attack_technique_factory import AttackTechniqueFactory + from pyrit.scenario.core.dataset_configuration import DatasetConfiguration + from pyrit.scenario.core.scenario_strategy import ScenarioStrategy + from pyrit.score import TrueFalseScorer + +logger = logging.getLogger(__name__) + + +class AdaptiveScenario(Scenario): + """ + Abstract base for adaptive (epsilon-greedy) scenarios. + + Subclasses must implement the standard ``Scenario`` class-method overrides + and implement ``_atomic_attack_prefix``. Selector wiring + and dispatcher construction are handled here. + """ + + VERSION: ClassVar[int] + + @classmethod + @abstractmethod + def _atomic_attack_prefix(cls) -> str: + """ + Return the prefix for per-objective atomic-attack names (e.g. ``"adaptive_text"``). + + Must be unique across adaptive subclasses — different modalities + emitting the same prefix would collide on ``atomic_attack_name`` and + merge their resume bookkeeping silently. + """ + raise NotImplementedError + + @classmethod + @abstractmethod + def get_strategy_class(cls) -> type[ScenarioStrategy]: + """Return the scenario's strategy enum (subclasses must override).""" + raise NotImplementedError + + @classmethod + @abstractmethod + def get_default_strategy(cls) -> ScenarioStrategy: + """Return the scenario's default strategy aggregate (subclasses must override).""" + raise NotImplementedError + + @classmethod + @abstractmethod + def default_dataset_config(cls) -> DatasetConfiguration: + """Return the scenario's default ``DatasetConfiguration`` (subclasses must override).""" + raise NotImplementedError + + def __init__( + self, + *, + objective_scorer: TrueFalseScorer | None = None, + selector: TechniqueSelector | None = None, + scenario_result_id: str | None = None, + ) -> None: + """ + Args: + objective_scorer (TrueFalseScorer | None): Scorer used to judge each + response. Defaults to the composite scorer from the base class. + selector (TechniqueSelector | None): Pre-built selector. When ``None`` + (default) an ``EpsilonGreedyTechniqueSelector`` is created + with default settings. + scenario_result_id (str | None): ID of an existing ``ScenarioResult`` to resume. + """ + if not objective_scorer: + objective_scorer = self._get_default_objective_scorer() + self._objective_scorer: TrueFalseScorer = objective_scorer + + self._selector: TechniqueSelector = selector if selector is not None else EpsilonGreedyTechniqueSelector() + + super().__init__( + version=self.VERSION, + strategy_class=self.get_strategy_class(), + default_strategy=self.get_default_strategy(), + default_dataset_config=self.default_dataset_config(), + objective_scorer=objective_scorer, + scenario_result_id=scenario_result_id, + ) + + def _get_attack_technique_factories(self) -> dict[str, AttackTechniqueFactory]: + """ + Build factories from the canonical scenario-techniques catalog, + augmented with any factory currently in the global + ``AttackTechniqueRegistry``. + + The catalog defines the deterministic baseline pool — it is also the + source of truth for the strategy enum's valid values, so iteration + order and presence of techniques do not depend on registry + initialization order. Registry-registered factories whose name + matches a catalog entry **override** the catalog default, letting + operators swap in tuned configurations (custom adversarial chat, + different converter chain, etc.) without editing core. Factories + registered only in the registry (no matching strategy enum value) + are returned too but the scenario will only consume those whose + names appear in ``self._scenario_strategies``. When the registry + has not been initialized yet, the catalog alone is used. + + Subclasses may override to further customize the pool. + + Returns: + dict[str, AttackTechniqueFactory]: Mapping of technique name to factory. + """ + # Local import: ``scenario_techniques`` imports ``pyrit.scenario.core``, + # which transitively re-imports this module, so a top-level import + # would form a cycle during ``pyrit.scenario`` package initialization. + from pyrit.setup.initializers.components.scenario_techniques import ( + build_scenario_technique_factories, + ) + + catalog = {factory.name: factory for factory in build_scenario_technique_factories()} + try: + registry_overrides = super()._get_attack_technique_factories() + except RuntimeError: + # Registry not initialized yet (e.g. bare CLI parse before + # ScenarioTechniqueInitializer has run). Catalog alone is the + # safe fallback and matches the strategy enum's value set. + registry_overrides = {} + return {**catalog, **registry_overrides} + + async def _get_atomic_attacks_async(self) -> list[AtomicAttack]: + """ + Build one ``AtomicAttack`` per (dataset, compatible seed group) pair. + + For each dataset, construct a single ``AdaptiveTechniqueDispatcher`` + shared across that dataset's seed groups. For each seed group, ask + the dispatcher to build its per-objective ``SequentialAttack`` and + wrap it in its own ``AtomicAttack``. All dispatchers across all + datasets share one ``TechniqueSelector`` instance so learning + accumulates globally; selection is committed up-front during + scenario initialization, before any execution starts. + + When ``self._include_baseline`` is true (the default under + ``BASELINE_ATTACK_POLICY = Enabled``), a baseline ``AtomicAttack`` + named ``"baseline"`` is prepended at index 0. + + Returns: + list[AtomicAttack]: One ``AtomicAttack`` per compatible + seed group across all datasets, with the baseline (when + enabled) prepended at index 0. + + Raises: + ValueError: If ``self._objective_target`` is not set, or if + ``_build_techniques_dict`` finds no usable techniques. + """ + if self._objective_target is None: + raise ValueError("objective_target must be set before creating attacks") + + techniques = self._build_techniques_dict(objective_target=self._objective_target) + + seed_groups_by_dataset = self._dataset_config.get_seed_attack_groups() + atomic_attacks: list[AtomicAttack] = [] + for dataset_name, seed_groups in seed_groups_by_dataset.items(): + atomic_attacks.extend( + await self._build_atomics_for_dataset_async( + dataset_name=dataset_name, + seed_groups=seed_groups, + techniques=techniques, + selector=self._selector, + ) + ) + + if self._include_baseline: + all_seed_groups = [g for groups in seed_groups_by_dataset.values() for g in groups] + atomic_attacks.insert(0, self._build_baseline_atomic_attack(seed_groups=all_seed_groups)) + + return atomic_attacks + + def _build_techniques_dict( + self, + *, + objective_target: PromptTarget, + ) -> dict[str, TechniqueBundle]: + """ + Resolve selected strategies into a ``{eval_hash: TechniqueBundle}`` map. + + Each bundle carries the inner attack strategy along with the factory's + ``seed_technique`` and ``adversarial_chat`` so the dispatcher can + reproduce the static ``AtomicAttack`` execution path per attempt. + + Technique keys are eval hashes derived from the inner attack strategy's + identifier (run through ``AtomicAttackEvaluationIdentifier`` so seeds, + scorers, and operational target params are excluded). The same hash is + auto-stamped on every persisted ``AttackResultEntry.atomic_attack_identifier`` + by the executor, which lets the selector aggregate historical success + rates by behavioral configuration via + ``MemoryInterface.get_attack_results(atomic_attack_eval_hashes=...)``. + + For factories whose attack class narrows ``attack_scoring_config`` to a + specific subtype (e.g. ``TAPAttackScoringConfig`` for TAP), this method + builds the matching subtype using the scenario's objective scorer. + Techniques whose factory rejects the scenario scorer at construction + time (e.g. TAP also requires a ``FloatScaleThresholdScorer`` at runtime) + are dropped with a warning so the rest of the pool continues to run. + + Returns: + dict[str, TechniqueBundle]: Mapping from technique eval hash to its + bundle, in the order selected strategies were resolved. + + Raises: + ValueError: If no techniques remain after filtering. Includes the + requested techniques and skip reasons. + """ + selected_techniques = sorted({s.value for s in self._scenario_strategies}) + factories = self._get_attack_technique_factories() + + techniques: dict[str, TechniqueBundle] = {} + skipped_no_factory: list[str] = [] + skipped_incompatible: dict[str, str] = {} + for technique_name in selected_techniques: + factory = factories.get(technique_name) + if factory is None: + skipped_no_factory.append(technique_name) + logger.warning(f"No factory for technique '{technique_name}', skipping.") + continue + scoring_config = self._build_scoring_config_for_factory(factory=factory) + if scoring_config is None: + required_type = factory.scoring_config_type + required_name = required_type.__name__ if required_type is not None else "AttackScoringConfig" + reason = f"scenario scorer is incompatible with required {required_name}" + skipped_incompatible[technique_name] = reason + logger.warning(f"Skipping technique '{technique_name}': {reason}") + continue + try: + technique = factory.create( + objective_target=objective_target, + attack_scoring_config=scoring_config, + ) + except (TypeError, ValueError) as exc: + skipped_incompatible[technique_name] = str(exc) + logger.warning(f"Skipping technique '{technique_name}': {type(exc).__name__}: {exc}") + continue + eval_hash = compute_inner_attack_eval_hash(attack=technique.attack) + adversarial_chat = factory.adversarial_chat + if adversarial_chat is None and factory.uses_adversarial: + adversarial_chat = get_default_adversarial_target() + techniques[eval_hash] = TechniqueBundle( + attack=technique.attack, + name=technique_name, + seed_technique=technique.seed_technique, + adversarial_chat=adversarial_chat, + ) + + if not techniques: + details: list[str] = [] + if skipped_no_factory: + details.append(f"no factory registered: {sorted(skipped_no_factory)}") + if skipped_incompatible: + details.append(f"incompatible with scenario scorer: {sorted(skipped_incompatible)}") + suffix = f" ({'; '.join(details)})" if details else "" + raise ValueError( + f"{type(self).__name__}: no usable techniques after resolving strategies. " + f"Check the --strategies selection.{suffix}" + ) + + return techniques + + def _build_scoring_config_for_factory(self, *, factory: AttackTechniqueFactory) -> AttackScoringConfig | None: + """ + Build the most specific scoring config the factory's attack class accepts. + + When the attack's constructor narrows ``attack_scoring_config`` to a + subtype of ``AttackScoringConfig`` (e.g. TAP requires + ``TAPAttackScoringConfig``), construct that subtype directly so the + factory does not have to fall back to its WARN policy and silently + substitute an internal default scorer. + + Returns ``None`` when the required subtype itself rejects the + scenario's ``objective_scorer`` (e.g. TAP requires a + ``FloatScaleThresholdScorer`` while the scenario provides a + ``TrueFalseScorer``). ``None`` signals that no scenario-scorer- + preserving config exists for this technique — the caller drops the + technique rather than relying on the factory's override policy to + react (under WARN/SKIP the factory would silently substitute its + internal default scorer, masking the incompatibility). + + Returns: + AttackScoringConfig | None: The most specific config that could + be built, or ``None`` if the technique is incompatible with + the scenario scorer. + """ + required = factory.scoring_config_type + if required is None or required is AttackScoringConfig: + return AttackScoringConfig(objective_scorer=self._objective_scorer) + try: + return required(objective_scorer=self._objective_scorer) + except (TypeError, ValueError): + return None + + async def _build_atomics_for_dataset_async( + self, + *, + dataset_name: str, + seed_groups: list[SeedAttackGroup], + techniques: dict[str, TechniqueBundle], + selector: TechniqueSelector, + ) -> list[AtomicAttack]: + """ + Build one ``AtomicAttack`` per seed group with at least one + compatible technique. + + A single ``AdaptiveTechniqueDispatcher`` is constructed for this + dataset and used to build a fresh ``SequentialAttack`` per seed + group. Each returned atomic carries one seed group and one + pre-built attack whose children were selected up-front via the + dispatcher. + + Seed groups for which no technique in the pool is compatible are + dropped here with a warning. + + Returns: + list[AtomicAttack]: One atomic per compatible seed group. + Empty list when every seed group is incompatible with + every technique. + + Raises: + ValueError: If ``self._objective_target`` is not set + (defensive guard; ``_get_atomic_attacks_async`` enforces + this earlier). + """ + if self._objective_target is None: # pragma: no cover - defensive + raise ValueError("objective_target must be set before creating attacks") + + dispatcher = AdaptiveTechniqueDispatcher( + objective_target=self._objective_target, + techniques=techniques, + selector=selector, + objective_scorer=self._objective_scorer, + max_attempts_per_objective=self.params.get("max_attempts_per_objective", 3), + scenario_result_id=self._scenario_result_id, + ) + + atomics: list[AtomicAttack] = [] + for seed_group in seed_groups: + compatible = dispatcher.compatible_techniques(seed_group=seed_group) + if not compatible: + logger.warning( + "AdaptiveScenario: no compatible techniques for seed group in dataset '%s' " + "(objective=%r); skipping.", + dataset_name, + seed_group.objective.value, + ) + continue + + attack = await dispatcher.build_attack_async(seed_group=seed_group, compatible=compatible) + objective_sha = to_sha256(seed_group.objective.value) + atomic_attack_name = f"{self._atomic_attack_prefix()}_{dataset_name}::{objective_sha}" + atomics.append( + AtomicAttack( + atomic_attack_name=atomic_attack_name, + attack_technique=AttackTechnique(attack=attack), + seed_groups=[seed_group], + objective_scorer=self._objective_scorer, + memory_labels=dict(self._memory_labels), + display_group=dataset_name, + ) + ) + + return atomics diff --git a/pyrit/scenario/scenarios/adaptive/dispatcher.py b/pyrit/scenario/scenarios/adaptive/dispatcher.py new file mode 100644 index 0000000000..7005e89d76 --- /dev/null +++ b/pyrit/scenario/scenarios/adaptive/dispatcher.py @@ -0,0 +1,234 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT license. + +""" +``AdaptiveTechniqueDispatcher`` — selects inner techniques per objective via a +``TechniqueSelector`` and builds a ``SequentialAttack`` to run them. + +The dispatcher is a plain class, not an ``AttackStrategy``. It does not +execute anything and does not persist anything. ``AdaptiveScenario`` calls +``build_attack_async`` once per ``SeedAttackGroup`` during scenario +initialization, wraps each returned attack in its own ``AtomicAttack``, and +hands them to the scenario base for execution. + +The returned attack is a plain ``SequentialAttack`` with +``SequenceCompletionPolicy.FIRST_SUCCESS``. The per-attempt dispatch trail +(which technique ran, with what outcome, in what order) is not stamped onto +the envelope — every child ``AttackResult`` in +``SequentialAttackResult.child_attack_results`` already carries its own +``outcome`` and its own ``atomic_attack_identifier.eval_hash``. Callers that +want a human-readable technique label per child read it directly from the +child via ``child.get_attack_strategy_identifier().unique_name`` (the +executor auto-stamps ``class_name`` and ``unique_name`` on every persisted +row), so there is no separate ``{eval_hash: name}`` map to consult. +""" + +from __future__ import annotations + +import logging +from dataclasses import dataclass +from typing import TYPE_CHECKING, Any + +from pyrit.executor.attack.compound.sequential_attack import ( + SequenceCompletionPolicy, + SequentialAttack, + SequentialChildAttack, +) + +if TYPE_CHECKING: + from pyrit.executor.attack.core.attack_strategy import AttackStrategy + from pyrit.models import AttackResult, SeedAttackGroup, SeedAttackTechniqueGroup + from pyrit.prompt_target import PromptTarget + from pyrit.scenario.scenarios.adaptive.selectors import TechniqueSelector + from pyrit.score import TrueFalseScorer + +logger = logging.getLogger(__name__) + + +# Memory-label key stamped onto persisted prompt rows so adaptive attempts +# can be filtered/grouped after a run. +ADAPTIVE_ATTEMPT_LABEL: str = "_adaptive_attempt" +"""1-based attempt index within the per-objective loop.""" + + +@dataclass(frozen=True) +class TechniqueBundle: + """ + Per-technique bundle consumed by the dispatcher. + + Carries the inner attack strategy alongside the factory-supplied + ``seed_technique`` (if any) and ``adversarial_chat`` (required when the + seed_technique contains a simulated-conversation config). ``name`` is the + factory-registration key; the dispatcher does not consume it, but it is + convenient for diagnostics and is preserved here so callers/tests can + cross-check which factory each bundle came from. + + Notebook/report code that wants a human-readable label for a persisted + child ``AttackResult`` should read it from the child itself via + ``child.get_attack_strategy_identifier()`` — the executor already stamps + ``class_name`` and ``unique_name`` on every row, so there is no need to + publish a separate ``{eval_hash: name}`` map. + """ + + attack: AttackStrategy[Any, AttackResult] + name: str = "" + seed_technique: SeedAttackTechniqueGroup | None = None + adversarial_chat: PromptTarget | None = None + + +class AdaptiveTechniqueDispatcher: + """ + Selects inner techniques per objective and builds a ``SequentialAttack``. + + Not an ``AttackStrategy``: the dispatcher does not execute anything + and does not persist anything. It is a small factory used by + ``AdaptiveScenario`` at initialization to translate one + ``SeedAttackGroup`` (one objective) into one ready-to-run attack. + + For each call: query the selector for the top + ``max_attempts_per_objective`` techniques compatible with the seed + group, then construct a ``SequentialAttack`` (with + ``SequenceCompletionPolicy.FIRST_SUCCESS``) whose children are the + chosen techniques in priority order. The selector is shared by + reference across all calls in a scenario so learning accumulates + across objectives — though all selections are committed up-front + during scenario initialization (see + ``AdaptiveScenario._get_atomic_attacks_async``). + """ + + def __init__( + self, + *, + objective_target: PromptTarget, + techniques: dict[str, TechniqueBundle], + selector: TechniqueSelector, + objective_scorer: TrueFalseScorer | None = None, + max_attempts_per_objective: int = 3, + scenario_result_id: str | None = None, + ) -> None: + """ + Args: + objective_target (PromptTarget): The target inner attacks run against. + techniques (dict[str, TechniqueBundle]): Mapping from + technique eval hash to its bundle. Must be non-empty. + selector (TechniqueSelector): Stateless technique selector. + objective_scorer (TrueFalseScorer | None): Scorer forwarded + to inner attacks that generate simulated conversations. + max_attempts_per_objective (int): Maximum attempts per + objective; must be >= 1. Defaults to 3. + scenario_result_id (str | None): Passed to the selector to + scope memory queries to this scenario run. Defaults to + ``None``. + + Raises: + ValueError: If ``techniques`` is empty or + ``max_attempts_per_objective`` < 1. + """ + if not techniques: + raise ValueError("techniques must contain at least one attack technique") + if max_attempts_per_objective < 1: + raise ValueError(f"max_attempts_per_objective must be >= 1, got {max_attempts_per_objective}") + self._objective_target = objective_target + self._techniques = techniques + self._selector = selector + self._objective_scorer = objective_scorer + self._max_attempts = max_attempts_per_objective + self._scenario_result_id = scenario_result_id + + def compatible_techniques(self, *, seed_group: SeedAttackGroup) -> list[str]: + """ + Return technique hashes whose ``seed_technique`` is compatible with ``seed_group``. + + Techniques with no ``seed_technique`` are universally compatible. + Used by ``AdaptiveScenario`` to drop seed groups with no usable + techniques before building atomic attacks. + + Returns: + list[str]: Technique eval hashes in declaration order. + """ + return [ + name + for name, bundle in self._techniques.items() + if bundle.seed_technique is None or seed_group.is_compatible_with_technique(technique=bundle.seed_technique) + ] + + async def build_attack_async( + self, + *, + seed_group: SeedAttackGroup, + compatible: list[str] | None = None, + ) -> SequentialAttack: + """ + Build a ``SequentialAttack`` for one ``SeedAttackGroup``. + + Queries the selector for the top + ``max_attempts_per_objective`` techniques (filtered by per-call + seed-group compatibility) and wraps them in a + ``SequentialAttack`` with + ``SequenceCompletionPolicy.FIRST_SUCCESS``. + + Args: + seed_group (SeedAttackGroup): The seed group for the + objective this attack will run against. Must carry a + non-None objective. + compatible (list[str] | None): Precomputed result of + ``compatible_techniques(seed_group=...)``. When ``None`` + (default) the dispatcher computes it itself. Callers that + already filter empty pools out via ``compatible_techniques`` + should pass the result through to avoid re-scanning the + technique map. + + Returns: + SequentialAttack: The ready-to-run attack. Each child's + identity is captured by its own + ``atomic_attack_identifier.eval_hash`` after execution; + callers wanting the friendly technique name read it + directly from the child via + ``child.get_attack_strategy_identifier().unique_name``. + + Raises: + ValueError: If ``seed_group.objective`` is not initialized, + or if no techniques in the pool are compatible with the + seed group. + """ + if seed_group.objective is None: + raise ValueError("seed_group.objective is not initialized") + + if compatible is None: + compatible = self.compatible_techniques(seed_group=seed_group) + if not compatible: + raise ValueError( + f"AdaptiveTechniqueDispatcher: no compatible techniques for seed group " + f"(objective={seed_group.objective.value!r})." + ) + + chosen_hashes = await self._selector.select_async( + technique_identifiers=compatible, + objective=seed_group.objective.value, + num_top_techniques=self._max_attempts, + scenario_result_id=self._scenario_result_id, + ) + + child_attacks: list[SequentialChildAttack] = [] + for attempt_idx, chosen in enumerate(chosen_hashes): + bundle = self._techniques[chosen] + execution_group = ( + seed_group.with_technique(technique=bundle.seed_technique) + if bundle.seed_technique is not None + else seed_group + ) + child_attacks.append( + SequentialChildAttack( + strategy=bundle.attack, + seed_group=execution_group, + adversarial_chat=bundle.adversarial_chat, + objective_scorer=self._objective_scorer, + memory_labels={ADAPTIVE_ATTEMPT_LABEL: str(attempt_idx + 1)}, + ) + ) + + return SequentialAttack( + objective_target=self._objective_target, + child_attacks=child_attacks, + completion_policy=SequenceCompletionPolicy.FIRST_SUCCESS, + ) diff --git a/pyrit/scenario/scenarios/adaptive/selectors/__init__.py b/pyrit/scenario/scenarios/adaptive/selectors/__init__.py new file mode 100644 index 0000000000..709146fb43 --- /dev/null +++ b/pyrit/scenario/scenarios/adaptive/selectors/__init__.py @@ -0,0 +1,18 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT license. + +"""Selector protocol and selector implementations.""" + +from pyrit.scenario.scenarios.adaptive.selectors.epsilon_greedy import ( + EpsilonGreedyTechniqueSelector, +) +from pyrit.scenario.scenarios.adaptive.selectors.technique_selector import ( + SelectorScope, + TechniqueSelector, +) + +__all__ = [ + "EpsilonGreedyTechniqueSelector", + "SelectorScope", + "TechniqueSelector", +] diff --git a/pyrit/scenario/scenarios/adaptive/selectors/epsilon_greedy.py b/pyrit/scenario/scenarios/adaptive/selectors/epsilon_greedy.py new file mode 100644 index 0000000000..a415182733 --- /dev/null +++ b/pyrit/scenario/scenarios/adaptive/selectors/epsilon_greedy.py @@ -0,0 +1,175 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT license. + +"""Epsilon-greedy technique selector for adaptive scenarios.""" + +from __future__ import annotations + +import hashlib +import logging +import random +import struct +from typing import TYPE_CHECKING + +from pyrit.analytics.technique_analysis import compute_technique_stats +from pyrit.scenario.scenarios.adaptive.selectors.technique_selector import SelectorScope + +if TYPE_CHECKING: + from collections.abc import Sequence + + from pyrit.analytics.result_analysis import AttackStats + +logger = logging.getLogger(__name__) + + +def _derive_rng(random_seed: int | None, decision_key: str) -> random.Random: + """ + Derive a per-decision ``Random`` from ``(random_seed, decision_key)``. + + Returns: + random.Random: A fresh ``random.Random`` seeded deterministically from the + inputs when ``random_seed`` is not None, or an unseeded ``Random`` otherwise. + """ + if random_seed is None: + return random.Random() + digest = hashlib.sha256(f"{random_seed}|{decision_key}".encode()).digest() + derived_seed = struct.unpack(" None: + """ + Args: + epsilon (float): Exploration probability in [0.0, 1.0]. Defaults to 0.2. + scope (SelectorScope | None): Filter describing which historical + ``AttackResult`` rows to use when estimating success rates. + Defaults to ``SelectorScope.all_runs()`` (all history). + random_seed (int | None): Base seed for deterministic per-decision RNG + derivation. Defaults to ``None`` (non-deterministic). + + Raises: + ValueError: If ``epsilon`` is outside [0.0, 1.0]. + """ + if not 0.0 <= epsilon <= 1.0: + raise ValueError(f"epsilon must be in [0.0, 1.0], got {epsilon}") + + self._epsilon = epsilon + self._scope = scope if scope is not None else SelectorScope.all_runs() + self._seed = random_seed + + async def select_async( + self, + *, + technique_identifiers: Sequence[str], + objective: str, + num_top_techniques: int = 1, + scenario_result_id: str | None = None, + ) -> Sequence[str]: + """ + Return up to ``num_top_techniques`` techniques in priority order. + + Args: + technique_identifiers (Sequence[str]): Available technique names. + objective (str): The objective text for scoping the per-decision RNG. + num_top_techniques (int): Max techniques to return. Defaults to 1. + scenario_result_id (str | None): The current scenario run ID, supplied + by the dispatcher. Folded into the per-decision RNG key so distinct + runs diverge while resumes (same ``scenario_result_id``) reproduce + the original picks; also forwarded to memory only when the + configured ``scope.current_run_only`` is ``True``. Defaults to + ``None``. + + Returns: + Sequence[str]: Techniques in priority order. Fewer than + ``num_top_techniques`` if not enough techniques are available. + + Raises: + ValueError: If ``technique_identifiers`` is empty. + """ + technique_list = list(technique_identifiers) + if not technique_list: + raise ValueError("technique_identifiers must contain at least one entry") + + num_top_techniques = min(num_top_techniques, len(technique_list)) + + # Fold scenario_result_id into the RNG key so two different scenario runs + # over the same objective explore differently. + decision_key = f"{objective}|{scenario_result_id or ''}" + rng = _derive_rng(self._seed, decision_key) + + effective_run_id = scenario_result_id if self._scope.current_run_only else None + stats = compute_technique_stats( + technique_eval_hashes=technique_list, + scenario_result_id=effective_run_id, + targeted_harm_categories=self._scope.targeted_harm_categories, + ) + + chosen: list[str] = [] + remaining = list(technique_list) + + for _ in range(num_top_techniques): + if not remaining: + break + + if rng.random() < self._epsilon: + pick = rng.choice(remaining) + else: + estimates = {t: self._estimate(technique=t, stats=stats) for t in remaining} + best = max(estimates.values()) + winners = [t for t, v in estimates.items() if v >= best - self._TIE_TOL] + pick = rng.choice(winners) + + chosen.append(pick) + remaining.remove(pick) + + return chosen + + @staticmethod + def _estimate(*, technique: str, stats: dict[str, AttackStats]) -> float: + """ + Laplace-smoothed success-rate estimate for a technique. + + Unseen techniques get ``(0 + 1) / (0 + 1) = 1.0`` (optimistic init). + + Args: + technique (str): The technique name. + stats (dict[str, AttackStats]): Pre-computed stats from memory. + + Returns: + float: Estimated success rate in ``(0, 1]``. + """ + technique_stats = stats.get(technique) + if technique_stats is None or technique_stats.total_decided == 0: + return 1.0 + return (technique_stats.successes + 1) / (technique_stats.total_decided + 1) diff --git a/pyrit/scenario/scenarios/adaptive/selectors/technique_selector.py b/pyrit/scenario/scenarios/adaptive/selectors/technique_selector.py new file mode 100644 index 0000000000..eada0fb5ed --- /dev/null +++ b/pyrit/scenario/scenarios/adaptive/selectors/technique_selector.py @@ -0,0 +1,103 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT license. + +"""Technique selector protocol for adaptive scenarios.""" + +from __future__ import annotations + +from dataclasses import dataclass +from typing import TYPE_CHECKING, Protocol, runtime_checkable + +if TYPE_CHECKING: + from collections.abc import Sequence + + +@dataclass(frozen=True) +class SelectorScope: + """ + Filter describing which historical ``AttackResult`` rows a selector + queries when estimating technique success rates. + + All fields default to "no restriction"; combine fields to narrow the + scope (e.g. current run only, same harm category). Filter values flow + through ``compute_technique_stats`` to + ``MemoryInterface.get_attack_results``. + + The scope is held by the selector at construction time. The per-call + ``scenario_result_id`` is supplied by the dispatcher and is forwarded + to memory only when ``current_run_only`` is set; otherwise the selector + queries across all runs. + + Per-technique disambiguation uses ``atomic_attack_identifier.eval_hash`` + (auto-stamped on every persisted attack result), which already encodes + the attack class plus its behavior-relevant params. Class-based + narrowing is therefore unnecessary at this layer. + """ + + current_run_only: bool = False + """Restrict to the dispatcher-supplied ``scenario_result_id`` for the + in-flight run. When ``False`` (default), query across all runs.""" + + targeted_harm_categories: Sequence[str] | None = None + """Filter to results whose prompts targeted these harm categories. + ``None`` means no harm-category filter.""" + + @classmethod + def all_runs(cls) -> SelectorScope: + """ + Build a scope that queries across all historical scenario runs (the default). + + Returns: + SelectorScope: A scope with no restrictions. + """ + return cls() + + @classmethod + def current_run(cls) -> SelectorScope: + """ + Build a scope restricted to the dispatcher-supplied scenario run. + + Returns: + SelectorScope: A scope with ``current_run_only=True``. + """ + return cls(current_run_only=True) + + +@runtime_checkable +class TechniqueSelector(Protocol): + """ + Protocol for adaptive technique selectors. + + Selectors are **stateless** — they query memory for historical success + rates rather than maintaining internal counts. Calling ``select_async`` + with the same arguments twice should yield the same answer + (deterministic given memory contents). + """ + + async def select_async( + self, + *, + technique_identifiers: Sequence[str], + objective: str, + num_top_techniques: int = 1, + scenario_result_id: str | None = None, + ) -> Sequence[str]: + """ + Return techniques in priority order (try first, try second, …). + + Args: + technique_identifiers (Sequence[str]): Available technique eval + hashes. + objective (str): The objective text for this selection. + num_top_techniques (int): Max techniques to return. Defaults to 1. + scenario_result_id (str | None): The current scenario run ID, + provided by the dispatcher. Selectors forward this to + memory only when their ``SelectorScope`` has + ``current_run_only=True``. + + Returns: + Sequence[str]: Up to ``num_top_techniques`` technique eval hashes + in priority order. Fewer if not enough techniques are + available. + """ + ... # pragma: no cover diff --git a/pyrit/scenario/scenarios/adaptive/text_adaptive.py b/pyrit/scenario/scenarios/adaptive/text_adaptive.py new file mode 100644 index 0000000000..1be215a536 --- /dev/null +++ b/pyrit/scenario/scenarios/adaptive/text_adaptive.py @@ -0,0 +1,174 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT license. + +""" +``TextAdaptive`` — text adaptive scenario. + +Picks attack techniques per-objective using an epsilon-greedy selector +informed by observed success rates. Runs up to ``max_attempts_per_objective`` +techniques per objective and stops early on success. ``prompt_sending`` is +excluded from the adaptive technique pool and runs as the baseline comparison +instead. +""" + +from __future__ import annotations + +import logging +from typing import TYPE_CHECKING, ClassVar + +from pyrit.common import apply_defaults +from pyrit.common.parameter import Parameter +from pyrit.registry.object_registries.attack_technique_registry import ( + AttackTechniqueRegistry, +) +from pyrit.registry.tag_query import TagQuery +from pyrit.scenario.core.dataset_configuration import DatasetConfiguration +from pyrit.scenario.scenarios.adaptive.adaptive_scenario import AdaptiveScenario + +if TYPE_CHECKING: + from pyrit.scenario.core.scenario_strategy import ScenarioStrategy + from pyrit.scenario.scenarios.adaptive.selectors import TechniqueSelector + from pyrit.score import TrueFalseScorer + +logger = logging.getLogger(__name__) + +# Techniques excluded from the adaptive technique pool. These run as the +# baseline comparison rather than as adversarial moves the selector chooses. +_EXCLUDED_TECHNIQUES = frozenset({"prompt_sending"}) + + +def _build_text_adaptive_strategy() -> type[ScenarioStrategy]: + """ + Build the strategy enum from the core scenario-techniques catalog, + excluding techniques that run as baseline. + + Returns: + type[ScenarioStrategy]: The dynamically-built strategy enum class. + + Logs a warning if any name in ``_EXCLUDED_TECHNIQUES`` is not present + in the current catalog. The exclusion is defensive — when the catalog + does not contain the named technique, the filter is a no-op, and we + surface that so a stale entry in the exclusion list (or a renamed + catalog entry) doesn't silently break the intended exclusion. + """ + # Local import: ``scenario_techniques`` imports ``pyrit.scenario.core``, + # which transitively re-imports this module, so a top-level import would + # form a cycle during ``pyrit.scenario`` package initialization. + from pyrit.setup.initializers.components.scenario_techniques import ( + build_scenario_technique_factories, + ) + + all_factories = list(build_scenario_technique_factories()) + catalog_names = {factory.name for factory in all_factories} + unmatched = _EXCLUDED_TECHNIQUES - catalog_names + if unmatched: + logger.warning( + "TextAdaptive: _EXCLUDED_TECHNIQUES entries %s are not in the current " + "scenario-techniques catalog %s; the exclusion is a no-op for those entries. " + "Remove stale entries or update the catalog.", + sorted(unmatched), + sorted(catalog_names), + ) + + factories = [factory for factory in all_factories if factory.name not in _EXCLUDED_TECHNIQUES] + + return AttackTechniqueRegistry.build_strategy_class_from_factories( # type: ignore[return-value, ty:invalid-return-type] + class_name="TextAdaptiveStrategy", + factories=factories, + aggregate_tags={ + "default": TagQuery.any_of("default"), + "single_turn": TagQuery.any_of("single_turn"), + "multi_turn": TagQuery.any_of("multi_turn"), + }, + ) + + +class TextAdaptive(AdaptiveScenario): + """ + Adaptive text-attack scenario. + + Selects techniques per-objective via an epsilon-greedy selector over the + set of selected strategies. ``prompt_sending`` runs as the baseline + comparison and is excluded from the adaptive technique pool. + """ + + _cached_strategy_class: ClassVar[type[ScenarioStrategy] | None] = None + + VERSION: int = 1 + + @classmethod + def _atomic_attack_prefix(cls) -> str: + """Return the prefix for per-objective atomic-attack names.""" + return "adaptive_text" + + @classmethod + def get_strategy_class(cls) -> type[ScenarioStrategy]: + """Return the strategy enum for this scenario, building it once on first access.""" + if cls._cached_strategy_class is None: + cls._cached_strategy_class = _build_text_adaptive_strategy() + return cls._cached_strategy_class + + @classmethod + def get_default_strategy(cls) -> ScenarioStrategy: + """Return the default strategy aggregate (resolves to every ``default``-tagged technique).""" + strategy_class = cls.get_strategy_class() + return strategy_class("default") + + @classmethod + def required_datasets(cls) -> list[str]: + """Return the dataset names this scenario expects when no override is provided.""" + return [ + "airt_hate", + "airt_fairness", + "airt_violence", + "airt_sexual", + "airt_harassment", + "airt_misinformation", + "airt_leakage", + ] + + @classmethod + def default_dataset_config(cls) -> DatasetConfiguration: + """Return the default ``DatasetConfiguration`` (required datasets, capped at 4 per dataset).""" + return DatasetConfiguration(dataset_names=cls.required_datasets(), max_dataset_size=4) + + @classmethod + def supported_parameters(cls) -> list[Parameter]: + """ + Declare custom parameters this scenario accepts from the CLI / config file. + + Returns: + list[Parameter]: Parameters configurable per-run. + """ + return [ + Parameter( + name="max_attempts_per_objective", + description="Max techniques tried per objective. Defaults to 3.", + param_type=int, + default=3, + ), + ] + + @apply_defaults + def __init__( + self, + *, + objective_scorer: TrueFalseScorer | None = None, + selector: TechniqueSelector | None = None, + scenario_result_id: str | None = None, + ) -> None: + """ + Args: + objective_scorer (TrueFalseScorer | None): Scorer used to judge each + response. Defaults to the composite scorer from the base class. + selector (TechniqueSelector | None): Pre-built selector. When ``None`` + (default) an ``EpsilonGreedyTechniqueSelector`` is created + with default settings. Pass a custom instance to tune + ``epsilon`` or ``random_seed``. + scenario_result_id (str | None): ID of an existing ``ScenarioResult`` to resume. + """ + super().__init__( + objective_scorer=objective_scorer, + selector=selector, + scenario_result_id=scenario_result_id, + ) diff --git a/tests/unit/analytics/test_technique_analysis.py b/tests/unit/analytics/test_technique_analysis.py new file mode 100644 index 0000000000..04b1d94890 --- /dev/null +++ b/tests/unit/analytics/test_technique_analysis.py @@ -0,0 +1,146 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT license. + +from unittest.mock import MagicMock, patch + +import pytest + +from pyrit.analytics.technique_analysis import compute_technique_stats +from pyrit.models import AttackOutcome + + +def _make_result(*, eval_hash: str | None, outcome: AttackOutcome) -> MagicMock: + r = MagicMock() + if eval_hash is None: + r.atomic_attack_identifier = None + else: + identifier = MagicMock() + identifier.eval_hash = eval_hash + r.atomic_attack_identifier = identifier + r.outcome = outcome + return r + + +@pytest.fixture(autouse=True) +def _patch_memory(): + mock_memory = MagicMock() + mock_memory.get_attack_results.return_value = [] + with patch("pyrit.analytics.technique_analysis.CentralMemory") as cm: + cm.get_memory_instance.return_value = mock_memory + yield mock_memory + + +class TestComputeTechniqueStats: + def test_empty_results_returns_empty(self, _patch_memory): + stats = compute_technique_stats(technique_eval_hashes=["a", "b"]) + assert stats == {} + + def test_empty_hashes_short_circuits(self, _patch_memory): + stats = compute_technique_stats(technique_eval_hashes=[]) + assert stats == {} + _patch_memory.get_attack_results.assert_not_called() + + def test_counts_successes_and_failures(self, _patch_memory): + _patch_memory.get_attack_results.return_value = [ + _make_result(eval_hash="a", outcome=AttackOutcome.SUCCESS), + _make_result(eval_hash="a", outcome=AttackOutcome.SUCCESS), + _make_result(eval_hash="a", outcome=AttackOutcome.FAILURE), + _make_result(eval_hash="b", outcome=AttackOutcome.FAILURE), + ] + + stats = compute_technique_stats(technique_eval_hashes=["a", "b"]) + + assert stats["a"].successes == 2 + assert stats["a"].failures == 1 + assert stats["a"].total_decided == 3 + assert stats["b"].successes == 0 + assert stats["b"].failures == 1 + + def test_counts_errors_and_undetermined(self, _patch_memory): + _patch_memory.get_attack_results.return_value = [ + _make_result(eval_hash="a", outcome=AttackOutcome.ERROR), + _make_result(eval_hash="a", outcome=AttackOutcome.UNDETERMINED), + ] + + stats = compute_technique_stats(technique_eval_hashes=["a"]) + + assert stats["a"].errors == 1 + assert stats["a"].undetermined == 1 + + def test_ignores_hashes_not_in_requested_list(self, _patch_memory): + _patch_memory.get_attack_results.return_value = [ + _make_result(eval_hash="a", outcome=AttackOutcome.SUCCESS), + _make_result(eval_hash="c", outcome=AttackOutcome.SUCCESS), + ] + + stats = compute_technique_stats(technique_eval_hashes=["a", "b"]) + + assert "a" in stats + assert "c" not in stats + + def test_skips_results_without_eval_hash(self, _patch_memory): + _patch_memory.get_attack_results.return_value = [ + _make_result(eval_hash="a", outcome=AttackOutcome.SUCCESS), + _make_result(eval_hash=None, outcome=AttackOutcome.SUCCESS), + ] + + stats = compute_technique_stats(technique_eval_hashes=["a"]) + + assert stats["a"].successes == 1 + + def test_passes_eval_hashes_to_memory_query(self, _patch_memory): + compute_technique_stats(technique_eval_hashes=["x", "y"]) + + call_kwargs = _patch_memory.get_attack_results.call_args[1] + assert call_kwargs["atomic_attack_eval_hashes"] == ["x", "y"] + assert call_kwargs["scenario_result_id"] is None + assert call_kwargs["targeted_harm_categories"] is None + + def test_passes_scenario_result_id_to_memory_query(self, _patch_memory): + compute_technique_stats(technique_eval_hashes=["x"], scenario_result_id="run-123") + + call_kwargs = _patch_memory.get_attack_results.call_args[1] + assert call_kwargs["scenario_result_id"] == "run-123" + + def test_omits_hashes_with_no_history(self, _patch_memory): + _patch_memory.get_attack_results.return_value = [ + _make_result(eval_hash="a", outcome=AttackOutcome.SUCCESS), + ] + + stats = compute_technique_stats(technique_eval_hashes=["a", "b"]) + + assert "a" in stats + assert "b" not in stats + + def test_success_rate_computed(self, _patch_memory): + _patch_memory.get_attack_results.return_value = [ + _make_result(eval_hash="a", outcome=AttackOutcome.SUCCESS), + _make_result(eval_hash="a", outcome=AttackOutcome.SUCCESS), + _make_result(eval_hash="a", outcome=AttackOutcome.FAILURE), + _make_result(eval_hash="a", outcome=AttackOutcome.FAILURE), + ] + + stats = compute_technique_stats(technique_eval_hashes=["a"]) + + assert stats["a"].success_rate == pytest.approx(0.5) + + def test_passes_harm_categories_to_memory_query(self, _patch_memory): + compute_technique_stats( + technique_eval_hashes=["x"], + targeted_harm_categories=["misinformation", "hate"], + ) + + call_kwargs = _patch_memory.get_attack_results.call_args[1] + assert call_kwargs["targeted_harm_categories"] == ["misinformation", "hate"] + + def test_injected_memory_bypasses_central_memory(self, _patch_memory): + injected = MagicMock() + injected.get_attack_results.return_value = [ + _make_result(eval_hash="a", outcome=AttackOutcome.SUCCESS), + ] + + stats = compute_technique_stats(technique_eval_hashes=["a"], memory=injected) + + injected.get_attack_results.assert_called_once() + _patch_memory.get_attack_results.assert_not_called() + assert stats["a"].successes == 1 diff --git a/tests/unit/cli/test_api_client.py b/tests/unit/cli/test_api_client.py index a11df194fe..95c232273e 100644 --- a/tests/unit/cli/test_api_client.py +++ b/tests/unit/cli/test_api_client.py @@ -59,6 +59,26 @@ async def test_async_context_manager_opens_and_closes(mock_httpx_client): # After exit, close was called mock_httpx_client.aclose.assert_awaited_once() assert c._client is None + # Default request_timeout (60s) propagates to the httpx client constructor. + fake_async_client_cls.assert_called_once_with(base_url="http://localhost:8000", timeout=60.0) + + +async def test_async_context_manager_passes_custom_request_timeout(mock_httpx_client): + c = PyRITApiClient(base_url="http://localhost:8000", request_timeout=120.0) + fake_async_client_cls = MagicMock(return_value=mock_httpx_client) + with patch("httpx.AsyncClient", fake_async_client_cls): + async with c: + pass + fake_async_client_cls.assert_called_once_with(base_url="http://localhost:8000", timeout=120.0) + + +async def test_async_context_manager_uses_default_when_request_timeout_is_none(mock_httpx_client): + c = PyRITApiClient(base_url="http://localhost:8000", request_timeout=None) + fake_async_client_cls = MagicMock(return_value=mock_httpx_client) + with patch("httpx.AsyncClient", fake_async_client_cls): + async with c: + pass + fake_async_client_cls.assert_called_once_with(base_url="http://localhost:8000", timeout=60.0) async def test_close_async_is_noop_when_already_closed(): @@ -199,10 +219,27 @@ async def test_start_scenario_run_async(client, mock_httpx_client): async def test_get_scenario_run_async(client, mock_httpx_client): + import httpx as _httpx + mock_httpx_client.get.return_value = _make_response(json_data={"status": "RUNNING"}) result = await client.get_scenario_run_async(scenario_result_id="abc") assert result == {"status": "RUNNING"} - mock_httpx_client.get.assert_awaited_once_with("/api/scenarios/runs/abc", params=None) + # Polling uses read=None so a busy server doesn't trip the client default + # timeout while a scenario is executing. + mock_httpx_client.get.assert_awaited_once() + args, kwargs = mock_httpx_client.get.call_args + assert args == ("/api/scenarios/runs/abc",) + assert kwargs["params"] is None + timeout = kwargs["timeout"] + assert isinstance(timeout, _httpx.Timeout) + assert timeout.read is None + assert timeout.connect == 10.0 + + +async def test_get_scenario_run_async_wraps_connect_error(client, mock_httpx_client): + mock_httpx_client.get.side_effect = httpx.ConnectError("nope") + with pytest.raises(ServerNotAvailableError, match="Cannot connect"): + await client.get_scenario_run_async(scenario_result_id="abc") async def test_get_scenario_run_results_async(client, mock_httpx_client): diff --git a/tests/unit/memory/memory_interface/test_interface_attack_results.py b/tests/unit/memory/memory_interface/test_interface_attack_results.py index ea55436865..cdb611e64b 100644 --- a/tests/unit/memory/memory_interface/test_interface_attack_results.py +++ b/tests/unit/memory/memory_interface/test_interface_attack_results.py @@ -1360,6 +1360,62 @@ def test_get_attack_results_attack_classes_empty_returns_all(sqlite_instance: Me assert len(results) == 2 +def _eval_hash_for(class_name: str) -> str: + from pyrit.models.identifiers.evaluation_identifier import AtomicAttackEvaluationIdentifier + + return AtomicAttackEvaluationIdentifier( + build_atomic_attack_identifier( + attack_identifier=ComponentIdentifier( + class_name=class_name, + class_module="pyrit.attacks", + ), + ), + ).eval_hash + + +def test_get_attack_results_by_atomic_attack_eval_hashes_single(sqlite_instance: MemoryInterface): + """Filter by a single eval_hash; only matching rows are returned.""" + ar1 = _make_attack_result_with_identifier("conv_1", "CrescendoAttack") + ar2 = _make_attack_result_with_identifier("conv_2", "ManualAttack") + sqlite_instance.add_attack_results_to_memory(attack_results=[ar1, ar2]) + + target_hash = _eval_hash_for("CrescendoAttack") + results = sqlite_instance.get_attack_results(atomic_attack_eval_hashes=[target_hash]) + assert len(results) == 1 + assert results[0].conversation_id == "conv_1" + + +def test_get_attack_results_by_atomic_attack_eval_hashes_multi_uses_or(sqlite_instance: MemoryInterface): + """Multiple eval_hashes OR-combine — matches any of the listed hashes.""" + ar1 = _make_attack_result_with_identifier("conv_1", "CrescendoAttack") + ar2 = _make_attack_result_with_identifier("conv_2", "ManualAttack") + ar3 = _make_attack_result_with_identifier("conv_3", "TreeOfAttacksAttack") + sqlite_instance.add_attack_results_to_memory(attack_results=[ar1, ar2, ar3]) + + hashes = [_eval_hash_for("CrescendoAttack"), _eval_hash_for("ManualAttack")] + results = sqlite_instance.get_attack_results(atomic_attack_eval_hashes=hashes) + assert {r.conversation_id for r in results} == {"conv_1", "conv_2"} + + +def test_get_attack_results_atomic_attack_eval_hashes_empty_returns_all(sqlite_instance: MemoryInterface): + """atomic_attack_eval_hashes=[] behaves like None (no filter applied).""" + ar1 = _make_attack_result_with_identifier("conv_1", "CrescendoAttack") + ar2 = _make_attack_result_with_identifier("conv_2", "ManualAttack") + sqlite_instance.add_attack_results_to_memory(attack_results=[ar1, ar2]) + + results = sqlite_instance.get_attack_results(atomic_attack_eval_hashes=[]) + assert len(results) == 2 + + +def test_get_attack_results_atomic_attack_eval_hashes_no_match(sqlite_instance: MemoryInterface): + """A non-matching eval_hash returns no rows.""" + ar1 = _make_attack_result_with_identifier("conv_1", "CrescendoAttack") + sqlite_instance.add_attack_results_to_memory(attack_results=[ar1]) + + results = sqlite_instance.get_attack_results(atomic_attack_eval_hashes=["deadbeef" * 8]) + assert len(results) == 0 + + def test_get_attack_results_converter_classes_none_returns_all(sqlite_instance: MemoryInterface): """Test that converter_classes=None (omitted) returns all attacks unfiltered.""" ar1 = _make_attack_result_with_identifier("conv_1", "Attack", ["Base64Converter"]) diff --git a/tests/unit/models/identifiers/test_evaluation_identifier.py b/tests/unit/models/identifiers/test_evaluation_identifier.py index bf344fe878..46a6898fed 100644 --- a/tests/unit/models/identifiers/test_evaluation_identifier.py +++ b/tests/unit/models/identifiers/test_evaluation_identifier.py @@ -586,6 +586,77 @@ def test_scorer_eval_hash_matches_with_and_without_round_robin(self): assert eval_direct == eval_rr +class TestComputeInnerAttackEvalHash: + """``compute_inner_attack_eval_hash`` should match what the executor stamps.""" + + def _attack_with_identifier(self, identifier: ComponentIdentifier): + from unittest.mock import MagicMock + + attack = MagicMock() + attack.get_identifier.return_value = identifier + return attack + + def test_matches_manual_two_step_composition(self): + """Helper equals the executor recipe (build_atomic_attack_identifier + AtomicAttackEvaluationIdentifier).""" + from pyrit.models.identifiers import ( + AtomicAttackEvaluationIdentifier, + build_atomic_attack_identifier, + compute_inner_attack_eval_hash, + ) + + inner_id = ComponentIdentifier( + class_name="PromptSendingAttack", + class_module="pyrit.executor.attack.single_turn.prompt_sending", + ) + attack = self._attack_with_identifier(inner_id) + + expected = AtomicAttackEvaluationIdentifier( + build_atomic_attack_identifier(attack_identifier=inner_id), + ).eval_hash + assert compute_inner_attack_eval_hash(attack=attack) == expected + + def test_differs_when_attack_class_differs(self): + from pyrit.models.identifiers import compute_inner_attack_eval_hash + + a = self._attack_with_identifier( + ComponentIdentifier(class_name="A", class_module="m"), + ) + b = self._attack_with_identifier( + ComponentIdentifier(class_name="B", class_module="m"), + ) + assert compute_inner_attack_eval_hash(attack=a) != compute_inner_attack_eval_hash(attack=b) + + def test_stable_across_calls_for_same_attack(self): + from pyrit.models.identifiers import compute_inner_attack_eval_hash + + attack = self._attack_with_identifier( + ComponentIdentifier(class_name="Same", class_module="m"), + ) + assert compute_inner_attack_eval_hash(attack=attack) == compute_inner_attack_eval_hash(attack=attack) + + def test_matches_persisted_row_eval_hash(self): + """Whatever the helper returns, persisting an attack result with the same + identifier must yield an entry with the same eval_hash.""" + from pyrit.memory.memory_models import AttackResultEntry + from pyrit.models import AttackResult + from pyrit.models.identifiers import build_atomic_attack_identifier, compute_inner_attack_eval_hash + + inner_id = ComponentIdentifier( + class_name="MyAttack", + class_module="pyrit.attacks", + ) + attack = self._attack_with_identifier(inner_id) + predicted = compute_inner_attack_eval_hash(attack=attack) + + result = AttackResult( + conversation_id="conv_1", + objective="o", + atomic_attack_identifier=build_atomic_attack_identifier(attack_identifier=inner_id), + ) + entry = AttackResultEntry(entry=result) + assert entry.atomic_attack_identifier["eval_hash"] == predicted + + # --------------------------------------------------------------------------- # OWN_RULE / leaf-entity eval-hash tests # --------------------------------------------------------------------------- diff --git a/tests/unit/models/test_import_boundary.py b/tests/unit/models/test_import_boundary.py index b4c33795e2..2c95c1e59d 100644 --- a/tests/unit/models/test_import_boundary.py +++ b/tests/unit/models/test_import_boundary.py @@ -59,6 +59,9 @@ "pyrit.models.data_type_serializer": { "pyrit.memory": "phase-9", }, + "pyrit.models.identifiers.evaluation_identifier": { + "pyrit.executor.attack.core.attack_strategy": "phase-7", + }, "pyrit.models.scenario_result": { "pyrit.score.scorer_evaluation.scorer_metrics": "phase-7", "pyrit.score.scorer_evaluation.scorer_metrics_io": "phase-7", diff --git a/tests/unit/scenario/scenarios/adaptive/test_dispatcher.py b/tests/unit/scenario/scenarios/adaptive/test_dispatcher.py new file mode 100644 index 0000000000..0ab722f224 --- /dev/null +++ b/tests/unit/scenario/scenarios/adaptive/test_dispatcher.py @@ -0,0 +1,265 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT license. + +from unittest.mock import MagicMock + +import pytest + +from pyrit.executor.attack.compound.sequential_attack import ( + SequenceCompletionPolicy, + SequentialAttack, +) +from pyrit.models import AttackOutcome, SeedAttackGroup, SeedObjective +from pyrit.scenario.scenarios.adaptive.dispatcher import ( + ADAPTIVE_ATTEMPT_LABEL, + AdaptiveTechniqueDispatcher, + TechniqueBundle, +) + + +def _make_bundle(*, name: str, outcomes: list[AttackOutcome], seed_technique=None) -> TechniqueBundle: + """Build a TechniqueBundle whose attack stub yields the given outcomes in order.""" + attack = MagicMock(name=f"attack-{name}") + attack._outcomes = outcomes + attack._name = name + return TechniqueBundle(attack=attack, name=name, seed_technique=seed_technique) + + +class _StubSelector: + """A deterministic selector stub that returns techniques in the order given.""" + + def __init__(self, *, technique_order: list[str]): + self._order = technique_order + + async def select_async( + self, + *, + technique_identifiers, + objective: str, + num_top_techniques: int = 1, + scenario_result_id: str | None = None, + ): + return self._order[:num_top_techniques] + + +@pytest.fixture +def selector(): + return _StubSelector(technique_order=["a", "b", "c"]) + + +@pytest.fixture +def target() -> MagicMock: + return MagicMock(name="objective_target") + + +@pytest.fixture +def seed_group() -> SeedAttackGroup: + return SeedAttackGroup(seeds=[SeedObjective(value="obj")]) + + +class TestDispatcherInit: + @pytest.mark.usefixtures("patch_central_database") + def test_init_rejects_empty_techniques(self, target, selector): + with pytest.raises(ValueError, match="techniques"): + AdaptiveTechniqueDispatcher( + objective_target=target, + techniques={}, + selector=selector, + ) + + @pytest.mark.parametrize("bad_max", [0, -1]) + @pytest.mark.usefixtures("patch_central_database") + def test_init_rejects_invalid_max_attempts(self, target, selector, bad_max): + with pytest.raises(ValueError, match="max_attempts_per_objective"): + AdaptiveTechniqueDispatcher( + objective_target=target, + techniques={"a": _make_bundle(name="a", outcomes=[AttackOutcome.SUCCESS])}, + selector=selector, + max_attempts_per_objective=bad_max, + ) + + +@pytest.mark.usefixtures("patch_central_database") +class TestCompatibleTechniques: + def test_returns_all_when_no_seed_technique(self, target, selector, seed_group): + bundles = { + "a": _make_bundle(name="a", outcomes=[AttackOutcome.SUCCESS]), + "b": _make_bundle(name="b", outcomes=[AttackOutcome.SUCCESS]), + } + dispatcher = AdaptiveTechniqueDispatcher( + objective_target=target, + techniques=bundles, + selector=selector, + ) + assert dispatcher.compatible_techniques(seed_group=seed_group) == ["a", "b"] + + +@pytest.mark.usefixtures("patch_central_database") +class TestBuildAttackAsync: + async def test_builds_sequential_attack(self, target, seed_group): + bundles = { + "a": _make_bundle(name="a", outcomes=[AttackOutcome.SUCCESS]), + "b": _make_bundle(name="b", outcomes=[AttackOutcome.SUCCESS]), + } + selector = _StubSelector(technique_order=["a", "b"]) + dispatcher = AdaptiveTechniqueDispatcher( + objective_target=target, + techniques=bundles, + selector=selector, + max_attempts_per_objective=2, + ) + + attack = await dispatcher.build_attack_async(seed_group=seed_group) + + # plain SequentialAttack (no adaptive subclass) + assert isinstance(attack, SequentialAttack) + assert type(attack) is SequentialAttack + assert len(attack._child_attacks) == 2 + # children in selection order + assert attack._child_attacks[0].strategy is bundles["a"].attack + assert attack._child_attacks[1].strategy is bundles["b"].attack + # 1-based per-attempt label stamped on each child + assert attack._child_attacks[0].memory_labels[ADAPTIVE_ATTEMPT_LABEL] == "1" + assert attack._child_attacks[1].memory_labels[ADAPTIVE_ATTEMPT_LABEL] == "2" + # default policy is FIRST_SUCCESS + assert attack._completion_policy is SequenceCompletionPolicy.FIRST_SUCCESS + + async def test_raises_when_no_compatible_techniques(self, target): + # bundle with an incompatible seed technique + incompatible_technique = MagicMock(name="incompatible_seed_technique") + bundle = _make_bundle(name="a", outcomes=[AttackOutcome.SUCCESS], seed_technique=incompatible_technique) + bundles = {"a": bundle} + selector = _StubSelector(technique_order=["a"]) + dispatcher = AdaptiveTechniqueDispatcher( + objective_target=target, + techniques=bundles, + selector=selector, + ) + seed_group = MagicMock(name="seed_group") + seed_group.objective = MagicMock(value="obj") + seed_group.is_compatible_with_technique.return_value = False + + with pytest.raises(ValueError, match="no compatible techniques"): + await dispatcher.build_attack_async(seed_group=seed_group) + + async def test_raises_when_seed_group_has_no_objective(self, target, selector): + bundles = {"a": _make_bundle(name="a", outcomes=[AttackOutcome.SUCCESS])} + dispatcher = AdaptiveTechniqueDispatcher( + objective_target=target, + techniques=bundles, + selector=selector, + ) + sg = MagicMock(name="seed_group") + sg.objective = None + with pytest.raises(ValueError, match="objective is not initialized"): + await dispatcher.build_attack_async(seed_group=sg) + + async def test_respects_max_attempts(self, target, seed_group): + bundles = { + "a": _make_bundle(name="a", outcomes=[AttackOutcome.SUCCESS]), + "b": _make_bundle(name="b", outcomes=[AttackOutcome.SUCCESS]), + "c": _make_bundle(name="c", outcomes=[AttackOutcome.SUCCESS]), + } + selector = _StubSelector(technique_order=["a", "b", "c"]) + dispatcher = AdaptiveTechniqueDispatcher( + objective_target=target, + techniques=bundles, + selector=selector, + max_attempts_per_objective=2, + ) + + attack = await dispatcher.build_attack_async(seed_group=seed_group) + # selector receives num_top_techniques=2, returns 2 items + assert len(attack._child_attacks) == 2 + + async def test_merges_seed_technique_into_child_seed_group(self, target): + """When a bundle declares a seed_technique it is merged into the seed group for that child.""" + seed_technique = MagicMock(name="seed_technique") + + outer_sg = MagicMock(name="seed_group") + outer_sg.objective = MagicMock(value="obj") + outer_sg.is_compatible_with_technique.return_value = True + merged_sg = MagicMock(name="merged_seed_group") + outer_sg.with_technique.return_value = merged_sg + + bundle = _make_bundle(name="a", outcomes=[AttackOutcome.SUCCESS], seed_technique=seed_technique) + bundles = {"a": bundle} + selector = _StubSelector(technique_order=["a"]) + dispatcher = AdaptiveTechniqueDispatcher( + objective_target=target, + techniques=bundles, + selector=selector, + ) + + attack = await dispatcher.build_attack_async(seed_group=outer_sg) + # The merged seed group is forwarded to the child attack. + outer_sg.with_technique.assert_called_once_with(technique=seed_technique) + assert attack._child_attacks[0].seed_group is merged_sg + + +@pytest.mark.usefixtures("patch_central_database") +class TestEvalHashRoundTrip: + """ + Pin the load-bearing invariant that ``compute_inner_attack_eval_hash`` + (used by ``AdaptiveScenario._build_techniques_dict`` to key the + ``techniques`` dict and by the selector to look up historical stats) + equals the ``eval_hash`` the executor stamps on persisted child rows. + + If the prediction helper and the write path ever drift (e.g. a new + field is added to the eval-hash rule on one side only), the selector + silently reads zero history for every technique and epsilon-greedy + degrades to random with no error. This test runs a real + ``PromptSendingAttack`` through the dispatcher's ``SequentialAttack`` + end-to-end and asserts the round-trip holds. + """ + + async def test_predicted_hash_matches_persisted_row(self, sqlite_instance): + from pyrit.executor.attack.single_turn.prompt_sending import PromptSendingAttack + from pyrit.memory.memory_models import AttackResultEntry + from pyrit.models import SeedAttackGroup, SeedObjective + from pyrit.models.identifiers import compute_inner_attack_eval_hash + from tests.unit.mocks import MockPromptTarget + + live_target = MockPromptTarget() + attack = PromptSendingAttack(objective_target=live_target) + predicted_hash = compute_inner_attack_eval_hash(attack=attack) + + bundles = {predicted_hash: TechniqueBundle(attack=attack, name="prompt_sending")} + dispatcher = AdaptiveTechniqueDispatcher( + objective_target=live_target, + techniques=bundles, + selector=_StubSelector(technique_order=[predicted_hash]), + max_attempts_per_objective=1, + ) + + sg = SeedAttackGroup(seeds=[SeedObjective(value="say hello")]) + sequential = await dispatcher.build_attack_async(seed_group=sg) + await sequential.execute_async(objective="say hello") + + with sqlite_instance.get_session() as session: + rows = session.query(AttackResultEntry).all() + + # Drill into the persisted envelope to find rows whose inner attack is PromptSendingAttack, + # then assert the eval_hash on those rows matches what the selector predicted. + matching_rows = [ + r + for r in rows + if r.atomic_attack_identifier + and r.atomic_attack_identifier.get("children", {}) + .get("attack_technique", {}) + .get("children", {}) + .get("attack", {}) + .get("class_name") + == "PromptSendingAttack" + ] + assert matching_rows, ( + f"Expected at least one persisted row whose inner attack is PromptSendingAttack; " + f"found rows: {[(r.id, r.atomic_attack_identifier) for r in rows]}" + ) + for row in matching_rows: + stamped_hash = row.atomic_attack_identifier["eval_hash"] + assert stamped_hash == predicted_hash, ( + f"Selector-side eval_hash ({predicted_hash}) drifted from executor-stamped " + f"eval_hash ({stamped_hash}) on persisted row {row.id}. " + f"compute_inner_attack_eval_hash and build_atomic_attack_identifier must agree." + ) diff --git a/tests/unit/scenario/scenarios/adaptive/test_epsilon_greedy.py b/tests/unit/scenario/scenarios/adaptive/test_epsilon_greedy.py new file mode 100644 index 0000000000..21144721f5 --- /dev/null +++ b/tests/unit/scenario/scenarios/adaptive/test_epsilon_greedy.py @@ -0,0 +1,168 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT license. + +from unittest.mock import patch + +import pytest + +from pyrit.analytics.result_analysis import AttackStats +from pyrit.scenario.scenarios.adaptive.selectors import ( + EpsilonGreedyTechniqueSelector, + SelectorScope, +) + +TECHNIQUES = ["a", "b", "c", "d"] + +_COMPUTE_PATH = "pyrit.scenario.scenarios.adaptive.selectors.epsilon_greedy.compute_technique_stats" + + +def _seeded_selector(*, epsilon: float = 0.0, random_seed: int = 0) -> EpsilonGreedyTechniqueSelector: + return EpsilonGreedyTechniqueSelector(epsilon=epsilon, random_seed=random_seed) + + +def _empty_rates(*args, **kwargs) -> dict[str, AttackStats]: + """Return empty stats (all techniques unseen).""" + return {} + + +def _rates_with_winner(winner: str, *, successes: int = 5, failures: int = 0): + """Return stats where one technique has a clear win record and others have failures.""" + + def _compute(*args, **kwargs): + stats = {} + total = successes + failures + stats[winner] = AttackStats( + success_rate=successes / total if total else None, + total_decided=total, + successes=successes, + failures=failures, + undetermined=0, + errors=0, + ) + for t in TECHNIQUES: + if t != winner: + stats[t] = AttackStats( + success_rate=0.0, + total_decided=5, + successes=0, + failures=5, + undetermined=0, + errors=0, + ) + return stats + + return _compute + + +class TestEpsilonGreedyTechniqueSelectorInit: + def test_init_defaults(self): + EpsilonGreedyTechniqueSelector() + + @pytest.mark.parametrize("bad_epsilon", [-0.1, 1.1, 2.0, -1.0]) + def test_init_rejects_out_of_range_epsilon(self, bad_epsilon): + with pytest.raises(ValueError, match="epsilon"): + EpsilonGreedyTechniqueSelector(epsilon=bad_epsilon) + + +class TestEpsilonGreedyTechniqueSelectorSelect: + @patch( + "pyrit.scenario.scenarios.adaptive.selectors.epsilon_greedy.compute_technique_stats", + side_effect=_empty_rates, + ) + async def test_select_empty_techniques_raises(self, _mock): + selector = _seeded_selector() + with pytest.raises(ValueError, match="technique_identifiers"): + await selector.select_async(technique_identifiers=[], objective="obj") + + @patch( + "pyrit.scenario.scenarios.adaptive.selectors.epsilon_greedy.compute_technique_stats", + side_effect=_empty_rates, + ) + async def test_select_all_unseen_ties_resolved_randomly(self, _mock): + winners = set() + for s in range(50): + sel = _seeded_selector(random_seed=s) + result = await sel.select_async(technique_identifiers=TECHNIQUES, objective="obj") + winners.add(result[0]) + assert len(winners) > 1 + assert winners.issubset(set(TECHNIQUES)) + + @patch( + "pyrit.scenario.scenarios.adaptive.selectors.epsilon_greedy.compute_technique_stats", + side_effect=_rates_with_winner("b"), + ) + async def test_select_exploits_clear_winner(self, _mock): + selector = _seeded_selector() + for _ in range(20): + result = await selector.select_async(technique_identifiers=TECHNIQUES, objective="obj") + assert result[0] == "b" + + @patch( + "pyrit.scenario.scenarios.adaptive.selectors.epsilon_greedy.compute_technique_stats", + side_effect=_empty_rates, + ) + async def test_select_epsilon_one_is_pure_random(self, _mock): + selector = _seeded_selector(epsilon=1.0) + picks = set() + for i in range(200): + result = await selector.select_async(technique_identifiers=TECHNIQUES, objective=f"obj-{i}") + picks.add(result[0]) + assert picks == set(TECHNIQUES) + + @patch( + "pyrit.scenario.scenarios.adaptive.selectors.epsilon_greedy.compute_technique_stats", + side_effect=_empty_rates, + ) + async def test_select_returns_multiple_techniques(self, _mock): + selector = _seeded_selector() + result = await selector.select_async(technique_identifiers=TECHNIQUES, objective="obj", num_top_techniques=3) + assert len(result) == 3 + assert len(set(result)) == 3 # no duplicates + + @patch( + "pyrit.scenario.scenarios.adaptive.selectors.epsilon_greedy.compute_technique_stats", + side_effect=_empty_rates, + ) + async def test_select_caps_at_available_techniques(self, _mock): + selector = _seeded_selector() + result = await selector.select_async(technique_identifiers=["a", "b"], objective="obj", num_top_techniques=5) + assert len(result) == 2 + + +class TestEpsilonGreedySelectorScope: + @patch(_COMPUTE_PATH, side_effect=_empty_rates) + async def test_default_scope_passes_none_scenario_result_id(self, mock_compute): + selector = _seeded_selector() + await selector.select_async(technique_identifiers=TECHNIQUES, objective="obj", scenario_result_id="run-1") + + # Default scope is all_runs(): the per-call scenario_result_id is dropped. + assert mock_compute.call_args.kwargs["scenario_result_id"] is None + assert mock_compute.call_args.kwargs["targeted_harm_categories"] is None + + @patch(_COMPUTE_PATH, side_effect=_empty_rates) + async def test_current_run_scope_forwards_scenario_result_id(self, mock_compute): + selector = EpsilonGreedyTechniqueSelector(epsilon=0.0, random_seed=0, scope=SelectorScope.current_run()) + await selector.select_async(technique_identifiers=TECHNIQUES, objective="obj", scenario_result_id="run-42") + + assert mock_compute.call_args.kwargs["scenario_result_id"] == "run-42" + + @patch(_COMPUTE_PATH, side_effect=_empty_rates) + async def test_scope_filter_fields_forwarded(self, mock_compute): + scope = SelectorScope( + targeted_harm_categories=["misinformation"], + ) + selector = EpsilonGreedyTechniqueSelector(epsilon=0.0, random_seed=0, scope=scope) + await selector.select_async(technique_identifiers=TECHNIQUES, objective="obj") + + kwargs = mock_compute.call_args.kwargs + assert kwargs["targeted_harm_categories"] == ["misinformation"] + + +class TestEpsilonGreedyEstimate: + def test_estimate_unseen_is_one(self): + assert EpsilonGreedyTechniqueSelector._estimate(technique="a", stats={}) == pytest.approx(1.0) + + def test_estimate_with_data(self): + stats = {"a": AttackStats(success_rate=0.6, total_decided=5, successes=3, failures=2, undetermined=0, errors=0)} + # (3 + 1) / (5 + 1) = 4/6 ≈ 0.6667 + assert EpsilonGreedyTechniqueSelector._estimate(technique="a", stats=stats) == pytest.approx(4 / 6) diff --git a/tests/unit/scenario/scenarios/adaptive/test_selector_scope.py b/tests/unit/scenario/scenarios/adaptive/test_selector_scope.py new file mode 100644 index 0000000000..f024d2c5a5 --- /dev/null +++ b/tests/unit/scenario/scenarios/adaptive/test_selector_scope.py @@ -0,0 +1,55 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT license. + +import dataclasses + +import pytest + +from pyrit.scenario.scenarios.adaptive.selectors import SelectorScope + + +class TestSelectorScopeDefaults: + def test_default_constructs_all_runs(self): + scope = SelectorScope() + assert scope.current_run_only is False + assert scope.targeted_harm_categories is None + + def test_all_runs_classmethod_equivalent_to_default(self): + assert SelectorScope.all_runs() == SelectorScope() + + def test_current_run_classmethod_sets_flag(self): + scope = SelectorScope.current_run() + assert scope.current_run_only is True + assert scope.targeted_harm_categories is None + + +class TestSelectorScopeFrozen: + def test_assigning_field_raises(self): + scope = SelectorScope() + with pytest.raises(dataclasses.FrozenInstanceError): + scope.current_run_only = True # type: ignore[misc] + + def test_assigning_new_field_raises(self): + scope = SelectorScope() + with pytest.raises(dataclasses.FrozenInstanceError): + scope.targeted_harm_categories = ("a",) # type: ignore[misc] + + +class TestSelectorScopeCombinations: + def test_fields_combine(self): + scope = SelectorScope( + current_run_only=True, + targeted_harm_categories=["misinformation"], + ) + assert scope.current_run_only is True + assert scope.targeted_harm_categories == ["misinformation"] + + def test_equality_value_based(self): + a = SelectorScope(targeted_harm_categories=("y",)) + b = SelectorScope(targeted_harm_categories=("y",)) + assert a == b + + def test_inequality_when_fields_differ(self): + a = SelectorScope.all_runs() + b = SelectorScope.current_run() + assert a != b diff --git a/tests/unit/scenario/scenarios/adaptive/test_technique_selector.py b/tests/unit/scenario/scenarios/adaptive/test_technique_selector.py new file mode 100644 index 0000000000..5167cf3310 --- /dev/null +++ b/tests/unit/scenario/scenarios/adaptive/test_technique_selector.py @@ -0,0 +1,13 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT license. + +from pyrit.scenario.scenarios.adaptive.selectors import ( + EpsilonGreedyTechniqueSelector, + TechniqueSelector, +) + + +class TestTechniqueSelectorProtocol: + def test_implements_protocol(self): + selector = EpsilonGreedyTechniqueSelector() + assert isinstance(selector, TechniqueSelector) diff --git a/tests/unit/scenario/scenarios/adaptive/test_text_adaptive.py b/tests/unit/scenario/scenarios/adaptive/test_text_adaptive.py new file mode 100644 index 0000000000..848851ae0a --- /dev/null +++ b/tests/unit/scenario/scenarios/adaptive/test_text_adaptive.py @@ -0,0 +1,553 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT license. + +"""Tests for the ``TextAdaptive`` scenario.""" + +from __future__ import annotations + +import uuid +import warnings +from unittest.mock import MagicMock, patch + +import pytest + +from pyrit.models import SeedAttackGroup, SeedObjective +from pyrit.models.identifiers import ComponentIdentifier +from pyrit.prompt_target import PromptTarget +from pyrit.registry.object_registries.attack_technique_registry import AttackTechniqueRegistry +from pyrit.scenario.core.dataset_configuration import DatasetConfiguration +from pyrit.scenario.core.scenario import BaselineAttackPolicy +from pyrit.scenario.scenarios.adaptive.dispatcher import ( + AdaptiveTechniqueDispatcher, +) +from pyrit.scenario.scenarios.adaptive.text_adaptive import TextAdaptive +from pyrit.score import TrueFalseScorer + +_MOCK_MANY_SHOT_EXAMPLES = [{"question": f"q{i}", "answer": f"a{i}"} for i in range(100)] + + +def _mock_id(name: str) -> ComponentIdentifier: + return ComponentIdentifier(class_name=name, class_module="test") + + +@pytest.fixture +def mock_objective_target() -> MagicMock: + mock = MagicMock(spec=PromptTarget) + mock.get_identifier.return_value = _mock_id("MockObjectiveTarget") + return mock + + +@pytest.fixture +def mock_objective_scorer() -> MagicMock: + mock = MagicMock(spec=TrueFalseScorer) + mock.get_identifier.return_value = _mock_id("MockObjectiveScorer") + return mock + + +@pytest.fixture(autouse=True) +def reset_technique_registry(): + """Reset registries and the cached strategy class between tests.""" + from pyrit.registry import TargetRegistry + + AttackTechniqueRegistry.reset_instance() + TargetRegistry.reset_instance() + TextAdaptive._cached_strategy_class = None + yield + AttackTechniqueRegistry.reset_instance() + TargetRegistry.reset_instance() + TextAdaptive._cached_strategy_class = None + + +@pytest.fixture(autouse=True) +def patch_many_shot_load(): + with patch( + "pyrit.executor.attack.single_turn.many_shot_jailbreak.load_many_shot_jailbreaking_dataset", + return_value=_MOCK_MANY_SHOT_EXAMPLES, + ): + yield + + +@pytest.fixture +def mock_runtime_env(): + with patch.dict( + "os.environ", + { + "OPENAI_CHAT_ENDPOINT": "https://test.openai.azure.com/", + "OPENAI_CHAT_KEY": "test-key", + "OPENAI_CHAT_MODEL": "gpt-4", + }, + ): + yield + + +def _make_seed_group(*, value: str, harm_categories: list[str] | None = None) -> SeedAttackGroup: + return SeedAttackGroup(seeds=[SeedObjective(value=value, harm_categories=harm_categories)]) + + +def _make_fake_factory(*, seed_technique=None, adversarial_chat=None, scoring_config_type=None) -> MagicMock: + """Return a stub attack-technique factory that produces a fake ``AttackTechnique``. + + Mocks the surface ``AdaptiveScenario._build_techniques_dict`` consumes + (``factory.create(...)``, ``factory.adversarial_chat``, and + ``factory.scoring_config_type``). Each call assigns a unique fake + attack identifier (via a fresh UUID) so the bundle dict keys (eval + hashes) don't collide across calls — no shared mutable test state, so + test execution order doesn't shift hash values. + """ + fake_id = uuid.uuid4().hex[:8] + + fake_technique = MagicMock() + fake_attack = MagicMock(name=f"fake-attack-strategy-{fake_id}") + fake_attack.get_identifier.return_value = ComponentIdentifier( + class_name=f"FakeAttack{fake_id}", + class_module="test_text_adaptive", + ) + fake_technique.attack = fake_attack + fake_technique.seed_technique = seed_technique + factory = MagicMock() + factory.create.return_value = fake_technique + factory.adversarial_chat = adversarial_chat + factory.scoring_config_type = scoring_config_type + return factory + + +FIXTURES = ["patch_central_database", "mock_runtime_env"] + + +@pytest.mark.usefixtures(*FIXTURES) +class TestTextAdaptiveBasics: + def test_version(self): + assert TextAdaptive.VERSION == 1 + + def test_baseline_enabled(self): + assert TextAdaptive.BASELINE_ATTACK_POLICY is BaselineAttackPolicy.Enabled + + def test_default_dataset_config(self): + config = TextAdaptive.default_dataset_config() + assert isinstance(config, DatasetConfiguration) + assert config.max_dataset_size == 4 + + def test_required_datasets_non_empty(self): + assert len(TextAdaptive.required_datasets()) > 0 + + def test_get_strategy_class_is_cached(self): + cls_a = TextAdaptive.get_strategy_class() + cls_b = TextAdaptive.get_strategy_class() + assert cls_a is cls_b + + def test_get_default_strategy(self): + strat = TextAdaptive.get_default_strategy() + # The default aggregate must resolve to something runnable. + assert strat is not None + + @patch("pyrit.scenario.core.scenario.Scenario._get_default_objective_scorer") + def test_init_stores_adaptive_params(self, mock_get_scorer, mock_objective_scorer): + mock_get_scorer.return_value = mock_objective_scorer + scenario = TextAdaptive() + scenario.set_params_from_args( + args={ + "max_attempts_per_objective": 7, + } + ) + assert scenario.params["max_attempts_per_objective"] == 7 + + +@pytest.mark.usefixtures(*FIXTURES) +class TestTextAdaptiveAtomicAttacks: + """Tests for ``_get_atomic_attacks_async`` overriding.""" + + async def _build_scenario_and_attacks( + self, + *, + mock_objective_target, + mock_objective_scorer, + seed_groups: dict[str, list[SeedAttackGroup]], + **scenario_kwargs, + ): + with patch.object(DatasetConfiguration, "get_seed_attack_groups", return_value=seed_groups): + scenario = TextAdaptive( + objective_scorer=mock_objective_scorer, + **scenario_kwargs, + ) + await scenario.initialize_async( + objective_target=mock_objective_target, + include_baseline=False, + ) + return scenario, await scenario._get_atomic_attacks_async() + + async def test_one_atomic_per_objective(self, mock_objective_target, mock_objective_scorer): + groups = { + "violence": [ + _make_seed_group(value="obj-v1", harm_categories=["violence"]), + _make_seed_group(value="obj-v2", harm_categories=["violence"]), + ], + "hate": [ + _make_seed_group(value="obj-h1", harm_categories=["hate"]), + ], + } + _scenario, attacks = await self._build_scenario_and_attacks( + mock_objective_target=mock_objective_target, + mock_objective_scorer=mock_objective_scorer, + seed_groups=groups, + ) + # One atomic per objective; each carries exactly one seed group. + assert len(attacks) == 3 + for a in attacks: + assert len(a.seed_groups) == 1 + + async def test_dispatchers_share_one_selector(self, mock_objective_target, mock_objective_scorer): + """All per-dataset dispatchers share one TechniqueSelector instance so + learning accumulates globally (selection is committed up-front but the + selector is still shared by reference across constructions). + """ + selectors_seen: list = [] + real_init = AdaptiveTechniqueDispatcher.__init__ + + def _spy_init(self, *args, **kwargs): + selectors_seen.append(kwargs["selector"]) + return real_init(self, *args, **kwargs) + + groups = { + "violence": [_make_seed_group(value="obj-v1", harm_categories=["violence"])], + "hate": [_make_seed_group(value="obj-h1", harm_categories=["hate"])], + } + with patch.object(DatasetConfiguration, "get_seed_attack_groups", return_value=groups): + scenario = TextAdaptive(objective_scorer=mock_objective_scorer) + await scenario.initialize_async( + objective_target=mock_objective_target, + include_baseline=False, + ) + # Only spy on the explicit invocation so the initialize_async call + # doesn't double-count dispatchers. + with patch.object(AdaptiveTechniqueDispatcher, "__init__", _spy_init): + await scenario._get_atomic_attacks_async() + + # One dispatcher per dataset; all share the same selector identity. + assert len(selectors_seen) == 2 + assert len({id(s) for s in selectors_seen}) == 1 + + async def test_atomic_names_contain_dataset_and_objective_hash(self, mock_objective_target, mock_objective_scorer): + groups = { + "violence": [_make_seed_group(value=f"obj-{i}", harm_categories=["violence"]) for i in range(5)], + "hate": [_make_seed_group(value=f"hate-{i}", harm_categories=["hate"]) for i in range(3)], + } + _scenario, attacks = await self._build_scenario_and_attacks( + mock_objective_target=mock_objective_target, + mock_objective_scorer=mock_objective_scorer, + seed_groups=groups, + ) + names = [atomic.atomic_attack_name for atomic in attacks] + # All names unique; each name embeds its dataset name. + assert len(set(names)) == len(names) == 8 + for atomic in attacks: + assert any(ds in atomic.atomic_attack_name for ds in groups) + + async def test_display_group_is_dataset_name(self, mock_objective_target, mock_objective_scorer): + groups = { + "violence": [_make_seed_group(value="obj-v", harm_categories=["violence"])], + "hate": [_make_seed_group(value="obj-h", harm_categories=["hate"])], + } + _scenario, attacks = await self._build_scenario_and_attacks( + mock_objective_target=mock_objective_target, + mock_objective_scorer=mock_objective_scorer, + seed_groups=groups, + ) + display_groups = {atomic.display_group for atomic in attacks} + assert display_groups == {"violence", "hate"} + + async def test_no_usable_techniques_raises(self, mock_objective_target, mock_objective_scorer): + groups = {"violence": [_make_seed_group(value="obj")]} + with patch.object(DatasetConfiguration, "get_seed_attack_groups", return_value=groups): + scenario = TextAdaptive(objective_scorer=mock_objective_scorer) + await scenario.initialize_async( + objective_target=mock_objective_target, + include_baseline=False, + ) + # Force the factory map to be empty. + with patch.object(scenario, "_get_attack_technique_factories", return_value={}): + with pytest.raises(ValueError, match="no usable techniques"): + await scenario._get_atomic_attacks_async() + + async def test_techniques_with_seed_technique_are_kept(self, mock_objective_target, mock_objective_scorer): + """Factories that declare a ``seed_technique`` participate in the pool + (the old behavior silently dropped them with a warning). + """ + groups = {"violence": [_make_seed_group(value="obj")]} + plain_factory = _make_fake_factory(seed_technique=None) + seeded_factory = _make_fake_factory(seed_technique=MagicMock(name="seed_technique")) + + with ( + patch.object(DatasetConfiguration, "get_seed_attack_groups", return_value=groups), + patch.object(SeedAttackGroup, "is_compatible_with_technique", return_value=True), + ): + scenario = TextAdaptive(objective_scorer=mock_objective_scorer) + strategy_class = scenario.get_strategy_class() + factories = {"role_play": plain_factory, "many_shot": seeded_factory} + with patch.object(scenario, "_get_attack_technique_factories", return_value=factories): + await scenario.initialize_async( + objective_target=mock_objective_target, + include_baseline=False, + scenario_strategies=[strategy_class("role_play"), strategy_class("many_shot")], + ) + attacks = scenario._atomic_attacks + techniques = scenario._build_techniques_dict(objective_target=mock_objective_target) + + # One atomic for the single objective. + assert len(attacks) == 1 + # Both factories survive in the technique pool; in particular the + # seeded one is no longer silently dropped. + technique_names = {b.name for b in techniques.values()} + assert "role_play" in technique_names + assert "many_shot" in technique_names + + async def test_incompatible_seed_technique_is_filtered_per_objective( + self, mock_objective_target, mock_objective_scorer + ): + """When one technique's ``seed_technique`` is incompatible but another + is universally compatible, the objective still produces an atomic that + uses only the compatible technique. + """ + groups = {"violence": [_make_seed_group(value="obj")]} + plain_factory = _make_fake_factory(seed_technique=None) + incompatible_factory = _make_fake_factory(seed_technique=MagicMock(name="incompatible_seed_technique")) + + # Only the plain factory (no seed_technique) is compatible. + with ( + patch.object(DatasetConfiguration, "get_seed_attack_groups", return_value=groups), + patch.object(SeedAttackGroup, "is_compatible_with_technique", return_value=False), + ): + scenario = TextAdaptive(objective_scorer=mock_objective_scorer) + strategy_class = scenario.get_strategy_class() + factories = {"role_play": plain_factory, "many_shot": incompatible_factory} + with patch.object(scenario, "_get_attack_technique_factories", return_value=factories): + await scenario.initialize_async( + objective_target=mock_objective_target, + include_baseline=False, + scenario_strategies=[strategy_class("role_play"), strategy_class("many_shot")], + ) + attacks = scenario._atomic_attacks + techniques = scenario._build_techniques_dict(objective_target=mock_objective_target) + + # Atomic survives because the plain factory keeps the compatible pool non-empty. + assert len(attacks) == 1 + assert len(attacks[0].seed_groups) == 1 + # Both factories live in the pool; per-objective compatibility filtering + # inside the dispatcher (``AdaptiveTechniqueDispatcher.compatible_techniques``) + # then drops the incompatible one before selection. + technique_names = {b.name for b in techniques.values()} + assert "role_play" in technique_names + assert "many_shot" in technique_names + + async def test_objective_skipped_when_no_compatible_techniques( + self, mock_objective_target, mock_objective_scorer, caplog + ): + """When every technique requires an incompatible seed_technique, the + objective is dropped with a warning rather than producing an atomic + attack with an empty technique pool. + """ + groups = { + "violence": [_make_seed_group(value="obj-keep")], + "hate": [_make_seed_group(value="obj-skip")], + } + seeded_factory = _make_fake_factory(seed_technique=MagicMock(name="seed_technique")) + + # is_compatible_with_technique returns True for "obj-keep", False for "obj-skip". + def _selective_compat(self_group, *, technique): + return self_group.objective.value == "obj-keep" + + with ( + patch.object(DatasetConfiguration, "get_seed_attack_groups", return_value=groups), + patch.object(SeedAttackGroup, "is_compatible_with_technique", _selective_compat), + ): + scenario = TextAdaptive(objective_scorer=mock_objective_scorer) + strategy_class = scenario.get_strategy_class() + with patch.object( + scenario, + "_get_attack_technique_factories", + return_value={"role_play": seeded_factory}, + ): + import logging + + with caplog.at_level(logging.WARNING): + await scenario.initialize_async( + objective_target=mock_objective_target, + include_baseline=False, + scenario_strategies=[strategy_class("role_play")], + ) + attacks = scenario._atomic_attacks + + # Only the compatible objective produced an atomic attack. + assert len(attacks) == 1 + # Skip was logged with the affected objective value. + assert any("obj-skip" in record.getMessage() for record in caplog.records) + + async def test_factory_with_narrowed_scoring_config_type_receives_subtype( + self, mock_objective_target, mock_objective_scorer + ): + """When a factory's attack class narrows ``attack_scoring_config`` to a + subtype, the scenario builds and passes that subtype to ``create``.""" + from pyrit.executor.attack import AttackScoringConfig + + class NarrowScoringConfig(AttackScoringConfig): + pass + + groups = {"violence": [_make_seed_group(value="obj")]} + narrow_factory = _make_fake_factory(scoring_config_type=NarrowScoringConfig) + with ( + patch.object(DatasetConfiguration, "get_seed_attack_groups", return_value=groups), + patch.object(SeedAttackGroup, "is_compatible_with_technique", return_value=True), + ): + scenario = TextAdaptive(objective_scorer=mock_objective_scorer) + strategy_class = scenario.get_strategy_class() + with patch.object( + scenario, + "_get_attack_technique_factories", + return_value={"role_play": narrow_factory}, + ): + await scenario.initialize_async( + objective_target=mock_objective_target, + include_baseline=False, + scenario_strategies=[strategy_class("role_play")], + ) + + narrow_factory.create.assert_called_once() + kwargs = narrow_factory.create.call_args.kwargs + passed_config = kwargs["attack_scoring_config"] + assert isinstance(passed_config, NarrowScoringConfig) + assert passed_config.objective_scorer is mock_objective_scorer + + async def test_factory_with_incompatible_narrowed_scoring_config_is_skipped( + self, mock_objective_target, mock_objective_scorer, caplog + ): + """When the narrowed ``attack_scoring_config`` subtype rejects the + scenario's objective scorer, the technique is skipped with a warning + rather than silently falling back to the base config (which could let + a WARN-policy factory substitute its internal default scorer).""" + import logging + + from pyrit.executor.attack import AttackScoringConfig + + class StrictScoringConfig(AttackScoringConfig): + def __init__(self, *, objective_scorer): + raise ValueError("StrictScoringConfig requires FloatScaleThresholdScorer") + + groups = {"violence": [_make_seed_group(value="obj")]} + good_factory = _make_fake_factory() + strict_factory = _make_fake_factory(scoring_config_type=StrictScoringConfig) + + with ( + patch.object(DatasetConfiguration, "get_seed_attack_groups", return_value=groups), + patch.object(SeedAttackGroup, "is_compatible_with_technique", return_value=True), + ): + scenario = TextAdaptive(objective_scorer=mock_objective_scorer) + strategy_class = scenario.get_strategy_class() + factories = {"role_play": good_factory, "tap": strict_factory} + with patch.object(scenario, "_get_attack_technique_factories", return_value=factories): + with caplog.at_level(logging.WARNING): + await scenario.initialize_async( + objective_target=mock_objective_target, + include_baseline=False, + scenario_strategies=[strategy_class("role_play"), strategy_class("tap")], + ) + techniques = scenario._build_techniques_dict(objective_target=mock_objective_target) + + # Strict factory's create is never called — incompatibility surfaces + # before construction, not via the factory's override policy. + strict_factory.create.assert_not_called() + # Only the compatible technique remains in the pool. + technique_names = {b.name for b in techniques.values()} + assert technique_names == {"role_play"} + # The skip reason mentions the required config type so operators can + # diagnose the mismatch. + assert any("tap" in r.getMessage() and "StrictScoringConfig" in r.getMessage() for r in caplog.records) + + async def test_factory_create_failure_skips_technique(self, mock_objective_target, mock_objective_scorer, caplog): + """A factory whose ``create`` raises ``ValueError`` (e.g. the attack + rejects the scenario's objective scorer) is logged and skipped, while + sibling techniques still build successfully. + """ + import logging + + groups = {"violence": [_make_seed_group(value="obj")]} + good_factory = _make_fake_factory() + bad_factory = _make_fake_factory() + bad_factory.create.side_effect = ValueError("requires FloatScaleThresholdScorer") + + with ( + patch.object(DatasetConfiguration, "get_seed_attack_groups", return_value=groups), + patch.object(SeedAttackGroup, "is_compatible_with_technique", return_value=True), + ): + scenario = TextAdaptive(objective_scorer=mock_objective_scorer) + strategy_class = scenario.get_strategy_class() + factories = {"role_play": good_factory, "tap": bad_factory} + with patch.object(scenario, "_get_attack_technique_factories", return_value=factories): + with caplog.at_level(logging.WARNING): + await scenario.initialize_async( + objective_target=mock_objective_target, + include_baseline=False, + scenario_strategies=[strategy_class("role_play"), strategy_class("tap")], + ) + attacks = scenario._atomic_attacks + techniques = scenario._build_techniques_dict(objective_target=mock_objective_target) + + assert len(attacks) == 1 + technique_names = {b.name for b in techniques.values()} + assert technique_names == {"role_play"} + assert any("tap" in r.getMessage() and "Skipping" in r.getMessage() for r in caplog.records) + + async def test_all_factories_failing_raises_with_reason(self, mock_objective_target, mock_objective_scorer): + """When every technique's ``create`` fails, ``_build_techniques_dict`` + raises a ``ValueError`` that names the incompatible technique(s).""" + groups = {"violence": [_make_seed_group(value="obj")]} + bad_factory = _make_fake_factory() + bad_factory.create.side_effect = ValueError("requires FloatScaleThresholdScorer") + + with ( + patch.object(DatasetConfiguration, "get_seed_attack_groups", return_value=groups), + patch.object(SeedAttackGroup, "is_compatible_with_technique", return_value=True), + ): + scenario = TextAdaptive(objective_scorer=mock_objective_scorer) + strategy_class = scenario.get_strategy_class() + with patch.object( + scenario, + "_get_attack_technique_factories", + return_value={"tap": bad_factory}, + ): + with pytest.raises(ValueError, match="incompatible with scenario scorer.*tap"): + await scenario.initialize_async( + objective_target=mock_objective_target, + include_baseline=False, + scenario_strategies=[strategy_class("tap")], + ) + + +@pytest.mark.usefixtures(*FIXTURES) +class TestTextAdaptiveBaselinePolicy: + async def test_initialize_async_accepts_explicit_baseline(self, mock_objective_target, mock_objective_scorer): + groups = {"violence": [_make_seed_group(value="obj", harm_categories=["violence"])]} + with patch.object(DatasetConfiguration, "get_seed_attack_groups", return_value=groups): + scenario = TextAdaptive(objective_scorer=mock_objective_scorer) + # Baseline is Enabled by default, so explicit include_baseline=True must not raise. + await scenario.initialize_async( + objective_target=mock_objective_target, + include_baseline=True, + ) + + async def test_baseline_emitted_at_index_zero_by_default(self, mock_objective_target, mock_objective_scorer): + """ + Under ``BASELINE_ATTACK_POLICY = Enabled`` (the default), the override + must explicitly prepend a baseline atomic at index 0 so the legacy + deprecation-rescue branch in ``Scenario.initialize_async`` (slated + for removal in 0.16.0) is bypassed and no DeprecationWarning fires. + """ + groups = {"violence": [_make_seed_group(value="obj", harm_categories=["violence"])]} + with patch.object(DatasetConfiguration, "get_seed_attack_groups", return_value=groups): + scenario = TextAdaptive(objective_scorer=mock_objective_scorer) + with warnings.catch_warnings(): + warnings.simplefilter("error", DeprecationWarning) + await scenario.initialize_async(objective_target=mock_objective_target) + + assert scenario._atomic_attacks, "expected at least one atomic attack" + assert scenario._atomic_attacks[0].atomic_attack_name == "baseline", ( + f"baseline must be prepended at index 0; got {[a.atomic_attack_name for a in scenario._atomic_attacks]}" + )