nvidia-nat 1.1.0a20251020__py3-none-any.whl

nat/tool/memory_tools/delete_memory_tool.py

@@ -0,0 +1,66 @@
+# 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 logging
+
+from pydantic import Field
+
+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.component_ref import MemoryRef
+from nat.data_models.function import FunctionBaseConfig
+from nat.memory.models import DeleteMemoryInput
+
+logger = logging.getLogger(__name__)
+
+
+class DeleteToolConfig(FunctionBaseConfig, name="delete_memory"):
+    """Function to delete memory from a hosted memory platform."""
+
+    description: str = Field(default="Tool to delete a memory from a hosted memory platform.",
+                             description="The description of this function's use for tool calling agents.")
+    memory: MemoryRef = Field(default=MemoryRef("saas_memory"),
+                              description=("Instance name of the memory client instance from the workflow "
+                                           "configuration object."))
+
+
+@register_function(config_type=DeleteToolConfig)
+async def delete_memory_tool(config: DeleteToolConfig, builder: Builder):
+    """
+    Function to delete memory from a hosted memory platform.
+    """
+
+    from langchain_core.tools import ToolException
+
+    # First, retrieve the memory client
+    memory_editor = await builder.get_memory_client(config.memory)
+
+    async def _arun(user_id: str) -> str:
+        """
+        Asynchronous execution of deletion of memories.
+        """
+
+        try:
+
+            await memory_editor.remove_items(user_id=user_id, )
+
+            return "Memories deleted!"
+
+        except Exception as e:
+
+            raise ToolException(f"Error deleting memory: {e}") from e
+
+    yield FunctionInfo.from_fn(_arun, description=config.description, input_schema=DeleteMemoryInput)

nat/tool/memory_tools/get_memory_tool.py

@@ -0,0 +1,72 @@
+# 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 logging
+
+from pydantic import Field
+
+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.component_ref import MemoryRef
+from nat.data_models.function import FunctionBaseConfig
+from nat.memory.models import SearchMemoryInput
+
+logger = logging.getLogger(__name__)
+
+
+class GetToolConfig(FunctionBaseConfig, name="get_memory"):
+    """Function to get memory to a hosted memory platform."""
+
+    description: str = Field(default=("Tool to retrieve a memory about a user's "
+                                      "interactions to help answer questions in a personalized way."),
+                             description="The description of this function's use for tool calling agents.")
+    memory: MemoryRef = Field(default=MemoryRef("saas_memory"),
+                              description=("Instance name of the memory client instance from the workflow "
+                                           "configuration object."))
+
+
+@register_function(config_type=GetToolConfig)
+async def get_memory_tool(config: GetToolConfig, builder: Builder):
+    """
+    Function to get memory to a hosted memory platform.
+    """
+
+    import json
+
+    from langchain_core.tools import ToolException
+
+    # First, retrieve the memory client
+    memory_editor = await builder.get_memory_client(config.memory)
+
+    async def _arun(search_input: SearchMemoryInput) -> str:
+        """
+        Asynchronous execution of collection of memories.
+        """
+        try:
+            memories = await memory_editor.search(
+                query=search_input.query,
+                top_k=search_input.top_k,
+                user_id=search_input.user_id,
+            )
+
+            memory_str = f"Memories as a JSON: \n{json.dumps([mem.model_dump(mode='json') for mem in memories])}"
+            return memory_str
+
+        except Exception as e:
+
+            raise ToolException(f"Error retrieving memory: {e}") from e
+
+    yield FunctionInfo.from_fn(_arun, description=config.description)

nat/tool/nvidia_rag.py

