From f95ed72df5b61be9ead26607a5416a294d0b545f Mon Sep 17 00:00:00 2001 From: Victor Valbuena Date: Tue, 7 Apr 2026 23:35:17 +0000 Subject: [PATCH 1/9] . --- pyrit/setup/initializers/airt.py | 64 +++++++++++++++++++---- tests/unit/setup/test_airt_initializer.py | 38 +++++++++++--- 2 files changed, 85 insertions(+), 17 deletions(-) diff --git a/pyrit/setup/initializers/airt.py b/pyrit/setup/initializers/airt.py index 96740565d8..52d9962754 100644 --- a/pyrit/setup/initializers/airt.py +++ b/pyrit/setup/initializers/airt.py @@ -43,12 +43,15 @@ class AIRTInitializer(PyRITInitializer): - Converter targets with Azure OpenAI configuration - Composite harm and objective scorers - Adversarial target configurations for attacks + - Use of an Azure SQL database Required Environment Variables: - AZURE_OPENAI_GPT4O_UNSAFE_CHAT_ENDPOINT: Azure OpenAI endpoint for converters and targets - AZURE_OPENAI_GPT4O_UNSAFE_CHAT_MODEL: Azure OpenAI model name for converters and targets - AZURE_OPENAI_GPT4O_UNSAFE_CHAT_ENDPOINT2: Azure OpenAI endpoint for scoring - AZURE_OPENAI_GPT4O_UNSAFE_CHAT_MODEL2: Azure OpenAI model name for scoring + - AZURE_SQL_DB_CONNECTION_STRING: Azure SQL database connection string + - AZURE_STORAGE_ACCOUNT_DB_DATA_CONTAINER_URL: Azure SQL database location Optional Environment Variables: - AZURE_OPENAI_GPT4O_UNSAFE_CHAT_KEY: API key for converter endpoint. If not set, Entra ID auth is used. @@ -90,6 +93,8 @@ def required_env_vars(self) -> list[str]: "AZURE_OPENAI_GPT4O_UNSAFE_CHAT_ENDPOINT2", "AZURE_OPENAI_GPT4O_UNSAFE_CHAT_MODEL2", "AZURE_CONTENT_SAFETY_API_ENDPOINT", + "AZURE_SQL_DB_CONNECTION_STRING", + "AZURE_STORAGE_ACCOUNT_DB_DATA_CONTAINER_URL" ] async def initialize_async(self) -> None: @@ -102,9 +107,14 @@ async def initialize_async(self) -> None: 3. Adversarial target configurations 4. Default values for all attack types """ + # Ensure op_name, username, and email are populated from GLOBAL_MEMORY_LABELS. + self._validate_memory_labels() + # Get environment variables (validated by validate() method) - converter_endpoint = os.getenv("AZURE_OPENAI_GPT4O_UNSAFE_CHAT_ENDPOINT") - converter_model_name = os.getenv("AZURE_OPENAI_GPT4O_UNSAFE_CHAT_MODEL") + converter_endpoint = os.getenv( + "AZURE_OPENAI_GPT4O_UNSAFE_CHAT_ENDPOINT") + converter_model_name = os.getenv( + "AZURE_OPENAI_GPT4O_UNSAFE_CHAT_MODEL") scorer_endpoint = os.getenv("AZURE_OPENAI_GPT4O_UNSAFE_CHAT_ENDPOINT2") scorer_model_name = os.getenv("AZURE_OPENAI_GPT4O_UNSAFE_CHAT_MODEL2") @@ -114,8 +124,10 @@ async def initialize_async(self) -> None: # model name can be empty in certain cases (e.g., custom model deployments that don't need model name) # Check for API keys first, fall back to Entra auth if not set - converter_api_key = os.getenv("AZURE_OPENAI_GPT4O_UNSAFE_CHAT_KEY") or get_azure_openai_auth(converter_endpoint) - scorer_api_key = os.getenv("AZURE_OPENAI_GPT4O_UNSAFE_CHAT_KEY2") or get_azure_openai_auth(scorer_endpoint) + converter_api_key = os.getenv( + "AZURE_OPENAI_GPT4O_UNSAFE_CHAT_KEY") or get_azure_openai_auth(converter_endpoint) + scorer_api_key = os.getenv( + "AZURE_OPENAI_GPT4O_UNSAFE_CHAT_KEY2") or get_azure_openai_auth(scorer_endpoint) content_safety_api_key_str = os.getenv("AZURE_CONTENT_SAFETY_API_KEY") content_safety_api_key: str | Callable[[], str] = ( content_safety_api_key_str @@ -150,7 +162,8 @@ def _setup_converter_target(self, *, endpoint: str, api_key: str, model_name: st temperature=1.1, ) - set_global_variable(name="default_converter_target", value=default_converter_target) + set_global_variable(name="default_converter_target", + value=default_converter_target) set_default_value( class_type=PromptConverter, parameter_name="converter_target", @@ -187,7 +200,8 @@ def _setup_scorers( TrueFalseInverterScorer( scorer=SelfAskRefusalScorer(chat_target=scorer_target), ), - FloatScaleThresholdScorer(scorer=SelfAskScaleScorer(chat_target=scorer_target), threshold=0.7), + FloatScaleThresholdScorer(scorer=SelfAskScaleScorer( + chat_target=scorer_target), threshold=0.7), ], ) @@ -201,16 +215,20 @@ def _setup_scorers( TrueFalseInverterScorer( scorer=SelfAskRefusalScorer(chat_target=scorer_target), ), - FloatScaleThresholdScorer(scorer=SelfAskScaleScorer(chat_target=scorer_target), threshold=0.7), + FloatScaleThresholdScorer(scorer=SelfAskScaleScorer( + chat_target=scorer_target), threshold=0.7), ], ) # Set global variables - set_global_variable(name="default_harm_scorer", value=default_harm_scorer) - set_global_variable(name="default_objective_scorer", value=default_objective_scorer) + set_global_variable(name="default_harm_scorer", + value=default_harm_scorer) + set_global_variable(name="default_objective_scorer", + value=default_objective_scorer) # Configure default attack scoring configuration - default_objective_scorer_config = AttackScoringConfig(objective_scorer=default_objective_scorer) + default_objective_scorer_config = AttackScoringConfig( + objective_scorer=default_objective_scorer) # Set default values for various attack types attack_classes = [ @@ -239,7 +257,8 @@ def _setup_adversarial_targets(self, *, endpoint: str, api_key: str, model_name: ) # Set global variable for easy access - set_global_variable(name="adversarial_config", value=adversarial_config) + set_global_variable(name="adversarial_config", + value=adversarial_config) # Set default adversarial configurations for various attack types attack_classes = [ @@ -255,3 +274,26 @@ def _setup_adversarial_targets(self, *, endpoint: str, api_key: str, model_name: parameter_name="attack_adversarial_config", value=adversarial_config, ) + + def _validate_memory_labels(self) -> None: + """ + Check that mandatory global memory labels (username, email, and op_name) + are populated. Note that this is a separate path than Initializer.validate + since the presence of GLOBAL_MEMORY_LABELS doesn't indicate it has the + necessary fields for AIRTInitializer. + + Raises: + ValueError: If mandatory global memory labels are missing. + """ + labels = os.environ.get("GLOBAL_MEMORY_LABELS") + if not labels: + raise ValueError( + "Error: GLOBAL_MEMORY_LABELS was not set! Please add it to `.env.local` before running the initializer.") + + missing_fields = list(set(labels) - {"op_name", "username", "email"}) + + if missing_fields: + raise ValueError( + f"Error: AIRTInitializer was called, but the following fields were not found: \ + {missing_fields}. Please add these to GLOBAL_MEMORY_LABELS in `.env.local` \ + before running the initializer.") diff --git a/tests/unit/setup/test_airt_initializer.py b/tests/unit/setup/test_airt_initializer.py index 2a8606cded..4618c625c3 100644 --- a/tests/unit/setup/test_airt_initializer.py +++ b/tests/unit/setup/test_airt_initializer.py @@ -79,8 +79,10 @@ async def test_initialize_runs_without_error(self): """Test that initialize runs without errors when no API keys are set (Entra auth fallback).""" init = AIRTInitializer() with ( - patch("pyrit.setup.initializers.airt.get_azure_openai_auth", return_value="mock_token"), - patch("pyrit.setup.initializers.airt.get_azure_token_provider", return_value="mock_token_provider"), + patch("pyrit.setup.initializers.airt.get_azure_openai_auth", + return_value="mock_token"), + patch("pyrit.setup.initializers.airt.get_azure_token_provider", + return_value="mock_token_provider"), ): await init.initialize_async() @@ -114,8 +116,10 @@ async def test_get_info_after_initialize_has_populated_data(self): """Test that get_info_async() returns populated data after initialization.""" init = AIRTInitializer() with ( - patch("pyrit.setup.initializers.airt.get_azure_openai_auth", return_value="mock_token"), - patch("pyrit.setup.initializers.airt.get_azure_token_provider", return_value="mock_token_provider"), + patch("pyrit.setup.initializers.airt.get_azure_openai_auth", + return_value="mock_token"), + patch("pyrit.setup.initializers.airt.get_azure_token_provider", + return_value="mock_token_provider"), ): await init.initialize_async() # get_info_async re-runs initialize_async internally, so patches must still be active @@ -129,7 +133,8 @@ async def test_get_info_after_initialize_has_populated_data(self): # Verify default_values list is populated and not empty assert isinstance(info["default_values"], list) - assert len(info["default_values"]) > 0, "default_values should be populated after initialization" + assert len(info["default_values"] + ) > 0, "default_values should be populated after initialization" # Verify expected default values are present default_values_str = str(info["default_values"]) @@ -139,7 +144,8 @@ async def test_get_info_after_initialize_has_populated_data(self): # Verify global_variables list is populated and not empty assert isinstance(info["global_variables"], list) - assert len(info["global_variables"]) > 0, "global_variables should be populated after initialization" + assert len(info["global_variables"] + ) > 0, "global_variables should be populated after initialization" # Verify expected global variables are present assert "default_converter_target" in info["global_variables"] @@ -174,6 +180,26 @@ def test_validate_missing_multiple_env_vars_raises_error(self): assert "AZURE_OPENAI_GPT4O_UNSAFE_CHAT_ENDPOINT" in error_message assert "AZURE_OPENAI_GPT4O_UNSAFE_CHAT_MODEL" in error_message + def test_validate_missing_memory_labels_raises_error(self): + """Test that GLOBAL_MEMORY_LABELS not being set raises an error.""" + del os.environ["GLOBAL_MEMORY_LABELS"] + init = AIRTInitializer() + with pytest.raises(ValueError) as exc_info: + init.validate() + + error_message = str(exc_info.value) + assert "GLOBAL_MEMORY_LABELS" in error_message + + def test_validate_missing_op_name_raises_error(self): + """Test that op_name not being set raises an error.""" + os.environ["GLOBAL_MEMORY_LABELS"] = "{'test_key': 'test_value'}" + init = AIRTInitializer() + with pytest.raises(ValueError) as exc_info: + init.validate() + + error_message = str(exc_info.value) + assert "op_name" in error_message + class TestAIRTInitializerGetInfo: """Tests for AIRTInitializer.get_info method - basic functionality.""" From 3363011b9c595815a27e4a78a37eb776b5a6233e Mon Sep 17 00:00:00 2001 From: Victor Valbuena Date: Wed, 8 Apr 2026 00:14:47 +0000 Subject: [PATCH 2/9] . --- pyrit/setup/initializers/airt.py | 4 +++- tests/unit/setup/test_airt_initializer.py | 10 ++++++++-- 2 files changed, 11 insertions(+), 3 deletions(-) diff --git a/pyrit/setup/initializers/airt.py b/pyrit/setup/initializers/airt.py index 52d9962754..936e7e5cd8 100644 --- a/pyrit/setup/initializers/airt.py +++ b/pyrit/setup/initializers/airt.py @@ -290,7 +290,9 @@ def _validate_memory_labels(self) -> None: raise ValueError( "Error: GLOBAL_MEMORY_LABELS was not set! Please add it to `.env.local` before running the initializer.") - missing_fields = list(set(labels) - {"op_name", "username", "email"}) + required_fields = ["op_name", "username", "email"] + missing_fields = [ + field for field in required_fields if field not in labels] if missing_fields: raise ValueError( diff --git a/tests/unit/setup/test_airt_initializer.py b/tests/unit/setup/test_airt_initializer.py index 4618c625c3..9811b15c1a 100644 --- a/tests/unit/setup/test_airt_initializer.py +++ b/tests/unit/setup/test_airt_initializer.py @@ -41,6 +41,9 @@ def setup_method(self) -> None: os.environ["AZURE_OPENAI_GPT4O_UNSAFE_CHAT_ENDPOINT2"] = "https://test-scorer.openai.azure.com" os.environ["AZURE_OPENAI_GPT4O_UNSAFE_CHAT_MODEL2"] = "gpt-4" os.environ["AZURE_CONTENT_SAFETY_API_ENDPOINT"] = "https://test-safety.cognitiveservices.azure.com" + os.environ["AZURE_SQL_DB_CONNECTION_STRING"] = "Server=test.database.windows.net;Database=testdb" + os.environ["AZURE_STORAGE_ACCOUNT_DB_DATA_CONTAINER_URL"] = "https://teststorage.blob.core.windows.net/data" + os.environ["GLOBAL_MEMORY_LABELS"] = "{'op_name': 'test_op', 'username': 'test_user', 'email': 'test@test.com'}" # Clean up globals for attr in [ "default_converter_target", @@ -61,6 +64,9 @@ def teardown_method(self) -> None: "AZURE_OPENAI_GPT4O_UNSAFE_CHAT_ENDPOINT2", "AZURE_OPENAI_GPT4O_UNSAFE_CHAT_MODEL2", "AZURE_CONTENT_SAFETY_API_ENDPOINT", + "AZURE_SQL_DB_CONNECTION_STRING", + "AZURE_STORAGE_ACCOUNT_DB_DATA_CONTAINER_URL", + "GLOBAL_MEMORY_LABELS", ]: if var in os.environ: del os.environ[var] @@ -185,7 +191,7 @@ def test_validate_missing_memory_labels_raises_error(self): del os.environ["GLOBAL_MEMORY_LABELS"] init = AIRTInitializer() with pytest.raises(ValueError) as exc_info: - init.validate() + init._validate_memory_labels() error_message = str(exc_info.value) assert "GLOBAL_MEMORY_LABELS" in error_message @@ -195,7 +201,7 @@ def test_validate_missing_op_name_raises_error(self): os.environ["GLOBAL_MEMORY_LABELS"] = "{'test_key': 'test_value'}" init = AIRTInitializer() with pytest.raises(ValueError) as exc_info: - init.validate() + init._validate_memory_labels() error_message = str(exc_info.value) assert "op_name" in error_message From 57a26fc5eab44bd634413c0fd52a2e39aad9f9ee Mon Sep 17 00:00:00 2001 From: Victor Valbuena Date: Wed, 8 Apr 2026 23:19:23 +0000 Subject: [PATCH 3/9] . --- pyrit/setup/configuration_loader.py | 2 ++ pyrit/setup/initializers/airt.py | 44 ++++++++++++++--------- tests/unit/setup/test_airt_initializer.py | 23 ++++-------- 3 files changed, 35 insertions(+), 34 deletions(-) diff --git a/pyrit/setup/configuration_loader.py b/pyrit/setup/configuration_loader.py index 5e2ce65b9b..3c67b2f4b7 100644 --- a/pyrit/setup/configuration_loader.py +++ b/pyrit/setup/configuration_loader.py @@ -70,6 +70,8 @@ class ConfigurationLoader(YamlLoadable): env_files: List of environment file paths to load. None means "use defaults (.env, .env.local)", [] means "load nothing". silent: Whether to suppress initialization messages. + operator: Name for the current operator, e.g. a team or username. + operation: Name for the current operation. Example YAML configuration: memory_db_type: sqlite diff --git a/pyrit/setup/initializers/airt.py b/pyrit/setup/initializers/airt.py index 936e7e5cd8..5e57369523 100644 --- a/pyrit/setup/initializers/airt.py +++ b/pyrit/setup/initializers/airt.py @@ -8,8 +8,12 @@ AIRT configuration including converters, scorers, and targets using Azure OpenAI. """ +import json import os from collections.abc import Callable +from pathlib import Path + +import yaml from pyrit.auth import get_azure_openai_auth, get_azure_token_provider from pyrit.common.apply_defaults import set_default_value, set_global_variable @@ -108,7 +112,7 @@ async def initialize_async(self) -> None: 4. Default values for all attack types """ # Ensure op_name, username, and email are populated from GLOBAL_MEMORY_LABELS. - self._validate_memory_labels() + self._validate_operation_fields() # Get environment variables (validated by validate() method) converter_endpoint = os.getenv( @@ -275,27 +279,33 @@ def _setup_adversarial_targets(self, *, endpoint: str, api_key: str, model_name: value=adversarial_config, ) - def _validate_memory_labels(self) -> None: + def _validate_operation_fields(self) -> None: """ - Check that mandatory global memory labels (username, email, and op_name) - are populated. Note that this is a separate path than Initializer.validate - since the presence of GLOBAL_MEMORY_LABELS doesn't indicate it has the - necessary fields for AIRTInitializer. + Check that mandatory global memory labels (operation, operator) + are populated. Raises: ValueError: If mandatory global memory labels are missing. """ - labels = os.environ.get("GLOBAL_MEMORY_LABELS") - if not labels: - raise ValueError( - "Error: GLOBAL_MEMORY_LABELS was not set! Please add it to `.env.local` before running the initializer.") + config_path = Path.home() / ".pyrit" / ".pyrit_conf" + with open(config_path) as f: + data = yaml.load(f, Loader=yaml.SafeLoader) - required_fields = ["op_name", "username", "email"] - missing_fields = [ - field for field in required_fields if field not in labels] + if "operator" not in data: + raise ValueError( + "Error: `operator` was not set in .pyrit_conf. This is a required value for the AIRTInitializer.") - if missing_fields: + if "operation" not in data: raise ValueError( - f"Error: AIRTInitializer was called, but the following fields were not found: \ - {missing_fields}. Please add these to GLOBAL_MEMORY_LABELS in `.env.local` \ - before running the initializer.") + "Error: `operation` was not set in .pyrit_conf. This is a required value for the AIRTInitializer.") + + labels = os.environ.get("GLOBAL_MEMORY_LABELS") + labels = json.loads(labels) if labels else {} + + if "username" not in labels: + labels["username"] = data["operator"] + + if "op_name" not in labels: + labels["op_name"] = data["operation"] + + os.environ["GLOBAL_MEMORY_LABELS"] = json.dumps(labels) diff --git a/tests/unit/setup/test_airt_initializer.py b/tests/unit/setup/test_airt_initializer.py index 9811b15c1a..8338769eb5 100644 --- a/tests/unit/setup/test_airt_initializer.py +++ b/tests/unit/setup/test_airt_initializer.py @@ -186,25 +186,14 @@ def test_validate_missing_multiple_env_vars_raises_error(self): assert "AZURE_OPENAI_GPT4O_UNSAFE_CHAT_ENDPOINT" in error_message assert "AZURE_OPENAI_GPT4O_UNSAFE_CHAT_MODEL" in error_message - def test_validate_missing_memory_labels_raises_error(self): - """Test that GLOBAL_MEMORY_LABELS not being set raises an error.""" - del os.environ["GLOBAL_MEMORY_LABELS"] - init = AIRTInitializer() - with pytest.raises(ValueError) as exc_info: - init._validate_memory_labels() + def test_validate_missing_operator_raises_error(self): + pass - error_message = str(exc_info.value) - assert "GLOBAL_MEMORY_LABELS" in error_message + def test_validate_missing_operation_raises_error(self): + pass - def test_validate_missing_op_name_raises_error(self): - """Test that op_name not being set raises an error.""" - os.environ["GLOBAL_MEMORY_LABELS"] = "{'test_key': 'test_value'}" - init = AIRTInitializer() - with pytest.raises(ValueError) as exc_info: - init._validate_memory_labels() - - error_message = str(exc_info.value) - assert "op_name" in error_message + def test_validate_db_connection_raises_error(self): + pass class TestAIRTInitializerGetInfo: From c84f4c088f1b6a6e63aa10c48c2c66d42042c4c9 Mon Sep 17 00:00:00 2001 From: Victor Valbuena Date: Wed, 8 Apr 2026 23:21:56 +0000 Subject: [PATCH 4/9] . --- tests/unit/setup/test_airt_initializer.py | 48 +++++++++++++++-------- 1 file changed, 32 insertions(+), 16 deletions(-) diff --git a/tests/unit/setup/test_airt_initializer.py b/tests/unit/setup/test_airt_initializer.py index 8338769eb5..95fe0f09ea 100644 --- a/tests/unit/setup/test_airt_initializer.py +++ b/tests/unit/setup/test_airt_initializer.py @@ -3,9 +3,10 @@ import os import sys -from unittest.mock import patch +from unittest.mock import mock_open, patch import pytest +import yaml from pyrit.common.apply_defaults import reset_default_values from pyrit.setup.initializers import AIRTInitializer @@ -85,10 +86,8 @@ async def test_initialize_runs_without_error(self): """Test that initialize runs without errors when no API keys are set (Entra auth fallback).""" init = AIRTInitializer() with ( - patch("pyrit.setup.initializers.airt.get_azure_openai_auth", - return_value="mock_token"), - patch("pyrit.setup.initializers.airt.get_azure_token_provider", - return_value="mock_token_provider"), + patch("pyrit.setup.initializers.airt.get_azure_openai_auth", return_value="mock_token"), + patch("pyrit.setup.initializers.airt.get_azure_token_provider", return_value="mock_token_provider"), ): await init.initialize_async() @@ -122,10 +121,8 @@ async def test_get_info_after_initialize_has_populated_data(self): """Test that get_info_async() returns populated data after initialization.""" init = AIRTInitializer() with ( - patch("pyrit.setup.initializers.airt.get_azure_openai_auth", - return_value="mock_token"), - patch("pyrit.setup.initializers.airt.get_azure_token_provider", - return_value="mock_token_provider"), + patch("pyrit.setup.initializers.airt.get_azure_openai_auth", return_value="mock_token"), + patch("pyrit.setup.initializers.airt.get_azure_token_provider", return_value="mock_token_provider"), ): await init.initialize_async() # get_info_async re-runs initialize_async internally, so patches must still be active @@ -139,8 +136,7 @@ async def test_get_info_after_initialize_has_populated_data(self): # Verify default_values list is populated and not empty assert isinstance(info["default_values"], list) - assert len(info["default_values"] - ) > 0, "default_values should be populated after initialization" + assert len(info["default_values"]) > 0, "default_values should be populated after initialization" # Verify expected default values are present default_values_str = str(info["default_values"]) @@ -150,8 +146,7 @@ async def test_get_info_after_initialize_has_populated_data(self): # Verify global_variables list is populated and not empty assert isinstance(info["global_variables"], list) - assert len(info["global_variables"] - ) > 0, "global_variables should be populated after initialization" + assert len(info["global_variables"]) > 0, "global_variables should be populated after initialization" # Verify expected global variables are present assert "default_converter_target" in info["global_variables"] @@ -187,13 +182,34 @@ def test_validate_missing_multiple_env_vars_raises_error(self): assert "AZURE_OPENAI_GPT4O_UNSAFE_CHAT_MODEL" in error_message def test_validate_missing_operator_raises_error(self): - pass + """Test that _validate_operation_fields raises error when operator is missing from .pyrit_conf.""" + conf_data = yaml.dump({"operation": "test_op"}) + init = AIRTInitializer() + with ( + patch("builtins.open", mock_open(read_data=conf_data)), + pytest.raises(ValueError, match="operator"), + ): + init._validate_operation_fields() def test_validate_missing_operation_raises_error(self): - pass + """Test that _validate_operation_fields raises error when operation is missing from .pyrit_conf.""" + conf_data = yaml.dump({"operator": "test_user"}) + init = AIRTInitializer() + with ( + patch("builtins.open", mock_open(read_data=conf_data)), + pytest.raises(ValueError, match="operation"), + ): + init._validate_operation_fields() def test_validate_db_connection_raises_error(self): - pass + """Test that validate raises error when AZURE_SQL_DB_CONNECTION_STRING is missing.""" + del os.environ["AZURE_SQL_DB_CONNECTION_STRING"] + init = AIRTInitializer() + with pytest.raises(ValueError) as exc_info: + init.validate() + + error_message = str(exc_info.value) + assert "AZURE_SQL_DB_CONNECTION_STRING" in error_message class TestAIRTInitializerGetInfo: From ffc4955dd7fd2478bfac3165b2aa11afa9934ce7 Mon Sep 17 00:00:00 2001 From: Victor Valbuena Date: Wed, 8 Apr 2026 23:34:28 +0000 Subject: [PATCH 5/9] precommit --- pyrit/setup/configuration_loader.py | 2 +- pyrit/setup/initializers/airt.py | 45 ++++++++++++----------------- 2 files changed, 19 insertions(+), 28 deletions(-) diff --git a/pyrit/setup/configuration_loader.py b/pyrit/setup/configuration_loader.py index 3c67b2f4b7..6f720b3035 100644 --- a/pyrit/setup/configuration_loader.py +++ b/pyrit/setup/configuration_loader.py @@ -71,7 +71,7 @@ class ConfigurationLoader(YamlLoadable): None means "use defaults (.env, .env.local)", [] means "load nothing". silent: Whether to suppress initialization messages. operator: Name for the current operator, e.g. a team or username. - operation: Name for the current operation. + operation: Name for the current operation. Example YAML configuration: memory_db_type: sqlite diff --git a/pyrit/setup/initializers/airt.py b/pyrit/setup/initializers/airt.py index 5e57369523..51ca7289da 100644 --- a/pyrit/setup/initializers/airt.py +++ b/pyrit/setup/initializers/airt.py @@ -98,7 +98,7 @@ def required_env_vars(self) -> list[str]: "AZURE_OPENAI_GPT4O_UNSAFE_CHAT_MODEL2", "AZURE_CONTENT_SAFETY_API_ENDPOINT", "AZURE_SQL_DB_CONNECTION_STRING", - "AZURE_STORAGE_ACCOUNT_DB_DATA_CONTAINER_URL" + "AZURE_STORAGE_ACCOUNT_DB_DATA_CONTAINER_URL", ] async def initialize_async(self) -> None: @@ -115,10 +115,8 @@ async def initialize_async(self) -> None: self._validate_operation_fields() # Get environment variables (validated by validate() method) - converter_endpoint = os.getenv( - "AZURE_OPENAI_GPT4O_UNSAFE_CHAT_ENDPOINT") - converter_model_name = os.getenv( - "AZURE_OPENAI_GPT4O_UNSAFE_CHAT_MODEL") + converter_endpoint = os.getenv("AZURE_OPENAI_GPT4O_UNSAFE_CHAT_ENDPOINT") + converter_model_name = os.getenv("AZURE_OPENAI_GPT4O_UNSAFE_CHAT_MODEL") scorer_endpoint = os.getenv("AZURE_OPENAI_GPT4O_UNSAFE_CHAT_ENDPOINT2") scorer_model_name = os.getenv("AZURE_OPENAI_GPT4O_UNSAFE_CHAT_MODEL2") @@ -128,10 +126,8 @@ async def initialize_async(self) -> None: # model name can be empty in certain cases (e.g., custom model deployments that don't need model name) # Check for API keys first, fall back to Entra auth if not set - converter_api_key = os.getenv( - "AZURE_OPENAI_GPT4O_UNSAFE_CHAT_KEY") or get_azure_openai_auth(converter_endpoint) - scorer_api_key = os.getenv( - "AZURE_OPENAI_GPT4O_UNSAFE_CHAT_KEY2") or get_azure_openai_auth(scorer_endpoint) + converter_api_key = os.getenv("AZURE_OPENAI_GPT4O_UNSAFE_CHAT_KEY") or get_azure_openai_auth(converter_endpoint) + scorer_api_key = os.getenv("AZURE_OPENAI_GPT4O_UNSAFE_CHAT_KEY2") or get_azure_openai_auth(scorer_endpoint) content_safety_api_key_str = os.getenv("AZURE_CONTENT_SAFETY_API_KEY") content_safety_api_key: str | Callable[[], str] = ( content_safety_api_key_str @@ -166,8 +162,7 @@ def _setup_converter_target(self, *, endpoint: str, api_key: str, model_name: st temperature=1.1, ) - set_global_variable(name="default_converter_target", - value=default_converter_target) + set_global_variable(name="default_converter_target", value=default_converter_target) set_default_value( class_type=PromptConverter, parameter_name="converter_target", @@ -204,8 +199,7 @@ def _setup_scorers( TrueFalseInverterScorer( scorer=SelfAskRefusalScorer(chat_target=scorer_target), ), - FloatScaleThresholdScorer(scorer=SelfAskScaleScorer( - chat_target=scorer_target), threshold=0.7), + FloatScaleThresholdScorer(scorer=SelfAskScaleScorer(chat_target=scorer_target), threshold=0.7), ], ) @@ -219,20 +213,16 @@ def _setup_scorers( TrueFalseInverterScorer( scorer=SelfAskRefusalScorer(chat_target=scorer_target), ), - FloatScaleThresholdScorer(scorer=SelfAskScaleScorer( - chat_target=scorer_target), threshold=0.7), + FloatScaleThresholdScorer(scorer=SelfAskScaleScorer(chat_target=scorer_target), threshold=0.7), ], ) # Set global variables - set_global_variable(name="default_harm_scorer", - value=default_harm_scorer) - set_global_variable(name="default_objective_scorer", - value=default_objective_scorer) + set_global_variable(name="default_harm_scorer", value=default_harm_scorer) + set_global_variable(name="default_objective_scorer", value=default_objective_scorer) # Configure default attack scoring configuration - default_objective_scorer_config = AttackScoringConfig( - objective_scorer=default_objective_scorer) + default_objective_scorer_config = AttackScoringConfig(objective_scorer=default_objective_scorer) # Set default values for various attack types attack_classes = [ @@ -261,8 +251,7 @@ def _setup_adversarial_targets(self, *, endpoint: str, api_key: str, model_name: ) # Set global variable for easy access - set_global_variable(name="adversarial_config", - value=adversarial_config) + set_global_variable(name="adversarial_config", value=adversarial_config) # Set default adversarial configurations for various attack types attack_classes = [ @@ -293,14 +282,16 @@ def _validate_operation_fields(self) -> None: if "operator" not in data: raise ValueError( - "Error: `operator` was not set in .pyrit_conf. This is a required value for the AIRTInitializer.") + "Error: `operator` was not set in .pyrit_conf. This is a required value for the AIRTInitializer." + ) if "operation" not in data: raise ValueError( - "Error: `operation` was not set in .pyrit_conf. This is a required value for the AIRTInitializer.") + "Error: `operation` was not set in .pyrit_conf. This is a required value for the AIRTInitializer." + ) - labels = os.environ.get("GLOBAL_MEMORY_LABELS") - labels = json.loads(labels) if labels else {} + raw_labels = os.environ.get("GLOBAL_MEMORY_LABELS") + labels = dict(json.loads(raw_labels)) if raw_labels else {} if "username" not in labels: labels["username"] = data["operator"] From 84dfda75fe23b35abf975f613e6f877ec7aae248 Mon Sep 17 00:00:00 2001 From: Victor Valbuena Date: Wed, 8 Apr 2026 23:50:24 +0000 Subject: [PATCH 6/9] unit test --- tests/unit/setup/test_airt_initializer.py | 20 ++++++++++++++++---- 1 file changed, 16 insertions(+), 4 deletions(-) diff --git a/tests/unit/setup/test_airt_initializer.py b/tests/unit/setup/test_airt_initializer.py index 95fe0f09ea..5f059302c9 100644 --- a/tests/unit/setup/test_airt_initializer.py +++ b/tests/unit/setup/test_airt_initializer.py @@ -3,6 +3,7 @@ import os import sys +from pathlib import Path from unittest.mock import mock_open, patch import pytest @@ -12,6 +13,17 @@ from pyrit.setup.initializers import AIRTInitializer +@pytest.fixture +def patch_pyrit_conf(tmp_path): + """Create a temporary .pyrit_conf file and patch _validate_operation_fields to read from it.""" + conf_file = tmp_path / ".pyrit_conf" + conf_file.write_text(yaml.dump({"operator": "test_user", "operation": "test_op"})) + with patch.object(Path, "home", return_value=tmp_path): + (tmp_path / ".pyrit").mkdir(exist_ok=True) + (tmp_path / ".pyrit" / ".pyrit_conf").write_text(yaml.dump({"operator": "test_user", "operation": "test_op"})) + yield + + class TestAIRTInitializer: """Tests for AIRTInitializer class - basic functionality.""" @@ -44,7 +56,7 @@ def setup_method(self) -> None: os.environ["AZURE_CONTENT_SAFETY_API_ENDPOINT"] = "https://test-safety.cognitiveservices.azure.com" os.environ["AZURE_SQL_DB_CONNECTION_STRING"] = "Server=test.database.windows.net;Database=testdb" os.environ["AZURE_STORAGE_ACCOUNT_DB_DATA_CONTAINER_URL"] = "https://teststorage.blob.core.windows.net/data" - os.environ["GLOBAL_MEMORY_LABELS"] = "{'op_name': 'test_op', 'username': 'test_user', 'email': 'test@test.com'}" + os.environ["GLOBAL_MEMORY_LABELS"] = '{"op_name": "test_op", "username": "test_user", "email": "test@test.com"}' # Clean up globals for attr in [ "default_converter_target", @@ -82,7 +94,7 @@ def teardown_method(self) -> None: delattr(sys.modules["__main__"], attr) @pytest.mark.asyncio - async def test_initialize_runs_without_error(self): + async def test_initialize_runs_without_error(self, patch_pyrit_conf): """Test that initialize runs without errors when no API keys are set (Entra auth fallback).""" init = AIRTInitializer() with ( @@ -92,7 +104,7 @@ async def test_initialize_runs_without_error(self): await init.initialize_async() @pytest.mark.asyncio - async def test_initialize_uses_api_keys_when_set(self): + async def test_initialize_uses_api_keys_when_set(self, patch_pyrit_conf): """Test that initialize uses API keys from env vars when they are set.""" os.environ["AZURE_OPENAI_GPT4O_UNSAFE_CHAT_KEY"] = "converter-key" os.environ["AZURE_OPENAI_GPT4O_UNSAFE_CHAT_KEY2"] = "scorer-key" @@ -117,7 +129,7 @@ async def test_initialize_uses_api_keys_when_set(self): del os.environ[var] @pytest.mark.asyncio - async def test_get_info_after_initialize_has_populated_data(self): + async def test_get_info_after_initialize_has_populated_data(self, patch_pyrit_conf): """Test that get_info_async() returns populated data after initialization.""" init = AIRTInitializer() with ( From e330f1ffab4c8fa32bc980d660a26bc2f36585d1 Mon Sep 17 00:00:00 2001 From: Victor Valbuena Date: Thu, 9 Apr 2026 18:19:41 +0000 Subject: [PATCH 7/9] . --- pyrit/setup/initializers/airt.py | 5 ++--- 1 file changed, 2 insertions(+), 3 deletions(-) diff --git a/pyrit/setup/initializers/airt.py b/pyrit/setup/initializers/airt.py index 51ca7289da..251ff7941d 100644 --- a/pyrit/setup/initializers/airt.py +++ b/pyrit/setup/initializers/airt.py @@ -11,12 +11,12 @@ import json import os from collections.abc import Callable -from pathlib import Path import yaml from pyrit.auth import get_azure_openai_auth, get_azure_token_provider from pyrit.common.apply_defaults import set_default_value, set_global_variable +from pyrit.common.path import DEFAULT_CONFIG_PATH from pyrit.executor.attack import ( AttackAdversarialConfig, AttackScoringConfig, @@ -276,8 +276,7 @@ def _validate_operation_fields(self) -> None: Raises: ValueError: If mandatory global memory labels are missing. """ - config_path = Path.home() / ".pyrit" / ".pyrit_conf" - with open(config_path) as f: + with open(DEFAULT_CONFIG_PATH) as f: data = yaml.load(f, Loader=yaml.SafeLoader) if "operator" not in data: From 2475596999c5172711e93fc66145ef6a7a927920 Mon Sep 17 00:00:00 2001 From: Victor Valbuena Date: Thu, 9 Apr 2026 19:10:46 +0000 Subject: [PATCH 8/9] removing op_name and username --- .env_local_example | 4 ++-- build_scripts/env_local_integration_test | 2 +- doc/code/memory/4_manually_working_with_memory.md | 2 +- pyrit/setup/initializers/airt.py | 10 +++++----- tests/unit/setup/test_airt_initializer.py | 4 +++- 5 files changed, 12 insertions(+), 10 deletions(-) diff --git a/.env_local_example b/.env_local_example index 5126c95c90..626476ca16 100644 --- a/.env_local_example +++ b/.env_local_example @@ -13,11 +13,11 @@ OPENAI_CHAT_MODEL="gpt-4o" ############## # The below GLOBAL_MEMORY_LABELS will be applied to all prompts sent via attacks and can be altered whenever needed. -# Example recommended labels are shown below: `username`, `op_name`. Others that may be useful include: +# Example recommended labels are shown below: `operator`, `operation`. Others that may be useful include: # `language`, `harm_category`, `stage`, or `technique. For the above labels, please stick to the exact spelling, # spacing, and casing for better standardization throughout the database. ############## -GLOBAL_MEMORY_LABELS = {"username": "username"} +GLOBAL_MEMORY_LABELS = {"operator": "operator", "operation": "operation"} ############## # Set optional OPENAI_CHAT_ADDITIONAL_REQUEST_HEADERS to include additional HTTP headers in a dictionary format for API requests, e.g., {'key1': 'value1'}. diff --git a/build_scripts/env_local_integration_test b/build_scripts/env_local_integration_test index d9873e4354..b61cfe7d20 100644 --- a/build_scripts/env_local_integration_test +++ b/build_scripts/env_local_integration_test @@ -22,7 +22,7 @@ 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} -GLOBAL_MEMORY_LABELS={"username": "integration-test", "op_name": "integration-test"} +GLOBAL_MEMORY_LABELS={"operator": "integration-test", "operation": "integration-test"} ############## # Set optional OPENAI_CHAT_ADDITIONAL_REQUEST_HEADERS to include additional HTTP headers in a dictionary format for API requests, e.g., {'key1': 'value1'}. diff --git a/doc/code/memory/4_manually_working_with_memory.md b/doc/code/memory/4_manually_working_with_memory.md index 8c27665f62..87a5831d92 100644 --- a/doc/code/memory/4_manually_working_with_memory.md +++ b/doc/code/memory/4_manually_working_with_memory.md @@ -32,7 +32,7 @@ This is especially nice with scoring. There are countless ways to do this, but t ![scoring_2.png](../../../assets/scoring_3_pivot.png) ## Using AzureSQL Query Editor to Query and Export Data -If you are using an AzureSQL Database, you can use the Query Editor to run SQL queries to retrieve desired data. Memory labels (`labels`) may be an especially useful column to query on for finding data pertaining to a specific operation, user, harm_category, etc. Memory labels are a free-from dictionary for tagging prompts with whatever information you'd like (e.g. `op_name`, `username`, `harm_category`). (For more information on memory labels, see the [Memory Labels Guide](../memory/5_memory_labels.ipynb).) An example is shown below: +If you are using an AzureSQL Database, you can use the Query Editor to run SQL queries to retrieve desired data. Memory labels (`labels`) may be an especially useful column to query on for finding data pertaining to a specific operation, user, harm_category, etc. Memory labels are a free-from dictionary for tagging prompts with whatever information you'd like (e.g. `operation`, `operator`, `harm_category`). (For more information on memory labels, see the [Memory Labels Guide](../memory/5_memory_labels.ipynb).) An example is shown below: 1. Write a SQL query in the Query Editor. You can either write these manually or use the "Open Query" option to load one in. The image below shows a query that gathers prompt entries with their corresponding scores for a specific operation (using the `labels` column) with a "float_scale" `score_type`. diff --git a/pyrit/setup/initializers/airt.py b/pyrit/setup/initializers/airt.py index 251ff7941d..c899797b92 100644 --- a/pyrit/setup/initializers/airt.py +++ b/pyrit/setup/initializers/airt.py @@ -111,7 +111,7 @@ async def initialize_async(self) -> None: 3. Adversarial target configurations 4. Default values for all attack types """ - # Ensure op_name, username, and email are populated from GLOBAL_MEMORY_LABELS. + # Ensure operator, operation, and email are populated from GLOBAL_MEMORY_LABELS. self._validate_operation_fields() # Get environment variables (validated by validate() method) @@ -292,10 +292,10 @@ def _validate_operation_fields(self) -> None: raw_labels = os.environ.get("GLOBAL_MEMORY_LABELS") labels = dict(json.loads(raw_labels)) if raw_labels else {} - if "username" not in labels: - labels["username"] = data["operator"] + if "operator" not in labels: + labels["operator"] = data["operator"] - if "op_name" not in labels: - labels["op_name"] = data["operation"] + if "operation" not in labels: + labels["operation"] = data["operation"] os.environ["GLOBAL_MEMORY_LABELS"] = json.dumps(labels) diff --git a/tests/unit/setup/test_airt_initializer.py b/tests/unit/setup/test_airt_initializer.py index 5f059302c9..1cccab0678 100644 --- a/tests/unit/setup/test_airt_initializer.py +++ b/tests/unit/setup/test_airt_initializer.py @@ -56,7 +56,9 @@ def setup_method(self) -> None: os.environ["AZURE_CONTENT_SAFETY_API_ENDPOINT"] = "https://test-safety.cognitiveservices.azure.com" os.environ["AZURE_SQL_DB_CONNECTION_STRING"] = "Server=test.database.windows.net;Database=testdb" os.environ["AZURE_STORAGE_ACCOUNT_DB_DATA_CONTAINER_URL"] = "https://teststorage.blob.core.windows.net/data" - os.environ["GLOBAL_MEMORY_LABELS"] = '{"op_name": "test_op", "username": "test_user", "email": "test@test.com"}' + os.environ["GLOBAL_MEMORY_LABELS"] = ( + '{"operation": "test_op", "operator": "test_user", "email": "test@test.com"}' + ) # Clean up globals for attr in [ "default_converter_target", From 5ba739cf62ab0b17619bae9a88706ad7dd80dfd5 Mon Sep 17 00:00:00 2001 From: Victor Valbuena Date: Thu, 9 Apr 2026 21:33:51 +0000 Subject: [PATCH 9/9] . --- tests/unit/setup/test_airt_initializer.py | 23 +++++++++++------------ 1 file changed, 11 insertions(+), 12 deletions(-) diff --git a/tests/unit/setup/test_airt_initializer.py b/tests/unit/setup/test_airt_initializer.py index 1cccab0678..61f74cbe57 100644 --- a/tests/unit/setup/test_airt_initializer.py +++ b/tests/unit/setup/test_airt_initializer.py @@ -3,8 +3,7 @@ import os import sys -from pathlib import Path -from unittest.mock import mock_open, patch +from unittest.mock import patch import pytest import yaml @@ -15,12 +14,10 @@ @pytest.fixture def patch_pyrit_conf(tmp_path): - """Create a temporary .pyrit_conf file and patch _validate_operation_fields to read from it.""" + """Create a temporary .pyrit_conf file and patch DEFAULT_CONFIG_PATH to point to it.""" conf_file = tmp_path / ".pyrit_conf" conf_file.write_text(yaml.dump({"operator": "test_user", "operation": "test_op"})) - with patch.object(Path, "home", return_value=tmp_path): - (tmp_path / ".pyrit").mkdir(exist_ok=True) - (tmp_path / ".pyrit" / ".pyrit_conf").write_text(yaml.dump({"operator": "test_user", "operation": "test_op"})) + with patch("pyrit.setup.initializers.airt.DEFAULT_CONFIG_PATH", conf_file): yield @@ -195,22 +192,24 @@ def test_validate_missing_multiple_env_vars_raises_error(self): assert "AZURE_OPENAI_GPT4O_UNSAFE_CHAT_ENDPOINT" in error_message assert "AZURE_OPENAI_GPT4O_UNSAFE_CHAT_MODEL" in error_message - def test_validate_missing_operator_raises_error(self): + def test_validate_missing_operator_raises_error(self, tmp_path): """Test that _validate_operation_fields raises error when operator is missing from .pyrit_conf.""" - conf_data = yaml.dump({"operation": "test_op"}) + conf_file = tmp_path / ".pyrit_conf" + conf_file.write_text(yaml.dump({"operation": "test_op"})) init = AIRTInitializer() with ( - patch("builtins.open", mock_open(read_data=conf_data)), + patch("pyrit.setup.initializers.airt.DEFAULT_CONFIG_PATH", conf_file), pytest.raises(ValueError, match="operator"), ): init._validate_operation_fields() - def test_validate_missing_operation_raises_error(self): + def test_validate_missing_operation_raises_error(self, tmp_path): """Test that _validate_operation_fields raises error when operation is missing from .pyrit_conf.""" - conf_data = yaml.dump({"operator": "test_user"}) + conf_file = tmp_path / ".pyrit_conf" + conf_file.write_text(yaml.dump({"operator": "test_user"})) init = AIRTInitializer() with ( - patch("builtins.open", mock_open(read_data=conf_data)), + patch("pyrit.setup.initializers.airt.DEFAULT_CONFIG_PATH", conf_file), pytest.raises(ValueError, match="operation"), ): init._validate_operation_fields()