nvidia-nat 1.4.0a20251120__py3-none-any.whl → 1.4.0a20260113__py3-none-any.whl

nat/middleware/middleware.py

@@ -1,4 +1,4 @@
-# SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: Apache-2.0
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
@@ -24,23 +24,26 @@ from __future__ import annotations
 
 import dataclasses
 from abc import ABC
+from abc import abstractmethod
 from collections.abc import AsyncIterator
 from collections.abc import Awaitable
 from collections.abc import Callable
 from typing import Any
 
 from pydantic import BaseModel
+from pydantic import ConfigDict
+from pydantic import Field
 
 #: Type alias for single-output invocation callables.
-CallNext = Callable[[Any], Awaitable[Any]]
+CallNext = Callable[..., Awaitable[Any]]
 
 #: Type alias for streaming invocation callables.
-CallNextStream = Callable[[Any], AsyncIterator[Any]]
+CallNextStream = Callable[..., AsyncIterator[Any]]
 
 
 @dataclasses.dataclass(frozen=True, kw_only=True)
 class FunctionMiddlewareContext:
-    """Context information about the function being wrapped by middleware.
+    """Static metadata about the function being wrapped by middleware.
 
     Middleware receives this context object which describes the function they
     are wrapping. This allows middleware to make decisions based on the
@@ -66,31 +69,58 @@ class FunctionMiddlewareContext:
     """Schema describing streaming outputs or :class:`types.NoneType` when absent."""
 
 
+class InvocationContext(BaseModel):
+    """Unified context for pre-invoke and post-invoke phases.
+
+    Used for both phases of middleware execution:
+    - Pre-invoke: output is None, modify modified_args/modified_kwargs to transform inputs
+    - Post-invoke: output contains the function result, modify output to transform results
+
+    This unified context simplifies the middleware interface by using a single
+    context type for both hooks.
+    """
+
+    model_config = ConfigDict(validate_assignment=True)
+
+    # Frozen fields - cannot be modified after creation
+    function_context: FunctionMiddlewareContext = Field(
+        frozen=True, description="Static metadata about the function being invoked (frozen).")
+    original_args: tuple[Any, ...] = Field(
+        frozen=True, description="The original function input arguments before any middleware processing.")
+    original_kwargs: dict[str, Any] = Field(
+        frozen=True, description="The original function input keyword arguments before any middleware processing.")
+
+    # Mutable fields - modify these to transform inputs/outputs
+    modified_args: tuple[Any, ...] = Field(description="Modified args after middleware processing.")
+    modified_kwargs: dict[str, Any] = Field(description="Modified kwargs after middleware processing.")
+    output: Any = Field(default=None, description="Function output. None pre-invoke, result post-invoke.")
+
+
 class Middleware(ABC):
-    """Base class for middleware-style wrapping.
+    """Base class for middleware-style wrapping with pre/post-invoke hooks.
 
     Middleware works like middleware in web frameworks:
 
-    1. **Preprocess**: Inspect and optionally modify inputs
+    1. **Preprocess**: Inspect and optionally modify inputs (via pre_invoke)
     2. **Call Next**: Delegate to the next middleware or the target itself
-    3. **Postprocess**: Process, transform, or augment the output
+    3. **Postprocess**: Process, transform, or augment the output (via post_invoke)
     4. **Continue**: Return or yield the final result
 
     Example::
 
-        class LoggingMiddleware(Middleware):
-            async def middleware_invoke(self, value, call_next, context):
-                # 1. Preprocess
-                print(f"Input: {value}")
-
-                # 2. Call next middleware/target
-                result = await call_next(value)
+        class LoggingMiddleware(FunctionMiddleware):
+            @property
+            def enabled(self) -> bool:
+                return True
 
-                # 3. Postprocess
-                print(f"Output: {result}")
+            async def pre_invoke(self, context: InvocationContext) -> InvocationContext | None:
+                print(f"Current args: {context.modified_args}")
+                print(f"Original args: {context.original_args}")
+                return None  # Pass through unchanged
 
-                # 4. Continue
-                return result
+            async def post_invoke(self, context: InvocationContext) -> InvocationContext | None:
+                print(f"Output: {context.output}")
+                return None  # Pass through unchanged
 
     Attributes:
         is_final: If True, this middleware terminates the chain. No subsequent
@@ -101,6 +131,78 @@ class Middleware(ABC):
     def __init__(self, *, is_final: bool = False) -> None:
         self._is_final = is_final
 
