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

nat/middleware/cache_middleware.py

@@ -0,0 +1,256 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025, 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.
+"""Cache middleware for function memoization with similarity matching.
+
+This module provides a cache middleware that memoizes function calls based on
+input similarity. It demonstrates the middleware pattern by:
+
+1. Preprocessing: Serializing and checking the cache for similar inputs
+2. Calling next: Delegating to the next middleware/function if no cache hit
+3. Postprocessing: Caching the result for future use
+4. Continuing: Returning the result (cached or fresh)
+
+The cache supports exact matching for maximum performance and fuzzy matching
+using Python's built-in difflib for similarity computation.
+"""
+
+from __future__ import annotations
+
+import json
+import logging
+from collections.abc import AsyncIterator
+from typing import Any
+from typing import Literal
+
+from pydantic import Field
+
+from nat.builder.context import Context
+from nat.builder.context import ContextState
+from nat.data_models.middleware import FunctionMiddlewareBaseConfig
+from nat.middleware.function_middleware import CallNext
+from nat.middleware.function_middleware import CallNextStream
+from nat.middleware.function_middleware import FunctionMiddleware
+from nat.middleware.function_middleware import FunctionMiddlewareContext
+
+logger = logging.getLogger(__name__)
+
+
+class CacheMiddleware(FunctionMiddleware):
+    """Cache middleware that memoizes function outputs based on input similarity.
+
+    This middleware demonstrates the four-phase middleware pattern:
+
+    1. **Preprocess**: Serialize input and check cache for similar entries
+    2. **Call Next**: Delegate to next middleware/function if cache miss
+    3. **Postprocess**: Store the result in cache for future use
+    4. **Continue**: Return the result (from cache or fresh)
+
+    The cache serializes function inputs to strings and performs similarity
+    matching against previously seen inputs. If a similar input is found above
+    the configured threshold, it returns the cached output without calling the
+    next middleware or function.
+
+    Args:
+        enabled_mode: Either "always" to always cache, or "eval" to only
+            cache when Context.is_evaluating is True.
+        similarity_threshold: Float between 0 and 1. If 1.0, performs
+            exact string matching. Otherwise uses difflib for similarity
+            computation.
+    """
+
+    def __init__(self, *, enabled_mode: str, similarity_threshold: float) -> None:
+        """Initialize the cache middleware.
+
+        Args:
+            enabled_mode: Either "always" or "eval". If "eval", only caches
+                when Context.is_evaluating is True.
+            similarity_threshold: Similarity threshold between 0 and 1.
+                If 1.0, performs exact matching. Otherwise uses fuzzy matching.
+        """
+        super().__init__(is_final=True)
+        self._enabled_mode = enabled_mode
+        self._similarity_threshold = similarity_threshold
+        self._cache: dict[str, Any] = {}
+
+    def _should_cache(self) -> bool:
+        """Check if caching should be enabled based on the current context."""
+        if self._enabled_mode == "always":
+            return True
+
+        # Get the current context and check if we're in evaluation mode
+        try:
+            context_state = ContextState.get()
+            context = Context(context_state)
+            return context.is_evaluating
+        except Exception:
+            logger.warning("Failed to get context for cache decision", exc_info=True)
+            return False
+
+    def _serialize_input(self, value: Any) -> str | None:
+        """Serialize the input value to a string for caching.
+
+        Args:
+            value: The input value to serialize.
+
+        Returns:
+            String representation of the input, or None if serialization
+            fails.
+        """
+        try:
+            # Try JSON serialization first for best results
+            return json.dumps(value, sort_keys=True, default=str)
+        except Exception:
+            logger.debug("Failed to serialize input for caching", exc_info=True)
+            return None
+
+    def _find_similar_key(self, input_str: str) -> str | None:
+        """Find a cached key that is similar to the input string.
+
+        Args:
+            input_str: The serialized input string to match.
+
+        Returns:
+            The most similar cached key if above threshold, None otherwise.
+        """
+        if self._similarity_threshold == 1.0:
+            # Exact matching - fast path
+            return input_str if input_str in self._cache else None
+
+        # Fuzzy matching using difflib
+        import difflib
+
+        best_match = None
+        best_ratio = 0.0
+
+        for cached_key in self._cache:
+            # Use SequenceMatcher for similarity computation
+            matcher = difflib.SequenceMatcher(None, input_str, cached_key)
+            ratio = matcher.ratio()
+
+            if ratio >= self._similarity_threshold and ratio > best_ratio:
+                best_ratio = ratio
+                best_match = cached_key
+
+        return best_match
+
+    async def function_middleware_invoke(self, value: Any, call_next: CallNext,
+                                         context: FunctionMiddlewareContext) -> Any:
+        """Cache middleware for single-output invocations.
+
+        Implements the four-phase middleware pattern:
+
+        1. **Preprocess**: Check if caching is enabled and serialize input
+        2. **Call Next**: Delegate to next middleware/function if cache miss
+        3. **Postprocess**: Store the result in cache
+        4. **Continue**: Return the result (cached or fresh)
+
+        Args:
+            value: The input value to process
+            call_next: Callable to invoke the next middleware or function
+            context: Metadata about the function being wrapped
+
+        Returns:
+            The cached output if found, otherwise the fresh output
+        """
+        # Phase 1: Preprocess - check if caching should be enabled
+        if not self._should_cache():
+            return await call_next(value)
+
+        # Phase 1: Preprocess - serialize the input
+        input_str = self._serialize_input(value)
+        if input_str is None:
+            # Can't serialize, pass through to next middleware/function
+            logger.debug("Could not serialize input for function %s, bypassing cache", context.name)
+            return await call_next(value)
+
+        # Phase 1: Preprocess - look for a similar cached input
+        similar_key = self._find_similar_key(input_str)
+        if similar_key is not None:
+            # Cache hit - short-circuit and return cached output
+            logger.debug("Cache hit for function %s with similarity %.2f",
+                         context.name,
+                         1.0 if similar_key == input_str else self._similarity_threshold)
+            # Phase 4: Continue - return cached result
+            return self._cache[similar_key]
+
+        # Phase 2: Call next - no cache hit, call next middleware/function
+        logger.debug("Cache miss for function %s", context.name)
+        result = await call_next(value)
+
+        # Phase 3: Postprocess - cache the result for future use
+        self._cache[input_str] = result
+        logger.debug("Cached result for function %s", context.name)
+
+        # Phase 4: Continue - return the fresh result
+        return result
+
+    async def function_middleware_stream(self,
+                                         value: Any,
+                                         call_next: CallNextStream,
+                                         context: FunctionMiddlewareContext) -> AsyncIterator[Any]:
+        """Cache middleware for streaming invocations - bypasses caching.
+
+        Streaming results are not cached as they would need to be buffered
+        entirely in memory, which would defeat the purpose of streaming.
+
+        This method demonstrates the middleware pattern for streams:
+
+        1. **Preprocess**: Log that we're bypassing cache
+        2. **Call Next**: Get stream from next middleware/function
+        3. **Process Chunks**: Yield each chunk as it arrives
+        4. **Continue**: Complete the stream
+
+        Args:
+            value: The input value to process
+            call_next: Callable to invoke the next middleware or function stream
+            context: Metadata about the function being wrapped
+
+        Yields:
+            Chunks from the stream (unmodified)
+        """
+        # Phase 1: Preprocess - log that we're bypassing cache for streams
+        logger.debug("Streaming call for function %s, bypassing cache", context.name)
+
+        # Phase 2-3: Call next and process chunks - yield chunks as they arrive
+        async for chunk in call_next(value):
+            yield chunk
+
+        # Phase 4: Continue - stream is complete (implicit)
+
+
+class CacheMiddlewareConfig(FunctionMiddlewareBaseConfig, name="cache"):
+    """Configuration for cache middleware.
+
+    The cache middleware memoizes function outputs based on input similarity,
+    with support for both exact and fuzzy matching.
+
+    Args:
+        enabled_mode: Controls when caching is active:
+            - "always": Cache is always enabled
+            - "eval": Cache only active when Context.is_evaluating is True
+        similarity_threshold: Float between 0 and 1 for input matching:
+            - 1.0: Exact string matching (fastest)
+            - < 1.0: Fuzzy matching using difflib similarity
+    """
+
+    enabled_mode: Literal["always", "eval"] = Field(
+        default="eval", description="When caching is enabled: 'always' or 'eval' (only during evaluation)")
+    similarity_threshold: float = Field(default=1.0,
+                                        ge=0.0,
+                                        le=1.0,
+                                        description="Similarity threshold between 0 and 1. Use 1.0 for exact matching")
+
+
+__all__ = ["CacheMiddleware", "CacheMiddlewareConfig"]

