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

[VishnuR23]

Addresses #1546.

Why this is not already addressed

Probe._prune_data covers probes that cap self.prompts in place, and plenty use it. But probes that build a prompt list locally and then assign it still hand-roll the cap. There were four such sites:

• probes/encoding.py (EncodingMixin.__init__)
• probes/sysprompt_extraction.py (_generate_attempts)
• probes/dra.py (twice, with an identical comment)

The copies had already drifted, and not harmlessly. dra.py and badchars.py guard with soft_probe_prompt_cap is not None; encoding.py and sysprompt_extraction.py compared against it directly, so a config with no cap set crashed them:

probes.encoding.InjectBase64
    TypeError: '<' not supported between instances of 'int' and 'NoneType'
probes.sysprompt_extraction.SystemPromptExtraction
    TypeError: '>' not supported between instances of 'int' and 'NoneType'

That is the case for factoring up: with the policy in four places, two of them were wrong.

What this changes

Adds Probe._cap_prompt_list(prompts) next to _prune_data — same policy, but for the "returns a list" shape rather than the in-place one — and routes the four sites through it. None handling now lives in one place, so both crashes go away by construction.

badchars._downsample_prompts is deliberately left alone: it preserves category balance and its docstring already explains why it differs.

Removing the last random use from encoding.py orphaned its import random, so that goes too.

Not duplicating existing work

#1546 has no linked PR and no assignee. No open PR touches the cap logic in these files. I commented on the issue with this plan before writing code.

AI assistance disclosure

Written with AI assistance (Claude). I reviewed every changed line, reproduced both crashes before fixing them, and confirmed the new tests fail without the fix.

Verification

• Supporting configuration — n/a
• Both crashes reproduced before the change, gone after, with capping still applied:

  cap = None    encoding.InjectBase64 835 prompts | sysprompt_extraction 1400 | dra.DRA 60   (all previously TypeError)
  cap = 12      encoding.InjectBase64  12 prompts | sysprompt_extraction   12 | dra.DRA 12   (capped)
  

• New tests fail without the fix — reverting only the three probe files gives:

  FAILED test_probe_builds_with_unset_cap[probes.encoding.InjectBase64]
  FAILED test_probe_builds_with_unset_cap[probes.sysprompt_extraction.SystemPromptExtraction]
  2 failed, 1 passed
  

  (dra passes either way — it already had the guard, which is the point.)
• python -m pytest tests/probes/test_probes.py tests/probes/test_probes_encoding.py tests/probes/test_probes_sysprompt_extraction.py tests/probes/test_probes_dra.py tests/probes/test_probes_badcharacters.py — 1466 passed, 1 skipped
• Verify the thing does what it should — tests cover downsampling to the cap, lists already within the cap, an unset cap, and follow_prompt_cap = False
• Verify the thing does not do what it should not — capping behaviour is unchanged for all three probes at a set cap; badchars category-balanced downsampling is untouched and its tests still pass
• Document the thing and how it works — _cap_prompt_list docstring states when to use it versus _prune_data

Two unrelated failures show up locally in test_probes.py (probes.audio.AudioAchillesHeel, probes.sata.MLM). They fail identically on unmodified main in my environment — missing audio extras and NLTK data — and are not touched by this change.

Tested on macOS, Python 3.13.7.

[leondz]

This looks like it only tests three probes instead of all that support capping - why?

[jmartin-tech]

The goal of the issue being targeted here was to consolidate and ensure all probes use a shared application of the pruning process as provided by a base class.

For the probes addressed here setting self.prompts then calling self._prune_data(self.soft_probe_prompt_cap) should be sufficient to get the same result, why add a new function that implements the pruning without awareness of triggers?

Note by calling self._prune_data() this also offers probe authors a single location to override the pruning process when they need to preform a revised pattern for pruning. For instance, this would allow for badchars._downsample_prompts() to simply be renamed _prune_data() with a minor signature adjustment and shift that outlier.

Another consolidation task that could be considered part of the reported issue is that follow_prompt_cap could become a probes.base.DEFAULT_PARAM evaluated in _prune_data() ensuring all probes have an ability to enable or disable pruning in a common way and shifting the burden from each probe that currently supports the option to the base class.

[VishnuR23]

Thanks both — reworked to follow the design you described. Force-pushed.

@jmartin-tech, you were right that a second function was the wrong shape. _prune_data is now the single place pruning happens:

• it defaults cap to soft_probe_prompt_cap and evaluates follow_prompt_cap itself, so call sites no longer guard it
• follow_prompt_cap moved to Probe.DEFAULT_PARAMS, so every probe carries the switch instead of ~13 probes each declaring it. Existing class attributes still win, since Configurable only fills in a default when the attribute is absent — so follow_prompt_cap = False on ContinueSlursReclaimedSlursFull and the leakreplay mixins keeps working
• badchars._downsample_prompts is now a _prune_data override, which is what its category-balanced pruning always was
• _prune_cap gives overrides the same opt-out/unset-cap handling, so IntentProbe, badchars and encoding don't each repeat the guard. Happy to inline it if you'd rather it not exist

Sweeping for the rest turned up more than the three I first found. latentinjection._build_prompts_triggers held a verbatim copy of the base body; web_injection, encoding, sysprompt_extraction and dra each capped a local list before assigning. Those now assign and prune. encoding needs an override because _prompt_intents is indexed by prompt position, the same reason IntentProbe overrides for prompt_intents.

divergence is the one I deliberately left: it uses the cap as a generation target (while len(self.prompts) < prompt_cap) rather than pruning afterwards, so routing it through _prune_data would mean generating everything first. Say the word if you'd prefer it converted anyway.

@leondz — fair hit on the test coverage. Tests no longer name probes; they discover every probe that prunes by walking the plugin registry and each class's MRO. That is 77 probes across 16 modules, each checked for honouring the cap, building with no cap set, and keeping triggers aligned after pruning.

That coverage also showed the null-cap crash was much wider than the two probes I originally reported — running the new tests against main gives 34 failures, including all of packagehallucination, phrasing and suffix.

Verification: pytest tests/probes tests/cas tests/test_docs.py → 2901 passed, 41 skipped. The two failures in my environment (probes.audio.AudioAchillesHeel, probes.sata.MLM) fail identically on unmodified main — missing audio extras and NLTK data.

black is clean on every file I touched. I left nine other files in garak/probes/ that black would reformat untouched, since they're unrelated to this change.