# How to Extend the VSS Agent with Custom NAT Tools and Function Registration

> Extend the NVIDIA AI Toolkit NAT agent by registering custom tools and functions. Learn how to create Python modules, import them, and expose them to the framework.

- Repository: [NVIDIA AI Blueprints/video-search-and-summarization](https://github.com/NVIDIA-AI-Blueprints/video-search-and-summarization)
- Tags: how-to-guide
- Published: 2026-05-15

---

**To extend the VSS agent with custom NAT tools, create a Python module with the `@builder` decorator in `agent/src/vss_agents/tools/`, import it into [`register.py`](https://github.com/NVIDIA-AI-Blueprints/video-search-and-summarization/blob/main/register.py), and add the module name to `__all__` to expose it to the NVIDIA AI Toolkit (NAT) framework.**

The Video Search & Summarization (VSS) agent from the `NVIDIA-AI-Blueprints/video-search-and-summarization` repository is built on NAT, which loads capabilities through a plugin architecture. All agent tools reside as Python modules under `agent/src/vss_agents/tools/` and register automatically through a central shim. By following the NAT builder pattern, you can add custom logic that becomes immediately available to LLM workflows and downstream pipelines.

## Understanding the NAT Tool Architecture

The VSS agent discovers tools through an entry-point group called `vss_tool_plugins` defined in [`pyproject.toml`](https://github.com/NVIDIA-AI-Blueprints/video-search-and-summarization/blob/main/pyproject.toml). When the agent starts, NAT scans this entry point and loads every module listed in the `__all__` array of [`agent/src/vss_agents/tools/register.py`](https://github.com/NVIDIA-AI-Blueprints/video-search-and-summarization/blob/main/agent/src/vss_agents/tools/register.py). Each tool module must expose at least one function decorated with `@builder` from `nvidia_nats.client`, which registers the function as a callable tool with defined input and output schemas.

## Step 1: Create a Custom Tool Module

Create a new Python file inside `agent/src/vss_agents/tools/`. A valid NAT tool uses Pydantic models for validation and the `@builder` decorator to register the function. The decorator requires a unique name, an input schema, and an output schema.

```python

# agent/src/vss_agents/tools/my_echo_tool.py

from nvidia_nats.client import builder, InputSchema, OutputSchema
from pydantic import BaseModel

class EchoInput(BaseModel):
    """Input schema for the echo tool."""
    text: str

class EchoOutput(BaseModel):
    """Output schema – the echoed text in upper case."""
    echoed: str

@builder(name="my_echo", input_schema=EchoInput, output_schema=EchoOutput)
def my_echo(request: EchoInput) -> EchoOutput:
    """Return the supplied text transformed to upper-case."""
    return EchoOutput(echoed=request.text.upper())

```

**Key requirements for the `@builder` decorator:**
- **name** – The unique string identifier used to invoke the tool (e.g., `"my_echo"`).
- **input_schema** – A Pydantic BaseModel defining the expected request payload.
- **output_schema** – A Pydantic BaseModel defining the JSON-serializable return structure.

## Step 2: Register the Tool in the Shim

Open [`agent/src/vss_agents/tools/register.py`](https://github.com/NVIDIA-AI-Blueprints/video-search-and-summarization/blob/main/agent/src/vss_agents/tools/register.py) and import your new module. Add the module name to the `__all__` list to ensure NAT exposes it during agent initialization.

```python

# agent/src/vss_agents/tools/register.py

from . import attribute_search

# ... existing imports ...

from . import vss_summarize
from . import my_echo_tool  # <-- NEW IMPORT

__all__ = [
    "attribute_search",
    # ... existing entries ...

    "vss_summarize",
    "my_echo_tool",  # <-- NEW EXPORT

]

```

The `__all__` array acts as the discovery manifest. NAT reads this list to determine which tool modules to load when the `vss_tool_plugins` entry point is processed.

## Step 3: Configure the Entry Point (Optional)

For external NAT-based services to discover your tool, declare it in the project root's [`pyproject.toml`](https://github.com/NVIDIA-AI-Blueprints/video-search-and-summarization/blob/main/pyproject.toml) under the `vss_tool_plugins` entry-point group.

```toml
[project.entry-points."vss_tool_plugins"]
attribute_search = "vss_agents.tools.attribute_search"

# ... existing tools ...

vss_summarize = "vss_agents.tools.vss_summarize"
my_echo_tool = "vss_agents.tools.my_echo_tool"  # <-- NEW ENTRY

```

This step is optional if you only need internal agent access, but required for standalone discovery by external NAT clients.

## Step 4: Invoke the Custom Tool

After restarting the agent server, NAT loads the new function automatically. You can invoke it via the NAT client SDK from external applications or resolve it internally from other tools.

**Using the NAT client SDK:**

```python
from nvidia_nats.client import NATClient

client = NATClient()
result = client.run_function(
    "my_echo",  # The name defined in the @builder decorator

    {"text": "Hello, VSS!"}  # Payload matching EchoInput

)
print(result)  # {"echoed": "HELLO, VSS!"}

```

**Using the internal resolver from another tool:**

```python
from nvidia_nats.client import resolve

echo_fn = resolve("my_echo")
output = echo_fn({"text": "pipeline test"})

# output == {"echoed": "PIPELINE TEST"}

```

The `resolve` helper allows synchronous calls between tools without overhead, while `NATClient` enables remote invocation.

## Testing Your Custom Tool

The repository includes a unit-test suite under `agent/tests/unit_test/tools/`. Create a corresponding test file to validate your tool's schema and execution path.

```python

# agent/tests/unit_test/tools/test_my_echo_tool.py

def test_my_echo():
    from vss_agents.tools.my_echo_tool import my_echo, EchoInput
    out = my_echo(EchoInput(text="test"))
    assert out.echoed == "TEST"

```

Run `pytest` from the agent directory to verify that the new tool integrates correctly with existing functionality without breaking the pipeline.

## Summary

- **Create** a new Python module in `agent/src/vss_agents/tools/` using the `@builder` decorator and Pydantic schemas to define inputs and outputs.
- **Register** the module by importing it into [`agent/src/vss_agents/tools/register.py`](https://github.com/NVIDIA-AI-Blueprints/video-search-and-summarization/blob/main/agent/src/vss_agents/tools/register.py) and adding it to the `__all__` list.
- **Configure** the `vss_tool_plugins` entry point in [`pyproject.toml`](https://github.com/NVIDIA-AI-Blueprints/video-search-and-summarization/blob/main/pyproject.toml) if external discovery is required.
- **Invoke** the tool via `NATClient.run_function()` for remote calls or `resolve()` for internal tool-to-tool communication.
- **Test** your implementation by adding unit tests in `agent/tests/unit_test/tools/` to ensure schema compliance and correctness.

## Frequently Asked Questions

### What is the `@builder` decorator in NAT?

The `@builder` decorator is provided by the `nvidia_nats.client` module and registers a Python function as a callable tool within the NAT framework. It requires a unique name, an input Pydantic model, and an output Pydantic model, enabling automatic validation and documentation generation.

### Do I need to modify [`pyproject.toml`](https://github.com/NVIDIA-AI-Blueprints/video-search-and-summarization/blob/main/pyproject.toml) for every new tool?

You only need to modify [`pyproject.toml`](https://github.com/NVIDIA-AI-Blueprints/video-search-and-summarization/blob/main/pyproject.toml) if you require the tool to be discoverable by external NAT-based services or standalone clients. For tools used exclusively within the VSS agent, registering in [`register.py`](https://github.com/NVIDIA-AI-Blueprints/video-search-and-summarization/blob/main/register.py) and adding to `__all__` is sufficient.

### Can custom tools perform I/O operations like writing to S3?

Yes, the function body inside a `@builder` decorated tool can execute any Python code, including I/O operations such as writing to S3, calling external APIs, or querying databases. However, ensure the function handles errors gracefully since NAT expects a JSON-serializable output matching the defined output schema.

### How does the VSS agent discover tools at runtime?

The VSS agent discovers tools through the `vss_tool_plugins` entry-point group defined in [`pyproject.toml`](https://github.com/NVIDIA-AI-Blueprints/video-search-and-summarization/blob/main/pyproject.toml). NAT scans this entry point and loads each module listed in the `__all__` array of [`agent/src/vss_agents/tools/register.py`](https://github.com/NVIDIA-AI-Blueprints/video-search-and-summarization/blob/main/agent/src/vss_agents/tools/register.py), automatically registering any functions decorated with `@builder`.