nvidia-nat 1.1.0a20251020__py3-none-any.whl

nat/llm/openai_llm.py

@@ -0,0 +1,54 @@
+# SPDX-FileCopyrightText: Copyright (c) 2024-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 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.optimizable import OptimizableMixin
+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,
+                        OptimizableMixin,
+                        TemperatureMixin,
+                        TopPMixin,
+                        ThinkingMixin,
+                        name="openai"):
+    """An OpenAI LLM provider to be used with an LLM client."""
+
+    model_config = ConfigDict(protected_namespaces=(), extra="allow")
+
+    api_key: str | None = Field(default=None, description="OpenAI API key to interact with hosted model.")
+    base_url: str | None = Field(default=None, description="Base url to the hosted model.")
+    model_name: str = Field(validation_alias=AliasChoices("model_name", "model"),
+                            serialization_alias="model",
+                            description="The OpenAI hosted model name.")
+    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):
+
+    yield LLMProviderInfo(config=config, description="An OpenAI model for use with an LLM client.")

nat/llm/register.py

@@ -0,0 +1,27 @@
+# SPDX-FileCopyrightText: Copyright (c) 2024-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.
+
+# 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/__init__.py

@@ -0,0 +1,14 @@
+# 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.

nat/llm/utils/env_config_value.py

@@ -0,0 +1,93 @@
+# SPDX-FileCopyrightText: Copyright (c) 2024-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 os
+from abc import ABC
+from enum import Enum
+
+
+class EnvConfigValueSource(Enum):
+    ENV_DEFAULT = 1
+    CONSTRUCTOR = 2
+    ENV_OVERRIDE = 3
+
+
+class EnvConfigValue(ABC):
+    """
+    A wrapper for a string used as a configuration value which can be loaded from the system environment or injected via
+    the constructor. This class should be subclassed and the class fields `_ENV_KEY` and `_ENV_KEY_OVERRIDE` can be set
+    to enable environment-loading functionality. Convienience properties are available to check from where the value was
+    loaded.
+    """
+
+    _ENV_KEY: str | None = None
+    _ENV_KEY_OVERRIDE: str | None = None
+    _ALLOW_NONE: bool = False
+
+    def __init__(self, value: str | None = None, use_env: bool = True):
+        """
+        Parameters
+        ----------
+        value : str, optional
+            The value to be contained in the EnvConfigValue. If the value is `None`, an attempt will be made to load it
+            from the environment using `_ENV_KEY`. if the `_ENV_KEY_OVERRIDE` field is not `None`, an attempt will be
+            made to load that environment variable in place of the passed-in value.
+        use_env : bool
+            If False, all environment-loading logic will be bypassed and the passed-in value will be used as-is.
+            defaults to True.
+        """
+
+        self._source = EnvConfigValueSource.CONSTRUCTOR
+
+        if use_env:
+            if value is None and self.__class__._ENV_KEY is not None:
+                value = os.environ.get(self.__class__._ENV_KEY, None)
+                self._source = EnvConfigValueSource.ENV_DEFAULT
+
+            if self.__class__._ENV_KEY_OVERRIDE is not None and self.__class__._ENV_KEY_OVERRIDE in os.environ:
+                value = os.environ[self.__class__._ENV_KEY_OVERRIDE]
+                self._source = EnvConfigValueSource.ENV_OVERRIDE
+
+            if not self.__class__._ALLOW_NONE and value is None:
+
+                message = ("value must not be None, but provided value was None and no environment-based default or "
+                           "override was found.")
+
+                if self.__class__._ENV_KEY is None:
+                    raise ValueError(message)
+
+                raise ValueError(
+                    f"{message} Try passing a value to the constructor, or setting the `{self.__class__._ENV_KEY}` "
+                    "environment variable.")
+
+        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
+
+        self._value = value
+        self._use_env = use_env
+
+    @property
+    def source(self) -> EnvConfigValueSource:
+        return self._source
+
+    @property
+    def use_env(self) -> bool:
+        return self._use_env
+
+    @property
+    def value(self) -> str | None:
+        return self._value

nat/llm/utils/error.py

@@ -0,0 +1,17 @@
+# SPDX-FileCopyrightText: Copyright (c) 2024-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_ERROR_MESSAGE = (
+    "{package} not found. Install it and other additional dependencies by running the following command:\n")

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/memory/__init__.py

@@ -0,0 +1,20 @@
+# SPDX-FileCopyrightText: Copyright (c) 2024-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.
+"""
+NAT Memory Module
+
+This package provides foundational classes and interfaces
+for managing text-based memory in NAT's LLM-based agents.
+"""

nat/memory/interfaces.py

