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

nat/test/tool_test_runner.py

@@ -14,18 +14,33 @@
 # limitations under the License.
 
 import asyncio
+import inspect
 import logging
 import typing
+from collections.abc import Sequence
 from contextlib import asynccontextmanager
 from unittest.mock import AsyncMock
 from unittest.mock import MagicMock
 
+from nat.authentication.interfaces import AuthProviderBase
 from nat.builder.builder import Builder
 from nat.builder.function import Function
+from nat.builder.function import FunctionGroup
 from nat.builder.function_info import FunctionInfo
 from nat.cli.type_registry import GlobalTypeRegistry
+from nat.data_models.authentication import AuthProviderBaseConfig
+from nat.data_models.embedder import EmbedderBaseConfig
 from nat.data_models.function import FunctionBaseConfig
+from nat.data_models.function import FunctionGroupBaseConfig
+from nat.data_models.function_dependencies import FunctionDependencies
+from nat.data_models.llm import LLMBaseConfig
+from nat.data_models.memory import MemoryBaseConfig
 from nat.data_models.object_store import ObjectStoreBaseConfig
+from nat.data_models.retriever import RetrieverBaseConfig
+from nat.data_models.ttc_strategy import TTCStrategyBaseConfig
+from nat.experimental.test_time_compute.models.stage_enums import PipelineTypeEnum
+from nat.experimental.test_time_compute.models.stage_enums import StageTypeEnum
+from nat.memory.interfaces import MemoryEditor
 from nat.object_store.interfaces import ObjectStore
 from nat.runtime.loader import PluginTypes
 from nat.runtime.loader import discover_and_register_plugins
@@ -46,6 +61,10 @@ class MockBuilder(Builder):
         """Add a mock function that returns a fixed response."""
         self._mocks[name] = mock_response
 
+    def mock_function_group(self, name: str, mock_response: typing.Any):
+        """Add a mock function group that returns a fixed response."""
+        self._mocks[name] = mock_response
+
     def mock_llm(self, name: str, mock_response: typing.Any):
         """Add a mock LLM that returns a fixed response."""
         self._mocks[f"llm_{name}"] = mock_response
@@ -70,14 +89,16 @@ class MockBuilder(Builder):
         """Add a mock TTC strategy that returns a fixed response."""
         self._mocks[f"ttc_strategy_{name}"] = mock_response
 
-    async def add_ttc_strategy(self, name: str, config):
+    def mock_auth_provider(self, name: str, mock_response: typing.Any):
+        """Add a mock auth provider that returns a fixed response."""
+        self._mocks[f"auth_provider_{name}"] = mock_response
+
+    async def add_ttc_strategy(self, name: str, config: TTCStrategyBaseConfig) -> None:
         """Mock implementation (no‑op)."""
         pass
 
-    async def get_ttc_strategy(self,
-                               strategy_name: str,
-                               pipeline_type: typing.Any = None,
-                               stage_type: typing.Any = None):
+    async def get_ttc_strategy(self, strategy_name: str, pipeline_type: PipelineTypeEnum,
+                               stage_type: StageTypeEnum) -> typing.Any:
         """Return a mock TTC strategy if one is configured."""
         key = f"ttc_strategy_{strategy_name}"
         if key in self._mocks:
@@ -90,16 +111,29 @@ class MockBuilder(Builder):
 
     async def get_ttc_strategy_config(self,
                                       strategy_name: str,
-                                      pipeline_type: typing.Any = None,
-                                      stage_type: typing.Any = None):
+                                      pipeline_type: PipelineTypeEnum,
+                                      stage_type: StageTypeEnum) -> TTCStrategyBaseConfig:
         """Mock implementation."""
+        return TTCStrategyBaseConfig()
+
+    async def add_auth_provider(self, name: str, config: AuthProviderBaseConfig) -> None:
+        """Mock implementation (no‑op)."""
         pass
 
+    async def get_auth_provider(self, auth_provider_name: str) -> AuthProviderBase:
+        """Return a mock auth provider if one is configured."""
+        key = f"auth_provider_{auth_provider_name}"
+        if key in self._mocks:
+            mock_auth = MagicMock()
+            mock_auth.authenticate = AsyncMock(return_value=self._mocks[key])
+            return mock_auth
+        raise ValueError(f"Auth provider '{auth_provider_name}' not mocked. Use mock_auth_provider() to add it.")
+
     async def add_function(self, name: str, config: FunctionBaseConfig) -> Function:
         """Mock implementation - not used in tool testing."""
         raise NotImplementedError("Mock implementation does not support add_function")
 
