diff --git a/pyrit/datasets/jailbreak/text_jailbreak.py b/pyrit/datasets/jailbreak/text_jailbreak.py index 68fce5d5d..6e5083bd4 100644 --- a/pyrit/datasets/jailbreak/text_jailbreak.py +++ b/pyrit/datasets/jailbreak/text_jailbreak.py @@ -222,7 +222,7 @@ def get_jailbreak_templates(cls, num_templates: Optional[int] = None) -> list[st ValueError: If no jailbreak templates are found in the jailbreak directory. ValueError: If n is larger than the number of templates that exist. """ - jailbreak_template_names = [str(f.stem) + ".yaml" for f in JAILBREAK_TEMPLATES_PATH.glob("*.yaml")] + jailbreak_template_names = sorted(cls._get_template_cache().keys()) if not jailbreak_template_names: raise ValueError("No jailbreak templates found in the jailbreak directory") @@ -232,7 +232,7 @@ def get_jailbreak_templates(cls, num_templates: Optional[int] = None) -> list[st f"Attempted to pull {num_templates} jailbreaks from a dataset" f" with only {len(jailbreak_template_names)} jailbreaks!" ) - jailbreak_template_names = random.choices(jailbreak_template_names, k=num_templates) + jailbreak_template_names = random.sample(jailbreak_template_names, k=num_templates) return jailbreak_template_names def get_jailbreak_system_prompt(self) -> str: diff --git a/tests/unit/datasets/test_jailbreak_text.py b/tests/unit/datasets/test_jailbreak_text.py index 658eb0f96..9deb982a2 100644 --- a/tests/unit/datasets/test_jailbreak_text.py +++ b/tests/unit/datasets/test_jailbreak_text.py @@ -7,7 +7,7 @@ import pytest -from pyrit.common.path import DATASETS_PATH +from pyrit.common.path import DATASETS_PATH, JAILBREAK_TEMPLATES_PATH from pyrit.datasets import TextJailBreak from pyrit.models import SeedPrompt @@ -82,6 +82,13 @@ def test_get_file_name_subdirectory(): assert "{{ prompt }}" not in result +def test_get_jailbreak_templates_includes_subdirectory_templates(): + """Subdirectory templates must appear in the listing, not just top-level ones.""" + templates = TextJailBreak.get_jailbreak_templates() + top_level_count = len(list(JAILBREAK_TEMPLATES_PATH.glob("*.yaml"))) + assert len(templates) > top_level_count, "Subdirectory templates should be included in the listing" + + def test_all_templates_render_without_syntax_errors(jailbreak_dir): """Test that all jailbreak templates can be successfully rendered with a test prompt.""" yaml_files = [f for f in jailbreak_dir.rglob("*.yaml") if "multi_parameter" not in f.parts] diff --git a/tests/unit/scenarios/test_jailbreak.py b/tests/unit/scenarios/test_jailbreak.py index 9c4e4f9a1..78f830f8b 100644 --- a/tests/unit/scenarios/test_jailbreak.py +++ b/tests/unit/scenarios/test_jailbreak.py @@ -7,6 +7,8 @@ import pytest +from pyrit.common.path import JAILBREAK_TEMPLATES_PATH +from pyrit.datasets import TextJailBreak from pyrit.executor.attack.single_turn.many_shot_jailbreak import ManyShotJailbreakAttack from pyrit.executor.attack.single_turn.prompt_sending import PromptSendingAttack from pyrit.executor.attack.single_turn.role_play import RolePlayAttack @@ -178,6 +180,19 @@ def test_init_raises_exception_when_both_num_and_which_jailbreaks(self, mock_ran with pytest.raises(ValueError): Jailbreak(num_templates=mock_random_num_templates, jailbreak_names=mock_templates) + def test_init_accepts_subdirectory_jailbreak_names(self, mock_objective_scorer, mock_memory_seed_groups): + """Test that explicit jailbreak names can reference templates stored in subdirectories.""" + # Pick a template that lives in a subdirectory (not top-level) + all_templates = TextJailBreak.get_jailbreak_templates() + top_level_names = {f.name for f in JAILBREAK_TEMPLATES_PATH.glob("*.yaml")} + subdir_templates = [t for t in all_templates if t not in top_level_names] + assert subdir_templates, "Expected at least one subdirectory template to exist" + subdir_name = subdir_templates[0] + + with patch.object(Jailbreak, "_resolve_seed_groups", return_value=mock_memory_seed_groups): + scenario = Jailbreak(objective_scorer=mock_objective_scorer, jailbreak_names=[subdir_name]) + assert scenario._jailbreaks == [subdir_name] + @pytest.mark.asyncio async def test_init_raises_exception_when_no_datasets_available(self, mock_objective_target, mock_objective_scorer): """Test that initialization raises ValueError when datasets are not available in memory."""