nat/middleware/function_middleware.py

@@ -0,0 +1,186 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025, 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.
+"""Function-specific middleware for the NeMo Agent toolkit.
+
+This module provides function-specific middleware implementations that extend
+the base Middleware class. FunctionMiddleware is a specialized middleware type
+designed specifically for wrapping function calls with dedicated methods
+for function-specific preprocessing and postprocessing.
+
+Middleware is configured at registration time and is bound to instances when they
+are constructed by the workflow builder.
+
+Middleware executes in the order provided and can optionally be marked as *final*.
+A final middleware terminates the chain, preventing subsequent middleware or the
+wrapped target from running unless the final middleware explicitly delegates to
+the next callable.
+"""
+
+from __future__ import annotations
+
+from collections.abc import AsyncIterator
+from collections.abc import Sequence
+from typing import Any
+
+from nat.middleware.middleware import CallNext
+from nat.middleware.middleware import CallNextStream
+from nat.middleware.middleware import FunctionMiddlewareContext
+from nat.middleware.middleware import Middleware
+
+
+class FunctionMiddleware(Middleware):
+    """Specialized middleware for function-specific wrapping.
+
+    This class extends the base Middleware class and provides function-specific
+    wrapping methods. Functions that use this middleware type will call
+    ``function_middleware_invoke`` and ``function_middleware_stream`` instead of
+    the base ``middleware_invoke`` and ``middleware_stream`` methods.
+    """
+
+    async def middleware_invoke(self, value: Any, call_next: CallNext, context: FunctionMiddlewareContext) -> Any:
+        """Delegate to function_middleware_invoke for function-specific handling."""
+        return await self.function_middleware_invoke(value, call_next, context)
+
+    async def middleware_stream(self, value: Any, call_next: CallNextStream,
+                                context: FunctionMiddlewareContext) -> AsyncIterator[Any]:
+        """Delegate to function_middleware_stream for function-specific handling."""
+        async for chunk in self.function_middleware_stream(value, call_next, context):
+            yield chunk
+
+    async def function_middleware_invoke(self, value: Any, call_next: CallNext,
+                                         context: FunctionMiddlewareContext) -> Any:
+        """Function-specific middleware for single-output invocations.
+
+        Args:
+            value: The input value to process
+            call_next: Callable to invoke the next middleware or function
+            context: Metadata about the function being wrapped
+
+        Returns:
+            The (potentially modified) output from the function
+
+        The default implementation simply delegates to ``call_next``. Override this
+        in subclasses to add function-specific preprocessing and postprocessing.
+        """
+        return await call_next(value)
+
+    async def function_middleware_stream(self,
+                                         value: Any,
+                                         call_next: CallNextStream,
+                                         context: FunctionMiddlewareContext) -> AsyncIterator[Any]:
+        """Function-specific middleware for streaming invocations.
+
+        Args:
+            value: The input value to process
+            call_next: Callable to invoke the next middleware or function stream
+            context: Metadata about the function being wrapped
+
+        Yields:
+            Chunks from the stream (potentially modified)
+
+        The default implementation forwards to ``call_next`` untouched. Override this
+        in subclasses to add function-specific preprocessing and chunk transformations.
+        """
+        async for chunk in call_next(value):
+            yield chunk
+
+
+class FunctionMiddlewareChain:
+    """Utility that composes middleware-style callables.
+
+    This class builds a chain of middleware that executes in order,
+    with each middleware able to preprocess inputs, call the next middleware,
+    and postprocess outputs.
+    """
+
+    def __init__(self, *, middleware: Sequence[Middleware], context: FunctionMiddlewareContext) -> None:
+        self._middleware = tuple(middleware)
+        self._context = context
+
+    def build_single(self, final_call: CallNext) -> CallNext:
+        """Build the middleware chain for single-output invocations.
+
+        Args:
+            final_call: The final function to call (the actual function implementation)
+
+        Returns:
+            A callable that executes the entire middleware chain
+        """
+        call = final_call
+
+        for mw in reversed(self._middleware):
+            call_next = call
+
+            async def wrapped(value: Any, *, _middleware: Middleware = mw, _call_next: CallNext = call_next) -> Any:
+                return await _middleware.middleware_invoke(value, _call_next, self._context)
+
+            call = wrapped
+
+        return call
+
+    def build_stream(self, final_call: CallNextStream) -> CallNextStream:
+        """Build the middleware chain for streaming invocations.
+
+        Args:
+            final_call: The final function to call (the actual function implementation)
+
+        Returns:
+            A callable that executes the entire middleware chain
+        """
+        call = final_call
+
+        for mw in reversed(self._middleware):
+            call_next = call
+
+            async def wrapped(value: Any,
+                              *,
+                              _middleware: Middleware = mw,
+                              _call_next: CallNextStream = call_next) -> AsyncIterator[Any]:
+                async for chunk in _middleware.middleware_stream(value, _call_next, self._context):
+                    yield chunk
+
+            call = wrapped
+
+        return call
+
+
+def validate_middleware(middleware: Sequence[Middleware] | None) -> tuple[Middleware, ...]:
+    """Validate a sequence of middleware, enforcing ordering guarantees."""
+
+    if not middleware:
+        return tuple()
+
+    final_found = False
+    for idx, mw in enumerate(middleware):
+        if not isinstance(mw, Middleware):
+            raise TypeError("All middleware must be instances of Middleware")
+
+        if mw.is_final:
+            if final_found:
+                raise ValueError("Only one final Middleware may be specified per function")
+
+            if idx != len(middleware) - 1:
+                raise ValueError("A final Middleware must be the last middleware in the chain")
+
+            final_found = True
+
+    return tuple(middleware)
+
+
+__all__ = [
+    "FunctionMiddleware",
+    "FunctionMiddlewareChain",
+    "validate_middleware",
+]