-    def get_function(self, name: str) -> Function:
+    async def get_function(self, name: str) -> Function:
         """Return a mock function if one is configured."""
         if name in self._mocks:
             mock_fn = AsyncMock()
@@ -109,25 +143,49 @@ class MockBuilder(Builder):
 
     def get_function_config(self, name: str) -> FunctionBaseConfig:
         """Mock implementation."""
-        pass
+        return FunctionBaseConfig()
+
+    async def add_function_group(self, name: str, config: FunctionGroupBaseConfig) -> FunctionGroup:
+        """Mock implementation - not used in tool testing."""
+        raise NotImplementedError("Mock implementation does not support add_function_group")
+
+    async def get_function_group(self, name: str) -> FunctionGroup:
+        """Return a mock function group if one is configured."""
+        if name in self._mocks:
+            mock_fn_group = MagicMock(spec=FunctionGroup)
+            mock_fn_group.ainvoke = AsyncMock(return_value=self._mocks[name])
+            return mock_fn_group
+        raise ValueError(f"Function group '{name}' not mocked. Use mock_function_group() to add it.")
+
+    def get_function_group_config(self, name: str) -> FunctionGroupBaseConfig:
+        """Mock implementation."""
+        return FunctionGroupBaseConfig()
 
     async def set_workflow(self, config: FunctionBaseConfig) -> Function:
         """Mock implementation."""
-        pass
+        mock_fn = AsyncMock()
+        mock_fn.ainvoke = AsyncMock(return_value="mock_workflow_result")
+        return mock_fn
 
     def get_workflow(self) -> Function:
         """Mock implementation."""
-        pass
+        mock_fn = AsyncMock()
+        mock_fn.ainvoke = AsyncMock(return_value="mock_workflow_result")
+        return mock_fn
 
     def get_workflow_config(self) -> FunctionBaseConfig:
         """Mock implementation."""
-        pass
+        return FunctionBaseConfig()
+
+    async def get_tools(self, tool_names: Sequence[str], wrapper_type) -> list[typing.Any]:
+        """Mock implementation."""
+        return []
 
-    def get_tool(self, fn_name: str, wrapper_type):
+    async def get_tool(self, fn_name: str, wrapper_type) -> typing.Any:
         """Mock implementation."""
         pass
 
-    async def add_llm(self, name: str, config):
+    async def add_llm(self, name: str, config) -> None:
         """Mock implementation."""
         pass
 
@@ -141,11 +199,11 @@ class MockBuilder(Builder):
             return mock_llm
         raise ValueError(f"LLM '{llm_name}' not mocked. Use mock_llm() to add it.")
 
-    def get_llm_config(self, llm_name: str):
+    def get_llm_config(self, llm_name: str) -> LLMBaseConfig:
         """Mock implementation."""
-        pass
+        return LLMBaseConfig()
 
-    async def add_embedder(self, name: str, config):
+    async def add_embedder(self, name: str, config) -> None:
         """Mock implementation."""
         pass
 
@@ -159,15 +217,14 @@ class MockBuilder(Builder):
             return mock_embedder
         raise ValueError(f"Embedder '{embedder_name}' not mocked. Use mock_embedder() to add it.")
 
-    def get_embedder_config(self, embedder_name: str):
+    def get_embedder_config(self, embedder_name: str) -> EmbedderBaseConfig:
         """Mock implementation."""
-        pass
+        return EmbedderBaseConfig()
 
-    async def add_memory_client(self, name: str, config):
-        """Mock implementation."""
-        pass
+    async def add_memory_client(self, name: str, config) -> MemoryEditor:
+        return MagicMock(spec=MemoryEditor)
 
-    def get_memory_client(self, memory_name: str):
+    async def get_memory_client(self, memory_name: str) -> MemoryEditor:
         """Return a mock memory client if one is configured."""
         key = f"memory_{memory_name}"
         if key in self._mocks:
@@ -177,11 +234,11 @@ class MockBuilder(Builder):
             return mock_memory
         raise ValueError(f"Memory client '{memory_name}' not mocked. Use mock_memory_client() to add it.")
 
-    def get_memory_client_config(self, memory_name: str):
+    def get_memory_client_config(self, memory_name: str) -> MemoryBaseConfig:
         """Mock implementation."""
-        pass
+        return MemoryBaseConfig()
 
-    async def add_retriever(self, name: str, config):
+    async def add_retriever(self, name: str, config) -> None:
         """Mock implementation."""
         pass
 
