diff --git a/pyrit/scenario/scenarios/airt/psychosocial.py b/pyrit/scenario/scenarios/airt/psychosocial.py index 904fcf0a86..99f0c7510b 100644 --- a/pyrit/scenario/scenarios/airt/psychosocial.py +++ b/pyrit/scenario/scenarios/airt/psychosocial.py @@ -460,24 +460,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 e26e4f1c7f..3367a6667f 100644 --- a/tests/unit/scenario/airt/test_psychosocial.py +++ b/tests/unit/scenario/airt/test_psychosocial.py @@ -674,6 +674,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_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( + 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_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( + 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"):