diff --git a/pyrit/backend/services/scenario_run_service.py b/pyrit/backend/services/scenario_run_service.py index b770ac9722..d35b14b688 100644 --- a/pyrit/backend/services/scenario_run_service.py +++ b/pyrit/backend/services/scenario_run_service.py @@ -28,7 +28,7 @@ TargetRegistry, ) from pyrit.scenario import Scenario -from pyrit.scenario.core import DatasetAttackConfiguration +from pyrit.scenario.core import DatasetAttackConfiguration, build_dataset_filters if TYPE_CHECKING: from pyrit.prompt_converter import PromptConverter @@ -302,15 +302,21 @@ def _build_init_kwargs( if request.labels: init_kwargs["memory_labels"] = request.labels + dataset_filters = build_dataset_filters(parameters=request.dataset_parameters) + # Resolve strategies and dataset config from a temporary instance of the # scenario. The downstream _initialize_scenario_async builds its own # instance (so scenario_result_id can be passed), so this is a cheap # throwaway used only for introspection. Introspection is required - # whenever the caller wants to override strategies, dataset names, or - # the sample cap, because each of those needs the scenario's own - # strategy enum or dataset-config subclass to be resolved correctly. + # whenever the caller wants to override strategies, dataset names, the + # sample cap, or dataset filters, because each of those needs the + # scenario's own strategy enum or dataset-config subclass to be resolved + # correctly. needs_introspection = ( - bool(request.strategies) or bool(request.dataset_names) or request.max_dataset_size is not None + bool(request.strategies) + or bool(request.dataset_names) + or request.max_dataset_size is not None + or bool(dataset_filters) ) if not needs_introspection: return init_kwargs @@ -334,7 +340,7 @@ def _build_init_kwargs( if strategy_converters: init_kwargs["strategy_converters"] = strategy_converters - if request.dataset_names or request.max_dataset_size is not None: + if request.dataset_names or request.max_dataset_size is not None or dataset_filters: default_config = introspection_instance._default_dataset_config if request.dataset_names: @@ -345,6 +351,7 @@ def _build_init_kwargs( init_kwargs["dataset_config"] = default_config_class( dataset_names=request.dataset_names, max_dataset_size=request.max_dataset_size, + filters=dataset_filters or None, ) except TypeError as exc: # The subclass __init__ takes extra required kwargs we cannot @@ -354,7 +361,7 @@ def _build_init_kwargs( # define a no-extra-required-args constructor or surface the # incompatibility through their own initialize_async validation. logger.warning( - "Cannot construct %s(dataset_names=..., max_dataset_size=...) (%s). " + "Cannot construct %s(dataset_names=..., max_dataset_size=..., filters=...) (%s). " "Falling back to a generic DatasetAttackConfiguration; scenario-specific " "dataset-config behavior may be lost.", default_config_class.__name__, @@ -363,12 +370,17 @@ def _build_init_kwargs( init_kwargs["dataset_config"] = DatasetAttackConfiguration( dataset_names=request.dataset_names, max_dataset_size=request.max_dataset_size, + filters=dataset_filters or None, ) - elif request.max_dataset_size is not None: + else: # Reuse the scenario's default dataset config (preserves subtype + # the scenario's own default dataset names) and override only the - # sample cap. Safe because the introspection instance is throwaway. - default_config.max_dataset_size = request.max_dataset_size + # sample cap and/or filters. Safe because the introspection instance + # is throwaway. + if request.max_dataset_size is not None: + default_config.max_dataset_size = request.max_dataset_size + if dataset_filters: + default_config.update_filters(filters=dataset_filters) init_kwargs["dataset_config"] = default_config return init_kwargs diff --git a/pyrit/cli/_cli_args.py b/pyrit/cli/_cli_args.py index 21fc0e37db..4fd1966a5c 100644 --- a/pyrit/cli/_cli_args.py +++ b/pyrit/cli/_cli_args.py @@ -249,9 +249,44 @@ def parse_memory_labels(json_string: str) -> dict[str, str]: return labels +def parse_dataset_parameter(arg: str) -> tuple[str, str]: + """ + Parse a single ``KEY=VALUE`` dataset-parameter token from the CLI. + + Note: The ``arg`` parameter is positional (not keyword-only) so it can be used directly + as an argparse ``type=`` callable and an ``_ArgSpec`` parser. This mirrors + ``_parse_initializer_arg`` and is an intentional exception to the keyword-only style rule + for argparse compatibility. + + Args: + arg (str): The raw ``KEY=VALUE`` token. + + Returns: + tuple[str, str]: The (key, value) pair. The value keeps its raw string form so the + server can coerce and validate it. + + Raises: + ValueError: If the token is not in ``KEY=VALUE`` form or the key is empty. Argparse + converts this into a clean CLI error; the shell catches it directly. + """ + if "=" not in arg: + raise ValueError(f"Dataset parameter must be in KEY=VALUE form, got: {arg!r}") + key, _, value = arg.partition("=") + key = key.strip() + if not key: + raise ValueError(f"Dataset parameter key cannot be empty in: {arg!r}") + return key, value + + # --------------------------------------------------------------------------- # Shared argument help text # --------------------------------------------------------------------------- + +# Dataset-filter keys advertised in --dataset-parameters help. Kept as a static list here (this +# module is on the fast --help path and stays import-light); a unit test asserts it matches the +# authoritative registry in pyrit.scenario.core.dataset_configuration. +_ADVERTISED_DATASET_FILTER_KEYS: tuple[str, ...] = ("harm_categories", "data_types") + ARG_HELP = { "config_file": CONFIG_FILE_HELP, "initializers": ( @@ -275,6 +310,10 @@ def parse_memory_labels(json_string: str) -> dict[str, str]: "Creates a new dataset config; fetches all items unless --max-dataset-size is also specified", "max_dataset_size": "Maximum number of items to use from the dataset (must be >= 1). " "Limits new datasets if --dataset-names provided, otherwise overrides scenario's default limit", + "dataset_parameters": "Dataset seed filters as KEY=VALUE tokens " + "(e.g., harm_categories=cyber data_types=text). Accepted keys: " + ", ".join(_ADVERTISED_DATASET_FILTER_KEYS) + ". " + "Keys filter seeds before sizing. " + "List values may be comma-separated (e.g., harm_categories=cyber,violence)", "target": "Name of a registered target from the TargetRegistry to use as the objective target. " "Targets are registered by initializers (e.g., 'target' initializer). " "Use --list-targets to see available target names after initializers have run", @@ -409,6 +448,12 @@ class _ArgSpec: result_key="max_dataset_size", parser=lambda v: validate_integer(v, name="--max-dataset-size", min_value=1), ) +_DATASET_PARAMETERS_ARG = _ArgSpec( + flags=["--dataset-parameters"], + result_key="dataset_parameters", + multi_value=True, + parser=parse_dataset_parameter, +) _TARGET_ARG = _ArgSpec( flags=["--target"], result_key="target", @@ -422,6 +467,7 @@ class _ArgSpec: _MEMORY_LABELS_ARG, _DATASET_NAMES_ARG, _MAX_DATASET_SIZE_ARG, + _DATASET_PARAMETERS_ARG, _TARGET_ARG, ] diff --git a/pyrit/cli/pyrit_scan.py b/pyrit/cli/pyrit_scan.py index 2b0f3b1a7c..6897d26d3f 100644 --- a/pyrit/cli/pyrit_scan.py +++ b/pyrit/cli/pyrit_scan.py @@ -24,6 +24,7 @@ _parse_initializer_arg, build_parameters_from_api, non_negative_int, + parse_dataset_parameter, positive_int, validate_log_level_argparse, ) @@ -267,6 +268,13 @@ def _build_base_parser(*, add_help: bool = True) -> ArgumentParser: type=positive_int, help=ARG_HELP["max_dataset_size"], ) + run_group.add_argument( + "--dataset-parameters", + type=parse_dataset_parameter, + nargs="+", + metavar="KEY=VALUE", + help=ARG_HELP["dataset_parameters"], + ) return parser @@ -644,6 +652,8 @@ def _build_run_request(*, parsed_args: Namespace, scenario_name: str) -> RunScen kwargs["dataset_names"] = parsed_args.dataset_names if parsed_args.max_dataset_size is not None: kwargs["max_dataset_size"] = parsed_args.max_dataset_size + if parsed_args.dataset_parameters: + kwargs["dataset_parameters"] = dict(parsed_args.dataset_parameters) if parsed_args.memory_labels: kwargs["labels"] = parse_memory_labels(json_string=parsed_args.memory_labels) diff --git a/pyrit/cli/pyrit_shell.py b/pyrit/cli/pyrit_shell.py index 52d303f9b3..ff9807a257 100644 --- a/pyrit/cli/pyrit_shell.py +++ b/pyrit/cli/pyrit_shell.py @@ -432,6 +432,8 @@ def do_run(self, line: str) -> None: request_kwargs["dataset_names"] = args["dataset_names"] if args.get("max_dataset_size") is not None: request_kwargs["max_dataset_size"] = args["max_dataset_size"] + if args.get("dataset_parameters"): + request_kwargs["dataset_parameters"] = dict(args["dataset_parameters"]) if args.get("memory_labels"): request_kwargs["labels"] = args["memory_labels"] diff --git a/pyrit/models/catalog/scenario.py b/pyrit/models/catalog/scenario.py index 5c510fab15..956482541f 100644 --- a/pyrit/models/catalog/scenario.py +++ b/pyrit/models/catalog/scenario.py @@ -50,6 +50,12 @@ class RunScenarioRequest(BaseModel): strategies: list[str] | None = Field(None, description="Strategy names to use (uses scenario default if omitted)") dataset_names: list[str] | None = Field(None, description="Dataset names to use (uses scenario default if omitted)") max_dataset_size: int | None = Field(None, ge=1, description="Maximum items per dataset") + dataset_parameters: dict[str, Any] | None = Field( + None, + description=( + "Dataset seed filters keyed by field, applied before sampling. Accepted keys: harm_categories, data_types." + ), + ) max_concurrency: int = Field(10, ge=1, le=100, description="Maximum concurrent operations") max_retries: int = Field(0, ge=0, le=20, description="Maximum retry attempts on failure") labels: dict[str, str] | None = Field(None, description="Labels to attach to memory entries") diff --git a/pyrit/scenario/core/__init__.py b/pyrit/scenario/core/__init__.py index 3faa4a4e52..224c3d3eb0 100644 --- a/pyrit/scenario/core/__init__.py +++ b/pyrit/scenario/core/__init__.py @@ -15,6 +15,7 @@ DatasetConstraintError, DatasetSourceKind, ResolvedDataset, + build_dataset_filters, require_nonempty, ) from pyrit.scenario.core.scenario import BaselineAttackPolicy, Scenario @@ -34,6 +35,7 @@ "INLINE_DATASET_NAME", "Parameter", "ResolvedDataset", + "build_dataset_filters", "require_nonempty", "Scenario", "ScenarioCompositeStrategy", diff --git a/pyrit/scenario/core/dataset_configuration.py b/pyrit/scenario/core/dataset_configuration.py index 020a5a6964..390ae25d19 100644 --- a/pyrit/scenario/core/dataset_configuration.py +++ b/pyrit/scenario/core/dataset_configuration.py @@ -31,7 +31,7 @@ from dataclasses import dataclass from enum import Enum from functools import cached_property -from typing import TYPE_CHECKING, TypeVar +from typing import TYPE_CHECKING, Any, TypeVar from pyrit.common.deprecation import print_deprecation_message from pyrit.memory import CentralMemory @@ -226,6 +226,60 @@ def _validate(resolved: ResolvedDataset) -> None: return _validate +# Authoritative set of dataset seed filters exposed via ``--dataset-parameters``. Each entry +# is used verbatim as a ``MemoryInterface.get_seeds`` keyword argument, so a filter key IS the +# get_seeds kwarg. Every exposed filter must be a list-valued (Sequence) get_seeds parameter -- +# a guard test enforces this and keeps the CLI's advertised keys in sync. Adding a filterable +# field is a one-line change here. +DATASET_FILTERS: frozenset[str] = frozenset({"harm_categories", "data_types"}) + + +def _coerce_filter_values(value: Any) -> list[str]: + """ + Coerce a single filter value into a list of non-empty strings. + + Accepts a list/tuple as-is (stringified) or splits a comma-separated string. + + Args: + value (Any): The raw filter value (scalar, comma-separated string, or sequence). + + Returns: + list[str]: The cleaned list of values. + """ + if isinstance(value, (list, tuple)): + return [str(item).strip() for item in value if str(item).strip()] + return [part.strip() for part in str(value).split(",") if part.strip()] + + +def build_dataset_filters(*, parameters: dict[str, Any] | None) -> dict[str, Any]: + """ + Map a user-facing dataset-parameter bag onto ``MemoryInterface.get_seeds`` filters. + + Each key must be one of ``DATASET_FILTERS`` (the exposed ``get_seeds`` kwargs); its value is + coerced into a list (a single token or a comma-separated string). Any other key is rejected + loudly. + + Args: + parameters (dict[str, Any] | None): The raw dataset-parameter bag (e.g. from the CLI + ``--dataset-parameters`` flag), keyed by ``get_seeds`` kwarg name. + + Returns: + dict[str, Any]: The filters ready to pass as ``**filters`` to ``get_seeds``. + + Raises: + ValueError: If a key is not present in ``DATASET_FILTERS``. + """ + filters: dict[str, Any] = {} + for key, raw_value in (parameters or {}).items(): + if key not in DATASET_FILTERS: + raise ValueError(f"Unknown dataset parameter '{key}'. Allowed: {', '.join(sorted(DATASET_FILTERS))}.") + values = _coerce_filter_values(raw_value) + if values: + filters[key] = values + + return filters + + def restrict_dataset_names(allowed: set[str]) -> Callable[[ResolvedDataset], None]: """ Build a validator that requires every contributing dataset name to be in ``allowed``. @@ -293,6 +347,7 @@ def __init__( seed_groups: list[SeedGroup] | None = None, dataset_names: list[str] | None = None, max_dataset_size: int | None = None, + filters: dict[str, Any] | None = None, validators: Sequence[Callable[[ResolvedDataset], None]] | None = None, auto_fetch: bool = True, ) -> None: @@ -306,6 +361,9 @@ def __init__( dataset_names (list[str] | None): Names of datasets to load from memory. max_dataset_size (int | None): If set, randomly samples up to this many items from the resolved dataset (without replacement). + filters (dict[str, Any] | None): Filters passed to ``MemoryInterface.get_seeds`` + when resolving named datasets (e.g. ``{"harm_categories": ["cyber"]}``). + Applied before ``max_dataset_size`` sampling; ignored for inline seeds. validators (Sequence[Callable[[ResolvedDataset], None]] | None): Constraint callbacks run against the resolved dataset; each raises on violation. These are appended to the subclass's ``_default_validators``. @@ -331,6 +389,7 @@ def __init__( self._seed_groups = list(seed_groups) if seed_groups is not None else None self._dataset_names = list(dataset_names) if dataset_names is not None else None self.max_dataset_size = max_dataset_size + self._filters: dict[str, Any] = dict(filters or {}) self._validators: list[Callable[[ResolvedDataset], None]] = [ *self._default_validators(), *(list(validators) if validators else []), @@ -389,6 +448,28 @@ def source_kind(self) -> DatasetSourceKind: return DatasetSourceKind.INLINE return DatasetSourceKind.MEMORY + @property + def filters(self) -> dict[str, Any]: + """ + The ``get_seeds`` filters applied when resolving named datasets. + + Returns: + dict[str, Any]: A copy of the configured filters. + """ + return dict(self._filters) + + def update_filters(self, *, filters: dict[str, Any]) -> None: + """ + Merge additional ``get_seeds`` filters into this configuration (run-time override). + + Used when a run overrides dataset selection without rebuilding the configuration -- + the provided filters take precedence over any already configured with the same key. + + Args: + filters (dict[str, Any]): Filters to merge, keyed by ``get_seeds`` kwarg name. + """ + self._filters = {**self._filters, **filters} + # ========================================================================= # Resolution helpers # ========================================================================= @@ -426,7 +507,7 @@ async def _collect_seeds_for_dataset_async(self, *, dataset_name: str) -> list[S DatasetConstraintError: If the dataset yields no seeds even after auto-fetch, or if auto-fetch itself fails (the provider error is chained as the cause). """ - found = list(self._memory.get_seeds(dataset_name=dataset_name)) + found = list(self._memory.get_seeds(dataset_name=dataset_name, **self._filters)) if not found and self._auto_fetch: try: await self._fetch_dataset_async(dataset_name=dataset_name) @@ -434,8 +515,12 @@ async def _collect_seeds_for_dataset_async(self, *, dataset_name: str) -> list[S raise DatasetConstraintError( f"Dataset '{dataset_name}' could not be loaded: auto-fetch from the registered provider failed." ) from exc - found = list(self._memory.get_seeds(dataset_name=dataset_name)) + found = list(self._memory.get_seeds(dataset_name=dataset_name, **self._filters)) if not found: + if self._filters and self._memory.get_seeds(dataset_name=dataset_name): + raise DatasetConstraintError( + f"Dataset '{dataset_name}' has seeds, but none match the configured filters {self._filters}." + ) hint = ( "auto-fetch from the registered provider did not populate it" if self._auto_fetch @@ -867,6 +952,7 @@ def per_dataset( dataset_names: Sequence[str], max_dataset_size: int | None = None, auto_fetch: bool = True, + filters: dict[str, Any] | None = None, validators: Sequence[Callable[[ResolvedDataset], None]] | None = None, ) -> CompoundDatasetAttackConfiguration: """ @@ -879,6 +965,7 @@ def per_dataset( dataset_names (Sequence[str]): The dataset names; one child is built per name. max_dataset_size (int | None): Per-dataset cap applied to each child. auto_fetch (bool): Passed to each child (fetch missing datasets into memory). + filters (dict[str, Any] | None): ``get_seeds`` filters applied to each child. validators (Sequence[Callable[[ResolvedDataset], None]] | None): Applied to each child. Returns: @@ -895,6 +982,7 @@ def per_dataset( dataset_names=[name], max_dataset_size=max_dataset_size, auto_fetch=auto_fetch, + filters=filters, validators=validators, ) for name in dataset_names @@ -928,6 +1016,20 @@ def source_kind(self) -> DatasetSourceKind: return DatasetSourceKind.INLINE return DatasetSourceKind.MEMORY + def update_filters(self, *, filters: dict[str, Any]) -> None: + """ + Merge filters into the compound and propagate them to every child configuration. + + The children run the actual ``get_seeds`` queries, so run-time filter overrides must + reach each child to take effect. + + Args: + filters (dict[str, Any]): Filters to merge, keyed by ``get_seeds`` kwarg name. + """ + super().update_filters(filters=filters) + for child in self._configurations: + child.update_filters(filters=filters) + async def get_seed_attack_groups_async(self) -> list[SeedAttackGroup]: """ Concatenate every child's flat result, then validate and apply the global cap. diff --git a/tests/unit/backend/test_scenario_run_service.py b/tests/unit/backend/test_scenario_run_service.py index 89180bfd80..f6c2017069 100644 --- a/tests/unit/backend/test_scenario_run_service.py +++ b/tests/unit/backend/test_scenario_run_service.py @@ -65,6 +65,7 @@ def _make_request( scenario_result_id: str | None = None, dataset_names: list[str] | None = None, max_dataset_size: int | None = None, + dataset_parameters: dict[str, Any] | None = None, ) -> RunScenarioRequest: """Create a RunScenarioRequest for testing.""" return RunScenarioRequest( @@ -75,6 +76,7 @@ def _make_request( scenario_result_id=scenario_result_id, dataset_names=dataset_names, max_dataset_size=max_dataset_size, + dataset_parameters=dataset_parameters, ) @@ -389,6 +391,56 @@ def __init__(self, *, required_extra: str, **kwargs: Any) -> None: for record in caplog.records ) + async def test_start_run_dataset_parameters_filters_new_config(self, mock_all_registries) -> None: + """``dataset_parameters`` with ``dataset_names`` builds a config carrying the filters.""" + + class _MarkerDatasetConfiguration(DatasetConfiguration): + pass + + scenario_instance = mock_all_registries["scenario_instance"] + scenario_instance._default_dataset_config = _MarkerDatasetConfiguration(dataset_names=["original"]) + + service = ScenarioRunService() + await service.start_run_async( + request=_make_request( + dataset_names=["custom"], + max_dataset_size=7, + dataset_parameters={"harm_categories": "cyber"}, + ) + ) + + init_call = mock_all_registries["scenario_registry"].create_and_initialize_async.await_args + built_config = init_call.kwargs["dataset_config"] + assert built_config.dataset_names == ["custom"] + assert built_config.max_dataset_size == 7 + assert built_config.filters == {"harm_categories": ["cyber"]} + + async def test_start_run_dataset_parameters_updates_default_config(self, mock_all_registries) -> None: + """``dataset_parameters`` with no ``dataset_names`` merges filters into the default config.""" + default_config = DatasetAttackConfiguration(dataset_names=["original"]) + scenario_instance = mock_all_registries["scenario_instance"] + scenario_instance._default_dataset_config = default_config + + service = ScenarioRunService() + await service.start_run_async(request=_make_request(dataset_parameters={"harm_categories": "cyber"})) + + init_call = mock_all_registries["scenario_registry"].create_and_initialize_async.await_args + built_config = init_call.kwargs["dataset_config"] + assert built_config is default_config + assert built_config.filters == {"harm_categories": ["cyber"]} + + async def test_start_run_dataset_parameters_rejects_max_dataset_size(self, mock_all_registries) -> None: + """``max_dataset_size`` in ``dataset_parameters`` is rejected as an unknown filter key.""" + service = ScenarioRunService() + with pytest.raises(ValueError, match="Unknown dataset parameter 'max_dataset_size'"): + await service.start_run_async(request=_make_request(dataset_parameters={"max_dataset_size": "9"})) + + async def test_start_run_dataset_parameters_unknown_key_raises(self, mock_all_registries) -> None: + """An unknown dataset parameter surfaces as a validation error.""" + service = ScenarioRunService() + with pytest.raises(ValueError, match="Unknown dataset parameter 'bogus'"): + await service.start_run_async(request=_make_request(dataset_parameters={"bogus": "x"})) + async def test_start_run_dataset_names_introspection_failure_raises(self, mock_memory) -> None: """Passing ``dataset_names`` against a non-no-arg-instantiable scenario fails fast.""" # Mirrors test_start_run_scenario_not_no_arg_instantiable_raises but for the dataset_names path. diff --git a/tests/unit/cli/test_dataset_filter_help.py b/tests/unit/cli/test_dataset_filter_help.py new file mode 100644 index 0000000000..e6d64c621b --- /dev/null +++ b/tests/unit/cli/test_dataset_filter_help.py @@ -0,0 +1,12 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT license. + +"""Guard that the CLI's advertised dataset-filter keys stay in sync with the resolver.""" + +from pyrit.cli._cli_args import _ADVERTISED_DATASET_FILTER_KEYS +from pyrit.scenario.core.dataset_configuration import DATASET_FILTERS + + +def test_cli_advertised_filters_match_dataset_configuration() -> None: + # The static CLI list must equal the exact filter kwargs the resolver accepts. + assert set(_ADVERTISED_DATASET_FILTER_KEYS) == set(DATASET_FILTERS) diff --git a/tests/unit/cli/test_pyrit_scan.py b/tests/unit/cli/test_pyrit_scan.py index 7ec86a0814..7d76a81d5b 100644 --- a/tests/unit/cli/test_pyrit_scan.py +++ b/tests/unit/cli/test_pyrit_scan.py @@ -84,6 +84,16 @@ def test_parse_args_with_memory_labels(self): args = pyrit_scan.parse_args(["test_scenario", "--memory-labels", '{"key":"value"}']) assert args.memory_labels == '{"key":"value"}' + def test_parse_args_with_dataset_parameters(self): + args = pyrit_scan.parse_args( + ["test_scenario", "--dataset-parameters", "harm_categories=cyber", "data_types=text"] + ) + assert args.dataset_parameters == [("harm_categories", "cyber"), ("data_types", "text")] + + def test_parse_args_dataset_parameter_without_equals_errors(self): + with pytest.raises(SystemExit): + pyrit_scan.parse_args(["test_scenario", "--dataset-parameters", "harm_categories"]) + def test_parse_args_complex_command(self): args = pyrit_scan.parse_args( [ @@ -606,6 +616,7 @@ def test_includes_initializer_args(self): max_retries=None, dataset_names=None, max_dataset_size=None, + dataset_parameters=None, memory_labels=None, ) request = pyrit_scan._build_run_request(parsed_args=parsed, scenario_name="s") @@ -621,6 +632,7 @@ def test_populates_optional_fields(self): max_retries=2, dataset_names=["d1"], max_dataset_size=10, + dataset_parameters=None, memory_labels='{"key":"value"}', ) request = pyrit_scan._build_run_request(parsed_args=parsed, scenario_name="s") @@ -631,6 +643,21 @@ def test_populates_optional_fields(self): assert request.max_dataset_size == 10 assert request.labels == {"key": "value"} + def test_populates_dataset_parameters(self): + parsed = Namespace( + target="t", + initializers=None, + scenario_strategies=None, + max_concurrency=None, + max_retries=None, + dataset_names=None, + max_dataset_size=None, + dataset_parameters=[("harm_categories", "cyber"), ("data_types", "text")], + memory_labels=None, + ) + request = pyrit_scan._build_run_request(parsed_args=parsed, scenario_name="s") + assert request.dataset_parameters == {"harm_categories": "cyber", "data_types": "text"} + def test_includes_scenario_declared_params(self): parsed = Namespace( target=None, @@ -640,6 +667,7 @@ def test_includes_scenario_declared_params(self): max_retries=None, dataset_names=None, max_dataset_size=None, + dataset_parameters=None, memory_labels=None, scenario__max_turns="7", ) diff --git a/tests/unit/cli/test_pyrit_shell.py b/tests/unit/cli/test_pyrit_shell.py index f13c237ec3..879e6b1855 100644 --- a/tests/unit/cli/test_pyrit_shell.py +++ b/tests/unit/cli/test_pyrit_shell.py @@ -545,6 +545,7 @@ def test_run_completed_path_with_results(self, shell, capsys): "memory_labels": {"k": "v"}, "dataset_names": ["d1"], "max_dataset_size": 5, + "dataset_parameters": [("harm_categories", "cyber"), ("data_types", "text")], }, ), patch("pyrit.cli._output.print_scenario_result_async", new_callable=AsyncMock), @@ -562,6 +563,7 @@ def test_run_completed_path_with_results(self, shell, capsys): assert sent.labels == {"k": "v"} assert sent.dataset_names == ["d1"] assert sent.max_dataset_size == 5 + assert sent.dataset_parameters == {"harm_categories": "cyber", "data_types": "text"} def test_run_failed_status_calls_summary(self, shell): s, client = shell diff --git a/tests/unit/scenario/core/test_dataset_configuration.py b/tests/unit/scenario/core/test_dataset_configuration.py index a1e94f41b1..f0b9535a6e 100644 --- a/tests/unit/scenario/core/test_dataset_configuration.py +++ b/tests/unit/scenario/core/test_dataset_configuration.py @@ -9,6 +9,7 @@ from pyrit.models import SeedAttackGroup, SeedGroup, SeedObjective, SeedPrompt from pyrit.scenario.core.dataset_configuration import ( + DATASET_FILTERS, INLINE_DATASET_NAME, CompoundDatasetAttackConfiguration, DatasetAttackConfiguration, @@ -16,6 +17,7 @@ DatasetConstraintError, DatasetSourceKind, ResolvedDataset, + build_dataset_filters, forbid_inline_seeds, require_harm_categories, require_inline_seeds, @@ -595,3 +597,117 @@ async def test_inline_children_combine(self) -> None: ) groups = await config.get_seed_attack_groups_async() assert sorted(g.objective.value for g in groups) == ["a", "b"] + + +class TestBuildDatasetFilters: + """``build_dataset_filters`` maps the user-facing bag onto ``get_seeds`` filters.""" + + def test_none_returns_empty(self) -> None: + assert build_dataset_filters(parameters=None) == {} + + def test_empty_returns_empty(self) -> None: + assert build_dataset_filters(parameters={}) == {} + + def test_max_dataset_size_is_rejected(self) -> None: + with pytest.raises(ValueError, match="Unknown dataset parameter 'max_dataset_size'"): + build_dataset_filters(parameters={"max_dataset_size": "13"}) + + def test_single_value_maps_to_list(self) -> None: + assert build_dataset_filters(parameters={"harm_categories": "cyber"}) == {"harm_categories": ["cyber"]} + + def test_comma_splits_into_list(self) -> None: + assert build_dataset_filters(parameters={"harm_categories": "cyber,violence"}) == { + "harm_categories": ["cyber", "violence"] + } + + def test_multiple_filters_combine(self) -> None: + assert build_dataset_filters(parameters={"harm_categories": "cyber", "data_types": "text,image_path"}) == { + "harm_categories": ["cyber"], + "data_types": ["text", "image_path"], + } + + def test_unknown_key_raises(self) -> None: + with pytest.raises(ValueError, match="Unknown dataset parameter 'bogus'"): + build_dataset_filters(parameters={"bogus": "x"}) + + def test_unexposed_get_seeds_kwarg_is_rejected(self) -> None: + # ``authors`` is a real get_seeds kwarg but is intentionally NOT exposed as a filter. + with pytest.raises(ValueError, match="Unknown dataset parameter 'authors'"): + build_dataset_filters(parameters={"authors": "jones"}) + + +class TestDatasetFilterRegistry: + """The exposed filter set is intentional and stays in contract with ``get_seeds``.""" + + def test_exposed_filters_are_frozen(self) -> None: + # Adding/removing a filter must be a deliberate edit to this expected set. + assert {"harm_categories", "data_types"} == DATASET_FILTERS + + def test_every_filter_is_a_sequence_get_seeds_param(self) -> None: + # Each exposed key must be a real get_seeds parameter AND list-valued, since + # build_dataset_filters always coerces values into a list. + import typing + from collections.abc import Sequence + + from pyrit.memory.memory_interface import MemoryInterface + + hints = typing.get_type_hints(MemoryInterface.get_seeds) + + def _allows_sequence(annotation: object) -> bool: + for candidate in (annotation, *typing.get_args(annotation)): + origin = typing.get_origin(candidate) or candidate + if isinstance(origin, type) and origin is not str and issubclass(origin, Sequence): + return True + return False + + for name in DATASET_FILTERS: + assert name in hints, f"'{name}' is not a MemoryInterface.get_seeds parameter" + assert _allows_sequence(hints[name]), f"'{name}' must be a Sequence-typed get_seeds parameter" + + +class TestDatasetConfigurationFilters: + """Filters are threaded into ``get_seeds`` and applied before sampling.""" + + async def test_filters_passed_to_get_seeds(self, mock_memory: MagicMock) -> None: + mock_memory.get_seeds.return_value = make_objectives("a", "b") + config = DatasetAttackConfiguration(dataset_names=["d1"], filters={"harm_categories": ["cyber"]}) + await config.get_seed_attack_groups_async() + mock_memory.get_seeds.assert_called_with(dataset_name="d1", harm_categories=["cyber"]) + + async def test_filter_removing_all_seeds_raises_specific_error(self, mock_memory: MagicMock) -> None: + def _get_seeds(*, dataset_name, **filters): + return [] if filters else make_objectives("a", "b") + + mock_memory.get_seeds.side_effect = _get_seeds + config = DatasetAttackConfiguration( + dataset_names=["d1"], filters={"harm_categories": ["missing"]}, auto_fetch=False + ) + with pytest.raises(DatasetConstraintError, match="none match the configured filters"): + await config.get_seed_attack_groups_async() + + async def test_update_filters_merges(self, mock_memory: MagicMock) -> None: + mock_memory.get_seeds.return_value = make_objectives("a") + config = DatasetAttackConfiguration(dataset_names=["d1"], filters={"harm_categories": ["a"]}) + config.update_filters(filters={"authors": ["jones"]}) + await config.get_seed_attack_groups_async() + mock_memory.get_seeds.assert_called_with(dataset_name="d1", harm_categories=["a"], authors=["jones"]) + + def test_filters_property_returns_copy(self) -> None: + config = DatasetAttackConfiguration(dataset_names=["d1"], filters={"harm_categories": ["a"]}) + config.filters["authors"] = ["mutated"] + assert config.filters == {"harm_categories": ["a"]} + + async def test_per_dataset_threads_filters_to_children(self, mock_memory: MagicMock) -> None: + mock_memory.get_seeds.return_value = make_objectives("a") + config = CompoundDatasetAttackConfiguration.per_dataset( + dataset_names=["d1"], filters={"harm_categories": ["cyber"]} + ) + await config.get_seed_attack_groups_async() + mock_memory.get_seeds.assert_called_with(dataset_name="d1", harm_categories=["cyber"]) + + async def test_compound_update_filters_propagates_to_children(self, mock_memory: MagicMock) -> None: + mock_memory.get_seeds.return_value = make_objectives("a") + config = CompoundDatasetAttackConfiguration.per_dataset(dataset_names=["d1"]) + config.update_filters(filters={"harm_categories": ["cyber"]}) + await config.get_seed_attack_groups_async() + mock_memory.get_seeds.assert_called_with(dataset_name="d1", harm_categories=["cyber"])