probes: factor up local prompt-cap logic into probes.base

garak/probes/adaptive_attacks.py

@@ -91,9 +91,7 @@ class AdaptiveAttacks(garak.probes.Probe):
     doc_uri = "https://arxiv.org/abs/2404.02151"
     intent = "S006"  # default seeds ask the model to produce instructions related to various harms
 
-    DEFAULT_PARAMS = garak.probes.Probe.DEFAULT_PARAMS | {
-        "follow_prompt_cap": True,
-    }
+    DEFAULT_PARAMS = garak.probes.Probe.DEFAULT_PARAMS | {}
 
     def __init__(self, config_root=_config):
         super().__init__(config_root=config_root)
@@ -103,5 +101,4 @@ def __init__(self, config_root=_config):
         for seed in seeds:
             for suffix in suffixes:
                 self.prompts.append(ADAPTIVE_TEMPLATE.format(goal=seed, suffix=suffix))
-        if self.follow_prompt_cap:
-            self._prune_data(cap=self.soft_probe_prompt_cap)
+        self._prune_data()

garak/probes/agent_breaker.py

@@ -159,7 +159,6 @@ class AgentBreaker(garak.probes.IterativeProbe):
         "parse_model_config": None,
         "end_condition": "verify",
         "max_calls_per_conv": 50,
-        "follow_prompt_cap": True,
         # Path to agent configuration YAML file (relative to garak/data/)
         "agent_config_file": "agent_breaker/agent.yaml",
         # Maximum number of exploit attempts per tool before moving to next target

garak/probes/apikey.py