+    # ==================== Abstract Members ====================
+
+    @property
+    @abstractmethod
+    def enabled(self) -> bool:
+        """Whether this middleware should execute.
+        """
+        ...
+
+    @abstractmethod
+    async def pre_invoke(self, context: InvocationContext) -> InvocationContext | None:
+        """Transform inputs before execution.
+
+        Called by specialized middleware invoke methods (e.g., function_middleware_invoke).
+        Use to validate, transform, or augment inputs. At this phase, context.output is None.
+
+        Args:
+            context: Invocation context (Pydantic model) containing:
+                - function_context: Static function metadata (frozen)
+                - original_args: What entered the middleware chain (frozen)
+                - original_kwargs: What entered the middleware chain (frozen)
+                - modified_args: Current args (mutable)
+                - modified_kwargs: Current kwargs (mutable)
+                - output: None (function not yet called)
+
+        Returns:
+            InvocationContext: Return the (modified) context to signal changes
+            None: Pass through unchanged (framework uses current context state)
+
+        Note:
+            Frozen fields (original_args, original_kwargs) cannot be modified.
+            Attempting to modify them raises ValidationError.
+
+        Raises:
+            Any exception to abort execution
+        """
+        ...
+
+    @abstractmethod
+    async def post_invoke(self, context: InvocationContext) -> InvocationContext | None:
+        """Transform output after execution.
+
+        Called by specialized middleware invoke methods (e.g., function_middleware_invoke).
+        For streaming, called per-chunk. Use to validate, transform, or augment outputs.
+
+        Args:
+            context: Invocation context (Pydantic model) containing:
+                - function_context: Static function metadata (frozen)
+                - original_args: What entered the middleware chain (frozen)
+                - original_kwargs: What entered the middleware chain (frozen)
+                - modified_args: What the function received (mutable)
+                - modified_kwargs: What the function received (mutable)
+                - output: Current output value (mutable)
+
+        Returns:
+            InvocationContext: Return the (modified) context to signal changes
+            None: Pass through unchanged (framework uses current context.output)
+
+        Example::
+
+            async def post_invoke(self, context: InvocationContext) -> InvocationContext | None:
+                # Wrap the output
+                context.output = {"result": context.output, "processed": True}
+                return context  # Signal modification
+
+        Raises:
+            Any exception to abort and propagate error
+        """
+        ...
+
+    # ==================== Properties ====================
+
     @property
     def is_final(self) -> bool:
         """Whether this middleware terminates the chain.
@@ -111,13 +213,20 @@ class Middleware(ABC):
 
         return self._is_final
 
-    async def middleware_invoke(self, value: Any, call_next: CallNext, context: FunctionMiddlewareContext) -> Any:
+    # ==================== Default Invoke Methods ====================
+
+    async def middleware_invoke(self,
+                                value: Any,
+                                call_next: CallNext,
+                                context: FunctionMiddlewareContext,
+                                **kwargs: Any) -> Any:
         """Middleware for single-output invocations.
 
         Args:
             value: The input value to process
             call_next: Callable to invoke the next middleware or target
             context: Metadata about the target being wrapped
+            kwargs: Additional function arguments
 
         Returns:
             The (potentially modified) output from the target
@@ -125,12 +234,12 @@ class Middleware(ABC):
         The default implementation simply delegates to ``call_next``. Override this
         to add preprocessing, postprocessing, or to short-circuit execution::
 
-            async def middleware_invoke(self, value, call_next, context):
+            async def middleware_invoke(self, value, call_next, context, **kwargs):
                 # Preprocess: modify input
                 modified_input = transform(value)
 
                 # Call next: delegate to next middleware/target
-                result = await call_next(modified_input)
+                result = await call_next(modified_input, **kwargs)
 
                 # Postprocess: modify output
                 modified_result = transform_output(result)
@@ -140,16 +249,20 @@ class Middleware(ABC):
         """
 
         del context  # Unused by the default implementation.
-        return await call_next(value)
+        return await call_next(value, **kwargs)
 
-    async def middleware_stream(self, value: Any, call_next: CallNextStream,
-                                context: FunctionMiddlewareContext) -> AsyncIterator[Any]:
+    async def middleware_stream(self,
+                                value: Any,
+                                call_next: CallNextStream,
+                                context: FunctionMiddlewareContext,
+                                **kwargs: Any) -> AsyncIterator[Any]:
         """Middleware for streaming invocations.
 
         Args:
             value: The input value to process
             call_next: Callable to invoke the next middleware or target stream
             context: Metadata about the target being wrapped
