Red Teaming functionality using function middleware

[kshiarl1]

Description

Closes

Red Teaming functionality using function middleware.

This PR introduces the necessary ingredients to perform AI red teaming on an agent workflow.
These are:

• Intercepting function and adding attack payloads (prompt injections)
• Evaluating the effect of these payloads on other parts of the workflow.
• Attack scenario definition and orchestration.

Key Features

Implements RedTeamingMiddleware.

• Special middleware that can replace inputs and output strings and floats to functions.
• Configurable payloads to replace i/o with.
• Can search within function schema to find the right fields to introduce the payload.

Implements RedTeamingEvaluationRunner.

• Allows the user to run red teaming using different threat scenarios.
• Threat scenarios are implemented via the RedTeamingEvaluationRunnerConfig.
• Processes, summarizes and saves results for ease of use post evaluation.

Implements a RedTeamingEvaluator.

• Can filter intermediate steps to find specific function inputs / outputs.
• Uses LLM as a judge to determine whether an attack is successful.
• Accepts scenario specific instructions for evaluation.

By Submitting this PR I confirm:

• I am familiar with the Contributing Guidelines.
• We require that all contributors "sign-off" on their commits. This certifies that the contribution is your original work, or you have rights to submit it under the same license, or a compatible license.
  • Any contribution which contains commits that are not Signed-Off will not be accepted.
• When the PR is ready for review, new or existing tests cover these changes.
• When the PR is ready for review, the documentation is up to date with these changes.

[copy-pr-bot]

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

[coderabbitai]

Important

Review skipped

Auto reviews are disabled on base/target branches other than the default branch.

🗂️ Base branches to auto review (2)

• develop
• release/.*

Please check the settings in the CodeRabbit UI or the .coderabbit.yaml file in this repository. To trigger a single review, invoke the @coderabbitai review command.

You can disable this status message by setting the reviews.review_status to false in the CodeRabbit configuration file.

✨ Finishing touches

🧪 Generate unit tests (beta)

• Create PR with unit tests
• Post copyable unit tests in a comment

---

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❤️ Share

• X
• Mastodon
• Reddit
• LinkedIn

_{Comment @coderabbitai help to get the list of available commands and usage tips.}

[ericevans-nv]

Please revert the TypeVar and Generic[EvaluatorTemplateItem] changes to EvalOutput. This added complexity is unused as the simpler non-generic class works fine since subclasses like RedTeamingEvalOutputItem already inherit from EvalOutputItem.

[kshiarl1]

This does not work. Here is the problem. I am happy to discuss a different solution but the problem is there and it prevent extensibility of the whole evaluation framework so it needs to be solved:

If EvalOutput is constructed with RedTeamingEvalOutputItem instead of EvalOutputItem and then model_dump_json is called on EvalOutput (in order to save the result) it will dump only the fields in EvalOutputItem. This model_dump_json operation happens in EvaluationRun.

NeMo-Agent-Toolkit/src/nat/eval/evaluate.py

Line 351 in 161c463

# create json content using the evaluation results

[kshiarl1]

The simplest solution is to change the typing inside EvalOutput to be list[Any] instead of list[EvalOutputItem]

[kshiarl1]

If you would like to verify yourself that there is a problem try running this standalone script:

from pydantic import BaseModel
from typing import List

class EvalOutputItem(BaseModel):
base_field: str

class NewEvalOutputItem(EvalOutputItem):
new_field: str # This gets lost during serialization

class EvalOutput(BaseModel):
items: List[EvalOutputItem] # Type annotation limits serialization

#This works fine
new_item = NewEvalOutputItem(base_field="test", new_field="extra")
eval_output = EvalOutput(items=[new_item])

#But this loses the new_field
json_str = eval_output.model_dump_json()
#Only serializes EvalOutputItem fields, not NewEvalOutputItem fields
print(json_str)

[ericevans-nv]

I think what we are looking for here is:

eval_output_items: list[SerializeAsAny[EvalOutputItem]] 

Section: Serializing with duck typing 🦆¶

[ericevans-nv]

I think what we are looking for here is:

eval_output_items: list[SerializeAsAny[EvalOutputItem]] 

Section: Serializing with duck typing 🦆¶