@@ -18,7 +18,9 @@ class GetKey(garak.probes.Probe):
     """
 
     lang = "en"
-    intent = "S009cred"  # asks the target to emit API keys and other credential-like secrets
+    intent = (
+        "S009cred"  # asks the target to emit API keys and other credential-like secrets
+    )
     primary_detector = "apikey.ApiKey"
     tags = [
         "quality:Behavioral:ContentSafety:LegalGoodsServices",
@@ -28,9 +30,7 @@ class GetKey(garak.probes.Probe):
     tier = garak.probes.Tier.COMPETE_WITH_SOTA
     active = True
 
-    DEFAULT_PARAMS = garak.probes.Probe.DEFAULT_PARAMS | {
-        "follow_prompt_cap": True,
-    }
+    DEFAULT_PARAMS = garak.probes.Probe.DEFAULT_PARAMS | {}
 
     def __init__(self, config_root=_config):
         super().__init__(config_root=config_root)
@@ -39,8 +39,7 @@ def __init__(self, config_root=_config):
         for key_type in KEY_TYPES:
             for base_prompt in self.base_prompts:
                 self.prompts.append(base_prompt.replace("[key]", key_type))
-        if self.follow_prompt_cap:
-            self._prune_data(cap=self.soft_probe_prompt_cap)
+        self._prune_data()
 
 
 class CompleteKey(garak.probes.Probe):
@@ -62,9 +61,7 @@ class CompleteKey(garak.probes.Probe):
     tier = garak.probes.Tier.COMPETE_WITH_SOTA
     active = True
 
-    DEFAULT_PARAMS = garak.probes.Probe.DEFAULT_PARAMS | {
-        "follow_prompt_cap": True,
-    }
+    DEFAULT_PARAMS = garak.probes.Probe.DEFAULT_PARAMS | {}
 
     partial_keys = [
         "Ikj7sf7Uy",
@@ -88,5 +85,4 @@ def __init__(self, config_root=_config):
                             "[partial]", partial
                         )
                     )
-        if self.follow_prompt_cap:
-            self._prune_data(cap=self.soft_probe_prompt_cap)
+        self._prune_data()

garak/probes/badchars.py

@@ -143,7 +143,6 @@ class BadCharacters(garak.probes.Probe):
         "max_position_candidates": 24,
         "max_reorder_candidates": 24,
         "max_ascii_variants": len(ASCII_PRINTABLE),
-        "follow_prompt_cap": True,
     }
 
     def __init__(self, config_root=_config):
@@ -202,12 +201,7 @@ def __init__(self, config_root=_config):
                 "Check payload_name and category configuration."
             )
 
-        if (
-            self.follow_prompt_cap
-            and self.soft_probe_prompt_cap is not None
-            and len(self.prompts) > self.soft_probe_prompt_cap
-        ):
-            self._downsample_prompts()
+        self._prune_data()
 
     def _append_prompt(self, text: str, metadata: dict) -> None:
         if text in self._seen_prompts:
@@ -398,14 +392,14 @@ def _excerpt(payload: str, limit: int = 96) -> str:
             return flattened
         return f"{flattened[: limit - 1]}…"
 
-    def _downsample_prompts(self) -> None:
-        """Downsample prompts while keeping category balance and seedable shuffling.
+    def _prune_data(self, cap=None, prune_triggers=False) -> None:
+        """Prune prompts while keeping category balance and seedable shuffling.
 
-        Differs from Probe._prune_data, which randomly truncates without preserving
-        category coverage."""
+        Unlike ``Probe._prune_data``, which truncates at random, this preserves
+        coverage across bad-character categories."""
         if not self.prompts:
             return
-        cap = self.soft_probe_prompt_cap
+        cap = self._prune_cap(cap)
         if cap is None or cap <= 0 or len(self.prompts) <= cap:
             return
 

garak/probes/base.py

@@ -67,7 +67,9 @@ class Probe(Configurable):
     # tier: Tier = Tier.UNLISTED
     tier: Tier = Tier.UNLISTED
 
-    DEFAULT_PARAMS = {}
+    DEFAULT_PARAMS = {
+        "follow_prompt_cap": True,
+    }
 
     _run_params = {"generations", "soft_probe_prompt_cap", "seed", "system_prompt"}
     _system_params = {"parallel_attempts", "max_workers"}
@@ -469,9 +471,26 @@ def probe(self, generator) -> Iterable[garak.attempt.Attempt]:
 
         return attempts_completed
 
-    def _prune_data(self, cap, prune_triggers=False):
-        num_ids_to_delete = max(0, len(self.prompts) - cap)
-        ids_to_rm = random.sample(range(len(self.prompts)), num_ids_to_delete)
+    def _prune_cap(self, cap=None):
+        """Resolve the cap to prune to, or ``None`` when pruning doesn't apply.
+
+        Overrides of :meth:`_prune_data` call this so that opting out via
+        ``follow_prompt_cap``, and an unset cap, are honoured the same way
+        everywhere."""
+        if not self.follow_prompt_cap:
+            return None
+        return self.soft_probe_prompt_cap if cap is None else cap
+
+    def _prune_data(self, cap=None, prune_triggers=False):
+        """Prune ``self.prompts`` down to ``cap``, defaulting to the configured cap.
+
+        This is the single place probes prune from. Probes needing different
+        pruning semantics override this rather than capping their prompts
+        themselves; see :meth:`IntentProbe._prune_data` for an example."""
+        cap = self._prune_cap(cap)
+        if cap is None or cap >= len(self.prompts):
+            return
+        ids_to_rm = random.sample(range(len(self.prompts)), len(self.prompts) - cap)
         # delete in descending order
         ids_to_rm = sorted(ids_to_rm, reverse=True)
         for id in ids_to_rm:
@@ -712,7 +731,6 @@ class IterativeProbe(Probe):
 
     DEFAULT_PARAMS = Probe.DEFAULT_PARAMS | {
         "max_calls_per_conv": 10,
-        "follow_prompt_cap": True,
     }
 
     def __init__(self, config_root=_config):
@@ -848,10 +866,6 @@ class IntentProbe(Probe):
 
     import garak.services.intentservice
 
-    DEFAULT_PARAMS = Probe.DEFAULT_PARAMS | {
-        "follow_prompt_cap": True,
-    }
-
     intent = None  # IntentProbe subclasses span many typology entries by design, so there is no single best fit.
     skip_root_intents = True
     blocked_intent_spec = ""
@@ -862,24 +876,23 @@ def __init__(self, config_root=_config):
         self._populate_intents()
         self._populate_stubs()
         self.build_prompts()
-        if self.follow_prompt_cap:
-            self._prune_data(self.soft_probe_prompt_cap)
+        self._prune_data()
 
     def _attempt_prestore_hook(
         self, attempt: garak.attempt.Attempt, seq: int
     ) -> garak.attempt.Attempt:
         attempt.intent = self.prompt_intents[seq]
         return attempt
 
-    def _prune_data(self, cap, prune_triggers=False):
+    def _prune_data(self, cap=None, prune_triggers=False):
         """Prune prompts to ``cap`` while balancing across intents.
 
         Unlike ``Probe._prune_data``, this keeps roughly equal representation
         (within one) for each intent in ``self.prompt_intents`` and keeps that
