diff --git a/pyrit/executor/attack/single_turn/many_shot_jailbreak.py b/pyrit/executor/attack/single_turn/many_shot_jailbreak.py index 91ffcb6c44..d9b9d29f8d 100644 --- a/pyrit/executor/attack/single_turn/many_shot_jailbreak.py +++ b/pyrit/executor/attack/single_turn/many_shot_jailbreak.py @@ -65,12 +65,12 @@ def __init__( prompt_normalizer (PromptNormalizer, Optional): Normalizer for handling prompts. max_attempts_on_failure (int, Optional): Maximum number of attempts to retry on failure. Defaults to 0. example_count (int): The number of examples to include from many_shot_examples or the Many - Shot Jailbreaking dataset. Defaults to the first 100. + Shot Jailbreaking dataset. Must be a positive integer. Defaults to the first 100. many_shot_examples (list[dict[str, str]], Optional): The many shot jailbreaking examples to use. If not provided, takes the first `example_count` examples from Many Shot Jailbreaking dataset. Raises: - ValueError: If many_shot_examples is empty. + ValueError: If example_count is not a positive integer or many_shot_examples is empty. """ ManyShotJailbreakAttack.validate_constructor_parameters( { @@ -106,12 +106,16 @@ def validate_constructor_parameters(cls, params: dict[str, Any]) -> None: Validate the supplied example selection without loading the default dataset. Raises: - ValueError: If the supplied example selection is empty. + ValueError: If ``example_count`` is not a positive integer or the supplied example + selection is empty. """ super().validate_constructor_parameters(params) examples = params.get("many_shot_examples") count = params.get("example_count", 100) - if count == 0 or (examples is not None and not examples[:count]): + # The count is a slice bound, so a negative value would drop examples from the end. + if isinstance(count, bool) or not isinstance(count, int) or count < 1: + raise ValueError(f"example_count must be a positive integer, got {count!r}.") + if examples is not None and not examples: raise ValueError("Many shot examples must be provided.") async def _perform_async(self, *, context: SingleTurnAttackContext[Any]) -> AttackResult: diff --git a/tests/unit/backend/test_technique_service.py b/tests/unit/backend/test_technique_service.py index e9c6582fe4..ce3cf9c570 100644 --- a/tests/unit/backend/test_technique_service.py +++ b/tests/unit/backend/test_technique_service.py @@ -496,6 +496,7 @@ async def register_then_initialize_async(*, request: RunScenarioRequest, init_kw ("PromptSendingAttack", {"max_attempts_on_failure": -1}), ("ManyShotJailbreakAttack", {"max_attempts_on_failure": -1}), ("ManyShotJailbreakAttack", {"example_count": 0}), + ("ManyShotJailbreakAttack", {"example_count": -1}), ("ManyShotJailbreakAttack", {"many_shot_examples": []}), ("CrescendoAttack", {"max_backtracks": -1}), ("CrescendoAttack", {"max_turns": 0}), diff --git a/tests/unit/executor/attack/single_turn/test_many_shot_jailbreak.py b/tests/unit/executor/attack/single_turn/test_many_shot_jailbreak.py index f56fa46833..eee0962c22 100644 --- a/tests/unit/executor/attack/single_turn/test_many_shot_jailbreak.py +++ b/tests/unit/executor/attack/single_turn/test_many_shot_jailbreak.py @@ -188,6 +188,27 @@ def test_init_raises_error_with_empty_custom_examples(self, mock_from_yaml, mock with pytest.raises(ValueError, match="Many shot examples must be provided"): ManyShotJailbreakAttack(objective_target=mock_objective_target, many_shot_examples=[]) + @pytest.mark.parametrize("example_count", [0, -1, -3, True, 2.5, "5", None]) + @patch("pyrit.executor.attack.single_turn.many_shot_jailbreak.SeedPrompt.from_yaml_file") + def test_init_rejects_example_count_that_is_not_a_positive_integer( + self, mock_from_yaml, mock_objective_target, mock_template, sample_many_shot_examples, example_count + ): + """A negative count would slice examples off the end instead of selecting from the start.""" + mock_from_yaml.return_value = mock_template + + with pytest.raises(ValueError, match="example_count must be a positive integer"): + ManyShotJailbreakAttack( + objective_target=mock_objective_target, + example_count=example_count, + many_shot_examples=sample_many_shot_examples, + ) + + @pytest.mark.parametrize("example_count", [0, -1, True, "5"]) + def test_validate_constructor_parameters_rejects_invalid_example_count(self, example_count): + """Deferred construction (e.g. the REST technique API) must reject the same counts.""" + with pytest.raises(ValueError, match="example_count must be a positive integer"): + ManyShotJailbreakAttack.validate_constructor_parameters({"example_count": example_count}) + @pytest.mark.usefixtures("patch_central_database") class TestManyShotJailbreakAttackParamsType: