diff --git a/pyrit/datasets/jailbreak/text_jailbreak.py b/pyrit/datasets/jailbreak/text_jailbreak.py index 5937bc6431..2e988fdfe2 100644 --- a/pyrit/datasets/jailbreak/text_jailbreak.py +++ b/pyrit/datasets/jailbreak/text_jailbreak.py @@ -211,20 +211,24 @@ def get_jailbreak_templates(cls, num_templates: int | None = None) -> list[str]: Retrieve all jailbreaks from the JAILBREAK_TEMPLATES_PATH. Args: - num_templates (int, optional): Number of jailbreak templates to return. None to get all. + num_templates (int | None): Number of jailbreak templates to return. None to get all. Returns: list[str]: List of jailbreak template file names. Raises: + ValueError: If ``num_templates`` is provided and is less than 1. ValueError: If no jailbreak templates are found in the jailbreak directory. ValueError: If n is larger than the number of templates that exist. """ + if num_templates is not None and num_templates < 1: + raise ValueError("num_templates must be at least 1") + jailbreak_template_names = sorted(cls._get_template_cache().keys()) if not jailbreak_template_names: raise ValueError("No jailbreak templates found in the jailbreak directory") - if num_templates: + if num_templates is not None: if num_templates > len(jailbreak_template_names): raise ValueError( f"Attempted to pull {num_templates} jailbreaks from a dataset" diff --git a/pyrit/scenario/scenarios/airt/jailbreak.py b/pyrit/scenario/scenarios/airt/jailbreak.py index 7c0957982b..7cdc9c1656 100644 --- a/pyrit/scenario/scenarios/airt/jailbreak.py +++ b/pyrit/scenario/scenarios/airt/jailbreak.py @@ -220,12 +220,16 @@ def set_params_from_args(self, *, args: dict[str, Any]) -> None: args (dict[str, Any]): Raw scenario run parameters. Raises: - ValueError: If ``num_jailbreak_attempts`` is less than one. + ValueError: If ``num_jailbreak_attempts`` is less than one, or if ``num_jailbreaks`` + is provided and is less than one. """ super().set_params_from_args(args=args) num_attempts = self.params["num_jailbreak_attempts"] if num_attempts < 1: raise ValueError("num_jailbreak_attempts must be at least 1") + num_jailbreaks = self.params.get("num_jailbreaks") + if num_jailbreaks is not None and num_jailbreaks < 1: + raise ValueError("num_jailbreaks must be at least 1") @apply_defaults def __init__( diff --git a/tests/unit/datasets/test_jailbreak_text.py b/tests/unit/datasets/test_jailbreak_text.py index cf7ffeebac..0b78f1bf57 100644 --- a/tests/unit/datasets/test_jailbreak_text.py +++ b/tests/unit/datasets/test_jailbreak_text.py @@ -89,6 +89,30 @@ def test_get_jailbreak_templates_includes_subdirectory_templates(): assert len(templates) > top_level_count, "Subdirectory templates should be included in the listing" +@pytest.mark.parametrize("num_templates", [1, 3]) +def test_get_jailbreak_templates_returns_requested_count(num_templates: int) -> None: + all_templates = TextJailBreak.get_jailbreak_templates() + templates = TextJailBreak.get_jailbreak_templates(num_templates=num_templates) + assert len(templates) == num_templates + assert set(templates).issubset(all_templates) + + +@pytest.mark.parametrize("num_templates", [0, -1]) +def test_get_jailbreak_templates_rejects_invalid_num_templates(num_templates: int) -> None: + with pytest.raises(ValueError, match="num_templates must be at least 1"): + TextJailBreak.get_jailbreak_templates(num_templates=num_templates) + + +def test_get_jailbreak_templates_none_returns_all() -> None: + assert TextJailBreak.get_jailbreak_templates(num_templates=None) == sorted(TextJailBreak._get_template_cache()) + + +def test_get_jailbreak_templates_rejects_count_above_population() -> None: + num_templates = len(TextJailBreak.get_jailbreak_templates()) + 1 + with pytest.raises(ValueError, match=f"Attempted to pull {num_templates} jailbreaks"): + TextJailBreak.get_jailbreak_templates(num_templates=num_templates) + + 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/scenario/airt/test_jailbreak.py b/tests/unit/scenario/airt/test_jailbreak.py index c6a68bf2c1..a501c0d050 100644 --- a/tests/unit/scenario/airt/test_jailbreak.py +++ b/tests/unit/scenario/airt/test_jailbreak.py @@ -208,6 +208,13 @@ def test_rejects_non_positive_num_jailbreak_attempts(self, mock_objective_scorer with pytest.raises(ValueError, match="num_jailbreak_attempts must be at least 1"): scenario.set_params_from_args(args={"num_jailbreak_attempts": num_attempts}) + @pytest.mark.parametrize("num_jailbreaks", [0, -1]) + def test_rejects_non_positive_num_jailbreaks(self, mock_objective_scorer, num_jailbreaks: int) -> None: + scenario = Jailbreak(objective_scorer=mock_objective_scorer) + + with pytest.raises(ValueError, match="num_jailbreaks must be at least 1"): + scenario.set_params_from_args(args={"num_jailbreaks": num_jailbreaks}) + async def test_default_draws_random_template_sample( self, mock_objective_target, mock_objective_scorer, mock_memory_seed_groups ):