nvidia-nat 1.2.1rc1__py3-none-any.whl → 1.3.0__py3-none-any.whl

nat/llm/litellm_llm.py

@@ -0,0 +1,69 @@
+# 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 collections.abc import AsyncIterator
+
+from pydantic import AliasChoices
+from pydantic import ConfigDict
+from pydantic import Field
+
+from nat.builder.builder import Builder
+from nat.builder.llm import LLMProviderInfo
+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 LiteLlmModelConfig(
+        LLMBaseConfig,
+        RetryMixin,
+        TemperatureMixin,
+        TopPMixin,
+        ThinkingMixin,
+        name="litellm",
+):
+    """A LiteLlm provider to be used with an LLM client."""
+
+    model_config = ConfigDict(protected_namespaces=(), extra="allow")
+
+    api_key: str | None = Field(default=None, description="API key to interact with hosted model.")
+    base_url: str | None = Field(default=None,
+                                 description="Base url to the hosted model.",
+                                 validation_alias=AliasChoices("base_url", "api_base"),
+                                 serialization_alias="api_base")
+    model_name: str = Field(validation_alias=AliasChoices("model_name", "model"),
+                            serialization_alias="model",
+                            description="The LiteLlm hosted model name.")
+    seed: int | None = Field(default=None, description="Random seed to set for generation.")
+
+
+@register_llm_provider(config_type=LiteLlmModelConfig)
+async def litellm_model(
+    config: LiteLlmModelConfig,
+    _builder: Builder,
+) -> AsyncIterator[LLMProviderInfo]:
+    """Litellm model provider.
+
+    Args:
+        config (LiteLlmModelConfig): The LiteLlm model configuration.
+        _builder (Builder): The NAT builder instance.
+
+    Returns:
+        AsyncIterator[LLMProviderInfo]: An async iterator that yields an LLMProviderInfo object.
+    """
+    yield LLMProviderInfo(config=config, description="A LiteLlm model for use with an LLM client.")

nat/llm/nim_llm.py

@@ -22,25 +22,37 @@ from nat.builder.builder import Builder
 from nat.builder.llm import LLMProviderInfo
 from nat.cli.register_workflow import register_llm_provider
 from nat.data_models.llm import LLMBaseConfig
+from nat.data_models.optimizable import OptimizableField
+from nat.data_models.optimizable import OptimizableMixin
+from nat.data_models.optimizable import SearchSpace
 from nat.data_models.retry_mixin import RetryMixin
-
-
-class NIMModelConfig(LLMBaseConfig, RetryMixin, name="nim"):
+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,
+                     OptimizableMixin,
+                     TemperatureMixin,
+                     TopPMixin,
+                     ThinkingMixin,
+                     name="nim"):
     """An NVIDIA Inference Microservice (NIM) llm provider to be used with an LLM client."""
 
-    model_config = ConfigDict(protected_namespaces=())
+    model_config = ConfigDict(protected_namespaces=(), extra="allow")
 
     api_key: str | None = Field(default=None, description="NVIDIA API key to interact with hosted NIM.")
     base_url: str | None = Field(default=None, description="Base url to the hosted NIM.")
     model_name: str = Field(validation_alias=AliasChoices("model_name", "model"),
                             serialization_alias="model",
                             description="The model name for the hosted NIM.")
-    temperature: float = Field(default=0.0, description="Sampling temperature in [0, 1].")
-    top_p: float = Field(default=1.0, description="Top-p for distribution sampling.")
-    max_tokens: PositiveInt = Field(default=300, description="Maximum number of tokens to generate.")
+    max_tokens: PositiveInt = OptimizableField(default=300,
+                                               description="Maximum number of tokens to generate.",
+                                               space=SearchSpace(high=2176, low=128, step=512))
 
 
 @register_llm_provider(config_type=NIMModelConfig)
-async def nim_model(llm_config: NIMModelConfig, builder: Builder):
+async def nim_model(llm_config: NIMModelConfig, _builder: Builder):
 
     yield LLMProviderInfo(config=llm_config, description="A NIM model for use with an LLM client.")

nat/llm/openai_llm.py

@@ -21,10 +21,20 @@ from nat.builder.builder import Builder
 from nat.builder.llm import LLMProviderInfo
 from nat.cli.register_workflow import register_llm_provider
 from nat.data_models.llm import LLMBaseConfig
+from nat.data_models.optimizable import OptimizableMixin
 from nat.data_models.retry_mixin import RetryMixin
