nvidia-nat 1.3.0a20250827__py3-none-any.whl → 1.3.0a20250828__py3-none-any.whl

nat/llm/aws_bedrock_llm.py

@@ -23,9 +23,11 @@ from nat.cli.register_workflow import register_llm_provider
 from nat.data_models.llm import LLMBaseConfig
 from nat.data_models.retry_mixin import RetryMixin
 from nat.data_models.temperature_mixin import TemperatureMixin
+from nat.data_models.thinking_mixin import ThinkingMixin
+from nat.data_models.top_p_mixin import TopPMixin
 
 
-class AWSBedrockModelConfig(LLMBaseConfig, RetryMixin, TemperatureMixin, name="aws_bedrock"):
+class AWSBedrockModelConfig(LLMBaseConfig, RetryMixin, TemperatureMixin, TopPMixin, ThinkingMixin, name="aws_bedrock"):
     """An AWS Bedrock llm provider to be used with an LLM client."""
 
     model_config = ConfigDict(protected_namespaces=())
@@ -34,14 +36,13 @@ class AWSBedrockModelConfig(LLMBaseConfig, RetryMixin, TemperatureMixin, name="a
     model_name: str = Field(validation_alias=AliasChoices("model_name", "model"),
                             serialization_alias="model",
                             description="The model name for the hosted AWS Bedrock.")
-    max_tokens: int | None = Field(default=1024,
-                                   gt=0,
-                                   description="Maximum number of tokens to generate."
-                                   "This field is ONLY required when using AWS Bedrock with Langchain.")
-    context_size: int | None = Field(default=1024,
-                                     gt=0,
-                                     description="Maximum number of tokens to generate."
-                                     "This field is ONLY required when using AWS Bedrock with LlamaIndex.")
+    max_tokens: int | None = Field(default=1024, gt=0, description="Maximum number of tokens to generate.")
+    context_size: int | None = Field(
+        default=1024,
+        gt=0,
+        description="The maximum number of tokens available for input. This is only required for LlamaIndex. "
+        "This field is ignored for Langchain.",
+    )
 
     # Client parameters
     region_name: str | None = Field(default="None", description="AWS region to use.")

nat/llm/azure_openai_llm.py

@@ -23,10 +23,18 @@ from nat.cli.register_workflow import register_llm_provider
 from nat.data_models.llm import LLMBaseConfig
 from nat.data_models.retry_mixin import RetryMixin
 from nat.data_models.temperature_mixin import TemperatureMixin
+from nat.data_models.thinking_mixin import ThinkingMixin
 from nat.data_models.top_p_mixin import TopPMixin
 
 
-class AzureOpenAIModelConfig(LLMBaseConfig, RetryMixin, TemperatureMixin, TopPMixin, name="azure_openai"):
+class AzureOpenAIModelConfig(
+        LLMBaseConfig,
+        RetryMixin,
+        TemperatureMixin,
+        TopPMixin,
+        ThinkingMixin,
+        name="azure_openai",
+):
     """An Azure OpenAI LLM provider to be used with an LLM client."""
 
     model_config = ConfigDict(protected_namespaces=(), extra="allow")

nat/llm/nim_llm.py

@@ -24,10 +24,11 @@ from nat.cli.register_workflow import register_llm_provider
 from nat.data_models.llm import LLMBaseConfig
 from nat.data_models.retry_mixin import RetryMixin
 from nat.data_models.temperature_mixin import TemperatureMixin
+from nat.data_models.thinking_mixin import ThinkingMixin
 from nat.data_models.top_p_mixin import TopPMixin
 
 
-class NIMModelConfig(LLMBaseConfig, RetryMixin, TemperatureMixin, TopPMixin, name="nim"):
+class NIMModelConfig(LLMBaseConfig, RetryMixin, TemperatureMixin, TopPMixin, ThinkingMixin, name="nim"):
     """An NVIDIA Inference Microservice (NIM) llm provider to be used with an LLM client."""
 
     model_config = ConfigDict(protected_namespaces=())

nat/llm/openai_llm.py

@@ -23,10 +23,11 @@ from nat.cli.register_workflow import register_llm_provider
 from nat.data_models.llm import LLMBaseConfig
 from nat.data_models.retry_mixin import RetryMixin
 from nat.data_models.temperature_mixin import TemperatureMixin
+from nat.data_models.thinking_mixin import ThinkingMixin
 from nat.data_models.top_p_mixin import TopPMixin
 
 
-class OpenAIModelConfig(LLMBaseConfig, RetryMixin, TemperatureMixin, TopPMixin, name="openai"):
+class OpenAIModelConfig(LLMBaseConfig, RetryMixin, TemperatureMixin, TopPMixin, ThinkingMixin, name="openai"):
     """An OpenAI LLM provider to be used with an LLM client."""
 
     model_config = ConfigDict(protected_namespaces=(), extra="allow")

nat/llm/utils/thinking.py

@@ -0,0 +1,215 @@
+# 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.
+
+import functools
+import inspect
+import logging
+import types
+from abc import abstractmethod
+from collections.abc import AsyncGenerator
+from collections.abc import Iterable
+from dataclasses import dataclass
+from typing import Any
+from typing import Callable
+from typing import TypeVar
+
+ModelType = TypeVar("ModelType")
+MessagesType = TypeVar("MessagesType")
+
+logger = logging.getLogger(__name__)
+
+
+class FunctionArgumentWrapper:
+    """
+    Wrapper for the arguments and keyword arguments of a function.
+
+    The arguments and keyword arguments are stored in the args and kwargs attributes, respectively.
+    """
+
+    def __init__(self, *args: Any, **kwargs: Any):
+        """
+        Initialize the FunctionArgumentWrapper.
+
+        Args:
+            args: The arguments to the function.
+            kwargs: The keyword arguments to the function.
+        """
+        self.args = args
+        self.kwargs = kwargs
+
+    def __repr__(self) -> str:
+        return f"FunctionArgumentWrapper(args={self.args}, kwargs={self.kwargs})"
+
+
+@dataclass
+class BaseThinkingInjector:
+    """
+    Base class for thinking injectors.
+
+    Args:
+        system_prompt: The system prompt to inject.
+        function_names: The function names to inject the system prompt into.
+    """
+
+    system_prompt: str
+    function_names: list[str]
+
+    @abstractmethod
+    def inject(self, *args, **kwargs) -> FunctionArgumentWrapper:
+        """
+        Inject the system prompt into the arguments.
+
+        Args:
+            args: The arguments to inject the system prompt into.
+            kwargs: The keyword arguments to inject the system prompt into.
+
+        Returns:
+            FunctionArgumentWrapper: An object that contains the transformed args and kwargs.
+        """
+        pass
+
+
+def _make_thinking_decorator(injector: BaseThinkingInjector):
+
+    def decorate(fn: Callable[..., Any]) -> Callable[..., Any]:
+
+        async def _call_async(obj: object, *call_args, **call_kwargs) -> Any:
+            new_args = injector.inject(*call_args, **call_kwargs)
+            return await fn(obj, *new_args.args, **new_args.kwargs)
+
+        async def _agen(obj: object, *call_args, **call_kwargs) -> AsyncGenerator[Any, None]:
+            new_args = injector.inject(*call_args, **call_kwargs)
+            async for item in fn(obj, *new_args.args, **new_args.kwargs):
+                yield item
+
+        def _gen(obj: object, *call_args, **call_kwargs) -> Iterable[Any]:
+            new_args = injector.inject(*call_args, **call_kwargs)
+            yield from fn(obj, *new_args.args, **new_args.kwargs)
+            return
+
+        def _sync(obj: object, *call_args, **call_kwargs) -> Any:
+            new_args = injector.inject(*call_args, **call_kwargs)
+            return fn(obj, *new_args.args, **new_args.kwargs)
+
+        # Decide which wrapper to return
+        if inspect.iscoroutinefunction(fn):
+            wrapper = _call_async
+        elif inspect.isasyncgenfunction(fn):
+            wrapper = _agen
+        elif inspect.isgeneratorfunction(fn):
+            wrapper = _gen
+        else:
+            wrapper = _sync
+
+        return functools.wraps(fn)(wrapper)
+
+    return decorate
+
+
+def patch_with_thinking(obj: ModelType, injector: BaseThinkingInjector) -> ModelType:
+    """
+    Patch the given object with a decorator that injects a system prompt into the supplied messages.
+    There is an assumption that the first non-object argument is the messages.
+
+    Args:
+        obj: The object to patch.
+        injector: The injector to use.
+
+    Returns:
+        The patched object.
+
+    Examples:
+        >>> from nat.llm.utils.thinking import BaseThinkingInjector
+        >>> from nat.llm.utils.thinking import FunctionArgumentWrapper
+        >>> from nat.llm.utils.thinking import patch_with_thinking
+        >>>
+        >>> class MockClass:
+        ...     def sync_method(self, *args, **kwargs):
+        ...         return (args, kwargs)
+        ...
+        >>> mock_obj_1 = MockClass()
+        >>> class AddThinking(BaseThinkingInjector):
+        ...     def inject(self, x: str, *args, **kwargs) -> FunctionArgumentWrapper:
+        ...         return FunctionArgumentWrapper(("thinking " + x), *args, **kwargs)
+        >>>
+        >>> patched_obj = patch_with_thinking(mock_obj_1, AddThinking(
+        ...     system_prompt="thinking",
+        ...     function_names=["sync_method"],
+        ... ))
+        >>> patched_obj.sync_method("test", 1, 2, 3, foo="bar")
+        (('thinking test', 1, 2, 3), {'foo': 'bar'})
+        >>>
+        >>> mock_obj_2 = MockClass()
+        >>> class AddThinkingWithArgs(BaseThinkingInjector):
+        ...     def inject(self, *args, **kwargs) -> FunctionArgumentWrapper:
+        ...         return FunctionArgumentWrapper("thinking", *args, **kwargs)
+        >>>
+        >>> patched_obj = patch_with_thinking(mock_obj_2, AddThinkingWithArgs(
+        ...     system_prompt="thinking",
+        ...     function_names=["sync_method"],
+        ... ))
+        >>> patched_obj.sync_method("test", 1, 2, 3, foo="bar")
+        (('thinking', 'test', 1, 2, 3), {'foo': 'bar'})
+        >>>
+        >>> mock_obj_3 = MockClass()
+        >>> class AddThinkingWithKwargs(BaseThinkingInjector):
+        ...     def inject(self, *args, **kwargs) -> FunctionArgumentWrapper:
+        ...         return FunctionArgumentWrapper(*args, thinking=True, **kwargs)
+        >>>
+        >>> patched_obj = patch_with_thinking(mock_obj_3, AddThinkingWithKwargs(
+        ...     system_prompt="thinking",
+        ...     function_names=["sync_method"],
+        ... ))
+        >>> patched_obj.sync_method("test", 1, 2, 3, foo="bar")
+        (('test', 1, 2, 3), {'thinking': True, 'foo': 'bar'})
+    """
+
+    decorator = _make_thinking_decorator(injector)
+
+    cls = obj if inspect.isclass(obj) else type(obj)
+    cls_name = getattr(cls, "__name__", str(cls))
+
+    for name, _ in inspect.getmembers(cls, callable):
+        if name not in injector.function_names:
+            continue
+
+        descriptor = inspect.getattr_static(cls, name)
+        original = descriptor.__func__ if isinstance(descriptor, types.MethodType) else descriptor
+        wrapped = decorator(original)
+
+        try:  # instance‑level first
+            if not inspect.isclass(obj):
+                object.__setattr__(obj, name, types.MethodType(wrapped, obj))
+                continue
+        except Exception as exc:
+            logger.info(
+                "Instance‑level patch failed for %s.%s (%s); "
+                "falling back to class‑level patch.",
+                cls_name,
+                name,
+                exc,
+            )
+
+        try:  # class‑level fallback
+            setattr(cls, name, wrapped)
+        except Exception as exc:
+            logger.info(
+                "Cannot patch method %s.%s with thinking: %s",
+                cls_name,
+                name,
+                exc,
+            )
+
+    return obj

nat/observability/processor/falsy_batch_filter_processor.py

@@ -0,0 +1,55 @@
+# 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.
+
+import logging
+from typing import TypeVar
+
+from nat.observability.processor.processor import Processor
+from nat.utils.type_utils import override
+
+logger = logging.getLogger(__name__)
+
+FalsyT = TypeVar("FalsyT")
+
+
+class FalsyBatchFilterProcessor(Processor[list[FalsyT], list[FalsyT]]):
+    """Processor that filters out falsy items from a batch."""
+
+    @override
+    async def process(self, item: list[FalsyT]) -> list[FalsyT]:
+        """Filter out falsy items from a batch.
+
+        Args:
+            item (list[FalsyT]): The batch of items to filter.
+
+        Returns:
+            list[FalsyT]: The filtered batch.
+        """
+        return [i for i in item if i]
+
+
+class DictBatchFilterProcessor(FalsyBatchFilterProcessor[dict]):
+    """Processor that filters out empty dict items from a batch."""
+    pass
+
+
+class ListBatchFilterProcessor(FalsyBatchFilterProcessor[list]):
+    """Processor that filters out empty list items from a batch."""
+    pass
+
+
+class SetBatchFilterProcessor(FalsyBatchFilterProcessor[set]):
+    """Processor that filters out empty set items from a batch."""
+    pass

nat/observability/processor/processor_factory.py

@@ -0,0 +1,70 @@
+# 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.
+
+from typing import Any
+
+from nat.observability.processor.processor import Processor
+
+
+def processor_factory(processor_class: type, from_type: type[Any], to_type: type[Any]) -> type[Processor]:
+    """Create a concrete processor class from a processor class and types.
+
+    Args:
+        processor_class (type): The processor class to create a concrete instance of
+        from_type (type[Any]): The type of the input data
+        to_type (type[Any]): The type of the output data
+
+    Returns:
+        type[Processor]: The concrete processor class
+    """
+
+    class ConcreteProcessor(processor_class[from_type, to_type]):  # type: ignore
+        pass
+
+    return ConcreteProcessor
+
+
+def processor_factory_from_type(processor_class: type, from_type: type[Any]) -> type[Processor]:
+    """Create a concrete processor class from a processor class and input type.
+
+    Args:
+        processor_class (type): The processor class to create a concrete instance of
+        from_type (type[Any]): The type of the input data
+
+    Returns:
+        type[Processor]: The concrete processor class
+    """
+
+    class ConcreteProcessor(processor_class[from_type]):  # type: ignore
+        pass
+
+    return ConcreteProcessor
+
+
+def processor_factory_to_type(processor_class: type, to_type: type[Any]) -> type[Processor]:
+    """Create a concrete processor class from a processor class and output type.
+
+    Args:
+        processor_class (type): The processor class to create a concrete instance of
+        to_type (type[Any]): The type of the output data
+
+    Returns:
+        type[Processor]: The concrete processor class
+    """
+
+    class ConcreteProcessor(processor_class[to_type]):  # type: ignore
+        pass
+
+    return ConcreteProcessor

nat/profiler/decorators/function_tracking.py

@@ -16,7 +16,9 @@
 import functools
 import inspect
 import uuid
+from collections.abc import Callable
 from typing import Any
+from typing import cast
 
 from pydantic import BaseModel
 
@@ -252,3 +254,126 @@ def track_function(func: Any = None, *, metadata: dict[str, Any] | None = None):
         return result
 
     return sync_wrapper
+
+
+def track_unregistered_function(func: Callable[..., Any] | None = None,
+                                *,
+                                name: str | None = None,
+                                metadata: dict[str, Any] | None = None) -> Callable[..., Any]:
+    """
+    Decorator that wraps any function with scope management and automatic tracking.
+
+    - Sets active function context using the function name
+    - Leverages Context.push_active_function for built-in tracking
+    - Avoids duplicate tracking entries by relying on the library's built-in systems
+    - Supports sync/async functions and generators
+
+    Args:
+        func: The function to wrap (auto-detected when used without parentheses)
+        name: Custom name to use for tracking instead of func.__name__
+        metadata: Additional metadata to include in tracking
+    """
+
+    # If called with parameters: @track_unregistered_function(name="...", metadata={...})
+    if func is None:
+
+        def decorator_wrapper(actual_func: Callable[..., Any]) -> Callable[..., Any]:
+            # Cast to ensure type checker understands this returns a callable
+            return cast(Callable[..., Any], track_unregistered_function(actual_func, name=name, metadata=metadata))
+
+        return decorator_wrapper
+
+    # Direct decoration: @track_unregistered_function or recursive call with actual function
+    function_name: str = name if name else func.__name__
+
+    # --- Validate metadata ---
+    if metadata is not None:
+        if not isinstance(metadata, dict):
+            raise TypeError("metadata must be a dict[str, Any].")
+        if any(not isinstance(k, str) for k in metadata.keys()):
+            raise TypeError("All metadata keys must be strings.")
+
+    trace_metadata = TraceMetadata(provided_metadata=metadata)
+
+    # --- Now detect the function type and wrap accordingly ---
+    if inspect.isasyncgenfunction(func):
+        # ---------------------
+        # ASYNC GENERATOR
+        # ---------------------
+
+        @functools.wraps(func)
+        async def async_gen_wrapper(*args, **kwargs):
+            context = Context.get()
+            input_data = (
+                *args,
+                kwargs,
+            )
+            # Only do context management - let push_active_function handle tracking
+            with context.push_active_function(function_name, input_data=input_data, metadata=trace_metadata) as manager:
+                final_outputs = []
+                async for item in func(*args, **kwargs):
+                    final_outputs.append(item)
+                    yield item
+
+                manager.set_output(final_outputs)
+
+        return async_gen_wrapper
+
+    if inspect.iscoroutinefunction(func):
+        # ---------------------
+        # ASYNC FUNCTION
+        # ---------------------
+        @functools.wraps(func)
+        async def async_wrapper(*args, **kwargs):
+            context = Context.get()
+            input_data = (
+                *args,
+                kwargs,
+            )
+
+            # Only do context management - let push_active_function handle tracking
+            with context.push_active_function(function_name, input_data=input_data, metadata=trace_metadata) as manager:
+                result = await func(*args, **kwargs)
+                manager.set_output(result)
+                return result
+
+        return async_wrapper
+
+    if inspect.isgeneratorfunction(func):
+        # ---------------------
+        # SYNC GENERATOR
+        # ---------------------
+        @functools.wraps(func)
+        def sync_gen_wrapper(*args, **kwargs):
+            context = Context.get()
+            input_data = (
+                *args,
+                kwargs,
+            )
+
+            # Only do context management - let push_active_function handle tracking
+            with context.push_active_function(function_name, input_data=input_data, metadata=trace_metadata) as manager:
+                final_outputs = []
+                for item in func(*args, **kwargs):
+                    final_outputs.append(item)
+                    yield item
+
+                manager.set_output(final_outputs)
+
+        return sync_gen_wrapper
+
+    @functools.wraps(func)
+    def sync_wrapper(*args, **kwargs):
+        context = Context.get()
+        input_data = (
+            *args,
+            kwargs,
+        )
+
+        # Only do context management - let push_active_function handle tracking
+        with context.push_active_function(function_name, input_data=input_data, metadata=trace_metadata) as manager:
+            result = func(*args, **kwargs)
+            manager.set_output(result)
+            return result
+
+    return sync_wrapper

{nvidia_nat-1.3.0a20250827.dist-info → nvidia_nat-1.3.0a20250828.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: nvidia-nat
-Version: 1.3.0a20250827
+Version: 1.3.0a20250828
 Summary: NVIDIA NeMo Agent toolkit
 Author: NVIDIA Corporation
 Maintainer: NVIDIA Corporation
@@ -243,6 +243,8 @@ Provides-Extra: agno
 Requires-Dist: nvidia-nat-agno; extra == "agno"
 Provides-Extra: crewai
 Requires-Dist: nvidia-nat-crewai; extra == "crewai"
+Provides-Extra: data-flywheel
+Requires-Dist: nvidia-nat-data-flywheel; extra == "data-flywheel"
 Provides-Extra: ingestion
 Requires-Dist: nvidia-nat-ingestion; extra == "ingestion"
 Provides-Extra: langchain