fix(divergence): guard max_tokens override when the generator lacks it

[adityasingh2400]

Fixes #1096.

Refocused per review. The reported problem is that the divergence probe's _generator_precall_hook reads generator.max_tokens unconditionally, so generators that do not carry max_tokens (the function generators, which by design only receive the prompt and static kwargs) raise AttributeError when run with this probe.

This guards the attribute the same way the promptinject probe does, checking that max_tokens is present on the generator before overriding it, and restoring it afterward only when it was actually changed. Generators without max_tokens are left untouched.

The earlier approach of having the function generators inherit the base DEFAULT_PARAMS is dropped here. As discussed, that changes how function generators are configured and is better evaluated as a separate enhancement PR on its own merits.

[jmartin-tech]

Is this how we want the function generators to work? Note that this would not pass these options to the function as only kwargs are passed by design. The current docs define that only the prompt as a str and the kwargs dictionary provided as static config are passed to these generators.

As noted in the reporting issue, there are currently two probes that expect to set generator options in pre or post inference stages. The promptinject probe guards to ensure the attribute is on the generator. I would suggest that divergence should also guard the expectation that max_tokens will exist on a the generator though I do think there might be some value in ensuring the base.Probe.DEFAULT_PARAMS are made available to function generators.

I could seen enhancing the function generators to inspect the signature of the function passed and add any generator attributes with matching names into the kwargs dictionary passed at inference time.

[erickgalinkin]

Is this how we want the function generators to work? Note that this would not pass these options to the function as only kwargs are passed by design. The current docs define that only the prompt as a str and the kwargs dictionary provided as static config are passed to these generators.

As noted in the reporting issue, there are currently two probes that expect to set generator options in pre or post inference stages. The promptinject probe guards to ensure the attribute is on the generator. I would suggest that divergence should also guard the expectation that max_tokens will exist on a the generator though I do think there might be some value in ensuring the base.Probe.DEFAULT_PARAMS are made available to function generators.

I could seen enhancing the function generators to inspect the signature of the function passed and add any generator attributes with matching names into the kwargs dictionary passed at inference time.

+1 on all of this. I think it changes the way that function is fundamentally different from other generators. If we want this functionality, then we need to be thoughtful about how the relevant function is invoked and if it would even accept our default params.

[adityasingh2400]

Good points, and I agree the current change is not the right shape on its own. Since function generators only receive the prompt and the static kwargs by design, inheriting DEFAULT_PARAMS on the generator does not actually make max_tokens reach the function, so it does not fully fix the divergence case and it blurs how function differs from the other generators.

I think the cleanest split is what you both described:

1. Make divergence guard the expectation that max_tokens exists on the generator, the same way promptinject already guards its attribute. That is the direct correctness fix for the reported issue and it does not change the function contract.
2. Optionally, enhance the function generators to inspect the passed function signature and copy any matching generator attributes (including base.Probe.DEFAULT_PARAMS names) into the kwargs at inference time. That keeps the prompt plus kwargs contract intact while letting a function opt into known params by naming them.

I am happy to rework this PR down to item 1 as the focused fix, and to add item 2 if you want function generators to surface those params. Which would you prefer, or both?

[jmartin-tech]

Refocus the PR on a defined target, the reported issue seem to focus on an expectation that all probes will have some parameters available which is simply not true today. Accounting for that gap should be the first iteration.

If you would like to address enhancing the capabilities of function generators that can be a separate enhancement PR that can be evaluated on its own merits.

[adityasingh2400]

Done in 3d0a50e, refocused as you suggested.

• The divergence probe now guards max_tokens in _generator_precall_hook, checking that it is on the generator before overriding (via dir, the same way promptinject guards its parameters). Generators that do not expose max_tokens are left untouched instead of raising AttributeError.
• Reverted the function generator DEFAULT_PARAMS change so this PR is purely the probe-side guard. Happy to open a separate enhancement PR for the function generator capabilities if that is wanted.
• Added two divergence tests covering the missing and present max_tokens paths.

[jmartin-tech]

DCO sign-off is required in all commit messages of the branch.

Commit sha: 3d0a50e, Author: Aditya Singh, Committer: Aditya Singh; The sign-off is missing.

[adityasingh2400]

Thanks for flagging. I added the DCO sign-off to all commits on the branch and force-pushed, so every commit now carries the Signed-off-by trailer and the DCO check is green. Let me know if anything else is needed.