From 2759075af5d2b5e5ec08bb75e6773a7c38ac30fa Mon Sep 17 00:00:00 2001 From: Shivam Sharma <192242476+ShivamSharma43@users.noreply.github.com> Date: Tue, 27 Oct 2026 18:59:15 +0530 Subject: [PATCH] FIX Reject max_examples below 1 in the Garak and ML-VLSU loaders Both loaders checked the cap only after adding an item, so max_examples=0 or a negative value still loaded one seed (Garak) or one text and image pair (ML-VLSU). Raise ValueError in __init__ instead, as the ComicJailbreak loader does. Fixes #3591 Co-Authored-By: Claude Opus 5.5 --- pyrit/datasets/seed_datasets/remote/garak_dataset.py | 5 +++++ .../seed_datasets/remote/vlsu_multimodal_dataset.py | 4 +++- .../datasets/test_garak_package_hallucination_dataset.py | 6 ++++++ tests/unit/datasets/test_vlsu_multimodal_dataset.py | 6 ++++++ 4 files changed, 20 insertions(+), 1 deletion(-) diff --git a/pyrit/datasets/seed_datasets/remote/garak_dataset.py b/pyrit/datasets/seed_datasets/remote/garak_dataset.py index edaff80eee..4ebc5ec50f 100644 --- a/pyrit/datasets/seed_datasets/remote/garak_dataset.py +++ b/pyrit/datasets/seed_datasets/remote/garak_dataset.py @@ -91,7 +91,12 @@ def __init__(self, *, max_examples: int | None = None) -> None: the complete reference list rely on this. A small value is used by the integration tests to keep the multi-million-row package registries fast. + + Raises: + ValueError: If max_examples is less than 1. """ + if max_examples is not None and max_examples < 1: + raise ValueError(f"max_examples must be at least 1 or None, got {max_examples}") self._max_examples = max_examples @property diff --git a/pyrit/datasets/seed_datasets/remote/vlsu_multimodal_dataset.py b/pyrit/datasets/seed_datasets/remote/vlsu_multimodal_dataset.py index d8012b66f8..8db352d6ca 100644 --- a/pyrit/datasets/seed_datasets/remote/vlsu_multimodal_dataset.py +++ b/pyrit/datasets/seed_datasets/remote/vlsu_multimodal_dataset.py @@ -122,8 +122,10 @@ def __init__( or quick validations. Raises: - ValueError: If any of the specified categories are invalid. + ValueError: If any of the specified categories are invalid, or if max_examples is less than 1. """ + if max_examples is not None and max_examples < 1: + raise ValueError(f"max_examples must be at least 1 or None, got {max_examples}") if unsafe_grades is None: unsafe_grades = ["unsafe", "borderline"] self.source = source diff --git a/tests/unit/datasets/test_garak_package_hallucination_dataset.py b/tests/unit/datasets/test_garak_package_hallucination_dataset.py index 27b430bfd9..0e8ed72d8c 100644 --- a/tests/unit/datasets/test_garak_package_hallucination_dataset.py +++ b/tests/unit/datasets/test_garak_package_hallucination_dataset.py @@ -125,6 +125,12 @@ async def test_max_examples_caps_seed_count(): assert [s.value for s in dataset.seeds] == ["pkg0", "pkg1"] +@pytest.mark.parametrize("max_examples", [0, -1]) +def test_max_examples_below_one_raises(max_examples): + with pytest.raises(ValueError, match="max_examples must be at least 1"): + _GarakNpmDataset(max_examples=max_examples) + + async def test_max_examples_none_loads_all(mock_npm_rows): loader = _GarakNpmDataset(max_examples=None) diff --git a/tests/unit/datasets/test_vlsu_multimodal_dataset.py b/tests/unit/datasets/test_vlsu_multimodal_dataset.py index 13b1582f94..206375d471 100644 --- a/tests/unit/datasets/test_vlsu_multimodal_dataset.py +++ b/tests/unit/datasets/test_vlsu_multimodal_dataset.py @@ -54,6 +54,12 @@ def test_init_with_empty_categories_raises(self): with pytest.raises(ValueError, match="`categories` must be a non-empty list"): _VLSUMultimodalDataset(categories=[]) + @pytest.mark.parametrize("max_examples", [0, -1]) + def test_init_with_max_examples_below_one_raises(self, max_examples): + """Test that max_examples below 1 raises ValueError instead of still loading one pair.""" + with pytest.raises(ValueError, match="max_examples must be at least 1"): + _VLSUMultimodalDataset(max_examples=max_examples) + def test_init_with_unsafe_grades(self): """Test initialization with custom unsafe grades.""" dataset = _VLSUMultimodalDataset(unsafe_grades=["unsafe"])