[ericevans-nv]

The hard-coded wrapper_type=LLMFrameworkEnum.LANGCHAIN is fragile, the register function shouldn't need to know what framework the evaluator uses internally. Consider following the pattern from the DynamicFunctionMiddleware implementation, which passes the builder instance and patches get_llm to discover and wrap components at runtime. This decouples the registration from the implementation details and allows the evaluator to request whatever framework it needs without the register function making assumptions.

[kshiarl1]

I took this from here:

NeMo-Agent-Toolkit/src/nat/eval/tunable_rag_evaluator/register.py

Line 40 in 161c463

async def register_tunable_rag_evaluator(config: TunableRagEvaluatorConfig, builder: EvalBuilder):

and it is also the case here:

NeMo-Agent-Toolkit/src/nat/eval/rag_evaluator/register.py

Line 122 in 161c463

llm = await builder.get_llm(config.llm_name, wrapper_type=LLMFrameworkEnum.LANGCHAIN)

so basically all LLMJudge evaluators in NAT use the same wrapper.

could you point me out to what you mean exactly?

[ericevans-nv]

Avoid injecting hardcoded names into already defined components. It's acceptable to add a new component with its full configuration (e.g., adding a new evaluator or LLM), but if a component is already configured by the user, the runner should only modify its parameters—not override its name or type.

[kshiarl1]

This is not yet in the full configuration. The RedTeamingEvaluationRunner constructs the full configuration it just forces the name.

[ericevans-nv]

The base workflow file should remain the source of truth for type information. The current RedTeamingRunnerConfig creates a parallel type system by defining fields like evaluator_llm: LLMBaseConfig and requiring rebuild_annotations() to stay in sync with the workflow's type registry—this is fragile and overcomplicated. Instead, the scenarios schema should work with the base workflow through mutations (path-value changes) rather than duplicating its type definitions. If a user has an LLM configured in their workflow, the scenario should reference it by name instead of injecting a new one with a hardcoded key. This would eliminate the need for rebuild_annotations(), remove type coupling between scenarios and workflows, respect user-configured values, and avoid duplicating type registration efforts. While this approach works for the red teaming evaluator since you own that component, we should not couple concerns this way for other components. Let's discuss a better pattern for converting scenarios to workflow configurations.

[kshiarl1]

Happy do discuss a different solution. I feel like the current solution is best for the user as it makes the components of red teaming much more explicit and explicitly ties the evaluator with the scenarios (as is the case). I don't see how this overcomplicates things as these operations are strictly performed for the red teaming workflow that can only be done using the red-team command.

[ericevans-nv]

RedTeamingRunner should inherit from MultiEvaluationRunner rather than being a standalone class that wraps it internally. The current design instantiates MultiEvaluationRunner inside run() and delegates to run_all(), which is essentially inheritance without the benefits. By extending MultiEvaluationRunner, the red teaming runner can override init to handle scenario-based config transformation, call super().run_all() for execution, and add pre/post-processing in an overridden run_all() method. This makes the relationship explicit, reduces indirection, and follows standard patterns for specialized runners.

[kshiarl1]

I could make the inheritance for the RedTeamingEvaluationRunner work. But because injeriting the config is not recommended in my opinion it might be better to not do this. See below.

[ericevans-nv]

Similarly, RedTeamingRunnerConfig should extend MultiEvaluationRunConfig. Currently, the runner manually transforms RedTeamingRunnerConfig into a MultiEvaluationRunConfig before passing it to MultiEvaluationRunner. If the config class extended MultiEvaluationRunConfig, the inheritance hierarchy would be consistent between the configs and runners, and the transformation could be simplified or handled through inheritance rather than manual conversion.

[kshiarl1]

Are you sure about this? The whole point of RedTeamingRunnerConfig is that it can be loaded and instantiated from yaml / json...If i inherit then I will also inherit the mandatory configs field that does not exist in the stored config. The two configs have nothing to do with each other.

[ericevans-nv]

Please make the update I mentioned to use SerializeAsAny instead of the generics, and we can address the LLM-wrapper related changes later. The current changes are fine to merge just don’t pull dev in yet. Once everyone gets their first PR in, we can update the epic branch. Comment /merge no squash to retain your commits.