@@ -0,0 +1,95 @@
+# 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 json
+import logging
+
+from pydantic import Field
+
+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.function import FunctionBaseConfig
+
+logger = logging.getLogger(__name__)
+
+
+class NVIDIARAGToolConfig(FunctionBaseConfig, name="nvidia_rag"):
+    """
+    Tool used to search the NVIDIA Developer database for documents across a variety of NVIDIA asset types.
+    """
+    base_url: str = Field(description="The base url to the RAG service.")
+    timeout: int = Field(default=60, description="The timeout configuration to use when sending requests.")
+    document_separator: str = Field(default="\n\n", description="The delimiter to use between retrieved documents.")
+    document_prompt: str = Field(default=("-------\n\n" + "Title: {document_title}\n"
+                                          "Text: {page_content}\nSource URL: {document_url}"),
+                                 description="The prompt to use to retrieve documents from the RAG service")
+    top_k: int = Field(default=4, description="The number of results to return from the RAG service.")
+    collection_name: str = Field(default="nvidia_api_catalog",
+                                 description=("The name of the collection to use when retrieving documents."))
+
+
+@register_function(config_type=NVIDIARAGToolConfig)
+async def nvidia_rag_tool(config: NVIDIARAGToolConfig, builder: Builder):
+    import httpx
+    from langchain.prompts import PromptTemplate
+    from langchain_core.documents import Document
+    from langchain_core.prompts import aformat_document
+
+    document_prompt = PromptTemplate.from_template(config.document_prompt)
+
+    async with httpx.AsyncClient(headers={
+            "accept": "application/json", "Content-Type": "application/json"
+    },
+                                 timeout=config.timeout) as client:
+
+        async def runnable(query: str) -> str:
+
+            try:
+                url = f"{config.base_url}/search"
+
+                payload = {"query": query, "top_k": config.top_k, "collection_name": config.collection_name}
+
+                logger.debug("Sending request to the RAG endpoint %s.", url)
+                response = await client.post(url, content=json.dumps(payload))
+
+                response.raise_for_status()
+
+                output = response.json()
+
+                docs = [
+                    Document(
+                        page_content=ret["content"],
+                        metadata={
+                            "document_title": ret["filename"],
+                            "document_url": "nemo_framework",
+                            "document_full_text": ret["content"],
+                            "score_rerank": ret["score"]
+                        },
+                        type="Document",
+                    ) for ret in output["chunks"]
+                ]
+
+                parsed_output = config.document_separator.join(
+                    [await aformat_document(doc, document_prompt) for doc in docs])
+                return parsed_output
+            except Exception as e:
+                logger.exception("Error while running the tool")
+                return f"Error while running the tool: {e}"
+
+        yield FunctionInfo.from_fn(
+            runnable,
+            description=("Search the NVIDIA Developer database for documents across a variety of "
+                         "NVIDIA asset types"))

nat/tool/register.py

@@ -0,0 +1,31 @@
+# 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
+
+# Import any tools which need to be automatically registered here
+from . import chat_completion
+from . import datetime_tools
+from . import document_search
+from . import github_tools
+from . import nvidia_rag
+from . import retriever
+from . import server_tools
+from .code_execution import register
+from .github_tools import github_tool
+from .github_tools import github_files_tool
+from .memory_tools import add_memory_tool
+from .memory_tools import delete_memory_tool
+from .memory_tools import get_memory_tool

nat/tool/retriever.py

