From f3368ea8ca2b56507c6b5f336e1c1bf97f3b40fa Mon Sep 17 00:00:00 2001 From: Utkarsh Bahuguna Date: Sun, 11 Oct 2026 00:12:35 +0530 Subject: [PATCH 1/2] run only the psychosocial sub-harms whose datasets were picked --- pyrit/scenario/scenarios/airt/psychosocial.py | 28 ++++++++++++----- tests/unit/scenario/airt/test_psychosocial.py | 31 +++++++++++++++++++ 2 files changed, 51 insertions(+), 8 deletions(-) diff --git a/pyrit/scenario/scenarios/airt/psychosocial.py b/pyrit/scenario/scenarios/airt/psychosocial.py index 1f54d65918..82d4b28567 100644 --- a/pyrit/scenario/scenarios/airt/psychosocial.py +++ b/pyrit/scenario/scenarios/airt/psychosocial.py @@ -459,24 +459,36 @@ def _selected_sub_harms(self) -> list[_SubHarm]: """ Resolve the ``sub_harm`` run parameter into the ordered list of sub-harm configs. - Accepts ``None`` / ``"all"`` (both sub-harms) or a single sub-harm name. Order follows - ``_SUB_HARMS`` for deterministic results. + Accepts ``None`` / ``"all"`` (both sub-harms) or a single sub-harm name. Only sub-harms whose + dataset is in the dataset configuration are kept, so selecting a subset of the datasets also + narrows the run. Order follows ``_SUB_HARMS`` for deterministic results. Returns: list[_SubHarm]: The selected sub-harms. Raises: ValueError: If an unknown sub-harm name is requested. + DatasetConstraintError: If none of the requested sub-harms has its dataset configured. """ requested = self.params.get("sub_harm") if not requested or requested == "all": - return list(_SUB_HARMS) - name = str(requested) - if name not in _SUB_HARMS_BY_NAME: - raise ValueError( - f"Unknown psychosocial sub_harm '{name}'. Valid values: {sorted(_SUB_HARMS_BY_NAME)} (or 'all')." + requested_harms = list(_SUB_HARMS) + else: + name = str(requested) + if name not in _SUB_HARMS_BY_NAME: + raise ValueError( + f"Unknown psychosocial sub_harm '{name}'. Valid values: {sorted(_SUB_HARMS_BY_NAME)} (or 'all')." + ) + requested_harms = [_SUB_HARMS_BY_NAME[name]] + configured = set(self._dataset_config.dataset_names) + selected = [harm for harm in requested_harms if harm.dataset_name in configured] + if not selected: + raise DatasetConstraintError( + f"Psychosocial sub_harm '{requested or 'all'}' needs " + f"{', '.join(repr(harm.dataset_name) for harm in requested_harms)}, but the selected datasets are " + f"{sorted(configured)}." ) - return [_SUB_HARMS_BY_NAME[name]] + return selected def _validate_runtime_configuration(self) -> None: config = self._dataset_config diff --git a/tests/unit/scenario/airt/test_psychosocial.py b/tests/unit/scenario/airt/test_psychosocial.py index 816fb42e0d..ec6280f518 100644 --- a/tests/unit/scenario/airt/test_psychosocial.py +++ b/tests/unit/scenario/airt/test_psychosocial.py @@ -621,6 +621,37 @@ async def test_single_sub_harm_hard_binds_dataset_config(self, mock_objective_ta await scenario.initialize_async() assert list(scenario._dataset_config.dataset_names) == ["airt_licensed_therapist"] + async def test_dataset_subset_narrows_the_sub_harms(self, mock_objective_target): + scenario = _scenario_with_mock_scorers() + with _patch_base_seed_groups(_make_seed_groups()): + scenario.set_params_from_args( + args={ + "objective_target": mock_objective_target, + "dataset_config": DatasetAttackConfiguration( + sources=[DatasetSource(name="airt_licensed_therapist")] + ), + } + ) + estimate = await scenario.get_run_size_estimate_async(target_is_configured=True) + await scenario.initialize_async() + assert list(scenario._dataset_config.dataset_names) == ["airt_licensed_therapist"] + assert {a.display_group for a in scenario._atomic_attacks} == {"licensed_therapist"} + assert [dataset.name for dataset in estimate.datasets] == ["airt_licensed_therapist"] + + async def test_sub_harm_without_its_dataset_is_rejected(self, mock_objective_target): + scenario = _scenario_with_mock_scorers() + with pytest.raises(DatasetConstraintError, match="needs 'airt_imminent_crisis'"): + scenario.set_params_from_args( + args={ + "objective_target": mock_objective_target, + "sub_harm": "imminent_crisis", + "dataset_config": DatasetAttackConfiguration( + sources=[DatasetSource(name="airt_licensed_therapist")] + ), + } + ) + await scenario.initialize_async() + async def test_unrelated_source_is_rejected(self, mock_objective_target): scenario = _scenario_with_mock_scorers() with pytest.raises(DatasetConstraintError, match="sub-harm datasets"): From cdddbd0b59b310e093f1282ff3db7335175a0616 Mon Sep 17 00:00:00 2001 From: Utkarsh Bahuguna Date: Sun, 11 Oct 2026 10:43:11 +0530 Subject: [PATCH 2/2] tighten tests: async names, typed args, public paths --- tests/unit/scenario/airt/test_psychosocial.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/tests/unit/scenario/airt/test_psychosocial.py b/tests/unit/scenario/airt/test_psychosocial.py index ec6280f518..4ff02ddef3 100644 --- a/tests/unit/scenario/airt/test_psychosocial.py +++ b/tests/unit/scenario/airt/test_psychosocial.py @@ -621,7 +621,7 @@ async def test_single_sub_harm_hard_binds_dataset_config(self, mock_objective_ta await scenario.initialize_async() assert list(scenario._dataset_config.dataset_names) == ["airt_licensed_therapist"] - async def test_dataset_subset_narrows_the_sub_harms(self, mock_objective_target): + async def test_dataset_subset_narrows_the_sub_harms_async(self, *, mock_objective_target: PromptTarget) -> None: scenario = _scenario_with_mock_scorers() with _patch_base_seed_groups(_make_seed_groups()): scenario.set_params_from_args( @@ -638,7 +638,7 @@ async def test_dataset_subset_narrows_the_sub_harms(self, mock_objective_target) assert {a.display_group for a in scenario._atomic_attacks} == {"licensed_therapist"} assert [dataset.name for dataset in estimate.datasets] == ["airt_licensed_therapist"] - async def test_sub_harm_without_its_dataset_is_rejected(self, mock_objective_target): + async def test_sub_harm_without_its_dataset_is_rejected_async(self, *, mock_objective_target: PromptTarget) -> None: scenario = _scenario_with_mock_scorers() with pytest.raises(DatasetConstraintError, match="needs 'airt_imminent_crisis'"): scenario.set_params_from_args(