-
-
-class OpenAIModelConfig(LLMBaseConfig, RetryMixin, name="openai"):
+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,
+                        OptimizableMixin,
+                        TemperatureMixin,
+                        TopPMixin,
+                        ThinkingMixin,
+                        name="openai"):
     """An OpenAI LLM provider to be used with an LLM client."""
 
     model_config = ConfigDict(protected_namespaces=(), extra="allow")
@@ -34,13 +44,11 @@ class OpenAIModelConfig(LLMBaseConfig, RetryMixin, name="openai"):
     model_name: str = Field(validation_alias=AliasChoices("model_name", "model"),
                             serialization_alias="model",
                             description="The OpenAI hosted model name.")
-    temperature: float = Field(default=0.0, description="Sampling temperature in [0, 1].")
-    top_p: float = Field(default=1.0, description="Top-p for distribution sampling.")
     seed: int | None = Field(default=None, description="Random seed to set for generation.")
     max_retries: int = Field(default=10, description="The max number of retries for the request.")
 
 
 @register_llm_provider(config_type=OpenAIModelConfig)
-async def openai_llm(config: OpenAIModelConfig, builder: Builder):
+async def openai_llm(config: OpenAIModelConfig, _builder: Builder):
 
     yield LLMProviderInfo(config=config, description="An OpenAI model for use with an LLM client.")

nat/llm/register.py

@@ -13,11 +13,15 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-# pylint: disable=unused-import
 # flake8: noqa
 # isort:skip_file
+"""Register LLM providers via import side effects.
 
+This module is imported by the NeMo Agent Toolkit runtime to ensure providers are registered and discoverable.
+"""
 # Import any providers which need to be automatically registered here
 from . import aws_bedrock_llm
+from . import azure_openai_llm
+from . import litellm_llm
 from . import nim_llm
 from . import openai_llm

nat/llm/utils/env_config_value.py

@@ -72,9 +72,8 @@ class EnvConfigValue(ABC):
                     f"{message} Try passing a value to the constructor, or setting the `{self.__class__._ENV_KEY}` "
                     "environment variable.")
 
-        else:
-            if not self.__class__._ALLOW_NONE and value is None:
-                raise ValueError("value must not be none")
+        elif not self.__class__._ALLOW_NONE and value is None:
+            raise ValueError("value must not be none")
 
         assert isinstance(value, str) or value is None
 

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 Callable
+from collections.abc import Iterable
+from dataclasses import dataclass
+from typing import Any
+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/meta/pypi.md

@@ -23,19 +23,19 @@ NeMo Agent toolkit is a flexible library designed to seamlessly integrate your e
 
 ## Key Features
 