@@ -0,0 +1,95 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import logging
+
+from pydantic import BaseModel
+from pydantic import Field
+
+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.component_ref import RetrieverRef
+from nat.data_models.function import FunctionBaseConfig
+from nat.retriever.interface import Retriever
+from nat.retriever.models import RetrieverError
+from nat.retriever.models import RetrieverOutput
+
+logger = logging.getLogger(__name__)
+
+
+class RetrieverConfig(FunctionBaseConfig, name="nat_retriever"):
+    """
+    Retriever tool which provides a common interface for different vectorstores. Its
+    configuration uses clients, which are the vectorstore-specific implementaiton of the retriever interface.
+    """
+    retriever: RetrieverRef = Field(description="The retriever instance name from the workflow configuration object.")
+    raise_errors: bool = Field(
+        default=True,
+        description="If true the tool will raise exceptions, otherwise it will log them as warnings and return []",
+    )
+    topic: str | None = Field(default=None, description="Used to provide a more detailed tool description to the agent")
+    description: str | None = Field(default=None, description="If present it will be used as the tool description")
+
+
+def _get_description_from_config(config: RetrieverConfig) -> str:
+    """
+    Generate a description of what the tool will do based on how it is configured.
+    """
+    description = "Retrieve document chunks{topic} which can be used to answer the provided question."
+
+    _topic = f" related to {config.topic}" if config.topic else ""
+
+    return description.format(topic=_topic) if not config.description else config.description
+
+
+@register_function(config_type=RetrieverConfig)
+async def retriever_tool(config: RetrieverConfig, builder: Builder):
+    """
+    Configure a NAT Retriever Tool which supports different clients such as Milvus and Nemo Retriever.
+
+    Args:
+        config: A config object with required parameters 'client' and 'client_config'
+        builder: A workflow builder object
+    """
+
+    class RetrieverInputSchema(BaseModel):
+        query: str = Field(description="The query to be searched in the configured data store")
+
+    client: Retriever = await builder.get_retriever(config.retriever)
+
+    async def _retrieve(query: str) -> RetrieverOutput:
+        try:
+            retrieved_context = await client.search(query=query)
+            logger.info("Retrieved %s records for query %s.", len(retrieved_context), query)
+            return retrieved_context
+
+        except RetrieverError as e:
+            if config.raise_errors:
+                logger.error("Retriever threw an error: %s.", e)
+                raise
+            logger.exception("Retriever threw an error: %s. Returning an empty response.", e)
+            return RetrieverOutput(results=[])
+
+    yield FunctionInfo.from_fn(
+        fn=_retrieve,
+        input_schema=RetrieverInputSchema,
+        description=_get_description_from_config(config),
+    )
+
+
+# Compatibility aliases with previous releases
+AIQRetrieverConfig = RetrieverConfig
+aiq_retriever_tool = retriever_tool

nat/tool/server_tools.py

@@ -0,0 +1,66 @@
+# 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 nat.builder.builder import Builder
+from nat.builder.function_info import FunctionInfo
+from nat.cli.register_workflow import register_function
+from nat.data_models.function import FunctionBaseConfig
+
+
+class RequestAttributesTool(FunctionBaseConfig, name="current_request_attributes"):
+    """
+    A simple tool that demonstrates how to retrieve user-defined request attributes from HTTP requests
+    within workflow tools. Please refer to the 'general' section of the configuration file located in the
+    'examples/getting_started/simple_web_query/configs/config-metadata.yml' directory to see how to define a
+    custom route using a YAML file and associate it with a corresponding function to acquire request attributes.
+    """
+    pass
+
+
+@register_function(config_type=RequestAttributesTool)
+async def current_request_attributes(config: RequestAttributesTool, builder: Builder):
+
+    from starlette.datastructures import Headers
+    from starlette.datastructures import QueryParams
+
+    async def _get_request_attributes(unused: str) -> str:
+
+        from nat.builder.context import Context
+        nat_context = Context.get()
+
+        method: str | None = nat_context.metadata.method
+        url_path: str | None = nat_context.metadata.url_path
+        url_scheme: str | None = nat_context.metadata.url_scheme
+        headers: Headers | None = nat_context.metadata.headers
+        query_params: QueryParams | None = nat_context.metadata.query_params
+        path_params: dict[str, str] | None = nat_context.metadata.path_params
+        client_host: str | None = nat_context.metadata.client_host
+        client_port: int | None = nat_context.metadata.client_port
+        cookies: dict[str, str] | None = nat_context.metadata.cookies
+        conversation_id: str | None = nat_context.conversation_id
+
+        return (f"Method: {method}, "
+                f"URL Path: {url_path}, "
+                f"URL Scheme: {url_scheme}, "
+                f"Headers: {dict(headers) if headers is not None else 'None'}, "
+                f"Query Params: {dict(query_params) if query_params is not None else 'None'}, "
+                f"Path Params: {path_params}, "
+                f"Client Host: {client_host}, "
+                f"Client Port: {client_port}, "
+                f"Cookies: {cookies}, "
+                f"Conversation Id: {conversation_id}")
+
+    yield FunctionInfo.from_fn(_get_request_attributes,
+                               description="Returns the acquired user defined request attributes.")

