diff --git a/.pyrit_conf_example b/.pyrit_conf_example index ed45b30ce..e136b7a91 100644 --- a/.pyrit_conf_example +++ b/.pyrit_conf_example @@ -24,8 +24,8 @@ memory_db_type: sqlite # Available initializers: # - simple: Basic OpenAI configuration (requires OPENAI_CHAT_* env vars) # - airt: AI Red Team setup with Azure OpenAI (requires AZURE_OPENAI_* env vars) -# - targets: Registers available prompt targets into the TargetRegistry -# - scorers: Registers pre-configured scorers into the ScorerRegistry +# - target: Registers available prompt targets into the TargetRegistry +# - scorer: Registers pre-configured scorers into the ScorerRegistry # - load_default_datasets: Loads default datasets for all registered scenarios # - objective_list: Sets default objectives for scenarios # @@ -39,7 +39,7 @@ memory_db_type: sqlite # Example: # initializers: # - simple -# - name: targets +# - name: target # args: # tags: # - default @@ -47,8 +47,8 @@ memory_db_type: sqlite initializers: - name: simple - name: load_default_datasets - - name: scorers - - name: targets + - name: scorer + - name: target args: tags: - default diff --git a/build_scripts/env_local_integration_test b/build_scripts/env_local_integration_test index 8d709f4aa..d9873e435 100644 --- a/build_scripts/env_local_integration_test +++ b/build_scripts/env_local_integration_test @@ -17,7 +17,7 @@ OPENAI_TTS_KEY=${OPENAI_TTS_KEY2} AZURE_SQL_DB_CONNECTION_STRING=${AZURE_SQL_DB_CONNECTION_STRING_TEST} AZURE_STORAGE_ACCOUNT_DB_DATA_CONTAINER_URL=${AZURE_STORAGE_ACCOUNT_DB_DATA_CONTAINER_URL_TEST} -# E2E scenario test variables (used by openai_objective_target initializer) +# E2E scenario test variables (used by target initializer) DEFAULT_OPENAI_FRONTEND_ENDPOINT=${AZURE_OPENAI_INTEGRATION_TEST_ENDPOINT} DEFAULT_OPENAI_FRONTEND_KEY=${AZURE_OPENAI_INTEGRATION_TEST_KEY} DEFAULT_OPENAI_FRONTEND_MODEL=${AZURE_OPENAI_INTEGRATION_TEST_MODEL} diff --git a/doc/code/front_end/1_pyrit_scan.ipynb b/doc/code/front_end/1_pyrit_scan.ipynb index e4df18870..2a4d6af9c 100644 --- a/doc/code/front_end/1_pyrit_scan.ipynb +++ b/doc/code/front_end/1_pyrit_scan.ipynb @@ -27,59 +27,65 @@ "output_type": "stream", "text": [ "Starting PyRIT...\n", - "usage: pyrit_scan [-h] [--log-level LOG_LEVEL] [--list-scenarios]\n", - " [--list-initializers] [--database DATABASE]\n", + "usage: pyrit_scan [-h] [--config-file CONFIG_FILE] [--log-level LOG_LEVEL]\n", + " [--list-scenarios] [--list-initializers] [--list-targets]\n", " [--initializers INITIALIZERS [INITIALIZERS ...]]\n", " [--initialization-scripts INITIALIZATION_SCRIPTS [INITIALIZATION_SCRIPTS ...]]\n", - " [--env-files ENV_FILES [ENV_FILES ...]]\n", " [--strategies SCENARIO_STRATEGIES [SCENARIO_STRATEGIES ...]]\n", " [--max-concurrency MAX_CONCURRENCY]\n", " [--max-retries MAX_RETRIES] [--memory-labels MEMORY_LABELS]\n", " [--dataset-names DATASET_NAMES [DATASET_NAMES ...]]\n", - " [--max-dataset-size MAX_DATASET_SIZE]\n", + " [--max-dataset-size MAX_DATASET_SIZE] [--target TARGET]\n", " [scenario_name]\n", "\n", "PyRIT Scanner - Run security scenarios against AI systems\n", "\n", "Examples:\n", - " # List available scenarios and initializers\n", + " # List available scenarios, initializers, and targets\n", " pyrit_scan --list-scenarios\n", " pyrit_scan --list-initializers\n", + " pyrit_scan --list-targets --initializers target\n", "\n", - " # Run a scenario with built-in initializers\n", - " pyrit_scan foundry --initializers openai_objective_target load_default_datasets\n", + " # Run a scenario with a target and initializers\n", + " pyrit_scan foundry.red_team_agent --target my_target --initializers target load_default_datasets\n", + "\n", + " # Run with a configuration file (recommended for complex setups)\n", + " pyrit_scan foundry.red_team_agent --target my_target --config-file ./my_config.yaml\n", "\n", " # Run with custom initialization scripts\n", - " pyrit_scan garak.encoding --initialization-scripts ./my_config.py\n", + " pyrit_scan garak.encoding --target my_target --initialization-scripts ./my_config.py\n", "\n", " # Run specific strategies or options\n", - " pyrit scan foundry --strategies base64 rot13 --initializers openai_objective_target\n", - " pyrit_scan foundry --initializers openai_objective_target --max-concurrency 10 --max-retries 3\n", - " pyrit_scan garak.encoding --initializers openai_objective_target --memory-labels '{\"run_id\":\"test123\"}'\n", + " pyrit_scan foundry.red_team_agent --target my_target --strategies base64 rot13 --initializers target\n", + " pyrit_scan foundry.red_team_agent --target my_target --initializers target --max-concurrency 10 --max-retries 3\n", "\n", "positional arguments:\n", " scenario_name Name of the scenario to run\n", "\n", "options:\n", " -h, --help show this help message and exit\n", + " --config-file CONFIG_FILE\n", + " Path to a YAML configuration file. Allows specifying\n", + " database, initializers (with args), initialization\n", + " scripts, and env files. CLI arguments override config\n", + " file values. If not specified, ~/.pyrit/.pyrit_conf is\n", + " loaded if it exists.\n", " --log-level LOG_LEVEL\n", " Logging level (DEBUG, INFO, WARNING, ERROR, CRITICAL)\n", " (default: WARNING)\n", " --list-scenarios List all available scenarios and exit\n", " --list-initializers List all available scenario initializers and exit\n", - " --database DATABASE Database type to use for memory storage (InMemory,\n", - " SQLite, AzureSQL) (default: SQLite)\n", + " --list-targets List all available targets from the TargetRegistry and\n", + " exit. Requires initializers that register targets\n", + " (e.g., --initializers target)\n", " --initializers INITIALIZERS [INITIALIZERS ...]\n", - " Built-in initializer names to run before the scenario\n", - " (e.g., openai_objective_target)\n", + " Built-in initializer names to run before the scenario.\n", + " Supports optional params with name:key=val syntax\n", + " (e.g., target:tags=default,scorer dataset:mode=strict)\n", " --initialization-scripts INITIALIZATION_SCRIPTS [INITIALIZATION_SCRIPTS ...]\n", " Paths to custom Python initialization scripts to run\n", " before the scenario\n", - " --env-files ENV_FILES [ENV_FILES ...]\n", - " Paths to environment files to load in order (e.g.,\n", - " .env.production .env.local). Later files override\n", - " earlier ones.\n", - " --strategies, -s SCENARIO_STRATEGIES [SCENARIO_STRATEGIES ...]\n", + " --strategies SCENARIO_STRATEGIES [SCENARIO_STRATEGIES ...], -s SCENARIO_STRATEGIES [SCENARIO_STRATEGIES ...]\n", " List of strategy names to run (e.g., base64 rot13)\n", " --max-concurrency MAX_CONCURRENCY\n", " Maximum number of concurrent attack executions (must\n", @@ -98,7 +104,12 @@ " --max-dataset-size MAX_DATASET_SIZE\n", " Maximum number of items to use from the dataset (must\n", " be >= 1). Limits new datasets if --dataset-names\n", - " provided, otherwise overrides scenario's default limit\n" + " provided, otherwise overrides scenario's default limit\n", + " --target TARGET Name of a registered target from the TargetRegistry to\n", + " use as the objective target. Targets are registered by\n", + " initializers (e.g., 'target' initializer). Use --list-\n", + " targets to see available target names after\n", + " initializers have run\n" ] } ], @@ -127,9 +138,9 @@ "output_type": "stream", "text": [ "Starting PyRIT...\n", - "Found default environment files: ['./.pyrit/.env', './.pyrit/.env.local']\n", + "Loading default configuration file: ./.pyrit/.pyrit_conf\n", + "Found default environment files: ['./.pyrit/.env']\n", "Loaded environment file: ./.pyrit/.env\n", - "Loaded environment file: ./.pyrit/.env.local\n", "\n", "Available Scenarios:\n", "================================================================================\n", @@ -164,6 +175,66 @@ " Default Datasets (1, max 4 per dataset):\n", " airt_malware\n", "\u001b[1m\u001b[36m\n", + " airt.jailbreak\u001b[0m\n", + " Class: Jailbreak\n", + " Description:\n", + " Jailbreak scenario implementation for PyRIT. This scenario tests how\n", + " vulnerable models are to jailbreak attacks by applying various\n", + " single-turn jailbreak templates to a set of test prompts. The responses\n", + " are scored to determine if the jailbreak was successful.\n", + " Aggregate Strategies:\n", + " - all, simple, complex\n", + " Available Strategies (4):\n", + " prompt_sending, many_shot, skeleton, role_play\n", + " Default Strategy: simple\n", + " Default Datasets (1, max 4 per dataset):\n", + " airt_harms\n", + "\u001b[1m\u001b[36m\n", + " airt.leakage\u001b[0m\n", + " Class: Leakage\n", + " Description:\n", + " Leakage scenario implementation for PyRIT. This scenario tests how\n", + " susceptible models are to leaking training data, PII, intellectual\n", + " property, or other confidential information. The LeakageScenario class\n", + " contains different attack variations designed to extract sensitive\n", + " information from models.\n", + " Aggregate Strategies:\n", + " - all, single_turn, multi_turn, ip, sensitive_data\n", + " Available Strategies (4):\n", + " first_letter, image, role_play, crescendo\n", + " Default Strategy: all\n", + " Default Datasets (1, max 4 per dataset):\n", + " airt_leakage\n", + "\u001b[1m\u001b[36m\n", + " airt.psychosocial\u001b[0m\n", + " Class: Psychosocial\n", + " Description:\n", + " Psychosocial Harms Scenario implementation for PyRIT. This scenario\n", + " contains various psychosocial harm-based checks that evaluate model\n", + " behavior for responses that may harm the well-being or mental health of\n", + " users. Examples include mistreating users in a mental health crisis or\n", + " aiding them in self-harm. All datasets in this scenario share the\n", + " unified 'psychosocial' harm category, but each dataset (e.g.,\n", + " imminent_crisis, dependency) can have unique: - Scoring rubrics - Seed\n", + " objectives The scenario tests whether models: - Recognize mental health\n", + " crises and respond appropriately - Avoid facilitating self-harm or\n", + " harmful behaviors - Provide crisis resources and encourage professional\n", + " help Example usage with custom configurations: # Define custom\n", + " configurations per subharm category custom_configs = {\n", + " \"airt_imminent_crisis\": SubharmConfig(\n", + " crescendo_system_prompt_path=\"path/to/custom_escalation.yaml\",\n", + " scoring_rubric_path=\"path/to/custom_rubric.yaml\", ), } scenario =\n", + " Psychosocial(subharm_configs=custom_configs) await\n", + " scenario.initialize_async( objective_target=target_llm,\n", + " scenario_strategies=[PsychosocialStrategy.ImminentCrisis], )\n", + " Aggregate Strategies:\n", + " - all\n", + " Available Strategies (2):\n", + " imminent_crisis, licensed_therapist\n", + " Default Strategy: all\n", + " Default Datasets (1, max 4 per dataset):\n", + " airt_imminent_crisis\n", + "\u001b[1m\u001b[36m\n", " airt.scam\u001b[0m\n", " Class: Scam\n", " Description:\n", @@ -227,7 +298,7 @@ "\n", "================================================================================\n", "\n", - "Total scenarios: 5\n" + "Total scenarios: 8\n" ] } ], @@ -252,7 +323,7 @@ "\n", "PyRITInitializers are how you can configure the CLI scanner. PyRIT includes several built-in initializers you can use with the `--initializers` flag.\n", "\n", - "The `--list-initializers` command shows all available initializers. Initializers are referenced by their filename (e.g., `objective_target`, `objective_list`, `simple`) regardless of which subdirectory they're in.\n", + "The `--list-initializers` command shows all available initializers. Initializers are referenced by their filename (e.g., `target`, `objective_list`, `simple`) regardless of which subdirectory they're in.\n", "\n", "List the available initializers using the --list-initializers flag." ] @@ -268,10 +339,27 @@ "output_type": "stream", "text": [ "Starting PyRIT...\n", + "Loading default configuration file: ./.pyrit/.pyrit_conf\n", + "Found default environment files: ['./.pyrit/.env']\n", + "Loaded environment file: ./.pyrit/.env\n", "\n", "Available Initializers:\n", "================================================================================\n", "\u001b[1m\u001b[36m\n", + " airt\u001b[0m\n", + " Class: AIRTInitializer\n", + " Name: AIRT Default Configuration\n", + " Execution Order: 1\n", + " Required Environment Variables:\n", + " - AZURE_OPENAI_GPT4O_UNSAFE_CHAT_ENDPOINT\n", + " - AZURE_OPENAI_GPT4O_UNSAFE_CHAT_MODEL\n", + " - AZURE_OPENAI_GPT4O_UNSAFE_CHAT_ENDPOINT2\n", + " - AZURE_OPENAI_GPT4O_UNSAFE_CHAT_MODEL2\n", + " - AZURE_CONTENT_SAFETY_API_ENDPOINT\n", + " Description:\n", + " AI Red Team setup with Azure OpenAI converters, composite harm/objective\n", + " scorers, and adversarial targets\n", + "\u001b[1m\u001b[36m\n", " load_default_datasets\u001b[0m\n", " Class: LoadDefaultDatasets\n", " Name: Default Dataset Loader for Scenarios\n", @@ -283,7 +371,7 @@ " customized in memory. Note: if you are using persistent memory, avoid\n", " calling this every time as datasets can take time to load.\n", "\u001b[1m\u001b[36m\n", - " objective_list\u001b[0m\n", + " scenario_objective_list\u001b[0m\n", " Class: ScenarioObjectiveListInitializer\n", " Name: Simple Objective List Configuration for Scenarios\n", " Execution Order: 10\n", @@ -291,22 +379,43 @@ " Description:\n", " Simple Objective List Configuration for Scenarios\n", "\u001b[1m\u001b[36m\n", - " openai_objective_target\u001b[0m\n", - " Class: ScenarioObjectiveTargetInitializer\n", - " Name: Simple Objective Target Configuration for Scenarios\n", - " Execution Order: 10\n", + " scorer\u001b[0m\n", + " Class: ScorerInitializer\n", + " Name: Scorer Initializer\n", + " Execution Order: 2\n", + " Required Environment Variables: None\n", + " Supported Parameters:\n", + " - tags [default: ['default']]: Tags for filtering (e.g., ['default'])\n", + " Description:\n", + " Instantiates a collection of scorers using targets from the\n", + " TargetRegistry and adds them to the ScorerRegistry\n", + "\u001b[1m\u001b[36m\n", + " simple\u001b[0m\n", + " Class: SimpleInitializer\n", + " Name: Simple Complete Configuration\n", + " Execution Order: 1\n", " Required Environment Variables:\n", - " - DEFAULT_OPENAI_FRONTEND_ENDPOINT\n", - " - DEFAULT_OPENAI_FRONTEND_KEY\n", + " - OPENAI_CHAT_ENDPOINT\n", + " - OPENAI_CHAT_MODEL\n", + " Description:\n", + " Complete simple setup with basic OpenAI converters, objective scorer (no\n", + " harm detection), and adversarial targets. Only requires\n", + " OPENAI_CHAT_ENDPOINT and OPENAI_CHAT_MODEL environment variables.\n", + "\u001b[1m\u001b[36m\n", + " target\u001b[0m\n", + " Class: TargetInitializer\n", + " Name: Target Initializer\n", + " Execution Order: 1\n", + " Required Environment Variables: None\n", + " Supported Parameters:\n", + " - tags [default: ['default']]: Target tags to register (e.g., ['default'], ['default', 'scorer'], or ['all'])\n", " Description:\n", - " This configuration sets up a simple objective target for scenarios using\n", - " OpenAIChatTarget with basic settings. It initializes an openAI chat\n", - " target using the OPENAI_CLI_ENDPOINT and OPENAI_CLI_KEY environment\n", - " variables.\n", + " Instantiates a collection of targets from available environment\n", + " variables and adds them to the TargetRegistry\n", "\n", "================================================================================\n", "\n", - "Total initializers: 3\n" + "Total initializers: 6\n" ] } ], @@ -324,13 +433,13 @@ "You need a single scenario to run, you need two things:\n", "\n", "1. A Scenario. Many are defined in `pyrit.scenario.scenarios`. But you can also define your own in initialization_scripts.\n", - "2. Initializers (which can be supplied via `--initializers` or `--initialization-scripts`). Scenarios often don't need many arguments, but they can be configured in different ways. And at the very least, most need an `objective_target` (the thing you're running a scan against).\n", + "2. Initializers (which can be supplied via `--initializers` or `--initialization-scripts` or `initializers` section of config file (see [here](../../getting_started/pyrit_conf.md))). Scenarios often don't need many arguments, but they can be configured in different ways. And at the very least, most need an `objective_target` (the thing you're running a scan against) which you can configure by using the `--target` flag if your initializer registers targets (e.g. `target` initializer)\n", "3. Scenario Strategies (optional). These are supplied by the `--scenario-strategies` flag and tell the scenario what to test, but they are always optional. Also note you can obtain these by running `--list-scenarios`\n", "\n", "Basic usage will look something like:\n", "\n", "```shell\n", - "pyrit_scan --initializers --scenario-strategies \n", + "pyrit_scan --target --initializers --scenario-strategies \n", "```\n", "\n", "You can also override scenario parameters directly from the CLI:\n", @@ -342,10 +451,10 @@ "Or concretely:\n", "\n", "```shell\n", - "!pyrit_scan foundry.red_team_agent --initializers simple openai_objective_target --scenario-strategies base64\n", + "!pyrit_scan foundry.red_team_agent --target openai_chat --initializers load_default_datasets target --scenario-strategies base64\n", "```\n", "\n", - "Example with a basic configuration that runs the Foundry scenario against the objective target defined in `openai_objective_target` (which just is an OpenAIChatTarget with `DEFAULT_OPENAI_FRONTEND_ENDPOINT` and `DEFAULT_OPENAI_FRONTEND_KEY`)." + "Example with a basic configuration that runs the Foundry scenario against the objective target defined in the `target` initializer." ] }, { @@ -359,35 +468,100 @@ "output_type": "stream", "text": [ "Starting PyRIT...\n", - "Found default environment files: ['./.pyrit/.env', './.pyrit/.env.local']\n", + "Loading default configuration file: ./.pyrit/.pyrit_conf\n", + "Found default environment files: ['./.pyrit/.env']\n", "Loaded environment file: ./.pyrit/.env\n", - "Loaded environment file: ./.pyrit/.env.local\n", - "Running 1 initializer(s)...\n", - "Found default environment files: ['./.pyrit/.env', './.pyrit/.env.local']\n", - "Loaded environment file: ./.pyrit/.env\n", - "Loaded environment file: ./.pyrit/.env.local\n", + "Running 2 initializer(s)...\n", "\n", "Running scenario: foundry.red_team_agent\n", "\n", "\u001b[36m====================================================================================================\u001b[0m\n", + "\u001b[1m\u001b[36m 📊 SCENARIO RESULTS: RedTeamAgent \u001b[0m\n", + "\u001b[36m====================================================================================================\u001b[0m\n", + "\n", + "\u001b[1m\u001b[36m▼ Scenario Information\u001b[0m\n", + "\u001b[36m────────────────────────────────────────────────────────────────────────────────────────────────────\u001b[0m\n", + "\u001b[1m 📋 Scenario Details\u001b[0m\n", + "\u001b[36m • Name: RedTeamAgent\u001b[0m\n", + "\u001b[36m • Scenario Version: 1\u001b[0m\n", + "\u001b[36m • PyRIT Version: 0.12.1.dev0\u001b[0m\n", + "\u001b[36m • Description:\u001b[0m\n", + "\u001b[36m RedTeamAgent is a preconfigured scenario that automatically generates multiple AtomicAttack instances based on\u001b[0m\n", + "\u001b[36m the specified attack strategies. It supports both single-turn attacks (with various converters) and multi-turn\u001b[0m\n", + "\u001b[36m attacks (Crescendo, RedTeaming), making it easy to quickly test a target against multiple attack vectors. The\u001b[0m\n", + "\u001b[36m scenario can expand difficulty levels (EASY, MODERATE, DIFFICULT) into their constituent attack strategies, or\u001b[0m\n", + "\u001b[36m you can specify individual strategies directly. This scenario is designed for use with the Foundry AI Red\u001b[0m\n", + "\u001b[36m Teaming Agent library, providing a consistent PyRIT contract for their integration.\u001b[0m\n", + "\n", + "\u001b[1m 🎯 Target Information\u001b[0m\n", + "\u001b[36m • Target Type: OpenAIChatTarget\u001b[0m\n", + "\u001b[36m • Target Model: gpt-4o\u001b[0m\n", + "\u001b[36m • Target Endpoint: https://pyrit-japan-test.openai.azure.com/openai/v1\u001b[0m\n", + "\n", + "\u001b[1m 📊 Scorer Information\u001b[0m\n", + "\u001b[37m ▸ Scorer Identifier\u001b[0m\n", + "\u001b[36m • Scorer Type: TrueFalseInverterScorer\u001b[0m\n", + "\u001b[36m • scorer_type: true_false\u001b[0m\n", + "\u001b[36m • score_aggregator: OR_\u001b[0m\n", + "\u001b[36m └─ Composite of 1 scorer(s):\u001b[0m\n", + "\u001b[36m • Scorer Type: SelfAskRefusalScorer\u001b[0m\n", + "\u001b[36m • scorer_type: true_false\u001b[0m\n", + "\u001b[36m • score_aggregator: OR_\u001b[0m\n", + "\u001b[36m • model_name: gpt-4o\u001b[0m\n", "\n", - "Error: 'charmap' codec can't encode character '\\U0001f4ca' in position 43: character maps to \n" + "\u001b[37m ▸ Performance Metrics\u001b[0m\n", + "\u001b[36m • Accuracy: 84.84%\u001b[0m\n", + "\u001b[36m • Accuracy Std Error: ±0.0185\u001b[0m\n", + "\u001b[36m • F1 Score: 0.8606\u001b[0m\n", + "\u001b[36m • Precision: 0.7928\u001b[0m\n", + "\u001b[32m • Recall: 0.9412\u001b[0m\n", + "\u001b[36m • Average Score Time: 1.27s\u001b[0m\n", + "\n", + "\u001b[1m\u001b[36m▼ Overall Statistics\u001b[0m\n", + "\u001b[36m────────────────────────────────────────────────────────────────────────────────────────────────────\u001b[0m\n", + "\u001b[1m 📈 Summary\u001b[0m\n", + "\u001b[32m • Total Strategies: 2\u001b[0m\n", + "\u001b[32m • Total Attack Results: 8\u001b[0m\n", + "\u001b[32m • Overall Success Rate: 12%\u001b[0m\n", + "\u001b[32m • Unique Objectives: 8\u001b[0m\n", + "\n", + "\u001b[1m\u001b[36m▼ Per-Strategy Breakdown\u001b[0m\n", + "\u001b[36m────────────────────────────────────────────────────────────────────────────────────────────────────\u001b[0m\n", + "\n", + "\u001b[1m 🔸 Strategy: baseline\u001b[0m\n", + "\u001b[33m • Number of Results: 4\u001b[0m\n", + "\u001b[32m • Success Rate: 0%\u001b[0m\n", + "\n", + "\u001b[1m 🔸 Strategy: base64\u001b[0m\n", + "\u001b[33m • Number of Results: 4\u001b[0m\n", + "\u001b[36m • Success Rate: 25%\u001b[0m\n", + "\n", + "\u001b[36m====================================================================================================\u001b[0m\n", + "\n" ] }, { "name": "stderr", "output_type": "stream", "text": [ + "\n", + "Loading datasets - this can take a few minutes: 0%| | 0/58 [00:00\n" ] } ], "source": [ - "!pyrit_scan foundry.red_team_agent --initializers openai_objective_target --strategies base64" + "!pyrit_scan foundry.red_team_agent --target openai_chat --initializers load_default_datasets target --strategies base64" ] }, { @@ -398,17 +572,17 @@ "Or with all options and multiple initializers and multiple strategies:\n", "\n", "```shell\n", - "pyrit_scan foundry.red_team_agent --database InMemory --initializers simple objective_target objective_list --scenario-strategies easy crescendo\n", + "pyrit_scan foundry.red_team_agent --target openai_chat --initializers load_default_datasets target --strategies easy crescendo\n", "```\n", "\n", "You can also override scenario execution parameters:\n", "\n", "```shell\n", "# Override concurrency and retry settings\n", - "pyrit_scan foundry.red_team_agent --initializers simple objective_target --max-concurrency 10 --max-retries 3\n", + "pyrit_scan foundry.red_team_agent --target openai_chat --initializers load_default_datasets target --max-concurrency 10 --max-retries 3\n", "\n", "# Add custom memory labels for tracking (must be valid JSON)\n", - "pyrit_scan foundry.red_team_agent --initializers simple objective_target --memory-labels '{\"experiment\": \"test1\", \"version\": \"v2\", \"researcher\": \"alice\"}'\n", + "pyrit_scan foundry.red_team_agent --target openai_chat --initializers load_default_datasets target --memory-labels '{\"experiment\": \"test1\", \"version\": \"v2\", \"researcher\": \"alice\"}'\n", "```\n", "\n", "Available CLI parameter overrides:\n", @@ -437,21 +611,22 @@ "cell_type": "code", "execution_count": null, "id": "10", - "metadata": {}, + "metadata": { + "lines_to_next_cell": 2 + }, "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" + "Found default environment files: ['./.pyrit/.env']\n", + "Loaded environment file: ./.pyrit/.env\n" ] }, { "data": { "text/plain": [ - "<__main__.MyCustomScenario at 0x13c63b4c2f0>" + "<__main__.MyCustomScenario at 0x1ec016f9c90>" ] }, "execution_count": null, @@ -545,7 +720,7 @@ "name": "python", "nbconvert_exporter": "python", "pygments_lexer": "ipython3", - "version": "3.13.5" + "version": "3.11.9" } }, "nbformat": 4, diff --git a/doc/code/front_end/1_pyrit_scan.py b/doc/code/front_end/1_pyrit_scan.py index 3e325bc5b..6e632f9ad 100644 --- a/doc/code/front_end/1_pyrit_scan.py +++ b/doc/code/front_end/1_pyrit_scan.py @@ -5,11 +5,7 @@ # extension: .py # format_name: percent # format_version: '1.3' -# jupytext_version: 1.17.3 -# kernelspec: -# display_name: pyrit-dev -# language: python -# name: python3 +# jupytext_version: 1.19.1 # --- # %% [markdown] @@ -47,7 +43,7 @@ # # PyRITInitializers are how you can configure the CLI scanner. PyRIT includes several built-in initializers you can use with the `--initializers` flag. # -# The `--list-initializers` command shows all available initializers. Initializers are referenced by their filename (e.g., `objective_target`, `objective_list`, `simple`) regardless of which subdirectory they're in. +# The `--list-initializers` command shows all available initializers. Initializers are referenced by their filename (e.g., `target`, `objective_list`, `simple`) regardless of which subdirectory they're in. # # List the available initializers using the --list-initializers flag. @@ -60,13 +56,13 @@ # You need a single scenario to run, you need two things: # # 1. A Scenario. Many are defined in `pyrit.scenario.scenarios`. But you can also define your own in initialization_scripts. -# 2. Initializers (which can be supplied via `--initializers` or `--initialization-scripts`). Scenarios often don't need many arguments, but they can be configured in different ways. And at the very least, most need an `objective_target` (the thing you're running a scan against). +# 2. Initializers (which can be supplied via `--initializers` or `--initialization-scripts` or `initializers` section of config file (see [here](../../getting_started/pyrit_conf.md))). Scenarios often don't need many arguments, but they can be configured in different ways. And at the very least, most need an `objective_target` (the thing you're running a scan against) which you can configure by using the `--target` flag if your initializer registers targets (e.g. `target` initializer) # 3. Scenario Strategies (optional). These are supplied by the `--scenario-strategies` flag and tell the scenario what to test, but they are always optional. Also note you can obtain these by running `--list-scenarios` # # Basic usage will look something like: # # ```shell -# pyrit_scan --initializers --scenario-strategies +# pyrit_scan --target --initializers --scenario-strategies # ``` # # You can also override scenario parameters directly from the CLI: @@ -78,29 +74,29 @@ # Or concretely: # # ```shell -# !pyrit_scan foundry.red_team_agent --initializers simple openai_objective_target --scenario-strategies base64 +# !pyrit_scan foundry.red_team_agent --target openai_chat --initializers load_default_datasets target --scenario-strategies base64 # ``` # -# Example with a basic configuration that runs the Foundry scenario against the objective target defined in `openai_objective_target` (which just is an OpenAIChatTarget with `DEFAULT_OPENAI_FRONTEND_ENDPOINT` and `DEFAULT_OPENAI_FRONTEND_KEY`). +# Example with a basic configuration that runs the Foundry scenario against the objective target defined in the `target` initializer. # %% -# !pyrit_scan foundry.red_team_agent --initializers openai_objective_target --strategies base64 +# !pyrit_scan foundry.red_team_agent --target openai_chat --initializers load_default_datasets target --strategies base64 # %% [markdown] # Or with all options and multiple initializers and multiple strategies: # # ```shell -# pyrit_scan foundry.red_team_agent --database InMemory --initializers simple objective_target objective_list --scenario-strategies easy crescendo +# pyrit_scan foundry.red_team_agent --target openai_chat --initializers load_default_datasets target --strategies easy crescendo # ``` # # You can also override scenario execution parameters: # # ```shell # # Override concurrency and retry settings -# pyrit_scan foundry.red_team_agent --initializers simple objective_target --max-concurrency 10 --max-retries 3 +# pyrit_scan foundry.red_team_agent --target openai_chat --initializers load_default_datasets target --max-concurrency 10 --max-retries 3 # # # Add custom memory labels for tracking (must be valid JSON) -# pyrit_scan foundry.red_team_agent --initializers simple objective_target --memory-labels '{"experiment": "test1", "version": "v2", "researcher": "alice"}' +# pyrit_scan foundry.red_team_agent --target openai_chat --initializers load_default_datasets target --memory-labels '{"experiment": "test1", "version": "v2", "researcher": "alice"}' # ``` # # Available CLI parameter overrides: @@ -175,6 +171,7 @@ async def _get_atomic_attacks_async(self): await initialize_pyrit_async(memory_db_type="InMemory") # type: ignore MyCustomScenario() + # %% [markdown] # Then discover and run it: # diff --git a/doc/code/front_end/2_pyrit_shell.md b/doc/code/front_end/2_pyrit_shell.md index b7c22e231..3c9396bb7 100644 --- a/doc/code/front_end/2_pyrit_shell.md +++ b/doc/code/front_end/2_pyrit_shell.md @@ -13,14 +13,15 @@ pyrit_shell With startup options: ```bash -# Set default database for all runs -pyrit_shell --database InMemory +# Load configuration file (if not provided, defaults to ~/.pyrit/.pyrit_conf if it exists) +# to set database preference, initializers, labels, env_file, and more. +pyrit_shell --config-file ./.pyrit_conf # Set default log level pyrit_shell --log-level DEBUG # Load initializers at startup -pyrit_shell --initializers openai_objective_target load_default_datasets +pyrit_shell --initializers load_default_datasets # Load custom initialization scripts pyrit_shell --initialization-scripts ./my_config.py @@ -28,12 +29,13 @@ pyrit_shell --initialization-scripts ./my_config.py ## Available Commands -Once in the shell, you have access to: +Once starting the shell, you will see the list of commands you have access to. Some of them are shown below: | Command | Description | |---------|-------------| | `list-scenarios` | List all available scenarios | | `list-initializers` | List all available initializers | +| `list-targets` | List all available targets from the registry | | `run [options]` | Run a scenario with optional parameters | | `scenario-history` | List all previous scenario runs in this session | | `print-scenario [N]` | Print detailed results for scenario run(s) | @@ -48,32 +50,32 @@ The `run` command executes scenarios with the same options as `pyrit_scan`: ### Basic Usage ```bash -pyrit> run foundry.red_team_agent --initializers openai_objective_target load_default_datasets +pyrit> run foundry.red_team_agent --target my_target --initializers target load_default_datasets ``` ### With Strategies ```bash -pyrit> run garak.encoding --initializers openai_objective_target load_default_datasets --strategies base64 rot13 +pyrit> run garak.encoding --target my_target --initializers target load_default_datasets --strategies base64 rot13 -pyrit> run foundry.red_team_agent --initializers openai_objective_target load_default_datasets -s jailbreak crescendo +pyrit> run foundry.red_team_agent --target my_target --initializers target load_default_datasets -s jailbreak crescendo ``` ### With Runtime Parameters ```bash # Set concurrency and retries -pyrit> run foundry.red_team_agent --initializers openai_objective_target load_default_datasets --max-concurrency 10 --max-retries 3 +pyrit> run foundry.red_team_agent --target my_target --initializers target load_default_datasets --max-concurrency 10 --max-retries 3 # Add memory labels for tracking -pyrit> run garak.encoding --initializers openai_objective_target load_default_datasets --memory-labels '{"experiment":"test1","version":"v2"}' +pyrit> run garak.encoding --target my_target --initializers target load_default_datasets --memory-labels '{"experiment":"test1","version":"v2"}' ``` ### Override Defaults Per-Run ```bash -# Override database and log level for this run only -pyrit> run garak.encoding --initializers openai_objective_target load_default_datasets --database InMemory --log-level DEBUG +# Override log level for this run only +pyrit> run garak.encoding --target my_target --initializers target load_default_datasets --log-level DEBUG ``` ### Run Command Options @@ -85,7 +87,6 @@ pyrit> run garak.encoding --initializers openai_objective_target load_default_da --max-concurrency Maximum concurrent operations --max-retries Maximum retry attempts --memory-labels JSON string of labels ---database Override default database (InMemory, SQLite, AzureSQL) --log-level Override default log level (DEBUG, INFO, WARNING, ERROR, CRITICAL) ``` @@ -114,9 +115,9 @@ pyrit> scenario-history Scenario Run History: ================================================================================ -1) foundry.red_team_agent --initializers openai_objective_target load_default_datasets --strategies base64 -2) garak.encoding --initializers openai_objective_target load_default_datasets --strategies rot13 -3) foundry.red_team_agent --initializers openai_objective_target load_default_datasets -s jailbreak +1) foundry.red_team_agent --initializers target load_default_datasets --strategies base64 +2) garak.encoding --initializers target load_default_datasets --strategies rot13 +3) foundry.red_team_agent --initializers target load_default_datasets -s jailbreak ================================================================================ Total runs: 3 @@ -130,7 +131,7 @@ The shell excels at interactive testing workflows: ```bash # Start shell with defaults -pyrit_shell --database InMemory --initializers openai_objective_target load_default_datasets +pyrit_shell --initializers target load_default_datasets # Quick exploration pyrit> list-scenarios @@ -161,7 +162,7 @@ pyrit> print-scenario 2 2. **Use short strategy aliases** with `-s`: ```bash - pyrit> run foundry.red_team_agent --initializers openai_objective_target load_default_datasets -s base64 rot13 + pyrit> run foundry.red_team_agent --initializers target load_default_datasets -s base64 rot13 ``` 3. **Review history regularly** to track what you've tested: diff --git a/doc/code/registry/1_class_registry.ipynb b/doc/code/registry/1_class_registry.ipynb index beb2243a9..26d562131 100644 --- a/doc/code/registry/1_class_registry.ipynb +++ b/doc/code/registry/1_class_registry.ipynb @@ -20,14 +20,12 @@ "name": "stdout", "output_type": "stream", "text": [ - "Available scenarios: ['airt.content_harms', 'airt.cyber', 'airt.scam', 'foundry.red_team_agent', 'garak.encoding']...\n", + "Available scenarios: ['airt.content_harms', 'airt.cyber', 'airt.jailbreak', 'airt.leakage', 'airt.psychosocial']...\n", "\n", - "airt.content_harms:\n", - " Class: ContentHarms\n", + "ContentHarms:\n", " Description: Content Harms Scenario implementation for PyRIT. This scenario contains various ...\n", "\n", - "airt.cyber:\n", - " Class: Cyber\n", + "Cyber:\n", " Description: Cyber scenario implementation for PyRIT. This scenario tests how willing models ...\n" ] } @@ -74,8 +72,6 @@ } ], "source": [ - "# Get a scenario class\n", - "\n", "scenario_class = registry.get_class(\"garak.encoding\")\n", "\n", "print(f\"Got class: {scenario_class}\")\n", @@ -102,56 +98,15 @@ "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" - ] - }, - { - "name": "stderr", - "output_type": "stream", - "text": [ - "\r\n", - "Loading datasets - this can take a few minutes: 0%| | 0/45 [00:00 [opts] - Execute a security scenario", " • scenario-history - View your session history", " • print-scenario [N] - Display detailed results", " • help [command] - Get help on any command", + " • clear - Clear the screen", " • exit - Quit the shell", ] cmd_section: list[tuple[str, ColorRole]] = [ @@ -296,7 +298,7 @@ def add(line: str, role: ColorRole, segments: Optional[list[tuple[int, int, Colo quick_start = [ "Quick Start:", " pyrit> list-scenarios", - " pyrit> run foundry --target my_target --initializers targets load_default_datasets", + " pyrit> run foundry.red_team_agent --target my_target --initializers target load_default_datasets", ] for qs in quick_start: full_line = _box_line(" " + qs) diff --git a/pyrit/cli/_cli_args.py b/pyrit/cli/_cli_args.py index 623ce7732..1264956cc 100644 --- a/pyrit/cli/_cli_args.py +++ b/pyrit/cli/_cli_args.py @@ -287,7 +287,7 @@ def parse_memory_labels(json_string: str) -> dict[str, str]: "max_dataset_size": "Maximum number of items to use from the dataset (must be >= 1). " "Limits new datasets if --dataset-names provided, otherwise overrides scenario's default limit", "target": "Name of a registered target from the TargetRegistry to use as the objective target. " - "Targets are registered by initializers (e.g., 'targets' initializer). " + "Targets are registered by initializers (e.g., 'target' initializer). " "Use --list-targets to see available target names after initializers have run", } diff --git a/pyrit/cli/frontend_core.py b/pyrit/cli/frontend_core.py index e70fb4272..3db555201 100644 --- a/pyrit/cli/frontend_core.py +++ b/pyrit/cli/frontend_core.py @@ -17,7 +17,6 @@ import logging import sys -from pathlib import Path from typing import TYPE_CHECKING, Any, Optional from pyrit.cli._cli_args import ARG_HELP as ARG_HELP @@ -63,6 +62,7 @@ def cprint(text: str, color: str = None, attrs: list = None) -> None: # type: i if TYPE_CHECKING: from collections.abc import Sequence + from pathlib import Path from pyrit.models.scenario_result import ScenarioResult from pyrit.registry import ( @@ -234,22 +234,18 @@ async def list_scenarios_async(*, context: FrontendCore) -> list[ScenarioMetadat async def list_initializers_async( - *, context: FrontendCore, discovery_path: Optional[Path] = None + *, + context: FrontendCore, ) -> Sequence[InitializerMetadata]: """ List metadata for all available initializers. Args: context: PyRIT context with loaded registries. - discovery_path: Optional path to discover initializers from. Returns: Sequence of initializer metadata dictionaries describing each initializer class. """ - if discovery_path: - registry = InitializerRegistry(discovery_path=discovery_path) - return registry.list_metadata() - if not context._initialized: await context.initialize_async() return context.initializer_registry.list_metadata() @@ -321,7 +317,7 @@ async def run_scenario_async( scenario_name: Name of the scenario to run. context: PyRIT context with loaded registries. target_name: Name of a registered target from the TargetRegistry to use as the - objective target. Targets are registered by initializers (e.g., the 'targets' + objective target. Targets are registered by initializers (e.g., the 'target' initializer). Use --list-targets to see available names after initializers run. scenario_strategies: Optional list of strategy names. max_concurrency: Max concurrent operations. @@ -384,7 +380,7 @@ async def run_scenario_async( raise ValueError( f"Target '{target_name}' not found. The target registry is empty.\n" "Targets are registered by initializers. Make sure to include an initializer " - "that registers targets (e.g., --initializers targets)." + "that registers targets (e.g., --initializers target)." ) raise ValueError( f"Target '{target_name}' not found in registry.\nAvailable targets: {', '.join(available_names)}" @@ -512,7 +508,7 @@ def format_scenario_metadata(*, scenario_metadata: ScenarioMetadata) -> None: Args: scenario_metadata: Dataclass containing scenario metadata. """ - _print_header(text=scenario_metadata.snake_class_name) + _print_header(text=scenario_metadata.registry_name) print(f" Class: {scenario_metadata.class_name}") description = scenario_metadata.class_description @@ -554,7 +550,7 @@ def format_initializer_metadata(*, initializer_metadata: InitializerMetadata) -> Args: initializer_metadata: Dataclass containing initializer metadata. """ - _print_header(text=initializer_metadata.snake_class_name) + _print_header(text=initializer_metadata.registry_name) print(f" Class: {initializer_metadata.class_name}") print(f" Name: {initializer_metadata.display_name}") print(f" Execution Order: {initializer_metadata.execution_order}") @@ -594,17 +590,6 @@ def resolve_initialization_scripts(script_paths: list[str]) -> list[Path]: return InitializerRegistry.resolve_script_paths(script_paths=script_paths) -def get_default_initializer_discovery_path() -> Path: - """ - Get the default path for discovering initializers. - - Returns: - Path to the scenarios initializers directory. - """ - pyrit_path = Path(__file__).parent.parent.resolve() - return pyrit_path / "setup" / "initializers" / "scenarios" - - async def print_scenarios_list_async(*, context: FrontendCore) -> int: """ Print a formatted list of all available scenarios. @@ -630,18 +615,17 @@ async def print_scenarios_list_async(*, context: FrontendCore) -> int: return 0 -async def print_initializers_list_async(*, context: FrontendCore, discovery_path: Optional[Path] = None) -> int: +async def print_initializers_list_async(*, context: FrontendCore) -> int: """ Print a formatted list of all available initializers. Args: context: PyRIT context with loaded registries. - discovery_path: Optional path to discover initializers from. Returns: Exit code (0 for success). """ - initializers = await list_initializers_async(context=context, discovery_path=discovery_path) + initializers = await list_initializers_async(context=context) if not initializers: print("No initializers found.") @@ -661,7 +645,7 @@ async def print_targets_list_async(*, context: FrontendCore) -> int: Print a formatted list of all available targets from the TargetRegistry. Targets are registered by initializers, so this requires initializers to run first. - If no targets are found, prints a hint about using the 'targets' initializer. + If no targets are found, prints a hint about using the 'target' initializer. Args: context: PyRIT context with loaded registries. @@ -675,7 +659,7 @@ async def print_targets_list_async(*, context: FrontendCore) -> int: print("\nNo targets found in registry.") print( "\nTargets are registered by initializers. Include an initializer that registers " - "targets, for example:\n --initializers targets\n" + "targets, for example:\n --initializers target\n" ) return 0 diff --git a/pyrit/cli/pyrit_backend.py b/pyrit/cli/pyrit_backend.py index bde9ce23b..fa965cd84 100644 --- a/pyrit/cli/pyrit_backend.py +++ b/pyrit/cli/pyrit_backend.py @@ -232,8 +232,7 @@ def main(*, args: Optional[list[str]] = None) -> int: # Handle list-initializers command if parsed_args.list_initializers: context = frontend_core.FrontendCore(config_file=parsed_args.config_file, log_level=parsed_args.log_level) - scenarios_path = frontend_core.get_default_initializer_discovery_path() - return asyncio.run(frontend_core.print_initializers_list_async(context=context, discovery_path=scenarios_path)) + return asyncio.run(frontend_core.print_initializers_list_async(context=context)) # Run the server try: diff --git a/pyrit/cli/pyrit_scan.py b/pyrit/cli/pyrit_scan.py index cba8c1e89..aefdfa5f2 100644 --- a/pyrit/cli/pyrit_scan.py +++ b/pyrit/cli/pyrit_scan.py @@ -32,20 +32,20 @@ def parse_args(args: Optional[list[str]] = None) -> Namespace: # List available scenarios, initializers, and targets pyrit_scan --list-scenarios pyrit_scan --list-initializers - pyrit_scan --list-targets --initializers targets + pyrit_scan --list-targets --initializers target # Run a scenario with a target and initializers - pyrit_scan foundry --target my_target --initializers targets load_default_datasets + pyrit_scan foundry.red_team_agent --target my_target --initializers target load_default_datasets # Run with a configuration file (recommended for complex setups) - pyrit_scan foundry --target my_target --config-file ./my_config.yaml + pyrit_scan foundry.red_team_agent --target my_target --config-file ./my_config.yaml # Run with custom initialization scripts pyrit_scan garak.encoding --target my_target --initialization-scripts ./my_config.py # Run specific strategies or options - pyrit_scan foundry --target my_target --strategies base64 rot13 --initializers targets - pyrit_scan foundry --target my_target --initializers targets --max-concurrency 10 --max-retries 3 + pyrit_scan foundry.red_team_agent --target my_target --strategies base64 rot13 --initializers target + pyrit_scan foundry.red_team_agent --target my_target --initializers target --max-concurrency 10 --max-retries 3 """, formatter_class=RawDescriptionHelpFormatter, ) @@ -79,7 +79,7 @@ def parse_args(args: Optional[list[str]] = None) -> Namespace: "--list-targets", action="store_true", help="List all available targets from the TargetRegistry and exit. " - "Requires initializers that register targets (e.g., --initializers targets)", + "Requires initializers that register targets (e.g., --initializers target)", ) parser.add_argument( @@ -189,14 +189,20 @@ def main(args: Optional[list[str]] = None) -> int: return asyncio.run(frontend_core.print_scenarios_list_async(context=context)) if parsed_args.list_initializers: - # Discover from scenarios directory - scenarios_path = frontend_core.get_default_initializer_discovery_path() + context = frontend_core.FrontendCore( + config_file=parsed_args.config_file, + log_level=parsed_args.log_level, + ) + return asyncio.run(frontend_core.print_initializers_list_async(context=context)) + if parsed_args.list_targets: + # Need initializers to populate target registry context = frontend_core.FrontendCore( config_file=parsed_args.config_file, + initializer_names=parsed_args.initializers, log_level=parsed_args.log_level, ) - return asyncio.run(frontend_core.print_initializers_list_async(context=context, discovery_path=scenarios_path)) + return asyncio.run(frontend_core.print_targets_list_async(context=context)) if parsed_args.list_targets: # Need initializers to populate target registry diff --git a/pyrit/cli/pyrit_shell.py b/pyrit/cli/pyrit_shell.py index 26217facb..f19602bee 100644 --- a/pyrit/cli/pyrit_shell.py +++ b/pyrit/cli/pyrit_shell.py @@ -199,9 +199,7 @@ def do_list_initializers(self, arg: str) -> None: """List all available initializers.""" self._ensure_initialized() try: - # Discover from scenarios directory by default (same as scan) - discovery_path = self._fc.get_default_initializer_discovery_path() - asyncio.run(self._fc.print_initializers_list_async(context=self.context, discovery_path=discovery_path)) + asyncio.run(self._fc.print_initializers_list_async(context=self.context)) except Exception as e: print(f"Error listing initializers: {e}") @@ -231,22 +229,22 @@ def do_run(self, line: str) -> None: --log-level Override default log level (DEBUG, INFO, WARNING, ERROR, CRITICAL) Examples: - run garak.encoding --target my_target --initializers targets \ + run garak.encoding --target my_target --initializers target \ load_default_datasets - run garak.encoding --target my_target --initializers targets \ + run garak.encoding --target my_target --initializers target \ load_default_datasets --strategies base64 rot13 - run foundry --target my_target --initializers targets:tags=default,scorer \ + run foundry.red_team_agent --target my_target --initializers target:tags=default,scorer \ dataset:mode=strict --strategies base64 - run foundry --target my_target --initializers targets \ + run foundry.red_team_agent --target my_target --initializers target \ load_default_datasets --max-concurrency 10 --max-retries 3 - run garak.encoding --target my_target --initializers targets \ + run garak.encoding --target my_target --initializers target \ load_default_datasets \ --memory-labels '{"run_id":"test123","env":"dev"}' - run foundry --target my_target --initializers targets \ + run foundry.red_team_agent --target my_target --initializers target \ load_default_datasets -s jailbreak crescendo - run garak.encoding --target my_target --initializers targets \ + run garak.encoding --target my_target --initializers target \ load_default_datasets --log-level DEBUG - run foundry --target my_target --initialization-scripts ./my_custom_init.py -s all + run foundry.red_team_agent --target my_target --initialization-scripts ./my_custom_init.py -s all Note: --target is required for every run. @@ -274,7 +272,7 @@ def do_run(self, line: str) -> None: " --log-level Override default log level (DEBUG, INFO, WARNING, ERROR, CRITICAL)" ) print("\nExample:") - print(" run foundry --target my_target --initializers targets load_default_datasets") + print(" run foundry.red_team_agent --target my_target --initializers target load_default_datasets") print("\nType 'help run' for more details and examples") return @@ -321,6 +319,8 @@ def do_run(self, line: str) -> None: ) # Store the command and result in history self._scenario_history.append((line, result)) + except KeyboardInterrupt: + print("\n\nScenario interrupted. Returning to shell.") except ValueError as e: print(f"Error: {e}") except Exception as e: @@ -429,20 +429,21 @@ def do_help(self, arg: str) -> None: print("=" * 70) print(" --target (REQUIRED)") print(f" {ARG_HELP['target']}") - print(" Example: run foundry --target my_target --initializers targets load_default_datasets") + print(" Example: run foundry.red_team_agent --target my_target") + print(" --initializers target load_default_datasets") print() print(" --initializers [ ...]") print(f" {ARG_HELP['initializers']}") - print(" Example: run foundry --target my_target --initializers targets load_default_datasets") - print(" With params: run foundry --target my_target --initializers targets:tags=default,scorer") - print( - " Multiple with params: run foundry --target my_target" - " --initializers targets:tags=default,scorer dataset:mode=strict" - ) + print(" Example: run foundry.red_team_agent --target my_target") + print(" --initializers target load_default_datasets") + print(" With params: run foundry.red_team_agent --target my_target") + print(" --initializers target:tags=default,scorer") + print(" Multiple with params: run foundry.red_team_agent --target my_target") + print(" --initializers target:tags=default,scorer dataset:mode=strict") print() print(" --initialization-scripts [ ...] (Alternative to --initializers)") print(f" {ARG_HELP['initialization_scripts']}") - print(" Example: run foundry --initialization-scripts ./my_init.py") + print(" Example: run foundry.red_team_agent --initialization-scripts ./my_init.py") print() print(" --strategies, -s [ ...]") print(f" {ARG_HELP['scenario_strategies']}") @@ -456,7 +457,7 @@ def do_help(self, arg: str) -> None: print() print(" --memory-labels ") print(f" {ARG_HELP['memory_labels']}") - print(' Example: run foundry --memory-labels \'{"env":"test"}\'') + print(' Example: run foundry.red_team_agent --memory-labels \'{"env":"test"}\'') print() print(" --log-level Override (DEBUG, INFO, WARNING, ERROR, CRITICAL)") print() diff --git a/pyrit/memory/sqlite_memory.py b/pyrit/memory/sqlite_memory.py index d59a6571d..7bd05b4f8 100644 --- a/pyrit/memory/sqlite_memory.py +++ b/pyrit/memory/sqlite_memory.py @@ -15,6 +15,7 @@ from sqlalchemy.exc import SQLAlchemyError from sqlalchemy.orm import joinedload, sessionmaker from sqlalchemy.orm.session import Session +from sqlalchemy.pool import StaticPool from sqlalchemy.sql.expression import TextClause from pyrit.common.path import DB_DATA_PATH @@ -84,6 +85,14 @@ def _create_engine(self, *, has_echo: bool) -> Engine: Creates an engine bound to the specified database file. The `has_echo` parameter controls the verbosity of SQL execution logging. + For in-memory databases (``db_path=":memory:"``), a ``StaticPool`` is used so + that a single shared connection backs all threads. SQLAlchemy's default pool + for ``:memory:`` is ``SingletonThreadPool``, which gives each thread its own + connection — and therefore its own *separate* in-memory database. That causes + tables created on one thread (e.g. a background initialisation thread) to be + invisible from another thread (e.g. the main thread), resulting in + "no such table" errors. + Args: has_echo (bool): Flag to enable detailed SQL execution logging. @@ -94,8 +103,17 @@ def _create_engine(self, *, has_echo: bool) -> Engine: SQLAlchemyError: If there's an issue creating the engine. """ try: - # Create the SQLAlchemy engine. - engine = create_engine(f"sqlite:///{self.db_path}", echo=has_echo) + extra_kwargs: dict[str, Any] = {} + + if self.db_path == ":memory:": + # Use StaticPool so every checkout returns the same underlying + # DBAPI connection, keeping all threads on a single in-memory + # database. ``check_same_thread=False`` is required because + # the connection will be shared across threads. + extra_kwargs["poolclass"] = StaticPool + extra_kwargs["connect_args"] = {"check_same_thread": False} + + engine = create_engine(f"sqlite:///{self.db_path}", echo=has_echo, **extra_kwargs) logger.info(f"Engine created successfully for database: {self.db_path}") return engine except SQLAlchemyError as e: diff --git a/pyrit/registry/base.py b/pyrit/registry/base.py index 51bf59229..02a973a86 100644 --- a/pyrit/registry/base.py +++ b/pyrit/registry/base.py @@ -12,8 +12,6 @@ from dataclasses import dataclass from typing import Any, Optional, Protocol, TypeVar, runtime_checkable -from pyrit.identifiers.class_name_utils import class_name_to_snake_case - # Type variable for metadata (invariant for Protocol compatibility) MetadataT = TypeVar("MetadataT") @@ -30,20 +28,14 @@ class ClassRegistryEntry: class_name (str): Python class name (e.g., "ContentHarmsScenario"). class_module (str): Full module path (e.g., "pyrit.scenario.scenarios.content_harms"). class_description (str): Human-readable description, typically from the class docstring. + registry_name (str): The suffix-stripped snake_case key used in the registry + (e.g., "content_harms" for ContentHarmsScenario). """ class_name: str class_module: str class_description: str = "" - - @property - def snake_class_name(self) -> str: - """ - Snake_case version of class_name (e.g., "content_harms_scenario"). - - Used by CLI formatting and as registry display keys. - """ - return class_name_to_snake_case(self.class_name) + registry_name: str = "" @runtime_checkable diff --git a/pyrit/registry/class_registries/initializer_registry.py b/pyrit/registry/class_registries/initializer_registry.py index cea7e1620..4a2332053 100644 --- a/pyrit/registry/class_registries/initializer_registry.py +++ b/pyrit/registry/class_registries/initializer_registry.py @@ -16,6 +16,7 @@ from pathlib import Path from typing import TYPE_CHECKING, Optional +from pyrit.identifiers.class_name_utils import class_name_to_snake_case from pyrit.registry.base import ClassRegistryEntry from pyrit.registry.class_registries.base_class_registry import ( BaseClassRegistry, @@ -93,9 +94,6 @@ def __init__(self, *, discovery_path: Optional[Path] = None, lazy_discovery: boo # At this point _discovery_path is guaranteed to be a Path assert self._discovery_path is not None - # Track file paths for collision detection and resolution - self._initializer_paths: dict[str, Path] = {} - super().__init__(lazy_discovery=lazy_discovery) def _discover(self) -> None: @@ -113,14 +111,12 @@ def _discover(self) -> None: if discovery_path.is_file(): self._process_file(file_path=discovery_path, base_class=PyRITInitializer) else: - for file_stem, file_path, initializer_class in discover_in_directory( + for _file_stem, _file_path, initializer_class in discover_in_directory( directory=discovery_path, base_class=PyRITInitializer, # type: ignore[type-abstract] recursive=True, ): self._register_initializer( - short_name=file_stem, - file_path=file_path, initializer_class=initializer_class, ) @@ -136,16 +132,6 @@ def _process_file(self, *, file_path: Path, base_class: type) -> None: short_name = file_path.stem - # Check for name collision - if short_name in self._initializer_paths: - existing_path = self._initializer_paths[short_name] - logger.error( - f"Initializer name collision: '{short_name}' found in both " - f"'{file_path}' and '{existing_path}'. " - f"Initializer filenames must be unique across all directories." - ) - return - try: spec = importlib.util.spec_from_file_location(f"initializer.{short_name}", file_path) if not spec or not spec.loader: @@ -163,8 +149,6 @@ def _process_file(self, *, file_path: Path, base_class: type) -> None: and not inspect.isabstract(attr) ): self._register_initializer( - short_name=short_name, - file_path=file_path, initializer_class=attr, # type: ignore[arg-type] ) @@ -174,32 +158,30 @@ def _process_file(self, *, file_path: Path, base_class: type) -> None: def _register_initializer( self, *, - short_name: str, - file_path: Path, initializer_class: type[PyRITInitializer], ) -> None: """ Register an initializer class. Args: - short_name: The short name for the initializer (filename without extension). - file_path: The path to the file containing the initializer. initializer_class: The initializer class to register. """ - # Check for name collision - if short_name in self._initializer_paths: - existing_path = self._initializer_paths[short_name] - logger.error( - f"Initializer name collision: '{short_name}' found in both '{file_path}' and '{existing_path}'." - ) - return - try: - # Create the entry + # Convert class name to snake_case for registry name + registry_name = class_name_to_snake_case(initializer_class.__name__, suffix="Initializer") + + # Check for registry key collision + if registry_name in self._class_entries: + logger.warning( + f"Initializer registry name collision: '{registry_name}' " + f"conflicts with an already-registered initializer. Original " + f"initializer is kept: {self._class_entries[registry_name].registered_class.__name__}" + ) + return + entry = ClassEntry(registered_class=initializer_class) - self._class_entries[short_name] = entry - self._initializer_paths[short_name] = file_path - logger.debug(f"Registered initializer: {short_name} ({initializer_class.__name__})") + self._class_entries[registry_name] = entry + logger.debug(f"Registered initializer: {registry_name} ({initializer_class.__name__})") except Exception as e: logger.warning(f"Failed to register initializer {initializer_class.__name__}: {e}") @@ -223,6 +205,7 @@ def _build_metadata(self, name: str, entry: ClassEntry[PyRITInitializer]) -> Ini class_name=initializer_class.__name__, class_module=initializer_class.__module__, class_description=instance.description, + registry_name=name, display_name=instance.name, required_env_vars=tuple(instance.required_env_vars), execution_order=instance.execution_order, @@ -236,44 +219,12 @@ def _build_metadata(self, name: str, entry: ClassEntry[PyRITInitializer]) -> Ini class_name=initializer_class.__name__, class_module=initializer_class.__module__, class_description="Error loading initializer metadata", + registry_name=name, display_name=name, required_env_vars=(), execution_order=100, ) - def resolve_initializer_paths(self, *, initializer_names: list[str]) -> list[Path]: - """ - Resolve initializer names to their file paths. - - Args: - initializer_names: List of initializer names to resolve. - - Returns: - List of resolved file paths. - - Raises: - ValueError: If any initializer name is not found or has no file path. - """ - self._ensure_discovered() - resolved_paths = [] - - for initializer_name in initializer_names: - if initializer_name not in self._class_entries: - available = ", ".join(sorted(self.get_names())) - raise ValueError( - f"Built-in initializer '{initializer_name}' not found.\n" - f"Available initializers: {available}\n" - f"Use 'pyrit_scan --list-initializers' to see detailed information." - ) - - initializer_file = self._initializer_paths.get(initializer_name) - if initializer_file is None: - raise ValueError(f"Could not locate file for initializer '{initializer_name}'.") - - resolved_paths.append(initializer_file) - - return resolved_paths - @staticmethod def resolve_script_paths(*, script_paths: list[str]) -> list[Path]: """ diff --git a/pyrit/registry/class_registries/scenario_registry.py b/pyrit/registry/class_registries/scenario_registry.py index 8d89b8036..6f6d949ad 100644 --- a/pyrit/registry/class_registries/scenario_registry.py +++ b/pyrit/registry/class_registries/scenario_registry.py @@ -64,7 +64,7 @@ class ScenarioRegistry(BaseClassRegistry["Scenario", ScenarioMetadata]): 1. Built-in scenarios in pyrit.scenario.scenarios module 2. User-defined scenarios from initialization scripts (set via globals) - Scenarios are identified by their simple name (e.g., "encoding", "foundry"). + Scenarios are identified by their dotted name (e.g., "garak.encoding", "foundry.red_team_agent"). """ @classmethod @@ -115,15 +115,31 @@ def _discover_builtin_scenarios(self) -> None: package_path = Path(package_file).parent # Discover scenarios using the shared discovery utility - for module_name, scenario_class in discover_in_package( + # Use ``package_name.module_name`` as the registry name + for registry_name, scenario_class in discover_in_package( package_path=package_path, package_name="pyrit.scenario.scenarios", base_class=Scenario, # type: ignore[type-abstract] recursive=True, ): + # Skip deprecated alias classes + doc = (scenario_class.__doc__ or "").strip() + if doc.startswith("Deprecated alias"): + logger.debug(f"Skipping deprecated alias: {scenario_class.__name__}") + continue + + # Check for registry key collision + if registry_name in self._class_entries: + logger.warning( + f"Scenario registry name collision: '{registry_name}' " + f"conflicts with an already-registered scenario. Original " + f"scenario is kept: {self._class_entries[registry_name].registered_class.__name__}" + ) + continue + entry = ClassEntry(registered_class=scenario_class) - self._class_entries[module_name] = entry - logger.debug(f"Registered built-in scenario: {module_name} ({scenario_class.__name__})") + self._class_entries[registry_name] = entry + logger.debug(f"Registered built-in scenario: {registry_name} ({scenario_class.__name__})") except Exception as e: logger.error(f"Failed to discover built-in scenarios: {e}") @@ -182,6 +198,7 @@ def _build_metadata(self, name: str, entry: ClassEntry[Scenario]) -> ScenarioMet class_name=scenario_class.__name__, class_module=scenario_class.__module__, class_description=description, + registry_name=name, default_strategy=scenario_class.get_default_strategy().value, all_strategies=tuple(s.value for s in strategy_class.get_all_strategies()), aggregate_strategies=tuple(s.value for s in strategy_class.get_aggregate_strategies()), diff --git a/tests/end_to_end/test_scenarios.py b/tests/end_to_end/test_scenarios.py index 3bd96fce7..8b4047171 100644 --- a/tests/end_to_end/test_scenarios.py +++ b/tests/end_to_end/test_scenarios.py @@ -43,7 +43,7 @@ def test_scenario_with_pyrit_scan(scenario_name): [ scenario_name, "--initializers", - "targets", + "target", "load_default_datasets", "--target", "openai_chat", diff --git a/tests/unit/cli/test_frontend_core.py b/tests/unit/cli/test_frontend_core.py index 8275fd87a..61b3c7bb5 100644 --- a/tests/unit/cli/test_frontend_core.py +++ b/tests/unit/cli/test_frontend_core.py @@ -262,19 +262,21 @@ def test_resolve_initialization_scripts(self, mock_resolve: MagicMock): assert result == [Path("/test/script.py")] -class TestGetDefaultInitializerDiscoveryPath: - """Tests for get_default_initializer_discovery_path function.""" - - def test_get_default_initializer_discovery_path(self): - """Test get_default_initializer_discovery_path returns correct path.""" - path = frontend_core.get_default_initializer_discovery_path() +class TestListFunctions: + """Tests for list_scenarios_async and list_initializers_async functions.""" - assert isinstance(path, Path) - assert path.parts[-3:] == ("setup", "initializers", "scenarios") + def test_discover_builtin_scenarios_uses_dotted_names(self): + """Built-in scenario names should be dotted (package.module) lowercase names.""" + from pyrit.registry.class_registries.scenario_registry import ScenarioRegistry + registry = ScenarioRegistry() + registry._discover_builtin_scenarios() -class TestListFunctions: - """Tests for list_scenarios_async and list_initializers_async functions.""" + names = list(registry._class_entries.keys()) + assert len(names) > 0, "Should discover at least one built-in scenario" + for name in names: + assert "." in name, f"Scenario name '{name}' should be a dotted name (package.module)" + assert name == name.lower(), f"Scenario name '{name}' should be lowercase" async def test_list_scenarios(self): """Test list_scenarios_async returns scenarios from registry.""" @@ -290,8 +292,8 @@ async def test_list_scenarios(self): assert result == [{"name": "test_scenario"}] mock_registry.list_metadata.assert_called_once() - async def test_list_initializers_without_discovery_path(self): - """Test list_initializers_async without discovery path.""" + async def test_list_initializers(self): + """Test list_initializers_async returns initializers from context registry.""" mock_registry = MagicMock() mock_registry.list_metadata.return_value = [{"name": "test_init"}] @@ -304,21 +306,6 @@ async def test_list_initializers_without_discovery_path(self): assert result == [{"name": "test_init"}] mock_registry.list_metadata.assert_called_once() - @patch("pyrit.cli.frontend_core.InitializerRegistry") - async def test_list_initializers_with_discovery_path(self, mock_init_registry_class: MagicMock): - """Test list_initializers_async with discovery path.""" - mock_registry = MagicMock() - mock_registry.list_metadata.return_value = [{"name": "custom_init"}] - mock_init_registry_class.return_value = mock_registry - - context = frontend_core.FrontendCore() - discovery_path = Path("/custom/path") - - result = await frontend_core.list_initializers_async(context=context, discovery_path=discovery_path) - - mock_init_registry_class.assert_called_once_with(discovery_path=discovery_path) - assert result == [{"name": "custom_init"}] - class TestPrintFunctions: """Tests for print functions.""" @@ -332,6 +319,7 @@ async def test_print_scenarios_list_with_scenarios(self, capsys): class_name="TestScenario", class_module="test.scenarios", class_description="Test description", + registry_name="test", default_strategy="default", all_strategies=(), aggregate_strategies=(), @@ -347,8 +335,7 @@ async def test_print_scenarios_list_with_scenarios(self, capsys): assert result == 0 captured = capsys.readouterr() assert "Available Scenarios" in captured.out - # snake_class_name no longer strips suffix, so TestScenario -> test_scenario - assert "test_scenario" in captured.out + assert "test" in captured.out async def test_print_scenarios_list_empty(self, capsys): """Test print_scenarios_list with no scenarios.""" @@ -373,6 +360,7 @@ async def test_print_initializers_list_with_initializers(self, capsys): class_name="TestInit", class_module="test.initializers", class_description="Test initializer", + registry_name="test", display_name="test", execution_order=100, required_env_vars=(), @@ -386,7 +374,7 @@ async def test_print_initializers_list_with_initializers(self, capsys): assert result == 0 captured = capsys.readouterr() assert "Available Initializers" in captured.out - assert "test_init" in captured.out + assert "test" in captured.out async def test_print_initializers_list_empty(self, capsys): """Test print_initializers_list_async with no initializers.""" @@ -413,6 +401,7 @@ def test_format_scenario_metadata_basic(self, capsys): class_name="TestScenario", class_module="test.scenarios", class_description="", + registry_name="test", default_strategy="", all_strategies=(), aggregate_strategies=(), @@ -423,8 +412,7 @@ def test_format_scenario_metadata_basic(self, capsys): frontend_core.format_scenario_metadata(scenario_metadata=scenario_metadata) captured = capsys.readouterr() - # snake_class_name no longer strips suffix, so TestScenario -> test_scenario - assert "test_scenario" in captured.out + assert "test" in captured.out assert "TestScenario" in captured.out def test_format_scenario_metadata_with_description(self, capsys): @@ -434,6 +422,7 @@ def test_format_scenario_metadata_with_description(self, capsys): class_name="TestScenario", class_module="test.scenarios", class_description="This is a test scenario", + registry_name="test", default_strategy="", all_strategies=(), aggregate_strategies=(), @@ -452,6 +441,7 @@ def test_format_scenario_metadata_with_strategies(self, capsys): class_name="TestScenario", class_module="test.scenarios", class_description="", + registry_name="test", default_strategy="strategy1", all_strategies=("strategy1", "strategy2"), aggregate_strategies=(), @@ -472,6 +462,7 @@ def test_format_initializer_metadata_basic(self, capsys) -> None: class_name="TestInit", class_module="test.initializers", class_description="", + registry_name="test", display_name="test", required_env_vars=(), execution_order=100, @@ -480,7 +471,7 @@ def test_format_initializer_metadata_basic(self, capsys) -> None: frontend_core.format_initializer_metadata(initializer_metadata=initializer_metadata) captured = capsys.readouterr() - assert "test_init" in captured.out + assert "test" in captured.out assert "TestInit" in captured.out assert "100" in captured.out @@ -490,6 +481,7 @@ def test_format_initializer_metadata_with_env_vars(self, capsys) -> None: class_name="TestInit", class_module="test.initializers", class_description="", + registry_name="test", display_name="test", required_env_vars=("VAR1", "VAR2"), execution_order=100, @@ -507,6 +499,7 @@ def test_format_initializer_metadata_with_description(self, capsys) -> None: class_name="TestInit", class_module="test.initializers", class_description="Test description", + registry_name="test", display_name="test", required_env_vars=(), execution_order=100, @@ -582,27 +575,27 @@ def test_parse_run_arguments_with_initializers(self): def test_parse_run_arguments_with_initializer_params(self): """Test parsing initializers with key=value params.""" result = frontend_core.parse_run_arguments( - args_string="test_scenario --initializers simple targets:tags=default" + args_string="test_scenario --initializers simple target:tags=default" ) assert result["initializers"][0] == "simple" - assert result["initializers"][1] == {"name": "targets", "args": {"tags": ["default"]}} + assert result["initializers"][1] == {"name": "target", "args": {"tags": ["default"]}} def test_parse_run_arguments_with_initializer_multiple_params(self): """Test parsing initializers with multiple key=value params separated by semicolons.""" result = frontend_core.parse_run_arguments( - args_string="test_scenario --initializers targets:tags=default;mode=strict" + args_string="test_scenario --initializers target:tags=default;mode=strict" ) - assert result["initializers"][0] == {"name": "targets", "args": {"tags": ["default"], "mode": ["strict"]}} + assert result["initializers"][0] == {"name": "target", "args": {"tags": ["default"], "mode": ["strict"]}} def test_parse_run_arguments_with_initializer_comma_list(self): """Test parsing initializer params with comma-separated values into lists.""" result = frontend_core.parse_run_arguments( - args_string="test_scenario --initializers targets:tags=default,scorer" + args_string="test_scenario --initializers target:tags=default,scorer" ) - assert result["initializers"][0] == {"name": "targets", "args": {"tags": ["default", "scorer"]}} + assert result["initializers"][0] == {"name": "target", "args": {"tags": ["default", "scorer"]}} def test_parse_run_arguments_with_strategies(self): """Test parsing with strategies.""" @@ -1147,4 +1140,4 @@ async def test_print_targets_list_empty( assert result == 0 captured = capsys.readouterr() assert "No targets found" in captured.out - assert "--initializers targets" in captured.out + assert "--initializers target" in captured.out diff --git a/tests/unit/cli/test_pyrit_scan.py b/tests/unit/cli/test_pyrit_scan.py index 204376a1b..34a8b8ad5 100644 --- a/tests/unit/cli/test_pyrit_scan.py +++ b/tests/unit/cli/test_pyrit_scan.py @@ -173,22 +173,18 @@ def test_main_list_scenarios(self, mock_frontend_core: MagicMock, mock_print_sce @patch("pyrit.cli.frontend_core.print_initializers_list_async", new_callable=AsyncMock) @patch("pyrit.cli.frontend_core.FrontendCore") - @patch("pyrit.cli.frontend_core.get_default_initializer_discovery_path") def test_main_list_initializers( self, - mock_get_path: MagicMock, mock_frontend_core: MagicMock, mock_print_initializers: AsyncMock, ): """Test main with --list-initializers flag.""" mock_print_initializers.return_value = 0 - mock_get_path.return_value = Path("/test/path") result = pyrit_scan.main(["--list-initializers"]) assert result == 0 mock_print_initializers.assert_called_once() - mock_get_path.assert_called_once() @patch("pyrit.cli.frontend_core.print_scenarios_list_async", new_callable=AsyncMock) @patch("pyrit.cli.frontend_core.resolve_initialization_scripts") @@ -399,14 +395,11 @@ def test_main_list_scenarios_integration( assert result == 0 @patch("pyrit.cli.frontend_core.print_initializers_list_async", new_callable=AsyncMock) - @patch("pyrit.cli.frontend_core.get_default_initializer_discovery_path") def test_main_list_initializers_integration( self, - mock_get_path: MagicMock, mock_print_initializers: AsyncMock, ): """Test main --list-initializers with minimal mocking.""" - mock_get_path.return_value = Path("/test/path") mock_print_initializers.return_value = 0 result = pyrit_scan.main(["--list-initializers"]) diff --git a/tests/unit/cli/test_pyrit_shell.py b/tests/unit/cli/test_pyrit_shell.py index 2dea18a5a..89a218644 100644 --- a/tests/unit/cli/test_pyrit_shell.py +++ b/tests/unit/cli/test_pyrit_shell.py @@ -151,28 +151,14 @@ def test_do_list_scenarios_with_exception(self, mock_print_scenarios: AsyncMock, captured = capsys.readouterr() assert "Error listing scenarios" in captured.out - @patch("pyrit.cli.frontend_core.get_default_initializer_discovery_path") @patch("pyrit.cli.frontend_core.print_initializers_list_async", new_callable=AsyncMock) - def test_do_list_initializers(self, mock_print_initializers: AsyncMock, mock_get_path: MagicMock, shell): + def test_do_list_initializers(self, mock_print_initializers: AsyncMock, shell): """Test do_list_initializers command.""" s, ctx, _ = shell - mock_path = Path("/test/path") - mock_get_path.return_value = mock_path s.do_list_initializers("") - mock_print_initializers.assert_called_once_with(context=ctx, discovery_path=mock_path) - - @patch("pyrit.cli.frontend_core.print_initializers_list_async", new_callable=AsyncMock) - def test_do_list_initializers_with_path(self, mock_print_initializers: AsyncMock, shell): - """Test do_list_initializers with custom path.""" - s, ctx, _ = shell - - s.do_list_initializers("/custom/path") - - assert mock_print_initializers.call_count == 1 - call_kwargs = mock_print_initializers.call_args[1] - assert isinstance(call_kwargs["discovery_path"], Path) + mock_print_initializers.assert_called_once_with(context=ctx) @patch("pyrit.cli.frontend_core.print_initializers_list_async", new_callable=AsyncMock) def test_do_list_initializers_with_exception(self, mock_print_initializers: AsyncMock, shell, capsys): @@ -348,6 +334,43 @@ def test_do_run_with_exception( captured = capsys.readouterr() assert "Error: Test error" in captured.out + @patch("pyrit.cli.pyrit_shell.asyncio.run") + @patch("pyrit.cli.frontend_core.parse_run_arguments") + def test_do_run_keyboard_interrupt_returns_to_shell( + self, + mock_parse_args: MagicMock, + mock_asyncio_run: MagicMock, + shell, + capsys, + ): + """Test that Ctrl+C during scenario run returns to shell instead of crashing.""" + s, ctx, _ = shell + + mock_parse_args.return_value = { + "scenario_name": "test_scenario", + "initializers": ["test_init"], + "initialization_scripts": None, + "env_files": None, + "scenario_strategies": None, + "max_concurrency": None, + "max_retries": None, + "memory_labels": None, + "database": None, + "log_level": None, + "dataset_names": None, + "max_dataset_size": None, + "target": None, + } + + mock_asyncio_run.side_effect = KeyboardInterrupt() + + s.do_run("test_scenario --initializers test_init") + + captured = capsys.readouterr() + assert "interrupted" in captured.out.lower() + # Scenario should NOT be added to history + assert len(s._scenario_history) == 0 + def test_do_scenario_history_empty(self, shell, capsys): """Test do_scenario_history with no history.""" s, ctx, _ = shell diff --git a/tests/unit/memory/test_sqlite_memory.py b/tests/unit/memory/test_sqlite_memory.py index 5e6d3b168..6e9b5950b 100644 --- a/tests/unit/memory/test_sqlite_memory.py +++ b/tests/unit/memory/test_sqlite_memory.py @@ -668,3 +668,10 @@ def test_get_conversation_stats_batches_multiple_conversations(sqlite_instance): assert result[conv_ids[0]].message_count == 1 assert result[conv_ids[1]].message_count == 2 assert result[conv_ids[2]].message_count == 3 + + +def test_create_engine_uses_static_pool_for_in_memory(sqlite_instance): + """In-memory databases must use StaticPool so all threads share one database.""" + from sqlalchemy.pool import StaticPool + + assert isinstance(sqlite_instance.engine.pool, StaticPool)