diff --git a/pyrit/datasets/seed_datasets/remote/red_team_social_bias_dataset.py b/pyrit/datasets/seed_datasets/remote/red_team_social_bias_dataset.py index 79a6d9eac6..2480494ed0 100644 --- a/pyrit/datasets/seed_datasets/remote/red_team_social_bias_dataset.py +++ b/pyrit/datasets/seed_datasets/remote/red_team_social_bias_dataset.py @@ -121,6 +121,8 @@ async def _fetch_dataset_async(self, *, cache: bool = True) -> SeedDataset: data_type="text", dataset_name=self.dataset_name, prompt_group_id=group_id, + # Only user turns are kept, and a group with several sequences needs a role on each + role="user", sequence=i, harm_categories=harm_categories, description=description, diff --git a/tests/unit/datasets/test_red_team_social_bias_dataset.py b/tests/unit/datasets/test_red_team_social_bias_dataset.py index f78e53e2b7..2d77faffc2 100644 --- a/tests/unit/datasets/test_red_team_social_bias_dataset.py +++ b/tests/unit/datasets/test_red_team_social_bias_dataset.py @@ -72,6 +72,26 @@ async def test_fetch_dataset_multi_turn_linked(mock_social_bias_data): assert multi_turn_prompts[1].value == "Second turn" +async def test_fetch_dataset_multi_turn_prompts_form_seed_groups(mock_social_bias_data): + loader = _RedTeamSocialBiasDataset() + + with patch.object( + loader, "_fetch_from_huggingface_async", new_callable=AsyncMock, return_value=mock_social_bias_data + ): + dataset = await loader.fetch_dataset_async() + + # Each multi-turn prompt is a user turn, so the group passes SeedGroup's role validation + multi_turn_prompts = [s for s in dataset.seeds if s.prompt_group_id is not None] + assert [p.role for p in multi_turn_prompts] == ["user", "user"] + assert [p.sequence for p in multi_turn_prompts] == [0, 1] + + seed_groups = dataset.seed_groups + assert len(seed_groups) == 2 + multi_turn_group = next(g for g in seed_groups if len(g.prompts) == 2) + assert [p.value for p in multi_turn_group.prompts] == ["First turn", "Second turn"] + assert [p.role for p in multi_turn_group.prompts] == ["user", "user"] + + def test_dataset_name(): loader = _RedTeamSocialBiasDataset() assert loader.dataset_name == "red_team_social_bias"