@@ -0,0 +1,183 @@
+# SPDX-FileCopyrightText: Copyright (c) 2024-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 abc import ABC
+from abc import abstractmethod
+from collections.abc import Callable
+
+from .models import MemoryItem
+
+
+class MemoryEditor(ABC):
+    """
+    Abstract interface for editing and
+    retrieving memory items.
+
+    A MemoryEditor is responsible for adding, searching, and
+    removing MemoryItems.
+
+    Implementations may integrate with
+    vector stores or other indexing backends.
+    """
+
+    @abstractmethod
+    async def add_items(self, items: list[MemoryItem]) -> None:
+        """
+        Insert multiple MemoryItems into the memory.
+
+        Args:
+            items (list[MemoryItem]): The items to be added.
+        """
+        raise NotImplementedError
+
+    @abstractmethod
+    async def search(self, query: str, top_k: int = 5, **kwargs) -> list[MemoryItem]:
+        """
+        Retrieve items relevant to the given query.
+        Relevance criteria depend on implementation.
+
+        Args:
+            query (str): The query string to match.
+            top_k (int): Maximum number of items to return.
+            kwargs (dict): Keyword arguments to pass to the search method.
+
+        Returns:
+            list[MemoryItem]: The most relevant MemoryItems.
+        """
+        raise NotImplementedError
+
+    @abstractmethod
+    async def remove_items(self, **kwargs) -> None:
+        """
+        Remove items. Additional parameters
+        needed for deletion can be specified in keyword arguments.
+
+        Args:
+            kwargs (dict): Keyword arguments to pass to the remove-items method.
+        """
+        raise NotImplementedError
+
+
+class MemoryIOBase(ABC):
+    """
+    Base abstract class for I/O operations
+    on memory, providing a common interface for
+
+    MemoryReader and MemoryWriter to interact
+    with a MemoryEditor.
+
+    Concrete subclasses should hold a
+    reference to a MemoryEditor instance.
+    """
+
+    def __init__(self, editor: MemoryEditor) -> None:
+        self._editor = editor
+
+
+class MemoryReader(MemoryIOBase):
+    """
+    Responsible for retrieving MemoryItems
+    from the MemoryEditor based on context or queries.
+    """
+
+    @abstractmethod
+    async def retrieve(self, context: str, top_k: int = 5) -> list[MemoryItem]:
+        """
+        Retrieve a subset of
+        MemoryItems relevant to the provided context.
+
+        Args:
+            context (str): A string representing
+            the current user context or query.
+            top_k (int): Maximum number of items to return.
+
+        Returns:
+            list[MemoryItem]: Relevant MemoryItems.
+        """
+        raise NotImplementedError
+
+
+class MemoryWriter(MemoryIOBase):
+    """
+    Responsible for converting new observations
+    (textual inputs) into MemoryItems andstoring
+    them via the MemoryEditor.
+    """
+
+    @abstractmethod
+    async def write(self, observation: str, context: str | None = None) -> list[MemoryItem]:
+        """
+        Process the given observation and store the resulting MemoryItems.
+
+        Args:
+            observation (str): The new textual input to record.
+            context (Optional[str]): Additional
+            context that might influence how the observation is stored.
+
+        Returns:
+            list[MemoryItem]: The newly created MemoryItems.
+        """
+        raise NotImplementedError
+
+
+class MemoryManager(ABC):
+    """
+    Manages the lifecycle of the stored
+    memory by applying policies such as summarization,
+    reflection, forgetting, and mergingn
+    to ensure long-term coherence and relevance.
+    """
+
+    @abstractmethod
+    async def summarize(self) -> None:
+        """
+        Summarize long or numerous MemoryItems into a more compact form.
+        This may remove the original items and store a new summary item.
+        """
+        raise NotImplementedError
+
+    @abstractmethod
+    async def reflect(self) -> None:
+        """
+        Generate higher-level insights or
+        abstractions from existing MemoryItems.
+        This may call out to an LLM or other
+        logic to produce conceptual memory.
+        """
+        raise NotImplementedError
+
+    @abstractmethod
+    async def forget(self, criteria: Callable[[MemoryItem], bool]) -> None:
+        """
+        Remove MemoryItems that are no
+        longer relevant or have low importance.
+
+        Args:
+            criteria (Callable[[MemoryItem], bool]): A function that
+            returns True for items to forget.
+        """
+        raise NotImplementedError
+
+    @abstractmethod
+    async def merge(self, criteria: Callable[[MemoryItem, MemoryItem], bool]) -> None:
+        """
+        Merge similar or redundant MemoryItems
+        into a smaller set of more concise items.
+
+        Args:
+            criteria (Callable[[MemoryItem, MemoryItem], bool]): A function
+            that determines which items can be merged.
+        """
+        raise NotImplementedError