-- [**Framework Agnostic:**](https://docs.nvidia.com/nemo/agent-toolkit/1.2/extend/plugins.html) Works with any agentic framework, so you can use your current technology stack without replatforming.
-- [**Reusability:**](https://docs.nvidia.com/nemo/agent-toolkit/1.2/extend/sharing-components.html) Every agent, tool, or workflow can be combined and repurposed, allowing developers to leverage existing work in new scenarios.
-- [**Rapid Development:**](https://docs.nvidia.com/nemo/agent-toolkit/1.2/tutorials/index.html) Start with a pre-built agent, tool, or workflow, and customize it to your needs.
-- [**Profiling:**](https://docs.nvidia.com/nemo/agent-toolkit/1.2/workflows/profiler.html) Profile entire workflows down to the tool and agent level, track input/output tokens and timings, and identify bottlenecks.
-- [**Observability:**](https://docs.nvidia.com/nemo/agent-toolkit/1.2/workflows/observe/observe-workflow-with-phoenix.html) Monitor and debug your workflows with any OpenTelemetry-compatible observability tool, with examples using [Phoenix](https://docs.nvidia.com/nemo/agent-toolkit/1.2/workflows/observe/observe-workflow-with-phoenix.html) and [W&B Weave](https://docs.nvidia.com/nemo/agent-toolkit/1.2/workflows/observe/observe-workflow-with-weave.html).
-- [**Evaluation System:**](https://docs.nvidia.com/nemo/agent-toolkit/1.2/workflows/evaluate.html) Validate and maintain accuracy of agentic workflows with built-in evaluation tools.
-- [**User Interface:**](https://docs.nvidia.com/nemo/agent-toolkit/1.2/quick-start/launching-ui.html) Use the NeMo Agent toolkit UI chat interface to interact with your agents, visualize output, and debug workflows.
-- [**MCP Compatibility**](https://docs.nvidia.com/nemo/agent-toolkit/1.2/workflows/mcp/mcp-client.html) Compatible with Model Context Protocol (MCP), allowing tools served by MCP Servers to be used as NeMo Agent toolkit functions.
+- [**Framework Agnostic:**](https://docs.nvidia.com/nemo/agent-toolkit/1.3/extend/plugins.html) Works with any agentic framework, so you can use your current technology stack without replatforming.
+- [**Reusability:**](https://docs.nvidia.com/nemo/agent-toolkit/1.3/extend/sharing-components.html) Every agent, tool, or workflow can be combined and repurposed, allowing developers to leverage existing work in new scenarios.
+- [**Rapid Development:**](https://docs.nvidia.com/nemo/agent-toolkit/1.3/tutorials/index.html) Start with a pre-built agent, tool, or workflow, and customize it to your needs.
+- [**Profiling:**](https://docs.nvidia.com/nemo/agent-toolkit/1.3/workflows/profiler.html) Profile entire workflows down to the tool and agent level, track input/output tokens and timings, and identify bottlenecks.
+- [**Observability:**](https://docs.nvidia.com/nemo/agent-toolkit/1.3/workflows/observe/observe-workflow-with-phoenix.html) Monitor and debug your workflows with any OpenTelemetry-compatible observability tool, with examples using [Phoenix](https://docs.nvidia.com/nemo/agent-toolkit/1.3/workflows/observe/observe-workflow-with-phoenix.html) and [W&B Weave](https://docs.nvidia.com/nemo/agent-toolkit/1.3/workflows/observe/observe-workflow-with-weave.html).
+- [**Evaluation System:**](https://docs.nvidia.com/nemo/agent-toolkit/1.3/workflows/evaluate.html) Validate and maintain accuracy of agentic workflows with built-in evaluation tools.
+- [**User Interface:**](https://docs.nvidia.com/nemo/agent-toolkit/1.3/quick-start/launching-ui.html) Use the NeMo Agent toolkit UI chat interface to interact with your agents, visualize output, and debug workflows.
+- [**MCP Compatibility**](https://docs.nvidia.com/nemo/agent-toolkit/1.3/workflows/mcp/mcp-client.html) Compatible with Model Context Protocol (MCP), allowing tools served by MCP Servers to be used as NeMo Agent toolkit functions.
 
 With NeMo Agent toolkit, you can move quickly, experiment freely, and ensure reliability across all your agent-driven projects.
 
 ## Links
- * [Documentation](https://docs.nvidia.com/nemo/agent-toolkit/1.2/index.html): Explore the full documentation for NeMo Agent toolkit.
+ * [Documentation](https://docs.nvidia.com/nemo/agent-toolkit/1.3/index.html): Explore the full documentation for NeMo Agent toolkit.
 
 ## First time user?
  If this is your first time using NeMo Agent toolkit, it is recommended to install the latest version from the [source repository](https://github.com/NVIDIA/NeMo-Agent-Toolkit?tab=readme-ov-file#quick-start) on GitHub. This package is intended for users who are familiar with NeMo Agent toolkit applications and need to add NeMo Agent toolkit as a dependency to their project.

nat/object_store/register.py

@@ -13,7 +13,6 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-# pylint: disable=unused-import
 # flake8: noqa
 # isort:skip_file
 

nat/observability/exporter/base_exporter.py

@@ -357,7 +357,7 @@ class BaseExporter(Exporter):
                 except Exception as e:
                     logger.warning("Error while canceling task %s: %s", task.get_name(), e)
 
-    async def _wait_for_tasks(self, timeout: float = 5.0):
+    async def wait_for_tasks(self, timeout: float = 5.0):
         """Wait for all tracked tasks to complete with a timeout.
 
         Note: This method is NOT called during normal stop() operation for performance.
@@ -372,10 +372,10 @@ class BaseExporter(Exporter):
         try:
             # Wait for all tasks to complete with a timeout
             await asyncio.wait_for(asyncio.gather(*self._tasks, return_exceptions=True), timeout=timeout)
-        except asyncio.TimeoutError:
+        except TimeoutError:
             logger.warning("%s: Some tasks did not complete within %s seconds", self.name, timeout)
         except Exception as e:
-            logger.error("%s: Error while waiting for tasks: %s", self.name, e, exc_info=True)
+            logger.exception("%s: Error while waiting for tasks: %s", self.name, e)
 
     @override
     async def stop(self):

nat/observability/exporter/file_exporter.py

@@ -24,7 +24,7 @@ from nat.observability.processor.intermediate_step_serializer import Intermediat
 logger = logging.getLogger(__name__)
 
 
-class FileExporter(FileExportMixin, RawExporter[IntermediateStep, str]):  # pylint: disable=R0901
+class FileExporter(FileExportMixin, RawExporter[IntermediateStep, str]):
     """A File exporter that exports telemetry traces to a local file."""
 
     def __init__(self, context_state: ContextState | None = None, **file_kwargs):