From ffad824996602ef0ad572c97b52b529942c0e84f Mon Sep 17 00:00:00 2001 From: Nina Chikanov Date: Tue, 23 Dec 2025 23:11:17 +0000 Subject: [PATCH] Update to Messages in AtomicAttack & add integration tests for scenarios notebooks. Rename other notebook test files to match test_notebooks_*.py format --- pyrit/scenario/core/atomic_attack.py | 51 +++---------------- ...otebooks.py => test_notebooks_executor.py} | 0 .../scenarios/test_notebooks_scenarios.py | 35 +++++++++++++ ..._notebooks.py => test_notebooks_scorer.py} | 0 4 files changed, 41 insertions(+), 45 deletions(-) rename tests/integration/executors/{test_executor_notebooks.py => test_notebooks_executor.py} (100%) create mode 100644 tests/integration/scenarios/test_notebooks_scenarios.py rename tests/integration/score/{test_scorer_notebooks.py => test_notebooks_scorer.py} (100%) diff --git a/pyrit/scenario/core/atomic_attack.py b/pyrit/scenario/core/atomic_attack.py index d4f3f6f60..5300c2eca 100644 --- a/pyrit/scenario/core/atomic_attack.py +++ b/pyrit/scenario/core/atomic_attack.py @@ -23,7 +23,7 @@ SingleTurnAttackContext, ) from pyrit.executor.attack.core.attack_executor import AttackExecutorResult -from pyrit.models import AttackResult, Message, SeedGroup +from pyrit.models import AttackResult, Message logger = logging.getLogger(__name__) @@ -93,7 +93,7 @@ def __init__( attack: AttackStrategy, objectives: List[str], prepended_conversations: Optional[List[List[Message]]] = None, - seed_groups: Optional[List[SeedGroup]] = None, + messages: Optional[List[Message]] = None, custom_prompts: Optional[List[str]] = None, memory_labels: Optional[Dict[str, str]] = None, **attack_execute_params: Any, @@ -109,7 +109,7 @@ def __init__( prepended_conversations (Optional[List[List[Message]]]): Optional list of conversation histories to prepend to each attack execution. This will be used for all objectives. - seed_groups (Optional[List[SeedGroup]]): List of seed groups + messages (Optional[List[Message]]): List of messages for single-turn attacks. Only valid for single-turn attacks. custom_prompts (Optional[List[str]]): List of custom prompts for multi-turn attacks. Only valid for multi-turn attacks. @@ -120,8 +120,6 @@ def __init__( Raises: ValueError: If objectives list is empty. - TypeError: If seed_groups is provided for multi-turn attacks or - custom_prompts is provided for single-turn attacks. """ self.atomic_attack_name = atomic_attack_name @@ -134,16 +132,9 @@ def __init__( # Determine context type once during initialization self._context_type: Literal["single_turn", "multi_turn", "unknown"] = self._determine_context_type(attack) - # Validate attack context type and parameters - self._validate_parameters( - seed_groups=seed_groups, - custom_prompts=custom_prompts, - ) - self._objectives = objectives self._prepended_conversations = prepended_conversations - self._seed_groups = seed_groups - self._custom_prompts = custom_prompts + self._messages = messages self._memory_labels = memory_labels or {} self._attack_execute_params = attack_execute_params @@ -181,36 +172,6 @@ def _determine_context_type(self, attack: AttackStrategy) -> Literal["single_tur return "multi_turn" return "unknown" - def _validate_parameters( - self, - *, - seed_groups: Optional[List[SeedGroup]], - custom_prompts: Optional[List[str]], - ) -> None: - """ - Validate that parameters match the attack context type. - - Args: - seed_groups (Optional[List[SeedGroup]]): Seed groups parameter. - custom_prompts (Optional[List[str]]): Custom prompts parameter. - - Raises: - TypeError: If parameters don't match the attack context type. - """ - # Validate seed_groups is only used with single-turn attacks - if seed_groups is not None and self._context_type != "single_turn": - raise TypeError( - f"seed_groups can only be used with single-turn attacks. " - f"Attack {self._attack.__class__.__name__} uses {self._context_type} context" - ) - - # Validate custom_prompts is only used with multi-turn attacks - if custom_prompts is not None and self._context_type != "multi_turn": - raise TypeError( - f"custom_prompts can only be used with multi-turn attacks. " - f"Attack {self._attack.__class__.__name__} uses {self._context_type} context" - ) - async def run_async( self, *, @@ -268,7 +229,7 @@ async def run_async( results = await executor.execute_single_turn_attacks_async( attack=self._attack, objectives=self._objectives, - seed_groups=self._seed_groups, + messages=self._messages, prepended_conversations=prepended_conversations, memory_labels=merged_memory_labels, return_partial_on_failure=return_partial_on_failure, @@ -278,7 +239,7 @@ async def run_async( results = await executor.execute_multi_turn_attacks_async( attack=self._attack, objectives=self._objectives, - custom_prompts=self._custom_prompts, + messages=self._messages, prepended_conversations=prepended_conversations, memory_labels=merged_memory_labels, return_partial_on_failure=return_partial_on_failure, diff --git a/tests/integration/executors/test_executor_notebooks.py b/tests/integration/executors/test_notebooks_executor.py similarity index 100% rename from tests/integration/executors/test_executor_notebooks.py rename to tests/integration/executors/test_notebooks_executor.py diff --git a/tests/integration/scenarios/test_notebooks_scenarios.py b/tests/integration/scenarios/test_notebooks_scenarios.py new file mode 100644 index 000000000..4e0004b8d --- /dev/null +++ b/tests/integration/scenarios/test_notebooks_scenarios.py @@ -0,0 +1,35 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT license. + +import os +import pathlib + +import nbformat +import pytest +from nbconvert.preprocessors import ExecutePreprocessor + +from pyrit.common import path + +nb_directory_path = pathlib.Path(path.DOCS_CODE_PATH, "scenarios").resolve() +folder_paths = [d for d in nb_directory_path.iterdir() if d.is_dir()] + + +@pytest.mark.parametrize( + "file_path", + [ + os.path.join(dir_path, file) + for dir_path in folder_paths + for file in os.listdir(dir_path) + if file.endswith(".ipynb") + ], +) +def test_execute_notebooks(file_path): + nb_path = pathlib.Path(file_path).resolve() + print(nb_path) + with open(nb_path, encoding="utf-8") as f: + nb = nbformat.read(f, as_version=4) + + ep = ExecutePreprocessor(timeout=900) + + # Execute notebook, test will throw exception if any cell fails + ep.preprocess(nb, {"metadata": {"path": nb_path.parent}}) diff --git a/tests/integration/score/test_scorer_notebooks.py b/tests/integration/score/test_notebooks_scorer.py similarity index 100% rename from tests/integration/score/test_scorer_notebooks.py rename to tests/integration/score/test_notebooks_scorer.py