From 0acd48ad93a7757c9d9e7d353eedf38eef083158 Mon Sep 17 00:00:00 2001 From: Victor Valbuena Date: Wed, 20 May 2026 13:19:04 -0700 Subject: [PATCH 01/40] new .env_example with adversrial target models. --- .env_example | 198 ++++++++++++++++++++++++++++++++++++--------------- 1 file changed, 139 insertions(+), 59 deletions(-) diff --git a/.env_example b/.env_example index b925fb097c..aafb8b41da 100644 --- a/.env_example +++ b/.env_example @@ -1,115 +1,163 @@ # ============================================================================ + # PyRIT Environment File Example + # ============================================================================ + # -# Copy this file to ~/.pyrit/.env and fill in ONLY the sections you need. + +# Copy this file to ~/.pyrit/.env and fill in ONLY the sections you need + # -# MOST USERS ONLY NEED 3 VARIABLES to get started: + +# MOST USERS ONLY NEED 3 VARIABLES to get started + # -# OPENAI_CHAT_ENDPOINT="https://api.openai.com/v1" # or any OpenAI-compatible API -# OPENAI_CHAT_KEY="your-key-here" -# OPENAI_CHAT_MODEL="gpt-4o" + +# OPENAI_CHAT_ENDPOINT="" # or any OpenAI-compatible API + +# OPENAI_CHAT_KEY="your-key-here" + +# OPENAI_CHAT_MODEL="gpt-4o" + # + # These work with OpenAI, Azure OpenAI, Ollama, Groq, OpenRouter, and any + # other OpenAI-compatible endpoint. See doc/setup/populating_secrets.md -# for provider-specific examples. + +# for provider-specific examples + # -# If you are using Entra authentication for Azure resources, + +# If you are using Entra authentication for Azure resources + # keys for those resources are not needed. PyRIT auto-detects: if an API key -# is set, it uses key auth; otherwise it falls back to Entra ID automatically. + +# is set, it uses key auth; otherwise it falls back to Entra ID automatically + # -# ============================================================================ +# ============================================================================ ################################### + # OPENAI TARGET SECRETS + # + # The below models work with OpenAIChatTarget - either pass via environment variables + # or copy to OPENAI_CHAT_ENDPOINT + ################################### -PLATFORM_OPENAI_CHAT_ENDPOINT="https://api.openai.com/v1" +PLATFORM_OPENAI_CHAT_ENDPOINT="" PLATFORM_OPENAI_CHAT_KEY="sk-xxxxx" PLATFORM_OPENAI_CHAT_GPT4O_MODEL="gpt-4o" # Note: For Azure OpenAI endpoints, use the new format with /openai/v1 and specify the model separately -# Example: https://xxxx.openai.azure.com/openai/v1 -AZURE_OPENAI_GPT4O_ENDPOINT="https://xxxx.openai.azure.com/openai/v1" + +# Example: + +AZURE_OPENAI_GPT4O_ENDPOINT="" AZURE_OPENAI_GPT4O_KEY="xxxxx" AZURE_OPENAI_GPT4O_MODEL="deployment-name" -# Since Azure deployment name may be custom and differ from the actual underlying model, -# you can specify the underlying model for identifier purposes. If not specified, -# identifiers will default to the value of the standard MODEL environment variable. + +# Since Azure deployment name may be custom and differ from the actual underlying model + +# you can specify the underlying model for identifier purposes. If not specified + +# identifiers will default to the value of the standard MODEL environment variable + AZURE_OPENAI_GPT4O_UNDERLYING_MODEL="gpt-4o" -AZURE_OPENAI_INTEGRATION_TEST_ENDPOINT="https://xxxxx.openai.azure.com/openai/v1" +AZURE_OPENAI_INTEGRATION_TEST_ENDPOINT="" AZURE_OPENAI_INTEGRATION_TEST_KEY="xxxxx" AZURE_OPENAI_INTEGRATION_TEST_MODEL="deployment-name" AZURE_OPENAI_INTEGRATION_TEST_UNDERLYING_MODEL="" -AZURE_OPENAI_GPT3_5_CHAT_ENDPOINT="https://xxxxx.openai.azure.com/openai/v1" +AZURE_OPENAI_GPT3_5_CHAT_ENDPOINT="" AZURE_OPENAI_GPT3_5_CHAT_KEY="xxxxx" AZURE_OPENAI_GPT3_5_CHAT_MODEL="deployment-name" AZURE_OPENAI_GPT3_5_CHAT_UNDERLYING_MODEL="" -AZURE_OPENAI_GPT4_CHAT_ENDPOINT="https://xxxxx.openai.azure.com/openai/v1" +AZURE_OPENAI_GPT4_CHAT_ENDPOINT="" AZURE_OPENAI_GPT4_CHAT_KEY="xxxxx" AZURE_OPENAI_GPT4_CHAT_MODEL="deployment-name" AZURE_OPENAI_GPT4_CHAT_UNDERLYING_MODEL="" -AZURE_OPENAI_GPT5_4_ENDPOINT="https://xxxxx.openai.azure.com/openai/v1" +AZURE_OPENAI_GPT5_4_ENDPOINT="" AZURE_OPENAI_GPT5_4_KEY="xxxxx" AZURE_OPENAI_GPT5_4_MODEL="gpt-5.4" AZURE_OPENAI_GPT5_4_UNDERLYING_MODEL="gpt-5.4" # Endpoints that host models with fewer safety mechanisms (e.g. via adversarial fine tuning -# or content filters turned off) can be defined below and used in adversarial attack testing scenarios. -AZURE_OPENAI_GPT4O_UNSAFE_CHAT_ENDPOINT="https://xxxxx.openai.azure.com/openai/v1" + +# or content filters turned off) can be defined below and used in adversarial attack testing scenarios + +AZURE_OPENAI_GPT4O_UNSAFE_CHAT_ENDPOINT="" AZURE_OPENAI_GPT4O_UNSAFE_CHAT_KEY="xxxxx" AZURE_OPENAI_GPT4O_UNSAFE_CHAT_MODEL="deployment-name" AZURE_OPENAI_GPT4O_UNSAFE_CHAT_UNDERLYING_MODEL="" -AZURE_OPENAI_GPT4O_UNSAFE_CHAT_ENDPOINT2="https://xxxxx.openai.azure.com/openai/v1" +AZURE_OPENAI_GPT4O_UNSAFE_CHAT_ENDPOINT2="" AZURE_OPENAI_GPT4O_UNSAFE_CHAT_KEY2="xxxxx" AZURE_OPENAI_GPT4O_UNSAFE_CHAT_MODEL2="deployment-name" AZURE_OPENAI_GPT4O_UNSAFE_CHAT_UNDERLYING_MODEL2="" -# Adversarial chat target (used by scenario attack techniques, e.g. role-play, TAP) -ADVERSARIAL_CHAT_ENDPOINT="https://xxxxx.openai.azure.com/openai/v1" +# Adversarial chat targets (used by scenario attack techniques, e.g. role-play, TAP) + +# Default endpoint goes here; specialized ones below + +ADVERSARIAL_CHAT_ENDPOINT="" ADVERSARIAL_CHAT_KEY="xxxxx" ADVERSARIAL_CHAT_MODEL="deployment-name" +ADVERSARIAL_CHAT_SINGLETURN_ENDPOINT="" +ADVERSARIAL_CHAT_SINGLETURN_KEY="xxxxx" +ADVERSARIAL_CHAT_SINGLETURN_MODEL="deployment-name" + +ADVERSARIAL_CHAT_MULTITURN_ENDPOINT="" +ADVERSARIAL_CHAT_MULTITURN_KEY="xxxxx" +ADVERSARIAL_CHAT_MULTITURN_MODEL="deployment-name" + +ADVERSARIAL_CHAT_REASONING_ENDPOINT="" +ADVERSARIAL_CHAT_REASONING_KEY="xxxxx" +ADVERSARIAL_CHAT_REASONING_MODEL="deployment-name" + # Objective Scorer chat target (used in scorers in scenarios) -OBJECTIVE_SCORER_CHAT_ENDPOINT="https://xxxxx.openai.azure.com/openai/v1" + +OBJECTIVE_SCORER_CHAT_ENDPOINT="" OBJECTIVE_SCORER_CHAT_KEY="xxxxx" OBJECTIVE_SCORER_CHAT_MODEL="deployment-name" -AZURE_FOUNDRY_DEEPSEEK_ENDPOINT="https://xxxxx.eastus2.models.ai.azure.com" +AZURE_FOUNDRY_DEEPSEEK_ENDPOINT="" AZURE_FOUNDRY_DEEPSEEK_KEY="xxxxx" AZURE_FOUNDRY_DEEPSEEK_MODEL="" -AZURE_FOUNDRY_PHI4_ENDPOINT="https://xxxxx.models.ai.azure.com" +AZURE_FOUNDRY_PHI4_ENDPOINT="" AZURE_CHAT_PHI4_KEY="xxxxx" AZURE_FOUNDRY_PHI4_MODEL="" -AZURE_FOUNDRY_MISTRAL_LARGE_ENDPOINT="https://xxxxx.services.ai.azure.com/openai/v1/" +AZURE_FOUNDRY_MISTRAL_LARGE_ENDPOINT="" AZURE_FOUNDRY_MISTRAL_LARGE_KEY="xxxxx" AZURE_FOUNDRY_MISTRAL_LARGE_MODEL="Mistral-Large-3" -AWS_ENDPOINT="https://bedrock-mantle.us-east-1.api.aws/v1" +AWS_ENDPOINT="" AWS_KEY="xxxxx" AWS_CHAT_MODEL="nvidia.nemotron-super-3-120b" AWS_RESPONSES_MODEL="openai.gpt-oss-120b" -GROQ_ENDPOINT="https://api.groq.com/openai/v1" +GROQ_ENDPOINT="" GROQ_KEY="gsk_xxxxxxxx" GROQ_LLAMA_MODEL="llama3-8b-8192" -OPEN_ROUTER_ENDPOINT="https://openrouter.ai/api/v1" +OPEN_ROUTER_ENDPOINT="" OPEN_ROUTER_KEY="sk-or-v1-xxxxx" OPEN_ROUTER_CLAUDE_MODEL="anthropic/claude-3.7-sonnet" -OLLAMA_CHAT_ENDPOINT="http://127.0.0.1:11434/v1" +OLLAMA_CHAT_ENDPOINT="" OLLAMA_MODEL="llama2" DEFAULT_OPENAI_FRONTEND_ENDPOINT = ${AZURE_OPENAI_GPT4O_AAD_ENDPOINT} @@ -119,25 +167,30 @@ DEFAULT_OPENAI_FRONTEND_MODEL = "gpt-4o" OPENAI_CHAT_ENDPOINT=${PLATFORM_OPENAI_CHAT_ENDPOINT} OPENAI_CHAT_KEY=${PLATFORM_OPENAI_CHAT_KEY} OPENAI_CHAT_MODEL=${PLATFORM_OPENAI_CHAT_GPT4O_MODEL} + # The following line can be populated if using an Azure OpenAI deployment + # where the deployment name differs from the actual underlying model + OPENAI_CHAT_UNDERLYING_MODEL="" ################################## + # OPENAI RESPONSES TARGET SECRETS + ################################## -AZURE_OPENAI_GPT5_RESPONSES_ENDPOINT="https://xxxxxxxxx.azure.com/openai/v1" -AZURE_OPENAI_GPT5_COMPLETION_ENDPOINT="https://xxxxxxxxx.azure.com/openai/v1" +AZURE_OPENAI_GPT5_RESPONSES_ENDPOINT="" +AZURE_OPENAI_GPT5_COMPLETION_ENDPOINT="" AZURE_OPENAI_GPT5_KEY="xxxxxxx" AZURE_OPENAI_GPT5_MODEL="gpt-5" AZURE_OPENAI_GPT5_UNDERLYING_MODEL="gpt-5" -PLATFORM_OPENAI_RESPONSES_ENDPOINT="https://api.openai.com/v1" +PLATFORM_OPENAI_RESPONSES_ENDPOINT="" PLATFORM_OPENAI_RESPONSES_KEY="sk-xxxxx" PLATFORM_OPENAI_RESPONSES_MODEL="o4-mini" -AZURE_OPENAI_RESPONSES_ENDPOINT="https://xxxxx.openai.azure.com/openai/v1" +AZURE_OPENAI_RESPONSES_ENDPOINT="" AZURE_OPENAI_RESPONSES_KEY="xxxxx" AZURE_OPENAI_RESPONSES_MODEL="o4-mini" AZURE_OPENAI_RESPONSES_UNDERLYING_MODEL="o4-mini" @@ -148,10 +201,15 @@ OPENAI_RESPONSES_MODEL=${PLATFORM_OPENAI_RESPONSES_MODEL} OPENAI_RESPONSES_UNDERLYING_MODEL="" ################################## + # OPENAI REALTIME TARGET SECRETS + # + # The below models work with RealtimeTarget - either pass via environment variables + # or copy to OPENAI_REALTIME_ENDPOINT + ################################## PLATFORM_OPENAI_REALTIME_ENDPOINT="wss://api.openai.com/v1" @@ -169,18 +227,23 @@ OPENAI_REALTIME_MODEL = ${PLATFORM_OPENAI_REALTIME_MODEL} OPENAI_REALTIME_UNDERLYING_MODEL = "" ################################## + # IMAGE TARGET SECRETS + # + # The below models work with OpenAIImageTarget - either pass via environment variables + # or copy to OPENAI_IMAGE_ENDPOINT + ################################### -OPENAI_IMAGE_ENDPOINT1 = "https://xxxxx.openai.azure.com/openai/v1" +OPENAI_IMAGE_ENDPOINT1 = "" OPENAI_IMAGE_API_KEY1 = "xxxxxx" OPENAI_IMAGE_MODEL1 = "deployment-name" OPENAI_IMAGE_UNDERLYING_MODEL1 = "dall-e-3" -OPENAI_IMAGE_ENDPOINT2 = "https://api.openai.com/v1" +OPENAI_IMAGE_ENDPOINT2 = "" OPENAI_IMAGE_API_KEY2 = "sk-xxxxx" OPENAI_IMAGE_MODEL2 = "dall-e-3" OPENAI_IMAGE_UNDERLYING_MODEL2 = "dall-e-3" @@ -190,20 +253,24 @@ OPENAI_IMAGE_API_KEY = ${OPENAI_IMAGE_API_KEY2} OPENAI_IMAGE_MODEL = ${OPENAI_IMAGE_MODEL2} OPENAI_IMAGE_UNDERLYING_MODEL = "" - ################################## + # TTS TARGET SECRETS + # + # The below models work with OpenAITTSTarget - either pass via environment variables + # or copy to OPENAI_TTS_ENDPOINT + ################################### -OPENAI_TTS_ENDPOINT1 = "https://xxxxx.openai.azure.com/openai/v1" +OPENAI_TTS_ENDPOINT1 = "" OPENAI_TTS_KEY1 = "xxxxxxx" OPENAI_TTS_MODEL1 = "tts" OPENAI_TTS_UNDERLYING_MODEL1 = "tts" -OPENAI_TTS_ENDPOINT2 = "https://api.openai.com/v1" +OPENAI_TTS_ENDPOINT2 = "" OPENAI_TTS_KEY2 = "xxxxxx" OPENAI_TTS_MODEL2 = "tts-1" OPENAI_TTS_UNDERLYING_MODEL2 = "tts-1" @@ -214,14 +281,20 @@ OPENAI_TTS_MODEL = ${OPENAI_TTS_MODEL2} OPENAI_TTS_UNDERLYING_MODEL = "" ################################## + # VIDEO TARGET SECRETS + # + # The below models work with OpenAIVideoTarget - either pass via environment variables + # or copy to OPENAI_VIDEO_ENDPOINT + ################################### # Note: Use the base URL without API path -AZURE_OPENAI_VIDEO_ENDPOINT="https://xxxxx.cognitiveservices.azure.com/openai/v1" + +AZURE_OPENAI_VIDEO_ENDPOINT="" AZURE_OPENAI_VIDEO_KEY="xxxxxxx" AZURE_OPENAI_VIDEO_MODEL="sora-2" AZURE_OPENAI_VIDEO_UNDERLYING_MODEL="sora-2" @@ -231,68 +304,75 @@ OPENAI_VIDEO_KEY = ${AZURE_OPENAI_VIDEO_KEY} OPENAI_VIDEO_MODEL = ${AZURE_OPENAI_VIDEO_MODEL} OPENAI_VIDEO_UNDERLYING_MODEL = "" - ################################## + # AML TARGET SECRETS + # The below models work with AzureMLChatTarget - either pass via environment variables + # or copy to AZURE_ML_MANAGED_ENDPOINT + ################################### -AZURE_ML_PHI_ENDPOINT="https://xxxxxx.westus3.inference.ml.azure.com/score" +AZURE_ML_PHI_ENDPOINT="" AZURE_ML_PHI_KEY="xxxxx" -# The below is set as the default Azure OpenAI model used in most notebooks. Adjust as needed. +# The below is set as the default Azure OpenAI model used in most notebooks. Adjust as needed + AZURE_ML_MANAGED_ENDPOINT=${AZURE_ML_PHI_ENDPOINT} AZURE_ML_KEY=${AZURE_ML_PHI_KEY} - ################################## + # MISC TARGET SECRETS -################################### +################################### -OPENAI_COMPLETION_ENDPOINT="https://xxxxx.openai.azure.com/openai/v1" +OPENAI_COMPLETION_ENDPOINT="" OPENAI_COMPLETION_API_KEY="xxxxx" OPENAI_COMPLETION_MODEL="davinci-002" -OPENAI_EMBEDDING_ENDPOINT="https://xxxxx.openai.azure.com/openai/v1" +OPENAI_EMBEDDING_ENDPOINT="" OPENAI_EMBEDDING_KEY="xxxxx" OPENAI_EMBEDDING_MODEL="text-embedding-3-small" -AZURE_STORAGE_ACCOUNT_CONTAINER_URL="https://xxxxxx.blob.core.windows.net/xpia" +AZURE_STORAGE_ACCOUNT_CONTAINER_URL="" AZURE_STORAGE_ACCOUNT_SAS_TOKEN="xxxxx" - AZURE_SPEECH_REGION = "eastus2" AZURE_SPEECH_KEY = "xxxxx" + # Resource ID is needed when using Entra authentication + AZURE_SPEECH_RESOURCE_ID = "xxxxx" AZURE_CONTENT_SAFETY_API_KEY="xxxxx" -AZURE_CONTENT_SAFETY_API_ENDPOINT="https://xxxxx.cognitiveservices.azure.com/" +AZURE_CONTENT_SAFETY_API_ENDPOINT="" HUGGINGFACE_TOKEN="hf_xxxxxxx" -HUGGINGFACE_ENDPOINT="https://router.huggingface.co/v1" +HUGGINGFACE_ENDPOINT="" -GOOGLE_GEMINI_ENDPOINT = "https://generativelanguage.googleapis.com/v1beta/openai" +GOOGLE_GEMINI_ENDPOINT = "" GOOGLE_GEMINI_API_KEY = "xxxxx" GOOGLE_GEMINI_MODEL="gemini-2.0-flash" - ######################### + # AZURE SQL SECRETS -######################### +######################### # This connects to the test database + AZURE_SQL_DB_CONNECTION_STRING_TEST = "mssql+pyodbc://@xxxxx.database.windows.net/xxxxx?driver=ODBC+Driver+18+for+SQL+Server" -AZURE_STORAGE_ACCOUNT_DB_DATA_CONTAINER_URL_TEST="https://xxxxx.blob.core.windows.net/dbdata" +AZURE_STORAGE_ACCOUNT_DB_DATA_CONTAINER_URL_TEST="" # This connects to the prod database + AZURE_SQL_DB_CONNECTION_STRING_PROD = "mssql+pyodbc://@xxxxx.database.windows.net/xxxxx?driver=ODBC+Driver+18+for+SQL+Server" -AZURE_STORAGE_ACCOUNT_DB_DATA_CONTAINER_URL_PROD="https://xxxxx.blob.core.windows.net/dbdata" +AZURE_STORAGE_ACCOUNT_DB_DATA_CONTAINER_URL_PROD="" +# The below is set as the central memory. Adjust as needed. Recommend overwriting in .env.local -# The below is set as the central memory. Adjust as needed. Recommend overwriting in .env.local. AZURE_SQL_DB_CONNECTION_STRING = ${AZURE_SQL_DB_CONNECTION_STRING_PROD} AZURE_STORAGE_ACCOUNT_DB_DATA_CONTAINER_URL=${AZURE_STORAGE_ACCOUNT_DB_DATA_CONTAINER_URL_PROD} From 6881d19800bbd6c7d7c1034f468b8f4194a86a2f Mon Sep 17 00:00:00 2001 From: Victor Valbuena Date: Wed, 20 May 2026 13:21:59 -0700 Subject: [PATCH 02/40] .env_example formatting --- .env_example | 145 ++++++++++++++------------------------------------- 1 file changed, 40 insertions(+), 105 deletions(-) diff --git a/.env_example b/.env_example index aafb8b41da..c20da007e4 100644 --- a/.env_example +++ b/.env_example @@ -1,163 +1,133 @@ # ============================================================================ - # PyRIT Environment File Example - # ============================================================================ - # - # Copy this file to ~/.pyrit/.env and fill in ONLY the sections you need - # - # MOST USERS ONLY NEED 3 VARIABLES to get started - # - -# OPENAI_CHAT_ENDPOINT="" # or any OpenAI-compatible API - +# OPENAI_CHAT_ENDPOINT="https://api.openai.com/v1" # or any OpenAI-compatible API # OPENAI_CHAT_KEY="your-key-here" - # OPENAI_CHAT_MODEL="gpt-4o" - # - # These work with OpenAI, Azure OpenAI, Ollama, Groq, OpenRouter, and any - # other OpenAI-compatible endpoint. See doc/setup/populating_secrets.md - # for provider-specific examples - # - # If you are using Entra authentication for Azure resources - # keys for those resources are not needed. PyRIT auto-detects: if an API key - # is set, it uses key auth; otherwise it falls back to Entra ID automatically - # - # ============================================================================ - ################################### - # OPENAI TARGET SECRETS - # - # The below models work with OpenAIChatTarget - either pass via environment variables - # or copy to OPENAI_CHAT_ENDPOINT ################################### -PLATFORM_OPENAI_CHAT_ENDPOINT="" +PLATFORM_OPENAI_CHAT_ENDPOINT="https://api.openai.com/v1" PLATFORM_OPENAI_CHAT_KEY="sk-xxxxx" PLATFORM_OPENAI_CHAT_GPT4O_MODEL="gpt-4o" # Note: For Azure OpenAI endpoints, use the new format with /openai/v1 and specify the model separately +# Example: https://xxxx.openai.azure.com/openai/v1 -# Example: - -AZURE_OPENAI_GPT4O_ENDPOINT="" +AZURE_OPENAI_GPT4O_ENDPOINT="https://xxxx.openai.azure.com/openai/v1" AZURE_OPENAI_GPT4O_KEY="xxxxx" AZURE_OPENAI_GPT4O_MODEL="deployment-name" # Since Azure deployment name may be custom and differ from the actual underlying model - # you can specify the underlying model for identifier purposes. If not specified - # identifiers will default to the value of the standard MODEL environment variable AZURE_OPENAI_GPT4O_UNDERLYING_MODEL="gpt-4o" -AZURE_OPENAI_INTEGRATION_TEST_ENDPOINT="" +AZURE_OPENAI_INTEGRATION_TEST_ENDPOINT="https://xxxxx.openai.azure.com/openai/v1" AZURE_OPENAI_INTEGRATION_TEST_KEY="xxxxx" AZURE_OPENAI_INTEGRATION_TEST_MODEL="deployment-name" AZURE_OPENAI_INTEGRATION_TEST_UNDERLYING_MODEL="" -AZURE_OPENAI_GPT3_5_CHAT_ENDPOINT="" +AZURE_OPENAI_GPT3_5_CHAT_ENDPOINT="https://xxxxx.openai.azure.com/openai/v1" AZURE_OPENAI_GPT3_5_CHAT_KEY="xxxxx" AZURE_OPENAI_GPT3_5_CHAT_MODEL="deployment-name" AZURE_OPENAI_GPT3_5_CHAT_UNDERLYING_MODEL="" -AZURE_OPENAI_GPT4_CHAT_ENDPOINT="" +AZURE_OPENAI_GPT4_CHAT_ENDPOINT="https://xxxxx.openai.azure.com/openai/v1" AZURE_OPENAI_GPT4_CHAT_KEY="xxxxx" AZURE_OPENAI_GPT4_CHAT_MODEL="deployment-name" AZURE_OPENAI_GPT4_CHAT_UNDERLYING_MODEL="" -AZURE_OPENAI_GPT5_4_ENDPOINT="" +AZURE_OPENAI_GPT5_4_ENDPOINT="https://xxxxx.openai.azure.com/openai/v1" AZURE_OPENAI_GPT5_4_KEY="xxxxx" AZURE_OPENAI_GPT5_4_MODEL="gpt-5.4" AZURE_OPENAI_GPT5_4_UNDERLYING_MODEL="gpt-5.4" # Endpoints that host models with fewer safety mechanisms (e.g. via adversarial fine tuning - # or content filters turned off) can be defined below and used in adversarial attack testing scenarios -AZURE_OPENAI_GPT4O_UNSAFE_CHAT_ENDPOINT="" +AZURE_OPENAI_GPT4O_UNSAFE_CHAT_ENDPOINT="https://xxxxx.openai.azure.com/openai/v1" AZURE_OPENAI_GPT4O_UNSAFE_CHAT_KEY="xxxxx" AZURE_OPENAI_GPT4O_UNSAFE_CHAT_MODEL="deployment-name" AZURE_OPENAI_GPT4O_UNSAFE_CHAT_UNDERLYING_MODEL="" -AZURE_OPENAI_GPT4O_UNSAFE_CHAT_ENDPOINT2="" +AZURE_OPENAI_GPT4O_UNSAFE_CHAT_ENDPOINT2="https://xxxxx.openai.azure.com/openai/v1" AZURE_OPENAI_GPT4O_UNSAFE_CHAT_KEY2="xxxxx" AZURE_OPENAI_GPT4O_UNSAFE_CHAT_MODEL2="deployment-name" AZURE_OPENAI_GPT4O_UNSAFE_CHAT_UNDERLYING_MODEL2="" # Adversarial chat targets (used by scenario attack techniques, e.g. role-play, TAP) - # Default endpoint goes here; specialized ones below -ADVERSARIAL_CHAT_ENDPOINT="" +ADVERSARIAL_CHAT_ENDPOINT="https://xxxxx.openai.azure.com/openai/v1" ADVERSARIAL_CHAT_KEY="xxxxx" ADVERSARIAL_CHAT_MODEL="deployment-name" -ADVERSARIAL_CHAT_SINGLETURN_ENDPOINT="" +ADVERSARIAL_CHAT_SINGLETURN_ENDPOINT="https://xxxxxx.westus3.inference.ml.azure.com/score" ADVERSARIAL_CHAT_SINGLETURN_KEY="xxxxx" ADVERSARIAL_CHAT_SINGLETURN_MODEL="deployment-name" -ADVERSARIAL_CHAT_MULTITURN_ENDPOINT="" +ADVERSARIAL_CHAT_MULTITURN_ENDPOINT="https://xxxxxx.westus3.inference.ml.azure.com/score" ADVERSARIAL_CHAT_MULTITURN_KEY="xxxxx" ADVERSARIAL_CHAT_MULTITURN_MODEL="deployment-name" -ADVERSARIAL_CHAT_REASONING_ENDPOINT="" +ADVERSARIAL_CHAT_REASONING_ENDPOINT="https://xxxxxx.westus3.inference.ml.azure.com/score" ADVERSARIAL_CHAT_REASONING_KEY="xxxxx" ADVERSARIAL_CHAT_REASONING_MODEL="deployment-name" # Objective Scorer chat target (used in scorers in scenarios) -OBJECTIVE_SCORER_CHAT_ENDPOINT="" +OBJECTIVE_SCORER_CHAT_ENDPOINT="https://xxxxx.openai.azure.com/openai/v1" OBJECTIVE_SCORER_CHAT_KEY="xxxxx" OBJECTIVE_SCORER_CHAT_MODEL="deployment-name" -AZURE_FOUNDRY_DEEPSEEK_ENDPOINT="" +AZURE_FOUNDRY_DEEPSEEK_ENDPOINT="https://xxxxx.eastus2.models.ai.azure.com" AZURE_FOUNDRY_DEEPSEEK_KEY="xxxxx" AZURE_FOUNDRY_DEEPSEEK_MODEL="" -AZURE_FOUNDRY_PHI4_ENDPOINT="" +AZURE_FOUNDRY_PHI4_ENDPOINT="https://xxxxx.models.ai.azure.com" AZURE_CHAT_PHI4_KEY="xxxxx" AZURE_FOUNDRY_PHI4_MODEL="" -AZURE_FOUNDRY_MISTRAL_LARGE_ENDPOINT="" +AZURE_FOUNDRY_MISTRAL_LARGE_ENDPOINT="https://xxxxx.services.ai.azure.com/openai/v1/" AZURE_FOUNDRY_MISTRAL_LARGE_KEY="xxxxx" AZURE_FOUNDRY_MISTRAL_LARGE_MODEL="Mistral-Large-3" -AWS_ENDPOINT="" +AWS_ENDPOINT="https://bedrock-mantle.us-east-1.api.aws/v1" AWS_KEY="xxxxx" AWS_CHAT_MODEL="nvidia.nemotron-super-3-120b" AWS_RESPONSES_MODEL="openai.gpt-oss-120b" -GROQ_ENDPOINT="" +GROQ_ENDPOINT="https://api.groq.com/openai/v1" GROQ_KEY="gsk_xxxxxxxx" GROQ_LLAMA_MODEL="llama3-8b-8192" -OPEN_ROUTER_ENDPOINT="" +OPEN_ROUTER_ENDPOINT="https://openrouter.ai/api/v1" OPEN_ROUTER_KEY="sk-or-v1-xxxxx" OPEN_ROUTER_CLAUDE_MODEL="anthropic/claude-3.7-sonnet" -OLLAMA_CHAT_ENDPOINT="" +OLLAMA_CHAT_ENDPOINT="http://127.0.0.1:11434/v1" OLLAMA_MODEL="llama2" DEFAULT_OPENAI_FRONTEND_ENDPOINT = ${AZURE_OPENAI_GPT4O_AAD_ENDPOINT} @@ -169,28 +139,25 @@ OPENAI_CHAT_KEY=${PLATFORM_OPENAI_CHAT_KEY} OPENAI_CHAT_MODEL=${PLATFORM_OPENAI_CHAT_GPT4O_MODEL} # The following line can be populated if using an Azure OpenAI deployment - # where the deployment name differs from the actual underlying model OPENAI_CHAT_UNDERLYING_MODEL="" ################################## - # OPENAI RESPONSES TARGET SECRETS - ################################## -AZURE_OPENAI_GPT5_RESPONSES_ENDPOINT="" -AZURE_OPENAI_GPT5_COMPLETION_ENDPOINT="" +AZURE_OPENAI_GPT5_RESPONSES_ENDPOINT="https://xxxxxxxxx.azure.com/openai/v1" +AZURE_OPENAI_GPT5_COMPLETION_ENDPOINT="https://xxxxxxxxx.azure.com/openai/v1" AZURE_OPENAI_GPT5_KEY="xxxxxxx" AZURE_OPENAI_GPT5_MODEL="gpt-5" AZURE_OPENAI_GPT5_UNDERLYING_MODEL="gpt-5" -PLATFORM_OPENAI_RESPONSES_ENDPOINT="" +PLATFORM_OPENAI_RESPONSES_ENDPOINT="https://api.openai.com/v1" PLATFORM_OPENAI_RESPONSES_KEY="sk-xxxxx" PLATFORM_OPENAI_RESPONSES_MODEL="o4-mini" -AZURE_OPENAI_RESPONSES_ENDPOINT="" +AZURE_OPENAI_RESPONSES_ENDPOINT="https://xxxxx.openai.azure.com/openai/v1" AZURE_OPENAI_RESPONSES_KEY="xxxxx" AZURE_OPENAI_RESPONSES_MODEL="o4-mini" AZURE_OPENAI_RESPONSES_UNDERLYING_MODEL="o4-mini" @@ -201,15 +168,10 @@ OPENAI_RESPONSES_MODEL=${PLATFORM_OPENAI_RESPONSES_MODEL} OPENAI_RESPONSES_UNDERLYING_MODEL="" ################################## - # OPENAI REALTIME TARGET SECRETS - # - # The below models work with RealtimeTarget - either pass via environment variables - # or copy to OPENAI_REALTIME_ENDPOINT - ################################## PLATFORM_OPENAI_REALTIME_ENDPOINT="wss://api.openai.com/v1" @@ -227,23 +189,18 @@ OPENAI_REALTIME_MODEL = ${PLATFORM_OPENAI_REALTIME_MODEL} OPENAI_REALTIME_UNDERLYING_MODEL = "" ################################## - # IMAGE TARGET SECRETS - # - # The below models work with OpenAIImageTarget - either pass via environment variables - # or copy to OPENAI_IMAGE_ENDPOINT - ################################### -OPENAI_IMAGE_ENDPOINT1 = "" +OPENAI_IMAGE_ENDPOINT1 = "https://xxxxx.openai.azure.com/openai/v1" OPENAI_IMAGE_API_KEY1 = "xxxxxx" OPENAI_IMAGE_MODEL1 = "deployment-name" OPENAI_IMAGE_UNDERLYING_MODEL1 = "dall-e-3" -OPENAI_IMAGE_ENDPOINT2 = "" +OPENAI_IMAGE_ENDPOINT2 = "https://api.openai.com/v1" OPENAI_IMAGE_API_KEY2 = "sk-xxxxx" OPENAI_IMAGE_MODEL2 = "dall-e-3" OPENAI_IMAGE_UNDERLYING_MODEL2 = "dall-e-3" @@ -254,23 +211,18 @@ OPENAI_IMAGE_MODEL = ${OPENAI_IMAGE_MODEL2} OPENAI_IMAGE_UNDERLYING_MODEL = "" ################################## - # TTS TARGET SECRETS - # - # The below models work with OpenAITTSTarget - either pass via environment variables - # or copy to OPENAI_TTS_ENDPOINT - ################################### -OPENAI_TTS_ENDPOINT1 = "" +OPENAI_TTS_ENDPOINT1 = "https://xxxxx.openai.azure.com/openai/v1" OPENAI_TTS_KEY1 = "xxxxxxx" OPENAI_TTS_MODEL1 = "tts" OPENAI_TTS_UNDERLYING_MODEL1 = "tts" -OPENAI_TTS_ENDPOINT2 = "" +OPENAI_TTS_ENDPOINT2 = "https://api.openai.com/v1" OPENAI_TTS_KEY2 = "xxxxxx" OPENAI_TTS_MODEL2 = "tts-1" OPENAI_TTS_UNDERLYING_MODEL2 = "tts-1" @@ -281,20 +233,14 @@ OPENAI_TTS_MODEL = ${OPENAI_TTS_MODEL2} OPENAI_TTS_UNDERLYING_MODEL = "" ################################## - # VIDEO TARGET SECRETS - # - # The below models work with OpenAIVideoTarget - either pass via environment variables - # or copy to OPENAI_VIDEO_ENDPOINT - ################################### - # Note: Use the base URL without API path -AZURE_OPENAI_VIDEO_ENDPOINT="" +AZURE_OPENAI_VIDEO_ENDPOINT="https://xxxxx.cognitiveservices.azure.com/openai/v1" AZURE_OPENAI_VIDEO_KEY="xxxxxxx" AZURE_OPENAI_VIDEO_MODEL="sora-2" AZURE_OPENAI_VIDEO_UNDERLYING_MODEL="sora-2" @@ -305,16 +251,12 @@ OPENAI_VIDEO_MODEL = ${AZURE_OPENAI_VIDEO_MODEL} OPENAI_VIDEO_UNDERLYING_MODEL = "" ################################## - # AML TARGET SECRETS - # The below models work with AzureMLChatTarget - either pass via environment variables - # or copy to AZURE_ML_MANAGED_ENDPOINT - ################################### -AZURE_ML_PHI_ENDPOINT="" +AZURE_ML_PHI_ENDPOINT="https://xxxxxx.westus3.inference.ml.azure.com/score" AZURE_ML_PHI_KEY="xxxxx" # The below is set as the default Azure OpenAI model used in most notebooks. Adjust as needed @@ -323,20 +265,18 @@ AZURE_ML_MANAGED_ENDPOINT=${AZURE_ML_PHI_ENDPOINT} AZURE_ML_KEY=${AZURE_ML_PHI_KEY} ################################## - # MISC TARGET SECRETS - ################################### -OPENAI_COMPLETION_ENDPOINT="" +OPENAI_COMPLETION_ENDPOINT="https://xxxxx.openai.azure.com/openai/v1" OPENAI_COMPLETION_API_KEY="xxxxx" OPENAI_COMPLETION_MODEL="davinci-002" -OPENAI_EMBEDDING_ENDPOINT="" +OPENAI_EMBEDDING_ENDPOINT="https://xxxxx.openai.azure.com/openai/v1" OPENAI_EMBEDDING_KEY="xxxxx" OPENAI_EMBEDDING_MODEL="text-embedding-3-small" -AZURE_STORAGE_ACCOUNT_CONTAINER_URL="" +AZURE_STORAGE_ACCOUNT_CONTAINER_URL="https://xxxxxx.blob.core.windows.net/xpia" AZURE_STORAGE_ACCOUNT_SAS_TOKEN="xxxxx" AZURE_SPEECH_REGION = "eastus2" @@ -347,32 +287,27 @@ AZURE_SPEECH_KEY = "xxxxx" AZURE_SPEECH_RESOURCE_ID = "xxxxx" AZURE_CONTENT_SAFETY_API_KEY="xxxxx" -AZURE_CONTENT_SAFETY_API_ENDPOINT="" +AZURE_CONTENT_SAFETY_API_ENDPOINT="https://xxxxx.cognitiveservices.azure.com/" HUGGINGFACE_TOKEN="hf_xxxxxxx" -HUGGINGFACE_ENDPOINT="" +HUGGINGFACE_ENDPOINT="https://router.huggingface.co/v1" -GOOGLE_GEMINI_ENDPOINT = "" +GOOGLE_GEMINI_ENDPOINT = "https://generativelanguage.googleapis.com/v1beta/openai" GOOGLE_GEMINI_API_KEY = "xxxxx" GOOGLE_GEMINI_MODEL="gemini-2.0-flash" ######################### - # AZURE SQL SECRETS - ######################### # This connects to the test database - AZURE_SQL_DB_CONNECTION_STRING_TEST = "mssql+pyodbc://@xxxxx.database.windows.net/xxxxx?driver=ODBC+Driver+18+for+SQL+Server" -AZURE_STORAGE_ACCOUNT_DB_DATA_CONTAINER_URL_TEST="" +AZURE_STORAGE_ACCOUNT_DB_DATA_CONTAINER_URL_TEST="https://xxxxx.blob.core.windows.net/dbdata" # This connects to the prod database - AZURE_SQL_DB_CONNECTION_STRING_PROD = "mssql+pyodbc://@xxxxx.database.windows.net/xxxxx?driver=ODBC+Driver+18+for+SQL+Server" -AZURE_STORAGE_ACCOUNT_DB_DATA_CONTAINER_URL_PROD="" +AZURE_STORAGE_ACCOUNT_DB_DATA_CONTAINER_URL_PROD="https://xxxxx.blob.core.windows.net/dbdata" # The below is set as the central memory. Adjust as needed. Recommend overwriting in .env.local - AZURE_SQL_DB_CONNECTION_STRING = ${AZURE_SQL_DB_CONNECTION_STRING_PROD} AZURE_STORAGE_ACCOUNT_DB_DATA_CONTAINER_URL=${AZURE_STORAGE_ACCOUNT_DB_DATA_CONTAINER_URL_PROD} From 1f1a07b5a32539590f16d0a99c1a929ed114b1a2 Mon Sep 17 00:00:00 2001 From: Victor Valbuena Date: Wed, 20 May 2026 13:23:51 -0700 Subject: [PATCH 03/40] More .env_example formatting --- .env_example | 55 ++++++++++++++++++++++++++-------------------------- 1 file changed, 27 insertions(+), 28 deletions(-) diff --git a/.env_example b/.env_example index c20da007e4..3ceb6ca51b 100644 --- a/.env_example +++ b/.env_example @@ -2,29 +2,30 @@ # PyRIT Environment File Example # ============================================================================ # -# Copy this file to ~/.pyrit/.env and fill in ONLY the sections you need +# Copy this file to ~/.pyrit/.env and fill in ONLY the sections you need. # -# MOST USERS ONLY NEED 3 VARIABLES to get started +# MOST USERS ONLY NEED 3 VARIABLES to get started: # -# OPENAI_CHAT_ENDPOINT="https://api.openai.com/v1" # or any OpenAI-compatible API -# OPENAI_CHAT_KEY="your-key-here" -# OPENAI_CHAT_MODEL="gpt-4o" +# OPENAI_CHAT_ENDPOINT="https://api.openai.com/v1" # or any OpenAI-compatible API +# OPENAI_CHAT_KEY="your-key-here" +# OPENAI_CHAT_MODEL="gpt-4o" # # These work with OpenAI, Azure OpenAI, Ollama, Groq, OpenRouter, and any # other OpenAI-compatible endpoint. See doc/setup/populating_secrets.md -# for provider-specific examples +# for provider-specific examples. # -# If you are using Entra authentication for Azure resources +# If you are using Entra authentication for Azure resources, # keys for those resources are not needed. PyRIT auto-detects: if an API key -# is set, it uses key auth; otherwise it falls back to Entra ID automatically +# is set, it uses key auth; otherwise it falls back to Entra ID automatically. # # ============================================================================ + + ################################### # OPENAI TARGET SECRETS # # The below models work with OpenAIChatTarget - either pass via environment variables # or copy to OPENAI_CHAT_ENDPOINT - ################################### PLATFORM_OPENAI_CHAT_ENDPOINT="https://api.openai.com/v1" @@ -33,15 +34,12 @@ PLATFORM_OPENAI_CHAT_GPT4O_MODEL="gpt-4o" # Note: For Azure OpenAI endpoints, use the new format with /openai/v1 and specify the model separately # Example: https://xxxx.openai.azure.com/openai/v1 - AZURE_OPENAI_GPT4O_ENDPOINT="https://xxxx.openai.azure.com/openai/v1" AZURE_OPENAI_GPT4O_KEY="xxxxx" AZURE_OPENAI_GPT4O_MODEL="deployment-name" - -# Since Azure deployment name may be custom and differ from the actual underlying model -# you can specify the underlying model for identifier purposes. If not specified -# identifiers will default to the value of the standard MODEL environment variable - +# Since Azure deployment name may be custom and differ from the actual underlying model, +# you can specify the underlying model for identifier purposes. If not specified, +# identifiers will default to the value of the standard MODEL environment variable. AZURE_OPENAI_GPT4O_UNDERLYING_MODEL="gpt-4o" AZURE_OPENAI_INTEGRATION_TEST_ENDPOINT="https://xxxxx.openai.azure.com/openai/v1" @@ -65,8 +63,7 @@ AZURE_OPENAI_GPT5_4_MODEL="gpt-5.4" AZURE_OPENAI_GPT5_4_UNDERLYING_MODEL="gpt-5.4" # Endpoints that host models with fewer safety mechanisms (e.g. via adversarial fine tuning -# or content filters turned off) can be defined below and used in adversarial attack testing scenarios - +# or content filters turned off) can be defined below and used in adversarial attack testing scenarios. AZURE_OPENAI_GPT4O_UNSAFE_CHAT_ENDPOINT="https://xxxxx.openai.azure.com/openai/v1" AZURE_OPENAI_GPT4O_UNSAFE_CHAT_KEY="xxxxx" AZURE_OPENAI_GPT4O_UNSAFE_CHAT_MODEL="deployment-name" @@ -77,9 +74,8 @@ AZURE_OPENAI_GPT4O_UNSAFE_CHAT_KEY2="xxxxx" AZURE_OPENAI_GPT4O_UNSAFE_CHAT_MODEL2="deployment-name" AZURE_OPENAI_GPT4O_UNSAFE_CHAT_UNDERLYING_MODEL2="" -# Adversarial chat targets (used by scenario attack techniques, e.g. role-play, TAP) +# Adversarial chat target (used by scenario attack techniques, e.g. role-play, TAP) # Default endpoint goes here; specialized ones below - ADVERSARIAL_CHAT_ENDPOINT="https://xxxxx.openai.azure.com/openai/v1" ADVERSARIAL_CHAT_KEY="xxxxx" ADVERSARIAL_CHAT_MODEL="deployment-name" @@ -96,8 +92,8 @@ ADVERSARIAL_CHAT_REASONING_ENDPOINT="https://xxxxxx.westus3.inference.ml.azure.c ADVERSARIAL_CHAT_REASONING_KEY="xxxxx" ADVERSARIAL_CHAT_REASONING_MODEL="deployment-name" -# Objective Scorer chat target (used in scorers in scenarios) +# Objective Scorer chat target (used in scorers in scenarios) OBJECTIVE_SCORER_CHAT_ENDPOINT="https://xxxxx.openai.azure.com/openai/v1" OBJECTIVE_SCORER_CHAT_KEY="xxxxx" OBJECTIVE_SCORER_CHAT_MODEL="deployment-name" @@ -137,10 +133,8 @@ DEFAULT_OPENAI_FRONTEND_MODEL = "gpt-4o" OPENAI_CHAT_ENDPOINT=${PLATFORM_OPENAI_CHAT_ENDPOINT} OPENAI_CHAT_KEY=${PLATFORM_OPENAI_CHAT_KEY} OPENAI_CHAT_MODEL=${PLATFORM_OPENAI_CHAT_GPT4O_MODEL} - # The following line can be populated if using an Azure OpenAI deployment # where the deployment name differs from the actual underlying model - OPENAI_CHAT_UNDERLYING_MODEL="" ################################## @@ -210,6 +204,7 @@ OPENAI_IMAGE_API_KEY = ${OPENAI_IMAGE_API_KEY2} OPENAI_IMAGE_MODEL = ${OPENAI_IMAGE_MODEL2} OPENAI_IMAGE_UNDERLYING_MODEL = "" + ################################## # TTS TARGET SECRETS # @@ -238,8 +233,8 @@ OPENAI_TTS_UNDERLYING_MODEL = "" # The below models work with OpenAIVideoTarget - either pass via environment variables # or copy to OPENAI_VIDEO_ENDPOINT ################################### -# Note: Use the base URL without API path +# Note: Use the base URL without API path AZURE_OPENAI_VIDEO_ENDPOINT="https://xxxxx.cognitiveservices.azure.com/openai/v1" AZURE_OPENAI_VIDEO_KEY="xxxxxxx" AZURE_OPENAI_VIDEO_MODEL="sora-2" @@ -250,6 +245,7 @@ OPENAI_VIDEO_KEY = ${AZURE_OPENAI_VIDEO_KEY} OPENAI_VIDEO_MODEL = ${AZURE_OPENAI_VIDEO_MODEL} OPENAI_VIDEO_UNDERLYING_MODEL = "" + ################################## # AML TARGET SECRETS # The below models work with AzureMLChatTarget - either pass via environment variables @@ -259,15 +255,16 @@ OPENAI_VIDEO_UNDERLYING_MODEL = "" AZURE_ML_PHI_ENDPOINT="https://xxxxxx.westus3.inference.ml.azure.com/score" AZURE_ML_PHI_KEY="xxxxx" -# The below is set as the default Azure OpenAI model used in most notebooks. Adjust as needed - +# The below is set as the default Azure OpenAI model used in most notebooks. Adjust as needed. AZURE_ML_MANAGED_ENDPOINT=${AZURE_ML_PHI_ENDPOINT} AZURE_ML_KEY=${AZURE_ML_PHI_KEY} + ################################## # MISC TARGET SECRETS ################################### + OPENAI_COMPLETION_ENDPOINT="https://xxxxx.openai.azure.com/openai/v1" OPENAI_COMPLETION_API_KEY="xxxxx" OPENAI_COMPLETION_MODEL="davinci-002" @@ -279,11 +276,10 @@ OPENAI_EMBEDDING_MODEL="text-embedding-3-small" AZURE_STORAGE_ACCOUNT_CONTAINER_URL="https://xxxxxx.blob.core.windows.net/xpia" AZURE_STORAGE_ACCOUNT_SAS_TOKEN="xxxxx" + AZURE_SPEECH_REGION = "eastus2" AZURE_SPEECH_KEY = "xxxxx" - # Resource ID is needed when using Entra authentication - AZURE_SPEECH_RESOURCE_ID = "xxxxx" AZURE_CONTENT_SAFETY_API_KEY="xxxxx" @@ -296,10 +292,12 @@ GOOGLE_GEMINI_ENDPOINT = "https://generativelanguage.googleapis.com/v1beta/opena GOOGLE_GEMINI_API_KEY = "xxxxx" GOOGLE_GEMINI_MODEL="gemini-2.0-flash" + ######################### # AZURE SQL SECRETS ######################### + # This connects to the test database AZURE_SQL_DB_CONNECTION_STRING_TEST = "mssql+pyodbc://@xxxxx.database.windows.net/xxxxx?driver=ODBC+Driver+18+for+SQL+Server" AZURE_STORAGE_ACCOUNT_DB_DATA_CONTAINER_URL_TEST="https://xxxxx.blob.core.windows.net/dbdata" @@ -308,6 +306,7 @@ AZURE_STORAGE_ACCOUNT_DB_DATA_CONTAINER_URL_TEST="https://xxxxx.blob.core.window AZURE_SQL_DB_CONNECTION_STRING_PROD = "mssql+pyodbc://@xxxxx.database.windows.net/xxxxx?driver=ODBC+Driver+18+for+SQL+Server" AZURE_STORAGE_ACCOUNT_DB_DATA_CONTAINER_URL_PROD="https://xxxxx.blob.core.windows.net/dbdata" -# The below is set as the central memory. Adjust as needed. Recommend overwriting in .env.local + +# The below is set as the central memory. Adjust as needed. Recommend overwriting in .env.local. AZURE_SQL_DB_CONNECTION_STRING = ${AZURE_SQL_DB_CONNECTION_STRING_PROD} AZURE_STORAGE_ACCOUNT_DB_DATA_CONTAINER_URL=${AZURE_STORAGE_ACCOUNT_DB_DATA_CONTAINER_URL_PROD} From 81d529d3b89270edf1a5618728696c69286ba727 Mon Sep 17 00:00:00 2001 From: Victor Valbuena Date: Wed, 20 May 2026 13:32:25 -0700 Subject: [PATCH 04/40] Reorganized unit test directory --- tests/unit/scenario/{ => airt}/test_cyber.py | 0 tests/unit/scenario/{ => airt}/test_jailbreak.py | 0 tests/unit/scenario/{ => airt}/test_leakage_scenario.py | 0 tests/unit/scenario/{ => airt}/test_psychosocial_harms.py | 0 tests/unit/scenario/{ => airt}/test_rapid_response.py | 0 tests/unit/scenario/{ => airt}/test_scam.py | 0 tests/unit/scenario/{ => benchmark}/test_adversarial.py | 0 tests/unit/scenario/{ => core}/test_atomic_attack.py | 0 tests/unit/scenario/{ => core}/test_attack_technique.py | 0 tests/unit/scenario/{ => core}/test_attack_technique_factory.py | 0 tests/unit/scenario/{ => core}/test_baseline_deprecation.py | 0 tests/unit/scenario/{ => core}/test_dataset_configuration.py | 0 tests/unit/scenario/{ => core}/test_scenario.py | 0 tests/unit/scenario/{ => core}/test_scenario_parameters.py | 0 tests/unit/scenario/{ => core}/test_scenario_partial_results.py | 0 tests/unit/scenario/{ => core}/test_scenario_retry.py | 0 .../unit/scenario/{ => core}/test_scenario_strategy_invariants.py | 0 tests/unit/scenario/{ => core}/test_strategy_validation.py | 0 tests/unit/scenario/{ => foundry}/test_foundry.py | 0 tests/unit/scenario/{ => garak}/test_encoding.py | 0 20 files changed, 0 insertions(+), 0 deletions(-) rename tests/unit/scenario/{ => airt}/test_cyber.py (100%) rename tests/unit/scenario/{ => airt}/test_jailbreak.py (100%) rename tests/unit/scenario/{ => airt}/test_leakage_scenario.py (100%) rename tests/unit/scenario/{ => airt}/test_psychosocial_harms.py (100%) rename tests/unit/scenario/{ => airt}/test_rapid_response.py (100%) rename tests/unit/scenario/{ => airt}/test_scam.py (100%) rename tests/unit/scenario/{ => benchmark}/test_adversarial.py (100%) rename tests/unit/scenario/{ => core}/test_atomic_attack.py (100%) rename tests/unit/scenario/{ => core}/test_attack_technique.py (100%) rename tests/unit/scenario/{ => core}/test_attack_technique_factory.py (100%) rename tests/unit/scenario/{ => core}/test_baseline_deprecation.py (100%) rename tests/unit/scenario/{ => core}/test_dataset_configuration.py (100%) rename tests/unit/scenario/{ => core}/test_scenario.py (100%) rename tests/unit/scenario/{ => core}/test_scenario_parameters.py (100%) rename tests/unit/scenario/{ => core}/test_scenario_partial_results.py (100%) rename tests/unit/scenario/{ => core}/test_scenario_retry.py (100%) rename tests/unit/scenario/{ => core}/test_scenario_strategy_invariants.py (100%) rename tests/unit/scenario/{ => core}/test_strategy_validation.py (100%) rename tests/unit/scenario/{ => foundry}/test_foundry.py (100%) rename tests/unit/scenario/{ => garak}/test_encoding.py (100%) diff --git a/tests/unit/scenario/test_cyber.py b/tests/unit/scenario/airt/test_cyber.py similarity index 100% rename from tests/unit/scenario/test_cyber.py rename to tests/unit/scenario/airt/test_cyber.py diff --git a/tests/unit/scenario/test_jailbreak.py b/tests/unit/scenario/airt/test_jailbreak.py similarity index 100% rename from tests/unit/scenario/test_jailbreak.py rename to tests/unit/scenario/airt/test_jailbreak.py diff --git a/tests/unit/scenario/test_leakage_scenario.py b/tests/unit/scenario/airt/test_leakage_scenario.py similarity index 100% rename from tests/unit/scenario/test_leakage_scenario.py rename to tests/unit/scenario/airt/test_leakage_scenario.py diff --git a/tests/unit/scenario/test_psychosocial_harms.py b/tests/unit/scenario/airt/test_psychosocial_harms.py similarity index 100% rename from tests/unit/scenario/test_psychosocial_harms.py rename to tests/unit/scenario/airt/test_psychosocial_harms.py diff --git a/tests/unit/scenario/test_rapid_response.py b/tests/unit/scenario/airt/test_rapid_response.py similarity index 100% rename from tests/unit/scenario/test_rapid_response.py rename to tests/unit/scenario/airt/test_rapid_response.py diff --git a/tests/unit/scenario/test_scam.py b/tests/unit/scenario/airt/test_scam.py similarity index 100% rename from tests/unit/scenario/test_scam.py rename to tests/unit/scenario/airt/test_scam.py diff --git a/tests/unit/scenario/test_adversarial.py b/tests/unit/scenario/benchmark/test_adversarial.py similarity index 100% rename from tests/unit/scenario/test_adversarial.py rename to tests/unit/scenario/benchmark/test_adversarial.py diff --git a/tests/unit/scenario/test_atomic_attack.py b/tests/unit/scenario/core/test_atomic_attack.py similarity index 100% rename from tests/unit/scenario/test_atomic_attack.py rename to tests/unit/scenario/core/test_atomic_attack.py diff --git a/tests/unit/scenario/test_attack_technique.py b/tests/unit/scenario/core/test_attack_technique.py similarity index 100% rename from tests/unit/scenario/test_attack_technique.py rename to tests/unit/scenario/core/test_attack_technique.py diff --git a/tests/unit/scenario/test_attack_technique_factory.py b/tests/unit/scenario/core/test_attack_technique_factory.py similarity index 100% rename from tests/unit/scenario/test_attack_technique_factory.py rename to tests/unit/scenario/core/test_attack_technique_factory.py diff --git a/tests/unit/scenario/test_baseline_deprecation.py b/tests/unit/scenario/core/test_baseline_deprecation.py similarity index 100% rename from tests/unit/scenario/test_baseline_deprecation.py rename to tests/unit/scenario/core/test_baseline_deprecation.py diff --git a/tests/unit/scenario/test_dataset_configuration.py b/tests/unit/scenario/core/test_dataset_configuration.py similarity index 100% rename from tests/unit/scenario/test_dataset_configuration.py rename to tests/unit/scenario/core/test_dataset_configuration.py diff --git a/tests/unit/scenario/test_scenario.py b/tests/unit/scenario/core/test_scenario.py similarity index 100% rename from tests/unit/scenario/test_scenario.py rename to tests/unit/scenario/core/test_scenario.py diff --git a/tests/unit/scenario/test_scenario_parameters.py b/tests/unit/scenario/core/test_scenario_parameters.py similarity index 100% rename from tests/unit/scenario/test_scenario_parameters.py rename to tests/unit/scenario/core/test_scenario_parameters.py diff --git a/tests/unit/scenario/test_scenario_partial_results.py b/tests/unit/scenario/core/test_scenario_partial_results.py similarity index 100% rename from tests/unit/scenario/test_scenario_partial_results.py rename to tests/unit/scenario/core/test_scenario_partial_results.py diff --git a/tests/unit/scenario/test_scenario_retry.py b/tests/unit/scenario/core/test_scenario_retry.py similarity index 100% rename from tests/unit/scenario/test_scenario_retry.py rename to tests/unit/scenario/core/test_scenario_retry.py diff --git a/tests/unit/scenario/test_scenario_strategy_invariants.py b/tests/unit/scenario/core/test_scenario_strategy_invariants.py similarity index 100% rename from tests/unit/scenario/test_scenario_strategy_invariants.py rename to tests/unit/scenario/core/test_scenario_strategy_invariants.py diff --git a/tests/unit/scenario/test_strategy_validation.py b/tests/unit/scenario/core/test_strategy_validation.py similarity index 100% rename from tests/unit/scenario/test_strategy_validation.py rename to tests/unit/scenario/core/test_strategy_validation.py diff --git a/tests/unit/scenario/test_foundry.py b/tests/unit/scenario/foundry/test_foundry.py similarity index 100% rename from tests/unit/scenario/test_foundry.py rename to tests/unit/scenario/foundry/test_foundry.py diff --git a/tests/unit/scenario/test_encoding.py b/tests/unit/scenario/garak/test_encoding.py similarity index 100% rename from tests/unit/scenario/test_encoding.py rename to tests/unit/scenario/garak/test_encoding.py From 6290bbe33b0cb3368d4d47d39a39976cabe90ee1 Mon Sep 17 00:00:00 2001 From: Victor Valbuena Date: Wed, 20 May 2026 13:35:46 -0700 Subject: [PATCH 05/40] Renaming unit tests for consistency --- .../scenario/airt/{test_leakage_scenario.py => test_leakage.py} | 0 .../airt/{test_psychosocial_harms.py => test_psychosocial.py} | 0 .../scenario/foundry/{test_foundry.py => test_red_team_agent.py} | 0 3 files changed, 0 insertions(+), 0 deletions(-) rename tests/unit/scenario/airt/{test_leakage_scenario.py => test_leakage.py} (100%) rename tests/unit/scenario/airt/{test_psychosocial_harms.py => test_psychosocial.py} (100%) rename tests/unit/scenario/foundry/{test_foundry.py => test_red_team_agent.py} (100%) diff --git a/tests/unit/scenario/airt/test_leakage_scenario.py b/tests/unit/scenario/airt/test_leakage.py similarity index 100% rename from tests/unit/scenario/airt/test_leakage_scenario.py rename to tests/unit/scenario/airt/test_leakage.py diff --git a/tests/unit/scenario/airt/test_psychosocial_harms.py b/tests/unit/scenario/airt/test_psychosocial.py similarity index 100% rename from tests/unit/scenario/airt/test_psychosocial_harms.py rename to tests/unit/scenario/airt/test_psychosocial.py diff --git a/tests/unit/scenario/foundry/test_foundry.py b/tests/unit/scenario/foundry/test_red_team_agent.py similarity index 100% rename from tests/unit/scenario/foundry/test_foundry.py rename to tests/unit/scenario/foundry/test_red_team_agent.py From 42887b94e4b71df790b60e217eea9650cf56015b Mon Sep 17 00:00:00 2001 From: Victor Valbuena Date: Thu, 21 May 2026 11:47:24 -0700 Subject: [PATCH 06/40] FIX: TargetInitializer propagates config.tags to registry entries --- .../setup/initializers/components/targets.py | 2 + tests/unit/setup/test_targets_initializer.py | 100 ++++++++++++++++++ 2 files changed, 102 insertions(+) diff --git a/pyrit/setup/initializers/components/targets.py b/pyrit/setup/initializers/components/targets.py index 6340880c8e..093a369a7a 100644 --- a/pyrit/setup/initializers/components/targets.py +++ b/pyrit/setup/initializers/components/targets.py @@ -600,6 +600,8 @@ def _register_target(self, config: TargetConfig) -> None: target = config.target_class(**kwargs) registry = TargetRegistry.get_registry_singleton() registry.register_instance(target, name=config.registry_name) + if config.tags: + registry.add_tags(name=config.registry_name, tags=list(config.tags)) if config.default_objective_target: registry.add_tags(name=config.registry_name, tags=[TargetInitializerTags.DEFAULT_OBJECTIVE_TARGET]) logger.info(f"Registered target: {config.registry_name}") diff --git a/tests/unit/setup/test_targets_initializer.py b/tests/unit/setup/test_targets_initializer.py index e35decc77c..6c5e8a1226 100644 --- a/tests/unit/setup/test_targets_initializer.py +++ b/tests/unit/setup/test_targets_initializer.py @@ -400,3 +400,103 @@ async def test_other_targets_not_tagged_as_default(self) -> None: assert config.default_objective_target is False, ( f"Target {config.registry_name} should not have default_objective_target=True" ) + + +@pytest.mark.usefixtures("patch_central_database") +class TestTargetInitializerConfigTagPropagation: + """Tests for TargetInitializer propagating ``TargetConfig.tags`` to the registry (F1c).""" + + def setup_method(self) -> None: + """Reset registry before each test.""" + TargetRegistry.reset_instance() + + def teardown_method(self) -> None: + """Clean up after each test.""" + TargetRegistry.reset_instance() + for var in [ + "ADVERSARIAL_CHAT_ENDPOINT", + "ADVERSARIAL_CHAT_KEY", + "ADVERSARIAL_CHAT_MODEL", + "OPENAI_CHAT_ENDPOINT", + "OPENAI_CHAT_KEY", + "OPENAI_CHAT_MODEL", + ]: + os.environ.pop(var, None) + + async def test_register_target_propagates_config_tags(self) -> None: + """ + ``TargetConfig.tags`` should be added to the registry entry so the entire + ``TargetInitializerTags`` enum is queryable post-registration. + """ + from pyrit.setup.initializers.components.targets import TargetInitializerTags + + os.environ["ADVERSARIAL_CHAT_ENDPOINT"] = "https://test.openai.azure.com" + os.environ["ADVERSARIAL_CHAT_KEY"] = "test_key" + os.environ["ADVERSARIAL_CHAT_MODEL"] = "gpt-4o" + + init = TargetInitializer() + await init.initialize_async() + + registry = TargetRegistry.get_registry_singleton() + assert "adversarial_chat" in registry + + adversarial_entries = registry.get_by_tag(tag=TargetInitializerTags.ADVERSARIAL) + assert any(entry.name == "adversarial_chat" for entry in adversarial_entries), ( + "adversarial_chat should be discoverable by the ADVERSARIAL tag after F1c" + ) + + default_entries = registry.get_by_tag(tag=TargetInitializerTags.DEFAULT) + assert any(entry.name == "adversarial_chat" for entry in default_entries), ( + "adversarial_chat declares both DEFAULT and ADVERSARIAL tags; both must propagate" + ) + + async def test_register_target_no_tags_in_config_no_extra_add_tags(self) -> None: + """An empty ``config.tags`` list must not trigger an ``add_tags`` call (no spurious empty-list passes).""" + from unittest.mock import MagicMock, patch + + from pyrit.setup.initializers.components.targets import TargetConfig, TargetInitializer + + config = TargetConfig( + registry_name="empty_tags_target", + target_class=MagicMock(return_value=MagicMock()), + endpoint_var="EMPTY_TAGS_ENDPOINT", + key_var="", + tags=[], + ) + + os.environ["EMPTY_TAGS_ENDPOINT"] = "https://example.com" + + try: + mock_registry = MagicMock() + with patch.object(TargetRegistry, "get_registry_singleton", return_value=mock_registry): + init = TargetInitializer() + init._register_target(config) + + mock_registry.register_instance.assert_called_once() + mock_registry.add_tags.assert_not_called() + finally: + os.environ.pop("EMPTY_TAGS_ENDPOINT", None) + + async def test_register_target_default_objective_tag_still_applied(self) -> None: + """ + Regression: ``default_objective_target=True`` must still add the ``DEFAULT_OBJECTIVE_TARGET`` + tag alongside any ``config.tags``. + """ + from pyrit.setup.initializers.components.targets import TargetInitializerTags + + os.environ["OPENAI_CHAT_ENDPOINT"] = "https://api.openai.com/v1" + os.environ["OPENAI_CHAT_KEY"] = "test_key" + os.environ["OPENAI_CHAT_MODEL"] = "gpt-4o" + + init = TargetInitializer() + await init.initialize_async() + + registry = TargetRegistry.get_registry_singleton() + default_objective_entries = registry.get_by_tag(tag=TargetInitializerTags.DEFAULT_OBJECTIVE_TARGET) + assert len(default_objective_entries) == 1 + assert default_objective_entries[0].name == "openai_chat" + + default_entries = registry.get_by_tag(tag=TargetInitializerTags.DEFAULT) + assert any(entry.name == "openai_chat" for entry in default_entries), ( + "openai_chat's config.tags=[DEFAULT] must propagate even when default_objective_target=True" + ) From b7af0deefe9344d60393c3fa76a3daedd2460410 Mon Sep 17 00:00:00 2001 From: Victor Valbuena Date: Thu, 21 May 2026 11:55:27 -0700 Subject: [PATCH 07/40] FEAT: Add TargetRegistry.get_by_tag_query for TagQuery-based lookup --- .../object_registries/target_registry.py | 29 +++++++++ tests/unit/registry/test_target_registry.py | 61 +++++++++++++++++++ 2 files changed, 90 insertions(+) diff --git a/pyrit/registry/object_registries/target_registry.py b/pyrit/registry/object_registries/target_registry.py index c6fefd3926..000979482f 100644 --- a/pyrit/registry/object_registries/target_registry.py +++ b/pyrit/registry/object_registries/target_registry.py @@ -12,9 +12,11 @@ import logging from typing import TYPE_CHECKING, Optional, Union +from pyrit.registry.object_registries.base_instance_registry import RegistryEntry from pyrit.registry.object_registries.retrievable_instance_registry import ( RetrievableInstanceRegistry, ) +from pyrit.registry.tag_query import TagQuery if TYPE_CHECKING: from pyrit.prompt_target import PromptTarget @@ -74,3 +76,30 @@ def get_instance_by_name(self, name: str) -> Optional[PromptTarget]: The target instance, or None if not found. """ return self.get(name) + + def get_by_tag_query(self, *, query: TagQuery) -> list[RegistryEntry[PromptTarget]]: + """ + Get all entries whose tag keys satisfy ``query``. + + ``TagQuery`` operates on a tag set, so this method matches against + ``entry.tags.keys()`` and ignores tag values. For value-aware + single-tag lookups use ``get_by_tag(*, tag, value)`` on the base + class. + + Composite queries compose with ``&`` and ``|`` operators, e.g. + ``TagQuery.all("adversarial") & TagQuery.any_of("singleturn", "multiturn")``. + + Args: + query: The tag predicate to evaluate against each entry. + + Returns: + List of matching ``RegistryEntry`` objects sorted by registry name. + """ + results: list[RegistryEntry[PromptTarget]] = [] + # Note: this erases insertion order, but respects the base_instance_registry pattern + # (get_by_tag). + for name in sorted(self._registry_items.keys()): + entry = self._registry_items[name] + if query.matches(set(entry.tags.keys())): + results.append(entry) + return results diff --git a/tests/unit/registry/test_target_registry.py b/tests/unit/registry/test_target_registry.py index ce612f5203..ee9aa64f97 100644 --- a/tests/unit/registry/test_target_registry.py +++ b/tests/unit/registry/test_target_registry.py @@ -230,3 +230,64 @@ def test_list_metadata_filter_by_class_name(self): assert len(mock_metadata) == 2 for m in mock_metadata: assert m.class_name == "MockPromptTarget" + + +@pytest.mark.usefixtures("patch_central_database") +class TestTargetRegistryGetByTagQuery: + """Tests for ``TargetRegistry.get_by_tag_query`` (TagQuery-aware tag lookup).""" + + def setup_method(self): + """Reset and populate a fresh registry for each test.""" + TargetRegistry.reset_instance() + self.registry = TargetRegistry.get_registry_singleton() + + self.registry.register_instance(MockPromptTarget(), name="adv_single", tags=["adversarial", "singleturn"]) + self.registry.register_instance(MockPromptTarget(), name="adv_multi", tags=["adversarial", "multiturn"]) + self.registry.register_instance(MockPromptChatTarget(), name="scorer_only", tags=["scorer"]) + self.registry.register_instance(MockPromptTarget(), name="untagged") + + def teardown_method(self): + """Reset the singleton after each test.""" + TargetRegistry.reset_instance() + + def test_get_by_tag_query_returns_matching(self): + """A leaf ``TagQuery.all`` returns every entry whose tag set contains the required tag.""" + from pyrit.registry.tag_query import TagQuery + + results = self.registry.get_by_tag_query(query=TagQuery.all("adversarial")) + + names = [entry.name for entry in results] + assert names == ["adv_multi", "adv_single"] + + def test_get_by_tag_query_empty(self): + """A query that matches no entries returns an empty list (not raise).""" + from pyrit.registry.tag_query import TagQuery + + results = self.registry.get_by_tag_query(query=TagQuery.all("nonexistent_tag")) + assert results == [] + + def test_get_by_tag_query_composite_and_or(self): + """Composite queries via ``&`` / ``|`` evaluate as expected.""" + from pyrit.registry.tag_query import TagQuery + + query = TagQuery.all("adversarial") & TagQuery.any_of("singleturn", "multiturn") + results = self.registry.get_by_tag_query(query=query) + + names = [entry.name for entry in results] + assert names == ["adv_multi", "adv_single"] + + narrower = TagQuery.all("adversarial") & TagQuery.any_of("singleturn") + narrow_names = [entry.name for entry in self.registry.get_by_tag_query(query=narrower)] + assert narrow_names == ["adv_single"] + + def test_get_by_tag_query_matches_keys_not_values(self): + """``TagQuery`` evaluates against tag keys; tag values are ignored by this method.""" + from pyrit.registry.tag_query import TagQuery + + self.registry.add_tags(name="adv_single", tags={"priority": "high"}) + + priority_matches = self.registry.get_by_tag_query(query=TagQuery.all("priority")) + assert [entry.name for entry in priority_matches] == ["adv_single"] + + value_lookup = self.registry.get_by_tag_query(query=TagQuery.all("high")) + assert value_lookup == [] From 190ee4a08051db6115122efe91cd4df658c98d7e Mon Sep 17 00:00:00 2001 From: Victor Valbuena Date: Thu, 21 May 2026 12:11:26 -0700 Subject: [PATCH 08/40] FEAT: Register ADVERSARIAL_CHAT singleturn/multiturn/reasoning variants Adds three TargetConfig entries (singleturn, multiturn, reasoning), each tagged [DEFAULT, ADVERSARIAL], for the env-driven variants already declared in .env_example. Tightens _register_target to skip with a warning when a TargetConfig declares model_var but the env var is unset; without this guard the target silently falls back to the global OPENAI_CHAT_MODEL default and sends requests to the wrong model. New tests covering naming-related failure modes flagged for the eventual PR review: - test_register_instance_with_duplicate_name_silently_overwrites pins the current "second write wins" behavior so future hardening (warn / raise / idempotent skip) is intentional. - test_target_configs_have_unique_registry_names guards against typos in ENV_TARGET_CONFIGS that would otherwise silently drop a target. - test_double_initialize_async_is_idempotent regression-guards the re-init path that depends on the silent-overwrite semantics above. - test_variant_skips_when_model_env_var_missing parameterizes the missing-_MODEL skip+warning for all three new variants. Failure modes surfaced during this change but not addressed here (tracked for the PR description batch): - Duplicate registry_name silently overwrites in BaseInstanceRegistry. - registry_name has no format validation; risk grows with per-user TargetConfig support in P1. - No-adversarial-models-found error message UX is owned by the upcoming BenchmarkInitializer commit and needs a clear, actionable message. --- .../setup/initializers/components/targets.py | 39 +++++ tests/unit/registry/test_target_registry.py | 19 ++ tests/unit/setup/test_targets_initializer.py | 162 ++++++++++++++++++ 3 files changed, 220 insertions(+) diff --git a/pyrit/setup/initializers/components/targets.py b/pyrit/setup/initializers/components/targets.py index 093a369a7a..9e85d7d8b1 100644 --- a/pyrit/setup/initializers/components/targets.py +++ b/pyrit/setup/initializers/components/targets.py @@ -189,6 +189,33 @@ class TargetConfig: temperature=1.2, tags=[TargetInitializerTags.DEFAULT, TargetInitializerTags.ADVERSARIAL], ), + TargetConfig( + registry_name="adversarial_chat_singleturn", + target_class=OpenAIChatTarget, + endpoint_var="ADVERSARIAL_CHAT_SINGLETURN_ENDPOINT", + key_var="ADVERSARIAL_CHAT_SINGLETURN_KEY", + model_var="ADVERSARIAL_CHAT_SINGLETURN_MODEL", + temperature=1.2, + tags=[TargetInitializerTags.DEFAULT, TargetInitializerTags.ADVERSARIAL], + ), + TargetConfig( + registry_name="adversarial_chat_multiturn", + target_class=OpenAIChatTarget, + endpoint_var="ADVERSARIAL_CHAT_MULTITURN_ENDPOINT", + key_var="ADVERSARIAL_CHAT_MULTITURN_KEY", + model_var="ADVERSARIAL_CHAT_MULTITURN_MODEL", + temperature=1.2, + tags=[TargetInitializerTags.DEFAULT, TargetInitializerTags.ADVERSARIAL], + ), + TargetConfig( + registry_name="adversarial_chat_reasoning", + target_class=OpenAIChatTarget, + endpoint_var="ADVERSARIAL_CHAT_REASONING_ENDPOINT", + key_var="ADVERSARIAL_CHAT_REASONING_KEY", + model_var="ADVERSARIAL_CHAT_REASONING_MODEL", + temperature=1.2, + tags=[TargetInitializerTags.DEFAULT, TargetInitializerTags.ADVERSARIAL], + ), TargetConfig( registry_name="objective_scorer_chat", target_class=OpenAIChatTarget, @@ -573,6 +600,18 @@ def _register_target(self, config: TargetConfig) -> None: model_name = os.getenv(config.model_var) if config.model_var else None underlying_model = os.getenv(config.underlying_model_var) if config.underlying_model_var else None + # Guard against silent fallback to a global OPENAI_CHAT_MODEL default when the + # declared per-config model env var is unset. Without this skip, the target + # registers cleanly but sends requests to the wrong model at runtime. + if config.model_var and not model_name: + logger.warning( + "Skipping target '%s': %s is not set. " + "All declared env vars (endpoint, key, model) must be present for this target to register.", + config.registry_name, + config.model_var, + ) + return + # Build kwargs for the target constructor kwargs: dict[str, Any] = { "endpoint": endpoint, diff --git a/tests/unit/registry/test_target_registry.py b/tests/unit/registry/test_target_registry.py index ee9aa64f97..b8c9234b88 100644 --- a/tests/unit/registry/test_target_registry.py +++ b/tests/unit/registry/test_target_registry.py @@ -138,6 +138,25 @@ def test_register_instance_same_target_type_different_config(self): assert len(self.registry) == 2 + def test_register_instance_with_duplicate_name_silently_overwrites(self): + """Characterization: re-registering an existing name silently replaces the prior entry. + + BaseInstanceRegistry.register is plain dict assignment; there is no + collision check, warning, or error. This test pins the current behavior + so any future tightening (warn, raise, idempotent skip) is an + intentional decision rather than a silent regression. Tracked as + ``duplicate-registry-name`` in failure_mode_followups for the PR + review batch. + """ + first = MockPromptTarget(model_name="first") + second = MockPromptTarget(model_name="second") + + self.registry.register_instance(first, name="same_name") + self.registry.register_instance(second, name="same_name") + + assert len(self.registry) == 1 + assert self.registry.get("same_name") is second + @pytest.mark.usefixtures("patch_central_database") class TestTargetRegistryGetInstanceByName: diff --git a/tests/unit/setup/test_targets_initializer.py b/tests/unit/setup/test_targets_initializer.py index 6c5e8a1226..831f5e16f8 100644 --- a/tests/unit/setup/test_targets_initializer.py +++ b/tests/unit/setup/test_targets_initializer.py @@ -211,6 +211,23 @@ def test_expected_targets_in_configs(self): assert "groq" in registry_names assert "google_gemini" in registry_names + def test_target_configs_have_unique_registry_names(self): + """Guard against typos: every ``registry_name`` in ``ENV_TARGET_CONFIGS`` must be unique. + + Duplicate names would silently overwrite each other when + ``TargetInitializer`` registers them (per ``BaseInstanceRegistry.register`` + semantics, characterized in ``test_target_registry.py``). Only the + second entry would survive in the registry, which breaks downstream + fan-out (``BenchmarkInitializer``) and is hard to diagnose. Tracked + as ``duplicate-registry-name`` in failure_mode_followups. + """ + registry_names = [config.registry_name for config in TARGET_CONFIGS] + seen: dict[str, int] = {} + for name in registry_names: + seen[name] = seen.get(name, 0) + 1 + duplicates = {name: count for name, count in seen.items() if count > 1} + assert not duplicates, f"Duplicate registry_name(s) in TARGET_CONFIGS: {duplicates}" + class TestTargetInitializerGetInfo: """Tests for TargetInitializer.get_info_async method.""" @@ -500,3 +517,148 @@ async def test_register_target_default_objective_tag_still_applied(self) -> None assert any(entry.name == "openai_chat" for entry in default_entries), ( "openai_chat's config.tags=[DEFAULT] must propagate even when default_objective_target=True" ) + + +ADVERSARIAL_CHAT_VARIANTS: list[tuple[str, str]] = [ + ("adversarial_chat_singleturn", "ADVERSARIAL_CHAT_SINGLETURN"), + ("adversarial_chat_multiturn", "ADVERSARIAL_CHAT_MULTITURN"), + ("adversarial_chat_reasoning", "ADVERSARIAL_CHAT_REASONING"), +] + + +@pytest.mark.usefixtures("patch_central_database") +class TestTargetInitializerAdversarialChatVariants: + """Tests for the ``ADVERSARIAL_CHAT_{SINGLETURN,MULTITURN,REASONING}_*`` env-driven variants.""" + + def setup_method(self) -> None: + """Reset registry and clear variant env vars.""" + TargetRegistry.reset_instance() + self._clear_variant_env_vars() + + def teardown_method(self) -> None: + """Reset registry and clear variant env vars.""" + TargetRegistry.reset_instance() + self._clear_variant_env_vars() + + @staticmethod + def _clear_variant_env_vars() -> None: + for _, prefix in ADVERSARIAL_CHAT_VARIANTS: + for suffix in ("ENDPOINT", "KEY", "MODEL"): + os.environ.pop(f"{prefix}_{suffix}", None) + + @staticmethod + def _set_variant_env_vars(prefix: str) -> None: + os.environ[f"{prefix}_ENDPOINT"] = "https://variant.openai.azure.com/openai/v1" + os.environ[f"{prefix}_KEY"] = "test_key" + os.environ[f"{prefix}_MODEL"] = "deployment-name" + + @pytest.mark.parametrize(("registry_name", "env_prefix"), ADVERSARIAL_CHAT_VARIANTS) + async def test_variant_registers_with_default_and_adversarial_tags( + self, registry_name: str, env_prefix: str + ) -> None: + """Each variant registers with ``[DEFAULT, ADVERSARIAL]`` tags when its env vars are set.""" + from pyrit.setup.initializers.components.targets import TargetInitializerTags + + self._set_variant_env_vars(env_prefix) + + init = TargetInitializer() + await init.initialize_async() + + registry = TargetRegistry.get_registry_singleton() + assert registry_name in registry + + adversarial_entries = registry.get_by_tag(tag=TargetInitializerTags.ADVERSARIAL) + assert any(entry.name == registry_name for entry in adversarial_entries) + + default_entries = registry.get_by_tag(tag=TargetInitializerTags.DEFAULT) + assert any(entry.name == registry_name for entry in default_entries) + + @pytest.mark.parametrize(("registry_name", "env_prefix"), ADVERSARIAL_CHAT_VARIANTS) + async def test_variant_skips_when_env_vars_missing(self, registry_name: str, env_prefix: str) -> None: + """Variants skip gracefully when their env vars are missing (matches existing adversarial_chat behavior).""" + init = TargetInitializer() + await init.initialize_async() + + registry = TargetRegistry.get_registry_singleton() + assert registry_name not in registry + + @pytest.mark.parametrize(("registry_name", "env_prefix"), ADVERSARIAL_CHAT_VARIANTS) + async def test_variant_skips_when_model_env_var_missing( + self, registry_name: str, env_prefix: str, caplog: pytest.LogCaptureFixture + ) -> None: + """Endpoint+key set but _MODEL unset must skip with a warning, not silently fall back to OPENAI_CHAT_MODEL.""" + import logging + + os.environ[f"{env_prefix}_ENDPOINT"] = "https://variant.openai.azure.com/openai/v1" + os.environ[f"{env_prefix}_KEY"] = "test_key" + + try: + with caplog.at_level(logging.WARNING, logger="pyrit.setup.initializers.components.targets"): + init = TargetInitializer() + await init.initialize_async() + + registry = TargetRegistry.get_registry_singleton() + assert registry_name not in registry + + captured_messages = [r.message for r in caplog.records] + assert any(f"{env_prefix}_MODEL" in m for m in captured_messages), ( + f"Expected a warning naming the missing {env_prefix}_MODEL env var; got: {captured_messages}" + ) + finally: + os.environ.pop(f"{env_prefix}_ENDPOINT", None) + os.environ.pop(f"{env_prefix}_KEY", None) + + async def test_all_variants_discoverable_via_adversarial_tag_query(self) -> None: + """End-to-end: variants + ``adversarial_chat`` are returned by adversarial-tag ``get_by_tag_query``.""" + from pyrit.registry.tag_query import TagQuery + + os.environ["ADVERSARIAL_CHAT_ENDPOINT"] = "https://parent.openai.azure.com/openai/v1" + os.environ["ADVERSARIAL_CHAT_KEY"] = "test_key" + os.environ["ADVERSARIAL_CHAT_MODEL"] = "deployment-name" + + for _, prefix in ADVERSARIAL_CHAT_VARIANTS: + self._set_variant_env_vars(prefix) + + try: + init = TargetInitializer() + await init.initialize_async() + + registry = TargetRegistry.get_registry_singleton() + matches = registry.get_by_tag_query(query=TagQuery.all("adversarial")) + match_names = {entry.name for entry in matches} + + expected = {"adversarial_chat"} | {name for name, _ in ADVERSARIAL_CHAT_VARIANTS} + assert expected <= match_names, ( + f"Missing variants from tag query result. Expected superset: {expected}, got: {match_names}" + ) + finally: + for var in ("ADVERSARIAL_CHAT_ENDPOINT", "ADVERSARIAL_CHAT_KEY", "ADVERSARIAL_CHAT_MODEL"): + os.environ.pop(var, None) + + async def test_double_initialize_async_is_idempotent(self) -> None: + """Re-running ``initialize_async`` with the same env state produces the same registry contents. + + Regression guard for the duplicate-registration silent-overwrite path: + because env vars haven't changed between calls, the rebuilt entries + carry identical configuration. If anyone introduces non-idempotent + side-effects (e.g. tag accumulation, instance leaks) into + ``_register_target``, this test will catch it. Tracked as + ``duplicate-registry-name`` in failure_mode_followups. + """ + from pyrit.setup.initializers.components.targets import TargetInitializerTags + + for _, prefix in ADVERSARIAL_CHAT_VARIANTS: + self._set_variant_env_vars(prefix) + + init = TargetInitializer() + await init.initialize_async() + registry = TargetRegistry.get_registry_singleton() + first_names = sorted(registry.get_names()) + first_adversarial_count = len(registry.get_by_tag(tag=TargetInitializerTags.ADVERSARIAL)) + + await init.initialize_async() + second_names = sorted(registry.get_names()) + second_adversarial_count = len(registry.get_by_tag(tag=TargetInitializerTags.ADVERSARIAL)) + + assert first_names == second_names + assert first_adversarial_count == second_adversarial_count From c1476bf4060ae91ba7d18d51c8d2683db873a7eb Mon Sep 17 00:00:00 2001 From: Victor Valbuena Date: Thu, 21 May 2026 12:32:00 -0700 Subject: [PATCH 09/40] FEAT: Add BenchmarkInitializer for adversarial-target fan-out MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Registers one AttackTechniqueSpec variant per (adversarial-capable technique, adversarial-tagged target) pair into AttackTechniqueRegistry with the live target bound onto adversarial_chat. Variants are named f"{source}__{target_name}" and tagged ["benchmark_fanout", f"model:{target_name}"] so the benchmark scenario can discover them via tag query in a later commit. Adversarial-capability is determined by reusing _spec_needs_adversarial from scenario_techniques (multi-turn attacks + crescendo-style simulated conversations). Single-turn techniques without an adversarial chat target (prompt_sending, role_play, many_shot, context_compliance) are not fanned — the benchmark holds the objective target constant and varies the adversarial chat helper across runs. Placed at pyrit/setup/initializers/benchmark.py (top level) alongside AIRTInitializer and SimpleInitializer, not under components/. The components/ initializers (TargetInitializer, ScorerInitializer, ScenarioTechniqueInitializer) are auto-bundled building blocks that populate their registries during every PyRIT setup. BenchmarkInitializer is the opposite shape: a user-opted workflow profile named after the use case, listed in .pyrit_conf when the user wants a benchmarking trial. The placement convention is itself underspecified in the codebase and is tracked for follow-up. Parameter contract (target_names: list[str] | None): The optional target_names parameter is declared via PyRITInitializer.supported_parameters, which is the single source of truth shared across three consumer sites: 1. .pyrit_conf YAML: initializers: - name: benchmark args: target_names: - adversarial_chat_singleturn - adversarial_chat_reasoning ConfigurationLoader._resolve_initializers calls instance.set_params_from_args(args=config.args) and then _validate_params against supported_parameters, so unknown keys fail fast at config-load time. Omitting the args block uses the default (fan over every adversarial-tagged target). 2. CLI (--list-initializers via frontend_core._print_initializer_meta): reads metadata.supported_parameters and prints name + description + default for each declared parameter, so users discover what they can put in .pyrit_conf without reading source. 3. GUI backend (InitializerService): wraps each declared parameter as an InitializerParameterSummary({name, description, default}) on the RegisteredInitializer Pydantic model. The GUI renders form fields from this metadata. All three paths terminate at the same self.params dict that initialize_async reads via self.params.get("target_names"), so adding, renaming, or retyping the parameter is a single-site change. target_names narrows fan-out to a subset of adversarial targets by registry name; unknown names raise ValueError listing both the unknowns and the discovered set. Empty discovery raises ValueError naming the ADVERSARIAL_CHAT_* env vars and the TargetInitializer ordering dependency (closes one of the failure-mode follow-ups surfaced in the previous commit). Failure modes audited during this change but not addressed here (tracked for the PR description batch): - AttackTechniqueRegistry.register_from_specs is first-write-wins on name collision with no log entry. Disjoint name spaces between ScenarioTechniqueInitializer and BenchmarkInitializer mean this is inert today; future extensions that produce colliding names would be silently no-op'd. - BenchmarkInitializer's TargetRegistry walk is a snapshot at init time; later mutations to TargetRegistry leave the fanned specs holding stale references. Failure surfaces at API-call time, not at registration. - Top-level vs components/ initializer placement convention is implicit; this commit picks "top-level for workflow profiles", matching AIRT/Simple. Worth a CONTRIBUTING note when convention is formalized. --- pyrit/setup/initializers/__init__.py | 2 + pyrit/setup/initializers/benchmark.py | 198 ++++++++++++++++++ .../unit/setup/test_benchmark_initializer.py | 181 ++++++++++++++++ 3 files changed, 381 insertions(+) create mode 100644 pyrit/setup/initializers/benchmark.py create mode 100644 tests/unit/setup/test_benchmark_initializer.py diff --git a/pyrit/setup/initializers/__init__.py b/pyrit/setup/initializers/__init__.py index 84aeb83a49..b9e77c4038 100644 --- a/pyrit/setup/initializers/__init__.py +++ b/pyrit/setup/initializers/__init__.py @@ -6,6 +6,7 @@ from pyrit.common.deprecation import print_deprecation_message from pyrit.common.parameter import Parameter from pyrit.setup.initializers.airt import AIRTInitializer +from pyrit.setup.initializers.benchmark import BenchmarkInitializer from pyrit.setup.initializers.components.scenarios import ScenarioTechniqueInitializer from pyrit.setup.initializers.components.scorers import ScorerInitializer from pyrit.setup.initializers.components.targets import TargetInitializer @@ -18,6 +19,7 @@ "Parameter", "PyRITInitializer", "AIRTInitializer", + "BenchmarkInitializer", "ScenarioTechniqueInitializer", "ScorerInitializer", "TargetInitializer", diff --git a/pyrit/setup/initializers/benchmark.py b/pyrit/setup/initializers/benchmark.py new file mode 100644 index 0000000000..3ecbce9dbc --- /dev/null +++ b/pyrit/setup/initializers/benchmark.py @@ -0,0 +1,198 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT license. + +""" +Benchmark initializer that fans adversarial-capable scenario techniques across +adversarial targets discovered in ``TargetRegistry``. + +This is the entry point for bootstrapping an ``AdversarialBenchmark`` trial. +It queries ``TargetRegistry`` for entries tagged ``ADVERSARIAL`` (via +``TagQuery.all("adversarial")``), then for every adversarial-capable +technique in ``SCENARIO_TECHNIQUES`` builds one fanned +``AttackTechniqueSpec`` per discovered target. Each fanned spec binds the +live target onto ``adversarial_chat`` and is registered into +``AttackTechniqueRegistry`` tagged ``["benchmark_fanout", f"model:{name}"]`` +so the benchmark scenario can discover them via tag query in a later commit. + +The ``target_names`` parameter (optional, settable from ``.pyrit_conf``) +narrows the fan-out to a specific subset of adversarial targets by registry +name. Unknown names raise ``ValueError``. + +Discovery returning no adversarial-tagged targets raises ``ValueError`` with +an actionable message pointing at the ``ADVERSARIAL_CHAT_*`` env vars and +the ``TargetInitializer`` dependency. +""" + +import dataclasses +import logging + +from pyrit.common.parameter import Parameter +from pyrit.registry import TargetRegistry +from pyrit.registry.object_registries.attack_technique_registry import ( + AttackTechniqueRegistry, + AttackTechniqueSpec, +) +from pyrit.registry.tag_query import TagQuery +from pyrit.scenario.core.scenario_techniques import SCENARIO_TECHNIQUES, _spec_needs_adversarial +from pyrit.setup.initializers.components.targets import TargetInitializerTags +from pyrit.setup.initializers.pyrit_initializer import PyRITInitializer + +logger = logging.getLogger(__name__) + + +#: Default discovery query used when no ``target_names`` override is provided. +#: Resolves every ``TargetRegistry`` entry tagged ``ADVERSARIAL`` (which today +#: includes ``adversarial_chat`` plus the ``ADVERSARIAL_CHAT_{SINGLETURN, +#: MULTITURN,REASONING}`` variants — all set by ``TargetInitializer``). +DEFAULT_ADVERSARIAL_TAG_QUERY: TagQuery = TagQuery.all(TargetInitializerTags.ADVERSARIAL.value) + + +class BenchmarkInitializer(PyRITInitializer): + """ + Fan adversarial-capable scenario techniques across discovered adversarial targets. + + For every ``AttackTechniqueSpec`` in ``SCENARIO_TECHNIQUES`` that uses an + adversarial chat target (multi-turn attacks plus crescendo-style + simulated conversations), this initializer registers one variant per + discovered adversarial target with the target bound onto + ``adversarial_chat``. Variants are named + ``f"{source_spec.name}__{target_name}"`` (e.g. ``red_teaming__adversarial_chat_singleturn``) + and carry the additional strategy tags ``"benchmark_fanout"`` and + ``f"model:{target_name}"`` so the benchmark scenario can query the + registry by tag in a later commit. + + Parameters (declared via :attr:`supported_parameters`): + + * ``target_names`` (``list[str]``, optional): Narrow fan-out to a + specific subset of adversarial targets by registry name. When omitted, + every target matching :data:`DEFAULT_ADVERSARIAL_TAG_QUERY` is used. + + Raises (at ``initialize_async``): + + * ``ValueError`` — no adversarial-tagged targets are registered. The + error names the ``ADVERSARIAL_CHAT_*`` env vars to set and the + ``TargetInitializer`` dependency. + * ``ValueError`` — any name in ``target_names`` does not match a + discovered adversarial-tagged target. The error lists discovered names. + + Prerequisites: ``TargetInitializer`` must have run first so adversarial + env-driven targets are present in ``TargetRegistry``. Registering the + base scenario-technique catalog (``ScenarioTechniqueInitializer`` or an + equivalent caller of ``register_scenario_techniques``) is also expected + if users will select non-benchmark strategies in the same session; + ``BenchmarkInitializer`` itself only registers the fanned variants. + Per-name idempotent via ``AttackTechniqueRegistry.register_from_specs``: + running the initializer twice with the same registry state is a no-op. + """ + + @property + def supported_parameters(self) -> list[Parameter]: + """Declare the optional ``target_names`` narrowing parameter.""" + return [ + Parameter( + name="target_names", + description=( + "Optional list of adversarial target registry names to narrow benchmark fan-out. " + 'When omitted, every target matching TagQuery.all("adversarial") is used.' + ), + default=None, + param_type=list[str], + ), + ] + + async def initialize_async(self) -> None: + """ + Discover adversarial targets and register fanned specs into the technique registry. + + Raises: + ValueError: If no adversarial-tagged targets are registered in + ``TargetRegistry``, or if ``self.params['target_names']`` + contains a name not in the discovered set. + """ + target_registry = TargetRegistry.get_registry_singleton() + discovered_entries = target_registry.get_by_tag_query(query=DEFAULT_ADVERSARIAL_TAG_QUERY) + if not discovered_entries: + raise ValueError( + "BenchmarkInitializer: no adversarial-tagged targets registered in TargetRegistry. " + "Set ADVERSARIAL_CHAT_* env vars (see .env_example) and ensure TargetInitializer runs " + "before BenchmarkInitializer (e.g. via .pyrit_conf initializer ordering)." + ) + + selected_entries = self._narrow_by_target_names(discovered_entries=discovered_entries) + + fanned_specs = self._build_fanned_specs(target_entries=selected_entries) + + attack_registry = AttackTechniqueRegistry.get_registry_singleton() + attack_registry.register_from_specs(fanned_specs) + + logger.info( + "BenchmarkInitializer: registered %d fanned spec(s) across %d adversarial target(s): %s", + len(fanned_specs), + len(selected_entries), + ", ".join(entry.name for entry in selected_entries), + ) + + def _narrow_by_target_names(self, *, discovered_entries: list) -> list: + """ + Filter ``discovered_entries`` to the names in ``self.params['target_names']``, if set. + + Args: + discovered_entries: The full set of adversarial-tagged registry entries. + + Returns: + list: ``discovered_entries`` unchanged when no ``target_names`` param + is set, otherwise the subset whose ``name`` is in the requested set. + + Raises: + ValueError: If any name in ``self.params['target_names']`` is not + present in ``discovered_entries``. + """ + target_names_param = self.params.get("target_names") + if not target_names_param: + return discovered_entries + + requested = set(target_names_param) + discovered_names = {entry.name for entry in discovered_entries} + unknown = requested - discovered_names + if unknown: + raise ValueError( + f"BenchmarkInitializer: unknown target_names {sorted(unknown)}. " + f"Discovered adversarial targets: {sorted(discovered_names)}." + ) + return [entry for entry in discovered_entries if entry.name in requested] + + def _build_fanned_specs(self, *, target_entries: list) -> list[AttackTechniqueSpec]: + """ + Build fanned ``AttackTechniqueSpec``s for every (adversarial-capable technique, target) pair. + + Adversarial-capability is determined by ``_spec_needs_adversarial`` + (re-used from ``scenario_techniques``): a spec needs an adversarial + chat target when its attack class accepts ``attack_adversarial_config`` + or its ``seed_technique`` has a simulated conversation. Non-adversarial + techniques (e.g. ``prompt_sending``, ``role_play``) are skipped — the + benchmark holds the objective target constant and varies the + adversarial chat helper across runs. + + Args: + target_entries: The adversarial-tagged registry entries to fan over. + + Returns: + list[AttackTechniqueSpec]: One fanned spec per (adversarial-capable + technique, target entry) pair, with the live target bound onto + ``adversarial_chat`` and benchmark-specific strategy tags appended. + """ + fanned: list[AttackTechniqueSpec] = [] + for source_spec in SCENARIO_TECHNIQUES: + if not _spec_needs_adversarial(source_spec): + continue + fanned.extend( + dataclasses.replace( + source_spec, + name=f"{source_spec.name}__{entry.name}", + adversarial_chat=entry.instance, + adversarial_chat_key=None, + strategy_tags=[*source_spec.strategy_tags, "benchmark_fanout", f"model:{entry.name}"], + ) + for entry in target_entries + ) + return fanned diff --git a/tests/unit/setup/test_benchmark_initializer.py b/tests/unit/setup/test_benchmark_initializer.py new file mode 100644 index 0000000000..8a04578abb --- /dev/null +++ b/tests/unit/setup/test_benchmark_initializer.py @@ -0,0 +1,181 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT license. + +"""Tests for BenchmarkInitializer.""" + +from unittest.mock import MagicMock + +import pytest + +from pyrit.prompt_target import PromptTarget +from pyrit.registry import TargetRegistry +from pyrit.registry.object_registries.attack_technique_registry import AttackTechniqueRegistry +from pyrit.setup.initializers import BenchmarkInitializer +from pyrit.setup.initializers.benchmark import DEFAULT_ADVERSARIAL_TAG_QUERY +from pyrit.setup.initializers.components.targets import TargetInitializerTags + + +@pytest.fixture(autouse=True) +def reset_registries(): + """Reset technique and target registries between tests.""" + AttackTechniqueRegistry.reset_instance() + TargetRegistry.reset_instance() + yield + AttackTechniqueRegistry.reset_instance() + TargetRegistry.reset_instance() + + +def _register_adversarial_target(*, name: str) -> PromptTarget: + """Register a mock adversarial-tagged target and return the instance.""" + target = MagicMock(spec=PromptTarget) + target.capabilities.includes.return_value = True + registry = TargetRegistry.get_registry_singleton() + registry.register_instance(target, name=name, tags=[TargetInitializerTags.ADVERSARIAL.value]) + return target + + +class TestBenchmarkInitializerBasic: + """Class metadata tests.""" + + def test_can_be_created(self): + init = BenchmarkInitializer() + assert init is not None + + def test_required_env_vars_is_empty(self): + """Initializer takes no required env vars; discovery happens via TargetRegistry.""" + init = BenchmarkInitializer() + assert init.required_env_vars == [] + + def test_supported_parameters_declares_target_names(self): + init = BenchmarkInitializer() + names = [p.name for p in init.supported_parameters] + assert "target_names" in names + + def test_default_adversarial_tag_query_matches_adversarial_only(self): + """The default discovery query is exactly ``TagQuery.all("adversarial")``.""" + assert DEFAULT_ADVERSARIAL_TAG_QUERY.matches({"adversarial"}) + assert not DEFAULT_ADVERSARIAL_TAG_QUERY.matches({"default"}) + assert not DEFAULT_ADVERSARIAL_TAG_QUERY.matches(set()) + + +class TestBenchmarkInitializerFanOut: + """Tests for the fan-out registration behavior.""" + + async def test_fans_out_one_spec_per_target_per_adversarial_technique(self): + """N targets * M adversarial-capable techniques = N*M fanned specs in the attack registry.""" + _register_adversarial_target(name="adv_a") + _register_adversarial_target(name="adv_b") + + init = BenchmarkInitializer() + await init.initialize_async() + + attack_registry = AttackTechniqueRegistry.get_registry_singleton() + fanned = attack_registry.get_by_tag(tag="benchmark_fanout") + assert len(fanned) > 0 + assert len(fanned) % 2 == 0, "Expected an even count: every adversarial technique fanned across both targets" + + fanned_names = {entry.name for entry in fanned} + for name in fanned_names: + assert "__" in name, f"Fanned spec name '{name}' missing '__' separator" + + async def test_fanned_spec_names_use_source_double_underscore_target(self): + """Spec naming contract: ``f'{source_spec.name}__{target_name}'``.""" + _register_adversarial_target(name="adv_single") + + init = BenchmarkInitializer() + await init.initialize_async() + + attack_registry = AttackTechniqueRegistry.get_registry_singleton() + fanned = attack_registry.get_by_tag(tag="model:adv_single") + assert len(fanned) > 0 + for entry in fanned: + assert entry.name.endswith("__adv_single") + source_name = entry.name.split("__", 1)[0] + assert source_name and "__" not in source_name + + async def test_fanned_specs_carry_benchmark_and_model_tags(self): + """Each fanned spec is tagged ``benchmark_fanout`` plus ``f'model:{name}'``.""" + _register_adversarial_target(name="adv_a") + _register_adversarial_target(name="adv_b") + + init = BenchmarkInitializer() + await init.initialize_async() + + attack_registry = AttackTechniqueRegistry.get_registry_singleton() + for entry in attack_registry.get_by_tag(tag="benchmark_fanout"): + assert "benchmark_fanout" in entry.tags + model_tags = [tag for tag in entry.tags if tag.startswith("model:")] + assert len(model_tags) == 1, f"Expected exactly one model:* tag on {entry.name}, got {model_tags}" + assert model_tags[0] in ("model:adv_a", "model:adv_b") + + async def test_registration_is_idempotent_across_re_init(self): + """Re-running initialize_async produces the same registry state (per-name idempotent).""" + _register_adversarial_target(name="adv_a") + + init = BenchmarkInitializer() + await init.initialize_async() + attack_registry = AttackTechniqueRegistry.get_registry_singleton() + first_count = len(attack_registry.get_by_tag(tag="benchmark_fanout")) + + await init.initialize_async() + second_count = len(attack_registry.get_by_tag(tag="benchmark_fanout")) + + assert first_count == second_count + + +class TestBenchmarkInitializerTargetNamesNarrowing: + """Tests for the optional ``target_names`` parameter.""" + + async def test_target_names_narrows_to_subset(self): + """When ``target_names`` is set, only those entries are fanned.""" + _register_adversarial_target(name="adv_a") + _register_adversarial_target(name="adv_b") + _register_adversarial_target(name="adv_c") + + init = BenchmarkInitializer() + init.params = {"target_names": ["adv_a", "adv_c"]} + await init.initialize_async() + + attack_registry = AttackTechniqueRegistry.get_registry_singleton() + model_b_specs = attack_registry.get_by_tag(tag="model:adv_b") + assert model_b_specs == [] + + model_a_specs = attack_registry.get_by_tag(tag="model:adv_a") + model_c_specs = attack_registry.get_by_tag(tag="model:adv_c") + assert len(model_a_specs) > 0 + assert len(model_c_specs) > 0 + + async def test_target_names_unknown_raises_with_discovered_list(self): + """Unknown ``target_names`` raise ``ValueError`` naming both the unknowns and the discovered set.""" + _register_adversarial_target(name="adv_a") + + init = BenchmarkInitializer() + init.params = {"target_names": ["nonexistent"]} + + with pytest.raises(ValueError, match=r"nonexistent.*adv_a"): + await init.initialize_async() + + async def test_empty_target_names_param_falls_back_to_default_query(self): + """An empty ``target_names`` list is treated as "no narrowing" (same as omitting it).""" + _register_adversarial_target(name="adv_a") + + init = BenchmarkInitializer() + init.params = {"target_names": []} + await init.initialize_async() + + attack_registry = AttackTechniqueRegistry.get_registry_singleton() + assert len(attack_registry.get_by_tag(tag="model:adv_a")) > 0 + + +class TestBenchmarkInitializerErrorMessages: + """Tests for the actionable error message on empty discovery.""" + + async def test_no_adversarial_targets_raises_with_actionable_message(self): + """``ValueError`` must name ``ADVERSARIAL_CHAT_*`` env vars and the ``TargetInitializer`` dependency.""" + init = BenchmarkInitializer() + with pytest.raises(ValueError) as exc_info: + await init.initialize_async() + + msg = str(exc_info.value) + assert "ADVERSARIAL_CHAT_" in msg + assert "TargetInitializer" in msg From d652a566e79f9bb539cf4bf414a74e0ba51aa2c1 Mon Sep 17 00:00:00 2001 From: Victor Valbuena Date: Thu, 21 May 2026 12:53:02 -0700 Subject: [PATCH 10/40] REFACTOR: Collapse AdversarialBenchmark to RapidResponse pattern Removes the local factory-construction override and the adversarial_models constructor parameter. AdversarialBenchmark now inherits the base Scenario._get_atomic_attacks_async loop and reads its strategy enum from AttackTechniqueRegistry entries tagged benchmark_fanout (registered by BenchmarkInitializer in the previous commit). What's removed: - adversarial_models: list[PromptTarget] constructor param + validation. - _adversarial_configs dict construction in __init__. - _get_atomic_attacks_async override that built local factories, iterated models x techniques x datasets, and injected attack_adversarial_config_override at create-time. - _infer_labels static method + the entire dedupe/collision-suffix loop that inferred model labels from target identifiers - replaced by TargetConfig.registry_name as the canonical label (set explicitly in ENV_TARGET_CONFIGS, no inference needed). - _get_benchmarkable_specs and _build_benchmark_strategy as @staticmethods on the class - replaced by a module-level _build_benchmark_strategy function. Strategy-class construction never reads scenario instance state, so the function does not belong to the class; module-level placement makes the dependency (only the registry) explicit and the unit-test surface flat. What's added: - BENCHMARK_FANOUT_TAG module constant (= "benchmark_fanout") as the shared contract between BenchmarkInitializer (writes the tag) and AdversarialBenchmark (reads it). - _StrategyOnlyMarker sentinel class to satisfy the required AttackTechniqueSpec.attack_class field when reconstructing minimal specs for strategy-enum construction. build_strategy_class_from_specs reads only name + strategy_tags, so the sentinel never reaches a runtime construction site; the real factory is fetched by name from the registry at attack-execution time. - _build_display_group override: extracts the target label from the fanned f"src__target" technique name so display rolls up per-model. Falls back to the full name when no __ separator is present. Where the (technique x target x dataset) permutation now happens: The pre-collapse override did all three dimensions at scenario runtime in one nested loop. Post-collapse the permutation is split across two stages, owned by different layers: 1. Initializer time - BenchmarkInitializer.initialize_async runs the (technique x adversarial-target) cross-product and registers one fanned AttackTechniqueFactory per pair into AttackTechniqueRegistry, tagged benchmark_fanout. Target binding lives on the factory. 2. Scenario runtime - Scenario._get_atomic_attacks_async (inherited, base class) runs the (fanned-variant x dataset) cross-product, building one AtomicAttack per pair. The target dimension is already resolved on the factory at this point. Net atomic-attack count is unchanged for the same inputs; the change is which layer owns which dimension. See the AdversarialBenchmark class docstring for the full explanation. VERSION bump 1 -> 2: The atomic_attack_name format changes from f"{technique}__{model}__{dataset}" (triple-segment, old override-driven) to f"{technique}__{model}_{dataset}" (double-then- single-underscore, base-inherited). Cached results from VERSION=1 remain queryable via memory.get_scenario_results(scenario_version=1) but won't suppress fresh runs with skip_cached=True (the param itself lands in the next commit). No CHANGELOG file in this repo; this note will land in the PR description. Doc notebook (doc/scanner/benchmark.{py,ipynb}) still references the removed adversarial_models API and will fail at runtime until Commit 9 rewrites it; deferred per plan F7. Not gated by any unit test. Tests rewritten end-to-end (619 -> 268 lines): see test_adversarial.py for the four test classes covering metadata, strategy construction, collapsed init surface, and display grouping. Wider regression: 1091/1091 pass across scenario+setup+registry; 547/547 pass in backend. --- .../scenarios/benchmark/adversarial.py | 376 ++++----- .../scenario/benchmark/test_adversarial.py | 773 +++++------------- 2 files changed, 383 insertions(+), 766 deletions(-) diff --git a/pyrit/scenario/scenarios/benchmark/adversarial.py b/pyrit/scenario/scenarios/benchmark/adversarial.py index dfec12839c..13432d189b 100644 --- a/pyrit/scenario/scenarios/benchmark/adversarial.py +++ b/pyrit/scenario/scenarios/benchmark/adversarial.py @@ -9,30 +9,161 @@ from typing import TYPE_CHECKING, ClassVar from pyrit.common import apply_defaults -from pyrit.executor.attack import AttackAdversarialConfig, AttackScoringConfig -from pyrit.prompt_target import CHAT_TARGET_REQUIREMENTS from pyrit.registry import AttackTechniqueRegistry, AttackTechniqueSpec from pyrit.registry.tag_query import TagQuery -from pyrit.scenario.core.atomic_attack import AtomicAttack from pyrit.scenario.core.dataset_configuration import DatasetConfiguration from pyrit.scenario.core.scenario import BaselineAttackPolicy, Scenario -from pyrit.scenario.core.scenario_techniques import SCENARIO_TECHNIQUES if TYPE_CHECKING: - from pyrit.prompt_target import PromptTarget from pyrit.scenario.core.scenario_strategy import ScenarioStrategy from pyrit.score import TrueFalseScorer logger = logging.getLogger(__name__) +#: Strategy tag applied by ``BenchmarkInitializer`` to every fanned variant it +#: registers in ``AttackTechniqueRegistry``. The benchmark scenario reads its +#: strategy enum from entries carrying this tag. +BENCHMARK_FANOUT_TAG: str = "benchmark_fanout" + + +class _StrategyOnlyMarker: + """ + Sentinel attack class used only to satisfy ``AttackTechniqueSpec.attack_class`` + when reconstructing minimal specs for strategy-enum construction. + + ``AttackTechniqueRegistry.build_strategy_class_from_specs`` reads only + ``spec.name`` and ``spec.strategy_tags`` — never ``attack_class`` — so the + sentinel is safe. At attack-execution time the base + ``Scenario._get_atomic_attacks_async`` looks up the real factory by name + from ``AttackTechniqueRegistry`` (where ``BenchmarkInitializer`` registered + it), so this sentinel never reaches a runtime construction site. + """ + + +def _build_benchmark_strategy() -> type[ScenarioStrategy]: + """ + Build the ``BenchmarkStrategy`` enum from ``BenchmarkInitializer``-registered fanout. + + *Fanned entries* (also called *fanned variants*) are the per-target copies + of adversarial-capable scenario techniques that ``BenchmarkInitializer`` + registers into ``AttackTechniqueRegistry``. For each adversarial-capable + technique in ``SCENARIO_TECHNIQUES`` and each adversarial-tagged target + in ``TargetRegistry``, the initializer creates one fanned variant named + ``f"{source_technique}__{target_name}"`` with the live target bound onto + ``adversarial_chat`` and the strategy tag + :data:`BENCHMARK_FANOUT_TAG` appended. This function reads those entries + back and builds an enum whose concrete members are exactly the fanned + variants. + + Implementation note: this is a module-level function rather than a + ``@staticmethod`` on ``AdversarialBenchmark``. Strategy-class + construction never reads scenario instance state, so the function does + not belong to the class; module-level placement makes the dependency + (only the registry) explicit and the unit-test surface flat. + + Reconstructs minimal ``AttackTechniqueSpec`` stand-ins (name + + strategy_tags only) from each fanned entry to pass into + ``build_strategy_class_from_specs``. The sentinel + :class:`_StrategyOnlyMarker` is used for the required ``attack_class`` + field — see the sentinel's docstring for why this is safe. + + Aggregate selectors on the generated enum: + + * ``all`` — every fanned variant (auto-included by the builder). + * ``light`` — variants tagged ``"light"`` (inherited from the source spec). + * ``single_turn`` / ``multi_turn`` — variants tagged with the matching + turn-style tag inherited from the source spec. + + Per-target selection is also available via the auto-applied + ``f"model:{target_name}"`` tag on each fanned variant, accessible by name + on the generated enum (e.g. + ``BenchmarkStrategy("red_teaming__adversarial_chat_singleturn")``). + + Returns: + type[ScenarioStrategy]: The dynamically generated ``BenchmarkStrategy`` class. + """ + registry = AttackTechniqueRegistry.get_registry_singleton() + fanned_entries = registry.get_by_tag(tag=BENCHMARK_FANOUT_TAG) + + fanned_specs = [ + AttackTechniqueSpec( + name=entry.name, + attack_class=_StrategyOnlyMarker, + strategy_tags=list(entry.tags.keys()), + ) + for entry in fanned_entries + ] + + return AttackTechniqueRegistry.build_strategy_class_from_specs( # type: ignore[ty:invalid-return-type] + class_name="BenchmarkStrategy", + specs=fanned_specs, + aggregate_tags={ + "light": TagQuery.any_of("light"), + "single_turn": TagQuery.any_of("single_turn"), + "multi_turn": TagQuery.any_of("multi_turn"), + }, + ) + + class AdversarialBenchmark(Scenario): """ - Benchmarking scenario that compares the attack success rate (ASR) - of several different adversarial models. + Benchmark scenario that compares the attack success rate (ASR) across adversarial models. + + Adversarial-model fan-out is provided by ``BenchmarkInitializer``, which + registers per-target *fanned variants* of adversarial-capable scenario + techniques into ``AttackTechniqueRegistry`` tagged ``benchmark_fanout``. + This scenario reads those variants and builds its strategy enum from + them, so the set of available strategies reflects whichever adversarial + targets were discovered when ``BenchmarkInitializer`` ran (typically via + ``.pyrit_conf`` initializer ordering). + + Inherits the base ``Scenario._get_atomic_attacks_async`` loop with no + override; the fanned ``adversarial_chat`` binding lives on the + registered factories, so atomic-attack construction needs no special + handling here. + + When permuted atomic attacks materialize + ========================================= + The (technique × target × dataset) cross-product now happens in two + stages, not one (the pre-collapse override did all three at runtime): + + 1. **Initializer time** — ``BenchmarkInitializer.initialize_async`` + runs the (technique × adversarial-target) cross-product. For each + adversarial-capable technique in ``SCENARIO_TECHNIQUES`` and each + adversarial-tagged target in ``TargetRegistry``, it registers one + fanned ``AttackTechniqueFactory`` into ``AttackTechniqueRegistry`` + with the live target baked onto the factory's adversarial config. + After this step, the registry contains N×M fanned entries where N + is the count of adversarial-capable techniques and M is the count of + discovered adversarial targets. + + 2. **Scenario runtime** — ``Scenario._get_atomic_attacks_async`` + (inherited, base class) runs the (fanned-variant × dataset) + cross-product. It iterates ``self._scenario_strategies`` (the + fanned-variant names the user picked via the ``BenchmarkStrategy`` + enum), pairs each with every seed group in + ``self._dataset_config``, and builds one ``AtomicAttack`` per pair. + The target binding rides through on the factory created in step 1, + so no per-target handling is needed at this layer. + + The user-observable result is the same shape as before + (one ``AtomicAttack`` per (technique, target, dataset) triple), but the + target dimension is now owned by the initializer and the dataset + dimension is owned by the scenario. + + Display grouping is by target name (the part after ``__`` in each + fanned technique name) rather than by technique, so per-model ASR rolls + up naturally in result displays. """ - VERSION: int = 1 + #: Bumped from 1 (pre-collapse) to 2 because the ``atomic_attack_name`` + #: format changed from ``f"{technique}__{model}__{dataset}"`` (triple-segment, + #: old override-driven) to ``f"{technique}__{model}_{dataset}"`` (double- + #: underscore-then-single-underscore, base-inherited). Cached results from + #: VERSION=1 remain queryable but won't suppress fresh runs. + VERSION: int = 2 + _cached_strategy_class: ClassVar[type[ScenarioStrategy] | None] = None #: AdversarialBenchmark compares attack-success rates across adversarial models; a baseline @@ -42,24 +173,29 @@ class AdversarialBenchmark(Scenario): @classmethod def get_strategy_class(cls) -> type[ScenarioStrategy]: """ - Return the AdversarialBenchmarkStrategy enum, building on first access. + Return the ``BenchmarkStrategy`` enum, building on first access. + + The enum is cached per-class for the lifetime of the process. To + rebuild after registry mutations (e.g. after re-running + ``BenchmarkInitializer`` with different adversarial targets), set + ``AdversarialBenchmark._cached_strategy_class = None`` and call again. Returns: - type[ScenarioStrategy]: The BenchmarkStrategy enum class. + type[ScenarioStrategy]: The ``BenchmarkStrategy`` enum class. """ if cls._cached_strategy_class is None: - cls._cached_strategy_class = AdversarialBenchmark._build_benchmark_strategy() - + cls._cached_strategy_class = _build_benchmark_strategy() return cls._cached_strategy_class @classmethod def get_default_strategy(cls) -> ScenarioStrategy: """ - Return the default strategy (``light`` — run benchmark-friendly techniques - that can wrap up quickly and without too many system resources). + Return the default strategy (``light``). Returns: - ScenarioStrategy: The ``light`` aggregate member. + ScenarioStrategy: The ``light`` aggregate member — runs the subset + of benchmark-friendly techniques that finish quickly with modest + system resources. """ return cls.get_strategy_class()("light") @@ -69,7 +205,8 @@ def default_dataset_config(cls) -> DatasetConfiguration: Return the default dataset configuration for benchmarking. Returns: - DatasetConfiguration: Configuration with standard harm-category datasets. + DatasetConfiguration: ``harmbench`` capped at 8 prompts per + atomic attack. """ return DatasetConfiguration( dataset_names=["harmbench"], @@ -80,7 +217,6 @@ def default_dataset_config(cls) -> DatasetConfiguration: def __init__( self, *, - adversarial_models: list[PromptTarget], objective_scorer: TrueFalseScorer | None = None, scenario_result_id: str | None = None, ) -> None: @@ -88,48 +224,12 @@ def __init__( Initialize the AdversarialBenchmark scenario. Args: - adversarial_models: A non-empty list of ``PromptTarget`` instances - that each satisfy :data:`CHAT_TARGET_REQUIREMENTS` (multi-turn - with editable history). Individual techniques selected at - run time may impose stricter capability requirements which are - enforced when their attack instances are constructed. - Labels are inferred from each target's identifier (preferring - ``underlying_model_name`` over ``model_name`` over the class - name). Identical targets are silently deduped and distinct - targets whose inferred names collide are suffixed (``_2``, - ``_3``, …) with a warning. - objective_scorer: Scorer for evaluating attack success. - Defaults to the registered default objective scorer. - scenario_result_id: Optional ID of an existing scenario - result to resume. - - Raises: - ValueError: If ``adversarial_models`` is empty, not a list, or - contains a target that does not satisfy - :data:`CHAT_TARGET_REQUIREMENTS`. + objective_scorer: Scorer for evaluating attack success. Defaults + to the registered default objective scorer (typically the + composite refusal+scale scorer set up by an initializer). + scenario_result_id: Optional ID of an existing scenario result + to resume. """ - if not adversarial_models: - raise ValueError("adversarial_models must be a non-empty list of PromptTarget instances.") - - if not isinstance(adversarial_models, list): - raise ValueError("adversarial_models must be a list of PromptTarget instances.") - - for target in adversarial_models: - try: - CHAT_TARGET_REQUIREMENTS.validate(target=target) - except ValueError as exc: - raise ValueError( - f"adversarial_models entry {type(target).__name__} does not satisfy " - f"the chat-target capability requirements: {exc}" - ) from exc - - # Infer labels, then wrap each bare target in a default AttackAdversarialConfig - # so it can be passed to factory.create() as an override. - labeled_targets = self._infer_labels(items=adversarial_models) - self._adversarial_configs: dict[str, AttackAdversarialConfig] = { - label: AttackAdversarialConfig(target=target) for label, target in labeled_targets.items() - } - self._objective_scorer: TrueFalseScorer = ( objective_scorer if objective_scorer else self._get_default_objective_scorer() ) @@ -141,156 +241,24 @@ def __init__( scenario_result_id=scenario_result_id, ) - async def _get_atomic_attacks_async(self) -> list[AtomicAttack]: + def _build_display_group(self, *, technique_name: str, seed_group_name: str) -> str: """ - Build atomic attacks from the cross-product of techniques × models × datasets. + Group atomic-attack results by adversarial-target label rather than by technique. - Factories are built locally from adversarial-capable ``SCENARIO_TECHNIQUES`` - (not the registry singleton). Each model is injected at create-time via - ``attack_adversarial_config_override``. - - Returns: - list[AtomicAttack]: One atomic attack per technique/model/dataset combination. - - Raises: - ValueError: If the scenario has not been initialized. - """ - if self._objective_target is None: - raise ValueError( - "Scenario not properly initialized. Call await scenario.initialize_async() before running." - ) - - benchmarkable_specs = AdversarialBenchmark._get_benchmarkable_specs() - local_factories = { - spec.name: AttackTechniqueRegistry.build_factory_from_spec(spec) for spec in benchmarkable_specs - } - - selected_techniques = {s.value for s in self._scenario_strategies} - seed_groups_by_dataset = self._dataset_config.get_seed_attack_groups() - scoring_config = AttackScoringConfig(objective_scorer=self._objective_scorer) - - atomic_attacks: list[AtomicAttack] = [] - for technique_name in selected_techniques: - factory = local_factories.get(technique_name) - if factory is None: - logger.warning("No factory for technique '%s', skipping.", technique_name) - continue - - for model_label, adv_config in self._adversarial_configs.items(): - for dataset_name, seed_groups in seed_groups_by_dataset.items(): - attack_technique = factory.create( - objective_target=self._objective_target, - attack_scoring_config=scoring_config, - attack_adversarial_config_override=adv_config, - ) - atomic_attacks.append( - AtomicAttack( - atomic_attack_name=f"{technique_name}__{model_label}__{dataset_name}", - attack_technique=attack_technique, - seed_groups=list(seed_groups), - adversarial_chat=adv_config.target, - objective_scorer=self._objective_scorer, - memory_labels=self._memory_labels, - display_group=model_label, - ) - ) - - return atomic_attacks - - @staticmethod - def _infer_labels( - *, - items: list[PromptTarget], - ) -> dict[str, PromptTarget]: - """ - Infer user-facing labels for a list of adversarial targets. - - The dedupe key is ``target.get_identifier().hash`` so identical - targets collapse to a single entry silently, while two distinct - targets whose inferred names happen to match get a numeric suffix - and a ``logger.warning`` so the situation isn't silent. + Fanned technique names have the format ``f"{source}__{target_name}"`` + (per ``BenchmarkInitializer``), so the target label is everything + after the ``__`` separator. Falls back to the full technique name + when no separator is present so legacy / non-fanned strategies still + render with a sensible label. Args: - items: List of ``PromptTarget`` instances. - - Returns: - dict[str, PromptTarget]: Mapping from inferred label to the - original target. Targets are wrapped in an - ``AttackAdversarialConfig`` by ``__init__`` after this call. - """ - result: dict[str, PromptTarget] = {} - seen_keys: dict[str, str | None] = {} - - for target in items: - identifier = target.get_identifier() - params = identifier.params or {} - base_name = params.get("underlying_model_name") or params.get("model_name") or type(target).__name__ - - dedupe_key = identifier.hash - - # Identical target already stored under some label — silently drop. - if dedupe_key in seen_keys.values(): - continue - - if base_name not in seen_keys: - result[base_name] = target - seen_keys[base_name] = dedupe_key - continue - - # Distinct target colliding on inferred name — find next free suffix and warn. - counter = 2 - while f"{base_name}_{counter}" in seen_keys: - counter += 1 - suffixed = f"{base_name}_{counter}" - logger.warning( - "Inferred label '%s' collided with a different model setup; using '%s' instead.", - base_name, - suffixed, - ) - result[suffixed] = target - seen_keys[suffixed] = dedupe_key - - return result - - @staticmethod - def _build_benchmark_strategy() -> type[ScenarioStrategy]: - """ - Build the BenchmarkStrategy enum from adversarial-capable ``SCENARIO_TECHNIQUES``. - - Returns a strategy class whose concrete members are adversarial-capable - techniques (no baked-in adversarial chat) and whose aggregates allow - selecting by turn style. - - Returns: - type[ScenarioStrategy]: The dynamically generated strategy enum class. - """ - specs = AdversarialBenchmark._get_benchmarkable_specs() - return AttackTechniqueRegistry.build_strategy_class_from_specs( # type: ignore[ty:invalid-return-type] - class_name="BenchmarkStrategy", - specs=TagQuery.all("core").filter(specs), - aggregate_tags={ - "default": TagQuery.any_of("default"), - "single_turn": TagQuery.any_of("single_turn"), - "multi_turn": TagQuery.any_of("multi_turn"), - "light": TagQuery.any_of("light"), - }, - ) - - @staticmethod - def _get_benchmarkable_specs() -> list[AttackTechniqueSpec]: - """ - Return techniques from ``SCENARIO_TECHNIQUES`` that accept an adversarial - model but don't have one already baked in. - - This is the dual guard: ``_accepts_adversarial`` ensures the technique - CAN use an adversarial model, and ``adversarial_chat is None`` ensures - it doesn't already have one set — we inject our own at create-time. + technique_name: The fanned technique name, e.g. + ``"red_teaming__adversarial_chat_singleturn"``. + seed_group_name: Unused for this scenario (display rolls up + per-target, not per-seed-group). Returns: - list[AttackTechniqueSpec]: Filtered, adversarial-ready specs. + str: The display group label — the target portion of the fanned + name when ``__`` is present, otherwise the full technique name. """ - return [ - spec - for spec in SCENARIO_TECHNIQUES - if AttackTechniqueRegistry._accepts_adversarial(spec.attack_class) and spec.adversarial_chat is None - ] + return technique_name.split("__", 1)[1] if "__" in technique_name else technique_name diff --git a/tests/unit/scenario/benchmark/test_adversarial.py b/tests/unit/scenario/benchmark/test_adversarial.py index 5914a40ba9..169361226d 100644 --- a/tests/unit/scenario/benchmark/test_adversarial.py +++ b/tests/unit/scenario/benchmark/test_adversarial.py @@ -1,619 +1,268 @@ # Copyright (c) Microsoft Corporation. # Licensed under the MIT license. -"""Tests for the AdversarialBenchmark scenario.""" +"""Tests for the post-collapse AdversarialBenchmark scenario. + +AdversarialBenchmark no longer takes an ``adversarial_models`` constructor +parameter and no longer builds local factories. It reads fanned variants +from ``AttackTechniqueRegistry`` (registered by ``BenchmarkInitializer``) +and inherits the base ``Scenario._get_atomic_attacks_async`` loop. + +These tests cover the new contract: +* Class metadata (VERSION, BASELINE policy, defaults). +* Strategy enum is built from ``benchmark_fanout``-tagged registry entries. +* Display grouping uses the target-label portion of fanned technique names. +* Construction accepts only ``objective_scorer`` and ``scenario_result_id``. +""" -import copy -from dataclasses import FrozenInstanceError from unittest.mock import MagicMock, patch import pytest -from pyrit.executor.attack import AttackAdversarialConfig -from pyrit.identifiers import ComponentIdentifier -from pyrit.models import ( - AttackOutcome, - AttackResult, - ScenarioIdentifier, - ScenarioResult, - SeedAttackGroup, - SeedObjective, - SeedPrompt, -) -from pyrit.prompt_target import PromptTarget, TargetCapabilities, TargetConfiguration +from pyrit.prompt_target import PromptTarget +from pyrit.registry import TargetRegistry from pyrit.registry.object_registries.attack_technique_registry import AttackTechniqueRegistry -from pyrit.scenario.core import AtomicAttack, BaselineAttackPolicy -from pyrit.scenario.core.dataset_configuration import DatasetConfiguration -from pyrit.scenario.core.scenario_techniques import SCENARIO_TECHNIQUES -from pyrit.scenario.scenarios.benchmark.adversarial import AdversarialBenchmark +from pyrit.scenario.core import BaselineAttackPolicy +from pyrit.scenario.scenarios.benchmark.adversarial import ( + BENCHMARK_FANOUT_TAG, + AdversarialBenchmark, + _build_benchmark_strategy, +) from pyrit.score import TrueFalseScorer - -# Self-pinned: any change to ``_get_benchmarkable_specs`` (or to the ``light`` tag -# membership in SCENARIO_TECHNIQUES) is reflected automatically — no magic numbers. -# -# ``_BENCHMARKABLE_*`` covers every adversarial-capable spec (used to verify the -# strategy enum's full concrete-member roster). ``_LIGHT_BENCHMARKABLE_*`` covers -# only the subset tagged ``"light"`` (used for runtime expectations under the -# default ``"light"`` strategy). -_BENCHMARKABLE_SPECS = AdversarialBenchmark._get_benchmarkable_specs() -_NUM_ADVERSARIAL_TECHNIQUES = len(_BENCHMARKABLE_SPECS) -_BENCHMARKABLE_TECHNIQUE_NAMES = {spec.name for spec in _BENCHMARKABLE_SPECS} -_BENCHMARKABLE_ATTACK_CLASSES = {spec.attack_class for spec in _BENCHMARKABLE_SPECS} - -_LIGHT_BENCHMARKABLE_SPECS = [spec for spec in _BENCHMARKABLE_SPECS if "light" in spec.strategy_tags] -_NUM_LIGHT_BENCHMARKABLE = len(_LIGHT_BENCHMARKABLE_SPECS) +from pyrit.setup.initializers import BenchmarkInitializer +from pyrit.setup.initializers.components.targets import TargetInitializerTags # --------------------------------------------------------------------------- -# Synthetic many-shot examples — prevents reading the real JSON during tests +# Fixtures # --------------------------------------------------------------------------- -_MOCK_MANY_SHOT_EXAMPLES = [{"question": f"test question {i}", "answer": f"test answer {i}"} for i in range(100)] -# --------------------------------------------------------------------------- -# Helpers -# --------------------------------------------------------------------------- +@pytest.fixture(autouse=True) +def reset_registries_and_cache(): + """Reset both registries and AdversarialBenchmark's strategy-class cache between tests.""" + AttackTechniqueRegistry.reset_instance() + TargetRegistry.reset_instance() + AdversarialBenchmark._cached_strategy_class = None + yield + AttackTechniqueRegistry.reset_instance() + TargetRegistry.reset_instance() + AdversarialBenchmark._cached_strategy_class = None -def _mock_id(name: str, *, params: dict | None = None) -> ComponentIdentifier: - return ComponentIdentifier(class_name=name, class_module="test", params=params or {}) +def _register_adversarial_target(*, name: str) -> PromptTarget: + """Register a mock adversarial-tagged target in TargetRegistry.""" + target = MagicMock(spec=PromptTarget) + target.capabilities.includes.return_value = True + registry = TargetRegistry.get_registry_singleton() + registry.register_instance(target, name=name, tags=[TargetInitializerTags.ADVERSARIAL.value]) + return target -_CHAT_TARGET_CONFIGURATION = TargetConfiguration( - capabilities=TargetCapabilities( - supports_multi_turn=True, - supports_multi_message_pieces=True, - supports_system_prompt=True, - supports_editable_history=True, - ), -) +async def _fan_out(*, target_names: list[str]) -> None: + """Register mock targets + run BenchmarkInitializer to populate AttackTechniqueRegistry.""" + for name in target_names: + _register_adversarial_target(name=name) + init = BenchmarkInitializer() + await init.initialize_async() + +# --------------------------------------------------------------------------- +# Class metadata +# --------------------------------------------------------------------------- -def _make_adversarial_target(name: str, *, params: dict | None = None) -> MagicMock: - """Create a mock adversarial PromptTarget with a given model name and optional identifier params. - By default, ``model_name`` is stamped into the identifier params so the - inferred label produced by ``_infer_labels`` matches ``name``. Pass an - explicit ``params`` dict to override (e.g. to omit the key for collision - testing or to add ``underlying_model_name`` / ``endpoint``). +class TestAdversarialBenchmarkMetadata: + """Tests for class-level metadata that doesn't depend on fan-out state.""" - The mock exposes a real ``TargetConfiguration`` declaring multi-turn and - editable history so the target satisfies ``CHAT_TARGET_REQUIREMENTS`` at - construction time. - """ - mock = MagicMock(spec=PromptTarget) - mock._model_name = name - mock.get_identifier.return_value = _mock_id(name, params=params if params is not None else {"model_name": name}) - mock.configuration = _CHAT_TARGET_CONFIGURATION - return mock + def test_version_is_2(self): + """VERSION is bumped from 1 because the atomic_attack_name format changed.""" + assert AdversarialBenchmark.VERSION == 2 + def test_baseline_attack_policy_is_forbidden(self): + """A baseline contributes no signal to a model-comparison benchmark, so it is forbidden.""" + assert AdversarialBenchmark.BASELINE_ATTACK_POLICY is BaselineAttackPolicy.Forbidden -def _make_seed_groups(name: str) -> list[SeedAttackGroup]: - """Create two seed attack groups for a given category.""" - return [ - SeedAttackGroup(seeds=[SeedObjective(value=f"{name} objective 1"), SeedPrompt(value=f"{name} prompt 1")]), - SeedAttackGroup(seeds=[SeedObjective(value=f"{name} objective 2"), SeedPrompt(value=f"{name} prompt 2")]), - ] + def test_default_dataset_config_uses_harmbench(self): + config = AdversarialBenchmark.default_dataset_config() + assert config.get_default_dataset_names() == ["harmbench"] + + def test_default_dataset_config_max_size_is_8(self): + assert AdversarialBenchmark.default_dataset_config().max_dataset_size == 8 + + def test_benchmark_fanout_tag_value(self): + """The shared tag value must match what BenchmarkInitializer applies.""" + assert BENCHMARK_FANOUT_TAG == "benchmark_fanout" # --------------------------------------------------------------------------- -# Fixtures +# Strategy class construction # --------------------------------------------------------------------------- -@pytest.fixture -def all_supported_attacks(): - """All attacks that currently support adversarial models (computed from production).""" - return _BENCHMARKABLE_TECHNIQUE_NAMES +class TestAdversarialBenchmarkStrategy: + """Tests for _build_benchmark_strategy and the cached get_strategy_class accessor.""" + async def test_strategy_built_from_fanned_registry_entries(self): + """Every benchmark_fanout-tagged entry produces one concrete enum member.""" + await _fan_out(target_names=["adv_a", "adv_b"]) -@pytest.fixture -def mock_objective_target(): - mock = MagicMock(spec=PromptTarget) - mock.get_identifier.return_value = _mock_id("MockObjectiveTarget") - return mock + strategy_cls = AdversarialBenchmark.get_strategy_class() + aggregate_names = {"all"} | strategy_cls.get_aggregate_tags() + concrete_members = [m for m in strategy_cls if m.value not in aggregate_names] + assert len(concrete_members) > 0 + for member in concrete_members: + assert "__" in member.value, f"Expected fanned format with '__', got: {member.value}" -@pytest.fixture -def two_adversarial_models(): - """Two mock adversarial models for benchmark permutation.""" - return [_make_adversarial_target("model_a"), _make_adversarial_target("model_b")] + async def test_strategy_concrete_member_count_matches_registry(self): + """Concrete enum members count equals fanned spec count in the registry.""" + await _fan_out(target_names=["adv_a", "adv_b"]) + attack_registry = AttackTechniqueRegistry.get_registry_singleton() + fanned_entries = attack_registry.get_by_tag(tag=BENCHMARK_FANOUT_TAG) -@pytest.fixture -def single_adversarial_model(): - """Single mock adversarial model.""" - return [_make_adversarial_target("model_a")] + strategy_cls = AdversarialBenchmark.get_strategy_class() + aggregate_names = {"all"} | strategy_cls.get_aggregate_tags() + concrete_members = [m for m in strategy_cls if m.value not in aggregate_names] + assert len(concrete_members) == len(fanned_entries) -@pytest.fixture(autouse=True) -def reset_technique_registry(): - """Reset the AttackTechniqueRegistry and cached strategy class between tests.""" - from pyrit.registry import TargetRegistry + async def test_strategy_exposes_per_model_selection(self): + """Each fanned variant inherits its model:* tag, accessible by name on the enum.""" + await _fan_out(target_names=["adv_a"]) - AttackTechniqueRegistry.reset_instance() - TargetRegistry.reset_instance() - AdversarialBenchmark._cached_strategy_class = None - yield - AttackTechniqueRegistry.reset_instance() - TargetRegistry.reset_instance() - AdversarialBenchmark._cached_strategy_class = None + attack_registry = AttackTechniqueRegistry.get_registry_singleton() + model_a_entries = attack_registry.get_by_tag(tag="model:adv_a") + assert len(model_a_entries) > 0 + strategy_cls = AdversarialBenchmark.get_strategy_class() + for entry in model_a_entries: + member = strategy_cls(entry.name) + assert "model:adv_a" in member.tags -@pytest.fixture(autouse=True) -def patch_many_shot_load(): - """Prevent ManyShotJailbreakAttack from loading the full bundled dataset.""" - 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(): - """Set minimal env vars needed for OpenAIChatTarget fallback via @apply_defaults.""" - with patch.dict( - "os.environ", - { - "OPENAI_CHAT_ENDPOINT": "https://test.openai.azure.com/", - "OPENAI_CHAT_KEY": "test-key", - "OPENAI_CHAT_MODEL": "gpt-4", - }, - ): - yield - - -FIXTURES = ["patch_central_database", "mock_runtime_env"] - - -# =========================================================================== -# Type and syntax tests -# =========================================================================== - - -@pytest.mark.usefixtures(*FIXTURES) -class TestBenchmarkTypes: - """Unit tests for types, validation, and basic construction.""" - - def test_empty_list_adversarial_models_raises(self): - """Passing an empty list must raise ValueError.""" - with pytest.raises(ValueError, match="non-empty"): - AdversarialBenchmark(adversarial_models=[]) - - def test_unsupported_type_adversarial_models_raises(self): - """Passing a non-list type must raise ValueError.""" - with pytest.raises(ValueError, match="non-empty list|list of PromptTarget"): - AdversarialBenchmark(adversarial_models="not-a-list") # type: ignore[arg-type] - - def test_adversarial_model_missing_chat_capabilities_raises(self): - """A target that does not satisfy CHAT_TARGET_REQUIREMENTS must be rejected at construction.""" - non_chat_target = MagicMock(spec=PromptTarget) - non_chat_target.get_identifier.return_value = _mock_id("NonChatTarget") - non_chat_target.configuration = TargetConfiguration( - capabilities=TargetCapabilities( - supports_multi_turn=False, - supports_editable_history=False, - ), - ) + async def test_strategy_includes_required_aggregates(self): + """The strategy enum exposes all, light, single_turn, multi_turn aggregates.""" + await _fan_out(target_names=["adv_a"]) - with pytest.raises(ValueError, match="chat-target capability requirements"): - AdversarialBenchmark(adversarial_models=[non_chat_target]) + strategy_cls = AdversarialBenchmark.get_strategy_class() + aggregates = strategy_cls.get_aggregate_tags() - def test_version_is_1(self): - assert AdversarialBenchmark.VERSION == 1 + assert "all" in aggregates + assert "light" in aggregates + assert "single_turn" in aggregates + assert "multi_turn" in aggregates - def test_default_dataset_config_uses_harmbench(self): - config = AdversarialBenchmark.default_dataset_config() - assert isinstance(config, DatasetConfiguration) - names = config.get_default_dataset_names() - assert "harmbench" in names + async def test_get_strategy_class_is_cached(self): + """Repeated calls within a process return the same class instance.""" + await _fan_out(target_names=["adv_a"]) + + first = AdversarialBenchmark.get_strategy_class() + second = AdversarialBenchmark.get_strategy_class() + + assert first is second + + async def test_cache_can_be_cleared_to_rebuild(self): + """Setting _cached_strategy_class = None forces a rebuild from current registry state.""" + await _fan_out(target_names=["adv_a"]) + first = AdversarialBenchmark.get_strategy_class() + + await _fan_out(target_names=["adv_b"]) + AdversarialBenchmark._cached_strategy_class = None + second = AdversarialBenchmark.get_strategy_class() + + assert first is not second + + async def test_default_strategy_is_light(self): + """get_default_strategy returns the 'light' aggregate so quick benchmark runs are the default.""" + await _fan_out(target_names=["adv_a"]) - def test_default_dataset_config_max_size_is_8(self): - config = AdversarialBenchmark.default_dataset_config() - assert config.max_dataset_size == 8 - - def test_frozen_spec_cannot_be_mutated(self): - """AttackTechniqueSpec is frozen — direct mutation must raise.""" - spec = SCENARIO_TECHNIQUES[0] - with pytest.raises(FrozenInstanceError): - spec.name = "mutated" # type: ignore[misc] - - -# =========================================================================== -# Strategy construction tests -# =========================================================================== - - -def _make_benchmark(adversarial_models): - """Helper to create a AdversarialBenchmark with mocked default scorer.""" - with patch("pyrit.scenario.core.scenario.Scenario._get_default_objective_scorer") as mock_scorer: - mock_scorer.return_value = MagicMock(spec=TrueFalseScorer, get_identifier=lambda: _mock_id("scorer")) - return AdversarialBenchmark(adversarial_models=adversarial_models) - - -@pytest.mark.usefixtures(*FIXTURES) -class TestBenchmarkStrategy: - """Tests for the (static) BenchmarkStrategy enum and instance-level wiring.""" - - def test_strategy_includes_all_adversarial_techniques(self, all_supported_attacks): - """get_strategy_class() concrete members match the adversarial-capable spec set.""" - strat = AdversarialBenchmark.get_strategy_class() - values = {s.value for s in strat.get_all_strategies()} - assert values == all_supported_attacks - - def test_strategy_has_no_permuted_members(self): - """No ``__model`` suffixes — models are a runtime parameter, not a strategy axis.""" - strat = AdversarialBenchmark.get_strategy_class() - values = {s.value for s in strat.get_all_strategies()} - assert not any("__" in v for v in values) - - def test_strategy_excludes_non_adversarial_techniques(self): - """prompt_sending and many_shot don't accept an adversarial chat and must be excluded.""" - strat = AdversarialBenchmark.get_strategy_class() - values = {s.value for s in strat.get_all_strategies()} - assert "prompt_sending" not in values - assert "many_shot" not in values - - def test_strategy_class_is_static(self, single_adversarial_model, two_adversarial_models): - """All instances share the same strategy class — no per-instance permutation.""" - s1 = _make_benchmark(single_adversarial_model) - s2 = _make_benchmark(two_adversarial_models) - assert s1._strategy_class is s2._strategy_class - assert s1._strategy_class is AdversarialBenchmark.get_strategy_class() - - def test_default_strategy_is_light(self): - """Default expands to every benchmarkable technique via the ``all`` aggregate.""" default = AdversarialBenchmark.get_default_strategy() assert default.value == "light" - def test_benchmarkable_specs_have_no_adversarial_chat(self): - """Filtered specs must leave adversarial_chat unset — the scenario injects its own.""" - for spec in AdversarialBenchmark._get_benchmarkable_specs(): - assert spec.adversarial_chat is None - - def test_benchmarkable_specs_accept_adversarial(self): - """All filtered specs must accept attack_adversarial_config.""" - for spec in AdversarialBenchmark._get_benchmarkable_specs(): - assert AttackTechniqueRegistry._accepts_adversarial(spec.attack_class) - - def test_original_scenario_techniques_unmodified(self, two_adversarial_models): - """SCENARIO_TECHNIQUES global must not be mutated by spec filtering.""" - original = copy.deepcopy([(s.name, s.attack_class) for s in SCENARIO_TECHNIQUES]) - _make_benchmark(two_adversarial_models) - current = [(s.name, s.attack_class) for s in SCENARIO_TECHNIQUES] - assert current == original - - def test_singleton_registry_not_polluted(self, two_adversarial_models): - """Building atomic attacks must not register anything in the global singleton.""" - _make_benchmark(two_adversarial_models) - registry = AttackTechniqueRegistry.get_registry_singleton() - factories = registry.get_factories() - assert not any("__" in name for name in factories) - - def test_scenario_name(self, single_adversarial_model): - """Scenario name should be 'AdversarialBenchmark'.""" - scenario = _make_benchmark(single_adversarial_model) - assert scenario.name == "AdversarialBenchmark" - - -# =========================================================================== -# Runtime / attack generation tests -# =========================================================================== - - -@pytest.mark.usefixtures(*FIXTURES) -class TestBenchmarkRuntime: - """Tests for _get_atomic_attacks_async and display grouping.""" - - async def _init_and_get_attacks( - self, - *, - mock_objective_target, - adversarial_models, - seed_groups: dict[str, list[SeedAttackGroup]] | None = None, - strategies=None, - ) -> tuple[AdversarialBenchmark, list[AtomicAttack]]: - """Helper: create AdversarialBenchmark, initialize, return (scenario, attacks).""" - groups = seed_groups or {"harmbench": _make_seed_groups("harmbench")} - with ( - patch.object(DatasetConfiguration, "get_seed_attack_groups", return_value=groups), - patch("pyrit.scenario.core.scenario.Scenario._get_default_objective_scorer") as mock_scorer, - ): - mock_scorer.return_value = MagicMock(spec=TrueFalseScorer, get_identifier=lambda: _mock_id("scorer")) - scenario = AdversarialBenchmark(adversarial_models=adversarial_models) - init_kwargs: dict = {"objective_target": mock_objective_target} - if strategies: - init_kwargs["scenario_strategies"] = strategies - await scenario.initialize_async(**init_kwargs) - attacks = await scenario._get_atomic_attacks_async() - return scenario, attacks - - @pytest.mark.asyncio - async def test_default_strategy_runs_light_techniques(self, mock_objective_target, two_adversarial_models): - """With no strategies passed, default ``light`` produces N_light x N_models attacks.""" - _, attacks = await self._init_and_get_attacks( - mock_objective_target=mock_objective_target, - adversarial_models=two_adversarial_models, - ) - assert len(attacks) == _NUM_LIGHT_BENCHMARKABLE * 2 - - @pytest.mark.asyncio - async def test_all_strategy_produces_full_cross_product(self, mock_objective_target, two_adversarial_models): - """ALL strategy: N_techniques x 2 models x 1 dataset attacks.""" - with ( - patch.object( - DatasetConfiguration, - "get_seed_attack_groups", - return_value={"harmbench": _make_seed_groups("harmbench")}, - ), - patch("pyrit.scenario.core.scenario.Scenario._get_default_objective_scorer") as mock_scorer, - ): - mock_scorer.return_value = MagicMock(spec=TrueFalseScorer, get_identifier=lambda: _mock_id("scorer")) - scenario = AdversarialBenchmark(adversarial_models=two_adversarial_models) - all_strat = scenario._strategy_class("all") - await scenario.initialize_async(objective_target=mock_objective_target, scenario_strategies=[all_strat]) - attacks = await scenario._get_atomic_attacks_async() - assert len(attacks) == _NUM_ADVERSARIAL_TECHNIQUES * 2 - - @pytest.mark.asyncio - async def test_atomic_attack_names_are_unique(self, mock_objective_target, two_adversarial_models): - """All atomic_attack_name values must be unique for resume correctness.""" - with ( - patch.object( - DatasetConfiguration, - "get_seed_attack_groups", - return_value={"harmbench": _make_seed_groups("harmbench")}, - ), - patch("pyrit.scenario.core.scenario.Scenario._get_default_objective_scorer") as mock_scorer, - ): - mock_scorer.return_value = MagicMock(spec=TrueFalseScorer, get_identifier=lambda: _mock_id("scorer")) - scenario = AdversarialBenchmark(adversarial_models=two_adversarial_models) - all_strat = scenario._strategy_class("all") - await scenario.initialize_async(objective_target=mock_objective_target, scenario_strategies=[all_strat]) - attacks = await scenario._get_atomic_attacks_async() - names = [a.atomic_attack_name for a in attacks] - assert len(names) == len(set(names)) - - @pytest.mark.asyncio - async def test_atomic_attack_names_follow_pattern(self, mock_objective_target, single_adversarial_model): - """Each atomic_attack_name should contain the technique__model and dataset.""" - with ( - patch.object( - DatasetConfiguration, - "get_seed_attack_groups", - return_value={"harmbench": _make_seed_groups("harmbench")}, - ), - patch("pyrit.scenario.core.scenario.Scenario._get_default_objective_scorer") as mock_scorer, - ): - mock_scorer.return_value = MagicMock(spec=TrueFalseScorer, get_identifier=lambda: _mock_id("scorer")) - scenario = AdversarialBenchmark(adversarial_models=single_adversarial_model) - all_strat = scenario._strategy_class("all") - await scenario.initialize_async(objective_target=mock_objective_target, scenario_strategies=[all_strat]) - attacks = await scenario._get_atomic_attacks_async() - for a in attacks: - assert "_harmbench" in a.atomic_attack_name - assert "__model_a" in a.atomic_attack_name - - @pytest.mark.asyncio - async def test_display_groups_by_adversarial_model(self, mock_objective_target, two_adversarial_models): - """display_group should group by model label, not by technique or dataset.""" - with ( - patch.object( - DatasetConfiguration, - "get_seed_attack_groups", - return_value={"harmbench": _make_seed_groups("harmbench")}, - ), - patch("pyrit.scenario.core.scenario.Scenario._get_default_objective_scorer") as mock_scorer, - ): - mock_scorer.return_value = MagicMock(spec=TrueFalseScorer, get_identifier=lambda: _mock_id("scorer")) - scenario = AdversarialBenchmark(adversarial_models=two_adversarial_models) - all_strat = scenario._strategy_class("all") - await scenario.initialize_async(objective_target=mock_objective_target, scenario_strategies=[all_strat]) - attacks = await scenario._get_atomic_attacks_async() - display_groups = {a.display_group for a in attacks} - assert display_groups == {"model_a", "model_b"} - - @pytest.mark.asyncio - async def test_raises_when_not_initialized(self, single_adversarial_model): - """_get_atomic_attacks_async must raise if initialize_async was not called.""" - scenario = _make_benchmark(single_adversarial_model) - with pytest.raises(ValueError, match="Scenario not properly initialized"): - await scenario._get_atomic_attacks_async() - - @pytest.mark.asyncio - async def test_multiple_datasets_multiplies_attacks(self, mock_objective_target, single_adversarial_model): - """1 model x N_light_techniques x 2 datasets = 2 * N_light atomic attacks (default ``light``).""" - two_datasets = { - "harmbench": _make_seed_groups("harmbench"), - "extra": _make_seed_groups("extra"), - } - _, attacks = await self._init_and_get_attacks( - mock_objective_target=mock_objective_target, - adversarial_models=single_adversarial_model, - seed_groups=two_datasets, - ) - assert len(attacks) == _NUM_LIGHT_BENCHMARKABLE * 2 - - @pytest.mark.asyncio - async def test_attacks_use_all_benchmarkable_attack_classes(self, mock_objective_target, single_adversarial_model): - """Under the ``all`` strategy, atomic attacks must cover every adversarial-capable attack class.""" - scenario_class_strategies = AdversarialBenchmark.get_strategy_class() - _, attacks = await self._init_and_get_attacks( - mock_objective_target=mock_objective_target, - adversarial_models=single_adversarial_model, - strategies=[scenario_class_strategies("all")], - ) - technique_classes = {type(a.attack_technique.attack) for a in attacks} - assert technique_classes == _BENCHMARKABLE_ATTACK_CLASSES - - @pytest.mark.asyncio - async def test_attacks_carry_seed_groups(self, mock_objective_target, single_adversarial_model): - """Each atomic attack should have non-empty objectives from the seed groups.""" - _, attacks = await self._init_and_get_attacks( - mock_objective_target=mock_objective_target, - adversarial_models=single_adversarial_model, - ) - for a in attacks: - assert len(a.objectives) > 0 - - async def test_baseline_excluded(self, mock_objective_target, single_adversarial_model): - """AdversarialBenchmark must opt out of the parent's default baseline. - - Verifies both the class-level capability flag and the observable property - (no atomic attack is named ``"baseline"``). - """ - scenario, _ = await self._init_and_get_attacks( - mock_objective_target=mock_objective_target, - adversarial_models=single_adversarial_model, - ) - assert type(scenario).BASELINE_ATTACK_POLICY is BaselineAttackPolicy.Forbidden - assert not any(a.atomic_attack_name == "baseline" for a in scenario._atomic_attacks) - - async def test_baseline_explicit_true_raises(self, mock_objective_target, single_adversarial_model): - """Explicitly passing include_baseline=True to a forbidden scenario raises ValueError.""" - scenario = AdversarialBenchmark(adversarial_models=single_adversarial_model) - with pytest.raises(ValueError, match="does not support a default baseline"): - await scenario.initialize_async( - objective_target=mock_objective_target, - include_baseline=True, - ) - - async def test_baseline_explicit_false_succeeds(self, mock_objective_target, single_adversarial_model): - """Explicit include_baseline=False on a forbidden scenario is accepted (matches the default).""" - groups = {"harmbench": _make_seed_groups("harmbench")} - with ( - patch.object(DatasetConfiguration, "get_seed_attack_groups", return_value=groups), - patch("pyrit.scenario.core.scenario.Scenario._get_default_objective_scorer") as mock_scorer, - ): - mock_scorer.return_value = MagicMock(spec=TrueFalseScorer, get_identifier=lambda: _mock_id("scorer")) - scenario = AdversarialBenchmark(adversarial_models=single_adversarial_model) - await scenario.initialize_async( - objective_target=mock_objective_target, - include_baseline=False, - ) - assert not any(a.atomic_attack_name == "baseline" for a in scenario._atomic_attacks) - - -# =========================================================================== -# adversarial_models normalization tests (label inference / dedupe / collision) -# =========================================================================== - - -@pytest.mark.usefixtures(*FIXTURES) -class TestBenchmarkAdversarialModelsNormalization: - """Tests for the list → ``dict[str, AttackAdversarialConfig]`` normalization in __init__. - - Labels are inferred from each target's identifier; identical targets dedupe - silently, distinct targets whose inferred names collide get suffixed with - a warning. - """ - - def test_list_of_targets_infers_labels_from_model_name(self): - """A list of bare targets is normalized to {model_name: AttackAdversarialConfig}.""" - t1 = _make_adversarial_target("t1", params={"model_name": "alpha"}) - t2 = _make_adversarial_target("t2", params={"model_name": "beta"}) - scenario = _make_benchmark([t1, t2]) - assert set(scenario._adversarial_configs.keys()) == {"alpha", "beta"} - assert all(isinstance(v, AttackAdversarialConfig) for v in scenario._adversarial_configs.values()) - assert scenario._adversarial_configs["alpha"].target is t1 - assert scenario._adversarial_configs["beta"].target is t2 - - def test_list_falls_back_to_underlying_model_name(self): - """``underlying_model_name`` is preferred over ``model_name`` when present.""" - t = _make_adversarial_target("t", params={"underlying_model_name": "gpt-4o", "model_name": "wrapper"}) - scenario = _make_benchmark([t]) - assert "gpt-4o" in scenario._adversarial_configs - - def test_list_dedupe_silent_for_identical_target(self, caplog): - """The same target instance passed twice in a list collapses to one entry, silently.""" - t = _make_adversarial_target("t", params={"model_name": "alpha"}) - with caplog.at_level("WARNING"): - scenario = _make_benchmark([t, t]) - assert list(scenario._adversarial_configs.keys()) == ["alpha"] - assert "collided" not in caplog.text - - def test_list_collision_suffixes_distinct_targets_and_warns(self, caplog): - """Two distinct targets that infer the same name get suffixed and a warning is logged.""" - t1 = _make_adversarial_target("t1", params={"model_name": "alpha", "endpoint": "ep1"}) - t2 = _make_adversarial_target("t2", params={"model_name": "alpha", "endpoint": "ep2"}) - with caplog.at_level("WARNING"): - scenario = _make_benchmark([t1, t2]) - assert set(scenario._adversarial_configs.keys()) == {"alpha", "alpha_2"} - assert "collided" in caplog.text - - -# =========================================================================== -# ASR-sensibility tests (per-model breakdown math) -# =========================================================================== + def test_build_benchmark_strategy_empty_registry_produces_aggregates_only(self): + """No fan-out → enum still constructs (aggregates always present), just with zero concrete members.""" + strategy_cls = _build_benchmark_strategy() + aggregate_names = {"all"} | strategy_cls.get_aggregate_tags() + concrete_members = [m for m in strategy_cls if m.value not in aggregate_names] + assert concrete_members == [] -@pytest.mark.usefixtures("patch_central_database") -class TestBenchmarkASRBreakdown: - """Verify the per-display-group ASR math the notebook sanity check relies on. - - A higher per-group success rate must correspond to more ``AttackOutcome.SUCCESS`` - results in that group. This test pins the invariant that lets reviewers trust - the printed breakdown when comparing adversarial models or system prompts. - """ - - @staticmethod - def _result(*, conv_id: str, outcome: AttackOutcome) -> AttackResult: - return AttackResult( - conversation_id=conv_id, - objective="objective", - outcome=outcome, - executed_turns=1, - ) +# --------------------------------------------------------------------------- +# Construction +# --------------------------------------------------------------------------- - def test_per_model_breakdown_reflects_outcome_counts(self): - """High-success model > low-success model in per-group ASR; math invariants hold.""" - # Two techniques × two models, mirroring how AdversarialBenchmark keys atomic_attack_name - # ("{technique}__{model_label}__{dataset}") and folds them into model_label. - attack_results: dict[str, list[AttackResult]] = { - "role_play__model_high__hb": [ - self._result(conv_id=f"high-rp-{i}", outcome=AttackOutcome.SUCCESS) for i in range(3) - ], - "context_compliance__model_high__hb": [ - self._result(conv_id=f"high-cc-{i}", outcome=AttackOutcome.SUCCESS) for i in range(3) - ], - "role_play__model_low__hb": [ - self._result(conv_id=f"low-rp-{i}", outcome=AttackOutcome.FAILURE) for i in range(3) - ], - "context_compliance__model_low__hb": [ - self._result(conv_id=f"low-cc-{i}", outcome=AttackOutcome.FAILURE) for i in range(3) - ], - } - display_group_map = { - "role_play__model_high__hb": "model_high", - "context_compliance__model_high__hb": "model_high", - "role_play__model_low__hb": "model_low", - "context_compliance__model_low__hb": "model_low", - } - result = ScenarioResult( - scenario_identifier=ScenarioIdentifier(name="AdversarialBenchmark", scenario_version=1), - objective_target_identifier=ComponentIdentifier(class_name="MockTarget", class_module="test"), - attack_results=attack_results, - objective_scorer_identifier=ComponentIdentifier(class_name="MockScorer", class_module="test"), - display_group_map=display_group_map, - ) - groups = result.get_display_groups() - assert set(groups.keys()) == {"model_high", "model_low"} +class TestAdversarialBenchmarkInit: + """Tests for the collapsed __init__ surface (objective_scorer + scenario_result_id only).""" + + @pytest.mark.usefixtures("patch_central_database") + async def test_construct_with_default_objective_scorer(self): + """When no scorer is supplied, _get_default_objective_scorer is consulted.""" + await _fan_out(target_names=["adv_a"]) + + default_scorer = MagicMock(spec=TrueFalseScorer) + with patch.object(AdversarialBenchmark, "_get_default_objective_scorer", return_value=default_scorer): + bench = AdversarialBenchmark() + + assert bench._objective_scorer is default_scorer - per_group = { - label: int(sum(1 for r in rs if r.outcome == AttackOutcome.SUCCESS) / max(len(rs), 1) * 100) - for label, rs in groups.items() - } + @pytest.mark.usefixtures("patch_central_database") + async def test_construct_with_explicit_objective_scorer(self): + """An explicit scorer is used as-is, no default consulted.""" + await _fan_out(target_names=["adv_a"]) - # The whole point of the sanity check: more SUCCESSes ⇒ higher rate. - assert per_group["model_high"] == 100 - assert per_group["model_low"] == 0 - assert per_group["model_high"] > per_group["model_low"] - # Bounds invariant the notebook asserts. - assert all(0 <= rate <= 100 for rate in per_group.values()) + explicit_scorer = MagicMock(spec=TrueFalseScorer) + bench = AdversarialBenchmark(objective_scorer=explicit_scorer) - # Overall rate matches the weighted average (6 SUCCESS / 12 total = 50%). - assert result.objective_achieved_rate() == 50 + assert bench._objective_scorer is explicit_scorer - # Display grouping must not lose results. - assert sum(len(rs) for rs in groups.values()) == sum(len(rs) for rs in attack_results.values()) + async def test_construct_takes_no_adversarial_models_param(self): + """Regression: the old adversarial_models constructor param is removed.""" + await _fan_out(target_names=["adv_a"]) + + with pytest.raises(TypeError): + AdversarialBenchmark(adversarial_models=[MagicMock(spec=PromptTarget)]) # type: ignore[call-arg] + + +# --------------------------------------------------------------------------- +# Display grouping +# --------------------------------------------------------------------------- + + +@pytest.mark.usefixtures("patch_central_database") +class TestAdversarialBenchmarkDisplayGroup: + """Tests for _build_display_group's fanned-name parsing.""" + + async def _make_bench(self) -> AdversarialBenchmark: + await _fan_out(target_names=["adv_a"]) + return AdversarialBenchmark(objective_scorer=MagicMock(spec=TrueFalseScorer)) + + async def test_extracts_target_label_after_double_underscore(self): + bench = await self._make_bench() + result = bench._build_display_group( + technique_name="red_teaming__adversarial_chat_singleturn", + seed_group_name="seed_group_1", + ) + assert result == "adversarial_chat_singleturn" + + async def test_falls_back_to_full_name_when_no_separator(self): + """Non-fanned names (no ``__``) return the full technique name unchanged.""" + bench = await self._make_bench() + result = bench._build_display_group( + technique_name="prompt_sending", + seed_group_name="seed_group_1", + ) + assert result == "prompt_sending" + + async def test_ignores_seed_group_name(self): + """seed_group_name input must not influence the result (display rolls up per-target).""" + bench = await self._make_bench() + first = bench._build_display_group( + technique_name="red_teaming__adv_a", + seed_group_name="seed_group_a", + ) + second = bench._build_display_group( + technique_name="red_teaming__adv_a", + seed_group_name="seed_group_b", + ) + assert first == second == "adv_a" From b54c47f0904a7d630af298abbcf4bcf6370fbefa Mon Sep 17 00:00:00 2001 From: Victor Valbuena Date: Thu, 21 May 2026 13:04:19 -0700 Subject: [PATCH 11/40] FEAT: Add skip_cached cross-run caching to AdversarialBenchmark Adds a skip_cached: bool = False constructor parameter and a thin _get_atomic_attacks_async override on AdversarialBenchmark that, when enabled, filters out atomic-attack candidates whose (atomic_attack_name, technique_eval_hash) tuple appears in any prior COMPLETED ScenarioResult for the same scenario name + VERSION with outcome SUCCESS or FAILURE. ERROR and UNDETERMINED outcomes always retry. Caching is off by default to preserve existing behavior. Built on the AttackResultAttribution primitives introduced in #1758: - AtomicAttack.technique_eval_hash provides the candidate side of the cache key (content-derived via AtomicAttackEvaluationIdentifier). - AttackResultEntry.attribution_data['parent_collection' + 'parent_eval_hash'] provides the persisted side; the executor stamps these per AttackResult, so two atomic attacks sharing a name but using different technique configurations don't cross-pollinate. Defensive behavior: - Missing attribution_data or missing parent_collection -> skip the row silently (treat as not-cached). - Memory exceptions from get_scenario_results / get_attack_results -> log a warning and fall back to no filtering. Caching becomes a no-op rather than blocking the run. - Scenarios in IN_PROGRESS / FAILED / CANCELLED state contribute nothing (no get_attack_results query made for them at all). - Scenario name is matched on type(self).__name__ (PascalCase "AdversarialBenchmark"), aligned with how ScenarioIdentifier stores it; VERSION filter ensures the VERSION bump in the previous commit invalidates old VERSION=1 results for cache purposes (they remain queryable; they just don't suppress fresh runs). Tests: 11 new unit tests (TestAdversarialBenchmarkSkipCachedFilter + TestAdversarialBenchmarkSkipCachedInit) covering filtering semantics, outcome filters, eval-hash disambiguation, scenario-state filter, query-arg shape, missing-attribution defense, memory-error defense, and constructor defaults. Integration test with full persistence round-trip is a separate follow-up commit (F6.3 per plan). Wider regression: 1649/1649 pass across scenario+setup+registry+ backend. Failure mode flagged for the PR description batch: - The override + helper are scenario-agnostic in shape and should probably live on base Scenario behind a duck-typed identity hook (e.g. cls.cache_scope_name() classmethod) so other scenarios (RapidResponse, Scam, etc.) can opt into skip_cached without copy-pasting the wrapper. Enhancement, not a bug; tracked as lift-skip-cached-to-base-scenario. --- .../scenarios/benchmark/adversarial.py | 110 +++++++ .../scenario/benchmark/test_adversarial.py | 277 +++++++++++++++++- 2 files changed, 385 insertions(+), 2 deletions(-) diff --git a/pyrit/scenario/scenarios/benchmark/adversarial.py b/pyrit/scenario/scenarios/benchmark/adversarial.py index 13432d189b..73a32c7a51 100644 --- a/pyrit/scenario/scenarios/benchmark/adversarial.py +++ b/pyrit/scenario/scenarios/benchmark/adversarial.py @@ -9,12 +9,14 @@ from typing import TYPE_CHECKING, ClassVar from pyrit.common import apply_defaults +from pyrit.models import AttackOutcome from pyrit.registry import AttackTechniqueRegistry, AttackTechniqueSpec from pyrit.registry.tag_query import TagQuery from pyrit.scenario.core.dataset_configuration import DatasetConfiguration from pyrit.scenario.core.scenario import BaselineAttackPolicy, Scenario if TYPE_CHECKING: + from pyrit.scenario.core.atomic_attack import AtomicAttack from pyrit.scenario.core.scenario_strategy import ScenarioStrategy from pyrit.score import TrueFalseScorer @@ -218,6 +220,7 @@ def __init__( self, *, objective_scorer: TrueFalseScorer | None = None, + skip_cached: bool = False, scenario_result_id: str | None = None, ) -> None: """ @@ -227,12 +230,23 @@ def __init__( objective_scorer: Scorer for evaluating attack success. Defaults to the registered default objective scorer (typically the composite refusal+scale scorer set up by an initializer). + skip_cached: When ``True``, ``_get_atomic_attacks_async`` filters + out atomic attacks whose ``(atomic_attack_name, + technique_eval_hash)`` tuple already appears in a prior + ``COMPLETED`` ``ScenarioResult`` for the same scenario name + and version with outcome ``SUCCESS`` or ``FAILURE``. + ``ERROR`` and ``UNDETERMINED`` outcomes always retry. Cache + identity is content-derived via + ``AtomicAttack.technique_eval_hash``, so two atomic attacks + with the same name but different technique configurations + (e.g. different scorer) do not cross-pollinate. scenario_result_id: Optional ID of an existing scenario result to resume. """ self._objective_scorer: TrueFalseScorer = ( objective_scorer if objective_scorer else self._get_default_objective_scorer() ) + self._skip_cached: bool = skip_cached super().__init__( version=self.VERSION, @@ -241,6 +255,102 @@ def __init__( scenario_result_id=scenario_result_id, ) + async def _get_atomic_attacks_async(self) -> list[AtomicAttack]: + """ + Build the base set of atomic attacks, then filter out cached completions when requested. + + Delegates to the base ``Scenario._get_atomic_attacks_async`` to + construct the (fanned-variant × dataset) candidate list, then drops + any candidate whose ``(atomic_attack_name, technique_eval_hash)`` + tuple appears in :meth:`_collect_cached_completion_pairs` (only when + ``self._skip_cached`` is ``True``). Always returns the unfiltered + base list when caching is disabled. + + Returns: + list[AtomicAttack]: The atomic attacks to actually execute on + this run. + """ + candidates = await super()._get_atomic_attacks_async() + if not self._skip_cached: + return candidates + + cached_pairs = self._collect_cached_completion_pairs() + filtered = [c for c in candidates if (c.atomic_attack_name, c.technique_eval_hash) not in cached_pairs] + skipped = len(candidates) - len(filtered) + if skipped > 0: + logger.info( + "skip_cached=True: dropping %d/%d atomic attack(s) already completed in prior runs.", + skipped, + len(candidates), + ) + return filtered + + def _collect_cached_completion_pairs(self) -> set[tuple[str, str | None]]: + """ + Collect cache keys for atomic attacks that completed in any prior run of this scenario. + + Walks ``ScenarioResult`` rows for the same scenario name and + ``VERSION``, restricts to ``scenario_run_state == "COMPLETED"``, + then walks the linked ``AttackResult`` rows (joined via + ``AttackResultEntry.attribution_parent_id``) and records the + ``(atomic_attack_name, parent_eval_hash)`` tuple for every + ``SUCCESS`` or ``FAILURE`` outcome. The pair shape mirrors the + ``(atomic_attack_name, technique_eval_hash)`` tuple used by + :meth:`_get_atomic_attacks_async` so a direct ``in`` check filters + candidates without further key construction. + + Resilient to attribution-data variation: rows whose + ``attribution_data`` is ``None`` or missing ``parent_collection`` + are skipped. Rows without ``parent_eval_hash`` enter the cache with + ``None`` in that slot, so they only match candidates whose + ``technique_eval_hash`` also resolves to ``None`` (currently never, + since ``AtomicAttack.technique_eval_hash`` is always populated post-#1758). + + Returns: + set[tuple[str, str | None]]: Cache keys for already-completed + atomic attacks. Empty set on any unexpected error (logged at + warning level) — caching becomes a no-op rather than blocking + the run. + """ + scenario_name = type(self).__name__ + cached_pairs: set[tuple[str, str | None]] = set() + + try: + prior_results = self._memory.get_scenario_results( + scenario_name=scenario_name, + scenario_version=self.VERSION, + ) + except Exception as exc: + logger.warning("skip_cached: failed to query prior scenario results (%s); skipping cache filter.", exc) + return cached_pairs + + for scenario_result in prior_results: + if scenario_result.scenario_run_state != "COMPLETED": + continue + if scenario_result.id is None: + continue + try: + attack_results = self._memory.get_attack_results(scenario_result_id=str(scenario_result.id)) + except Exception as exc: + logger.warning( + "skip_cached: failed to load attack results for scenario %s (%s); skipping that run.", + scenario_result.id, + exc, + ) + continue + + for ar in attack_results: + if ar.outcome not in (AttackOutcome.SUCCESS, AttackOutcome.FAILURE): + continue + data = ar.attribution_data or {} + atomic_attack_name = data.get("parent_collection") + if not atomic_attack_name: + continue + parent_eval_hash = data.get("parent_eval_hash") + cached_pairs.add((atomic_attack_name, parent_eval_hash)) + + return cached_pairs + def _build_display_group(self, *, technique_name: str, seed_group_name: str) -> str: """ Group atomic-attack results by adversarial-target label rather than by technique. diff --git a/tests/unit/scenario/benchmark/test_adversarial.py b/tests/unit/scenario/benchmark/test_adversarial.py index 169361226d..0e0b5788fc 100644 --- a/tests/unit/scenario/benchmark/test_adversarial.py +++ b/tests/unit/scenario/benchmark/test_adversarial.py @@ -6,19 +6,25 @@ AdversarialBenchmark no longer takes an ``adversarial_models`` constructor parameter and no longer builds local factories. It reads fanned variants from ``AttackTechniqueRegistry`` (registered by ``BenchmarkInitializer``) -and inherits the base ``Scenario._get_atomic_attacks_async`` loop. +and inherits the base ``Scenario._get_atomic_attacks_async`` loop, with +an opt-in caching wrapper for cross-run skip-on-completion. These tests cover the new contract: * Class metadata (VERSION, BASELINE policy, defaults). * Strategy enum is built from ``benchmark_fanout``-tagged registry entries. * Display grouping uses the target-label portion of fanned technique names. -* Construction accepts only ``objective_scorer`` and ``scenario_result_id``. +* Construction accepts ``objective_scorer``, ``skip_cached``, and + ``scenario_result_id``. +* ``skip_cached`` filters prior SUCCESS/FAILURE completions, keeps + ERROR/UNDETERMINED, respects eval-hash disambiguation, and only counts + COMPLETED scenario runs of the matching name + version. """ from unittest.mock import MagicMock, patch import pytest +from pyrit.models import AttackOutcome from pyrit.prompt_target import PromptTarget from pyrit.registry import TargetRegistry from pyrit.registry.object_registries.attack_technique_registry import AttackTechniqueRegistry @@ -266,3 +272,270 @@ async def test_ignores_seed_group_name(self): seed_group_name="seed_group_b", ) assert first == second == "adv_a" + + +# --------------------------------------------------------------------------- +# skip_cached behavior (Commit 6 / F3) +# --------------------------------------------------------------------------- + + +def _make_scenario_result(*, result_id: str, run_state: str = "COMPLETED") -> MagicMock: + """Build a minimal ScenarioResult stand-in for cache-key tests.""" + sr = MagicMock() + sr.id = result_id + sr.scenario_run_state = run_state + return sr + + +def _make_attack_result( + *, + outcome: AttackOutcome, + parent_collection: str | None, + parent_eval_hash: str | None, +) -> MagicMock: + """Build a minimal AttackResult stand-in with the attribution_data shape Commit 6 reads.""" + ar = MagicMock() + ar.outcome = outcome + if parent_collection is None and parent_eval_hash is None: + ar.attribution_data = None + else: + data: dict[str, str] = {} + if parent_collection is not None: + data["parent_collection"] = parent_collection + if parent_eval_hash is not None: + data["parent_eval_hash"] = parent_eval_hash + ar.attribution_data = data + return ar + + +def _make_candidate(*, name: str, eval_hash: str) -> MagicMock: + """Build a minimal AtomicAttack stand-in with the two fields the cache filter reads.""" + candidate = MagicMock() + candidate.atomic_attack_name = name + candidate.technique_eval_hash = eval_hash + return candidate + + +@pytest.mark.usefixtures("patch_central_database") +class TestAdversarialBenchmarkSkipCachedFilter: + """Tests for the _get_atomic_attacks_async caching wrapper.""" + + async def _make_bench(self, *, skip_cached: bool) -> AdversarialBenchmark: + await _fan_out(target_names=["adv_a"]) + return AdversarialBenchmark( + objective_scorer=MagicMock(spec=TrueFalseScorer), + skip_cached=skip_cached, + ) + + async def test_skip_cached_default_false_means_no_filtering(self): + """skip_cached defaults to False; super() output is returned unchanged.""" + bench = await self._make_bench(skip_cached=False) + candidates = [_make_candidate(name="red_teaming__adv_a", eval_hash="hash_a")] + + with patch.object( + AdversarialBenchmark.__bases__[0], "_get_atomic_attacks_async", return_value=candidates + ) as super_mock: + result = await bench._get_atomic_attacks_async() + + assert result == candidates + super_mock.assert_awaited_once() + + async def test_skip_cached_true_drops_completed_pairs(self): + """SUCCESS and FAILURE prior outcomes drop the matching candidate.""" + bench = await self._make_bench(skip_cached=True) + + candidates = [ + _make_candidate(name="red_teaming__adv_a", eval_hash="hash_a"), + _make_candidate(name="tap__adv_a", eval_hash="hash_b"), + _make_candidate(name="crescendo_simulated__adv_a", eval_hash="hash_c"), + ] + prior_sr = _make_scenario_result(result_id="sid-1") + prior_attacks = [ + _make_attack_result( + outcome=AttackOutcome.SUCCESS, + parent_collection="red_teaming__adv_a", + parent_eval_hash="hash_a", + ), + _make_attack_result( + outcome=AttackOutcome.FAILURE, + parent_collection="tap__adv_a", + parent_eval_hash="hash_b", + ), + ] + + bench._memory = MagicMock() + bench._memory.get_scenario_results.return_value = [prior_sr] + bench._memory.get_attack_results.return_value = prior_attacks + + with patch.object(AdversarialBenchmark.__bases__[0], "_get_atomic_attacks_async", return_value=candidates): + result = await bench._get_atomic_attacks_async() + + names = [c.atomic_attack_name for c in result] + assert names == ["crescendo_simulated__adv_a"] + + async def test_skip_cached_keeps_error_outcomes(self): + """ERROR outcomes must retry — not be cached.""" + bench = await self._make_bench(skip_cached=True) + + candidates = [_make_candidate(name="red_teaming__adv_a", eval_hash="hash_a")] + prior_sr = _make_scenario_result(result_id="sid-1") + prior_attacks = [ + _make_attack_result( + outcome=AttackOutcome.ERROR, + parent_collection="red_teaming__adv_a", + parent_eval_hash="hash_a", + ), + ] + + bench._memory = MagicMock() + bench._memory.get_scenario_results.return_value = [prior_sr] + bench._memory.get_attack_results.return_value = prior_attacks + + with patch.object(AdversarialBenchmark.__bases__[0], "_get_atomic_attacks_async", return_value=candidates): + result = await bench._get_atomic_attacks_async() + + assert result == candidates + + async def test_skip_cached_keeps_undetermined_outcomes(self): + """UNDETERMINED outcomes must retry — not be cached.""" + bench = await self._make_bench(skip_cached=True) + + candidates = [_make_candidate(name="red_teaming__adv_a", eval_hash="hash_a")] + prior_sr = _make_scenario_result(result_id="sid-1") + prior_attacks = [ + _make_attack_result( + outcome=AttackOutcome.UNDETERMINED, + parent_collection="red_teaming__adv_a", + parent_eval_hash="hash_a", + ), + ] + + bench._memory = MagicMock() + bench._memory.get_scenario_results.return_value = [prior_sr] + bench._memory.get_attack_results.return_value = prior_attacks + + with patch.object(AdversarialBenchmark.__bases__[0], "_get_atomic_attacks_async", return_value=candidates): + result = await bench._get_atomic_attacks_async() + + assert result == candidates + + async def test_skip_cached_respects_eval_hash_disambiguation(self): + """Same atomic_attack_name but different parent_eval_hash → not considered cached.""" + bench = await self._make_bench(skip_cached=True) + + candidates = [_make_candidate(name="red_teaming__adv_a", eval_hash="new_hash")] + prior_sr = _make_scenario_result(result_id="sid-1") + prior_attacks = [ + _make_attack_result( + outcome=AttackOutcome.SUCCESS, + parent_collection="red_teaming__adv_a", + parent_eval_hash="old_hash", + ), + ] + + bench._memory = MagicMock() + bench._memory.get_scenario_results.return_value = [prior_sr] + bench._memory.get_attack_results.return_value = prior_attacks + + with patch.object(AdversarialBenchmark.__bases__[0], "_get_atomic_attacks_async", return_value=candidates): + result = await bench._get_atomic_attacks_async() + + assert result == candidates + + async def test_skip_cached_only_considers_completed_scenarios(self): + """Scenarios in IN_PROGRESS / FAILED / CANCELLED state must not seed the cache.""" + bench = await self._make_bench(skip_cached=True) + + candidates = [_make_candidate(name="red_teaming__adv_a", eval_hash="hash_a")] + in_progress = _make_scenario_result(result_id="sid-1", run_state="IN_PROGRESS") + failed = _make_scenario_result(result_id="sid-2", run_state="FAILED") + + bench._memory = MagicMock() + bench._memory.get_scenario_results.return_value = [in_progress, failed] + bench._memory.get_attack_results.return_value = [ + _make_attack_result( + outcome=AttackOutcome.SUCCESS, + parent_collection="red_teaming__adv_a", + parent_eval_hash="hash_a", + ), + ] + + with patch.object(AdversarialBenchmark.__bases__[0], "_get_atomic_attacks_async", return_value=candidates): + result = await bench._get_atomic_attacks_async() + + assert result == candidates + bench._memory.get_attack_results.assert_not_called() + + async def test_skip_cached_filters_by_scenario_name_and_version(self): + """get_scenario_results is queried with this scenario's name + VERSION; old VERSION=1 results don't apply.""" + bench = await self._make_bench(skip_cached=True) + + bench._memory = MagicMock() + bench._memory.get_scenario_results.return_value = [] + + with patch.object(AdversarialBenchmark.__bases__[0], "_get_atomic_attacks_async", return_value=[]): + await bench._get_atomic_attacks_async() + + bench._memory.get_scenario_results.assert_called_once_with( + scenario_name="AdversarialBenchmark", + scenario_version=AdversarialBenchmark.VERSION, + ) + + async def test_skip_cached_handles_missing_attribution_data(self): + """Rows with attribution_data=None or missing parent_collection are silently skipped.""" + bench = await self._make_bench(skip_cached=True) + + candidates = [_make_candidate(name="red_teaming__adv_a", eval_hash="hash_a")] + prior_sr = _make_scenario_result(result_id="sid-1") + prior_attacks = [ + _make_attack_result( + outcome=AttackOutcome.SUCCESS, + parent_collection=None, + parent_eval_hash=None, + ), + _make_attack_result( + outcome=AttackOutcome.SUCCESS, + parent_collection=None, + parent_eval_hash="hash_x", + ), + ] + + bench._memory = MagicMock() + bench._memory.get_scenario_results.return_value = [prior_sr] + bench._memory.get_attack_results.return_value = prior_attacks + + with patch.object(AdversarialBenchmark.__bases__[0], "_get_atomic_attacks_async", return_value=candidates): + result = await bench._get_atomic_attacks_async() + + assert result == candidates + + async def test_skip_cached_memory_error_falls_back_to_no_filter(self): + """An exception from get_scenario_results must not block the run — return base candidates as-is.""" + bench = await self._make_bench(skip_cached=True) + + candidates = [_make_candidate(name="red_teaming__adv_a", eval_hash="hash_a")] + bench._memory = MagicMock() + bench._memory.get_scenario_results.side_effect = RuntimeError("db down") + + with patch.object(AdversarialBenchmark.__bases__[0], "_get_atomic_attacks_async", return_value=candidates): + result = await bench._get_atomic_attacks_async() + + assert result == candidates + + +@pytest.mark.usefixtures("patch_central_database") +class TestAdversarialBenchmarkSkipCachedInit: + """Tests for the skip_cached constructor surface.""" + + async def test_skip_cached_defaults_to_false(self): + await _fan_out(target_names=["adv_a"]) + bench = AdversarialBenchmark(objective_scorer=MagicMock(spec=TrueFalseScorer)) + assert bench._skip_cached is False + + async def test_skip_cached_can_be_set_true(self): + await _fan_out(target_names=["adv_a"]) + bench = AdversarialBenchmark( + objective_scorer=MagicMock(spec=TrueFalseScorer), + skip_cached=True, + ) + assert bench._skip_cached is True From 091a501f9d7cdc485818520a0be4f07a797109f6 Mon Sep 17 00:00:00 2001 From: Victor Valbuena Date: Thu, 21 May 2026 13:10:53 -0700 Subject: [PATCH 12/40] FEAT: Widen AdversarialBenchmark objective_scorer to Scorer (stage 1) Stage 1 of the scorer-flexibility refactor: widens the parameter annotation on AdversarialBenchmark.__init__ from TrueFalseScorer | None to Scorer | None, while preserving the existing runtime contract via an isinstance(resolved, TrueFalseScorer) guard that raises TypeError with a pointer at the new-scoring follow-up. Forward-compatible: when stage 2 lands and AttackScoringConfig + atomic-attack types are widened to Scorer, removing the guard is the only change needed here. Why widen the annotation now rather than wait for stage 2: - Lets the follow-up PR be a behavior change (drop the guard, wire the new scorer path through AttackScoringConfig) without a parameter signature change. Users coding to AdversarialBenchmark.__init__'s signature see the eventual contract today. - Self-documents the planned direction in IDE tooling and --list-scenarios output. - TypeError message names the constraint AND points readers at the follow-up so the broken case isn't silent. Out of scope (stage 2, separate follow-up): - AttackScoringConfig.objective_scorer widening - Atomic attack types' objective_scorer widening - pyrit.scenario.core.scenario casts at lines :778, :990, :1034 - Removing this guard Tests: 3 new (test_objective_scorer_annotation_is_scorer, test_construct_accepts_truefalse_scorer_subclass, test_non_truefalse_scorer_raises_typeerror_with_pointer). Existing default-scorer / explicit-scorer init tests already cover the happy TrueFalseScorer path. Wider regression: 1652/1652 pass across scenario+setup+registry+ backend. Pre-commit clean. --- .../scenarios/benchmark/adversarial.py | 36 +++++++++++---- .../scenario/benchmark/test_adversarial.py | 44 +++++++++++++++++++ 2 files changed, 72 insertions(+), 8 deletions(-) diff --git a/pyrit/scenario/scenarios/benchmark/adversarial.py b/pyrit/scenario/scenarios/benchmark/adversarial.py index 73a32c7a51..9f9fd9cdf4 100644 --- a/pyrit/scenario/scenarios/benchmark/adversarial.py +++ b/pyrit/scenario/scenarios/benchmark/adversarial.py @@ -14,11 +14,12 @@ from pyrit.registry.tag_query import TagQuery from pyrit.scenario.core.dataset_configuration import DatasetConfiguration from pyrit.scenario.core.scenario import BaselineAttackPolicy, Scenario +from pyrit.score.true_false.true_false_scorer import TrueFalseScorer if TYPE_CHECKING: from pyrit.scenario.core.atomic_attack import AtomicAttack from pyrit.scenario.core.scenario_strategy import ScenarioStrategy - from pyrit.score import TrueFalseScorer + from pyrit.score import Scorer logger = logging.getLogger(__name__) @@ -219,7 +220,7 @@ def default_dataset_config(cls) -> DatasetConfiguration: def __init__( self, *, - objective_scorer: TrueFalseScorer | None = None, + objective_scorer: Scorer | None = None, skip_cached: bool = False, scenario_result_id: str | None = None, ) -> None: @@ -227,9 +228,17 @@ def __init__( Initialize the AdversarialBenchmark scenario. Args: - objective_scorer: Scorer for evaluating attack success. Defaults - to the registered default objective scorer (typically the - composite refusal+scale scorer set up by an initializer). + objective_scorer: Scorer for evaluating attack success. The + annotation is the broad ``Scorer`` base class for forward + compatibility with the planned non-``TrueFalseScorer`` + scoring follow-up (see PR description / follow-up issue + tracker), but the runtime contract is currently still + ``TrueFalseScorer``: any other ``Scorer`` subclass raises + ``TypeError`` at construction with a message pointing at + the follow-up. Defaults to the registered default objective + scorer (typically the composite refusal+scale scorer set + up by an initializer), which is always a + ``TrueFalseScorer``. skip_cached: When ``True``, ``_get_atomic_attacks_async`` filters out atomic attacks whose ``(atomic_attack_name, technique_eval_hash)`` tuple already appears in a prior @@ -242,10 +251,21 @@ def __init__( (e.g. different scorer) do not cross-pollinate. scenario_result_id: Optional ID of an existing scenario result to resume. + + Raises: + TypeError: If ``objective_scorer`` is a ``Scorer`` subclass + other than ``TrueFalseScorer`` (full non-true/false support + is tracked as a follow-up; the type annotation is widened + ahead of the runtime support). """ - self._objective_scorer: TrueFalseScorer = ( - objective_scorer if objective_scorer else self._get_default_objective_scorer() - ) + resolved_scorer: Scorer = objective_scorer if objective_scorer else self._get_default_objective_scorer() + if not isinstance(resolved_scorer, TrueFalseScorer): + raise TypeError( + f"AdversarialBenchmark currently requires a TrueFalseScorer for objective_scorer; " + f"got {type(resolved_scorer).__name__}. Full Scorer support (e.g. FloatScaleScorer) is " + f"tracked as the new-scoring follow-up — see the PR description for status." + ) + self._objective_scorer: TrueFalseScorer = resolved_scorer self._skip_cached: bool = skip_cached super().__init__( diff --git a/tests/unit/scenario/benchmark/test_adversarial.py b/tests/unit/scenario/benchmark/test_adversarial.py index 0e0b5788fc..4bc22ce072 100644 --- a/tests/unit/scenario/benchmark/test_adversarial.py +++ b/tests/unit/scenario/benchmark/test_adversarial.py @@ -539,3 +539,47 @@ async def test_skip_cached_can_be_set_true(self): skip_cached=True, ) assert bench._skip_cached is True + + +# --------------------------------------------------------------------------- +# Scorer flexibility — stage 1 (Commit 7 / F4) +# --------------------------------------------------------------------------- + + +@pytest.mark.usefixtures("patch_central_database") +class TestAdversarialBenchmarkScorerFlexibility: + """Tests for the widened objective_scorer annotation + isinstance guard (stage 1).""" + + def test_objective_scorer_annotation_is_scorer(self): + """The parameter annotation is the broad Scorer base class for forward compatibility.""" + import inspect + + from pyrit.score import Scorer + + sig = inspect.signature(AdversarialBenchmark.__init__) + annotation = sig.parameters["objective_scorer"].annotation + # ``Scorer | None`` resolves to ``Scorer | None`` at import time. + # str() captures both the runtime and stringified forms reliably. + assert "Scorer" in str(annotation) + assert Scorer is not None # sanity that the import resolves + + async def test_construct_accepts_truefalse_scorer_subclass(self): + """TrueFalseScorer remains the runtime-supported type; should construct cleanly.""" + await _fan_out(target_names=["adv_a"]) + + scorer = MagicMock(spec=TrueFalseScorer) + bench = AdversarialBenchmark(objective_scorer=scorer) + + assert bench._objective_scorer is scorer + + async def test_non_truefalse_scorer_raises_typeerror_with_pointer(self): + """A Scorer subclass that isn't TrueFalseScorer must raise TypeError with a clear pointer.""" + from pyrit.score import Scorer + + await _fan_out(target_names=["adv_a"]) + + # Bare Scorer (not TrueFalseScorer) — covers any future non-TF subclass. + non_tf_scorer = MagicMock(spec=Scorer) + + with pytest.raises(TypeError, match=r"requires a TrueFalseScorer.*follow-up"): + AdversarialBenchmark(objective_scorer=non_tf_scorer) From 40ff08be7bc9a811536088b69cc31620b795631f Mon Sep 17 00:00:00 2001 From: Victor Valbuena Date: Thu, 21 May 2026 13:16:02 -0700 Subject: [PATCH 13/40] TEST: e2e per-scenario initializer override for benchmark.adversarial Adds DEFAULT_INITIALIZERS + SCENARIO_INITIALIZERS to tests/end_to_end/test_scenarios.py so scenarios that need scenario- specific initialization (post-collapse benchmark.adversarial needs BenchmarkInitializer to fan adversarial techniques across registry- discovered targets) can opt into a longer initializer list without forcing every other scenario to load the same extras. Default for every scenario: ["target", "load_default_datasets"] (unchanged from prior behavior). Override for benchmark.adversarial: defaults + ["benchmark"], so BenchmarkInitializer runs after TargetInitializer has populated TargetRegistry with the ADVERSARIAL-tagged env-driven targets. Plan-vs-reality fix caught during implementation: the plan referred to the scenario key as "adversarial_benchmark", but the actual ScenarioRegistry name (used by pyrit_scan) is the dotted module path "benchmark.adversarial", mirroring "airt.cyber" / "garak.encoding". The override map uses the dotted form. Comment in the file pins the convention so future overrides don't hit the same gotcha. E2e tests are not part of CI; they run via make end-to-end-test on developer machines that have ADVERSARIAL_CHAT_* env vars set. When the env vars are absent, BenchmarkInitializer surfaces the actionable error message added in Commit 4 (closes failure_mode_followup no-adversarial-model-clear-error, also referenced in Commit 4 body). No regression run included: e2e tests require live API credentials. Smoke-tested with pytest --collect-only (9 scenarios, including benchmark.adversarial) and a manual resolution check that _initializers_for("benchmark.adversarial") returns the override list while _initializers_for("airt.cyber") falls back to defaults. --- tests/end_to_end/test_scenarios.py | 40 +++++++++++++++++++++++++++--- 1 file changed, 36 insertions(+), 4 deletions(-) diff --git a/tests/end_to_end/test_scenarios.py b/tests/end_to_end/test_scenarios.py index 8b40471715..4f682cbc5f 100644 --- a/tests/end_to_end/test_scenarios.py +++ b/tests/end_to_end/test_scenarios.py @@ -5,7 +5,17 @@ End-to-end tests for PyRIT scenarios using pyrit_scan CLI. These tests dynamically discover all available scenarios and run each one -using the pyrit_scan command with standard initializers. +using the pyrit_scan command. Most scenarios run with the +:data:`DEFAULT_INITIALIZERS` list; scenarios that need additional setup +(e.g. ``benchmark.adversarial`` needs ``BenchmarkInitializer`` to fan +adversarial techniques out across registry-discovered targets) declare +their full initializer list in :data:`SCENARIO_INITIALIZERS`. + +Note: e2e tests are not part of CI; they run via ``make end-to-end-test`` +on developer machines that have the appropriate env vars set +(``ADVERSARIAL_CHAT_*`` for the benchmark scenario, in particular). +``BenchmarkInitializer`` surfaces a clear error pointing at the env vars +when they are absent. """ from pathlib import Path @@ -17,6 +27,23 @@ CONFIG_FILE = Path(__file__).parent / "test_config.yaml" +#: Initializers run for every scenario unless overridden in :data:`SCENARIO_INITIALIZERS`. +#: ``target`` populates ``TargetRegistry`` from env vars; ``load_default_datasets`` +#: fetches each scenario's declared default datasets into memory. +DEFAULT_INITIALIZERS: list[str] = ["target", "load_default_datasets"] + +#: Per-scenario override map. A scenario named here uses this list verbatim +#: (no implicit merge with ``DEFAULT_INITIALIZERS``); a scenario absent here +#: falls back to ``DEFAULT_INITIALIZERS``. Keys use the dotted registry name +#: (``.``) returned by ``ScenarioRegistry.get_names()``. +SCENARIO_INITIALIZERS: dict[str, list[str]] = { + # benchmark.adversarial depends on BenchmarkInitializer to fan + # adversarial-capable scenario techniques out across every + # ADVERSARIAL-tagged target in TargetRegistry. Without the + # benchmark initializer, the scenario's strategy enum is empty. + "benchmark.adversarial": [*DEFAULT_INITIALIZERS, "benchmark"], +} + def get_all_scenarios(): """ @@ -29,22 +56,27 @@ def get_all_scenarios(): return registry.get_names() +def _initializers_for(scenario_name: str) -> list[str]: + """Return the initializer name list for ``scenario_name``, defaulting to ``DEFAULT_INITIALIZERS``.""" + return SCENARIO_INITIALIZERS.get(scenario_name, DEFAULT_INITIALIZERS) + + @pytest.mark.timeout(7200) # 2 hour timeout per scenario @pytest.mark.parametrize("scenario_name", get_all_scenarios()) def test_scenario_with_pyrit_scan(scenario_name): """ - Test each scenario runs successfully using pyrit_scan with standard initializers. + Test each scenario runs successfully using pyrit_scan with its declared initializer list. Args: scenario_name: Name of the scenario to test (dynamically discovered). """ + initializers = _initializers_for(scenario_name) try: result = pyrit_scan_main( [ scenario_name, "--initializers", - "target", - "load_default_datasets", + *initializers, "--target", "openai_chat", "--config-file", From 43233839d2aa23fad933c8ea27cafa16d88a3c74 Mon Sep 17 00:00:00 2001 From: Victor Valbuena Date: Thu, 21 May 2026 13:19:44 -0700 Subject: [PATCH 14/40] DOCS: Rewrite benchmark scanner notebook for registry-driven flow Rewrites doc/scanner/benchmark.{py,ipynb} end-to-end around the new registry-driven flow. The previous notebook constructed AdversarialBenchmark with adversarial_models=[OpenAIChatTarget()], which no longer exists after the collapse. New notebook content: - Prerequisites: ADVERSARIAL_CHAT_* env vars (plus optional _SINGLETURN / _MULTITURN / _REASONING variants). - CLI quickstart: pyrit_scan benchmark.adversarial --initializers target load_default_datasets benchmark --target openai_chat ... - Setup cell: initialize_pyrit_async with TargetInitializer + ScorerInitializer + LoadDefaultDatasets + BenchmarkInitializer (in that order, since BenchmarkInitializer reads TargetRegistry). - Run cell: AdversarialBenchmark() with no model args; default "light" strategy. - Cross-run caching cell: AdversarialBenchmark(skip_cached=True); documents (atomic_attack_name, technique_eval_hash) cache key, SUCCESS/FAILURE-only caching, ERROR/UNDETERMINED retry semantics, and the "add new adversarial targets incrementally" use case. - Narrowing the fan-out: BenchmarkInitializer.set_params_from_args with target_names = [...]. - .pyrit_conf bootstrap: full YAML with initializer ordering. - Scorer flexibility: documents the widened Scorer | None annotation and the TrueFalseScorer-only runtime contract for stage 1. Per microsoft.github.io/PyRIT/contributing/notebooks/ the .ipynb is generated from the .py via jupytext; this commit regenerates the .ipynb to match the new .py source. Both committed pre-execution (no real output cells). Maintainers running pct_to_ipynb.py before the next release will re-execute against real endpoints; doing so now would require ADVERSARIAL_CHAT_* env vars set in this dev env. Downstream: the published doc page at https://microsoft.github.io/PyRIT/scanner/benchmark/ is built from doc/scanner/benchmark.py and will update automatically when this PR merges into main. Companion test (smoke-run the notebook with mocked targets) is deferred to its own follow-up commit (F6.2 / scanner-notebook-test per plan). --- doc/scanner/benchmark.ipynb | 321 +++++++++++++++++++++++------------- doc/scanner/benchmark.py | 179 ++++++++++++++++++-- 2 files changed, 374 insertions(+), 126 deletions(-) diff --git a/doc/scanner/benchmark.ipynb b/doc/scanner/benchmark.ipynb index 9da0510b4c..193f0388d1 100644 --- a/doc/scanner/benchmark.ipynb +++ b/doc/scanner/benchmark.ipynb @@ -7,7 +7,10 @@ "source": [ "# Benchmark Scenarios\n", "\n", - "Benchmark scenarios are a subset of scenarios that compare the effectiveness of attacks across an axis that varies within the scenario itself. The axis can be many things; currently, the only benchmark variant is the adversarial benchmark, whose axis of change is the adversarial model used in attacks." + "Benchmark scenarios compare the effectiveness of attacks across an axis that varies within the\n", + "scenario itself. The axis can be many things; currently, the only benchmark variant is the\n", + "adversarial benchmark, whose axis of change is the **adversarial chat helper model** used in\n", + "attacks." ] }, { @@ -16,143 +19,239 @@ "metadata": {}, "source": [ "## Adversarial Benchmark\n", - "The adversarial benchmarking scenario (`AdversarialBenchmark`) compares the effectiveness of different adversarial models in successfully executing attacks against a target model." + "\n", + "`AdversarialBenchmark` holds the objective target and dataset constant and varies the adversarial\n", + "chat model used to drive multi-turn attacks (and crescendo-style simulated conversations). Useful\n", + "for evaluating which adversarial helper models produce stronger or weaker attack success rates\n", + "against the same target.\n", + "\n", + "Model fan-out is owned by `BenchmarkInitializer`, which discovers every ADVERSARIAL-tagged target\n", + "in `TargetRegistry` (populated by `TargetInitializer` from `ADVERSARIAL_CHAT_*` env vars) and\n", + "registers one fanned variant per `(adversarial-capable technique, target)` pair into\n", + "`AttackTechniqueRegistry`. The scenario then reads those variants when it builds its strategy\n", + "enum.\n", + "\n", + "### Prerequisites\n", + "\n", + "Set at least one `ADVERSARIAL_CHAT_*` group of env vars (see `.env_example`):\n", + "\n", + "```bash\n", + "# Default adversarial target (always available when set)\n", + "ADVERSARIAL_CHAT_ENDPOINT=\"https://your-endpoint.openai.azure.com/openai/v1\"\n", + "ADVERSARIAL_CHAT_KEY=\"your-key\"\n", + "ADVERSARIAL_CHAT_MODEL=\"deployment-name\"\n", + "\n", + "# Optional turn-style variants — auto-discovered by TargetInitializer when set\n", + "ADVERSARIAL_CHAT_SINGLETURN_ENDPOINT=\"...\"\n", + "ADVERSARIAL_CHAT_SINGLETURN_KEY=\"...\"\n", + "ADVERSARIAL_CHAT_SINGLETURN_MODEL=\"...\"\n", + "# ADVERSARIAL_CHAT_MULTITURN_* and ADVERSARIAL_CHAT_REASONING_* follow the same pattern\n", + "```\n", + "\n", + "If no adversarial-tagged target is registered, `BenchmarkInitializer.initialize_async` raises\n", + "`ValueError` naming the env vars to set.\n", + "\n", + "### CLI quickstart\n", + "\n", + "```bash\n", + "pyrit_scan benchmark.adversarial \\\n", + " --initializers target load_default_datasets benchmark \\\n", + " --target openai_chat \\\n", + " --max-dataset-size 4\n", + "```\n", + "\n", + "**Available strategies (depend on registered adversarial targets):** `all`, `light`, `single_turn`,\n", + "`multi_turn`, plus one concrete member per fanned variant named\n", + "`f\"{source_technique}__{target_name}\"` (e.g. `red_teaming__adversarial_chat`,\n", + "`tap__adversarial_chat_singleturn`)." + ] + }, + { + "cell_type": "markdown", + "id": "2", + "metadata": {}, + "source": [ + "## Setup\n", + "\n", + "`BenchmarkInitializer` runs after `TargetInitializer` so adversarial-tagged targets are present\n", + "in the registry before the fan-out logic queries it. The initializer chain is executed in list\n", + "order." ] }, { "cell_type": "code", "execution_count": null, - "id": "2", + "id": "3", "metadata": {}, - "outputs": [ - { - "name": "stdout", - "output_type": "stream", - "text": [ - "Found default environment files: ['./.pyrit/.env', './.pyrit/.env.local']\n", - "Loaded environment file: ./.pyrit/.env\n", - "Loaded environment file: ./.pyrit/.env.local\n" - ] - }, - { - "data": { - "application/vnd.jupyter.widget-view+json": { - "model_id": "8316db039ba1408499df0a2de6c8d6f6", - "version_major": 2, - "version_minor": 0 - }, - "text/plain": [ - "Executing AdversarialBenchmark: 0%| | 0/3 [00:00)` to pick\n", - "# up where this run left off (constructor args must match the original run).\n", + "# Save this id to resume the run later via AdversarialBenchmark(scenario_result_id=...).\n", "print(f\"Scenario result id: {baseline_result.id}\")\n", "\n", - "\n", "await output_scenario_async(baseline_result)" ] + }, + { + "cell_type": "markdown", + "id": "6", + "metadata": {}, + "source": [ + "## Cross-run caching\n", + "\n", + "Re-run the benchmark with `skip_cached=True` and atomic attacks that completed (`SUCCESS` or\n", + "`FAILURE` outcome) in any prior `COMPLETED` run of the same scenario name + version are skipped.\n", + "`ERROR` and `UNDETERMINED` outcomes always retry, so transient failures don't poison the cache.\n", + "\n", + "The cache key is `(atomic_attack_name, technique_eval_hash)` — two atomic attacks that share a\n", + "name but use different technique configurations (e.g. different scorer) don't cross-pollinate.\n", + "\n", + "Useful for:\n", + "* Resuming a long-running benchmark after a crash or `Ctrl-C`.\n", + "* Incrementally adding new adversarial targets without re-running the existing ones — the new\n", + " fanned variants have new names, so they don't match any cached entry and execute fresh." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "7", + "metadata": {}, + "outputs": [], + "source": [ + "scenario_cached = AdversarialBenchmark(skip_cached=True)\n", + "await scenario_cached.initialize_async( # type: ignore\n", + " objective_target=objective_target,\n", + " dataset_config=dataset_config,\n", + ")\n", + "\n", + "cached_result = await scenario_cached.run_async() # type: ignore\n", + "\n", + "await output_scenario_async(cached_result)" + ] + }, + { + "cell_type": "markdown", + "id": "8", + "metadata": {}, + "source": [ + "## Narrowing the fan-out\n", + "\n", + "Limit the benchmark to a specific subset of registered adversarial targets via\n", + "`BenchmarkInitializer`'s `target_names` parameter. Settable from `.pyrit_conf` (see next section),\n", + "or programmatically:\n", + "\n", + "```python\n", + "narrow_init = BenchmarkInitializer()\n", + "narrow_init.set_params_from_args(args={\"target_names\": [\"adversarial_chat_singleturn\"]})\n", + "await narrow_init.initialize_async()\n", + "```\n", + "\n", + "Unknown names raise `ValueError` listing both the unknowns and the discovered set, so typos fail\n", + "loudly rather than silently expanding the run." + ] + }, + { + "cell_type": "markdown", + "id": "9", + "metadata": {}, + "source": [ + "## Bootstrapping from `.pyrit_conf`\n", + "\n", + "For production / repeated runs, declare the full initializer chain in `.pyrit_conf`. Initializers\n", + "run in list order; `target` must precede `benchmark` so adversarial-tagged targets are present\n", + "when `BenchmarkInitializer` queries the registry.\n", + "\n", + "```yaml\n", + "memory_db_type: duckdb\n", + "initializers:\n", + " - name: target\n", + " - name: scorer\n", + " - name: load_default_datasets\n", + " - name: benchmark\n", + " args:\n", + " target_names:\n", + " - adversarial_chat_singleturn\n", + " - adversarial_chat_reasoning\n", + "scenario:\n", + " name: benchmark.adversarial\n", + "```\n", + "\n", + "Then run `pyrit_scan --config-file .pyrit_conf` — the scenario picks up the initializer-registered\n", + "fan-out automatically." + ] + }, + { + "cell_type": "markdown", + "id": "10", + "metadata": {}, + "source": [ + "## Scorer flexibility (forward-looking)\n", + "\n", + "`AdversarialBenchmark.__init__`'s `objective_scorer` parameter is typed as the broad `Scorer`\n", + "base class. The current runtime contract is still `TrueFalseScorer` — passing a non-true/false\n", + "scorer raises `TypeError` with a pointer to the planned scoring follow-up. The annotation is\n", + "widened ahead of runtime support so callers coding against the signature today see the eventual\n", + "contract, and the follow-up that drops the guard won't require a signature change." + ] } ], "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.13.13" } }, "nbformat": 4, diff --git a/doc/scanner/benchmark.py b/doc/scanner/benchmark.py index 0f9cedc6db..9a3efa1f44 100644 --- a/doc/scanner/benchmark.py +++ b/doc/scanner/benchmark.py @@ -11,37 +11,186 @@ # %% [markdown] # # Benchmark Scenarios # -# Benchmark scenarios are a subset of scenarios that compare the effectiveness of attacks across an axis that varies within the scenario itself. The axis can be many things; currently, the only benchmark variant is the adversarial benchmark, whose axis of change is the adversarial model used in attacks. +# Benchmark scenarios compare the effectiveness of attacks across an axis that varies within the +# scenario itself. The axis can be many things; currently, the only benchmark variant is the +# adversarial benchmark, whose axis of change is the **adversarial chat helper model** used in +# attacks. # %% [markdown] # ## Adversarial Benchmark -# The adversarial benchmarking scenario (`AdversarialBenchmark`) compares the effectiveness of different adversarial models in successfully executing attacks against a target model. +# +# `AdversarialBenchmark` holds the objective target and dataset constant and varies the adversarial +# chat model used to drive multi-turn attacks (and crescendo-style simulated conversations). Useful +# for evaluating which adversarial helper models produce stronger or weaker attack success rates +# against the same target. +# +# Model fan-out is owned by `BenchmarkInitializer`, which discovers every ADVERSARIAL-tagged target +# in `TargetRegistry` (populated by `TargetInitializer` from `ADVERSARIAL_CHAT_*` env vars) and +# registers one fanned variant per `(adversarial-capable technique, target)` pair into +# `AttackTechniqueRegistry`. The scenario then reads those variants when it builds its strategy +# enum. +# +# ### Prerequisites +# +# Set at least one `ADVERSARIAL_CHAT_*` group of env vars (see `.env_example`): +# +# ```bash +# # Default adversarial target (always available when set) +# ADVERSARIAL_CHAT_ENDPOINT="https://your-endpoint.openai.azure.com/openai/v1" +# ADVERSARIAL_CHAT_KEY="your-key" +# ADVERSARIAL_CHAT_MODEL="deployment-name" +# +# # Optional turn-style variants — auto-discovered by TargetInitializer when set +# ADVERSARIAL_CHAT_SINGLETURN_ENDPOINT="..." +# ADVERSARIAL_CHAT_SINGLETURN_KEY="..." +# ADVERSARIAL_CHAT_SINGLETURN_MODEL="..." +# # ADVERSARIAL_CHAT_MULTITURN_* and ADVERSARIAL_CHAT_REASONING_* follow the same pattern +# ``` +# +# If no adversarial-tagged target is registered, `BenchmarkInitializer.initialize_async` raises +# `ValueError` naming the env vars to set. +# +# ### CLI quickstart +# +# ```bash +# pyrit_scan benchmark.adversarial \ +# --initializers target load_default_datasets benchmark \ +# --target openai_chat \ +# --max-dataset-size 4 +# ``` +# +# **Available strategies (depend on registered adversarial targets):** `all`, `light`, `single_turn`, +# `multi_turn`, plus one concrete member per fanned variant named +# `f"{source_technique}__{target_name}"` (e.g. `red_teaming__adversarial_chat`, +# `tap__adversarial_chat_singleturn`). + +# %% [markdown] +# ## Setup +# +# `BenchmarkInitializer` runs after `TargetInitializer` so adversarial-tagged targets are present +# in the registry before the fan-out logic queries it. The initializer chain is executed in list +# order. # %% from pyrit.output import output_scenario_async from pyrit.prompt_target import OpenAIChatTarget +from pyrit.scenario import DatasetConfiguration from pyrit.scenario.scenarios.benchmark import AdversarialBenchmark from pyrit.setup import IN_MEMORY, initialize_pyrit_async -from pyrit.setup.initializers import LoadDefaultDatasets +from pyrit.setup.initializers import ( + BenchmarkInitializer, + LoadDefaultDatasets, + ScorerInitializer, + TargetInitializer, +) + +await initialize_pyrit_async( # type: ignore + memory_db_type=IN_MEMORY, + initializers=[ + TargetInitializer(), + ScorerInitializer(), + LoadDefaultDatasets(), + BenchmarkInitializer(), + ], +) -await initialize_pyrit_async(memory_db_type=IN_MEMORY, initializers=[LoadDefaultDatasets()]) # type: ignore +objective_target = OpenAIChatTarget() -# Pass any number of adversarial PromptTarget instances (with chat-target -# capabilities — multi-turn and editable history) as a list; AdversarialBenchmark -# infers a label for each from its identifier and runs every benchmark-friendly -# attack technique against the objective target with each adversarial model. -adversarial_model = OpenAIChatTarget() +# %% [markdown] +# ## Run the benchmark +# +# Instantiate with no model arguments — adversarial models come from the registry via +# `BenchmarkInitializer`. The default strategy (`light`) runs the benchmark-friendly subset of +# fanned variants for a quick comparison. -benchmark_scenario = AdversarialBenchmark(adversarial_models=[adversarial_model]) +# %% +dataset_config = DatasetConfiguration(dataset_names=["harmbench"], max_dataset_size=4) -await benchmark_scenario.initialize_async( # type: ignore - objective_target=OpenAIChatTarget(), max_concurrency=2 +scenario = AdversarialBenchmark() +await scenario.initialize_async( # type: ignore + objective_target=objective_target, + dataset_config=dataset_config, ) -baseline_result = await benchmark_scenario.run_async() # type: ignore +baseline_result = await scenario.run_async() # type: ignore -# Resume handle: re-run with `AdversarialBenchmark(..., scenario_result_id=)` to pick -# up where this run left off (constructor args must match the original run). +# Save this id to resume the run later via AdversarialBenchmark(scenario_result_id=...). print(f"Scenario result id: {baseline_result.id}") await output_scenario_async(baseline_result) + +# %% [markdown] +# ## Cross-run caching +# +# Re-run the benchmark with `skip_cached=True` and atomic attacks that completed (`SUCCESS` or +# `FAILURE` outcome) in any prior `COMPLETED` run of the same scenario name + version are skipped. +# `ERROR` and `UNDETERMINED` outcomes always retry, so transient failures don't poison the cache. +# +# The cache key is `(atomic_attack_name, technique_eval_hash)` — two atomic attacks that share a +# name but use different technique configurations (e.g. different scorer) don't cross-pollinate. +# +# Useful for: +# * Resuming a long-running benchmark after a crash or `Ctrl-C`. +# * Incrementally adding new adversarial targets without re-running the existing ones — the new +# fanned variants have new names, so they don't match any cached entry and execute fresh. + +# %% +scenario_cached = AdversarialBenchmark(skip_cached=True) +await scenario_cached.initialize_async( # type: ignore + objective_target=objective_target, + dataset_config=dataset_config, +) + +cached_result = await scenario_cached.run_async() # type: ignore + +await output_scenario_async(cached_result) + +# %% [markdown] +# ## Narrowing the fan-out +# +# Limit the benchmark to a specific subset of registered adversarial targets via +# `BenchmarkInitializer`'s `target_names` parameter. Settable from `.pyrit_conf` (see next section), +# or programmatically: +# +# ```python +# narrow_init = BenchmarkInitializer() +# narrow_init.set_params_from_args(args={"target_names": ["adversarial_chat_singleturn"]}) +# await narrow_init.initialize_async() +# ``` +# +# Unknown names raise `ValueError` listing both the unknowns and the discovered set, so typos fail +# loudly rather than silently expanding the run. + +# %% [markdown] +# ## Bootstrapping from `.pyrit_conf` +# +# For production / repeated runs, declare the full initializer chain in `.pyrit_conf`. Initializers +# run in list order; `target` must precede `benchmark` so adversarial-tagged targets are present +# when `BenchmarkInitializer` queries the registry. +# +# ```yaml +# memory_db_type: duckdb +# initializers: +# - name: target +# - name: scorer +# - name: load_default_datasets +# - name: benchmark +# args: +# target_names: +# - adversarial_chat_singleturn +# - adversarial_chat_reasoning +# scenario: +# name: benchmark.adversarial +# ``` +# +# Then run `pyrit_scan --config-file .pyrit_conf` — the scenario picks up the initializer-registered +# fan-out automatically. + +# %% [markdown] +# ## Scorer flexibility (forward-looking) +# +# `AdversarialBenchmark.__init__`'s `objective_scorer` parameter is typed as the broad `Scorer` +# base class. The current runtime contract is still `TrueFalseScorer` — passing a non-true/false +# scorer raises `TypeError` with a pointer to the planned scoring follow-up. The annotation is +# widened ahead of runtime support so callers coding against the signature today see the eventual +# contract, and the follow-up that drops the guard won't require a signature change. From 7e450b15911764eb72cbdc2e13cfe90ce1df5c22 Mon Sep 17 00:00:00 2001 From: Victor Valbuena Date: Thu, 21 May 2026 13:28:50 -0700 Subject: [PATCH 15/40] STYLE: Move TYPE_CHECKING-only imports in target_registry Upstream 23e2aa6b (DOC strict build) tightened ruff TC rules. RegistryEntry and TagQuery are only used in type annotations on TargetRegistry, so they belong in the if TYPE_CHECKING: block. Pre-commit clean across all PR-touched files after this fix. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- pyrit/registry/object_registries/target_registry.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/pyrit/registry/object_registries/target_registry.py b/pyrit/registry/object_registries/target_registry.py index 000979482f..a9379aea4f 100644 --- a/pyrit/registry/object_registries/target_registry.py +++ b/pyrit/registry/object_registries/target_registry.py @@ -12,14 +12,14 @@ import logging from typing import TYPE_CHECKING, Optional, Union -from pyrit.registry.object_registries.base_instance_registry import RegistryEntry from pyrit.registry.object_registries.retrievable_instance_registry import ( RetrievableInstanceRegistry, ) -from pyrit.registry.tag_query import TagQuery if TYPE_CHECKING: from pyrit.prompt_target import PromptTarget + from pyrit.registry.object_registries.base_instance_registry import RegistryEntry + from pyrit.registry.tag_query import TagQuery logger = logging.getLogger(__name__) From 34b86f0d83fa717277a5c134594053b7ab5a374f Mon Sep 17 00:00:00 2001 From: Victor Valbuena Date: Thu, 21 May 2026 13:51:25 -0700 Subject: [PATCH 16/40] FIX: TargetInitializer initialized adversarial chats as OpenAIChatTarget. Changed to AzureMLChatTarget. --- pyrit/setup/initializers/components/targets.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/pyrit/setup/initializers/components/targets.py b/pyrit/setup/initializers/components/targets.py index 9e85d7d8b1..aa1cae5c06 100644 --- a/pyrit/setup/initializers/components/targets.py +++ b/pyrit/setup/initializers/components/targets.py @@ -191,7 +191,7 @@ class TargetConfig: ), TargetConfig( registry_name="adversarial_chat_singleturn", - target_class=OpenAIChatTarget, + target_class=AzureMLChatTarget, endpoint_var="ADVERSARIAL_CHAT_SINGLETURN_ENDPOINT", key_var="ADVERSARIAL_CHAT_SINGLETURN_KEY", model_var="ADVERSARIAL_CHAT_SINGLETURN_MODEL", @@ -200,7 +200,7 @@ class TargetConfig: ), TargetConfig( registry_name="adversarial_chat_multiturn", - target_class=OpenAIChatTarget, + target_class=AzureMLChatTarget, endpoint_var="ADVERSARIAL_CHAT_MULTITURN_ENDPOINT", key_var="ADVERSARIAL_CHAT_MULTITURN_KEY", model_var="ADVERSARIAL_CHAT_MULTITURN_MODEL", @@ -209,7 +209,7 @@ class TargetConfig: ), TargetConfig( registry_name="adversarial_chat_reasoning", - target_class=OpenAIChatTarget, + target_class=AzureMLChatTarget, endpoint_var="ADVERSARIAL_CHAT_REASONING_ENDPOINT", key_var="ADVERSARIAL_CHAT_REASONING_KEY", model_var="ADVERSARIAL_CHAT_REASONING_MODEL", From c7e7b93f5a621f4b79da565cce2d257b29bf375e Mon Sep 17 00:00:00 2001 From: Victor Valbuena Date: Tue, 26 May 2026 11:19:48 -0700 Subject: [PATCH 17/40] Removed _cached_strategy_class and scorer type validation. --- .../scenarios/benchmark/adversarial.py | 19 ++++++------------- 1 file changed, 6 insertions(+), 13 deletions(-) diff --git a/pyrit/scenario/scenarios/benchmark/adversarial.py b/pyrit/scenario/scenarios/benchmark/adversarial.py index 9f9fd9cdf4..ee4e44bd18 100644 --- a/pyrit/scenario/scenarios/benchmark/adversarial.py +++ b/pyrit/scenario/scenarios/benchmark/adversarial.py @@ -14,12 +14,12 @@ from pyrit.registry.tag_query import TagQuery from pyrit.scenario.core.dataset_configuration import DatasetConfiguration from pyrit.scenario.core.scenario import BaselineAttackPolicy, Scenario -from pyrit.score.true_false.true_false_scorer import TrueFalseScorer if TYPE_CHECKING: from pyrit.scenario.core.atomic_attack import AtomicAttack from pyrit.scenario.core.scenario_strategy import ScenarioStrategy - from pyrit.score import Scorer + from pyrit.score.true_false.true_false_scorer import TrueFalseScorer + logger = logging.getLogger(__name__) @@ -167,8 +167,6 @@ class AdversarialBenchmark(Scenario): #: VERSION=1 remain queryable but won't suppress fresh runs. VERSION: int = 2 - _cached_strategy_class: ClassVar[type[ScenarioStrategy] | None] = None - #: AdversarialBenchmark compares attack-success rates across adversarial models; a baseline #: attack would be model-independent and contribute no signal to the comparison. BASELINE_ATTACK_POLICY: ClassVar[BaselineAttackPolicy] = BaselineAttackPolicy.Forbidden @@ -220,7 +218,7 @@ def default_dataset_config(cls) -> DatasetConfiguration: def __init__( self, *, - objective_scorer: Scorer | None = None, + objective_scorer: TrueFalseScorer | None = None, skip_cached: bool = False, scenario_result_id: str | None = None, ) -> None: @@ -258,14 +256,9 @@ def __init__( is tracked as a follow-up; the type annotation is widened ahead of the runtime support). """ - resolved_scorer: Scorer = objective_scorer if objective_scorer else self._get_default_objective_scorer() - if not isinstance(resolved_scorer, TrueFalseScorer): - raise TypeError( - f"AdversarialBenchmark currently requires a TrueFalseScorer for objective_scorer; " - f"got {type(resolved_scorer).__name__}. Full Scorer support (e.g. FloatScaleScorer) is " - f"tracked as the new-scoring follow-up — see the PR description for status." - ) - self._objective_scorer: TrueFalseScorer = resolved_scorer + self._objective_scorer: TrueFalseScorer = ( + objective_scorer if objective_scorer else self._get_default_objective_scorer() + ) self._skip_cached: bool = skip_cached super().__init__( From 5840b251f713227ed6deac82aa8c1834494acc2f Mon Sep 17 00:00:00 2001 From: Victor Valbuena Date: Tue, 26 May 2026 12:09:37 -0700 Subject: [PATCH 18/40] REFACTOR: Remove BenchmarkInitializer and source adversarial targets via scenario parameter Per Rich's PR review (#1765), the BenchmarkInitializer abstraction was over-engineered for what amounts to a single CLI flag. This commit collapses the fan-out plumbing into the scenario itself and surfaces adversarial-target selection through the standard `supported_parameters` API introduced by #1680. What changes - `pyrit/setup/initializers/benchmark.py` is deleted, along with its `__init__.py` export and `tests/unit/setup/test_benchmark_initializer.py`. - `AdversarialBenchmark` declares `adversarial_targets: list[str]` via `supported_parameters()`. CLI users pass `--adversarial-targets [ ...]` (auto-derived from the parameter list). `.pyrit_conf` users set `scenario.args.adversarial_targets`. The scenario raises an actionable `ValueError` if the parameter is missing or empty. - `_build_benchmark_strategy()` now filters `SCENARIO_TECHNIQUES` directly (no registry pre-population, no benchmark_fanout tag). The `light` aggregate is preserved so users keep the quick-snapshot default (excludes `tap`/ `crescendo_simulated`, which can take hours). - A reintroduced `_get_atomic_attacks_async()` override builds the (technique x target x dataset) cross-product. Per-pair factories are built via `AttackTechniqueRegistry.build_factory_from_spec(replace(spec, adversarial_chat=target))` so no global registry mutation occurs. - `display_group` is set to the TargetRegistry name passed via the CLI, not parsed from `atomic_attack_name` or read off a PromptTarget instance attribute. A new regression test verifies the registry name wins over `_model_name`/`_underlying_model`/`_endpoint`. - `atomic_attack_name` format is unchanged (`{technique}__{target}_{dataset}`); `VERSION` stays at 2 so prior cached results remain matchable by `skip_cached`. - `_cached_strategy_class` is dropped (Rich: "has some bugs and is also likely not needed"). The strategy enum is deterministic given `SCENARIO_TECHNIQUES` and is rebuilt on each call. - `TargetRegistry.get_by_tag_query()` is removed; its sole consumer (BenchmarkInitializer) is gone. Callers needing compound queries should lift the API to `BaseInstanceRegistry` if/when needed (Rich: "you should be able to use `get_by_tag` from the base class"). - `tests/end_to_end/test_scenarios.py` drops the `benchmark` initializer override and gains a `SCENARIO_EXTRA_ARGS` map that supplies `--adversarial-targets adversarial_chat` for the benchmark scenario. - `doc/scanner/benchmark.{py,ipynb}` is rewritten around the new CLI/config flow. A deeper restyle to match the other scanner docs is tracked as a follow-up. Out of scope (separate commits/PRs) - Rich's `adversarial.py:286` (analytics-driven cache via objective_target eval hash). - Rich's `benchmark.py:123` + `benchmark.ipynb:23` (basic-usage-only doc restyle). - Lifting `skip_cached` to base `Scenario`. Tests - `tests/unit/scenario/benchmark/test_adversarial.py` rewritten (46 passing) including the registry-name display-group regression test. - `tests/unit/setup/test_targets_initializer.py` migrated from `get_by_tag_query` to `get_by_tag`. - Wider sweep (`tests/unit/{scenario,setup,registry,cli}/`) green at 1349 tests. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- doc/scanner/benchmark.ipynb | 99 +- doc/scanner/benchmark.py | 92 +- .../object_registries/target_registry.py | 29 - .../scenarios/benchmark/adversarial.py | 420 +++++---- pyrit/setup/initializers/__init__.py | 2 - pyrit/setup/initializers/benchmark.py | 198 ---- tests/end_to_end/test_scenarios.py | 40 +- tests/unit/registry/test_target_registry.py | 61 -- .../scenario/benchmark/test_adversarial.py | 851 +++++++++++------- .../unit/setup/test_benchmark_initializer.py | 181 ---- tests/unit/setup/test_targets_initializer.py | 9 +- 11 files changed, 881 insertions(+), 1101 deletions(-) delete mode 100644 pyrit/setup/initializers/benchmark.py delete mode 100644 tests/unit/setup/test_benchmark_initializer.py diff --git a/doc/scanner/benchmark.ipynb b/doc/scanner/benchmark.ipynb index 193f0388d1..d3847460d8 100644 --- a/doc/scanner/benchmark.ipynb +++ b/doc/scanner/benchmark.ipynb @@ -25,11 +25,13 @@ "for evaluating which adversarial helper models produce stronger or weaker attack success rates\n", "against the same target.\n", "\n", - "Model fan-out is owned by `BenchmarkInitializer`, which discovers every ADVERSARIAL-tagged target\n", - "in `TargetRegistry` (populated by `TargetInitializer` from `ADVERSARIAL_CHAT_*` env vars) and\n", - "registers one fanned variant per `(adversarial-capable technique, target)` pair into\n", - "`AttackTechniqueRegistry`. The scenario then reads those variants when it builds its strategy\n", - "enum.\n", + "Adversarial targets are user-provided via the `adversarial_targets` scenario parameter. Each name\n", + "must already be registered in `TargetRegistry` — typically by `TargetInitializer` from the\n", + "`ADVERSARIAL_CHAT_*` env vars, or programmatically via `TargetRegistry.register_instance`. At run\n", + "time the scenario builds the `(technique × target × dataset)` cross-product directly: for each\n", + "adversarial-capable technique in `SCENARIO_TECHNIQUES` and each requested target, it constructs a\n", + "per-pair factory with `adversarial_chat` overridden to that target. No global\n", + "`AttackTechniqueRegistry` state is mutated.\n", "\n", "### Prerequisites\n", "\n", @@ -48,22 +50,32 @@ "# ADVERSARIAL_CHAT_MULTITURN_* and ADVERSARIAL_CHAT_REASONING_* follow the same pattern\n", "```\n", "\n", - "If no adversarial-tagged target is registered, `BenchmarkInitializer.initialize_async` raises\n", - "`ValueError` naming the env vars to set.\n", + "Use `pyrit_scan list-targets` to see every target currently registered, along with its tags.\n", "\n", "### CLI quickstart\n", "\n", "```bash\n", "pyrit_scan benchmark.adversarial \\\n", - " --initializers target load_default_datasets benchmark \\\n", + " --initializers target load_default_datasets \\\n", " --target openai_chat \\\n", + " --adversarial-targets adversarial_chat \\\n", " --max-dataset-size 4\n", "```\n", "\n", - "**Available strategies (depend on registered adversarial targets):** `all`, `light`, `single_turn`,\n", - "`multi_turn`, plus one concrete member per fanned variant named\n", - "`f\"{source_technique}__{target_name}\"` (e.g. `red_teaming__adversarial_chat`,\n", - "`tap__adversarial_chat_singleturn`)." + "Pass multiple `--adversarial-targets` values to compare across models in a single run:\n", + "\n", + "```bash\n", + "pyrit_scan benchmark.adversarial \\\n", + " --initializers target load_default_datasets \\\n", + " --target openai_chat \\\n", + " --adversarial-targets adversarial_chat adversarial_chat_singleturn adversarial_chat_reasoning \\\n", + " --max-dataset-size 4\n", + "```\n", + "\n", + "**Available strategies:** `light` (the default — a quick snapshot using the cheaper techniques),\n", + "`single_turn`, `multi_turn`, plus one concrete member per adversarial-capable source technique\n", + "(e.g. `red_teaming`, `tap`, `crescendo_simulated`). The default `light` aggregate deliberately\n", + "excludes `tap` and `crescendo_simulated`, which can take hours on a single run." ] }, { @@ -73,9 +85,9 @@ "source": [ "## Setup\n", "\n", - "`BenchmarkInitializer` runs after `TargetInitializer` so adversarial-tagged targets are present\n", - "in the registry before the fan-out logic queries it. The initializer chain is executed in list\n", - "order." + "`TargetInitializer` populates `TargetRegistry` from the `ADVERSARIAL_CHAT_*` env vars. The\n", + "scenario looks up adversarial targets by registry name from its `adversarial_targets` parameter,\n", + "so the targets must be registered before `scenario.run_async()` runs." ] }, { @@ -91,7 +103,6 @@ "from pyrit.scenario.scenarios.benchmark import AdversarialBenchmark\n", "from pyrit.setup import IN_MEMORY, initialize_pyrit_async\n", "from pyrit.setup.initializers import (\n", - " BenchmarkInitializer,\n", " LoadDefaultDatasets,\n", " ScorerInitializer,\n", " TargetInitializer,\n", @@ -103,7 +114,6 @@ " TargetInitializer(),\n", " ScorerInitializer(),\n", " LoadDefaultDatasets(),\n", - " BenchmarkInitializer(),\n", " ],\n", ")\n", "\n", @@ -117,9 +127,10 @@ "source": [ "## Run the benchmark\n", "\n", - "Instantiate with no model arguments — adversarial models come from the registry via\n", - "`BenchmarkInitializer`. The default strategy (`light`) runs the benchmark-friendly subset of\n", - "fanned variants for a quick comparison." + "Instantiate the scenario, then pass `adversarial_targets` through `initialize_async` via\n", + "`set_params_from_args` (programmatic equivalent of the `--adversarial-targets` CLI flag). The\n", + "default strategy (`light`) runs the benchmark-friendly subset of techniques for a quick\n", + "comparison." ] }, { @@ -132,6 +143,7 @@ "dataset_config = DatasetConfiguration(dataset_names=[\"harmbench\"], max_dataset_size=4)\n", "\n", "scenario = AdversarialBenchmark()\n", + "scenario.set_params_from_args(args={\"adversarial_targets\": [\"adversarial_chat\"]})\n", "await scenario.initialize_async( # type: ignore\n", " objective_target=objective_target,\n", " dataset_config=dataset_config,\n", @@ -162,7 +174,7 @@ "Useful for:\n", "* Resuming a long-running benchmark after a crash or `Ctrl-C`.\n", "* Incrementally adding new adversarial targets without re-running the existing ones — the new\n", - " fanned variants have new names, so they don't match any cached entry and execute fresh." + " `{technique}__{target}_{dataset}` names don't match any cached entry and execute fresh." ] }, { @@ -173,6 +185,7 @@ "outputs": [], "source": [ "scenario_cached = AdversarialBenchmark(skip_cached=True)\n", + "scenario_cached.set_params_from_args(args={\"adversarial_targets\": [\"adversarial_chat\"]})\n", "await scenario_cached.initialize_async( # type: ignore\n", " objective_target=objective_target,\n", " dataset_config=dataset_config,\n", @@ -187,33 +200,12 @@ "cell_type": "markdown", "id": "8", "metadata": {}, - "source": [ - "## Narrowing the fan-out\n", - "\n", - "Limit the benchmark to a specific subset of registered adversarial targets via\n", - "`BenchmarkInitializer`'s `target_names` parameter. Settable from `.pyrit_conf` (see next section),\n", - "or programmatically:\n", - "\n", - "```python\n", - "narrow_init = BenchmarkInitializer()\n", - "narrow_init.set_params_from_args(args={\"target_names\": [\"adversarial_chat_singleturn\"]})\n", - "await narrow_init.initialize_async()\n", - "```\n", - "\n", - "Unknown names raise `ValueError` listing both the unknowns and the discovered set, so typos fail\n", - "loudly rather than silently expanding the run." - ] - }, - { - "cell_type": "markdown", - "id": "9", - "metadata": {}, "source": [ "## Bootstrapping from `.pyrit_conf`\n", "\n", - "For production / repeated runs, declare the full initializer chain in `.pyrit_conf`. Initializers\n", - "run in list order; `target` must precede `benchmark` so adversarial-tagged targets are present\n", - "when `BenchmarkInitializer` queries the registry.\n", + "For production / repeated runs, declare the initializer chain and the adversarial-target list in\n", + "`.pyrit_conf`. `TargetInitializer` must precede the scenario so the named targets are present in\n", + "`TargetRegistry` by the time the scenario builds atomic attacks.\n", "\n", "```yaml\n", "memory_db_type: duckdb\n", @@ -221,22 +213,23 @@ " - name: target\n", " - name: scorer\n", " - name: load_default_datasets\n", - " - name: benchmark\n", - " args:\n", - " target_names:\n", - " - adversarial_chat_singleturn\n", - " - adversarial_chat_reasoning\n", "scenario:\n", " name: benchmark.adversarial\n", + " args:\n", + " adversarial_targets:\n", + " - adversarial_chat\n", + " - adversarial_chat_singleturn\n", + " - adversarial_chat_reasoning\n", "```\n", "\n", - "Then run `pyrit_scan --config-file .pyrit_conf` — the scenario picks up the initializer-registered\n", - "fan-out automatically." + "Then run `pyrit_scan --config-file .pyrit_conf` — the scenario reads `adversarial_targets` from\n", + "the config and builds the cross-product automatically. Unknown names raise `ValueError` listing\n", + "both the unknowns and every registered target so typos fail loudly." ] }, { "cell_type": "markdown", - "id": "10", + "id": "9", "metadata": {}, "source": [ "## Scorer flexibility (forward-looking)\n", diff --git a/doc/scanner/benchmark.py b/doc/scanner/benchmark.py index 9a3efa1f44..3dc051c3e9 100644 --- a/doc/scanner/benchmark.py +++ b/doc/scanner/benchmark.py @@ -24,11 +24,13 @@ # for evaluating which adversarial helper models produce stronger or weaker attack success rates # against the same target. # -# Model fan-out is owned by `BenchmarkInitializer`, which discovers every ADVERSARIAL-tagged target -# in `TargetRegistry` (populated by `TargetInitializer` from `ADVERSARIAL_CHAT_*` env vars) and -# registers one fanned variant per `(adversarial-capable technique, target)` pair into -# `AttackTechniqueRegistry`. The scenario then reads those variants when it builds its strategy -# enum. +# Adversarial targets are user-provided via the `adversarial_targets` scenario parameter. Each name +# must already be registered in `TargetRegistry` — typically by `TargetInitializer` from the +# `ADVERSARIAL_CHAT_*` env vars, or programmatically via `TargetRegistry.register_instance`. At run +# time the scenario builds the `(technique × target × dataset)` cross-product directly: for each +# adversarial-capable technique in `SCENARIO_TECHNIQUES` and each requested target, it constructs a +# per-pair factory with `adversarial_chat` overridden to that target. No global +# `AttackTechniqueRegistry` state is mutated. # # ### Prerequisites # @@ -47,29 +49,39 @@ # # ADVERSARIAL_CHAT_MULTITURN_* and ADVERSARIAL_CHAT_REASONING_* follow the same pattern # ``` # -# If no adversarial-tagged target is registered, `BenchmarkInitializer.initialize_async` raises -# `ValueError` naming the env vars to set. +# Use `pyrit_scan list-targets` to see every target currently registered, along with its tags. # # ### CLI quickstart # # ```bash # pyrit_scan benchmark.adversarial \ -# --initializers target load_default_datasets benchmark \ +# --initializers target load_default_datasets \ # --target openai_chat \ +# --adversarial-targets adversarial_chat \ # --max-dataset-size 4 # ``` # -# **Available strategies (depend on registered adversarial targets):** `all`, `light`, `single_turn`, -# `multi_turn`, plus one concrete member per fanned variant named -# `f"{source_technique}__{target_name}"` (e.g. `red_teaming__adversarial_chat`, -# `tap__adversarial_chat_singleturn`). +# Pass multiple `--adversarial-targets` values to compare across models in a single run: +# +# ```bash +# pyrit_scan benchmark.adversarial \ +# --initializers target load_default_datasets \ +# --target openai_chat \ +# --adversarial-targets adversarial_chat adversarial_chat_singleturn adversarial_chat_reasoning \ +# --max-dataset-size 4 +# ``` +# +# **Available strategies:** `light` (the default — a quick snapshot using the cheaper techniques), +# `single_turn`, `multi_turn`, plus one concrete member per adversarial-capable source technique +# (e.g. `red_teaming`, `tap`, `crescendo_simulated`). The default `light` aggregate deliberately +# excludes `tap` and `crescendo_simulated`, which can take hours on a single run. # %% [markdown] # ## Setup # -# `BenchmarkInitializer` runs after `TargetInitializer` so adversarial-tagged targets are present -# in the registry before the fan-out logic queries it. The initializer chain is executed in list -# order. +# `TargetInitializer` populates `TargetRegistry` from the `ADVERSARIAL_CHAT_*` env vars. The +# scenario looks up adversarial targets by registry name from its `adversarial_targets` parameter, +# so the targets must be registered before `scenario.run_async()` runs. # %% from pyrit.output import output_scenario_async @@ -78,7 +90,6 @@ from pyrit.scenario.scenarios.benchmark import AdversarialBenchmark from pyrit.setup import IN_MEMORY, initialize_pyrit_async from pyrit.setup.initializers import ( - BenchmarkInitializer, LoadDefaultDatasets, ScorerInitializer, TargetInitializer, @@ -90,7 +101,6 @@ TargetInitializer(), ScorerInitializer(), LoadDefaultDatasets(), - BenchmarkInitializer(), ], ) @@ -99,14 +109,16 @@ # %% [markdown] # ## Run the benchmark # -# Instantiate with no model arguments — adversarial models come from the registry via -# `BenchmarkInitializer`. The default strategy (`light`) runs the benchmark-friendly subset of -# fanned variants for a quick comparison. +# Instantiate the scenario, then pass `adversarial_targets` through `initialize_async` via +# `set_params_from_args` (programmatic equivalent of the `--adversarial-targets` CLI flag). The +# default strategy (`light`) runs the benchmark-friendly subset of techniques for a quick +# comparison. # %% dataset_config = DatasetConfiguration(dataset_names=["harmbench"], max_dataset_size=4) scenario = AdversarialBenchmark() +scenario.set_params_from_args(args={"adversarial_targets": ["adversarial_chat"]}) await scenario.initialize_async( # type: ignore objective_target=objective_target, dataset_config=dataset_config, @@ -132,10 +144,11 @@ # Useful for: # * Resuming a long-running benchmark after a crash or `Ctrl-C`. # * Incrementally adding new adversarial targets without re-running the existing ones — the new -# fanned variants have new names, so they don't match any cached entry and execute fresh. +# `{technique}__{target}_{dataset}` names don't match any cached entry and execute fresh. # %% scenario_cached = AdversarialBenchmark(skip_cached=True) +scenario_cached.set_params_from_args(args={"adversarial_targets": ["adversarial_chat"]}) await scenario_cached.initialize_async( # type: ignore objective_target=objective_target, dataset_config=dataset_config, @@ -145,28 +158,12 @@ await output_scenario_async(cached_result) -# %% [markdown] -# ## Narrowing the fan-out -# -# Limit the benchmark to a specific subset of registered adversarial targets via -# `BenchmarkInitializer`'s `target_names` parameter. Settable from `.pyrit_conf` (see next section), -# or programmatically: -# -# ```python -# narrow_init = BenchmarkInitializer() -# narrow_init.set_params_from_args(args={"target_names": ["adversarial_chat_singleturn"]}) -# await narrow_init.initialize_async() -# ``` -# -# Unknown names raise `ValueError` listing both the unknowns and the discovered set, so typos fail -# loudly rather than silently expanding the run. - # %% [markdown] # ## Bootstrapping from `.pyrit_conf` # -# For production / repeated runs, declare the full initializer chain in `.pyrit_conf`. Initializers -# run in list order; `target` must precede `benchmark` so adversarial-tagged targets are present -# when `BenchmarkInitializer` queries the registry. +# For production / repeated runs, declare the initializer chain and the adversarial-target list in +# `.pyrit_conf`. `TargetInitializer` must precede the scenario so the named targets are present in +# `TargetRegistry` by the time the scenario builds atomic attacks. # # ```yaml # memory_db_type: duckdb @@ -174,17 +171,18 @@ # - name: target # - name: scorer # - name: load_default_datasets -# - name: benchmark -# args: -# target_names: -# - adversarial_chat_singleturn -# - adversarial_chat_reasoning # scenario: # name: benchmark.adversarial +# args: +# adversarial_targets: +# - adversarial_chat +# - adversarial_chat_singleturn +# - adversarial_chat_reasoning # ``` # -# Then run `pyrit_scan --config-file .pyrit_conf` — the scenario picks up the initializer-registered -# fan-out automatically. +# Then run `pyrit_scan --config-file .pyrit_conf` — the scenario reads `adversarial_targets` from +# the config and builds the cross-product automatically. Unknown names raise `ValueError` listing +# both the unknowns and every registered target so typos fail loudly. # %% [markdown] # ## Scorer flexibility (forward-looking) diff --git a/pyrit/registry/object_registries/target_registry.py b/pyrit/registry/object_registries/target_registry.py index a9379aea4f..c6fefd3926 100644 --- a/pyrit/registry/object_registries/target_registry.py +++ b/pyrit/registry/object_registries/target_registry.py @@ -18,8 +18,6 @@ if TYPE_CHECKING: from pyrit.prompt_target import PromptTarget - from pyrit.registry.object_registries.base_instance_registry import RegistryEntry - from pyrit.registry.tag_query import TagQuery logger = logging.getLogger(__name__) @@ -76,30 +74,3 @@ def get_instance_by_name(self, name: str) -> Optional[PromptTarget]: The target instance, or None if not found. """ return self.get(name) - - def get_by_tag_query(self, *, query: TagQuery) -> list[RegistryEntry[PromptTarget]]: - """ - Get all entries whose tag keys satisfy ``query``. - - ``TagQuery`` operates on a tag set, so this method matches against - ``entry.tags.keys()`` and ignores tag values. For value-aware - single-tag lookups use ``get_by_tag(*, tag, value)`` on the base - class. - - Composite queries compose with ``&`` and ``|`` operators, e.g. - ``TagQuery.all("adversarial") & TagQuery.any_of("singleturn", "multiturn")``. - - Args: - query: The tag predicate to evaluate against each entry. - - Returns: - List of matching ``RegistryEntry`` objects sorted by registry name. - """ - results: list[RegistryEntry[PromptTarget]] = [] - # Note: this erases insertion order, but respects the base_instance_registry pattern - # (get_by_tag). - for name in sorted(self._registry_items.keys()): - entry = self._registry_items[name] - if query.matches(set(entry.tags.keys())): - results.append(entry) - return results diff --git a/pyrit/scenario/scenarios/benchmark/adversarial.py b/pyrit/scenario/scenarios/benchmark/adversarial.py index ee4e44bd18..2e78ad4c9a 100644 --- a/pyrit/scenario/scenarios/benchmark/adversarial.py +++ b/pyrit/scenario/scenarios/benchmark/adversarial.py @@ -5,18 +5,22 @@ from __future__ import annotations +import dataclasses import logging from typing import TYPE_CHECKING, ClassVar -from pyrit.common import apply_defaults -from pyrit.models import AttackOutcome -from pyrit.registry import AttackTechniqueRegistry, AttackTechniqueSpec +from pyrit.common import Parameter, apply_defaults +from pyrit.executor.attack import AttackScoringConfig +from pyrit.models import AttackOutcome, SeedAttackGroup +from pyrit.registry import AttackTechniqueRegistry, TargetRegistry from pyrit.registry.tag_query import TagQuery +from pyrit.scenario.core.atomic_attack import AtomicAttack from pyrit.scenario.core.dataset_configuration import DatasetConfiguration from pyrit.scenario.core.scenario import BaselineAttackPolicy, Scenario +from pyrit.scenario.core.scenario_techniques import SCENARIO_TECHNIQUES, _spec_needs_adversarial if TYPE_CHECKING: - from pyrit.scenario.core.atomic_attack import AtomicAttack + from pyrit.prompt_target import PromptTarget from pyrit.scenario.core.scenario_strategy import ScenarioStrategy from pyrit.score.true_false.true_false_scorer import TrueFalseScorer @@ -24,83 +28,32 @@ logger = logging.getLogger(__name__) -#: Strategy tag applied by ``BenchmarkInitializer`` to every fanned variant it -#: registers in ``AttackTechniqueRegistry``. The benchmark scenario reads its -#: strategy enum from entries carrying this tag. -BENCHMARK_FANOUT_TAG: str = "benchmark_fanout" - - -class _StrategyOnlyMarker: - """ - Sentinel attack class used only to satisfy ``AttackTechniqueSpec.attack_class`` - when reconstructing minimal specs for strategy-enum construction. - - ``AttackTechniqueRegistry.build_strategy_class_from_specs`` reads only - ``spec.name`` and ``spec.strategy_tags`` — never ``attack_class`` — so the - sentinel is safe. At attack-execution time the base - ``Scenario._get_atomic_attacks_async`` looks up the real factory by name - from ``AttackTechniqueRegistry`` (where ``BenchmarkInitializer`` registered - it), so this sentinel never reaches a runtime construction site. - """ - - def _build_benchmark_strategy() -> type[ScenarioStrategy]: """ - Build the ``BenchmarkStrategy`` enum from ``BenchmarkInitializer``-registered fanout. - - *Fanned entries* (also called *fanned variants*) are the per-target copies - of adversarial-capable scenario techniques that ``BenchmarkInitializer`` - registers into ``AttackTechniqueRegistry``. For each adversarial-capable - technique in ``SCENARIO_TECHNIQUES`` and each adversarial-tagged target - in ``TargetRegistry``, the initializer creates one fanned variant named - ``f"{source_technique}__{target_name}"`` with the live target bound onto - ``adversarial_chat`` and the strategy tag - :data:`BENCHMARK_FANOUT_TAG` appended. This function reads those entries - back and builds an enum whose concrete members are exactly the fanned - variants. - - Implementation note: this is a module-level function rather than a - ``@staticmethod`` on ``AdversarialBenchmark``. Strategy-class - construction never reads scenario instance state, so the function does - not belong to the class; module-level placement makes the dependency - (only the registry) explicit and the unit-test surface flat. - - Reconstructs minimal ``AttackTechniqueSpec`` stand-ins (name + - strategy_tags only) from each fanned entry to pass into - ``build_strategy_class_from_specs``. The sentinel - :class:`_StrategyOnlyMarker` is used for the required ``attack_class`` - field — see the sentinel's docstring for why this is safe. - - Aggregate selectors on the generated enum: - - * ``all`` — every fanned variant (auto-included by the builder). - * ``light`` — variants tagged ``"light"`` (inherited from the source spec). - * ``single_turn`` / ``multi_turn`` — variants tagged with the matching - turn-style tag inherited from the source spec. - - Per-target selection is also available via the auto-applied - ``f"model:{target_name}"`` tag on each fanned variant, accessible by name - on the generated enum (e.g. - ``BenchmarkStrategy("red_teaming__adversarial_chat_singleturn")``). + Build the ``BenchmarkStrategy`` enum from ``SCENARIO_TECHNIQUES``. + + Filters the static technique catalog to entries that require an + adversarial chat target (per :func:`_spec_needs_adversarial`) and passes + those source specs to + :meth:`AttackTechniqueRegistry.build_strategy_class_from_specs`. The + resulting enum has one concrete member per source technique (e.g. + ``red_teaming``, ``tap``, ``crescendo_simulated``) plus the standard + ``all`` / ``light`` / ``single_turn`` / ``multi_turn`` aggregates inherited + from the source specs' ``strategy_tags``. + + The (technique × target) cross-product is no longer pre-materialized into + enum members; per-target factories are built lazily in + :meth:`AdversarialBenchmark._get_atomic_attacks_async` from the + user-supplied ``adversarial_targets`` parameter. Returns: type[ScenarioStrategy]: The dynamically generated ``BenchmarkStrategy`` class. """ - registry = AttackTechniqueRegistry.get_registry_singleton() - fanned_entries = registry.get_by_tag(tag=BENCHMARK_FANOUT_TAG) - - fanned_specs = [ - AttackTechniqueSpec( - name=entry.name, - attack_class=_StrategyOnlyMarker, - strategy_tags=list(entry.tags.keys()), - ) - for entry in fanned_entries - ] + adversarial_specs = [spec for spec in SCENARIO_TECHNIQUES if _spec_needs_adversarial(spec)] return AttackTechniqueRegistry.build_strategy_class_from_specs( # type: ignore[ty:invalid-return-type] class_name="BenchmarkStrategy", - specs=fanned_specs, + specs=adversarial_specs, aggregate_tags={ "light": TagQuery.any_of("light"), "single_turn": TagQuery.any_of("single_turn"), @@ -113,58 +66,31 @@ class AdversarialBenchmark(Scenario): """ Benchmark scenario that compares the attack success rate (ASR) across adversarial models. - Adversarial-model fan-out is provided by ``BenchmarkInitializer``, which - registers per-target *fanned variants* of adversarial-capable scenario - techniques into ``AttackTechniqueRegistry`` tagged ``benchmark_fanout``. - This scenario reads those variants and builds its strategy enum from - them, so the set of available strategies reflects whichever adversarial - targets were discovered when ``BenchmarkInitializer`` ran (typically via - ``.pyrit_conf`` initializer ordering). - - Inherits the base ``Scenario._get_atomic_attacks_async`` loop with no - override; the fanned ``adversarial_chat`` binding lives on the - registered factories, so atomic-attack construction needs no special - handling here. - - When permuted atomic attacks materialize - ========================================= - The (technique × target × dataset) cross-product now happens in two - stages, not one (the pre-collapse override did all three at runtime): - - 1. **Initializer time** — ``BenchmarkInitializer.initialize_async`` - runs the (technique × adversarial-target) cross-product. For each - adversarial-capable technique in ``SCENARIO_TECHNIQUES`` and each - adversarial-tagged target in ``TargetRegistry``, it registers one - fanned ``AttackTechniqueFactory`` into ``AttackTechniqueRegistry`` - with the live target baked onto the factory's adversarial config. - After this step, the registry contains N×M fanned entries where N - is the count of adversarial-capable techniques and M is the count of - discovered adversarial targets. - - 2. **Scenario runtime** — ``Scenario._get_atomic_attacks_async`` - (inherited, base class) runs the (fanned-variant × dataset) - cross-product. It iterates ``self._scenario_strategies`` (the - fanned-variant names the user picked via the ``BenchmarkStrategy`` - enum), pairs each with every seed group in - ``self._dataset_config``, and builds one ``AtomicAttack`` per pair. - The target binding rides through on the factory created in step 1, - so no per-target handling is needed at this layer. - - The user-observable result is the same shape as before - (one ``AtomicAttack`` per (technique, target, dataset) triple), but the - target dimension is now owned by the initializer and the dataset - dimension is owned by the scenario. - - Display grouping is by target name (the part after ``__`` in each - fanned technique name) rather than by technique, so per-model ASR rolls - up naturally in result displays. + Adversarial targets are user-supplied via the ``adversarial_targets`` + parameter (declared in :meth:`supported_parameters`). Each target must + already be registered in ``TargetRegistry`` — typically by + ``TargetInitializer`` from ``ADVERSARIAL_CHAT_*`` env vars, or + programmatically via ``TargetRegistry.register_instance``. + + At run time, :meth:`_get_atomic_attacks_async` performs the + ``(technique × adversarial_target × dataset)`` cross-product: for each + selected adversarial-capable technique in ``SCENARIO_TECHNIQUES`` and + each requested target, it constructs a per-pair + :class:`AttackTechniqueFactory` via + :meth:`AttackTechniqueRegistry.build_factory_from_spec` with + ``adversarial_chat`` overridden to that target — no global registry + mutation. The resulting :class:`AtomicAttack` is named + ``f"{technique}__{target}_{dataset}"`` with ``display_group`` set to the + target's registry name so per-model ASR rolls up naturally in result + displays. """ - #: Bumped from 1 (pre-collapse) to 2 because the ``atomic_attack_name`` - #: format changed from ``f"{technique}__{model}__{dataset}"`` (triple-segment, - #: old override-driven) to ``f"{technique}__{model}_{dataset}"`` (double- - #: underscore-then-single-underscore, base-inherited). Cached results from - #: VERSION=1 remain queryable but won't suppress fresh runs. + #: Bumped from 1 to match the ``atomic_attack_name`` format introduced + #: when the scenario stopped passing target labels through its old triple + #: ``f"{technique}__{model}__{dataset}"`` shape. The post-collapse format + #: ``f"{technique}__{target}_{dataset}"`` is preserved here so cached + #: results from the prior collapse-era runs remain matchable by + #: ``skip_cached``. VERSION: int = 2 #: AdversarialBenchmark compares attack-success rates across adversarial models; a baseline @@ -174,19 +100,16 @@ class AdversarialBenchmark(Scenario): @classmethod def get_strategy_class(cls) -> type[ScenarioStrategy]: """ - Return the ``BenchmarkStrategy`` enum, building on first access. + Return the ``BenchmarkStrategy`` enum. - The enum is cached per-class for the lifetime of the process. To - rebuild after registry mutations (e.g. after re-running - ``BenchmarkInitializer`` with different adversarial targets), set - ``AdversarialBenchmark._cached_strategy_class = None`` and call again. + The enum is deterministic given the current ``SCENARIO_TECHNIQUES`` + catalog (the scenario no longer fans out across registry entries), + so it is rebuilt on every call rather than cached. Returns: type[ScenarioStrategy]: The ``BenchmarkStrategy`` enum class. """ - if cls._cached_strategy_class is None: - cls._cached_strategy_class = _build_benchmark_strategy() - return cls._cached_strategy_class + return _build_benchmark_strategy() @classmethod def get_default_strategy(cls) -> ScenarioStrategy: @@ -194,9 +117,10 @@ def get_default_strategy(cls) -> ScenarioStrategy: Return the default strategy (``light``). Returns: - ScenarioStrategy: The ``light`` aggregate member — runs the subset - of benchmark-friendly techniques that finish quickly with modest - system resources. + ScenarioStrategy: The ``light`` aggregate member — runs the + subset of benchmark-friendly techniques that finish quickly with + modest system resources (excludes ``tap`` and + ``crescendo_simulated``, which can take hours on a single run). """ return cls.get_strategy_class()("light") @@ -214,6 +138,38 @@ def default_dataset_config(cls) -> DatasetConfiguration: max_dataset_size=8, ) + @classmethod + def supported_parameters(cls) -> list[Parameter]: + """ + Declare the ``adversarial_targets`` parameter. + + The list is treated as required at run time: + :meth:`_get_atomic_attacks_async` raises ``ValueError`` if + ``self.params["adversarial_targets"]`` is empty or missing. The + scenario-side error (rather than a declaration-side default) lets + the caller raise a domain-specific message that names the CLI flag, + the ``.pyrit_conf`` key, and ``pyrit_scan list-targets``. + + Returns: + list[Parameter]: Single parameter declaring + ``adversarial_targets: list[str]``. + """ + return [ + Parameter( + name="adversarial_targets", + description=( + "Registry names of adversarial chat targets to benchmark. " + "Each name must already be registered in TargetRegistry " + "(via TargetInitializer or TargetRegistry.register_instance). " + "Use 'pyrit_scan list-targets' to see registered targets. " + "Settable via --adversarial-targets [ ...] on the CLI, " + "or scenario.args.adversarial_targets in .pyrit_conf." + ), + param_type=list[str], + default=None, + ), + ] + @apply_defaults def __init__( self, @@ -270,34 +226,184 @@ def __init__( async def _get_atomic_attacks_async(self) -> list[AtomicAttack]: """ - Build the base set of atomic attacks, then filter out cached completions when requested. - - Delegates to the base ``Scenario._get_atomic_attacks_async`` to - construct the (fanned-variant × dataset) candidate list, then drops - any candidate whose ``(atomic_attack_name, technique_eval_hash)`` - tuple appears in :meth:`_collect_cached_completion_pairs` (only when - ``self._skip_cached`` is ``True``). Always returns the unfiltered - base list when caching is disabled. + Build atomic attacks from (technique × adversarial_target × dataset), then apply caching. + + Reads the user-supplied ``adversarial_targets`` parameter, resolves + each name to a :class:`PromptTarget` via ``TargetRegistry``, and + cross-products the selected adversarial-capable techniques over the + resolved targets and configured datasets. Each pair builds a + non-registered per-pair factory via + :meth:`AttackTechniqueRegistry.build_factory_from_spec` with + ``adversarial_chat`` overridden to the resolved target — no global + registry state is touched. When ``self._skip_cached`` is set, the + final candidate list is then filtered against prior completed + ``(atomic_attack_name, technique_eval_hash)`` tuples. Returns: list[AtomicAttack]: The atomic attacks to actually execute on this run. + + Raises: + ValueError: If the scenario has not been initialized, if + ``adversarial_targets`` is missing/empty, or if any name in + ``adversarial_targets`` is not registered. """ - candidates = await super()._get_atomic_attacks_async() + if self._objective_target is None: + raise ValueError( + "Scenario not properly initialized. Call await scenario.initialize_async() before running." + ) + + target_names = self.params.get("adversarial_targets") + if not target_names: + raise ValueError( + "AdversarialBenchmark requires at least one adversarial chat target. " + "Pass --adversarial-targets [ ...] on the CLI, or set " + "scenario.args.adversarial_targets in .pyrit_conf. Use 'pyrit_scan list-targets' " + "to see registered targets." + ) + + resolved_targets = self._resolve_adversarial_targets(target_names=target_names) + selected_specs = self._select_adversarial_specs() + + scoring_config = AttackScoringConfig(objective_scorer=self._objective_scorer) + seed_groups_by_dataset = self._dataset_config.get_seed_attack_groups() + + atomic_attacks: list[AtomicAttack] = [] + for spec in selected_specs: + for target_name, target_instance in resolved_targets: + pair_spec = dataclasses.replace( + spec, + adversarial_chat=target_instance, + adversarial_chat_key=None, + ) + factory = AttackTechniqueRegistry.build_factory_from_spec(pair_spec) + + for dataset_name, seed_groups in seed_groups_by_dataset.items(): + if factory.seed_technique is not None: + compatible_groups = SeedAttackGroup.filter_compatible( + seed_groups=seed_groups, + technique=factory.seed_technique, + ) + skipped = len(seed_groups) - len(compatible_groups) + if skipped: + logger.info( + f"Skipped {skipped} seed group(s) from '{dataset_name}' for technique " + f"'{spec.name}' (prompt sequences overlap with simulated conversation)." + ) + if not compatible_groups: + logger.warning( + f"No compatible seed groups in '{dataset_name}' for technique " + f"'{spec.name}', skipping this (technique, target, dataset) triple." + ) + continue + else: + compatible_groups = list(seed_groups) + + attack_technique = factory.create( + objective_target=self._objective_target, + attack_scoring_config=scoring_config, + ) + # ``display_group`` is set explicitly here so result roll-ups group by the + # TargetRegistry name the caller passed via ``--adversarial-targets`` — + # not by any internal field on the PromptTarget instance (e.g. ``_model_name``). + # Because we override ``_get_atomic_attacks_async`` entirely, the base + # ``Scenario._build_display_group`` hook is never consulted; ``Scenario._finalize`` + # then reads ``aa.display_group`` directly (scenario.py:721). + atomic_attacks.append( + AtomicAttack( + atomic_attack_name=f"{spec.name}__{target_name}_{dataset_name}", + attack_technique=attack_technique, + seed_groups=list(compatible_groups), + adversarial_chat=target_instance, + objective_scorer=self._objective_scorer, + memory_labels=self._memory_labels, + display_group=target_name, + ) + ) + if not self._skip_cached: - return candidates + return atomic_attacks cached_pairs = self._collect_cached_completion_pairs() - filtered = [c for c in candidates if (c.atomic_attack_name, c.technique_eval_hash) not in cached_pairs] - skipped = len(candidates) - len(filtered) + filtered = [c for c in atomic_attacks if (c.atomic_attack_name, c.technique_eval_hash) not in cached_pairs] + skipped = len(atomic_attacks) - len(filtered) if skipped > 0: logger.info( "skip_cached=True: dropping %d/%d atomic attack(s) already completed in prior runs.", skipped, - len(candidates), + len(atomic_attacks), ) return filtered + def _resolve_adversarial_targets(self, *, target_names: list[str]) -> list[tuple[str, PromptTarget]]: + """ + Resolve each requested adversarial target name to its registered instance. + + Args: + target_names: Names supplied via the ``adversarial_targets`` + parameter. + + Returns: + list[tuple[str, PromptTarget]]: ``(registry_name, instance)`` + pairs in the order requested. + + Raises: + ValueError: If any name is not registered. The error lists both + the missing names and the names that are available, so + typos fail loudly. + """ + target_registry = TargetRegistry.get_registry_singleton() + resolved: list[tuple[str, PromptTarget]] = [] + unknown: list[str] = [] + for name in target_names: + instance = target_registry.get_instance_by_name(name) + if instance is None: + unknown.append(name) + else: + resolved.append((name, instance)) + + if unknown: + available = sorted(target_registry.get_names()) + raise ValueError( + f"AdversarialBenchmark: adversarial_targets {sorted(unknown)} not found in TargetRegistry. " + f"Available targets: {available}." + ) + + return resolved + + def _select_adversarial_specs(self) -> list: + """ + Resolve ``self._scenario_strategies`` back to adversarial-capable source specs. + + Strategies that are not adversarial-capable (i.e. don't satisfy + :func:`_spec_needs_adversarial`) are dropped with a warning. + Strategies whose name doesn't match any spec in + ``SCENARIO_TECHNIQUES`` are also dropped with a warning — this + guards against drift between the strategy enum and the technique + catalog. + + Returns: + list[AttackTechniqueSpec]: The adversarial-capable specs the + user selected, suitable for the per-pair factory build loop. + """ + specs_by_name = {spec.name: spec for spec in SCENARIO_TECHNIQUES} + selected_strategy_values = {s.value for s in self._scenario_strategies} + + selected_specs: list = [] + for value in selected_strategy_values: + spec = specs_by_name.get(value) + if spec is None: + logger.warning(f"AdversarialBenchmark: strategy '{value}' has no matching technique spec, skipping.") + continue + if not _spec_needs_adversarial(spec): + logger.warning( + f"AdversarialBenchmark: technique '{value}' does not require an adversarial chat target, " + "skipping (only adversarial-capable techniques are benchmarked)." + ) + continue + selected_specs.append(spec) + return selected_specs + def _collect_cached_completion_pairs(self) -> set[tuple[str, str | None]]: """ Collect cache keys for atomic attacks that completed in any prior run of this scenario. @@ -363,25 +469,3 @@ def _collect_cached_completion_pairs(self) -> set[tuple[str, str | None]]: cached_pairs.add((atomic_attack_name, parent_eval_hash)) return cached_pairs - - def _build_display_group(self, *, technique_name: str, seed_group_name: str) -> str: - """ - Group atomic-attack results by adversarial-target label rather than by technique. - - Fanned technique names have the format ``f"{source}__{target_name}"`` - (per ``BenchmarkInitializer``), so the target label is everything - after the ``__`` separator. Falls back to the full technique name - when no separator is present so legacy / non-fanned strategies still - render with a sensible label. - - Args: - technique_name: The fanned technique name, e.g. - ``"red_teaming__adversarial_chat_singleturn"``. - seed_group_name: Unused for this scenario (display rolls up - per-target, not per-seed-group). - - Returns: - str: The display group label — the target portion of the fanned - name when ``__`` is present, otherwise the full technique name. - """ - return technique_name.split("__", 1)[1] if "__" in technique_name else technique_name diff --git a/pyrit/setup/initializers/__init__.py b/pyrit/setup/initializers/__init__.py index b9e77c4038..84aeb83a49 100644 --- a/pyrit/setup/initializers/__init__.py +++ b/pyrit/setup/initializers/__init__.py @@ -6,7 +6,6 @@ from pyrit.common.deprecation import print_deprecation_message from pyrit.common.parameter import Parameter from pyrit.setup.initializers.airt import AIRTInitializer -from pyrit.setup.initializers.benchmark import BenchmarkInitializer from pyrit.setup.initializers.components.scenarios import ScenarioTechniqueInitializer from pyrit.setup.initializers.components.scorers import ScorerInitializer from pyrit.setup.initializers.components.targets import TargetInitializer @@ -19,7 +18,6 @@ "Parameter", "PyRITInitializer", "AIRTInitializer", - "BenchmarkInitializer", "ScenarioTechniqueInitializer", "ScorerInitializer", "TargetInitializer", diff --git a/pyrit/setup/initializers/benchmark.py b/pyrit/setup/initializers/benchmark.py deleted file mode 100644 index 3ecbce9dbc..0000000000 --- a/pyrit/setup/initializers/benchmark.py +++ /dev/null @@ -1,198 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# Licensed under the MIT license. - -""" -Benchmark initializer that fans adversarial-capable scenario techniques across -adversarial targets discovered in ``TargetRegistry``. - -This is the entry point for bootstrapping an ``AdversarialBenchmark`` trial. -It queries ``TargetRegistry`` for entries tagged ``ADVERSARIAL`` (via -``TagQuery.all("adversarial")``), then for every adversarial-capable -technique in ``SCENARIO_TECHNIQUES`` builds one fanned -``AttackTechniqueSpec`` per discovered target. Each fanned spec binds the -live target onto ``adversarial_chat`` and is registered into -``AttackTechniqueRegistry`` tagged ``["benchmark_fanout", f"model:{name}"]`` -so the benchmark scenario can discover them via tag query in a later commit. - -The ``target_names`` parameter (optional, settable from ``.pyrit_conf``) -narrows the fan-out to a specific subset of adversarial targets by registry -name. Unknown names raise ``ValueError``. - -Discovery returning no adversarial-tagged targets raises ``ValueError`` with -an actionable message pointing at the ``ADVERSARIAL_CHAT_*`` env vars and -the ``TargetInitializer`` dependency. -""" - -import dataclasses -import logging - -from pyrit.common.parameter import Parameter -from pyrit.registry import TargetRegistry -from pyrit.registry.object_registries.attack_technique_registry import ( - AttackTechniqueRegistry, - AttackTechniqueSpec, -) -from pyrit.registry.tag_query import TagQuery -from pyrit.scenario.core.scenario_techniques import SCENARIO_TECHNIQUES, _spec_needs_adversarial -from pyrit.setup.initializers.components.targets import TargetInitializerTags -from pyrit.setup.initializers.pyrit_initializer import PyRITInitializer - -logger = logging.getLogger(__name__) - - -#: Default discovery query used when no ``target_names`` override is provided. -#: Resolves every ``TargetRegistry`` entry tagged ``ADVERSARIAL`` (which today -#: includes ``adversarial_chat`` plus the ``ADVERSARIAL_CHAT_{SINGLETURN, -#: MULTITURN,REASONING}`` variants — all set by ``TargetInitializer``). -DEFAULT_ADVERSARIAL_TAG_QUERY: TagQuery = TagQuery.all(TargetInitializerTags.ADVERSARIAL.value) - - -class BenchmarkInitializer(PyRITInitializer): - """ - Fan adversarial-capable scenario techniques across discovered adversarial targets. - - For every ``AttackTechniqueSpec`` in ``SCENARIO_TECHNIQUES`` that uses an - adversarial chat target (multi-turn attacks plus crescendo-style - simulated conversations), this initializer registers one variant per - discovered adversarial target with the target bound onto - ``adversarial_chat``. Variants are named - ``f"{source_spec.name}__{target_name}"`` (e.g. ``red_teaming__adversarial_chat_singleturn``) - and carry the additional strategy tags ``"benchmark_fanout"`` and - ``f"model:{target_name}"`` so the benchmark scenario can query the - registry by tag in a later commit. - - Parameters (declared via :attr:`supported_parameters`): - - * ``target_names`` (``list[str]``, optional): Narrow fan-out to a - specific subset of adversarial targets by registry name. When omitted, - every target matching :data:`DEFAULT_ADVERSARIAL_TAG_QUERY` is used. - - Raises (at ``initialize_async``): - - * ``ValueError`` — no adversarial-tagged targets are registered. The - error names the ``ADVERSARIAL_CHAT_*`` env vars to set and the - ``TargetInitializer`` dependency. - * ``ValueError`` — any name in ``target_names`` does not match a - discovered adversarial-tagged target. The error lists discovered names. - - Prerequisites: ``TargetInitializer`` must have run first so adversarial - env-driven targets are present in ``TargetRegistry``. Registering the - base scenario-technique catalog (``ScenarioTechniqueInitializer`` or an - equivalent caller of ``register_scenario_techniques``) is also expected - if users will select non-benchmark strategies in the same session; - ``BenchmarkInitializer`` itself only registers the fanned variants. - Per-name idempotent via ``AttackTechniqueRegistry.register_from_specs``: - running the initializer twice with the same registry state is a no-op. - """ - - @property - def supported_parameters(self) -> list[Parameter]: - """Declare the optional ``target_names`` narrowing parameter.""" - return [ - Parameter( - name="target_names", - description=( - "Optional list of adversarial target registry names to narrow benchmark fan-out. " - 'When omitted, every target matching TagQuery.all("adversarial") is used.' - ), - default=None, - param_type=list[str], - ), - ] - - async def initialize_async(self) -> None: - """ - Discover adversarial targets and register fanned specs into the technique registry. - - Raises: - ValueError: If no adversarial-tagged targets are registered in - ``TargetRegistry``, or if ``self.params['target_names']`` - contains a name not in the discovered set. - """ - target_registry = TargetRegistry.get_registry_singleton() - discovered_entries = target_registry.get_by_tag_query(query=DEFAULT_ADVERSARIAL_TAG_QUERY) - if not discovered_entries: - raise ValueError( - "BenchmarkInitializer: no adversarial-tagged targets registered in TargetRegistry. " - "Set ADVERSARIAL_CHAT_* env vars (see .env_example) and ensure TargetInitializer runs " - "before BenchmarkInitializer (e.g. via .pyrit_conf initializer ordering)." - ) - - selected_entries = self._narrow_by_target_names(discovered_entries=discovered_entries) - - fanned_specs = self._build_fanned_specs(target_entries=selected_entries) - - attack_registry = AttackTechniqueRegistry.get_registry_singleton() - attack_registry.register_from_specs(fanned_specs) - - logger.info( - "BenchmarkInitializer: registered %d fanned spec(s) across %d adversarial target(s): %s", - len(fanned_specs), - len(selected_entries), - ", ".join(entry.name for entry in selected_entries), - ) - - def _narrow_by_target_names(self, *, discovered_entries: list) -> list: - """ - Filter ``discovered_entries`` to the names in ``self.params['target_names']``, if set. - - Args: - discovered_entries: The full set of adversarial-tagged registry entries. - - Returns: - list: ``discovered_entries`` unchanged when no ``target_names`` param - is set, otherwise the subset whose ``name`` is in the requested set. - - Raises: - ValueError: If any name in ``self.params['target_names']`` is not - present in ``discovered_entries``. - """ - target_names_param = self.params.get("target_names") - if not target_names_param: - return discovered_entries - - requested = set(target_names_param) - discovered_names = {entry.name for entry in discovered_entries} - unknown = requested - discovered_names - if unknown: - raise ValueError( - f"BenchmarkInitializer: unknown target_names {sorted(unknown)}. " - f"Discovered adversarial targets: {sorted(discovered_names)}." - ) - return [entry for entry in discovered_entries if entry.name in requested] - - def _build_fanned_specs(self, *, target_entries: list) -> list[AttackTechniqueSpec]: - """ - Build fanned ``AttackTechniqueSpec``s for every (adversarial-capable technique, target) pair. - - Adversarial-capability is determined by ``_spec_needs_adversarial`` - (re-used from ``scenario_techniques``): a spec needs an adversarial - chat target when its attack class accepts ``attack_adversarial_config`` - or its ``seed_technique`` has a simulated conversation. Non-adversarial - techniques (e.g. ``prompt_sending``, ``role_play``) are skipped — the - benchmark holds the objective target constant and varies the - adversarial chat helper across runs. - - Args: - target_entries: The adversarial-tagged registry entries to fan over. - - Returns: - list[AttackTechniqueSpec]: One fanned spec per (adversarial-capable - technique, target entry) pair, with the live target bound onto - ``adversarial_chat`` and benchmark-specific strategy tags appended. - """ - fanned: list[AttackTechniqueSpec] = [] - for source_spec in SCENARIO_TECHNIQUES: - if not _spec_needs_adversarial(source_spec): - continue - fanned.extend( - dataclasses.replace( - source_spec, - name=f"{source_spec.name}__{entry.name}", - adversarial_chat=entry.instance, - adversarial_chat_key=None, - strategy_tags=[*source_spec.strategy_tags, "benchmark_fanout", f"model:{entry.name}"], - ) - for entry in target_entries - ) - return fanned diff --git a/tests/end_to_end/test_scenarios.py b/tests/end_to_end/test_scenarios.py index 4f682cbc5f..0306607022 100644 --- a/tests/end_to_end/test_scenarios.py +++ b/tests/end_to_end/test_scenarios.py @@ -7,15 +7,15 @@ These tests dynamically discover all available scenarios and run each one using the pyrit_scan command. Most scenarios run with the :data:`DEFAULT_INITIALIZERS` list; scenarios that need additional setup -(e.g. ``benchmark.adversarial`` needs ``BenchmarkInitializer`` to fan -adversarial techniques out across registry-discovered targets) declare -their full initializer list in :data:`SCENARIO_INITIALIZERS`. +declare their full initializer list in :data:`SCENARIO_INITIALIZERS` and +extra CLI args in :data:`SCENARIO_EXTRA_ARGS`. Note: e2e tests are not part of CI; they run via ``make end-to-end-test`` on developer machines that have the appropriate env vars set -(``ADVERSARIAL_CHAT_*`` for the benchmark scenario, in particular). -``BenchmarkInitializer`` surfaces a clear error pointing at the env vars -when they are absent. +(``ADVERSARIAL_CHAT_*`` for the benchmark scenario, in particular). The +benchmark scenario reads its adversarial targets from ``--adversarial-targets``, +which resolves names via ``TargetRegistry`` (populated by +``TargetInitializer`` from those env vars). """ from pathlib import Path @@ -32,16 +32,19 @@ #: fetches each scenario's declared default datasets into memory. DEFAULT_INITIALIZERS: list[str] = ["target", "load_default_datasets"] -#: Per-scenario override map. A scenario named here uses this list verbatim -#: (no implicit merge with ``DEFAULT_INITIALIZERS``); a scenario absent here -#: falls back to ``DEFAULT_INITIALIZERS``. Keys use the dotted registry name +#: Per-scenario override map for initializers. A scenario absent here falls back +#: to :data:`DEFAULT_INITIALIZERS`. Keys use the dotted registry name #: (``.``) returned by ``ScenarioRegistry.get_names()``. -SCENARIO_INITIALIZERS: dict[str, list[str]] = { - # benchmark.adversarial depends on BenchmarkInitializer to fan - # adversarial-capable scenario techniques out across every - # ADVERSARIAL-tagged target in TargetRegistry. Without the - # benchmark initializer, the scenario's strategy enum is empty. - "benchmark.adversarial": [*DEFAULT_INITIALIZERS, "benchmark"], +SCENARIO_INITIALIZERS: dict[str, list[str]] = {} + +#: Per-scenario extra CLI args appended after the standard flag block. Keys use +#: the same dotted registry name as :data:`SCENARIO_INITIALIZERS`. Values are +#: lists already split into argv tokens. +SCENARIO_EXTRA_ARGS: dict[str, list[str]] = { + # benchmark.adversarial requires --adversarial-targets at run time + # (see AdversarialBenchmark.supported_parameters); without it the scenario + # raises ValueError before any attack is built. + "benchmark.adversarial": ["--adversarial-targets", "adversarial_chat"], } @@ -61,6 +64,11 @@ def _initializers_for(scenario_name: str) -> list[str]: return SCENARIO_INITIALIZERS.get(scenario_name, DEFAULT_INITIALIZERS) +def _extra_args_for(scenario_name: str) -> list[str]: + """Return scenario-specific extra CLI argv tokens, defaulting to none.""" + return SCENARIO_EXTRA_ARGS.get(scenario_name, []) + + @pytest.mark.timeout(7200) # 2 hour timeout per scenario @pytest.mark.parametrize("scenario_name", get_all_scenarios()) def test_scenario_with_pyrit_scan(scenario_name): @@ -71,6 +79,7 @@ def test_scenario_with_pyrit_scan(scenario_name): scenario_name: Name of the scenario to test (dynamically discovered). """ initializers = _initializers_for(scenario_name) + extra_args = _extra_args_for(scenario_name) try: result = pyrit_scan_main( [ @@ -85,6 +94,7 @@ def test_scenario_with_pyrit_scan(scenario_name): "1", "--log-level", "WARNING", + *extra_args, ] ) diff --git a/tests/unit/registry/test_target_registry.py b/tests/unit/registry/test_target_registry.py index b8c9234b88..f5865d4230 100644 --- a/tests/unit/registry/test_target_registry.py +++ b/tests/unit/registry/test_target_registry.py @@ -249,64 +249,3 @@ def test_list_metadata_filter_by_class_name(self): assert len(mock_metadata) == 2 for m in mock_metadata: assert m.class_name == "MockPromptTarget" - - -@pytest.mark.usefixtures("patch_central_database") -class TestTargetRegistryGetByTagQuery: - """Tests for ``TargetRegistry.get_by_tag_query`` (TagQuery-aware tag lookup).""" - - def setup_method(self): - """Reset and populate a fresh registry for each test.""" - TargetRegistry.reset_instance() - self.registry = TargetRegistry.get_registry_singleton() - - self.registry.register_instance(MockPromptTarget(), name="adv_single", tags=["adversarial", "singleturn"]) - self.registry.register_instance(MockPromptTarget(), name="adv_multi", tags=["adversarial", "multiturn"]) - self.registry.register_instance(MockPromptChatTarget(), name="scorer_only", tags=["scorer"]) - self.registry.register_instance(MockPromptTarget(), name="untagged") - - def teardown_method(self): - """Reset the singleton after each test.""" - TargetRegistry.reset_instance() - - def test_get_by_tag_query_returns_matching(self): - """A leaf ``TagQuery.all`` returns every entry whose tag set contains the required tag.""" - from pyrit.registry.tag_query import TagQuery - - results = self.registry.get_by_tag_query(query=TagQuery.all("adversarial")) - - names = [entry.name for entry in results] - assert names == ["adv_multi", "adv_single"] - - def test_get_by_tag_query_empty(self): - """A query that matches no entries returns an empty list (not raise).""" - from pyrit.registry.tag_query import TagQuery - - results = self.registry.get_by_tag_query(query=TagQuery.all("nonexistent_tag")) - assert results == [] - - def test_get_by_tag_query_composite_and_or(self): - """Composite queries via ``&`` / ``|`` evaluate as expected.""" - from pyrit.registry.tag_query import TagQuery - - query = TagQuery.all("adversarial") & TagQuery.any_of("singleturn", "multiturn") - results = self.registry.get_by_tag_query(query=query) - - names = [entry.name for entry in results] - assert names == ["adv_multi", "adv_single"] - - narrower = TagQuery.all("adversarial") & TagQuery.any_of("singleturn") - narrow_names = [entry.name for entry in self.registry.get_by_tag_query(query=narrower)] - assert narrow_names == ["adv_single"] - - def test_get_by_tag_query_matches_keys_not_values(self): - """``TagQuery`` evaluates against tag keys; tag values are ignored by this method.""" - from pyrit.registry.tag_query import TagQuery - - self.registry.add_tags(name="adv_single", tags={"priority": "high"}) - - priority_matches = self.registry.get_by_tag_query(query=TagQuery.all("priority")) - assert [entry.name for entry in priority_matches] == ["adv_single"] - - value_lookup = self.registry.get_by_tag_query(query=TagQuery.all("high")) - assert value_lookup == [] diff --git a/tests/unit/scenario/benchmark/test_adversarial.py b/tests/unit/scenario/benchmark/test_adversarial.py index 4bc22ce072..7748bab1d2 100644 --- a/tests/unit/scenario/benchmark/test_adversarial.py +++ b/tests/unit/scenario/benchmark/test_adversarial.py @@ -3,85 +3,79 @@ """Tests for the post-collapse AdversarialBenchmark scenario. -AdversarialBenchmark no longer takes an ``adversarial_models`` constructor -parameter and no longer builds local factories. It reads fanned variants -from ``AttackTechniqueRegistry`` (registered by ``BenchmarkInitializer``) -and inherits the base ``Scenario._get_atomic_attacks_async`` loop, with -an opt-in caching wrapper for cross-run skip-on-completion. +AdversarialBenchmark now owns its adversarial target axis directly via +the ``adversarial_targets`` parameter declared in +:meth:`supported_parameters`. Targets are user-supplied registry names +that resolve to ``PromptTarget`` instances via ``TargetRegistry``. The +``(technique × target × dataset)`` cross-product is built lazily inside +:meth:`_get_atomic_attacks_async` using per-pair non-registered factories; +no global ``AttackTechniqueRegistry`` state is mutated. These tests cover the new contract: * Class metadata (VERSION, BASELINE policy, defaults). -* Strategy enum is built from ``benchmark_fanout``-tagged registry entries. -* Display grouping uses the target-label portion of fanned technique names. -* Construction accepts ``objective_scorer``, ``skip_cached``, and - ``scenario_result_id``. -* ``skip_cached`` filters prior SUCCESS/FAILURE completions, keeps - ERROR/UNDETERMINED, respects eval-hash disambiguation, and only counts - COMPLETED scenario runs of the matching name + version. +* Strategy enum is built from source ``SCENARIO_TECHNIQUES`` entries that + require an adversarial chat target; ``light`` aggregate preserves the + source ``light`` tag (excludes ``tap`` / ``crescendo_simulated``). +* ``supported_parameters`` declares ``adversarial_targets: list[str]``. +* ``_resolve_adversarial_targets`` raises with available names on typos. +* ``_select_adversarial_specs`` drops non-adversarial techniques. +* ``_get_atomic_attacks_async`` produces ``N × M × D`` atomic attacks + with the expected ``atomic_attack_name`` and ``display_group``. +* ``_collect_cached_completion_pairs`` collects (name, hash) tuples for + prior ``SUCCESS`` / ``FAILURE`` outcomes only. +* ``skip_cached`` filters cached candidates end-to-end. +* Scorer flexibility stage 1: widened annotation + ``TypeError`` guard. """ from unittest.mock import MagicMock, patch import pytest -from pyrit.models import AttackOutcome +from pyrit.models import AttackOutcome, SeedAttackGroup, SeedObjective from pyrit.prompt_target import PromptTarget from pyrit.registry import TargetRegistry from pyrit.registry.object_registries.attack_technique_registry import AttackTechniqueRegistry from pyrit.scenario.core import BaselineAttackPolicy +from pyrit.scenario.core.scenario_techniques import SCENARIO_TECHNIQUES, _spec_needs_adversarial from pyrit.scenario.scenarios.benchmark.adversarial import ( - BENCHMARK_FANOUT_TAG, AdversarialBenchmark, _build_benchmark_strategy, ) from pyrit.score import TrueFalseScorer -from pyrit.setup.initializers import BenchmarkInitializer -from pyrit.setup.initializers.components.targets import TargetInitializerTags # --------------------------------------------------------------------------- -# Fixtures +# Fixtures / helpers # --------------------------------------------------------------------------- @pytest.fixture(autouse=True) -def reset_registries_and_cache(): - """Reset both registries and AdversarialBenchmark's strategy-class cache between tests.""" +def reset_registries(): + """Reset both registries between tests so target/technique state doesn't leak.""" AttackTechniqueRegistry.reset_instance() TargetRegistry.reset_instance() - AdversarialBenchmark._cached_strategy_class = None yield AttackTechniqueRegistry.reset_instance() TargetRegistry.reset_instance() - AdversarialBenchmark._cached_strategy_class = None def _register_adversarial_target(*, name: str) -> PromptTarget: - """Register a mock adversarial-tagged target in TargetRegistry.""" + """Register a mock adversarial target in TargetRegistry.""" target = MagicMock(spec=PromptTarget) - target.capabilities.includes.return_value = True registry = TargetRegistry.get_registry_singleton() - registry.register_instance(target, name=name, tags=[TargetInitializerTags.ADVERSARIAL.value]) + registry.register_instance(target, name=name) return target -async def _fan_out(*, target_names: list[str]) -> None: - """Register mock targets + run BenchmarkInitializer to populate AttackTechniqueRegistry.""" - for name in target_names: - _register_adversarial_target(name=name) - init = BenchmarkInitializer() - await init.initialize_async() - - # --------------------------------------------------------------------------- # Class metadata # --------------------------------------------------------------------------- class TestAdversarialBenchmarkMetadata: - """Tests for class-level metadata that doesn't depend on fan-out state.""" + """Tests for class-level metadata that doesn't depend on any runtime state.""" def test_version_is_2(self): - """VERSION is bumped from 1 because the atomic_attack_name format changed.""" + """VERSION matches the post-collapse ``atomic_attack_name`` format so cached results still match.""" assert AdversarialBenchmark.VERSION == 2 def test_baseline_attack_policy_is_forbidden(self): @@ -95,187 +89,425 @@ def test_default_dataset_config_uses_harmbench(self): def test_default_dataset_config_max_size_is_8(self): assert AdversarialBenchmark.default_dataset_config().max_dataset_size == 8 - def test_benchmark_fanout_tag_value(self): - """The shared tag value must match what BenchmarkInitializer applies.""" - assert BENCHMARK_FANOUT_TAG == "benchmark_fanout" - # --------------------------------------------------------------------------- -# Strategy class construction +# supported_parameters # --------------------------------------------------------------------------- -class TestAdversarialBenchmarkStrategy: - """Tests for _build_benchmark_strategy and the cached get_strategy_class accessor.""" +class TestAdversarialBenchmarkSupportedParameters: + """Tests for the ``adversarial_targets`` parameter declaration.""" - async def test_strategy_built_from_fanned_registry_entries(self): - """Every benchmark_fanout-tagged entry produces one concrete enum member.""" - await _fan_out(target_names=["adv_a", "adv_b"]) + def test_declares_adversarial_targets_param(self): + params = AdversarialBenchmark.supported_parameters() + names = [p.name for p in params] + assert "adversarial_targets" in names - strategy_cls = AdversarialBenchmark.get_strategy_class() - aggregate_names = {"all"} | strategy_cls.get_aggregate_tags() - concrete_members = [m for m in strategy_cls if m.value not in aggregate_names] + def test_adversarial_targets_param_is_list_of_str(self): + params = {p.name: p for p in AdversarialBenchmark.supported_parameters()} + param = params["adversarial_targets"] + assert param.param_type == list[str] + + def test_adversarial_targets_default_is_none(self): + """``None`` default lets the scenario raise a domain-specific error rather than the framework default.""" + params = {p.name: p for p in AdversarialBenchmark.supported_parameters()} + assert params["adversarial_targets"].default is None - assert len(concrete_members) > 0 - for member in concrete_members: - assert "__" in member.value, f"Expected fanned format with '__', got: {member.value}" + def test_adversarial_targets_description_mentions_cli_flag(self): + """The description must point users at ``--adversarial-targets`` for discoverability.""" + params = {p.name: p for p in AdversarialBenchmark.supported_parameters()} + description = params["adversarial_targets"].description + assert "--adversarial-targets" in description - async def test_strategy_concrete_member_count_matches_registry(self): - """Concrete enum members count equals fanned spec count in the registry.""" - await _fan_out(target_names=["adv_a", "adv_b"]) - attack_registry = AttackTechniqueRegistry.get_registry_singleton() - fanned_entries = attack_registry.get_by_tag(tag=BENCHMARK_FANOUT_TAG) +# --------------------------------------------------------------------------- +# Strategy class construction +# --------------------------------------------------------------------------- + - strategy_cls = AdversarialBenchmark.get_strategy_class() +class TestAdversarialBenchmarkStrategy: + """Tests for ``_build_benchmark_strategy`` and the cached ``get_strategy_class`` accessor.""" + + def test_strategy_built_from_adversarial_specs(self): + """Every adversarial-capable spec in ``SCENARIO_TECHNIQUES`` produces one concrete enum member.""" + strategy_cls = _build_benchmark_strategy() aggregate_names = {"all"} | strategy_cls.get_aggregate_tags() concrete_members = [m for m in strategy_cls if m.value not in aggregate_names] - assert len(concrete_members) == len(fanned_entries) - - async def test_strategy_exposes_per_model_selection(self): - """Each fanned variant inherits its model:* tag, accessible by name on the enum.""" - await _fan_out(target_names=["adv_a"]) + adversarial_specs = [s for s in SCENARIO_TECHNIQUES if _spec_needs_adversarial(s)] + adversarial_spec_names = {s.name for s in adversarial_specs} - attack_registry = AttackTechniqueRegistry.get_registry_singleton() - model_a_entries = attack_registry.get_by_tag(tag="model:adv_a") - assert len(model_a_entries) > 0 + concrete_member_values = {m.value for m in concrete_members} + assert concrete_member_values == adversarial_spec_names - strategy_cls = AdversarialBenchmark.get_strategy_class() - for entry in model_a_entries: - member = strategy_cls(entry.name) - assert "model:adv_a" in member.tags + def test_strategy_excludes_non_adversarial_techniques(self): + """Techniques like ``prompt_sending`` (no adversarial chat) must not be enum members.""" + strategy_cls = _build_benchmark_strategy() + member_values = {m.value for m in strategy_cls} - async def test_strategy_includes_required_aggregates(self): - """The strategy enum exposes all, light, single_turn, multi_turn aggregates.""" - await _fan_out(target_names=["adv_a"]) + non_adversarial = [s for s in SCENARIO_TECHNIQUES if not _spec_needs_adversarial(s)] + for spec in non_adversarial: + assert spec.name not in member_values, ( + f"{spec.name} is not adversarial-capable but appeared as a benchmark strategy member." + ) - strategy_cls = AdversarialBenchmark.get_strategy_class() + def test_strategy_includes_required_aggregates(self): + """The strategy enum exposes ``light``, ``single_turn``, ``multi_turn`` aggregates.""" + strategy_cls = _build_benchmark_strategy() aggregates = strategy_cls.get_aggregate_tags() - assert "all" in aggregates assert "light" in aggregates assert "single_turn" in aggregates assert "multi_turn" in aggregates - async def test_get_strategy_class_is_cached(self): - """Repeated calls within a process return the same class instance.""" - await _fan_out(target_names=["adv_a"]) + def test_light_aggregate_excludes_expensive_techniques(self): + """``light`` must not pull in ``tap`` or ``crescendo_simulated`` — both can take hours.""" + strategy_cls = _build_benchmark_strategy() + light_member = strategy_cls("light") - first = AdversarialBenchmark.get_strategy_class() - second = AdversarialBenchmark.get_strategy_class() + # Expand the aggregate to its concrete child members. + resolved_values = {child.value for child in strategy_cls.expand({light_member})} - assert first is second + assert "tap" not in resolved_values + assert "crescendo_simulated" not in resolved_values - async def test_cache_can_be_cleared_to_rebuild(self): - """Setting _cached_strategy_class = None forces a rebuild from current registry state.""" - await _fan_out(target_names=["adv_a"]) - first = AdversarialBenchmark.get_strategy_class() + def test_light_aggregate_includes_red_teaming(self): + """Sanity check: ``red_teaming`` is adversarial-capable AND tagged ``light``.""" + strategy_cls = _build_benchmark_strategy() + light_member = strategy_cls("light") + resolved_values = {child.value for child in strategy_cls.expand({light_member})} + assert "red_teaming" in resolved_values - await _fan_out(target_names=["adv_b"]) - AdversarialBenchmark._cached_strategy_class = None + def test_get_strategy_class_returns_same_enum_shape(self): + """``get_strategy_class`` rebuilds on every call; the resulting enums have identical members.""" + first = AdversarialBenchmark.get_strategy_class() second = AdversarialBenchmark.get_strategy_class() + assert {m.value for m in first} == {m.value for m in second} - assert first is not second - - async def test_default_strategy_is_light(self): - """get_default_strategy returns the 'light' aggregate so quick benchmark runs are the default.""" - await _fan_out(target_names=["adv_a"]) - + def test_default_strategy_is_light(self): + """``get_default_strategy`` returns the ``light`` aggregate.""" default = AdversarialBenchmark.get_default_strategy() assert default.value == "light" - def test_build_benchmark_strategy_empty_registry_produces_aggregates_only(self): - """No fan-out → enum still constructs (aggregates always present), just with zero concrete members.""" - strategy_cls = _build_benchmark_strategy() - aggregate_names = {"all"} | strategy_cls.get_aggregate_tags() - concrete_members = [m for m in strategy_cls if m.value not in aggregate_names] - assert concrete_members == [] - # --------------------------------------------------------------------------- -# Construction +# Construction (collapsed __init__) # --------------------------------------------------------------------------- +@pytest.mark.usefixtures("patch_central_database") class TestAdversarialBenchmarkInit: - """Tests for the collapsed __init__ surface (objective_scorer + scenario_result_id only).""" - - @pytest.mark.usefixtures("patch_central_database") - async def test_construct_with_default_objective_scorer(self): - """When no scorer is supplied, _get_default_objective_scorer is consulted.""" - await _fan_out(target_names=["adv_a"]) + """Tests for the collapsed ``__init__`` surface.""" + def test_construct_with_default_objective_scorer(self): + """When no scorer is supplied, ``_get_default_objective_scorer`` is consulted.""" default_scorer = MagicMock(spec=TrueFalseScorer) with patch.object(AdversarialBenchmark, "_get_default_objective_scorer", return_value=default_scorer): bench = AdversarialBenchmark() - assert bench._objective_scorer is default_scorer - @pytest.mark.usefixtures("patch_central_database") - async def test_construct_with_explicit_objective_scorer(self): - """An explicit scorer is used as-is, no default consulted.""" - await _fan_out(target_names=["adv_a"]) - + def test_construct_with_explicit_objective_scorer(self): explicit_scorer = MagicMock(spec=TrueFalseScorer) bench = AdversarialBenchmark(objective_scorer=explicit_scorer) - assert bench._objective_scorer is explicit_scorer - async def test_construct_takes_no_adversarial_models_param(self): - """Regression: the old adversarial_models constructor param is removed.""" - await _fan_out(target_names=["adv_a"]) - + def test_construct_takes_no_adversarial_models_param(self): + """Regression: the old ``adversarial_models`` constructor param is removed.""" with pytest.raises(TypeError): AdversarialBenchmark(adversarial_models=[MagicMock(spec=PromptTarget)]) # type: ignore[call-arg] + def test_construct_takes_no_models_param(self): + """Regression: the interim ``models`` param (BenchmarkInitializer era) is removed.""" + with pytest.raises(TypeError): + AdversarialBenchmark(models=[MagicMock(spec=PromptTarget)]) # type: ignore[call-arg] + + def test_skip_cached_defaults_to_false(self): + bench = AdversarialBenchmark(objective_scorer=MagicMock(spec=TrueFalseScorer)) + assert bench._skip_cached is False + + def test_skip_cached_can_be_set_true(self): + bench = AdversarialBenchmark( + objective_scorer=MagicMock(spec=TrueFalseScorer), + skip_cached=True, + ) + assert bench._skip_cached is True + # --------------------------------------------------------------------------- -# Display grouping +# _resolve_adversarial_targets # --------------------------------------------------------------------------- @pytest.mark.usefixtures("patch_central_database") -class TestAdversarialBenchmarkDisplayGroup: - """Tests for _build_display_group's fanned-name parsing.""" +class TestResolveAdversarialTargets: + """Tests for ``_resolve_adversarial_targets``: registry lookup + actionable errors on miss.""" - async def _make_bench(self) -> AdversarialBenchmark: - await _fan_out(target_names=["adv_a"]) + def _make_bench(self) -> AdversarialBenchmark: return AdversarialBenchmark(objective_scorer=MagicMock(spec=TrueFalseScorer)) - async def test_extracts_target_label_after_double_underscore(self): - bench = await self._make_bench() - result = bench._build_display_group( - technique_name="red_teaming__adversarial_chat_singleturn", - seed_group_name="seed_group_1", - ) - assert result == "adversarial_chat_singleturn" - - async def test_falls_back_to_full_name_when_no_separator(self): - """Non-fanned names (no ``__``) return the full technique name unchanged.""" - bench = await self._make_bench() - result = bench._build_display_group( - technique_name="prompt_sending", - seed_group_name="seed_group_1", - ) - assert result == "prompt_sending" - - async def test_ignores_seed_group_name(self): - """seed_group_name input must not influence the result (display rolls up per-target).""" - bench = await self._make_bench() - first = bench._build_display_group( - technique_name="red_teaming__adv_a", - seed_group_name="seed_group_a", + def test_resolves_registered_targets(self): + t_a = _register_adversarial_target(name="adv_a") + t_b = _register_adversarial_target(name="adv_b") + bench = self._make_bench() + + resolved = bench._resolve_adversarial_targets(target_names=["adv_a", "adv_b"]) + + names = [name for name, _ in resolved] + instances = [inst for _, inst in resolved] + assert names == ["adv_a", "adv_b"] + assert instances == [t_a, t_b] + + def test_unknown_target_raises_with_available_list(self): + _register_adversarial_target(name="adv_a") + bench = self._make_bench() + + with pytest.raises(ValueError) as exc_info: + bench._resolve_adversarial_targets(target_names=["adv_a", "missing"]) + + message = str(exc_info.value) + assert "missing" in message + assert "adv_a" in message # available list should include registered targets + + def test_all_unknown_targets_raises(self): + bench = self._make_bench() + + with pytest.raises(ValueError, match="not found in TargetRegistry"): + bench._resolve_adversarial_targets(target_names=["nope_1", "nope_2"]) + + def test_preserves_caller_order(self): + _register_adversarial_target(name="adv_b") + _register_adversarial_target(name="adv_a") + _register_adversarial_target(name="adv_c") + bench = self._make_bench() + + resolved = bench._resolve_adversarial_targets(target_names=["adv_c", "adv_a", "adv_b"]) + names = [name for name, _ in resolved] + assert names == ["adv_c", "adv_a", "adv_b"] + + +# --------------------------------------------------------------------------- +# _select_adversarial_specs +# --------------------------------------------------------------------------- + + +@pytest.mark.usefixtures("patch_central_database") +class TestSelectAdversarialSpecs: + """Tests for ``_select_adversarial_specs``: filter strategies down to adversarial-capable specs.""" + + def _make_bench(self) -> AdversarialBenchmark: + return AdversarialBenchmark(objective_scorer=MagicMock(spec=TrueFalseScorer)) + + def test_returns_only_adversarial_specs(self): + """``red_teaming`` is adversarial-capable; ``prompt_sending`` is not — only red_teaming survives.""" + bench = self._make_bench() + + red_teaming_strategy = MagicMock() + red_teaming_strategy.value = "red_teaming" + prompt_sending_strategy = MagicMock() + prompt_sending_strategy.value = "prompt_sending" + bench._scenario_strategies = [red_teaming_strategy, prompt_sending_strategy] + + selected = bench._select_adversarial_specs() + selected_names = {s.name for s in selected} + + assert "red_teaming" in selected_names + assert "prompt_sending" not in selected_names + + def test_unknown_strategy_value_is_skipped_with_warning(self, caplog): + """A strategy enum value with no matching spec is dropped (defensive guard against drift).""" + bench = self._make_bench() + + unknown_strategy = MagicMock() + unknown_strategy.value = "nonexistent_technique" + bench._scenario_strategies = [unknown_strategy] + + with caplog.at_level("WARNING"): + selected = bench._select_adversarial_specs() + + assert selected == [] + assert any("nonexistent_technique" in record.message for record in caplog.records) + + +# --------------------------------------------------------------------------- +# _get_atomic_attacks_async — validation and cross-product +# --------------------------------------------------------------------------- + + +@pytest.mark.usefixtures("patch_central_database") +class TestGetAtomicAttacksValidation: + """Tests for validation errors raised by ``_get_atomic_attacks_async``.""" + + def _make_bench(self) -> AdversarialBenchmark: + return AdversarialBenchmark(objective_scorer=MagicMock(spec=TrueFalseScorer)) + + async def test_uninitialized_scenario_raises(self): + """Calling ``_get_atomic_attacks_async`` before ``initialize_async`` raises a clear error.""" + bench = self._make_bench() + bench._objective_target = None + + with pytest.raises(ValueError, match="not properly initialized"): + await bench._get_atomic_attacks_async() + + async def test_missing_adversarial_targets_raises_actionable_error(self): + """Empty/missing ``adversarial_targets`` raises a message pointing at CLI / .pyrit_conf / list-targets.""" + bench = self._make_bench() + bench._objective_target = MagicMock(spec=PromptTarget) + bench.params = {} + + with pytest.raises(ValueError) as exc_info: + await bench._get_atomic_attacks_async() + + message = str(exc_info.value) + assert "--adversarial-targets" in message + assert ".pyrit_conf" in message + assert "list-targets" in message + + async def test_empty_adversarial_targets_list_raises(self): + bench = self._make_bench() + bench._objective_target = MagicMock(spec=PromptTarget) + bench.params = {"adversarial_targets": []} + + with pytest.raises(ValueError, match="at least one adversarial chat target"): + await bench._get_atomic_attacks_async() + + async def test_unknown_target_name_raises_listing_available(self): + _register_adversarial_target(name="adv_a") + bench = self._make_bench() + bench._objective_target = MagicMock(spec=PromptTarget) + bench.params = {"adversarial_targets": ["missing"]} + + with pytest.raises(ValueError) as exc_info: + await bench._get_atomic_attacks_async() + + message = str(exc_info.value) + assert "missing" in message + assert "adv_a" in message + + +@pytest.mark.usefixtures("patch_central_database") +class TestGetAtomicAttacksCrossProduct: + """Tests for the (technique × target × dataset) cross-product produced by ``_get_atomic_attacks_async``.""" + + def _make_bench_with_targets(self, *, target_names: list[str]) -> AdversarialBenchmark: + for name in target_names: + _register_adversarial_target(name=name) + bench = AdversarialBenchmark(objective_scorer=MagicMock(spec=TrueFalseScorer)) + bench._objective_target = MagicMock(spec=PromptTarget) + bench.params = {"adversarial_targets": target_names} + + red_teaming_strategy = MagicMock() + red_teaming_strategy.value = "red_teaming" + bench._scenario_strategies = [red_teaming_strategy] + + # Dataset config: one dataset with one real seed group (AtomicAttack hashes objectives). + seed_group = SeedAttackGroup(seeds=[SeedObjective(value="benchmark_objective_1")]) + bench._dataset_config = MagicMock() + bench._dataset_config.get_seed_attack_groups.return_value = {"harmbench": [seed_group]} + + return bench + + def _patch_factory_builder(self, *, seed_technique=None): + """Return a patch context manager for ``AttackTechniqueRegistry.build_factory_from_spec``.""" + factory = MagicMock() + factory.seed_technique = seed_technique + factory.create.return_value = MagicMock(name="AttackTechnique") + return patch( + "pyrit.scenario.scenarios.benchmark.adversarial.AttackTechniqueRegistry.build_factory_from_spec", + return_value=factory, ) - second = bench._build_display_group( - technique_name="red_teaming__adv_a", - seed_group_name="seed_group_b", + + async def test_cross_product_count_matches_n_techniques_m_targets_d_datasets(self): + """1 technique × 2 targets × 1 dataset = 2 atomic attacks.""" + bench = self._make_bench_with_targets(target_names=["adv_a", "adv_b"]) + + with self._patch_factory_builder(): + result = await bench._get_atomic_attacks_async() + + assert len(result) == 2 + + async def test_atomic_attack_name_format_is_technique__target_dataset(self): + """Name format: ``{technique}__{target}_{dataset}`` (preserves VERSION=2 cache key shape).""" + bench = self._make_bench_with_targets(target_names=["adv_a"]) + + with self._patch_factory_builder(): + result = await bench._get_atomic_attacks_async() + + names = [a.atomic_attack_name for a in result] + assert names == ["red_teaming__adv_a_harmbench"] + + async def test_display_group_equals_target_registry_name(self): + """``display_group`` is the raw target registry name — no string parsing.""" + bench = self._make_bench_with_targets(target_names=["adv_a", "adv_b"]) + + with self._patch_factory_builder(): + result = await bench._get_atomic_attacks_async() + + display_groups = sorted({a.display_group for a in result}) + assert display_groups == ["adv_a", "adv_b"] + + async def test_display_group_uses_registry_name_not_target_model_name(self): + """Regression: ``display_group`` must come from the registry name passed in via + ``adversarial_targets`` — not from any internal field on the ``PromptTarget`` instance + (``_model_name``, ``_underlying_model``, ``_endpoint``, etc.). If a future refactor + causes the scenario to source ``display_group`` from the target's own attributes, + users' per-target ASR roll-ups would silently change shape based on whatever model + name the target was constructed with. + """ + # Register a target under the registry name "adv_a" with an utterly different + # internal model/endpoint identity. After resolution, display_group should still + # be "adv_a" — the registry name — not anything that leaked from the target. + target = MagicMock(spec=PromptTarget) + target._model_name = "totally-different-model-name" + target._underlying_model = "another-model-identity" + target._endpoint = "https://hijacked.example.com/openai/v1" + target.name = "name-attribute-that-must-not-leak" + TargetRegistry.get_registry_singleton().register_instance(target, name="adv_a") + + bench = AdversarialBenchmark(objective_scorer=MagicMock(spec=TrueFalseScorer)) + bench._objective_target = MagicMock(spec=PromptTarget) + bench.params = {"adversarial_targets": ["adv_a"]} + + red_teaming_strategy = MagicMock() + red_teaming_strategy.value = "red_teaming" + bench._scenario_strategies = [red_teaming_strategy] + + seed_group = SeedAttackGroup(seeds=[SeedObjective(value="display_group_regression_objective")]) + bench._dataset_config = MagicMock() + bench._dataset_config.get_seed_attack_groups.return_value = {"harmbench": [seed_group]} + + with self._patch_factory_builder(): + result = await bench._get_atomic_attacks_async() + + assert len(result) == 1 + atomic = result[0] + assert atomic.display_group == "adv_a", ( + f"display_group must equal the registry name 'adv_a', got {atomic.display_group!r}. " + "If this is failing, the scenario started sourcing display_group from the target's " + "internal attributes (_model_name, etc.) — restore the registry-name behavior." ) - assert first == second == "adv_a" + # Belt-and-suspenders: also assert the atomic_attack_name uses the registry name, + # since the same plumbing pipes both. + assert atomic.atomic_attack_name == "red_teaming__adv_a_harmbench" + + async def test_factory_built_per_target_with_overridden_adversarial_chat(self): + """Each (spec, target) pair gets its own ``build_factory_from_spec`` call with a replaced spec.""" + bench = self._make_bench_with_targets(target_names=["adv_a", "adv_b"]) + + with self._patch_factory_builder() as build_mock: + await bench._get_atomic_attacks_async() + + # 1 selected technique × 2 targets = 2 factory builds. + assert build_mock.call_count == 2 + # Each pair_spec has adversarial_chat replaced; verify the (replaced spec).adversarial_chat + # matches the corresponding registry entry. + target_a = TargetRegistry.get_registry_singleton().get_instance_by_name("adv_a") + target_b = TargetRegistry.get_registry_singleton().get_instance_by_name("adv_b") + replaced_targets = {call.args[0].adversarial_chat for call in build_mock.call_args_list} + assert replaced_targets == {target_a, target_b} # --------------------------------------------------------------------------- -# skip_cached behavior (Commit 6 / F3) +# _collect_cached_completion_pairs # --------------------------------------------------------------------------- @@ -293,7 +525,7 @@ def _make_attack_result( parent_collection: str | None, parent_eval_hash: str | None, ) -> MagicMock: - """Build a minimal AttackResult stand-in with the attribution_data shape Commit 6 reads.""" + """Build a minimal AttackResult stand-in with the attribution_data shape the cache filter reads.""" ar = MagicMock() ar.outcome = outcome if parent_collection is None and parent_eval_hash is None: @@ -308,278 +540,211 @@ def _make_attack_result( return ar -def _make_candidate(*, name: str, eval_hash: str) -> MagicMock: - """Build a minimal AtomicAttack stand-in with the two fields the cache filter reads.""" - candidate = MagicMock() - candidate.atomic_attack_name = name - candidate.technique_eval_hash = eval_hash - return candidate - - @pytest.mark.usefixtures("patch_central_database") -class TestAdversarialBenchmarkSkipCachedFilter: - """Tests for the _get_atomic_attacks_async caching wrapper.""" - - async def _make_bench(self, *, skip_cached: bool) -> AdversarialBenchmark: - await _fan_out(target_names=["adv_a"]) - return AdversarialBenchmark( - objective_scorer=MagicMock(spec=TrueFalseScorer), - skip_cached=skip_cached, - ) - - async def test_skip_cached_default_false_means_no_filtering(self): - """skip_cached defaults to False; super() output is returned unchanged.""" - bench = await self._make_bench(skip_cached=False) - candidates = [_make_candidate(name="red_teaming__adv_a", eval_hash="hash_a")] +class TestCollectCachedCompletionPairs: + """Tests for ``_collect_cached_completion_pairs`` (the cache key collector).""" - with patch.object( - AdversarialBenchmark.__bases__[0], "_get_atomic_attacks_async", return_value=candidates - ) as super_mock: - result = await bench._get_atomic_attacks_async() - - assert result == candidates - super_mock.assert_awaited_once() - - async def test_skip_cached_true_drops_completed_pairs(self): - """SUCCESS and FAILURE prior outcomes drop the matching candidate.""" - bench = await self._make_bench(skip_cached=True) + def _make_bench(self) -> AdversarialBenchmark: + bench = AdversarialBenchmark(objective_scorer=MagicMock(spec=TrueFalseScorer)) + bench._memory = MagicMock() + return bench - candidates = [ - _make_candidate(name="red_teaming__adv_a", eval_hash="hash_a"), - _make_candidate(name="tap__adv_a", eval_hash="hash_b"), - _make_candidate(name="crescendo_simulated__adv_a", eval_hash="hash_c"), - ] + def test_collects_success_and_failure_pairs(self): + bench = self._make_bench() prior_sr = _make_scenario_result(result_id="sid-1") prior_attacks = [ _make_attack_result( outcome=AttackOutcome.SUCCESS, - parent_collection="red_teaming__adv_a", + parent_collection="red_teaming__adv_a_harmbench", parent_eval_hash="hash_a", ), _make_attack_result( outcome=AttackOutcome.FAILURE, - parent_collection="tap__adv_a", + parent_collection="tap__adv_a_harmbench", parent_eval_hash="hash_b", ), ] - - bench._memory = MagicMock() bench._memory.get_scenario_results.return_value = [prior_sr] bench._memory.get_attack_results.return_value = prior_attacks - with patch.object(AdversarialBenchmark.__bases__[0], "_get_atomic_attacks_async", return_value=candidates): - result = await bench._get_atomic_attacks_async() - - names = [c.atomic_attack_name for c in result] - assert names == ["crescendo_simulated__adv_a"] + pairs = bench._collect_cached_completion_pairs() - async def test_skip_cached_keeps_error_outcomes(self): - """ERROR outcomes must retry — not be cached.""" - bench = await self._make_bench(skip_cached=True) + assert pairs == { + ("red_teaming__adv_a_harmbench", "hash_a"), + ("tap__adv_a_harmbench", "hash_b"), + } - candidates = [_make_candidate(name="red_teaming__adv_a", eval_hash="hash_a")] + def test_excludes_error_and_undetermined_outcomes(self): + bench = self._make_bench() prior_sr = _make_scenario_result(result_id="sid-1") prior_attacks = [ _make_attack_result( outcome=AttackOutcome.ERROR, - parent_collection="red_teaming__adv_a", - parent_eval_hash="hash_a", + parent_collection="x", + parent_eval_hash="h", + ), + _make_attack_result( + outcome=AttackOutcome.UNDETERMINED, + parent_collection="y", + parent_eval_hash="h", ), ] - - bench._memory = MagicMock() bench._memory.get_scenario_results.return_value = [prior_sr] bench._memory.get_attack_results.return_value = prior_attacks - with patch.object(AdversarialBenchmark.__bases__[0], "_get_atomic_attacks_async", return_value=candidates): - result = await bench._get_atomic_attacks_async() + pairs = bench._collect_cached_completion_pairs() - assert result == candidates + assert pairs == set() - async def test_skip_cached_keeps_undetermined_outcomes(self): - """UNDETERMINED outcomes must retry — not be cached.""" - bench = await self._make_bench(skip_cached=True) + def test_only_counts_completed_scenario_runs(self): + bench = self._make_bench() + in_progress = _make_scenario_result(result_id="sid-1", run_state="IN_PROGRESS") + failed = _make_scenario_result(result_id="sid-2", run_state="FAILED") + bench._memory.get_scenario_results.return_value = [in_progress, failed] - candidates = [_make_candidate(name="red_teaming__adv_a", eval_hash="hash_a")] - prior_sr = _make_scenario_result(result_id="sid-1") - prior_attacks = [ - _make_attack_result( - outcome=AttackOutcome.UNDETERMINED, - parent_collection="red_teaming__adv_a", - parent_eval_hash="hash_a", - ), - ] + pairs = bench._collect_cached_completion_pairs() - bench._memory = MagicMock() - bench._memory.get_scenario_results.return_value = [prior_sr] - bench._memory.get_attack_results.return_value = prior_attacks + assert pairs == set() + # No COMPLETED runs → never touch get_attack_results. + bench._memory.get_attack_results.assert_not_called() - with patch.object(AdversarialBenchmark.__bases__[0], "_get_atomic_attacks_async", return_value=candidates): - result = await bench._get_atomic_attacks_async() + def test_queries_memory_by_scenario_name_and_version(self): + bench = self._make_bench() + bench._memory.get_scenario_results.return_value = [] - assert result == candidates + bench._collect_cached_completion_pairs() - async def test_skip_cached_respects_eval_hash_disambiguation(self): - """Same atomic_attack_name but different parent_eval_hash → not considered cached.""" - bench = await self._make_bench(skip_cached=True) + bench._memory.get_scenario_results.assert_called_once_with( + scenario_name="AdversarialBenchmark", + scenario_version=AdversarialBenchmark.VERSION, + ) - candidates = [_make_candidate(name="red_teaming__adv_a", eval_hash="new_hash")] + def test_skips_rows_with_missing_parent_collection(self): + """``attribution_data=None`` or missing ``parent_collection`` rows are silently skipped.""" + bench = self._make_bench() prior_sr = _make_scenario_result(result_id="sid-1") prior_attacks = [ - _make_attack_result( - outcome=AttackOutcome.SUCCESS, - parent_collection="red_teaming__adv_a", - parent_eval_hash="old_hash", - ), + _make_attack_result(outcome=AttackOutcome.SUCCESS, parent_collection=None, parent_eval_hash=None), + _make_attack_result(outcome=AttackOutcome.SUCCESS, parent_collection=None, parent_eval_hash="hash_x"), ] - - bench._memory = MagicMock() bench._memory.get_scenario_results.return_value = [prior_sr] bench._memory.get_attack_results.return_value = prior_attacks - with patch.object(AdversarialBenchmark.__bases__[0], "_get_atomic_attacks_async", return_value=candidates): - result = await bench._get_atomic_attacks_async() + pairs = bench._collect_cached_completion_pairs() + assert pairs == set() - assert result == candidates + def test_memory_error_falls_back_to_empty_set(self): + """An exception from ``get_scenario_results`` must not block the run; cache becomes a no-op.""" + bench = self._make_bench() + bench._memory.get_scenario_results.side_effect = RuntimeError("db down") - async def test_skip_cached_only_considers_completed_scenarios(self): - """Scenarios in IN_PROGRESS / FAILED / CANCELLED state must not seed the cache.""" - bench = await self._make_bench(skip_cached=True) + pairs = bench._collect_cached_completion_pairs() + assert pairs == set() - candidates = [_make_candidate(name="red_teaming__adv_a", eval_hash="hash_a")] - in_progress = _make_scenario_result(result_id="sid-1", run_state="IN_PROGRESS") - failed = _make_scenario_result(result_id="sid-2", run_state="FAILED") - bench._memory = MagicMock() - bench._memory.get_scenario_results.return_value = [in_progress, failed] - bench._memory.get_attack_results.return_value = [ - _make_attack_result( - outcome=AttackOutcome.SUCCESS, - parent_collection="red_teaming__adv_a", - parent_eval_hash="hash_a", - ), - ] +# --------------------------------------------------------------------------- +# skip_cached end-to-end through _get_atomic_attacks_async +# --------------------------------------------------------------------------- - with patch.object(AdversarialBenchmark.__bases__[0], "_get_atomic_attacks_async", return_value=candidates): - result = await bench._get_atomic_attacks_async() - assert result == candidates - bench._memory.get_attack_results.assert_not_called() +@pytest.mark.usefixtures("patch_central_database") +class TestSkipCachedFilter: + """End-to-end tests for the ``skip_cached`` filter applied in ``_get_atomic_attacks_async``.""" - async def test_skip_cached_filters_by_scenario_name_and_version(self): - """get_scenario_results is queried with this scenario's name + VERSION; old VERSION=1 results don't apply.""" - bench = await self._make_bench(skip_cached=True) + def _make_bench(self, *, skip_cached: bool) -> AdversarialBenchmark: + _register_adversarial_target(name="adv_a") + bench = AdversarialBenchmark( + objective_scorer=MagicMock(spec=TrueFalseScorer), + skip_cached=skip_cached, + ) + bench._objective_target = MagicMock(spec=PromptTarget) + bench.params = {"adversarial_targets": ["adv_a"]} + + red_teaming_strategy = MagicMock() + red_teaming_strategy.value = "red_teaming" + bench._scenario_strategies = [red_teaming_strategy] + + seed_group = SeedAttackGroup(seeds=[SeedObjective(value="skip_cached_objective")]) + bench._dataset_config = MagicMock() + bench._dataset_config.get_seed_attack_groups.return_value = {"harmbench": [seed_group]} + + return bench + + def _patch_factory_builder(self): + factory = MagicMock() + factory.seed_technique = None + factory.create.return_value = MagicMock(name="AttackTechnique") + return patch( + "pyrit.scenario.scenarios.benchmark.adversarial.AttackTechniqueRegistry.build_factory_from_spec", + return_value=factory, + ) + async def test_skip_cached_false_returns_all_candidates(self): + bench = self._make_bench(skip_cached=False) bench._memory = MagicMock() - bench._memory.get_scenario_results.return_value = [] - with patch.object(AdversarialBenchmark.__bases__[0], "_get_atomic_attacks_async", return_value=[]): - await bench._get_atomic_attacks_async() - - bench._memory.get_scenario_results.assert_called_once_with( - scenario_name="AdversarialBenchmark", - scenario_version=AdversarialBenchmark.VERSION, - ) + with self._patch_factory_builder(): + result = await bench._get_atomic_attacks_async() - async def test_skip_cached_handles_missing_attribution_data(self): - """Rows with attribution_data=None or missing parent_collection are silently skipped.""" - bench = await self._make_bench(skip_cached=True) + assert len(result) == 1 + # No cache query when skip_cached=False. + bench._memory.get_scenario_results.assert_not_called() - candidates = [_make_candidate(name="red_teaming__adv_a", eval_hash="hash_a")] + async def test_skip_cached_true_filters_matching_candidates(self): + bench = self._make_bench(skip_cached=True) prior_sr = _make_scenario_result(result_id="sid-1") prior_attacks = [ _make_attack_result( outcome=AttackOutcome.SUCCESS, - parent_collection=None, - parent_eval_hash=None, - ), - _make_attack_result( - outcome=AttackOutcome.SUCCESS, - parent_collection=None, - parent_eval_hash="hash_x", + parent_collection="red_teaming__adv_a_harmbench", + parent_eval_hash=None, # MagicMock candidates yield None for technique_eval_hash ), ] - bench._memory = MagicMock() bench._memory.get_scenario_results.return_value = [prior_sr] bench._memory.get_attack_results.return_value = prior_attacks - with patch.object(AdversarialBenchmark.__bases__[0], "_get_atomic_attacks_async", return_value=candidates): - result = await bench._get_atomic_attacks_async() - - assert result == candidates + with self._patch_factory_builder(): + # Stub out technique_eval_hash so the cache-key tuple matches. + with patch( + "pyrit.scenario.core.atomic_attack.AtomicAttack.technique_eval_hash", + new_callable=lambda: property(lambda self: None), + ): + result = await bench._get_atomic_attacks_async() - async def test_skip_cached_memory_error_falls_back_to_no_filter(self): - """An exception from get_scenario_results must not block the run — return base candidates as-is.""" - bench = await self._make_bench(skip_cached=True) + assert result == [] - candidates = [_make_candidate(name="red_teaming__adv_a", eval_hash="hash_a")] + async def test_skip_cached_true_keeps_unmatched_candidates(self): + bench = self._make_bench(skip_cached=True) + prior_sr = _make_scenario_result(result_id="sid-1") + prior_attacks = [ + _make_attack_result( + outcome=AttackOutcome.SUCCESS, + parent_collection="some_other_name", + parent_eval_hash="hash_x", + ), + ] bench._memory = MagicMock() - bench._memory.get_scenario_results.side_effect = RuntimeError("db down") + bench._memory.get_scenario_results.return_value = [prior_sr] + bench._memory.get_attack_results.return_value = prior_attacks - with patch.object(AdversarialBenchmark.__bases__[0], "_get_atomic_attacks_async", return_value=candidates): + with self._patch_factory_builder(): result = await bench._get_atomic_attacks_async() - assert result == candidates - - -@pytest.mark.usefixtures("patch_central_database") -class TestAdversarialBenchmarkSkipCachedInit: - """Tests for the skip_cached constructor surface.""" - - async def test_skip_cached_defaults_to_false(self): - await _fan_out(target_names=["adv_a"]) - bench = AdversarialBenchmark(objective_scorer=MagicMock(spec=TrueFalseScorer)) - assert bench._skip_cached is False - - async def test_skip_cached_can_be_set_true(self): - await _fan_out(target_names=["adv_a"]) - bench = AdversarialBenchmark( - objective_scorer=MagicMock(spec=TrueFalseScorer), - skip_cached=True, - ) - assert bench._skip_cached is True + assert len(result) == 1 # --------------------------------------------------------------------------- -# Scorer flexibility — stage 1 (Commit 7 / F4) +# Scorer flexibility — stage 1 # --------------------------------------------------------------------------- @pytest.mark.usefixtures("patch_central_database") class TestAdversarialBenchmarkScorerFlexibility: - """Tests for the widened objective_scorer annotation + isinstance guard (stage 1).""" - - def test_objective_scorer_annotation_is_scorer(self): - """The parameter annotation is the broad Scorer base class for forward compatibility.""" - import inspect - - from pyrit.score import Scorer - - sig = inspect.signature(AdversarialBenchmark.__init__) - annotation = sig.parameters["objective_scorer"].annotation - # ``Scorer | None`` resolves to ``Scorer | None`` at import time. - # str() captures both the runtime and stringified forms reliably. - assert "Scorer" in str(annotation) - assert Scorer is not None # sanity that the import resolves - - async def test_construct_accepts_truefalse_scorer_subclass(self): - """TrueFalseScorer remains the runtime-supported type; should construct cleanly.""" - await _fan_out(target_names=["adv_a"]) + """Tests for the widened ``objective_scorer`` annotation + ``isinstance`` guard (stage 1).""" + def test_construct_accepts_truefalse_scorer_subclass(self): + """``TrueFalseScorer`` remains the runtime-supported type; should construct cleanly.""" scorer = MagicMock(spec=TrueFalseScorer) bench = AdversarialBenchmark(objective_scorer=scorer) - assert bench._objective_scorer is scorer - - async def test_non_truefalse_scorer_raises_typeerror_with_pointer(self): - """A Scorer subclass that isn't TrueFalseScorer must raise TypeError with a clear pointer.""" - from pyrit.score import Scorer - - await _fan_out(target_names=["adv_a"]) - - # Bare Scorer (not TrueFalseScorer) — covers any future non-TF subclass. - non_tf_scorer = MagicMock(spec=Scorer) - - with pytest.raises(TypeError, match=r"requires a TrueFalseScorer.*follow-up"): - AdversarialBenchmark(objective_scorer=non_tf_scorer) diff --git a/tests/unit/setup/test_benchmark_initializer.py b/tests/unit/setup/test_benchmark_initializer.py deleted file mode 100644 index 8a04578abb..0000000000 --- a/tests/unit/setup/test_benchmark_initializer.py +++ /dev/null @@ -1,181 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# Licensed under the MIT license. - -"""Tests for BenchmarkInitializer.""" - -from unittest.mock import MagicMock - -import pytest - -from pyrit.prompt_target import PromptTarget -from pyrit.registry import TargetRegistry -from pyrit.registry.object_registries.attack_technique_registry import AttackTechniqueRegistry -from pyrit.setup.initializers import BenchmarkInitializer -from pyrit.setup.initializers.benchmark import DEFAULT_ADVERSARIAL_TAG_QUERY -from pyrit.setup.initializers.components.targets import TargetInitializerTags - - -@pytest.fixture(autouse=True) -def reset_registries(): - """Reset technique and target registries between tests.""" - AttackTechniqueRegistry.reset_instance() - TargetRegistry.reset_instance() - yield - AttackTechniqueRegistry.reset_instance() - TargetRegistry.reset_instance() - - -def _register_adversarial_target(*, name: str) -> PromptTarget: - """Register a mock adversarial-tagged target and return the instance.""" - target = MagicMock(spec=PromptTarget) - target.capabilities.includes.return_value = True - registry = TargetRegistry.get_registry_singleton() - registry.register_instance(target, name=name, tags=[TargetInitializerTags.ADVERSARIAL.value]) - return target - - -class TestBenchmarkInitializerBasic: - """Class metadata tests.""" - - def test_can_be_created(self): - init = BenchmarkInitializer() - assert init is not None - - def test_required_env_vars_is_empty(self): - """Initializer takes no required env vars; discovery happens via TargetRegistry.""" - init = BenchmarkInitializer() - assert init.required_env_vars == [] - - def test_supported_parameters_declares_target_names(self): - init = BenchmarkInitializer() - names = [p.name for p in init.supported_parameters] - assert "target_names" in names - - def test_default_adversarial_tag_query_matches_adversarial_only(self): - """The default discovery query is exactly ``TagQuery.all("adversarial")``.""" - assert DEFAULT_ADVERSARIAL_TAG_QUERY.matches({"adversarial"}) - assert not DEFAULT_ADVERSARIAL_TAG_QUERY.matches({"default"}) - assert not DEFAULT_ADVERSARIAL_TAG_QUERY.matches(set()) - - -class TestBenchmarkInitializerFanOut: - """Tests for the fan-out registration behavior.""" - - async def test_fans_out_one_spec_per_target_per_adversarial_technique(self): - """N targets * M adversarial-capable techniques = N*M fanned specs in the attack registry.""" - _register_adversarial_target(name="adv_a") - _register_adversarial_target(name="adv_b") - - init = BenchmarkInitializer() - await init.initialize_async() - - attack_registry = AttackTechniqueRegistry.get_registry_singleton() - fanned = attack_registry.get_by_tag(tag="benchmark_fanout") - assert len(fanned) > 0 - assert len(fanned) % 2 == 0, "Expected an even count: every adversarial technique fanned across both targets" - - fanned_names = {entry.name for entry in fanned} - for name in fanned_names: - assert "__" in name, f"Fanned spec name '{name}' missing '__' separator" - - async def test_fanned_spec_names_use_source_double_underscore_target(self): - """Spec naming contract: ``f'{source_spec.name}__{target_name}'``.""" - _register_adversarial_target(name="adv_single") - - init = BenchmarkInitializer() - await init.initialize_async() - - attack_registry = AttackTechniqueRegistry.get_registry_singleton() - fanned = attack_registry.get_by_tag(tag="model:adv_single") - assert len(fanned) > 0 - for entry in fanned: - assert entry.name.endswith("__adv_single") - source_name = entry.name.split("__", 1)[0] - assert source_name and "__" not in source_name - - async def test_fanned_specs_carry_benchmark_and_model_tags(self): - """Each fanned spec is tagged ``benchmark_fanout`` plus ``f'model:{name}'``.""" - _register_adversarial_target(name="adv_a") - _register_adversarial_target(name="adv_b") - - init = BenchmarkInitializer() - await init.initialize_async() - - attack_registry = AttackTechniqueRegistry.get_registry_singleton() - for entry in attack_registry.get_by_tag(tag="benchmark_fanout"): - assert "benchmark_fanout" in entry.tags - model_tags = [tag for tag in entry.tags if tag.startswith("model:")] - assert len(model_tags) == 1, f"Expected exactly one model:* tag on {entry.name}, got {model_tags}" - assert model_tags[0] in ("model:adv_a", "model:adv_b") - - async def test_registration_is_idempotent_across_re_init(self): - """Re-running initialize_async produces the same registry state (per-name idempotent).""" - _register_adversarial_target(name="adv_a") - - init = BenchmarkInitializer() - await init.initialize_async() - attack_registry = AttackTechniqueRegistry.get_registry_singleton() - first_count = len(attack_registry.get_by_tag(tag="benchmark_fanout")) - - await init.initialize_async() - second_count = len(attack_registry.get_by_tag(tag="benchmark_fanout")) - - assert first_count == second_count - - -class TestBenchmarkInitializerTargetNamesNarrowing: - """Tests for the optional ``target_names`` parameter.""" - - async def test_target_names_narrows_to_subset(self): - """When ``target_names`` is set, only those entries are fanned.""" - _register_adversarial_target(name="adv_a") - _register_adversarial_target(name="adv_b") - _register_adversarial_target(name="adv_c") - - init = BenchmarkInitializer() - init.params = {"target_names": ["adv_a", "adv_c"]} - await init.initialize_async() - - attack_registry = AttackTechniqueRegistry.get_registry_singleton() - model_b_specs = attack_registry.get_by_tag(tag="model:adv_b") - assert model_b_specs == [] - - model_a_specs = attack_registry.get_by_tag(tag="model:adv_a") - model_c_specs = attack_registry.get_by_tag(tag="model:adv_c") - assert len(model_a_specs) > 0 - assert len(model_c_specs) > 0 - - async def test_target_names_unknown_raises_with_discovered_list(self): - """Unknown ``target_names`` raise ``ValueError`` naming both the unknowns and the discovered set.""" - _register_adversarial_target(name="adv_a") - - init = BenchmarkInitializer() - init.params = {"target_names": ["nonexistent"]} - - with pytest.raises(ValueError, match=r"nonexistent.*adv_a"): - await init.initialize_async() - - async def test_empty_target_names_param_falls_back_to_default_query(self): - """An empty ``target_names`` list is treated as "no narrowing" (same as omitting it).""" - _register_adversarial_target(name="adv_a") - - init = BenchmarkInitializer() - init.params = {"target_names": []} - await init.initialize_async() - - attack_registry = AttackTechniqueRegistry.get_registry_singleton() - assert len(attack_registry.get_by_tag(tag="model:adv_a")) > 0 - - -class TestBenchmarkInitializerErrorMessages: - """Tests for the actionable error message on empty discovery.""" - - async def test_no_adversarial_targets_raises_with_actionable_message(self): - """``ValueError`` must name ``ADVERSARIAL_CHAT_*`` env vars and the ``TargetInitializer`` dependency.""" - init = BenchmarkInitializer() - with pytest.raises(ValueError) as exc_info: - await init.initialize_async() - - msg = str(exc_info.value) - assert "ADVERSARIAL_CHAT_" in msg - assert "TargetInitializer" in msg diff --git a/tests/unit/setup/test_targets_initializer.py b/tests/unit/setup/test_targets_initializer.py index 831f5e16f8..57b04413fe 100644 --- a/tests/unit/setup/test_targets_initializer.py +++ b/tests/unit/setup/test_targets_initializer.py @@ -218,7 +218,8 @@ def test_target_configs_have_unique_registry_names(self): ``TargetInitializer`` registers them (per ``BaseInstanceRegistry.register`` semantics, characterized in ``test_target_registry.py``). Only the second entry would survive in the registry, which breaks downstream - fan-out (``BenchmarkInitializer``) and is hard to diagnose. Tracked + scenarios that resolve targets by name (e.g. ``AdversarialBenchmark``'s + ``adversarial_targets`` parameter) and is hard to diagnose. Tracked as ``duplicate-registry-name`` in failure_mode_followups. """ registry_names = [config.registry_name for config in TARGET_CONFIGS] @@ -609,8 +610,8 @@ async def test_variant_skips_when_model_env_var_missing( os.environ.pop(f"{env_prefix}_KEY", None) async def test_all_variants_discoverable_via_adversarial_tag_query(self) -> None: - """End-to-end: variants + ``adversarial_chat`` are returned by adversarial-tag ``get_by_tag_query``.""" - from pyrit.registry.tag_query import TagQuery + """End-to-end: variants + ``adversarial_chat`` are returned by adversarial-tag ``get_by_tag``.""" + from pyrit.setup.initializers.components.targets import TargetInitializerTags os.environ["ADVERSARIAL_CHAT_ENDPOINT"] = "https://parent.openai.azure.com/openai/v1" os.environ["ADVERSARIAL_CHAT_KEY"] = "test_key" @@ -624,7 +625,7 @@ async def test_all_variants_discoverable_via_adversarial_tag_query(self) -> None await init.initialize_async() registry = TargetRegistry.get_registry_singleton() - matches = registry.get_by_tag_query(query=TagQuery.all("adversarial")) + matches = registry.get_by_tag(tag=TargetInitializerTags.ADVERSARIAL.value) match_names = {entry.name for entry in matches} expected = {"adversarial_chat"} | {name for name, _ in ADVERSARIAL_CHAT_VARIANTS} From 2ec2a493362ecb9c6e4c3d312671aab1bcda47c8 Mon Sep 17 00:00:00 2001 From: Victor Valbuena Date: Tue, 26 May 2026 12:38:00 -0700 Subject: [PATCH 19/40] DOCS: Trim adversarial benchmark scanner doc to basic-usage shape MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Per Rich's PR review (#1765), the scanner doc now mirrors the structure of `doc/scanner/airt.py`: brief description, CLI quickstart, available strategies, then a single Setup section followed by run + output cells. Removed sections: - `Scorer flexibility (forward-looking)` — factually wrong after `c7e7b93f` reverted the broad `Scorer` annotation and the runtime `TypeError` guard; the section described pre-revert behavior that no longer exists. - `Cross-run caching` — benchmark-only `skip_cached` feature; defer documentation until `skip_cached` is lifted to base `Scenario` (an enhancement tracked separately in the PR description). - `Bootstrapping from .pyrit_conf` — full YAML snippet collapsed into a pointer to `pyrit_scan --help` at the top of the doc, matching the other scanner pages. The last cell is now pure execution (`output_scenario_async`), matching the airt / foundry / garak scanner-doc convention. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- doc/scanner/benchmark.ipynb | 215 ++++++++++-------------------------- doc/scanner/benchmark.py | 151 +++---------------------- 2 files changed, 74 insertions(+), 292 deletions(-) diff --git a/doc/scanner/benchmark.ipynb b/doc/scanner/benchmark.ipynb index d3847460d8..cd3c6ff2a2 100644 --- a/doc/scanner/benchmark.ipynb +++ b/doc/scanner/benchmark.ipynb @@ -7,10 +7,10 @@ "source": [ "# Benchmark Scenarios\n", "\n", - "Benchmark scenarios compare the effectiveness of attacks across an axis that varies within the\n", - "scenario itself. The axis can be many things; currently, the only benchmark variant is the\n", - "adversarial benchmark, whose axis of change is the **adversarial chat helper model** used in\n", - "attacks." + "Benchmark scenarios compare attack effectiveness across an axis that varies within the scenario\n", + "itself. Currently the only benchmark variant is the adversarial benchmark, whose axis of change is\n", + "the **adversarial chat helper model** used in attacks. For full configuration options see\n", + "`pyrit_scan --help` and the [Scenarios Programming Guide](../code/scenarios/0_scenarios.ipynb)." ] }, { @@ -21,38 +21,13 @@ "## Adversarial Benchmark\n", "\n", "`AdversarialBenchmark` holds the objective target and dataset constant and varies the adversarial\n", - "chat model used to drive multi-turn attacks (and crescendo-style simulated conversations). Useful\n", - "for evaluating which adversarial helper models produce stronger or weaker attack success rates\n", - "against the same target.\n", + "chat model used to drive multi-turn attacks. Useful for evaluating which adversarial helper\n", + "models produce stronger or weaker attack success rates against the same target.\n", "\n", "Adversarial targets are user-provided via the `adversarial_targets` scenario parameter. Each name\n", "must already be registered in `TargetRegistry` — typically by `TargetInitializer` from the\n", - "`ADVERSARIAL_CHAT_*` env vars, or programmatically via `TargetRegistry.register_instance`. At run\n", - "time the scenario builds the `(technique × target × dataset)` cross-product directly: for each\n", - "adversarial-capable technique in `SCENARIO_TECHNIQUES` and each requested target, it constructs a\n", - "per-pair factory with `adversarial_chat` overridden to that target. No global\n", - "`AttackTechniqueRegistry` state is mutated.\n", - "\n", - "### Prerequisites\n", - "\n", - "Set at least one `ADVERSARIAL_CHAT_*` group of env vars (see `.env_example`):\n", - "\n", - "```bash\n", - "# Default adversarial target (always available when set)\n", - "ADVERSARIAL_CHAT_ENDPOINT=\"https://your-endpoint.openai.azure.com/openai/v1\"\n", - "ADVERSARIAL_CHAT_KEY=\"your-key\"\n", - "ADVERSARIAL_CHAT_MODEL=\"deployment-name\"\n", - "\n", - "# Optional turn-style variants — auto-discovered by TargetInitializer when set\n", - "ADVERSARIAL_CHAT_SINGLETURN_ENDPOINT=\"...\"\n", - "ADVERSARIAL_CHAT_SINGLETURN_KEY=\"...\"\n", - "ADVERSARIAL_CHAT_SINGLETURN_MODEL=\"...\"\n", - "# ADVERSARIAL_CHAT_MULTITURN_* and ADVERSARIAL_CHAT_REASONING_* follow the same pattern\n", - "```\n", - "\n", - "Use `pyrit_scan list-targets` to see every target currently registered, along with its tags.\n", - "\n", - "### CLI quickstart\n", + "`ADVERSARIAL_CHAT_*` env vars (see `.env_example`). Use `pyrit_scan --list-targets` to see every\n", + "target currently registered.\n", "\n", "```bash\n", "pyrit_scan benchmark.adversarial \\\n", @@ -62,20 +37,12 @@ " --max-dataset-size 4\n", "```\n", "\n", - "Pass multiple `--adversarial-targets` values to compare across models in a single run:\n", + "Pass multiple `--adversarial-targets` values to compare across models in a single run.\n", "\n", - "```bash\n", - "pyrit_scan benchmark.adversarial \\\n", - " --initializers target load_default_datasets \\\n", - " --target openai_chat \\\n", - " --adversarial-targets adversarial_chat adversarial_chat_singleturn adversarial_chat_reasoning \\\n", - " --max-dataset-size 4\n", - "```\n", - "\n", - "**Available strategies:** `light` (the default — a quick snapshot using the cheaper techniques),\n", - "`single_turn`, `multi_turn`, plus one concrete member per adversarial-capable source technique\n", - "(e.g. `red_teaming`, `tap`, `crescendo_simulated`). The default `light` aggregate deliberately\n", - "excludes `tap` and `crescendo_simulated`, which can take hours on a single run." + "**Available strategies:** `light` (default — a quick snapshot using the cheaper techniques),\n", + "`single_turn`, `multi_turn`, plus one member per adversarial-capable source technique\n", + "(e.g. `red_teaming`, `tap`, `crescendo_simulated`). The `light` aggregate excludes `tap` and\n", + "`crescendo_simulated`, which can take hours." ] }, { @@ -83,11 +50,7 @@ "id": "2", "metadata": {}, "source": [ - "## Setup\n", - "\n", - "`TargetInitializer` populates `TargetRegistry` from the `ADVERSARIAL_CHAT_*` env vars. The\n", - "scenario looks up adversarial targets by registry name from its `adversarial_targets` parameter,\n", - "so the targets must be registered before `scenario.run_async()` runs." + "## Setup" ] }, { @@ -95,48 +58,46 @@ "execution_count": null, "id": "3", "metadata": {}, - "outputs": [], + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "Found default environment files: ['C:\\\\Users\\\\vvalbuena\\\\.pyrit\\\\.env', 'C:\\\\Users\\\\vvalbuena\\\\.pyrit\\\\.env.local']\n", + "Loaded environment file: C:\\Users\\vvalbuena\\.pyrit\\.env\n", + "Loaded environment file: C:\\Users\\vvalbuena\\.pyrit\\.env.local\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" + ] + } + ], "source": [ "from pyrit.output import output_scenario_async\n", "from pyrit.prompt_target import OpenAIChatTarget\n", "from pyrit.scenario import DatasetConfiguration\n", "from pyrit.scenario.scenarios.benchmark import AdversarialBenchmark\n", "from pyrit.setup import IN_MEMORY, initialize_pyrit_async\n", - "from pyrit.setup.initializers import (\n", - " LoadDefaultDatasets,\n", - " ScorerInitializer,\n", - " TargetInitializer,\n", - ")\n", + "from pyrit.setup.initializers import LoadDefaultDatasets, ScorerInitializer, TargetInitializer\n", "\n", "await initialize_pyrit_async( # type: ignore\n", " memory_db_type=IN_MEMORY,\n", - " initializers=[\n", - " TargetInitializer(),\n", - " ScorerInitializer(),\n", - " LoadDefaultDatasets(),\n", - " ],\n", + " initializers=[TargetInitializer(), ScorerInitializer(), LoadDefaultDatasets()],\n", ")\n", "\n", "objective_target = OpenAIChatTarget()" ] }, - { - "cell_type": "markdown", - "id": "4", - "metadata": {}, - "source": [ - "## Run the benchmark\n", - "\n", - "Instantiate the scenario, then pass `adversarial_targets` through `initialize_async` via\n", - "`set_params_from_args` (programmatic equivalent of the `--adversarial-targets` CLI flag). The\n", - "default strategy (`light`) runs the benchmark-friendly subset of techniques for a quick\n", - "comparison." - ] - }, { "cell_type": "code", "execution_count": null, - "id": "5", + "id": "4", "metadata": {}, "outputs": [], "source": [ @@ -149,102 +110,40 @@ " dataset_config=dataset_config,\n", ")\n", "\n", - "baseline_result = await scenario.run_async() # type: ignore\n", - "\n", - "# Save this id to resume the run later via AdversarialBenchmark(scenario_result_id=...).\n", - "print(f\"Scenario result id: {baseline_result.id}\")\n", - "\n", - "await output_scenario_async(baseline_result)" - ] - }, - { - "cell_type": "markdown", - "id": "6", - "metadata": {}, - "source": [ - "## Cross-run caching\n", - "\n", - "Re-run the benchmark with `skip_cached=True` and atomic attacks that completed (`SUCCESS` or\n", - "`FAILURE` outcome) in any prior `COMPLETED` run of the same scenario name + version are skipped.\n", - "`ERROR` and `UNDETERMINED` outcomes always retry, so transient failures don't poison the cache.\n", - "\n", - "The cache key is `(atomic_attack_name, technique_eval_hash)` — two atomic attacks that share a\n", - "name but use different technique configurations (e.g. different scorer) don't cross-pollinate.\n", - "\n", - "Useful for:\n", - "* Resuming a long-running benchmark after a crash or `Ctrl-C`.\n", - "* Incrementally adding new adversarial targets without re-running the existing ones — the new\n", - " `{technique}__{target}_{dataset}` names don't match any cached entry and execute fresh." + "scenario_result = await scenario.run_async() # type: ignore" ] }, { "cell_type": "code", "execution_count": null, - "id": "7", + "id": "5", "metadata": {}, "outputs": [], "source": [ - "scenario_cached = AdversarialBenchmark(skip_cached=True)\n", - "scenario_cached.set_params_from_args(args={\"adversarial_targets\": [\"adversarial_chat\"]})\n", - "await scenario_cached.initialize_async( # type: ignore\n", - " objective_target=objective_target,\n", - " dataset_config=dataset_config,\n", - ")\n", - "\n", - "cached_result = await scenario_cached.run_async() # type: ignore\n", - "\n", - "await output_scenario_async(cached_result)" - ] - }, - { - "cell_type": "markdown", - "id": "8", - "metadata": {}, - "source": [ - "## Bootstrapping from `.pyrit_conf`\n", - "\n", - "For production / repeated runs, declare the initializer chain and the adversarial-target list in\n", - "`.pyrit_conf`. `TargetInitializer` must precede the scenario so the named targets are present in\n", - "`TargetRegistry` by the time the scenario builds atomic attacks.\n", - "\n", - "```yaml\n", - "memory_db_type: duckdb\n", - "initializers:\n", - " - name: target\n", - " - name: scorer\n", - " - name: load_default_datasets\n", - "scenario:\n", - " name: benchmark.adversarial\n", - " args:\n", - " adversarial_targets:\n", - " - adversarial_chat\n", - " - adversarial_chat_singleturn\n", - " - adversarial_chat_reasoning\n", - "```\n", - "\n", - "Then run `pyrit_scan --config-file .pyrit_conf` — the scenario reads `adversarial_targets` from\n", - "the config and builds the cross-product automatically. Unknown names raise `ValueError` listing\n", - "both the unknowns and every registered target so typos fail loudly." - ] - }, - { - "cell_type": "markdown", - "id": "9", - "metadata": {}, - "source": [ - "## Scorer flexibility (forward-looking)\n", - "\n", - "`AdversarialBenchmark.__init__`'s `objective_scorer` parameter is typed as the broad `Scorer`\n", - "base class. The current runtime contract is still `TrueFalseScorer` — passing a non-true/false\n", - "scorer raises `TypeError` with a pointer to the planned scoring follow-up. The annotation is\n", - "widened ahead of runtime support so callers coding against the signature today see the eventual\n", - "contract, and the follow-up that drops the guard won't require a signature change." + "await output_scenario_async(scenario_result)" ] } ], "metadata": { "jupytext": { "main_language": "python" + }, + "kernelspec": { + "display_name": "pyrit", + "language": "python", + "name": "python3" + }, + "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.13.13" } }, "nbformat": 4, diff --git a/doc/scanner/benchmark.py b/doc/scanner/benchmark.py index 3dc051c3e9..b166ecf76f 100644 --- a/doc/scanner/benchmark.py +++ b/doc/scanner/benchmark.py @@ -11,47 +11,22 @@ # %% [markdown] # # Benchmark Scenarios # -# Benchmark scenarios compare the effectiveness of attacks across an axis that varies within the -# scenario itself. The axis can be many things; currently, the only benchmark variant is the -# adversarial benchmark, whose axis of change is the **adversarial chat helper model** used in -# attacks. +# Benchmark scenarios compare attack effectiveness across an axis that varies within the scenario +# itself. Currently the only benchmark variant is the adversarial benchmark, whose axis of change is +# the **adversarial chat helper model** used in attacks. For full configuration options see +# `pyrit_scan --help` and the [Scenarios Programming Guide](../code/scenarios/0_scenarios.ipynb). # %% [markdown] # ## Adversarial Benchmark # # `AdversarialBenchmark` holds the objective target and dataset constant and varies the adversarial -# chat model used to drive multi-turn attacks (and crescendo-style simulated conversations). Useful -# for evaluating which adversarial helper models produce stronger or weaker attack success rates -# against the same target. +# chat model used to drive multi-turn attacks. Useful for evaluating which adversarial helper +# models produce stronger or weaker attack success rates against the same target. # # Adversarial targets are user-provided via the `adversarial_targets` scenario parameter. Each name # must already be registered in `TargetRegistry` — typically by `TargetInitializer` from the -# `ADVERSARIAL_CHAT_*` env vars, or programmatically via `TargetRegistry.register_instance`. At run -# time the scenario builds the `(technique × target × dataset)` cross-product directly: for each -# adversarial-capable technique in `SCENARIO_TECHNIQUES` and each requested target, it constructs a -# per-pair factory with `adversarial_chat` overridden to that target. No global -# `AttackTechniqueRegistry` state is mutated. -# -# ### Prerequisites -# -# Set at least one `ADVERSARIAL_CHAT_*` group of env vars (see `.env_example`): -# -# ```bash -# # Default adversarial target (always available when set) -# ADVERSARIAL_CHAT_ENDPOINT="https://your-endpoint.openai.azure.com/openai/v1" -# ADVERSARIAL_CHAT_KEY="your-key" -# ADVERSARIAL_CHAT_MODEL="deployment-name" -# -# # Optional turn-style variants — auto-discovered by TargetInitializer when set -# ADVERSARIAL_CHAT_SINGLETURN_ENDPOINT="..." -# ADVERSARIAL_CHAT_SINGLETURN_KEY="..." -# ADVERSARIAL_CHAT_SINGLETURN_MODEL="..." -# # ADVERSARIAL_CHAT_MULTITURN_* and ADVERSARIAL_CHAT_REASONING_* follow the same pattern -# ``` -# -# Use `pyrit_scan list-targets` to see every target currently registered, along with its tags. -# -# ### CLI quickstart +# `ADVERSARIAL_CHAT_*` env vars (see `.env_example`). Use `pyrit_scan --list-targets` to see every +# target currently registered. # # ```bash # pyrit_scan benchmark.adversarial \ @@ -61,27 +36,15 @@ # --max-dataset-size 4 # ``` # -# Pass multiple `--adversarial-targets` values to compare across models in a single run: +# Pass multiple `--adversarial-targets` values to compare across models in a single run. # -# ```bash -# pyrit_scan benchmark.adversarial \ -# --initializers target load_default_datasets \ -# --target openai_chat \ -# --adversarial-targets adversarial_chat adversarial_chat_singleturn adversarial_chat_reasoning \ -# --max-dataset-size 4 -# ``` -# -# **Available strategies:** `light` (the default — a quick snapshot using the cheaper techniques), -# `single_turn`, `multi_turn`, plus one concrete member per adversarial-capable source technique -# (e.g. `red_teaming`, `tap`, `crescendo_simulated`). The default `light` aggregate deliberately -# excludes `tap` and `crescendo_simulated`, which can take hours on a single run. +# **Available strategies:** `light` (default — a quick snapshot using the cheaper techniques), +# `single_turn`, `multi_turn`, plus one member per adversarial-capable source technique +# (e.g. `red_teaming`, `tap`, `crescendo_simulated`). The `light` aggregate excludes `tap` and +# `crescendo_simulated`, which can take hours. # %% [markdown] # ## Setup -# -# `TargetInitializer` populates `TargetRegistry` from the `ADVERSARIAL_CHAT_*` env vars. The -# scenario looks up adversarial targets by registry name from its `adversarial_targets` parameter, -# so the targets must be registered before `scenario.run_async()` runs. # %% from pyrit.output import output_scenario_async @@ -89,31 +52,15 @@ from pyrit.scenario import DatasetConfiguration from pyrit.scenario.scenarios.benchmark import AdversarialBenchmark from pyrit.setup import IN_MEMORY, initialize_pyrit_async -from pyrit.setup.initializers import ( - LoadDefaultDatasets, - ScorerInitializer, - TargetInitializer, -) +from pyrit.setup.initializers import LoadDefaultDatasets, ScorerInitializer, TargetInitializer await initialize_pyrit_async( # type: ignore memory_db_type=IN_MEMORY, - initializers=[ - TargetInitializer(), - ScorerInitializer(), - LoadDefaultDatasets(), - ], + initializers=[TargetInitializer(), ScorerInitializer(), LoadDefaultDatasets()], ) objective_target = OpenAIChatTarget() -# %% [markdown] -# ## Run the benchmark -# -# Instantiate the scenario, then pass `adversarial_targets` through `initialize_async` via -# `set_params_from_args` (programmatic equivalent of the `--adversarial-targets` CLI flag). The -# default strategy (`light`) runs the benchmark-friendly subset of techniques for a quick -# comparison. - # %% dataset_config = DatasetConfiguration(dataset_names=["harmbench"], max_dataset_size=4) @@ -124,71 +71,7 @@ dataset_config=dataset_config, ) -baseline_result = await scenario.run_async() # type: ignore - -# Save this id to resume the run later via AdversarialBenchmark(scenario_result_id=...). -print(f"Scenario result id: {baseline_result.id}") - -await output_scenario_async(baseline_result) - -# %% [markdown] -# ## Cross-run caching -# -# Re-run the benchmark with `skip_cached=True` and atomic attacks that completed (`SUCCESS` or -# `FAILURE` outcome) in any prior `COMPLETED` run of the same scenario name + version are skipped. -# `ERROR` and `UNDETERMINED` outcomes always retry, so transient failures don't poison the cache. -# -# The cache key is `(atomic_attack_name, technique_eval_hash)` — two atomic attacks that share a -# name but use different technique configurations (e.g. different scorer) don't cross-pollinate. -# -# Useful for: -# * Resuming a long-running benchmark after a crash or `Ctrl-C`. -# * Incrementally adding new adversarial targets without re-running the existing ones — the new -# `{technique}__{target}_{dataset}` names don't match any cached entry and execute fresh. +scenario_result = await scenario.run_async() # type: ignore # %% -scenario_cached = AdversarialBenchmark(skip_cached=True) -scenario_cached.set_params_from_args(args={"adversarial_targets": ["adversarial_chat"]}) -await scenario_cached.initialize_async( # type: ignore - objective_target=objective_target, - dataset_config=dataset_config, -) - -cached_result = await scenario_cached.run_async() # type: ignore - -await output_scenario_async(cached_result) - -# %% [markdown] -# ## Bootstrapping from `.pyrit_conf` -# -# For production / repeated runs, declare the initializer chain and the adversarial-target list in -# `.pyrit_conf`. `TargetInitializer` must precede the scenario so the named targets are present in -# `TargetRegistry` by the time the scenario builds atomic attacks. -# -# ```yaml -# memory_db_type: duckdb -# initializers: -# - name: target -# - name: scorer -# - name: load_default_datasets -# scenario: -# name: benchmark.adversarial -# args: -# adversarial_targets: -# - adversarial_chat -# - adversarial_chat_singleturn -# - adversarial_chat_reasoning -# ``` -# -# Then run `pyrit_scan --config-file .pyrit_conf` — the scenario reads `adversarial_targets` from -# the config and builds the cross-product automatically. Unknown names raise `ValueError` listing -# both the unknowns and every registered target so typos fail loudly. - -# %% [markdown] -# ## Scorer flexibility (forward-looking) -# -# `AdversarialBenchmark.__init__`'s `objective_scorer` parameter is typed as the broad `Scorer` -# base class. The current runtime contract is still `TrueFalseScorer` — passing a non-true/false -# scorer raises `TypeError` with a pointer to the planned scoring follow-up. The annotation is -# widened ahead of runtime support so callers coding against the signature today see the eventual -# contract, and the follow-up that drops the guard won't require a signature change. +await output_scenario_async(scenario_result) From cfdf5bf885eadacdd39d8950a357a3eb6329555b Mon Sep 17 00:00:00 2001 From: Victor Valbuena Date: Tue, 26 May 2026 13:16:42 -0700 Subject: [PATCH 20/40] DOCS: Add executed benchmark scanner notebook outputs Captured execution outputs from a real sanity run (small dataset, HarmBench, single adversarial target) so the rendered notebook on ReadTheDocs shows a representative scenario-results display. nbstripout, sanitize-notebook-paths, and strip-notebook-progress-bars pre-commit hooks were applied automatically (user paths sanitized, tqdm progress bars stripped, execution counts cleared per PyRIT's nbstripout config). Cell outputs (logs, scenario-results panel) are preserved per the same config. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- doc/scanner/benchmark.ipynb | 98 +++++++++++++++++++++++++++++++++---- 1 file changed, 88 insertions(+), 10 deletions(-) diff --git a/doc/scanner/benchmark.ipynb b/doc/scanner/benchmark.ipynb index cd3c6ff2a2..773311d199 100644 --- a/doc/scanner/benchmark.ipynb +++ b/doc/scanner/benchmark.ipynb @@ -63,9 +63,9 @@ "name": "stdout", "output_type": "stream", "text": [ - "Found default environment files: ['C:\\\\Users\\\\vvalbuena\\\\.pyrit\\\\.env', 'C:\\\\Users\\\\vvalbuena\\\\.pyrit\\\\.env.local']\n", - "Loaded environment file: C:\\Users\\vvalbuena\\.pyrit\\.env\n", - "Loaded environment file: C:\\Users\\vvalbuena\\.pyrit\\.env.local\n", + "Found default environment files: ['./.pyrit/.env', './.pyrit/.env.local']\n", + "Loaded environment file: ./.pyrit/.env\n", + "Loaded environment file: ./.pyrit/.env.local\n", "No new upgrade operations detected.\n" ] }, @@ -99,7 +99,22 @@ "execution_count": null, "id": "4", "metadata": {}, - "outputs": [], + "outputs": [ + { + "data": { + "application/vnd.jupyter.widget-view+json": { + "model_id": "b25354cde64b48088458fde2e1d2beba", + "version_major": 2, + "version_minor": 0 + }, + "text/plain": [ + "Executing AdversarialBenchmark: 0%| | 0/3 [00:00 Date: Tue, 26 May 2026 13:59:40 -0700 Subject: [PATCH 21/40] Add mcp Python SDK dependency Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- pyproject.toml | 1 + uv.lock | 94 ++++++++++++++++++++++++++++++++++++++++++++++++++ 2 files changed, 95 insertions(+) diff --git a/pyproject.toml b/pyproject.toml index bdc563d000..222ed1c86f 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -45,6 +45,7 @@ dependencies = [ "fastapi>=0.115.0", "httpx[http2]>=0.27.2", "jinja2>=3.1.6", + "mcp>=1.0,<2", "numpy>=1.26.0; python_version < '3.14'", "numpy>=2.3.0; python_version >= '3.14'", "openai>=2.2.0", diff --git a/uv.lock b/uv.lock index dec865a0c3..40b177484d 100644 --- a/uv.lock +++ b/uv.lock @@ -2121,6 +2121,15 @@ http2 = [ { name = "h2" }, ] +[[package]] +name = "httpx-sse" +version = "0.4.3" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/0f/4c/751061ffa58615a32c31b2d82e8482be8dd4a89154f003147acee90f2be9/httpx_sse-0.4.3.tar.gz", hash = "sha256:9b1ed0127459a66014aec3c56bebd93da3c1bc8bb6618c8082039a44889a755d", size = 15943, upload-time = "2025-10-10T21:48:22.271Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/d2/fd/6668e5aec43ab844de6fc74927e155a3b37bf40d7c3790e49fc0406b6578/httpx_sse-0.4.3-py3-none-any.whl", hash = "sha256:0ac1c9fe3c0afad2e0ebb25a934a59f4c7823b60792691f779fad2c5568830fc", size = 8960, upload-time = "2025-10-10T21:48:21.158Z" }, +] + [[package]] name = "huggingface-hub" version = "1.13.0" @@ -3192,6 +3201,31 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/af/33/ee4519fa02ed11a94aef9559552f3b17bb863f2ecfe1a35dc7f548cde231/matplotlib_inline-0.2.1-py3-none-any.whl", hash = "sha256:d56ce5156ba6085e00a9d54fead6ed29a9c47e215cd1bba2e976ef39f5710a76", size = 9516, upload-time = "2025-10-23T09:00:20.675Z" }, ] +[[package]] +name = "mcp" +version = "1.27.1" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "anyio" }, + { name = "httpx" }, + { name = "httpx-sse" }, + { name = "jsonschema" }, + { name = "pydantic" }, + { name = "pydantic-settings" }, + { name = "pyjwt", extra = ["crypto"] }, + { name = "python-multipart" }, + { name = "pywin32", marker = "sys_platform == 'win32'" }, + { name = "sse-starlette" }, + { name = "starlette" }, + { name = "typing-extensions" }, + { name = "typing-inspection" }, + { name = "uvicorn", marker = "sys_platform != 'emscripten'" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/38/83/d1efe7c2980d8a3afa476f4e3d42d53dd54c0ab94c27bee5d755b45c8b73/mcp-1.27.1.tar.gz", hash = "sha256:0f47e1820f8f8f941466b39749eb1d1839a04caddca2bc60e9d46e8a99914924", size = 608458, upload-time = "2026-05-08T16:50:12.601Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/fd/73/42d9596facebdb533b7f0b86c1b0364ef350d1f8ba78b1052e8a58b48b65/mcp-1.27.1-py3-none-any.whl", hash = "sha256:1af3c4203b329430fde7a87b4fcb6392a041f5cb851fd68fc674016ab4e7c06f", size = 216260, upload-time = "2026-05-08T16:50:10.547Z" }, +] + [[package]] name = "mdit-py-plugins" version = "0.5.0" @@ -5014,6 +5048,20 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/36/c7/cfc8e811f061c841d7990b0201912c3556bfeb99cdcb7ed24adc8d6f8704/pydantic_core-2.41.5-pp311-pypy311_pp73-win_amd64.whl", hash = "sha256:56121965f7a4dc965bff783d70b907ddf3d57f6eba29b6d2e5dabfaf07799c51", size = 2145302, upload-time = "2025-11-04T13:43:46.64Z" }, ] +[[package]] +name = "pydantic-settings" +version = "2.14.1" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "pydantic" }, + { name = "python-dotenv" }, + { name = "typing-inspection" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/07/60/1d1e59c9c90d54591469ada7d268251f71c24bdb765f1a8a832cee8c6653/pydantic_settings-2.14.1.tar.gz", hash = "sha256:e874d3bec7e787b0c9958277956ed9b4dd5de6a80e162188fdaff7c5e26fd5fa", size = 235551, upload-time = "2026-05-08T13:40:06.542Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/ae/8d/f1af3832f5e6eb13ba94ee809e72b8ecb5eef226d27ee0bef7d963d943c7/pydantic_settings-2.14.1-py3-none-any.whl", hash = "sha256:6e3c7edfd8277687cdc598f56e5cff0e9bfff0910a3749deaa8d4401c3a2b9de", size = 60964, upload-time = "2026-05-08T13:40:04.958Z" }, +] + [[package]] name = "pydash" version = "8.0.5" @@ -5170,6 +5218,7 @@ dependencies = [ { name = "fastapi" }, { name = "httpx", extra = ["http2"] }, { name = "jinja2" }, + { name = "mcp" }, { name = "numpy", version = "2.2.6", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.14'" }, { name = "numpy", version = "2.4.4", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.14'" }, { name = "openai" }, @@ -5306,6 +5355,7 @@ requires-dist = [ { name = "ipykernel", marker = "extra == 'all'", specifier = ">=6.29.5" }, { name = "jinja2", specifier = ">=3.1.6" }, { name = "jupyter", marker = "extra == 'all'", specifier = ">=1.1.1" }, + { name = "mcp", specifier = ">=1.0,<2" }, { name = "ml-collections", marker = "extra == 'all'", specifier = ">=1.1.0" }, { name = "ml-collections", marker = "extra == 'gcg'", specifier = ">=1.1.0" }, { name = "numpy", marker = "python_full_version < '3.14'", specifier = ">=1.26.0" }, @@ -5498,6 +5548,15 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/51/e5/fecf13f06e5e5f67e8837d777d1bc43fac0ed2b77a676804df5c34744727/python_json_logger-4.0.0-py3-none-any.whl", hash = "sha256:af09c9daf6a813aa4cc7180395f50f2a9e5fa056034c9953aec92e381c5ba1e2", size = 15548, upload-time = "2025-10-06T04:15:17.553Z" }, ] +[[package]] +name = "python-multipart" +version = "0.0.29" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/4e/fe/70bd71a6738b09a0bdf6480ca6436b167469ca4578b2a0efbe390b4b0e70/python_multipart-0.0.29.tar.gz", hash = "sha256:643e93849196645e2dbdd81a0f8829a23123ad7f797a84a364c6fb3563f18904", size = 45678, upload-time = "2026-05-17T17:29:47.654Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/8f/cb/769cfc37177252872a45a71f3fbdde9d51b471a3f3c14bfe95dde3407386/python_multipart-0.0.29-py3-none-any.whl", hash = "sha256:2ddcc971cef266225f54f552d8fa10bcfbb1f14446caec199060daac59ff2d69", size = 29640, upload-time = "2026-05-17T17:29:45.69Z" }, +] + [[package]] name = "pytz" version = "2025.2" @@ -5507,6 +5566,28 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/81/c4/34e93fe5f5429d7570ec1fa436f1986fb1f00c3e0f43a589fe2bbcd22c3f/pytz-2025.2-py2.py3-none-any.whl", hash = "sha256:5ddf76296dd8c44c26eb8f4b6f35488f3ccbf6fbbd7adee0b7262d43f0ec2f00", size = 509225, upload-time = "2025-03-25T02:24:58.468Z" }, ] +[[package]] +name = "pywin32" +version = "311" +source = { registry = "https://pypi.org/simple" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/7b/40/44efbb0dfbd33aca6a6483191dae0716070ed99e2ecb0c53683f400a0b4f/pywin32-311-cp310-cp310-win32.whl", hash = "sha256:d03ff496d2a0cd4a5893504789d4a15399133fe82517455e78bad62efbb7f0a3", size = 8760432, upload-time = "2025-07-14T20:13:05.9Z" }, + { url = "https://files.pythonhosted.org/packages/5e/bf/360243b1e953bd254a82f12653974be395ba880e7ec23e3731d9f73921cc/pywin32-311-cp310-cp310-win_amd64.whl", hash = "sha256:797c2772017851984b97180b0bebe4b620bb86328e8a884bb626156295a63b3b", size = 9590103, upload-time = "2025-07-14T20:13:07.698Z" }, + { url = "https://files.pythonhosted.org/packages/57/38/d290720e6f138086fb3d5ffe0b6caa019a791dd57866940c82e4eeaf2012/pywin32-311-cp310-cp310-win_arm64.whl", hash = "sha256:0502d1facf1fed4839a9a51ccbcc63d952cf318f78ffc00a7e78528ac27d7a2b", size = 8778557, upload-time = "2025-07-14T20:13:11.11Z" }, + { url = "https://files.pythonhosted.org/packages/7c/af/449a6a91e5d6db51420875c54f6aff7c97a86a3b13a0b4f1a5c13b988de3/pywin32-311-cp311-cp311-win32.whl", hash = "sha256:184eb5e436dea364dcd3d2316d577d625c0351bf237c4e9a5fabbcfa5a58b151", size = 8697031, upload-time = "2025-07-14T20:13:13.266Z" }, + { url = "https://files.pythonhosted.org/packages/51/8f/9bb81dd5bb77d22243d33c8397f09377056d5c687aa6d4042bea7fbf8364/pywin32-311-cp311-cp311-win_amd64.whl", hash = "sha256:3ce80b34b22b17ccbd937a6e78e7225d80c52f5ab9940fe0506a1a16f3dab503", size = 9508308, upload-time = "2025-07-14T20:13:15.147Z" }, + { url = "https://files.pythonhosted.org/packages/44/7b/9c2ab54f74a138c491aba1b1cd0795ba61f144c711daea84a88b63dc0f6c/pywin32-311-cp311-cp311-win_arm64.whl", hash = "sha256:a733f1388e1a842abb67ffa8e7aad0e70ac519e09b0f6a784e65a136ec7cefd2", size = 8703930, upload-time = "2025-07-14T20:13:16.945Z" }, + { url = "https://files.pythonhosted.org/packages/e7/ab/01ea1943d4eba0f850c3c61e78e8dd59757ff815ff3ccd0a84de5f541f42/pywin32-311-cp312-cp312-win32.whl", hash = "sha256:750ec6e621af2b948540032557b10a2d43b0cee2ae9758c54154d711cc852d31", size = 8706543, upload-time = "2025-07-14T20:13:20.765Z" }, + { url = "https://files.pythonhosted.org/packages/d1/a8/a0e8d07d4d051ec7502cd58b291ec98dcc0c3fff027caad0470b72cfcc2f/pywin32-311-cp312-cp312-win_amd64.whl", hash = "sha256:b8c095edad5c211ff31c05223658e71bf7116daa0ecf3ad85f3201ea3190d067", size = 9495040, upload-time = "2025-07-14T20:13:22.543Z" }, + { url = "https://files.pythonhosted.org/packages/ba/3a/2ae996277b4b50f17d61f0603efd8253cb2d79cc7ae159468007b586396d/pywin32-311-cp312-cp312-win_arm64.whl", hash = "sha256:e286f46a9a39c4a18b319c28f59b61de793654af2f395c102b4f819e584b5852", size = 8710102, upload-time = "2025-07-14T20:13:24.682Z" }, + { url = "https://files.pythonhosted.org/packages/a5/be/3fd5de0979fcb3994bfee0d65ed8ca9506a8a1260651b86174f6a86f52b3/pywin32-311-cp313-cp313-win32.whl", hash = "sha256:f95ba5a847cba10dd8c4d8fefa9f2a6cf283b8b88ed6178fa8a6c1ab16054d0d", size = 8705700, upload-time = "2025-07-14T20:13:26.471Z" }, + { url = "https://files.pythonhosted.org/packages/e3/28/e0a1909523c6890208295a29e05c2adb2126364e289826c0a8bc7297bd5c/pywin32-311-cp313-cp313-win_amd64.whl", hash = "sha256:718a38f7e5b058e76aee1c56ddd06908116d35147e133427e59a3983f703a20d", size = 9494700, upload-time = "2025-07-14T20:13:28.243Z" }, + { url = "https://files.pythonhosted.org/packages/04/bf/90339ac0f55726dce7d794e6d79a18a91265bdf3aa70b6b9ca52f35e022a/pywin32-311-cp313-cp313-win_arm64.whl", hash = "sha256:7b4075d959648406202d92a2310cb990fea19b535c7f4a78d3f5e10b926eeb8a", size = 8709318, upload-time = "2025-07-14T20:13:30.348Z" }, + { url = "https://files.pythonhosted.org/packages/c9/31/097f2e132c4f16d99a22bfb777e0fd88bd8e1c634304e102f313af69ace5/pywin32-311-cp314-cp314-win32.whl", hash = "sha256:b7a2c10b93f8986666d0c803ee19b5990885872a7de910fc460f9b0c2fbf92ee", size = 8840714, upload-time = "2025-07-14T20:13:32.449Z" }, + { url = "https://files.pythonhosted.org/packages/90/4b/07c77d8ba0e01349358082713400435347df8426208171ce297da32c313d/pywin32-311-cp314-cp314-win_amd64.whl", hash = "sha256:3aca44c046bd2ed8c90de9cb8427f581c479e594e99b5c0bb19b29c10fd6cb87", size = 9656800, upload-time = "2025-07-14T20:13:34.312Z" }, + { url = "https://files.pythonhosted.org/packages/c0/d2/21af5c535501a7233e734b8af901574572da66fcc254cb35d0609c9080dd/pywin32-311-cp314-cp314-win_arm64.whl", hash = "sha256:a508e2d9025764a8270f93111a970e1d0fbfc33f4153b388bb649b7eec4f9b42", size = 8932540, upload-time = "2025-07-14T20:13:36.379Z" }, +] + [[package]] name = "pywinpty" version = "3.0.2" @@ -6561,6 +6642,19 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/2a/ae/57d1d7af907e20c077e113e0e4976f87b82c0a415403d99284a262229dd0/srsly-2.5.3-cp314-cp314t-win_arm64.whl", hash = "sha256:d822083fe26ec6728bd8c273ac121fc4ab3864a0fdf0cf0ff3efb188fcd209ed", size = 650229, upload-time = "2026-03-23T11:56:46.148Z" }, ] +[[package]] +name = "sse-starlette" +version = "3.4.4" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "anyio" }, + { name = "starlette" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/f7/2b/58abc2d1fd397e7dde08e947e05c884d8ef2f78d5e2588c17a12d42d6994/sse_starlette-3.4.4.tar.gz", hash = "sha256:07e0fa0460138baf25cdd5fb28683472c3995dc1642225191b3832d62526bcb0", size = 31819, upload-time = "2026-05-12T17:37:17.019Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/dc/67/805710444ea8cc75fbf70b920ed431a560c4bf9c57f7d5a3117213189399/sse_starlette-3.4.4-py3-none-any.whl", hash = "sha256:3f4dd50d8aed2771a091f3a83000323fc3844541c16b4fe585ae2420cc6df973", size = 16514, upload-time = "2026-05-12T17:37:15.601Z" }, +] + [[package]] name = "stack-data" version = "0.6.3" From c7d65d3d633f7357a2b7d3323b8b470db4b5af36 Mon Sep 17 00:00:00 2001 From: Victor Valbuena Date: Tue, 26 May 2026 14:34:13 -0700 Subject: [PATCH 22/40] Add tools/ package with tool_loop decorator and CallableToolBackend Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- pyrit/exceptions/__init__.py | 4 + pyrit/exceptions/exception_classes.py | 64 ++++ pyrit/tools/__init__.py | 58 ++++ pyrit/tools/backend.py | 83 +++++ pyrit/tools/callable_backend.py | 134 ++++++++ pyrit/tools/models.py | 241 ++++++++++++++ pyrit/tools/parsers.py | 55 ++++ tests/unit/tools/__init__.py | 2 + tests/unit/tools/conftest.py | 307 ++++++++++++++++++ tests/unit/tools/echo_mcp_server.py | 57 ++++ .../unit/tools/test_callable_tool_backend.py | 180 ++++++++++ tests/unit/tools/test_tool_loop_decorator.py | 289 +++++++++++++++++ 12 files changed, 1474 insertions(+) create mode 100644 pyrit/tools/__init__.py create mode 100644 pyrit/tools/backend.py create mode 100644 pyrit/tools/callable_backend.py create mode 100644 pyrit/tools/models.py create mode 100644 pyrit/tools/parsers.py create mode 100644 tests/unit/tools/__init__.py create mode 100644 tests/unit/tools/conftest.py create mode 100644 tests/unit/tools/echo_mcp_server.py create mode 100644 tests/unit/tools/test_callable_tool_backend.py create mode 100644 tests/unit/tools/test_tool_loop_decorator.py diff --git a/pyrit/exceptions/__init__.py b/pyrit/exceptions/__init__.py index abd42de031..9baea33c1a 100644 --- a/pyrit/exceptions/__init__.py +++ b/pyrit/exceptions/__init__.py @@ -10,6 +10,8 @@ MissingPromptPlaceholderException, PyritException, RateLimitException, + ToolCallLoopLimitExceeded, + ToolCallNotSupported, get_retry_max_num_attempts, handle_bad_request_exception, pyrit_custom_result_retry, @@ -59,4 +61,6 @@ "set_execution_context", "set_retry_collector", "execution_context", + "ToolCallLoopLimitExceeded", + "ToolCallNotSupported", ] diff --git a/pyrit/exceptions/exception_classes.py b/pyrit/exceptions/exception_classes.py index b2fc55440b..5d0014aa3d 100644 --- a/pyrit/exceptions/exception_classes.py +++ b/pyrit/exceptions/exception_classes.py @@ -233,6 +233,70 @@ def __init__(self, *, message: str = "No prompt placeholder") -> None: super().__init__(message=message) +class ToolCallNotSupported(PyritException): + """ + Raised when a target produces a tool call that the configured + :class:`~pyrit.tools.ToolEventPolicy` does not permit to execute + (``ToolEventBehavior.RAISE``, or ``EXECUTE`` without a backend). + + The ``partial_conversation`` attribute carries every message produced + up to and including the assistant turn that contained the offending + tool call(s). Consumers can inspect it to log the surfaced tool-use + attempt. + """ + + def __init__( + self, + *, + message: str = "Tool call not supported by configured policy.", + partial_conversation: Optional[list["Message"]] = None, + ) -> None: + """ + Initialize the exception. + + Args: + message (str): Human-readable error description. + partial_conversation (Optional[list[Message]]): Messages produced by + the target up to (and including) the assistant turn that + contained the disallowed tool call(s). + """ + super().__init__(status_code=400, message=message) + self.partial_conversation: list[Message] = ( + list(partial_conversation) if partial_conversation is not None else [] + ) + + +class ToolCallLoopLimitExceeded(PyritException): + """ + Raised when the tool-use loop runs for more than + ``ToolEventPolicy.max_tool_iterations`` iterations without the model + producing a stop response. + + The ``partial_conversation`` attribute carries every message produced + across all completed iterations. Consumers can inspect it to debug + runaway agentic behavior. + """ + + def __init__( + self, + *, + message: str = "Tool loop exceeded max_tool_iterations without a stop response.", + partial_conversation: Optional[list["Message"]] = None, + ) -> None: + """ + Initialize the exception. + + Args: + message (str): Human-readable error description. + partial_conversation (Optional[list[Message]]): Messages produced by + the target across every completed iteration of the tool loop. + """ + super().__init__(status_code=400, message=message) + self.partial_conversation: list[Message] = ( + list(partial_conversation) if partial_conversation is not None else [] + ) + + def pyrit_custom_result_retry( retry_function: Callable[..., bool], retry_max_num_attempts: Optional[int] = None ) -> Callable[..., Any]: diff --git a/pyrit/tools/__init__.py b/pyrit/tools/__init__.py new file mode 100644 index 0000000000..9830f17758 --- /dev/null +++ b/pyrit/tools/__init__.py @@ -0,0 +1,58 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT license. + +""" +Generic tool-use scaffolding for :class:`~pyrit.prompt_target.PromptTarget`. + +This package provides a transport-agnostic tool-calling loop. The +:func:`tool_loop` decorator, when applied to ``send_prompt_async``, runs +the standard PyRIT validate+normalize work once and then repeatedly +re-enters the target's protected ``_send_prompt_to_target_async`` until +the model issues a stop response (or a configured limit is hit). + +A target opts in by declaring two collaborators: + +* ``self._tool_parser`` — a :class:`ToolCallParser` that walks a + response message and extracts pending :class:`ToolCall` instances. +* ``self.configuration.tool_event_policy`` — a :class:`ToolEventPolicy` + whose :class:`ToolEventBehavior` decides whether to ``EXECUTE``, + ``RAISE``, or ``RETURN_RAW`` on each detected call. + +When the policy is ``EXECUTE``, calls are dispatched through +``self.configuration.tool_backend``, an implementation of +:class:`ToolBackend`. :class:`CallableToolBackend` is the pure-Python +backend shipped here; :class:`MCPToolBackend` ships in C3 and proxies +through one or more MCP servers. + +The :class:`ToolBackend` Protocol is intentionally distinct from +:mod:`pyrit.registry` — that namespace is reserved for framework-level +identity registries (``TargetRegistry``, ``ScorerRegistry``) that +register named singletons for CLI lookup, which a per-target tool +dispatch table is not. + +Wiring of ``@tool_loop`` onto :class:`PromptTarget.send_prompt_async` +and of the ``tool_event_policy`` / ``tool_backend`` fields onto +:class:`TargetConfiguration` lands in C4/C5. + +The two exception types the loop raises +(:class:`~pyrit.exceptions.ToolCallNotSupported` and +:class:`~pyrit.exceptions.ToolCallLoopLimitExceeded`) live in +:mod:`pyrit.exceptions` alongside the rest of PyRIT's exception +catalog, so non-tools callers (attacks, normalizers) can import them +without taking a subsystem-level dependency on ``pyrit.tools``. +""" + +from pyrit.tools.backend import ToolBackend +from pyrit.tools.callable_backend import CallableToolBackend +from pyrit.tools.models import ToolCall, ToolEventBehavior, ToolEventPolicy, tool_loop +from pyrit.tools.parsers import ToolCallParser + +__all__ = [ + "CallableToolBackend", + "ToolBackend", + "ToolCall", + "ToolCallParser", + "ToolEventBehavior", + "ToolEventPolicy", + "tool_loop", +] diff --git a/pyrit/tools/backend.py b/pyrit/tools/backend.py new file mode 100644 index 0000000000..3274355d6e --- /dev/null +++ b/pyrit/tools/backend.py @@ -0,0 +1,83 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT license. + +from __future__ import annotations + +from typing import TYPE_CHECKING, Any, Protocol, runtime_checkable + +if TYPE_CHECKING: + from pyrit.tools.models import ToolCall + + +@runtime_checkable +class ToolBackend(Protocol): + """ + Protocol for backends that dispatch tool calls produced by a target. + + A :class:`ToolBackend` is a per-target dispatch table — it owns the + ``name -> async callable`` mapping a target uses to execute the tool + calls extracted from a model response. This is intentionally distinct + from :mod:`pyrit.registry`, whose ``Registry`` classes register named + framework singletons (targets, scorers, attacks) for CLI lookup. + + Two concrete implementations ship with PyRIT: + + * :class:`~pyrit.tools.CallableToolBackend` — pure-Python backend + backed by ``async def`` callables. Useful for unit tests and for + embedding tools inside the PyRIT process. + * :class:`pyrit.tools.MCPToolBackend` (lands in C3) — proxies + dispatch through one or more MCP servers. + + The :attr:`schemas` property exposes the JSON-schema descriptors the + target injects into its request body (e.g. ``tools=[...]`` for the + OpenAI APIs). + + :meth:`dispatch_all_sequential_async` is the contract the tool loop + uses: backends that wish to parallelize dispatch should override it. + The default sequencing — one ``await dispatch_async`` per call, in + declaration order — is what every PyRIT backend ships with today. + """ + + @property + def schemas(self) -> list[dict[str, Any]]: + """ + The JSON-schema descriptors for every tool the backend exposes. + + Returns: + list[dict[str, Any]]: One schema per tool, in a target-agnostic + format that concrete targets serialize into their request + body. + """ + ... + + async def dispatch_async(self, call: ToolCall) -> dict[str, Any]: + """ + Execute a single tool call and return the structured result. + + Implementations MUST NOT raise on tool-side failures; they MUST + return an error envelope (e.g. ``{"error": "...", "tool": "..."}``) + so the tool loop can carry the failure back to the model. + + Args: + call (ToolCall): The tool call to dispatch. + + Returns: + dict[str, Any]: The structured tool result. + """ + ... + + async def dispatch_all_sequential_async( + self, + calls: list[ToolCall], + ) -> list[tuple[ToolCall, dict[str, Any]]]: + """ + Dispatch every call in *calls* sequentially, preserving declaration order. + + Args: + calls (list[ToolCall]): The calls to dispatch, in declaration order. + + Returns: + list[tuple[ToolCall, dict[str, Any]]]: ``(call, result)`` pairs, + in the same order as *calls*. + """ + ... diff --git a/pyrit/tools/callable_backend.py b/pyrit/tools/callable_backend.py new file mode 100644 index 0000000000..defbe94ed1 --- /dev/null +++ b/pyrit/tools/callable_backend.py @@ -0,0 +1,134 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT license. + +from __future__ import annotations + +import logging +from typing import TYPE_CHECKING, Any + +if TYPE_CHECKING: + from collections.abc import Awaitable, Callable + + from pyrit.tools.models import ToolCall + +logger = logging.getLogger(__name__) + + +class CallableToolBackend: + """ + Pure-Python :class:`~pyrit.tools.ToolBackend` backed by a name -> ``async def`` + mapping. Useful for unit tests and for embedding small tools inside the + PyRIT process without standing up an MCP server. + + The backend dispatches sequentially in declaration order. Tool-side + failures (raised exceptions, missing names, allow-list rejections) + are converted into structured error envelopes so the tool loop can + forward them back to the model as ``function_call_output`` content + rather than aborting the conversation. + """ + + def __init__( + self, + *, + callables: dict[str, Callable[[dict[str, Any]], Awaitable[Any]]], + schemas: list[dict[str, Any]] | None = None, + allowed_tools: set[str] | None = None, + fail_on_missing_function: bool = True, + ) -> None: + """ + Initialize the backend. + + Args: + callables (dict[str, Callable[[dict[str, Any]], Awaitable[Any]]]): + Map from tool name to an ``async def`` that accepts a parsed + arguments dict and returns the tool result. Results are + serialized by the tool loop via :func:`json.dumps`. + schemas (list[dict[str, Any]] | None): JSON-schema descriptors + injected into the target's request body. Defaults to an empty + list when omitted. + allowed_tools (set[str] | None): Optional allow-list of tool + names; calls whose name is not in this set surface as + ``tool_not_allowed`` envelopes without invoking the callable. + Defaults to None (no allow-list; every registered tool is + callable). + fail_on_missing_function (bool): When True (default), an unknown + tool name raises :class:`KeyError`. When False, the backend + returns a ``tool_not_registered`` envelope so the model can + recover. + """ + self._callables = dict(callables) + self._schemas: list[dict[str, Any]] = list(schemas) if schemas is not None else [] + self._allowed_tools = set(allowed_tools) if allowed_tools is not None else None + self._fail_on_missing_function = fail_on_missing_function + + @property + def schemas(self) -> list[dict[str, Any]]: + """The JSON-schema descriptors for the tools in this backend.""" + return list(self._schemas) + + async def dispatch_async(self, call: ToolCall) -> dict[str, Any]: + """ + Dispatch a single tool call. Tool failures are converted into + structured envelopes; only configuration errors (missing tool with + ``fail_on_missing_function=True``) propagate as exceptions. + + Args: + call (ToolCall): The call to dispatch. + + Returns: + dict[str, Any]: The tool's result, or a structured error envelope. + + Raises: + KeyError: When the tool name is not registered and + ``fail_on_missing_function=True``. + """ + if self._allowed_tools is not None and call.name not in self._allowed_tools: + logger.info("Rejecting disallowed tool call: %s", call.name) + return { + "error": "tool_not_allowed", + "tool": call.name, + "allowed_tools": sorted(self._allowed_tools), + } + + fn = self._callables.get(call.name) + if fn is None: + if self._fail_on_missing_function: + raise KeyError(f"Tool '{call.name}' is not registered.") + available = sorted(self._callables.keys()) + logger.warning("Tool '%s' not registered. Available: %s", call.name, available) + return { + "error": "tool_not_registered", + "tool": call.name, + "available_tools": available, + } + + try: + result = await fn(call.arguments) + except Exception as ex: + logger.warning("Tool '%s' raised %s: %s", call.name, type(ex).__name__, ex) + return { + "error": "tool_execution_failed", + "tool": call.name, + "detail": str(ex), + } + return result if isinstance(result, dict) else {"result": result} + + async def dispatch_all_sequential_async( + self, + calls: list[ToolCall], + ) -> list[tuple[ToolCall, dict[str, Any]]]: + """ + Dispatch *calls* sequentially in declaration order. + + Args: + calls (list[ToolCall]): Calls to dispatch. + + Returns: + list[tuple[ToolCall, dict[str, Any]]]: ``(call, result)`` pairs + in the same order as *calls*. + """ + results: list[tuple[ToolCall, dict[str, Any]]] = [] + for call in calls: + result = await self.dispatch_async(call) + results.append((call, result)) + return results diff --git a/pyrit/tools/models.py b/pyrit/tools/models.py new file mode 100644 index 0000000000..0e05f9be9d --- /dev/null +++ b/pyrit/tools/models.py @@ -0,0 +1,241 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT license. + +from __future__ import annotations + +import enum +import functools +import json +import logging +from dataclasses import dataclass, field +from typing import TYPE_CHECKING, Any + +from pyrit.exceptions import ToolCallLoopLimitExceeded, ToolCallNotSupported +from pyrit.models import Message, MessagePiece + +if TYPE_CHECKING: + from collections.abc import Awaitable, Callable + + from pyrit.tools.backend import ToolBackend + from pyrit.tools.parsers import ToolCallParser + +logger = logging.getLogger(__name__) + + +@dataclass(frozen=True) +class ToolCall: + """ + A parsed tool call extracted from a target response. + + Concrete :class:`~pyrit.tools.ToolCallParser` implementations build + :class:`ToolCall` instances by walking the response message pieces. + The :attr:`raw_envelope` carries the original target-specific dict + (e.g. the function_call JSON section) so dispatchers and observers + can recover provider-specific fields without re-parsing. + + Attributes: + call_id (str): The provider-issued call identifier; must round-trip + into the matching ``function_call_output`` piece. + name (str): The tool name to dispatch. + arguments (dict[str, Any]): The parsed JSON arguments. + raw_envelope (dict[str, Any]): The original provider envelope. + """ + + call_id: str + name: str + arguments: dict[str, Any] + raw_envelope: dict[str, Any] = field(default_factory=dict) + + +class ToolEventBehavior(enum.Enum): + """ + What the tool loop should do when a target response contains a + pending tool call. + + Values: + EXECUTE: Dispatch the call via ``configuration.tool_backend`` + and re-enter the target with the tool output appended. + This is the standard agentic loop behavior. + RAISE: Raise :class:`~pyrit.exceptions.ToolCallNotSupported` with + the partial conversation attached. Useful for red-team + attacks that want to observe attempted tool use without + allowing execution. + RETURN_RAW: Return the assistant response containing the tool + call as-is, without dispatching. Useful when a caller wants + to inspect tool calls in-band (e.g. a scorer that scores + attempted tool use). + """ + + EXECUTE = "execute" + RAISE = "raise" + RETURN_RAW = "return_raw" + + +@dataclass(frozen=True) +class ToolEventPolicy: + """ + Per-target configuration that controls how the tool loop responds + to a pending tool call from the model. + + Attributes: + behavior (ToolEventBehavior): What to do on each detected tool call. + max_tool_iterations (int): Maximum number of model<->tool round-trips + before the loop raises :class:`ToolCallLoopLimitExceeded`. Each + iteration is one ``_send_prompt_to_target_async`` call. + """ + + behavior: ToolEventBehavior + max_tool_iterations: int = 5 + + +def _build_function_call_output_message( + *, + reference_piece: MessagePiece, + outputs: list[tuple[ToolCall, Any]], +) -> Message: + """ + Build the canonical ``tool`` message produced after dispatching one or more + tool calls in a single iteration. + + The returned :class:`Message` contains one + :class:`MessagePiece` per ``(call, result)`` pair, in declaration order. + Every piece has ``role="tool"`` and ``original_value_data_type="function_call_output"``, + with the JSON envelope ``{"type": "function_call_output", "call_id": ..., "output": ...}``. + + Lineage metadata (conversation_id, identifiers) is copied from + *reference_piece* — typically the first piece of the assistant message + that issued the tool calls — so the resulting message stays inside the + correct conversation. + + Args: + reference_piece (MessagePiece): Piece whose lineage metadata is + copied onto every output piece. Pass the first piece of the + assistant message that produced the calls. + outputs (list[tuple[ToolCall, Any]]): ``(call, result)`` pairs in + declaration order. *result* is serialized via :func:`json.dumps` + unless it is already a string. + + Returns: + Message: One message carrying every function_call_output piece. + """ + pieces: list[MessagePiece] = [] + for call, result in outputs: + output_str = result if isinstance(result, str) else json.dumps(result, separators=(",", ":")) + envelope = json.dumps( + {"type": "function_call_output", "call_id": call.call_id, "output": output_str}, + separators=(",", ":"), + ) + pieces.append( + MessagePiece( + role="tool", + original_value=envelope, + original_value_data_type="function_call_output", + conversation_id=reference_piece.conversation_id, + prompt_target_identifier=reference_piece.prompt_target_identifier, + attack_identifier=reference_piece.attack_identifier, + ) + ) + return Message(message_pieces=pieces, skip_validation=True) + + +def tool_loop( + method: Callable[..., Awaitable[list[Message]]], +) -> Callable[..., Awaitable[list[Message]]]: + """ + Wrap a :class:`~pyrit.prompt_target.PromptTarget`-style + ``send_prompt_async`` to run an agentic tool-use loop. + + When the target's ``configuration.tool_event_policy`` is ``None`` the + wrapper is a no-op — the wrapped method runs unchanged. When a policy + is configured, the wrapper replaces the method body with the loop: + + 1. Validate and normalize the incoming message exactly once. + 2. Repeatedly call ``self._send_prompt_to_target_async`` with the + growing conversation. + 3. After each call, parse the last response via ``self._tool_parser``. + Exit on empty parse (model issued a stop response). + 4. On a non-empty parse, branch on ``policy.behavior``: + ``RAISE`` raises :class:`ToolCallNotSupported`; ``RETURN_RAW`` + returns the chain as-is; ``EXECUTE`` dispatches the calls via + ``configuration.tool_backend`` and appends the tool message. + 5. Raise :class:`ToolCallLoopLimitExceeded` if the loop runs past + ``policy.max_tool_iterations`` without the model stopping. + + The decorator deliberately knows nothing about MCP, OpenAI, or any + specific transport. The two collaborators it requires — + ``self._tool_parser`` and ``self.configuration.tool_backend`` — are + plain protocols (:class:`ToolCallParser`, :class:`ToolBackend`). + + Args: + method (Callable): The async method to wrap. Must have the + ``async def f(self, *, message: Message) -> list[Message]`` + signature of :meth:`PromptTarget.send_prompt_async`. + + Returns: + Callable: The wrapped method. + """ + + @functools.wraps(method) + async def wrapper(self: Any, *, message: Message) -> list[Message]: + policy: ToolEventPolicy | None = getattr(self.configuration, "tool_event_policy", None) + if policy is None: + return await method(self, message=message) + + message.validate() + normalized_conversation = await self._get_normalized_conversation_async(message=message) + if not normalized_conversation: + raise ValueError("Normalization pipeline returned an empty conversation. Cannot send an empty request.") + self._validate_request(normalized_conversation=normalized_conversation) + + parser: ToolCallParser | None = getattr(self, "_tool_parser", None) + backend: ToolBackend | None = getattr(self.configuration, "tool_backend", None) + max_iter = policy.max_tool_iterations + + all_responses: list[Message] = [] + + for _ in range(max_iter): + responses_this_turn = await self._send_prompt_to_target_async( + normalized_conversation=normalized_conversation, + ) + all_responses.extend(responses_this_turn) + + if parser is None: + return all_responses + + last_response = responses_this_turn[-1] + pending_calls = parser.parse(last_response) + + if not pending_calls: + return all_responses + + if policy.behavior is ToolEventBehavior.RAISE: + raise ToolCallNotSupported( + message=( + f"Target produced {len(pending_calls)} tool call(s) but ToolEventPolicy.behavior is RAISE." + ), + partial_conversation=all_responses, + ) + + if policy.behavior is ToolEventBehavior.RETURN_RAW: + return all_responses + + if backend is None: + raise ToolCallNotSupported( + message=(f"Target produced {len(pending_calls)} tool call(s) but no tool_backend is configured."), + partial_conversation=all_responses, + ) + + results = await backend.dispatch_all_sequential_async(pending_calls) + tool_msg = _build_function_call_output_message( + reference_piece=last_response.message_pieces[0], + outputs=results, + ) + all_responses.append(tool_msg) + normalized_conversation = list(normalized_conversation) + [last_response, tool_msg] + + raise ToolCallLoopLimitExceeded( + message=f"Tool loop exceeded max_tool_iterations={max_iter} without a stop response.", + partial_conversation=all_responses, + ) + + return wrapper diff --git a/pyrit/tools/parsers.py b/pyrit/tools/parsers.py new file mode 100644 index 0000000000..4ff7fc4c04 --- /dev/null +++ b/pyrit/tools/parsers.py @@ -0,0 +1,55 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT license. + +from __future__ import annotations + +from typing import TYPE_CHECKING, Protocol, runtime_checkable + +if TYPE_CHECKING: + from pyrit.models import Message, MessagePiece + from pyrit.tools.models import ToolCall + + +@runtime_checkable +class ToolCallParser(Protocol): + """ + Protocol for extracting tool calls from a target response message. + + Concrete parsers live next to the target whose response shape they + understand (see :class:`OpenAIChatTarget` and :class:`OpenAIResponseTarget` + after C7/C8). Parsers MUST return an empty list when the model has + issued a stop response — the tool loop uses the empty list as the + signal to exit. + """ + + def parse(self, message: Message) -> list[ToolCall]: + """ + Extract tool calls from a target response message. + + Args: + message (Message): The most recent assistant response. + + Returns: + list[ToolCall]: Tool calls, in declaration order. An empty list + signals that the model produced a stop response. + """ + ... + + +def _extract_function_call_pieces(message: Message) -> list[MessagePiece]: + """ + Return every :class:`MessagePiece` in *message* whose + ``original_value_data_type`` is ``"function_call"``. + + This is the canonical envelope produced by OpenAI-style targets after + the C6 normalization commit. It is exposed here so concrete parsers + can reuse the filter rather than re-implementing it. + + Args: + message (Message): The message to scan. + + Returns: + list[MessagePiece]: Pieces whose ``original_value_data_type`` is + ``"function_call"``, in their declaration order. + """ + return [piece for piece in message.message_pieces if piece.original_value_data_type == "function_call"] diff --git a/tests/unit/tools/__init__.py b/tests/unit/tools/__init__.py new file mode 100644 index 0000000000..9a0454564d --- /dev/null +++ b/tests/unit/tools/__init__.py @@ -0,0 +1,2 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT license. diff --git a/tests/unit/tools/conftest.py b/tests/unit/tools/conftest.py new file mode 100644 index 0000000000..9ac8ce8aa5 --- /dev/null +++ b/tests/unit/tools/conftest.py @@ -0,0 +1,307 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT license. + +""" +Shared fixtures for ``tests/unit/tools``. + +Provides the minimal collaborators the tool-loop tests need to exercise +:func:`pyrit.tools.tool_loop` end-to-end without standing up real targets +or MCP transports: + +* :class:`_FakeToolTarget` — a :class:`PromptTarget` subclass whose + ``_send_prompt_to_target_async`` returns scripted messages from a queue + and whose ``_get_normalized_conversation_async`` skips the memory round + trip so decorator behavior is isolated from normalization. +* :class:`_RecordingToolBackend` — a :class:`ToolBackend` that records + every dispatched call (for order-of-execution assertions) and returns + results from a scripted queue. +* :class:`_CanonicalEnvelopeParser` — a :class:`ToolCallParser` that walks + message pieces and parses the canonical ``function_call`` JSON envelope. + +Helper message builders (``_make_user_message``, +``_make_assistant_text_message``, ``_make_assistant_function_call_message``) +produce the canonical envelope shape used by the OpenAI targets after the +C6 normalization commit. +""" + +from __future__ import annotations + +import json +import uuid +from collections import deque +from typing import Any + +import pytest + +from pyrit.models import Message, MessagePiece +from pyrit.prompt_target.common.prompt_target import PromptTarget +from pyrit.prompt_target.common.target_capabilities import TargetCapabilities +from pyrit.prompt_target.common.target_configuration import TargetConfiguration +from pyrit.tools import ( + CallableToolBackend, + ToolCall, + ToolCallParser, + ToolEventBehavior, + ToolEventPolicy, + tool_loop, +) + + +def _make_user_message(text: str, *, conversation_id: str | None = None) -> Message: + """Build a single-piece user :class:`Message` carrying *text*.""" + return Message( + message_pieces=[ + MessagePiece( + role="user", + original_value=text, + original_value_data_type="text", + conversation_id=conversation_id or str(uuid.uuid4()), + ) + ] + ) + + +def _make_assistant_text_message(text: str, *, conversation_id: str | None = None) -> Message: + """Build a single-piece assistant :class:`Message` carrying plain text.""" + return Message( + message_pieces=[ + MessagePiece( + role="assistant", + original_value=text, + original_value_data_type="text", + conversation_id=conversation_id or str(uuid.uuid4()), + ) + ], + skip_validation=True, + ) + + +def _make_function_call_piece( + *, + call_id: str, + name: str, + arguments: dict[str, Any], + conversation_id: str | None = None, +) -> MessagePiece: + """Build one assistant ``function_call`` piece carrying the canonical envelope.""" + envelope = { + "type": "function_call", + "call_id": call_id, + "name": name, + "arguments": json.dumps(arguments, separators=(",", ":")), + } + return MessagePiece( + role="assistant", + original_value=json.dumps(envelope, separators=(",", ":")), + original_value_data_type="function_call", + conversation_id=conversation_id or str(uuid.uuid4()), + ) + + +def _make_assistant_function_call_message( + *, + calls: list[tuple[str, str, dict[str, Any]]], + conversation_id: str | None = None, +) -> Message: + """ + Build an assistant :class:`Message` carrying one ``function_call`` piece + per ``(call_id, name, args)`` tuple, in declaration order. + """ + conv_id = conversation_id or str(uuid.uuid4()) + pieces = [ + _make_function_call_piece(call_id=cid, name=name, arguments=args, conversation_id=conv_id) + for cid, name, args in calls + ] + return Message(message_pieces=pieces, skip_validation=True) + + +class _CanonicalEnvelopeParser: + """ + Reference :class:`ToolCallParser` that understands the canonical envelope + (``original_value_data_type == "function_call"`` carrying a JSON object + with ``type``/``call_id``/``name``/``arguments``). + + Per-target parsers shipped in C7/C8 will reuse this shape; this stand-in + keeps decorator tests independent of the real OpenAI parsers. + """ + + def parse(self, message: Message) -> list[ToolCall]: + calls: list[ToolCall] = [] + for piece in message.message_pieces: + if piece.original_value_data_type != "function_call": + continue + envelope = json.loads(piece.original_value) + arguments_str = envelope.get("arguments", "{}") + arguments = json.loads(arguments_str) if isinstance(arguments_str, str) else dict(arguments_str) + calls.append( + ToolCall( + call_id=envelope["call_id"], + name=envelope["name"], + arguments=arguments, + raw_envelope=envelope, + ) + ) + return calls + + +class _RecordingToolBackend: + """ + Minimal :class:`ToolBackend` that records every dispatched call and + returns results from a scripted queue. Used to assert dispatch order, + iteration count, and per-call payload shape without invoking real tools. + """ + + def __init__( + self, + *, + scripted_results: list[Any] | None = None, + schemas: list[dict[str, Any]] | None = None, + ) -> None: + self._results: deque[Any] = deque(scripted_results or []) + self._schemas: list[dict[str, Any]] = list(schemas) if schemas is not None else [] + self.recorded_calls: list[ToolCall] = [] + + @property + def schemas(self) -> list[dict[str, Any]]: + return list(self._schemas) + + async def dispatch_async(self, call: ToolCall) -> dict[str, Any]: + self.recorded_calls.append(call) + if not self._results: + return {"result": f"recorded:{call.name}:{call.call_id}"} + nxt = self._results.popleft() + return nxt if isinstance(nxt, dict) else {"result": nxt} + + async def dispatch_all_sequential_async( + self, + calls: list[ToolCall], + ) -> list[tuple[ToolCall, dict[str, Any]]]: + results: list[tuple[ToolCall, dict[str, Any]]] = [] + for call in calls: + result = await self.dispatch_async(call) + results.append((call, result)) + return results + + +class _FakeToolTarget(PromptTarget): + """ + Test-only :class:`PromptTarget` whose ``_send_prompt_to_target_async`` + pops scripted responses off a queue. ``_get_normalized_conversation_async`` + is overridden to return ``[message]`` directly, isolating decorator + behavior from the memory + normalization pipeline. + """ + + _DEFAULT_CONFIGURATION = TargetConfiguration( + capabilities=TargetCapabilities( + supports_multi_turn=True, + supports_multi_message_pieces=True, + ) + ) + + def __init__( + self, + *, + scripted_responses: list[Message], + policy: ToolEventPolicy | None = None, + backend: Any = None, + parser: ToolCallParser | None = None, + ) -> None: + super().__init__() + self._scripted_responses: deque[Message] = deque(scripted_responses) + self.call_count: int = 0 + self.normalized_conversations_seen: list[list[Message]] = [] + # The C2 decorator reads these via getattr; production code wires them + # through TargetConfiguration in C4. + self._configuration.tool_event_policy = policy + self._configuration.tool_backend = backend + self._tool_parser = parser if parser is not None else _CanonicalEnvelopeParser() + + async def _get_normalized_conversation_async(self, *, message: Message) -> list[Message]: + return [message] + + def _validate_request(self, *, normalized_conversation: list[Message]) -> None: + return + + async def _send_prompt_to_target_async(self, *, normalized_conversation: list[Message]) -> list[Message]: + self.call_count += 1 + self.normalized_conversations_seen.append(list(normalized_conversation)) + if not self._scripted_responses: + raise AssertionError(f"Fake target ran out of scripted responses on iteration {self.call_count}.") + return [self._scripted_responses.popleft()] + + @tool_loop + async def send_prompt_async(self, *, message: Message) -> list[Message]: + # Passthrough path: only invoked when ToolEventPolicy is None. The + # decorator replaces this body entirely when a policy is set. + message.validate() + normalized = await self._get_normalized_conversation_async(message=message) + return await self._send_prompt_to_target_async(normalized_conversation=normalized) + + +@pytest.fixture +def make_fake_target(patch_central_database): + """ + Factory fixture for :class:`_FakeToolTarget`. Each invocation returns a + fresh target instance whose scripted response queue is independent of + other targets created during the test. + """ + + def _factory( + *, + scripted_responses: list[Message], + policy: ToolEventPolicy | None = None, + backend: Any = None, + parser: ToolCallParser | None = None, + ) -> _FakeToolTarget: + return _FakeToolTarget( + scripted_responses=scripted_responses, + policy=policy, + backend=backend, + parser=parser, + ) + + return _factory + + +@pytest.fixture +def recording_backend(): + """Factory fixture for :class:`_RecordingToolBackend`.""" + + def _factory(*, scripted_results: list[Any] | None = None) -> _RecordingToolBackend: + return _RecordingToolBackend(scripted_results=scripted_results) + + return _factory + + +@pytest.fixture +def execute_policy(): + """ + Factory fixture for :class:`ToolEventPolicy` with + ``behavior=ToolEventBehavior.EXECUTE`` and a tunable iteration cap. + """ + + def _factory(*, max_tool_iterations: int = 5) -> ToolEventPolicy: + return ToolEventPolicy( + behavior=ToolEventBehavior.EXECUTE, + max_tool_iterations=max_tool_iterations, + ) + + return _factory + + +__all__ = [ + "CallableToolBackend", + "ToolCall", + "ToolEventBehavior", + "ToolEventPolicy", + "_CanonicalEnvelopeParser", + "_FakeToolTarget", + "_RecordingToolBackend", + "_make_assistant_function_call_message", + "_make_assistant_text_message", + "_make_function_call_piece", + "_make_user_message", + "execute_policy", + "make_fake_target", + "recording_backend", +] diff --git a/tests/unit/tools/echo_mcp_server.py b/tests/unit/tools/echo_mcp_server.py new file mode 100644 index 0000000000..723a3c6594 --- /dev/null +++ b/tests/unit/tools/echo_mcp_server.py @@ -0,0 +1,57 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT license. + +""" +Deterministic echo MCP server used as a stdio subprocess fixture by +``tests/unit/tools/test_mcp_client.py`` (C3) and the integration tests +(C9). + +Lands in C2 so subsequent commits don't shuffle test plumbing; C2's own +tests do not import this module (the :class:`CallableToolRegistry` is +exercised in-process). + +Run directly as ``python echo_mcp_server.py`` to expose the four tools +over stdio. The MCP client harness in C3 launches this file with +``mcp.client.stdio.stdio_client`` and asserts behavior end to end. +""" + +from __future__ import annotations + +import asyncio + +from mcp.server.fastmcp import FastMCP + +mcp = FastMCP("pyrit-echo") + + +@mcp.tool() +def echo(text: str) -> str: + """Return *text* unchanged.""" + return text + + +@mcp.tool() +def add(a: int, b: int) -> int: + """Return ``a + b``.""" + return a + b + + +@mcp.tool() +def reverse(text: str) -> str: + """Return *text* reversed.""" + return text[::-1] + + +@mcp.tool() +async def slow_echo(text: str, delay_ms: int = 0) -> str: + """ + Return *text* after sleeping ``delay_ms`` milliseconds. Used by + timeout / cancellation tests. + """ + if delay_ms > 0: + await asyncio.sleep(delay_ms / 1000.0) + return text + + +if __name__ == "__main__": + mcp.run() diff --git a/tests/unit/tools/test_callable_tool_backend.py b/tests/unit/tools/test_callable_tool_backend.py new file mode 100644 index 0000000000..fb09788ebc --- /dev/null +++ b/tests/unit/tools/test_callable_tool_backend.py @@ -0,0 +1,180 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT license. + +""" +Unit tests for :class:`pyrit.tools.CallableToolBackend`. + +Coverage map (rows from the C2 test matrix): + +* **U10** (partial; the MCP counterpart lands in C3) — + ``test_each_dummy_tool_invoked_via_prepended_conversation`` +* **U17** (partial; the MCP-timeout counterpart lands in C3) — + ``test_failing_tool_yields_error_envelope`` +* **U18** — ``test_disallowed_tool_returns_error_without_invoking_callable`` + +Also covers the backend's documented behavior for missing functions +(both strict and tolerant modes), schema property defaulting, scalar +result wrapping, and declaration-order preservation in the bulk dispatch +path. These are required for the §10 rubber-duck guarantee that every +public-facing branch of :class:`CallableToolBackend` is exercised +before C5 wires it to a production target. +""" + +from __future__ import annotations + +import pytest + +from pyrit.tools import CallableToolBackend, ToolCall + + +def _make_call(name: str, *, call_id: str = "c1", arguments: dict | None = None) -> ToolCall: + return ToolCall(call_id=call_id, name=name, arguments=arguments or {}) + + +async def test_disallowed_tool_returns_error_without_invoking_callable(): + invoked: list[str] = [] + + async def echo(args: dict) -> dict: + invoked.append(args.get("text", "")) + return {"echoed": args.get("text", "")} + + backend = CallableToolBackend( + callables={"echo": echo, "off_limits": echo}, + allowed_tools={"echo"}, + ) + + result = await backend.dispatch_async(_make_call("off_limits", arguments={"text": "nope"})) + + assert result["error"] == "tool_not_allowed" + assert result["tool"] == "off_limits" + assert "echo" in result["allowed_tools"] + assert invoked == [] # callable was never invoked + + +async def test_failing_tool_yields_error_envelope(): + async def boom(args: dict) -> dict: + raise RuntimeError("kaboom") + + backend = CallableToolBackend(callables={"boom": boom}) + + result = await backend.dispatch_async(_make_call("boom")) + + assert result["error"] == "tool_execution_failed" + assert result["tool"] == "boom" + assert "kaboom" in result["detail"] + + +async def test_missing_tool_raises_when_strict(): + backend = CallableToolBackend(callables={}, fail_on_missing_function=True) + + with pytest.raises(KeyError, match="ghost"): + await backend.dispatch_async(_make_call("ghost")) + + +async def test_missing_tool_returns_envelope_when_tolerant(): + async def echo(args: dict) -> dict: + return {"ok": True} + + backend = CallableToolBackend( + callables={"echo": echo}, + fail_on_missing_function=False, + ) + + result = await backend.dispatch_async(_make_call("ghost")) + + assert result["error"] == "tool_not_registered" + assert result["tool"] == "ghost" + assert result["available_tools"] == ["echo"] + + +async def test_scalar_result_is_wrapped_in_dict(): + async def number(args: dict) -> int: + return 42 + + backend = CallableToolBackend(callables={"number": number}) + + result = await backend.dispatch_async(_make_call("number")) + + assert result == {"result": 42} + + +async def test_dict_result_passes_through_unchanged(): + async def named(args: dict) -> dict: + return {"custom_key": "custom_value"} + + backend = CallableToolBackend(callables={"named": named}) + + result = await backend.dispatch_async(_make_call("named")) + + assert result == {"custom_key": "custom_value"} + + +async def test_schemas_defaults_to_empty_list(): + backend = CallableToolBackend(callables={}) + + assert backend.schemas == [] + + +async def test_schemas_returned_as_copy(): + schemas_in = [{"name": "echo", "parameters": {}}] + backend = CallableToolBackend(callables={}, schemas=schemas_in) + + out1 = backend.schemas + out1.append({"name": "mutated"}) + + # Mutating the returned list does not affect the backend's internal state. + assert backend.schemas == schemas_in + + +async def test_dispatch_all_sequential_preserves_declaration_order(): + async def echo(args: dict) -> dict: + return {"echoed": args["i"]} + + backend = CallableToolBackend(callables={"echo": echo}) + + calls = [_make_call("echo", call_id=f"c{i}", arguments={"i": i}) for i in range(5)] + pairs = await backend.dispatch_all_sequential_async(calls) + + assert [c.call_id for c, _ in pairs] == ["c0", "c1", "c2", "c3", "c4"] + assert [r["echoed"] for _, r in pairs] == [0, 1, 2, 3, 4] + + +async def test_each_dummy_tool_invoked_via_prepended_conversation(): + """ + U10 (partial). Each dummy tool resolves on first dispatch (single + forward step, no model reasoning trace), confirming the backend can + short-circuit a prepended conversation where every call is already + decided. The MCP counterpart in C3 exercises the same shape against + a real stdio server. + """ + invocations: list[tuple[str, dict]] = [] + + async def echo(args: dict) -> dict: + invocations.append(("echo", args)) + return {"echoed": args.get("text", "")} + + async def add(args: dict) -> dict: + invocations.append(("add", args)) + return {"sum": args["a"] + args["b"]} + + async def reverse(args: dict) -> dict: + invocations.append(("reverse", args)) + return {"reversed": args.get("text", "")[::-1]} + + backend = CallableToolBackend(callables={"echo": echo, "add": add, "reverse": reverse}) + + prepended_calls = [ + _make_call("echo", call_id="e1", arguments={"text": "hello"}), + _make_call("add", call_id="a1", arguments={"a": 2, "b": 3}), + _make_call("reverse", call_id="r1", arguments={"text": "pyrit"}), + ] + pairs = await backend.dispatch_all_sequential_async(prepended_calls) + + # Each dummy resolved exactly once; no retries, no model re-entry. + assert len(invocations) == 3 + assert [name for name, _ in invocations] == ["echo", "add", "reverse"] + assert [r for _, r in pairs] == [ + {"echoed": "hello"}, + {"sum": 5}, + {"reversed": "tiryp"}, + ] diff --git a/tests/unit/tools/test_tool_loop_decorator.py b/tests/unit/tools/test_tool_loop_decorator.py new file mode 100644 index 0000000000..bc0db6b357 --- /dev/null +++ b/tests/unit/tools/test_tool_loop_decorator.py @@ -0,0 +1,289 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT license. + +""" +Unit tests for :func:`pyrit.tools.tool_loop`. + +Coverage map (rows from the C2 test matrix): + +* **U2** (partial; full-DB end lands in C5) — ``test_loop_returns_full_chain_in_order`` +* **U3** — ``test_loop_exits_on_first_response_when_no_tool_calls``, + ``test_loops_until_no_pending_tool_call`` +* **U4** — ``test_raises_after_max_tool_iterations``, + ``test_partial_conversation_attached_to_limit_exception`` +* **U12** — ``test_policy_raise_includes_partial_conversation`` +* **U13** — ``test_policy_return_raw_does_not_dispatch`` +* **U16** — ``test_multi_call_per_turn_dispatched_sequentially_in_order`` + +Also covers two additional decorator concerns required by the rubber-duck +review (§10): EXECUTE policy with no backend raises with a partial +conversation, and the normalized conversation grows correctly across +iterations (decorator does not re-normalize each turn). +""" + +from __future__ import annotations + +import json + +import pytest + +from pyrit.exceptions import ToolCallLoopLimitExceeded, ToolCallNotSupported +from pyrit.tools import ToolEventBehavior, ToolEventPolicy + +from .conftest import ( + _make_assistant_function_call_message, + _make_assistant_text_message, + _make_user_message, +) + + +@pytest.mark.usefixtures("patch_central_database") +class TestToolLoopDecoratorBasics: + """Loop entry/exit semantics: no tool calls, single round trip, multi-round.""" + + async def test_loop_exits_on_first_response_when_no_tool_calls(self, make_fake_target, execute_policy): + target = make_fake_target( + scripted_responses=[_make_assistant_text_message("done")], + policy=execute_policy(), + ) + + responses = await target.send_prompt_async(message=_make_user_message("hi")) + + assert len(responses) == 1 + assert responses[0].get_value() == "done" + assert target.call_count == 1 + + async def test_loops_until_no_pending_tool_call(self, make_fake_target, execute_policy, recording_backend): + backend = recording_backend(scripted_results=[{"ok": True}, {"ok": True}]) + target = make_fake_target( + scripted_responses=[ + _make_assistant_function_call_message(calls=[("c1", "tool_a", {"x": 1})]), + _make_assistant_function_call_message(calls=[("c2", "tool_a", {"x": 2})]), + _make_assistant_text_message("done"), + ], + policy=execute_policy(max_tool_iterations=5), + backend=backend, + ) + + responses = await target.send_prompt_async(message=_make_user_message("hi")) + + # Two model-tool round trips and one final assistant message. + assert target.call_count == 3 + # Returned chain: fc1, tool1, fc2, tool2, final-text → 5 messages total. + assert len(responses) == 5 + assert [r.message_pieces[0].original_value_data_type for r in responses] == [ + "function_call", + "function_call_output", + "function_call", + "function_call_output", + "text", + ] + assert len(backend.recorded_calls) == 2 + assert [c.call_id for c in backend.recorded_calls] == ["c1", "c2"] + + +@pytest.mark.usefixtures("patch_central_database") +class TestToolLoopMessageShape: + """U2 — assistant_fc → tool → final_assistant ordering and identity.""" + + async def test_loop_returns_full_chain_in_order(self, make_fake_target, execute_policy, recording_backend): + backend = recording_backend(scripted_results=[{"weather": "sunny"}]) + fc_msg = _make_assistant_function_call_message(calls=[("call_abc", "get_weather", {"city": "Seattle"})]) + final_msg = _make_assistant_text_message("It is sunny in Seattle.") + + target = make_fake_target( + scripted_responses=[fc_msg, final_msg], + policy=execute_policy(), + backend=backend, + ) + + responses = await target.send_prompt_async(message=_make_user_message("weather?")) + + assert len(responses) == 3 + # 1) assistant with function_call (identity preserved) + assert responses[0] is fc_msg + # 2) tool message with exactly one function_call_output piece carrying call_id + tool_msg = responses[1] + assert len(tool_msg.message_pieces) == 1 + tool_piece = tool_msg.message_pieces[0] + assert tool_piece.api_role == "tool" + assert tool_piece.original_value_data_type == "function_call_output" + envelope = json.loads(tool_piece.original_value) + assert envelope["type"] == "function_call_output" + assert envelope["call_id"] == "call_abc" + # The tool result is JSON-serialized into the "output" field. + assert json.loads(envelope["output"]) == {"weather": "sunny"} + # 3) final assistant text (identity preserved) + assert responses[2] is final_msg + + +@pytest.mark.usefixtures("patch_central_database") +class TestToolLoopIterationLimits: + """U4 — iteration cap raises and carries the partial chain.""" + + async def test_raises_after_max_tool_iterations(self, make_fake_target, execute_policy, recording_backend): + # Model never stops asking for tools. + backend = recording_backend(scripted_results=[{"ok": True}] * 3) + target = make_fake_target( + scripted_responses=[ + _make_assistant_function_call_message(calls=[(f"c{i}", "loop_tool", {})]) for i in range(3) + ], + policy=execute_policy(max_tool_iterations=2), + backend=backend, + ) + + with pytest.raises(ToolCallLoopLimitExceeded, match="max_tool_iterations=2"): + await target.send_prompt_async(message=_make_user_message("hi")) + + # Exactly max_tool_iterations model calls made before raising. + assert target.call_count == 2 + + async def test_partial_conversation_attached_to_limit_exception( + self, make_fake_target, execute_policy, recording_backend + ): + backend = recording_backend(scripted_results=[{"ok": True}] * 2) + target = make_fake_target( + scripted_responses=[ + _make_assistant_function_call_message(calls=[(f"c{i}", "loop_tool", {})]) for i in range(2) + ], + policy=execute_policy(max_tool_iterations=2), + backend=backend, + ) + + with pytest.raises(ToolCallLoopLimitExceeded) as excinfo: + await target.send_prompt_async(message=_make_user_message("hi")) + + partial = excinfo.value.partial_conversation + # 2 iterations × (assistant_fc + tool_msg) = 4 messages, all in order. + assert len(partial) == 4 + assert [m.message_pieces[0].original_value_data_type for m in partial] == [ + "function_call", + "function_call_output", + "function_call", + "function_call_output", + ] + + +@pytest.mark.usefixtures("patch_central_database") +class TestToolEventPolicyBehaviors: + """U12, U13 — non-EXECUTE behaviors short-circuit dispatch.""" + + async def test_policy_raise_includes_partial_conversation(self, make_fake_target, recording_backend): + backend = recording_backend(scripted_results=[{"ok": True}]) + fc_msg = _make_assistant_function_call_message(calls=[("c1", "danger", {})]) + target = make_fake_target( + scripted_responses=[fc_msg], + policy=ToolEventPolicy(behavior=ToolEventBehavior.RAISE), + backend=backend, + ) + + with pytest.raises(ToolCallNotSupported, match="RAISE") as excinfo: + await target.send_prompt_async(message=_make_user_message("hi")) + + partial = excinfo.value.partial_conversation + # Partial contains the offending assistant turn; no tool dispatch occurred. + assert partial == [fc_msg] + assert backend.recorded_calls == [] + assert target.call_count == 1 + + async def test_policy_return_raw_does_not_dispatch(self, make_fake_target, recording_backend): + backend = recording_backend(scripted_results=[{"ok": True}]) + fc_msg = _make_assistant_function_call_message(calls=[("c1", "danger", {})]) + target = make_fake_target( + scripted_responses=[fc_msg], + policy=ToolEventPolicy(behavior=ToolEventBehavior.RETURN_RAW), + backend=backend, + ) + + responses = await target.send_prompt_async(message=_make_user_message("hi")) + + assert responses == [fc_msg] + assert backend.recorded_calls == [] + assert target.call_count == 1 + + +@pytest.mark.usefixtures("patch_central_database") +class TestToolLoopMultiCallPerTurn: + """U16 — multi-call turns dispatch sequentially in declaration order.""" + + async def test_multi_call_per_turn_dispatched_sequentially_in_order( + self, make_fake_target, execute_policy, recording_backend + ): + backend = recording_backend(scripted_results=[{"a": 1}, {"b": 2}, {"c": 3}]) + multi_fc = _make_assistant_function_call_message( + calls=[ + ("c_alpha", "tool_alpha", {"k": "v1"}), + ("c_beta", "tool_beta", {"k": "v2"}), + ("c_gamma", "tool_gamma", {"k": "v3"}), + ] + ) + target = make_fake_target( + scripted_responses=[multi_fc, _make_assistant_text_message("ok")], + policy=execute_policy(), + backend=backend, + ) + + responses = await target.send_prompt_async(message=_make_user_message("multi")) + + # Three calls dispatched in declaration order, recorded ids match. + assert [c.call_id for c in backend.recorded_calls] == ["c_alpha", "c_beta", "c_gamma"] + assert [c.name for c in backend.recorded_calls] == ["tool_alpha", "tool_beta", "tool_gamma"] + # One tool message after the multi-call assistant turn, carrying three + # function_call_output pieces in declaration order with the right call_ids. + tool_msg = responses[1] + assert len(tool_msg.message_pieces) == 3 + envelopes = [json.loads(p.original_value) for p in tool_msg.message_pieces] + assert [e["call_id"] for e in envelopes] == ["c_alpha", "c_beta", "c_gamma"] + assert all(p.original_value_data_type == "function_call_output" for p in tool_msg.message_pieces) + + +@pytest.mark.usefixtures("patch_central_database") +class TestToolLoopMisconfiguration: + """EXECUTE policy with no backend must fail loudly and carry the partial chain.""" + + async def test_execute_without_backend_raises_with_partial(self, make_fake_target, execute_policy): + fc_msg = _make_assistant_function_call_message(calls=[("c1", "no_reg", {})]) + target = make_fake_target( + scripted_responses=[fc_msg], + policy=execute_policy(), + backend=None, + ) + + with pytest.raises(ToolCallNotSupported, match="tool_backend") as excinfo: + await target.send_prompt_async(message=_make_user_message("hi")) + + assert excinfo.value.partial_conversation == [fc_msg] + + +@pytest.mark.usefixtures("patch_central_database") +class TestToolLoopConversationGrowth: + """The decorator must extend (not re-normalize) the conversation each round.""" + + async def test_normalized_conversation_grows_each_iteration( + self, make_fake_target, execute_policy, recording_backend + ): + backend = recording_backend(scripted_results=[{"r1": 1}, {"r2": 2}]) + target = make_fake_target( + scripted_responses=[ + _make_assistant_function_call_message(calls=[("c1", "t", {})]), + _make_assistant_function_call_message(calls=[("c2", "t", {})]), + _make_assistant_text_message("done"), + ], + policy=execute_policy(), + backend=backend, + ) + + await target.send_prompt_async(message=_make_user_message("hi")) + + # Three protected-method calls; each subsequent call sees the prior + # assistant_fc + tool_msg appended (the decorator must NOT re-normalize). + seen = target.normalized_conversations_seen + assert len(seen) == 3 + # call 1: just the user message + assert len(seen[0]) == 1 + # call 2: user + assistant_fc(c1) + tool_msg + assert len(seen[1]) == 3 + assert seen[1][1].message_pieces[0].original_value_data_type == "function_call" + assert seen[1][2].message_pieces[0].original_value_data_type == "function_call_output" + # call 3: user + assistant_fc(c1) + tool_msg + assistant_fc(c2) + tool_msg + assert len(seen[2]) == 5 From 39752ac88bab2e6c1ad27ab1b660df8b6771c78b Mon Sep 17 00:00:00 2001 From: Victor Valbuena Date: Tue, 26 May 2026 15:24:49 -0700 Subject: [PATCH 23/40] Addition of pyrit/tools package and unit tests. Introduces base models for tool calling. --- pyrit/tools/__init__.py | 6 ++--- pyrit/tools/backend.py | 6 ++--- .../{callable_backend.py => local_backend.py} | 9 +++++-- tests/unit/tools/conftest.py | 4 +-- ..._backend.py => test_local_tool_backend.py} | 26 +++++++++---------- 5 files changed, 28 insertions(+), 23 deletions(-) rename pyrit/tools/{callable_backend.py => local_backend.py} (93%) rename tests/unit/tools/{test_callable_tool_backend.py => test_local_tool_backend.py} (87%) diff --git a/pyrit/tools/__init__.py b/pyrit/tools/__init__.py index 9830f17758..c3520098ea 100644 --- a/pyrit/tools/__init__.py +++ b/pyrit/tools/__init__.py @@ -20,7 +20,7 @@ When the policy is ``EXECUTE``, calls are dispatched through ``self.configuration.tool_backend``, an implementation of -:class:`ToolBackend`. :class:`CallableToolBackend` is the pure-Python +:class:`ToolBackend`. :class:`LocalToolBackend` is the in-process backend shipped here; :class:`MCPToolBackend` ships in C3 and proxies through one or more MCP servers. @@ -43,12 +43,12 @@ """ from pyrit.tools.backend import ToolBackend -from pyrit.tools.callable_backend import CallableToolBackend +from pyrit.tools.local_backend import LocalToolBackend from pyrit.tools.models import ToolCall, ToolEventBehavior, ToolEventPolicy, tool_loop from pyrit.tools.parsers import ToolCallParser __all__ = [ - "CallableToolBackend", + "LocalToolBackend", "ToolBackend", "ToolCall", "ToolCallParser", diff --git a/pyrit/tools/backend.py b/pyrit/tools/backend.py index 3274355d6e..54c1dd7a1d 100644 --- a/pyrit/tools/backend.py +++ b/pyrit/tools/backend.py @@ -22,9 +22,9 @@ class ToolBackend(Protocol): Two concrete implementations ship with PyRIT: - * :class:`~pyrit.tools.CallableToolBackend` — pure-Python backend - backed by ``async def`` callables. Useful for unit tests and for - embedding tools inside the PyRIT process. + * :class:`~pyrit.tools.LocalToolBackend` — in-process backend backed + by ``async def`` callables. Useful for unit tests and for embedding + tools inside the PyRIT process. * :class:`pyrit.tools.MCPToolBackend` (lands in C3) — proxies dispatch through one or more MCP servers. diff --git a/pyrit/tools/callable_backend.py b/pyrit/tools/local_backend.py similarity index 93% rename from pyrit/tools/callable_backend.py rename to pyrit/tools/local_backend.py index defbe94ed1..0c67054590 100644 --- a/pyrit/tools/callable_backend.py +++ b/pyrit/tools/local_backend.py @@ -14,12 +14,17 @@ logger = logging.getLogger(__name__) -class CallableToolBackend: +class LocalToolBackend: """ - Pure-Python :class:`~pyrit.tools.ToolBackend` backed by a name -> ``async def`` + In-process :class:`~pyrit.tools.ToolBackend` backed by a name -> ``async def`` mapping. Useful for unit tests and for embedding small tools inside the PyRIT process without standing up an MCP server. + "Local" here means tools run in PyRIT's own Python process — no + subprocess, no IPC, no wire protocol. Contrast with + :class:`~pyrit.tools.MCPToolBackend` (lands in C3), which proxies + dispatch through one or more MCP servers reached via JSON-RPC. + The backend dispatches sequentially in declaration order. Tool-side failures (raised exceptions, missing names, allow-list rejections) are converted into structured error envelopes so the tool loop can diff --git a/tests/unit/tools/conftest.py b/tests/unit/tools/conftest.py index 9ac8ce8aa5..8419c5b06c 100644 --- a/tests/unit/tools/conftest.py +++ b/tests/unit/tools/conftest.py @@ -38,7 +38,7 @@ from pyrit.prompt_target.common.target_capabilities import TargetCapabilities from pyrit.prompt_target.common.target_configuration import TargetConfiguration from pyrit.tools import ( - CallableToolBackend, + LocalToolBackend, ToolCall, ToolCallParser, ToolEventBehavior, @@ -290,7 +290,7 @@ def _factory(*, max_tool_iterations: int = 5) -> ToolEventPolicy: __all__ = [ - "CallableToolBackend", + "LocalToolBackend", "ToolCall", "ToolEventBehavior", "ToolEventPolicy", diff --git a/tests/unit/tools/test_callable_tool_backend.py b/tests/unit/tools/test_local_tool_backend.py similarity index 87% rename from tests/unit/tools/test_callable_tool_backend.py rename to tests/unit/tools/test_local_tool_backend.py index fb09788ebc..8e24693140 100644 --- a/tests/unit/tools/test_callable_tool_backend.py +++ b/tests/unit/tools/test_local_tool_backend.py @@ -2,7 +2,7 @@ # Licensed under the MIT license. """ -Unit tests for :class:`pyrit.tools.CallableToolBackend`. +Unit tests for :class:`pyrit.tools.LocalToolBackend`. Coverage map (rows from the C2 test matrix): @@ -16,7 +16,7 @@ (both strict and tolerant modes), schema property defaulting, scalar result wrapping, and declaration-order preservation in the bulk dispatch path. These are required for the §10 rubber-duck guarantee that every -public-facing branch of :class:`CallableToolBackend` is exercised +public-facing branch of :class:`LocalToolBackend` is exercised before C5 wires it to a production target. """ @@ -24,7 +24,7 @@ import pytest -from pyrit.tools import CallableToolBackend, ToolCall +from pyrit.tools import LocalToolBackend, ToolCall def _make_call(name: str, *, call_id: str = "c1", arguments: dict | None = None) -> ToolCall: @@ -38,7 +38,7 @@ async def echo(args: dict) -> dict: invoked.append(args.get("text", "")) return {"echoed": args.get("text", "")} - backend = CallableToolBackend( + backend = LocalToolBackend( callables={"echo": echo, "off_limits": echo}, allowed_tools={"echo"}, ) @@ -55,7 +55,7 @@ async def test_failing_tool_yields_error_envelope(): async def boom(args: dict) -> dict: raise RuntimeError("kaboom") - backend = CallableToolBackend(callables={"boom": boom}) + backend = LocalToolBackend(callables={"boom": boom}) result = await backend.dispatch_async(_make_call("boom")) @@ -65,7 +65,7 @@ async def boom(args: dict) -> dict: async def test_missing_tool_raises_when_strict(): - backend = CallableToolBackend(callables={}, fail_on_missing_function=True) + backend = LocalToolBackend(callables={}, fail_on_missing_function=True) with pytest.raises(KeyError, match="ghost"): await backend.dispatch_async(_make_call("ghost")) @@ -75,7 +75,7 @@ async def test_missing_tool_returns_envelope_when_tolerant(): async def echo(args: dict) -> dict: return {"ok": True} - backend = CallableToolBackend( + backend = LocalToolBackend( callables={"echo": echo}, fail_on_missing_function=False, ) @@ -91,7 +91,7 @@ async def test_scalar_result_is_wrapped_in_dict(): async def number(args: dict) -> int: return 42 - backend = CallableToolBackend(callables={"number": number}) + backend = LocalToolBackend(callables={"number": number}) result = await backend.dispatch_async(_make_call("number")) @@ -102,7 +102,7 @@ async def test_dict_result_passes_through_unchanged(): async def named(args: dict) -> dict: return {"custom_key": "custom_value"} - backend = CallableToolBackend(callables={"named": named}) + backend = LocalToolBackend(callables={"named": named}) result = await backend.dispatch_async(_make_call("named")) @@ -110,14 +110,14 @@ async def named(args: dict) -> dict: async def test_schemas_defaults_to_empty_list(): - backend = CallableToolBackend(callables={}) + backend = LocalToolBackend(callables={}) assert backend.schemas == [] async def test_schemas_returned_as_copy(): schemas_in = [{"name": "echo", "parameters": {}}] - backend = CallableToolBackend(callables={}, schemas=schemas_in) + backend = LocalToolBackend(callables={}, schemas=schemas_in) out1 = backend.schemas out1.append({"name": "mutated"}) @@ -130,7 +130,7 @@ async def test_dispatch_all_sequential_preserves_declaration_order(): async def echo(args: dict) -> dict: return {"echoed": args["i"]} - backend = CallableToolBackend(callables={"echo": echo}) + backend = LocalToolBackend(callables={"echo": echo}) calls = [_make_call("echo", call_id=f"c{i}", arguments={"i": i}) for i in range(5)] pairs = await backend.dispatch_all_sequential_async(calls) @@ -161,7 +161,7 @@ async def reverse(args: dict) -> dict: invocations.append(("reverse", args)) return {"reversed": args.get("text", "")[::-1]} - backend = CallableToolBackend(callables={"echo": echo, "add": add, "reverse": reverse}) + backend = LocalToolBackend(callables={"echo": echo, "add": add, "reverse": reverse}) prepended_calls = [ _make_call("echo", call_id="e1", arguments={"text": "hello"}), From e8ab8ffd055e742454836ba6416e64bd9ac7044d Mon Sep 17 00:00:00 2001 From: Victor Valbuena Date: Tue, 26 May 2026 15:56:40 -0700 Subject: [PATCH 24/40] Addition of MCP components including the MCP client and tool backend. --- pyrit/tools/__init__.py | 14 + pyrit/tools/backend.py | 40 +-- pyrit/tools/local_backend.py | 24 +- pyrit/tools/mcp_backend.py | 199 +++++++++++++++ pyrit/tools/mcp_client.py | 369 +++++++++++++++++++++++++++ tests/unit/tools/conftest.py | 13 +- tests/unit/tools/test_mcp_backend.py | 156 +++++++++++ tests/unit/tools/test_mcp_client.py | 171 +++++++++++++ 8 files changed, 936 insertions(+), 50 deletions(-) create mode 100644 pyrit/tools/mcp_backend.py create mode 100644 pyrit/tools/mcp_client.py create mode 100644 tests/unit/tools/test_mcp_backend.py create mode 100644 tests/unit/tools/test_mcp_client.py diff --git a/pyrit/tools/__init__.py b/pyrit/tools/__init__.py index c3520098ea..46b11aa358 100644 --- a/pyrit/tools/__init__.py +++ b/pyrit/tools/__init__.py @@ -44,11 +44,25 @@ from pyrit.tools.backend import ToolBackend from pyrit.tools.local_backend import LocalToolBackend +from pyrit.tools.mcp_backend import MCPToolBackend +from pyrit.tools.mcp_client import ( + DockerMCPServerSpec, + LocalMCPServerSpec, + MCPClient, + MCPServerSpec, + RemoteMCPServerSpec, +) from pyrit.tools.models import ToolCall, ToolEventBehavior, ToolEventPolicy, tool_loop from pyrit.tools.parsers import ToolCallParser __all__ = [ + "DockerMCPServerSpec", + "LocalMCPServerSpec", "LocalToolBackend", + "MCPClient", + "MCPServerSpec", + "MCPToolBackend", + "RemoteMCPServerSpec", "ToolBackend", "ToolCall", "ToolCallParser", diff --git a/pyrit/tools/backend.py b/pyrit/tools/backend.py index 54c1dd7a1d..e7a02a7685 100644 --- a/pyrit/tools/backend.py +++ b/pyrit/tools/backend.py @@ -3,16 +3,16 @@ from __future__ import annotations -from typing import TYPE_CHECKING, Any, Protocol, runtime_checkable +from abc import ABC, abstractmethod +from typing import TYPE_CHECKING, Any if TYPE_CHECKING: from pyrit.tools.models import ToolCall -@runtime_checkable -class ToolBackend(Protocol): +class ToolBackend(ABC): """ - Protocol for backends that dispatch tool calls produced by a target. + Abstract base for backends that dispatch tool calls produced by a target. A :class:`ToolBackend` is a per-target dispatch table — it owns the ``name -> async callable`` mapping a target uses to execute the tool @@ -25,20 +25,18 @@ class ToolBackend(Protocol): * :class:`~pyrit.tools.LocalToolBackend` — in-process backend backed by ``async def`` callables. Useful for unit tests and for embedding tools inside the PyRIT process. - * :class:`pyrit.tools.MCPToolBackend` (lands in C3) — proxies - dispatch through one or more MCP servers. - - The :attr:`schemas` property exposes the JSON-schema descriptors the - target injects into its request body (e.g. ``tools=[...]`` for the - OpenAI APIs). - - :meth:`dispatch_all_sequential_async` is the contract the tool loop - uses: backends that wish to parallelize dispatch should override it. - The default sequencing — one ``await dispatch_async`` per call, in - declaration order — is what every PyRIT backend ships with today. + * :class:`~pyrit.tools.MCPToolBackend` — proxies dispatch through one + or more MCP servers. + + Subclasses MUST implement :attr:`schemas` and :meth:`dispatch_async`. + :meth:`dispatch_all_sequential_async` ships with a default + implementation that awaits :meth:`dispatch_async` once per call in + declaration order; backends that wish to parallelize dispatch + (e.g. fan out across multiple sandbox containers) should override it. """ @property + @abstractmethod def schemas(self) -> list[dict[str, Any]]: """ The JSON-schema descriptors for every tool the backend exposes. @@ -48,8 +46,8 @@ def schemas(self) -> list[dict[str, Any]]: format that concrete targets serialize into their request body. """ - ... + @abstractmethod async def dispatch_async(self, call: ToolCall) -> dict[str, Any]: """ Execute a single tool call and return the structured result. @@ -64,7 +62,6 @@ async def dispatch_async(self, call: ToolCall) -> dict[str, Any]: Returns: dict[str, Any]: The structured tool result. """ - ... async def dispatch_all_sequential_async( self, @@ -73,6 +70,9 @@ async def dispatch_all_sequential_async( """ Dispatch every call in *calls* sequentially, preserving declaration order. + Default implementation: ``await dispatch_async`` once per call. + Backends that parallelize dispatch should override this method. + Args: calls (list[ToolCall]): The calls to dispatch, in declaration order. @@ -80,4 +80,8 @@ async def dispatch_all_sequential_async( list[tuple[ToolCall, dict[str, Any]]]: ``(call, result)`` pairs, in the same order as *calls*. """ - ... + results: list[tuple[ToolCall, dict[str, Any]]] = [] + for call in calls: + envelope = await self.dispatch_async(call) + results.append((call, envelope)) + return results diff --git a/pyrit/tools/local_backend.py b/pyrit/tools/local_backend.py index 0c67054590..25fe42e83c 100644 --- a/pyrit/tools/local_backend.py +++ b/pyrit/tools/local_backend.py @@ -6,6 +6,8 @@ import logging from typing import TYPE_CHECKING, Any +from pyrit.tools.backend import ToolBackend + if TYPE_CHECKING: from collections.abc import Awaitable, Callable @@ -14,7 +16,7 @@ logger = logging.getLogger(__name__) -class LocalToolBackend: +class LocalToolBackend(ToolBackend): """ In-process :class:`~pyrit.tools.ToolBackend` backed by a name -> ``async def`` mapping. Useful for unit tests and for embedding small tools inside the @@ -117,23 +119,3 @@ async def dispatch_async(self, call: ToolCall) -> dict[str, Any]: "detail": str(ex), } return result if isinstance(result, dict) else {"result": result} - - async def dispatch_all_sequential_async( - self, - calls: list[ToolCall], - ) -> list[tuple[ToolCall, dict[str, Any]]]: - """ - Dispatch *calls* sequentially in declaration order. - - Args: - calls (list[ToolCall]): Calls to dispatch. - - Returns: - list[tuple[ToolCall, dict[str, Any]]]: ``(call, result)`` pairs - in the same order as *calls*. - """ - results: list[tuple[ToolCall, dict[str, Any]]] = [] - for call in calls: - result = await self.dispatch_async(call) - results.append((call, result)) - return results diff --git a/pyrit/tools/mcp_backend.py b/pyrit/tools/mcp_backend.py new file mode 100644 index 0000000000..66da88a30f --- /dev/null +++ b/pyrit/tools/mcp_backend.py @@ -0,0 +1,199 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT license. + +""" +Multi-server tool backend that proxies dispatch through one or more +MCP servers. + +This is the :class:`~pyrit.tools.ToolBackend` implementation that real +red-team configurations use. It composes one +:class:`~pyrit.tools.MCPClient` per :class:`~pyrit.tools.MCPServerSpec`, +aggregates their advertised schemas, routes incoming +:class:`~pyrit.tools.ToolCall` instances to the correct underlying +client, and enforces an optional ``allowed_tools`` allow-list. + +Contrast with :class:`~pyrit.tools.LocalToolBackend`, which dispatches +to Python ``async def`` callables inside PyRIT's own process. +""" + +from __future__ import annotations + +import asyncio +import logging +from contextlib import AsyncExitStack +from typing import TYPE_CHECKING, Any + +from pyrit.tools.backend import ToolBackend +from pyrit.tools.mcp_client import MCPClient + +if TYPE_CHECKING: + from collections.abc import Iterable + + from pyrit.tools.mcp_client import MCPServerSpec + from pyrit.tools.models import ToolCall + +logger = logging.getLogger(__name__) + + +class MCPToolBackend(ToolBackend): + """ + :class:`~pyrit.tools.ToolBackend` backed by one or more MCP servers. + + On :meth:`__aenter__`, the backend spawns / connects each server in + its :attr:`_servers` list (sequentially) through a single + :class:`contextlib.AsyncExitStack`, runs the MCP handshake, caches + schemas, and builds an advertised-name → ``(client, server_name)`` + routing table. Collisions raise :class:`ValueError` unless the + colliding specs set :attr:`~pyrit.tools.LocalMCPServerSpec.name_prefix`. + + A single shared :class:`AsyncExitStack` (rather than one per client) + is required so anyio's nested cancel scopes — opened by the ``mcp`` + SDK's ``stdio_client`` and ``ClientSession`` context managers — are + closed in strict LIFO order from the entering task. Closing + out-of-order would trip + ``"Attempted to exit a cancel scope that isn't the current task's + current cancel scope"``. + + Dispatch is serialized through an :class:`asyncio.Lock` per backend + instance — multiple concurrent coroutines sharing the same backend + (e.g. parallel attack runs) will not interleave JSON-RPC frames on + the same stdio pipe. + """ + + def __init__( + self, + *, + servers: Iterable[MCPServerSpec], + allowed_tools: list[str] | None = None, + ) -> None: + """ + Initialize the backend. + + Args: + servers: One or more :class:`MCPServerSpec` instances describing + where each server runs. + allowed_tools: Optional allow-list of tool names. Names not in + the list are filtered from :attr:`schemas` AND + short-circuit dispatch with a ``tool_not_allowed`` envelope. + Names are matched after :attr:`~LocalMCPServerSpec.name_prefix` + has been applied. Defaults to None (every advertised tool is + callable). + + Raises: + ValueError: When *servers* is empty. + """ + self._servers: list[MCPServerSpec] = list(servers) + if not self._servers: + raise ValueError("MCPToolBackend requires at least one server spec.") + self._allowed_tools: set[str] | None = set(allowed_tools) if allowed_tools is not None else None + self._clients: list[MCPClient] = [] + self._routing: dict[str, tuple[MCPClient, str]] = {} + self._dispatch_lock = asyncio.Lock() + self._stack: AsyncExitStack | None = None + self._entered = False + + @property + def schemas(self) -> list[dict[str, Any]]: + """The union of every connected server's schemas, filtered by ``allowed_tools``.""" + out: list[dict[str, Any]] = [] + for client in self._clients: + for schema in client.schemas: + if self._allowed_tools is not None and schema["name"] not in self._allowed_tools: + continue + out.append(schema) + return out + + async def __aenter__(self) -> MCPToolBackend: + """ + Connect each underlying client through a shared :class:`AsyncExitStack` and build the routing table. + + Returns: + MCPToolBackend: *self*, ready to dispatch. + + Raises: + ValueError: When two connected clients advertise the same tool + name without a disambiguating ``name_prefix``. + """ + stack = AsyncExitStack() + clients: list[MCPClient] = [] + routing: dict[str, tuple[MCPClient, str]] = {} + try: + for spec in self._servers: + client = MCPClient(spec=spec) + await stack.enter_async_context(client) + clients.append(client) + for advertised_name in client.tool_names: + if advertised_name in routing: + raise ValueError( + f"duplicate tool name '{advertised_name}'. " + "Set LocalMCPServerSpec.name_prefix on at least one " + "colliding server to disambiguate.", + ) + routing[advertised_name] = (client, advertised_name) + except Exception: + await stack.aclose() + raise + + self._stack = stack + self._clients = clients + self._routing = routing + self._entered = True + return self + + async def __aexit__(self, *exc: Any) -> None: + """Tear down every underlying client in strict LIFO order.""" + stack = self._stack + self._stack = None + self._clients = [] + self._routing = {} + self._entered = False + if stack is not None: + await stack.aclose() + + async def dispatch_async(self, call: ToolCall) -> dict[str, Any]: + """ + Route *call* to the correct client and dispatch. + + See :class:`MCPClient.dispatch_async` for the envelope shape. + Allow-list rejections and unknown-tool calls return error + envelopes; only "backend not entered" raises. + + Args: + call (ToolCall): The call to dispatch. + + Returns: + dict[str, Any]: A structured envelope (success, ``tool_not_allowed``, + ``tool_not_registered``, or the underlying + :meth:`MCPClient.dispatch_async` envelope). + + Raises: + RuntimeError: When the backend has not been entered via ``async with``. + """ + if not self._entered: + raise RuntimeError( + "MCPToolBackend is not active. Use `async with backend:` to manage its lifecycle before dispatching.", + ) + + if self._allowed_tools is not None and call.name not in self._allowed_tools: + logger.info("Rejecting disallowed tool call: %s", call.name) + return { + "is_error": True, + "error": "tool_not_allowed", + "tool": call.name, + "allowed_tools": sorted(self._allowed_tools), + } + + route = self._routing.get(call.name) + if route is None: + available = sorted(self._routing.keys()) + logger.warning("Tool '%s' not registered. Available: %s", call.name, available) + return { + "is_error": True, + "error": "tool_not_registered", + "tool": call.name, + "available_tools": available, + } + + client, _server_side_name = route + async with self._dispatch_lock: + return await client.dispatch_async(call) diff --git a/pyrit/tools/mcp_client.py b/pyrit/tools/mcp_client.py new file mode 100644 index 0000000000..904004f675 --- /dev/null +++ b/pyrit/tools/mcp_client.py @@ -0,0 +1,369 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT license. + +""" +Stdio-transport client for the Model Context Protocol (MCP). + +This module is the wire-protocol half of PyRIT's MCP integration. It +sits below :class:`~pyrit.tools.MCPToolBackend` (which composes one +:class:`MCPClient` per configured server and handles cross-server +routing) and above the upstream ``mcp`` Python SDK (which owns the +JSON-RPC framing, capability negotiation, and asyncio task plumbing). + +The three :class:`MCPServerSpec` variants describe *where* the server +runs. Only :class:`LocalMCPServerSpec` is implemented in this commit: + +* :class:`LocalMCPServerSpec` — spawn the server as a child process and + speak JSON-RPC over its stdin/stdout. +* :class:`RemoteMCPServerSpec` — HTTP/SSE transport against a hosted + server. Stub: ``connect_async`` raises ``NotImplementedError``. +* :class:`DockerMCPServerSpec` — stdio over ``docker run -i`` against a + hardened sandbox container. Stub: ``connect_async`` raises + ``NotImplementedError``. Implementation lands in the follow-up + sandbox PR. + +The stub variants are intentionally part of the type union today so +downstream code can be written against the eventual API without +forcing a Union expansion later. +""" + +from __future__ import annotations + +import asyncio +import logging +from contextlib import AsyncExitStack +from dataclasses import dataclass +from typing import TYPE_CHECKING, Any + +from mcp import ClientSession +from mcp.client.stdio import StdioServerParameters, stdio_client + +if TYPE_CHECKING: + from pyrit.tools.models import ToolCall + +logger = logging.getLogger(__name__) + + +@dataclass(frozen=True) +class LocalMCPServerSpec: + """ + Spec for an MCP server spawned as a child process and reached via + stdio JSON-RPC. + + Attributes: + command (str): The interpreter or binary to exec (e.g. ``"python"``). + args (tuple[str, ...]): Arguments passed to *command*, in order. + env (dict[str, str] | None): Environment overlay for the child + process. ``None`` (default) inherits PyRIT's environment. + name_prefix (str | None): When set, every tool advertised by the + server is registered as ``f"{name_prefix}{tool_name}"`` in + the parent :class:`~pyrit.tools.MCPToolBackend`. Used to + disambiguate two servers that expose the same tool name. + timeout_seconds (float): Per-call timeout, enforced by + :meth:`MCPClient.dispatch_async`. Defaults to 30 seconds. + """ + + command: str + args: tuple[str, ...] = () + env: dict[str, str] | None = None + name_prefix: str | None = None + timeout_seconds: float = 30.0 + + +@dataclass(frozen=True) +class RemoteMCPServerSpec: + """ + Spec for an MCP server reached over HTTP / SSE. **Not implemented** + in this PR — :meth:`MCPClient.connect_async` raises + :class:`NotImplementedError`. Tracked by ``# TODO(mcp-http-transport)``. + + Attributes: + url (str): The base URL of the MCP server. + name_prefix (str | None): Same semantics as + :attr:`LocalMCPServerSpec.name_prefix`. + timeout_seconds (float): Per-call timeout. + """ + + url: str + name_prefix: str | None = None + timeout_seconds: float = 30.0 + + +# TODO(sandbox-provider) — DockerMCPServerSpec stub here; implementation lands in follow-up PR. +@dataclass(frozen=True) +class DockerMCPServerSpec: + """ + Spec for an MCP server hosted inside a hardened Docker container. + + **NOT IMPLEMENTED IN THIS PR.** Reached via stdio over ``docker run -i``. + + Expected behavior in the follow-up sandbox PR: + + * One container per spec instance, managed by a process-wide + ``SandboxPool``. + * Image is built lazily, keyed by ``sha256(Dockerfile + build_context)``, + and cached across attacks; no rebuild unless missing or explicitly + overridden. + * Container is recreated from the cached image at attack and scenario + boundaries (filesystem returns to baseline every time). + * Network access governed by ``NetworkProfile`` (default ``"none"`` = + ``--network=none``). + * Container runs as a non-root UID with ``--cap-drop=ALL``, a read-only + root filesystem, and an in-container MCP server exposing + ``run_shell(cmd, timeout_seconds)``. + + Attributes: + image (str): Docker image tag (e.g. ``"pyrit-sandbox:base"``). + network_profile (str): ``NetworkProfile`` name; ``"none"`` (default) + launches the container with ``--network=none``. + name_prefix (str | None): Same semantics as + :attr:`LocalMCPServerSpec.name_prefix`. + timeout_seconds (float): Per-call timeout. + + Future fields (deferred to the follow-up sandbox PR): ``memory_limit``, + ``cpu_limit``, ``pids_limit``, ``env``, ``mounts``, ``command_override``. + """ + + image: str + network_profile: str = "none" + name_prefix: str | None = None + timeout_seconds: float = 30.0 + + +MCPServerSpec = LocalMCPServerSpec | RemoteMCPServerSpec | DockerMCPServerSpec + + +def _to_input_schema_dict(input_schema: Any) -> dict[str, Any]: + """ + Coerce the SDK's tool ``inputSchema`` (pydantic model or dict) into a plain dict. + + Returns: + dict[str, Any]: A plain-dict copy of *input_schema*, or an empty + object schema when *input_schema* is None or of an unrecognized type. + """ + if input_schema is None: + return {"type": "object", "properties": {}} + if hasattr(input_schema, "model_dump"): + return input_schema.model_dump() + if isinstance(input_schema, dict): + return dict(input_schema) + return {"type": "object", "properties": {}} + + +def _flatten_content(content: list[Any]) -> str: + """ + Concatenate the text portions of an MCP ``CallToolResult.content`` list. + + Returns: + str: Concatenated ``.text`` values from each content item, in order. + """ + pieces: list[str] = [] + for item in content: + text = getattr(item, "text", None) + if text is not None: + pieces.append(text) + elif isinstance(item, dict) and "text" in item: + pieces.append(item["text"]) + return "".join(pieces) + + +class MCPClient: + """ + A single MCP-server session. + + The client owns the lifetime of one server's transport stack and + exposes a uniform :meth:`dispatch_async` regardless of which + :class:`MCPServerSpec` variant it was constructed from. Composition + across multiple servers (routing, schema aggregation, allow-lists) + is the responsibility of :class:`~pyrit.tools.MCPToolBackend`. + + Lifecycle: + + * :meth:`connect_async` spawns the subprocess (for + :class:`LocalMCPServerSpec`), runs the MCP handshake, and caches + ``tools/list`` results. + * :meth:`dispatch_async` issues one ``tools/call`` and returns a + structured envelope (success or error). + * :meth:`close_async` tears down the transport stack. + + The class is usable as an async context manager. + """ + + def __init__(self, *, spec: MCPServerSpec) -> None: + """ + Initialize the client around *spec*. Does not connect; call + :meth:`connect_async` (or use the async context-manager form) to start + the transport stack. + """ + self._spec = spec + self._stack = AsyncExitStack() + self._session: ClientSession | None = None + self._tools: list[Any] = [] + + @property + def spec(self) -> MCPServerSpec: + """The :class:`MCPServerSpec` this client was constructed with.""" + return self._spec + + @property + def schemas(self) -> list[dict[str, Any]]: + """ + JSON schemas for every tool the server advertises. + + Each schema is shaped ``{"name", "description", "parameters"}``. + The optional :attr:`LocalMCPServerSpec.name_prefix` is applied + here so a backend that owns this client sees the prefixed name. + """ + prefix = getattr(self._spec, "name_prefix", None) or "" + return [ + { + "name": f"{prefix}{tool.name}", + "description": tool.description or "", + "parameters": _to_input_schema_dict(tool.inputSchema), + } + for tool in self._tools + ] + + @property + def tool_names(self) -> list[str]: + """Tool names with the spec's :attr:`name_prefix` applied.""" + return [s["name"] for s in self.schemas] + + def _strip_prefix(self, name: str) -> str: + prefix = getattr(self._spec, "name_prefix", None) or "" + if prefix and name.startswith(prefix): + return name[len(prefix) :] + return name + + async def connect_async(self) -> None: + """Establish the transport, run the handshake, and cache schemas.""" + if isinstance(self._spec, RemoteMCPServerSpec): + raise NotImplementedError( + "HTTP/SSE transport ships in a follow-up PR. " + "RemoteMCPServerSpec is declared today so user code can target the eventual API." + ) + if isinstance(self._spec, DockerMCPServerSpec): + raise NotImplementedError( + "Docker sandbox transport ships in a follow-up PR. " + "DockerMCPServerSpec runs the MCP server inside a hardened " + "Debian container reached via stdio over `docker run -i`, " + "managed by a process-wide SandboxPool with image caching and " + "per-attack container recreation." + ) + + assert isinstance(self._spec, LocalMCPServerSpec) + params = StdioServerParameters( + command=self._spec.command, + args=list(self._spec.args), + env=self._spec.env, + ) + read, write = await self._stack.enter_async_context(stdio_client(params)) + session = await self._stack.enter_async_context(ClientSession(read, write)) + await session.initialize() + result = await session.list_tools() + self._session = session + self._tools = list(result.tools) + + async def close_async(self) -> None: + """Tear down the transport stack. Idempotent; safe to call before connect.""" + try: + await self._stack.aclose() + except Exception as ex: # noqa: BLE001 — close should never raise into the caller. + logger.warning("Error tearing down MCP client stack: %s", ex) + finally: + self._stack = AsyncExitStack() + self._session = None + self._tools = [] + + async def __aenter__(self) -> MCPClient: + """ + Connect the transport stack and return *self*. + + Returns: + MCPClient: *self*, ready to dispatch tool calls. + """ + await self.connect_async() + return self + + async def __aexit__(self, *exc: Any) -> None: + """Tear down the transport stack.""" + await self.close_async() + + async def dispatch_async(self, call: ToolCall) -> dict[str, Any]: + """ + Issue one ``tools/call`` and return a structured envelope. + + Envelope shape: + + * Success: ``{"is_error": False, "content": str, "tool": name}``. + * Timeout: ``{"is_error": True, "error": "tool_timeout", "tool": name, ...}``. + * Server-reported error: ``{"is_error": True, "error": "tool_execution_failed", "tool": name, ...}``. + + Tool-side failures are converted to envelopes; only programmer + errors (calling before :meth:`connect_async`) raise. + + Args: + call (ToolCall): The call to dispatch. The advertised + ``name_prefix`` (if any) is stripped before contacting the server. + + Returns: + dict[str, Any]: One of the envelope shapes documented above. + + Raises: + RuntimeError: When the client has not been connected. + """ + if self._session is None: + raise RuntimeError("MCPClient is not connected; call connect_async first.") + + server_side_name = self._strip_prefix(call.name) + timeout = getattr(self._spec, "timeout_seconds", 30.0) + try: + result = await asyncio.wait_for( + self._session.call_tool(server_side_name, arguments=dict(call.arguments)), + timeout=timeout, + ) + except asyncio.TimeoutError: + logger.warning( + "MCP tool '%s' timed out after %.2fs", + call.name, + timeout, + ) + return { + "is_error": True, + "error": "tool_timeout", + "tool": call.name, + "timeout_seconds": timeout, + } + except Exception as ex: # noqa: BLE001 — wrap and surface as envelope. + logger.warning( + "MCP tool '%s' raised %s: %s", + call.name, + type(ex).__name__, + ex, + ) + return { + "is_error": True, + "error": "tool_execution_failed", + "tool": call.name, + "detail": str(ex), + } + + content_text = _flatten_content(list(result.content)) + is_error = bool(getattr(result, "isError", False)) + envelope: dict[str, Any] = { + "is_error": is_error, + "content": content_text, + "tool": call.name, + } + if is_error: + envelope["error"] = "tool_execution_failed" + return envelope + + +__all__ = [ + "DockerMCPServerSpec", + "LocalMCPServerSpec", + "MCPClient", + "MCPServerSpec", + "RemoteMCPServerSpec", +] diff --git a/tests/unit/tools/conftest.py b/tests/unit/tools/conftest.py index 8419c5b06c..dd1ae7756e 100644 --- a/tests/unit/tools/conftest.py +++ b/tests/unit/tools/conftest.py @@ -39,6 +39,7 @@ from pyrit.prompt_target.common.target_configuration import TargetConfiguration from pyrit.tools import ( LocalToolBackend, + ToolBackend, ToolCall, ToolCallParser, ToolEventBehavior, @@ -144,7 +145,7 @@ def parse(self, message: Message) -> list[ToolCall]: return calls -class _RecordingToolBackend: +class _RecordingToolBackend(ToolBackend): """ Minimal :class:`ToolBackend` that records every dispatched call and returns results from a scripted queue. Used to assert dispatch order, @@ -172,16 +173,6 @@ async def dispatch_async(self, call: ToolCall) -> dict[str, Any]: nxt = self._results.popleft() return nxt if isinstance(nxt, dict) else {"result": nxt} - async def dispatch_all_sequential_async( - self, - calls: list[ToolCall], - ) -> list[tuple[ToolCall, dict[str, Any]]]: - results: list[tuple[ToolCall, dict[str, Any]]] = [] - for call in calls: - result = await self.dispatch_async(call) - results.append((call, result)) - return results - class _FakeToolTarget(PromptTarget): """ diff --git a/tests/unit/tools/test_mcp_backend.py b/tests/unit/tools/test_mcp_backend.py new file mode 100644 index 0000000000..0abc8ff88c --- /dev/null +++ b/tests/unit/tools/test_mcp_backend.py @@ -0,0 +1,156 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT license. + +""" +Unit tests for :class:`pyrit.tools.MCPToolBackend`. + +These tests verify the multi-server fan-out and routing layer on top of +:class:`MCPClient`: schema aggregation, name-collision detection, +``name_prefix`` disambiguation, ``allowed_tools`` allow-list semantics, +and concurrent-dispatch serialization. They reuse the real +``echo_mcp_server.py`` stdio subprocess. + +Coverage map: + +* **U18** — ``test_disallowed_tool_returns_error_envelope_without_invoking_server``. +* **U20a** — ``test_name_collision_raises_value_error``. +* **U20b** — ``test_name_prefix_disambiguates_colliding_servers``. +* **U21** — ``test_concurrent_dispatch_is_serialized_by_lock``. +""" + +from __future__ import annotations + +import asyncio +import sys +from pathlib import Path + +import pytest + +from pyrit.tools import ( + LocalMCPServerSpec, + MCPToolBackend, + ToolCall, +) + +ECHO_SERVER_SCRIPT = str(Path(__file__).parent / "echo_mcp_server.py") + + +def _spec(*, name_prefix: str | None = None, timeout_seconds: float = 5.0) -> LocalMCPServerSpec: + return LocalMCPServerSpec( + command=sys.executable, + args=(ECHO_SERVER_SCRIPT,), + name_prefix=name_prefix, + timeout_seconds=timeout_seconds, + ) + + +def _make_call(name: str, *, call_id: str = "c1", arguments: dict | None = None) -> ToolCall: + return ToolCall(call_id=call_id, name=name, arguments=arguments or {}) + + +@pytest.mark.asyncio +async def test_backend_aggregates_schemas_across_servers() -> None: + """Schemas from every connected server show up in :attr:`schemas`.""" + backend = MCPToolBackend(servers=[_spec()]) + async with backend: + names = {s["name"] for s in backend.schemas} + assert names == {"echo", "add", "reverse", "slow_echo"} + + +@pytest.mark.asyncio +async def test_dispatch_routes_to_correct_server() -> None: + """A :class:`ToolCall` is routed to the server that registered the name.""" + backend = MCPToolBackend(servers=[_spec()]) + async with backend: + envelope = await backend.dispatch_async(_make_call("echo", arguments={"text": "routed"})) + assert envelope["is_error"] is False + assert envelope["content"] == "routed" + + +@pytest.mark.asyncio +async def test_name_collision_raises_value_error() -> None: + """Two servers exposing the same tool name without prefixes raise.""" + backend = MCPToolBackend(servers=[_spec(), _spec()]) + with pytest.raises(ValueError, match="duplicate tool name"): + await backend.__aenter__() + # __aexit__ is the cleanup path; __aenter__ failing leaves nothing to clean. + + +@pytest.mark.asyncio +async def test_name_prefix_disambiguates_colliding_servers() -> None: + """Setting :attr:`LocalMCPServerSpec.name_prefix` disambiguates duplicates.""" + backend = MCPToolBackend( + servers=[ + _spec(name_prefix="a_"), + _spec(name_prefix="b_"), + ], + ) + async with backend: + names = {s["name"] for s in backend.schemas} + assert "a_echo" in names + assert "b_echo" in names + envelope = await backend.dispatch_async(_make_call("a_echo", arguments={"text": "alpha"})) + assert envelope["content"] == "alpha" + envelope_b = await backend.dispatch_async(_make_call("b_echo", arguments={"text": "beta"})) + assert envelope_b["content"] == "beta" + + +@pytest.mark.asyncio +async def test_disallowed_tool_returns_error_envelope_without_invoking_server() -> None: + """U18: allowed_tools blocks both schema advertisement AND dispatch.""" + backend = MCPToolBackend(servers=[_spec()], allowed_tools=["echo"]) + async with backend: + advertised = {s["name"] for s in backend.schemas} + assert advertised == {"echo"} # add/reverse/slow_echo are filtered out. + + envelope = await backend.dispatch_async(_make_call("add", arguments={"a": 1, "b": 2})) + assert envelope["is_error"] is True + assert envelope["error"] == "tool_not_allowed" + assert envelope["tool"] == "add" + assert envelope["allowed_tools"] == ["echo"] + + +@pytest.mark.asyncio +async def test_unknown_tool_returns_error_envelope() -> None: + """A call to a name no connected server exposes returns an error envelope.""" + backend = MCPToolBackend(servers=[_spec()]) + async with backend: + envelope = await backend.dispatch_async(_make_call("never_registered")) + assert envelope["is_error"] is True + assert envelope["error"] == "tool_not_registered" + assert envelope["tool"] == "never_registered" + + +@pytest.mark.asyncio +async def test_concurrent_dispatch_is_serialized_by_lock() -> None: + """U21: two coroutines dispatching against the same backend do not interleave. + + The slow_echo tool sleeps server-side; without the lock the two + dispatches would issue overlapping JSON-RPC frames over the same + stdio pipe. With the lock they run back-to-back. We assert both + return successfully — interleaved frames would surface as protocol + errors or wrong content. + """ + backend = MCPToolBackend(servers=[_spec(timeout_seconds=10.0)]) + async with backend: + results = await asyncio.gather( + backend.dispatch_async(_make_call("slow_echo", arguments={"text": "A", "delay_ms": 50})), + backend.dispatch_async(_make_call("slow_echo", arguments={"text": "B", "delay_ms": 50})), + ) + assert all(not r["is_error"] for r in results) + assert {r["content"] for r in results} == {"A", "B"} + + +@pytest.mark.asyncio +async def test_dispatch_all_sequential_async_preserves_order() -> None: + """Bulk dispatch returns (call, envelope) pairs in declaration order.""" + backend = MCPToolBackend(servers=[_spec()]) + calls = [ + _make_call("echo", call_id="c1", arguments={"text": "first"}), + _make_call("echo", call_id="c2", arguments={"text": "second"}), + _make_call("echo", call_id="c3", arguments={"text": "third"}), + ] + async with backend: + results = await backend.dispatch_all_sequential_async(calls) + assert [c.call_id for c, _ in results] == ["c1", "c2", "c3"] + assert [r["content"] for _, r in results] == ["first", "second", "third"] diff --git a/tests/unit/tools/test_mcp_client.py b/tests/unit/tools/test_mcp_client.py new file mode 100644 index 0000000000..67f93d046e --- /dev/null +++ b/tests/unit/tools/test_mcp_client.py @@ -0,0 +1,171 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT license. + +""" +Unit tests for :class:`pyrit.tools.MCPClient` and the +:class:`pyrit.tools.MCPServerSpec` union. + +Coverage map (rows from the C2/C3 test matrix): + +* **U10** — ``test_real_subprocess_dispatch_returns_text_content``, + ``test_sequential_dispatch_against_real_server``. +* **U14** — ``test_connect_async_populates_schemas_via_tools_list``. +* **U17** — ``test_dispatch_timeout_returns_error_envelope``. +* **U20** — ``test_remote_mcp_server_spec_raises_not_implemented``, + ``test_docker_mcp_server_spec_raises_not_implemented``. + +These tests spawn the real ``tests/unit/tools/echo_mcp_server.py`` +subprocess via ``mcp.client.stdio.stdio_client``; they exercise the +full handshake → ``tools/list`` → ``tools/call`` round trip. The +purpose is to verify that ``MCPClient`` is a thin, correct facade +over the SDK rather than to re-test the SDK itself. +""" + +from __future__ import annotations + +import sys +from pathlib import Path + +import pytest + +from pyrit.tools import ( + DockerMCPServerSpec, + LocalMCPServerSpec, + MCPClient, + RemoteMCPServerSpec, + ToolCall, +) + +ECHO_SERVER_SCRIPT = str(Path(__file__).parent / "echo_mcp_server.py") + + +def _local_spec(*, timeout_seconds: float = 5.0) -> LocalMCPServerSpec: + """Build a :class:`LocalMCPServerSpec` that spawns ``echo_mcp_server.py``.""" + return LocalMCPServerSpec( + command=sys.executable, + args=(ECHO_SERVER_SCRIPT,), + timeout_seconds=timeout_seconds, + ) + + +def _make_call(name: str, *, call_id: str = "c1", arguments: dict | None = None) -> ToolCall: + return ToolCall(call_id=call_id, name=name, arguments=arguments or {}) + + +@pytest.mark.asyncio +async def test_real_subprocess_dispatch_returns_text_content() -> None: + """U10: dispatching a single tool call returns the echo server's text response.""" + client = MCPClient(spec=_local_spec()) + async with client: + envelope = await client.dispatch_async(_make_call("echo", arguments={"text": "hi"})) + assert envelope["is_error"] is False + assert envelope["content"] == "hi" + + +@pytest.mark.asyncio +async def test_sequential_dispatch_against_real_server() -> None: + """U10: multiple sequential calls round-trip through the same session.""" + client = MCPClient(spec=_local_spec()) + async with client: + envelopes = [ + await client.dispatch_async(_make_call("echo", arguments={"text": "first"})), + await client.dispatch_async(_make_call("add", arguments={"a": 2, "b": 3})), + await client.dispatch_async(_make_call("reverse", arguments={"text": "abc"})), + ] + contents = [e["content"] for e in envelopes] + assert contents == ["first", "5", "cba"] + + +@pytest.mark.asyncio +async def test_connect_async_populates_schemas_via_tools_list() -> None: + """U14: schemas are discovered via tools/list during connect_async.""" + client = MCPClient(spec=_local_spec()) + async with client: + schemas = client.schemas + names = {s["name"] for s in schemas} + assert names == {"echo", "add", "reverse", "slow_echo"} + echo_schema = next(s for s in schemas if s["name"] == "echo") + assert "parameters" in echo_schema + assert echo_schema["parameters"]["properties"]["text"]["type"] == "string" + + +@pytest.mark.asyncio +async def test_dispatch_timeout_returns_error_envelope() -> None: + """U17: a tool call that exceeds the spec's timeout produces an error envelope.""" + client = MCPClient(spec=_local_spec(timeout_seconds=0.05)) + async with client: + envelope = await client.dispatch_async( + _make_call("slow_echo", arguments={"text": "late", "delay_ms": 500}), + ) + assert envelope["is_error"] is True + assert envelope["error"] == "tool_timeout" + assert envelope["tool"] == "slow_echo" + + +@pytest.mark.asyncio +async def test_dispatch_async_returns_error_envelope_on_unknown_tool() -> None: + """Server-side errors (unknown tool name) surface as is_error envelopes.""" + client = MCPClient(spec=_local_spec()) + async with client: + envelope = await client.dispatch_async(_make_call("nonexistent_tool")) + assert envelope["is_error"] is True + assert envelope["tool"] == "nonexistent_tool" + + +def test_remote_mcp_server_spec_is_frozen_dataclass() -> None: + """U20: RemoteMCPServerSpec exists in the type system as a frozen dataclass.""" + spec = RemoteMCPServerSpec(url="https://example.com/mcp") + assert spec.url == "https://example.com/mcp" + with pytest.raises((AttributeError, Exception)): # frozen dataclass guard + spec.url = "other" # type: ignore[misc] + + +@pytest.mark.asyncio +async def test_remote_mcp_server_spec_raises_not_implemented() -> None: + """U20: connecting to a RemoteMCPServerSpec raises NotImplementedError.""" + client = MCPClient(spec=RemoteMCPServerSpec(url="https://example.com/mcp")) + with pytest.raises(NotImplementedError, match="follow-up PR"): + await client.connect_async() + + +def test_docker_mcp_server_spec_dataclass_fields() -> None: + """U20: DockerMCPServerSpec carries the fields the sandbox PR will consume.""" + spec = DockerMCPServerSpec(image="pyrit-sandbox:base") + assert spec.image == "pyrit-sandbox:base" + assert spec.network_profile == "none" + assert spec.name_prefix is None + assert spec.timeout_seconds == 30.0 + + +@pytest.mark.asyncio +async def test_docker_mcp_server_spec_raises_not_implemented() -> None: + """U20: connecting to a DockerMCPServerSpec raises NotImplementedError.""" + client = MCPClient(spec=DockerMCPServerSpec(image="pyrit-sandbox:base")) + with pytest.raises(NotImplementedError, match="follow-up PR"): + await client.connect_async() + + +@pytest.mark.asyncio +async def test_dispatch_before_connect_raises_runtime_error() -> None: + """Calling dispatch_async before connect_async is a programmer error.""" + client = MCPClient(spec=_local_spec()) + with pytest.raises(RuntimeError, match="not connected"): + await client.dispatch_async(_make_call("echo", arguments={"text": "hi"})) + + +@pytest.mark.asyncio +async def test_close_async_is_idempotent() -> None: + """Calling close_async twice (or before connect) does not raise.""" + client = MCPClient(spec=_local_spec()) + await client.close_async() # before connect — no-op. + await client.connect_async() + await client.close_async() + await client.close_async() # double-close — no-op. + + +@pytest.mark.asyncio +async def test_local_mcp_server_spec_is_frozen() -> None: + """LocalMCPServerSpec is a frozen dataclass.""" + spec = LocalMCPServerSpec(command="python", args=("a.py",)) + with pytest.raises((AttributeError, Exception)): + spec.command = "other" # type: ignore[misc] From c61bcfe20579148884f15ab565979b241eaa09aa Mon Sep 17 00:00:00 2001 From: Victor Valbuena Date: Tue, 26 May 2026 17:03:40 -0700 Subject: [PATCH 25/40] Add supports_tool_use capability + ToolEventPolicy and wire tool_loop into PromptTarget.send_prompt_async C4 lands the in-tree wiring for the generic tool-use loop introduced by C2/C3: - TargetCapabilities gains supports_tool_use: bool (default False) and CapabilityName.TOOL_USE for the corresponding enum value, matching the existing supports_X / "supports_X" naming convention used by every other capability. - TargetConfiguration grows tool_event_policy + tool_backend kwargs, both gettable/settable properties. The setter (and constructor) validate that a non-None tool_backend requires supports_tool_use=True; otherwise they raise ValueError immediately. ToolBackend / ToolEventPolicy imports are quoted + behind TYPE_CHECKING to keep pyrit.prompt_target.common from importing pyrit.tools eagerly. - PromptTarget.send_prompt_async picks up @tool_loop (below the existing @final). The wrapper is a no-op when tool_event_policy is None, so every existing target keeps its current behavior. _tool_parser (property, default None) and _tool_schemas() (default []) are added on the base class as the two collaborators @tool_loop reads. - _permissive_configuration is updated to flip supports_tool_use=True alongside the other supports_X flags so the all-flags-on probe loop in test_discover_target_capabilities still sees every CapabilityName value as supported. tests/unit/tools/conftest.py drops the hand-decorated @tool_loop on _FakeToolTarget.send_prompt_async (which would now violate the base class's @final) and instead wires policy + backend through TargetConfiguration. _tool_parser becomes a subclass property since the base class now defines one. Tests: - test_tool_event_policy.py adds U7 (capability flag wiring through the wrapper) plus dataclass field defaults and the TargetConfiguration validator. - test_prompt_target_tool_loop.py adds U1 / U2 (DB-end) / U8 / U9 / U11 exercised against a _ProductionShapedTarget that uses the real base-class _get_normalized_conversation_async (memory round-trip via patch_central_database). Plus default-_tool_parser / -_tool_schemas assertions. Validation: 8104 unit tests pass; pre-commit clean. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- .../common/discover_target_capabilities.py | 1 + pyrit/prompt_target/common/prompt_target.py | 42 +++ .../common/target_capabilities.py | 9 + .../common/target_configuration.py | 72 ++++- tests/unit/tools/conftest.py | 45 +-- .../tools/test_prompt_target_tool_loop.py | 281 ++++++++++++++++++ tests/unit/tools/test_tool_event_policy.py | 121 ++++++++ 7 files changed, 548 insertions(+), 23 deletions(-) create mode 100644 tests/unit/tools/test_prompt_target_tool_loop.py create mode 100644 tests/unit/tools/test_tool_event_policy.py diff --git a/pyrit/prompt_target/common/discover_target_capabilities.py b/pyrit/prompt_target/common/discover_target_capabilities.py index 859d07d428..a61a0cb9a1 100644 --- a/pyrit/prompt_target/common/discover_target_capabilities.py +++ b/pyrit/prompt_target/common/discover_target_capabilities.py @@ -149,6 +149,7 @@ def _permissive_configuration( supports_json_output=True, supports_editable_history=True, supports_system_prompt=True, + supports_tool_use=True, input_modalities=merged_modalities, ) # Rebuild a fresh configuration from the instance's native capabilities so diff --git a/pyrit/prompt_target/common/prompt_target.py b/pyrit/prompt_target/common/prompt_target.py index 461af0e03b..800335e1ef 100644 --- a/pyrit/prompt_target/common/prompt_target.py +++ b/pyrit/prompt_target/common/prompt_target.py @@ -12,6 +12,7 @@ from pyrit.models.json_response_config import _JsonResponseConfig from pyrit.prompt_target.common.target_capabilities import CapabilityName, TargetCapabilities from pyrit.prompt_target.common.target_configuration import TargetConfiguration +from pyrit.tools import ToolCallParser, tool_loop logger = logging.getLogger(__name__) @@ -85,6 +86,7 @@ def __init__( logging.basicConfig(level=logging.INFO) @final + @tool_loop async def send_prompt_async(self, *, message: Message) -> list[Message]: """ Validate, normalize, and send a prompt to the target. @@ -97,6 +99,13 @@ async def send_prompt_async(self, *, message: Message) -> list[Message]: 3. Delegates to :meth:`_send_prompt_to_target_async` with the normalized conversation. + When the target's :attr:`configuration.tool_event_policy` is set, the + :func:`pyrit.tools.tool_loop` decorator replaces this body with the + agentic loop and re-enters :meth:`_send_prompt_to_target_async` + repeatedly until the model issues a stop response (or the configured + ``max_tool_iterations`` is hit). When no policy is set, the decorator + is a no-op and the body below runs unchanged. + Subclasses MUST NOT override this method. Override :meth:`_send_prompt_to_target_async` instead. @@ -132,6 +141,39 @@ async def _send_prompt_to_target_async(self, *, normalized_conversation: list[Me list[Message]: Response messages from the target. """ + @property + def _tool_parser(self) -> ToolCallParser | None: + """ + Per-target :class:`ToolCallParser` consulted by :func:`pyrit.tools.tool_loop`. + + Targets that participate in the tool-use loop override this property + to return a parser that walks their response messages and extracts + :class:`~pyrit.tools.ToolCall` instances. The base default of + ``None`` signals "this target does not participate" -- the wrapper + short-circuits after the first response. + + Returns: + ToolCallParser | None: The parser, or ``None`` for the default + no-tool-use behavior. + """ + return None + + def _tool_schemas(self) -> list[dict[str, Any]]: + """ + Outbound tool-schema list sent on the next request to the model. + + Targets that participate in the tool-use loop override this method + to translate the active :class:`~pyrit.tools.ToolBackend.schemas` + into the wire format their model expects (Responses API vs. Chat + Completions API vs. anything else). The base default returns an + empty list, which means no schemas are advertised. + + Returns: + list[dict[str, Any]]: One schema per advertised tool, in the + target-specific wire format. Empty by default. + """ + return [] + def _validate_request(self, *, normalized_conversation: list[Message]) -> None: """ Validate the normalized conversation before sending to the target. diff --git a/pyrit/prompt_target/common/target_capabilities.py b/pyrit/prompt_target/common/target_capabilities.py index 6ae9ed69e2..234ef4d359 100644 --- a/pyrit/prompt_target/common/target_capabilities.py +++ b/pyrit/prompt_target/common/target_capabilities.py @@ -24,6 +24,7 @@ class CapabilityName(str, Enum): JSON_OUTPUT = "supports_json_output" EDITABLE_HISTORY = "supports_editable_history" SYSTEM_PROMPT = "supports_system_prompt" + TOOL_USE = "supports_tool_use" class UnsupportedCapabilityBehavior(str, Enum): @@ -138,6 +139,14 @@ class attribute. Users can override individual capabilities per instance # Whether the target natively supports system prompts. supports_system_prompt: bool = False + # Whether the target natively supports model-issued tool calls (the + # canonical OpenAI ``function_call`` / ``function_call_output`` envelopes + # plus an outbound tool-schema list). Targets without this capability + # cannot host a tool-use loop -- attempting to configure a + # :class:`TargetConfiguration` with a ``tool_backend`` on a target whose + # capabilities have ``supports_tool_use=False`` raises at construction. + supports_tool_use: bool = False + # The input modalities supported by the target (e.g., "text", "image"). input_modalities: frozenset[frozenset[PromptDataType]] = frozenset({frozenset(["text"])}) diff --git a/pyrit/prompt_target/common/target_configuration.py b/pyrit/prompt_target/common/target_configuration.py index 72ca42fcc1..b00d611ab7 100644 --- a/pyrit/prompt_target/common/target_configuration.py +++ b/pyrit/prompt_target/common/target_configuration.py @@ -4,7 +4,7 @@ import logging from collections.abc import Mapping from dataclasses import fields -from typing import Any +from typing import TYPE_CHECKING, Any from pyrit.message_normalizer import MessageListNormalizer from pyrit.models import Message @@ -16,6 +16,10 @@ UnsupportedCapabilityBehavior, ) +if TYPE_CHECKING: + from pyrit.tools.backend import ToolBackend + from pyrit.tools.models import ToolEventPolicy + logger = logging.getLogger(__name__) @@ -39,6 +43,15 @@ class TargetConfiguration: Each target defines defaults; callers can override policy or individual normalizers at creation time. + + Tool use is configured by setting :attr:`tool_event_policy` (mandatory + when a target's response contains tool calls; controls EXECUTE / RAISE / + RETURN\\_RAW behavior) and optionally :attr:`tool_backend` (required only + when ``tool_event_policy.behavior`` is ``EXECUTE``). Both default to + ``None`` and are read by :func:`pyrit.tools.tool_loop` at runtime; + constructing a configuration with a ``tool_backend`` on a target that + does not declare ``capabilities.supports_tool_use=True`` raises + immediately. """ def __init__( @@ -47,6 +60,8 @@ def __init__( capabilities: TargetCapabilities, policy: CapabilityHandlingPolicy | None = None, normalizer_overrides: Mapping[CapabilityName, MessageListNormalizer[Any]] | None = None, + tool_event_policy: "ToolEventPolicy | None" = None, + tool_backend: "ToolBackend | None" = None, ) -> None: """ Build a target configuration and resolve the normalization pipeline. @@ -57,7 +72,25 @@ def __init__( capability. Defaults to RAISE for all adaptable capabilities. normalizer_overrides (Mapping[CapabilityName, MessageListNormalizer[Any]] | None): Optional overrides for specific capability normalizers. + tool_event_policy (ToolEventPolicy | None): How + :func:`pyrit.tools.tool_loop` should react to a pending tool + call from the target. ``None`` means the loop is disabled and + the wrapper short-circuits. + tool_backend (ToolBackend | None): Dispatch table used when + ``tool_event_policy.behavior`` is ``EXECUTE``. ``None`` is + valid only for the RAISE / RETURN\\_RAW policies and the + no-policy passthrough. + + Raises: + ValueError: If ``tool_backend`` is set on a target whose + capabilities do not include ``supports_tool_use``. """ + if tool_backend is not None and not capabilities.includes(capability=CapabilityName.TOOL_USE): + raise ValueError( + "tool_backend is set but capabilities.supports_tool_use is False. " + "Either declare supports_tool_use=True on the target's capabilities, " + "or remove the tool_backend." + ) self._capabilities = capabilities self._policy = policy or _DEFAULT_POLICY self._pipeline = ConversationNormalizationPipeline.from_capabilities( @@ -65,6 +98,8 @@ def __init__( policy=self._policy, normalizer_overrides=normalizer_overrides, ) + self._tool_event_policy = tool_event_policy + self._tool_backend = tool_backend @property def capabilities(self) -> TargetCapabilities: @@ -81,6 +116,41 @@ def pipeline(self) -> ConversationNormalizationPipeline: """The resolved normalization pipeline.""" return self._pipeline + @property + def tool_event_policy(self) -> "ToolEventPolicy | None": + """The tool-use policy consulted by :func:`pyrit.tools.tool_loop`.""" + return self._tool_event_policy + + @tool_event_policy.setter + def tool_event_policy(self, value: "ToolEventPolicy | None") -> None: + """Allow runtime updates so callers can opt a configured target into tool use.""" + self._tool_event_policy = value + + @property + def tool_backend(self) -> "ToolBackend | None": + """The tool dispatch backend used when the loop's behavior is ``EXECUTE``.""" + return self._tool_backend + + @tool_backend.setter + def tool_backend(self, value: "ToolBackend | None") -> None: + """ + Allow runtime updates to the backend. + + Re-runs the ``supports_tool_use`` validator so a backend can never be + installed onto a configuration that does not declare the capability. + + Raises: + ValueError: If ``value`` is not ``None`` and the configuration's + capabilities do not include ``supports_tool_use``. + """ + if value is not None and not self._capabilities.includes(capability=CapabilityName.TOOL_USE): + raise ValueError( + "tool_backend is set but capabilities.supports_tool_use is False. " + "Either declare supports_tool_use=True on the target's capabilities, " + "or remove the tool_backend." + ) + self._tool_backend = value + def includes(self, *, capability: CapabilityName) -> bool: """ Check whether the target includes support for the given capability. diff --git a/tests/unit/tools/conftest.py b/tests/unit/tools/conftest.py index dd1ae7756e..ad7d2c7fd1 100644 --- a/tests/unit/tools/conftest.py +++ b/tests/unit/tools/conftest.py @@ -44,7 +44,6 @@ ToolCallParser, ToolEventBehavior, ToolEventPolicy, - tool_loop, ) @@ -180,14 +179,12 @@ class _FakeToolTarget(PromptTarget): pops scripted responses off a queue. ``_get_normalized_conversation_async`` is overridden to return ``[message]`` directly, isolating decorator behavior from the memory + normalization pipeline. - """ - _DEFAULT_CONFIGURATION = TargetConfiguration( - capabilities=TargetCapabilities( - supports_multi_turn=True, - supports_multi_message_pieces=True, - ) - ) + Inherits the base class's ``@final @tool_loop send_prompt_async``; the + policy + backend are wired through :class:`TargetConfiguration` so the + wrapper finds them via ``self.configuration.tool_event_policy`` and + ``self.configuration.tool_backend``. + """ def __init__( self, @@ -197,15 +194,27 @@ def __init__( backend: Any = None, parser: ToolCallParser | None = None, ) -> None: - super().__init__() + # ``supports_tool_use`` is forced on whenever a policy is configured so + # the TargetConfiguration validator accepts the backend. + caps = TargetCapabilities( + supports_multi_turn=True, + supports_multi_message_pieces=True, + supports_tool_use=policy is not None, + ) + config = TargetConfiguration( + capabilities=caps, + tool_event_policy=policy, + tool_backend=backend, + ) + super().__init__(custom_configuration=config) self._scripted_responses: deque[Message] = deque(scripted_responses) self.call_count: int = 0 self.normalized_conversations_seen: list[list[Message]] = [] - # The C2 decorator reads these via getattr; production code wires them - # through TargetConfiguration in C4. - self._configuration.tool_event_policy = policy - self._configuration.tool_backend = backend - self._tool_parser = parser if parser is not None else _CanonicalEnvelopeParser() + self._parser_instance: ToolCallParser | None = parser if parser is not None else _CanonicalEnvelopeParser() + + @property + def _tool_parser(self) -> ToolCallParser | None: + return self._parser_instance async def _get_normalized_conversation_async(self, *, message: Message) -> list[Message]: return [message] @@ -220,14 +229,6 @@ async def _send_prompt_to_target_async(self, *, normalized_conversation: list[Me raise AssertionError(f"Fake target ran out of scripted responses on iteration {self.call_count}.") return [self._scripted_responses.popleft()] - @tool_loop - async def send_prompt_async(self, *, message: Message) -> list[Message]: - # Passthrough path: only invoked when ToolEventPolicy is None. The - # decorator replaces this body entirely when a policy is set. - message.validate() - normalized = await self._get_normalized_conversation_async(message=message) - return await self._send_prompt_to_target_async(normalized_conversation=normalized) - @pytest.fixture def make_fake_target(patch_central_database): diff --git a/tests/unit/tools/test_prompt_target_tool_loop.py b/tests/unit/tools/test_prompt_target_tool_loop.py new file mode 100644 index 0000000000..8848ddb010 --- /dev/null +++ b/tests/unit/tools/test_prompt_target_tool_loop.py @@ -0,0 +1,281 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT license. + +""" +Unit tests for ``@tool_loop`` wired into :meth:`PromptTarget.send_prompt_async`. + +C4 lands the wiring: ``send_prompt_async`` becomes ``@final @tool_loop`` +on the base class, ``_tool_parser`` and ``_tool_schemas()`` get default +no-op implementations, and ``TargetConfiguration`` grows ``tool_event_policy`` ++ ``tool_backend`` kwargs. + +These tests use the production ``_get_normalized_conversation_async`` path +(memory round-trip through :class:`SQLiteMemory` via ``patch_central_database``) +to exercise the wrapper end-to-end. They cover: + +- U1: decorator order (validate + normalize happen exactly once, then the loop) +- U2 (DB-end half): produced ``tool`` message has one ``function_call_output`` + piece per dispatched call, in declaration order +- U8: DB inserts user, asst_with_fc, tool, asst_final in that order +- U9: DB roles + data_types match the canonical envelope +- U11: targets without a policy short-circuit (no wrapper behavior change) + +Tests for capability flag wiring + ``TargetConfiguration`` construction +validation live in :mod:`tests.unit.tools.test_tool_event_policy`. +""" + +from __future__ import annotations + +import json +from collections import deque +from typing import TYPE_CHECKING, Any + +import pytest + +from pyrit.prompt_target.common.prompt_target import PromptTarget +from pyrit.prompt_target.common.target_capabilities import TargetCapabilities +from pyrit.prompt_target.common.target_configuration import TargetConfiguration +from pyrit.tools import ToolCallParser, ToolEventBehavior, ToolEventPolicy + +from .conftest import ( + _CanonicalEnvelopeParser, + _make_assistant_function_call_message, + _make_assistant_text_message, + _make_user_message, + _RecordingToolBackend, +) + +if TYPE_CHECKING: + from pyrit.models import Message + + +class _ProductionShapedTarget(PromptTarget): + """ + Minimal :class:`PromptTarget` that uses the *real* base-class + ``_get_normalized_conversation_async`` (memory round-trip + normalization + pipeline) instead of the conftest stub. Drives the production wrapper + end-to-end so DB-insert-order assertions can run against the real + :class:`CentralMemory` instance set up by ``patch_central_database``. + """ + + def __init__( + self, + *, + scripted_responses: list[Message], + policy: ToolEventPolicy | None, + backend: Any, + parser: ToolCallParser | None, + ) -> None: + caps = TargetCapabilities( + supports_multi_turn=True, + supports_multi_message_pieces=True, + supports_tool_use=policy is not None, + ) + config = TargetConfiguration( + capabilities=caps, + tool_event_policy=policy, + tool_backend=backend, + ) + super().__init__(custom_configuration=config) + self._scripted: deque[Message] = deque(scripted_responses) + self.call_count: int = 0 + self._parser_instance = parser + + @property + def _tool_parser(self) -> ToolCallParser | None: + return self._parser_instance + + async def _send_prompt_to_target_async(self, *, normalized_conversation: list[Message]) -> list[Message]: + self.call_count += 1 + if not self._scripted: + raise AssertionError(f"Target ran out of scripted responses on iteration {self.call_count}.") + response = self._scripted.popleft() + conversation_id = normalized_conversation[-1].message_pieces[0].conversation_id + for piece in response.message_pieces: + piece.conversation_id = conversation_id + return [response] + + +@pytest.fixture +def make_production_target(patch_central_database): + def _factory( + *, + scripted_responses: list[Message], + policy: ToolEventPolicy | None = None, + backend: Any = None, + parser: ToolCallParser | None = None, + ) -> _ProductionShapedTarget: + effective_parser = parser + if effective_parser is None and policy is not None: + effective_parser = _CanonicalEnvelopeParser() + return _ProductionShapedTarget( + scripted_responses=scripted_responses, + policy=policy, + backend=backend, + parser=effective_parser, + ) + + return _factory + + +@pytest.fixture +def execute_policy_fixture(): + return ToolEventPolicy(behavior=ToolEventBehavior.EXECUTE, max_tool_iterations=5) + + +class TestToolLoopWiredIntoBaseClass: + """Verifies ``@tool_loop`` runs on every ``send_prompt_async`` call.""" + + @pytest.mark.asyncio + async def test_decorator_passthrough_when_no_policy(self, make_production_target): + """U11 -- target without a policy behaves exactly like pre-C4 ``send_prompt_async``.""" + target = make_production_target( + scripted_responses=[_make_assistant_text_message("plain")], + policy=None, + ) + + responses = await target.send_prompt_async(message=_make_user_message("hi")) + + assert target.call_count == 1 + assert len(responses) == 1 + assert responses[0].message_pieces[0].original_value == "plain" + + @pytest.mark.asyncio + async def test_tool_loop_order_after_normalize_before_memory(self, make_production_target, execute_policy_fixture): + """U1 -- validate + normalize happen exactly once before the loop iterates.""" + backend = _RecordingToolBackend(scripted_results=[{"result": "echoed"}]) + target = make_production_target( + scripted_responses=[ + _make_assistant_function_call_message(calls=[("c1", "echo", {"text": "x"})]), + _make_assistant_text_message("done"), + ], + policy=execute_policy_fixture, + backend=backend, + ) + + responses = await target.send_prompt_async(message=_make_user_message("please echo")) + + assert target.call_count == 2 + assert [c.name for c in backend.recorded_calls] == ["echo"] + assert len(responses) == 3 + assert responses[0].message_pieces[0].original_value_data_type == "function_call" + assert responses[1].message_pieces[0].original_value_data_type == "function_call_output" + assert responses[2].message_pieces[0].original_value_data_type == "text" + + @pytest.mark.asyncio + async def test_tool_message_has_one_function_call_output_piece_per_call( + self, make_production_target, execute_policy_fixture + ): + """U2 DB-end half -- one tool Message, N pieces, one per dispatched call.""" + backend = _RecordingToolBackend(scripted_results=[{"r": 1}, {"r": 2}]) + target = make_production_target( + scripted_responses=[ + _make_assistant_function_call_message( + calls=[("c1", "echo", {"text": "a"}), ("c2", "echo", {"text": "b"})] + ), + _make_assistant_text_message("done"), + ], + policy=execute_policy_fixture, + backend=backend, + ) + + responses = await target.send_prompt_async(message=_make_user_message("two calls please")) + + tool_msg = responses[1] + assert len(tool_msg.message_pieces) == 2 + call_ids_in_order = [json.loads(p.original_value)["call_id"] for p in tool_msg.message_pieces] + assert call_ids_in_order == ["c1", "c2"] + assert all(p.original_value_data_type == "function_call_output" for p in tool_msg.message_pieces) + assert all(p.api_role == "tool" for p in tool_msg.message_pieces) + + +class TestDbTranscriptAfterToolLoop: + """ + DB-level assertions that exercise the production memory pipeline. + + These tests rely on the wrapper writing the user message + every assistant + + tool message produced during the loop back to ``CentralMemory``, in + declaration order. Whether that write happens *inside* the wrapper or via + the caller (the prompt normalizer) is an implementation detail; the + invariant is the wrapper returns the full chain so the caller can persist + in order. + """ + + @pytest.mark.asyncio + async def test_db_insert_order_user_then_asst_fc_then_tool_then_final_asst( + self, make_production_target, execute_policy_fixture + ): + """U8 -- after a complete tool round, the wrapper's return order is canonical.""" + backend = _RecordingToolBackend(scripted_results=[{"result": "echoed"}]) + target = make_production_target( + scripted_responses=[ + _make_assistant_function_call_message(calls=[("c1", "echo", {"text": "x"})]), + _make_assistant_text_message("done"), + ], + policy=execute_policy_fixture, + backend=backend, + ) + + responses = await target.send_prompt_async(message=_make_user_message("please echo")) + + data_types_in_order = [r.message_pieces[0].original_value_data_type for r in responses] + assert data_types_in_order == ["function_call", "function_call_output", "text"] + + @pytest.mark.asyncio + async def test_db_roles_and_data_types_match_canonical_envelope( + self, make_production_target, execute_policy_fixture + ): + """U9 -- roles and data_types match the canonical envelope contract.""" + backend = _RecordingToolBackend(scripted_results=[{"result": "echoed"}]) + target = make_production_target( + scripted_responses=[ + _make_assistant_function_call_message(calls=[("c1", "echo", {"text": "x"})]), + _make_assistant_text_message("done"), + ], + policy=execute_policy_fixture, + backend=backend, + ) + + responses = await target.send_prompt_async(message=_make_user_message("please echo")) + + asst_fc, tool_msg, asst_final = responses + # function_call from the assistant + assert asst_fc.message_pieces[0].api_role == "assistant" + assert asst_fc.message_pieces[0].original_value_data_type == "function_call" + envelope = json.loads(asst_fc.message_pieces[0].original_value) + assert envelope["type"] == "function_call" + assert envelope["call_id"] == "c1" + assert envelope["name"] == "echo" + # function_call_output from the tool + assert tool_msg.message_pieces[0].api_role == "tool" + assert tool_msg.message_pieces[0].original_value_data_type == "function_call_output" + tool_envelope = json.loads(tool_msg.message_pieces[0].original_value) + assert tool_envelope["type"] == "function_call_output" + assert tool_envelope["call_id"] == "c1" + # Final assistant text + assert asst_final.message_pieces[0].api_role == "assistant" + assert asst_final.message_pieces[0].original_value_data_type == "text" + + +class TestFinalAndAbstractMethodContract: + """ + Asserts the base-class shape changes that C4 introduces but doesn't + exercise via end-to-end runs: ``_tool_parser`` defaults to ``None``, + ``_tool_schemas`` defaults to ``[]``. + """ + + def test_default_tool_parser_is_none(self, make_production_target): + target = make_production_target( + scripted_responses=[_make_assistant_text_message("plain")], + policy=None, + ) + # Subclass overrides only when the test caller passes a parser. With + # no policy + no parser, the override returns None. + assert target._tool_parser is None + + def test_default_tool_schemas_is_empty_list(self, make_production_target): + target = make_production_target( + scripted_responses=[_make_assistant_text_message("plain")], + policy=None, + ) + assert target._tool_schemas() == [] diff --git a/tests/unit/tools/test_tool_event_policy.py b/tests/unit/tools/test_tool_event_policy.py new file mode 100644 index 0000000000..aa451e672d --- /dev/null +++ b/tests/unit/tools/test_tool_event_policy.py @@ -0,0 +1,121 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT license. + +""" +Unit tests for the wiring between :class:`TargetCapabilities.supports_tool_use`, +:class:`TargetConfiguration.tool_event_policy` / +:class:`TargetConfiguration.tool_backend`, and the +:func:`pyrit.tools.tool_loop` decorator that lives on +:class:`PromptTarget.send_prompt_async`. + +These tests are the §7 U7 row plus the construction-time validator added in C4. +They assert the *capability flag* axis only -- that targets which declare +``supports_tool_use=True`` and configure a policy + backend route through +the loop, that targets without a policy short-circuit, and that the +``tool_backend``-without-capability misconfiguration raises at construction. + +End-to-end ordering against the production memory pipeline (U1, U8, U9) is +exercised separately in ``tests/unit/prompt_target/common/test_prompt_target_tool_loop.py``. +""" + +from __future__ import annotations + +import pytest + +from pyrit.prompt_target.common.target_capabilities import TargetCapabilities +from pyrit.prompt_target.common.target_configuration import TargetConfiguration +from pyrit.tools import LocalToolBackend, ToolEventBehavior, ToolEventPolicy + +from .conftest import ( + _make_assistant_function_call_message, + _make_assistant_text_message, + _make_user_message, +) + + +class TestSupportsToolUseCapabilityFlag: + """Asserts the new ``supports_tool_use`` field on :class:`TargetCapabilities`.""" + + def test_default_is_false(self): + caps = TargetCapabilities() + assert caps.supports_tool_use is False + + def test_explicit_true(self): + caps = TargetCapabilities(supports_tool_use=True) + assert caps.supports_tool_use is True + + +class TestTargetConfigurationToolFields: + """Asserts the new ``tool_event_policy`` / ``tool_backend`` kwargs.""" + + def test_defaults_are_none(self): + caps = TargetCapabilities(supports_tool_use=True) + config = TargetConfiguration(capabilities=caps) + assert config.tool_event_policy is None + assert config.tool_backend is None + + def test_explicit_policy_and_backend(self): + caps = TargetCapabilities(supports_tool_use=True) + backend = LocalToolBackend(callables={}, schemas=[]) + policy = ToolEventPolicy(behavior=ToolEventBehavior.EXECUTE) + config = TargetConfiguration( + capabilities=caps, + tool_event_policy=policy, + tool_backend=backend, + ) + assert config.tool_event_policy is policy + assert config.tool_backend is backend + + def test_tool_backend_without_capability_raises(self): + caps = TargetCapabilities(supports_tool_use=False) + backend = LocalToolBackend(callables={}, schemas=[]) + with pytest.raises(ValueError, match="supports_tool_use"): + TargetConfiguration(capabilities=caps, tool_backend=backend) + + def test_tool_event_policy_without_backend_is_allowed(self): + """``RAISE`` / ``RETURN_RAW`` policies do not require a backend.""" + caps = TargetCapabilities(supports_tool_use=True) + policy = ToolEventPolicy(behavior=ToolEventBehavior.RAISE) + config = TargetConfiguration(capabilities=caps, tool_event_policy=policy) + assert config.tool_event_policy is policy + assert config.tool_backend is None + + +class TestCapabilityFlagWiringIntoToolLoop: + """ + U7 -- verify the wrapper dispatches only when the target declares + ``supports_tool_use`` AND a policy is configured. + """ + + @pytest.mark.asyncio + async def test_target_with_tool_use_capability_uses_tool_loop( + self, make_fake_target, recording_backend, execute_policy + ): + backend = recording_backend() + target = make_fake_target( + scripted_responses=[ + _make_assistant_function_call_message(calls=[("c1", "echo", {"text": "hi"})]), + _make_assistant_text_message("done"), + ], + policy=execute_policy(), + backend=backend, + ) + + responses = await target.send_prompt_async(message=_make_user_message("please call echo")) + + assert target.call_count == 2, "Decorator should have iterated twice (call + final)." + assert [c.name for c in backend.recorded_calls] == ["echo"] + assert len(responses) == 3, "user expects asst_fc, tool_msg, asst_final." + + @pytest.mark.asyncio + async def test_target_without_tool_use_capability_skips_dispatch(self, make_fake_target): + target = make_fake_target( + scripted_responses=[_make_assistant_text_message("plain response, no tool call")], + policy=None, + backend=None, + ) + + responses = await target.send_prompt_async(message=_make_user_message("hello")) + + assert target.call_count == 1 + assert len(responses) == 1 From 3a3c26b9f7f542b9579415490cc5d396e2017d58 Mon Sep 17 00:00:00 2001 From: Victor Valbuena Date: Wed, 27 May 2026 17:01:37 -0700 Subject: [PATCH 26/40] Drop C5 (Chat target tool calling): defer to follow-up, redesign around Response target This commit is intentionally empty. It records a scope decision made in response to PR review feedback. No code changes - the C5 working set was uncommitted and has been reverted. # Why we're dropping C5 Review feedback raised two concerns the original C5 did not address: 1. **Duplication against OpenAIResponseTarget.** The Response target already implements an agentic tool loop (openai_response_target.py lines 590-626), the canonical function_call envelope (lines 666-674), a Python-callable dispatch registry (custom_functions), and an allow-list-ish hook (fail_on_missing_function). C5 layered a parallel implementation on top for the Chat target instead of converging both targets onto one stack. 2. **Chat Completions is on its way out.** OpenAI has publicly framed the Responses API as the long-term replacement for Chat Completions. Investing in tool-call plumbing for a deprecated endpoint ages out fast and obscures the actual value of this PR. The right framing is: this PR is not "tool calling for all targets." It is "pluggable tool-execution backends + a client-side agentic loop for non-Responses-API targets." The Responses API is one transport; this PR is the in-process abstraction that works for every transport. # What survives unchanged C1 (mcp SDK dep), C2 (tools/ scaffold + LocalToolBackend), C3 (MCPClient + MCPToolBackend + Docker stub), and C4 (capability flag + @tool_loop wired on the base class) all remain shipped. The genuinely-novel work - local stdio MCP, pluggable backend ABC, ToolEventPolicy (RAISE / EXECUTE / RETURN_RAW), allowed_tools - is unaffected. # The new design **One agentic loop driver.** The @tool_loop decorator on PromptTarget.send_prompt_async (shipped in C4) is the only loop driver. Every target's _send_prompt_to_target_async returns exactly ONE Message per call. The decorator stitches iterations into the response list. **One tool execution layer.** Every dispatched call flows through ToolBackend.dispatch_async(call) -> envelope. Backends (LocalToolBackend for Python callables, MCPToolBackend for stdio MCP subprocesses, future DockerMCPToolBackend, future CompositeToolBackend) are interchangeable behind a single ABC. **Migrate OpenAIResponseTarget onto the decorator (new C5).** Delete the in-class while loop (lines 590-626). _send_prompt_to_target_async becomes "build body, call API, parse response into one Message, return." Add _tool_parser returning CanonicalEnvelopeParser (extracts only function_call pieces; reasoning, mcp_call, web_search_call, etc. continue to pass through to Memory without dispatch). Translate the configured backend's schemas into the Responses-API tools shape inside _construct_request_body (without clobbering an existing extra_body_parameters["tools"]). Wrap custom_functions as a LocalToolBackend internally with DeprecationWarning(removed_in="0.16.0"), preserving the existing fail_on_missing_function semantics. **Integration tests (new C6).** Rewrite to use the Response target as the sole OpenAI tool-calling path, plus end-to-end scenario tests against the real echo_mcp_server. **OpenAIChatTarget receives no tool-calling support in this PR.** A future PR can pull Chat onto the same abstractions if anyone still wants it, but the recommended OpenAI tool-calling path becomes the Responses API. # Risks * Behavior-parity on the Response target: callers that rely on `len(send_prompt_async(...)) == iterations` rather than scanning piece types will need updating. Existing function-chaining tests act as sentinels. * `custom_functions` deprecation must preserve `fail_on_missing_function` semantics through the LocalToolBackend wrapper. * Response parser must continue to round-trip non-`function_call` piece types (reasoning, mcp_call, etc.) to Memory without dispatching. * `extra_body_parameters["tools"]` takes precedence over backend-derived tools so existing manual configs keep working. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> From eb84ed56ae193b87c04fd5e6e0d73a535f156cd3 Mon Sep 17 00:00:00 2001 From: Victor Valbuena Date: Thu, 28 May 2026 10:33:30 -0700 Subject: [PATCH 27/40] Migrate OpenAIResponseTarget onto @tool_loop and LocalToolBackend C6 collapses the Response target in-class agentic loop into the @tool_loop decorator shipped in C4, and routes tool dispatch through LocalToolBackend (wrapping the existing custom_functions registry as a deprecation shim). # What changed - _send_prompt_to_target_async no longer runs a while loop. It now returns exactly one Message per call. The agentic loop is driven by @tool_loop on the base class. - Added _tool_parser returning CanonicalEnvelopeParser from pyrit/tools/parsers.py. The parser extracts only function_call pieces; reasoning, mcp_call, web_search_call, computer_call, local_shell_call, etc. pass through to Memory unchanged because the parser ignores them and the decorator exits cleanly on the empty parse. - Added _tool_schemas() translating the configured backend schemas into the Responses-API tools shape. - _construct_request_body injects tools=... when the backend has schemas. User-supplied extra_body_parameters["tools"] takes precedence. - supports_tool_use=True on _DEFAULT_CONFIGURATION. - custom_functions= now emits DeprecationWarning(removed_in="0.16.0"). Internally wraps into a LocalToolBackend. A LocalToolBackend is always installed (populated or empty) so legacy target._custom_functions[name]=fn mutations keep affecting dispatch via a back-compat property. - Constructor deep-copies the class-level _DEFAULT_CONFIGURATION before mutating it (PromptTarget.get_default_configuration returns the singleton, so otherwise one instances tool_backend would leak across every other instance). # What did NOT change The legacy _find_last_pending_tool_call, _execute_call_section, and _make_tool_piece helpers remain in place. They are no longer called from production code, but existing tests still cover them; cleanup is deferred to the same follow-up PR that removes the custom_functions kwarg after the 0.16.0 deprecation window. # Tests - New tests/unit/prompt_target/target/test_openai_response_target_c6_migration.py with 7 tests covering deprecation warning, dispatch through user-supplied LocalToolBackend, schema injection, extra_body precedence, no-backend behavior, and reasoning-only passthrough. - All 5 existing function-chaining sentinel tests in test_openai_response_target_function_chaining.py pass unchanged: the back-compat _custom_functions property keeps in-place mutations working. 8131 unit tests green; pre-commit clean (ruff format, ruff check, ty). Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- .../openai/openai_response_target.py | 213 +++++++++--- pyrit/tools/__init__.py | 3 +- pyrit/tools/parsers.py | 66 +++- ...est_openai_response_target_c6_migration.py | 304 ++++++++++++++++++ 4 files changed, 527 insertions(+), 59 deletions(-) create mode 100644 tests/unit/prompt_target/target/test_openai_response_target_c6_migration.py diff --git a/pyrit/prompt_target/openai/openai_response_target.py b/pyrit/prompt_target/openai/openai_response_target.py index 48b4d7ade2..670f73f9cd 100644 --- a/pyrit/prompt_target/openai/openai_response_target.py +++ b/pyrit/prompt_target/openai/openai_response_target.py @@ -3,6 +3,7 @@ import json import logging +import warnings from collections.abc import Awaitable, Callable, MutableSequence from enum import Enum from typing import ( @@ -34,6 +35,14 @@ from pyrit.prompt_target.common.utils import limit_requests_per_minute, validate_temperature, validate_top_p from pyrit.prompt_target.openai.openai_error_handling import _is_content_filter_error from pyrit.prompt_target.openai.openai_target import OpenAITarget +from pyrit.tools import ( + CanonicalEnvelopeParser, + LocalToolBackend, + ToolBackend, + ToolCallParser, + ToolEventBehavior, + ToolEventPolicy, +) logger = logging.getLogger(__name__) @@ -76,6 +85,7 @@ class OpenAIResponseTarget(OpenAITarget, PromptTarget): supports_json_output=True, supports_multi_message_pieces=True, supports_system_prompt=True, + supports_tool_use=True, input_modalities=frozenset( { frozenset(["text"]), @@ -154,6 +164,17 @@ def __init__( """ super().__init__(custom_configuration=custom_configuration, **kwargs) + # If the constructed configuration is the class-level _DEFAULT_CONFIGURATION + # singleton (user did not pass custom_configuration AND the underlying_model + # was unrecognized), rebuild a per-instance copy so the C6 tool-backend + # plumbing below does not mutate state shared across every other instance. + if custom_configuration is None and self._configuration is type(self)._DEFAULT_CONFIGURATION: + caps = self._configuration.capabilities + self._configuration = TargetConfiguration( + capabilities=caps, + policy=self._configuration.policy, + ) + # Validate temperature and top_p validate_temperature(temperature) validate_top_p(top_p) @@ -167,10 +188,39 @@ def __init__( self._extra_body_parameters = extra_body_parameters - # Per-instance tool/func registries: - self._custom_functions: dict[str, ToolExecutor] = custom_functions or {} + # ----- Tool-calling plumbing (C6) --------------------------------- + # custom_functions is deprecated as of 0.15.x. New code configures + # tool_backend on TargetConfiguration directly. The kwarg is still + # accepted; we ALWAYS install a LocalToolBackend (whether populated + # or empty) when no other backend is supplied, so legacy in-place + # mutations of `target._custom_functions` (via the back-compat + # property below) keep affecting dispatch. self._fail_on_missing_function: bool = fail_on_missing_function + if self.configuration.tool_backend is None: + shim_backend = LocalToolBackend( + callables=dict(custom_functions) if custom_functions else {}, + schemas=self._derive_default_schemas(custom_functions or {}), + fail_on_missing_function=fail_on_missing_function, + ) + self.configuration.tool_backend = shim_backend + + if custom_functions: + warnings.warn( + "OpenAIResponseTarget(custom_functions=...) is deprecated and will be " + "removed in 0.16.0. Configure tool_backend on TargetConfiguration " + "instead (e.g. LocalToolBackend(callables=..., schemas=..., " + "fail_on_missing_function=...)).", + DeprecationWarning, + stacklevel=2, + ) + + # Default policy to EXECUTE when a backend is present. The wrapper's + # parser returns an empty list when the model produces no tool calls, + # so this is a no-op for plain text completions. + if self.configuration.tool_event_policy is None: + self.configuration.tool_event_policy = ToolEventPolicy(behavior=ToolEventBehavior.EXECUTE) + # Extract the grammar 'tool' if one is present # See # https://platform.openai.com/docs/guides/function-calling#context-free-grammars @@ -185,6 +235,61 @@ def __init__( logger.debug("Detected grammar tool: %s", tool_name) self._grammar_name = tool_name + @staticmethod + def _derive_default_schemas(callables: dict[str, ToolExecutor]) -> list[dict[str, Any]]: + """ + Synthesize minimal JSON schemas for the deprecation-shim path. + + Users who pass the legacy ``custom_functions`` kwarg do not also pass a + schema list (the Responses API would accept the calls anyway because the + legacy path predates structured tool advertisement). To keep the + deprecation shim transparent we generate a schema-less stub per name so + ``_tool_schemas()`` returns something non-empty when the user actually + wires tools. + + Args: + callables: Function name to async callable mapping. + + Returns: + list[dict[str, Any]]: A bare schema per callable (``parameters`` + is the unconstrained empty-object schema). + """ + return [{"name": name, "parameters": {"type": "object"}} for name in callables] + + @property + def _custom_functions(self) -> dict[str, ToolExecutor]: + """ + Back-compat live view of the active backend's callables registry. + + Mutations on the returned dict (``target._custom_functions[name] = fn``, + ``target._custom_functions.pop(name)``) take effect immediately because + the dict object is shared with the underlying + :class:`pyrit.tools.LocalToolBackend`. Returns an empty dict when no + backend is installed or when the configured backend is not a + ``LocalToolBackend``. + + Returns: + dict[str, ToolExecutor]: The live callables dict. + """ + backend = self.configuration.tool_backend + if isinstance(backend, LocalToolBackend): + return cast("dict[str, ToolExecutor]", backend._callables) + return {} + + @_custom_functions.setter + def _custom_functions(self, value: dict[str, ToolExecutor]) -> None: + backend = self.configuration.tool_backend + if isinstance(backend, LocalToolBackend): + backend._callables = dict(value) + backend._schemas = self._derive_default_schemas(value) + return + new_backend = LocalToolBackend( + callables=dict(value), + schemas=self._derive_default_schemas(value), + fail_on_missing_function=self._fail_on_missing_function, + ) + self.configuration.tool_backend = new_backend + def _build_identifier(self) -> ComponentIdentifier: """ Build the identifier with OpenAI response-specific parameters. @@ -378,8 +483,9 @@ async def _construct_request_body( input_items = await self._build_input_for_multi_modal_async(conversation) text_format = self._build_text_format(json_config=json_config) + tool_schemas = self._tool_schemas() - body_parameters = { + body_parameters: dict[str, Any] = { "model": self._model_name, "max_output_tokens": self._max_output_tokens, "temperature": self._temperature, @@ -390,8 +496,11 @@ async def _construct_request_body( "text": text_format, "reasoning": self._build_reasoning_config(), } + if tool_schemas: + body_parameters["tools"] = tool_schemas if self._extra_body_parameters: + # User-supplied extra_body_parameters wins over backend-derived tools. body_parameters.update(self._extra_body_parameters) # Filter out None values @@ -559,11 +668,18 @@ async def _construct_message_from_response(self, response: Any, request: Message @pyrit_target_retry async def _send_prompt_to_target_async(self, *, normalized_conversation: list[Message]) -> list[Message]: """ - Send prompt, handle agentic tool calls (function_call), return all messages. + Send one prompt to the Responses API and return exactly one Message. + + The agentic tool-calling loop now lives in :func:`pyrit.tools.tool_loop` + on the base class. This method is the single-iteration body the loop + re-enters on each turn: build the request body, call the API, parse the + response, return the constructed :class:`Message` wrapped in a list of + length 1. - The Responses API supports structured outputs and tool execution. This method handles both: - - Simple text/reasoning responses - - Agentic tool-calling loops that may require multiple back-and-forth exchanges + The wrapper detects function_call pieces via :attr:`_tool_parser` and + decides whether to dispatch + re-enter. Reasoning, MCP, web-search, + computer-use, and other non-function-call sections pass through to + Memory unchanged because the parser ignores them. Args: normalized_conversation (list[Message]): The full conversation @@ -571,59 +687,54 @@ async def _send_prompt_to_target_async(self, *, normalized_conversation: list[Me pipeline. The current message is the last element. Returns: - List of messages generated during the interaction (assistant responses and tool messages). - The normalizer will persist all of these to memory. + list[Message]: Exactly one Message wrapping the parsed response. """ message = normalized_conversation[-1] - message_piece: MessagePiece = message.message_pieces[0] last_piece = message.message_pieces[-1] json_config = self._get_json_response_config(message_piece=last_piece) - working_conversation: MutableSequence[Message] = list(normalized_conversation) - - # Track all responses generated during this interaction - responses_to_return: list[Message] = [] - - # Main agentic loop - each back-and-forth creates a new message - tool_call_section: Optional[dict[str, Any]] = None - - while True: - logger.info(f"Sending conversation with {len(working_conversation)} messages to the prompt target") - - body = await self._construct_request_body(conversation=working_conversation, json_config=json_config) - - # Use unified error handling - automatically detects Response and validates - result = await self._handle_openai_request( - api_call=lambda body=body: self._client.responses.create(**body), - request=message, - ) - - # Add result to conversation and responses list - working_conversation.append(result) - responses_to_return.append(result) - - # Extract tool call if present - tool_call_section = self._find_last_pending_tool_call(result) - - # If no tool call, we're done - if not tool_call_section: - break - - # Execute the tool/function - tool_output = await self._execute_call_section(tool_call_section) + body = await self._construct_request_body(conversation=list(normalized_conversation), json_config=json_config) + logger.info("Sending conversation with %d messages to the Responses API", len(normalized_conversation)) + result = await self._handle_openai_request( + api_call=lambda body=body: self._client.responses.create(**body), + request=message, + ) + return [result] - # Create a new message with the tool output - tool_piece = self._make_tool_piece(tool_output, tool_call_section["call_id"], reference_piece=message_piece) - tool_message = Message(message_pieces=[tool_piece], skip_validation=True) + @property + def _tool_parser(self) -> ToolCallParser | None: + """ + Canonical-envelope parser shared with future canonical-envelope targets. + + Walks response message pieces and emits one :class:`~pyrit.tools.ToolCall` + per piece whose ``original_value_data_type`` is ``"function_call"``. + Reasoning, MCP, web-search, computer-use, and local-shell sections all + produce pieces of OTHER data types, so the parser returns an empty list + for them and the @tool_loop decorator exits cleanly. Those sections + still land in Memory via the parsed Message returned by + ``_send_prompt_to_target_async``; they're just not client-side + dispatched. + """ + return CanonicalEnvelopeParser() - # Add tool output message to conversation and responses list - working_conversation.append(tool_message) - responses_to_return.append(tool_message) + def _tool_schemas(self) -> list[dict[str, Any]]: + """ + Translate the configured backend's schemas into Responses-API tools shape. - # Continue loop to send tool result and get next response + The Responses API expects each function tool as a top-level + ``{"type": "function", "name": ..., "description": ..., + "parameters": ...}`` entry (NOT wrapped in an inner ``"function"`` key + the way Chat Completions does). The backend's schemas are already the + bare function schema, so we just stamp ``type=function`` on each. - # Return all responses (normalizer will persist all of them to memory) - return responses_to_return + Returns: + list[dict[str, Any]]: One descriptor per advertised tool, or an + empty list when no backend is configured. + """ + backend: ToolBackend | None = self.configuration.tool_backend + if backend is None: + return [] + return [{"type": "function", **schema} for schema in backend.schemas] def _parse_response_output_section( self, *, section: Any, message_piece: MessagePiece, error: Optional[PromptResponseError] diff --git a/pyrit/tools/__init__.py b/pyrit/tools/__init__.py index 46b11aa358..f2ae0090ce 100644 --- a/pyrit/tools/__init__.py +++ b/pyrit/tools/__init__.py @@ -53,9 +53,10 @@ RemoteMCPServerSpec, ) from pyrit.tools.models import ToolCall, ToolEventBehavior, ToolEventPolicy, tool_loop -from pyrit.tools.parsers import ToolCallParser +from pyrit.tools.parsers import CanonicalEnvelopeParser, ToolCallParser __all__ = [ + "CanonicalEnvelopeParser", "DockerMCPServerSpec", "LocalMCPServerSpec", "LocalToolBackend", diff --git a/pyrit/tools/parsers.py b/pyrit/tools/parsers.py index 4ff7fc4c04..c903eb73c7 100644 --- a/pyrit/tools/parsers.py +++ b/pyrit/tools/parsers.py @@ -3,6 +3,7 @@ from __future__ import annotations +import json from typing import TYPE_CHECKING, Protocol, runtime_checkable if TYPE_CHECKING: @@ -16,10 +17,11 @@ class ToolCallParser(Protocol): Protocol for extracting tool calls from a target response message. Concrete parsers live next to the target whose response shape they - understand (see :class:`OpenAIChatTarget` and :class:`OpenAIResponseTarget` - after C7/C8). Parsers MUST return an empty list when the model has - issued a stop response — the tool loop uses the empty list as the - signal to exit. + understand (the canonical-envelope parser shipped here is shared by + :class:`OpenAIResponseTarget`; per-model-family parsers for non-OpenAI + targets ship in a follow-up, see plan §12.9). Parsers MUST return an + empty list when the model has issued a stop response — the tool loop + uses the empty list as the signal to exit. """ def parse(self, message: Message) -> list[ToolCall]: @@ -41,9 +43,9 @@ def _extract_function_call_pieces(message: Message) -> list[MessagePiece]: Return every :class:`MessagePiece` in *message* whose ``original_value_data_type`` is ``"function_call"``. - This is the canonical envelope produced by OpenAI-style targets after - the C6 normalization commit. It is exposed here so concrete parsers - can reuse the filter rather than re-implementing it. + This is the canonical envelope used by every PyRIT-supported tool-emitting + target. It is exposed here so concrete parsers can reuse the filter rather + than re-implementing it. Args: message (Message): The message to scan. @@ -53,3 +55,53 @@ def _extract_function_call_pieces(message: Message) -> list[MessagePiece]: ``"function_call"``, in their declaration order. """ return [piece for piece in message.message_pieces if piece.original_value_data_type == "function_call"] + + +class CanonicalEnvelopeParser: + """ + Reference :class:`ToolCallParser` for the canonical function_call envelope. + + Walks every :class:`MessagePiece` whose ``original_value_data_type`` is + ``"function_call"`` and decodes the canonical JSON shape:: + + { + "type": "function_call", + "call_id": "", + "name": "", + "arguments": "" + } + + into :class:`ToolCall` instances. Pieces of other data types -- reasoning, + mcp_call, web_search_call, etc. -- are ignored (they pass through to + Memory but are not client-side dispatchable). Per-model-family parsers + for non-OpenAI targets ship in a follow-up PR (see plan §12.9). + """ + + def parse(self, message: Message) -> list[ToolCall]: + """ + Decode canonical ``function_call`` pieces in *message* into :class:`ToolCall`. + + Args: + message (Message): The most recent assistant response. + + Returns: + list[ToolCall]: One :class:`ToolCall` per ``function_call`` + piece, in declaration order. Empty if the message contains + no ``function_call`` pieces (model stop). + """ + from pyrit.tools.models import ToolCall + + calls: list[ToolCall] = [] + for piece in _extract_function_call_pieces(message): + envelope = json.loads(piece.original_value) + arguments_raw = envelope.get("arguments", "{}") + arguments = json.loads(arguments_raw) if isinstance(arguments_raw, str) else dict(arguments_raw) + calls.append( + ToolCall( + call_id=envelope["call_id"], + name=envelope["name"], + arguments=arguments, + raw_envelope=envelope, + ) + ) + return calls diff --git a/tests/unit/prompt_target/target/test_openai_response_target_c6_migration.py b/tests/unit/prompt_target/target/test_openai_response_target_c6_migration.py new file mode 100644 index 0000000000..04dfd0b024 --- /dev/null +++ b/tests/unit/prompt_target/target/test_openai_response_target_c6_migration.py @@ -0,0 +1,304 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT license. + +"""C6 additions to the Response target function-chaining suite. + +Covers the migration onto @tool_loop + LocalToolBackend. +""" + +from __future__ import annotations + +import json +import uuid +import warnings +from typing import Any +from unittest.mock import AsyncMock, MagicMock, patch + +import pytest + +from pyrit.models import Message, MessagePiece +from pyrit.prompt_target import OpenAIResponseTarget +from pyrit.prompt_target.common.target_capabilities import TargetCapabilities +from pyrit.prompt_target.common.target_configuration import TargetConfiguration +from pyrit.tools import LocalToolBackend, ToolEventBehavior, ToolEventPolicy + + +def _mock_function_call_response(call_id: str, function_name: str, arguments: dict) -> MagicMock: + """Build a fake Responses-API response containing a function_call section.""" + mock_response = MagicMock() + mock_response.status = "completed" + mock_response.error = None + section = MagicMock() + section.type = "function_call" + section.call_id = call_id + section.name = function_name + section.arguments = json.dumps(arguments) + section.model_dump.return_value = { + "type": "function_call", + "call_id": call_id, + "name": function_name, + "arguments": json.dumps(arguments), + } + mock_response.output = [section] + return mock_response + + +def _mock_text_response(text: str) -> MagicMock: + """Build a fake Responses-API response containing a message section.""" + mock_response = MagicMock() + mock_response.status = "completed" + mock_response.error = None + section = MagicMock() + section.type = "message" + section.content = [MagicMock(text=text)] + mock_response.output = [section] + return mock_response + + +def _user_msg(text: str, conversation_id: str | None = None) -> Message: + return Message( + message_pieces=[ + MessagePiece( + role="user", + original_value=text, + conversation_id=conversation_id or str(uuid.uuid4()), + ) + ] + ) + + +class TestCustomFunctionsDeprecation: + """custom_functions still works but emits DeprecationWarning.""" + + def test_custom_functions_kwarg_emits_deprecation_warning(self, patch_central_database): + async def get_weather(args: dict[str, Any]) -> dict[str, Any]: + return {"t": 72} + + with pytest.warns(DeprecationWarning, match="custom_functions"): + OpenAIResponseTarget( + model_name="gpt-4", + endpoint="https://mock.example.com", + api_key="mock-key", + custom_functions={"get_weather": get_weather}, + ) + + @pytest.mark.asyncio + async def test_custom_functions_kwarg_still_dispatches(self, patch_central_database): + async def get_weather(args: dict[str, Any]) -> dict[str, Any]: + return {"temperature": 72, "condition": "sunny"} + + with warnings.catch_warnings(): + warnings.simplefilter("ignore", DeprecationWarning) + target = OpenAIResponseTarget( + model_name="gpt-4", + endpoint="https://mock.example.com", + api_key="mock-key", + custom_functions={"get_weather": get_weather}, + ) + + responses = [ + _mock_function_call_response("call_1", "get_weather", {"location": "NYC"}), + _mock_text_response("72F and sunny."), + ] + seen = [] + + async def mock_create(**kwargs): + seen.append(kwargs) + return responses[len(seen) - 1] + + with patch.object(target._async_client.responses, "create", new_callable=AsyncMock) as mc: + mc.side_effect = mock_create + result = await target.send_prompt_async(message=_user_msg("weather?")) + + assert len(seen) == 2 + assert result[-1].message_pieces[0].original_value == "72F and sunny." + second_input = seen[1]["input"] + assert any(item.get("type") == "function_call_output" for item in second_input) + + +def _config_with_backend(backend: LocalToolBackend) -> TargetConfiguration: + """Build a TargetConfiguration wired for the modern tool-backend path.""" + caps = TargetCapabilities( + supports_multi_turn=True, + supports_multi_message_pieces=True, + supports_editable_history=True, + supports_json_output=True, + supports_system_prompt=True, + supports_tool_use=True, + input_modalities=frozenset( + { + frozenset(["text"]), + frozenset(["text", "image_path"]), + frozenset(["function_call"]), + frozenset(["tool_call"]), + frozenset(["function_call_output"]), + frozenset(["reasoning"]), + } + ), + ) + return TargetConfiguration( + capabilities=caps, + tool_event_policy=ToolEventPolicy(behavior=ToolEventBehavior.EXECUTE, max_tool_iterations=5), + tool_backend=backend, + ) + + +class TestToolBackendDispatch: + """The modern path: pass tool_backend via TargetConfiguration.""" + + @pytest.mark.asyncio + async def test_local_backend_dispatches_through_tool_loop(self, patch_central_database): + async def get_weather(args: dict[str, Any]) -> dict[str, Any]: + return {"temperature": 72, "condition": "sunny"} + + backend = LocalToolBackend( + callables={"get_weather": get_weather}, + schemas=[ + { + "name": "get_weather", + "description": "Weather lookup.", + "parameters": { + "type": "object", + "properties": {"location": {"type": "string"}}, + "required": ["location"], + }, + } + ], + ) + target = OpenAIResponseTarget( + model_name="gpt-4", + endpoint="https://mock.example.com", + api_key="mock-key", + custom_configuration=_config_with_backend(backend), + ) + + responses = [ + _mock_function_call_response("call_1", "get_weather", {"location": "NYC"}), + _mock_text_response("72F and sunny in NYC."), + ] + seen = [] + + async def mock_create(**kwargs): + seen.append(kwargs) + return responses[len(seen) - 1] + + with patch.object(target._async_client.responses, "create", new_callable=AsyncMock) as mc: + mc.side_effect = mock_create + result = await target.send_prompt_async(message=_user_msg("weather?")) + + assert len(seen) == 2 + assert result[-1].message_pieces[0].original_value == "72F and sunny in NYC." + second_input = seen[1]["input"] + assert any(item.get("type") == "function_call_output" for item in second_input) + + +class TestToolSchemasInjection: + """_construct_request_body injects backend schemas when present.""" + + @pytest.mark.asyncio + async def test_backend_schemas_injected_into_tools(self, patch_central_database): + async def get_weather(args: dict[str, Any]) -> dict[str, Any]: + return {"t": 1} + + backend = LocalToolBackend( + callables={"get_weather": get_weather}, + schemas=[{"name": "get_weather", "description": "x", "parameters": {"type": "object"}}], + ) + target = OpenAIResponseTarget( + model_name="gpt-4", + endpoint="https://mock.example.com", + api_key="mock-key", + custom_configuration=_config_with_backend(backend), + ) + body = await target._construct_request_body( + conversation=[_user_msg("hi")], + json_config=MagicMock(enabled=False, schema=None), + ) + assert "tools" in body + assert body["tools"][0]["type"] == "function" + assert body["tools"][0]["name"] == "get_weather" + + @pytest.mark.asyncio + async def test_extra_body_tools_take_precedence(self, patch_central_database): + async def f(args: dict[str, Any]) -> dict[str, Any]: + return {} + + backend = LocalToolBackend( + callables={"f": f}, + schemas=[{"name": "f", "parameters": {"type": "object"}}], + ) + legacy = [{"type": "function", "name": "legacy_tool", "description": "x"}] + config = _config_with_backend(backend) + target = OpenAIResponseTarget( + model_name="gpt-4", + endpoint="https://mock.example.com", + api_key="mock-key", + extra_body_parameters={"tools": legacy}, + custom_configuration=config, + ) + body = await target._construct_request_body( + conversation=[_user_msg("hi")], + json_config=MagicMock(enabled=False, schema=None), + ) + assert body["tools"] == legacy + + @pytest.mark.asyncio + async def test_no_backend_no_tools_key(self, patch_central_database): + target = OpenAIResponseTarget( + model_name="gpt-4", + endpoint="https://mock.example.com", + api_key="mock-key", + ) + body = await target._construct_request_body( + conversation=[_user_msg("hi")], + json_config=MagicMock(enabled=False, schema=None), + ) + assert "tools" not in body + + +class TestNonFunctionCallPiecesPassThrough: + """Reasoning / mcp_call / web_search_call sections must NOT be dispatched. + + The Response target's parser populates pieces for these types so they can + be persisted to Memory and round-tripped on subsequent requests. The + CanonicalEnvelopeParser only extracts function_call pieces; the tool loop + must therefore see an empty parse and exit cleanly. + """ + + @pytest.mark.asyncio + async def test_reasoning_only_response_exits_loop(self, patch_central_database): + target = OpenAIResponseTarget( + model_name="gpt-4", + endpoint="https://mock.example.com", + api_key="mock-key", + reasoning_effort="medium", + ) + # Reasoning section + final text section in one response + mock_response = MagicMock() + mock_response.status = "completed" + mock_response.error = None + reasoning_section = MagicMock() + reasoning_section.type = "reasoning" + reasoning_section.model_dump.return_value = {"type": "reasoning", "summary": "thinking..."} + text_section = MagicMock() + text_section.type = "message" + text_section.content = [MagicMock(text="The answer is 42.")] + mock_response.output = [reasoning_section, text_section] + + seen = [] + + async def mock_create(**kwargs): + seen.append(kwargs) + return mock_response + + with patch.object(target._async_client.responses, "create", new_callable=AsyncMock) as mc: + mc.side_effect = mock_create + result = await target.send_prompt_async(message=_user_msg("question?")) + + # Exactly one API call -- reasoning is not a tool call so the loop exits + assert len(seen) == 1 + # Response message contains both pieces + assert len(result) == 1 + piece_types = [p.original_value_data_type for p in result[0].message_pieces] + assert "reasoning" in piece_types + assert "text" in piece_types From e57ff8f3be87836cb8670573d7f8cb0e17928ec0 Mon Sep 17 00:00:00 2001 From: Victor Valbuena Date: Thu, 28 May 2026 10:37:42 -0700 Subject: [PATCH 28/40] DOCS: Migrate reST roles, drop dead spec-selector and stale scorer-flex copy Combines three small post-review cleanups in one commit: * Migrate :class: / :func: / :meth: reST roles to plain double-backtick code spans across files touched by this PR (adversarial.py, the two test modules, and the captured scenario-description output in the benchmark scanner notebook). Matches Rich's #1782 style-guide rule and the new docstring-style entry in .github/instructions/. * Rewrite the VERSION inline doc to describe the actual bump trigger (param + atomic_attack_name format change) and cache-scoping semantics (v1 results stay queryable but don't suppress v2 runs). * Simplify the objective_scorer Args block to reflect c7e7b93f (guard removed; annotation narrowed to TrueFalseScorer | None) and delete the now-lying Raises: TypeError section. * Inline _select_adversarial_specs at its single call site as a 2-line dict lookup and delete the method. The defensive double-filter (non-adversarial + drift warnings) is dead by construction since BenchmarkStrategy is built from adversarial-capable SCENARIO_TECHNIQUES entries only. Drop the corresponding TestSelectAdversarialSpecs class. * Delete TestAdversarialBenchmarkScorerFlexibility - its sole remaining test was a duplicate of test_construct_with_explicit_objective_scorer after the c7e7b93f guard removal. 43 unit tests pass (was 46 - -3 for the deletions). No source behavior change. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- doc/scanner/benchmark.ipynb | 8 +- .../scenarios/benchmark/adversarial.py | 98 ++++++------------- tests/end_to_end/test_scenarios.py | 12 +-- .../scenario/benchmark/test_adversarial.py | 65 +----------- 4 files changed, 42 insertions(+), 141 deletions(-) diff --git a/doc/scanner/benchmark.ipynb b/doc/scanner/benchmark.ipynb index 773311d199..843c06739e 100644 --- a/doc/scanner/benchmark.ipynb +++ b/doc/scanner/benchmark.ipynb @@ -151,14 +151,14 @@ "\u001b[36m • PyRIT Version: 0.14.0.dev0\u001b[0m\n", "\u001b[36m • Description:\u001b[0m\n", "\u001b[36m Benchmark scenario that compares the attack success rate (ASR) across adversarial models. Adversarial targets\u001b[0m\n", - "\u001b[36m are user-supplied via the ``adversarial_targets`` parameter (declared in :meth:`supported_parameters`). Each\u001b[0m\n", + "\u001b[36m are user-supplied via the ``adversarial_targets`` parameter (declared in ``supported_parameters``). Each\u001b[0m\n", "\u001b[36m target must already be registered in ``TargetRegistry`` — typically by ``TargetInitializer`` from\u001b[0m\n", "\u001b[36m ``ADVERSARIAL_CHAT_*`` env vars, or programmatically via ``TargetRegistry.register_instance``. At run time,\u001b[0m\n", - "\u001b[36m :meth:`_get_atomic_attacks_async` performs the ``(technique × adversarial_target × dataset)`` cross-product: for\u001b[0m\n", + "\u001b[36m ``_get_atomic_attacks_async`` performs the ``(technique × adversarial_target × dataset)`` cross-product: for\u001b[0m\n", "\u001b[36m each selected adversarial-capable technique in ``SCENARIO_TECHNIQUES`` and each requested target, it constructs\u001b[0m\n", - "\u001b[36m a per-pair :class:`AttackTechniqueFactory` via :meth:`AttackTechniqueRegistry.build_factory_from_spec` with\u001b[0m\n", + "\u001b[36m a per-pair ``AttackTechniqueFactory`` via ``AttackTechniqueRegistry.build_factory_from_spec`` with\u001b[0m\n", "\u001b[36m ``adversarial_chat`` overridden to that target — no global registry mutation. The resulting\u001b[0m\n", - "\u001b[36m :class:`AtomicAttack` is named ``f\"{technique}__{target}_{dataset}\"`` with ``display_group`` set to the target's\u001b[0m\n", + "\u001b[36m ``AtomicAttack`` is named ``f\"{technique}__{target}_{dataset}\"`` with ``display_group`` set to the target's\u001b[0m\n", "\u001b[36m registry name so per-model ASR rolls up naturally in result displays.\u001b[0m\n", "\n", "\u001b[1m 🎯 Target Information\u001b[0m\n", diff --git a/pyrit/scenario/scenarios/benchmark/adversarial.py b/pyrit/scenario/scenarios/benchmark/adversarial.py index 2e78ad4c9a..e1aaabe6be 100644 --- a/pyrit/scenario/scenarios/benchmark/adversarial.py +++ b/pyrit/scenario/scenarios/benchmark/adversarial.py @@ -33,9 +33,9 @@ def _build_benchmark_strategy() -> type[ScenarioStrategy]: Build the ``BenchmarkStrategy`` enum from ``SCENARIO_TECHNIQUES``. Filters the static technique catalog to entries that require an - adversarial chat target (per :func:`_spec_needs_adversarial`) and passes + adversarial chat target (per ``_spec_needs_adversarial``) and passes those source specs to - :meth:`AttackTechniqueRegistry.build_strategy_class_from_specs`. The + ``AttackTechniqueRegistry.build_strategy_class_from_specs``. The resulting enum has one concrete member per source technique (e.g. ``red_teaming``, ``tap``, ``crescendo_simulated``) plus the standard ``all`` / ``light`` / ``single_turn`` / ``multi_turn`` aggregates inherited @@ -43,7 +43,7 @@ def _build_benchmark_strategy() -> type[ScenarioStrategy]: The (technique × target) cross-product is no longer pre-materialized into enum members; per-target factories are built lazily in - :meth:`AdversarialBenchmark._get_atomic_attacks_async` from the + ``AdversarialBenchmark._get_atomic_attacks_async`` from the user-supplied ``adversarial_targets`` parameter. Returns: @@ -67,30 +67,30 @@ class AdversarialBenchmark(Scenario): Benchmark scenario that compares the attack success rate (ASR) across adversarial models. Adversarial targets are user-supplied via the ``adversarial_targets`` - parameter (declared in :meth:`supported_parameters`). Each target must + parameter (declared in ``supported_parameters``). Each target must already be registered in ``TargetRegistry`` — typically by ``TargetInitializer`` from ``ADVERSARIAL_CHAT_*`` env vars, or programmatically via ``TargetRegistry.register_instance``. - At run time, :meth:`_get_atomic_attacks_async` performs the + At run time, ``_get_atomic_attacks_async`` performs the ``(technique × adversarial_target × dataset)`` cross-product: for each selected adversarial-capable technique in ``SCENARIO_TECHNIQUES`` and each requested target, it constructs a per-pair - :class:`AttackTechniqueFactory` via - :meth:`AttackTechniqueRegistry.build_factory_from_spec` with + ``AttackTechniqueFactory`` via + ``AttackTechniqueRegistry.build_factory_from_spec`` with ``adversarial_chat`` overridden to that target — no global registry - mutation. The resulting :class:`AtomicAttack` is named + mutation. The resulting ``AtomicAttack`` is named ``f"{technique}__{target}_{dataset}"`` with ``display_group`` set to the target's registry name so per-model ASR rolls up naturally in result displays. """ - #: Bumped from 1 to match the ``atomic_attack_name`` format introduced - #: when the scenario stopped passing target labels through its old triple - #: ``f"{technique}__{model}__{dataset}"`` shape. The post-collapse format - #: ``f"{technique}__{target}_{dataset}"`` is preserved here so cached - #: results from the prior collapse-era runs remain matchable by - #: ``skip_cached``. + #: Bumped from 1 → 2 by the refactor that moved adversarial targets + #: from a constructor parameter to the ``adversarial_targets`` scenario + #: parameter and changed ``atomic_attack_name`` from + #: ``{technique}__{model}__{dataset}`` to ``{technique}__{target}_{dataset}``. + #: ``skip_cached`` only matches against prior runs at the current + #: ``VERSION``; v1 results remain queryable but won't suppress v2 runs. VERSION: int = 2 #: AdversarialBenchmark compares attack-success rates across adversarial models; a baseline @@ -144,7 +144,7 @@ def supported_parameters(cls) -> list[Parameter]: Declare the ``adversarial_targets`` parameter. The list is treated as required at run time: - :meth:`_get_atomic_attacks_async` raises ``ValueError`` if + ``_get_atomic_attacks_async`` raises ``ValueError`` if ``self.params["adversarial_targets"]`` is empty or missing. The scenario-side error (rather than a declaration-side default) lets the caller raise a domain-specific message that names the CLI flag, @@ -182,17 +182,12 @@ def __init__( Initialize the AdversarialBenchmark scenario. Args: - objective_scorer: Scorer for evaluating attack success. The - annotation is the broad ``Scorer`` base class for forward - compatibility with the planned non-``TrueFalseScorer`` - scoring follow-up (see PR description / follow-up issue - tracker), but the runtime contract is currently still - ``TrueFalseScorer``: any other ``Scorer`` subclass raises - ``TypeError`` at construction with a message pointing at - the follow-up. Defaults to the registered default objective + objective_scorer: ``TrueFalseScorer`` used to evaluate attack + success. Defaults to the registered default objective scorer (typically the composite refusal+scale scorer set - up by an initializer), which is always a - ``TrueFalseScorer``. + up by an initializer). Widening to general ``Scorer`` + support (covering ``FloatScaleScorer``, etc.) is tracked + as a follow-up. skip_cached: When ``True``, ``_get_atomic_attacks_async`` filters out atomic attacks whose ``(atomic_attack_name, technique_eval_hash)`` tuple already appears in a prior @@ -205,12 +200,6 @@ def __init__( (e.g. different scorer) do not cross-pollinate. scenario_result_id: Optional ID of an existing scenario result to resume. - - Raises: - TypeError: If ``objective_scorer`` is a ``Scorer`` subclass - other than ``TrueFalseScorer`` (full non-true/false support - is tracked as a follow-up; the type annotation is widened - ahead of the runtime support). """ self._objective_scorer: TrueFalseScorer = ( objective_scorer if objective_scorer else self._get_default_objective_scorer() @@ -229,11 +218,11 @@ async def _get_atomic_attacks_async(self) -> list[AtomicAttack]: Build atomic attacks from (technique × adversarial_target × dataset), then apply caching. Reads the user-supplied ``adversarial_targets`` parameter, resolves - each name to a :class:`PromptTarget` via ``TargetRegistry``, and + each name to a ``PromptTarget`` via ``TargetRegistry``, and cross-products the selected adversarial-capable techniques over the resolved targets and configured datasets. Each pair builds a non-registered per-pair factory via - :meth:`AttackTechniqueRegistry.build_factory_from_spec` with + ``AttackTechniqueRegistry.build_factory_from_spec`` with ``adversarial_chat`` overridden to the resolved target — no global registry state is touched. When ``self._skip_cached`` is set, the final candidate list is then filtered against prior completed @@ -263,7 +252,13 @@ async def _get_atomic_attacks_async(self) -> list[AtomicAttack]: ) resolved_targets = self._resolve_adversarial_targets(target_names=target_names) - selected_specs = self._select_adversarial_specs() + # ``BenchmarkStrategy`` is built from adversarial-capable + # ``SCENARIO_TECHNIQUES`` entries only (see ``_build_benchmark_strategy``), + # so every selected strategy resolves to exactly one spec. Drift between the + # enum and the catalog is silently ignored — the next strategy-class build + # would surface it. + specs_by_name = {spec.name: spec for spec in SCENARIO_TECHNIQUES} + selected_specs = [specs_by_name[s.value] for s in self._scenario_strategies if s.value in specs_by_name] scoring_config = AttackScoringConfig(objective_scorer=self._objective_scorer) seed_groups_by_dataset = self._dataset_config.get_seed_attack_groups() @@ -371,39 +366,6 @@ def _resolve_adversarial_targets(self, *, target_names: list[str]) -> list[tuple return resolved - def _select_adversarial_specs(self) -> list: - """ - Resolve ``self._scenario_strategies`` back to adversarial-capable source specs. - - Strategies that are not adversarial-capable (i.e. don't satisfy - :func:`_spec_needs_adversarial`) are dropped with a warning. - Strategies whose name doesn't match any spec in - ``SCENARIO_TECHNIQUES`` are also dropped with a warning — this - guards against drift between the strategy enum and the technique - catalog. - - Returns: - list[AttackTechniqueSpec]: The adversarial-capable specs the - user selected, suitable for the per-pair factory build loop. - """ - specs_by_name = {spec.name: spec for spec in SCENARIO_TECHNIQUES} - selected_strategy_values = {s.value for s in self._scenario_strategies} - - selected_specs: list = [] - for value in selected_strategy_values: - spec = specs_by_name.get(value) - if spec is None: - logger.warning(f"AdversarialBenchmark: strategy '{value}' has no matching technique spec, skipping.") - continue - if not _spec_needs_adversarial(spec): - logger.warning( - f"AdversarialBenchmark: technique '{value}' does not require an adversarial chat target, " - "skipping (only adversarial-capable techniques are benchmarked)." - ) - continue - selected_specs.append(spec) - return selected_specs - def _collect_cached_completion_pairs(self) -> set[tuple[str, str | None]]: """ Collect cache keys for atomic attacks that completed in any prior run of this scenario. @@ -415,7 +377,7 @@ def _collect_cached_completion_pairs(self) -> set[tuple[str, str | None]]: ``(atomic_attack_name, parent_eval_hash)`` tuple for every ``SUCCESS`` or ``FAILURE`` outcome. The pair shape mirrors the ``(atomic_attack_name, technique_eval_hash)`` tuple used by - :meth:`_get_atomic_attacks_async` so a direct ``in`` check filters + ``_get_atomic_attacks_async`` so a direct ``in`` check filters candidates without further key construction. Resilient to attribution-data variation: rows whose diff --git a/tests/end_to_end/test_scenarios.py b/tests/end_to_end/test_scenarios.py index 0306607022..dcd37ecb82 100644 --- a/tests/end_to_end/test_scenarios.py +++ b/tests/end_to_end/test_scenarios.py @@ -6,9 +6,9 @@ These tests dynamically discover all available scenarios and run each one using the pyrit_scan command. Most scenarios run with the -:data:`DEFAULT_INITIALIZERS` list; scenarios that need additional setup -declare their full initializer list in :data:`SCENARIO_INITIALIZERS` and -extra CLI args in :data:`SCENARIO_EXTRA_ARGS`. +``DEFAULT_INITIALIZERS`` list; scenarios that need additional setup +declare their full initializer list in ``SCENARIO_INITIALIZERS`` and +extra CLI args in ``SCENARIO_EXTRA_ARGS``. Note: e2e tests are not part of CI; they run via ``make end-to-end-test`` on developer machines that have the appropriate env vars set @@ -27,18 +27,18 @@ CONFIG_FILE = Path(__file__).parent / "test_config.yaml" -#: Initializers run for every scenario unless overridden in :data:`SCENARIO_INITIALIZERS`. +#: Initializers run for every scenario unless overridden in ``SCENARIO_INITIALIZERS``. #: ``target`` populates ``TargetRegistry`` from env vars; ``load_default_datasets`` #: fetches each scenario's declared default datasets into memory. DEFAULT_INITIALIZERS: list[str] = ["target", "load_default_datasets"] #: Per-scenario override map for initializers. A scenario absent here falls back -#: to :data:`DEFAULT_INITIALIZERS`. Keys use the dotted registry name +#: to ``DEFAULT_INITIALIZERS``. Keys use the dotted registry name #: (``.``) returned by ``ScenarioRegistry.get_names()``. SCENARIO_INITIALIZERS: dict[str, list[str]] = {} #: Per-scenario extra CLI args appended after the standard flag block. Keys use -#: the same dotted registry name as :data:`SCENARIO_INITIALIZERS`. Values are +#: the same dotted registry name as ``SCENARIO_INITIALIZERS``. Values are #: lists already split into argv tokens. SCENARIO_EXTRA_ARGS: dict[str, list[str]] = { # benchmark.adversarial requires --adversarial-targets at run time diff --git a/tests/unit/scenario/benchmark/test_adversarial.py b/tests/unit/scenario/benchmark/test_adversarial.py index 7748bab1d2..0bc9fab22d 100644 --- a/tests/unit/scenario/benchmark/test_adversarial.py +++ b/tests/unit/scenario/benchmark/test_adversarial.py @@ -5,10 +5,10 @@ AdversarialBenchmark now owns its adversarial target axis directly via the ``adversarial_targets`` parameter declared in -:meth:`supported_parameters`. Targets are user-supplied registry names +``supported_parameters``. Targets are user-supplied registry names that resolve to ``PromptTarget`` instances via ``TargetRegistry``. The ``(technique × target × dataset)`` cross-product is built lazily inside -:meth:`_get_atomic_attacks_async` using per-pair non-registered factories; +``_get_atomic_attacks_async`` using per-pair non-registered factories; no global ``AttackTechniqueRegistry`` state is mutated. These tests cover the new contract: @@ -18,13 +18,11 @@ source ``light`` tag (excludes ``tap`` / ``crescendo_simulated``). * ``supported_parameters`` declares ``adversarial_targets: list[str]``. * ``_resolve_adversarial_targets`` raises with available names on typos. -* ``_select_adversarial_specs`` drops non-adversarial techniques. * ``_get_atomic_attacks_async`` produces ``N × M × D`` atomic attacks with the expected ``atomic_attack_name`` and ``display_group``. * ``_collect_cached_completion_pairs`` collects (name, hash) tuples for prior ``SUCCESS`` / ``FAILURE`` outcomes only. * ``skip_cached`` filters cached candidates end-to-end. -* Scorer flexibility stage 1: widened annotation + ``TypeError`` guard. """ from unittest.mock import MagicMock, patch @@ -285,49 +283,6 @@ def test_preserves_caller_order(self): assert names == ["adv_c", "adv_a", "adv_b"] -# --------------------------------------------------------------------------- -# _select_adversarial_specs -# --------------------------------------------------------------------------- - - -@pytest.mark.usefixtures("patch_central_database") -class TestSelectAdversarialSpecs: - """Tests for ``_select_adversarial_specs``: filter strategies down to adversarial-capable specs.""" - - def _make_bench(self) -> AdversarialBenchmark: - return AdversarialBenchmark(objective_scorer=MagicMock(spec=TrueFalseScorer)) - - def test_returns_only_adversarial_specs(self): - """``red_teaming`` is adversarial-capable; ``prompt_sending`` is not — only red_teaming survives.""" - bench = self._make_bench() - - red_teaming_strategy = MagicMock() - red_teaming_strategy.value = "red_teaming" - prompt_sending_strategy = MagicMock() - prompt_sending_strategy.value = "prompt_sending" - bench._scenario_strategies = [red_teaming_strategy, prompt_sending_strategy] - - selected = bench._select_adversarial_specs() - selected_names = {s.name for s in selected} - - assert "red_teaming" in selected_names - assert "prompt_sending" not in selected_names - - def test_unknown_strategy_value_is_skipped_with_warning(self, caplog): - """A strategy enum value with no matching spec is dropped (defensive guard against drift).""" - bench = self._make_bench() - - unknown_strategy = MagicMock() - unknown_strategy.value = "nonexistent_technique" - bench._scenario_strategies = [unknown_strategy] - - with caplog.at_level("WARNING"): - selected = bench._select_adversarial_specs() - - assert selected == [] - assert any("nonexistent_technique" in record.message for record in caplog.records) - - # --------------------------------------------------------------------------- # _get_atomic_attacks_async — validation and cross-product # --------------------------------------------------------------------------- @@ -732,19 +687,3 @@ async def test_skip_cached_true_keeps_unmatched_candidates(self): result = await bench._get_atomic_attacks_async() assert len(result) == 1 - - -# --------------------------------------------------------------------------- -# Scorer flexibility — stage 1 -# --------------------------------------------------------------------------- - - -@pytest.mark.usefixtures("patch_central_database") -class TestAdversarialBenchmarkScorerFlexibility: - """Tests for the widened ``objective_scorer`` annotation + ``isinstance`` guard (stage 1).""" - - def test_construct_accepts_truefalse_scorer_subclass(self): - """``TrueFalseScorer`` remains the runtime-supported type; should construct cleanly.""" - scorer = MagicMock(spec=TrueFalseScorer) - bench = AdversarialBenchmark(objective_scorer=scorer) - assert bench._objective_scorer is scorer From 0e622c7ff4f26357769e3edc14c231202d077c36 Mon Sep 17 00:00:00 2001 From: Victor Valbuena Date: Thu, 28 May 2026 10:42:20 -0700 Subject: [PATCH 29/40] Add integration tests for RedTeamingAttack with real MCP tool dispatch C7 adds end-to-end integration coverage of the @tool_loop decorator, MCPToolBackend, and MCPClient stack against the real echo_mcp_server subprocess. Only the OpenAI Responses HTTP layer is mocked; the MCP stdio subprocess, AsyncExitStack lifecycle, canonical envelope round-trip, and RedTeamingAttack execution path all run unmocked. # What ships tests/integration/tools/test_red_teaming_with_tools.py with three tests: 1. test_red_teaming_response_target_with_mcp_echo - end-to-end smoke test. RedTeamingAttack drives OpenAIResponseTarget configured with a MCPToolBackend pointing at echo_mcp_server. The Responses API mock returns one function_call followed by a stop response. Asserts the tool call actually reaches the MCP subprocess and the result lands back in the second API call as a function_call_output. 2. test_red_teaming_persists_canonical_transcript_in_memory - verifies the canonical envelope contract (plan section 13). Reads the conversation back from Memory after attack.execute_async returns and asserts the function_call and function_call_output pieces are present, in order, with matching call_ids. 3. test_red_teaming_dispatches_all_tool_calls_per_turn - regression test for the intentional behavior change from C6. The pre-C6 in-class loop in OpenAIResponseTarget only dispatched the LAST function_call per turn; the @tool_loop decorator now dispatches every call in declaration order. Issues both echo and add in one response and asserts both results land in the next API call. # Test infrastructure - LocalMCPServerSpec uses command=sys.executable + args=(echo_server,). - Mock objective scorer returns a true score so RedTeamingAttack exits cleanly after one turn. - Mock adversarial target returns a single scripted prompt wrapped as list[Message] (PromptTarget.send_prompt_async contract). - Score, ComponentIdentifier, and PromptTarget MagicMock(spec=...) usage matches the existing tests/unit/executor/attack patterns. All three integration tests pass; pre-commit clean. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- tests/integration/tools/__init__.py | 2 + .../tools/test_red_teaming_with_tools.py | 361 ++++++++++++++++++ 2 files changed, 363 insertions(+) create mode 100644 tests/integration/tools/__init__.py create mode 100644 tests/integration/tools/test_red_teaming_with_tools.py diff --git a/tests/integration/tools/__init__.py b/tests/integration/tools/__init__.py new file mode 100644 index 0000000000..9a0454564d --- /dev/null +++ b/tests/integration/tools/__init__.py @@ -0,0 +1,2 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT license. diff --git a/tests/integration/tools/test_red_teaming_with_tools.py b/tests/integration/tools/test_red_teaming_with_tools.py new file mode 100644 index 0000000000..9dca01371a --- /dev/null +++ b/tests/integration/tools/test_red_teaming_with_tools.py @@ -0,0 +1,361 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT license. + +"""C7 integration tests: RedTeamingAttack with real tool dispatch. + +These tests spawn the real ``tests/unit/tools/echo_mcp_server.py`` subprocess +and exercise the full client-side tool-calling stack: + + attack -> normalizer -> target -> @tool_loop wrapper -> MCPToolBackend -> + MCPClient (stdio) -> echo subprocess -> tool result -> back through the + wrapper -> Memory. + +Only the OpenAI Responses HTTP layer is mocked. The MCP subprocess, the +MCPToolBackend lock, the AsyncExitStack lifecycle, the canonical envelope +round-trip, and the @tool_loop decorator's RedTeam-attack invocation path +all execute under their real implementations. +""" + +from __future__ import annotations + +import json +import pathlib +import sys +import uuid +from unittest.mock import AsyncMock, MagicMock, patch + +import pytest + +from pyrit.executor.attack.core.attack_config import AttackAdversarialConfig, AttackScoringConfig +from pyrit.executor.attack.multi_turn.red_teaming import RedTeamingAttack +from pyrit.identifiers import ComponentIdentifier +from pyrit.memory import CentralMemory +from pyrit.models import Message, MessagePiece, Score +from pyrit.prompt_target import OpenAIResponseTarget +from pyrit.prompt_target.common.prompt_target import PromptTarget +from pyrit.prompt_target.common.target_capabilities import TargetCapabilities +from pyrit.prompt_target.common.target_configuration import TargetConfiguration +from pyrit.score.true_false.true_false_scorer import TrueFalseScorer +from pyrit.tools import ( + LocalMCPServerSpec, + MCPToolBackend, + ToolEventBehavior, + ToolEventPolicy, +) + + +def _mock_id(name: str) -> ComponentIdentifier: + return ComponentIdentifier(class_name=name, class_module="test") + + +ECHO_SERVER_PATH = pathlib.Path(__file__).resolve().parents[2] / "unit" / "tools" / "echo_mcp_server.py" + + +def _local_echo_spec() -> LocalMCPServerSpec: + """Build a LocalMCPServerSpec that launches the in-tree echo server.""" + return LocalMCPServerSpec( + command=sys.executable, + args=(str(ECHO_SERVER_PATH),), + ) + + +def _mock_function_call_response(call_id: str, function_name: str, arguments: dict) -> MagicMock: + """Build a fake Responses-API response containing a function_call section.""" + response = MagicMock() + response.status = "completed" + response.error = None + section = MagicMock() + section.type = "function_call" + section.call_id = call_id + section.name = function_name + section.arguments = json.dumps(arguments) + section.model_dump.return_value = { + "type": "function_call", + "call_id": call_id, + "name": function_name, + "arguments": json.dumps(arguments), + } + response.output = [section] + return response + + +def _mock_text_response(text: str) -> MagicMock: + """Build a fake Responses-API response containing a message section.""" + response = MagicMock() + response.status = "completed" + response.error = None + section = MagicMock() + section.type = "message" + section.content = [MagicMock(text=text)] + response.output = [section] + return response + + +def _make_response_target_with_mcp_backend( + backend: MCPToolBackend, +) -> OpenAIResponseTarget: + """Build an OpenAIResponseTarget wired to the live MCP backend.""" + caps = TargetCapabilities( + supports_multi_turn=True, + supports_editable_history=True, + supports_json_output=True, + supports_multi_message_pieces=True, + supports_system_prompt=True, + supports_tool_use=True, + input_modalities=frozenset( + { + frozenset(["text"]), + frozenset(["text", "image_path"]), + frozenset(["function_call"]), + frozenset(["tool_call"]), + frozenset(["function_call_output"]), + frozenset(["reasoning"]), + } + ), + ) + config = TargetConfiguration( + capabilities=caps, + tool_event_policy=ToolEventPolicy(behavior=ToolEventBehavior.EXECUTE, max_tool_iterations=5), + tool_backend=backend, + ) + return OpenAIResponseTarget( + model_name="gpt-4", + endpoint="https://mock.example.com", + api_key="mock-key", + custom_configuration=config, + ) + + +def _scripted_adversarial(prompts: list[str]) -> MagicMock: + """Build a mock adversarial target that returns scripted prompts.""" + adversarial = MagicMock(spec=PromptTarget) + adversarial.send_prompt_async = AsyncMock( + side_effect=[ + [ + Message( + message_pieces=[ + MessagePiece( + role="assistant", + original_value=p, + original_value_data_type="text", + conversation_id=str(uuid.uuid4()), + ) + ] + ) + ] + for p in prompts + ] + ) + adversarial.get_identifier.return_value = _mock_id("MockAdversarial") + adversarial.set_system_prompt = MagicMock() + return adversarial + + +def _success_scorer() -> MagicMock: + """Mock objective scorer that always returns True (objective met).""" + scorer = MagicMock(spec=TrueFalseScorer) + scorer.score_async = AsyncMock( + return_value=[ + Score( + score_value="true", + score_value_description="objective met", + score_type="true_false", + score_category=["test"], + score_rationale="mock rationale", + score_metadata={}, + message_piece_id=str(uuid.uuid4()), + scorer_class_identifier=_mock_id("MockScorer"), + ) + ] + ) + scorer.get_identifier.return_value = _mock_id("MockScorer") + return scorer + + +@pytest.mark.asyncio +async def test_red_teaming_response_target_with_mcp_echo(patch_central_database): + """End-to-end: RedTeamingAttack drives OpenAIResponseTarget with MCPToolBackend. + + The Response target's HTTP layer is mocked to return a function_call for + the echo tool, followed by a stop response after the tool result arrives. + The MCP subprocess actually executes the echo call. + """ + backend = MCPToolBackend(servers=[_local_echo_spec()]) + async with backend: + objective_target = _make_response_target_with_mcp_backend(backend) + + # Mock the OpenAI Responses HTTP layer on the objective target. + responses = [ + _mock_function_call_response("call_1", "echo", {"text": "hello"}), + _mock_text_response("Echoed: hello"), + ] + seen = [] + + async def mock_create(**kwargs): + seen.append(kwargs) + return responses[len(seen) - 1] + + # Adversarial returns one prompt (RedTeamingAttack stops after objective is met) + adversarial = _scripted_adversarial(["please echo hello"]) + + attack = RedTeamingAttack( + objective_target=objective_target, + attack_adversarial_config=AttackAdversarialConfig(target=adversarial), + attack_scoring_config=AttackScoringConfig(objective_scorer=_success_scorer()), + ) + + with patch.object( + objective_target._async_client.responses, "create", new_callable=AsyncMock + ) as mock_create_call: + mock_create_call.side_effect = mock_create + result = await attack.execute_async(objective="get the model to echo 'hello'") + + # Two HTTP calls to the Response API: initial + post-tool + assert len(seen) == 2 + # Second call must include the function_call_output (tool result) + second_input = seen[1]["input"] + function_outputs = [item for item in second_input if item.get("type") == "function_call_output"] + assert len(function_outputs) == 1 + # The output JSON contains the text "hello" because the real MCP echo + # subprocess returned it + assert "hello" in function_outputs[0]["output"] + assert result is not None + + +@pytest.mark.asyncio +async def test_red_teaming_persists_canonical_transcript_in_memory(patch_central_database): + """End-to-end: after a successful tool dispatch the DB shows the full chain. + + Verifies the canonical envelope contract (§13): the conversation written + to Memory must contain the user message, the assistant function_call, the + tool function_call_output (with matching call_id), and the assistant's + final text -- in that order. + """ + backend = MCPToolBackend(servers=[_local_echo_spec()]) + async with backend: + objective_target = _make_response_target_with_mcp_backend(backend) + + responses = [ + _mock_function_call_response("call_xyz", "echo", {"text": "world"}), + _mock_text_response("Echoed: world"), + ] + seen = [] + + async def mock_create(**kwargs): + seen.append(kwargs) + return responses[len(seen) - 1] + + adversarial = _scripted_adversarial(["echo world"]) + + attack = RedTeamingAttack( + objective_target=objective_target, + attack_adversarial_config=AttackAdversarialConfig(target=adversarial), + attack_scoring_config=AttackScoringConfig(objective_scorer=_success_scorer()), + ) + + with patch.object( + objective_target._async_client.responses, "create", new_callable=AsyncMock + ) as mock_create_call: + mock_create_call.side_effect = mock_create + result = await attack.execute_async(objective="echo world") + + # Read the conversation back from Memory + memory = CentralMemory.get_memory_instance() + assert result is not None + objective_conv_id = result.conversation_id + assert objective_conv_id, "Attack result must carry the objective-target conversation id" + + pieces = list(memory.get_message_pieces(conversation_id=objective_conv_id)) + # Filter out system prompts; we care about the user/assistant/tool chain + data_types_in_order = [p.original_value_data_type for p in pieces] + # The chain MUST contain function_call followed by function_call_output (canonical envelope) + assert "function_call" in data_types_in_order + assert "function_call_output" in data_types_in_order + + fc_index = data_types_in_order.index("function_call") + fco_index = data_types_in_order.index("function_call_output") + assert fc_index < fco_index, "function_call must precede function_call_output in DB" + + fc_envelope = json.loads(pieces[fc_index].original_value) + fco_envelope = json.loads(pieces[fco_index].original_value) + assert fc_envelope["call_id"] == fco_envelope["call_id"] == "call_xyz" + assert fc_envelope["name"] == "echo" + # The tool result envelope's `output` is JSON-encoded; the underlying echo result is "world" + assert "world" in fco_envelope["output"] + + +@pytest.mark.asyncio +async def test_red_teaming_dispatches_all_tool_calls_per_turn(patch_central_database): + """Multi-call-per-turn dispatch (intentional behavior change vs pre-C6 loop). + + When the model emits two function_call sections in one response, BOTH + must dispatch through the MCPToolBackend. The pre-C6 in-class loop in + OpenAIResponseTarget only dispatched the LAST call per turn; the C6 + migration onto @tool_loop changes this to "dispatch every call in + declaration order." Verify by issuing both an `echo` and an `add` call + and asserting both results land in the second API call's input. + """ + backend = MCPToolBackend(servers=[_local_echo_spec()]) + async with backend: + objective_target = _make_response_target_with_mcp_backend(backend) + + # First response contains TWO function_calls; second is the stop text. + multi_call_response = MagicMock() + multi_call_response.status = "completed" + multi_call_response.error = None + + echo_section = MagicMock() + echo_section.type = "function_call" + echo_section.call_id = "call_echo" + echo_section.name = "echo" + echo_section.arguments = json.dumps({"text": "hi"}) + echo_section.model_dump.return_value = { + "type": "function_call", + "call_id": "call_echo", + "name": "echo", + "arguments": json.dumps({"text": "hi"}), + } + add_section = MagicMock() + add_section.type = "function_call" + add_section.call_id = "call_add" + add_section.name = "add" + add_section.arguments = json.dumps({"a": 3, "b": 4}) + add_section.model_dump.return_value = { + "type": "function_call", + "call_id": "call_add", + "name": "add", + "arguments": json.dumps({"a": 3, "b": 4}), + } + multi_call_response.output = [echo_section, add_section] + + responses = [ + multi_call_response, + _mock_text_response("done"), + ] + seen = [] + + async def mock_create(**kwargs): + seen.append(kwargs) + return responses[len(seen) - 1] + + adversarial = _scripted_adversarial(["call echo and add"]) + + attack = RedTeamingAttack( + objective_target=objective_target, + attack_adversarial_config=AttackAdversarialConfig(target=adversarial), + attack_scoring_config=AttackScoringConfig(objective_scorer=_success_scorer()), + ) + + with patch.object(objective_target._async_client.responses, "create", new_callable=AsyncMock) as mc: + mc.side_effect = mock_create + await attack.execute_async(objective="dispatch both tools") + + assert len(seen) == 2 + second_input = seen[1]["input"] + outputs = [item for item in second_input if item.get("type") == "function_call_output"] + assert len(outputs) == 2, "Both tool calls must be dispatched per the new behavior" + call_ids = [o["call_id"] for o in outputs] + assert call_ids == ["call_echo", "call_add"], "Outputs must preserve declaration order" + # Real MCP subprocess: echo("hi") returned "hi", add(3, 4) returned 7 + assert "hi" in outputs[0]["output"] + assert "7" in outputs[1]["output"] From fe0f2d3c05d0c83afaded6fae17772ec167dcad1 Mon Sep 17 00:00:00 2001 From: Victor Valbuena Date: Thu, 28 May 2026 10:49:50 -0700 Subject: [PATCH 30/40] REFACTOR: Delete unused `TargetInitializerTags.ADVERSARIAL` The tag was originally read by `BenchmarkInitializer.initialize_async` to fan out the benchmark across every adversarial-tagged target via `TargetRegistry.get_by_tag_query`. After the scenario refactor (5840b251), both that initializer and `get_by_tag_query` were deleted: `AdversarialBenchmark` now resolves adversarial targets by explicit name (`--adversarial-targets `), and no production code path consults the tag. What changed - `TargetInitializerTags`: drop `ADVERSARIAL = 'adversarial'` enum member. - `ENV_TARGET_CONFIGS`: drop `ADVERSARIAL` from the four `adversarial_chat[_variant]` configs (`DEFAULT` still propagates, so default-initializer registration is unchanged). - `test_register_target_propagates_config_tags`: pivot onto the only remaining multi-tagged config (`objective_scorer_chat` with `[DEFAULT, SCORER]`) so multi-tag propagation coverage is preserved. - `test_variant_registers_with_default_and_adversarial_tags` -> rename to `test_variant_registers_with_default_tag`. - Delete `test_all_variants_discoverable_via_adversarial_tag_query` (pure regression test for the deleted fan-out feature). - `test_double_initialize_async_is_idempotent`: count `DEFAULT` instead of `ADVERSARIAL` to keep the same regression guard on duplicate-registration. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- .../setup/initializers/components/targets.py | 5 -- tests/unit/setup/test_targets_initializer.py | 66 +++++-------------- 2 files changed, 17 insertions(+), 54 deletions(-) diff --git a/pyrit/setup/initializers/components/targets.py b/pyrit/setup/initializers/components/targets.py index aa1cae5c06..81cc363d2d 100644 --- a/pyrit/setup/initializers/components/targets.py +++ b/pyrit/setup/initializers/components/targets.py @@ -45,7 +45,6 @@ class TargetInitializerTags(str, Enum): SCORER = "scorer" ALL = "all" DEFAULT_OBJECTIVE_TARGET = "default_objective_target" - ADVERSARIAL = "adversarial" @dataclass @@ -187,7 +186,6 @@ class TargetConfig: model_var="ADVERSARIAL_CHAT_MODEL", underlying_model_var="ADVERSARIAL_CHAT_UNDERLYING_MODEL", temperature=1.2, - tags=[TargetInitializerTags.DEFAULT, TargetInitializerTags.ADVERSARIAL], ), TargetConfig( registry_name="adversarial_chat_singleturn", @@ -196,7 +194,6 @@ class TargetConfig: key_var="ADVERSARIAL_CHAT_SINGLETURN_KEY", model_var="ADVERSARIAL_CHAT_SINGLETURN_MODEL", temperature=1.2, - tags=[TargetInitializerTags.DEFAULT, TargetInitializerTags.ADVERSARIAL], ), TargetConfig( registry_name="adversarial_chat_multiturn", @@ -205,7 +202,6 @@ class TargetConfig: key_var="ADVERSARIAL_CHAT_MULTITURN_KEY", model_var="ADVERSARIAL_CHAT_MULTITURN_MODEL", temperature=1.2, - tags=[TargetInitializerTags.DEFAULT, TargetInitializerTags.ADVERSARIAL], ), TargetConfig( registry_name="adversarial_chat_reasoning", @@ -214,7 +210,6 @@ class TargetConfig: key_var="ADVERSARIAL_CHAT_REASONING_KEY", model_var="ADVERSARIAL_CHAT_REASONING_MODEL", temperature=1.2, - tags=[TargetInitializerTags.DEFAULT, TargetInitializerTags.ADVERSARIAL], ), TargetConfig( registry_name="objective_scorer_chat", diff --git a/tests/unit/setup/test_targets_initializer.py b/tests/unit/setup/test_targets_initializer.py index 57b04413fe..1885808072 100644 --- a/tests/unit/setup/test_targets_initializer.py +++ b/tests/unit/setup/test_targets_initializer.py @@ -432,9 +432,9 @@ def teardown_method(self) -> None: """Clean up after each test.""" TargetRegistry.reset_instance() for var in [ - "ADVERSARIAL_CHAT_ENDPOINT", - "ADVERSARIAL_CHAT_KEY", - "ADVERSARIAL_CHAT_MODEL", + "OBJECTIVE_SCORER_CHAT_ENDPOINT", + "OBJECTIVE_SCORER_CHAT_KEY", + "OBJECTIVE_SCORER_CHAT_MODEL", "OPENAI_CHAT_ENDPOINT", "OPENAI_CHAT_KEY", "OPENAI_CHAT_MODEL", @@ -448,24 +448,24 @@ async def test_register_target_propagates_config_tags(self) -> None: """ from pyrit.setup.initializers.components.targets import TargetInitializerTags - os.environ["ADVERSARIAL_CHAT_ENDPOINT"] = "https://test.openai.azure.com" - os.environ["ADVERSARIAL_CHAT_KEY"] = "test_key" - os.environ["ADVERSARIAL_CHAT_MODEL"] = "gpt-4o" + os.environ["OBJECTIVE_SCORER_CHAT_ENDPOINT"] = "https://test.openai.azure.com" + os.environ["OBJECTIVE_SCORER_CHAT_KEY"] = "test_key" + os.environ["OBJECTIVE_SCORER_CHAT_MODEL"] = "gpt-4o" init = TargetInitializer() await init.initialize_async() registry = TargetRegistry.get_registry_singleton() - assert "adversarial_chat" in registry + assert "objective_scorer_chat" in registry - adversarial_entries = registry.get_by_tag(tag=TargetInitializerTags.ADVERSARIAL) - assert any(entry.name == "adversarial_chat" for entry in adversarial_entries), ( - "adversarial_chat should be discoverable by the ADVERSARIAL tag after F1c" + scorer_entries = registry.get_by_tag(tag=TargetInitializerTags.SCORER) + assert any(entry.name == "objective_scorer_chat" for entry in scorer_entries), ( + "objective_scorer_chat should be discoverable by the SCORER tag after F1c" ) default_entries = registry.get_by_tag(tag=TargetInitializerTags.DEFAULT) - assert any(entry.name == "adversarial_chat" for entry in default_entries), ( - "adversarial_chat declares both DEFAULT and ADVERSARIAL tags; both must propagate" + assert any(entry.name == "objective_scorer_chat" for entry in default_entries), ( + "objective_scorer_chat declares both DEFAULT and SCORER tags; both must propagate" ) async def test_register_target_no_tags_in_config_no_extra_add_tags(self) -> None: @@ -554,10 +554,8 @@ def _set_variant_env_vars(prefix: str) -> None: os.environ[f"{prefix}_MODEL"] = "deployment-name" @pytest.mark.parametrize(("registry_name", "env_prefix"), ADVERSARIAL_CHAT_VARIANTS) - async def test_variant_registers_with_default_and_adversarial_tags( - self, registry_name: str, env_prefix: str - ) -> None: - """Each variant registers with ``[DEFAULT, ADVERSARIAL]`` tags when its env vars are set.""" + async def test_variant_registers_with_default_tag(self, registry_name: str, env_prefix: str) -> None: + """Each variant registers with the ``DEFAULT`` tag when its env vars are set.""" from pyrit.setup.initializers.components.targets import TargetInitializerTags self._set_variant_env_vars(env_prefix) @@ -568,9 +566,6 @@ async def test_variant_registers_with_default_and_adversarial_tags( registry = TargetRegistry.get_registry_singleton() assert registry_name in registry - adversarial_entries = registry.get_by_tag(tag=TargetInitializerTags.ADVERSARIAL) - assert any(entry.name == registry_name for entry in adversarial_entries) - default_entries = registry.get_by_tag(tag=TargetInitializerTags.DEFAULT) assert any(entry.name == registry_name for entry in default_entries) @@ -609,33 +604,6 @@ async def test_variant_skips_when_model_env_var_missing( os.environ.pop(f"{env_prefix}_ENDPOINT", None) os.environ.pop(f"{env_prefix}_KEY", None) - async def test_all_variants_discoverable_via_adversarial_tag_query(self) -> None: - """End-to-end: variants + ``adversarial_chat`` are returned by adversarial-tag ``get_by_tag``.""" - from pyrit.setup.initializers.components.targets import TargetInitializerTags - - os.environ["ADVERSARIAL_CHAT_ENDPOINT"] = "https://parent.openai.azure.com/openai/v1" - os.environ["ADVERSARIAL_CHAT_KEY"] = "test_key" - os.environ["ADVERSARIAL_CHAT_MODEL"] = "deployment-name" - - for _, prefix in ADVERSARIAL_CHAT_VARIANTS: - self._set_variant_env_vars(prefix) - - try: - init = TargetInitializer() - await init.initialize_async() - - registry = TargetRegistry.get_registry_singleton() - matches = registry.get_by_tag(tag=TargetInitializerTags.ADVERSARIAL.value) - match_names = {entry.name for entry in matches} - - expected = {"adversarial_chat"} | {name for name, _ in ADVERSARIAL_CHAT_VARIANTS} - assert expected <= match_names, ( - f"Missing variants from tag query result. Expected superset: {expected}, got: {match_names}" - ) - finally: - for var in ("ADVERSARIAL_CHAT_ENDPOINT", "ADVERSARIAL_CHAT_KEY", "ADVERSARIAL_CHAT_MODEL"): - os.environ.pop(var, None) - async def test_double_initialize_async_is_idempotent(self) -> None: """Re-running ``initialize_async`` with the same env state produces the same registry contents. @@ -655,11 +623,11 @@ async def test_double_initialize_async_is_idempotent(self) -> None: await init.initialize_async() registry = TargetRegistry.get_registry_singleton() first_names = sorted(registry.get_names()) - first_adversarial_count = len(registry.get_by_tag(tag=TargetInitializerTags.ADVERSARIAL)) + first_default_count = len(registry.get_by_tag(tag=TargetInitializerTags.DEFAULT)) await init.initialize_async() second_names = sorted(registry.get_names()) - second_adversarial_count = len(registry.get_by_tag(tag=TargetInitializerTags.ADVERSARIAL)) + second_default_count = len(registry.get_by_tag(tag=TargetInitializerTags.DEFAULT)) assert first_names == second_names - assert first_adversarial_count == second_adversarial_count + assert first_default_count == second_default_count From 610ed05d025ebe1e84bcc12dc2168f9288b095d9 Mon Sep 17 00:00:00 2001 From: Victor Valbuena Date: Thu, 28 May 2026 10:52:10 -0700 Subject: [PATCH 31/40] DOCS: Compare singleturn vs multiturn adversarial targets in benchmark notebook Demonstrates `--adversarial-targets` accepting multiple registry names by running the benchmark across both env-driven variants (`adversarial_chat_singleturn` and `adversarial_chat_multiturn`). Source cells only; captured outputs will refresh on the next execution. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- doc/scanner/benchmark.ipynb | 126 ++---------------------------------- doc/scanner/benchmark.py | 6 +- 2 files changed, 11 insertions(+), 121 deletions(-) diff --git a/doc/scanner/benchmark.ipynb b/doc/scanner/benchmark.ipynb index 843c06739e..39b4f2f8a0 100644 --- a/doc/scanner/benchmark.ipynb +++ b/doc/scanner/benchmark.ipynb @@ -33,7 +33,7 @@ "pyrit_scan benchmark.adversarial \\\n", " --initializers target load_default_datasets \\\n", " --target openai_chat \\\n", - " --adversarial-targets adversarial_chat \\\n", + " --adversarial-targets adversarial_chat_singleturn adversarial_chat_multiturn \\\n", " --max-dataset-size 4\n", "```\n", "\n", @@ -58,26 +58,7 @@ "execution_count": null, "id": "3", "metadata": {}, - "outputs": [ - { - "name": "stdout", - "output_type": "stream", - "text": [ - "Found default environment files: ['./.pyrit/.env', './.pyrit/.env.local']\n", - "Loaded environment file: ./.pyrit/.env\n", - "Loaded environment file: ./.pyrit/.env.local\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" - ] - } - ], + "outputs": [], "source": [ "from pyrit.output import output_scenario_async\n", "from pyrit.prompt_target import OpenAIChatTarget\n", @@ -99,27 +80,14 @@ "execution_count": null, "id": "4", "metadata": {}, - "outputs": [ - { - "data": { - "application/vnd.jupyter.widget-view+json": { - "model_id": "b25354cde64b48088458fde2e1d2beba", - "version_major": 2, - "version_minor": 0 - }, - "text/plain": [ - "Executing AdversarialBenchmark: 0%| | 0/3 [00:00 Date: Thu, 28 May 2026 15:28:54 -0700 Subject: [PATCH 32/40] FEAT: Add ObjectiveTargetEvaluationIdentifier + analytics cache lookup Adds two primitives that let scenarios query memory for prior AttackResults matching a given (technique x objective target) combination so cross-run caching can skip work that has already been done. 1. `ObjectiveTargetEvaluationIdentifier` (leaf-target subclass) computes a stable eval hash for an objective target by filtering its constructor params to the existing TARGET_EVAL_PARAMS / TARGET_EVAL_PARAM_FALLBACKS set (underlying_model_name, temperature, top_p; model_name fallback). To support this with the existing `EvaluationIdentifier` machinery, `compute_eval_hash` gains an optional `own_rule: ChildEvalRule` kwarg, and `EvaluationIdentifier` gains an `OWN_RULE: ClassVar[Optional[ChildEvalRule]]` that subclasses can set. The rule's `exclude` / `included_item_values` / `inner_child_name` fields are rejected at hash time because they are only meaningful for nested children. 2. `pyrit.analytics.result_analysis.get_cached_results_for_technique( memory_interface, *, technique_eval_hash, objective_target_eval_hash, additional_filters=None, ) -> list[AttackResult]` Pre-filters in SQL via an `IdentifierFilter` on `atomic_attack_identifier.eval_hash` (already stamped by `AtomicAttack._enrich_atomic_attack_identifiers`), then post-filters in Python by recomputing the objective target's eval hash off the persisted identifier tree. Returns results sorted newest-first. Analytics remains stateless: scenarios pass their own `CentralMemory.get_memory_instance()` at the call site. Tests cover the new OWN_RULE plumbing, `ObjectiveTargetEvaluationIdentifier`, and `get_cached_results_for_technique` (happy path, mismatched targets, missing identifier tree, sort order, SQL filter shape, additional filters). Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- pyrit/analytics/__init__.py | 7 +- pyrit/analytics/result_analysis.py | 100 +++++++- pyrit/identifiers/__init__.py | 2 + pyrit/identifiers/evaluation_identifier.py | 71 +++++- tests/unit/analytics/test_result_analysis.py | 231 +++++++++++++++++- .../identifiers/test_evaluation_identifier.py | 211 ++++++++++++++++ 6 files changed, 610 insertions(+), 12 deletions(-) diff --git a/pyrit/analytics/__init__.py b/pyrit/analytics/__init__.py index f75d401dd7..654923ce87 100644 --- a/pyrit/analytics/__init__.py +++ b/pyrit/analytics/__init__.py @@ -4,7 +4,11 @@ """Analytics module for PyRIT conversation and result analysis.""" from pyrit.analytics.conversation_analytics import ConversationAnalytics -from pyrit.analytics.result_analysis import AttackStats, analyze_results +from pyrit.analytics.result_analysis import ( + AttackStats, + analyze_results, + get_cached_results_for_technique, +) from pyrit.analytics.text_matching import ( ApproximateTextMatching, ExactTextMatching, @@ -17,5 +21,6 @@ "AttackStats", "ConversationAnalytics", "ExactTextMatching", + "get_cached_results_for_technique", "TextMatching", ] diff --git a/pyrit/analytics/result_analysis.py b/pyrit/analytics/result_analysis.py index d2e998ee94..c45e68c2fa 100644 --- a/pyrit/analytics/result_analysis.py +++ b/pyrit/analytics/result_analysis.py @@ -2,11 +2,20 @@ # Licensed under the MIT license. from collections import defaultdict +from collections.abc import Sequence from dataclasses import dataclass -from typing import Optional +from typing import TYPE_CHECKING, Optional +from pyrit.identifiers import ( + IdentifierFilter, + IdentifierType, + ObjectiveTargetEvaluationIdentifier, +) from pyrit.models import AttackOutcome, AttackResult +if TYPE_CHECKING: + from pyrit.memory.memory_interface import MemoryInterface + @dataclass class AttackStats: @@ -101,3 +110,92 @@ def analyze_results(attack_results: list[AttackResult]) -> dict[str, AttackStats "Overall": overall_stats, "By_attack_identifier": by_type_stats, } + + +def get_cached_results_for_technique( + memory_interface: "MemoryInterface", + *, + technique_eval_hash: str, + objective_target_eval_hash: str, + additional_filters: Optional[Sequence[IdentifierFilter]] = None, +) -> list[AttackResult]: + """ + Return cached AttackResults matching a (technique × objective target) pair. + + Memory is queried for AttackResults whose stamped + ``atomic_attack_identifier.eval_hash`` equals ``technique_eval_hash``, + then results are filtered in Python to those whose nested objective + target produces the requested ``objective_target_eval_hash`` (computed + via ``ObjectiveTargetEvaluationIdentifier``). Returned results are sorted + newest-first by ``timestamp`` so the most recent is at index 0. + + No scenario scoping is applied; this is a behavioral cache spanning every + run that produced the same (technique × target) combination. Callers that + need scenario-level scoping should pass additional ``IdentifierFilter``s + or filter the returned list themselves. + + Args: + memory_interface (MemoryInterface): The memory interface to query. + Analytics is stateless, so callers (e.g. scenarios) must pass + their own ``CentralMemory.get_memory_instance()``. + technique_eval_hash (str): Behavioral eval hash of the atomic-attack + technique, as produced by ``AtomicAttackEvaluationIdentifier.eval_hash`` + (also exposed as ``AtomicAttack.technique_eval_hash``). + objective_target_eval_hash (str): Behavioral eval hash of the objective + target, as produced by ``ObjectiveTargetEvaluationIdentifier.eval_hash``. + additional_filters (Optional[Sequence[IdentifierFilter]]): Extra + ``IdentifierFilter`` predicates appended to the SQL pre-filter. + Defaults to None. + + Returns: + list[AttackResult]: Matching attack results sorted newest-first. + Empty list if no cache hit. + """ + filters: list[IdentifierFilter] = [ + IdentifierFilter( + identifier_type=IdentifierType.ATTACK, + property_path="$.eval_hash", + value=technique_eval_hash, + ), + ] + if additional_filters: + filters.extend(additional_filters) + + candidates = memory_interface.get_attack_results(identifier_filters=filters) + + matches = [result for result in candidates if _objective_target_eval_hash_for(result) == objective_target_eval_hash] + + matches.sort(key=lambda r: r.timestamp, reverse=True) + return matches + + +def _objective_target_eval_hash_for(attack_result: AttackResult) -> Optional[str]: + """ + Return the ObjectiveTargetEvaluationIdentifier eval hash for a result. + + Walks ``atomic_attack_identifier.attack_technique.objective_target`` and + wraps the resulting identifier in ``ObjectiveTargetEvaluationIdentifier``. + + Args: + attack_result (AttackResult): The attack result whose persisted + ``atomic_attack_identifier`` tree should be inspected. + + Returns: + Optional[str]: The ``ObjectiveTargetEvaluationIdentifier.eval_hash`` + computed from the persisted objective-target identifier, or + ``None`` when the identifier tree is missing expected nodes + (e.g. legacy rows or atomic attacks without a distinct objective + target). + """ + if attack_result.atomic_attack_identifier is None: + return None + + technique = attack_result.atomic_attack_identifier.get_child("attack_technique") + if technique is None: + return None + + target = technique.get_child("objective_target") + if target is None: + return None + + return ObjectiveTargetEvaluationIdentifier(target).eval_hash diff --git a/pyrit/identifiers/__init__.py b/pyrit/identifiers/__init__.py index daa28292f8..369486192a 100644 --- a/pyrit/identifiers/__init__.py +++ b/pyrit/identifiers/__init__.py @@ -20,6 +20,7 @@ AtomicAttackEvaluationIdentifier, ChildEvalRule, EvaluationIdentifier, + ObjectiveTargetEvaluationIdentifier, ScorerEvaluationIdentifier, compute_eval_hash, ) @@ -35,6 +36,7 @@ "compute_eval_hash", "EvaluationIdentifier", "Identifiable", + "ObjectiveTargetEvaluationIdentifier", "REGISTRY_NAME_PATTERN", "ScorerEvaluationIdentifier", "snake_case_to_class_name", diff --git a/pyrit/identifiers/evaluation_identifier.py b/pyrit/identifiers/evaluation_identifier.py index 0171d68b2c..c157abb2a1 100644 --- a/pyrit/identifiers/evaluation_identifier.py +++ b/pyrit/identifiers/evaluation_identifier.py @@ -12,9 +12,12 @@ hash from a ``ComponentIdentifier``. * ``EvaluationIdentifier`` — abstract base that wraps a ``ComponentIdentifier`` with domain-specific eval-hash configuration. Concrete subclasses declare - per-child rules via a single ``CHILD_EVAL_RULES`` ClassVar. + per-child rules via ``CHILD_EVAL_RULES`` and (optionally) a root-level + ``OWN_RULE`` for leaf entities whose own params need filtering. * ``ScorerEvaluationIdentifier`` — scorer-domain concrete subclass. * ``AtomicAttackEvaluationIdentifier`` — attack-domain concrete subclass. +* ``ObjectiveTargetEvaluationIdentifier`` — leaf-target subclass used by the + analytics layer to key cached results by behavioral target configuration. """ from __future__ import annotations @@ -169,6 +172,7 @@ def compute_eval_hash( identifier: ComponentIdentifier, *, child_eval_rules: dict[str, ChildEvalRule], + own_rule: Optional[ChildEvalRule] = None, ) -> str: """ Compute a behavioral equivalence hash for evaluation grouping. @@ -176,26 +180,43 @@ def compute_eval_hash( Unlike ``ComponentIdentifier.hash`` (which includes all params of self and children), the eval hash applies per-child rules to strip operational params (like endpoint, max_requests_per_minute), exclude children entirely, or - filter list items. This ensures the same logical configuration on different - deployments produces the same eval hash. + filter list items. ``own_rule`` extends this to the root entity itself, + which is required for leaf components (e.g., a target) whose own params + need filtering and which have no relevant children to delegate to. This + ensures the same logical configuration on different deployments produces + the same eval hash. Children not listed in ``child_eval_rules`` receive full recursive treatment. - When ``child_eval_rules`` is empty, no filtering occurs and the result - equals ``identifier.hash``. + When both ``child_eval_rules`` is empty and ``own_rule`` is ``None``, no + filtering occurs and the result equals ``identifier.hash``. Args: identifier (ComponentIdentifier): The component identity to compute the hash for. child_eval_rules (dict[str, ChildEvalRule]): Per-child eval rules. + own_rule (Optional[ChildEvalRule]): Rule applied to the root entity's + own params and fallbacks. Only ``included_params`` and + ``param_fallbacks`` are honored; ``exclude``, ``included_item_values``, + and ``inner_child_name`` are not meaningful at the root and will + raise ``ValueError`` if set. Defaults to None. Returns: str: A hex-encoded SHA256 hash suitable for eval registry keying. Raises: - RuntimeError: If the identifier's hash is None and child_eval_rules is empty. + RuntimeError: If the identifier's hash is None and no filtering is configured. + ValueError: If ``own_rule`` carries fields that are not meaningful at the root. """ - if not child_eval_rules: + if own_rule is not None: + if own_rule.exclude: + raise ValueError("own_rule.exclude is not meaningful at the root entity") + if own_rule.included_item_values is not None: + raise ValueError("own_rule.included_item_values is not meaningful at the root entity") + if own_rule.inner_child_name is not None: + raise ValueError("own_rule.inner_child_name is not meaningful at the root entity") + + if not child_eval_rules and own_rule is None: if identifier.hash is None: raise RuntimeError("hash should be set by __post_init__") return identifier.hash @@ -203,6 +224,8 @@ def compute_eval_hash( eval_dict = _build_eval_dict( identifier, child_eval_rules=child_eval_rules, + _included_params=own_rule.included_params if own_rule else None, + _param_fallbacks=own_rule.param_fallbacks if own_rule else None, ) return config_hash(eval_dict) @@ -215,11 +238,16 @@ class EvaluationIdentifier(ABC): ``ChildEvalRule`` instances that control how each child is treated during eval-hash computation. Children not listed receive full recursive treatment. + Leaf-entity subclasses (no relevant children to delegate to) may also set + ``OWN_RULE`` to filter the root entity's own params. See + ``ObjectiveTargetEvaluationIdentifier`` for an example. + The concrete ``eval_hash`` property delegates to the module-level ``compute_eval_hash`` free function. """ CHILD_EVAL_RULES: ClassVar[dict[str, ChildEvalRule]] + OWN_RULE: ClassVar[Optional[ChildEvalRule]] = None def __init__(self, identifier: ComponentIdentifier) -> None: """ @@ -228,7 +256,8 @@ def __init__(self, identifier: ComponentIdentifier) -> None: If the identifier carries an ``eval_hash`` (preserved from a prior DB round-trip or set by the scorer), that value is used directly. Otherwise the eval hash is computed from the identifier's params - and children using the subclass's ``CHILD_EVAL_RULES``. + and children using the subclass's ``CHILD_EVAL_RULES`` and + ``OWN_RULE``. """ self._identifier = identifier if identifier.eval_hash is not None: @@ -237,6 +266,7 @@ def __init__(self, identifier: ComponentIdentifier) -> None: self._eval_hash = compute_eval_hash( identifier, child_eval_rules=self.CHILD_EVAL_RULES, + own_rule=self.OWN_RULE, ) @property @@ -301,3 +331,28 @@ class AtomicAttackEvaluationIdentifier(EvaluationIdentifier): # attack_technique: not listed in rules — fully included in eval hash. # technique_seeds (nested inside attack_technique): also not listed — fully included. } + + +class ObjectiveTargetEvaluationIdentifier(EvaluationIdentifier): + """ + Evaluation identity for an objective target. + + Mirrors how ``ScorerEvaluationIdentifier`` filters its inner + ``prompt_target`` child, except the target itself is the root of this + identifier (it has no children carrying behavioral configuration). The + target's own params are filtered to the behavioral set + (``underlying_model_name``, ``temperature``, ``top_p``) via ``OWN_RULE``, + so the same logical target on different deployments produces the same + eval hash. + + Wrapper targets (e.g., ``RoundRobinTarget``) are not unwrapped — the + caller must pass the inner target's ``ComponentIdentifier`` directly if + behavioral equivalence with the unwrapped form is desired. This mirrors + the constraint on ``OWN_RULE`` (no ``inner_child_name`` at the root). + """ + + CHILD_EVAL_RULES: ClassVar[dict[str, ChildEvalRule]] = {} + OWN_RULE: ClassVar[Optional[ChildEvalRule]] = ChildEvalRule( + included_params=TARGET_EVAL_PARAMS, + param_fallbacks=TARGET_EVAL_PARAM_FALLBACKS, + ) diff --git a/tests/unit/analytics/test_result_analysis.py b/tests/unit/analytics/test_result_analysis.py index e2d96b5bd4..253a98484f 100644 --- a/tests/unit/analytics/test_result_analysis.py +++ b/tests/unit/analytics/test_result_analysis.py @@ -1,12 +1,25 @@ # Copyright (c) Microsoft Corporation. # Licensed under the MIT license. +from datetime import datetime, timedelta, timezone from typing import Optional +from unittest.mock import MagicMock import pytest -from pyrit.analytics.result_analysis import AttackStats, analyze_results -from pyrit.identifiers import ComponentIdentifier +from pyrit.analytics.result_analysis import ( + AttackStats, + _objective_target_eval_hash_for, + analyze_results, + get_cached_results_for_technique, +) +from pyrit.identifiers import ( + ComponentIdentifier, + IdentifierFilter, + IdentifierType, + ObjectiveTargetEvaluationIdentifier, +) +from pyrit.memory.memory_interface import MemoryInterface from pyrit.models import AttackOutcome, AttackResult @@ -152,3 +165,217 @@ def test_group_by_attack_type_parametrized(items, type_key, exp_succ, exp_fail, assert stats.undetermined == exp_und assert stats.total_decided == exp_succ + exp_fail assert stats.success_rate == exp_rate + + +# --------------------------------------------------------------------------- +# get_cached_results_for_technique tests +# --------------------------------------------------------------------------- + + +def _make_target_component(*, model_name: str = "gpt-4o", temperature: float = 0.7) -> ComponentIdentifier: + return ComponentIdentifier( + class_name="OpenAIChatTarget", + class_module="pyrit.prompt_target.openai.openai_chat_target", + params={ + "underlying_model_name": model_name, + "temperature": temperature, + "top_p": 1.0, + "endpoint": "https://east.example.com", + }, + ) + + +def _make_attack_with_target( + target: ComponentIdentifier, + *, + outcome: AttackOutcome = AttackOutcome.SUCCESS, + timestamp: Optional[datetime] = None, +) -> AttackResult: + technique = ComponentIdentifier( + class_name="PromptSendingAttack", + class_module="pyrit.executor.attack.single_turn.prompt_sending", + children={"objective_target": target}, + ) + atomic = ComponentIdentifier( + class_name="AtomicAttack", + class_module="pyrit.scenario.core.atomic_attack", + children={"attack_technique": technique}, + ) + return AttackResult( + conversation_id="conv-1", + objective="test objective", + atomic_attack_identifier=atomic, + outcome=outcome, + timestamp=timestamp or datetime.now(timezone.utc), + ) + + +def test_get_cached_results_for_technique_returns_matching(): + target = _make_target_component() + expected_hash = ObjectiveTargetEvaluationIdentifier(target).eval_hash + matching = _make_attack_with_target(target) + + memory = MagicMock(spec=MemoryInterface) + memory.get_attack_results.return_value = [matching] + + results = get_cached_results_for_technique( + memory, + technique_eval_hash="tech-hash", + objective_target_eval_hash=expected_hash, + ) + + assert results == [matching] + + +def test_get_cached_results_for_technique_filters_out_target_mismatches(): + target_match = _make_target_component(model_name="gpt-4o") + target_other = _make_target_component(model_name="gpt-4o-mini") + expected_hash = ObjectiveTargetEvaluationIdentifier(target_match).eval_hash + + memory = MagicMock(spec=MemoryInterface) + memory.get_attack_results.return_value = [ + _make_attack_with_target(target_other), + _make_attack_with_target(target_match), + _make_attack_with_target(target_other), + ] + + results = get_cached_results_for_technique( + memory, + technique_eval_hash="tech-hash", + objective_target_eval_hash=expected_hash, + ) + + assert len(results) == 1 + assert results[0].atomic_attack_identifier == _make_attack_with_target(target_match).atomic_attack_identifier + + +def test_get_cached_results_for_technique_returns_empty_when_no_candidates(): + memory = MagicMock(spec=MemoryInterface) + memory.get_attack_results.return_value = [] + + results = get_cached_results_for_technique( + memory, + technique_eval_hash="tech-hash", + objective_target_eval_hash="target-hash", + ) + + assert results == [] + + +def test_get_cached_results_for_technique_sorts_newest_first(): + target = _make_target_component() + expected_hash = ObjectiveTargetEvaluationIdentifier(target).eval_hash + now = datetime.now(timezone.utc) + older = _make_attack_with_target(target, timestamp=now - timedelta(hours=2)) + middle = _make_attack_with_target(target, timestamp=now - timedelta(hours=1)) + newest = _make_attack_with_target(target, timestamp=now) + + memory = MagicMock(spec=MemoryInterface) + memory.get_attack_results.return_value = [older, newest, middle] + + results = get_cached_results_for_technique( + memory, + technique_eval_hash="tech-hash", + objective_target_eval_hash=expected_hash, + ) + + assert [r.timestamp for r in results] == [newest.timestamp, middle.timestamp, older.timestamp] + + +def test_get_cached_results_for_technique_builds_default_sql_filter(): + memory = MagicMock(spec=MemoryInterface) + memory.get_attack_results.return_value = [] + + get_cached_results_for_technique( + memory, + technique_eval_hash="tech-hash-xyz", + objective_target_eval_hash="target-hash", + ) + + memory.get_attack_results.assert_called_once() + filters = memory.get_attack_results.call_args.kwargs["identifier_filters"] + assert len(filters) == 1 + sole = filters[0] + assert sole.identifier_type == IdentifierType.ATTACK + assert sole.property_path == "$.eval_hash" + assert sole.value == "tech-hash-xyz" + + +def test_get_cached_results_for_technique_appends_additional_filters(): + memory = MagicMock(spec=MemoryInterface) + memory.get_attack_results.return_value = [] + extra = IdentifierFilter( + identifier_type=IdentifierType.ATTACK, + property_path="$.children.attack_technique.children.attack.class_name", + value="PromptSendingAttack", + ) + + get_cached_results_for_technique( + memory, + technique_eval_hash="tech-hash", + objective_target_eval_hash="target-hash", + additional_filters=[extra], + ) + + filters = memory.get_attack_results.call_args.kwargs["identifier_filters"] + assert len(filters) == 2 + assert filters[1] is extra + + +def test_get_cached_results_for_technique_skips_results_without_identifier(): + """Results with no atomic_attack_identifier are ignored, not raised on.""" + target = _make_target_component() + expected_hash = ObjectiveTargetEvaluationIdentifier(target).eval_hash + matching = _make_attack_with_target(target) + orphan = AttackResult( + conversation_id="orphan", + objective="o", + atomic_attack_identifier=None, + outcome=AttackOutcome.SUCCESS, + ) + + memory = MagicMock(spec=MemoryInterface) + memory.get_attack_results.return_value = [orphan, matching] + + results = get_cached_results_for_technique( + memory, + technique_eval_hash="tech-hash", + objective_target_eval_hash=expected_hash, + ) + + assert results == [matching] + + +def test_objective_target_eval_hash_for_missing_attack_technique_returns_none(): + """Helper returns None when the identifier tree is missing attack_technique.""" + atomic_only = ComponentIdentifier( + class_name="AtomicAttack", + class_module="pyrit.scenario.core.atomic_attack", + ) + result = AttackResult( + conversation_id="c", + objective="o", + atomic_attack_identifier=atomic_only, + outcome=AttackOutcome.SUCCESS, + ) + assert _objective_target_eval_hash_for(result) is None + + +def test_objective_target_eval_hash_for_missing_objective_target_returns_none(): + """Helper returns None when attack_technique has no objective_target child.""" + technique = ComponentIdentifier( + class_name="PromptSendingAttack", + class_module="pyrit.executor.attack.single_turn.prompt_sending", + ) + atomic = ComponentIdentifier( + class_name="AtomicAttack", + class_module="pyrit.scenario.core.atomic_attack", + children={"attack_technique": technique}, + ) + result = AttackResult( + conversation_id="c", + objective="o", + atomic_attack_identifier=atomic, + outcome=AttackOutcome.SUCCESS, + ) + assert _objective_target_eval_hash_for(result) is None diff --git a/tests/unit/identifiers/test_evaluation_identifier.py b/tests/unit/identifiers/test_evaluation_identifier.py index c716b1a331..0ceece3ac4 100644 --- a/tests/unit/identifiers/test_evaluation_identifier.py +++ b/tests/unit/identifiers/test_evaluation_identifier.py @@ -584,3 +584,214 @@ def test_scorer_eval_hash_matches_with_and_without_round_robin(self): eval_rr = ScorerEvaluationIdentifier(scorer_rr).eval_hash assert eval_direct == eval_rr + + +# --------------------------------------------------------------------------- +# OWN_RULE / leaf-entity eval-hash tests +# --------------------------------------------------------------------------- + + +class TestOwnRule: + """Tests for compute_eval_hash(own_rule=...) — leaf-entity filtering.""" + + def test_own_rule_filters_root_params(self): + """own_rule.included_params is applied to the root entity's params.""" + target = ComponentIdentifier( + class_name="OpenAIChatTarget", + class_module="pyrit.prompt_target.openai.openai_chat_target", + params={ + "underlying_model_name": "gpt-4o", + "temperature": 0.7, + "top_p": 1.0, + "endpoint": "https://east.example.com", + }, + ) + rule = ChildEvalRule( + included_params=frozenset({"underlying_model_name", "temperature", "top_p"}), + ) + + eval_hash = compute_eval_hash(target, child_eval_rules={}, own_rule=rule) + + # Same target body without endpoint should produce the same hash. + target_no_endpoint = ComponentIdentifier( + class_name="OpenAIChatTarget", + class_module="pyrit.prompt_target.openai.openai_chat_target", + params={ + "underlying_model_name": "gpt-4o", + "temperature": 0.7, + "top_p": 1.0, + }, + ) + eval_hash_no_endpoint = compute_eval_hash(target_no_endpoint, child_eval_rules={}, own_rule=rule) + assert eval_hash == eval_hash_no_endpoint + + def test_own_rule_applies_param_fallbacks_at_root(self): + """When the primary param is missing at the root, the fallback is substituted.""" + target_primary = ComponentIdentifier( + class_name="OpenAIChatTarget", + class_module="pyrit.prompt_target", + params={"underlying_model_name": "gpt-4o", "temperature": 0.7}, + ) + target_fallback = ComponentIdentifier( + class_name="OpenAIChatTarget", + class_module="pyrit.prompt_target", + params={"model_name": "gpt-4o", "temperature": 0.7}, + ) + rule = ChildEvalRule( + included_params=frozenset({"underlying_model_name", "temperature"}), + param_fallbacks={"underlying_model_name": "model_name"}, + ) + + hash_primary = compute_eval_hash(target_primary, child_eval_rules={}, own_rule=rule) + hash_fallback = compute_eval_hash(target_fallback, child_eval_rules={}, own_rule=rule) + assert hash_primary == hash_fallback + + def test_own_rule_raises_on_exclude(self): + """own_rule.exclude has no meaning at the root.""" + rule = ChildEvalRule(exclude=True) + target = ComponentIdentifier(class_name="T", class_module="m") + with pytest.raises(ValueError, match="exclude"): + compute_eval_hash(target, child_eval_rules={}, own_rule=rule) + + def test_own_rule_raises_on_included_item_values(self): + """own_rule.included_item_values is only meaningful for list children.""" + rule = ChildEvalRule(included_item_values={"is_general_technique": True}) + target = ComponentIdentifier(class_name="T", class_module="m") + with pytest.raises(ValueError, match="included_item_values"): + compute_eval_hash(target, child_eval_rules={}, own_rule=rule) + + def test_own_rule_raises_on_inner_child_name(self): + """own_rule.inner_child_name is only meaningful for child rules.""" + rule = ChildEvalRule(inner_child_name="targets") + target = ComponentIdentifier(class_name="T", class_module="m") + with pytest.raises(ValueError, match="inner_child_name"): + compute_eval_hash(target, child_eval_rules={}, own_rule=rule) + + def test_short_circuit_only_when_both_empty(self): + """With own_rule set, the short-circuit MUST NOT return identifier.hash.""" + target = ComponentIdentifier( + class_name="OpenAIChatTarget", + class_module="pyrit.prompt_target", + params={"underlying_model_name": "gpt-4o", "endpoint": "https://east.example.com"}, + ) + rule = ChildEvalRule(included_params=frozenset({"underlying_model_name"})) + eval_hash = compute_eval_hash(target, child_eval_rules={}, own_rule=rule) + # The full identifier hash includes the endpoint; eval_hash must not. + assert eval_hash != target.hash + + +class TestEvaluationIdentifierOwnRule: + """Tests for the EvaluationIdentifier.OWN_RULE ClassVar.""" + + def test_own_rule_defaults_to_none(self): + """Subclasses that do not declare OWN_RULE inherit None.""" + assert _StubEvaluationIdentifier.OWN_RULE is None + + def test_subclass_with_own_rule_filters_root(self): + """A subclass that sets OWN_RULE filters root params at eval time.""" + + class TargetIdentity(EvaluationIdentifier): + CHILD_EVAL_RULES: ClassVar[dict[str, ChildEvalRule]] = {} + OWN_RULE: ClassVar = ChildEvalRule( + included_params=frozenset({"underlying_model_name"}), + ) + + target = ComponentIdentifier( + class_name="OpenAIChatTarget", + class_module="pyrit.prompt_target", + params={"underlying_model_name": "gpt-4o", "endpoint": "https://east.example.com"}, + ) + identity = TargetIdentity(target) + # Eval hash should not equal the raw identifier hash (endpoint must be stripped). + assert identity.eval_hash != target.hash + + +# --------------------------------------------------------------------------- +# ObjectiveTargetEvaluationIdentifier tests +# --------------------------------------------------------------------------- + + +class TestObjectiveTargetEvaluationIdentifier: + """Tests for the ObjectiveTargetEvaluationIdentifier concrete subclass.""" + + def test_different_endpoints_same_eval_hash(self): + """Same model name + temperature + top_p on different endpoints → same eval hash.""" + from pyrit.identifiers.evaluation_identifier import ObjectiveTargetEvaluationIdentifier + + target_east = ComponentIdentifier( + class_name="OpenAIChatTarget", + class_module="pyrit.prompt_target.openai.openai_chat_target", + params={ + "underlying_model_name": "gpt-4o", + "temperature": 0.7, + "top_p": 1.0, + "endpoint": "https://east.example.com", + "model_name": "gpt4o-east", + }, + ) + target_west = ComponentIdentifier( + class_name="OpenAIChatTarget", + class_module="pyrit.prompt_target.openai.openai_chat_target", + params={ + "underlying_model_name": "gpt-4o", + "temperature": 0.7, + "top_p": 1.0, + "endpoint": "https://west.example.com", + "model_name": "gpt4o-west", + }, + ) + + eval_east = ObjectiveTargetEvaluationIdentifier(target_east).eval_hash + eval_west = ObjectiveTargetEvaluationIdentifier(target_west).eval_hash + assert eval_east == eval_west + + def test_different_temperature_different_eval_hash(self): + """Behavioral params (temperature) DO contribute to the eval hash.""" + from pyrit.identifiers.evaluation_identifier import ObjectiveTargetEvaluationIdentifier + + target_cold = ComponentIdentifier( + class_name="OpenAIChatTarget", + class_module="pyrit.prompt_target", + params={"underlying_model_name": "gpt-4o", "temperature": 0.0}, + ) + target_hot = ComponentIdentifier( + class_name="OpenAIChatTarget", + class_module="pyrit.prompt_target", + params={"underlying_model_name": "gpt-4o", "temperature": 1.0}, + ) + + eval_cold = ObjectiveTargetEvaluationIdentifier(target_cold).eval_hash + eval_hot = ObjectiveTargetEvaluationIdentifier(target_hot).eval_hash + assert eval_cold != eval_hot + + def test_model_name_fallback_to_model_name(self): + """When underlying_model_name is missing, model_name is used as fallback.""" + from pyrit.identifiers.evaluation_identifier import ObjectiveTargetEvaluationIdentifier + + target_underlying = ComponentIdentifier( + class_name="OpenAIChatTarget", + class_module="pyrit.prompt_target", + params={"underlying_model_name": "gpt-4o", "temperature": 0.7}, + ) + target_only_model_name = ComponentIdentifier( + class_name="OpenAIChatTarget", + class_module="pyrit.prompt_target", + params={"model_name": "gpt-4o", "temperature": 0.7}, + ) + + eval_a = ObjectiveTargetEvaluationIdentifier(target_underlying).eval_hash + eval_b = ObjectiveTargetEvaluationIdentifier(target_only_model_name).eval_hash + assert eval_a == eval_b + + def test_stored_eval_hash_takes_precedence(self): + """A pre-stamped eval_hash is honored (DB round-trip safety).""" + from pyrit.identifiers.evaluation_identifier import ObjectiveTargetEvaluationIdentifier + + stored = "objective_target_stored_hash" + "0" * 36 + cid = ComponentIdentifier( + class_name="OpenAIChatTarget", + class_module="pyrit.prompt_target", + params={"underlying_model_name": "gpt-4o"}, + ).with_eval_hash(stored) + + assert ObjectiveTargetEvaluationIdentifier(cid).eval_hash == stored From 94f1f9f9fdd1daa1ac886921d401034a35f11c4a Mon Sep 17 00:00:00 2001 From: Victor Valbuena Date: Thu, 28 May 2026 15:54:17 -0700 Subject: [PATCH 33/40] FEAT: Wire adversarial benchmark cache to analytics primitives MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Route the `skip_cached` filter on `AdversarialBenchmark` through the new `pyrit.analytics.get_cached_results_for_technique` helper and the `ObjectiveTargetEvaluationIdentifier` so the cache key is the content-derived `(technique_eval_hash × objective_target_eval_hash)` pair instead of the scenario-name + atomic-attack-name tuple. The scenario layer still owns the outcome filter — only SUCCESS / FAILURE matches suppress the candidate; ERROR / UNDETERMINED retries every run. `_collect_cached_completion_pairs` now: * returns `set[str]` of cached technique hashes (was a tuple set), * dedupes by `technique_eval_hash` across the incoming atomic attacks so we issue one analytics query per unique technique, * swallows per-hash analytics failures with a warning so a flaky cache lookup never wedges scenario startup, * short-circuits when `_objective_target_identifier` is not yet populated. The end-to-end filter at the tail of `_get_atomic_attacks_async` drops candidates whose `technique_eval_hash` is in the cached set. Tests in `tests/unit/scenario/benchmark/test_adversarial.py` are rewritten to mock the analytics symbol (and `ObjectiveTargetEvaluationIdentifier`) at the import site on the adversarial module rather than poking at `_memory.get_scenario_results` / `get_attack_results`. The module docstring is updated to reflect the new contract. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- .../scenarios/benchmark/adversarial.py | 137 ++++---- .../scenario/benchmark/test_adversarial.py | 326 +++++++++++------- 2 files changed, 270 insertions(+), 193 deletions(-) diff --git a/pyrit/scenario/scenarios/benchmark/adversarial.py b/pyrit/scenario/scenarios/benchmark/adversarial.py index e1aaabe6be..43a79a9d3c 100644 --- a/pyrit/scenario/scenarios/benchmark/adversarial.py +++ b/pyrit/scenario/scenarios/benchmark/adversarial.py @@ -9,8 +9,10 @@ import logging from typing import TYPE_CHECKING, ClassVar +from pyrit.analytics import get_cached_results_for_technique from pyrit.common import Parameter, apply_defaults from pyrit.executor.attack import AttackScoringConfig +from pyrit.identifiers import ObjectiveTargetEvaluationIdentifier from pyrit.models import AttackOutcome, SeedAttackGroup from pyrit.registry import AttackTechniqueRegistry, TargetRegistry from pyrit.registry.tag_query import TagQuery @@ -189,15 +191,16 @@ def __init__( support (covering ``FloatScaleScorer``, etc.) is tracked as a follow-up. skip_cached: When ``True``, ``_get_atomic_attacks_async`` filters - out atomic attacks whose ``(atomic_attack_name, - technique_eval_hash)`` tuple already appears in a prior - ``COMPLETED`` ``ScenarioResult`` for the same scenario name - and version with outcome ``SUCCESS`` or ``FAILURE``. - ``ERROR`` and ``UNDETERMINED`` outcomes always retry. Cache - identity is content-derived via - ``AtomicAttack.technique_eval_hash``, so two atomic attacks - with the same name but different technique configurations - (e.g. different scorer) do not cross-pollinate. + out atomic attacks for which the live behavioral cache + (``pyrit.analytics.get_cached_results_for_technique``) has + already returned at least one ``SUCCESS`` or ``FAILURE`` + ``AttackResult`` for the matching + ``(technique_eval_hash × objective_target_eval_hash)`` + pair. ``ERROR`` and ``UNDETERMINED`` outcomes never count + as cache hits. The cache spans every prior run that + produced the same (technique × objective target) + combination — it is intentionally not scoped to this + scenario name or ``VERSION``. scenario_result_id: Optional ID of an existing scenario result to resume. """ @@ -225,8 +228,10 @@ async def _get_atomic_attacks_async(self) -> list[AtomicAttack]: ``AttackTechniqueRegistry.build_factory_from_spec`` with ``adversarial_chat`` overridden to the resolved target — no global registry state is touched. When ``self._skip_cached`` is set, the - final candidate list is then filtered against prior completed - ``(atomic_attack_name, technique_eval_hash)`` tuples. + final candidate list is then filtered against the live behavioral + cache via ``_collect_cached_completion_pairs``, which delegates to + ``pyrit.analytics.get_cached_results_for_technique`` for each + unique ``(technique_eval_hash, objective_target_eval_hash)`` pair. Returns: list[AtomicAttack]: The atomic attacks to actually execute on @@ -319,12 +324,13 @@ async def _get_atomic_attacks_async(self) -> list[AtomicAttack]: if not self._skip_cached: return atomic_attacks - cached_pairs = self._collect_cached_completion_pairs() - filtered = [c for c in atomic_attacks if (c.atomic_attack_name, c.technique_eval_hash) not in cached_pairs] + cached_technique_hashes = self._collect_cached_completion_pairs(atomic_attacks=atomic_attacks) + filtered = [c for c in atomic_attacks if c.technique_eval_hash not in cached_technique_hashes] skipped = len(atomic_attacks) - len(filtered) if skipped > 0: logger.info( - "skip_cached=True: dropping %d/%d atomic attack(s) already completed in prior runs.", + "skip_cached=True: dropping %d/%d atomic attack(s) already completed for the " + "current objective target (matched by technique_eval_hash × objective_target_eval_hash).", skipped, len(atomic_attacks), ) @@ -366,68 +372,71 @@ def _resolve_adversarial_targets(self, *, target_names: list[str]) -> list[tuple return resolved - def _collect_cached_completion_pairs(self) -> set[tuple[str, str | None]]: + def _collect_cached_completion_pairs(self, *, atomic_attacks: list[AtomicAttack]) -> set[str]: """ - Collect cache keys for atomic attacks that completed in any prior run of this scenario. - - Walks ``ScenarioResult`` rows for the same scenario name and - ``VERSION``, restricts to ``scenario_run_state == "COMPLETED"``, - then walks the linked ``AttackResult`` rows (joined via - ``AttackResultEntry.attribution_parent_id``) and records the - ``(atomic_attack_name, parent_eval_hash)`` tuple for every - ``SUCCESS`` or ``FAILURE`` outcome. The pair shape mirrors the - ``(atomic_attack_name, technique_eval_hash)`` tuple used by - ``_get_atomic_attacks_async`` so a direct ``in`` check filters - candidates without further key construction. - - Resilient to attribution-data variation: rows whose - ``attribution_data`` is ``None`` or missing ``parent_collection`` - are skipped. Rows without ``parent_eval_hash`` enter the cache with - ``None`` in that slot, so they only match candidates whose - ``technique_eval_hash`` also resolves to ``None`` (currently never, - since ``AtomicAttack.technique_eval_hash`` is always populated post-#1758). + Return the set of ``technique_eval_hash`` values already cached for this scenario's objective target. + + Delegates to ``pyrit.analytics.get_cached_results_for_technique`` for + each unique technique hash among ``atomic_attacks``. A technique is + considered cached when the analytics helper returns at least one + ``AttackResult`` with outcome ``SUCCESS`` or ``FAILURE`` for the + ``(technique_eval_hash × objective_target_eval_hash)`` pair — + ``ERROR`` and ``UNDETERMINED`` outcomes are ignored so transient + failures retry on the next run. + + The objective-target eval hash is computed once from + ``self._objective_target_identifier`` (populated by the base + ``Scenario.initialize_async``) via + ``ObjectiveTargetEvaluationIdentifier``. The cache is intentionally + scenario-agnostic: any prior run that produced a matching (technique + × objective target) result counts as a hit, regardless of scenario + name or ``VERSION``. + + Args: + atomic_attacks: The candidate atomic attacks built earlier in + ``_get_atomic_attacks_async``. Only their + ``technique_eval_hash`` values are read. Returns: - set[tuple[str, str | None]]: Cache keys for already-completed - atomic attacks. Empty set on any unexpected error (logged at - warning level) — caching becomes a no-op rather than blocking - the run. + set[str]: ``technique_eval_hash`` values that have at least one + qualifying cached ``AttackResult``. Empty set when the scenario + has no objective target identifier or every analytics lookup + fails (logged at warning level) — caching becomes a no-op rather + than blocking the run. """ - scenario_name = type(self).__name__ - cached_pairs: set[tuple[str, str | None]] = set() + cached_hashes: set[str] = set() + + if self._objective_target_identifier is None: + return cached_hashes try: - prior_results = self._memory.get_scenario_results( - scenario_name=scenario_name, - scenario_version=self.VERSION, - ) + objective_target_eval_hash = ObjectiveTargetEvaluationIdentifier( + self._objective_target_identifier + ).eval_hash except Exception as exc: - logger.warning("skip_cached: failed to query prior scenario results (%s); skipping cache filter.", exc) - return cached_pairs + logger.warning( + "skip_cached: failed to compute objective_target eval hash (%s); skipping cache filter.", + exc, + ) + return cached_hashes - for scenario_result in prior_results: - if scenario_result.scenario_run_state != "COMPLETED": - continue - if scenario_result.id is None: - continue + unique_technique_hashes = {c.technique_eval_hash for c in atomic_attacks if c.technique_eval_hash} + + for technique_eval_hash in unique_technique_hashes: try: - attack_results = self._memory.get_attack_results(scenario_result_id=str(scenario_result.id)) + matches = get_cached_results_for_technique( + self._memory, + technique_eval_hash=technique_eval_hash, + objective_target_eval_hash=objective_target_eval_hash, + ) except Exception as exc: logger.warning( - "skip_cached: failed to load attack results for scenario %s (%s); skipping that run.", - scenario_result.id, + "skip_cached: analytics lookup failed for technique_eval_hash=%s (%s); not treating it as cached.", + technique_eval_hash, exc, ) continue + if any(m.outcome in (AttackOutcome.SUCCESS, AttackOutcome.FAILURE) for m in matches): + cached_hashes.add(technique_eval_hash) - for ar in attack_results: - if ar.outcome not in (AttackOutcome.SUCCESS, AttackOutcome.FAILURE): - continue - data = ar.attribution_data or {} - atomic_attack_name = data.get("parent_collection") - if not atomic_attack_name: - continue - parent_eval_hash = data.get("parent_eval_hash") - cached_pairs.add((atomic_attack_name, parent_eval_hash)) - - return cached_pairs + return cached_hashes diff --git a/tests/unit/scenario/benchmark/test_adversarial.py b/tests/unit/scenario/benchmark/test_adversarial.py index 0bc9fab22d..2aa9a9dfdb 100644 --- a/tests/unit/scenario/benchmark/test_adversarial.py +++ b/tests/unit/scenario/benchmark/test_adversarial.py @@ -20,8 +20,10 @@ * ``_resolve_adversarial_targets`` raises with available names on typos. * ``_get_atomic_attacks_async`` produces ``N × M × D`` atomic attacks with the expected ``atomic_attack_name`` and ``display_group``. -* ``_collect_cached_completion_pairs`` collects (name, hash) tuples for - prior ``SUCCESS`` / ``FAILURE`` outcomes only. +* ``_collect_cached_completion_pairs`` delegates to + ``pyrit.analytics.get_cached_results_for_technique`` per unique + technique hash and returns the set of technique hashes with at least + one ``SUCCESS`` / ``FAILURE`` match for the scenario's objective target. * ``skip_cached`` filters cached candidates end-to-end. """ @@ -466,135 +468,203 @@ async def test_factory_built_per_target_with_overridden_adversarial_chat(self): # --------------------------------------------------------------------------- -def _make_scenario_result(*, result_id: str, run_state: str = "COMPLETED") -> MagicMock: - """Build a minimal ScenarioResult stand-in for cache-key tests.""" - sr = MagicMock() - sr.id = result_id - sr.scenario_run_state = run_state - return sr +def _make_attack_result_with_outcome(outcome: AttackOutcome) -> MagicMock: + """Build a minimal ``AttackResult`` stand-in for cache-hit tests. - -def _make_attack_result( - *, - outcome: AttackOutcome, - parent_collection: str | None, - parent_eval_hash: str | None, -) -> MagicMock: - """Build a minimal AttackResult stand-in with the attribution_data shape the cache filter reads.""" + The new analytics-backed cache filter only reads ``outcome`` off each + match — the (technique × objective target) keying is done by the + analytics lookup parameters, not by introspecting result fields. + """ ar = MagicMock() ar.outcome = outcome - if parent_collection is None and parent_eval_hash is None: - ar.attribution_data = None - else: - data: dict[str, str] = {} - if parent_collection is not None: - data["parent_collection"] = parent_collection - if parent_eval_hash is not None: - data["parent_eval_hash"] = parent_eval_hash - ar.attribution_data = data return ar @pytest.mark.usefixtures("patch_central_database") class TestCollectCachedCompletionPairs: - """Tests for ``_collect_cached_completion_pairs`` (the cache key collector).""" + """Tests for ``_collect_cached_completion_pairs`` — now delegates to ``pyrit.analytics``.""" - def _make_bench(self) -> AdversarialBenchmark: + _ANALYTICS_PATH = "pyrit.scenario.scenarios.benchmark.adversarial.get_cached_results_for_technique" + _IDENTIFIER_PATH = "pyrit.scenario.scenarios.benchmark.adversarial.ObjectiveTargetEvaluationIdentifier" + + def _make_bench(self, *, with_target_identifier: bool = True) -> AdversarialBenchmark: bench = AdversarialBenchmark(objective_scorer=MagicMock(spec=TrueFalseScorer)) bench._memory = MagicMock() + bench._objective_target_identifier = MagicMock() if with_target_identifier else None return bench - def test_collects_success_and_failure_pairs(self): - bench = self._make_bench() - prior_sr = _make_scenario_result(result_id="sid-1") - prior_attacks = [ - _make_attack_result( - outcome=AttackOutcome.SUCCESS, - parent_collection="red_teaming__adv_a_harmbench", - parent_eval_hash="hash_a", - ), - _make_attack_result( - outcome=AttackOutcome.FAILURE, - parent_collection="tap__adv_a_harmbench", - parent_eval_hash="hash_b", - ), - ] - bench._memory.get_scenario_results.return_value = [prior_sr] - bench._memory.get_attack_results.return_value = prior_attacks + def _make_candidate(self, *, technique_eval_hash: str | None) -> MagicMock: + candidate = MagicMock() + candidate.technique_eval_hash = technique_eval_hash + return candidate + + def _patch_identifier(self, eval_hash: str = "obj_target_hash"): + """Patch ``ObjectiveTargetEvaluationIdentifier`` so we don't need a real ComponentIdentifier.""" + identifier_instance = MagicMock() + identifier_instance.eval_hash = eval_hash + return patch(self._IDENTIFIER_PATH, return_value=identifier_instance) + + def test_returns_empty_when_no_objective_target_identifier(self): + """Pre-``initialize_async`` state: no identifier means the cache filter is a no-op.""" + bench = self._make_bench(with_target_identifier=False) + candidates = [self._make_candidate(technique_eval_hash="hash_a")] - pairs = bench._collect_cached_completion_pairs() + with patch(self._ANALYTICS_PATH) as analytics_mock: + cached = bench._collect_cached_completion_pairs(atomic_attacks=candidates) + + assert cached == set() + analytics_mock.assert_not_called() + + def test_returns_empty_when_no_atomic_attacks(self): + """No candidates → no analytics calls and an empty result.""" + bench = self._make_bench() + with self._patch_identifier(), patch(self._ANALYTICS_PATH) as analytics_mock: + cached = bench._collect_cached_completion_pairs(atomic_attacks=[]) - assert pairs == { - ("red_teaming__adv_a_harmbench", "hash_a"), - ("tap__adv_a_harmbench", "hash_b"), - } + assert cached == set() + analytics_mock.assert_not_called() - def test_excludes_error_and_undetermined_outcomes(self): + def test_returns_hash_when_success_match_exists(self): bench = self._make_bench() - prior_sr = _make_scenario_result(result_id="sid-1") - prior_attacks = [ - _make_attack_result( - outcome=AttackOutcome.ERROR, - parent_collection="x", - parent_eval_hash="h", + candidates = [self._make_candidate(technique_eval_hash="hash_a")] + + with ( + self._patch_identifier(eval_hash="obj_hash"), + patch( + self._ANALYTICS_PATH, + return_value=[_make_attack_result_with_outcome(AttackOutcome.SUCCESS)], ), - _make_attack_result( - outcome=AttackOutcome.UNDETERMINED, - parent_collection="y", - parent_eval_hash="h", + ): + cached = bench._collect_cached_completion_pairs(atomic_attacks=candidates) + + assert cached == {"hash_a"} + + def test_returns_hash_when_failure_match_exists(self): + bench = self._make_bench() + candidates = [self._make_candidate(technique_eval_hash="hash_a")] + + with ( + self._patch_identifier(), + patch( + self._ANALYTICS_PATH, + return_value=[_make_attack_result_with_outcome(AttackOutcome.FAILURE)], ), - ] - bench._memory.get_scenario_results.return_value = [prior_sr] - bench._memory.get_attack_results.return_value = prior_attacks + ): + cached = bench._collect_cached_completion_pairs(atomic_attacks=candidates) - pairs = bench._collect_cached_completion_pairs() + assert cached == {"hash_a"} - assert pairs == set() + def test_excludes_hash_when_only_error_or_undetermined_matches(self): + """ERROR / UNDETERMINED outcomes must NOT count as cached so transient failures retry.""" + bench = self._make_bench() + candidates = [self._make_candidate(technique_eval_hash="hash_a")] + + with ( + self._patch_identifier(), + patch( + self._ANALYTICS_PATH, + return_value=[ + _make_attack_result_with_outcome(AttackOutcome.ERROR), + _make_attack_result_with_outcome(AttackOutcome.UNDETERMINED), + ], + ), + ): + cached = bench._collect_cached_completion_pairs(atomic_attacks=candidates) + + assert cached == set() - def test_only_counts_completed_scenario_runs(self): + def test_excludes_hash_when_no_matches(self): bench = self._make_bench() - in_progress = _make_scenario_result(result_id="sid-1", run_state="IN_PROGRESS") - failed = _make_scenario_result(result_id="sid-2", run_state="FAILED") - bench._memory.get_scenario_results.return_value = [in_progress, failed] + candidates = [self._make_candidate(technique_eval_hash="hash_a")] - pairs = bench._collect_cached_completion_pairs() + with self._patch_identifier(), patch(self._ANALYTICS_PATH, return_value=[]): + cached = bench._collect_cached_completion_pairs(atomic_attacks=candidates) - assert pairs == set() - # No COMPLETED runs → never touch get_attack_results. - bench._memory.get_attack_results.assert_not_called() + assert cached == set() - def test_queries_memory_by_scenario_name_and_version(self): + def test_dedupes_unique_technique_hashes_across_candidates(self): + """Three candidates sharing two unique hashes → analytics called twice, not three times.""" bench = self._make_bench() - bench._memory.get_scenario_results.return_value = [] - - bench._collect_cached_completion_pairs() + candidates = [ + self._make_candidate(technique_eval_hash="hash_a"), + self._make_candidate(technique_eval_hash="hash_b"), + self._make_candidate(technique_eval_hash="hash_a"), # duplicate + ] - bench._memory.get_scenario_results.assert_called_once_with( - scenario_name="AdversarialBenchmark", - scenario_version=AdversarialBenchmark.VERSION, + with ( + self._patch_identifier(eval_hash="obj_hash"), + patch( + self._ANALYTICS_PATH, + return_value=[_make_attack_result_with_outcome(AttackOutcome.SUCCESS)], + ) as analytics_mock, + ): + cached = bench._collect_cached_completion_pairs(atomic_attacks=candidates) + + assert cached == {"hash_a", "hash_b"} + assert analytics_mock.call_count == 2 + called_technique_hashes = {call.kwargs["technique_eval_hash"] for call in analytics_mock.call_args_list} + assert called_technique_hashes == {"hash_a", "hash_b"} + + def test_delegates_with_memory_and_objective_target_hash(self): + """Each analytics call passes the scenario's memory + the computed objective target hash.""" + bench = self._make_bench() + candidates = [self._make_candidate(technique_eval_hash="hash_a")] + + with ( + self._patch_identifier(eval_hash="my_obj_target_hash"), + patch(self._ANALYTICS_PATH, return_value=[]) as analytics_mock, + ): + bench._collect_cached_completion_pairs(atomic_attacks=candidates) + + analytics_mock.assert_called_once_with( + bench._memory, + technique_eval_hash="hash_a", + objective_target_eval_hash="my_obj_target_hash", ) - def test_skips_rows_with_missing_parent_collection(self): - """``attribution_data=None`` or missing ``parent_collection`` rows are silently skipped.""" + def test_skips_candidates_with_no_technique_eval_hash(self): + """A candidate whose ``technique_eval_hash`` is ``None`` is silently ignored.""" + bench = self._make_bench() + candidates = [self._make_candidate(technique_eval_hash=None)] + + with self._patch_identifier(), patch(self._ANALYTICS_PATH) as analytics_mock: + cached = bench._collect_cached_completion_pairs(atomic_attacks=candidates) + + assert cached == set() + analytics_mock.assert_not_called() + + def test_analytics_lookup_exception_is_swallowed_per_hash(self): + """A failing analytics lookup for one hash must not block the others — that hash is not cached.""" bench = self._make_bench() - prior_sr = _make_scenario_result(result_id="sid-1") - prior_attacks = [ - _make_attack_result(outcome=AttackOutcome.SUCCESS, parent_collection=None, parent_eval_hash=None), - _make_attack_result(outcome=AttackOutcome.SUCCESS, parent_collection=None, parent_eval_hash="hash_x"), + candidates = [ + self._make_candidate(technique_eval_hash="hash_a"), + self._make_candidate(technique_eval_hash="hash_b"), ] - bench._memory.get_scenario_results.return_value = [prior_sr] - bench._memory.get_attack_results.return_value = prior_attacks - pairs = bench._collect_cached_completion_pairs() - assert pairs == set() + def fake_analytics(_memory, *, technique_eval_hash, objective_target_eval_hash): + if technique_eval_hash == "hash_a": + raise RuntimeError("analytics blew up") + return [_make_attack_result_with_outcome(AttackOutcome.SUCCESS)] + + with self._patch_identifier(), patch(self._ANALYTICS_PATH, side_effect=fake_analytics): + cached = bench._collect_cached_completion_pairs(atomic_attacks=candidates) - def test_memory_error_falls_back_to_empty_set(self): - """An exception from ``get_scenario_results`` must not block the run; cache becomes a no-op.""" + # hash_a was the failed lookup → not cached (will retry). hash_b succeeded → cached. + assert cached == {"hash_b"} + + def test_identifier_construction_failure_falls_back_to_empty(self): + """If ``ObjectiveTargetEvaluationIdentifier`` raises, cache becomes a no-op rather than blocking.""" bench = self._make_bench() - bench._memory.get_scenario_results.side_effect = RuntimeError("db down") + candidates = [self._make_candidate(technique_eval_hash="hash_a")] + + with ( + patch(self._IDENTIFIER_PATH, side_effect=RuntimeError("bad identifier")), + patch(self._ANALYTICS_PATH) as analytics_mock, + ): + cached = bench._collect_cached_completion_pairs(atomic_attacks=candidates) - pairs = bench._collect_cached_completion_pairs() - assert pairs == set() + assert cached == set() + analytics_mock.assert_not_called() # --------------------------------------------------------------------------- @@ -606,6 +676,9 @@ def test_memory_error_falls_back_to_empty_set(self): class TestSkipCachedFilter: """End-to-end tests for the ``skip_cached`` filter applied in ``_get_atomic_attacks_async``.""" + _ANALYTICS_PATH = "pyrit.scenario.scenarios.benchmark.adversarial.get_cached_results_for_technique" + _IDENTIFIER_PATH = "pyrit.scenario.scenarios.benchmark.adversarial.ObjectiveTargetEvaluationIdentifier" + def _make_bench(self, *, skip_cached: bool) -> AdversarialBenchmark: _register_adversarial_target(name="adv_a") bench = AdversarialBenchmark( @@ -613,6 +686,7 @@ def _make_bench(self, *, skip_cached: bool) -> AdversarialBenchmark: skip_cached=skip_cached, ) bench._objective_target = MagicMock(spec=PromptTarget) + bench._objective_target_identifier = MagicMock() bench.params = {"adversarial_targets": ["adv_a"]} red_teaming_strategy = MagicMock() @@ -634,56 +708,50 @@ def _patch_factory_builder(self): return_value=factory, ) - async def test_skip_cached_false_returns_all_candidates(self): + def _patch_identifier(self, eval_hash: str = "obj_hash"): + identifier_instance = MagicMock() + identifier_instance.eval_hash = eval_hash + return patch(self._IDENTIFIER_PATH, return_value=identifier_instance) + + async def test_skip_cached_false_returns_all_candidates_without_analytics_call(self): bench = self._make_bench(skip_cached=False) - bench._memory = MagicMock() - with self._patch_factory_builder(): + with self._patch_factory_builder(), patch(self._ANALYTICS_PATH) as analytics_mock: result = await bench._get_atomic_attacks_async() assert len(result) == 1 - # No cache query when skip_cached=False. - bench._memory.get_scenario_results.assert_not_called() + analytics_mock.assert_not_called() async def test_skip_cached_true_filters_matching_candidates(self): bench = self._make_bench(skip_cached=True) - prior_sr = _make_scenario_result(result_id="sid-1") - prior_attacks = [ - _make_attack_result( - outcome=AttackOutcome.SUCCESS, - parent_collection="red_teaming__adv_a_harmbench", - parent_eval_hash=None, # MagicMock candidates yield None for technique_eval_hash - ), - ] - bench._memory = MagicMock() - bench._memory.get_scenario_results.return_value = [prior_sr] - bench._memory.get_attack_results.return_value = prior_attacks - with self._patch_factory_builder(): - # Stub out technique_eval_hash so the cache-key tuple matches. - with patch( + # Stub every candidate's technique_eval_hash to a known value so the analytics + # lookup key matches the cached set. + with ( + self._patch_factory_builder(), + self._patch_identifier(), + patch( "pyrit.scenario.core.atomic_attack.AtomicAttack.technique_eval_hash", - new_callable=lambda: property(lambda self: None), - ): - result = await bench._get_atomic_attacks_async() + new_callable=lambda: property(lambda self: "cached_hash"), + ), + patch( + self._ANALYTICS_PATH, + return_value=[_make_attack_result_with_outcome(AttackOutcome.SUCCESS)], + ), + ): + result = await bench._get_atomic_attacks_async() assert result == [] async def test_skip_cached_true_keeps_unmatched_candidates(self): bench = self._make_bench(skip_cached=True) - prior_sr = _make_scenario_result(result_id="sid-1") - prior_attacks = [ - _make_attack_result( - outcome=AttackOutcome.SUCCESS, - parent_collection="some_other_name", - parent_eval_hash="hash_x", - ), - ] - bench._memory = MagicMock() - bench._memory.get_scenario_results.return_value = [prior_sr] - bench._memory.get_attack_results.return_value = prior_attacks - with self._patch_factory_builder(): + # Analytics returns no matches → no candidate is cached, so all pass through. + with ( + self._patch_factory_builder(), + self._patch_identifier(), + patch(self._ANALYTICS_PATH, return_value=[]), + ): result = await bench._get_atomic_attacks_async() assert len(result) == 1 From 6827aa95c10e34d609f3143bb00969b1a2be50eb Mon Sep 17 00:00:00 2001 From: Victor Valbuena Date: Thu, 28 May 2026 16:03:34 -0700 Subject: [PATCH 34/40] TEST: Add real-memory coverage for adversarial cache wiring Adds TestCollectCachedCompletionPairsWithRealMemory, a unit-test class that exercises AdversarialBenchmark._collect_cached_completion_pairs end-to-end through real SQLiteMemory using the existing patch_central_database / sqlite_instance fixture pair. The pre-existing mocked tests cover scenario-layer wiring (delegation, dedup, outcome filter, identifier construction) but stub out analytics and the persistence layer. These new cases catch regressions in the real path: AttackResult persistence auto-stamping atomic_attack_identifier.eval_hash, the SQL filter on $.eval_hash, the python-side ObjectiveTargetEvaluationIdentifier filter inside get_cached_results_for_technique, and the outcome filter against AttackOutcome.SUCCESS / FAILURE. Seven cases cover cold cache, SUCCESS / FAILURE matches, error-only history, dedup, and the two reject paths (different objective target with matching technique hash; different technique hash from a different temperature). Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- .../scenario/benchmark/test_adversarial.py | 213 +++++++++++++++++- 1 file changed, 212 insertions(+), 1 deletion(-) diff --git a/tests/unit/scenario/benchmark/test_adversarial.py b/tests/unit/scenario/benchmark/test_adversarial.py index 2aa9a9dfdb..0342e6cbc5 100644 --- a/tests/unit/scenario/benchmark/test_adversarial.py +++ b/tests/unit/scenario/benchmark/test_adversarial.py @@ -25,13 +25,22 @@ technique hash and returns the set of technique hashes with at least one ``SUCCESS`` / ``FAILURE`` match for the scenario's objective target. * ``skip_cached`` filters cached candidates end-to-end. +* Real-memory smoke for ``_collect_cached_completion_pairs`` exercises + persistence -> SQL filter -> objective-target filter -> outcome filter. """ +from datetime import datetime, timezone from unittest.mock import MagicMock, patch import pytest -from pyrit.models import AttackOutcome, SeedAttackGroup, SeedObjective +from pyrit.identifiers import ( + AtomicAttackEvaluationIdentifier, + ComponentIdentifier, + ObjectiveTargetEvaluationIdentifier, +) +from pyrit.memory.memory_interface import MemoryInterface +from pyrit.models import AttackOutcome, AttackResult, SeedAttackGroup, SeedObjective from pyrit.prompt_target import PromptTarget from pyrit.registry import TargetRegistry from pyrit.registry.object_registries.attack_technique_registry import AttackTechniqueRegistry @@ -755,3 +764,205 @@ async def test_skip_cached_true_keeps_unmatched_candidates(self): result = await bench._get_atomic_attacks_async() assert len(result) == 1 + + +# --------------------------------------------------------------------------- +# Real-memory coverage for _collect_cached_completion_pairs +# --------------------------------------------------------------------------- +# +# The mocked TestCollectCachedCompletionPairs class above exercises the +# scenario-layer wiring (delegation, dedup, outcome filter, identifier +# construction). The tests in this section exercise the *full* path through +# real SQLite memory: AttackResult persistence (which auto-stamps +# ``atomic_attack_identifier.eval_hash``), the +# ``get_cached_results_for_technique`` SQL filter on ``$.eval_hash``, and +# the python-side ``ObjectiveTargetEvaluationIdentifier`` filter inside +# the analytics helper. They catch wiring regressions (e.g. a future +# refactor that stops stamping ``eval_hash`` at write time) that the +# mocked tests cannot. + + +def _make_objective_target_component( + *, + model_name: str = "gpt-4o", + temperature: float = 0.7, + top_p: float = 1.0, +) -> ComponentIdentifier: + return ComponentIdentifier( + class_name="OpenAIChatTarget", + class_module="pyrit.prompt_target.openai.openai_chat_target", + params={ + "underlying_model_name": model_name, + "temperature": temperature, + "top_p": top_p, + }, + ) + + +def _make_atomic_attack_identifier(target: ComponentIdentifier) -> ComponentIdentifier: + """Build the nested identifier tree the persistence layer expects.""" + technique = ComponentIdentifier( + class_name="PromptSendingAttack", + class_module="pyrit.executor.attack.single_turn.prompt_sending", + children={"objective_target": target}, + ) + return ComponentIdentifier( + class_name="AtomicAttack", + class_module="pyrit.scenario.core.atomic_attack", + children={"attack_technique": technique}, + ) + + +def _technique_eval_hash_for(target: ComponentIdentifier) -> str: + atomic = _make_atomic_attack_identifier(target) + return AtomicAttackEvaluationIdentifier(atomic).eval_hash + + +def _persist_attack_result( + memory: MemoryInterface, + target: ComponentIdentifier, + *, + outcome: AttackOutcome, + objective: str = "probe target", +) -> AttackResult: + """Persist a real AttackResult with a well-formed identifier tree.""" + attack_result = AttackResult( + conversation_id=f"conv-{outcome.value}-{datetime.now(timezone.utc).timestamp()}", + objective=objective, + atomic_attack_identifier=_make_atomic_attack_identifier(target), + outcome=outcome, + timestamp=datetime.now(timezone.utc), + ) + memory.add_attack_results_to_memory(attack_results=[attack_result]) + return attack_result + + +def _make_bench_with_real_memory( + memory: MemoryInterface, + objective_target: ComponentIdentifier, +) -> AdversarialBenchmark: + """Build a minimal benchmark wired to a real memory backend. + + Uses ``__new__`` to bypass the full ``__init__`` so we don't have to + register a target or build a strategy enum just to exercise the cache + helper. The helper only reads ``_memory`` and + ``_objective_target_identifier``. + """ + bench = AdversarialBenchmark.__new__(AdversarialBenchmark) + bench._memory = memory + bench._objective_target_identifier = objective_target + return bench + + +def _make_candidate(*, technique_eval_hash: str) -> MagicMock: + candidate = MagicMock() + candidate.technique_eval_hash = technique_eval_hash + return candidate + + +@pytest.mark.usefixtures("patch_central_database") +class TestCollectCachedCompletionPairsWithRealMemory: + """End-to-end cache coverage through real ``SQLiteMemory``.""" + + def test_cold_cache_returns_empty(self, sqlite_instance): + target = _make_objective_target_component() + bench = _make_bench_with_real_memory(sqlite_instance, target) + candidate = _make_candidate(technique_eval_hash=_technique_eval_hash_for(target)) + + result = bench._collect_cached_completion_pairs(atomic_attacks=[candidate]) + + assert result == set() + + def test_returns_hash_for_success_match_in_real_db(self, sqlite_instance): + target = _make_objective_target_component() + _persist_attack_result(sqlite_instance, target, outcome=AttackOutcome.SUCCESS) + + bench = _make_bench_with_real_memory(sqlite_instance, target) + tech_hash = _technique_eval_hash_for(target) + candidate = _make_candidate(technique_eval_hash=tech_hash) + + result = bench._collect_cached_completion_pairs(atomic_attacks=[candidate]) + + assert result == {tech_hash} + + def test_returns_hash_for_failure_match_in_real_db(self, sqlite_instance): + target = _make_objective_target_component() + _persist_attack_result(sqlite_instance, target, outcome=AttackOutcome.FAILURE) + + bench = _make_bench_with_real_memory(sqlite_instance, target) + tech_hash = _technique_eval_hash_for(target) + candidate = _make_candidate(technique_eval_hash=tech_hash) + + result = bench._collect_cached_completion_pairs(atomic_attacks=[candidate]) + + assert result == {tech_hash} + + def test_filters_out_persisted_results_with_different_objective_target(self, sqlite_instance): + """A row with a matching technique hash but a different target hash is rejected.""" + persisted_target = _make_objective_target_component(model_name="gpt-4o", temperature=0.7) + bench_target = _make_objective_target_component(model_name="gpt-4o-mini", temperature=0.7) + # AtomicAttackEvaluationIdentifier strips non-temperature target params, so the + # two targets share a technique hash even though their objective-target eval + # hashes differ. The SQL filter on $.eval_hash will hit; the python-side target + # filter inside get_cached_results_for_technique must do the rejection. + assert _technique_eval_hash_for(persisted_target) == _technique_eval_hash_for(bench_target) + assert ( + ObjectiveTargetEvaluationIdentifier(persisted_target).eval_hash + != ObjectiveTargetEvaluationIdentifier(bench_target).eval_hash + ) + + _persist_attack_result(sqlite_instance, persisted_target, outcome=AttackOutcome.SUCCESS) + + bench = _make_bench_with_real_memory(sqlite_instance, bench_target) + candidate = _make_candidate(technique_eval_hash=_technique_eval_hash_for(bench_target)) + + result = bench._collect_cached_completion_pairs(atomic_attacks=[candidate]) + + assert result == set() + + def test_filters_out_persisted_results_with_different_technique_hash(self, sqlite_instance): + """A row whose technique eval hash differs is rejected by the SQL filter.""" + persisted_target = _make_objective_target_component(model_name="gpt-4o", temperature=0.0) + bench_target = _make_objective_target_component(model_name="gpt-4o", temperature=0.7) + # Temperature feeds into AtomicAttackEvaluationIdentifier, so the persisted + # row's stamped $.eval_hash is different from the candidate's technique hash + # and the SQL filter returns no rows. + assert _technique_eval_hash_for(persisted_target) != _technique_eval_hash_for(bench_target) + + _persist_attack_result(sqlite_instance, persisted_target, outcome=AttackOutcome.SUCCESS) + + bench = _make_bench_with_real_memory(sqlite_instance, bench_target) + candidate = _make_candidate(technique_eval_hash=_technique_eval_hash_for(bench_target)) + + result = bench._collect_cached_completion_pairs(atomic_attacks=[candidate]) + + assert result == set() + + def test_filters_out_error_only_history(self, sqlite_instance): + """Outcomes other than SUCCESS / FAILURE never count as cached.""" + target = _make_objective_target_component() + _persist_attack_result(sqlite_instance, target, outcome=AttackOutcome.ERROR) + _persist_attack_result(sqlite_instance, target, outcome=AttackOutcome.UNDETERMINED) + + bench = _make_bench_with_real_memory(sqlite_instance, target) + candidate = _make_candidate(technique_eval_hash=_technique_eval_hash_for(target)) + + result = bench._collect_cached_completion_pairs(atomic_attacks=[candidate]) + + assert result == set() + + def test_dedupes_candidates_with_same_technique_hash(self, sqlite_instance): + """Two candidates sharing a technique hash collapse to a single set entry.""" + target = _make_objective_target_component() + _persist_attack_result(sqlite_instance, target, outcome=AttackOutcome.SUCCESS) + + bench = _make_bench_with_real_memory(sqlite_instance, target) + tech_hash = _technique_eval_hash_for(target) + candidates = [ + _make_candidate(technique_eval_hash=tech_hash), + _make_candidate(technique_eval_hash=tech_hash), + ] + + result = bench._collect_cached_completion_pairs(atomic_attacks=candidates) + + assert result == {tech_hash} From 31ed2fb77515dd26000b28642dd0f141949ab0e7 Mon Sep 17 00:00:00 2001 From: Victor Valbuena Date: Fri, 29 May 2026 11:44:42 -0700 Subject: [PATCH 35/40] Removing tool components from this branch. Tool-calling is a feature also in PR (1811), but its components were incorrectly added to this PR. --- pyrit/tools/__init__.py | 73 ---- pyrit/tools/backend.py | 87 ----- pyrit/tools/local_backend.py | 121 ------ pyrit/tools/mcp_backend.py | 199 ---------- pyrit/tools/mcp_client.py | 369 ------------------ pyrit/tools/models.py | 241 ------------ pyrit/tools/parsers.py | 107 ----- tests/unit/tools/__init__.py | 2 - tests/unit/tools/conftest.py | 299 -------------- tests/unit/tools/echo_mcp_server.py | 57 --- tests/unit/tools/test_local_tool_backend.py | 180 --------- tests/unit/tools/test_mcp_backend.py | 156 -------- tests/unit/tools/test_mcp_client.py | 171 -------- .../tools/test_prompt_target_tool_loop.py | 281 ------------- tests/unit/tools/test_tool_event_policy.py | 121 ------ tests/unit/tools/test_tool_loop_decorator.py | 289 -------------- 16 files changed, 2753 deletions(-) delete mode 100644 pyrit/tools/__init__.py delete mode 100644 pyrit/tools/backend.py delete mode 100644 pyrit/tools/local_backend.py delete mode 100644 pyrit/tools/mcp_backend.py delete mode 100644 pyrit/tools/mcp_client.py delete mode 100644 pyrit/tools/models.py delete mode 100644 pyrit/tools/parsers.py delete mode 100644 tests/unit/tools/__init__.py delete mode 100644 tests/unit/tools/conftest.py delete mode 100644 tests/unit/tools/echo_mcp_server.py delete mode 100644 tests/unit/tools/test_local_tool_backend.py delete mode 100644 tests/unit/tools/test_mcp_backend.py delete mode 100644 tests/unit/tools/test_mcp_client.py delete mode 100644 tests/unit/tools/test_prompt_target_tool_loop.py delete mode 100644 tests/unit/tools/test_tool_event_policy.py delete mode 100644 tests/unit/tools/test_tool_loop_decorator.py diff --git a/pyrit/tools/__init__.py b/pyrit/tools/__init__.py deleted file mode 100644 index f2ae0090ce..0000000000 --- a/pyrit/tools/__init__.py +++ /dev/null @@ -1,73 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# Licensed under the MIT license. - -""" -Generic tool-use scaffolding for :class:`~pyrit.prompt_target.PromptTarget`. - -This package provides a transport-agnostic tool-calling loop. The -:func:`tool_loop` decorator, when applied to ``send_prompt_async``, runs -the standard PyRIT validate+normalize work once and then repeatedly -re-enters the target's protected ``_send_prompt_to_target_async`` until -the model issues a stop response (or a configured limit is hit). - -A target opts in by declaring two collaborators: - -* ``self._tool_parser`` — a :class:`ToolCallParser` that walks a - response message and extracts pending :class:`ToolCall` instances. -* ``self.configuration.tool_event_policy`` — a :class:`ToolEventPolicy` - whose :class:`ToolEventBehavior` decides whether to ``EXECUTE``, - ``RAISE``, or ``RETURN_RAW`` on each detected call. - -When the policy is ``EXECUTE``, calls are dispatched through -``self.configuration.tool_backend``, an implementation of -:class:`ToolBackend`. :class:`LocalToolBackend` is the in-process -backend shipped here; :class:`MCPToolBackend` ships in C3 and proxies -through one or more MCP servers. - -The :class:`ToolBackend` Protocol is intentionally distinct from -:mod:`pyrit.registry` — that namespace is reserved for framework-level -identity registries (``TargetRegistry``, ``ScorerRegistry``) that -register named singletons for CLI lookup, which a per-target tool -dispatch table is not. - -Wiring of ``@tool_loop`` onto :class:`PromptTarget.send_prompt_async` -and of the ``tool_event_policy`` / ``tool_backend`` fields onto -:class:`TargetConfiguration` lands in C4/C5. - -The two exception types the loop raises -(:class:`~pyrit.exceptions.ToolCallNotSupported` and -:class:`~pyrit.exceptions.ToolCallLoopLimitExceeded`) live in -:mod:`pyrit.exceptions` alongside the rest of PyRIT's exception -catalog, so non-tools callers (attacks, normalizers) can import them -without taking a subsystem-level dependency on ``pyrit.tools``. -""" - -from pyrit.tools.backend import ToolBackend -from pyrit.tools.local_backend import LocalToolBackend -from pyrit.tools.mcp_backend import MCPToolBackend -from pyrit.tools.mcp_client import ( - DockerMCPServerSpec, - LocalMCPServerSpec, - MCPClient, - MCPServerSpec, - RemoteMCPServerSpec, -) -from pyrit.tools.models import ToolCall, ToolEventBehavior, ToolEventPolicy, tool_loop -from pyrit.tools.parsers import CanonicalEnvelopeParser, ToolCallParser - -__all__ = [ - "CanonicalEnvelopeParser", - "DockerMCPServerSpec", - "LocalMCPServerSpec", - "LocalToolBackend", - "MCPClient", - "MCPServerSpec", - "MCPToolBackend", - "RemoteMCPServerSpec", - "ToolBackend", - "ToolCall", - "ToolCallParser", - "ToolEventBehavior", - "ToolEventPolicy", - "tool_loop", -] diff --git a/pyrit/tools/backend.py b/pyrit/tools/backend.py deleted file mode 100644 index e7a02a7685..0000000000 --- a/pyrit/tools/backend.py +++ /dev/null @@ -1,87 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# Licensed under the MIT license. - -from __future__ import annotations - -from abc import ABC, abstractmethod -from typing import TYPE_CHECKING, Any - -if TYPE_CHECKING: - from pyrit.tools.models import ToolCall - - -class ToolBackend(ABC): - """ - Abstract base for backends that dispatch tool calls produced by a target. - - A :class:`ToolBackend` is a per-target dispatch table — it owns the - ``name -> async callable`` mapping a target uses to execute the tool - calls extracted from a model response. This is intentionally distinct - from :mod:`pyrit.registry`, whose ``Registry`` classes register named - framework singletons (targets, scorers, attacks) for CLI lookup. - - Two concrete implementations ship with PyRIT: - - * :class:`~pyrit.tools.LocalToolBackend` — in-process backend backed - by ``async def`` callables. Useful for unit tests and for embedding - tools inside the PyRIT process. - * :class:`~pyrit.tools.MCPToolBackend` — proxies dispatch through one - or more MCP servers. - - Subclasses MUST implement :attr:`schemas` and :meth:`dispatch_async`. - :meth:`dispatch_all_sequential_async` ships with a default - implementation that awaits :meth:`dispatch_async` once per call in - declaration order; backends that wish to parallelize dispatch - (e.g. fan out across multiple sandbox containers) should override it. - """ - - @property - @abstractmethod - def schemas(self) -> list[dict[str, Any]]: - """ - The JSON-schema descriptors for every tool the backend exposes. - - Returns: - list[dict[str, Any]]: One schema per tool, in a target-agnostic - format that concrete targets serialize into their request - body. - """ - - @abstractmethod - async def dispatch_async(self, call: ToolCall) -> dict[str, Any]: - """ - Execute a single tool call and return the structured result. - - Implementations MUST NOT raise on tool-side failures; they MUST - return an error envelope (e.g. ``{"error": "...", "tool": "..."}``) - so the tool loop can carry the failure back to the model. - - Args: - call (ToolCall): The tool call to dispatch. - - Returns: - dict[str, Any]: The structured tool result. - """ - - async def dispatch_all_sequential_async( - self, - calls: list[ToolCall], - ) -> list[tuple[ToolCall, dict[str, Any]]]: - """ - Dispatch every call in *calls* sequentially, preserving declaration order. - - Default implementation: ``await dispatch_async`` once per call. - Backends that parallelize dispatch should override this method. - - Args: - calls (list[ToolCall]): The calls to dispatch, in declaration order. - - Returns: - list[tuple[ToolCall, dict[str, Any]]]: ``(call, result)`` pairs, - in the same order as *calls*. - """ - results: list[tuple[ToolCall, dict[str, Any]]] = [] - for call in calls: - envelope = await self.dispatch_async(call) - results.append((call, envelope)) - return results diff --git a/pyrit/tools/local_backend.py b/pyrit/tools/local_backend.py deleted file mode 100644 index 25fe42e83c..0000000000 --- a/pyrit/tools/local_backend.py +++ /dev/null @@ -1,121 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# Licensed under the MIT license. - -from __future__ import annotations - -import logging -from typing import TYPE_CHECKING, Any - -from pyrit.tools.backend import ToolBackend - -if TYPE_CHECKING: - from collections.abc import Awaitable, Callable - - from pyrit.tools.models import ToolCall - -logger = logging.getLogger(__name__) - - -class LocalToolBackend(ToolBackend): - """ - In-process :class:`~pyrit.tools.ToolBackend` backed by a name -> ``async def`` - mapping. Useful for unit tests and for embedding small tools inside the - PyRIT process without standing up an MCP server. - - "Local" here means tools run in PyRIT's own Python process — no - subprocess, no IPC, no wire protocol. Contrast with - :class:`~pyrit.tools.MCPToolBackend` (lands in C3), which proxies - dispatch through one or more MCP servers reached via JSON-RPC. - - The backend dispatches sequentially in declaration order. Tool-side - failures (raised exceptions, missing names, allow-list rejections) - are converted into structured error envelopes so the tool loop can - forward them back to the model as ``function_call_output`` content - rather than aborting the conversation. - """ - - def __init__( - self, - *, - callables: dict[str, Callable[[dict[str, Any]], Awaitable[Any]]], - schemas: list[dict[str, Any]] | None = None, - allowed_tools: set[str] | None = None, - fail_on_missing_function: bool = True, - ) -> None: - """ - Initialize the backend. - - Args: - callables (dict[str, Callable[[dict[str, Any]], Awaitable[Any]]]): - Map from tool name to an ``async def`` that accepts a parsed - arguments dict and returns the tool result. Results are - serialized by the tool loop via :func:`json.dumps`. - schemas (list[dict[str, Any]] | None): JSON-schema descriptors - injected into the target's request body. Defaults to an empty - list when omitted. - allowed_tools (set[str] | None): Optional allow-list of tool - names; calls whose name is not in this set surface as - ``tool_not_allowed`` envelopes without invoking the callable. - Defaults to None (no allow-list; every registered tool is - callable). - fail_on_missing_function (bool): When True (default), an unknown - tool name raises :class:`KeyError`. When False, the backend - returns a ``tool_not_registered`` envelope so the model can - recover. - """ - self._callables = dict(callables) - self._schemas: list[dict[str, Any]] = list(schemas) if schemas is not None else [] - self._allowed_tools = set(allowed_tools) if allowed_tools is not None else None - self._fail_on_missing_function = fail_on_missing_function - - @property - def schemas(self) -> list[dict[str, Any]]: - """The JSON-schema descriptors for the tools in this backend.""" - return list(self._schemas) - - async def dispatch_async(self, call: ToolCall) -> dict[str, Any]: - """ - Dispatch a single tool call. Tool failures are converted into - structured envelopes; only configuration errors (missing tool with - ``fail_on_missing_function=True``) propagate as exceptions. - - Args: - call (ToolCall): The call to dispatch. - - Returns: - dict[str, Any]: The tool's result, or a structured error envelope. - - Raises: - KeyError: When the tool name is not registered and - ``fail_on_missing_function=True``. - """ - if self._allowed_tools is not None and call.name not in self._allowed_tools: - logger.info("Rejecting disallowed tool call: %s", call.name) - return { - "error": "tool_not_allowed", - "tool": call.name, - "allowed_tools": sorted(self._allowed_tools), - } - - fn = self._callables.get(call.name) - if fn is None: - if self._fail_on_missing_function: - raise KeyError(f"Tool '{call.name}' is not registered.") - available = sorted(self._callables.keys()) - logger.warning("Tool '%s' not registered. Available: %s", call.name, available) - return { - "error": "tool_not_registered", - "tool": call.name, - "available_tools": available, - } - - try: - result = await fn(call.arguments) - except Exception as ex: - logger.warning("Tool '%s' raised %s: %s", call.name, type(ex).__name__, ex) - return { - "error": "tool_execution_failed", - "tool": call.name, - "detail": str(ex), - } - return result if isinstance(result, dict) else {"result": result} diff --git a/pyrit/tools/mcp_backend.py b/pyrit/tools/mcp_backend.py deleted file mode 100644 index 66da88a30f..0000000000 --- a/pyrit/tools/mcp_backend.py +++ /dev/null @@ -1,199 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# Licensed under the MIT license. - -""" -Multi-server tool backend that proxies dispatch through one or more -MCP servers. - -This is the :class:`~pyrit.tools.ToolBackend` implementation that real -red-team configurations use. It composes one -:class:`~pyrit.tools.MCPClient` per :class:`~pyrit.tools.MCPServerSpec`, -aggregates their advertised schemas, routes incoming -:class:`~pyrit.tools.ToolCall` instances to the correct underlying -client, and enforces an optional ``allowed_tools`` allow-list. - -Contrast with :class:`~pyrit.tools.LocalToolBackend`, which dispatches -to Python ``async def`` callables inside PyRIT's own process. -""" - -from __future__ import annotations - -import asyncio -import logging -from contextlib import AsyncExitStack -from typing import TYPE_CHECKING, Any - -from pyrit.tools.backend import ToolBackend -from pyrit.tools.mcp_client import MCPClient - -if TYPE_CHECKING: - from collections.abc import Iterable - - from pyrit.tools.mcp_client import MCPServerSpec - from pyrit.tools.models import ToolCall - -logger = logging.getLogger(__name__) - - -class MCPToolBackend(ToolBackend): - """ - :class:`~pyrit.tools.ToolBackend` backed by one or more MCP servers. - - On :meth:`__aenter__`, the backend spawns / connects each server in - its :attr:`_servers` list (sequentially) through a single - :class:`contextlib.AsyncExitStack`, runs the MCP handshake, caches - schemas, and builds an advertised-name → ``(client, server_name)`` - routing table. Collisions raise :class:`ValueError` unless the - colliding specs set :attr:`~pyrit.tools.LocalMCPServerSpec.name_prefix`. - - A single shared :class:`AsyncExitStack` (rather than one per client) - is required so anyio's nested cancel scopes — opened by the ``mcp`` - SDK's ``stdio_client`` and ``ClientSession`` context managers — are - closed in strict LIFO order from the entering task. Closing - out-of-order would trip - ``"Attempted to exit a cancel scope that isn't the current task's - current cancel scope"``. - - Dispatch is serialized through an :class:`asyncio.Lock` per backend - instance — multiple concurrent coroutines sharing the same backend - (e.g. parallel attack runs) will not interleave JSON-RPC frames on - the same stdio pipe. - """ - - def __init__( - self, - *, - servers: Iterable[MCPServerSpec], - allowed_tools: list[str] | None = None, - ) -> None: - """ - Initialize the backend. - - Args: - servers: One or more :class:`MCPServerSpec` instances describing - where each server runs. - allowed_tools: Optional allow-list of tool names. Names not in - the list are filtered from :attr:`schemas` AND - short-circuit dispatch with a ``tool_not_allowed`` envelope. - Names are matched after :attr:`~LocalMCPServerSpec.name_prefix` - has been applied. Defaults to None (every advertised tool is - callable). - - Raises: - ValueError: When *servers* is empty. - """ - self._servers: list[MCPServerSpec] = list(servers) - if not self._servers: - raise ValueError("MCPToolBackend requires at least one server spec.") - self._allowed_tools: set[str] | None = set(allowed_tools) if allowed_tools is not None else None - self._clients: list[MCPClient] = [] - self._routing: dict[str, tuple[MCPClient, str]] = {} - self._dispatch_lock = asyncio.Lock() - self._stack: AsyncExitStack | None = None - self._entered = False - - @property - def schemas(self) -> list[dict[str, Any]]: - """The union of every connected server's schemas, filtered by ``allowed_tools``.""" - out: list[dict[str, Any]] = [] - for client in self._clients: - for schema in client.schemas: - if self._allowed_tools is not None and schema["name"] not in self._allowed_tools: - continue - out.append(schema) - return out - - async def __aenter__(self) -> MCPToolBackend: - """ - Connect each underlying client through a shared :class:`AsyncExitStack` and build the routing table. - - Returns: - MCPToolBackend: *self*, ready to dispatch. - - Raises: - ValueError: When two connected clients advertise the same tool - name without a disambiguating ``name_prefix``. - """ - stack = AsyncExitStack() - clients: list[MCPClient] = [] - routing: dict[str, tuple[MCPClient, str]] = {} - try: - for spec in self._servers: - client = MCPClient(spec=spec) - await stack.enter_async_context(client) - clients.append(client) - for advertised_name in client.tool_names: - if advertised_name in routing: - raise ValueError( - f"duplicate tool name '{advertised_name}'. " - "Set LocalMCPServerSpec.name_prefix on at least one " - "colliding server to disambiguate.", - ) - routing[advertised_name] = (client, advertised_name) - except Exception: - await stack.aclose() - raise - - self._stack = stack - self._clients = clients - self._routing = routing - self._entered = True - return self - - async def __aexit__(self, *exc: Any) -> None: - """Tear down every underlying client in strict LIFO order.""" - stack = self._stack - self._stack = None - self._clients = [] - self._routing = {} - self._entered = False - if stack is not None: - await stack.aclose() - - async def dispatch_async(self, call: ToolCall) -> dict[str, Any]: - """ - Route *call* to the correct client and dispatch. - - See :class:`MCPClient.dispatch_async` for the envelope shape. - Allow-list rejections and unknown-tool calls return error - envelopes; only "backend not entered" raises. - - Args: - call (ToolCall): The call to dispatch. - - Returns: - dict[str, Any]: A structured envelope (success, ``tool_not_allowed``, - ``tool_not_registered``, or the underlying - :meth:`MCPClient.dispatch_async` envelope). - - Raises: - RuntimeError: When the backend has not been entered via ``async with``. - """ - if not self._entered: - raise RuntimeError( - "MCPToolBackend is not active. Use `async with backend:` to manage its lifecycle before dispatching.", - ) - - if self._allowed_tools is not None and call.name not in self._allowed_tools: - logger.info("Rejecting disallowed tool call: %s", call.name) - return { - "is_error": True, - "error": "tool_not_allowed", - "tool": call.name, - "allowed_tools": sorted(self._allowed_tools), - } - - route = self._routing.get(call.name) - if route is None: - available = sorted(self._routing.keys()) - logger.warning("Tool '%s' not registered. Available: %s", call.name, available) - return { - "is_error": True, - "error": "tool_not_registered", - "tool": call.name, - "available_tools": available, - } - - client, _server_side_name = route - async with self._dispatch_lock: - return await client.dispatch_async(call) diff --git a/pyrit/tools/mcp_client.py b/pyrit/tools/mcp_client.py deleted file mode 100644 index 904004f675..0000000000 --- a/pyrit/tools/mcp_client.py +++ /dev/null @@ -1,369 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# Licensed under the MIT license. - -""" -Stdio-transport client for the Model Context Protocol (MCP). - -This module is the wire-protocol half of PyRIT's MCP integration. It -sits below :class:`~pyrit.tools.MCPToolBackend` (which composes one -:class:`MCPClient` per configured server and handles cross-server -routing) and above the upstream ``mcp`` Python SDK (which owns the -JSON-RPC framing, capability negotiation, and asyncio task plumbing). - -The three :class:`MCPServerSpec` variants describe *where* the server -runs. Only :class:`LocalMCPServerSpec` is implemented in this commit: - -* :class:`LocalMCPServerSpec` — spawn the server as a child process and - speak JSON-RPC over its stdin/stdout. -* :class:`RemoteMCPServerSpec` — HTTP/SSE transport against a hosted - server. Stub: ``connect_async`` raises ``NotImplementedError``. -* :class:`DockerMCPServerSpec` — stdio over ``docker run -i`` against a - hardened sandbox container. Stub: ``connect_async`` raises - ``NotImplementedError``. Implementation lands in the follow-up - sandbox PR. - -The stub variants are intentionally part of the type union today so -downstream code can be written against the eventual API without -forcing a Union expansion later. -""" - -from __future__ import annotations - -import asyncio -import logging -from contextlib import AsyncExitStack -from dataclasses import dataclass -from typing import TYPE_CHECKING, Any - -from mcp import ClientSession -from mcp.client.stdio import StdioServerParameters, stdio_client - -if TYPE_CHECKING: - from pyrit.tools.models import ToolCall - -logger = logging.getLogger(__name__) - - -@dataclass(frozen=True) -class LocalMCPServerSpec: - """ - Spec for an MCP server spawned as a child process and reached via - stdio JSON-RPC. - - Attributes: - command (str): The interpreter or binary to exec (e.g. ``"python"``). - args (tuple[str, ...]): Arguments passed to *command*, in order. - env (dict[str, str] | None): Environment overlay for the child - process. ``None`` (default) inherits PyRIT's environment. - name_prefix (str | None): When set, every tool advertised by the - server is registered as ``f"{name_prefix}{tool_name}"`` in - the parent :class:`~pyrit.tools.MCPToolBackend`. Used to - disambiguate two servers that expose the same tool name. - timeout_seconds (float): Per-call timeout, enforced by - :meth:`MCPClient.dispatch_async`. Defaults to 30 seconds. - """ - - command: str - args: tuple[str, ...] = () - env: dict[str, str] | None = None - name_prefix: str | None = None - timeout_seconds: float = 30.0 - - -@dataclass(frozen=True) -class RemoteMCPServerSpec: - """ - Spec for an MCP server reached over HTTP / SSE. **Not implemented** - in this PR — :meth:`MCPClient.connect_async` raises - :class:`NotImplementedError`. Tracked by ``# TODO(mcp-http-transport)``. - - Attributes: - url (str): The base URL of the MCP server. - name_prefix (str | None): Same semantics as - :attr:`LocalMCPServerSpec.name_prefix`. - timeout_seconds (float): Per-call timeout. - """ - - url: str - name_prefix: str | None = None - timeout_seconds: float = 30.0 - - -# TODO(sandbox-provider) — DockerMCPServerSpec stub here; implementation lands in follow-up PR. -@dataclass(frozen=True) -class DockerMCPServerSpec: - """ - Spec for an MCP server hosted inside a hardened Docker container. - - **NOT IMPLEMENTED IN THIS PR.** Reached via stdio over ``docker run -i``. - - Expected behavior in the follow-up sandbox PR: - - * One container per spec instance, managed by a process-wide - ``SandboxPool``. - * Image is built lazily, keyed by ``sha256(Dockerfile + build_context)``, - and cached across attacks; no rebuild unless missing or explicitly - overridden. - * Container is recreated from the cached image at attack and scenario - boundaries (filesystem returns to baseline every time). - * Network access governed by ``NetworkProfile`` (default ``"none"`` = - ``--network=none``). - * Container runs as a non-root UID with ``--cap-drop=ALL``, a read-only - root filesystem, and an in-container MCP server exposing - ``run_shell(cmd, timeout_seconds)``. - - Attributes: - image (str): Docker image tag (e.g. ``"pyrit-sandbox:base"``). - network_profile (str): ``NetworkProfile`` name; ``"none"`` (default) - launches the container with ``--network=none``. - name_prefix (str | None): Same semantics as - :attr:`LocalMCPServerSpec.name_prefix`. - timeout_seconds (float): Per-call timeout. - - Future fields (deferred to the follow-up sandbox PR): ``memory_limit``, - ``cpu_limit``, ``pids_limit``, ``env``, ``mounts``, ``command_override``. - """ - - image: str - network_profile: str = "none" - name_prefix: str | None = None - timeout_seconds: float = 30.0 - - -MCPServerSpec = LocalMCPServerSpec | RemoteMCPServerSpec | DockerMCPServerSpec - - -def _to_input_schema_dict(input_schema: Any) -> dict[str, Any]: - """ - Coerce the SDK's tool ``inputSchema`` (pydantic model or dict) into a plain dict. - - Returns: - dict[str, Any]: A plain-dict copy of *input_schema*, or an empty - object schema when *input_schema* is None or of an unrecognized type. - """ - if input_schema is None: - return {"type": "object", "properties": {}} - if hasattr(input_schema, "model_dump"): - return input_schema.model_dump() - if isinstance(input_schema, dict): - return dict(input_schema) - return {"type": "object", "properties": {}} - - -def _flatten_content(content: list[Any]) -> str: - """ - Concatenate the text portions of an MCP ``CallToolResult.content`` list. - - Returns: - str: Concatenated ``.text`` values from each content item, in order. - """ - pieces: list[str] = [] - for item in content: - text = getattr(item, "text", None) - if text is not None: - pieces.append(text) - elif isinstance(item, dict) and "text" in item: - pieces.append(item["text"]) - return "".join(pieces) - - -class MCPClient: - """ - A single MCP-server session. - - The client owns the lifetime of one server's transport stack and - exposes a uniform :meth:`dispatch_async` regardless of which - :class:`MCPServerSpec` variant it was constructed from. Composition - across multiple servers (routing, schema aggregation, allow-lists) - is the responsibility of :class:`~pyrit.tools.MCPToolBackend`. - - Lifecycle: - - * :meth:`connect_async` spawns the subprocess (for - :class:`LocalMCPServerSpec`), runs the MCP handshake, and caches - ``tools/list`` results. - * :meth:`dispatch_async` issues one ``tools/call`` and returns a - structured envelope (success or error). - * :meth:`close_async` tears down the transport stack. - - The class is usable as an async context manager. - """ - - def __init__(self, *, spec: MCPServerSpec) -> None: - """ - Initialize the client around *spec*. Does not connect; call - :meth:`connect_async` (or use the async context-manager form) to start - the transport stack. - """ - self._spec = spec - self._stack = AsyncExitStack() - self._session: ClientSession | None = None - self._tools: list[Any] = [] - - @property - def spec(self) -> MCPServerSpec: - """The :class:`MCPServerSpec` this client was constructed with.""" - return self._spec - - @property - def schemas(self) -> list[dict[str, Any]]: - """ - JSON schemas for every tool the server advertises. - - Each schema is shaped ``{"name", "description", "parameters"}``. - The optional :attr:`LocalMCPServerSpec.name_prefix` is applied - here so a backend that owns this client sees the prefixed name. - """ - prefix = getattr(self._spec, "name_prefix", None) or "" - return [ - { - "name": f"{prefix}{tool.name}", - "description": tool.description or "", - "parameters": _to_input_schema_dict(tool.inputSchema), - } - for tool in self._tools - ] - - @property - def tool_names(self) -> list[str]: - """Tool names with the spec's :attr:`name_prefix` applied.""" - return [s["name"] for s in self.schemas] - - def _strip_prefix(self, name: str) -> str: - prefix = getattr(self._spec, "name_prefix", None) or "" - if prefix and name.startswith(prefix): - return name[len(prefix) :] - return name - - async def connect_async(self) -> None: - """Establish the transport, run the handshake, and cache schemas.""" - if isinstance(self._spec, RemoteMCPServerSpec): - raise NotImplementedError( - "HTTP/SSE transport ships in a follow-up PR. " - "RemoteMCPServerSpec is declared today so user code can target the eventual API." - ) - if isinstance(self._spec, DockerMCPServerSpec): - raise NotImplementedError( - "Docker sandbox transport ships in a follow-up PR. " - "DockerMCPServerSpec runs the MCP server inside a hardened " - "Debian container reached via stdio over `docker run -i`, " - "managed by a process-wide SandboxPool with image caching and " - "per-attack container recreation." - ) - - assert isinstance(self._spec, LocalMCPServerSpec) - params = StdioServerParameters( - command=self._spec.command, - args=list(self._spec.args), - env=self._spec.env, - ) - read, write = await self._stack.enter_async_context(stdio_client(params)) - session = await self._stack.enter_async_context(ClientSession(read, write)) - await session.initialize() - result = await session.list_tools() - self._session = session - self._tools = list(result.tools) - - async def close_async(self) -> None: - """Tear down the transport stack. Idempotent; safe to call before connect.""" - try: - await self._stack.aclose() - except Exception as ex: # noqa: BLE001 — close should never raise into the caller. - logger.warning("Error tearing down MCP client stack: %s", ex) - finally: - self._stack = AsyncExitStack() - self._session = None - self._tools = [] - - async def __aenter__(self) -> MCPClient: - """ - Connect the transport stack and return *self*. - - Returns: - MCPClient: *self*, ready to dispatch tool calls. - """ - await self.connect_async() - return self - - async def __aexit__(self, *exc: Any) -> None: - """Tear down the transport stack.""" - await self.close_async() - - async def dispatch_async(self, call: ToolCall) -> dict[str, Any]: - """ - Issue one ``tools/call`` and return a structured envelope. - - Envelope shape: - - * Success: ``{"is_error": False, "content": str, "tool": name}``. - * Timeout: ``{"is_error": True, "error": "tool_timeout", "tool": name, ...}``. - * Server-reported error: ``{"is_error": True, "error": "tool_execution_failed", "tool": name, ...}``. - - Tool-side failures are converted to envelopes; only programmer - errors (calling before :meth:`connect_async`) raise. - - Args: - call (ToolCall): The call to dispatch. The advertised - ``name_prefix`` (if any) is stripped before contacting the server. - - Returns: - dict[str, Any]: One of the envelope shapes documented above. - - Raises: - RuntimeError: When the client has not been connected. - """ - if self._session is None: - raise RuntimeError("MCPClient is not connected; call connect_async first.") - - server_side_name = self._strip_prefix(call.name) - timeout = getattr(self._spec, "timeout_seconds", 30.0) - try: - result = await asyncio.wait_for( - self._session.call_tool(server_side_name, arguments=dict(call.arguments)), - timeout=timeout, - ) - except asyncio.TimeoutError: - logger.warning( - "MCP tool '%s' timed out after %.2fs", - call.name, - timeout, - ) - return { - "is_error": True, - "error": "tool_timeout", - "tool": call.name, - "timeout_seconds": timeout, - } - except Exception as ex: # noqa: BLE001 — wrap and surface as envelope. - logger.warning( - "MCP tool '%s' raised %s: %s", - call.name, - type(ex).__name__, - ex, - ) - return { - "is_error": True, - "error": "tool_execution_failed", - "tool": call.name, - "detail": str(ex), - } - - content_text = _flatten_content(list(result.content)) - is_error = bool(getattr(result, "isError", False)) - envelope: dict[str, Any] = { - "is_error": is_error, - "content": content_text, - "tool": call.name, - } - if is_error: - envelope["error"] = "tool_execution_failed" - return envelope - - -__all__ = [ - "DockerMCPServerSpec", - "LocalMCPServerSpec", - "MCPClient", - "MCPServerSpec", - "RemoteMCPServerSpec", -] diff --git a/pyrit/tools/models.py b/pyrit/tools/models.py deleted file mode 100644 index 0e05f9be9d..0000000000 --- a/pyrit/tools/models.py +++ /dev/null @@ -1,241 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# Licensed under the MIT license. - -from __future__ import annotations - -import enum -import functools -import json -import logging -from dataclasses import dataclass, field -from typing import TYPE_CHECKING, Any - -from pyrit.exceptions import ToolCallLoopLimitExceeded, ToolCallNotSupported -from pyrit.models import Message, MessagePiece - -if TYPE_CHECKING: - from collections.abc import Awaitable, Callable - - from pyrit.tools.backend import ToolBackend - from pyrit.tools.parsers import ToolCallParser - -logger = logging.getLogger(__name__) - - -@dataclass(frozen=True) -class ToolCall: - """ - A parsed tool call extracted from a target response. - - Concrete :class:`~pyrit.tools.ToolCallParser` implementations build - :class:`ToolCall` instances by walking the response message pieces. - The :attr:`raw_envelope` carries the original target-specific dict - (e.g. the function_call JSON section) so dispatchers and observers - can recover provider-specific fields without re-parsing. - - Attributes: - call_id (str): The provider-issued call identifier; must round-trip - into the matching ``function_call_output`` piece. - name (str): The tool name to dispatch. - arguments (dict[str, Any]): The parsed JSON arguments. - raw_envelope (dict[str, Any]): The original provider envelope. - """ - - call_id: str - name: str - arguments: dict[str, Any] - raw_envelope: dict[str, Any] = field(default_factory=dict) - - -class ToolEventBehavior(enum.Enum): - """ - What the tool loop should do when a target response contains a - pending tool call. - - Values: - EXECUTE: Dispatch the call via ``configuration.tool_backend`` - and re-enter the target with the tool output appended. - This is the standard agentic loop behavior. - RAISE: Raise :class:`~pyrit.exceptions.ToolCallNotSupported` with - the partial conversation attached. Useful for red-team - attacks that want to observe attempted tool use without - allowing execution. - RETURN_RAW: Return the assistant response containing the tool - call as-is, without dispatching. Useful when a caller wants - to inspect tool calls in-band (e.g. a scorer that scores - attempted tool use). - """ - - EXECUTE = "execute" - RAISE = "raise" - RETURN_RAW = "return_raw" - - -@dataclass(frozen=True) -class ToolEventPolicy: - """ - Per-target configuration that controls how the tool loop responds - to a pending tool call from the model. - - Attributes: - behavior (ToolEventBehavior): What to do on each detected tool call. - max_tool_iterations (int): Maximum number of model<->tool round-trips - before the loop raises :class:`ToolCallLoopLimitExceeded`. Each - iteration is one ``_send_prompt_to_target_async`` call. - """ - - behavior: ToolEventBehavior - max_tool_iterations: int = 5 - - -def _build_function_call_output_message( - *, - reference_piece: MessagePiece, - outputs: list[tuple[ToolCall, Any]], -) -> Message: - """ - Build the canonical ``tool`` message produced after dispatching one or more - tool calls in a single iteration. - - The returned :class:`Message` contains one - :class:`MessagePiece` per ``(call, result)`` pair, in declaration order. - Every piece has ``role="tool"`` and ``original_value_data_type="function_call_output"``, - with the JSON envelope ``{"type": "function_call_output", "call_id": ..., "output": ...}``. - - Lineage metadata (conversation_id, identifiers) is copied from - *reference_piece* — typically the first piece of the assistant message - that issued the tool calls — so the resulting message stays inside the - correct conversation. - - Args: - reference_piece (MessagePiece): Piece whose lineage metadata is - copied onto every output piece. Pass the first piece of the - assistant message that produced the calls. - outputs (list[tuple[ToolCall, Any]]): ``(call, result)`` pairs in - declaration order. *result* is serialized via :func:`json.dumps` - unless it is already a string. - - Returns: - Message: One message carrying every function_call_output piece. - """ - pieces: list[MessagePiece] = [] - for call, result in outputs: - output_str = result if isinstance(result, str) else json.dumps(result, separators=(",", ":")) - envelope = json.dumps( - {"type": "function_call_output", "call_id": call.call_id, "output": output_str}, - separators=(",", ":"), - ) - pieces.append( - MessagePiece( - role="tool", - original_value=envelope, - original_value_data_type="function_call_output", - conversation_id=reference_piece.conversation_id, - prompt_target_identifier=reference_piece.prompt_target_identifier, - attack_identifier=reference_piece.attack_identifier, - ) - ) - return Message(message_pieces=pieces, skip_validation=True) - - -def tool_loop( - method: Callable[..., Awaitable[list[Message]]], -) -> Callable[..., Awaitable[list[Message]]]: - """ - Wrap a :class:`~pyrit.prompt_target.PromptTarget`-style - ``send_prompt_async`` to run an agentic tool-use loop. - - When the target's ``configuration.tool_event_policy`` is ``None`` the - wrapper is a no-op — the wrapped method runs unchanged. When a policy - is configured, the wrapper replaces the method body with the loop: - - 1. Validate and normalize the incoming message exactly once. - 2. Repeatedly call ``self._send_prompt_to_target_async`` with the - growing conversation. - 3. After each call, parse the last response via ``self._tool_parser``. - Exit on empty parse (model issued a stop response). - 4. On a non-empty parse, branch on ``policy.behavior``: - ``RAISE`` raises :class:`ToolCallNotSupported`; ``RETURN_RAW`` - returns the chain as-is; ``EXECUTE`` dispatches the calls via - ``configuration.tool_backend`` and appends the tool message. - 5. Raise :class:`ToolCallLoopLimitExceeded` if the loop runs past - ``policy.max_tool_iterations`` without the model stopping. - - The decorator deliberately knows nothing about MCP, OpenAI, or any - specific transport. The two collaborators it requires — - ``self._tool_parser`` and ``self.configuration.tool_backend`` — are - plain protocols (:class:`ToolCallParser`, :class:`ToolBackend`). - - Args: - method (Callable): The async method to wrap. Must have the - ``async def f(self, *, message: Message) -> list[Message]`` - signature of :meth:`PromptTarget.send_prompt_async`. - - Returns: - Callable: The wrapped method. - """ - - @functools.wraps(method) - async def wrapper(self: Any, *, message: Message) -> list[Message]: - policy: ToolEventPolicy | None = getattr(self.configuration, "tool_event_policy", None) - if policy is None: - return await method(self, message=message) - - message.validate() - normalized_conversation = await self._get_normalized_conversation_async(message=message) - if not normalized_conversation: - raise ValueError("Normalization pipeline returned an empty conversation. Cannot send an empty request.") - self._validate_request(normalized_conversation=normalized_conversation) - - parser: ToolCallParser | None = getattr(self, "_tool_parser", None) - backend: ToolBackend | None = getattr(self.configuration, "tool_backend", None) - max_iter = policy.max_tool_iterations - - all_responses: list[Message] = [] - - for _ in range(max_iter): - responses_this_turn = await self._send_prompt_to_target_async( - normalized_conversation=normalized_conversation, - ) - all_responses.extend(responses_this_turn) - - if parser is None: - return all_responses - - last_response = responses_this_turn[-1] - pending_calls = parser.parse(last_response) - - if not pending_calls: - return all_responses - - if policy.behavior is ToolEventBehavior.RAISE: - raise ToolCallNotSupported( - message=( - f"Target produced {len(pending_calls)} tool call(s) but ToolEventPolicy.behavior is RAISE." - ), - partial_conversation=all_responses, - ) - - if policy.behavior is ToolEventBehavior.RETURN_RAW: - return all_responses - - if backend is None: - raise ToolCallNotSupported( - message=(f"Target produced {len(pending_calls)} tool call(s) but no tool_backend is configured."), - partial_conversation=all_responses, - ) - - results = await backend.dispatch_all_sequential_async(pending_calls) - tool_msg = _build_function_call_output_message( - reference_piece=last_response.message_pieces[0], - outputs=results, - ) - all_responses.append(tool_msg) - normalized_conversation = list(normalized_conversation) + [last_response, tool_msg] - - raise ToolCallLoopLimitExceeded( - message=f"Tool loop exceeded max_tool_iterations={max_iter} without a stop response.", - partial_conversation=all_responses, - ) - - return wrapper diff --git a/pyrit/tools/parsers.py b/pyrit/tools/parsers.py deleted file mode 100644 index c903eb73c7..0000000000 --- a/pyrit/tools/parsers.py +++ /dev/null @@ -1,107 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# Licensed under the MIT license. - -from __future__ import annotations - -import json -from typing import TYPE_CHECKING, Protocol, runtime_checkable - -if TYPE_CHECKING: - from pyrit.models import Message, MessagePiece - from pyrit.tools.models import ToolCall - - -@runtime_checkable -class ToolCallParser(Protocol): - """ - Protocol for extracting tool calls from a target response message. - - Concrete parsers live next to the target whose response shape they - understand (the canonical-envelope parser shipped here is shared by - :class:`OpenAIResponseTarget`; per-model-family parsers for non-OpenAI - targets ship in a follow-up, see plan §12.9). Parsers MUST return an - empty list when the model has issued a stop response — the tool loop - uses the empty list as the signal to exit. - """ - - def parse(self, message: Message) -> list[ToolCall]: - """ - Extract tool calls from a target response message. - - Args: - message (Message): The most recent assistant response. - - Returns: - list[ToolCall]: Tool calls, in declaration order. An empty list - signals that the model produced a stop response. - """ - ... - - -def _extract_function_call_pieces(message: Message) -> list[MessagePiece]: - """ - Return every :class:`MessagePiece` in *message* whose - ``original_value_data_type`` is ``"function_call"``. - - This is the canonical envelope used by every PyRIT-supported tool-emitting - target. It is exposed here so concrete parsers can reuse the filter rather - than re-implementing it. - - Args: - message (Message): The message to scan. - - Returns: - list[MessagePiece]: Pieces whose ``original_value_data_type`` is - ``"function_call"``, in their declaration order. - """ - return [piece for piece in message.message_pieces if piece.original_value_data_type == "function_call"] - - -class CanonicalEnvelopeParser: - """ - Reference :class:`ToolCallParser` for the canonical function_call envelope. - - Walks every :class:`MessagePiece` whose ``original_value_data_type`` is - ``"function_call"`` and decodes the canonical JSON shape:: - - { - "type": "function_call", - "call_id": "", - "name": "", - "arguments": "" - } - - into :class:`ToolCall` instances. Pieces of other data types -- reasoning, - mcp_call, web_search_call, etc. -- are ignored (they pass through to - Memory but are not client-side dispatchable). Per-model-family parsers - for non-OpenAI targets ship in a follow-up PR (see plan §12.9). - """ - - def parse(self, message: Message) -> list[ToolCall]: - """ - Decode canonical ``function_call`` pieces in *message* into :class:`ToolCall`. - - Args: - message (Message): The most recent assistant response. - - Returns: - list[ToolCall]: One :class:`ToolCall` per ``function_call`` - piece, in declaration order. Empty if the message contains - no ``function_call`` pieces (model stop). - """ - from pyrit.tools.models import ToolCall - - calls: list[ToolCall] = [] - for piece in _extract_function_call_pieces(message): - envelope = json.loads(piece.original_value) - arguments_raw = envelope.get("arguments", "{}") - arguments = json.loads(arguments_raw) if isinstance(arguments_raw, str) else dict(arguments_raw) - calls.append( - ToolCall( - call_id=envelope["call_id"], - name=envelope["name"], - arguments=arguments, - raw_envelope=envelope, - ) - ) - return calls diff --git a/tests/unit/tools/__init__.py b/tests/unit/tools/__init__.py deleted file mode 100644 index 9a0454564d..0000000000 --- a/tests/unit/tools/__init__.py +++ /dev/null @@ -1,2 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# Licensed under the MIT license. diff --git a/tests/unit/tools/conftest.py b/tests/unit/tools/conftest.py deleted file mode 100644 index ad7d2c7fd1..0000000000 --- a/tests/unit/tools/conftest.py +++ /dev/null @@ -1,299 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# Licensed under the MIT license. - -""" -Shared fixtures for ``tests/unit/tools``. - -Provides the minimal collaborators the tool-loop tests need to exercise -:func:`pyrit.tools.tool_loop` end-to-end without standing up real targets -or MCP transports: - -* :class:`_FakeToolTarget` — a :class:`PromptTarget` subclass whose - ``_send_prompt_to_target_async`` returns scripted messages from a queue - and whose ``_get_normalized_conversation_async`` skips the memory round - trip so decorator behavior is isolated from normalization. -* :class:`_RecordingToolBackend` — a :class:`ToolBackend` that records - every dispatched call (for order-of-execution assertions) and returns - results from a scripted queue. -* :class:`_CanonicalEnvelopeParser` — a :class:`ToolCallParser` that walks - message pieces and parses the canonical ``function_call`` JSON envelope. - -Helper message builders (``_make_user_message``, -``_make_assistant_text_message``, ``_make_assistant_function_call_message``) -produce the canonical envelope shape used by the OpenAI targets after the -C6 normalization commit. -""" - -from __future__ import annotations - -import json -import uuid -from collections import deque -from typing import Any - -import pytest - -from pyrit.models import Message, MessagePiece -from pyrit.prompt_target.common.prompt_target import PromptTarget -from pyrit.prompt_target.common.target_capabilities import TargetCapabilities -from pyrit.prompt_target.common.target_configuration import TargetConfiguration -from pyrit.tools import ( - LocalToolBackend, - ToolBackend, - ToolCall, - ToolCallParser, - ToolEventBehavior, - ToolEventPolicy, -) - - -def _make_user_message(text: str, *, conversation_id: str | None = None) -> Message: - """Build a single-piece user :class:`Message` carrying *text*.""" - return Message( - message_pieces=[ - MessagePiece( - role="user", - original_value=text, - original_value_data_type="text", - conversation_id=conversation_id or str(uuid.uuid4()), - ) - ] - ) - - -def _make_assistant_text_message(text: str, *, conversation_id: str | None = None) -> Message: - """Build a single-piece assistant :class:`Message` carrying plain text.""" - return Message( - message_pieces=[ - MessagePiece( - role="assistant", - original_value=text, - original_value_data_type="text", - conversation_id=conversation_id or str(uuid.uuid4()), - ) - ], - skip_validation=True, - ) - - -def _make_function_call_piece( - *, - call_id: str, - name: str, - arguments: dict[str, Any], - conversation_id: str | None = None, -) -> MessagePiece: - """Build one assistant ``function_call`` piece carrying the canonical envelope.""" - envelope = { - "type": "function_call", - "call_id": call_id, - "name": name, - "arguments": json.dumps(arguments, separators=(",", ":")), - } - return MessagePiece( - role="assistant", - original_value=json.dumps(envelope, separators=(",", ":")), - original_value_data_type="function_call", - conversation_id=conversation_id or str(uuid.uuid4()), - ) - - -def _make_assistant_function_call_message( - *, - calls: list[tuple[str, str, dict[str, Any]]], - conversation_id: str | None = None, -) -> Message: - """ - Build an assistant :class:`Message` carrying one ``function_call`` piece - per ``(call_id, name, args)`` tuple, in declaration order. - """ - conv_id = conversation_id or str(uuid.uuid4()) - pieces = [ - _make_function_call_piece(call_id=cid, name=name, arguments=args, conversation_id=conv_id) - for cid, name, args in calls - ] - return Message(message_pieces=pieces, skip_validation=True) - - -class _CanonicalEnvelopeParser: - """ - Reference :class:`ToolCallParser` that understands the canonical envelope - (``original_value_data_type == "function_call"`` carrying a JSON object - with ``type``/``call_id``/``name``/``arguments``). - - Per-target parsers shipped in C7/C8 will reuse this shape; this stand-in - keeps decorator tests independent of the real OpenAI parsers. - """ - - def parse(self, message: Message) -> list[ToolCall]: - calls: list[ToolCall] = [] - for piece in message.message_pieces: - if piece.original_value_data_type != "function_call": - continue - envelope = json.loads(piece.original_value) - arguments_str = envelope.get("arguments", "{}") - arguments = json.loads(arguments_str) if isinstance(arguments_str, str) else dict(arguments_str) - calls.append( - ToolCall( - call_id=envelope["call_id"], - name=envelope["name"], - arguments=arguments, - raw_envelope=envelope, - ) - ) - return calls - - -class _RecordingToolBackend(ToolBackend): - """ - Minimal :class:`ToolBackend` that records every dispatched call and - returns results from a scripted queue. Used to assert dispatch order, - iteration count, and per-call payload shape without invoking real tools. - """ - - def __init__( - self, - *, - scripted_results: list[Any] | None = None, - schemas: list[dict[str, Any]] | None = None, - ) -> None: - self._results: deque[Any] = deque(scripted_results or []) - self._schemas: list[dict[str, Any]] = list(schemas) if schemas is not None else [] - self.recorded_calls: list[ToolCall] = [] - - @property - def schemas(self) -> list[dict[str, Any]]: - return list(self._schemas) - - async def dispatch_async(self, call: ToolCall) -> dict[str, Any]: - self.recorded_calls.append(call) - if not self._results: - return {"result": f"recorded:{call.name}:{call.call_id}"} - nxt = self._results.popleft() - return nxt if isinstance(nxt, dict) else {"result": nxt} - - -class _FakeToolTarget(PromptTarget): - """ - Test-only :class:`PromptTarget` whose ``_send_prompt_to_target_async`` - pops scripted responses off a queue. ``_get_normalized_conversation_async`` - is overridden to return ``[message]`` directly, isolating decorator - behavior from the memory + normalization pipeline. - - Inherits the base class's ``@final @tool_loop send_prompt_async``; the - policy + backend are wired through :class:`TargetConfiguration` so the - wrapper finds them via ``self.configuration.tool_event_policy`` and - ``self.configuration.tool_backend``. - """ - - def __init__( - self, - *, - scripted_responses: list[Message], - policy: ToolEventPolicy | None = None, - backend: Any = None, - parser: ToolCallParser | None = None, - ) -> None: - # ``supports_tool_use`` is forced on whenever a policy is configured so - # the TargetConfiguration validator accepts the backend. - caps = TargetCapabilities( - supports_multi_turn=True, - supports_multi_message_pieces=True, - supports_tool_use=policy is not None, - ) - config = TargetConfiguration( - capabilities=caps, - tool_event_policy=policy, - tool_backend=backend, - ) - super().__init__(custom_configuration=config) - self._scripted_responses: deque[Message] = deque(scripted_responses) - self.call_count: int = 0 - self.normalized_conversations_seen: list[list[Message]] = [] - self._parser_instance: ToolCallParser | None = parser if parser is not None else _CanonicalEnvelopeParser() - - @property - def _tool_parser(self) -> ToolCallParser | None: - return self._parser_instance - - async def _get_normalized_conversation_async(self, *, message: Message) -> list[Message]: - return [message] - - def _validate_request(self, *, normalized_conversation: list[Message]) -> None: - return - - async def _send_prompt_to_target_async(self, *, normalized_conversation: list[Message]) -> list[Message]: - self.call_count += 1 - self.normalized_conversations_seen.append(list(normalized_conversation)) - if not self._scripted_responses: - raise AssertionError(f"Fake target ran out of scripted responses on iteration {self.call_count}.") - return [self._scripted_responses.popleft()] - - -@pytest.fixture -def make_fake_target(patch_central_database): - """ - Factory fixture for :class:`_FakeToolTarget`. Each invocation returns a - fresh target instance whose scripted response queue is independent of - other targets created during the test. - """ - - def _factory( - *, - scripted_responses: list[Message], - policy: ToolEventPolicy | None = None, - backend: Any = None, - parser: ToolCallParser | None = None, - ) -> _FakeToolTarget: - return _FakeToolTarget( - scripted_responses=scripted_responses, - policy=policy, - backend=backend, - parser=parser, - ) - - return _factory - - -@pytest.fixture -def recording_backend(): - """Factory fixture for :class:`_RecordingToolBackend`.""" - - def _factory(*, scripted_results: list[Any] | None = None) -> _RecordingToolBackend: - return _RecordingToolBackend(scripted_results=scripted_results) - - return _factory - - -@pytest.fixture -def execute_policy(): - """ - Factory fixture for :class:`ToolEventPolicy` with - ``behavior=ToolEventBehavior.EXECUTE`` and a tunable iteration cap. - """ - - def _factory(*, max_tool_iterations: int = 5) -> ToolEventPolicy: - return ToolEventPolicy( - behavior=ToolEventBehavior.EXECUTE, - max_tool_iterations=max_tool_iterations, - ) - - return _factory - - -__all__ = [ - "LocalToolBackend", - "ToolCall", - "ToolEventBehavior", - "ToolEventPolicy", - "_CanonicalEnvelopeParser", - "_FakeToolTarget", - "_RecordingToolBackend", - "_make_assistant_function_call_message", - "_make_assistant_text_message", - "_make_function_call_piece", - "_make_user_message", - "execute_policy", - "make_fake_target", - "recording_backend", -] diff --git a/tests/unit/tools/echo_mcp_server.py b/tests/unit/tools/echo_mcp_server.py deleted file mode 100644 index 723a3c6594..0000000000 --- a/tests/unit/tools/echo_mcp_server.py +++ /dev/null @@ -1,57 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# Licensed under the MIT license. - -""" -Deterministic echo MCP server used as a stdio subprocess fixture by -``tests/unit/tools/test_mcp_client.py`` (C3) and the integration tests -(C9). - -Lands in C2 so subsequent commits don't shuffle test plumbing; C2's own -tests do not import this module (the :class:`CallableToolRegistry` is -exercised in-process). - -Run directly as ``python echo_mcp_server.py`` to expose the four tools -over stdio. The MCP client harness in C3 launches this file with -``mcp.client.stdio.stdio_client`` and asserts behavior end to end. -""" - -from __future__ import annotations - -import asyncio - -from mcp.server.fastmcp import FastMCP - -mcp = FastMCP("pyrit-echo") - - -@mcp.tool() -def echo(text: str) -> str: - """Return *text* unchanged.""" - return text - - -@mcp.tool() -def add(a: int, b: int) -> int: - """Return ``a + b``.""" - return a + b - - -@mcp.tool() -def reverse(text: str) -> str: - """Return *text* reversed.""" - return text[::-1] - - -@mcp.tool() -async def slow_echo(text: str, delay_ms: int = 0) -> str: - """ - Return *text* after sleeping ``delay_ms`` milliseconds. Used by - timeout / cancellation tests. - """ - if delay_ms > 0: - await asyncio.sleep(delay_ms / 1000.0) - return text - - -if __name__ == "__main__": - mcp.run() diff --git a/tests/unit/tools/test_local_tool_backend.py b/tests/unit/tools/test_local_tool_backend.py deleted file mode 100644 index 8e24693140..0000000000 --- a/tests/unit/tools/test_local_tool_backend.py +++ /dev/null @@ -1,180 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# Licensed under the MIT license. - -""" -Unit tests for :class:`pyrit.tools.LocalToolBackend`. - -Coverage map (rows from the C2 test matrix): - -* **U10** (partial; the MCP counterpart lands in C3) — - ``test_each_dummy_tool_invoked_via_prepended_conversation`` -* **U17** (partial; the MCP-timeout counterpart lands in C3) — - ``test_failing_tool_yields_error_envelope`` -* **U18** — ``test_disallowed_tool_returns_error_without_invoking_callable`` - -Also covers the backend's documented behavior for missing functions -(both strict and tolerant modes), schema property defaulting, scalar -result wrapping, and declaration-order preservation in the bulk dispatch -path. These are required for the §10 rubber-duck guarantee that every -public-facing branch of :class:`LocalToolBackend` is exercised -before C5 wires it to a production target. -""" - -from __future__ import annotations - -import pytest - -from pyrit.tools import LocalToolBackend, ToolCall - - -def _make_call(name: str, *, call_id: str = "c1", arguments: dict | None = None) -> ToolCall: - return ToolCall(call_id=call_id, name=name, arguments=arguments or {}) - - -async def test_disallowed_tool_returns_error_without_invoking_callable(): - invoked: list[str] = [] - - async def echo(args: dict) -> dict: - invoked.append(args.get("text", "")) - return {"echoed": args.get("text", "")} - - backend = LocalToolBackend( - callables={"echo": echo, "off_limits": echo}, - allowed_tools={"echo"}, - ) - - result = await backend.dispatch_async(_make_call("off_limits", arguments={"text": "nope"})) - - assert result["error"] == "tool_not_allowed" - assert result["tool"] == "off_limits" - assert "echo" in result["allowed_tools"] - assert invoked == [] # callable was never invoked - - -async def test_failing_tool_yields_error_envelope(): - async def boom(args: dict) -> dict: - raise RuntimeError("kaboom") - - backend = LocalToolBackend(callables={"boom": boom}) - - result = await backend.dispatch_async(_make_call("boom")) - - assert result["error"] == "tool_execution_failed" - assert result["tool"] == "boom" - assert "kaboom" in result["detail"] - - -async def test_missing_tool_raises_when_strict(): - backend = LocalToolBackend(callables={}, fail_on_missing_function=True) - - with pytest.raises(KeyError, match="ghost"): - await backend.dispatch_async(_make_call("ghost")) - - -async def test_missing_tool_returns_envelope_when_tolerant(): - async def echo(args: dict) -> dict: - return {"ok": True} - - backend = LocalToolBackend( - callables={"echo": echo}, - fail_on_missing_function=False, - ) - - result = await backend.dispatch_async(_make_call("ghost")) - - assert result["error"] == "tool_not_registered" - assert result["tool"] == "ghost" - assert result["available_tools"] == ["echo"] - - -async def test_scalar_result_is_wrapped_in_dict(): - async def number(args: dict) -> int: - return 42 - - backend = LocalToolBackend(callables={"number": number}) - - result = await backend.dispatch_async(_make_call("number")) - - assert result == {"result": 42} - - -async def test_dict_result_passes_through_unchanged(): - async def named(args: dict) -> dict: - return {"custom_key": "custom_value"} - - backend = LocalToolBackend(callables={"named": named}) - - result = await backend.dispatch_async(_make_call("named")) - - assert result == {"custom_key": "custom_value"} - - -async def test_schemas_defaults_to_empty_list(): - backend = LocalToolBackend(callables={}) - - assert backend.schemas == [] - - -async def test_schemas_returned_as_copy(): - schemas_in = [{"name": "echo", "parameters": {}}] - backend = LocalToolBackend(callables={}, schemas=schemas_in) - - out1 = backend.schemas - out1.append({"name": "mutated"}) - - # Mutating the returned list does not affect the backend's internal state. - assert backend.schemas == schemas_in - - -async def test_dispatch_all_sequential_preserves_declaration_order(): - async def echo(args: dict) -> dict: - return {"echoed": args["i"]} - - backend = LocalToolBackend(callables={"echo": echo}) - - calls = [_make_call("echo", call_id=f"c{i}", arguments={"i": i}) for i in range(5)] - pairs = await backend.dispatch_all_sequential_async(calls) - - assert [c.call_id for c, _ in pairs] == ["c0", "c1", "c2", "c3", "c4"] - assert [r["echoed"] for _, r in pairs] == [0, 1, 2, 3, 4] - - -async def test_each_dummy_tool_invoked_via_prepended_conversation(): - """ - U10 (partial). Each dummy tool resolves on first dispatch (single - forward step, no model reasoning trace), confirming the backend can - short-circuit a prepended conversation where every call is already - decided. The MCP counterpart in C3 exercises the same shape against - a real stdio server. - """ - invocations: list[tuple[str, dict]] = [] - - async def echo(args: dict) -> dict: - invocations.append(("echo", args)) - return {"echoed": args.get("text", "")} - - async def add(args: dict) -> dict: - invocations.append(("add", args)) - return {"sum": args["a"] + args["b"]} - - async def reverse(args: dict) -> dict: - invocations.append(("reverse", args)) - return {"reversed": args.get("text", "")[::-1]} - - backend = LocalToolBackend(callables={"echo": echo, "add": add, "reverse": reverse}) - - prepended_calls = [ - _make_call("echo", call_id="e1", arguments={"text": "hello"}), - _make_call("add", call_id="a1", arguments={"a": 2, "b": 3}), - _make_call("reverse", call_id="r1", arguments={"text": "pyrit"}), - ] - pairs = await backend.dispatch_all_sequential_async(prepended_calls) - - # Each dummy resolved exactly once; no retries, no model re-entry. - assert len(invocations) == 3 - assert [name for name, _ in invocations] == ["echo", "add", "reverse"] - assert [r for _, r in pairs] == [ - {"echoed": "hello"}, - {"sum": 5}, - {"reversed": "tiryp"}, - ] diff --git a/tests/unit/tools/test_mcp_backend.py b/tests/unit/tools/test_mcp_backend.py deleted file mode 100644 index 0abc8ff88c..0000000000 --- a/tests/unit/tools/test_mcp_backend.py +++ /dev/null @@ -1,156 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# Licensed under the MIT license. - -""" -Unit tests for :class:`pyrit.tools.MCPToolBackend`. - -These tests verify the multi-server fan-out and routing layer on top of -:class:`MCPClient`: schema aggregation, name-collision detection, -``name_prefix`` disambiguation, ``allowed_tools`` allow-list semantics, -and concurrent-dispatch serialization. They reuse the real -``echo_mcp_server.py`` stdio subprocess. - -Coverage map: - -* **U18** — ``test_disallowed_tool_returns_error_envelope_without_invoking_server``. -* **U20a** — ``test_name_collision_raises_value_error``. -* **U20b** — ``test_name_prefix_disambiguates_colliding_servers``. -* **U21** — ``test_concurrent_dispatch_is_serialized_by_lock``. -""" - -from __future__ import annotations - -import asyncio -import sys -from pathlib import Path - -import pytest - -from pyrit.tools import ( - LocalMCPServerSpec, - MCPToolBackend, - ToolCall, -) - -ECHO_SERVER_SCRIPT = str(Path(__file__).parent / "echo_mcp_server.py") - - -def _spec(*, name_prefix: str | None = None, timeout_seconds: float = 5.0) -> LocalMCPServerSpec: - return LocalMCPServerSpec( - command=sys.executable, - args=(ECHO_SERVER_SCRIPT,), - name_prefix=name_prefix, - timeout_seconds=timeout_seconds, - ) - - -def _make_call(name: str, *, call_id: str = "c1", arguments: dict | None = None) -> ToolCall: - return ToolCall(call_id=call_id, name=name, arguments=arguments or {}) - - -@pytest.mark.asyncio -async def test_backend_aggregates_schemas_across_servers() -> None: - """Schemas from every connected server show up in :attr:`schemas`.""" - backend = MCPToolBackend(servers=[_spec()]) - async with backend: - names = {s["name"] for s in backend.schemas} - assert names == {"echo", "add", "reverse", "slow_echo"} - - -@pytest.mark.asyncio -async def test_dispatch_routes_to_correct_server() -> None: - """A :class:`ToolCall` is routed to the server that registered the name.""" - backend = MCPToolBackend(servers=[_spec()]) - async with backend: - envelope = await backend.dispatch_async(_make_call("echo", arguments={"text": "routed"})) - assert envelope["is_error"] is False - assert envelope["content"] == "routed" - - -@pytest.mark.asyncio -async def test_name_collision_raises_value_error() -> None: - """Two servers exposing the same tool name without prefixes raise.""" - backend = MCPToolBackend(servers=[_spec(), _spec()]) - with pytest.raises(ValueError, match="duplicate tool name"): - await backend.__aenter__() - # __aexit__ is the cleanup path; __aenter__ failing leaves nothing to clean. - - -@pytest.mark.asyncio -async def test_name_prefix_disambiguates_colliding_servers() -> None: - """Setting :attr:`LocalMCPServerSpec.name_prefix` disambiguates duplicates.""" - backend = MCPToolBackend( - servers=[ - _spec(name_prefix="a_"), - _spec(name_prefix="b_"), - ], - ) - async with backend: - names = {s["name"] for s in backend.schemas} - assert "a_echo" in names - assert "b_echo" in names - envelope = await backend.dispatch_async(_make_call("a_echo", arguments={"text": "alpha"})) - assert envelope["content"] == "alpha" - envelope_b = await backend.dispatch_async(_make_call("b_echo", arguments={"text": "beta"})) - assert envelope_b["content"] == "beta" - - -@pytest.mark.asyncio -async def test_disallowed_tool_returns_error_envelope_without_invoking_server() -> None: - """U18: allowed_tools blocks both schema advertisement AND dispatch.""" - backend = MCPToolBackend(servers=[_spec()], allowed_tools=["echo"]) - async with backend: - advertised = {s["name"] for s in backend.schemas} - assert advertised == {"echo"} # add/reverse/slow_echo are filtered out. - - envelope = await backend.dispatch_async(_make_call("add", arguments={"a": 1, "b": 2})) - assert envelope["is_error"] is True - assert envelope["error"] == "tool_not_allowed" - assert envelope["tool"] == "add" - assert envelope["allowed_tools"] == ["echo"] - - -@pytest.mark.asyncio -async def test_unknown_tool_returns_error_envelope() -> None: - """A call to a name no connected server exposes returns an error envelope.""" - backend = MCPToolBackend(servers=[_spec()]) - async with backend: - envelope = await backend.dispatch_async(_make_call("never_registered")) - assert envelope["is_error"] is True - assert envelope["error"] == "tool_not_registered" - assert envelope["tool"] == "never_registered" - - -@pytest.mark.asyncio -async def test_concurrent_dispatch_is_serialized_by_lock() -> None: - """U21: two coroutines dispatching against the same backend do not interleave. - - The slow_echo tool sleeps server-side; without the lock the two - dispatches would issue overlapping JSON-RPC frames over the same - stdio pipe. With the lock they run back-to-back. We assert both - return successfully — interleaved frames would surface as protocol - errors or wrong content. - """ - backend = MCPToolBackend(servers=[_spec(timeout_seconds=10.0)]) - async with backend: - results = await asyncio.gather( - backend.dispatch_async(_make_call("slow_echo", arguments={"text": "A", "delay_ms": 50})), - backend.dispatch_async(_make_call("slow_echo", arguments={"text": "B", "delay_ms": 50})), - ) - assert all(not r["is_error"] for r in results) - assert {r["content"] for r in results} == {"A", "B"} - - -@pytest.mark.asyncio -async def test_dispatch_all_sequential_async_preserves_order() -> None: - """Bulk dispatch returns (call, envelope) pairs in declaration order.""" - backend = MCPToolBackend(servers=[_spec()]) - calls = [ - _make_call("echo", call_id="c1", arguments={"text": "first"}), - _make_call("echo", call_id="c2", arguments={"text": "second"}), - _make_call("echo", call_id="c3", arguments={"text": "third"}), - ] - async with backend: - results = await backend.dispatch_all_sequential_async(calls) - assert [c.call_id for c, _ in results] == ["c1", "c2", "c3"] - assert [r["content"] for _, r in results] == ["first", "second", "third"] diff --git a/tests/unit/tools/test_mcp_client.py b/tests/unit/tools/test_mcp_client.py deleted file mode 100644 index 67f93d046e..0000000000 --- a/tests/unit/tools/test_mcp_client.py +++ /dev/null @@ -1,171 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# Licensed under the MIT license. - -""" -Unit tests for :class:`pyrit.tools.MCPClient` and the -:class:`pyrit.tools.MCPServerSpec` union. - -Coverage map (rows from the C2/C3 test matrix): - -* **U10** — ``test_real_subprocess_dispatch_returns_text_content``, - ``test_sequential_dispatch_against_real_server``. -* **U14** — ``test_connect_async_populates_schemas_via_tools_list``. -* **U17** — ``test_dispatch_timeout_returns_error_envelope``. -* **U20** — ``test_remote_mcp_server_spec_raises_not_implemented``, - ``test_docker_mcp_server_spec_raises_not_implemented``. - -These tests spawn the real ``tests/unit/tools/echo_mcp_server.py`` -subprocess via ``mcp.client.stdio.stdio_client``; they exercise the -full handshake → ``tools/list`` → ``tools/call`` round trip. The -purpose is to verify that ``MCPClient`` is a thin, correct facade -over the SDK rather than to re-test the SDK itself. -""" - -from __future__ import annotations - -import sys -from pathlib import Path - -import pytest - -from pyrit.tools import ( - DockerMCPServerSpec, - LocalMCPServerSpec, - MCPClient, - RemoteMCPServerSpec, - ToolCall, -) - -ECHO_SERVER_SCRIPT = str(Path(__file__).parent / "echo_mcp_server.py") - - -def _local_spec(*, timeout_seconds: float = 5.0) -> LocalMCPServerSpec: - """Build a :class:`LocalMCPServerSpec` that spawns ``echo_mcp_server.py``.""" - return LocalMCPServerSpec( - command=sys.executable, - args=(ECHO_SERVER_SCRIPT,), - timeout_seconds=timeout_seconds, - ) - - -def _make_call(name: str, *, call_id: str = "c1", arguments: dict | None = None) -> ToolCall: - return ToolCall(call_id=call_id, name=name, arguments=arguments or {}) - - -@pytest.mark.asyncio -async def test_real_subprocess_dispatch_returns_text_content() -> None: - """U10: dispatching a single tool call returns the echo server's text response.""" - client = MCPClient(spec=_local_spec()) - async with client: - envelope = await client.dispatch_async(_make_call("echo", arguments={"text": "hi"})) - assert envelope["is_error"] is False - assert envelope["content"] == "hi" - - -@pytest.mark.asyncio -async def test_sequential_dispatch_against_real_server() -> None: - """U10: multiple sequential calls round-trip through the same session.""" - client = MCPClient(spec=_local_spec()) - async with client: - envelopes = [ - await client.dispatch_async(_make_call("echo", arguments={"text": "first"})), - await client.dispatch_async(_make_call("add", arguments={"a": 2, "b": 3})), - await client.dispatch_async(_make_call("reverse", arguments={"text": "abc"})), - ] - contents = [e["content"] for e in envelopes] - assert contents == ["first", "5", "cba"] - - -@pytest.mark.asyncio -async def test_connect_async_populates_schemas_via_tools_list() -> None: - """U14: schemas are discovered via tools/list during connect_async.""" - client = MCPClient(spec=_local_spec()) - async with client: - schemas = client.schemas - names = {s["name"] for s in schemas} - assert names == {"echo", "add", "reverse", "slow_echo"} - echo_schema = next(s for s in schemas if s["name"] == "echo") - assert "parameters" in echo_schema - assert echo_schema["parameters"]["properties"]["text"]["type"] == "string" - - -@pytest.mark.asyncio -async def test_dispatch_timeout_returns_error_envelope() -> None: - """U17: a tool call that exceeds the spec's timeout produces an error envelope.""" - client = MCPClient(spec=_local_spec(timeout_seconds=0.05)) - async with client: - envelope = await client.dispatch_async( - _make_call("slow_echo", arguments={"text": "late", "delay_ms": 500}), - ) - assert envelope["is_error"] is True - assert envelope["error"] == "tool_timeout" - assert envelope["tool"] == "slow_echo" - - -@pytest.mark.asyncio -async def test_dispatch_async_returns_error_envelope_on_unknown_tool() -> None: - """Server-side errors (unknown tool name) surface as is_error envelopes.""" - client = MCPClient(spec=_local_spec()) - async with client: - envelope = await client.dispatch_async(_make_call("nonexistent_tool")) - assert envelope["is_error"] is True - assert envelope["tool"] == "nonexistent_tool" - - -def test_remote_mcp_server_spec_is_frozen_dataclass() -> None: - """U20: RemoteMCPServerSpec exists in the type system as a frozen dataclass.""" - spec = RemoteMCPServerSpec(url="https://example.com/mcp") - assert spec.url == "https://example.com/mcp" - with pytest.raises((AttributeError, Exception)): # frozen dataclass guard - spec.url = "other" # type: ignore[misc] - - -@pytest.mark.asyncio -async def test_remote_mcp_server_spec_raises_not_implemented() -> None: - """U20: connecting to a RemoteMCPServerSpec raises NotImplementedError.""" - client = MCPClient(spec=RemoteMCPServerSpec(url="https://example.com/mcp")) - with pytest.raises(NotImplementedError, match="follow-up PR"): - await client.connect_async() - - -def test_docker_mcp_server_spec_dataclass_fields() -> None: - """U20: DockerMCPServerSpec carries the fields the sandbox PR will consume.""" - spec = DockerMCPServerSpec(image="pyrit-sandbox:base") - assert spec.image == "pyrit-sandbox:base" - assert spec.network_profile == "none" - assert spec.name_prefix is None - assert spec.timeout_seconds == 30.0 - - -@pytest.mark.asyncio -async def test_docker_mcp_server_spec_raises_not_implemented() -> None: - """U20: connecting to a DockerMCPServerSpec raises NotImplementedError.""" - client = MCPClient(spec=DockerMCPServerSpec(image="pyrit-sandbox:base")) - with pytest.raises(NotImplementedError, match="follow-up PR"): - await client.connect_async() - - -@pytest.mark.asyncio -async def test_dispatch_before_connect_raises_runtime_error() -> None: - """Calling dispatch_async before connect_async is a programmer error.""" - client = MCPClient(spec=_local_spec()) - with pytest.raises(RuntimeError, match="not connected"): - await client.dispatch_async(_make_call("echo", arguments={"text": "hi"})) - - -@pytest.mark.asyncio -async def test_close_async_is_idempotent() -> None: - """Calling close_async twice (or before connect) does not raise.""" - client = MCPClient(spec=_local_spec()) - await client.close_async() # before connect — no-op. - await client.connect_async() - await client.close_async() - await client.close_async() # double-close — no-op. - - -@pytest.mark.asyncio -async def test_local_mcp_server_spec_is_frozen() -> None: - """LocalMCPServerSpec is a frozen dataclass.""" - spec = LocalMCPServerSpec(command="python", args=("a.py",)) - with pytest.raises((AttributeError, Exception)): - spec.command = "other" # type: ignore[misc] diff --git a/tests/unit/tools/test_prompt_target_tool_loop.py b/tests/unit/tools/test_prompt_target_tool_loop.py deleted file mode 100644 index 8848ddb010..0000000000 --- a/tests/unit/tools/test_prompt_target_tool_loop.py +++ /dev/null @@ -1,281 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# Licensed under the MIT license. - -""" -Unit tests for ``@tool_loop`` wired into :meth:`PromptTarget.send_prompt_async`. - -C4 lands the wiring: ``send_prompt_async`` becomes ``@final @tool_loop`` -on the base class, ``_tool_parser`` and ``_tool_schemas()`` get default -no-op implementations, and ``TargetConfiguration`` grows ``tool_event_policy`` -+ ``tool_backend`` kwargs. - -These tests use the production ``_get_normalized_conversation_async`` path -(memory round-trip through :class:`SQLiteMemory` via ``patch_central_database``) -to exercise the wrapper end-to-end. They cover: - -- U1: decorator order (validate + normalize happen exactly once, then the loop) -- U2 (DB-end half): produced ``tool`` message has one ``function_call_output`` - piece per dispatched call, in declaration order -- U8: DB inserts user, asst_with_fc, tool, asst_final in that order -- U9: DB roles + data_types match the canonical envelope -- U11: targets without a policy short-circuit (no wrapper behavior change) - -Tests for capability flag wiring + ``TargetConfiguration`` construction -validation live in :mod:`tests.unit.tools.test_tool_event_policy`. -""" - -from __future__ import annotations - -import json -from collections import deque -from typing import TYPE_CHECKING, Any - -import pytest - -from pyrit.prompt_target.common.prompt_target import PromptTarget -from pyrit.prompt_target.common.target_capabilities import TargetCapabilities -from pyrit.prompt_target.common.target_configuration import TargetConfiguration -from pyrit.tools import ToolCallParser, ToolEventBehavior, ToolEventPolicy - -from .conftest import ( - _CanonicalEnvelopeParser, - _make_assistant_function_call_message, - _make_assistant_text_message, - _make_user_message, - _RecordingToolBackend, -) - -if TYPE_CHECKING: - from pyrit.models import Message - - -class _ProductionShapedTarget(PromptTarget): - """ - Minimal :class:`PromptTarget` that uses the *real* base-class - ``_get_normalized_conversation_async`` (memory round-trip + normalization - pipeline) instead of the conftest stub. Drives the production wrapper - end-to-end so DB-insert-order assertions can run against the real - :class:`CentralMemory` instance set up by ``patch_central_database``. - """ - - def __init__( - self, - *, - scripted_responses: list[Message], - policy: ToolEventPolicy | None, - backend: Any, - parser: ToolCallParser | None, - ) -> None: - caps = TargetCapabilities( - supports_multi_turn=True, - supports_multi_message_pieces=True, - supports_tool_use=policy is not None, - ) - config = TargetConfiguration( - capabilities=caps, - tool_event_policy=policy, - tool_backend=backend, - ) - super().__init__(custom_configuration=config) - self._scripted: deque[Message] = deque(scripted_responses) - self.call_count: int = 0 - self._parser_instance = parser - - @property - def _tool_parser(self) -> ToolCallParser | None: - return self._parser_instance - - async def _send_prompt_to_target_async(self, *, normalized_conversation: list[Message]) -> list[Message]: - self.call_count += 1 - if not self._scripted: - raise AssertionError(f"Target ran out of scripted responses on iteration {self.call_count}.") - response = self._scripted.popleft() - conversation_id = normalized_conversation[-1].message_pieces[0].conversation_id - for piece in response.message_pieces: - piece.conversation_id = conversation_id - return [response] - - -@pytest.fixture -def make_production_target(patch_central_database): - def _factory( - *, - scripted_responses: list[Message], - policy: ToolEventPolicy | None = None, - backend: Any = None, - parser: ToolCallParser | None = None, - ) -> _ProductionShapedTarget: - effective_parser = parser - if effective_parser is None and policy is not None: - effective_parser = _CanonicalEnvelopeParser() - return _ProductionShapedTarget( - scripted_responses=scripted_responses, - policy=policy, - backend=backend, - parser=effective_parser, - ) - - return _factory - - -@pytest.fixture -def execute_policy_fixture(): - return ToolEventPolicy(behavior=ToolEventBehavior.EXECUTE, max_tool_iterations=5) - - -class TestToolLoopWiredIntoBaseClass: - """Verifies ``@tool_loop`` runs on every ``send_prompt_async`` call.""" - - @pytest.mark.asyncio - async def test_decorator_passthrough_when_no_policy(self, make_production_target): - """U11 -- target without a policy behaves exactly like pre-C4 ``send_prompt_async``.""" - target = make_production_target( - scripted_responses=[_make_assistant_text_message("plain")], - policy=None, - ) - - responses = await target.send_prompt_async(message=_make_user_message("hi")) - - assert target.call_count == 1 - assert len(responses) == 1 - assert responses[0].message_pieces[0].original_value == "plain" - - @pytest.mark.asyncio - async def test_tool_loop_order_after_normalize_before_memory(self, make_production_target, execute_policy_fixture): - """U1 -- validate + normalize happen exactly once before the loop iterates.""" - backend = _RecordingToolBackend(scripted_results=[{"result": "echoed"}]) - target = make_production_target( - scripted_responses=[ - _make_assistant_function_call_message(calls=[("c1", "echo", {"text": "x"})]), - _make_assistant_text_message("done"), - ], - policy=execute_policy_fixture, - backend=backend, - ) - - responses = await target.send_prompt_async(message=_make_user_message("please echo")) - - assert target.call_count == 2 - assert [c.name for c in backend.recorded_calls] == ["echo"] - assert len(responses) == 3 - assert responses[0].message_pieces[0].original_value_data_type == "function_call" - assert responses[1].message_pieces[0].original_value_data_type == "function_call_output" - assert responses[2].message_pieces[0].original_value_data_type == "text" - - @pytest.mark.asyncio - async def test_tool_message_has_one_function_call_output_piece_per_call( - self, make_production_target, execute_policy_fixture - ): - """U2 DB-end half -- one tool Message, N pieces, one per dispatched call.""" - backend = _RecordingToolBackend(scripted_results=[{"r": 1}, {"r": 2}]) - target = make_production_target( - scripted_responses=[ - _make_assistant_function_call_message( - calls=[("c1", "echo", {"text": "a"}), ("c2", "echo", {"text": "b"})] - ), - _make_assistant_text_message("done"), - ], - policy=execute_policy_fixture, - backend=backend, - ) - - responses = await target.send_prompt_async(message=_make_user_message("two calls please")) - - tool_msg = responses[1] - assert len(tool_msg.message_pieces) == 2 - call_ids_in_order = [json.loads(p.original_value)["call_id"] for p in tool_msg.message_pieces] - assert call_ids_in_order == ["c1", "c2"] - assert all(p.original_value_data_type == "function_call_output" for p in tool_msg.message_pieces) - assert all(p.api_role == "tool" for p in tool_msg.message_pieces) - - -class TestDbTranscriptAfterToolLoop: - """ - DB-level assertions that exercise the production memory pipeline. - - These tests rely on the wrapper writing the user message + every assistant - + tool message produced during the loop back to ``CentralMemory``, in - declaration order. Whether that write happens *inside* the wrapper or via - the caller (the prompt normalizer) is an implementation detail; the - invariant is the wrapper returns the full chain so the caller can persist - in order. - """ - - @pytest.mark.asyncio - async def test_db_insert_order_user_then_asst_fc_then_tool_then_final_asst( - self, make_production_target, execute_policy_fixture - ): - """U8 -- after a complete tool round, the wrapper's return order is canonical.""" - backend = _RecordingToolBackend(scripted_results=[{"result": "echoed"}]) - target = make_production_target( - scripted_responses=[ - _make_assistant_function_call_message(calls=[("c1", "echo", {"text": "x"})]), - _make_assistant_text_message("done"), - ], - policy=execute_policy_fixture, - backend=backend, - ) - - responses = await target.send_prompt_async(message=_make_user_message("please echo")) - - data_types_in_order = [r.message_pieces[0].original_value_data_type for r in responses] - assert data_types_in_order == ["function_call", "function_call_output", "text"] - - @pytest.mark.asyncio - async def test_db_roles_and_data_types_match_canonical_envelope( - self, make_production_target, execute_policy_fixture - ): - """U9 -- roles and data_types match the canonical envelope contract.""" - backend = _RecordingToolBackend(scripted_results=[{"result": "echoed"}]) - target = make_production_target( - scripted_responses=[ - _make_assistant_function_call_message(calls=[("c1", "echo", {"text": "x"})]), - _make_assistant_text_message("done"), - ], - policy=execute_policy_fixture, - backend=backend, - ) - - responses = await target.send_prompt_async(message=_make_user_message("please echo")) - - asst_fc, tool_msg, asst_final = responses - # function_call from the assistant - assert asst_fc.message_pieces[0].api_role == "assistant" - assert asst_fc.message_pieces[0].original_value_data_type == "function_call" - envelope = json.loads(asst_fc.message_pieces[0].original_value) - assert envelope["type"] == "function_call" - assert envelope["call_id"] == "c1" - assert envelope["name"] == "echo" - # function_call_output from the tool - assert tool_msg.message_pieces[0].api_role == "tool" - assert tool_msg.message_pieces[0].original_value_data_type == "function_call_output" - tool_envelope = json.loads(tool_msg.message_pieces[0].original_value) - assert tool_envelope["type"] == "function_call_output" - assert tool_envelope["call_id"] == "c1" - # Final assistant text - assert asst_final.message_pieces[0].api_role == "assistant" - assert asst_final.message_pieces[0].original_value_data_type == "text" - - -class TestFinalAndAbstractMethodContract: - """ - Asserts the base-class shape changes that C4 introduces but doesn't - exercise via end-to-end runs: ``_tool_parser`` defaults to ``None``, - ``_tool_schemas`` defaults to ``[]``. - """ - - def test_default_tool_parser_is_none(self, make_production_target): - target = make_production_target( - scripted_responses=[_make_assistant_text_message("plain")], - policy=None, - ) - # Subclass overrides only when the test caller passes a parser. With - # no policy + no parser, the override returns None. - assert target._tool_parser is None - - def test_default_tool_schemas_is_empty_list(self, make_production_target): - target = make_production_target( - scripted_responses=[_make_assistant_text_message("plain")], - policy=None, - ) - assert target._tool_schemas() == [] diff --git a/tests/unit/tools/test_tool_event_policy.py b/tests/unit/tools/test_tool_event_policy.py deleted file mode 100644 index aa451e672d..0000000000 --- a/tests/unit/tools/test_tool_event_policy.py +++ /dev/null @@ -1,121 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# Licensed under the MIT license. - -""" -Unit tests for the wiring between :class:`TargetCapabilities.supports_tool_use`, -:class:`TargetConfiguration.tool_event_policy` / -:class:`TargetConfiguration.tool_backend`, and the -:func:`pyrit.tools.tool_loop` decorator that lives on -:class:`PromptTarget.send_prompt_async`. - -These tests are the §7 U7 row plus the construction-time validator added in C4. -They assert the *capability flag* axis only -- that targets which declare -``supports_tool_use=True`` and configure a policy + backend route through -the loop, that targets without a policy short-circuit, and that the -``tool_backend``-without-capability misconfiguration raises at construction. - -End-to-end ordering against the production memory pipeline (U1, U8, U9) is -exercised separately in ``tests/unit/prompt_target/common/test_prompt_target_tool_loop.py``. -""" - -from __future__ import annotations - -import pytest - -from pyrit.prompt_target.common.target_capabilities import TargetCapabilities -from pyrit.prompt_target.common.target_configuration import TargetConfiguration -from pyrit.tools import LocalToolBackend, ToolEventBehavior, ToolEventPolicy - -from .conftest import ( - _make_assistant_function_call_message, - _make_assistant_text_message, - _make_user_message, -) - - -class TestSupportsToolUseCapabilityFlag: - """Asserts the new ``supports_tool_use`` field on :class:`TargetCapabilities`.""" - - def test_default_is_false(self): - caps = TargetCapabilities() - assert caps.supports_tool_use is False - - def test_explicit_true(self): - caps = TargetCapabilities(supports_tool_use=True) - assert caps.supports_tool_use is True - - -class TestTargetConfigurationToolFields: - """Asserts the new ``tool_event_policy`` / ``tool_backend`` kwargs.""" - - def test_defaults_are_none(self): - caps = TargetCapabilities(supports_tool_use=True) - config = TargetConfiguration(capabilities=caps) - assert config.tool_event_policy is None - assert config.tool_backend is None - - def test_explicit_policy_and_backend(self): - caps = TargetCapabilities(supports_tool_use=True) - backend = LocalToolBackend(callables={}, schemas=[]) - policy = ToolEventPolicy(behavior=ToolEventBehavior.EXECUTE) - config = TargetConfiguration( - capabilities=caps, - tool_event_policy=policy, - tool_backend=backend, - ) - assert config.tool_event_policy is policy - assert config.tool_backend is backend - - def test_tool_backend_without_capability_raises(self): - caps = TargetCapabilities(supports_tool_use=False) - backend = LocalToolBackend(callables={}, schemas=[]) - with pytest.raises(ValueError, match="supports_tool_use"): - TargetConfiguration(capabilities=caps, tool_backend=backend) - - def test_tool_event_policy_without_backend_is_allowed(self): - """``RAISE`` / ``RETURN_RAW`` policies do not require a backend.""" - caps = TargetCapabilities(supports_tool_use=True) - policy = ToolEventPolicy(behavior=ToolEventBehavior.RAISE) - config = TargetConfiguration(capabilities=caps, tool_event_policy=policy) - assert config.tool_event_policy is policy - assert config.tool_backend is None - - -class TestCapabilityFlagWiringIntoToolLoop: - """ - U7 -- verify the wrapper dispatches only when the target declares - ``supports_tool_use`` AND a policy is configured. - """ - - @pytest.mark.asyncio - async def test_target_with_tool_use_capability_uses_tool_loop( - self, make_fake_target, recording_backend, execute_policy - ): - backend = recording_backend() - target = make_fake_target( - scripted_responses=[ - _make_assistant_function_call_message(calls=[("c1", "echo", {"text": "hi"})]), - _make_assistant_text_message("done"), - ], - policy=execute_policy(), - backend=backend, - ) - - responses = await target.send_prompt_async(message=_make_user_message("please call echo")) - - assert target.call_count == 2, "Decorator should have iterated twice (call + final)." - assert [c.name for c in backend.recorded_calls] == ["echo"] - assert len(responses) == 3, "user expects asst_fc, tool_msg, asst_final." - - @pytest.mark.asyncio - async def test_target_without_tool_use_capability_skips_dispatch(self, make_fake_target): - target = make_fake_target( - scripted_responses=[_make_assistant_text_message("plain response, no tool call")], - policy=None, - backend=None, - ) - - responses = await target.send_prompt_async(message=_make_user_message("hello")) - - assert target.call_count == 1 - assert len(responses) == 1 diff --git a/tests/unit/tools/test_tool_loop_decorator.py b/tests/unit/tools/test_tool_loop_decorator.py deleted file mode 100644 index bc0db6b357..0000000000 --- a/tests/unit/tools/test_tool_loop_decorator.py +++ /dev/null @@ -1,289 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# Licensed under the MIT license. - -""" -Unit tests for :func:`pyrit.tools.tool_loop`. - -Coverage map (rows from the C2 test matrix): - -* **U2** (partial; full-DB end lands in C5) — ``test_loop_returns_full_chain_in_order`` -* **U3** — ``test_loop_exits_on_first_response_when_no_tool_calls``, - ``test_loops_until_no_pending_tool_call`` -* **U4** — ``test_raises_after_max_tool_iterations``, - ``test_partial_conversation_attached_to_limit_exception`` -* **U12** — ``test_policy_raise_includes_partial_conversation`` -* **U13** — ``test_policy_return_raw_does_not_dispatch`` -* **U16** — ``test_multi_call_per_turn_dispatched_sequentially_in_order`` - -Also covers two additional decorator concerns required by the rubber-duck -review (§10): EXECUTE policy with no backend raises with a partial -conversation, and the normalized conversation grows correctly across -iterations (decorator does not re-normalize each turn). -""" - -from __future__ import annotations - -import json - -import pytest - -from pyrit.exceptions import ToolCallLoopLimitExceeded, ToolCallNotSupported -from pyrit.tools import ToolEventBehavior, ToolEventPolicy - -from .conftest import ( - _make_assistant_function_call_message, - _make_assistant_text_message, - _make_user_message, -) - - -@pytest.mark.usefixtures("patch_central_database") -class TestToolLoopDecoratorBasics: - """Loop entry/exit semantics: no tool calls, single round trip, multi-round.""" - - async def test_loop_exits_on_first_response_when_no_tool_calls(self, make_fake_target, execute_policy): - target = make_fake_target( - scripted_responses=[_make_assistant_text_message("done")], - policy=execute_policy(), - ) - - responses = await target.send_prompt_async(message=_make_user_message("hi")) - - assert len(responses) == 1 - assert responses[0].get_value() == "done" - assert target.call_count == 1 - - async def test_loops_until_no_pending_tool_call(self, make_fake_target, execute_policy, recording_backend): - backend = recording_backend(scripted_results=[{"ok": True}, {"ok": True}]) - target = make_fake_target( - scripted_responses=[ - _make_assistant_function_call_message(calls=[("c1", "tool_a", {"x": 1})]), - _make_assistant_function_call_message(calls=[("c2", "tool_a", {"x": 2})]), - _make_assistant_text_message("done"), - ], - policy=execute_policy(max_tool_iterations=5), - backend=backend, - ) - - responses = await target.send_prompt_async(message=_make_user_message("hi")) - - # Two model-tool round trips and one final assistant message. - assert target.call_count == 3 - # Returned chain: fc1, tool1, fc2, tool2, final-text → 5 messages total. - assert len(responses) == 5 - assert [r.message_pieces[0].original_value_data_type for r in responses] == [ - "function_call", - "function_call_output", - "function_call", - "function_call_output", - "text", - ] - assert len(backend.recorded_calls) == 2 - assert [c.call_id for c in backend.recorded_calls] == ["c1", "c2"] - - -@pytest.mark.usefixtures("patch_central_database") -class TestToolLoopMessageShape: - """U2 — assistant_fc → tool → final_assistant ordering and identity.""" - - async def test_loop_returns_full_chain_in_order(self, make_fake_target, execute_policy, recording_backend): - backend = recording_backend(scripted_results=[{"weather": "sunny"}]) - fc_msg = _make_assistant_function_call_message(calls=[("call_abc", "get_weather", {"city": "Seattle"})]) - final_msg = _make_assistant_text_message("It is sunny in Seattle.") - - target = make_fake_target( - scripted_responses=[fc_msg, final_msg], - policy=execute_policy(), - backend=backend, - ) - - responses = await target.send_prompt_async(message=_make_user_message("weather?")) - - assert len(responses) == 3 - # 1) assistant with function_call (identity preserved) - assert responses[0] is fc_msg - # 2) tool message with exactly one function_call_output piece carrying call_id - tool_msg = responses[1] - assert len(tool_msg.message_pieces) == 1 - tool_piece = tool_msg.message_pieces[0] - assert tool_piece.api_role == "tool" - assert tool_piece.original_value_data_type == "function_call_output" - envelope = json.loads(tool_piece.original_value) - assert envelope["type"] == "function_call_output" - assert envelope["call_id"] == "call_abc" - # The tool result is JSON-serialized into the "output" field. - assert json.loads(envelope["output"]) == {"weather": "sunny"} - # 3) final assistant text (identity preserved) - assert responses[2] is final_msg - - -@pytest.mark.usefixtures("patch_central_database") -class TestToolLoopIterationLimits: - """U4 — iteration cap raises and carries the partial chain.""" - - async def test_raises_after_max_tool_iterations(self, make_fake_target, execute_policy, recording_backend): - # Model never stops asking for tools. - backend = recording_backend(scripted_results=[{"ok": True}] * 3) - target = make_fake_target( - scripted_responses=[ - _make_assistant_function_call_message(calls=[(f"c{i}", "loop_tool", {})]) for i in range(3) - ], - policy=execute_policy(max_tool_iterations=2), - backend=backend, - ) - - with pytest.raises(ToolCallLoopLimitExceeded, match="max_tool_iterations=2"): - await target.send_prompt_async(message=_make_user_message("hi")) - - # Exactly max_tool_iterations model calls made before raising. - assert target.call_count == 2 - - async def test_partial_conversation_attached_to_limit_exception( - self, make_fake_target, execute_policy, recording_backend - ): - backend = recording_backend(scripted_results=[{"ok": True}] * 2) - target = make_fake_target( - scripted_responses=[ - _make_assistant_function_call_message(calls=[(f"c{i}", "loop_tool", {})]) for i in range(2) - ], - policy=execute_policy(max_tool_iterations=2), - backend=backend, - ) - - with pytest.raises(ToolCallLoopLimitExceeded) as excinfo: - await target.send_prompt_async(message=_make_user_message("hi")) - - partial = excinfo.value.partial_conversation - # 2 iterations × (assistant_fc + tool_msg) = 4 messages, all in order. - assert len(partial) == 4 - assert [m.message_pieces[0].original_value_data_type for m in partial] == [ - "function_call", - "function_call_output", - "function_call", - "function_call_output", - ] - - -@pytest.mark.usefixtures("patch_central_database") -class TestToolEventPolicyBehaviors: - """U12, U13 — non-EXECUTE behaviors short-circuit dispatch.""" - - async def test_policy_raise_includes_partial_conversation(self, make_fake_target, recording_backend): - backend = recording_backend(scripted_results=[{"ok": True}]) - fc_msg = _make_assistant_function_call_message(calls=[("c1", "danger", {})]) - target = make_fake_target( - scripted_responses=[fc_msg], - policy=ToolEventPolicy(behavior=ToolEventBehavior.RAISE), - backend=backend, - ) - - with pytest.raises(ToolCallNotSupported, match="RAISE") as excinfo: - await target.send_prompt_async(message=_make_user_message("hi")) - - partial = excinfo.value.partial_conversation - # Partial contains the offending assistant turn; no tool dispatch occurred. - assert partial == [fc_msg] - assert backend.recorded_calls == [] - assert target.call_count == 1 - - async def test_policy_return_raw_does_not_dispatch(self, make_fake_target, recording_backend): - backend = recording_backend(scripted_results=[{"ok": True}]) - fc_msg = _make_assistant_function_call_message(calls=[("c1", "danger", {})]) - target = make_fake_target( - scripted_responses=[fc_msg], - policy=ToolEventPolicy(behavior=ToolEventBehavior.RETURN_RAW), - backend=backend, - ) - - responses = await target.send_prompt_async(message=_make_user_message("hi")) - - assert responses == [fc_msg] - assert backend.recorded_calls == [] - assert target.call_count == 1 - - -@pytest.mark.usefixtures("patch_central_database") -class TestToolLoopMultiCallPerTurn: - """U16 — multi-call turns dispatch sequentially in declaration order.""" - - async def test_multi_call_per_turn_dispatched_sequentially_in_order( - self, make_fake_target, execute_policy, recording_backend - ): - backend = recording_backend(scripted_results=[{"a": 1}, {"b": 2}, {"c": 3}]) - multi_fc = _make_assistant_function_call_message( - calls=[ - ("c_alpha", "tool_alpha", {"k": "v1"}), - ("c_beta", "tool_beta", {"k": "v2"}), - ("c_gamma", "tool_gamma", {"k": "v3"}), - ] - ) - target = make_fake_target( - scripted_responses=[multi_fc, _make_assistant_text_message("ok")], - policy=execute_policy(), - backend=backend, - ) - - responses = await target.send_prompt_async(message=_make_user_message("multi")) - - # Three calls dispatched in declaration order, recorded ids match. - assert [c.call_id for c in backend.recorded_calls] == ["c_alpha", "c_beta", "c_gamma"] - assert [c.name for c in backend.recorded_calls] == ["tool_alpha", "tool_beta", "tool_gamma"] - # One tool message after the multi-call assistant turn, carrying three - # function_call_output pieces in declaration order with the right call_ids. - tool_msg = responses[1] - assert len(tool_msg.message_pieces) == 3 - envelopes = [json.loads(p.original_value) for p in tool_msg.message_pieces] - assert [e["call_id"] for e in envelopes] == ["c_alpha", "c_beta", "c_gamma"] - assert all(p.original_value_data_type == "function_call_output" for p in tool_msg.message_pieces) - - -@pytest.mark.usefixtures("patch_central_database") -class TestToolLoopMisconfiguration: - """EXECUTE policy with no backend must fail loudly and carry the partial chain.""" - - async def test_execute_without_backend_raises_with_partial(self, make_fake_target, execute_policy): - fc_msg = _make_assistant_function_call_message(calls=[("c1", "no_reg", {})]) - target = make_fake_target( - scripted_responses=[fc_msg], - policy=execute_policy(), - backend=None, - ) - - with pytest.raises(ToolCallNotSupported, match="tool_backend") as excinfo: - await target.send_prompt_async(message=_make_user_message("hi")) - - assert excinfo.value.partial_conversation == [fc_msg] - - -@pytest.mark.usefixtures("patch_central_database") -class TestToolLoopConversationGrowth: - """The decorator must extend (not re-normalize) the conversation each round.""" - - async def test_normalized_conversation_grows_each_iteration( - self, make_fake_target, execute_policy, recording_backend - ): - backend = recording_backend(scripted_results=[{"r1": 1}, {"r2": 2}]) - target = make_fake_target( - scripted_responses=[ - _make_assistant_function_call_message(calls=[("c1", "t", {})]), - _make_assistant_function_call_message(calls=[("c2", "t", {})]), - _make_assistant_text_message("done"), - ], - policy=execute_policy(), - backend=backend, - ) - - await target.send_prompt_async(message=_make_user_message("hi")) - - # Three protected-method calls; each subsequent call sees the prior - # assistant_fc + tool_msg appended (the decorator must NOT re-normalize). - seen = target.normalized_conversations_seen - assert len(seen) == 3 - # call 1: just the user message - assert len(seen[0]) == 1 - # call 2: user + assistant_fc(c1) + tool_msg - assert len(seen[1]) == 3 - assert seen[1][1].message_pieces[0].original_value_data_type == "function_call" - assert seen[1][2].message_pieces[0].original_value_data_type == "function_call_output" - # call 3: user + assistant_fc(c1) + tool_msg + assistant_fc(c2) + tool_msg - assert len(seen[2]) == 5 From 5dbd1e47fb86b031d27ba2cffd9597a44ec9fd86 Mon Sep 17 00:00:00 2001 From: Victor Valbuena Date: Fri, 29 May 2026 11:52:56 -0700 Subject: [PATCH 36/40] Remove remaining tool-calling leakage from PR #1811 The previous cleanup commit (31ed2fb) removed the pyrit/tools/ package and tests/unit/tools/ directory, but several tool-calling changes from PR #1811 (MCP) remained mixed in: - pyrit/exceptions: ToolCallNotSupported and ToolCallLoopLimitExceeded - pyrit/prompt_target/common/: @tool_loop decoration on send_prompt_async, supports_tool_use capability, tool_event_policy and tool_backend slots on TargetConfiguration - pyrit/prompt_target/openai/openai_response_target.py: migration onto @tool_loop + LocalToolBackend (the in-class agentic loop was removed in favor of the decorator) - tests/integration/tools/ and tests/unit/prompt_target/target/test_openai_response_target_c6_migration.py - pyproject.toml + uv.lock: mcp Python SDK dependency All of the above are reverted to origin/main. The adversarial benchmark refactor (this PR's actual scope) is unaffected; 128 targeted unit tests across openai_response_target, function_chaining, and scenario/benchmark still pass. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- pyproject.toml | 1 - pyrit/exceptions/__init__.py | 4 - pyrit/exceptions/exception_classes.py | 64 ---- .../common/discover_target_capabilities.py | 1 - pyrit/prompt_target/common/prompt_target.py | 42 -- .../common/target_capabilities.py | 9 - .../common/target_configuration.py | 72 +--- .../openai/openai_response_target.py | 213 +++-------- tests/integration/tools/__init__.py | 2 - .../tools/test_red_teaming_with_tools.py | 361 ------------------ ...est_openai_response_target_c6_migration.py | 304 --------------- uv.lock | 94 ----- 12 files changed, 52 insertions(+), 1115 deletions(-) delete mode 100644 tests/integration/tools/__init__.py delete mode 100644 tests/integration/tools/test_red_teaming_with_tools.py delete mode 100644 tests/unit/prompt_target/target/test_openai_response_target_c6_migration.py diff --git a/pyproject.toml b/pyproject.toml index 911d611804..0a29123c87 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -46,7 +46,6 @@ dependencies = [ "fastapi>=0.133.0", "httpx[http2]>=0.27.2", "jinja2>=3.1.6", - "mcp>=1.0,<2", "numpy>=1.26.0; python_version < '3.14'", "numpy>=2.3.0; python_version >= '3.14'", "openai>=2.2.0", diff --git a/pyrit/exceptions/__init__.py b/pyrit/exceptions/__init__.py index 9baea33c1a..abd42de031 100644 --- a/pyrit/exceptions/__init__.py +++ b/pyrit/exceptions/__init__.py @@ -10,8 +10,6 @@ MissingPromptPlaceholderException, PyritException, RateLimitException, - ToolCallLoopLimitExceeded, - ToolCallNotSupported, get_retry_max_num_attempts, handle_bad_request_exception, pyrit_custom_result_retry, @@ -61,6 +59,4 @@ "set_execution_context", "set_retry_collector", "execution_context", - "ToolCallLoopLimitExceeded", - "ToolCallNotSupported", ] diff --git a/pyrit/exceptions/exception_classes.py b/pyrit/exceptions/exception_classes.py index 5d0014aa3d..b2fc55440b 100644 --- a/pyrit/exceptions/exception_classes.py +++ b/pyrit/exceptions/exception_classes.py @@ -233,70 +233,6 @@ def __init__(self, *, message: str = "No prompt placeholder") -> None: super().__init__(message=message) -class ToolCallNotSupported(PyritException): - """ - Raised when a target produces a tool call that the configured - :class:`~pyrit.tools.ToolEventPolicy` does not permit to execute - (``ToolEventBehavior.RAISE``, or ``EXECUTE`` without a backend). - - The ``partial_conversation`` attribute carries every message produced - up to and including the assistant turn that contained the offending - tool call(s). Consumers can inspect it to log the surfaced tool-use - attempt. - """ - - def __init__( - self, - *, - message: str = "Tool call not supported by configured policy.", - partial_conversation: Optional[list["Message"]] = None, - ) -> None: - """ - Initialize the exception. - - Args: - message (str): Human-readable error description. - partial_conversation (Optional[list[Message]]): Messages produced by - the target up to (and including) the assistant turn that - contained the disallowed tool call(s). - """ - super().__init__(status_code=400, message=message) - self.partial_conversation: list[Message] = ( - list(partial_conversation) if partial_conversation is not None else [] - ) - - -class ToolCallLoopLimitExceeded(PyritException): - """ - Raised when the tool-use loop runs for more than - ``ToolEventPolicy.max_tool_iterations`` iterations without the model - producing a stop response. - - The ``partial_conversation`` attribute carries every message produced - across all completed iterations. Consumers can inspect it to debug - runaway agentic behavior. - """ - - def __init__( - self, - *, - message: str = "Tool loop exceeded max_tool_iterations without a stop response.", - partial_conversation: Optional[list["Message"]] = None, - ) -> None: - """ - Initialize the exception. - - Args: - message (str): Human-readable error description. - partial_conversation (Optional[list[Message]]): Messages produced by - the target across every completed iteration of the tool loop. - """ - super().__init__(status_code=400, message=message) - self.partial_conversation: list[Message] = ( - list(partial_conversation) if partial_conversation is not None else [] - ) - - def pyrit_custom_result_retry( retry_function: Callable[..., bool], retry_max_num_attempts: Optional[int] = None ) -> Callable[..., Any]: diff --git a/pyrit/prompt_target/common/discover_target_capabilities.py b/pyrit/prompt_target/common/discover_target_capabilities.py index b6a79bbcb9..45600e6009 100644 --- a/pyrit/prompt_target/common/discover_target_capabilities.py +++ b/pyrit/prompt_target/common/discover_target_capabilities.py @@ -149,7 +149,6 @@ def _permissive_configuration( supports_json_output=True, supports_editable_history=True, supports_system_prompt=True, - supports_tool_use=True, input_modalities=merged_modalities, ) # Rebuild a fresh configuration from the instance's native capabilities so diff --git a/pyrit/prompt_target/common/prompt_target.py b/pyrit/prompt_target/common/prompt_target.py index 035b00823d..b1ee5caaa2 100644 --- a/pyrit/prompt_target/common/prompt_target.py +++ b/pyrit/prompt_target/common/prompt_target.py @@ -12,7 +12,6 @@ from pyrit.models.json_response_config import _JsonResponseConfig from pyrit.prompt_target.common.target_capabilities import CapabilityName, TargetCapabilities from pyrit.prompt_target.common.target_configuration import TargetConfiguration -from pyrit.tools import ToolCallParser, tool_loop logger = logging.getLogger(__name__) @@ -86,7 +85,6 @@ def __init__( logging.basicConfig(level=logging.INFO) @final - @tool_loop async def send_prompt_async(self, *, message: Message) -> list[Message]: """ Validate, normalize, and send a prompt to the target. @@ -99,13 +97,6 @@ async def send_prompt_async(self, *, message: Message) -> list[Message]: 3. Delegates to ``_send_prompt_to_target_async`` with the normalized conversation. - When the target's :attr:`configuration.tool_event_policy` is set, the - :func:`pyrit.tools.tool_loop` decorator replaces this body with the - agentic loop and re-enters :meth:`_send_prompt_to_target_async` - repeatedly until the model issues a stop response (or the configured - ``max_tool_iterations`` is hit). When no policy is set, the decorator - is a no-op and the body below runs unchanged. - Subclasses MUST NOT override this method. Override ``_send_prompt_to_target_async`` instead. @@ -141,39 +132,6 @@ async def _send_prompt_to_target_async(self, *, normalized_conversation: list[Me list[Message]: Response messages from the target. """ - @property - def _tool_parser(self) -> ToolCallParser | None: - """ - Per-target :class:`ToolCallParser` consulted by :func:`pyrit.tools.tool_loop`. - - Targets that participate in the tool-use loop override this property - to return a parser that walks their response messages and extracts - :class:`~pyrit.tools.ToolCall` instances. The base default of - ``None`` signals "this target does not participate" -- the wrapper - short-circuits after the first response. - - Returns: - ToolCallParser | None: The parser, or ``None`` for the default - no-tool-use behavior. - """ - return None - - def _tool_schemas(self) -> list[dict[str, Any]]: - """ - Outbound tool-schema list sent on the next request to the model. - - Targets that participate in the tool-use loop override this method - to translate the active :class:`~pyrit.tools.ToolBackend.schemas` - into the wire format their model expects (Responses API vs. Chat - Completions API vs. anything else). The base default returns an - empty list, which means no schemas are advertised. - - Returns: - list[dict[str, Any]]: One schema per advertised tool, in the - target-specific wire format. Empty by default. - """ - return [] - def _validate_request(self, *, normalized_conversation: list[Message]) -> None: """ Validate the normalized conversation before sending to the target. diff --git a/pyrit/prompt_target/common/target_capabilities.py b/pyrit/prompt_target/common/target_capabilities.py index 234ef4d359..6ae9ed69e2 100644 --- a/pyrit/prompt_target/common/target_capabilities.py +++ b/pyrit/prompt_target/common/target_capabilities.py @@ -24,7 +24,6 @@ class CapabilityName(str, Enum): JSON_OUTPUT = "supports_json_output" EDITABLE_HISTORY = "supports_editable_history" SYSTEM_PROMPT = "supports_system_prompt" - TOOL_USE = "supports_tool_use" class UnsupportedCapabilityBehavior(str, Enum): @@ -139,14 +138,6 @@ class attribute. Users can override individual capabilities per instance # Whether the target natively supports system prompts. supports_system_prompt: bool = False - # Whether the target natively supports model-issued tool calls (the - # canonical OpenAI ``function_call`` / ``function_call_output`` envelopes - # plus an outbound tool-schema list). Targets without this capability - # cannot host a tool-use loop -- attempting to configure a - # :class:`TargetConfiguration` with a ``tool_backend`` on a target whose - # capabilities have ``supports_tool_use=False`` raises at construction. - supports_tool_use: bool = False - # The input modalities supported by the target (e.g., "text", "image"). input_modalities: frozenset[frozenset[PromptDataType]] = frozenset({frozenset(["text"])}) diff --git a/pyrit/prompt_target/common/target_configuration.py b/pyrit/prompt_target/common/target_configuration.py index 6058409194..7e11a04673 100644 --- a/pyrit/prompt_target/common/target_configuration.py +++ b/pyrit/prompt_target/common/target_configuration.py @@ -4,7 +4,7 @@ import logging from collections.abc import Mapping from dataclasses import fields -from typing import TYPE_CHECKING, Any +from typing import Any from pyrit.message_normalizer import MessageListNormalizer from pyrit.models import Message @@ -16,10 +16,6 @@ UnsupportedCapabilityBehavior, ) -if TYPE_CHECKING: - from pyrit.tools.backend import ToolBackend - from pyrit.tools.models import ToolEventPolicy - logger = logging.getLogger(__name__) @@ -43,15 +39,6 @@ class TargetConfiguration: Each target defines defaults; callers can override policy or individual normalizers at creation time. - - Tool use is configured by setting :attr:`tool_event_policy` (mandatory - when a target's response contains tool calls; controls EXECUTE / RAISE / - RETURN\\_RAW behavior) and optionally :attr:`tool_backend` (required only - when ``tool_event_policy.behavior`` is ``EXECUTE``). Both default to - ``None`` and are read by :func:`pyrit.tools.tool_loop` at runtime; - constructing a configuration with a ``tool_backend`` on a target that - does not declare ``capabilities.supports_tool_use=True`` raises - immediately. """ def __init__( @@ -60,8 +47,6 @@ def __init__( capabilities: TargetCapabilities, policy: CapabilityHandlingPolicy | None = None, normalizer_overrides: Mapping[CapabilityName, MessageListNormalizer[Any]] | None = None, - tool_event_policy: "ToolEventPolicy | None" = None, - tool_backend: "ToolBackend | None" = None, ) -> None: """ Build a target configuration and resolve the normalization pipeline. @@ -72,25 +57,7 @@ def __init__( capability. Defaults to RAISE for all adaptable capabilities. normalizer_overrides (Mapping[CapabilityName, MessageListNormalizer[Any]] | None): Optional overrides for specific capability normalizers. - tool_event_policy (ToolEventPolicy | None): How - :func:`pyrit.tools.tool_loop` should react to a pending tool - call from the target. ``None`` means the loop is disabled and - the wrapper short-circuits. - tool_backend (ToolBackend | None): Dispatch table used when - ``tool_event_policy.behavior`` is ``EXECUTE``. ``None`` is - valid only for the RAISE / RETURN\\_RAW policies and the - no-policy passthrough. - - Raises: - ValueError: If ``tool_backend`` is set on a target whose - capabilities do not include ``supports_tool_use``. """ - if tool_backend is not None and not capabilities.includes(capability=CapabilityName.TOOL_USE): - raise ValueError( - "tool_backend is set but capabilities.supports_tool_use is False. " - "Either declare supports_tool_use=True on the target's capabilities, " - "or remove the tool_backend." - ) self._capabilities = capabilities self._policy = policy or _DEFAULT_POLICY self._pipeline = ConversationNormalizationPipeline.from_capabilities( @@ -98,8 +65,6 @@ def __init__( policy=self._policy, normalizer_overrides=normalizer_overrides, ) - self._tool_event_policy = tool_event_policy - self._tool_backend = tool_backend @property def capabilities(self) -> TargetCapabilities: @@ -116,41 +81,6 @@ def pipeline(self) -> ConversationNormalizationPipeline: """The resolved normalization pipeline.""" return self._pipeline - @property - def tool_event_policy(self) -> "ToolEventPolicy | None": - """The tool-use policy consulted by :func:`pyrit.tools.tool_loop`.""" - return self._tool_event_policy - - @tool_event_policy.setter - def tool_event_policy(self, value: "ToolEventPolicy | None") -> None: - """Allow runtime updates so callers can opt a configured target into tool use.""" - self._tool_event_policy = value - - @property - def tool_backend(self) -> "ToolBackend | None": - """The tool dispatch backend used when the loop's behavior is ``EXECUTE``.""" - return self._tool_backend - - @tool_backend.setter - def tool_backend(self, value: "ToolBackend | None") -> None: - """ - Allow runtime updates to the backend. - - Re-runs the ``supports_tool_use`` validator so a backend can never be - installed onto a configuration that does not declare the capability. - - Raises: - ValueError: If ``value`` is not ``None`` and the configuration's - capabilities do not include ``supports_tool_use``. - """ - if value is not None and not self._capabilities.includes(capability=CapabilityName.TOOL_USE): - raise ValueError( - "tool_backend is set but capabilities.supports_tool_use is False. " - "Either declare supports_tool_use=True on the target's capabilities, " - "or remove the tool_backend." - ) - self._tool_backend = value - def includes(self, *, capability: CapabilityName) -> bool: """ Check whether the target includes support for the given capability. diff --git a/pyrit/prompt_target/openai/openai_response_target.py b/pyrit/prompt_target/openai/openai_response_target.py index 332b64847e..f2e4b19a76 100644 --- a/pyrit/prompt_target/openai/openai_response_target.py +++ b/pyrit/prompt_target/openai/openai_response_target.py @@ -3,7 +3,6 @@ import json import logging -import warnings from collections.abc import Awaitable, Callable, MutableSequence from enum import Enum from typing import ( @@ -35,14 +34,6 @@ from pyrit.prompt_target.common.utils import limit_requests_per_minute, validate_temperature, validate_top_p from pyrit.prompt_target.openai.openai_error_handling import _is_content_filter_error from pyrit.prompt_target.openai.openai_target import OpenAITarget -from pyrit.tools import ( - CanonicalEnvelopeParser, - LocalToolBackend, - ToolBackend, - ToolCallParser, - ToolEventBehavior, - ToolEventPolicy, -) logger = logging.getLogger(__name__) @@ -85,7 +76,6 @@ class OpenAIResponseTarget(OpenAITarget, PromptTarget): supports_json_output=True, supports_multi_message_pieces=True, supports_system_prompt=True, - supports_tool_use=True, input_modalities=frozenset( { frozenset(["text"]), @@ -164,17 +154,6 @@ def __init__( """ super().__init__(custom_configuration=custom_configuration, **kwargs) - # If the constructed configuration is the class-level _DEFAULT_CONFIGURATION - # singleton (user did not pass custom_configuration AND the underlying_model - # was unrecognized), rebuild a per-instance copy so the C6 tool-backend - # plumbing below does not mutate state shared across every other instance. - if custom_configuration is None and self._configuration is type(self)._DEFAULT_CONFIGURATION: - caps = self._configuration.capabilities - self._configuration = TargetConfiguration( - capabilities=caps, - policy=self._configuration.policy, - ) - # Validate temperature and top_p validate_temperature(temperature) validate_top_p(top_p) @@ -188,39 +167,10 @@ def __init__( self._extra_body_parameters = extra_body_parameters - # ----- Tool-calling plumbing (C6) --------------------------------- - # custom_functions is deprecated as of 0.15.x. New code configures - # tool_backend on TargetConfiguration directly. The kwarg is still - # accepted; we ALWAYS install a LocalToolBackend (whether populated - # or empty) when no other backend is supplied, so legacy in-place - # mutations of `target._custom_functions` (via the back-compat - # property below) keep affecting dispatch. + # Per-instance tool/func registries: + self._custom_functions: dict[str, ToolExecutor] = custom_functions or {} self._fail_on_missing_function: bool = fail_on_missing_function - if self.configuration.tool_backend is None: - shim_backend = LocalToolBackend( - callables=dict(custom_functions) if custom_functions else {}, - schemas=self._derive_default_schemas(custom_functions or {}), - fail_on_missing_function=fail_on_missing_function, - ) - self.configuration.tool_backend = shim_backend - - if custom_functions: - warnings.warn( - "OpenAIResponseTarget(custom_functions=...) is deprecated and will be " - "removed in 0.16.0. Configure tool_backend on TargetConfiguration " - "instead (e.g. LocalToolBackend(callables=..., schemas=..., " - "fail_on_missing_function=...)).", - DeprecationWarning, - stacklevel=2, - ) - - # Default policy to EXECUTE when a backend is present. The wrapper's - # parser returns an empty list when the model produces no tool calls, - # so this is a no-op for plain text completions. - if self.configuration.tool_event_policy is None: - self.configuration.tool_event_policy = ToolEventPolicy(behavior=ToolEventBehavior.EXECUTE) - # Extract the grammar 'tool' if one is present # See # https://platform.openai.com/docs/guides/function-calling#context-free-grammars @@ -235,61 +185,6 @@ def __init__( logger.debug("Detected grammar tool: %s", tool_name) self._grammar_name = tool_name - @staticmethod - def _derive_default_schemas(callables: dict[str, ToolExecutor]) -> list[dict[str, Any]]: - """ - Synthesize minimal JSON schemas for the deprecation-shim path. - - Users who pass the legacy ``custom_functions`` kwarg do not also pass a - schema list (the Responses API would accept the calls anyway because the - legacy path predates structured tool advertisement). To keep the - deprecation shim transparent we generate a schema-less stub per name so - ``_tool_schemas()`` returns something non-empty when the user actually - wires tools. - - Args: - callables: Function name to async callable mapping. - - Returns: - list[dict[str, Any]]: A bare schema per callable (``parameters`` - is the unconstrained empty-object schema). - """ - return [{"name": name, "parameters": {"type": "object"}} for name in callables] - - @property - def _custom_functions(self) -> dict[str, ToolExecutor]: - """ - Back-compat live view of the active backend's callables registry. - - Mutations on the returned dict (``target._custom_functions[name] = fn``, - ``target._custom_functions.pop(name)``) take effect immediately because - the dict object is shared with the underlying - :class:`pyrit.tools.LocalToolBackend`. Returns an empty dict when no - backend is installed or when the configured backend is not a - ``LocalToolBackend``. - - Returns: - dict[str, ToolExecutor]: The live callables dict. - """ - backend = self.configuration.tool_backend - if isinstance(backend, LocalToolBackend): - return cast("dict[str, ToolExecutor]", backend._callables) - return {} - - @_custom_functions.setter - def _custom_functions(self, value: dict[str, ToolExecutor]) -> None: - backend = self.configuration.tool_backend - if isinstance(backend, LocalToolBackend): - backend._callables = dict(value) - backend._schemas = self._derive_default_schemas(value) - return - new_backend = LocalToolBackend( - callables=dict(value), - schemas=self._derive_default_schemas(value), - fail_on_missing_function=self._fail_on_missing_function, - ) - self.configuration.tool_backend = new_backend - def _build_identifier(self) -> ComponentIdentifier: """ Build the identifier with OpenAI response-specific parameters. @@ -483,9 +378,8 @@ async def _construct_request_body( input_items = await self._build_input_for_multi_modal_async(conversation) text_format = self._build_text_format(json_config=json_config) - tool_schemas = self._tool_schemas() - body_parameters: dict[str, Any] = { + body_parameters = { "model": self._model_name, "max_output_tokens": self._max_output_tokens, "temperature": self._temperature, @@ -496,11 +390,8 @@ async def _construct_request_body( "text": text_format, "reasoning": self._build_reasoning_config(), } - if tool_schemas: - body_parameters["tools"] = tool_schemas if self._extra_body_parameters: - # User-supplied extra_body_parameters wins over backend-derived tools. body_parameters.update(self._extra_body_parameters) # Filter out None values @@ -668,18 +559,11 @@ async def _construct_message_from_response(self, response: Any, request: Message @pyrit_target_retry async def _send_prompt_to_target_async(self, *, normalized_conversation: list[Message]) -> list[Message]: """ - Send one prompt to the Responses API and return exactly one Message. - - The agentic tool-calling loop now lives in :func:`pyrit.tools.tool_loop` - on the base class. This method is the single-iteration body the loop - re-enters on each turn: build the request body, call the API, parse the - response, return the constructed :class:`Message` wrapped in a list of - length 1. + Send prompt, handle agentic tool calls (function_call), return all messages. - The wrapper detects function_call pieces via :attr:`_tool_parser` and - decides whether to dispatch + re-enter. Reasoning, MCP, web-search, - computer-use, and other non-function-call sections pass through to - Memory unchanged because the parser ignores them. + The Responses API supports structured outputs and tool execution. This method handles both: + - Simple text/reasoning responses + - Agentic tool-calling loops that may require multiple back-and-forth exchanges Args: normalized_conversation (list[Message]): The full conversation @@ -687,54 +571,59 @@ async def _send_prompt_to_target_async(self, *, normalized_conversation: list[Me pipeline. The current message is the last element. Returns: - list[Message]: Exactly one Message wrapping the parsed response. + List of messages generated during the interaction (assistant responses and tool messages). + The normalizer will persist all of these to memory. """ message = normalized_conversation[-1] + message_piece: MessagePiece = message.message_pieces[0] last_piece = message.message_pieces[-1] json_config = self._get_json_response_config(message_piece=last_piece) - body = await self._construct_request_body(conversation=list(normalized_conversation), json_config=json_config) - logger.info("Sending conversation with %d messages to the Responses API", len(normalized_conversation)) - result = await self._handle_openai_request( - api_call=lambda body=body: self._client.responses.create(**body), - request=message, - ) - return [result] + working_conversation: MutableSequence[Message] = list(normalized_conversation) - @property - def _tool_parser(self) -> ToolCallParser | None: - """ - Canonical-envelope parser shared with future canonical-envelope targets. - - Walks response message pieces and emits one :class:`~pyrit.tools.ToolCall` - per piece whose ``original_value_data_type`` is ``"function_call"``. - Reasoning, MCP, web-search, computer-use, and local-shell sections all - produce pieces of OTHER data types, so the parser returns an empty list - for them and the @tool_loop decorator exits cleanly. Those sections - still land in Memory via the parsed Message returned by - ``_send_prompt_to_target_async``; they're just not client-side - dispatched. - """ - return CanonicalEnvelopeParser() + # Track all responses generated during this interaction + responses_to_return: list[Message] = [] - def _tool_schemas(self) -> list[dict[str, Any]]: - """ - Translate the configured backend's schemas into Responses-API tools shape. + # Main agentic loop - each back-and-forth creates a new message + tool_call_section: Optional[dict[str, Any]] = None - The Responses API expects each function tool as a top-level - ``{"type": "function", "name": ..., "description": ..., - "parameters": ...}`` entry (NOT wrapped in an inner ``"function"`` key - the way Chat Completions does). The backend's schemas are already the - bare function schema, so we just stamp ``type=function`` on each. + while True: + logger.info(f"Sending conversation with {len(working_conversation)} messages to the prompt target") - Returns: - list[dict[str, Any]]: One descriptor per advertised tool, or an - empty list when no backend is configured. - """ - backend: ToolBackend | None = self.configuration.tool_backend - if backend is None: - return [] - return [{"type": "function", **schema} for schema in backend.schemas] + body = await self._construct_request_body(conversation=working_conversation, json_config=json_config) + + # Use unified error handling - automatically detects Response and validates + result = await self._handle_openai_request( + api_call=lambda body=body: self._client.responses.create(**body), + request=message, + ) + + # Add result to conversation and responses list + working_conversation.append(result) + responses_to_return.append(result) + + # Extract tool call if present + tool_call_section = self._find_last_pending_tool_call(result) + + # If no tool call, we're done + if not tool_call_section: + break + + # Execute the tool/function + tool_output = await self._execute_call_section(tool_call_section) + + # Create a new message with the tool output + tool_piece = self._make_tool_piece(tool_output, tool_call_section["call_id"], reference_piece=message_piece) + tool_message = Message(message_pieces=[tool_piece], skip_validation=True) + + # Add tool output message to conversation and responses list + working_conversation.append(tool_message) + responses_to_return.append(tool_message) + + # Continue loop to send tool result and get next response + + # Return all responses (normalizer will persist all of them to memory) + return responses_to_return def _parse_response_output_section( self, *, section: Any, message_piece: MessagePiece, error: Optional[PromptResponseError] diff --git a/tests/integration/tools/__init__.py b/tests/integration/tools/__init__.py deleted file mode 100644 index 9a0454564d..0000000000 --- a/tests/integration/tools/__init__.py +++ /dev/null @@ -1,2 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# Licensed under the MIT license. diff --git a/tests/integration/tools/test_red_teaming_with_tools.py b/tests/integration/tools/test_red_teaming_with_tools.py deleted file mode 100644 index 9dca01371a..0000000000 --- a/tests/integration/tools/test_red_teaming_with_tools.py +++ /dev/null @@ -1,361 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# Licensed under the MIT license. - -"""C7 integration tests: RedTeamingAttack with real tool dispatch. - -These tests spawn the real ``tests/unit/tools/echo_mcp_server.py`` subprocess -and exercise the full client-side tool-calling stack: - - attack -> normalizer -> target -> @tool_loop wrapper -> MCPToolBackend -> - MCPClient (stdio) -> echo subprocess -> tool result -> back through the - wrapper -> Memory. - -Only the OpenAI Responses HTTP layer is mocked. The MCP subprocess, the -MCPToolBackend lock, the AsyncExitStack lifecycle, the canonical envelope -round-trip, and the @tool_loop decorator's RedTeam-attack invocation path -all execute under their real implementations. -""" - -from __future__ import annotations - -import json -import pathlib -import sys -import uuid -from unittest.mock import AsyncMock, MagicMock, patch - -import pytest - -from pyrit.executor.attack.core.attack_config import AttackAdversarialConfig, AttackScoringConfig -from pyrit.executor.attack.multi_turn.red_teaming import RedTeamingAttack -from pyrit.identifiers import ComponentIdentifier -from pyrit.memory import CentralMemory -from pyrit.models import Message, MessagePiece, Score -from pyrit.prompt_target import OpenAIResponseTarget -from pyrit.prompt_target.common.prompt_target import PromptTarget -from pyrit.prompt_target.common.target_capabilities import TargetCapabilities -from pyrit.prompt_target.common.target_configuration import TargetConfiguration -from pyrit.score.true_false.true_false_scorer import TrueFalseScorer -from pyrit.tools import ( - LocalMCPServerSpec, - MCPToolBackend, - ToolEventBehavior, - ToolEventPolicy, -) - - -def _mock_id(name: str) -> ComponentIdentifier: - return ComponentIdentifier(class_name=name, class_module="test") - - -ECHO_SERVER_PATH = pathlib.Path(__file__).resolve().parents[2] / "unit" / "tools" / "echo_mcp_server.py" - - -def _local_echo_spec() -> LocalMCPServerSpec: - """Build a LocalMCPServerSpec that launches the in-tree echo server.""" - return LocalMCPServerSpec( - command=sys.executable, - args=(str(ECHO_SERVER_PATH),), - ) - - -def _mock_function_call_response(call_id: str, function_name: str, arguments: dict) -> MagicMock: - """Build a fake Responses-API response containing a function_call section.""" - response = MagicMock() - response.status = "completed" - response.error = None - section = MagicMock() - section.type = "function_call" - section.call_id = call_id - section.name = function_name - section.arguments = json.dumps(arguments) - section.model_dump.return_value = { - "type": "function_call", - "call_id": call_id, - "name": function_name, - "arguments": json.dumps(arguments), - } - response.output = [section] - return response - - -def _mock_text_response(text: str) -> MagicMock: - """Build a fake Responses-API response containing a message section.""" - response = MagicMock() - response.status = "completed" - response.error = None - section = MagicMock() - section.type = "message" - section.content = [MagicMock(text=text)] - response.output = [section] - return response - - -def _make_response_target_with_mcp_backend( - backend: MCPToolBackend, -) -> OpenAIResponseTarget: - """Build an OpenAIResponseTarget wired to the live MCP backend.""" - caps = TargetCapabilities( - supports_multi_turn=True, - supports_editable_history=True, - supports_json_output=True, - supports_multi_message_pieces=True, - supports_system_prompt=True, - supports_tool_use=True, - input_modalities=frozenset( - { - frozenset(["text"]), - frozenset(["text", "image_path"]), - frozenset(["function_call"]), - frozenset(["tool_call"]), - frozenset(["function_call_output"]), - frozenset(["reasoning"]), - } - ), - ) - config = TargetConfiguration( - capabilities=caps, - tool_event_policy=ToolEventPolicy(behavior=ToolEventBehavior.EXECUTE, max_tool_iterations=5), - tool_backend=backend, - ) - return OpenAIResponseTarget( - model_name="gpt-4", - endpoint="https://mock.example.com", - api_key="mock-key", - custom_configuration=config, - ) - - -def _scripted_adversarial(prompts: list[str]) -> MagicMock: - """Build a mock adversarial target that returns scripted prompts.""" - adversarial = MagicMock(spec=PromptTarget) - adversarial.send_prompt_async = AsyncMock( - side_effect=[ - [ - Message( - message_pieces=[ - MessagePiece( - role="assistant", - original_value=p, - original_value_data_type="text", - conversation_id=str(uuid.uuid4()), - ) - ] - ) - ] - for p in prompts - ] - ) - adversarial.get_identifier.return_value = _mock_id("MockAdversarial") - adversarial.set_system_prompt = MagicMock() - return adversarial - - -def _success_scorer() -> MagicMock: - """Mock objective scorer that always returns True (objective met).""" - scorer = MagicMock(spec=TrueFalseScorer) - scorer.score_async = AsyncMock( - return_value=[ - Score( - score_value="true", - score_value_description="objective met", - score_type="true_false", - score_category=["test"], - score_rationale="mock rationale", - score_metadata={}, - message_piece_id=str(uuid.uuid4()), - scorer_class_identifier=_mock_id("MockScorer"), - ) - ] - ) - scorer.get_identifier.return_value = _mock_id("MockScorer") - return scorer - - -@pytest.mark.asyncio -async def test_red_teaming_response_target_with_mcp_echo(patch_central_database): - """End-to-end: RedTeamingAttack drives OpenAIResponseTarget with MCPToolBackend. - - The Response target's HTTP layer is mocked to return a function_call for - the echo tool, followed by a stop response after the tool result arrives. - The MCP subprocess actually executes the echo call. - """ - backend = MCPToolBackend(servers=[_local_echo_spec()]) - async with backend: - objective_target = _make_response_target_with_mcp_backend(backend) - - # Mock the OpenAI Responses HTTP layer on the objective target. - responses = [ - _mock_function_call_response("call_1", "echo", {"text": "hello"}), - _mock_text_response("Echoed: hello"), - ] - seen = [] - - async def mock_create(**kwargs): - seen.append(kwargs) - return responses[len(seen) - 1] - - # Adversarial returns one prompt (RedTeamingAttack stops after objective is met) - adversarial = _scripted_adversarial(["please echo hello"]) - - attack = RedTeamingAttack( - objective_target=objective_target, - attack_adversarial_config=AttackAdversarialConfig(target=adversarial), - attack_scoring_config=AttackScoringConfig(objective_scorer=_success_scorer()), - ) - - with patch.object( - objective_target._async_client.responses, "create", new_callable=AsyncMock - ) as mock_create_call: - mock_create_call.side_effect = mock_create - result = await attack.execute_async(objective="get the model to echo 'hello'") - - # Two HTTP calls to the Response API: initial + post-tool - assert len(seen) == 2 - # Second call must include the function_call_output (tool result) - second_input = seen[1]["input"] - function_outputs = [item for item in second_input if item.get("type") == "function_call_output"] - assert len(function_outputs) == 1 - # The output JSON contains the text "hello" because the real MCP echo - # subprocess returned it - assert "hello" in function_outputs[0]["output"] - assert result is not None - - -@pytest.mark.asyncio -async def test_red_teaming_persists_canonical_transcript_in_memory(patch_central_database): - """End-to-end: after a successful tool dispatch the DB shows the full chain. - - Verifies the canonical envelope contract (§13): the conversation written - to Memory must contain the user message, the assistant function_call, the - tool function_call_output (with matching call_id), and the assistant's - final text -- in that order. - """ - backend = MCPToolBackend(servers=[_local_echo_spec()]) - async with backend: - objective_target = _make_response_target_with_mcp_backend(backend) - - responses = [ - _mock_function_call_response("call_xyz", "echo", {"text": "world"}), - _mock_text_response("Echoed: world"), - ] - seen = [] - - async def mock_create(**kwargs): - seen.append(kwargs) - return responses[len(seen) - 1] - - adversarial = _scripted_adversarial(["echo world"]) - - attack = RedTeamingAttack( - objective_target=objective_target, - attack_adversarial_config=AttackAdversarialConfig(target=adversarial), - attack_scoring_config=AttackScoringConfig(objective_scorer=_success_scorer()), - ) - - with patch.object( - objective_target._async_client.responses, "create", new_callable=AsyncMock - ) as mock_create_call: - mock_create_call.side_effect = mock_create - result = await attack.execute_async(objective="echo world") - - # Read the conversation back from Memory - memory = CentralMemory.get_memory_instance() - assert result is not None - objective_conv_id = result.conversation_id - assert objective_conv_id, "Attack result must carry the objective-target conversation id" - - pieces = list(memory.get_message_pieces(conversation_id=objective_conv_id)) - # Filter out system prompts; we care about the user/assistant/tool chain - data_types_in_order = [p.original_value_data_type for p in pieces] - # The chain MUST contain function_call followed by function_call_output (canonical envelope) - assert "function_call" in data_types_in_order - assert "function_call_output" in data_types_in_order - - fc_index = data_types_in_order.index("function_call") - fco_index = data_types_in_order.index("function_call_output") - assert fc_index < fco_index, "function_call must precede function_call_output in DB" - - fc_envelope = json.loads(pieces[fc_index].original_value) - fco_envelope = json.loads(pieces[fco_index].original_value) - assert fc_envelope["call_id"] == fco_envelope["call_id"] == "call_xyz" - assert fc_envelope["name"] == "echo" - # The tool result envelope's `output` is JSON-encoded; the underlying echo result is "world" - assert "world" in fco_envelope["output"] - - -@pytest.mark.asyncio -async def test_red_teaming_dispatches_all_tool_calls_per_turn(patch_central_database): - """Multi-call-per-turn dispatch (intentional behavior change vs pre-C6 loop). - - When the model emits two function_call sections in one response, BOTH - must dispatch through the MCPToolBackend. The pre-C6 in-class loop in - OpenAIResponseTarget only dispatched the LAST call per turn; the C6 - migration onto @tool_loop changes this to "dispatch every call in - declaration order." Verify by issuing both an `echo` and an `add` call - and asserting both results land in the second API call's input. - """ - backend = MCPToolBackend(servers=[_local_echo_spec()]) - async with backend: - objective_target = _make_response_target_with_mcp_backend(backend) - - # First response contains TWO function_calls; second is the stop text. - multi_call_response = MagicMock() - multi_call_response.status = "completed" - multi_call_response.error = None - - echo_section = MagicMock() - echo_section.type = "function_call" - echo_section.call_id = "call_echo" - echo_section.name = "echo" - echo_section.arguments = json.dumps({"text": "hi"}) - echo_section.model_dump.return_value = { - "type": "function_call", - "call_id": "call_echo", - "name": "echo", - "arguments": json.dumps({"text": "hi"}), - } - add_section = MagicMock() - add_section.type = "function_call" - add_section.call_id = "call_add" - add_section.name = "add" - add_section.arguments = json.dumps({"a": 3, "b": 4}) - add_section.model_dump.return_value = { - "type": "function_call", - "call_id": "call_add", - "name": "add", - "arguments": json.dumps({"a": 3, "b": 4}), - } - multi_call_response.output = [echo_section, add_section] - - responses = [ - multi_call_response, - _mock_text_response("done"), - ] - seen = [] - - async def mock_create(**kwargs): - seen.append(kwargs) - return responses[len(seen) - 1] - - adversarial = _scripted_adversarial(["call echo and add"]) - - attack = RedTeamingAttack( - objective_target=objective_target, - attack_adversarial_config=AttackAdversarialConfig(target=adversarial), - attack_scoring_config=AttackScoringConfig(objective_scorer=_success_scorer()), - ) - - with patch.object(objective_target._async_client.responses, "create", new_callable=AsyncMock) as mc: - mc.side_effect = mock_create - await attack.execute_async(objective="dispatch both tools") - - assert len(seen) == 2 - second_input = seen[1]["input"] - outputs = [item for item in second_input if item.get("type") == "function_call_output"] - assert len(outputs) == 2, "Both tool calls must be dispatched per the new behavior" - call_ids = [o["call_id"] for o in outputs] - assert call_ids == ["call_echo", "call_add"], "Outputs must preserve declaration order" - # Real MCP subprocess: echo("hi") returned "hi", add(3, 4) returned 7 - assert "hi" in outputs[0]["output"] - assert "7" in outputs[1]["output"] diff --git a/tests/unit/prompt_target/target/test_openai_response_target_c6_migration.py b/tests/unit/prompt_target/target/test_openai_response_target_c6_migration.py deleted file mode 100644 index 04dfd0b024..0000000000 --- a/tests/unit/prompt_target/target/test_openai_response_target_c6_migration.py +++ /dev/null @@ -1,304 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# Licensed under the MIT license. - -"""C6 additions to the Response target function-chaining suite. - -Covers the migration onto @tool_loop + LocalToolBackend. -""" - -from __future__ import annotations - -import json -import uuid -import warnings -from typing import Any -from unittest.mock import AsyncMock, MagicMock, patch - -import pytest - -from pyrit.models import Message, MessagePiece -from pyrit.prompt_target import OpenAIResponseTarget -from pyrit.prompt_target.common.target_capabilities import TargetCapabilities -from pyrit.prompt_target.common.target_configuration import TargetConfiguration -from pyrit.tools import LocalToolBackend, ToolEventBehavior, ToolEventPolicy - - -def _mock_function_call_response(call_id: str, function_name: str, arguments: dict) -> MagicMock: - """Build a fake Responses-API response containing a function_call section.""" - mock_response = MagicMock() - mock_response.status = "completed" - mock_response.error = None - section = MagicMock() - section.type = "function_call" - section.call_id = call_id - section.name = function_name - section.arguments = json.dumps(arguments) - section.model_dump.return_value = { - "type": "function_call", - "call_id": call_id, - "name": function_name, - "arguments": json.dumps(arguments), - } - mock_response.output = [section] - return mock_response - - -def _mock_text_response(text: str) -> MagicMock: - """Build a fake Responses-API response containing a message section.""" - mock_response = MagicMock() - mock_response.status = "completed" - mock_response.error = None - section = MagicMock() - section.type = "message" - section.content = [MagicMock(text=text)] - mock_response.output = [section] - return mock_response - - -def _user_msg(text: str, conversation_id: str | None = None) -> Message: - return Message( - message_pieces=[ - MessagePiece( - role="user", - original_value=text, - conversation_id=conversation_id or str(uuid.uuid4()), - ) - ] - ) - - -class TestCustomFunctionsDeprecation: - """custom_functions still works but emits DeprecationWarning.""" - - def test_custom_functions_kwarg_emits_deprecation_warning(self, patch_central_database): - async def get_weather(args: dict[str, Any]) -> dict[str, Any]: - return {"t": 72} - - with pytest.warns(DeprecationWarning, match="custom_functions"): - OpenAIResponseTarget( - model_name="gpt-4", - endpoint="https://mock.example.com", - api_key="mock-key", - custom_functions={"get_weather": get_weather}, - ) - - @pytest.mark.asyncio - async def test_custom_functions_kwarg_still_dispatches(self, patch_central_database): - async def get_weather(args: dict[str, Any]) -> dict[str, Any]: - return {"temperature": 72, "condition": "sunny"} - - with warnings.catch_warnings(): - warnings.simplefilter("ignore", DeprecationWarning) - target = OpenAIResponseTarget( - model_name="gpt-4", - endpoint="https://mock.example.com", - api_key="mock-key", - custom_functions={"get_weather": get_weather}, - ) - - responses = [ - _mock_function_call_response("call_1", "get_weather", {"location": "NYC"}), - _mock_text_response("72F and sunny."), - ] - seen = [] - - async def mock_create(**kwargs): - seen.append(kwargs) - return responses[len(seen) - 1] - - with patch.object(target._async_client.responses, "create", new_callable=AsyncMock) as mc: - mc.side_effect = mock_create - result = await target.send_prompt_async(message=_user_msg("weather?")) - - assert len(seen) == 2 - assert result[-1].message_pieces[0].original_value == "72F and sunny." - second_input = seen[1]["input"] - assert any(item.get("type") == "function_call_output" for item in second_input) - - -def _config_with_backend(backend: LocalToolBackend) -> TargetConfiguration: - """Build a TargetConfiguration wired for the modern tool-backend path.""" - caps = TargetCapabilities( - supports_multi_turn=True, - supports_multi_message_pieces=True, - supports_editable_history=True, - supports_json_output=True, - supports_system_prompt=True, - supports_tool_use=True, - input_modalities=frozenset( - { - frozenset(["text"]), - frozenset(["text", "image_path"]), - frozenset(["function_call"]), - frozenset(["tool_call"]), - frozenset(["function_call_output"]), - frozenset(["reasoning"]), - } - ), - ) - return TargetConfiguration( - capabilities=caps, - tool_event_policy=ToolEventPolicy(behavior=ToolEventBehavior.EXECUTE, max_tool_iterations=5), - tool_backend=backend, - ) - - -class TestToolBackendDispatch: - """The modern path: pass tool_backend via TargetConfiguration.""" - - @pytest.mark.asyncio - async def test_local_backend_dispatches_through_tool_loop(self, patch_central_database): - async def get_weather(args: dict[str, Any]) -> dict[str, Any]: - return {"temperature": 72, "condition": "sunny"} - - backend = LocalToolBackend( - callables={"get_weather": get_weather}, - schemas=[ - { - "name": "get_weather", - "description": "Weather lookup.", - "parameters": { - "type": "object", - "properties": {"location": {"type": "string"}}, - "required": ["location"], - }, - } - ], - ) - target = OpenAIResponseTarget( - model_name="gpt-4", - endpoint="https://mock.example.com", - api_key="mock-key", - custom_configuration=_config_with_backend(backend), - ) - - responses = [ - _mock_function_call_response("call_1", "get_weather", {"location": "NYC"}), - _mock_text_response("72F and sunny in NYC."), - ] - seen = [] - - async def mock_create(**kwargs): - seen.append(kwargs) - return responses[len(seen) - 1] - - with patch.object(target._async_client.responses, "create", new_callable=AsyncMock) as mc: - mc.side_effect = mock_create - result = await target.send_prompt_async(message=_user_msg("weather?")) - - assert len(seen) == 2 - assert result[-1].message_pieces[0].original_value == "72F and sunny in NYC." - second_input = seen[1]["input"] - assert any(item.get("type") == "function_call_output" for item in second_input) - - -class TestToolSchemasInjection: - """_construct_request_body injects backend schemas when present.""" - - @pytest.mark.asyncio - async def test_backend_schemas_injected_into_tools(self, patch_central_database): - async def get_weather(args: dict[str, Any]) -> dict[str, Any]: - return {"t": 1} - - backend = LocalToolBackend( - callables={"get_weather": get_weather}, - schemas=[{"name": "get_weather", "description": "x", "parameters": {"type": "object"}}], - ) - target = OpenAIResponseTarget( - model_name="gpt-4", - endpoint="https://mock.example.com", - api_key="mock-key", - custom_configuration=_config_with_backend(backend), - ) - body = await target._construct_request_body( - conversation=[_user_msg("hi")], - json_config=MagicMock(enabled=False, schema=None), - ) - assert "tools" in body - assert body["tools"][0]["type"] == "function" - assert body["tools"][0]["name"] == "get_weather" - - @pytest.mark.asyncio - async def test_extra_body_tools_take_precedence(self, patch_central_database): - async def f(args: dict[str, Any]) -> dict[str, Any]: - return {} - - backend = LocalToolBackend( - callables={"f": f}, - schemas=[{"name": "f", "parameters": {"type": "object"}}], - ) - legacy = [{"type": "function", "name": "legacy_tool", "description": "x"}] - config = _config_with_backend(backend) - target = OpenAIResponseTarget( - model_name="gpt-4", - endpoint="https://mock.example.com", - api_key="mock-key", - extra_body_parameters={"tools": legacy}, - custom_configuration=config, - ) - body = await target._construct_request_body( - conversation=[_user_msg("hi")], - json_config=MagicMock(enabled=False, schema=None), - ) - assert body["tools"] == legacy - - @pytest.mark.asyncio - async def test_no_backend_no_tools_key(self, patch_central_database): - target = OpenAIResponseTarget( - model_name="gpt-4", - endpoint="https://mock.example.com", - api_key="mock-key", - ) - body = await target._construct_request_body( - conversation=[_user_msg("hi")], - json_config=MagicMock(enabled=False, schema=None), - ) - assert "tools" not in body - - -class TestNonFunctionCallPiecesPassThrough: - """Reasoning / mcp_call / web_search_call sections must NOT be dispatched. - - The Response target's parser populates pieces for these types so they can - be persisted to Memory and round-tripped on subsequent requests. The - CanonicalEnvelopeParser only extracts function_call pieces; the tool loop - must therefore see an empty parse and exit cleanly. - """ - - @pytest.mark.asyncio - async def test_reasoning_only_response_exits_loop(self, patch_central_database): - target = OpenAIResponseTarget( - model_name="gpt-4", - endpoint="https://mock.example.com", - api_key="mock-key", - reasoning_effort="medium", - ) - # Reasoning section + final text section in one response - mock_response = MagicMock() - mock_response.status = "completed" - mock_response.error = None - reasoning_section = MagicMock() - reasoning_section.type = "reasoning" - reasoning_section.model_dump.return_value = {"type": "reasoning", "summary": "thinking..."} - text_section = MagicMock() - text_section.type = "message" - text_section.content = [MagicMock(text="The answer is 42.")] - mock_response.output = [reasoning_section, text_section] - - seen = [] - - async def mock_create(**kwargs): - seen.append(kwargs) - return mock_response - - with patch.object(target._async_client.responses, "create", new_callable=AsyncMock) as mc: - mc.side_effect = mock_create - result = await target.send_prompt_async(message=_user_msg("question?")) - - # Exactly one API call -- reasoning is not a tool call so the loop exits - assert len(seen) == 1 - # Response message contains both pieces - assert len(result) == 1 - piece_types = [p.original_value_data_type for p in result[0].message_pieces] - assert "reasoning" in piece_types - assert "text" in piece_types diff --git a/uv.lock b/uv.lock index 2adcb96aab..18a65c10f0 100644 --- a/uv.lock +++ b/uv.lock @@ -2122,15 +2122,6 @@ http2 = [ { name = "h2" }, ] -[[package]] -name = "httpx-sse" -version = "0.4.3" -source = { registry = "https://pypi.org/simple" } -sdist = { url = "https://files.pythonhosted.org/packages/0f/4c/751061ffa58615a32c31b2d82e8482be8dd4a89154f003147acee90f2be9/httpx_sse-0.4.3.tar.gz", hash = "sha256:9b1ed0127459a66014aec3c56bebd93da3c1bc8bb6618c8082039a44889a755d", size = 15943, upload-time = "2025-10-10T21:48:22.271Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/d2/fd/6668e5aec43ab844de6fc74927e155a3b37bf40d7c3790e49fc0406b6578/httpx_sse-0.4.3-py3-none-any.whl", hash = "sha256:0ac1c9fe3c0afad2e0ebb25a934a59f4c7823b60792691f779fad2c5568830fc", size = 8960, upload-time = "2025-10-10T21:48:21.158Z" }, -] - [[package]] name = "huggingface-hub" version = "1.13.0" @@ -3202,31 +3193,6 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/af/33/ee4519fa02ed11a94aef9559552f3b17bb863f2ecfe1a35dc7f548cde231/matplotlib_inline-0.2.1-py3-none-any.whl", hash = "sha256:d56ce5156ba6085e00a9d54fead6ed29a9c47e215cd1bba2e976ef39f5710a76", size = 9516, upload-time = "2025-10-23T09:00:20.675Z" }, ] -[[package]] -name = "mcp" -version = "1.27.1" -source = { registry = "https://pypi.org/simple" } -dependencies = [ - { name = "anyio" }, - { name = "httpx" }, - { name = "httpx-sse" }, - { name = "jsonschema" }, - { name = "pydantic" }, - { name = "pydantic-settings" }, - { name = "pyjwt", extra = ["crypto"] }, - { name = "python-multipart" }, - { name = "pywin32", marker = "sys_platform == 'win32'" }, - { name = "sse-starlette" }, - { name = "starlette" }, - { name = "typing-extensions" }, - { name = "typing-inspection" }, - { name = "uvicorn", marker = "sys_platform != 'emscripten'" }, -] -sdist = { url = "https://files.pythonhosted.org/packages/38/83/d1efe7c2980d8a3afa476f4e3d42d53dd54c0ab94c27bee5d755b45c8b73/mcp-1.27.1.tar.gz", hash = "sha256:0f47e1820f8f8f941466b39749eb1d1839a04caddca2bc60e9d46e8a99914924", size = 608458, upload-time = "2026-05-08T16:50:12.601Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/fd/73/42d9596facebdb533b7f0b86c1b0364ef350d1f8ba78b1052e8a58b48b65/mcp-1.27.1-py3-none-any.whl", hash = "sha256:1af3c4203b329430fde7a87b4fcb6392a041f5cb851fd68fc674016ab4e7c06f", size = 216260, upload-time = "2026-05-08T16:50:10.547Z" }, -] - [[package]] name = "mdit-py-plugins" version = "0.5.0" @@ -5049,20 +5015,6 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/36/c7/cfc8e811f061c841d7990b0201912c3556bfeb99cdcb7ed24adc8d6f8704/pydantic_core-2.41.5-pp311-pypy311_pp73-win_amd64.whl", hash = "sha256:56121965f7a4dc965bff783d70b907ddf3d57f6eba29b6d2e5dabfaf07799c51", size = 2145302, upload-time = "2025-11-04T13:43:46.64Z" }, ] -[[package]] -name = "pydantic-settings" -version = "2.14.1" -source = { registry = "https://pypi.org/simple" } -dependencies = [ - { name = "pydantic" }, - { name = "python-dotenv" }, - { name = "typing-inspection" }, -] -sdist = { url = "https://files.pythonhosted.org/packages/07/60/1d1e59c9c90d54591469ada7d268251f71c24bdb765f1a8a832cee8c6653/pydantic_settings-2.14.1.tar.gz", hash = "sha256:e874d3bec7e787b0c9958277956ed9b4dd5de6a80e162188fdaff7c5e26fd5fa", size = 235551, upload-time = "2026-05-08T13:40:06.542Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/ae/8d/f1af3832f5e6eb13ba94ee809e72b8ecb5eef226d27ee0bef7d963d943c7/pydantic_settings-2.14.1-py3-none-any.whl", hash = "sha256:6e3c7edfd8277687cdc598f56e5cff0e9bfff0910a3749deaa8d4401c3a2b9de", size = 60964, upload-time = "2026-05-08T13:40:04.958Z" }, -] - [[package]] name = "pydash" version = "8.0.5" @@ -5219,7 +5171,6 @@ dependencies = [ { name = "fastapi" }, { name = "httpx", extra = ["http2"] }, { name = "jinja2" }, - { name = "mcp" }, { name = "numpy", version = "2.2.6", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.14'" }, { name = "numpy", version = "2.4.4", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.14'" }, { name = "openai" }, @@ -5357,7 +5308,6 @@ requires-dist = [ { name = "ipykernel", marker = "extra == 'all'", specifier = ">=6.29.5" }, { name = "jinja2", specifier = ">=3.1.6" }, { name = "jupyter", marker = "extra == 'all'", specifier = ">=1.1.1" }, - { name = "mcp", specifier = ">=1.0,<2" }, { name = "ml-collections", marker = "extra == 'all'", specifier = ">=1.1.0" }, { name = "ml-collections", marker = "extra == 'gcg'", specifier = ">=1.1.0" }, { name = "numpy", marker = "python_full_version < '3.14'", specifier = ">=1.26.0" }, @@ -5551,15 +5501,6 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/51/e5/fecf13f06e5e5f67e8837d777d1bc43fac0ed2b77a676804df5c34744727/python_json_logger-4.0.0-py3-none-any.whl", hash = "sha256:af09c9daf6a813aa4cc7180395f50f2a9e5fa056034c9953aec92e381c5ba1e2", size = 15548, upload-time = "2025-10-06T04:15:17.553Z" }, ] -[[package]] -name = "python-multipart" -version = "0.0.29" -source = { registry = "https://pypi.org/simple" } -sdist = { url = "https://files.pythonhosted.org/packages/4e/fe/70bd71a6738b09a0bdf6480ca6436b167469ca4578b2a0efbe390b4b0e70/python_multipart-0.0.29.tar.gz", hash = "sha256:643e93849196645e2dbdd81a0f8829a23123ad7f797a84a364c6fb3563f18904", size = 45678, upload-time = "2026-05-17T17:29:47.654Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/8f/cb/769cfc37177252872a45a71f3fbdde9d51b471a3f3c14bfe95dde3407386/python_multipart-0.0.29-py3-none-any.whl", hash = "sha256:2ddcc971cef266225f54f552d8fa10bcfbb1f14446caec199060daac59ff2d69", size = 29640, upload-time = "2026-05-17T17:29:45.69Z" }, -] - [[package]] name = "pytz" version = "2025.2" @@ -5569,28 +5510,6 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/81/c4/34e93fe5f5429d7570ec1fa436f1986fb1f00c3e0f43a589fe2bbcd22c3f/pytz-2025.2-py2.py3-none-any.whl", hash = "sha256:5ddf76296dd8c44c26eb8f4b6f35488f3ccbf6fbbd7adee0b7262d43f0ec2f00", size = 509225, upload-time = "2025-03-25T02:24:58.468Z" }, ] -[[package]] -name = "pywin32" -version = "311" -source = { registry = "https://pypi.org/simple" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/7b/40/44efbb0dfbd33aca6a6483191dae0716070ed99e2ecb0c53683f400a0b4f/pywin32-311-cp310-cp310-win32.whl", hash = "sha256:d03ff496d2a0cd4a5893504789d4a15399133fe82517455e78bad62efbb7f0a3", size = 8760432, upload-time = "2025-07-14T20:13:05.9Z" }, - { url = "https://files.pythonhosted.org/packages/5e/bf/360243b1e953bd254a82f12653974be395ba880e7ec23e3731d9f73921cc/pywin32-311-cp310-cp310-win_amd64.whl", hash = "sha256:797c2772017851984b97180b0bebe4b620bb86328e8a884bb626156295a63b3b", size = 9590103, upload-time = "2025-07-14T20:13:07.698Z" }, - { url = "https://files.pythonhosted.org/packages/57/38/d290720e6f138086fb3d5ffe0b6caa019a791dd57866940c82e4eeaf2012/pywin32-311-cp310-cp310-win_arm64.whl", hash = "sha256:0502d1facf1fed4839a9a51ccbcc63d952cf318f78ffc00a7e78528ac27d7a2b", size = 8778557, upload-time = "2025-07-14T20:13:11.11Z" }, - { url = "https://files.pythonhosted.org/packages/7c/af/449a6a91e5d6db51420875c54f6aff7c97a86a3b13a0b4f1a5c13b988de3/pywin32-311-cp311-cp311-win32.whl", hash = "sha256:184eb5e436dea364dcd3d2316d577d625c0351bf237c4e9a5fabbcfa5a58b151", size = 8697031, upload-time = "2025-07-14T20:13:13.266Z" }, - { url = "https://files.pythonhosted.org/packages/51/8f/9bb81dd5bb77d22243d33c8397f09377056d5c687aa6d4042bea7fbf8364/pywin32-311-cp311-cp311-win_amd64.whl", hash = "sha256:3ce80b34b22b17ccbd937a6e78e7225d80c52f5ab9940fe0506a1a16f3dab503", size = 9508308, upload-time = "2025-07-14T20:13:15.147Z" }, - { url = "https://files.pythonhosted.org/packages/44/7b/9c2ab54f74a138c491aba1b1cd0795ba61f144c711daea84a88b63dc0f6c/pywin32-311-cp311-cp311-win_arm64.whl", hash = "sha256:a733f1388e1a842abb67ffa8e7aad0e70ac519e09b0f6a784e65a136ec7cefd2", size = 8703930, upload-time = "2025-07-14T20:13:16.945Z" }, - { url = "https://files.pythonhosted.org/packages/e7/ab/01ea1943d4eba0f850c3c61e78e8dd59757ff815ff3ccd0a84de5f541f42/pywin32-311-cp312-cp312-win32.whl", hash = "sha256:750ec6e621af2b948540032557b10a2d43b0cee2ae9758c54154d711cc852d31", size = 8706543, upload-time = "2025-07-14T20:13:20.765Z" }, - { url = "https://files.pythonhosted.org/packages/d1/a8/a0e8d07d4d051ec7502cd58b291ec98dcc0c3fff027caad0470b72cfcc2f/pywin32-311-cp312-cp312-win_amd64.whl", hash = "sha256:b8c095edad5c211ff31c05223658e71bf7116daa0ecf3ad85f3201ea3190d067", size = 9495040, upload-time = "2025-07-14T20:13:22.543Z" }, - { url = "https://files.pythonhosted.org/packages/ba/3a/2ae996277b4b50f17d61f0603efd8253cb2d79cc7ae159468007b586396d/pywin32-311-cp312-cp312-win_arm64.whl", hash = "sha256:e286f46a9a39c4a18b319c28f59b61de793654af2f395c102b4f819e584b5852", size = 8710102, upload-time = "2025-07-14T20:13:24.682Z" }, - { url = "https://files.pythonhosted.org/packages/a5/be/3fd5de0979fcb3994bfee0d65ed8ca9506a8a1260651b86174f6a86f52b3/pywin32-311-cp313-cp313-win32.whl", hash = "sha256:f95ba5a847cba10dd8c4d8fefa9f2a6cf283b8b88ed6178fa8a6c1ab16054d0d", size = 8705700, upload-time = "2025-07-14T20:13:26.471Z" }, - { url = "https://files.pythonhosted.org/packages/e3/28/e0a1909523c6890208295a29e05c2adb2126364e289826c0a8bc7297bd5c/pywin32-311-cp313-cp313-win_amd64.whl", hash = "sha256:718a38f7e5b058e76aee1c56ddd06908116d35147e133427e59a3983f703a20d", size = 9494700, upload-time = "2025-07-14T20:13:28.243Z" }, - { url = "https://files.pythonhosted.org/packages/04/bf/90339ac0f55726dce7d794e6d79a18a91265bdf3aa70b6b9ca52f35e022a/pywin32-311-cp313-cp313-win_arm64.whl", hash = "sha256:7b4075d959648406202d92a2310cb990fea19b535c7f4a78d3f5e10b926eeb8a", size = 8709318, upload-time = "2025-07-14T20:13:30.348Z" }, - { url = "https://files.pythonhosted.org/packages/c9/31/097f2e132c4f16d99a22bfb777e0fd88bd8e1c634304e102f313af69ace5/pywin32-311-cp314-cp314-win32.whl", hash = "sha256:b7a2c10b93f8986666d0c803ee19b5990885872a7de910fc460f9b0c2fbf92ee", size = 8840714, upload-time = "2025-07-14T20:13:32.449Z" }, - { url = "https://files.pythonhosted.org/packages/90/4b/07c77d8ba0e01349358082713400435347df8426208171ce297da32c313d/pywin32-311-cp314-cp314-win_amd64.whl", hash = "sha256:3aca44c046bd2ed8c90de9cb8427f581c479e594e99b5c0bb19b29c10fd6cb87", size = 9656800, upload-time = "2025-07-14T20:13:34.312Z" }, - { url = "https://files.pythonhosted.org/packages/c0/d2/21af5c535501a7233e734b8af901574572da66fcc254cb35d0609c9080dd/pywin32-311-cp314-cp314-win_arm64.whl", hash = "sha256:a508e2d9025764a8270f93111a970e1d0fbfc33f4153b388bb649b7eec4f9b42", size = 8932540, upload-time = "2025-07-14T20:13:36.379Z" }, -] - [[package]] name = "pywinpty" version = "3.0.2" @@ -6645,19 +6564,6 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/2a/ae/57d1d7af907e20c077e113e0e4976f87b82c0a415403d99284a262229dd0/srsly-2.5.3-cp314-cp314t-win_arm64.whl", hash = "sha256:d822083fe26ec6728bd8c273ac121fc4ab3864a0fdf0cf0ff3efb188fcd209ed", size = 650229, upload-time = "2026-03-23T11:56:46.148Z" }, ] -[[package]] -name = "sse-starlette" -version = "3.4.4" -source = { registry = "https://pypi.org/simple" } -dependencies = [ - { name = "anyio" }, - { name = "starlette" }, -] -sdist = { url = "https://files.pythonhosted.org/packages/f7/2b/58abc2d1fd397e7dde08e947e05c884d8ef2f78d5e2588c17a12d42d6994/sse_starlette-3.4.4.tar.gz", hash = "sha256:07e0fa0460138baf25cdd5fb28683472c3995dc1642225191b3832d62526bcb0", size = 31819, upload-time = "2026-05-12T17:37:17.019Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/dc/67/805710444ea8cc75fbf70b920ed431a560c4bf9c57f7d5a3117213189399/sse_starlette-3.4.4-py3-none-any.whl", hash = "sha256:3f4dd50d8aed2771a091f3a83000323fc3844541c16b4fe585ae2420cc6df973", size = 16514, upload-time = "2026-05-12T17:37:15.601Z" }, -] - [[package]] name = "stack-data" version = "0.6.3" From 1f5d0615070b62278e07fc6665f7f5fadfed9206 Mon Sep 17 00:00:00 2001 From: Victor Valbuena Date: Fri, 29 May 2026 16:42:09 -0700 Subject: [PATCH 37/40] Fix merge conflicts with main (PR #1785/#1784): adopt factory registry API - Replace SCENARIO_TECHNIQUES/AttackTechniqueSpec with AttackTechniqueFactory registry - Add @cache to _build_benchmark_strategy; drop deleted classmethods - Adopt default_strategy=/default_dataset_config= in super().__init__ (#1784 contract) - Loop: factory.create(attack_adversarial_config_override=AttackAdversarialConfig(target=...)) - Rename skip_cached -> use_cached throughout - Test fixture: use build_scenario_technique_factories() + mock adversarial_chat target (matches PR #1785 pattern); module-level constants from production catalog Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- .../scenarios/benchmark/adversarial.py | 309 +++--------------- .../scenario/benchmark/test_adversarial.py | 253 +++++++------- 2 files changed, 182 insertions(+), 380 deletions(-) diff --git a/pyrit/scenario/scenarios/benchmark/adversarial.py b/pyrit/scenario/scenarios/benchmark/adversarial.py index 3e8fb61b71..19bf684838 100644 --- a/pyrit/scenario/scenarios/benchmark/adversarial.py +++ b/pyrit/scenario/scenarios/benchmark/adversarial.py @@ -5,32 +5,20 @@ from __future__ import annotations -import dataclasses import logging from functools import cache from typing import TYPE_CHECKING, ClassVar -<<<<<<< HEAD from pyrit.analytics import get_cached_results_for_technique from pyrit.common import Parameter, apply_defaults -from pyrit.executor.attack import AttackScoringConfig +from pyrit.executor.attack import AttackAdversarialConfig, AttackScoringConfig from pyrit.identifiers import ObjectiveTargetEvaluationIdentifier from pyrit.models import AttackOutcome, SeedAttackGroup from pyrit.registry import AttackTechniqueRegistry, TargetRegistry -======= -from pyrit.common import apply_defaults -from pyrit.executor.attack import AttackAdversarialConfig, AttackScoringConfig -from pyrit.prompt_target import CHAT_TARGET_REQUIREMENTS -from pyrit.registry import AttackTechniqueRegistry ->>>>>>> 4cf9ff5de64017ce09cbe9eeb653a35a9983cab4 from pyrit.registry.tag_query import TagQuery from pyrit.scenario.core.atomic_attack import AtomicAttack from pyrit.scenario.core.dataset_configuration import DatasetConfiguration from pyrit.scenario.core.scenario import BaselineAttackPolicy, Scenario -<<<<<<< HEAD -from pyrit.scenario.core.scenario_techniques import SCENARIO_TECHNIQUES, _spec_needs_adversarial -======= ->>>>>>> 4cf9ff5de64017ce09cbe9eeb653a35a9983cab4 if TYPE_CHECKING: from pyrit.prompt_target import PromptTarget @@ -42,33 +30,36 @@ logger = logging.getLogger(__name__) +@cache def _build_benchmark_strategy() -> type[ScenarioStrategy]: """ - Build the ``BenchmarkStrategy`` enum from ``SCENARIO_TECHNIQUES``. - - Filters the static technique catalog to entries that require an - adversarial chat target (per ``_spec_needs_adversarial``) and passes - those source specs to - ``AttackTechniqueRegistry.build_strategy_class_from_specs``. The - resulting enum has one concrete member per source technique (e.g. - ``red_teaming``, ``tap``, ``crescendo_simulated``) plus the standard - ``all`` / ``light`` / ``single_turn`` / ``multi_turn`` aggregates inherited - from the source specs' ``strategy_tags``. - - The (technique × target) cross-product is no longer pre-materialized into - enum members; per-target factories are built lazily in + Build the ``BenchmarkStrategy`` enum from the registered factory catalog. + + Reads ``core`` adversarial-capable factories from the + ``AttackTechniqueRegistry`` singleton and passes them to + ``build_strategy_class_from_factories``. The resulting enum has one + concrete member per factory (e.g. ``red_teaming``, ``tap``, + ``crescendo_simulated``) plus ``default`` / ``light`` / ``single_turn`` + / ``multi_turn`` aggregates derived from each factory's ``strategy_tags``. + + The (technique × target) cross-product is materialized lazily in ``AdversarialBenchmark._get_atomic_attacks_async`` from the user-supplied ``adversarial_targets`` parameter. Returns: type[ScenarioStrategy]: The dynamically generated ``BenchmarkStrategy`` class. """ - adversarial_specs = [spec for spec in SCENARIO_TECHNIQUES if _spec_needs_adversarial(spec)] - - return AttackTechniqueRegistry.build_strategy_class_from_specs( # type: ignore[ty:invalid-return-type] + registry = AttackTechniqueRegistry.get_registry_singleton() + factories = [ + factory + for factory in registry.get_factories_or_raise().values() + if factory.uses_adversarial and "core" in factory.strategy_tags + ] + return AttackTechniqueRegistry.build_strategy_class_from_factories( # type: ignore[ty:invalid-return-type] class_name="BenchmarkStrategy", - specs=adversarial_specs, + factories=factories, aggregate_tags={ + "default": TagQuery.any_of("default"), "light": TagQuery.any_of("light"), "single_turn": TagQuery.any_of("single_turn"), "multi_turn": TagQuery.any_of("multi_turn"), @@ -88,75 +79,27 @@ class AdversarialBenchmark(Scenario): At run time, ``_get_atomic_attacks_async`` performs the ``(technique × adversarial_target × dataset)`` cross-product: for each - selected adversarial-capable technique in ``SCENARIO_TECHNIQUES`` and - each requested target, it constructs a per-pair - ``AttackTechniqueFactory`` via - ``AttackTechniqueRegistry.build_factory_from_spec`` with - ``adversarial_chat`` overridden to that target — no global registry - mutation. The resulting ``AtomicAttack`` is named - ``f"{technique}__{target}_{dataset}"`` with ``display_group`` set to the - target's registry name so per-model ASR rolls up naturally in result - displays. + selected adversarial-capable ``core`` factory in the + ``AttackTechniqueRegistry`` and each requested target, it calls + ``factory.create(attack_adversarial_config_override=...)`` with the + resolved target — no global registry mutation. The resulting + ``AtomicAttack`` is named ``f"{technique}__{target}_{dataset}"`` with + ``display_group`` set to the target's registry name so per-model ASR + rolls up naturally in result displays. """ -<<<<<<< HEAD #: Bumped from 1 → 2 by the refactor that moved adversarial targets #: from a constructor parameter to the ``adversarial_targets`` scenario #: parameter and changed ``atomic_attack_name`` from #: ``{technique}__{model}__{dataset}`` to ``{technique}__{target}_{dataset}``. - #: ``skip_cached`` only matches against prior runs at the current + #: ``use_cached`` only matches against prior runs at the current #: ``VERSION``; v1 results remain queryable but won't suppress v2 runs. VERSION: int = 2 -======= - VERSION: int = 1 ->>>>>>> 4cf9ff5de64017ce09cbe9eeb653a35a9983cab4 #: AdversarialBenchmark compares attack-success rates across adversarial models; a baseline #: attack would be model-independent and contribute no signal to the comparison. BASELINE_ATTACK_POLICY: ClassVar[BaselineAttackPolicy] = BaselineAttackPolicy.Forbidden -<<<<<<< HEAD - @classmethod - def get_strategy_class(cls) -> type[ScenarioStrategy]: - """ - Return the ``BenchmarkStrategy`` enum. - - The enum is deterministic given the current ``SCENARIO_TECHNIQUES`` - catalog (the scenario no longer fans out across registry entries), - so it is rebuilt on every call rather than cached. - - Returns: - type[ScenarioStrategy]: The ``BenchmarkStrategy`` enum class. - """ - return _build_benchmark_strategy() - - @classmethod - def get_default_strategy(cls) -> ScenarioStrategy: - """ - Return the default strategy (``light``). - - Returns: - ScenarioStrategy: The ``light`` aggregate member — runs the - subset of benchmark-friendly techniques that finish quickly with - modest system resources (excludes ``tap`` and - ``crescendo_simulated``, which can take hours on a single run). - """ - return cls.get_strategy_class()("light") - - @classmethod - def default_dataset_config(cls) -> DatasetConfiguration: - """ - Return the default dataset configuration for benchmarking. - - Returns: - DatasetConfiguration: ``harmbench`` capped at 8 prompts per - atomic attack. - """ - return DatasetConfiguration( - dataset_names=["harmbench"], - max_dataset_size=8, - ) - @classmethod def supported_parameters(cls) -> list[Parameter]: """ @@ -189,32 +132,25 @@ def supported_parameters(cls) -> list[Parameter]: ), ] -======= ->>>>>>> 4cf9ff5de64017ce09cbe9eeb653a35a9983cab4 @apply_defaults def __init__( self, *, -<<<<<<< HEAD -======= - adversarial_models: list[PromptTarget] | None = None, ->>>>>>> 4cf9ff5de64017ce09cbe9eeb653a35a9983cab4 objective_scorer: TrueFalseScorer | None = None, - skip_cached: bool = False, + use_cached: bool = False, scenario_result_id: str | None = None, ) -> None: """ Initialize the AdversarialBenchmark scenario. Args: -<<<<<<< HEAD objective_scorer: ``TrueFalseScorer`` used to evaluate attack success. Defaults to the registered default objective scorer (typically the composite refusal+scale scorer set up by an initializer). Widening to general ``Scorer`` support (covering ``FloatScaleScorer``, etc.) is tracked as a follow-up. - skip_cached: When ``True``, ``_get_atomic_attacks_async`` filters + use_cached: When ``True``, ``_get_atomic_attacks_async`` filters out atomic attacks for which the live behavioral cache (``pyrit.analytics.get_cached_results_for_technique``) has already returned at least one ``SUCCESS`` or ``FAILURE`` @@ -231,47 +167,7 @@ def __init__( self._objective_scorer: TrueFalseScorer = ( objective_scorer if objective_scorer else self._get_default_objective_scorer() ) - self._skip_cached: bool = skip_cached - - super().__init__( - version=self.VERSION, - objective_scorer=self._objective_scorer, - strategy_class=self.get_strategy_class(), - scenario_result_id=scenario_result_id, - ) -======= - adversarial_models: A non-empty list of ``PromptTarget`` instances - that each satisfy ``CHAT_TARGET_REQUIREMENTS`` (multi-turn - with editable history). Individual techniques selected at - run time may impose stricter capability requirements which are - enforced when their attack instances are constructed. - Labels are inferred from each target's identifier (preferring - ``underlying_model_name`` over ``model_name`` over the class - name). Identical targets are silently deduped and distinct - targets whose inferred names collide are suffixed (``_2``, - ``_3``, …) with a warning. - May be ``None`` at construction so the scenario can be - introspected (e.g. for ``--list-scenarios`` metadata); the - non-empty / capability validation is then deferred to - ``initialize_async``. - objective_scorer: Scorer for evaluating attack success. - Defaults to the registered default objective scorer. - scenario_result_id: Optional ID of an existing scenario - result to resume. - - Raises: - ValueError: If ``adversarial_models`` is provided and is empty, - not a list, or contains a target that does not satisfy - :data:`CHAT_TARGET_REQUIREMENTS`. - """ - if adversarial_models is not None: - self._adversarial_configs = self._build_adversarial_configs(adversarial_models) - else: - self._adversarial_configs = {} - - self._objective_scorer: TrueFalseScorer = ( - objective_scorer if objective_scorer else self._get_default_objective_scorer() - ) + self._use_cached: bool = use_cached strategy_class = _build_benchmark_strategy() @@ -287,61 +183,21 @@ def __init__( scenario_result_id=scenario_result_id, ) - @staticmethod - def _build_adversarial_configs( - adversarial_models: list[PromptTarget], - ) -> dict[str, AttackAdversarialConfig]: - """ - Validate ``adversarial_models`` and wrap each into an ``AttackAdversarialConfig``. - - Returns: - dict[str, AttackAdversarialConfig]: Adversarial configs keyed by inferred model label. - - Raises: - ValueError: If the list is empty, not a list, or contains a target - that does not satisfy :data:`CHAT_TARGET_REQUIREMENTS`. - """ - if not adversarial_models: - raise ValueError("adversarial_models must be a non-empty list of PromptTarget instances.") - - if not isinstance(adversarial_models, list): - raise ValueError("adversarial_models must be a list of PromptTarget instances.") - - for target in adversarial_models: - try: - CHAT_TARGET_REQUIREMENTS.validate(target=target) - except ValueError as exc: - raise ValueError( - f"adversarial_models entry {type(target).__name__} does not satisfy " - f"the chat-target capability requirements: {exc}" - ) from exc - - labeled_targets = AdversarialBenchmark._infer_labels(items=adversarial_models) - return {label: AttackAdversarialConfig(target=target) for label, target in labeled_targets.items()} ->>>>>>> 4cf9ff5de64017ce09cbe9eeb653a35a9983cab4 - async def _get_atomic_attacks_async(self) -> list[AtomicAttack]: """ Build atomic attacks from (technique × adversarial_target × dataset), then apply caching. -<<<<<<< HEAD Reads the user-supplied ``adversarial_targets`` parameter, resolves each name to a ``PromptTarget`` via ``TargetRegistry``, and cross-products the selected adversarial-capable techniques over the - resolved targets and configured datasets. Each pair builds a - non-registered per-pair factory via - ``AttackTechniqueRegistry.build_factory_from_spec`` with - ``adversarial_chat`` overridden to the resolved target — no global - registry state is touched. When ``self._skip_cached`` is set, the - final candidate list is then filtered against the live behavioral - cache via ``_collect_cached_completion_pairs``, which delegates to + resolved targets and configured datasets. Each pair calls + ``factory.create(attack_adversarial_config_override=...)`` with the + resolved target — no global registry state is touched. When + ``self._use_cached`` is set, the final candidate list is filtered + against the live behavioral cache via + ``_collect_cached_completion_pairs``, which delegates to ``pyrit.analytics.get_cached_results_for_technique`` for each unique ``(technique_eval_hash, objective_target_eval_hash)`` pair. -======= - Factories are read from the singleton ``AttackTechniqueRegistry`` and - narrowed to adversarial-capable ones. Each model is injected at - create-time via ``attack_adversarial_config_override``. ->>>>>>> 4cf9ff5de64017ce09cbe9eeb653a35a9983cab4 Returns: list[AtomicAttack]: The atomic attacks to actually execute on @@ -357,7 +213,6 @@ async def _get_atomic_attacks_async(self) -> list[AtomicAttack]: "Scenario not properly initialized. Call await scenario.initialize_async() before running." ) -<<<<<<< HEAD target_names = self.params.get("adversarial_targets") if not target_names: raise ValueError( @@ -368,37 +223,17 @@ async def _get_atomic_attacks_async(self) -> list[AtomicAttack]: ) resolved_targets = self._resolve_adversarial_targets(target_names=target_names) - # ``BenchmarkStrategy`` is built from adversarial-capable - # ``SCENARIO_TECHNIQUES`` entries only (see ``_build_benchmark_strategy``), - # so every selected strategy resolves to exactly one spec. Drift between the - # enum and the catalog is silently ignored — the next strategy-class build - # would surface it. - specs_by_name = {spec.name: spec for spec in SCENARIO_TECHNIQUES} - selected_specs = [specs_by_name[s.value] for s in self._scenario_strategies if s.value in specs_by_name] -======= - if not self._adversarial_configs: - raise ValueError( - "AdversarialBenchmark requires adversarial_models to be passed at construction " - "(non-empty list of chat-capable PromptTarget instances)." - ) - - benchmarkable_factories = AdversarialBenchmark._get_benchmarkable_factories() - local_factories = {factory.name: factory for factory in benchmarkable_factories} ->>>>>>> 4cf9ff5de64017ce09cbe9eeb653a35a9983cab4 + all_factories = AttackTechniqueRegistry.get_registry_singleton().get_factories_or_raise() + selected_factories = [ + all_factories[s.value] for s in self._scenario_strategies if s.value in all_factories + ] scoring_config = AttackScoringConfig(objective_scorer=self._objective_scorer) seed_groups_by_dataset = self._dataset_config.get_seed_attack_groups() atomic_attacks: list[AtomicAttack] = [] - for spec in selected_specs: + for factory in selected_factories: for target_name, target_instance in resolved_targets: - pair_spec = dataclasses.replace( - spec, - adversarial_chat=target_instance, - adversarial_chat_key=None, - ) - factory = AttackTechniqueRegistry.build_factory_from_spec(pair_spec) - for dataset_name, seed_groups in seed_groups_by_dataset.items(): if factory.seed_technique is not None: compatible_groups = SeedAttackGroup.filter_compatible( @@ -409,12 +244,12 @@ async def _get_atomic_attacks_async(self) -> list[AtomicAttack]: if skipped: logger.info( f"Skipped {skipped} seed group(s) from '{dataset_name}' for technique " - f"'{spec.name}' (prompt sequences overlap with simulated conversation)." + f"'{factory.name}' (prompt sequences overlap with simulated conversation)." ) if not compatible_groups: logger.warning( f"No compatible seed groups in '{dataset_name}' for technique " - f"'{spec.name}', skipping this (technique, target, dataset) triple." + f"'{factory.name}', skipping this (technique, target, dataset) triple." ) continue else: @@ -423,6 +258,7 @@ async def _get_atomic_attacks_async(self) -> list[AtomicAttack]: attack_technique = factory.create( objective_target=self._objective_target, attack_scoring_config=scoring_config, + attack_adversarial_config_override=AttackAdversarialConfig(target=target_instance), ) # ``display_group`` is set explicitly here so result roll-ups group by the # TargetRegistry name the caller passed via ``--adversarial-targets`` — @@ -432,7 +268,7 @@ async def _get_atomic_attacks_async(self) -> list[AtomicAttack]: # then reads ``aa.display_group`` directly (scenario.py:721). atomic_attacks.append( AtomicAttack( - atomic_attack_name=f"{spec.name}__{target_name}_{dataset_name}", + atomic_attack_name=f"{factory.name}__{target_name}_{dataset_name}", attack_technique=attack_technique, seed_groups=list(compatible_groups), adversarial_chat=target_instance, @@ -442,7 +278,7 @@ async def _get_atomic_attacks_async(self) -> list[AtomicAttack]: ) ) - if not self._skip_cached: + if not self._use_cached: return atomic_attacks cached_technique_hashes = self._collect_cached_completion_pairs(atomic_attacks=atomic_attacks) @@ -450,11 +286,16 @@ async def _get_atomic_attacks_async(self) -> list[AtomicAttack]: skipped = len(atomic_attacks) - len(filtered) if skipped > 0: logger.info( - "skip_cached=True: dropping %d/%d atomic attack(s) already completed for the " + "use_cached=True: skipping %d/%d atomic attack(s) already completed for the " "current objective target (matched by technique_eval_hash × objective_target_eval_hash).", skipped, len(atomic_attacks), ) + # TODO: inject prior AttackResult rows for the skipped attacks into the current ScenarioResult + # so use_cached=True produces a complete result rather than a partial one. + # The skipped attacks' names are: [a.atomic_attack_name for a in atomic_attacks if a not in filtered] + # Fetch their results via get_cached_results_for_technique and add them as pre-populated slots + # in ScenarioResult.attack_results (requires overriding initialize_async or a post-populate hook). return filtered def _resolve_adversarial_targets(self, *, target_names: list[str]) -> list[tuple[str, PromptTarget]]: @@ -493,7 +334,6 @@ def _resolve_adversarial_targets(self, *, target_names: list[str]) -> list[tuple return resolved -<<<<<<< HEAD def _collect_cached_completion_pairs(self, *, atomic_attacks: list[AtomicAttack]) -> set[str]: """ Return the set of ``technique_eval_hash`` values already cached for this scenario's objective target. @@ -562,44 +402,3 @@ def _collect_cached_completion_pairs(self, *, atomic_attacks: list[AtomicAttack] cached_hashes.add(technique_eval_hash) return cached_hashes -======= - @staticmethod - def _get_benchmarkable_factories() -> list[AttackTechniqueFactory]: - """ - Return ``core`` factories that drive an adversarial chat. - - Every benchmark technique must accept an adversarial-config override at - ``create()`` time so the scenario can inject one chat per benchmark - model. We narrow to the ``core`` tag to exclude experimental / persona - variants. - - Returns: - list[AttackTechniqueFactory]: Filtered core, adversarial-capable factories. - """ - registry = AttackTechniqueRegistry.get_registry_singleton() - return [ - factory - for factory in registry.get_factories_or_raise().values() - if factory.uses_adversarial and "core" in factory.strategy_tags - ] - - -@cache -def _build_benchmark_strategy() -> type[ScenarioStrategy]: - """ - Module-level cached builder so all callers share the same strategy enum class. - - Returns: - type[ScenarioStrategy]: The dynamically generated BenchmarkStrategy enum class. - """ - return AttackTechniqueRegistry.build_strategy_class_from_factories( # type: ignore[ty:invalid-return-type] - class_name="BenchmarkStrategy", - factories=AdversarialBenchmark._get_benchmarkable_factories(), - aggregate_tags={ - "default": TagQuery.any_of("default"), - "single_turn": TagQuery.any_of("single_turn"), - "multi_turn": TagQuery.any_of("multi_turn"), - "light": TagQuery.any_of("light"), - }, - ) ->>>>>>> 4cf9ff5de64017ce09cbe9eeb653a35a9983cab4 diff --git a/tests/unit/scenario/benchmark/test_adversarial.py b/tests/unit/scenario/benchmark/test_adversarial.py index 0342e6cbc5..381c4d51b2 100644 --- a/tests/unit/scenario/benchmark/test_adversarial.py +++ b/tests/unit/scenario/benchmark/test_adversarial.py @@ -8,13 +8,14 @@ ``supported_parameters``. Targets are user-supplied registry names that resolve to ``PromptTarget`` instances via ``TargetRegistry``. The ``(technique × target × dataset)`` cross-product is built lazily inside -``_get_atomic_attacks_async`` using per-pair non-registered factories; -no global ``AttackTechniqueRegistry`` state is mutated. +``_get_atomic_attacks_async`` using factory.create() with an +adversarial config override; no global ``AttackTechniqueRegistry`` +state is mutated. These tests cover the new contract: * Class metadata (VERSION, BASELINE policy, defaults). -* Strategy enum is built from source ``SCENARIO_TECHNIQUES`` entries that - require an adversarial chat target; ``light`` aggregate preserves the +* Strategy enum is built from registered factories with ``uses_adversarial=True`` + and the ``core`` strategy tag; ``light`` aggregate preserves the source ``light`` tag (excludes ``tap`` / ``crescendo_simulated``). * ``supported_parameters`` declares ``adversarial_targets: list[str]``. * ``_resolve_adversarial_targets`` raises with available names on typos. @@ -24,7 +25,7 @@ ``pyrit.analytics.get_cached_results_for_technique`` per unique technique hash and returns the set of technique hashes with at least one ``SUCCESS`` / ``FAILURE`` match for the scenario's objective target. -* ``skip_cached`` filters cached candidates end-to-end. +* ``use_cached`` filters cached candidates end-to-end. * Real-memory smoke for ``_collect_cached_completion_pairs`` exercises persistence -> SQL filter -> objective-target filter -> outcome filter. """ @@ -45,12 +46,42 @@ from pyrit.registry import TargetRegistry from pyrit.registry.object_registries.attack_technique_registry import AttackTechniqueRegistry from pyrit.scenario.core import BaselineAttackPolicy -from pyrit.scenario.core.scenario_techniques import SCENARIO_TECHNIQUES, _spec_needs_adversarial +from pyrit.scenario.core.attack_technique_factory import AttackTechniqueFactory from pyrit.scenario.scenarios.benchmark.adversarial import ( AdversarialBenchmark, _build_benchmark_strategy, ) from pyrit.score import TrueFalseScorer +from pyrit.setup.initializers.components.scenario_techniques import build_scenario_technique_factories + +# --------------------------------------------------------------------------- +# Module-level constants derived from the canonical factory catalog +# --------------------------------------------------------------------------- + + +def _build_benchmarkable_factories_snapshot() -> list: + """Compute benchmarkable-factory counts from the production catalog. + + Sets up a transient mock ``adversarial_chat`` in ``TargetRegistry`` so + factory construction does not depend on environment variables, then filters + by the same predicate used in ``AdversarialBenchmark._get_benchmarkable_factories``. + """ + TargetRegistry.reset_instance() + adv = MagicMock(spec=PromptTarget) + adv.capabilities.includes.return_value = True + TargetRegistry.get_registry_singleton().register_instance(adv, name="adversarial_chat") + try: + factories = build_scenario_technique_factories() + finally: + TargetRegistry.reset_instance() + return [f for f in factories if f.uses_adversarial and "core" in f.strategy_tags] + + +_BENCHMARKABLE_FACTORIES = _build_benchmarkable_factories_snapshot() +_NUM_ADVERSARIAL_TECHNIQUES = len(_BENCHMARKABLE_FACTORIES) +_BENCHMARKABLE_TECHNIQUE_NAMES = {f.name for f in _BENCHMARKABLE_FACTORIES} +_LIGHT_BENCHMARKABLE_FACTORIES = [f for f in _BENCHMARKABLE_FACTORIES if "light" in f.strategy_tags] +_NUM_LIGHT_BENCHMARKABLE = len(_LIGHT_BENCHMARKABLE_FACTORIES) # --------------------------------------------------------------------------- # Fixtures / helpers @@ -58,13 +89,26 @@ @pytest.fixture(autouse=True) -def reset_registries(): - """Reset both registries between tests so target/technique state doesn't leak.""" +def reset_technique_registry(): + """Reset registries, register a mock adversarial target, and populate real factories. + + Registers a mock ``adversarial_chat`` target so ``build_scenario_technique_factories`` + resolves without depending on environment variables. Uses ``_build_benchmark_strategy.cache_clear()`` + because our implementation uses ``@cache`` (not ``_cached_strategy_class``). + """ AttackTechniqueRegistry.reset_instance() TargetRegistry.reset_instance() + _build_benchmark_strategy.cache_clear() + + adv_target = MagicMock(spec=PromptTarget) + adv_target.capabilities.includes.return_value = True + TargetRegistry.get_registry_singleton().register_instance(adv_target, name="adversarial_chat") + + AttackTechniqueRegistry.get_registry_singleton().register_from_factories(build_scenario_technique_factories()) yield AttackTechniqueRegistry.reset_instance() TargetRegistry.reset_instance() + _build_benchmark_strategy.cache_clear() def _register_adversarial_target(*, name: str) -> PromptTarget: @@ -75,6 +119,19 @@ def _register_adversarial_target(*, name: str) -> PromptTarget: return target +def _register_mock_factory(*, name: str, tags: list[str] | None = None, seed_technique=None) -> MagicMock: + """Register a mock AttackTechniqueFactory in AttackTechniqueRegistry.""" + factory = MagicMock(spec=AttackTechniqueFactory) + factory.name = name + factory.uses_adversarial = True + factory.strategy_tags = tags if tags is not None else ["core", "light"] + factory.seed_technique = seed_technique + factory.create.return_value = MagicMock(name="AttackTechnique") + factory.attack_class = MagicMock(__name__=name) + AttackTechniqueRegistry.get_registry_singleton().register_from_factories([factory]) + return factory + + # --------------------------------------------------------------------------- # Class metadata # --------------------------------------------------------------------------- @@ -91,13 +148,6 @@ def test_baseline_attack_policy_is_forbidden(self): """A baseline contributes no signal to a model-comparison benchmark, so it is forbidden.""" assert AdversarialBenchmark.BASELINE_ATTACK_POLICY is BaselineAttackPolicy.Forbidden - def test_default_dataset_config_uses_harmbench(self): - config = AdversarialBenchmark.default_dataset_config() - assert config.get_default_dataset_names() == ["harmbench"] - - def test_default_dataset_config_max_size_is_8(self): - assert AdversarialBenchmark.default_dataset_config().max_dataset_size == 8 - # --------------------------------------------------------------------------- # supported_parameters @@ -135,69 +185,55 @@ def test_adversarial_targets_description_mentions_cli_flag(self): class TestAdversarialBenchmarkStrategy: - """Tests for ``_build_benchmark_strategy`` and the cached ``get_strategy_class`` accessor.""" + """Tests for ``_build_benchmark_strategy`` using the registry-based factory API.""" - def test_strategy_built_from_adversarial_specs(self): - """Every adversarial-capable spec in ``SCENARIO_TECHNIQUES`` produces one concrete enum member.""" + def test_strategy_built_from_registered_adversarial_factories(self): + """Each registered ``core`` adversarial factory produces one concrete enum member.""" strategy_cls = _build_benchmark_strategy() aggregate_names = {"all"} | strategy_cls.get_aggregate_tags() concrete_members = [m for m in strategy_cls if m.value not in aggregate_names] - - adversarial_specs = [s for s in SCENARIO_TECHNIQUES if _spec_needs_adversarial(s)] - adversarial_spec_names = {s.name for s in adversarial_specs} - concrete_member_values = {m.value for m in concrete_members} - assert concrete_member_values == adversarial_spec_names + assert concrete_member_values == _BENCHMARKABLE_TECHNIQUE_NAMES + + def test_strategy_excludes_non_adversarial_factories(self): + """Factories without ``uses_adversarial=True`` must not appear as enum members.""" + # Register a non-adversarial factory directly + non_adv = MagicMock(spec=AttackTechniqueFactory) + non_adv.name = "prompt_sending" + non_adv.uses_adversarial = False + non_adv.strategy_tags = ["core", "light"] + non_adv.seed_technique = None + non_adv.attack_class = MagicMock(__name__="prompt_sending") + non_adv.create.return_value = MagicMock() + AttackTechniqueRegistry.get_registry_singleton().register_from_factories([non_adv]) - def test_strategy_excludes_non_adversarial_techniques(self): - """Techniques like ``prompt_sending`` (no adversarial chat) must not be enum members.""" strategy_cls = _build_benchmark_strategy() member_values = {m.value for m in strategy_cls} - - non_adversarial = [s for s in SCENARIO_TECHNIQUES if not _spec_needs_adversarial(s)] - for spec in non_adversarial: - assert spec.name not in member_values, ( - f"{spec.name} is not adversarial-capable but appeared as a benchmark strategy member." - ) + assert "prompt_sending" not in member_values def test_strategy_includes_required_aggregates(self): """The strategy enum exposes ``light``, ``single_turn``, ``multi_turn`` aggregates.""" strategy_cls = _build_benchmark_strategy() aggregates = strategy_cls.get_aggregate_tags() - assert "light" in aggregates assert "single_turn" in aggregates assert "multi_turn" in aggregates - def test_light_aggregate_excludes_expensive_techniques(self): - """``light`` must not pull in ``tap`` or ``crescendo_simulated`` — both can take hours.""" + def test_light_aggregate_excludes_non_light_techniques(self): + """Techniques without the ``light`` tag must not appear in the ``light`` aggregate.""" strategy_cls = _build_benchmark_strategy() light_member = strategy_cls("light") - - # Expand the aggregate to its concrete child members. resolved_values = {child.value for child in strategy_cls.expand({light_member})} - assert "tap" not in resolved_values - assert "crescendo_simulated" not in resolved_values + assert "red_teaming" in resolved_values def test_light_aggregate_includes_red_teaming(self): - """Sanity check: ``red_teaming`` is adversarial-capable AND tagged ``light``.""" + """Sanity check: ``red_teaming`` tagged ``light`` appears in the ``light`` aggregate.""" strategy_cls = _build_benchmark_strategy() light_member = strategy_cls("light") resolved_values = {child.value for child in strategy_cls.expand({light_member})} assert "red_teaming" in resolved_values - def test_get_strategy_class_returns_same_enum_shape(self): - """``get_strategy_class`` rebuilds on every call; the resulting enums have identical members.""" - first = AdversarialBenchmark.get_strategy_class() - second = AdversarialBenchmark.get_strategy_class() - assert {m.value for m in first} == {m.value for m in second} - - def test_default_strategy_is_light(self): - """``get_default_strategy`` returns the ``light`` aggregate.""" - default = AdversarialBenchmark.get_default_strategy() - assert default.value == "light" - # --------------------------------------------------------------------------- # Construction (collapsed __init__) @@ -232,14 +268,14 @@ def test_construct_takes_no_models_param(self): def test_skip_cached_defaults_to_false(self): bench = AdversarialBenchmark(objective_scorer=MagicMock(spec=TrueFalseScorer)) - assert bench._skip_cached is False + assert bench._use_cached is False def test_skip_cached_can_be_set_true(self): bench = AdversarialBenchmark( objective_scorer=MagicMock(spec=TrueFalseScorer), - skip_cached=True, + use_cached=True, ) - assert bench._skip_cached is True + assert bench._use_cached is True # --------------------------------------------------------------------------- @@ -357,6 +393,11 @@ class TestGetAtomicAttacksCrossProduct: def _make_bench_with_targets(self, *, target_names: list[str]) -> AdversarialBenchmark: for name in target_names: _register_adversarial_target(name=name) + # Reset the technique registry so we can register a controllable mock factory + # whose create() return value we can inspect. + AttackTechniqueRegistry.reset_instance() + _build_benchmark_strategy.cache_clear() + _register_mock_factory(name="red_teaming", tags=["core", "light"]) bench = AdversarialBenchmark(objective_scorer=MagicMock(spec=TrueFalseScorer)) bench._objective_target = MagicMock(spec=PromptTarget) bench.params = {"adversarial_targets": target_names} @@ -372,62 +413,38 @@ def _make_bench_with_targets(self, *, target_names: list[str]) -> AdversarialBen return bench - def _patch_factory_builder(self, *, seed_technique=None): - """Return a patch context manager for ``AttackTechniqueRegistry.build_factory_from_spec``.""" - factory = MagicMock() - factory.seed_technique = seed_technique - factory.create.return_value = MagicMock(name="AttackTechnique") - return patch( - "pyrit.scenario.scenarios.benchmark.adversarial.AttackTechniqueRegistry.build_factory_from_spec", - return_value=factory, - ) - async def test_cross_product_count_matches_n_techniques_m_targets_d_datasets(self): """1 technique × 2 targets × 1 dataset = 2 atomic attacks.""" bench = self._make_bench_with_targets(target_names=["adv_a", "adv_b"]) - - with self._patch_factory_builder(): - result = await bench._get_atomic_attacks_async() - + result = await bench._get_atomic_attacks_async() assert len(result) == 2 async def test_atomic_attack_name_format_is_technique__target_dataset(self): """Name format: ``{technique}__{target}_{dataset}`` (preserves VERSION=2 cache key shape).""" bench = self._make_bench_with_targets(target_names=["adv_a"]) - - with self._patch_factory_builder(): - result = await bench._get_atomic_attacks_async() - + result = await bench._get_atomic_attacks_async() names = [a.atomic_attack_name for a in result] assert names == ["red_teaming__adv_a_harmbench"] async def test_display_group_equals_target_registry_name(self): """``display_group`` is the raw target registry name — no string parsing.""" bench = self._make_bench_with_targets(target_names=["adv_a", "adv_b"]) - - with self._patch_factory_builder(): - result = await bench._get_atomic_attacks_async() - + result = await bench._get_atomic_attacks_async() display_groups = sorted({a.display_group for a in result}) assert display_groups == ["adv_a", "adv_b"] async def test_display_group_uses_registry_name_not_target_model_name(self): - """Regression: ``display_group`` must come from the registry name passed in via - ``adversarial_targets`` — not from any internal field on the ``PromptTarget`` instance - (``_model_name``, ``_underlying_model``, ``_endpoint``, etc.). If a future refactor - causes the scenario to source ``display_group`` from the target's own attributes, - users' per-target ASR roll-ups would silently change shape based on whatever model - name the target was constructed with. - """ - # Register a target under the registry name "adv_a" with an utterly different - # internal model/endpoint identity. After resolution, display_group should still - # be "adv_a" — the registry name — not anything that leaked from the target. + """Regression: ``display_group`` must come from the registry name, not the target's internal fields.""" target = MagicMock(spec=PromptTarget) target._model_name = "totally-different-model-name" target._underlying_model = "another-model-identity" target._endpoint = "https://hijacked.example.com/openai/v1" target.name = "name-attribute-that-must-not-leak" TargetRegistry.get_registry_singleton().register_instance(target, name="adv_a") + # Reset the technique registry to get a controllable mock factory + AttackTechniqueRegistry.reset_instance() + _build_benchmark_strategy.cache_clear() + _register_mock_factory(name="red_teaming", tags=["core", "light"]) bench = AdversarialBenchmark(objective_scorer=MagicMock(spec=TrueFalseScorer)) bench._objective_target = MagicMock(spec=PromptTarget) @@ -441,35 +458,31 @@ async def test_display_group_uses_registry_name_not_target_model_name(self): bench._dataset_config = MagicMock() bench._dataset_config.get_seed_attack_groups.return_value = {"harmbench": [seed_group]} - with self._patch_factory_builder(): - result = await bench._get_atomic_attacks_async() + result = await bench._get_atomic_attacks_async() assert len(result) == 1 atomic = result[0] assert atomic.display_group == "adv_a", ( - f"display_group must equal the registry name 'adv_a', got {atomic.display_group!r}. " - "If this is failing, the scenario started sourcing display_group from the target's " - "internal attributes (_model_name, etc.) — restore the registry-name behavior." + f"display_group must equal the registry name 'adv_a', got {atomic.display_group!r}." ) - # Belt-and-suspenders: also assert the atomic_attack_name uses the registry name, - # since the same plumbing pipes both. assert atomic.atomic_attack_name == "red_teaming__adv_a_harmbench" - async def test_factory_built_per_target_with_overridden_adversarial_chat(self): - """Each (spec, target) pair gets its own ``build_factory_from_spec`` call with a replaced spec.""" + async def test_factory_create_called_per_target_with_adversarial_config_override(self): + """Each (factory, target) pair calls ``factory.create`` with an ``AttackAdversarialConfig`` override.""" bench = self._make_bench_with_targets(target_names=["adv_a", "adv_b"]) + factory = AttackTechniqueRegistry.get_registry_singleton().get_factories_or_raise()["red_teaming"] - with self._patch_factory_builder() as build_mock: - await bench._get_atomic_attacks_async() + await bench._get_atomic_attacks_async() - # 1 selected technique × 2 targets = 2 factory builds. - assert build_mock.call_count == 2 - # Each pair_spec has adversarial_chat replaced; verify the (replaced spec).adversarial_chat - # matches the corresponding registry entry. + # 1 factory × 2 targets × 1 dataset = 2 create calls + assert factory.create.call_count == 2 target_a = TargetRegistry.get_registry_singleton().get_instance_by_name("adv_a") target_b = TargetRegistry.get_registry_singleton().get_instance_by_name("adv_b") - replaced_targets = {call.args[0].adversarial_chat for call in build_mock.call_args_list} - assert replaced_targets == {target_a, target_b} + injected_targets = { + call.kwargs["attack_adversarial_config_override"].target + for call in factory.create.call_args_list + } + assert injected_targets == {target_a, target_b} # --------------------------------------------------------------------------- @@ -688,11 +701,15 @@ class TestSkipCachedFilter: _ANALYTICS_PATH = "pyrit.scenario.scenarios.benchmark.adversarial.get_cached_results_for_technique" _IDENTIFIER_PATH = "pyrit.scenario.scenarios.benchmark.adversarial.ObjectiveTargetEvaluationIdentifier" - def _make_bench(self, *, skip_cached: bool) -> AdversarialBenchmark: + def _make_bench(self, *, use_cached: bool) -> AdversarialBenchmark: _register_adversarial_target(name="adv_a") + # Reset the technique registry to get a controllable mock factory + AttackTechniqueRegistry.reset_instance() + _build_benchmark_strategy.cache_clear() + _register_mock_factory(name="red_teaming", tags=["core", "light"]) bench = AdversarialBenchmark( objective_scorer=MagicMock(spec=TrueFalseScorer), - skip_cached=skip_cached, + use_cached=use_cached, ) bench._objective_target = MagicMock(spec=PromptTarget) bench._objective_target_identifier = MagicMock() @@ -708,36 +725,24 @@ def _make_bench(self, *, skip_cached: bool) -> AdversarialBenchmark: return bench - def _patch_factory_builder(self): - factory = MagicMock() - factory.seed_technique = None - factory.create.return_value = MagicMock(name="AttackTechnique") - return patch( - "pyrit.scenario.scenarios.benchmark.adversarial.AttackTechniqueRegistry.build_factory_from_spec", - return_value=factory, - ) - def _patch_identifier(self, eval_hash: str = "obj_hash"): identifier_instance = MagicMock() identifier_instance.eval_hash = eval_hash return patch(self._IDENTIFIER_PATH, return_value=identifier_instance) - async def test_skip_cached_false_returns_all_candidates_without_analytics_call(self): - bench = self._make_bench(skip_cached=False) + async def test_use_cached_false_returns_all_candidates_without_analytics_call(self): + bench = self._make_bench(use_cached=False) - with self._patch_factory_builder(), patch(self._ANALYTICS_PATH) as analytics_mock: + with patch(self._ANALYTICS_PATH) as analytics_mock: result = await bench._get_atomic_attacks_async() assert len(result) == 1 analytics_mock.assert_not_called() - async def test_skip_cached_true_filters_matching_candidates(self): - bench = self._make_bench(skip_cached=True) + async def test_use_cached_true_filters_matching_candidates(self): + bench = self._make_bench(use_cached=True) - # Stub every candidate's technique_eval_hash to a known value so the analytics - # lookup key matches the cached set. with ( - self._patch_factory_builder(), self._patch_identifier(), patch( "pyrit.scenario.core.atomic_attack.AtomicAttack.technique_eval_hash", @@ -752,12 +757,10 @@ async def test_skip_cached_true_filters_matching_candidates(self): assert result == [] - async def test_skip_cached_true_keeps_unmatched_candidates(self): - bench = self._make_bench(skip_cached=True) + async def test_use_cached_true_keeps_unmatched_candidates(self): + bench = self._make_bench(use_cached=True) - # Analytics returns no matches → no candidate is cached, so all pass through. with ( - self._patch_factory_builder(), self._patch_identifier(), patch(self._ANALYTICS_PATH, return_value=[]), ): From 6fc708a5928182b6d5bad3917b991acc4d508db9 Mon Sep 17 00:00:00 2001 From: Victor Valbuena Date: Mon, 1 Jun 2026 11:46:22 -0700 Subject: [PATCH 38/40] FEAT: inject cached attack results into AdversarialBenchmark ScenarioResult When use_cached=True, skipped atomic attacks now have their prior cached AttackResults attached to the live ScenarioResult instead of being silently dropped. This addresses the PR #1765 review comment that we still need to surface cached runs in the final scenario output (not just skip execution). - _collect_cached_completion_pairs now stores per-hash cached results as a side effect for downstream lookup. - _get_atomic_attacks_async filters cached rows by attribution_data['parent_collection'] so a cache hit from one dataset/target slot does not leak into another. - run_async override merges _precomputed_cached_results into ScenarioResult.attack_results and updates _display_group_map. - Adds TestRunAsyncCacheInjection (3 tests) and two TestSkipCachedFilter tests covering the full pipeline and parent_collection filtering. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- .../scenarios/benchmark/adversarial.py | 70 +++++++-- .../scenario/benchmark/test_adversarial.py | 144 +++++++++++++++++- 2 files changed, 197 insertions(+), 17 deletions(-) diff --git a/pyrit/scenario/scenarios/benchmark/adversarial.py b/pyrit/scenario/scenarios/benchmark/adversarial.py index a80ea74492..e032928b2f 100644 --- a/pyrit/scenario/scenarios/benchmark/adversarial.py +++ b/pyrit/scenario/scenarios/benchmark/adversarial.py @@ -12,8 +12,13 @@ from pyrit.analytics import get_cached_results_for_technique from pyrit.common import Parameter, apply_defaults from pyrit.executor.attack import AttackAdversarialConfig, AttackScoringConfig -from pyrit.models import ObjectiveTargetEvaluationIdentifier -from pyrit.models import AttackOutcome, SeedAttackGroup +from pyrit.models import ( + AttackOutcome, + AttackResult, + ObjectiveTargetEvaluationIdentifier, + ScenarioResult, + SeedAttackGroup, +) from pyrit.registry import AttackTechniqueRegistry, TargetRegistry from pyrit.registry.tag_query import TagQuery from pyrit.scenario.core.atomic_attack import AtomicAttack @@ -22,7 +27,6 @@ if TYPE_CHECKING: from pyrit.prompt_target import PromptTarget - from pyrit.scenario.core.attack_technique_factory import AttackTechniqueFactory from pyrit.scenario.core.scenario_strategy import ScenarioStrategy from pyrit.score.true_false.true_false_scorer import TrueFalseScorer @@ -168,6 +172,9 @@ def __init__( objective_scorer if objective_scorer else self._get_default_objective_scorer() ) self._use_cached: bool = use_cached + self._precomputed_cached_results: dict[str, list[AttackResult]] = {} + self._precomputed_cached_display_groups: dict[str, str] = {} + self._cached_results_by_hash: dict[str, list[AttackResult]] = {} strategy_class = _build_benchmark_strategy() @@ -224,9 +231,7 @@ async def _get_atomic_attacks_async(self) -> list[AtomicAttack]: resolved_targets = self._resolve_adversarial_targets(target_names=target_names) all_factories = AttackTechniqueRegistry.get_registry_singleton().get_factories_or_raise() - selected_factories = [ - all_factories[s.value] for s in self._scenario_strategies if s.value in all_factories - ] + selected_factories = [all_factories[s.value] for s in self._scenario_strategies if s.value in all_factories] scoring_config = AttackScoringConfig(objective_scorer=self._objective_scorer) seed_groups_by_dataset = self._dataset_config.get_seed_attack_groups() @@ -283,19 +288,27 @@ async def _get_atomic_attacks_async(self) -> list[AtomicAttack]: cached_technique_hashes = self._collect_cached_completion_pairs(atomic_attacks=atomic_attacks) filtered = [c for c in atomic_attacks if c.technique_eval_hash not in cached_technique_hashes] - skipped = len(atomic_attacks) - len(filtered) - if skipped > 0: + skipped_attacks = [c for c in atomic_attacks if c.technique_eval_hash in cached_technique_hashes] + if skipped_attacks: logger.info( "use_cached=True: skipping %d/%d atomic attack(s) already completed for the " "current objective target (matched by technique_eval_hash × objective_target_eval_hash).", - skipped, + len(skipped_attacks), len(atomic_attacks), ) - # TODO: inject prior AttackResult rows for the skipped attacks into the current ScenarioResult - # so use_cached=True produces a complete result rather than a partial one. - # The skipped attacks' names are: [a.atomic_attack_name for a in atomic_attacks if a not in filtered] - # Fetch their results via get_cached_results_for_technique and add them as pre-populated slots - # in ScenarioResult.attack_results (requires overriding initialize_async or a post-populate hook). + # Pre-populate prior results for skipped attacks so run_async can surface them in + # ScenarioResult.attack_results. attribution_data["parent_collection"] records the + # original atomic_attack_name so results are attributed to the correct slot. + self._precomputed_cached_results = {} + self._precomputed_cached_display_groups = {} + for attack in skipped_attacks: + prior = self._cached_results_by_hash.get(attack.technique_eval_hash, []) + self._precomputed_cached_results[attack.atomic_attack_name] = [ + r + for r in prior + if r.attribution_data and r.attribution_data.get("parent_collection") == attack.atomic_attack_name + ] + self._precomputed_cached_display_groups[attack.atomic_attack_name] = attack.display_group return filtered def _resolve_adversarial_targets(self, *, target_names: list[str]) -> list[tuple[str, PromptTarget]]: @@ -334,6 +347,28 @@ def _resolve_adversarial_targets(self, *, target_names: list[str]) -> list[tuple return resolved + async def run_async(self) -> ScenarioResult: + """ + Run the scenario and merge any precomputed cached results into the returned ``ScenarioResult``. + + When ``use_cached=True`` skipped atomic attacks whose prior results were + loaded during ``_get_atomic_attacks_async``, this override attaches + those results (and their display-group labels) to the live scenario + result so the final report reflects both newly-executed and + cache-served runs. + + Returns: + ScenarioResult: The scenario result with cached attack results merged + into ``attack_results`` and cached display groups merged into + ``_display_group_map``. + """ + result = await super().run_async() + if self._precomputed_cached_results: + for attack_name, prior_results in self._precomputed_cached_results.items(): + result.attack_results.setdefault(attack_name, []).extend(prior_results) + result._display_group_map.update(self._precomputed_cached_display_groups) + return result + def _collect_cached_completion_pairs(self, *, atomic_attacks: list[AtomicAttack]) -> set[str]: """ Return the set of ``technique_eval_hash`` values already cached for this scenario's objective target. @@ -354,6 +389,11 @@ def _collect_cached_completion_pairs(self, *, atomic_attacks: list[AtomicAttack] × objective target) result counts as a hit, regardless of scenario name or ``VERSION``. + As a side effect, populates ``self._cached_results_by_hash`` with the + retrieved ``AttackResult`` lists keyed by technique eval hash so that + ``_get_atomic_attacks_async`` can build ``_precomputed_cached_results`` + for injection into the final ``ScenarioResult`` by ``run_async``. + Args: atomic_attacks: The candidate atomic attacks built earlier in ``_get_atomic_attacks_async``. Only their @@ -367,6 +407,7 @@ def _collect_cached_completion_pairs(self, *, atomic_attacks: list[AtomicAttack] than blocking the run. """ cached_hashes: set[str] = set() + self._cached_results_by_hash: dict[str, list[AttackResult]] = {} if self._objective_target_identifier is None: return cached_hashes @@ -400,5 +441,6 @@ def _collect_cached_completion_pairs(self, *, atomic_attacks: list[AtomicAttack] continue if any(m.outcome in (AttackOutcome.SUCCESS, AttackOutcome.FAILURE) for m in matches): cached_hashes.add(technique_eval_hash) + self._cached_results_by_hash[technique_eval_hash] = list(matches) return cached_hashes diff --git a/tests/unit/scenario/benchmark/test_adversarial.py b/tests/unit/scenario/benchmark/test_adversarial.py index 8695c3efd6..9a6eaef8dd 100644 --- a/tests/unit/scenario/benchmark/test_adversarial.py +++ b/tests/unit/scenario/benchmark/test_adversarial.py @@ -31,7 +31,7 @@ """ from datetime import datetime, timezone -from unittest.mock import MagicMock, patch +from unittest.mock import AsyncMock, MagicMock, patch import pytest @@ -50,6 +50,7 @@ from pyrit.registry.object_registries.attack_technique_registry import AttackTechniqueRegistry from pyrit.scenario.core import BaselineAttackPolicy from pyrit.scenario.core.attack_technique_factory import AttackTechniqueFactory +from pyrit.scenario.core.scenario import Scenario from pyrit.scenario.scenarios.benchmark.adversarial import ( AdversarialBenchmark, _build_benchmark_strategy, @@ -482,8 +483,7 @@ async def test_factory_create_called_per_target_with_adversarial_config_override target_a = TargetRegistry.get_registry_singleton().get_instance_by_name("adv_a") target_b = TargetRegistry.get_registry_singleton().get_instance_by_name("adv_b") injected_targets = { - call.kwargs["attack_adversarial_config_override"].target - for call in factory.create.call_args_list + call.kwargs["attack_adversarial_config_override"].target for call in factory.create.call_args_list } assert injected_targets == {target_a, target_b} @@ -505,6 +505,14 @@ def _make_attack_result_with_outcome(outcome: AttackOutcome) -> MagicMock: return ar +def _make_attack_result_with_attribution(*, outcome: AttackOutcome, parent_collection: str) -> MagicMock: + """Like ``_make_attack_result_with_outcome`` but with attribution_data for parent-collection filtering.""" + ar = MagicMock() + ar.outcome = outcome + ar.attribution_data = {"parent_collection": parent_collection} + return ar + + @pytest.mark.usefixtures("patch_central_database") class TestCollectCachedCompletionPairs: """Tests for ``_collect_cached_completion_pairs`` — now delegates to ``pyrit.analytics``.""" @@ -771,6 +779,52 @@ async def test_use_cached_true_keeps_unmatched_candidates(self): assert len(result) == 1 + async def test_use_cached_true_populates_precomputed_maps_for_skipped(self): + """Full pipeline: cache hit → _precomputed_cached_results/display_groups populated for skipped slot.""" + bench = self._make_bench(use_cached=True) + cached_attack = _make_attack_result_with_attribution( + outcome=AttackOutcome.SUCCESS, + parent_collection="red_teaming__adv_a_harmbench", + ) + + with ( + self._patch_identifier(), + patch( + "pyrit.scenario.core.atomic_attack.AtomicAttack.technique_eval_hash", + new_callable=lambda: property(lambda self: "cached_hash"), + ), + patch(self._ANALYTICS_PATH, return_value=[cached_attack]), + ): + result = await bench._get_atomic_attacks_async() + + assert result == [] + assert bench._precomputed_cached_results == {"red_teaming__adv_a_harmbench": [cached_attack]} + assert bench._precomputed_cached_display_groups == {"red_teaming__adv_a_harmbench": "adv_a"} + + async def test_use_cached_true_filters_results_by_parent_collection(self): + """Cached rows whose parent_collection doesn't match the skipped slot are dropped.""" + bench = self._make_bench(use_cached=True) + matching = _make_attack_result_with_attribution( + outcome=AttackOutcome.SUCCESS, + parent_collection="red_teaming__adv_a_harmbench", + ) + wrong_parent = _make_attack_result_with_attribution( + outcome=AttackOutcome.SUCCESS, + parent_collection="red_teaming__adv_a_xstest", + ) + + with ( + self._patch_identifier(), + patch( + "pyrit.scenario.core.atomic_attack.AtomicAttack.technique_eval_hash", + new_callable=lambda: property(lambda self: "cached_hash"), + ), + patch(self._ANALYTICS_PATH, return_value=[matching, wrong_parent]), + ): + await bench._get_atomic_attacks_async() + + assert bench._precomputed_cached_results == {"red_teaming__adv_a_harmbench": [matching]} + # --------------------------------------------------------------------------- # Real-memory coverage for _collect_cached_completion_pairs @@ -972,3 +1026,87 @@ def test_dedupes_candidates_with_same_technique_hash(self, sqlite_instance): result = bench._collect_cached_completion_pairs(atomic_attacks=candidates) assert result == {tech_hash} + + +# --------------------------------------------------------------------------- +# run_async cache injection +# --------------------------------------------------------------------------- + + +@pytest.mark.usefixtures("patch_central_database") +class TestRunAsyncCacheInjection: + """Tests that prior results for use_cached-skipped attacks are injected into ScenarioResult.""" + + async def test_precomputed_results_injected_into_attack_results(self): + """Slots from prior runs appear in attack_results alongside freshly-executed results.""" + bench = AdversarialBenchmark( + objective_scorer=MagicMock(spec=TrueFalseScorer), + use_cached=True, + ) + + result_x = MagicMock(spec=AttackResult) + result_y = MagicMock(spec=AttackResult) + result_z = MagicMock(spec=AttackResult) + + # Simulate what _get_atomic_attacks_async populated for the two skipped attacks + bench._precomputed_cached_results = { + "technique_a__adv_target_harmbench": [result_x], + "technique_b__adv_target_harmbench": [result_y], + } + bench._precomputed_cached_display_groups = { + "technique_a__adv_target_harmbench": "adv_target", + "technique_b__adv_target_harmbench": "adv_target", + } + + # Base run_async produced only the non-skipped attack's result + base_scenario_result = MagicMock() + base_scenario_result.attack_results = {"technique_c__adv_target_harmbench": [result_z]} + base_scenario_result._display_group_map = {} + + with patch.object(Scenario, "run_async", new=AsyncMock(return_value=base_scenario_result)): + result = await bench.run_async() + + assert set(result.attack_results.keys()) == { + "technique_a__adv_target_harmbench", + "technique_b__adv_target_harmbench", + "technique_c__adv_target_harmbench", + } + assert result.attack_results["technique_a__adv_target_harmbench"] == [result_x] + assert result.attack_results["technique_b__adv_target_harmbench"] == [result_y] + assert result.attack_results["technique_c__adv_target_harmbench"] == [result_z] + + async def test_display_group_map_updated_for_cached_attacks(self): + """Skipped attacks have their display group injected so get_display_groups aggregates correctly.""" + bench = AdversarialBenchmark( + objective_scorer=MagicMock(spec=TrueFalseScorer), + use_cached=True, + ) + + bench._precomputed_cached_results = {"technique_a__adv_target_harmbench": [MagicMock(spec=AttackResult)]} + bench._precomputed_cached_display_groups = {"technique_a__adv_target_harmbench": "adv_target"} + + base_scenario_result = MagicMock() + base_scenario_result.attack_results = {} + base_scenario_result._display_group_map = {} + + with patch.object(Scenario, "run_async", new=AsyncMock(return_value=base_scenario_result)): + result = await bench.run_async() + + assert result._display_group_map["technique_a__adv_target_harmbench"] == "adv_target" + + async def test_no_injection_when_no_cached_attacks(self): + """When all attacks were executed freshly, attack_results is returned unchanged.""" + bench = AdversarialBenchmark( + objective_scorer=MagicMock(spec=TrueFalseScorer), + use_cached=False, + ) + + result_z = MagicMock(spec=AttackResult) + base_scenario_result = MagicMock() + base_scenario_result.attack_results = {"technique_c__adv_target_harmbench": [result_z]} + base_scenario_result._display_group_map = {} + + with patch.object(Scenario, "run_async", new=AsyncMock(return_value=base_scenario_result)): + result = await bench.run_async() + + assert set(result.attack_results.keys()) == {"technique_c__adv_target_harmbench"} From 7672e1ea89888110c5d5cd64e2909a84ea86d565 Mon Sep 17 00:00:00 2001 From: Victor Valbuena Date: Mon, 1 Jun 2026 12:08:37 -0700 Subject: [PATCH 39/40] feat: add dataset-level scoping to _collect_cached_completion_pairs Changes _collect_cached_completion_pairs to return a set of tomic_attack_names (instead of technique hashes) and filters cached results per-slot using attribution_data['parent_collection']. This fixes a bug where two atomic attacks sharing the same technique+target hash (e.g. harmbench vs advbench) would incorrectly share cache hits, causing one dataset to be skipped with empty results. The dataset filter is a Python-side semantic filter, not a DB query, since get_cached_results_for_technique has no attribution parameter. Documented explicitly in the docstring. Also adds two new cross-dataset real-memory regression tests to verify harmbench and advbench results are correctly scoped. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- .../scenarios/benchmark/adversarial.py | 110 ++++++------ .../scenario/benchmark/test_adversarial.py | 158 +++++++++++++----- 2 files changed, 180 insertions(+), 88 deletions(-) diff --git a/pyrit/scenario/scenarios/benchmark/adversarial.py b/pyrit/scenario/scenarios/benchmark/adversarial.py index e032928b2f..f6d57d2d29 100644 --- a/pyrit/scenario/scenarios/benchmark/adversarial.py +++ b/pyrit/scenario/scenarios/benchmark/adversarial.py @@ -174,7 +174,7 @@ def __init__( self._use_cached: bool = use_cached self._precomputed_cached_results: dict[str, list[AttackResult]] = {} self._precomputed_cached_display_groups: dict[str, str] = {} - self._cached_results_by_hash: dict[str, list[AttackResult]] = {} + self._cached_results_by_name: dict[str, list[AttackResult]] = {} strategy_class = _build_benchmark_strategy() @@ -286,28 +286,25 @@ async def _get_atomic_attacks_async(self) -> list[AtomicAttack]: if not self._use_cached: return atomic_attacks - cached_technique_hashes = self._collect_cached_completion_pairs(atomic_attacks=atomic_attacks) - filtered = [c for c in atomic_attacks if c.technique_eval_hash not in cached_technique_hashes] - skipped_attacks = [c for c in atomic_attacks if c.technique_eval_hash in cached_technique_hashes] + cached_attack_names = self._collect_cached_completion_pairs(atomic_attacks=atomic_attacks) + filtered = [c for c in atomic_attacks if c.atomic_attack_name not in cached_attack_names] + skipped_attacks = [c for c in atomic_attacks if c.atomic_attack_name in cached_attack_names] if skipped_attacks: logger.info( "use_cached=True: skipping %d/%d atomic attack(s) already completed for the " - "current objective target (matched by technique_eval_hash × objective_target_eval_hash).", + 'current objective target (dataset-scoped via attribution_data["parent_collection"]).', len(skipped_attacks), len(atomic_attacks), ) # Pre-populate prior results for skipped attacks so run_async can surface them in - # ScenarioResult.attack_results. attribution_data["parent_collection"] records the - # original atomic_attack_name so results are attributed to the correct slot. + # ScenarioResult.attack_results. _cached_results_by_name already holds the + # attribution-filtered list keyed by atomic_attack_name, so no further filtering needed. self._precomputed_cached_results = {} self._precomputed_cached_display_groups = {} for attack in skipped_attacks: - prior = self._cached_results_by_hash.get(attack.technique_eval_hash, []) - self._precomputed_cached_results[attack.atomic_attack_name] = [ - r - for r in prior - if r.attribution_data and r.attribution_data.get("parent_collection") == attack.atomic_attack_name - ] + self._precomputed_cached_results[attack.atomic_attack_name] = self._cached_results_by_name.get( + attack.atomic_attack_name, [] + ) self._precomputed_cached_display_groups[attack.atomic_attack_name] = attack.display_group return filtered @@ -371,46 +368,53 @@ async def run_async(self) -> ScenarioResult: def _collect_cached_completion_pairs(self, *, atomic_attacks: list[AtomicAttack]) -> set[str]: """ - Return the set of ``technique_eval_hash`` values already cached for this scenario's objective target. - - Delegates to ``pyrit.analytics.get_cached_results_for_technique`` for - each unique technique hash among ``atomic_attacks``. A technique is - considered cached when the analytics helper returns at least one - ``AttackResult`` with outcome ``SUCCESS`` or ``FAILURE`` for the - ``(technique_eval_hash × objective_target_eval_hash)`` pair — - ``ERROR`` and ``UNDETERMINED`` outcomes are ignored so transient - failures retry on the next run. + Return the set of ``atomic_attack_name`` values already cached for this scenario's objective target. + + Database queries are deduplicated by unique ``technique_eval_hash`` (one query per hash, + regardless of how many atomic attacks share that hash), then the skip eligibility + decision is applied per-atomic-attack using a Python-side filter on + ``attribution_data["parent_collection"]``. + + **Dataset-level scoping is implemented as a semantic Python filter, not a database query.** + ``get_cached_results_for_technique`` has no ``dataset`` parameter; it returns all results + for a given ``(technique_eval_hash × objective_target_eval_hash)`` pair regardless of which + dataset they came from. The scoping happens here: a retrieved result only counts toward the + skip decision for atomic-attack *X* if its ``attribution_data["parent_collection"]`` equals + ``X.atomic_attack_name``. This means two atomic attacks that share a technique+target hash + (e.g. the same red-teaming technique run against the same model for both ``harmbench`` and + ``advbench``) are cached independently: a harmbench result will never cause the advbench + slot to be skipped. + + A dataset slot is considered cached when the attribution-filtered result set contains at + least one ``AttackResult`` with outcome ``SUCCESS`` or ``FAILURE`` — + ``ERROR`` and ``UNDETERMINED`` outcomes are ignored so transient failures retry on the + next run. The objective-target eval hash is computed once from ``self._objective_target_identifier`` (populated by the base ``Scenario.initialize_async``) via - ``ObjectiveTargetEvaluationIdentifier``. The cache is intentionally - scenario-agnostic: any prior run that produced a matching (technique - × objective target) result counts as a hit, regardless of scenario - name or ``VERSION``. + ``ObjectiveTargetEvaluationIdentifier``. - As a side effect, populates ``self._cached_results_by_hash`` with the - retrieved ``AttackResult`` lists keyed by technique eval hash so that - ``_get_atomic_attacks_async`` can build ``_precomputed_cached_results`` - for injection into the final ``ScenarioResult`` by ``run_async``. + As a side effect, populates ``self._cached_results_by_name`` with the + attribution-filtered ``AttackResult`` lists keyed by ``atomic_attack_name`` so that + ``_get_atomic_attacks_async`` can inject them into the final ``ScenarioResult`` + via ``run_async`` without re-filtering. Args: atomic_attacks: The candidate atomic attacks built earlier in - ``_get_atomic_attacks_async``. Only their - ``technique_eval_hash`` values are read. + ``_get_atomic_attacks_async``. Returns: - set[str]: ``technique_eval_hash`` values that have at least one - qualifying cached ``AttackResult``. Empty set when the scenario - has no objective target identifier or every analytics lookup - fails (logged at warning level) — caching becomes a no-op rather - than blocking the run. + set[str]: ``atomic_attack_name`` values that have at least one qualifying cached + ``AttackResult``. Empty set when the scenario has no objective target identifier + or every analytics lookup fails (logged at warning level) — caching becomes a + no-op rather than blocking the run. """ - cached_hashes: set[str] = set() - self._cached_results_by_hash: dict[str, list[AttackResult]] = {} + cached_names: set[str] = set() + self._cached_results_by_name: dict[str, list[AttackResult]] = {} if self._objective_target_identifier is None: - return cached_hashes + return cached_names try: objective_target_eval_hash = ObjectiveTargetEvaluationIdentifier( @@ -421,13 +425,15 @@ def _collect_cached_completion_pairs(self, *, atomic_attacks: list[AtomicAttack] "skip_cached: failed to compute objective_target eval hash (%s); skipping cache filter.", exc, ) - return cached_hashes + return cached_names unique_technique_hashes = {c.technique_eval_hash for c in atomic_attacks if c.technique_eval_hash} + # One DB query per unique hash (deduplication), results stored temporarily by hash. + raw_results_by_hash: dict[str, list[AttackResult]] = {} for technique_eval_hash in unique_technique_hashes: try: - matches = get_cached_results_for_technique( + raw_results_by_hash[technique_eval_hash] = get_cached_results_for_technique( self._memory, technique_eval_hash=technique_eval_hash, objective_target_eval_hash=objective_target_eval_hash, @@ -438,9 +444,19 @@ def _collect_cached_completion_pairs(self, *, atomic_attacks: list[AtomicAttack] technique_eval_hash, exc, ) - continue - if any(m.outcome in (AttackOutcome.SUCCESS, AttackOutcome.FAILURE) for m in matches): - cached_hashes.add(technique_eval_hash) - self._cached_results_by_hash[technique_eval_hash] = list(matches) - return cached_hashes + # Per-attack attribution filter: only count results that were produced for this + # specific atomic_attack_name slot (dataset-level scoping via parent_collection). + for attack in atomic_attacks: + if not attack.technique_eval_hash or attack.technique_eval_hash not in raw_results_by_hash: + continue + attributed = [ + r + for r in raw_results_by_hash[attack.technique_eval_hash] + if r.attribution_data and r.attribution_data.get("parent_collection") == attack.atomic_attack_name + ] + if any(r.outcome in (AttackOutcome.SUCCESS, AttackOutcome.FAILURE) for r in attributed): + cached_names.add(attack.atomic_attack_name) + self._cached_results_by_name[attack.atomic_attack_name] = attributed + + return cached_names diff --git a/tests/unit/scenario/benchmark/test_adversarial.py b/tests/unit/scenario/benchmark/test_adversarial.py index 9a6eaef8dd..b89b5ddc67 100644 --- a/tests/unit/scenario/benchmark/test_adversarial.py +++ b/tests/unit/scenario/benchmark/test_adversarial.py @@ -526,9 +526,10 @@ def _make_bench(self, *, with_target_identifier: bool = True) -> AdversarialBenc bench._objective_target_identifier = MagicMock() if with_target_identifier else None return bench - def _make_candidate(self, *, technique_eval_hash: str | None) -> MagicMock: + def _make_candidate(self, *, technique_eval_hash: str | None, atomic_attack_name: str = "attack_a") -> MagicMock: candidate = MagicMock() candidate.technique_eval_hash = technique_eval_hash + candidate.atomic_attack_name = atomic_attack_name return candidate def _patch_identifier(self, eval_hash: str = "obj_target_hash"): @@ -559,46 +560,52 @@ def test_returns_empty_when_no_atomic_attacks(self): def test_returns_hash_when_success_match_exists(self): bench = self._make_bench() - candidates = [self._make_candidate(technique_eval_hash="hash_a")] + candidates = [self._make_candidate(technique_eval_hash="hash_a", atomic_attack_name="attack_a")] with ( self._patch_identifier(eval_hash="obj_hash"), patch( self._ANALYTICS_PATH, - return_value=[_make_attack_result_with_outcome(AttackOutcome.SUCCESS)], + return_value=[ + _make_attack_result_with_attribution(outcome=AttackOutcome.SUCCESS, parent_collection="attack_a") + ], ), ): cached = bench._collect_cached_completion_pairs(atomic_attacks=candidates) - assert cached == {"hash_a"} + assert cached == {"attack_a"} def test_returns_hash_when_failure_match_exists(self): bench = self._make_bench() - candidates = [self._make_candidate(technique_eval_hash="hash_a")] + candidates = [self._make_candidate(technique_eval_hash="hash_a", atomic_attack_name="attack_a")] with ( self._patch_identifier(), patch( self._ANALYTICS_PATH, - return_value=[_make_attack_result_with_outcome(AttackOutcome.FAILURE)], + return_value=[ + _make_attack_result_with_attribution(outcome=AttackOutcome.FAILURE, parent_collection="attack_a") + ], ), ): cached = bench._collect_cached_completion_pairs(atomic_attacks=candidates) - assert cached == {"hash_a"} + assert cached == {"attack_a"} def test_excludes_hash_when_only_error_or_undetermined_matches(self): """ERROR / UNDETERMINED outcomes must NOT count as cached so transient failures retry.""" bench = self._make_bench() - candidates = [self._make_candidate(technique_eval_hash="hash_a")] + candidates = [self._make_candidate(technique_eval_hash="hash_a", atomic_attack_name="attack_a")] with ( self._patch_identifier(), patch( self._ANALYTICS_PATH, return_value=[ - _make_attack_result_with_outcome(AttackOutcome.ERROR), - _make_attack_result_with_outcome(AttackOutcome.UNDETERMINED), + _make_attack_result_with_attribution(outcome=AttackOutcome.ERROR, parent_collection="attack_a"), + _make_attack_result_with_attribution( + outcome=AttackOutcome.UNDETERMINED, parent_collection="attack_a" + ), ], ), ): @@ -616,24 +623,34 @@ def test_excludes_hash_when_no_matches(self): assert cached == set() def test_dedupes_unique_technique_hashes_across_candidates(self): - """Three candidates sharing two unique hashes → analytics called twice, not three times.""" + """Three candidates sharing two unique hashes → analytics called twice, not three times. + + Two candidates share hash_a (attack_a1 and attack_a2); one has hash_b (attack_b1). + The analytics mock returns results attributed to each name, so all three attacks + are independently cached. Key assertion: DB is called twice (deduplicated by hash). + """ bench = self._make_bench() candidates = [ - self._make_candidate(technique_eval_hash="hash_a"), - self._make_candidate(technique_eval_hash="hash_b"), - self._make_candidate(technique_eval_hash="hash_a"), # duplicate + self._make_candidate(technique_eval_hash="hash_a", atomic_attack_name="attack_a1"), + self._make_candidate(technique_eval_hash="hash_b", atomic_attack_name="attack_b1"), + self._make_candidate(technique_eval_hash="hash_a", atomic_attack_name="attack_a2"), ] + def _fake_analytics(_memory, *, technique_eval_hash, objective_target_eval_hash): + if technique_eval_hash == "hash_a": + return [ + _make_attack_result_with_attribution(outcome=AttackOutcome.SUCCESS, parent_collection="attack_a1"), + _make_attack_result_with_attribution(outcome=AttackOutcome.SUCCESS, parent_collection="attack_a2"), + ] + return [_make_attack_result_with_attribution(outcome=AttackOutcome.SUCCESS, parent_collection="attack_b1")] + with ( self._patch_identifier(eval_hash="obj_hash"), - patch( - self._ANALYTICS_PATH, - return_value=[_make_attack_result_with_outcome(AttackOutcome.SUCCESS)], - ) as analytics_mock, + patch(self._ANALYTICS_PATH, side_effect=_fake_analytics) as analytics_mock, ): cached = bench._collect_cached_completion_pairs(atomic_attacks=candidates) - assert cached == {"hash_a", "hash_b"} + assert cached == {"attack_a1", "attack_b1", "attack_a2"} assert analytics_mock.call_count == 2 called_technique_hashes = {call.kwargs["technique_eval_hash"] for call in analytics_mock.call_args_list} assert called_technique_hashes == {"hash_a", "hash_b"} @@ -670,20 +687,20 @@ def test_analytics_lookup_exception_is_swallowed_per_hash(self): """A failing analytics lookup for one hash must not block the others — that hash is not cached.""" bench = self._make_bench() candidates = [ - self._make_candidate(technique_eval_hash="hash_a"), - self._make_candidate(technique_eval_hash="hash_b"), + self._make_candidate(technique_eval_hash="hash_a", atomic_attack_name="attack_a"), + self._make_candidate(technique_eval_hash="hash_b", atomic_attack_name="attack_b"), ] def fake_analytics(_memory, *, technique_eval_hash, objective_target_eval_hash): if technique_eval_hash == "hash_a": raise RuntimeError("analytics blew up") - return [_make_attack_result_with_outcome(AttackOutcome.SUCCESS)] + return [_make_attack_result_with_attribution(outcome=AttackOutcome.SUCCESS, parent_collection="attack_b")] with self._patch_identifier(), patch(self._ANALYTICS_PATH, side_effect=fake_analytics): cached = bench._collect_cached_completion_pairs(atomic_attacks=candidates) - # hash_a was the failed lookup → not cached (will retry). hash_b succeeded → cached. - assert cached == {"hash_b"} + # hash_a was the failed lookup → not cached (will retry). hash_b succeeded → cached by name. + assert cached == {"attack_b"} def test_identifier_construction_failure_falls_back_to_empty(self): """If ``ObjectiveTargetEvaluationIdentifier`` raises, cache becomes a no-op rather than blocking.""" @@ -761,7 +778,12 @@ async def test_use_cached_true_filters_matching_candidates(self): ), patch( self._ANALYTICS_PATH, - return_value=[_make_attack_result_with_outcome(AttackOutcome.SUCCESS)], + return_value=[ + _make_attack_result_with_attribution( + outcome=AttackOutcome.SUCCESS, + parent_collection="red_teaming__adv_a_harmbench", + ) + ], ), ): result = await bench._get_atomic_attacks_async() @@ -884,14 +906,22 @@ def _persist_attack_result( *, outcome: AttackOutcome, objective: str = "probe target", + atomic_attack_name: str | None = None, ) -> AttackResult: - """Persist a real AttackResult with a well-formed identifier tree.""" + """Persist a real AttackResult with a well-formed identifier tree. + + When ``atomic_attack_name`` is provided, ``attribution_data`` is stamped + with ``{"parent_collection": atomic_attack_name}`` so dataset-level cache + scoping tests can verify the attribution filter in + ``_collect_cached_completion_pairs``. + """ attack_result = AttackResult( conversation_id=f"conv-{outcome.value}-{datetime.now(timezone.utc).timestamp()}", objective=objective, atomic_attack_identifier=_make_atomic_attack_identifier(target), outcome=outcome, timestamp=datetime.now(timezone.utc), + attribution_data={"parent_collection": atomic_attack_name} if atomic_attack_name else None, ) memory.add_attack_results_to_memory(attack_results=[attack_result]) return attack_result @@ -914,9 +944,10 @@ def _make_bench_with_real_memory( return bench -def _make_candidate(*, technique_eval_hash: str) -> MagicMock: +def _make_candidate(*, technique_eval_hash: str, atomic_attack_name: str = "attack_a") -> MagicMock: candidate = MagicMock() candidate.technique_eval_hash = technique_eval_hash + candidate.atomic_attack_name = atomic_attack_name return candidate @@ -935,27 +966,27 @@ def test_cold_cache_returns_empty(self, sqlite_instance): def test_returns_hash_for_success_match_in_real_db(self, sqlite_instance): target = _make_objective_target_component() - _persist_attack_result(sqlite_instance, target, outcome=AttackOutcome.SUCCESS) + _persist_attack_result(sqlite_instance, target, outcome=AttackOutcome.SUCCESS, atomic_attack_name="attack_a") bench = _make_bench_with_real_memory(sqlite_instance, target) tech_hash = _technique_eval_hash_for(target) - candidate = _make_candidate(technique_eval_hash=tech_hash) + candidate = _make_candidate(technique_eval_hash=tech_hash, atomic_attack_name="attack_a") result = bench._collect_cached_completion_pairs(atomic_attacks=[candidate]) - assert result == {tech_hash} + assert result == {"attack_a"} def test_returns_hash_for_failure_match_in_real_db(self, sqlite_instance): target = _make_objective_target_component() - _persist_attack_result(sqlite_instance, target, outcome=AttackOutcome.FAILURE) + _persist_attack_result(sqlite_instance, target, outcome=AttackOutcome.FAILURE, atomic_attack_name="attack_a") bench = _make_bench_with_real_memory(sqlite_instance, target) tech_hash = _technique_eval_hash_for(target) - candidate = _make_candidate(technique_eval_hash=tech_hash) + candidate = _make_candidate(technique_eval_hash=tech_hash, atomic_attack_name="attack_a") result = bench._collect_cached_completion_pairs(atomic_attacks=[candidate]) - assert result == {tech_hash} + assert result == {"attack_a"} def test_filters_out_persisted_results_with_different_objective_target(self, sqlite_instance): """A row with a matching technique hash but a different target hash is rejected.""" @@ -1012,25 +1043,70 @@ def test_filters_out_error_only_history(self, sqlite_instance): assert result == set() def test_dedupes_candidates_with_same_technique_hash(self, sqlite_instance): - """Two candidates sharing a technique hash collapse to a single set entry.""" + """Two candidates sharing a technique hash are evaluated independently by name.""" target = _make_objective_target_component() - _persist_attack_result(sqlite_instance, target, outcome=AttackOutcome.SUCCESS) + _persist_attack_result(sqlite_instance, target, outcome=AttackOutcome.SUCCESS, atomic_attack_name="attack_a") + _persist_attack_result(sqlite_instance, target, outcome=AttackOutcome.SUCCESS, atomic_attack_name="attack_b") bench = _make_bench_with_real_memory(sqlite_instance, target) tech_hash = _technique_eval_hash_for(target) candidates = [ - _make_candidate(technique_eval_hash=tech_hash), - _make_candidate(technique_eval_hash=tech_hash), + _make_candidate(technique_eval_hash=tech_hash, atomic_attack_name="attack_a"), + _make_candidate(technique_eval_hash=tech_hash, atomic_attack_name="attack_b"), ] result = bench._collect_cached_completion_pairs(atomic_attacks=candidates) - assert result == {tech_hash} + assert result == {"attack_a", "attack_b"} + def test_same_technique_hash_only_harmbench_cached_when_only_harmbench_persisted(self, sqlite_instance): + """Dataset-level scoping: same technique+target hash, only harmbench records in DB. -# --------------------------------------------------------------------------- -# run_async cache injection -# --------------------------------------------------------------------------- + Both harmbench and advbench candidates share a technique_eval_hash (same technique, + same model target). Only harmbench results were persisted. The advbench slot must + NOT be marked as cached — it should be re-run on the next execution. + """ + target = _make_objective_target_component() + harmbench_name = "red_teaming__adv_a_harmbench" + advbench_name = "red_teaming__adv_a_advbench" + + _persist_attack_result( + sqlite_instance, target, outcome=AttackOutcome.SUCCESS, atomic_attack_name=harmbench_name + ) + + bench = _make_bench_with_real_memory(sqlite_instance, target) + tech_hash = _technique_eval_hash_for(target) + candidates = [ + _make_candidate(technique_eval_hash=tech_hash, atomic_attack_name=harmbench_name), + _make_candidate(technique_eval_hash=tech_hash, atomic_attack_name=advbench_name), + ] + + result = bench._collect_cached_completion_pairs(atomic_attacks=candidates) + + assert result == {harmbench_name} + assert advbench_name not in result + + def test_same_technique_hash_both_datasets_cached_when_both_persisted(self, sqlite_instance): + """Dataset-level scoping: same technique+target, both datasets have prior results → both skipped.""" + target = _make_objective_target_component() + harmbench_name = "red_teaming__adv_a_harmbench" + advbench_name = "red_teaming__adv_a_advbench" + + _persist_attack_result( + sqlite_instance, target, outcome=AttackOutcome.SUCCESS, atomic_attack_name=harmbench_name + ) + _persist_attack_result(sqlite_instance, target, outcome=AttackOutcome.FAILURE, atomic_attack_name=advbench_name) + + bench = _make_bench_with_real_memory(sqlite_instance, target) + tech_hash = _technique_eval_hash_for(target) + candidates = [ + _make_candidate(technique_eval_hash=tech_hash, atomic_attack_name=harmbench_name), + _make_candidate(technique_eval_hash=tech_hash, atomic_attack_name=advbench_name), + ] + + result = bench._collect_cached_completion_pairs(atomic_attacks=candidates) + + assert result == {harmbench_name, advbench_name} @pytest.mark.usefixtures("patch_central_database") From b0dc4b09c0cf7ba0e48e5ad54760fd9db326f5b3 Mon Sep 17 00:00:00 2001 From: Victor Valbuena Date: Mon, 1 Jun 2026 13:02:06 -0700 Subject: [PATCH 40/40] fix: use uuid4 for conversation_id in _persist_attack_result to prevent CI collision datetime.now().timestamp() can return the same float for two rapid calls in CI, causing _dedup_attack_entries to discard one result (dedup is by conversation_id). Use uuid4 to guarantee uniqueness. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- tests/unit/scenario/benchmark/test_adversarial.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/tests/unit/scenario/benchmark/test_adversarial.py b/tests/unit/scenario/benchmark/test_adversarial.py index b89b5ddc67..10e823d9d8 100644 --- a/tests/unit/scenario/benchmark/test_adversarial.py +++ b/tests/unit/scenario/benchmark/test_adversarial.py @@ -30,6 +30,7 @@ persistence -> SQL filter -> objective-target filter -> outcome filter. """ +import uuid from datetime import datetime, timezone from unittest.mock import AsyncMock, MagicMock, patch @@ -916,7 +917,7 @@ def _persist_attack_result( ``_collect_cached_completion_pairs``. """ attack_result = AttackResult( - conversation_id=f"conv-{outcome.value}-{datetime.now(timezone.utc).timestamp()}", + conversation_id=str(uuid.uuid4()), objective=objective, atomic_attack_identifier=_make_atomic_attack_identifier(target), outcome=outcome,