nat/utils/callable_utils.py

@@ -0,0 +1,70 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import inspect
+from collections.abc import Callable
+from typing import Any
+
+
+async def ainvoke_any(func: Callable[..., Any], *args: Any, **kwargs: Any) -> Any:
+    """Execute any type of callable and return the result.
+
+    Handles synchronous functions, asynchronous functions, generators,
+    and async generators uniformly, returning the final result value.
+
+    Args:
+        func (Callable[..., Any]): The function to execute (sync/async function, generator, etc.)
+
+    Returns:
+        Any: The result of executing the callable
+    """
+    # Execute the function
+    result_value = func(*args, **kwargs)
+
+    # Handle different return types
+    if inspect.iscoroutine(result_value):
+        # Async function - await the coroutine
+        return await result_value
+
+    if inspect.isgenerator(result_value):
+        # Sync generator - consume until StopIteration and get return value
+        try:
+            while True:
+                next(result_value)
+        except StopIteration as e:
+            # Return the generator's return value, or None if not provided
+            return e.value
+
+    if inspect.isasyncgen(result_value):
+        # Async generator - consume all values and return the last one
+        last_value = None
+        async for value in result_value:
+            last_value = value
+        return last_value
+
+    # Direct value from sync function (most common case)
+    return result_value
+
+
+def is_async_callable(func: Callable[..., Any]) -> bool:
+    """Check if a function is async (coroutine function or async generator function).
+
+    Args:
+        func (Callable[..., Any]): The function to check
+
+    Returns:
+        bool: True if the function is async, False otherwise
+    """
+    return inspect.iscoroutinefunction(func) or inspect.isasyncgenfunction(func)

nat/utils/data_models/schema_validator.py

@@ -0,0 +1,58 @@
+# 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 yaml
+from pydantic import ValidationError
+
+from ..exception_handlers.schemas import schema_exception_handler
+from ..exception_handlers.schemas import yaml_exception_handler
+
+
+@schema_exception_handler
+def validate_schema(metadata, Schema):
+
+    try:
+        return Schema(**metadata)
+    except ValidationError as e:
+
+        raise e
+
+
+@yaml_exception_handler
+def validate_yaml(ctx, param, value):
+    """
+    Validate that the file is a valid YAML file
+
+    Parameters
+    ----------
+    ctx: Click context
+    param: Click parameter
+    value: Path to YAML file
+
+    Returns
+    -------
+    str: Path to valid YAML file
+
+    Raises
+    ------
+    ValueError: If file is invalid or unreadable
+    """
+    if value is None:
+        return None
+
+    with open(value, encoding="utf-8") as f:
+        yaml.safe_load(f)
+
+    return value

nat/utils/debugging_utils.py

@@ -0,0 +1,43 @@
+# 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.
+
+
+def is_debugger_attached() -> bool:
+    """
+    Check if a debugger is attached to the current process.
+
+    Returns
+    -------
+    bool
+        True if a debugger is attached, False otherwise
+    """
+    import sys
+
+    if "debugpy" in sys.modules:
+
+        import debugpy
+
+        return debugpy.is_client_connected()
+
+    trace_func = sys.gettrace()
+
+    # The presence of a trace function and pydevd means a debugger is attached
+    if (trace_func is not None):
+        trace_module = getattr(trace_func, "__module__", None)
+
+        if (trace_module is not None and trace_module.find("pydevd") != -1):
+            return True
+
+    return False