@@ -194,13 +251,13 @@ class MockBuilder(Builder):
             return mock_retriever
         raise ValueError(f"Retriever '{retriever_name}' not mocked. Use mock_retriever() to add it.")
 
-    async def get_retriever_config(self, retriever_name: str):
+    async def get_retriever_config(self, retriever_name: str) -> RetrieverBaseConfig:
         """Mock implementation."""
-        pass
+        return RetrieverBaseConfig()
 
-    async def add_object_store(self, name: str, config: ObjectStoreBaseConfig):
+    async def add_object_store(self, name: str, config: ObjectStoreBaseConfig) -> ObjectStore:
         """Mock implementation for object store."""
-        pass
+        return MagicMock(spec=ObjectStore)
 
     async def get_object_store_client(self, object_store_name: str) -> ObjectStore:
         """Return a mock object store client if one is configured."""
@@ -216,7 +273,7 @@ class MockBuilder(Builder):
 
     def get_object_store_config(self, object_store_name: str) -> ObjectStoreBaseConfig:
         """Mock implementation for object store config."""
-        pass
+        return ObjectStoreBaseConfig()
 
     def get_user_manager(self):
         """Mock implementation."""
@@ -224,9 +281,13 @@ class MockBuilder(Builder):
         mock_user.get_id = MagicMock(return_value="test_user")
         return mock_user
 
-    def get_function_dependencies(self, fn_name: str):
+    def get_function_dependencies(self, fn_name: str) -> FunctionDependencies:
         """Mock implementation."""
-        pass
+        return FunctionDependencies()
+
+    def get_function_group_dependencies(self, fn_name: str) -> FunctionDependencies:
+        """Mock implementation."""
+        return FunctionDependencies()
 
 
 class ToolTestRunner:
@@ -322,15 +383,19 @@ class ToolTestRunner:
 
             # Execute the tool
             if input_data is not None:
-                if asyncio.iscoroutinefunction(tool_function):
+                if isinstance(tool_function, Function):
+                    result = await tool_function.ainvoke(input_data)
+                elif asyncio.iscoroutinefunction(tool_function):
                     result = await tool_function(input_data)
                 else:
                     result = tool_function(input_data)
+            elif isinstance(tool_function, Function):
+                # Function objects require input, so pass None if no input_data
+                result = await tool_function.ainvoke(None)
+            elif asyncio.iscoroutinefunction(tool_function):
+                result = await tool_function()
             else:
-                if asyncio.iscoroutinefunction(tool_function):
-                    result = await tool_function()
-                else:
-                    result = tool_function()
+                result = tool_function()
 
             # Assert expected output if provided
             if expected_output is not None:
@@ -402,8 +467,8 @@ class ToolTestRunner:
             elif isinstance(tool_result, FunctionInfo):
                 if tool_result.single_fn:
                     tool_function = tool_result.single_fn
-                elif tool_result.streaming_fn:
-                    tool_function = tool_result.streaming_fn
+                elif tool_result.stream_fn:
+                    tool_function = tool_result.stream_fn
                 else:
                     raise ValueError("Tool function not found in FunctionInfo")
             elif callable(tool_result):
@@ -413,15 +478,17 @@ class ToolTestRunner:
 
             # Execute the tool
             if input_data is not None:
-                if asyncio.iscoroutinefunction(tool_function):
-                    result = await tool_function(input_data)
+                if isinstance(tool_function, Function):
+                    result = await tool_function.ainvoke(input_data)
                 else:
-                    result = tool_function(input_data)
+                    maybe_result = tool_function(input_data)
+                    result = await maybe_result if inspect.isawaitable(maybe_result) else maybe_result
+            elif isinstance(tool_function, Function):
+                # Function objects require input, so pass None if no input_data
+                result = await tool_function.ainvoke(None)
             else:
-                if asyncio.iscoroutinefunction(tool_function):
-                    result = await tool_function()
-                else:
-                    result = tool_function()
+                maybe_result = tool_function()
+                result = await maybe_result if inspect.isawaitable(maybe_result) else maybe_result
 
             # Assert expected output if provided
             if expected_output is not None:

nat/test/utils.py