nat/middleware/middleware.py

@@ -0,0 +1,184 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025, 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.
+"""Base middleware class for the NeMo Agent toolkit.
+
+This module provides the base Middleware class that defines the middleware pattern
+for wrapping and modifying function calls. Middleware works like middleware in
+web frameworks - they can modify inputs, call the next middleware in the chain,
+process outputs, and continue.
+"""
+
+from __future__ import annotations
+
+import dataclasses
+from abc import ABC
+from collections.abc import AsyncIterator
+from collections.abc import Awaitable
+from collections.abc import Callable
+from typing import Any
+
+from pydantic import BaseModel
+
+#: Type alias for single-output invocation callables.
+CallNext = Callable[[Any], Awaitable[Any]]
+
+#: Type alias for streaming invocation callables.
+CallNextStream = Callable[[Any], AsyncIterator[Any]]
+
+
+@dataclasses.dataclass(frozen=True, kw_only=True)
+class FunctionMiddlewareContext:
+    """Context information 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
+    function's name, configuration, schema, etc.
+    """
+
+    name: str
+    """Name of the function being wrapped."""
+
+    config: Any
+    """Configuration object for the function."""
+
+    description: str | None
+    """Optional description of the function."""
+
+    input_schema: type[BaseModel] | None
+    """Schema describing expected inputs or :class:`NoneType` when absent."""
+
+    single_output_schema: type[BaseModel] | type[None]
+    """Schema describing single outputs or :class:`types.NoneType` when absent."""
+
+    stream_output_schema: type[BaseModel] | type[None]
+    """Schema describing streaming outputs or :class:`types.NoneType` when absent."""
+
+
+class Middleware(ABC):
+    """Base class for middleware-style wrapping.
+
+    Middleware works like middleware in web frameworks:
+
+    1. **Preprocess**: Inspect and optionally modify inputs
+    2. **Call Next**: Delegate to the next middleware or the target itself
+    3. **Postprocess**: Process, transform, or augment the output
+    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)
+
+                # 3. Postprocess
+                print(f"Output: {result}")
+
+                # 4. Continue
+                return result
+
+    Attributes:
+        is_final: If True, this middleware terminates the chain. No subsequent
+            middleware or the target will be called unless this middleware
+            explicitly delegates to ``call_next``.
+    """
+
+    def __init__(self, *, is_final: bool = False) -> None:
+        self._is_final = is_final
+
+    @property
+    def is_final(self) -> bool:
+        """Whether this middleware terminates the chain.
+
+        A final middleware prevents subsequent middleware and the target
+        from running unless it explicitly calls ``call_next``.
+        """
+
+        return self._is_final
+
+    async def middleware_invoke(self, value: Any, call_next: CallNext, context: FunctionMiddlewareContext) -> 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
+
+        Returns:
+            The (potentially modified) output from the target
+
+        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):
+                # Preprocess: modify input
+                modified_input = transform(value)
+
+                # Call next: delegate to next middleware/target
+                result = await call_next(modified_input)
+
+                # Postprocess: modify output
+                modified_result = transform_output(result)
+
+                # Continue: return final result
+                return modified_result
+        """
+
+        del context  # Unused by the default implementation.
+        return await call_next(value)
+
+    async def middleware_stream(self, value: Any, call_next: CallNextStream,
+                                context: FunctionMiddlewareContext) -> 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
+
+        Yields:
+            Chunks from the stream (potentially modified)
+
+        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):
+                # 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):
+                    # Process each chunk
+                    modified_chunk = transform_chunk(chunk)
+                    yield modified_chunk
+
+                # Postprocess: cleanup after stream ends
+                await cleanup()
+        """
+
+        del context  # Unused by the default implementation.
+        async for chunk in call_next(value):
+            yield chunk
+
+
+__all__ = [
+    "CallNext",
+    "CallNextStream",
+    "Middleware",
+    "FunctionMiddlewareContext",
+]