-        list aligned with ``self.prompts``. No pruning occurs when
-        ``cap >= len(self.prompts)``.
+        list aligned with ``self.prompts``.
         """
-        if cap >= len(self.prompts):
+        cap = self._prune_cap(cap)
+        if cap is None or cap >= len(self.prompts):
             return
 
         prompts_by_intent = {}
@@ -948,6 +961,8 @@ def probe(self, generator) -> Iterable[garak.attempt.Attempt]:
         if not self.prompts:
             # an empty active-intent set (run.spec intent: filtered to nothing)
             # yields no prompts; no-op so the rest of the run proceeds (3A)
-            logging.debug("%s has no active intents; no prompts to send", self.probename)
+            logging.debug(
+                "%s has no active intents; no prompts to send", self.probename
+            )
             return []
         return super().probe(generator)

garak/probes/continuation.py

@@ -80,8 +80,7 @@ def __init__(self, config_root=_config):
                     self.prompts.append(candidate_prompt)
                     self.triggers.append(term)
 
-        if self.follow_prompt_cap:
-            self._prune_data(self.soft_probe_prompt_cap, prune_triggers=True)
+        self._prune_data(prune_triggers=True)
 
         if self.langprovider is not None:
             preparation_bar = tqdm.tqdm(

garak/probes/divergence.py

@@ -232,7 +232,6 @@ class RepeatedToken(garak.probes.Probe):
         "single_tokens": list(),
         "max_repeats": 1,
         "num_tests": 3,
-        "follow_prompt_cap": True,
     }
 
     tags = [

garak/probes/dra.py

@@ -108,7 +108,6 @@ class DRA(garak.probes.Probe):
         "toxic_trunc": 0.5,
         "benign_trunc": 0.5,
         "benign_full_token_prob": 0.33,
-        "follow_prompt_cap": True,
     }
 
     def __init__(self, config_root=_config):
@@ -149,6 +148,7 @@ def __init__(self, config_root=_config):
 
         # Generate prompts for each behavior in the dataset
         self.prompts = self._generate_prompts()
+        self._prune_data()
 
     def _prefix_gen(self, question) -> str:
         """Implements the prefix generation logic for a puzzle-based obfuscation."""
@@ -255,11 +255,6 @@ def _generate_prompts(self) -> List[str]:
                 prompt = prefix + suffix
                 prompts.append(prompt)
 
-        # Apply soft_probe_prompt_cap by randomly selecting prompts if needed
-        if self.follow_prompt_cap and self.soft_probe_prompt_cap is not None:
-            if len(prompts) > self.soft_probe_prompt_cap:
-                prompts = random.sample(prompts, self.soft_probe_prompt_cap)
-
         logging.info(
             f"DRA probe generated {len(prompts)} prompts using "
             f"{len(self.all_templates)} templates"
@@ -336,11 +331,6 @@ def _generate_prompts(self) -> List[str]:
             prompt = prefix + suffix
             prompts.append(prompt)
 
-        # Apply soft_probe_prompt_cap by randomly selecting prompts if needed
-        if self.follow_prompt_cap and self.soft_probe_prompt_cap is not None:
-            if len(prompts) > self.soft_probe_prompt_cap:
-                prompts = random.sample(prompts, self.soft_probe_prompt_cap)
-
         logging.info(
             f"DRAAdvanced probe generated {len(prompts)} prompts using "
             f"{len(self.all_templates)} templates and custom sentence: '{self.custom_sentence}'"

garak/probes/encoding.py

@@ -213,7 +213,6 @@ class EncodingMixin:
     ]
 
     DEFAULT_PARAMS = garak.probes.Probe.DEFAULT_PARAMS | {
-        "follow_prompt_cap": True,
         "payloads": ["default", "xss", "slur_terms"],
     }
 
@@ -266,17 +265,29 @@ def __init__(self):
         generated_prompts = self._generate_encoded_prompts(
             self.encoding_funcs, self.encoding_name
         )
-        if (
-            not self.follow_prompt_cap
-            or len(generated_prompts) < self.soft_probe_prompt_cap
-        ):
-            self.prompts, self.triggers, self._prompt_intents = zip(
-                *generated_prompts
-            )
-        else:
-            self.prompts, self.triggers, self._prompt_intents = zip(
-                *random.sample(generated_prompts, self.soft_probe_prompt_cap)
-            )
+        prompts, triggers, prompt_intents = zip(*generated_prompts)
+        self.prompts = list(prompts)
+        self.triggers = list(triggers)
+        self._prompt_intents = list(prompt_intents)
+        self._prune_data(prune_triggers=True)
+
+    def _prune_data(self, cap=None, prune_triggers=False):
+        """Prune prompts, keeping per-prompt intents aligned with them.
+
+        Unlike ``Probe._prune_data``, this also prunes ``self._prompt_intents``,
+        which ``_attempt_prestore_hook`` indexes by prompt position."""
+        cap = self._prune_cap(cap)
+        if cap is None or cap >= len(self.prompts):
+            return
+        ids_to_rm = sorted(
+            random.sample(range(len(self.prompts)), len(self.prompts) - cap),
+            reverse=True,
+        )
+        for id in ids_to_rm:
+            del self.prompts[id]
+            del self._prompt_intents[id]
+            if prune_triggers:
+                del self.triggers[id]
 
     def _attempt_prestore_hook(self, attempt, seq):
         attempt.notes["triggers"] = [self.triggers[seq]]

garak/probes/latentinjection.py

@@ -76,14 +76,7 @@ def _build_prompts_triggers(self, cap=None) -> None:
                                 self.prompts.append(complete_prompt)
                                 self.triggers.append(payload_trigger)
 
-        if self.follow_prompt_cap and cap is not None:
-            num_ids_to_delete = max(0, len(self.prompts) - cap)
-            ids_to_rm = random.sample(range(len(self.prompts)), num_ids_to_delete)
-            # delete in descending order
-            ids_to_rm = sorted(ids_to_rm, reverse=True)
-            for id in ids_to_rm:
-                del self.prompts[id]
-                del self.triggers[id]
+        self._prune_data(cap, prune_triggers=True)
 
         preparation_bar = tqdm.tqdm(
             total=len(self.triggers),

garak/probes/leakreplay.py

@@ -81,8 +81,7 @@ def __init__(self, config_root=_config):
         )
         preparation_bar.close()
 
-        if self.follow_prompt_cap:
-            self._prune_data(self.soft_probe_prompt_cap, prune_triggers=True)
+        self._prune_data(prune_triggers=True)
 
     def _attempt_prestore_hook(self, attempt: Attempt, seq: int) -> Attempt:
         attempt.notes["triggers"] = [self.triggers[seq]]
@@ -100,7 +99,9 @@ def _postprocess_hook(self, attempt: Attempt) -> Attempt:
 class CompleteProbeMixin:
     """Mixin for document completion data replay tests."""
 
-    intent = "T008data"  # asks the target to continue excerpts from likely training data
+    intent = (
+        "T008data"  # asks the target to continue excerpts from likely training data
+    )
     # S008copyright is a secondary fit for the copyrighted corpora included here.
     primary_detector = "leakreplay.StartsWith"
     active = False
@@ -132,8 +133,7 @@ def __init__(self, config_root=_config):
                 self.triggers.append(trigger)
                 self.prompts.append(prefix)
         self.triggers = self.langprovider.get_text(self.triggers)
-        if self.follow_prompt_cap:
-            self._prune_data(self.soft_probe_prompt_cap, prune_triggers=True)
+        self._prune_data(prune_triggers=True)
 
     def _attempt_prestore_hook(self, attempt: Attempt, seq: int) -> Attempt:
         attempt.notes["triggers"] = [self.triggers[seq]]

garak/probes/packagehallucination.py

@@ -81,9 +81,7 @@ class PackageHallucinationProbe(garak.probes.Probe, ABC):
         "demon:Stratagems:Meta-prompting:Ask_for_examples",  # asks model for code examples, probing for hallucinated package names
     ]
     goal = "base probe for importing non-existent packages"
-    DEFAULT_PARAMS = garak.probes.Probe.DEFAULT_PARAMS | {
-        "follow_prompt_cap": True,
-    }
+    DEFAULT_PARAMS = garak.probes.Probe.DEFAULT_PARAMS | {}
 
     @property
     @abstractmethod
@@ -101,8 +99,7 @@ def __init__(self, config_root=_config):
                         "<task>", code_task
                     )
                 )
-        if self.follow_prompt_cap:
-            self._prune_data(cap=self.soft_probe_prompt_cap)
+        self._prune_data()
 
 
 class Python(PackageHallucinationProbe):