+            kwargs: Additional function arguments
 
         Yields:
             Chunks from the stream (potentially modified)
@@ -157,12 +270,12 @@ class Middleware(ABC):
         The default implementation forwards to ``call_next`` untouched. Override this
         to add preprocessing, transform chunks, or perform cleanup::
 
-            async def middleware_stream(self, value, call_next, context):
+            async def middleware_stream(self, value, call_next, context, **kwargs):
                 # Preprocess: setup or modify input
                 modified_input = transform(value)
 
                 # Call next: get stream from next middleware/target
-                async for chunk in call_next(modified_input):
+                async for chunk in call_next(modified_input, **kwargs):
                     # Process each chunk
                     modified_chunk = transform_chunk(chunk)
                     yield modified_chunk
@@ -172,13 +285,14 @@ class Middleware(ABC):
         """
 
         del context  # Unused by the default implementation.
-        async for chunk in call_next(value):
+        async for chunk in call_next(value, **kwargs):
             yield chunk
 
 
 __all__ = [
     "CallNext",
     "CallNextStream",
-    "Middleware",
     "FunctionMiddlewareContext",
+    "InvocationContext",
+    "Middleware",
 ]

nat/middleware/red_teaming/__init__.py

@@ -0,0 +1,14 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.

nat/middleware/red_teaming/red_teaming_middleware.py

@@ -0,0 +1,344 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""Red teaming middleware for attacking agent functions.
+
+This module provides a middleware for red teaming and security testing that can
+intercept and modify function inputs or outputs with configurable attack payloads.
+
+The middleware supports:
+- Targeting specific functions or entire function groups
+- Field-level search within input/output schemas
+- Multiple attack modes (replace, append_start, append_middle, append_end)
+- Both regular and streaming function calls
+- Type-safe operations on strings, integers, and floats
+"""
+
+from __future__ import annotations
+
+import logging
+import random
+import re
+from typing import Any
+from typing import Literal
+from typing import cast
+
+from jsonpath_ng import parse
+from pydantic import BaseModel
+
+from nat.middleware.function_middleware import CallNext
+from nat.middleware.function_middleware import FunctionMiddleware
+from nat.middleware.function_middleware import FunctionMiddlewareContext
+
+logger = logging.getLogger(__name__)
+
+
+class RedTeamingMiddleware(FunctionMiddleware):
+    """Middleware for red teaming that intercepts and modifies function inputs/outputs.
+
+    This middleware enables systematic security testing by injecting attack payloads
+    into function inputs or outputs. It supports flexible targeting, field-level
+    modifications, and multiple attack modes.
+
+    Features:
+
+    * Target specific functions or entire function groups
+    * Search for specific fields in input/output schemas
+    * Apply attacks via replace or append modes
+    * Support for both regular and streaming calls
+    * Type-safe operations on strings, numbers
+
+    Example::
+
+        # In YAML config
+        middleware:
+          prompt_injection:
+            _type: red_teaming
+            attack_payload: "Ignore previous instructions"
+            target_function_or_group: my_llm.generate
+            payload_placement: append_start
+            target_location: input
+            target_field: prompt
+
+    Args:
+        attack_payload: The malicious payload to inject.
+        target_function_or_group: Function or group to target (None for all).
+        payload_placement: How to apply (replace, append_start, append_middle, append_end).
+        target_location: Whether to attack input or output.
+        target_field: Field name or path to attack (None for direct value).
+    """
+
+    def __init__(
+        self,
+        *,
+        attack_payload: str,
+        target_function_or_group: str | None = None,
+        payload_placement: Literal["replace", "append_start", "append_middle", "append_end"] = "append_end",
+        target_location: Literal["input", "output"] = "input",
+        target_field: str | None = None,
+        target_field_resolution_strategy: Literal["random", "first", "last", "all", "error"] = "error",
+        call_limit: int | None = None,
+    ) -> None:
+        """Initialize red teaming middleware.
+
+        Args:
+            attack_payload: The value to inject to the function input or output.
+            target_function_or_group: Optional function/group to target.
+            payload_placement: How to apply the payload (replace or append modes).
+            target_location: Whether to place the payload in the input or output.
+            target_field: JSONPath to the field to attack.
+            target_field_resolution_strategy: Strategy (random/first/last/all/error).
+            call_limit: Maximum number of times the middleware will apply a payload.
+        """
+        super().__init__(is_final=False)
+        self._attack_payload = attack_payload
+        self._target_function_or_group = target_function_or_group
+        self._payload_placement = payload_placement
+        self._target_location = target_location
+        self._target_field = target_field
+        self._target_field_resolution_strategy = target_field_resolution_strategy
+        self._call_count: int = 0  # Count the number of times the middleware has applied a payload
+        self._call_limit = call_limit
+        logger.info(
+            "RedTeamingMiddleware initialized: payload=%s, target=%s, placement=%s, location=%s, field=%s",
+            attack_payload,
+            target_function_or_group,
+            payload_placement,
+            target_location,
+            target_field,
+        )
+
+    def _should_apply_payload(self, context_name: str) -> bool:
+        """Check if this function should be attacked based on targeting configuration.
+
+        Args:
+            context_name: The name of the function from context (e.g., "calculator.add")
+
+        Returns:
+            True if the function should be attacked, False otherwise
+        """
+        # If no target specified, attack all functions
+        if self._target_function_or_group is None:
+            return True
+
+        target = self._target_function_or_group
+
+        # Group targeting - match if context starts with the group name
+        # Handle both "group.function" and just "function" in context
+        if "." in context_name and "." not in target:
+            context_group = context_name.split(".", 1)[0]
+            return context_group == target
+
+        if context_name == "<workflow>":
+            return target in {"<workflow>", "workflow"}
+
+        # If context has no dot, match if it equals the target exactly
+        return context_name == target
+
+    def _find_middle_sentence_index(self, text: str) -> int:
+        """Find the index to insert text at the middle sentence boundary.
+
+        Args:
+            text: The text to analyze
+
+        Returns:
+            The character index where the middle sentence ends
+        """
+        # Find all sentence boundaries using regex
+        # Match sentence-ending punctuation followed by space/newline or end of string
+        sentence_pattern = r"[.!?](?:\s+|$)"
+        matches = list(re.finditer(sentence_pattern, text))
+
+        if not matches:
+            # No sentence boundaries found, insert at middle character
+            return len(text) // 2
+
+        # Find the sentence boundary closest to the middle
+        text_midpoint = len(text) // 2
+        closest_match = min(matches, key=lambda m: abs(m.end() - text_midpoint))
+
+        return closest_match.end()
+
+    def _apply_payload_to_simple_type(self,
+                                      original_value: list | str | int | float,
+                                      attack_payload: str,
+                                      payload_placement: str) -> Any:
+        """Apply the attack payload to simple types (str, int, float) value.
+
+        Args:
+            original_value: The original value to attack
+            attack_payload: The payload to inject
+            payload_placement: How to apply the payload
+
+        Returns:
+            The modified value with attack applied
+
+        Raises:
+            ValueError: If attack cannot be applied due to type mismatch
+        """
+        # Determine actual type from value if not provided
+        value_type = type(original_value)
+
+        # Handle string attacks
+        if value_type is str or isinstance(original_value, str):
+            original_str = str(original_value)
+
+            if payload_placement == "replace":
+                return attack_payload
+            elif payload_placement == "append_start":
+                return f"{attack_payload}{original_str}"
+            elif payload_placement == "append_end":
+                return f"{original_str}{attack_payload}"
+            elif payload_placement == "append_middle":
+                insert_index = self._find_middle_sentence_index(original_str)
+                return f"{original_str[:insert_index]}{attack_payload}{original_str[insert_index:]}"
+            else:
+                raise ValueError(f"Unknown payload placement: {payload_placement}")
+
+        # Handle int/float attacks
+        if isinstance(original_value, int | float):
+            # For numbers, only replace is allowed
+            if payload_placement != "replace":
+                logger.warning(
+                    "Payload placement '%s' not supported for numeric types (int/float). "
+                    "Falling back to 'replace' mode for field with value %s",
+                    payload_placement,
+                    original_value,
+                )
+
+            # Convert payload to the appropriate numeric type
+            try:
+                if value_type is int or isinstance(original_value, int):
+                    return int(attack_payload)
+                return float(attack_payload)
+            except (ValueError, TypeError) as e:
+                raise ValueError(f"Cannot convert attack payload '{attack_payload}' to {value_type.__name__}") from e
+
+    def _resolve_multiple_field_matches(self, matches):
+        if self._target_field_resolution_strategy == "error":
+            raise ValueError(f"Multiple matches found for target_field: {self._target_field}")
+        elif self._target_field_resolution_strategy == "random":
+            return [random.choice(matches)]
+        elif self._target_field_resolution_strategy == "first":
+            return [matches[0]]
+        elif self._target_field_resolution_strategy == "last":
+            return [matches[-1]]
+        elif self._target_field_resolution_strategy == "all":
+            return matches
+        else:
+            raise ValueError(f"Unknown target_field_resolution_strategy: {self._target_field_resolution_strategy}")
+
+    def _apply_payload_to_complex_type(self, value: list | dict | BaseModel) -> list | dict | BaseModel:
+        if self._target_field is None:
+            if isinstance(value, BaseModel):
+                value_details = value.model_dump_json()
+            else:
+                value_details = ""
+            additional_info = ("Additional info: A pydantic BaseModel with fields:" +
+                               value_details if value_details else "")
+            raise ValueError("Applying an attack payload to complex type, requires a target_field. \n"
+                             f"Input value: {value}.: {value_details}. {additional_info} \n"
+                             "A target field can be specified in the middleware configuration as a jsonpath.")
+
+        # Convert BaseModel to dict for jsonpath processing
+        original_type = type(value)
+        is_basemodel = isinstance(value, BaseModel)
+        if is_basemodel:
+            value_to_modify = value.model_dump()
+        else:
+            value_to_modify = value
+
+        jsonpath_expr = parse(self._target_field)
+        matches = jsonpath_expr.find(value_to_modify)
+        if len(matches) == 0:
+            raise ValueError(f"No matches found for target_field: {self._target_field} in value: {value}")
+        if len(matches) > 1:
+            matches = self._resolve_multiple_field_matches(matches)
+        else:
+            matches = [matches[0]]
+        modified_values = [
+            self._apply_payload_to_simple_type(match.value, self._attack_payload, self._payload_placement)
+            for match in matches
+        ]
+        for match, modified_value in zip(matches, modified_values):
+            match.full_path.update(value_to_modify, modified_value)
+
+        # Reconstruct BaseModel if original was BaseModel
+        if is_basemodel:
+            assert isinstance(value_to_modify, dict)
+            return cast(type[BaseModel], original_type)(**value_to_modify)
+        return value_to_modify
+
+    def _apply_payload_to_function_value(self, value: Any) -> Any:
+        if self._call_limit is not None and self._call_count >= self._call_limit:
+            logger.warning("Call limit reached for red teaming middleware. "
+                           "Not applying attack payload to value: %s",
+                           value)
+            return value
+        if isinstance(value, list | dict | BaseModel):
+            modified_value = self._apply_payload_to_complex_type(value)
+        elif isinstance(value, str | int | float):
+            modified_value = self._apply_payload_to_simple_type(value, self._attack_payload, self._payload_placement)
+        else:
+            raise ValueError(f"Unsupported function input/output type: {type(value).__name__}")
+        self._call_count += 1
+        return modified_value
+
+    def _apply_payload_to_function_value_with_exception(self, value: Any, context: FunctionMiddlewareContext) -> Any:
+        try:
+            return self._apply_payload_to_function_value(value)
+        except Exception as e:
+            logger.error("Failed to apply red team attack to function %s: %s", context.name, e, exc_info=True)
+            raise
+
+    async def function_middleware_invoke(self,
+                                         *args: Any,
+                                         call_next: CallNext,
+                                         context: FunctionMiddlewareContext,
+                                         **kwargs: Any) -> Any:
+        """Invoke middleware for single-output functions.
+
+        Args:
+            args: Positional arguments passed to the function (first arg is typically the input value).
+            call_next: Callable to invoke next middleware/function.
+            context: Metadata about the function being wrapped.
+            kwargs: Keyword arguments passed to the function.
+
+        Returns:
+            The output value (potentially modified if attacking output).
+        """
+        value = args[0] if args else None
+
+        # Check if we should attack this function
+        if not self._should_apply_payload(context.name):
+            logger.debug("Skipping function %s (not targeted)", context.name)
+            return await call_next(value, *args[1:], **kwargs)
+
+        if self._target_location == "input":
+            # Attack the input before calling the function
+            modified_input = self._apply_payload_to_function_value_with_exception(value, context)
+            # Call next with modified input
+            return await call_next(modified_input, *args[1:], **kwargs)
+
+        elif self._target_location == "output":  # target_location == "output"
+            # Call function first, then attack the output
+            output = await call_next(value, *args[1:], **kwargs)
+            modified_output = self._apply_payload_to_function_value_with_exception(output, context)
+            return modified_output
+        else:
+            raise ValueError(f"Unknown target_location: {self._target_location}. "
+                             "Attack payloads can only be applied to function input or output.")
+
+
+__all__ = ["RedTeamingMiddleware"]