nat/middleware/register.py

@@ -0,0 +1,35 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025, 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.
+"""Registration module for built-in middleware."""
+
+from __future__ import annotations
+
+from nat.cli.register_workflow import register_middleware
+from nat.middleware.cache_middleware import CacheMiddleware
+from nat.middleware.cache_middleware import CacheMiddlewareConfig
+
+
+@register_middleware(config_type=CacheMiddlewareConfig)
+async def cache_middleware(config: CacheMiddlewareConfig, builder):
+    """Build a cache middleware from configuration.
+
+    Args:
+        config: The cache middleware configuration
+        builder: The workflow builder (unused but required by component pattern)
+
+    Yields:
+        A configured cache middleware instance
+    """
+    yield CacheMiddleware(enabled_mode=config.enabled_mode, similarity_threshold=config.similarity_threshold)

nat/profiler/decorators/framework_wrapper.py

@@ -34,6 +34,7 @@ _library_instrumented = {
     "semantic_kernel": False,
     "agno": False,
     "adk": False,
+    "strands": False,
 }
 
 callback_handler_var: ContextVar[Any | None] = ContextVar("callback_handler_var", default=None)
@@ -131,6 +132,21 @@ def set_framework_profiler_handler(
                     _library_instrumented["adk"] = True
                     logger.debug("ADK callback handler registered")
 
+            if (LLMFrameworkEnum.STRANDS in frameworks and not _library_instrumented["strands"]):
+                try:
+                    from nat.plugins.strands.strands_callback_handler import StrandsProfilerHandler
+                except ImportError as e:
+                    logger.warning(
+                        "Strands profiler not available. Install NAT with Strands extras: "
+                        "pip install \"nvidia-nat[strands]\". Error: %s",
+                        e,
+                    )
+                else:
+                    handler = StrandsProfilerHandler()
+                    handler.instrument()
+                    _library_instrumented["strands"] = True
+                    logger.debug("Strands callback handler registered")
+
             # IMPORTANT: actually call the wrapped function as an async context manager
             async with func(workflow_config, builder) as result:
                 yield result