nvidia-nat-test 1.4.0a20260117__py3-none-any.whl

nat/meta/pypi.md

@@ -0,0 +1,23 @@
+<!--
+SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+SPDX-License-Identifier: Apache-2.0
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+-->
+
+![NVIDIA NeMo Agent Toolkit](https://media.githubusercontent.com/media/NVIDIA/NeMo-Agent-Toolkit/refs/heads/main/docs/source/_static/banner.png "NeMo Agent toolkit banner image")
+
+# NVIDIA NeMo Agent Toolkit Subpackage
+This is a subpackage for NeMo Agent toolkit test utilities.
+
+For more information about the NVIDIA NeMo Agent toolkit, please visit the [NeMo Agent toolkit GitHub Repo](https://github.com/NVIDIA/NeMo-Agent-Toolkit).

nat/test/__init__.py

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

nat/test/embedder.py

@@ -0,0 +1,44 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from pydantic import ConfigDict
+
+from nat.builder.builder import Builder
+from nat.builder.embedder import EmbedderProviderInfo
+from nat.builder.framework_enum import LLMFrameworkEnum
+from nat.cli.register_workflow import register_embedder_client
+from nat.cli.register_workflow import register_embedder_provider
+from nat.data_models.embedder import EmbedderBaseConfig
+
+
+class EmbedderTestConfig(EmbedderBaseConfig, name="test_embedder"):
+    model_config = ConfigDict(protected_namespaces=())
+
+    model_name: str = "nvidia/nv-embedqa-e5-v5"
+    embedding_size: int = 768
+
+
+@register_embedder_provider(config_type=EmbedderTestConfig)
+async def embedder_test_provider(config: EmbedderTestConfig, builder: Builder):
+
+    yield EmbedderProviderInfo(config=config, description="Test embedder provider")
+
+
+@register_embedder_client(config_type=EmbedderTestConfig, wrapper_type=LLMFrameworkEnum.LANGCHAIN)
+async def embedder_langchain_test_client(config: EmbedderTestConfig, builder: Builder):
+
+    from langchain_community.embeddings import DeterministicFakeEmbedding
+
+    yield DeterministicFakeEmbedding(size=config.embedding_size)

nat/test/functions.py

@@ -0,0 +1,99 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from collections.abc import AsyncGenerator
+
+from nat.builder.builder import Builder
+from nat.builder.function_info import FunctionInfo
+from nat.cli.register_workflow import register_function
+from nat.data_models.api_server import ChatRequest
+from nat.data_models.api_server import ChatResponse
+from nat.data_models.api_server import ChatResponseChunk
+from nat.data_models.api_server import Usage
+from nat.data_models.function import FunctionBaseConfig
+
+
+class EchoFunctionConfig(FunctionBaseConfig, name="test_echo"):
+    use_openai_api: bool = False
+
+
+@register_function(config_type=EchoFunctionConfig)
+async def echo_function(config: EchoFunctionConfig, builder: Builder):
+
+    async def inner(message: str) -> str:
+        return message
+
+    async def inner_oai(message: ChatRequest) -> ChatResponse:
+        content = message.messages[0].content
+
+        # Create usage statistics for the response
+        prompt_tokens = sum(len(str(msg.content).split()) for msg in message.messages)
+        completion_tokens = len(content.split()) if content else 0
+        total_tokens = prompt_tokens + completion_tokens
+        usage = Usage(prompt_tokens=prompt_tokens, completion_tokens=completion_tokens, total_tokens=total_tokens)
+        return ChatResponse.from_string(content, usage=usage)
+
+    if (config.use_openai_api):
+        yield inner_oai
+    else:
+        yield inner
+
+
+class StreamingEchoFunctionConfig(FunctionBaseConfig, name="test_streaming_echo"):
+    use_openai_api: bool = False
+
+
+@register_function(config_type=StreamingEchoFunctionConfig)
+async def streaming_function(config: StreamingEchoFunctionConfig, builder: Builder):
+
+    def oai_to_list(message: ChatRequest) -> list[str]:
+        return [m.content for m in message.messages]
+
+    async def inner(message: list[str]) -> AsyncGenerator[str]:
+        for value in message:
+            yield value
+
+    async def inner_oai(message: ChatRequest) -> AsyncGenerator[ChatResponseChunk]:
+        for value in oai_to_list(message):
+            yield ChatResponseChunk.from_string(value)
+
+    yield FunctionInfo.from_fn(inner_oai if config.use_openai_api else inner, converters=[oai_to_list])
+
+
+class ConstantFunctionConfig(FunctionBaseConfig, name="test_constant"):
+    response: str
+
+
+@register_function(config_type=ConstantFunctionConfig)
+async def constant_function(config: ConstantFunctionConfig, builder: Builder):
+
+    async def inner() -> str:
+        return config.response
+
+    yield inner
+
+
+class StreamingConstantFunctionConfig(FunctionBaseConfig, name="test_streaming_constant"):
+    responses: list[str]
+
+
+@register_function(config_type=StreamingConstantFunctionConfig)
+async def streaming_constant_function(config: StreamingConstantFunctionConfig, builder: Builder):
+
+    async def inner() -> AsyncGenerator[str]:
+        for value in config.responses:
+            yield value
+
+    yield inner

nat/test/llm.py

@@ -0,0 +1,244 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+# pylint: disable=unused-argument,missing-class-docstring,missing-function-docstring,import-outside-toplevel
+# pylint: disable=too-few-public-methods
+
+import asyncio
+import time
+from collections.abc import AsyncGenerator
+from collections.abc import Iterator
+from itertools import cycle as iter_cycle
+from typing import Any
+
+from pydantic import Field
+
+from nat.builder.builder import Builder
+from nat.builder.framework_enum import LLMFrameworkEnum
+from nat.builder.llm import LLMProviderInfo
+from nat.cli.register_workflow import register_llm_client
+from nat.cli.register_workflow import register_llm_provider
+from nat.data_models.llm import LLMBaseConfig
+
+
+class MockLLMConfig(LLMBaseConfig, name="nat_test_llm"):
+    """Mock LLM configuration."""
+
+    response_seq: list[str] = Field(
+        default=[],
+        description="Returns the next element in order (wraps)",
+    )
+    delay_ms: int = Field(default=0, ge=0, description="Artificial per-call delay in milliseconds to mimic latency")
+
+
+class _ResponseChooser:
+    """
+    Helper class to choose the next response according to config using itertools.cycle and provide synchronous and
+    asynchronous sleep functions.
+    """
+
+    def __init__(self, response_seq: list[str], delay_ms: int):
+        self._cycler = iter_cycle(response_seq) if response_seq else None
+        self._delay_ms = delay_ms
+
+    def next_response(self) -> str:
+        """Return the next response in the cycle, or an empty string if no responses are configured."""
+        if self._cycler is None:
+            return ""
+        return next(self._cycler)
+
+    def sync_sleep(self) -> None:
+        time.sleep(self._delay_ms / 1000.0)
+
+    async def async_sleep(self) -> None:
+        await asyncio.sleep(self._delay_ms / 1000.0)
+
+
+@register_llm_provider(config_type=MockLLMConfig)
+async def test_llm_provider(config: MockLLMConfig, builder: Builder) -> AsyncGenerator[LLMProviderInfo, None]:
+    """Register the `nat_test_llm` provider for the NAT registry."""
+    yield LLMProviderInfo(config=config, description="Test LLM provider")
+
+
+@register_llm_client(config_type=MockLLMConfig, wrapper_type=LLMFrameworkEnum.LANGCHAIN)
+async def test_llm_langchain(config: MockLLMConfig, builder: Builder):
+    """LLM client for LangChain/LangGraph."""
+
+    chooser = _ResponseChooser(response_seq=config.response_seq, delay_ms=config.delay_ms)
+
+    class LangChainTestLLM:
+
+        def invoke(self, messages: Any, **_kwargs: Any) -> str:
+            chooser.sync_sleep()
+            return chooser.next_response()
+
+        async def ainvoke(self, messages: Any, **_kwargs: Any) -> str:
+            await chooser.async_sleep()
+            return chooser.next_response()
+
+        def stream(self, messages: Any, **_kwargs: Any) -> Iterator[str]:
+            chooser.sync_sleep()
+            yield chooser.next_response()
+
+        async def astream(self, messages: Any, **_kwargs: Any) -> AsyncGenerator[str]:
+            await chooser.async_sleep()
+            yield chooser.next_response()
+
+        def bind_tools(self, tools: Any, **_kwargs: Any) -> "LangChainTestLLM":
+            """Bind tools to the LLM. Returns self to maintain fluent interface."""
+            return self
+
+        def bind(self, **_kwargs: Any) -> "LangChainTestLLM":
+            """Bind additional parameters to the LLM. Returns self to maintain fluent interface."""
+            return self
+
+    yield LangChainTestLLM()
+
+
+@register_llm_client(config_type=MockLLMConfig, wrapper_type=LLMFrameworkEnum.LLAMA_INDEX)
+async def test_llm_llama_index(config: MockLLMConfig, builder: Builder):
+
+    try:
+        from llama_index.core.base.llms.types import ChatMessage
+        from llama_index.core.base.llms.types import ChatResponse
+    except ImportError as exc:
+        raise ImportError("llama_index is required for using the test_llm with llama_index. "
+                          "Please install the `nvidia-nat-llama-index` package. ") from exc
+
+    chooser = _ResponseChooser(response_seq=config.response_seq, delay_ms=config.delay_ms)
+
+    class LITestLLM:
+
+        def chat(self, messages: list[Any] | None = None, **_kwargs: Any) -> ChatResponse:
+            chooser.sync_sleep()
+            return ChatResponse(message=ChatMessage(chooser.next_response()))
+
+        async def achat(self, messages: list[Any] | None = None, **_kwargs: Any) -> ChatResponse:
+            await chooser.async_sleep()
+            return ChatResponse(message=ChatMessage(chooser.next_response()))
+
+        def stream_chat(self, messages: list[Any] | None = None, **_kwargs: Any) -> Iterator[ChatResponse]:
+            chooser.sync_sleep()
+            yield ChatResponse(message=ChatMessage(chooser.next_response()))
+
+        async def astream_chat(self,
+                               messages: list[Any] | None = None,
+                               **_kwargs: Any) -> AsyncGenerator[ChatResponse, None]:
+            await chooser.async_sleep()
+            yield ChatResponse(message=ChatMessage(chooser.next_response()))
+
+    yield LITestLLM()
+
+
+@register_llm_client(config_type=MockLLMConfig, wrapper_type=LLMFrameworkEnum.CREWAI)
+async def test_llm_crewai(config: MockLLMConfig, builder: Builder):
+    """LLM client for CrewAI."""
+
+    chooser = _ResponseChooser(response_seq=config.response_seq, delay_ms=config.delay_ms)
+
+    class CrewAITestLLM:
+
+        def call(self, messages: list[dict[str, str]] | None = None, **kwargs: Any) -> str:
+            chooser.sync_sleep()
+            return chooser.next_response()
+
+    yield CrewAITestLLM()
+
+
+@register_llm_client(config_type=MockLLMConfig, wrapper_type=LLMFrameworkEnum.SEMANTIC_KERNEL)
+async def test_llm_semantic_kernel(config: MockLLMConfig, builder: Builder):
+    """LLM client for SemanticKernel."""
+
+    try:
+        from semantic_kernel.contents.chat_message_content import ChatMessageContent
+        from semantic_kernel.contents.utils.author_role import AuthorRole
+    except ImportError as exc:
+        raise ImportError("Semantic Kernel is required for using the test_llm with semantic_kernel. "
+                          "Please install the `nvidia-nat-semantic-kernel` package. ") from exc
+
+    chooser = _ResponseChooser(response_seq=config.response_seq, delay_ms=config.delay_ms)
+
+    class SKTestLLM:
+
+        async def get_chat_message_contents(self, chat_history: Any, **_kwargs: Any) -> list[ChatMessageContent]:
+            await chooser.async_sleep()
+            text = chooser.next_response()
+            return [ChatMessageContent(role=AuthorRole.ASSISTANT, content=text)]
+
+        async def get_streaming_chat_message_contents(self, chat_history: Any,
+                                                      **_kwargs: Any) -> AsyncGenerator[ChatMessageContent, None]:
+            await chooser.async_sleep()
+            text = chooser.next_response()
+            yield ChatMessageContent(role=AuthorRole.ASSISTANT, content=text)
+
+    yield SKTestLLM()
+
+
+@register_llm_client(config_type=MockLLMConfig, wrapper_type=LLMFrameworkEnum.AGNO)
+async def test_llm_agno(config: MockLLMConfig, builder: Builder):
+    """LLM client for agno."""
+
+    chooser = _ResponseChooser(response_seq=config.response_seq, delay_ms=config.delay_ms)
+
+    class AgnoTestLLM:
+
+        def invoke(self, messages: Any | None = None, **_kwargs: Any) -> str:
+            chooser.sync_sleep()
+            return chooser.next_response()
+
+        async def ainvoke(self, messages: Any | None = None, **_kwargs: Any) -> str:
+            await chooser.async_sleep()
+            return chooser.next_response()
+
+        def invoke_stream(self, messages: Any | None = None, **_kwargs: Any) -> Iterator[str]:
+            chooser.sync_sleep()
+            yield chooser.next_response()
+
+        async def ainvoke_stream(self, messages: Any | None = None, **_kwargs: Any) -> AsyncGenerator[str, None]:
+            await chooser.async_sleep()
+            yield chooser.next_response()
+
+    yield AgnoTestLLM()
+
+
+@register_llm_client(config_type=MockLLMConfig, wrapper_type=LLMFrameworkEnum.ADK)
+async def test_llm_adk(config: MockLLMConfig, builder: Builder):
+    """LLM client for Google ADK."""
+
+    try:
+        from google.adk.models.base_llm import BaseLlm
+        from google.adk.models.llm_request import LlmRequest
+        from google.adk.models.llm_response import LlmResponse
+        from google.genai import types
+    except ImportError as exc:
+        raise ImportError("Google ADK is required for using the test_llm with ADK. "
+                          "Please install the `nvidia-nat-adk` package. ") from exc
+
+    chooser = _ResponseChooser(response_seq=config.response_seq, delay_ms=config.delay_ms)
+
+    class ADKTestLLM(BaseLlm):
+
+        async def generate_content_async(self,
+                                         llm_request: LlmRequest,
+                                         stream: bool = False) -> AsyncGenerator[LlmResponse, None]:
+            self._maybe_append_user_content(llm_request)
+            await chooser.async_sleep()
+            text = chooser.next_response()
+            yield LlmResponse(content=types.Content(role="model", parts=[types.Part.from_text(text=text)]))
+
+        def connect(self, *_args: Any, **_kwargs: Any) -> None:
+            return None
+
+    yield ADKTestLLM(model="nat_test_llm")

nat/test/memory.py

@@ -0,0 +1,41 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from nat.builder.builder import Builder
+from nat.cli.register_workflow import register_memory
+from nat.data_models.memory import MemoryBaseConfig
+from nat.memory.interfaces import MemoryEditor
+from nat.memory.models import MemoryItem
+
+
+class DummyMemoryConfig(MemoryBaseConfig, name="test_dummy"):
+    pass
+
+
+@register_memory(config_type=DummyMemoryConfig)
+async def echo_function(config: DummyMemoryConfig, builder: Builder):
+
+    class DummyMemoryEditor(MemoryEditor):
+
+        async def add_items(self, items: list[MemoryItem]) -> None:
+            pass
+
+        async def search(self, query: str, top_k: int = 5, **kwargs) -> list[MemoryItem]:
+            return []
+
+        async def remove_items(self, **kwargs) -> None:
+            pass
+
+    yield DummyMemoryEditor()

nat/test/object_store_tests.py

@@ -0,0 +1,117 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import uuid
+from abc import abstractmethod
+from contextlib import asynccontextmanager
+
+import pytest
+import pytest_asyncio
+
+from nat.data_models.object_store import KeyAlreadyExistsError
+from nat.data_models.object_store import NoSuchKeyError
+from nat.object_store.interfaces import ObjectStore
+from nat.object_store.models import ObjectStoreItem
+
+
+@pytest.mark.asyncio(loop_scope="class")
+class ObjectStoreTests:
+
+    @abstractmethod
+    @asynccontextmanager
+    async def _get_store(self):
+        yield
+
+    @pytest_asyncio.fixture(loop_scope="class", scope="class")
+    async def store(self):
+
+        async with self._get_store() as store:
+            yield store
+
+    async def test_create_object_store(self, store: ObjectStore):
+        assert isinstance(store, ObjectStore)
+
+    async def test_put_object(self, store: ObjectStore):
+
+        # Use a random key to avoid conflicts with other tests
+        key = f"test_key_{uuid.uuid4()}"
+
+        initial_item = ObjectStoreItem(data=b"test_value")
+        await store.put_object(key, initial_item)
+
+        # Try to put the same object again
+        with pytest.raises(KeyAlreadyExistsError):
+            await store.put_object(key, initial_item)
+
+    async def test_upsert_object(self, store: ObjectStore):
+        key = f"test_key_{uuid.uuid4()}"
+
+        initial_item = ObjectStoreItem(data=b"test_value", content_type="text/plain", metadata={"key": "value"})
+
+        await store.upsert_object(key, initial_item)
+
+        # Check that the object exists
+        retrieved_item = await store.get_object(key)
+        assert retrieved_item.data == initial_item.data
+        assert retrieved_item.content_type == initial_item.content_type
+        assert retrieved_item.metadata == initial_item.metadata
+
+        # Upsert the object with a new value
+        new_item = ObjectStoreItem(data=b"new_value", content_type="application/json", metadata={"key": "new_value"})
+        await store.upsert_object(key, new_item)
+
+        # Check that the object was updated
+        retrieved_item = await store.get_object(key)
+        assert retrieved_item.data == new_item.data
+        assert retrieved_item.content_type == new_item.content_type
+        assert retrieved_item.metadata == new_item.metadata
+
+    async def test_get_object(self, store: ObjectStore):
+
+        key = f"test_key_{uuid.uuid4()}"
+
+        initial_item = ObjectStoreItem(data=b"test_value", content_type="text/plain", metadata={"key": "value"})
+        await store.put_object(key, initial_item)
+
+        retrieved_item = await store.get_object(key)
+        assert retrieved_item.data == initial_item.data
+        assert retrieved_item.content_type == initial_item.content_type
+        assert retrieved_item.metadata == initial_item.metadata
+
+        # Try to get an object that doesn't exist
+        with pytest.raises(NoSuchKeyError):
+            await store.get_object(f"test_key_{uuid.uuid4()}")
+
+    async def test_delete_object(self, store: ObjectStore):
+
+        key = f"test_key_{uuid.uuid4()}"
+
+        initial_item = ObjectStoreItem(data=b"test_value")
+        await store.put_object(key, initial_item)
+
+        # Check that the object exists
+        retrieved_item = await store.get_object(key)
+        assert retrieved_item.data == initial_item.data
+
+        # Delete the object
+        await store.delete_object(key)
+
+        # Try to get the object again
+        with pytest.raises(NoSuchKeyError):
+            await store.get_object(key)
+
+        # Try to delete the object again
+        with pytest.raises(NoSuchKeyError):
+            await store.delete_object(key)