FIX: count languages in the Multilingual run size estimate

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]:
         """

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)