diff --git a/pyrit/scenario/scenarios/airt/multilingual.py b/pyrit/scenario/scenarios/airt/multilingual.py index 1493e8129e..8f30d5afec 100644 --- a/pyrit/scenario/scenarios/airt/multilingual.py +++ b/pyrit/scenario/scenarios/airt/multilingual.py @@ -13,7 +13,15 @@ from pyrit.common.path import DATASETS_PATH from pyrit.converter import RandomTranslationConverter, TranslationConverter from pyrit.executor.attack import PromptSendingAttack -from pyrit.models import Parameter, SeedDataset +from pyrit.models import ( + BoundedDatasetSize, + Parameter, + ScenarioRunSizeComponent, + ScenarioRunSizeEstimate, + ScenarioRunSizeEstimateStatus, + ScenarioRunSizeFactor, + SeedDataset, +) from pyrit.registry.components.attack_technique_registry import AttackTechniqueRegistry from pyrit.scenario.core import ( AtomicAttack, @@ -267,6 +275,22 @@ async def _resolve_languages_async(self) -> list[str]: if persisted: return _normalize_languages(list(persisted)) + selection = self._get_language_selection() + if isinstance(selection, list): + return selection + return _normalize_languages(random.sample(self._default_languages, selection)) + + def _get_language_selection(self) -> list[str] | int: + """ + Read the language parameters for a fresh run. + + Returns: + list[str] | int: The explicit normalized languages, or how many to draw at random. + + Raises: + ValueError: If both ``num_languages`` and ``languages`` are provided, + or if ``num_languages`` is out of bounds. + """ num_languages = self.params.get("num_languages") languages = self.params.get("languages") @@ -283,7 +307,65 @@ async def _resolve_languages_async(self) -> list[str]: count = int(num_languages) if num_languages is not None else _DEFAULT_NUM_LANGUAGES if count < 1 or count > len(self._default_languages): raise ValueError(f"num_languages must be between 1 and {len(self._default_languages)}.") - return _normalize_languages(random.sample(self._default_languages, count)) + return count + + async def _estimate_run_size_async(self, *, budget: BoundedDatasetSize) -> ScenarioRunSizeEstimate: + """ + Estimate the technique sweep across the language and translation-strategy axes. + + Args: + budget: Resolved logical-group size contract before technique expansion. + + Returns: + ScenarioRunSizeEstimate: Baseline plus one slice per selected language and strategy. + """ + seed_group_count, datasets = await self._get_dataset_size_for_estimate_async(budget=budget) + selection = self._get_language_selection() + language_count = len(selection) if isinstance(selection, list) else selection + strategies = set(self.params.get("translation_strategies") or [_TRANSLATION, _RANDOM_TRANSLATION]) + technique_count = len(self._scenario_techniques) + sweep_factors = [ + ScenarioRunSizeFactor(label="selected seed-group estimate", count=seed_group_count), + ScenarioRunSizeFactor(label="selected concrete techniques", count=technique_count), + ] + + components: list[ScenarioRunSizeComponent] = [] + if self._include_baseline: + components.append( + ScenarioRunSizeComponent( + label="Baseline", + count=seed_group_count, + factors=[ScenarioRunSizeFactor(label="selected logical seed groups", count=seed_group_count)], + is_baseline=True, + note="One unmodified prompt-sending unit per selected seed group.", + ) + ) + if _TRANSLATION in strategies: + components.append( + ScenarioRunSizeComponent( + label="Translation", + count=seed_group_count * technique_count * language_count, + factors=[*sweep_factors, ScenarioRunSizeFactor(label="languages", count=language_count)], + note="Each objective is translated into every selected language.", + ) + ) + if _RANDOM_TRANSLATION in strategies: + components.append( + ScenarioRunSizeComponent( + label="Random translation", + count=seed_group_count * technique_count, + factors=sweep_factors, + note="One slice that mixes words from the selected languages.", + ) + ) + + return ScenarioRunSizeEstimate( + status=ScenarioRunSizeEstimateStatus.Approximate, + total_attack_count=sum(component.count for component in components), + components=components, + datasets=datasets, + note="Counts planned outer execution units; retries and internal attack turns are excluded.", + ) def _build_initial_scenario_metadata(self) -> dict[str, Any]: """ diff --git a/tests/unit/scenario/airt/test_multilingual.py b/tests/unit/scenario/airt/test_multilingual.py index 152f7f3fe6..b9ee90fd8d 100644 --- a/tests/unit/scenario/airt/test_multilingual.py +++ b/tests/unit/scenario/airt/test_multilingual.py @@ -3,6 +3,7 @@ """Tests for the Multilingual scenario.""" +from typing import Any from unittest.mock import AsyncMock, MagicMock, call, patch import pytest @@ -517,3 +518,38 @@ async def test_converter_retry_exhaustion_retries_only_failed_language( ScenarioRunState.IN_PROGRESS, ScenarioRunState.COMPLETED, ] + + @pytest.mark.parametrize( + "params", + [ + {}, + {"num_languages": 3}, + {"languages": ["French", "german", "French"]}, + {"languages": ["French", "German"], "translation_strategies": [_TRANSLATION]}, + {"languages": ["French", "German"], "translation_strategies": [_RANDOM_TRANSLATION]}, + {"languages": ["French"], "include_baseline": False}, + {"scenario_techniques": ["prompt_sending", "base64"]}, + ], + ) + async def test_run_size_estimate_matches_the_initialized_run_async( + self, + *, + params: dict[str, Any], + mock_objective_target: PromptTarget, + mock_adversarial_chat: PromptTarget, + mock_objective_scorer: TrueFalseScorer, + ) -> None: + seed_groups = [AttackSeedGroup(seeds=[SeedObjective(value=f"objective {index}")]) for index in range(5)] + technique_class = _build_multilingual_technique() + if "scenario_techniques" in params: + params = { + **params, + "scenario_techniques": [technique_class(name) for name in params["scenario_techniques"]], + } + with _patch_seed_groups(seed_groups): + scenario = Multilingual(adversarial_chat=mock_adversarial_chat, objective_scorer=mock_objective_scorer) + scenario.set_params_from_args(args={"objective_target": mock_objective_target, **params}) + estimate = await scenario.get_run_size_estimate_async(target_is_configured=True) + await scenario.initialize_async() + + assert estimate.estimated_attack_count == sum(len(attack.seed_groups) for attack in scenario._atomic_attacks)