@@ -0,0 +1,155 @@
+# 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 importlib.resources
+import inspect
+import json
+import subprocess
+import typing
+from contextlib import asynccontextmanager
+from pathlib import Path
+
+if typing.TYPE_CHECKING:
+    from collections.abc import AsyncIterator
+
+    from httpx import AsyncClient
+
+    from nat.data_models.config import Config
+    from nat.front_ends.fastapi.fastapi_front_end_plugin_worker import FastApiFrontEndPluginWorker
+    from nat.utils.type_utils import StrPath
+
+
+def locate_repo_root() -> Path:
+    result = subprocess.run(["git", "rev-parse", "--show-toplevel"], check=False, capture_output=True, text=True)
+    assert result.returncode == 0, f"Failed to get git root: {result.stderr}"
+    return Path(result.stdout.strip())
+
+
+def locate_example_src_dir(example_config_class: type) -> Path:
+    """
+    Locate the example src directory for an example's config class.
+    """
+    package_name = inspect.getmodule(example_config_class).__package__
+    return importlib.resources.files(package_name)
+
+
+def locate_example_dir(example_config_class: type) -> Path:
+    """
+    Locate the example directory for an example's config class.
+    """
+    src_dir = locate_example_src_dir(example_config_class)
+    example_dir = src_dir.parent.parent
+    return example_dir
+
+
+def locate_example_config(example_config_class: type,
+                          config_file: str = "config.yml",
+                          assert_exists: bool = True) -> Path:
+    """
+    Locate the example config file for an example's config class, assumes the example contains a 'configs' directory
+    """
+    example_dir = locate_example_src_dir(example_config_class)
+    config_path = example_dir.joinpath("configs", config_file).absolute()
+    if assert_exists:
+        assert config_path.exists(), f"Config file {config_path} does not exist"
+
+    return config_path
+
+
+async def run_workflow(
+    *,
+    config: "Config | None" = None,
+    config_file: "StrPath | None" = None,
+    question: str,
+    expected_answer: str,
+    assert_expected_answer: bool = True,
+) -> str:
+    from nat.builder.workflow_builder import WorkflowBuilder
+    from nat.runtime.loader import load_config
+    from nat.runtime.session import SessionManager
+
+    if config is None:
+        assert config_file is not None, "Either config_file or config must be provided"
+        assert Path(config_file).exists(), f"Config file {config_file} does not exist"
+        config = load_config(config_file)
+
+    async with WorkflowBuilder.from_config(config=config) as workflow_builder:
+        workflow = SessionManager(await workflow_builder.build())
+        async with workflow.run(question) as runner:
+            result = await runner.result(to_type=str)
+
+    if assert_expected_answer:
+        assert expected_answer.lower() in result.lower(), f"Expected '{expected_answer}' in '{result}'"
+
+    return result
+
+
+@asynccontextmanager
+async def build_nat_client(
+        config: "Config",
+        worker_class: "type[FastApiFrontEndPluginWorker] | None" = None) -> "AsyncIterator[AsyncClient]":
+    """
+    Build a NAT client for testing purposes.
+
+    Creates a test client with an ASGI transport for the specified configuration.
+    The client is backed by a FastAPI application built from the provided worker class.
+
+    Args:
+        config: The NAT configuration to use for building the client.
+        worker_class: Optional worker class to use. Defaults to FastApiFrontEndPluginWorker.
+
+    Yields:
+        An AsyncClient instance configured for testing.
+    """
+    from asgi_lifespan import LifespanManager
+    from httpx import ASGITransport
+    from httpx import AsyncClient
+
+    from nat.front_ends.fastapi.fastapi_front_end_plugin_worker import FastApiFrontEndPluginWorker
+
+    if worker_class is None:
+        worker_class = FastApiFrontEndPluginWorker
+
+    worker = worker_class(config)
+    app = worker.build_app()
+
+    async with LifespanManager(app):
+        async with AsyncClient(transport=ASGITransport(app=app), base_url="http://test") as client:
+            yield client
+
+
+def validate_workflow_output(workflow_output_file: Path) -> None:
+    """
+    Validate the contents of the workflow output file.
+    WIP: output format should be published as a schema and this validation should be done against that schema.
+    """
+    # Ensure the workflow_output.json file was created
+    assert workflow_output_file.exists(), "The workflow_output.json file was not created"
+
+    # Read and validate the workflow_output.json file
+    try:
+        with open(workflow_output_file, encoding="utf-8") as f:
+            result_json = json.load(f)
+    except json.JSONDecodeError as err:
+        raise RuntimeError("Failed to parse workflow_output.json as valid JSON") from err
+
+    assert isinstance(result_json, list), "The workflow_output.json file is not a list"
+    assert len(result_json) > 0, "The workflow_output.json file is empty"
+    assert isinstance(result_json[0], dict), "The workflow_output.json file is not a list of dictionaries"
+
+    # Ensure required keys exist
+    required_keys = ["id", "question", "answer", "generated_answer", "intermediate_steps"]
+    for key in required_keys:
+        assert all(item.get(key) for item in result_json), f"The '{key}' key is missing in workflow_output.json"

