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

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, 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. 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. 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.

# 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 and import your new module. Add the module name to the __all__ list to ensure NAT exposes it during agent initialization.

# 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 under the vss_tool_plugins entry-point group.

[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:

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:

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.

# 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 and adding it to the __all__ list.
• Configure the vss_tool_plugins entry point in 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 for every new tool?

You only need to modify 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 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. NAT scans this entry point and loads each module listed in the __all__ array of agent/src/vss_agents/tools/register.py, automatically registering any functions decorated with @builder.

{
  "mcpServers": {
    "instagit": {
      "command": "npx",
      "args": ["-y", "instagit@latest"]
    }
  }
}