diff --git a/pyrit/datasets/seed_datasets/remote/comic_jailbreak_dataset.py b/pyrit/datasets/seed_datasets/remote/comic_jailbreak_dataset.py index 0d13f5a24c..c4b4d7d822 100644 --- a/pyrit/datasets/seed_datasets/remote/comic_jailbreak_dataset.py +++ b/pyrit/datasets/seed_datasets/remote/comic_jailbreak_dataset.py @@ -149,17 +149,21 @@ def __init__( source_type: The type of source ('public_url' or 'file'). templates: List of template names to include. If None, all 5 templates are used. max_examples: Maximum number of source goals to render. Each goal produces up to - ``len(templates)`` image+text pairs. If None, all goals are rendered. Useful for + ``len(templates)`` image+text pairs. Goals with no text for any of the selected + templates are skipped and do not count. If None, all goals are rendered. Useful for CI and quick validations where rendering all 300 goals × 5 templates is too slow. Raises: - ValueError: If any template name is invalid. + ValueError: If any template name is invalid, or if max_examples is less than 1. """ self.source = source self.source_type: Literal["public_url", "file"] = source_type self.templates = templates or list(self.TEMPLATE_NAMES) self.max_examples = max_examples + 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}") + invalid = set(self.templates) - set(self.TEMPLATE_NAMES) if invalid: raise ValueError( @@ -224,6 +228,7 @@ async def _fetch_dataset_async(self, *, cache: bool = True) -> SeedDataset: alias_overrides=self.HARM_CATEGORY_ALIAS_OVERRIDES, ) + rendered_templates = 0 for template_name in self.templates: col_name = template_name.capitalize() text_to_render = example.get(col_name, "").strip() @@ -248,6 +253,11 @@ async def _fetch_dataset_async(self, *, cache: bool = True) -> SeedDataset: behavior=example.get("Behavior", ""), ) seeds.extend(pair) + rendered_templates += 1 + + # A goal with no text for the selected templates rendered nothing, so it does not count + if not rendered_templates: + continue processed_goals += 1 if self.max_examples is not None and processed_goals >= self.max_examples: diff --git a/tests/unit/datasets/test_comic_jailbreak_dataset.py b/tests/unit/datasets/test_comic_jailbreak_dataset.py index ec2a01da0a..28933c2aa1 100644 --- a/tests/unit/datasets/test_comic_jailbreak_dataset.py +++ b/tests/unit/datasets/test_comic_jailbreak_dataset.py @@ -224,10 +224,32 @@ async def test_fetch_dataset_respects_max_examples(self): goals = {s.metadata["goal"] for s in dataset.seeds if isinstance(s, SeedPrompt)} assert goals == {"Goal 0", "Goal 1"} + async def test_fetch_dataset_max_examples_skips_goals_without_template_text(self): + """Goals with no text for the selected templates render nothing and do not count toward max_examples.""" + mock_data = [_make_example(Goal=f"Goal {i}", Code="") for i in range(3)] + [ + _make_example(Goal=f"Goal {i}", Code="print('hello')") for i in range(3, 6) + ] + loader = _ComicJailbreakDataset(templates=["code"], max_examples=2) + + with ( + patch.object(loader, "_fetch_from_url", return_value=mock_data), + patch.object(loader, "_fetch_template_async", new_callable=AsyncMock, return_value="/fake/template.png"), + patch.object(loader, "_render_comic_async", new_callable=AsyncMock, return_value="/fake/rendered.png"), + ): + dataset = await loader.fetch_dataset_async(cache=False) + + goals = {s.metadata["goal"] for s in dataset.seeds if isinstance(s, SeedPrompt)} + assert goals == {"Goal 3", "Goal 4"} + def test_init_default_max_examples_is_none(self): loader = _ComicJailbreakDataset() assert loader.max_examples is None + @pytest.mark.parametrize("max_examples", [0, -1]) + def test_init_rejects_max_examples_below_one(self, max_examples): + with pytest.raises(ValueError, match="max_examples must be at least 1"): + _ComicJailbreakDataset(max_examples=max_examples) + class TestComicJailbreakTemplates: """Tests for the COMIC_JAILBREAK_TEMPLATES constant."""