{nvidia_nat_test-1.2.1.dist-info → nvidia_nat_test-1.3.0.dist-info}/METADATA

@@ -1,14 +1,25 @@
 Metadata-Version: 2.4
 Name: nvidia-nat-test
-Version: 1.2.1
+Version: 1.3.0
 Summary: Testing utilities for NeMo Agent toolkit
+Author: NVIDIA Corporation
+Maintainer: NVIDIA Corporation
+License: Apache-2.0
+Project-URL: documentation, https://docs.nvidia.com/nemo/agent-toolkit/latest/
+Project-URL: source, https://github.com/NVIDIA/NeMo-Agent-Toolkit
 Keywords: ai,rag,agents
 Classifier: Programming Language :: Python
-Requires-Python: <3.13,>=3.11
+Classifier: Programming Language :: Python :: 3.11
+Classifier: Programming Language :: Python :: 3.12
+Classifier: Programming Language :: Python :: 3.13
+Requires-Python: <3.14,>=3.11
 Description-Content-Type: text/markdown
-Requires-Dist: nvidia-nat==v1.2.1
+License-File: LICENSE-3rd-party.txt
+License-File: LICENSE.md
+Requires-Dist: nvidia-nat==v1.3.0
 Requires-Dist: langchain-community~=0.3
 Requires-Dist: pytest~=8.3
+Dynamic: license-file
 
 <!--
 SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.

nvidia_nat_test-1.3.0.dist-info/RECORD

@@ -0,0 +1,18 @@
+nat/meta/pypi.md,sha256=LLKJHg5oN1-M9Pqfk3Bmphkk4O2TFsyiixuK5T0Y-gw,1100
+nat/test/__init__.py,sha256=_RnTJnsUucHvla_nYKqD4O4g8Bz0tcuDRzWk1bEhcy0,875
+nat/test/embedder.py,sha256=ClDyK1kna4hCBSlz71gK1B-ZjlwcBHTDQRekoNM81Bs,1809
+nat/test/functions.py,sha256=ZxXVzfaLBGOpR5qtmMrKU7q-M9-vVGGj3Xi5mrw4vHY,3557
+nat/test/llm.py,sha256=f6bz6arAQjhjuOKFrLfu_U1LbiyFzQmpM-q8b-WKSrU,9550
+nat/test/memory.py,sha256=xki_A2yiMhEZuQk60K7t04QRqf32nQqnfzD5Iv7fkvw,1456
+nat/test/object_store_tests.py,sha256=PyJioOtoSzILPq6LuD-sOZ_89PIcgXWZweoHBQpK2zQ,4281
+nat/test/plugin.py,sha256=VbKEduSIll8TlyKQtLX5LJ0LzCyXQDwCTxeirZKnrMY,22257
+nat/test/register.py,sha256=o1BEA5fyxyFyCxXhQ6ArmtuNpgRyTEfvw6HdBgECPLI,897
+nat/test/tool_test_runner.py,sha256=SxavwXHkvCQDl_PUiiiqgvGfexKJJTeBdI5i1qk6AzI,21712
+nat/test/utils.py,sha256=Lml187P9SUP3IB_HhBaU1XNhiljcpOFFZOAxgQR1vQo,5936
+nvidia_nat_test-1.3.0.dist-info/licenses/LICENSE-3rd-party.txt,sha256=fOk5jMmCX9YoKWyYzTtfgl-SUy477audFC5hNY4oP7Q,284609
+nvidia_nat_test-1.3.0.dist-info/licenses/LICENSE.md,sha256=QwcOLU5TJoTeUhuIXzhdCEEDDvorGiC6-3YTOl4TecE,11356
+nvidia_nat_test-1.3.0.dist-info/METADATA,sha256=xW11lwoKQIoD6bkD-SKBUZPnkdUZshmM9FgdObuApr0,1907
+nvidia_nat_test-1.3.0.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+nvidia_nat_test-1.3.0.dist-info/entry_points.txt,sha256=7dOP9XB6iMDqvav3gYx9VWUwA8RrFzhbAa8nGeC8e4Y,99
+nvidia_nat_test-1.3.0.dist-info/top_level.txt,sha256=8-CJ2cP6-f0ZReXe5Hzqp-5pvzzHz-5Ds5H2bGqh1-U,4
+nvidia_nat_test-1.3.0.dist-info/RECORD,,