From 62319e7cbcc6d334b814ead0a10176656c75a435 Mon Sep 17 00:00:00 2001 From: Bryan Bednarski Date: Mon, 18 May 2026 13:02:59 -0700 Subject: [PATCH 01/13] first pass at stabilizing the public API Signed-off-by: Bryan Bednarski --- .../build-workflows/advanced/middleware.md | 16 +-- .../function-groups.md | 12 +- docs/source/components/sharing-components.md | 2 +- .../adding-an-authentication-provider.md | 2 +- .../adding-an-llm-provider.md | 8 +- .../custom-dataset-loader.md | 4 +- .../custom-components/custom-evaluator.md | 8 +- .../custom-functions/function-groups.md | 77 ++++++------ .../custom-functions/functions.md | 20 +-- .../custom-functions/per-user-functions.md | 19 +-- .../extend/custom-components/finetuning.md | 12 +- .../extend/custom-components/object-store.md | 4 +- .../extend/custom-components/optimizer.md | 8 +- .../custom-components/telemetry-exporters.md | 16 +-- docs/source/extend/plugins.md | 36 +++--- .../testing/add-unit-tests-for-tools.md | 31 +++++ .../tutorials/add-tools-to-a-workflow.md | 10 +- docs/source/improve-workflows/evaluate.md | 2 +- docs/source/improve-workflows/optimizer.md | 4 +- .../improve-workflows/test-time-compute.md | 2 +- docs/source/index.md | 1 + .../workflow/templates/workflow.py.j2 | 10 +- .../tests/nat/tools/test_tool_test_runner.py | 40 ++++++ .../langchain/tools/exa_internet_search.py | 4 +- .../langchain/tools/tavily_internet_search.py | 4 +- .../src/nat/test/tool_test_runner.py | 114 ++++++++++++++++++ 26 files changed, 329 insertions(+), 137 deletions(-) diff --git a/docs/source/build-workflows/advanced/middleware.md b/docs/source/build-workflows/advanced/middleware.md index 4635ddb6ca..b4d7c322ae 100644 --- a/docs/source/build-workflows/advanced/middleware.md +++ b/docs/source/build-workflows/advanced/middleware.md @@ -244,8 +244,8 @@ Key benefits of extending `DynamicFunctionMiddleware`: Create a registration module following the idiomatic pattern: ```python -from nat.builder.builder import Builder -from nat.cli.register_workflow import register_middleware +from nat.plugin_api import Builder +from nat.plugin_api import register_middleware from .logging_middleware import LoggingMiddleware, LoggingMiddlewareConfig @@ -289,8 +289,8 @@ functions: Register your function without needing to specify middleware in the decorator: ```python -from nat.cli.register_workflow import register_function -from nat.builder.builder import Builder +from nat.plugin_api import register_function +from nat.plugin_api import Builder @register_function(config_type=MyAPIFunctionConfig) @@ -521,9 +521,9 @@ function_groups: ``` ```python -from nat.cli.register_workflow import register_function_group -from nat.builder.function import FunctionGroup -from nat.data_models.function import FunctionGroupBaseConfig +from nat.plugin_api import register_function_group +from nat.plugin_api import FunctionGroup +from nat.plugin_api import FunctionGroupBaseConfig class WeatherAPIGroupConfig(FunctionGroupBaseConfig, name="weather_api_group"): @@ -828,7 +828,7 @@ Solution: Ensure the register module is imported. NeMo Agent Toolkit automatical - {py:class}`~nat.middleware.function_middleware.FunctionMiddlewareChain`: Chain management - {py:class}`~nat.middleware.cache.cache_middleware_config.CacheMiddlewareConfig`: Cache configuration - {py:class}`~nat.middleware.cache.cache_middleware.CacheMiddleware`: Cache implementation -- {py:func}`~nat.cli.register_workflow.register_middleware`: Registration decorator +- {py:func}`~nat.plugin_api.register_middleware`: Registration decorator ## See Also diff --git a/docs/source/build-workflows/functions-and-function-groups/function-groups.md b/docs/source/build-workflows/functions-and-function-groups/function-groups.md index 1ea1a46acb..6886b0bd86 100644 --- a/docs/source/build-workflows/functions-and-function-groups/function-groups.md +++ b/docs/source/build-workflows/functions-and-function-groups/function-groups.md @@ -206,7 +206,7 @@ workflow: - **Multiple functions need the same connection** (database, API client, cache) - **Functions share configuration** (credentials, endpoints, settings) -- **You want to namespace related functions** (`math.add`, `math.multiply`) +- **You want to namespace related functions** (`math__add`, `math__multiply`) - **Functions need to share state** (session data, counters, caches) - **You have a family of operations** (CRUD operations, data transformations) @@ -495,7 +495,7 @@ workflow: llm_name: my_llm ``` -All functions in the `math` group (`math.add`, `math.multiply`) become available as tools for the agent. +All functions in the `math` group (`math__add`, `math__multiply`) become available as tools for the agent. #### Example 2: Including Specific Functions @@ -602,7 +602,7 @@ async with WorkflowBuilder() as builder: To wrap all accessible functions in a group for a specific agent framework: ```python -from nat.data_models.component_ref import FunctionGroupRef +from nat.plugin_api import FunctionGroupRef from nat.builder.framework_enum import LLMFrameworkEnum async with WorkflowBuilder() as builder: @@ -640,7 +640,7 @@ async with WorkflowBuilder() as builder: # Now only "add" functions are accessible accessible = await math_group.get_accessible_functions() - # Returns: ["math.add"] + # Returns: ["math__add"] ``` #### Per-Function Filters @@ -758,11 +758,11 @@ Instance names become part of function names, so keep them concise: ```python # Good group = FunctionGroup(config=config, instance_name="db") -# Results in: db.query, db.insert +# Results in: db__query, db__insert # Less ideal group = FunctionGroup(config=config, instance_name="database_operations") -# Results in: database_operations.query, database_operations.insert +# Results in: database_operations__query, database_operations__insert ``` #### Use Environment Variables for Secrets diff --git a/docs/source/components/sharing-components.md b/docs/source/components/sharing-components.md index fc8a959e6b..61b358fe7e 100644 --- a/docs/source/components/sharing-components.md +++ b/docs/source/components/sharing-components.md @@ -58,7 +58,7 @@ requirements. ```python from pydantic import Field -from nat.data_models.function import FunctionBaseConfig +from nat.plugin_api import FunctionBaseConfig class MyFnConfig(FunctionBaseConfig, name="my_fn_name"): # includes a name """The docstring should provide a description of the components utility.""" # includes a docstring diff --git a/docs/source/extend/custom-components/adding-an-authentication-provider.md b/docs/source/extend/custom-components/adding-an-authentication-provider.md index 1b699cf5e5..b7eb2d60bd 100644 --- a/docs/source/extend/custom-components/adding-an-authentication-provider.md +++ b/docs/source/extend/custom-components/adding-an-authentication-provider.md @@ -70,7 +70,7 @@ class OAuth2AuthCodeFlowProviderConfig(AuthProviderBaseConfig, name="oauth2_auth ``` ### Registering the Provider -An asynchronous function decorated with {py:func}`~nat.cli.register_workflow.register_auth_provider` is used to register the provider with NeMo Agent Toolkit by yielding an instance of +An asynchronous function decorated with {py:func}`~nat.plugin_api.register_auth_provider` is used to register the provider with NeMo Agent Toolkit by yielding an instance of {py:class}`~nat.authentication.interfaces.AuthProviderBase`. The `OAuth2AuthCodeFlowProviderConfig` from the previous section is registered as follows: diff --git a/docs/source/extend/custom-components/adding-an-llm-provider.md b/docs/source/extend/custom-components/adding-an-llm-provider.md index 9e79547ac1..ba851c9d55 100644 --- a/docs/source/extend/custom-components/adding-an-llm-provider.md +++ b/docs/source/extend/custom-components/adding-an-llm-provider.md @@ -118,10 +118,10 @@ class NIMModelConfig(LLMBaseConfig, ThinkingMixin, name="nim"): ``` ### Registering the Provider -An asynchronous function decorated with {py:deco}`nat.cli.register_workflow.register_llm_provider` is used to register the provider with NeMo Agent Toolkit by yielding an instance of {class}`nat.builder.llm.LLMProviderInfo`. +An asynchronous function decorated with {py:deco}`nat.plugin_api.register_llm_provider` is used to register the provider with NeMo Agent Toolkit by yielding an instance of {class}`nat.builder.llm.LLMProviderInfo`. :::{note} -Registering an embedder or retriever provider is similar; however, the function should be decorated with {py:deco}`nat.cli.register_workflow.register_embedder_provider` or {py:deco}`nat.cli.register_workflow.register_retriever_provider`. +Registering an embedder or retriever provider is similar; however, the function should be decorated with {py:deco}`nat.plugin_api.register_embedder_provider` or {py:deco}`nat.plugin_api.register_retriever_provider`. ::: @@ -146,10 +146,10 @@ async def openai_llm(config: OpenAIModelConfig, builder: Builder): ``` ## LLM Clients -As previously mentioned, each LLM client is specific to both the LLM API and the framework being used. The LLM client is registered by defining an asynchronous function decorated with {py:deco}`nat.cli.register_workflow.register_llm_client`. The `register_llm_client` decorator receives two required parameters: `config_type`, which is the configuration class of the provider, and `wrapper_type`, which identifies the framework being used. +As previously mentioned, each LLM client is specific to both the LLM API and the framework being used. The LLM client is registered by defining an asynchronous function decorated with {py:deco}`nat.plugin_api.register_llm_client`. The `register_llm_client` decorator receives two required parameters: `config_type`, which is the configuration class of the provider, and `wrapper_type`, which identifies the framework being used. :::{note} -Registering an embedder or retriever client is similar. However, the function should be decorated with {py:deco}`nat.cli.register_workflow.register_embedder_client` or {py:deco}`nat.cli.register_workflow.register_retriever_client`. +Registering an embedder or retriever client is similar. However, the function should be decorated with {py:deco}`nat.plugin_api.register_embedder_client` or {py:deco}`nat.plugin_api.register_retriever_client`. ::: The wrapped function in turn receives two required positional arguments: an instance of the configuration class of the provider, and an instance of {class}`nat.builder.builder.Builder`. The function should then yield a client suitable for the given provider and framework. The exact type is dictated by the framework itself and not by NeMo Agent Toolkit. diff --git a/docs/source/extend/custom-components/custom-dataset-loader.md b/docs/source/extend/custom-components/custom-dataset-loader.md index 85cdfc5f0a..59e0b58680 100644 --- a/docs/source/extend/custom-components/custom-dataset-loader.md +++ b/docs/source/extend/custom-components/custom-dataset-loader.md @@ -48,9 +48,9 @@ The following example shows how to define and register a custom dataset loader f import pandas as pd from pydantic import Field -from nat.builder.builder import EvalBuilder +from nat.plugin_api import EvalBuilder from nat.builder.dataset_loader import DatasetLoaderInfo -from nat.cli.register_workflow import register_dataset_loader +from nat.plugin_api import register_dataset_loader from nat.data_models.dataset_handler import EvalDatasetBaseConfig diff --git a/docs/source/extend/custom-components/custom-evaluator.md b/docs/source/extend/custom-components/custom-evaluator.md index 88b72e352e..f5577acaaa 100644 --- a/docs/source/extend/custom-components/custom-evaluator.md +++ b/docs/source/extend/custom-components/custom-evaluator.md @@ -48,9 +48,9 @@ The following example shows how to define and register a custom evaluator. The c ```python from pydantic import Field -from nat.builder.builder import EvalBuilder +from nat.plugin_api import EvalBuilder from nat.builder.evaluator import EvaluatorInfo -from nat.cli.register_workflow import register_evaluator +from nat.plugin_api import register_evaluator from nat.data_models.evaluator import EvaluatorBaseConfig @@ -165,9 +165,9 @@ from collections import Counter from pydantic import Field -from nat.builder.builder import EvalBuilder +from nat.plugin_api import EvalBuilder from nat.builder.evaluator import EvaluatorInfo -from nat.cli.register_workflow import register_evaluator +from nat.plugin_api import register_evaluator from nat.data_models.evaluator import EvaluatorBaseConfig from nat.plugins.eval.data_models.evaluator_io import EvalOutputItem from nat.plugins.eval.evaluator.atif_base_evaluator import AtifBaseEvaluator diff --git a/docs/source/extend/custom-components/custom-functions/function-groups.md b/docs/source/extend/custom-components/custom-functions/function-groups.md index b0c73bd370..f1faa9f2e0 100644 --- a/docs/source/extend/custom-components/custom-functions/function-groups.md +++ b/docs/source/extend/custom-components/custom-functions/function-groups.md @@ -31,18 +31,18 @@ Create a custom function group when you need to: - **Bundle related operations**: Group CRUD operations, file operations, or API endpoints that belong together - **Centralize configuration**: Manage credentials, endpoints, and settings in one place for multiple functions - **Create reusable components**: Package functionality that can be used across multiple workflows -- **Namespace functions**: Organize functions into logical groups, such as `db.query`, `db.insert`, `api.get`, and `api.post` +- **Namespace functions**: Organize functions into logical groups, such as `db__query`, `db__insert`, `api__get`, and `api__post` ## Step 1: Define the Configuration -Every function group needs a configuration class that inherits from {py:class}`~nat.data_models.function.FunctionGroupBaseConfig`. +Every function group needs a configuration class that inherits from {py:class}`~nat.plugin_api.FunctionGroupBaseConfig`. ### Minimal Configuration Start with the simplest possible configuration: ```python -from nat.data_models.function import FunctionGroupBaseConfig +from nat.plugin_api import FunctionGroupBaseConfig class MyGroupConfig(FunctionGroupBaseConfig, name="my_group"): """Configuration for my custom function group.""" @@ -57,7 +57,8 @@ Add fields for any settings your functions need to share: ```python from pydantic import Field -from nat.data_models.function import FunctionGroupBaseConfig + +from nat.plugin_api import FunctionGroupBaseConfig class DatabaseGroupConfig(FunctionGroupBaseConfig, name="database_group"): """Configuration for database operations.""" @@ -85,7 +86,7 @@ function_groups: ### Controlling Function Exposure -The {py:class}`~nat.data_models.function.FunctionGroupBaseConfig` configuration class has two optional fields: `include` and `exclude`. These fields are used to control which functions are exposed through the function group or excluded from the function group. +The {py:class}`~nat.plugin_api.FunctionGroupBaseConfig` configuration class has two optional fields: `include` and `exclude`. These fields are used to control which functions are exposed through the function group or excluded from the function group. If your function group is intended to override the default behavior of the function group, you can use the `include` field to specify which functions to expose and the `exclude` field to specify which functions to exclude. @@ -121,17 +122,17 @@ When to use `include`, `exclude`, or neither: ## Step 2: Register and Implement the Function Group -Use the {py:deco}`~nat.cli.register_workflow.register_function_group` decorator to register your function group builder. +Use the {py:deco}`~nat.plugin_api.register_function_group` decorator to register your function group builder. ### Basic Implementation Here's the simplest function group implementation: ```python -from nat.builder.workflow_builder import Builder -from nat.builder.function import FunctionGroup -from nat.cli.register_workflow import register_function_group -from nat.data_models.function import FunctionGroupBaseConfig +from nat.plugin_api import Builder +from nat.plugin_api import FunctionGroup +from nat.plugin_api import FunctionGroupBaseConfig +from nat.plugin_api import register_function_group class MyGroupConfig(FunctionGroupBaseConfig, name="my_group"): """Configuration for my custom function group.""" @@ -161,7 +162,7 @@ async def build_my_group(config: MyGroupConfig, _builder: Builder): **Key components**: - **Decorator**: `@register_function_group(config_type=MyGroupConfig)` registers the builder -- **Instance name**: `instance_name="my"` creates the namespace (`my.greet`, `my.farewell`) +- **Instance name**: `instance_name="my"` creates the namespace (`my__greet`, `my__farewell`) - **Function definitions**: Define async functions that implement your logic - **Add to group**: Use `group.add_function()` to register each function - **Yield**: `yield group` makes the group available to workflows @@ -172,7 +173,8 @@ Access configuration values in your functions to customize behavior: ```python import httpx -from nat.cli.register_workflow import register_function_group + +from nat.plugin_api import register_function_group @register_function_group(config_type=APIGroupConfig) async def build_api_group(config: APIGroupConfig, _builder: Builder): @@ -212,9 +214,9 @@ For functions that need shared resources (for example, connections and clients), ```python import asyncpg -from nat.cli.register_workflow import register_function_group -from nat.builder.workflow_builder import Builder -from nat.builder.function import FunctionGroup +from nat.plugin_api import Builder +from nat.plugin_api import FunctionGroup +from nat.plugin_api import register_function_group @register_function_group(config_type=DatabaseGroupConfig) async def build_database_group(config: DatabaseGroupConfig, _builder: Builder): @@ -276,7 +278,7 @@ After creating your function group, you can work with it programmatically in you ### Accessing Functions -Functions are referenced as `instance_name.function_name`: +Functions are referenced as `instance_name__function_name`: ```python from nat.builder.workflow_builder import WorkflowBuilder @@ -286,7 +288,7 @@ async with WorkflowBuilder() as builder: await builder.add_function_group("my", MyGroupConfig(include=["greet", "farewell"])) # Access individual function by fully qualified name - greet = await builder.get_function("my.greet") + greet = await builder.get_function("my__greet") result = await greet.ainvoke("World") print(result) # "Hello, World!" ``` @@ -321,30 +323,26 @@ async with WorkflowBuilder() as builder: ### Testing Your Function Group -Test individual functions through the group: +Test individual functions through the group with `nvidia-nat-test`: ```python -import pytest -from nat.builder.workflow_builder import WorkflowBuilder +from nat.test import ToolTestRunner + -@pytest.mark.asyncio async def test_my_function_group(): - async with WorkflowBuilder() as builder: - await builder.add_function_group("my", MyGroupConfig()) - my_group = await builder.get_function_group("my") - - # Test each function - all_funcs = await my_group.get_all_functions() - - # Test greet function - greet = all_funcs["greet"] - result = await greet.ainvoke("Alice") - assert result == "Hello, Alice!" - - # Test farewell function - farewell = all_funcs["farewell"] - result = await farewell.ainvoke("Bob") - assert result == "Goodbye, Bob!" + runner = ToolTestRunner() + + await runner.test_function_group( + config_type=MyGroupConfig, + expected_functions=["my__greet", "my__farewell"], + ) + + result = await runner.test_function_group_tool( + config_type=MyGroupConfig, + function_name="greet", + input_data="Alice", + ) + assert result == "Hello, Alice!" ``` ## Step 5: Advanced - Dynamic Filtering (Optional) @@ -374,8 +372,9 @@ Group-level filters receive a list of function names and return a filtered list: ```python from collections.abc import Sequence -from nat.cli.register_workflow import register_function_group -from nat.builder.function import FunctionGroup + +from nat.plugin_api import FunctionGroup +from nat.plugin_api import register_function_group class EnvironmentGroupConfig(FunctionGroupBaseConfig, name="env_group"): """Configuration with environment setting.""" diff --git a/docs/source/extend/custom-components/custom-functions/functions.md b/docs/source/extend/custom-components/custom-functions/functions.md index e9fc7589a5..8ae5fddd51 100644 --- a/docs/source/extend/custom-components/custom-functions/functions.md +++ b/docs/source/extend/custom-components/custom-functions/functions.md @@ -68,7 +68,7 @@ Both of these methods will result in a function that can be used in the same way ### Function Configuration Object -To use a function from a configuration file, it must be registered with NeMo Agent Toolkit. Registering a function is done with the {py:deco}`nat.cli.register_workflow.register_function` decorator. More information about registering components can be found in the [Plugin System](../../plugins.md) documentation. +To use a function from a configuration file, it must be registered with NeMo Agent Toolkit. Registering a function is done with the {py:deco}`nat.plugin_api.register_function` decorator. More information about registering components can be found in the [Plugin System](../../plugins.md) documentation. When registering a function, we first need to define the function configuration object. This object is used to configure the function and is passed to the function when it is invoked. Any options that are available to the function must be specified in the configuration object. @@ -82,7 +82,7 @@ class MyFunctionConfig(FunctionBaseConfig, name="my_function"): option3: dict[str, float] ``` -The configuration object must inherit from {py:class}`~nat.data_models.function.FunctionBaseConfig` and must have a `name` attribute. The `name` attribute is used to identify the function in the configuration file. +The configuration object must inherit from {py:class}`~nat.plugin_api.FunctionBaseConfig` and must have a `name` attribute. The `name` attribute is used to identify the function in the configuration file. Additionally, the configuration object can use Pydantic's features to provide validation and documentation for each of the options. For example, the following configuration will validate that `option2` is a positive integer, and documents all properties with a description and default value. @@ -101,7 +101,7 @@ This additional metadata will ensure that the configuration object is properly v With the configuration object defined, there are several options available to register the function: -* **Register a function from a callable using {py:class}`~nat.builder.function_info.FunctionInfo`**: +* **Register a function from a callable using {py:class}`~nat.plugin_api.FunctionInfo`**: ```python @register_function(config_type=MyFunctionConfig) @@ -194,7 +194,7 @@ With the configuration object defined, there are several options available to re ## Initialization and Cleanup -Its required to use an async context manager coroutine to register a function (it's not necessary to use `@asynccontextmanager`, since {py:deco}`nat.cli.register_workflow.register_function` does this for you). This is because the function may need to execute some initialization before construction or cleanup after it is used. For example, if the function needs to load a model, connect to a resource, or download data, this can be done in the register function. +Its required to use an async context manager coroutine to register a function (it's not necessary to use `@asynccontextmanager`, since {py:deco}`nat.plugin_api.register_function` does this for you). This is because the function may need to execute some initialization before construction or cleanup after it is used. For example, if the function needs to load a model, connect to a resource, or download data, this can be done in the register function. ```python @register_function(config_type=MyFunctionConfig) @@ -296,7 +296,7 @@ async def my_function(config: MyFunctionConfig, builder: Builder): ### Functions with Multiple Arguments -It is possible to create a function with a callable that has multiple arguments. When a function with multiple arguments is passed to {py:meth}`~nat.builder.function_info.FunctionInfo.from_fn`, the function will be wrapped with a lambda function which takes a single argument and passes it to the original function. For example, the following function takes two arguments, `input_data` and `repeat`: +It is possible to create a function with a callable that has multiple arguments. When a function with multiple arguments is passed to {py:meth}`~nat.plugin_api.FunctionInfo.from_fn`, the function will be wrapped with a lambda function which takes a single argument and passes it to the original function. For example, the following function takes two arguments, `input_data` and `repeat`: ```python async def multi_arg_function(input_data: list[float], repeat: int) -> list[float]: @@ -338,7 +338,7 @@ class MyFunction(Function[MyInput, MySingleOutput, MyStreamingOutput]): yield MyStreamingOutput(value, i) ``` -Similarly this can be accomplished using {py:meth}`~nat.builder.function_info.FunctionInfo.create` which is a more verbose version of {py:meth}`~nat.builder.function_info.FunctionInfo.from_fn`. +Similarly this can be accomplished using {py:meth}`~nat.plugin_api.FunctionInfo.create` which is a more verbose version of {py:meth}`~nat.plugin_api.FunctionInfo.from_fn`. ```python async def my_ainvoke(self, value: MyInput) -> MySingleOutput: @@ -358,7 +358,7 @@ assert function_info.single_output_type == MySingleOutput assert function_info.stream_output_type == MyStreamingOutput ``` -Finally, when using {py:meth}`~nat.builder.function_info.FunctionInfo.create` a conversion function can be provided to convert the single output to a streaming output, and a streaming output into a single output. This is useful when converting between streaming and single outputs is trivial and defining both methods would be overkill. For example, the following function converts a streaming output to a single output by joining the items with a comma: +Finally, when using {py:meth}`~nat.plugin_api.FunctionInfo.create` a conversion function can be provided to convert the single output to a streaming output, and a streaming output into a single output. This is useful when converting between streaming and single outputs is trivial and defining both methods would be overkill. For example, the following function converts a streaming output to a single output by joining the items with a comma: ```python # Define a conversion function to convert a streaming output to a single output @@ -412,7 +412,7 @@ Output schemas can also be overridden in a similar manner but for different purp ## Instantiating Functions -Once a function is registered, it can be instantiated using the {py:class}`~nat.builder.workflow_builder.WorkflowBuilder` class. The `WorkflowBuilder` class is used to create and manage all components in a workflow. When calling {py:meth}`~nat.builder.workflow_builder.WorkflowBuilder.add_function`, which function to create is determined by the type of the configuration object. The builder will match the configuration object type to the type used in the {py:deco}`nat.cli.register_workflow.register_function` decorator. +Once a function is registered, it can be instantiated by a workflow builder. The builder will match the configuration object type to the type used in the {py:deco}`nat.plugin_api.register_function` decorator. ```python @@ -526,7 +526,7 @@ async def ainvoke(value: typing.Any, to_type: type): ### Adding Custom Converters -Functions support custom type converters for complex conversion scenarios. To add a custom converter to a function, provide a list of converter callables to the {py:meth}`~nat.builder.function_info.FunctionInfo.from_fn` or {py:meth}`~nat.builder.function_info.FunctionInfo.create` methods when creating a function. A converter callable is any python function which takes a single value and returns a converted value. These functions must be annotated with the type it will convert from and the type it will convert to. +Functions support custom type converters for complex conversion scenarios. To add a custom converter to a function, provide a list of converter callables to the {py:meth}`~nat.plugin_api.FunctionInfo.from_fn` or {py:meth}`~nat.plugin_api.FunctionInfo.create` methods when creating a function. A converter callable is any python function which takes a single value and returns a converted value. These functions must be annotated with the type it will convert from and the type it will convert to. For example, the following converter will convert an `int` to a `str`: @@ -535,7 +535,7 @@ def my_converter(value: int) -> str: return str(value) ``` -This converter can then be passed to the {py:meth}`~nat.builder.function_info.FunctionInfo.from_fn` or {py:meth}`~nat.builder.function_info.FunctionInfo.create` methods when registering the function: +This converter can then be passed to the {py:meth}`~nat.plugin_api.FunctionInfo.from_fn` or {py:meth}`~nat.plugin_api.FunctionInfo.create` methods when registering the function: ```python @register_function(config_type=MyFunctionConfig) diff --git a/docs/source/extend/custom-components/custom-functions/per-user-functions.md b/docs/source/extend/custom-components/custom-functions/per-user-functions.md index 46fdc6154b..ffb55b3fe7 100644 --- a/docs/source/extend/custom-components/custom-functions/per-user-functions.md +++ b/docs/source/extend/custom-components/custom-functions/per-user-functions.md @@ -36,14 +36,15 @@ Per-user functions are useful when you need: ### The `@register_per_user_function` Decorator -To register a per-user function, use the {py:deco}`nat.cli.register_workflow.register_per_user_function` decorator. This decorator is similar to {py:deco}`nat.cli.register_workflow.register_function` but requires explicit schema definitions for input and output types. +To register a per-user function, use the {py:deco}`nat.plugin_api.register_per_user_function` decorator. This decorator is similar to {py:deco}`nat.plugin_api.register_function` but requires explicit schema definitions for input and output types. ```python from pydantic import BaseModel, Field -from nat.builder.builder import Builder -from nat.builder.function_info import FunctionInfo -from nat.cli.register_workflow import register_per_user_function -from nat.data_models.function import FunctionBaseConfig + +from nat.plugin_api import Builder +from nat.plugin_api import FunctionBaseConfig +from nat.plugin_api import FunctionInfo +from nat.plugin_api import register_per_user_function # Define input and output schemas @@ -114,12 +115,12 @@ async def with_simple_types(config, builder): ## Registering Per-User Function Groups -Function groups that need per-user state can be registered using the {py:deco}`nat.cli.register_workflow.register_per_user_function_group` decorator. +Function groups that need per-user state can be registered using the {py:deco}`nat.plugin_api.register_per_user_function_group` decorator. ```python -from nat.cli.register_workflow import register_per_user_function_group -from nat.data_models.function import FunctionGroupBaseConfig -from nat.builder.function import FunctionGroup +from nat.plugin_api import FunctionGroup +from nat.plugin_api import FunctionGroupBaseConfig +from nat.plugin_api import register_per_user_function_group class MyPerUserGroupConfig(FunctionGroupBaseConfig, name="my_per_user_group"): diff --git a/docs/source/extend/custom-components/finetuning.md b/docs/source/extend/custom-components/finetuning.md index 1cdd6a51f8..5e38b370d9 100644 --- a/docs/source/extend/custom-components/finetuning.md +++ b/docs/source/extend/custom-components/finetuning.md @@ -171,8 +171,8 @@ Implement the `TrajectoryBuilder` interface's methods. Create a registration module: ```python -from nat.builder.builder import Builder -from nat.cli.register_workflow import register_trajectory_builder +from nat.plugin_api import Builder +from nat.plugin_api import register_trajectory_builder from .my_trajectory_builder import MyTrajectoryBuilder, MyTrajectoryBuilderConfig @@ -297,8 +297,8 @@ Implement the `TrainerAdapter` interface's methods. #### Step 3: Register the Component ```python -from nat.builder.builder import Builder -from nat.cli.register_workflow import register_trainer_adapter +from nat.plugin_api import Builder +from nat.plugin_api import register_trainer_adapter from .my_trainer_adapter import MyTrainerAdapter, MyTrainerAdapterConfig @@ -411,8 +411,8 @@ define configuration, implement methods, and register the component. Once you have your `MyTrainer` and `MyTrainerConfig` implemented, register it as follows: ```python -from nat.builder.builder import Builder -from nat.cli.register_workflow import register_trainer +from nat.plugin_api import Builder +from nat.plugin_api import register_trainer from .my_trainer import MyTrainer, MyTrainerConfig diff --git a/docs/source/extend/custom-components/object-store.md b/docs/source/extend/custom-components/object-store.md index 9da2c78305..4547d49641 100644 --- a/docs/source/extend/custom-components/object-store.md +++ b/docs/source/extend/custom-components/object-store.md @@ -153,8 +153,8 @@ In the NeMo Agent Toolkit system, anything that extends {py:class}`~nat.data_mod 3. **Register your object store with NeMo Agent Toolkit** using the `@register_object_store` decorator: ```python - from nat.builder.builder import Builder - from nat.cli.register_workflow import register_object_store + from nat.plugin_api import Builder + from nat.plugin_api import register_object_store @register_object_store(config_type=MyCustomObjectStoreConfig) async def my_custom_object_store(config: MyCustomObjectStoreConfig, _builder: Builder): diff --git a/docs/source/extend/custom-components/optimizer.md b/docs/source/extend/custom-components/optimizer.md index 970716ac6f..1752c403eb 100644 --- a/docs/source/extend/custom-components/optimizer.md +++ b/docs/source/extend/custom-components/optimizer.md @@ -35,7 +35,7 @@ NeMo Agent Toolkit provides a pluggable optimizer system for tuning workflow par - {py:class}`~nat.plugins.config_optimizer.parameters.base.BaseParameterOptimizer`: Abstract base class for parameter optimization strategies. Requires implementing an async `run()` method that returns an optimized `Config`. * **Registration** - - {py:deco}`~nat.cli.register_workflow.register_optimizer`: Decorator that registers an optimizer strategy with the global type registry so the optimizer runtime can resolve the strategy from the type of `cfg.optimizer.numeric` or `cfg.optimizer.prompt`. + - {py:deco}`~nat.plugin_api.register_optimizer`: Decorator that registers an optimizer strategy with the global type registry so the optimizer runtime can resolve the strategy from the type of `cfg.optimizer.numeric` or `cfg.optimizer.prompt`. ## Adding a Custom Prompt Optimizer @@ -96,10 +96,10 @@ The `run()` method receives: ### 3. Register the Optimizer -Use the {py:deco}`~nat.cli.register_workflow.register_optimizer` decorator to register your strategy: +Use the {py:deco}`~nat.plugin_api.register_optimizer` decorator to register your strategy: ```python -from nat.cli.register_workflow import register_optimizer +from nat.plugin_api import register_optimizer @register_optimizer(config_type=IterativeRefinementPromptConfig) @@ -194,7 +194,7 @@ class RandomSearchOptimizer(BaseParameterOptimizer): ### 3. Register and Configure ```python -from nat.cli.register_workflow import register_optimizer +from nat.plugin_api import register_optimizer @register_optimizer(config_type=RandomSearchConfig) diff --git a/docs/source/extend/custom-components/telemetry-exporters.md b/docs/source/extend/custom-components/telemetry-exporters.md index 6d2652780a..91e310610c 100644 --- a/docs/source/extend/custom-components/telemetry-exporters.md +++ b/docs/source/extend/custom-components/telemetry-exporters.md @@ -88,8 +88,8 @@ Want to get started quickly? Here's a minimal working example that creates a con ```python from pydantic import Field -from nat.builder.builder import Builder -from nat.cli.register_workflow import register_telemetry_exporter +from nat.plugin_api import Builder +from nat.plugin_api import register_telemetry_exporter from nat.data_models.telemetry_exporter import TelemetryExporterBaseConfig from nat.observability.exporter.raw_exporter import RawExporter from nat.data_models.intermediate_step import IntermediateStep @@ -469,8 +469,8 @@ Create a registration function using the `@register_telemetry_exporter` decorato ```python import logging -from nat.builder.builder import Builder -from nat.cli.register_workflow import register_telemetry_exporter +from nat.plugin_api import Builder +from nat.plugin_api import register_telemetry_exporter logger = logging.getLogger(__name__) @@ -537,9 +537,9 @@ class MyCustomExporter(SpanExporter[Span, dict]): ```python from pydantic import Field -from nat.cli.register_workflow import register_telemetry_exporter +from nat.plugin_api import register_telemetry_exporter from nat.data_models.telemetry_exporter import TelemetryExporterBaseConfig -from nat.builder.builder import Builder +from nat.plugin_api import Builder # Configuration class can be in the same file as registration class MyTelemetryExporter(TelemetryExporterBaseConfig, name="my_exporter"): @@ -1423,8 +1423,8 @@ Here's a complete example of a custom telemetry exporter: import logging from pydantic import Field import aiohttp -from nat.builder.builder import Builder -from nat.cli.register_workflow import register_telemetry_exporter +from nat.plugin_api import Builder +from nat.plugin_api import register_telemetry_exporter from nat.data_models.telemetry_exporter import TelemetryExporterBaseConfig from nat.observability.exporter.span_exporter import SpanExporter from nat.observability.exporter.base_exporter import IsolatedAttribute diff --git a/docs/source/extend/plugins.md b/docs/source/extend/plugins.md index 92e02a1ce4..5b7b4175d8 100644 --- a/docs/source/extend/plugins.md +++ b/docs/source/extend/plugins.md @@ -28,28 +28,32 @@ These two concepts allow the library to be extended by installing any compatible NeMo Agent Toolkit utilizes the this plugin system for all first party components. This allows the library to be modular and extendable by default. Plugins from external libraries are treated exactly the same as first party plugins. +External plugin packages should import public plugin-authoring APIs from `nat.plugin_api`. This module is the stable +surface for decorators, function configuration bases, function groups, and common plugin helpers. See the +[Public Plugin API](./plugin-api.md) documentation for the compatibility contract. + ## Supported Plugin Types NeMo Agent Toolkit currently supports the following plugin types: - **CLI Commands**: CLI commands extend the `nat` command-line interface with plugin-specific commands. For example, the MCP and A2A plugins provide their own CLI commands for client operations and server management. To register a CLI command, add an entry point in the `nat.cli` group. -- **Dataset Loaders**: [Dataset loaders](../improve-workflows/evaluate.md#using-datasets) define how evaluation datasets are loaded and parsed. Built-in dataset loaders support `json`, `jsonl`, `csv`, `xls`, `parquet`, and `custom` formats. You can add support for additional dataset formats by creating a custom dataset loader plugin. To register a dataset loader, you can use the {py:deco}`nat.cli.register_workflow.register_dataset_loader` decorator. See the [Custom Dataset Loader](./custom-components/custom-dataset-loader.md) documentation for a step-by-step guide. -- **Embedder Clients**: [Embedder](../build-workflows/embedders.md) Clients are implementations of embedder providers, which are specific to a [LLM](../build-workflows/llms/index.md) framework. For example, when using the OpenAI embedder provider with the LangChain/LangGraph framework, the LangChain/LangGraph OpenAI embedder client needs to be registered. To register an embedder client, you can use the {py:deco}`nat.cli.register_workflow.register_embedder_client` decorator. -- **Embedder Providers**: Embedder Providers are services that provide a way to embed text. For example, OpenAI and NVIDIA NIMs are embedder providers. To register an embedder provider, you can use the {py:deco}`nat.cli.register_workflow.register_embedder_provider` decorator. -- **Evaluators**: [Evaluators](../improve-workflows/evaluate.md) are used by the evaluation framework to evaluate the performance of NeMo Agent Toolkit workflows. To register an evaluator, you can use the {py:deco}`nat.cli.register_workflow.register_evaluator` decorator. -- **Front Ends**: Front ends are the mechanism by which NeMo Agent Toolkit workflows are executed. Examples of front ends include a FastAPI server or a CLI. To register a front end, you can use the {py:deco}`nat.cli.register_workflow.register_front_end` decorator. -- **Functions**: [Functions](../build-workflows/functions-and-function-groups/functions.md) are one of the core building blocks of NeMo Agent Toolkit. They are used to define the tools and agents that can be used in a workflow. To register a function, you can use the {py:deco}`nat.cli.register_workflow.register_function` decorator. -- **LLM Clients**: LLM Clients are implementations of LLM providers that are specific to a LLM framework. For example, when using the NVIDIA NIMs LLM provider with the LangChain/LangGraph framework, the NVIDIA LangChain/LangGraph LLM client needs to be registered. To register an LLM client, you can use the {py:deco}`nat.cli.register_llm_client` decorator. -- **LLM Providers**: An LLM provider is a service that provides a way to interact with an LLM. For example, OpenAI and NVIDIA NIMs are LLM providers. To register an LLM provider, you can use the {py:deco}`nat.cli.register_workflow.register_llm_provider` decorator. -- **Logging Methods**: Logging methods control the destination and format of log messages. To register a logging method, you can use the {py:deco}`nat.cli.register_workflow.register_logging_method` decorator. -- **Memory**: [Memory](../build-workflows/memory.md) plugins are used to store and retrieve information from a database to be used by an LLM. Examples of memory plugins include Zep, Mem0 or MemMachine. To register a memory plugin, you can use the {py:deco}`nat.cli.register_workflow.register_memory` decorator. -- **Registry Handlers**: Registry handlers are used to register custom agent registries with NeMo Agent Toolkit. An agent registry is a collection of tools, agents, and workflows that can be used in a workflow. To register a registry handler, you can use the {py:deco}`nat.cli.register_workflow.register_registry_handler` decorator. -- **Retriever Clients**: [Retriever](../build-workflows/retrievers.md) clients are implementations of retriever providers, which are specific to a LLM framework. For example, when using the Milvus retriever provider with the LangChain/LangGraph framework, the LangChain/LangGraph Milvus retriever client needs to be registered. To register a retriever client, you can use the {py:deco}`nat.cli.register_workflow.register_retriever_client` decorator. -- **Retriever Providers**: Retriever providers are services that provide a way to retrieve information from a database. Examples of retriever providers include Chroma and Milvus. To register a retriever provider, you can use the {py:deco}`nat.cli.register_workflow.register_retriever_provider` decorator. -- **Telemetry Exporters**: [Telemetry exporters](../run-workflows/observe/observe.md) send telemetry data to a telemetry service. To register a telemetry exporter, you can use the {py:deco}`nat.cli.register_workflow.register_telemetry_exporter` decorator. -- **Tool Wrappers**: Tool wrappers are used to wrap functions in a way that is specific to a LLM framework. For example, when using the LangChain/LangGraph framework, NeMo Agent Toolkit functions need to be wrapped in `BaseTool` class to be compatible with LangChain/LangGraph. To register a tool wrapper, you can use the {py:deco}`nat.cli.register_workflow.register_tool_wrapper` decorator. -- **API Authentication Providers**: [API authentication providers](../components/auth/api-authentication.md) are services that provide a way to authenticate requests to an API provider. Examples of authentication providers include OAuth 2.0 Authorization Code Grant and API Key. To register an API authentication provider, you can use the {py:deco}`nat.cli.register_workflow.register_auth_provider` decorator. +- **Dataset Loaders**: [Dataset loaders](../improve-workflows/evaluate.md#using-datasets) define how evaluation datasets are loaded and parsed. Built-in dataset loaders support `json`, `jsonl`, `csv`, `xls`, `parquet`, and `custom` formats. You can add support for additional dataset formats by creating a custom dataset loader plugin. To register a dataset loader, you can use the {py:deco}`nat.plugin_api.register_dataset_loader` decorator. See the [Custom Dataset Loader](./custom-components/custom-dataset-loader.md) documentation for a step-by-step guide. +- **Embedder Clients**: [Embedder](../build-workflows/embedders.md) Clients are implementations of embedder providers, which are specific to a [LLM](../build-workflows/llms/index.md) framework. For example, when using the OpenAI embedder provider with the LangChain/LangGraph framework, the LangChain/LangGraph OpenAI embedder client needs to be registered. To register an embedder client, you can use the {py:deco}`nat.plugin_api.register_embedder_client` decorator. +- **Embedder Providers**: Embedder Providers are services that provide a way to embed text. For example, OpenAI and NVIDIA NIMs are embedder providers. To register an embedder provider, you can use the {py:deco}`nat.plugin_api.register_embedder_provider` decorator. +- **Evaluators**: [Evaluators](../improve-workflows/evaluate.md) are used by the evaluation framework to evaluate the performance of NeMo Agent Toolkit workflows. To register an evaluator, you can use the {py:deco}`nat.plugin_api.register_evaluator` decorator. +- **Front Ends**: Front ends are the mechanism by which NeMo Agent Toolkit workflows are executed. Examples of front ends include a FastAPI server or a CLI. To register a front end, you can use the {py:deco}`nat.plugin_api.register_front_end` decorator. +- **Functions**: [Functions](../build-workflows/functions-and-function-groups/functions.md) are one of the core building blocks of NeMo Agent Toolkit. They are used to define the tools and agents that can be used in a workflow. To register a function, you can use the {py:deco}`nat.plugin_api.register_function` decorator. +- **LLM Clients**: LLM Clients are implementations of LLM providers that are specific to a LLM framework. For example, when using the NVIDIA NIMs LLM provider with the LangChain/LangGraph framework, the NVIDIA LangChain/LangGraph LLM client needs to be registered. To register an LLM client, you can use the {py:deco}`nat.plugin_api.register_llm_client` decorator. +- **LLM Providers**: An LLM provider is a service that provides a way to interact with an LLM. For example, OpenAI and NVIDIA NIMs are LLM providers. To register an LLM provider, you can use the {py:deco}`nat.plugin_api.register_llm_provider` decorator. +- **Logging Methods**: Logging methods control the destination and format of log messages. To register a logging method, you can use the {py:deco}`nat.plugin_api.register_logging_method` decorator. +- **Memory**: [Memory](../build-workflows/memory.md) plugins are used to store and retrieve information from a database to be used by an LLM. Examples of memory plugins include Zep, Mem0 or MemMachine. To register a memory plugin, you can use the {py:deco}`nat.plugin_api.register_memory` decorator. +- **Registry Handlers**: Registry handlers are used to register custom agent registries with NeMo Agent Toolkit. An agent registry is a collection of tools, agents, and workflows that can be used in a workflow. To register a registry handler, you can use the {py:deco}`nat.plugin_api.register_registry_handler` decorator. +- **Retriever Clients**: [Retriever](../build-workflows/retrievers.md) clients are implementations of retriever providers, which are specific to a LLM framework. For example, when using the Milvus retriever provider with the LangChain/LangGraph framework, the LangChain/LangGraph Milvus retriever client needs to be registered. To register a retriever client, you can use the {py:deco}`nat.plugin_api.register_retriever_client` decorator. +- **Retriever Providers**: Retriever providers are services that provide a way to retrieve information from a database. Examples of retriever providers include Chroma and Milvus. To register a retriever provider, you can use the {py:deco}`nat.plugin_api.register_retriever_provider` decorator. +- **Telemetry Exporters**: [Telemetry exporters](../run-workflows/observe/observe.md) send telemetry data to a telemetry service. To register a telemetry exporter, you can use the {py:deco}`nat.plugin_api.register_telemetry_exporter` decorator. +- **Tool Wrappers**: Tool wrappers are used to wrap functions in a way that is specific to a LLM framework. For example, when using the LangChain/LangGraph framework, NeMo Agent Toolkit functions need to be wrapped in `BaseTool` class to be compatible with LangChain/LangGraph. To register a tool wrapper, you can use the {py:deco}`nat.plugin_api.register_tool_wrapper` decorator. +- **API Authentication Providers**: [API authentication providers](../components/auth/api-authentication.md) are services that provide a way to authenticate requests to an API provider. Examples of authentication providers include OAuth 2.0 Authorization Code Grant and API Key. To register an API authentication provider, you can use the {py:deco}`nat.plugin_api.register_auth_provider` decorator. ## Anatomy of a Plugin diff --git a/docs/source/extend/testing/add-unit-tests-for-tools.md b/docs/source/extend/testing/add-unit-tests-for-tools.md index f1f0f67ad0..1790a41a6a 100644 --- a/docs/source/extend/testing/add-unit-tests-for-tools.md +++ b/docs/source/extend/testing/add-unit-tests-for-tools.md @@ -91,6 +91,37 @@ async def test_tool_error_handling(): ## Advanced Usage +### Testing Function Groups + +Use `test_function_group` to validate the tools exposed by a function group, and `test_function_group_tool` to invoke one +tool from the group without creating a full workflow: + +```python +from nat.test import ToolTestRunner +from my_search_provider.register import SearchGroupConfig + + +async def test_search_group_exposes_tools(): + runner = ToolTestRunner() + + await runner.test_function_group( + config_type=SearchGroupConfig, + expected_functions=["search__query", "search__extract"] + ) + + +async def test_search_group_query_tool(): + runner = ToolTestRunner() + + result = await runner.test_function_group_tool( + config_type=SearchGroupConfig, + function_name="query", + input_kwargs={"query": "latest CUDA release"}, + ) + + assert "results" in result +``` + ### Testing Tools with Dependencies For tools that depend on [LLMs](../../build-workflows/llms/index.md), [memory](../../build-workflows/memory.md), [retrievers](../../build-workflows/retrievers.md), or other components, use the mocked dependencies context: diff --git a/docs/source/get-started/tutorials/add-tools-to-a-workflow.md b/docs/source/get-started/tutorials/add-tools-to-a-workflow.md index 3f1c477ed4..6c05573954 100644 --- a/docs/source/get-started/tutorials/add-tools-to-a-workflow.md +++ b/docs/source/get-started/tutorials/add-tools-to-a-workflow.md @@ -108,12 +108,14 @@ Workflow Result: ['To trace only specific parts of a LangChain application, you can either manually pass in a LangChainTracer instance as a callback or use the tracing_v2_enabled context manager. Additionally, you can configure a LangChainTracer instance to trace a specific invocation.'] ``` -## Alternate Method Using a Web Search Tool -Adding individual web pages to a workflow can be cumbersome, especially when dealing with multiple web pages. An alternative method is to use a web search tool. NeMo Agent Toolkit provides two web search tools: `tavily_internet_search` which utilizes the [Tavily Search API](https://tavily.com/), and `exa_internet_search` which utilizes the [Exa Search API](https://exa.ai/). +## Alternate Method Using a LangChain Web Search Tool +Adding individual web pages to a workflow can be cumbersome, especially when dealing with multiple web pages. An alternative method is to use a web search tool. The `nvidia-nat[langchain]` package provides two LangChain-backed web search tools: `tavily_internet_search` which uses the [Tavily Search API](https://tavily.com/), and `exa_internet_search` which uses the [Exa Search API](https://exa.ai/). + +These tools are useful when your workflow already uses the LangChain/LangGraph integration. Framework-agnostic provider packages should expose their tools through the public `nat.plugin_api` function or function group APIs instead. ### Using Tavily Search -The `tavily_internet_search` tool is part of the `nvidia-nat[langchain]` package, to install the package run: +The `tavily_internet_search` tool is part of the `nvidia-nat[langchain]` package and wraps the LangChain Tavily integration. To install the package run: ```bash # local package install from source uv pip install -e ".[langchain]" @@ -156,7 +158,7 @@ Workflow Result: ### Using Exa Search -The `exa_internet_search` tool is also part of the `nvidia-nat[langchain]` package. If you haven't already installed it: +The `exa_internet_search` tool is also part of the `nvidia-nat[langchain]` package and wraps the LangChain Exa integration. If you haven't already installed it: ```bash # local package install from source uv pip install -e ".[langchain]" diff --git a/docs/source/improve-workflows/evaluate.md b/docs/source/improve-workflows/evaluate.md index d695d813cf..c432e4c912 100644 --- a/docs/source/improve-workflows/evaluate.md +++ b/docs/source/improve-workflows/evaluate.md @@ -507,7 +507,7 @@ Callbacks are registered via the `@register_eval_callback(config_type=...)` deco For example, a provider registers its callback by decorating a factory function: ```python -from nat.cli.register_workflow import register_eval_callback +from nat.plugin_api import register_eval_callback @register_eval_callback(config_type=MyTelemetryExporter) def _build_my_eval_callback(config, **kwargs): diff --git a/docs/source/improve-workflows/optimizer.md b/docs/source/improve-workflows/optimizer.md index d6666a2b68..63cb2e9ace 100644 --- a/docs/source/improve-workflows/optimizer.md +++ b/docs/source/improve-workflows/optimizer.md @@ -196,7 +196,7 @@ Here's how you can define optimizable fields in your workflow's data models: ```python from pydantic import BaseModel -from nat.data_models.function import FunctionBaseConfig +from nat.plugin_api import FunctionBaseConfig from nat.data_models.optimizable import OptimizableField, SearchSpace, OptimizableMixin class SomeImageAgentConfig(FunctionBaseConfig, OptimizableMixin, name="some_image_agent_config"): @@ -700,7 +700,7 @@ Callbacks are registered via the `@register_optimizer_callback(config_type=...)` For example, a provider registers its callback by decorating a factory function: ```python -from nat.cli.register_workflow import register_optimizer_callback +from nat.plugin_api import register_optimizer_callback @register_optimizer_callback(config_type=MyTelemetryExporter) def _build_my_optimizer_callback(config, *, dataset_name=None, **kwargs): diff --git a/docs/source/improve-workflows/test-time-compute.md b/docs/source/improve-workflows/test-time-compute.md index 2c343d5bf1..068dd0947e 100644 --- a/docs/source/improve-workflows/test-time-compute.md +++ b/docs/source/improve-workflows/test-time-compute.md @@ -141,7 +141,7 @@ Follow the steps below to create and register a new strategy. 3. Register the strategy. ```python - from nat.cli.register_workflow import register_ttc_strategy + from nat.plugin_api import register_ttc_strategy @register_ttc_strategy(config_type=MyStrategyConfig) async def register_my_strategy(cfg: MyStrategyConfig, builder: Builder): diff --git a/docs/source/index.md b/docs/source/index.md index 61b07ef0cd..3f7e005c49 100644 --- a/docs/source/index.md +++ b/docs/source/index.md @@ -231,6 +231,7 @@ Sharing Components <./components/sharing-components.md> :caption: Extend Plugins <./extend/plugins.md> +Plugin API <./extend/plugin-api.md> Custom Components <./extend/custom-components/index.md> ./extend/testing/index.md ``` diff --git a/packages/nvidia_nat_core/src/nat/cli/commands/workflow/templates/workflow.py.j2 b/packages/nvidia_nat_core/src/nat/cli/commands/workflow/templates/workflow.py.j2 index 1d781432f8..3147851a30 100644 --- a/packages/nvidia_nat_core/src/nat/cli/commands/workflow/templates/workflow.py.j2 +++ b/packages/nvidia_nat_core/src/nat/cli/commands/workflow/templates/workflow.py.j2 @@ -2,11 +2,11 @@ import logging from pydantic import Field -from nat.builder.builder import Builder -from nat.builder.framework_enum import LLMFrameworkEnum -from nat.builder.function_info import FunctionInfo -from nat.cli.register_workflow import register_function -from nat.data_models.function import FunctionBaseConfig +from nat.plugin_api import Builder +from nat.plugin_api import FunctionBaseConfig +from nat.plugin_api import FunctionInfo +from nat.plugin_api import LLMFrameworkEnum +from nat.plugin_api import register_function logger = logging.getLogger(__name__) diff --git a/packages/nvidia_nat_core/tests/nat/tools/test_tool_test_runner.py b/packages/nvidia_nat_core/tests/nat/tools/test_tool_test_runner.py index 51a8bcd2df..cebfd92d72 100644 --- a/packages/nvidia_nat_core/tests/nat/tools/test_tool_test_runner.py +++ b/packages/nvidia_nat_core/tests/nat/tools/test_tool_test_runner.py @@ -14,9 +14,12 @@ # limitations under the License. from nat.builder.builder import Builder +from nat.builder.function import FunctionGroup from nat.builder.function_info import FunctionInfo from nat.cli.register_workflow import register_function +from nat.cli.register_workflow import register_function_group from nat.data_models.function import FunctionBaseConfig +from nat.data_models.function import FunctionGroupBaseConfig from nat.test.tool_test_runner import ToolTestRunner @@ -40,6 +43,21 @@ async def _calc_fn(input_data: str) -> str: yield FunctionInfo.from_fn(_calc_fn, description=_calc_fn.__doc__) +class SimpleCalculatorGroupConfig(FunctionGroupBaseConfig, name="test_simple_calculator_group"): + pass + + +@register_function_group(config_type=SimpleCalculatorGroupConfig) +async def simple_calculator_group(config: SimpleCalculatorGroupConfig, _builder: Builder): + + async def _add(lhs: int, rhs: int) -> int: + return lhs + rhs + + group = FunctionGroup(config=config, instance_name="calculator") + group.add_function("add", _add, description="Add two numbers.") + yield group + + # This test is to ensure ToolTestRunner is working correctly, and also a demonstration of how to test tools # in complete isolation without requiring spinning up entire workflows, agents, and external services. async def test_simple_calculator_tool(): @@ -125,3 +143,25 @@ async def test_tool_with_mocked_training_components(): trajectory_builder = await mock_builder.get_trajectory_builder("my_trajectory_builder") assert trajectory_builder is not None assert await trajectory_builder.build() == {"trajectories": []} + + +async def test_function_group_exposes_expected_tools(): + runner = ToolTestRunner() + + function_names = await runner.test_function_group(config_type=SimpleCalculatorGroupConfig, + expected_functions=["calculator__add"]) + + assert function_names == {"calculator__add"} + + +async def test_function_group_tool_with_kwargs(): + runner = ToolTestRunner() + + result = await runner.test_function_group_tool(config_type=SimpleCalculatorGroupConfig, + function_name="add", + input_kwargs={ + "lhs": 2, "rhs": 3 + }, + expected_output=5) + + assert result == 5 diff --git a/packages/nvidia_nat_langchain/src/nat/plugins/langchain/tools/exa_internet_search.py b/packages/nvidia_nat_langchain/src/nat/plugins/langchain/tools/exa_internet_search.py index 68363641c8..1d3b58d09e 100644 --- a/packages/nvidia_nat_langchain/src/nat/plugins/langchain/tools/exa_internet_search.py +++ b/packages/nvidia_nat_langchain/src/nat/plugins/langchain/tools/exa_internet_search.py @@ -34,7 +34,7 @@ # Internet Search tool class ExaInternetSearchToolConfig(FunctionBaseConfig, name="exa_internet_search"): """ - Tool that retrieves relevant contexts from web search (using Exa) for the given question. + LangChain-backed tool that retrieves relevant contexts from Exa web search for the given question. Requires an EXA_API_KEY. """ max_results: int = Field(default=5, ge=1, description="Maximum number of search results to return.") @@ -118,7 +118,7 @@ async def _exa_internet_search(question: str) -> str: await asyncio.sleep(2**attempt) return f"Web search failed after {tool_config.max_retries} attempts for: {question}" - # Create a Generic NAT tool that can be used with any supported LLM framework + # Create a NAT function backed by the LangChain Exa tool. yield FunctionInfo.from_fn( _exa_internet_search, description=_exa_internet_search.__doc__, diff --git a/packages/nvidia_nat_langchain/src/nat/plugins/langchain/tools/tavily_internet_search.py b/packages/nvidia_nat_langchain/src/nat/plugins/langchain/tools/tavily_internet_search.py index ccc63a45a1..d6a0aa2dd2 100644 --- a/packages/nvidia_nat_langchain/src/nat/plugins/langchain/tools/tavily_internet_search.py +++ b/packages/nvidia_nat_langchain/src/nat/plugins/langchain/tools/tavily_internet_search.py @@ -28,7 +28,7 @@ # Internet Search tool class TavilyInternetSearchToolConfig(FunctionBaseConfig, name="tavily_internet_search"): """ - Tool that retrieves relevant contexts from web search (using Tavily) for the given question. + LangChain-backed tool that retrieves relevant contexts from Tavily web search for the given question. Requires a TAVILY_API_KEY. """ max_results: int = 3 @@ -91,7 +91,7 @@ async def _tavily_internet_search(question: str) -> str: return f"Web search failed after {tool_config.max_retries} attempts for: {question}" await asyncio.sleep(2**attempt) - # Create a Generic NAT tool that can be used with any supported LLM framework + # Create a NAT function backed by the LangChain Tavily tool. yield FunctionInfo.from_fn( _tavily_internet_search, description=_tavily_internet_search.__doc__, diff --git a/packages/nvidia_nat_test/src/nat/test/tool_test_runner.py b/packages/nvidia_nat_test/src/nat/test/tool_test_runner.py index b44b17af10..24171b52b3 100644 --- a/packages/nvidia_nat_test/src/nat/test/tool_test_runner.py +++ b/packages/nvidia_nat_test/src/nat/test/tool_test_runner.py @@ -584,6 +584,120 @@ async def test_tool_with_builder( return result + def _resolve_function_group_tool(self, functions: dict[str, Function], function_name: str) -> Function: + if function_name in functions: + return functions[function_name] + + suffix = f"{FunctionGroup.SEPARATOR}{function_name}" + matches = [fn for name, fn in functions.items() if name.endswith(suffix)] + if len(matches) == 1: + return matches[0] + if len(matches) > 1: + raise ValueError( + f"Function name '{function_name}' is ambiguous. Use the fully qualified function group name.") + + raise ValueError(f"Function group tool '{function_name}' not found. Available tools: {sorted(functions)}") + + async def test_function_group( + self, + config_type: type[FunctionGroupBaseConfig], + config_params: dict[str, typing.Any] | None = None, + expected_functions: Sequence[str] | None = None, + builder: MockBuilder | None = None, + ) -> set[str]: + """ + Test a function group in isolation and optionally assert the exposed function names. + + Args: + config_type: The function group configuration class. + config_params: Parameters to pass to the config constructor. + expected_functions: Fully qualified function names expected from the group. + builder: Optional pre-configured MockBuilder with mocked dependencies. + + Returns: + The function group's accessible function names. + """ + config_params = config_params or {} + config = config_type(**config_params) + + registry = GlobalTypeRegistry.get() + try: + group_registration = registry.get_function_group(config_type) + except KeyError: + raise ValueError( + f"Function group {config_type} is not registered. Make sure it's imported and registered with " + "@register_function_group.") + + async with group_registration.build_fn(config, builder or MockBuilder()) as group_result: + if not isinstance(group_result, FunctionGroup): + raise ValueError(f"Unexpected function group result type: {type(group_result)}") + + functions = await group_result.get_accessible_functions() + function_names = set(functions) + if expected_functions is not None: + assert function_names == set(expected_functions), ( + f"Expected function group tools {sorted(expected_functions)}, got {sorted(function_names)}") + + return function_names + + async def test_function_group_tool( + self, + config_type: type[FunctionGroupBaseConfig], + function_name: str, + config_params: dict[str, typing.Any] | None = None, + input_data: typing.Any = None, + input_kwargs: dict[str, typing.Any] | None = None, + expected_output: typing.Any = None, + builder: MockBuilder | None = None, + ) -> typing.Any: + """ + Test one tool exposed by a function group. + + Args: + config_type: The function group configuration class. + function_name: Fully qualified function name, or the unqualified name when it is unique in the group. + config_params: Parameters to pass to the config constructor. + input_data: Positional input to pass to the function. + input_kwargs: Keyword input to pass to the function. + expected_output: Expected output for assertion. + builder: Optional pre-configured MockBuilder with mocked dependencies. + + Returns: + The function output. + """ + if input_data is not None and input_kwargs is not None: + raise ValueError("Use either input_data or input_kwargs, not both.") + + config_params = config_params or {} + config = config_type(**config_params) + + registry = GlobalTypeRegistry.get() + try: + group_registration = registry.get_function_group(config_type) + except KeyError: + raise ValueError( + f"Function group {config_type} is not registered. Make sure it's imported and registered with " + "@register_function_group.") + + async with group_registration.build_fn(config, builder or MockBuilder()) as group_result: + if not isinstance(group_result, FunctionGroup): + raise ValueError(f"Unexpected function group result type: {type(group_result)}") + + functions = await group_result.get_accessible_functions() + tool = self._resolve_function_group_tool(functions, function_name) + + if input_kwargs is not None: + result = await tool.acall_invoke(**input_kwargs) + elif input_data is not None: + result = await tool.acall_invoke(input_data) + else: + result = await tool.acall_invoke() + + if expected_output is not None: + assert result == expected_output, f"Expected {expected_output}, got {result}" + + return result + @asynccontextmanager async def with_mocked_dependencies(): From 636376a6e363f32876d2057333252aefa2cec50d Mon Sep 17 00:00:00 2001 From: Bryan Bednarski Date: Mon, 18 May 2026 13:49:48 -0700 Subject: [PATCH 02/13] add new plugin_api file and associated changes Signed-off-by: Bryan Bednarski --- .../function-groups.md | 2 +- docs/source/build-workflows/mcp-client.md | 2 +- examples/A2A/math_assistant_a2a/README.md | 6 +- .../configs/config-client.yml | 6 +- examples/dynamo_integration/README.md | 2 +- .../react_benchmark_agent/README.md | 8 +- .../react_benchmark_agent/DEVELOPER_NOTES.md | 3 +- .../evaluators/tsq_evaluator.py | 4 +- .../tests/test_tsq_formula.py | 8 +- .../nat/builder/per_user_workflow_builder.py | 2 + .../src/nat/builder/workflow_builder.py | 2 + .../nvidia_nat_core/src/nat/plugin_api.py | 132 ++++++++++++++++++ .../tests/nat/test_plugin_api.py | 80 +++++++++++ .../src/nat/plugins/eval/runtime/builder.py | 2 + 14 files changed, 238 insertions(+), 21 deletions(-) create mode 100644 packages/nvidia_nat_core/src/nat/plugin_api.py create mode 100644 packages/nvidia_nat_core/tests/nat/test_plugin_api.py diff --git a/docs/source/build-workflows/functions-and-function-groups/function-groups.md b/docs/source/build-workflows/functions-and-function-groups/function-groups.md index 6886b0bd86..ad289e7b0f 100644 --- a/docs/source/build-workflows/functions-and-function-groups/function-groups.md +++ b/docs/source/build-workflows/functions-and-function-groups/function-groups.md @@ -365,7 +365,7 @@ This will return a list of all accessible functions in the function group that a Functions inside a group are automatically namespaced by the group instance name. This creates a clear hierarchy and prevents naming conflicts. -To maintain compatibility with third-party libraries, the namespace separator switched from `.` (period) to `__` (double underscore). +Function groups expose functions with fully qualified names that use `__` (double underscore) between the group instance name and the function name. This keeps names compatible with third-party frameworks that reject or reinterpret `.` in tool names. **Pattern**: `instance_name__function_name` diff --git a/docs/source/build-workflows/mcp-client.md b/docs/source/build-workflows/mcp-client.md index ce31a08ad9..5bb40c6bc5 100644 --- a/docs/source/build-workflows/mcp-client.md +++ b/docs/source/build-workflows/mcp-client.md @@ -77,7 +77,7 @@ Example: workflows: _type: react_agent tool_names: - - mcp_tools.tool_a + - mcp_tools__tool_a ``` An additional case to note is when a function group is served by an MCP server, the tools within the function group must still be accessed by their full name. This is the same as the prior case, but there is an important difference. Consider the following example: diff --git a/examples/A2A/math_assistant_a2a/README.md b/examples/A2A/math_assistant_a2a/README.md index 5fa8874a3e..4801be24a2 100644 --- a/examples/A2A/math_assistant_a2a/README.md +++ b/examples/A2A/math_assistant_a2a/README.md @@ -202,9 +202,9 @@ workflow: _type: per_user_react_agent # Per-user ReAct agent tool_names: - calculator_a2a # Per-user A2A client - - mcp_time.get_current_time_mcp - - logic_evaluator.if_then_else - - logic_evaluator.evaluate_condition + - mcp_time__get_current_time_mcp + - logic_evaluator__if_then_else + - logic_evaluator__evaluate_condition llm_name: nim_llm ``` diff --git a/examples/A2A/math_assistant_a2a_protected/configs/config-client.yml b/examples/A2A/math_assistant_a2a_protected/configs/config-client.yml index 8c67043c7e..ee9cdb23a5 100644 --- a/examples/A2A/math_assistant_a2a_protected/configs/config-client.yml +++ b/examples/A2A/math_assistant_a2a_protected/configs/config-client.yml @@ -70,9 +70,9 @@ workflow: _type: per_user_react_agent tool_names: - calculator_a2a # A2A calculator functions (per-user with OAuth2) - - mcp_time.get_current_time_mcp # Local time function - - logic_evaluator.if_then_else # Conditional logic - - logic_evaluator.evaluate_condition # Comparison operations + - mcp_time__get_current_time_mcp # Local time function + - logic_evaluator__if_then_else # Conditional logic + - logic_evaluator__evaluate_condition # Comparison operations llm_name: nim_llm verbose: true retry_parsing_errors: true diff --git a/examples/dynamo_integration/README.md b/examples/dynamo_integration/README.md index 8d0a4b8a3f..8f0b694d65 100644 --- a/examples/dynamo_integration/README.md +++ b/examples/dynamo_integration/README.md @@ -301,7 +301,7 @@ external/dynamo/ # Dynamo backend (separate location) workflow: _type: react_agent llm_name: dynamo_llm - tool_names: [banking_tools.get_account_balance, ...] + tool_names: [banking_tools__get_account_balance, ...] ``` ### With Self-Evaluation Loop diff --git a/examples/dynamo_integration/react_benchmark_agent/README.md b/examples/dynamo_integration/react_benchmark_agent/README.md index 81fb2c549d..cde1050a38 100644 --- a/examples/dynamo_integration/react_benchmark_agent/README.md +++ b/examples/dynamo_integration/react_benchmark_agent/README.md @@ -332,9 +332,9 @@ workflow: _type: react_agent llm_name: dynamo_llm tool_names: [ - banking_tools.get_account_balance, - banking_tools.transfer_funds, - # ... all tools with banking_tools. prefix + banking_tools__get_account_balance, + banking_tools__transfer_funds, + # ... all tools with banking_tools__ prefix ] verbose: true max_tool_calls: 25 @@ -469,7 +469,7 @@ functions: react_workflow: _type: react_agent llm_name: dynamo_llm - tool_names: [banking_tools.get_account_balance, ...] + tool_names: [banking_tools__get_account_balance, ...] verbose: true max_tool_calls: 25 diff --git a/examples/dynamo_integration/react_benchmark_agent/src/react_benchmark_agent/DEVELOPER_NOTES.md b/examples/dynamo_integration/react_benchmark_agent/src/react_benchmark_agent/DEVELOPER_NOTES.md index 803106caa8..a793508127 100644 --- a/examples/dynamo_integration/react_benchmark_agent/src/react_benchmark_agent/DEVELOPER_NOTES.md +++ b/examples/dynamo_integration/react_benchmark_agent/src/react_benchmark_agent/DEVELOPER_NOTES.md @@ -203,7 +203,7 @@ class ReactBenchmarkAgentFunctionConfig(FunctionBaseConfig, name="react_benchmar - Loads tool schemas from `data/raw/banking/tools.json` - Creates stub functions for each tool via `create_tool_stub_function()` -- Registers them as a function group accessible by `banking_tools.` +- Registers them as a function group accessible by `banking_tools__` **`tool_intent_stubs.py`** (lines 79-136) @@ -511,4 +511,3 @@ def calculate_tool_accuracy(actual, expected): **File:** `evaluators/action_completion_evaluator.py` The AC evaluator measures whether the agent addressed all user goals. - diff --git a/examples/dynamo_integration/react_benchmark_agent/src/react_benchmark_agent/evaluators/tsq_evaluator.py b/examples/dynamo_integration/react_benchmark_agent/src/react_benchmark_agent/evaluators/tsq_evaluator.py index 44f747a060..8c4431f1c3 100644 --- a/examples/dynamo_integration/react_benchmark_agent/src/react_benchmark_agent/evaluators/tsq_evaluator.py +++ b/examples/dynamo_integration/react_benchmark_agent/src/react_benchmark_agent/evaluators/tsq_evaluator.py @@ -150,7 +150,7 @@ def normalize_tool_name(tool_name: str) -> str: Handles: - Case normalization (lowercase) - Underscore and dash removal - - Module prefix stripping (e.g., 'banking_tools.report_lost_stolen_card' -> 'reportloststolencard') + - Module prefix stripping (e.g., 'banking_tools__report_lost_stolen_card' -> 'reportloststolencard') Args: tool_name: Raw tool name from trajectory or expected list @@ -161,7 +161,7 @@ def normalize_tool_name(tool_name: str) -> str: if not tool_name: return "" - # Strip module prefix (e.g., "banking_tools.report_lost_stolen_card" -> "report_lost_stolen_card") + # Strip module prefix (e.g., "banking_tools__report_lost_stolen_card" -> "report_lost_stolen_card") if FunctionGroup.SEPARATOR in tool_name: _, tool_name = FunctionGroup.decompose(tool_name) diff --git a/examples/dynamo_integration/react_benchmark_agent/tests/test_tsq_formula.py b/examples/dynamo_integration/react_benchmark_agent/tests/test_tsq_formula.py index fec8e241ff..e2e4bd0510 100644 --- a/examples/dynamo_integration/react_benchmark_agent/tests/test_tsq_formula.py +++ b/examples/dynamo_integration/react_benchmark_agent/tests/test_tsq_formula.py @@ -358,7 +358,7 @@ def test_nested_payload_format(self): }, "payload": { "event_type": "TOOL_START", - "name": "banking_tools.report_lost_stolen_card", + "name": "banking_tools__report_lost_stolen_card", "data": { "input": { "input_params": { @@ -371,7 +371,7 @@ def test_nested_payload_format(self): tool_calls = self.extract_tool_calls_from_trajectory(trajectory) assert len(tool_calls) == 1 - assert tool_calls[0]["tool"] == "banking_tools.report_lost_stolen_card" + assert tool_calls[0]["tool"] == "banking_tools__report_lost_stolen_card" def test_flat_legacy_format(self): """Test extraction from flat structure (legacy format).""" @@ -470,7 +470,7 @@ def test_real_profiler_data_structure(self): "event_timestamp": 1764917512.0873613, "span_event_timestamp": None, "framework": None, - "name": "banking_tools.report_lost_stolen_card", + "name": "banking_tools__report_lost_stolen_card", "tags": None, "metadata": {}, "data": { @@ -487,7 +487,7 @@ def test_real_profiler_data_structure(self): tool_calls = self.extract_tool_calls_from_trajectory(trajectory) assert len(tool_calls) == 1 - assert tool_calls[0]["tool"] == "banking_tools.report_lost_stolen_card" + assert tool_calls[0]["tool"] == "banking_tools__report_lost_stolen_card" # Should extract nested input_params assert "card_type" in tool_calls[0]["parameters"] assert isinstance(tool_calls[0]["parameters"], dict) diff --git a/packages/nvidia_nat_core/src/nat/builder/per_user_workflow_builder.py b/packages/nvidia_nat_core/src/nat/builder/per_user_workflow_builder.py index 54da2d8dd7..ef0362ed94 100644 --- a/packages/nvidia_nat_core/src/nat/builder/per_user_workflow_builder.py +++ b/packages/nvidia_nat_core/src/nat/builder/per_user_workflow_builder.py @@ -227,6 +227,8 @@ async def add_function(self, name: str | FunctionRef, config: FunctionBaseConfig def _check_backwards_compatibility_function_name(self, name: str) -> str: if name in self._per_user_functions: return name + # TODO(#1952): In the next breaking release, remove dot separator compatibility and require + # FunctionGroup.SEPARATOR (`__`) for function group tool names. new_name = name.replace(FunctionGroup.LEGACY_SEPARATOR, FunctionGroup.SEPARATOR) if new_name in self._per_user_functions: logger.warning( diff --git a/packages/nvidia_nat_core/src/nat/builder/workflow_builder.py b/packages/nvidia_nat_core/src/nat/builder/workflow_builder.py index 7d10fc6054..7702adb624 100644 --- a/packages/nvidia_nat_core/src/nat/builder/workflow_builder.py +++ b/packages/nvidia_nat_core/src/nat/builder/workflow_builder.py @@ -706,6 +706,8 @@ async def add_function_group(self, name: str | FunctionGroupRef, config: Functio def _check_backwards_compatibility_function_name(self, name: str) -> str: if name in self._functions: return name + # TODO(#1952): In the next breaking release, remove dot separator compatibility and require + # FunctionGroup.SEPARATOR (`__`) for function group tool names. new_name = name.replace(FunctionGroup.LEGACY_SEPARATOR, FunctionGroup.SEPARATOR) if new_name in self._functions: logger.warning( diff --git a/packages/nvidia_nat_core/src/nat/plugin_api.py b/packages/nvidia_nat_core/src/nat/plugin_api.py new file mode 100644 index 0000000000..23c4e2e907 --- /dev/null +++ b/packages/nvidia_nat_core/src/nat/plugin_api.py @@ -0,0 +1,132 @@ +# 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. + +"""Stable public API for NeMo Agent Toolkit plugin authors. + +External plugin packages should prefer importing from this module instead of depending on implementation-oriented +modules such as ``nat.cli.register_workflow`` or ``nat.builder.function`` directly. +""" + +from nat.builder.builder import Builder +from nat.builder.builder import EvalBuilder +from nat.builder.framework_enum import LLMFrameworkEnum +from nat.builder.function import Function +from nat.builder.function import FunctionGroup +from nat.builder.function_info import FunctionInfo +from nat.cli.register_workflow import register_auth_provider +from nat.cli.register_workflow import register_dataset_loader +from nat.cli.register_workflow import register_embedder_client +from nat.cli.register_workflow import register_embedder_provider +from nat.cli.register_workflow import register_eval_callback +from nat.cli.register_workflow import register_evaluator +from nat.cli.register_workflow import register_front_end +from nat.cli.register_workflow import register_function +from nat.cli.register_workflow import register_function_group +from nat.cli.register_workflow import register_llm_client +from nat.cli.register_workflow import register_llm_provider +from nat.cli.register_workflow import register_logging_method +from nat.cli.register_workflow import register_memory +from nat.cli.register_workflow import register_middleware +from nat.cli.register_workflow import register_object_store +from nat.cli.register_workflow import register_optimizer +from nat.cli.register_workflow import register_optimizer_callback +from nat.cli.register_workflow import register_per_user_function +from nat.cli.register_workflow import register_per_user_function_group +from nat.cli.register_workflow import register_registry_handler +from nat.cli.register_workflow import register_retriever_client +from nat.cli.register_workflow import register_retriever_provider +from nat.cli.register_workflow import register_telemetry_exporter +from nat.cli.register_workflow import register_tool_wrapper +from nat.cli.register_workflow import register_trainer +from nat.cli.register_workflow import register_trainer_adapter +from nat.cli.register_workflow import register_trajectory_builder +from nat.cli.register_workflow import register_ttc_strategy +from nat.data_models.common import OptionalSecretStr +from nat.data_models.common import SerializableSecretStr +from nat.data_models.common import get_secret_value +from nat.data_models.common import set_secret_from_env +from nat.data_models.component_ref import AuthenticationRef +from nat.data_models.component_ref import ComponentRef +from nat.data_models.component_ref import EmbedderRef +from nat.data_models.component_ref import FunctionGroupRef +from nat.data_models.component_ref import FunctionRef +from nat.data_models.component_ref import LLMRef +from nat.data_models.component_ref import MemoryRef +from nat.data_models.component_ref import MiddlewareRef +from nat.data_models.component_ref import ObjectStoreRef +from nat.data_models.component_ref import RetrieverRef +from nat.data_models.component_ref import TrainerAdapterRef +from nat.data_models.component_ref import TrainerRef +from nat.data_models.component_ref import TrajectoryBuilderRef +from nat.data_models.component_ref import TTCStrategyRef +from nat.data_models.function import FunctionBaseConfig +from nat.data_models.function import FunctionGroupBaseConfig + +__all__ = [ + "AuthenticationRef", + "Builder", + "ComponentRef", + "EmbedderRef", + "EvalBuilder", + "Function", + "FunctionBaseConfig", + "FunctionGroup", + "FunctionGroupBaseConfig", + "FunctionGroupRef", + "FunctionInfo", + "FunctionRef", + "LLMFrameworkEnum", + "LLMRef", + "MemoryRef", + "MiddlewareRef", + "ObjectStoreRef", + "OptionalSecretStr", + "RetrieverRef", + "SerializableSecretStr", + "TTCStrategyRef", + "TrainerAdapterRef", + "TrainerRef", + "TrajectoryBuilderRef", + "get_secret_value", + "register_auth_provider", + "register_dataset_loader", + "register_embedder_client", + "register_embedder_provider", + "register_eval_callback", + "register_evaluator", + "register_front_end", + "register_function", + "register_function_group", + "register_llm_client", + "register_llm_provider", + "register_logging_method", + "register_memory", + "register_middleware", + "register_object_store", + "register_optimizer", + "register_optimizer_callback", + "register_per_user_function", + "register_per_user_function_group", + "register_registry_handler", + "register_retriever_client", + "register_retriever_provider", + "register_telemetry_exporter", + "register_tool_wrapper", + "register_trainer", + "register_trainer_adapter", + "register_trajectory_builder", + "register_ttc_strategy", + "set_secret_from_env", +] diff --git a/packages/nvidia_nat_core/tests/nat/test_plugin_api.py b/packages/nvidia_nat_core/tests/nat/test_plugin_api.py new file mode 100644 index 0000000000..53ce0d7586 --- /dev/null +++ b/packages/nvidia_nat_core/tests/nat/test_plugin_api.py @@ -0,0 +1,80 @@ +# 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. + +from pathlib import Path + +from nat import plugin_api +from nat.builder.builder import Builder +from nat.builder.function import FunctionGroup +from nat.cli.register_workflow import register_function +from nat.cli.register_workflow import register_function_group +from nat.data_models.function import FunctionBaseConfig +from nat.data_models.function import FunctionGroupBaseConfig + + +def test_plugin_api_exports_core_function_authoring_surface(): + assert plugin_api.Builder is Builder + assert plugin_api.FunctionBaseConfig is FunctionBaseConfig + assert plugin_api.FunctionGroupBaseConfig is FunctionGroupBaseConfig + assert plugin_api.FunctionGroup is FunctionGroup + assert plugin_api.register_function is register_function + assert plugin_api.register_function_group is register_function_group + + +def test_plugin_api_all_exports_are_importable(): + for name in plugin_api.__all__: + assert getattr(plugin_api, name) is not None + + +def test_plugin_authoring_docs_prefer_public_api_imports(): + repo_root = Path(__file__).parents[4] + paths = [ + repo_root / "docs/source/components/sharing-components.md", + repo_root / "docs/source/build-workflows/advanced/middleware.md", + repo_root / "docs/source/build-workflows/functions-and-function-groups/function-groups.md", + repo_root / "docs/source/extend/custom-components", + repo_root / "docs/source/extend/plugins.md", + repo_root / "docs/source/improve-workflows/evaluate.md", + repo_root / "docs/source/improve-workflows/optimizer.md", + repo_root / "docs/source/improve-workflows/test-time-compute.md", + repo_root / "packages/nvidia_nat_core/src/nat/cli/commands/workflow/templates/workflow.py.j2", + ] + denied_patterns = [ + "nat.cli.register_workflow", + "nat.cli.register_llm_client", + "from nat.builder.builder import Builder", + "from nat.builder.builder import EvalBuilder", + "from nat.builder.function import FunctionGroup", + "from nat.builder.function_info import FunctionInfo", + "from nat.data_models.component_ref import", + "from nat.data_models.function import FunctionBaseConfig", + "from nat.data_models.function import FunctionGroupBaseConfig", + ] + + files: list[Path] = [] + for path in paths: + if path.is_dir(): + files.extend(path.rglob("*.md")) + else: + files.append(path) + + violations = [] + for file_path in files: + text = file_path.read_text(encoding="utf-8") + for pattern in denied_patterns: + if pattern in text: + violations.append(f"{file_path.relative_to(repo_root)} contains {pattern!r}") + + assert not violations, "Plugin authoring docs should use nat.plugin_api:\n" + "\n".join(violations) diff --git a/packages/nvidia_nat_eval/src/nat/plugins/eval/runtime/builder.py b/packages/nvidia_nat_eval/src/nat/plugins/eval/runtime/builder.py index 5340b92e4e..cba469d280 100644 --- a/packages/nvidia_nat_eval/src/nat/plugins/eval/runtime/builder.py +++ b/packages/nvidia_nat_eval/src/nat/plugins/eval/runtime/builder.py @@ -138,6 +138,8 @@ async def get_all_tools(self, wrapper_type: LLMFrameworkEnum | str): async def get_tool(fn_name: str): # Maintain backwards compatibility with the old function group name format + # TODO(#1952): In the next breaking release, remove dot separator compatibility and require + # FunctionGroup.SEPARATOR (`__`) for function group tool names. new_fn_name = fn_name.replace(FunctionGroup.LEGACY_SEPARATOR, FunctionGroup.SEPARATOR) if (fn_name not in self._functions) and (new_fn_name in self._functions): logger.warning( From 231dce0089ad4e5fcde4ea81ff51a15691da30de Mon Sep 17 00:00:00 2001 From: Bryan Bednarski Date: Mon, 18 May 2026 19:21:01 -0700 Subject: [PATCH 03/13] expand public facade, tighten contract Signed-off-by: Bryan Bednarski --- .../build-workflows/advanced/middleware.md | 16 +- .../function-groups.md | 2 +- docs/source/build-workflows/retrievers.md | 2 + .../custom-components/adding-a-retriever.md | 8 + .../adding-an-authentication-provider.md | 7 +- .../adding-an-llm-provider.md | 24 ++- .../custom-dataset-loader.md | 4 +- .../custom-components/custom-evaluator.md | 9 +- .../extend/custom-components/finetuning.md | 11 +- .../source/extend/custom-components/memory.md | 41 ++--- .../extend/custom-components/object-store.md | 42 ++--- .../extend/custom-components/optimizer.md | 10 +- .../custom-components/telemetry-exporters.md | 10 +- .../existing-agents/langgraph.md | 2 +- .../nvidia_nat_core/src/nat/plugin_api.py | 80 +++++++++ .../tests/nat/test_plugin_api.py | 159 ++++++++++++++++-- 16 files changed, 335 insertions(+), 92 deletions(-) diff --git a/docs/source/build-workflows/advanced/middleware.md b/docs/source/build-workflows/advanced/middleware.md index b4d7c322ae..000bc09432 100644 --- a/docs/source/build-workflows/advanced/middleware.md +++ b/docs/source/build-workflows/advanced/middleware.md @@ -84,7 +84,7 @@ Create a configuration class inheriting from `DynamicMiddlewareConfig`: ```python from pydantic import Field -from nat.middleware.dynamic.dynamic_middleware_config import DynamicMiddlewareConfig + from nat.plugin_api import DynamicMiddlewareConfig class LoggingMiddlewareConfig(DynamicMiddlewareConfig, name="logging_middleware"): @@ -168,8 +168,8 @@ Create the middleware class inheriting from `DynamicFunctionMiddleware`: ```python import logging -from nat.middleware.dynamic.dynamic_function_middleware import DynamicFunctionMiddleware -from nat.middleware.middleware import InvocationContext +from nat.plugin_api import DynamicFunctionMiddleware +from nat.plugin_api import InvocationContext logger = logging.getLogger(__name__) @@ -429,6 +429,9 @@ async def caching_middleware(config: CachingMiddlewareConfig, builder: Builder): Final middleware can short-circuit execution: ```python +from nat.plugin_api import FunctionMiddleware +from nat.plugin_api import FunctionMiddlewareBaseConfig + class ValidationMiddlewareConfig(FunctionMiddlewareBaseConfig, name="validation"): strict_mode: bool = Field(default=True) @@ -603,7 +606,8 @@ Test middleware in isolation: ```python import pytest from unittest.mock import MagicMock -from nat.middleware.middleware import FunctionMiddlewareContext, InvocationContext +from nat.plugin_api import FunctionMiddlewareContext +from nat.plugin_api import InvocationContext @pytest.mark.asyncio @@ -823,8 +827,8 @@ Solution: Ensure the register module is imported. NeMo Agent Toolkit automatical ## API Reference -- {py:class}`~nat.middleware.function_middleware.FunctionMiddleware`: Base class -- {py:class}`~nat.middleware.function_middleware.FunctionMiddlewareContext`: Context info +- {py:class}`~nat.plugin_api.FunctionMiddleware`: Base class +- {py:class}`~nat.plugin_api.FunctionMiddlewareContext`: Context info - {py:class}`~nat.middleware.function_middleware.FunctionMiddlewareChain`: Chain management - {py:class}`~nat.middleware.cache.cache_middleware_config.CacheMiddlewareConfig`: Cache configuration - {py:class}`~nat.middleware.cache.cache_middleware.CacheMiddleware`: Cache implementation diff --git a/docs/source/build-workflows/functions-and-function-groups/function-groups.md b/docs/source/build-workflows/functions-and-function-groups/function-groups.md index ad289e7b0f..4b7182b42c 100644 --- a/docs/source/build-workflows/functions-and-function-groups/function-groups.md +++ b/docs/source/build-workflows/functions-and-function-groups/function-groups.md @@ -603,7 +603,7 @@ To wrap all accessible functions in a group for a specific agent framework: ```python from nat.plugin_api import FunctionGroupRef -from nat.builder.framework_enum import LLMFrameworkEnum +from nat.plugin_api import LLMFrameworkEnum async with WorkflowBuilder() as builder: await builder.add_function_group("math", MathGroupConfig(include=["add", "multiply"])) diff --git a/docs/source/build-workflows/retrievers.md b/docs/source/build-workflows/retrievers.md index c83d33174a..43c3aff072 100644 --- a/docs/source/build-workflows/retrievers.md +++ b/docs/source/build-workflows/retrievers.md @@ -81,6 +81,8 @@ Retrievers are configured similarly to other NeMo Agent Toolkit components, such Below is an example config object for the NeMo Retriever: ```python +from nat.plugin_api import RetrieverBaseConfig + class NemoRetrieverConfig(RetrieverBaseConfig, name="nemo_retriever"): """ Configuration for a Retriever which pulls data from a Nemo Retriever service. diff --git a/docs/source/extend/custom-components/adding-a-retriever.md b/docs/source/extend/custom-components/adding-a-retriever.md index 94655ff0b6..02c2d57ab2 100644 --- a/docs/source/extend/custom-components/adding-a-retriever.md +++ b/docs/source/extend/custom-components/adding-a-retriever.md @@ -41,6 +41,11 @@ class Retriever(ABC): Next, create the config for the provider and register it with NeMo Agent Toolkit: ```python +from nat.plugin_api import Builder +from nat.plugin_api import RetrieverBaseConfig +from nat.plugin_api import RetrieverProviderInfo +from nat.plugin_api import register_retriever_provider + class ExampleRetrieverConfig(RetrieverBaseConfig, name="example_retriever"): """ Configuration for a Retriever provider. The parameters will depend on the particular provider. These are examples. @@ -61,6 +66,9 @@ async def example_retriever(retriever_config: ExampleRetrieverConfig, builder: B Lastly, implement and register the retriever client: ```python +from nat.plugin_api import Builder +from nat.plugin_api import register_retriever_client + @register_retriever_client(config_type=ExampleRetrieverConfig, wrapper_type=None) async def nemo_retriever_client(config: ExampleRetrieverConfig, builder: Builder): from example_plugin.retriever import ExampleRetriever diff --git a/docs/source/extend/custom-components/adding-an-authentication-provider.md b/docs/source/extend/custom-components/adding-an-authentication-provider.md index b7eb2d60bd..c5e764e16f 100644 --- a/docs/source/extend/custom-components/adding-an-authentication-provider.md +++ b/docs/source/extend/custom-components/adding-an-authentication-provider.md @@ -42,12 +42,14 @@ from the clients that facilitate the authentication process. Authentication prov ## Extending an API Authentication Provider The first step in adding an authentication provider is to create a configuration model that inherits from the -{py:class}`~nat.data_models.authentication.AuthProviderBaseConfig` class and define the credentials required to +{py:class}`~nat.plugin_api.AuthProviderBaseConfig` class and define the credentials required to authenticate with the target API resource. The following example shows how to define and register a custom evaluator and can be found here: {py:class}`~nat.authentication.oauth2.oauth2_auth_code_flow_provider_config.OAuth2AuthCodeFlowProviderConfig` class: ```python +from nat.plugin_api import AuthProviderBaseConfig + class OAuth2AuthCodeFlowProviderConfig(AuthProviderBaseConfig, name="oauth2_auth_code_flow"): client_id: str = Field(description="The client ID for OAuth 2.0 authentication.") @@ -75,6 +77,9 @@ An asynchronous function decorated with {py:func}`~nat.plugin_api.register_auth_ The `OAuth2AuthCodeFlowProviderConfig` from the previous section is registered as follows: ```python +from nat.plugin_api import Builder +from nat.plugin_api import register_auth_provider + @register_auth_provider(config_type=OAuth2AuthCodeFlowProviderConfig) async def oauth2_client(authentication_provider: OAuth2AuthCodeFlowProviderConfig, builder: Builder): from nat.authentication.oauth2.oauth2_auth_code_flow_provider import OAuth2AuthCodeFlowProvider diff --git a/docs/source/extend/custom-components/adding-an-llm-provider.md b/docs/source/extend/custom-components/adding-an-llm-provider.md index ba851c9d55..f1e0162477 100644 --- a/docs/source/extend/custom-components/adding-an-llm-provider.md +++ b/docs/source/extend/custom-components/adding-an-llm-provider.md @@ -32,14 +32,16 @@ nat info components -t llm_client -q openai ## Provider Types -In NeMo Agent Toolkit, there are three provider types: `llm`, `embedder`, and `retriever`. The three provider types are defined by their respective base configuration classes: {class}`nat.data_models.llm.LLMBaseConfig`, {class}`nat.data_models.embedder.EmbedderBaseConfig`, and {class}`nat.data_models.retriever.RetrieverBaseConfig`. This guide focuses on adding an LLM provider. However, the process for adding an [embedder](../../build-workflows/embedders.md) or [retriever](../../build-workflows/retrievers.md) provider is similar. +In NeMo Agent Toolkit, there are three provider types: `llm`, `embedder`, and `retriever`. The three provider types are defined by their respective base configuration classes: {class}`nat.plugin_api.LLMBaseConfig`, {class}`nat.plugin_api.EmbedderBaseConfig`, and {class}`nat.plugin_api.RetrieverBaseConfig`. This guide focuses on adding an LLM provider. However, the process for adding an [embedder](../../build-workflows/embedders.md) or [retriever](../../build-workflows/retrievers.md) provider is similar. ## Defining an LLM Provider -The first step to adding an LLM provider is to subclass the {class}`nat.data_models.llm.LLMBaseConfig` class and add the configuration parameters needed to interact with the LLM API. Typically, this involves a `model_name` parameter and an `api_key` parameter; however, the exact parameters will depend on the API. The only requirement is a unique name for the provider. +The first step to adding an LLM provider is to subclass the {class}`nat.plugin_api.LLMBaseConfig` class and add the configuration parameters needed to interact with the LLM API. Typically, this involves a `model_name` parameter and an `api_key` parameter; however, the exact parameters will depend on the API. The only requirement is a unique name for the provider. Examine the previously mentioned {class}`nat.llm.openai_llm.OpenAIModelConfig` class: ```python +from nat.plugin_api import LLMBaseConfig + class OpenAIModelConfig(LLMBaseConfig, name="openai"): """An OpenAI LLM provider to be used with an LLM client.""" @@ -66,6 +68,7 @@ The {class}`nat.data_models.retry_mixin.RetryMixin` is a mixin that adds a `max_ ```python from nat.data_models.retry_mixin import RetryMixin +from nat.plugin_api import LLMBaseConfig class OpenAIModelConfig(LLMBaseConfig, RetryMixin, name="openai"): """An OpenAI LLM provider to be used with an LLM client.""" @@ -96,6 +99,7 @@ The {class}`nat.data_models.thinking_mixin.ThinkingMixin` is a mixin that adds a ```python from nat.data_models.thinking_mixin import ThinkingMixin +from nat.plugin_api import LLMBaseConfig class NIMModelConfig(LLMBaseConfig, ThinkingMixin, name="nim"): """An NIM LLM provider to be used with an LLM client.""" @@ -118,7 +122,7 @@ class NIMModelConfig(LLMBaseConfig, ThinkingMixin, name="nim"): ``` ### Registering the Provider -An asynchronous function decorated with {py:deco}`nat.plugin_api.register_llm_provider` is used to register the provider with NeMo Agent Toolkit by yielding an instance of {class}`nat.builder.llm.LLMProviderInfo`. +An asynchronous function decorated with {py:deco}`nat.plugin_api.register_llm_provider` is used to register the provider with NeMo Agent Toolkit by yielding an instance of {class}`nat.plugin_api.LLMProviderInfo`. :::{note} Registering an embedder or retriever provider is similar; however, the function should be decorated with {py:deco}`nat.plugin_api.register_embedder_provider` or {py:deco}`nat.plugin_api.register_retriever_provider`. @@ -128,6 +132,10 @@ Registering an embedder or retriever provider is similar; however, the function The `OpenAIModelConfig` from the previous section is registered as follows: `packages/nvidia_nat_core/src/nat/llm/openai_llm.py`: ```python +from nat.plugin_api import Builder +from nat.plugin_api import LLMProviderInfo +from nat.plugin_api import register_llm_provider + @register_llm_provider(config_type=OpenAIModelConfig) async def openai_llm(config: OpenAIModelConfig, builder: Builder): @@ -136,6 +144,10 @@ async def openai_llm(config: OpenAIModelConfig, builder: Builder): In the above example we didn't need to take any additional actions other than yielding the provider info. However, in some cases additional set up may be required, such as connecting to a cluster and performing validation could be performed in this method. In addition to this, any cleanup that needs to be done when the provider is no longer needed can be performed after the `yield` statement in the `finally` clause of a `try` statement. If this were needed we could update the above example as follows: ```python +from nat.plugin_api import Builder +from nat.plugin_api import LLMProviderInfo +from nat.plugin_api import register_llm_provider + @register_llm_provider(config_type=OpenAIModelConfig) async def openai_llm(config: OpenAIModelConfig, builder: Builder): # Perform any setup actions here and pre-flight checks here raising an exception if needed @@ -152,12 +164,16 @@ As previously mentioned, each LLM client is specific to both the LLM API and the Registering an embedder or retriever client is similar. However, the function should be decorated with {py:deco}`nat.plugin_api.register_embedder_client` or {py:deco}`nat.plugin_api.register_retriever_client`. ::: -The wrapped function in turn receives two required positional arguments: an instance of the configuration class of the provider, and an instance of {class}`nat.builder.builder.Builder`. The function should then yield a client suitable for the given provider and framework. The exact type is dictated by the framework itself and not by NeMo Agent Toolkit. +The wrapped function in turn receives two required positional arguments: an instance of the configuration class of the provider, and an instance of {class}`nat.plugin_api.Builder`. The function should then yield a client suitable for the given provider and framework. The exact type is dictated by the framework itself and not by NeMo Agent Toolkit. Since many frameworks provide clients for many of the common LLM APIs, in NeMo Agent Toolkit, the client registration functions are often simple factory methods. For example, the OpenAI client registration function for LangChain/LangGraph is as follows: `packages/nvidia_nat_langchain/src/nat/plugins/langchain/llm.py`: ```python +from nat.plugin_api import Builder +from nat.plugin_api import LLMFrameworkEnum +from nat.plugin_api import register_llm_client + @register_llm_client(config_type=OpenAIModelConfig, wrapper_type=LLMFrameworkEnum.LANGCHAIN) async def openai_langchain(llm_config: OpenAIModelConfig, builder: Builder): diff --git a/docs/source/extend/custom-components/custom-dataset-loader.md b/docs/source/extend/custom-components/custom-dataset-loader.md index 59e0b58680..b2fbff5341 100644 --- a/docs/source/extend/custom-components/custom-dataset-loader.md +++ b/docs/source/extend/custom-components/custom-dataset-loader.md @@ -48,10 +48,10 @@ The following example shows how to define and register a custom dataset loader f import pandas as pd from pydantic import Field +from nat.plugin_api import DatasetLoaderInfo from nat.plugin_api import EvalBuilder -from nat.builder.dataset_loader import DatasetLoaderInfo +from nat.plugin_api import EvalDatasetBaseConfig from nat.plugin_api import register_dataset_loader -from nat.data_models.dataset_handler import EvalDatasetBaseConfig class EvalDatasetTsvConfig(EvalDatasetBaseConfig, name="tsv"): diff --git a/docs/source/extend/custom-components/custom-evaluator.md b/docs/source/extend/custom-components/custom-evaluator.md index f5577acaaa..99df98ee73 100644 --- a/docs/source/extend/custom-components/custom-evaluator.md +++ b/docs/source/extend/custom-components/custom-evaluator.md @@ -49,9 +49,9 @@ The following example shows how to define and register a custom evaluator. The c from pydantic import Field from nat.plugin_api import EvalBuilder -from nat.builder.evaluator import EvaluatorInfo +from nat.plugin_api import EvaluatorBaseConfig +from nat.plugin_api import EvaluatorInfo from nat.plugin_api import register_evaluator -from nat.data_models.evaluator import EvaluatorBaseConfig class SimilarityEvaluatorConfig(EvaluatorBaseConfig, name="similarity"): @@ -166,9 +166,9 @@ from collections import Counter from pydantic import Field from nat.plugin_api import EvalBuilder -from nat.builder.evaluator import EvaluatorInfo +from nat.plugin_api import EvaluatorBaseConfig +from nat.plugin_api import EvaluatorInfo from nat.plugin_api import register_evaluator -from nat.data_models.evaluator import EvaluatorBaseConfig from nat.plugins.eval.data_models.evaluator_io import EvalOutputItem from nat.plugins.eval.evaluator.atif_base_evaluator import AtifBaseEvaluator from nat.plugins.eval.evaluator.atif_evaluator import AtifEvalSample @@ -308,4 +308,3 @@ The results of each evaluator is stored in a separate file with name `_ } ``` The contents of the file have been `snipped` for brevity. - diff --git a/docs/source/extend/custom-components/finetuning.md b/docs/source/extend/custom-components/finetuning.md index 5e38b370d9..dd1ad6f518 100644 --- a/docs/source/extend/custom-components/finetuning.md +++ b/docs/source/extend/custom-components/finetuning.md @@ -70,8 +70,9 @@ from abc import ABC, abstractmethod from typing import Any from nat.data_models.evaluate_runtime import EvaluationRunOutput -from nat.data_models.finetuning import FinetuneConfig, TrajectoryBuilderConfig, TrajectoryCollection from nat.data_models.evaluator import EvalOutputItem +from nat.data_models.finetuning import FinetuneConfig, TrajectoryCollection +from nat.plugin_api import TrajectoryBuilderConfig class TrajectoryBuilder(ABC): @@ -137,7 +138,7 @@ Create a configuration class that inherits from `TrajectoryBuilderConfig`: ```python from pydantic import Field -from nat.data_models.finetuning import TrajectoryBuilderConfig +from nat.plugin_api import TrajectoryBuilderConfig class MyTrajectoryBuilderConfig(TrajectoryBuilderConfig, name="my_traj_builder"): @@ -202,9 +203,9 @@ The `TrainerAdapter` bridges the gap between NeMo Agent Toolkit and external tra from abc import ABC, abstractmethod from typing import Any +from nat.plugin_api import TrainerAdapterConfig from nat.data_models.finetuning import ( FinetuneConfig, - TrainerAdapterConfig, TrainingJobRef, TrainingJobStatus, TrajectoryCollection, @@ -263,7 +264,7 @@ class TrainerAdapter(ABC): ```python from pydantic import BaseModel, Field -from nat.data_models.finetuning import TrainerAdapterConfig +from nat.plugin_api import TrainerAdapterConfig class MyBackendConfig(BaseModel): @@ -319,10 +320,10 @@ The `Trainer` orchestrates the complete finetuning workflow, coordinating the tr from abc import ABC, abstractmethod from typing import Any +from nat.plugin_api import TrainerConfig from nat.data_models.finetuning import ( FinetuneConfig, FinetuneRunConfig, - TrainerConfig, TrainingJobRef, TrainingJobStatus, TrajectoryCollection, diff --git a/docs/source/extend/custom-components/memory.md b/docs/source/extend/custom-components/memory.md index 16a3b58deb..9754e72c41 100644 --- a/docs/source/extend/custom-components/memory.md +++ b/docs/source/extend/custom-components/memory.md @@ -22,16 +22,15 @@ This documentation presumes familiarity with the NeMo Agent Toolkit [memory modu ## Key Memory Module Components * **Memory Data Models** - - **{py:class}`~nat.data_models.memory.MemoryBaseConfig`**: A Pydantic base class that all memory config classes must extend. This is used for specifying memory registration in the NeMo Agent Toolkit config file. - - **{py:class}`~nat.data_models.memory.MemoryBaseConfigT`**: A generic type alias for memory config classes. + - **{py:class}`~nat.plugin_api.MemoryBaseConfig`**: A Pydantic base class that all memory config classes must extend. This is used for specifying memory registration in the NeMo Agent Toolkit config file. * **Memory Interfaces** - - **{py:class}`~nat.memory.interfaces.MemoryEditor`** (abstract interface): The low-level API for adding, searching, and removing memory items. - - **{py:class}`~nat.memory.interfaces.MemoryReader`** and **{py:class}`~nat.memory.interfaces.MemoryWriter`** (abstract classes): Provide structured read/write logic on top of the `MemoryEditor`. - - **{py:class}`~nat.memory.interfaces.MemoryManager`** (abstract interface): Manages higher-level memory operations like summarization or reflection if needed. + - **{py:class}`~nat.plugin_api.MemoryEditor`** (abstract interface): The low-level API for adding, searching, and removing memory items. + - **{py:class}`~nat.plugin_api.MemoryReader`** and **{py:class}`~nat.plugin_api.MemoryWriter`** (abstract classes): Provide structured read/write logic on top of the `MemoryEditor`. + - **{py:class}`~nat.plugin_api.MemoryManager`** (abstract interface): Manages higher-level memory operations like summarization or reflection if needed. * **Memory Models** - - **{py:class}`~nat.memory.models.MemoryItem`**: The main object representing a piece of memory. It includes: + - **{py:class}`~nat.plugin_api.MemoryItem`**: The main object representing a piece of memory. It includes: ```python conversation: list[dict[str, str]] # user/assistant messages tags: list[str] = [] @@ -44,13 +43,13 @@ This documentation presumes familiarity with the NeMo Agent Toolkit [memory modu ## Adding a Memory Module -In the NeMo Agent Toolkit system, anything that extends {py:class}`~nat.data_models.memory.MemoryBaseConfig` and is declared with a `name="some_memory"` can be discovered as a *Memory type* by the NeMo Agent Toolkit global type registry. This allows you to define a custom memory class to handle your own backends (Redis, custom database, a vector store, etc.). Then your memory class can be selected in the NeMo Agent Toolkit config YAML via `_type: `. +In the NeMo Agent Toolkit system, anything that extends {py:class}`~nat.plugin_api.MemoryBaseConfig` and is declared with a `name="some_memory"` can be discovered as a *Memory type* by the NeMo Agent Toolkit global type registry. This allows you to define a custom memory class to handle your own backends (Redis, custom database, a vector store, etc.). Then your memory class can be selected in the NeMo Agent Toolkit config YAML via `_type: `. ### Basic Steps -1. **Create a config Class** that extends {py:class}`~nat.data_models.memory.MemoryBaseConfig`: +1. **Create a config Class** that extends {py:class}`~nat.plugin_api.MemoryBaseConfig`: ```python - from nat.data_models.memory import MemoryBaseConfig + from nat.plugin_api import MemoryBaseConfig class MyCustomMemoryConfig(MemoryBaseConfig, name="my_custom_memory"): # You can define any fields you want. For example: @@ -62,9 +61,10 @@ In the NeMo Agent Toolkit system, anything that extends {py:class}`~nat.data_mod The `name="my_custom_memory"` ensures that NeMo Agent Toolkit can recognize it when the user places `_type: my_custom_memory` in the memory config. ::: -2. **Implement a {py:class}`~nat.memory.interfaces.MemoryEditor`** that uses your backend**: +2. **Implement a {py:class}`~nat.plugin_api.MemoryEditor`** that uses your backend**: ```python - from nat.memory.interfaces import MemoryEditor, MemoryItem + from nat.plugin_api import MemoryEditor + from nat.plugin_api import MemoryItem class MyCustomMemoryEditor(MemoryEditor): def __init__(self, config: MyCustomMemoryConfig): @@ -104,22 +104,23 @@ In the NeMo Agent Toolkit system, anything that extends {py:class}`~nat.data_mod ## Bringing Your Own Memory Client Implementation A typical pattern is: -- You define a *config class* that extends {py:class}`~nat.data_models.memory.MemoryBaseConfig` (giving it a unique `_type` / name). -- You define the actual *runtime logic* in a "Memory Editor" or "Memory Client" class that implements {py:class}`~nat.memory.interfaces.MemoryEditor`. +- You define a *config class* that extends {py:class}`~nat.plugin_api.MemoryBaseConfig` (giving it a unique `_type` / name). +- You define the actual *runtime logic* in a "Memory Editor" or "Memory Client" class that implements {py:class}`~nat.plugin_api.MemoryEditor`. - You connect them together (for example, by implementing a small factory function or a method in the builder that says: "Given `MyCustomMemoryConfig`, return `MyCustomMemoryEditor(config)`"). ### Example: Minimal Skeleton ```python # my_custom_memory_config.py -from nat.data_models.memory import MemoryBaseConfig +from nat.plugin_api import MemoryBaseConfig class MyCustomMemoryConfig(MemoryBaseConfig, name="my_custom_memory"): url: str token: str # my_custom_memory_editor.py -from nat.memory.interfaces import MemoryEditor, MemoryItem +from nat.plugin_api import MemoryEditor +from nat.plugin_api import MemoryItem class MyCustomMemoryEditor(MemoryEditor): def __init__(self, cfg: MyCustomMemoryConfig): @@ -163,7 +164,7 @@ memories = await memory_client.search(query="What did user prefer last time?", t **Inside Tools**: Tools that read or write memory simply call the memory client. For example: ```python -from nat.memory.models import MemoryItem +from nat.plugin_api import MemoryItem from langchain_core.tools import ToolException async def add_memory_tool_action(item: MemoryItem, memory_name: str): @@ -226,8 +227,8 @@ For convenient memory persistence, you can use the [automatic memory wrapper](.. To **bring your own memory**: -1. **Implement** a custom {py:class}`~nat.data_models.memory.MemoryBaseConfig` (with a unique `_type`). -2. **Implement** a custom {py:class}`~nat.memory.interfaces.MemoryEditor` that can handle `add_items`, `search`, `remove_items` calls. +1. **Implement** a custom {py:class}`~nat.plugin_api.MemoryBaseConfig` (with a unique `_type`). +2. **Implement** a custom {py:class}`~nat.plugin_api.MemoryEditor` that can handle `add_items`, `search`, `remove_items` calls. 3. **Register** your config class so that the NeMo Agent Toolkit type registry is aware of `_type: `. 4. In your `.yml` config, specify: ```yaml @@ -242,8 +243,8 @@ To **bring your own memory**: ## Summary -- The **Memory** module in NeMo Agent Toolkit revolves around the {py:class}`~nat.memory.interfaces.MemoryEditor` interface and {py:class}`~nat.memory.models.MemoryItem` model. -- **Configuration** is done via a subclass of {py:class}`~nat.data_models.memory.MemoryBaseConfig` that is *discriminated* by the `_type` field in the YAML config. +- The **Memory** module in NeMo Agent Toolkit revolves around the {py:class}`~nat.plugin_api.MemoryEditor` interface and {py:class}`~nat.plugin_api.MemoryItem` model. +- **Configuration** is done via a subclass of {py:class}`~nat.plugin_api.MemoryBaseConfig` that is *discriminated* by the `_type` field in the YAML config. - **Registration** can be as simple as adding `name="my_custom_memory"` to your config class and letting NeMo Agent Toolkit discover it. - Tools and workflows then seamlessly **read/write** user memory by calling `builder.get_memory_client(...)`. diff --git a/docs/source/extend/custom-components/object-store.md b/docs/source/extend/custom-components/object-store.md index 4547d49641..d562de9f72 100644 --- a/docs/source/extend/custom-components/object-store.md +++ b/docs/source/extend/custom-components/object-store.md @@ -22,11 +22,10 @@ This documentation presumes familiarity with the NeMo Agent Toolkit [object stor ## Key Object Store Module Components * **Object Store Data Models** - - **{py:class}`~nat.data_models.object_store.ObjectStoreBaseConfig`**: A Pydantic base class that all object store config classes must extend. This is used for specifying object store registration in the NeMo Agent Toolkit config file. - - **{py:class}`~nat.data_models.object_store.ObjectStoreBaseConfigT`**: A generic type alias for object store config classes. + - **{py:class}`~nat.plugin_api.ObjectStoreBaseConfig`**: A Pydantic base class that all object store config classes must extend. This is used for specifying object store registration in the NeMo Agent Toolkit config file. * **Object Store Interfaces** - - **{py:class}`~nat.object_store.interfaces.ObjectStore`** (abstract interface): The core interface for object store operations, including put, upsert, get, and delete operations. + - **{py:class}`~nat.plugin_api.ObjectStore`** (abstract interface): The core interface for object store operations, including put, upsert, get, and delete operations. ```python class ObjectStore(ABC): @abstractmethod @@ -47,7 +46,7 @@ This documentation presumes familiarity with the NeMo Agent Toolkit [object stor ``` * **Object Store Models** - - **{py:class}`~nat.object_store.models.ObjectStoreItem`**: The main object representing an item in the object store. + - **{py:class}`~nat.plugin_api.ObjectStoreItem`**: The main object representing an item in the object store. ```python class ObjectStoreItem: data: bytes # The binary data to store @@ -56,18 +55,18 @@ This documentation presumes familiarity with the NeMo Agent Toolkit [object stor ``` * **Object Store Exceptions** - - **{py:class}`~nat.data_models.object_store.KeyAlreadyExistsError`**: Raised when trying to store an object with a key that already exists (for `put_object`) - - **{py:class}`~nat.data_models.object_store.NoSuchKeyError`**: Raised when trying to retrieve or delete an object with a non-existent key + - **{py:class}`~nat.plugin_api.KeyAlreadyExistsError`**: Raised when trying to store an object with a key that already exists (for `put_object`) + - **{py:class}`~nat.plugin_api.NoSuchKeyError`**: Raised when trying to retrieve or delete an object with a non-existent key ## Adding an Object Store Provider -In the NeMo Agent Toolkit system, anything that extends {py:class}`~nat.data_models.object_store.ObjectStoreBaseConfig` and is declared with a `name="some_object_store"` can be discovered as an *Object Store type* by the NeMo Agent Toolkit global type registry. This allows you to define a custom object store class to handle your own backends (for example, Redis, custom database, or cloud storage). Then your object store class can be selected in the NeMo Agent Toolkit config YAML using `_type: `. +In the NeMo Agent Toolkit system, anything that extends {py:class}`~nat.plugin_api.ObjectStoreBaseConfig` and is declared with a `name="some_object_store"` can be discovered as an *Object Store type* by the NeMo Agent Toolkit global type registry. This allows you to define a custom object store class to handle your own backends (for example, Redis, custom database, or cloud storage). Then your object store class can be selected in the NeMo Agent Toolkit config YAML using `_type: `. ### Basic Steps -1. **Create a config Class** that extends {py:class}`~nat.data_models.object_store.ObjectStoreBaseConfig`: +1. **Create a config Class** that extends {py:class}`~nat.plugin_api.ObjectStoreBaseConfig`: ```python - from nat.data_models.object_store import ObjectStoreBaseConfig + from nat.plugin_api import ObjectStoreBaseConfig class MyCustomObjectStoreConfig(ObjectStoreBaseConfig, name="my_custom_object_store"): # You can define any fields you want. For example: @@ -80,14 +79,15 @@ In the NeMo Agent Toolkit system, anything that extends {py:class}`~nat.data_mod The `name="my_custom_object_store"` ensures that NeMo Agent Toolkit can recognize it when the user places `_type: my_custom_object_store` in the object store config. ::: -2. **Implement an {py:class}`~nat.object_store.interfaces.ObjectStore`** that uses your backend: +2. **Implement an {py:class}`~nat.plugin_api.ObjectStore`** that uses your backend: It is recommended to have this implementation in a separate file from the config class and registration code. ```python - from nat.object_store.interfaces import ObjectStore - from nat.object_store.models import ObjectStoreItem - from nat.data_models.object_store import KeyAlreadyExistsError, NoSuchKeyError + from nat.plugin_api import KeyAlreadyExistsError + from nat.plugin_api import NoSuchKeyError + from nat.plugin_api import ObjectStore + from nat.plugin_api import ObjectStoreItem from nat.utils.type_utils import override class MyCustomObjectStore(ObjectStore): @@ -180,8 +180,8 @@ In the NeMo Agent Toolkit system, anything that extends {py:class}`~nat.data_mod ## Bringing Your Own Object Store Implementation A typical pattern is: -- You define a *config class* that extends {py:class}`~nat.data_models.object_store.ObjectStoreBaseConfig` (giving it a unique `_type` / name). -- You define the actual *runtime logic* in an "Object Store" class that implements {py:class}`~nat.object_store.interfaces.ObjectStore`. +- You define a *config class* that extends {py:class}`~nat.plugin_api.ObjectStoreBaseConfig` (giving it a unique `_type` / name). +- You define the actual *runtime logic* in an "Object Store" class that implements {py:class}`~nat.plugin_api.ObjectStore`. - You connect them together using the `@register_object_store` decorator. ### Example: Minimal Skeleton @@ -196,10 +196,10 @@ my_custom_object_store `my_custom_object_store.py` contents: ```python -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 +from nat.plugin_api import KeyAlreadyExistsError +from nat.plugin_api import NoSuchKeyError +from nat.plugin_api import ObjectStore +from nat.plugin_api import ObjectStoreItem from nat.utils.type_utils import override class MyCustomObjectStore(ObjectStore): @@ -232,7 +232,7 @@ class MyCustomObjectStore(ObjectStore): `object_store.py` contents: ```python -from nat.data_models.object_store import ObjectStoreBaseConfig +from nat.plugin_api import ObjectStoreBaseConfig class MyCustomObjectStoreConfig(ObjectStoreBaseConfig, name="my_custom_object_store"): url: str @@ -274,7 +274,7 @@ print(item.data.decode("utf-8")) **Inside Functions**: Functions that read or write to object stores simply call the object store client. For example: ```python -from nat.object_store.models import ObjectStoreItem +from nat.plugin_api import ObjectStoreItem from langchain_core.tools import ToolException async def store_file_tool_action(file_data: bytes, key: str, object_store_name: str): diff --git a/docs/source/extend/custom-components/optimizer.md b/docs/source/extend/custom-components/optimizer.md index 1752c403eb..3a5a116e3c 100644 --- a/docs/source/extend/custom-components/optimizer.md +++ b/docs/source/extend/custom-components/optimizer.md @@ -26,8 +26,8 @@ NeMo Agent Toolkit provides a pluggable optimizer system for tuning workflow par ## Key Interfaces * **Configuration Base Classes** - - {py:class}`~nat.data_models.optimizer.OptimizerStrategyBaseConfig`: Base class that all optimizer strategy configuration models must extend. Provides an `enabled` field and integrates with the NeMo Agent Toolkit type registry. - - {py:class}`~nat.data_models.optimizer.PromptOptimizationConfig`: Base for prompt optimization strategy configuration models. Adds `prompt_population_init_function` and `prompt_recombination_function` fields. + - {py:class}`~nat.plugin_api.OptimizerStrategyBaseConfig`: Base class that all optimizer strategy configuration models must extend. Provides an `enabled` field and integrates with the NeMo Agent Toolkit type registry. + - {py:class}`~nat.plugin_api.PromptOptimizationConfig`: Base for prompt optimization strategy configuration models. Adds `prompt_population_init_function` and `prompt_recombination_function` fields. - {py:class}`~nat.data_models.optimizer.OptunaParameterOptimizationConfig`: Built-in config for Optuna-based numeric parameter optimization. * **Optimizer ABCs** @@ -41,12 +41,12 @@ NeMo Agent Toolkit provides a pluggable optimizer system for tuning workflow par ### 1. Define a config class -Create a config class extending {py:class}`~nat.data_models.optimizer.PromptOptimizationConfig` with a unique `name`: +Create a config class extending {py:class}`~nat.plugin_api.PromptOptimizationConfig` with a unique `name`: ```python from pydantic import Field -from nat.data_models.optimizer import PromptOptimizationConfig +from nat.plugin_api import PromptOptimizationConfig class IterativeRefinementPromptConfig(PromptOptimizationConfig, name="iterative"): @@ -152,7 +152,7 @@ The pattern is the same, but parameter optimizers extend {py:class}`~nat.plugins ```python from pydantic import Field -from nat.data_models.optimizer import OptimizerStrategyBaseConfig +from nat.plugin_api import OptimizerStrategyBaseConfig class RandomSearchConfig(OptimizerStrategyBaseConfig, name="random_search"): diff --git a/docs/source/extend/custom-components/telemetry-exporters.md b/docs/source/extend/custom-components/telemetry-exporters.md index 91e310610c..59cdfca4b0 100644 --- a/docs/source/extend/custom-components/telemetry-exporters.md +++ b/docs/source/extend/custom-components/telemetry-exporters.md @@ -89,8 +89,8 @@ Want to get started quickly? Here's a minimal working example that creates a con from pydantic import Field from nat.plugin_api import Builder +from nat.plugin_api import TelemetryExporterBaseConfig from nat.plugin_api import register_telemetry_exporter -from nat.data_models.telemetry_exporter import TelemetryExporterBaseConfig from nat.observability.exporter.raw_exporter import RawExporter from nat.data_models.intermediate_step import IntermediateStep @@ -329,7 +329,7 @@ Create a configuration class that inherits from `TelemetryExporterBaseConfig`: ```python from pydantic import Field -from nat.data_models.telemetry_exporter import TelemetryExporterBaseConfig +from nat.plugin_api import TelemetryExporterBaseConfig class CustomTelemetryExporter(TelemetryExporterBaseConfig, name="custom"): """A simple custom telemetry exporter for sending traces to a custom service.""" @@ -537,9 +537,9 @@ class MyCustomExporter(SpanExporter[Span, dict]): ```python from pydantic import Field -from nat.plugin_api import register_telemetry_exporter -from nat.data_models.telemetry_exporter import TelemetryExporterBaseConfig from nat.plugin_api import Builder +from nat.plugin_api import TelemetryExporterBaseConfig +from nat.plugin_api import register_telemetry_exporter # Configuration class can be in the same file as registration class MyTelemetryExporter(TelemetryExporterBaseConfig, name="my_exporter"): @@ -1424,8 +1424,8 @@ import logging from pydantic import Field import aiohttp from nat.plugin_api import Builder +from nat.plugin_api import TelemetryExporterBaseConfig from nat.plugin_api import register_telemetry_exporter -from nat.data_models.telemetry_exporter import TelemetryExporterBaseConfig from nat.observability.exporter.span_exporter import SpanExporter from nat.observability.exporter.base_exporter import IsolatedAttribute from nat.data_models.span import Span diff --git a/docs/source/run-workflows/existing-agents/langgraph.md b/docs/source/run-workflows/existing-agents/langgraph.md index b5a68b4862..8a7c721e7a 100644 --- a/docs/source/run-workflows/existing-agents/langgraph.md +++ b/docs/source/run-workflows/existing-agents/langgraph.md @@ -123,7 +123,7 @@ agent = create_deep_agent( ```python from deepagents import create_deep_agent -from nat.builder.framework_enum import LLMFrameworkEnum +from nat.plugin_api import LLMFrameworkEnum from nat.builder.sync_builder import SyncBuilder # Get model from NeMo Agent Toolkit configuration diff --git a/packages/nvidia_nat_core/src/nat/plugin_api.py b/packages/nvidia_nat_core/src/nat/plugin_api.py index 23c4e2e907..0bef4e0830 100644 --- a/packages/nvidia_nat_core/src/nat/plugin_api.py +++ b/packages/nvidia_nat_core/src/nat/plugin_api.py @@ -21,10 +21,15 @@ from nat.builder.builder import Builder from nat.builder.builder import EvalBuilder +from nat.builder.dataset_loader import DatasetLoaderInfo +from nat.builder.embedder import EmbedderProviderInfo +from nat.builder.evaluator import EvaluatorInfo from nat.builder.framework_enum import LLMFrameworkEnum from nat.builder.function import Function from nat.builder.function import FunctionGroup from nat.builder.function_info import FunctionInfo +from nat.builder.llm import LLMProviderInfo +from nat.builder.retriever import RetrieverProviderInfo from nat.cli.register_workflow import register_auth_provider from nat.cli.register_workflow import register_dataset_loader from nat.cli.register_workflow import register_embedder_client @@ -53,6 +58,7 @@ from nat.cli.register_workflow import register_trainer_adapter from nat.cli.register_workflow import register_trajectory_builder from nat.cli.register_workflow import register_ttc_strategy +from nat.data_models.authentication import AuthProviderBaseConfig from nat.data_models.common import OptionalSecretStr from nat.data_models.common import SerializableSecretStr from nat.data_models.common import get_secret_value @@ -71,15 +77,60 @@ from nat.data_models.component_ref import TrainerRef from nat.data_models.component_ref import TrajectoryBuilderRef from nat.data_models.component_ref import TTCStrategyRef +from nat.data_models.dataset_handler import EvalDatasetBaseConfig +from nat.data_models.embedder import EmbedderBaseConfig +from nat.data_models.evaluator import EvaluatorBaseConfig +from nat.data_models.finetuning import TrainerAdapterConfig +from nat.data_models.finetuning import TrainerConfig +from nat.data_models.finetuning import TrajectoryBuilderConfig +from nat.data_models.front_end import FrontEndBaseConfig from nat.data_models.function import FunctionBaseConfig from nat.data_models.function import FunctionGroupBaseConfig +from nat.data_models.llm import LLMBaseConfig +from nat.data_models.logging import LoggingBaseConfig +from nat.data_models.memory import MemoryBaseConfig +from nat.data_models.middleware import FunctionMiddlewareBaseConfig +from nat.data_models.middleware import MiddlewareBaseConfig +from nat.data_models.object_store import KeyAlreadyExistsError +from nat.data_models.object_store import NoSuchKeyError +from nat.data_models.object_store import ObjectStoreBaseConfig +from nat.data_models.optimizer import OptimizerStrategyBaseConfig +from nat.data_models.optimizer import PromptOptimizationConfig +from nat.data_models.registry_handler import RegistryHandlerBaseConfig +from nat.data_models.retriever import RetrieverBaseConfig +from nat.data_models.telemetry_exporter import TelemetryExporterBaseConfig +from nat.data_models.ttc_strategy import TTCStrategyBaseConfig +from nat.memory.interfaces import MemoryEditor +from nat.memory.interfaces import MemoryManager +from nat.memory.interfaces import MemoryReader +from nat.memory.interfaces import MemoryWriter +from nat.memory.models import MemoryItem +from nat.middleware.dynamic.dynamic_function_middleware import DynamicFunctionMiddleware +from nat.middleware.dynamic.dynamic_middleware_config import DynamicMiddlewareConfig +from nat.middleware.function_middleware import FunctionMiddleware +from nat.middleware.middleware import FunctionMiddlewareContext +from nat.middleware.middleware import InvocationContext +from nat.object_store.interfaces import ObjectStore +from nat.object_store.models import ObjectStoreItem +# Public contract: keep this list exact and update docs/source/extend/plugin-api.md plus +# packages/nvidia_nat_core/tests/nat/test_plugin_api.py whenever symbols are added or removed. __all__ = [ + "AuthProviderBaseConfig", "AuthenticationRef", "Builder", "ComponentRef", + "DatasetLoaderInfo", + "DynamicFunctionMiddleware", + "DynamicMiddlewareConfig", + "EmbedderBaseConfig", + "EmbedderProviderInfo", "EmbedderRef", "EvalBuilder", + "EvalDatasetBaseConfig", + "EvaluatorBaseConfig", + "EvaluatorInfo", + "FrontEndBaseConfig", "Function", "FunctionBaseConfig", "FunctionGroup", @@ -87,17 +138,46 @@ "FunctionGroupRef", "FunctionInfo", "FunctionRef", + "FunctionMiddleware", + "FunctionMiddlewareBaseConfig", + "FunctionMiddlewareContext", + "InvocationContext", + "KeyAlreadyExistsError", + "LLMBaseConfig", "LLMFrameworkEnum", + "LLMProviderInfo", "LLMRef", + "LoggingBaseConfig", + "MemoryBaseConfig", + "MemoryEditor", + "MemoryItem", + "MemoryManager", + "MemoryReader", "MemoryRef", + "MemoryWriter", + "MiddlewareBaseConfig", "MiddlewareRef", + "NoSuchKeyError", + "ObjectStore", "ObjectStoreRef", + "ObjectStoreItem", + "ObjectStoreBaseConfig", + "OptimizerStrategyBaseConfig", "OptionalSecretStr", + "PromptOptimizationConfig", + "RegistryHandlerBaseConfig", + "RetrieverBaseConfig", + "RetrieverProviderInfo", "RetrieverRef", "SerializableSecretStr", + "TelemetryExporterBaseConfig", "TTCStrategyRef", + "TTCStrategyBaseConfig", + "TrainerAdapterConfig", "TrainerAdapterRef", + "TrainerConfig", "TrainerRef", + "TrajectoryBuilderConfig", "TrajectoryBuilderRef", "get_secret_value", "register_auth_provider", diff --git a/packages/nvidia_nat_core/tests/nat/test_plugin_api.py b/packages/nvidia_nat_core/tests/nat/test_plugin_api.py index 53ce0d7586..ae78b062e2 100644 --- a/packages/nvidia_nat_core/tests/nat/test_plugin_api.py +++ b/packages/nvidia_nat_core/tests/nat/test_plugin_api.py @@ -13,29 +13,115 @@ # See the License for the specific language governing permissions and # limitations under the License. +import importlib from pathlib import Path from nat import plugin_api -from nat.builder.builder import Builder -from nat.builder.function import FunctionGroup -from nat.cli.register_workflow import register_function -from nat.cli.register_workflow import register_function_group -from nat.data_models.function import FunctionBaseConfig -from nat.data_models.function import FunctionGroupBaseConfig +EXPECTED_PLUGIN_API_EXPORTS = { + "AuthProviderBaseConfig": ("nat.data_models.authentication", "AuthProviderBaseConfig"), + "AuthenticationRef": ("nat.data_models.component_ref", "AuthenticationRef"), + "Builder": ("nat.builder.builder", "Builder"), + "ComponentRef": ("nat.data_models.component_ref", "ComponentRef"), + "DatasetLoaderInfo": ("nat.builder.dataset_loader", "DatasetLoaderInfo"), + "DynamicFunctionMiddleware": ("nat.middleware.dynamic.dynamic_function_middleware", "DynamicFunctionMiddleware"), + "DynamicMiddlewareConfig": ("nat.middleware.dynamic.dynamic_middleware_config", "DynamicMiddlewareConfig"), + "EmbedderBaseConfig": ("nat.data_models.embedder", "EmbedderBaseConfig"), + "EmbedderProviderInfo": ("nat.builder.embedder", "EmbedderProviderInfo"), + "EmbedderRef": ("nat.data_models.component_ref", "EmbedderRef"), + "EvalBuilder": ("nat.builder.builder", "EvalBuilder"), + "EvalDatasetBaseConfig": ("nat.data_models.dataset_handler", "EvalDatasetBaseConfig"), + "EvaluatorBaseConfig": ("nat.data_models.evaluator", "EvaluatorBaseConfig"), + "EvaluatorInfo": ("nat.builder.evaluator", "EvaluatorInfo"), + "FrontEndBaseConfig": ("nat.data_models.front_end", "FrontEndBaseConfig"), + "Function": ("nat.builder.function", "Function"), + "FunctionBaseConfig": ("nat.data_models.function", "FunctionBaseConfig"), + "FunctionGroup": ("nat.builder.function", "FunctionGroup"), + "FunctionGroupBaseConfig": ("nat.data_models.function", "FunctionGroupBaseConfig"), + "FunctionGroupRef": ("nat.data_models.component_ref", "FunctionGroupRef"), + "FunctionInfo": ("nat.builder.function_info", "FunctionInfo"), + "FunctionRef": ("nat.data_models.component_ref", "FunctionRef"), + "FunctionMiddleware": ("nat.middleware.function_middleware", "FunctionMiddleware"), + "FunctionMiddlewareBaseConfig": ("nat.data_models.middleware", "FunctionMiddlewareBaseConfig"), + "FunctionMiddlewareContext": ("nat.middleware.middleware", "FunctionMiddlewareContext"), + "InvocationContext": ("nat.middleware.middleware", "InvocationContext"), + "KeyAlreadyExistsError": ("nat.data_models.object_store", "KeyAlreadyExistsError"), + "LLMBaseConfig": ("nat.data_models.llm", "LLMBaseConfig"), + "LLMFrameworkEnum": ("nat.builder.framework_enum", "LLMFrameworkEnum"), + "LLMProviderInfo": ("nat.builder.llm", "LLMProviderInfo"), + "LLMRef": ("nat.data_models.component_ref", "LLMRef"), + "LoggingBaseConfig": ("nat.data_models.logging", "LoggingBaseConfig"), + "MemoryBaseConfig": ("nat.data_models.memory", "MemoryBaseConfig"), + "MemoryEditor": ("nat.memory.interfaces", "MemoryEditor"), + "MemoryItem": ("nat.memory.models", "MemoryItem"), + "MemoryManager": ("nat.memory.interfaces", "MemoryManager"), + "MemoryReader": ("nat.memory.interfaces", "MemoryReader"), + "MemoryRef": ("nat.data_models.component_ref", "MemoryRef"), + "MemoryWriter": ("nat.memory.interfaces", "MemoryWriter"), + "MiddlewareBaseConfig": ("nat.data_models.middleware", "MiddlewareBaseConfig"), + "MiddlewareRef": ("nat.data_models.component_ref", "MiddlewareRef"), + "NoSuchKeyError": ("nat.data_models.object_store", "NoSuchKeyError"), + "ObjectStore": ("nat.object_store.interfaces", "ObjectStore"), + "ObjectStoreRef": ("nat.data_models.component_ref", "ObjectStoreRef"), + "ObjectStoreItem": ("nat.object_store.models", "ObjectStoreItem"), + "ObjectStoreBaseConfig": ("nat.data_models.object_store", "ObjectStoreBaseConfig"), + "OptimizerStrategyBaseConfig": ("nat.data_models.optimizer", "OptimizerStrategyBaseConfig"), + "OptionalSecretStr": ("nat.data_models.common", "OptionalSecretStr"), + "PromptOptimizationConfig": ("nat.data_models.optimizer", "PromptOptimizationConfig"), + "RegistryHandlerBaseConfig": ("nat.data_models.registry_handler", "RegistryHandlerBaseConfig"), + "RetrieverBaseConfig": ("nat.data_models.retriever", "RetrieverBaseConfig"), + "RetrieverProviderInfo": ("nat.builder.retriever", "RetrieverProviderInfo"), + "RetrieverRef": ("nat.data_models.component_ref", "RetrieverRef"), + "SerializableSecretStr": ("nat.data_models.common", "SerializableSecretStr"), + "TelemetryExporterBaseConfig": ("nat.data_models.telemetry_exporter", "TelemetryExporterBaseConfig"), + "TTCStrategyRef": ("nat.data_models.component_ref", "TTCStrategyRef"), + "TTCStrategyBaseConfig": ("nat.data_models.ttc_strategy", "TTCStrategyBaseConfig"), + "TrainerAdapterConfig": ("nat.data_models.finetuning", "TrainerAdapterConfig"), + "TrainerAdapterRef": ("nat.data_models.component_ref", "TrainerAdapterRef"), + "TrainerConfig": ("nat.data_models.finetuning", "TrainerConfig"), + "TrainerRef": ("nat.data_models.component_ref", "TrainerRef"), + "TrajectoryBuilderConfig": ("nat.data_models.finetuning", "TrajectoryBuilderConfig"), + "TrajectoryBuilderRef": ("nat.data_models.component_ref", "TrajectoryBuilderRef"), + "get_secret_value": ("nat.data_models.common", "get_secret_value"), + "register_auth_provider": ("nat.cli.register_workflow", "register_auth_provider"), + "register_dataset_loader": ("nat.cli.register_workflow", "register_dataset_loader"), + "register_embedder_client": ("nat.cli.register_workflow", "register_embedder_client"), + "register_embedder_provider": ("nat.cli.register_workflow", "register_embedder_provider"), + "register_eval_callback": ("nat.cli.register_workflow", "register_eval_callback"), + "register_evaluator": ("nat.cli.register_workflow", "register_evaluator"), + "register_front_end": ("nat.cli.register_workflow", "register_front_end"), + "register_function": ("nat.cli.register_workflow", "register_function"), + "register_function_group": ("nat.cli.register_workflow", "register_function_group"), + "register_llm_client": ("nat.cli.register_workflow", "register_llm_client"), + "register_llm_provider": ("nat.cli.register_workflow", "register_llm_provider"), + "register_logging_method": ("nat.cli.register_workflow", "register_logging_method"), + "register_memory": ("nat.cli.register_workflow", "register_memory"), + "register_middleware": ("nat.cli.register_workflow", "register_middleware"), + "register_object_store": ("nat.cli.register_workflow", "register_object_store"), + "register_optimizer": ("nat.cli.register_workflow", "register_optimizer"), + "register_optimizer_callback": ("nat.cli.register_workflow", "register_optimizer_callback"), + "register_per_user_function": ("nat.cli.register_workflow", "register_per_user_function"), + "register_per_user_function_group": ("nat.cli.register_workflow", "register_per_user_function_group"), + "register_registry_handler": ("nat.cli.register_workflow", "register_registry_handler"), + "register_retriever_client": ("nat.cli.register_workflow", "register_retriever_client"), + "register_retriever_provider": ("nat.cli.register_workflow", "register_retriever_provider"), + "register_telemetry_exporter": ("nat.cli.register_workflow", "register_telemetry_exporter"), + "register_tool_wrapper": ("nat.cli.register_workflow", "register_tool_wrapper"), + "register_trainer": ("nat.cli.register_workflow", "register_trainer"), + "register_trainer_adapter": ("nat.cli.register_workflow", "register_trainer_adapter"), + "register_trajectory_builder": ("nat.cli.register_workflow", "register_trajectory_builder"), + "register_ttc_strategy": ("nat.cli.register_workflow", "register_ttc_strategy"), + "set_secret_from_env": ("nat.data_models.common", "set_secret_from_env"), +} -def test_plugin_api_exports_core_function_authoring_surface(): - assert plugin_api.Builder is Builder - assert plugin_api.FunctionBaseConfig is FunctionBaseConfig - assert plugin_api.FunctionGroupBaseConfig is FunctionGroupBaseConfig - assert plugin_api.FunctionGroup is FunctionGroup - assert plugin_api.register_function is register_function - assert plugin_api.register_function_group is register_function_group +def test_plugin_api_exports_public_contract(): + assert len(plugin_api.__all__) == len(set(plugin_api.__all__)) + assert set(plugin_api.__all__) == set(EXPECTED_PLUGIN_API_EXPORTS) -def test_plugin_api_all_exports_are_importable(): - for name in plugin_api.__all__: - assert getattr(plugin_api, name) is not None + for public_name, (module_name, source_name) in EXPECTED_PLUGIN_API_EXPORTS.items(): + source_module = importlib.import_module(module_name) + assert getattr(plugin_api, public_name) is getattr(source_module, source_name) def test_plugin_authoring_docs_prefer_public_api_imports(): @@ -49,18 +135,59 @@ def test_plugin_authoring_docs_prefer_public_api_imports(): repo_root / "docs/source/improve-workflows/evaluate.md", repo_root / "docs/source/improve-workflows/optimizer.md", repo_root / "docs/source/improve-workflows/test-time-compute.md", + repo_root / "docs/source/run-workflows/existing-agents/langgraph.md", repo_root / "packages/nvidia_nat_core/src/nat/cli/commands/workflow/templates/workflow.py.j2", ] denied_patterns = [ "nat.cli.register_workflow", "nat.cli.register_llm_client", + "from nat.builder.dataset_loader import DatasetLoaderInfo", + "from nat.builder.embedder import EmbedderProviderInfo", + "from nat.builder.evaluator import EvaluatorInfo", + "from nat.builder.framework_enum import LLMFrameworkEnum", "from nat.builder.builder import Builder", "from nat.builder.builder import EvalBuilder", "from nat.builder.function import FunctionGroup", "from nat.builder.function_info import FunctionInfo", + "from nat.builder.llm import LLMProviderInfo", + "from nat.builder.retriever import RetrieverProviderInfo", + "from nat.data_models.authentication import AuthProviderBaseConfig", "from nat.data_models.component_ref import", + "from nat.data_models.dataset_handler import EvalDatasetBaseConfig", + "from nat.data_models.embedder import EmbedderBaseConfig", + "from nat.data_models.evaluator import EvaluatorBaseConfig", + "from nat.data_models.finetuning import TrainerAdapterConfig", + "from nat.data_models.finetuning import TrainerConfig", + "from nat.data_models.finetuning import TrajectoryBuilderConfig", + "from nat.data_models.front_end import FrontEndBaseConfig", "from nat.data_models.function import FunctionBaseConfig", "from nat.data_models.function import FunctionGroupBaseConfig", + "from nat.data_models.llm import LLMBaseConfig", + "from nat.data_models.logging import LoggingBaseConfig", + "from nat.data_models.memory import MemoryBaseConfig", + "from nat.data_models.middleware import FunctionMiddlewareBaseConfig", + "from nat.data_models.middleware import MiddlewareBaseConfig", + "from nat.data_models.object_store import KeyAlreadyExistsError", + "from nat.data_models.object_store import NoSuchKeyError", + "from nat.data_models.object_store import ObjectStoreBaseConfig", + "from nat.data_models.optimizer import OptimizerStrategyBaseConfig", + "from nat.data_models.optimizer import PromptOptimizationConfig", + "from nat.data_models.registry_handler import RegistryHandlerBaseConfig", + "from nat.data_models.retriever import RetrieverBaseConfig", + "from nat.data_models.telemetry_exporter import TelemetryExporterBaseConfig", + "from nat.data_models.ttc_strategy import TTCStrategyBaseConfig", + "from nat.memory.interfaces import MemoryEditor", + "from nat.memory.interfaces import MemoryManager", + "from nat.memory.interfaces import MemoryReader", + "from nat.memory.interfaces import MemoryWriter", + "from nat.memory.models import MemoryItem", + "from nat.middleware.dynamic.dynamic_function_middleware import DynamicFunctionMiddleware", + "from nat.middleware.dynamic.dynamic_middleware_config import DynamicMiddlewareConfig", + "from nat.middleware.function_middleware import FunctionMiddleware", + "from nat.middleware.middleware import FunctionMiddlewareContext", + "from nat.middleware.middleware import InvocationContext", + "from nat.object_store.interfaces import ObjectStore", + "from nat.object_store.models import ObjectStoreItem", ] files: list[Path] = [] From 5d50d652e10afb079d9585b03eee549e7a26e43b Mon Sep 17 00:00:00 2001 From: Bryan Bednarski Date: Mon, 18 May 2026 21:37:07 -0700 Subject: [PATCH 04/13] reduce set of public APIs being made available and add documentation on differed items Signed-off-by: Bryan Bednarski --- .../build-workflows/advanced/middleware.md | 2 +- .../extend/custom-components/finetuning.md | 16 +- .../extend/custom-components/optimizer.md | 18 +- docs/source/extend/plugin-api.md | 142 +++++++++++++++ docs/source/extend/plugins.md | 6 +- docs/source/improve-workflows/optimizer.md | 2 +- .../improve-workflows/test-time-compute.md | 2 +- .../nvidia_nat_core/src/nat/plugin_api.py | 44 ----- .../tests/nat/test_plugin_api.py | 163 ++++++++++++++---- 9 files changed, 295 insertions(+), 100 deletions(-) create mode 100644 docs/source/extend/plugin-api.md diff --git a/docs/source/build-workflows/advanced/middleware.md b/docs/source/build-workflows/advanced/middleware.md index 000bc09432..0805a335c6 100644 --- a/docs/source/build-workflows/advanced/middleware.md +++ b/docs/source/build-workflows/advanced/middleware.md @@ -84,7 +84,7 @@ Create a configuration class inheriting from `DynamicMiddlewareConfig`: ```python from pydantic import Field - from nat.plugin_api import DynamicMiddlewareConfig +from nat.plugin_api import DynamicMiddlewareConfig class LoggingMiddlewareConfig(DynamicMiddlewareConfig, name="logging_middleware"): diff --git a/docs/source/extend/custom-components/finetuning.md b/docs/source/extend/custom-components/finetuning.md index dd1ad6f518..24fa381c06 100644 --- a/docs/source/extend/custom-components/finetuning.md +++ b/docs/source/extend/custom-components/finetuning.md @@ -72,7 +72,7 @@ from typing import Any from nat.data_models.evaluate_runtime import EvaluationRunOutput from nat.data_models.evaluator import EvalOutputItem from nat.data_models.finetuning import FinetuneConfig, TrajectoryCollection -from nat.plugin_api import TrajectoryBuilderConfig +from nat.data_models.finetuning import TrajectoryBuilderConfig class TrajectoryBuilder(ABC): @@ -138,7 +138,7 @@ Create a configuration class that inherits from `TrajectoryBuilderConfig`: ```python from pydantic import Field -from nat.plugin_api import TrajectoryBuilderConfig +from nat.data_models.finetuning import TrajectoryBuilderConfig class MyTrajectoryBuilderConfig(TrajectoryBuilderConfig, name="my_traj_builder"): @@ -173,7 +173,7 @@ Create a registration module: ```python from nat.plugin_api import Builder -from nat.plugin_api import register_trajectory_builder +from nat.cli.register_workflow import register_trajectory_builder from .my_trajectory_builder import MyTrajectoryBuilder, MyTrajectoryBuilderConfig @@ -203,7 +203,7 @@ The `TrainerAdapter` bridges the gap between NeMo Agent Toolkit and external tra from abc import ABC, abstractmethod from typing import Any -from nat.plugin_api import TrainerAdapterConfig +from nat.data_models.finetuning import TrainerAdapterConfig from nat.data_models.finetuning import ( FinetuneConfig, TrainingJobRef, @@ -264,7 +264,7 @@ class TrainerAdapter(ABC): ```python from pydantic import BaseModel, Field -from nat.plugin_api import TrainerAdapterConfig +from nat.data_models.finetuning import TrainerAdapterConfig class MyBackendConfig(BaseModel): @@ -299,7 +299,7 @@ Implement the `TrainerAdapter` interface's methods. ```python from nat.plugin_api import Builder -from nat.plugin_api import register_trainer_adapter +from nat.cli.register_workflow import register_trainer_adapter from .my_trainer_adapter import MyTrainerAdapter, MyTrainerAdapterConfig @@ -320,7 +320,7 @@ The `Trainer` orchestrates the complete finetuning workflow, coordinating the tr from abc import ABC, abstractmethod from typing import Any -from nat.plugin_api import TrainerConfig +from nat.data_models.finetuning import TrainerConfig from nat.data_models.finetuning import ( FinetuneConfig, FinetuneRunConfig, @@ -413,7 +413,7 @@ Once you have your `MyTrainer` and `MyTrainerConfig` implemented, register it as ```python from nat.plugin_api import Builder -from nat.plugin_api import register_trainer +from nat.cli.register_workflow import register_trainer from .my_trainer import MyTrainer, MyTrainerConfig diff --git a/docs/source/extend/custom-components/optimizer.md b/docs/source/extend/custom-components/optimizer.md index 3a5a116e3c..970716ac6f 100644 --- a/docs/source/extend/custom-components/optimizer.md +++ b/docs/source/extend/custom-components/optimizer.md @@ -26,8 +26,8 @@ NeMo Agent Toolkit provides a pluggable optimizer system for tuning workflow par ## Key Interfaces * **Configuration Base Classes** - - {py:class}`~nat.plugin_api.OptimizerStrategyBaseConfig`: Base class that all optimizer strategy configuration models must extend. Provides an `enabled` field and integrates with the NeMo Agent Toolkit type registry. - - {py:class}`~nat.plugin_api.PromptOptimizationConfig`: Base for prompt optimization strategy configuration models. Adds `prompt_population_init_function` and `prompt_recombination_function` fields. + - {py:class}`~nat.data_models.optimizer.OptimizerStrategyBaseConfig`: Base class that all optimizer strategy configuration models must extend. Provides an `enabled` field and integrates with the NeMo Agent Toolkit type registry. + - {py:class}`~nat.data_models.optimizer.PromptOptimizationConfig`: Base for prompt optimization strategy configuration models. Adds `prompt_population_init_function` and `prompt_recombination_function` fields. - {py:class}`~nat.data_models.optimizer.OptunaParameterOptimizationConfig`: Built-in config for Optuna-based numeric parameter optimization. * **Optimizer ABCs** @@ -35,18 +35,18 @@ NeMo Agent Toolkit provides a pluggable optimizer system for tuning workflow par - {py:class}`~nat.plugins.config_optimizer.parameters.base.BaseParameterOptimizer`: Abstract base class for parameter optimization strategies. Requires implementing an async `run()` method that returns an optimized `Config`. * **Registration** - - {py:deco}`~nat.plugin_api.register_optimizer`: Decorator that registers an optimizer strategy with the global type registry so the optimizer runtime can resolve the strategy from the type of `cfg.optimizer.numeric` or `cfg.optimizer.prompt`. + - {py:deco}`~nat.cli.register_workflow.register_optimizer`: Decorator that registers an optimizer strategy with the global type registry so the optimizer runtime can resolve the strategy from the type of `cfg.optimizer.numeric` or `cfg.optimizer.prompt`. ## Adding a Custom Prompt Optimizer ### 1. Define a config class -Create a config class extending {py:class}`~nat.plugin_api.PromptOptimizationConfig` with a unique `name`: +Create a config class extending {py:class}`~nat.data_models.optimizer.PromptOptimizationConfig` with a unique `name`: ```python from pydantic import Field -from nat.plugin_api import PromptOptimizationConfig +from nat.data_models.optimizer import PromptOptimizationConfig class IterativeRefinementPromptConfig(PromptOptimizationConfig, name="iterative"): @@ -96,10 +96,10 @@ The `run()` method receives: ### 3. Register the Optimizer -Use the {py:deco}`~nat.plugin_api.register_optimizer` decorator to register your strategy: +Use the {py:deco}`~nat.cli.register_workflow.register_optimizer` decorator to register your strategy: ```python -from nat.plugin_api import register_optimizer +from nat.cli.register_workflow import register_optimizer @register_optimizer(config_type=IterativeRefinementPromptConfig) @@ -152,7 +152,7 @@ The pattern is the same, but parameter optimizers extend {py:class}`~nat.plugins ```python from pydantic import Field -from nat.plugin_api import OptimizerStrategyBaseConfig +from nat.data_models.optimizer import OptimizerStrategyBaseConfig class RandomSearchConfig(OptimizerStrategyBaseConfig, name="random_search"): @@ -194,7 +194,7 @@ class RandomSearchOptimizer(BaseParameterOptimizer): ### 3. Register and Configure ```python -from nat.plugin_api import register_optimizer +from nat.cli.register_workflow import register_optimizer @register_optimizer(config_type=RandomSearchConfig) diff --git a/docs/source/extend/plugin-api.md b/docs/source/extend/plugin-api.md new file mode 100644 index 0000000000..6c7798dc75 --- /dev/null +++ b/docs/source/extend/plugin-api.md @@ -0,0 +1,142 @@ + + +# Public Plugin API + +External plugin packages should import NeMo Agent Toolkit plugin-authoring APIs from `nat.plugin_api`. +This module is the stable public import surface for registering common plugin components and authoring functions or +function groups. + +```python +from nat.plugin_api import Builder +from nat.plugin_api import FunctionGroup +from nat.plugin_api import FunctionGroupBaseConfig +from nat.plugin_api import SerializableSecretStr +from nat.plugin_api import register_function_group +``` + +## Public Surface + +The following `nat.plugin_api` exports are intended for plugin authors: + +- Registration decorators for common external plugin components, such as `register_function`, + `register_function_group`, `register_llm_provider`, `register_embedder_provider`, `register_retriever_provider`, + `register_memory`, `register_object_store`, `register_middleware`, `register_telemetry_exporter`, and + `register_tool_wrapper`. +- Function authoring types, including `Builder`, `EvalBuilder`, `Function`, `FunctionInfo`, and `FunctionGroup`. +- Component configuration bases, including `FunctionBaseConfig`, `FunctionGroupBaseConfig`, `LLMBaseConfig`, + `EmbedderBaseConfig`, `RetrieverBaseConfig`, `MemoryBaseConfig`, `ObjectStoreBaseConfig`, + `MiddlewareBaseConfig`, `AuthProviderBaseConfig`, `EvaluatorBaseConfig`, `EvalDatasetBaseConfig`, and + `TelemetryExporterBaseConfig`. +- Registration return helpers, including `LLMProviderInfo`, `EmbedderProviderInfo`, `RetrieverProviderInfo`, + `EvaluatorInfo`, and `DatasetLoaderInfo`. +- Small implementation contracts needed by registered components, including `FunctionMiddleware`, + `DynamicFunctionMiddleware`, `MemoryEditor`, `ObjectStore`, and their associated context or value models. +- Component reference types, such as `FunctionRef`, `FunctionGroupRef`, `LLMRef`, `EmbedderRef`, `RetrieverRef`, + `MemoryRef`, `ObjectStoreRef`, `MiddlewareRef`, and `AuthenticationRef`. +- Framework wrapper identifiers, including `LLMFrameworkEnum`. +- Secret helpers, including `SerializableSecretStr`, `OptionalSecretStr`, `get_secret_value`, and `set_secret_from_env`. + +When a symbol is exported from `nat.plugin_api`, external packages can depend on that symbol's documented behavior across +minor and patch releases. Breaking changes to this public surface require a major release. + +Installed plugins execute as trusted Python code in the application environment. This public facade defines stable import +paths and authoring contracts; it does not make untrusted plugin packages safe to install or execute. + +The contract is intentionally explicit: adding a new public symbol requires adding it to `nat.plugin_api.__all__`, the +public API export test, and this documentation. Symbols that are not exported from `nat.plugin_api` should be treated as +implementation details unless a subsystem guide explicitly documents them as a specialized extension interface. Larger +subsystem-specific pipelines, such as telemetry processors or finetuning runtime interfaces, remain in their owning +modules until those contracts are promoted deliberately. + +## Surface Review + +The public facade is intentionally narrower than every component type supported by the runtime. The table below records +the current promotion decision for the major plugin-authoring surfaces. + +| Area | Public API status | Motivation | +| --- | --- | --- | +| Functions | Stable public | Core external plugin unit. Third-party tool and workflow packages need `register_function`, `FunctionBaseConfig`, `FunctionInfo`, and `Builder`. | +| Function groups | Stable public | Best fit for providers exposing multiple related tools. Supports external packages that share clients/resources and expose `group__function` names. | +| Builders | Stable public | Registered build functions receive a builder. Authors need a stable builder type without depending on `WorkflowBuilder`. | +| Config bases | Stable public | Public decorators require corresponding config base classes for typed YAML/discovery contracts. | +| Provider info objects | Stable public | LLM, embedder, retriever, dataset, and evaluator registrations yield these helper objects. | +| Component refs | Stable public | External configs need stable references to configured functions, LLMs, embedders, retrievers, memory, object stores, middleware, and auth providers. | +| Secrets | Stable public | External providers commonly need API keys and environment-backed secrets. Public helpers reduce raw-string credential patterns. | +| Registration decorators | Stable public | Decorators are the core plugin discovery and registration API. | +| LLM | Stable public | External LLM providers and framework clients are primary integration points. | +| Embedder | Stable public | External embedding providers and framework clients are expected provider plugins. | +| Retriever | Stable public | External retrieval providers and framework clients are expected provider plugins. | +| Evaluator and dataset loader | Stable public | Evaluation integrations and dataset loaders are documented plugin types with direct external authoring use cases. | +| Memory | Stable public, trusted plugin | External memory backends are documented integration points. They may handle user data, so plugins must be trusted. | +| Object store | Stable public, trusted plugin | External storage backends need the config base, object-store interface, item model, and standard errors. Plugins must be trusted. | +| Middleware | Stable public, trusted plugin | Middleware supports caching, policy, auth injection, redaction, and tracing. It can observe or alter calls, so plugins must be trusted. | +| Telemetry | Stable public, trusted plugin | External observability exporters are common integrations. They may receive traces or user data, so plugins must be trusted. | +| Auth provider | Stable public, trusted plugin | API integrations need auth providers. They handle credentials or tokens, so plugins must be trusted. | +| Front end | Deferred | Runtime hosting surfaces need a more explicit compatibility and security contract before being promoted through `nat.plugin_api`. | +| Logging | Deferred | External log sinks may exfiltrate sensitive logs. Keep the existing implementation API until the stable contract and trust guidance are clearer. | +| Registry handler | Deferred | Registry handlers influence component discovery/resolution. Keep out of the stable facade until that extension contract is reviewed. | +| Optimizer and optimizer callback | Deferred | Optimizer extension points are specialized subsystem APIs and are not required by common integration packages. | +| Trainer, trainer adapter, and trajectory builder | Deferred | Finetuning extensions are broad subsystem APIs. Keep them in owning modules until the finetuning compatibility contract is promoted deliberately. | +| TTC strategy | Deferred | Test-time compute is an advanced/experimental subsystem. Do not imply stable public facade support until that API matures. | + +Deferred surfaces remain available through their existing modules where those subsystem guides document them, but they are +not part of the stable `nat.plugin_api` facade. The deferred candidate list is also captured in +`packages/nvidia_nat_core/tests/nat/test_plugin_api.py` so CI can catch accidental promotion or stale candidate paths. + +## Private Implementation Modules + +Implementation modules remain importable for backward compatibility, but external plugin packages should not treat them +as stable API contracts. Prefer `nat.plugin_api` over direct imports from modules such as: + +- `nat.cli.register_workflow` +- `nat.cli.type_registry` +- `nat.builder.function` +- `nat.builder.function_base` +- `nat.builder.function_info` +- `nat.builder.workflow_builder` + +Concrete builders such as `WorkflowBuilder` are runtime implementation details. Plugin builders should type against +`Builder`, and plugin tests should use the utilities in `nvidia-nat-test` where possible. + +## Function Group Contract + +Function groups are the preferred pattern when one external service exposes multiple related tools. A function group: + +- Shares one configuration object across all functions in the group. +- Can share clients, connections, caches, and other resources. +- Exposes functions using `instance_name__function_name`. +- Supports `include` and `exclude` fields from `FunctionGroupBaseConfig` to control which functions are exposed through + workflow tool references. + +```python +class SearchConfig(FunctionGroupBaseConfig, name="search_provider"): + api_key: SerializableSecretStr + + +@register_function_group(config_type=SearchConfig) +async def build_search(config: SearchConfig, _builder: Builder): + group = FunctionGroup(config=config, instance_name="search") + + async def search(query: str) -> dict: + ... + + group.add_function("search", search, description=search.__doc__) + yield group +``` + +With this group configured as `search`, the function name is `search__search`. diff --git a/docs/source/extend/plugins.md b/docs/source/extend/plugins.md index 5b7b4175d8..820bd275b9 100644 --- a/docs/source/extend/plugins.md +++ b/docs/source/extend/plugins.md @@ -42,13 +42,13 @@ NeMo Agent Toolkit currently supports the following plugin types: - **Embedder Clients**: [Embedder](../build-workflows/embedders.md) Clients are implementations of embedder providers, which are specific to a [LLM](../build-workflows/llms/index.md) framework. For example, when using the OpenAI embedder provider with the LangChain/LangGraph framework, the LangChain/LangGraph OpenAI embedder client needs to be registered. To register an embedder client, you can use the {py:deco}`nat.plugin_api.register_embedder_client` decorator. - **Embedder Providers**: Embedder Providers are services that provide a way to embed text. For example, OpenAI and NVIDIA NIMs are embedder providers. To register an embedder provider, you can use the {py:deco}`nat.plugin_api.register_embedder_provider` decorator. - **Evaluators**: [Evaluators](../improve-workflows/evaluate.md) are used by the evaluation framework to evaluate the performance of NeMo Agent Toolkit workflows. To register an evaluator, you can use the {py:deco}`nat.plugin_api.register_evaluator` decorator. -- **Front Ends**: Front ends are the mechanism by which NeMo Agent Toolkit workflows are executed. Examples of front ends include a FastAPI server or a CLI. To register a front end, you can use the {py:deco}`nat.plugin_api.register_front_end` decorator. +- **Front Ends**: Front ends are the mechanism by which NeMo Agent Toolkit workflows are executed. Examples of front ends include a FastAPI server or a CLI. Front-end registration remains a specialized extension point and is not yet part of the stable `nat.plugin_api` facade. - **Functions**: [Functions](../build-workflows/functions-and-function-groups/functions.md) are one of the core building blocks of NeMo Agent Toolkit. They are used to define the tools and agents that can be used in a workflow. To register a function, you can use the {py:deco}`nat.plugin_api.register_function` decorator. - **LLM Clients**: LLM Clients are implementations of LLM providers that are specific to a LLM framework. For example, when using the NVIDIA NIMs LLM provider with the LangChain/LangGraph framework, the NVIDIA LangChain/LangGraph LLM client needs to be registered. To register an LLM client, you can use the {py:deco}`nat.plugin_api.register_llm_client` decorator. - **LLM Providers**: An LLM provider is a service that provides a way to interact with an LLM. For example, OpenAI and NVIDIA NIMs are LLM providers. To register an LLM provider, you can use the {py:deco}`nat.plugin_api.register_llm_provider` decorator. -- **Logging Methods**: Logging methods control the destination and format of log messages. To register a logging method, you can use the {py:deco}`nat.plugin_api.register_logging_method` decorator. +- **Logging Methods**: Logging methods control the destination and format of log messages. Logging method registration remains a specialized extension point and is not yet part of the stable `nat.plugin_api` facade. - **Memory**: [Memory](../build-workflows/memory.md) plugins are used to store and retrieve information from a database to be used by an LLM. Examples of memory plugins include Zep, Mem0 or MemMachine. To register a memory plugin, you can use the {py:deco}`nat.plugin_api.register_memory` decorator. -- **Registry Handlers**: Registry handlers are used to register custom agent registries with NeMo Agent Toolkit. An agent registry is a collection of tools, agents, and workflows that can be used in a workflow. To register a registry handler, you can use the {py:deco}`nat.plugin_api.register_registry_handler` decorator. +- **Registry Handlers**: Registry handlers are used to register custom agent registries with NeMo Agent Toolkit. An agent registry is a collection of tools, agents, and workflows that can be used in a workflow. Registry handler registration remains a specialized extension point and is not yet part of the stable `nat.plugin_api` facade. - **Retriever Clients**: [Retriever](../build-workflows/retrievers.md) clients are implementations of retriever providers, which are specific to a LLM framework. For example, when using the Milvus retriever provider with the LangChain/LangGraph framework, the LangChain/LangGraph Milvus retriever client needs to be registered. To register a retriever client, you can use the {py:deco}`nat.plugin_api.register_retriever_client` decorator. - **Retriever Providers**: Retriever providers are services that provide a way to retrieve information from a database. Examples of retriever providers include Chroma and Milvus. To register a retriever provider, you can use the {py:deco}`nat.plugin_api.register_retriever_provider` decorator. - **Telemetry Exporters**: [Telemetry exporters](../run-workflows/observe/observe.md) send telemetry data to a telemetry service. To register a telemetry exporter, you can use the {py:deco}`nat.plugin_api.register_telemetry_exporter` decorator. diff --git a/docs/source/improve-workflows/optimizer.md b/docs/source/improve-workflows/optimizer.md index 63cb2e9ace..bc84ec3dae 100644 --- a/docs/source/improve-workflows/optimizer.md +++ b/docs/source/improve-workflows/optimizer.md @@ -700,7 +700,7 @@ Callbacks are registered via the `@register_optimizer_callback(config_type=...)` For example, a provider registers its callback by decorating a factory function: ```python -from nat.plugin_api import register_optimizer_callback +from nat.cli.register_workflow import register_optimizer_callback @register_optimizer_callback(config_type=MyTelemetryExporter) def _build_my_optimizer_callback(config, *, dataset_name=None, **kwargs): diff --git a/docs/source/improve-workflows/test-time-compute.md b/docs/source/improve-workflows/test-time-compute.md index 068dd0947e..2c343d5bf1 100644 --- a/docs/source/improve-workflows/test-time-compute.md +++ b/docs/source/improve-workflows/test-time-compute.md @@ -141,7 +141,7 @@ Follow the steps below to create and register a new strategy. 3. Register the strategy. ```python - from nat.plugin_api import register_ttc_strategy + from nat.cli.register_workflow import register_ttc_strategy @register_ttc_strategy(config_type=MyStrategyConfig) async def register_my_strategy(cfg: MyStrategyConfig, builder: Builder): diff --git a/packages/nvidia_nat_core/src/nat/plugin_api.py b/packages/nvidia_nat_core/src/nat/plugin_api.py index 0bef4e0830..f617efad62 100644 --- a/packages/nvidia_nat_core/src/nat/plugin_api.py +++ b/packages/nvidia_nat_core/src/nat/plugin_api.py @@ -36,28 +36,19 @@ from nat.cli.register_workflow import register_embedder_provider from nat.cli.register_workflow import register_eval_callback from nat.cli.register_workflow import register_evaluator -from nat.cli.register_workflow import register_front_end from nat.cli.register_workflow import register_function from nat.cli.register_workflow import register_function_group from nat.cli.register_workflow import register_llm_client from nat.cli.register_workflow import register_llm_provider -from nat.cli.register_workflow import register_logging_method from nat.cli.register_workflow import register_memory from nat.cli.register_workflow import register_middleware from nat.cli.register_workflow import register_object_store -from nat.cli.register_workflow import register_optimizer -from nat.cli.register_workflow import register_optimizer_callback from nat.cli.register_workflow import register_per_user_function from nat.cli.register_workflow import register_per_user_function_group -from nat.cli.register_workflow import register_registry_handler from nat.cli.register_workflow import register_retriever_client from nat.cli.register_workflow import register_retriever_provider from nat.cli.register_workflow import register_telemetry_exporter from nat.cli.register_workflow import register_tool_wrapper -from nat.cli.register_workflow import register_trainer -from nat.cli.register_workflow import register_trainer_adapter -from nat.cli.register_workflow import register_trajectory_builder -from nat.cli.register_workflow import register_ttc_strategy from nat.data_models.authentication import AuthProviderBaseConfig from nat.data_models.common import OptionalSecretStr from nat.data_models.common import SerializableSecretStr @@ -73,33 +64,20 @@ from nat.data_models.component_ref import MiddlewareRef from nat.data_models.component_ref import ObjectStoreRef from nat.data_models.component_ref import RetrieverRef -from nat.data_models.component_ref import TrainerAdapterRef -from nat.data_models.component_ref import TrainerRef -from nat.data_models.component_ref import TrajectoryBuilderRef -from nat.data_models.component_ref import TTCStrategyRef from nat.data_models.dataset_handler import EvalDatasetBaseConfig from nat.data_models.embedder import EmbedderBaseConfig from nat.data_models.evaluator import EvaluatorBaseConfig -from nat.data_models.finetuning import TrainerAdapterConfig -from nat.data_models.finetuning import TrainerConfig -from nat.data_models.finetuning import TrajectoryBuilderConfig -from nat.data_models.front_end import FrontEndBaseConfig from nat.data_models.function import FunctionBaseConfig from nat.data_models.function import FunctionGroupBaseConfig from nat.data_models.llm import LLMBaseConfig -from nat.data_models.logging import LoggingBaseConfig from nat.data_models.memory import MemoryBaseConfig from nat.data_models.middleware import FunctionMiddlewareBaseConfig from nat.data_models.middleware import MiddlewareBaseConfig from nat.data_models.object_store import KeyAlreadyExistsError from nat.data_models.object_store import NoSuchKeyError from nat.data_models.object_store import ObjectStoreBaseConfig -from nat.data_models.optimizer import OptimizerStrategyBaseConfig -from nat.data_models.optimizer import PromptOptimizationConfig -from nat.data_models.registry_handler import RegistryHandlerBaseConfig from nat.data_models.retriever import RetrieverBaseConfig from nat.data_models.telemetry_exporter import TelemetryExporterBaseConfig -from nat.data_models.ttc_strategy import TTCStrategyBaseConfig from nat.memory.interfaces import MemoryEditor from nat.memory.interfaces import MemoryManager from nat.memory.interfaces import MemoryReader @@ -130,7 +108,6 @@ "EvalDatasetBaseConfig", "EvaluatorBaseConfig", "EvaluatorInfo", - "FrontEndBaseConfig", "Function", "FunctionBaseConfig", "FunctionGroup", @@ -147,7 +124,6 @@ "LLMFrameworkEnum", "LLMProviderInfo", "LLMRef", - "LoggingBaseConfig", "MemoryBaseConfig", "MemoryEditor", "MemoryItem", @@ -162,23 +138,12 @@ "ObjectStoreRef", "ObjectStoreItem", "ObjectStoreBaseConfig", - "OptimizerStrategyBaseConfig", "OptionalSecretStr", - "PromptOptimizationConfig", - "RegistryHandlerBaseConfig", "RetrieverBaseConfig", "RetrieverProviderInfo", "RetrieverRef", "SerializableSecretStr", "TelemetryExporterBaseConfig", - "TTCStrategyRef", - "TTCStrategyBaseConfig", - "TrainerAdapterConfig", - "TrainerAdapterRef", - "TrainerConfig", - "TrainerRef", - "TrajectoryBuilderConfig", - "TrajectoryBuilderRef", "get_secret_value", "register_auth_provider", "register_dataset_loader", @@ -186,27 +151,18 @@ "register_embedder_provider", "register_eval_callback", "register_evaluator", - "register_front_end", "register_function", "register_function_group", "register_llm_client", "register_llm_provider", - "register_logging_method", "register_memory", "register_middleware", "register_object_store", - "register_optimizer", - "register_optimizer_callback", "register_per_user_function", "register_per_user_function_group", - "register_registry_handler", "register_retriever_client", "register_retriever_provider", "register_telemetry_exporter", "register_tool_wrapper", - "register_trainer", - "register_trainer_adapter", - "register_trajectory_builder", - "register_ttc_strategy", "set_secret_from_env", ] diff --git a/packages/nvidia_nat_core/tests/nat/test_plugin_api.py b/packages/nvidia_nat_core/tests/nat/test_plugin_api.py index ae78b062e2..d52c168a46 100644 --- a/packages/nvidia_nat_core/tests/nat/test_plugin_api.py +++ b/packages/nvidia_nat_core/tests/nat/test_plugin_api.py @@ -14,6 +14,7 @@ # limitations under the License. import importlib +import re from pathlib import Path from nat import plugin_api @@ -33,7 +34,6 @@ "EvalDatasetBaseConfig": ("nat.data_models.dataset_handler", "EvalDatasetBaseConfig"), "EvaluatorBaseConfig": ("nat.data_models.evaluator", "EvaluatorBaseConfig"), "EvaluatorInfo": ("nat.builder.evaluator", "EvaluatorInfo"), - "FrontEndBaseConfig": ("nat.data_models.front_end", "FrontEndBaseConfig"), "Function": ("nat.builder.function", "Function"), "FunctionBaseConfig": ("nat.data_models.function", "FunctionBaseConfig"), "FunctionGroup": ("nat.builder.function", "FunctionGroup"), @@ -50,7 +50,6 @@ "LLMFrameworkEnum": ("nat.builder.framework_enum", "LLMFrameworkEnum"), "LLMProviderInfo": ("nat.builder.llm", "LLMProviderInfo"), "LLMRef": ("nat.data_models.component_ref", "LLMRef"), - "LoggingBaseConfig": ("nat.data_models.logging", "LoggingBaseConfig"), "MemoryBaseConfig": ("nat.data_models.memory", "MemoryBaseConfig"), "MemoryEditor": ("nat.memory.interfaces", "MemoryEditor"), "MemoryItem": ("nat.memory.models", "MemoryItem"), @@ -65,23 +64,12 @@ "ObjectStoreRef": ("nat.data_models.component_ref", "ObjectStoreRef"), "ObjectStoreItem": ("nat.object_store.models", "ObjectStoreItem"), "ObjectStoreBaseConfig": ("nat.data_models.object_store", "ObjectStoreBaseConfig"), - "OptimizerStrategyBaseConfig": ("nat.data_models.optimizer", "OptimizerStrategyBaseConfig"), "OptionalSecretStr": ("nat.data_models.common", "OptionalSecretStr"), - "PromptOptimizationConfig": ("nat.data_models.optimizer", "PromptOptimizationConfig"), - "RegistryHandlerBaseConfig": ("nat.data_models.registry_handler", "RegistryHandlerBaseConfig"), "RetrieverBaseConfig": ("nat.data_models.retriever", "RetrieverBaseConfig"), "RetrieverProviderInfo": ("nat.builder.retriever", "RetrieverProviderInfo"), "RetrieverRef": ("nat.data_models.component_ref", "RetrieverRef"), "SerializableSecretStr": ("nat.data_models.common", "SerializableSecretStr"), "TelemetryExporterBaseConfig": ("nat.data_models.telemetry_exporter", "TelemetryExporterBaseConfig"), - "TTCStrategyRef": ("nat.data_models.component_ref", "TTCStrategyRef"), - "TTCStrategyBaseConfig": ("nat.data_models.ttc_strategy", "TTCStrategyBaseConfig"), - "TrainerAdapterConfig": ("nat.data_models.finetuning", "TrainerAdapterConfig"), - "TrainerAdapterRef": ("nat.data_models.component_ref", "TrainerAdapterRef"), - "TrainerConfig": ("nat.data_models.finetuning", "TrainerConfig"), - "TrainerRef": ("nat.data_models.component_ref", "TrainerRef"), - "TrajectoryBuilderConfig": ("nat.data_models.finetuning", "TrajectoryBuilderConfig"), - "TrajectoryBuilderRef": ("nat.data_models.component_ref", "TrajectoryBuilderRef"), "get_secret_value": ("nat.data_models.common", "get_secret_value"), "register_auth_provider": ("nat.cli.register_workflow", "register_auth_provider"), "register_dataset_loader": ("nat.cli.register_workflow", "register_dataset_loader"), @@ -89,31 +77,113 @@ "register_embedder_provider": ("nat.cli.register_workflow", "register_embedder_provider"), "register_eval_callback": ("nat.cli.register_workflow", "register_eval_callback"), "register_evaluator": ("nat.cli.register_workflow", "register_evaluator"), - "register_front_end": ("nat.cli.register_workflow", "register_front_end"), "register_function": ("nat.cli.register_workflow", "register_function"), "register_function_group": ("nat.cli.register_workflow", "register_function_group"), "register_llm_client": ("nat.cli.register_workflow", "register_llm_client"), "register_llm_provider": ("nat.cli.register_workflow", "register_llm_provider"), - "register_logging_method": ("nat.cli.register_workflow", "register_logging_method"), "register_memory": ("nat.cli.register_workflow", "register_memory"), "register_middleware": ("nat.cli.register_workflow", "register_middleware"), "register_object_store": ("nat.cli.register_workflow", "register_object_store"), - "register_optimizer": ("nat.cli.register_workflow", "register_optimizer"), - "register_optimizer_callback": ("nat.cli.register_workflow", "register_optimizer_callback"), "register_per_user_function": ("nat.cli.register_workflow", "register_per_user_function"), "register_per_user_function_group": ("nat.cli.register_workflow", "register_per_user_function_group"), - "register_registry_handler": ("nat.cli.register_workflow", "register_registry_handler"), "register_retriever_client": ("nat.cli.register_workflow", "register_retriever_client"), "register_retriever_provider": ("nat.cli.register_workflow", "register_retriever_provider"), "register_telemetry_exporter": ("nat.cli.register_workflow", "register_telemetry_exporter"), "register_tool_wrapper": ("nat.cli.register_workflow", "register_tool_wrapper"), - "register_trainer": ("nat.cli.register_workflow", "register_trainer"), - "register_trainer_adapter": ("nat.cli.register_workflow", "register_trainer_adapter"), - "register_trajectory_builder": ("nat.cli.register_workflow", "register_trajectory_builder"), - "register_ttc_strategy": ("nat.cli.register_workflow", "register_ttc_strategy"), "set_secret_from_env": ("nat.data_models.common", "set_secret_from_env"), } +DEFERRED_PLUGIN_API_CANDIDATES = { + "FrontEndBaseConfig": { + "source": ("nat.data_models.front_end", "FrontEndBaseConfig"), + "reason": "runtime hosting surface; needs explicit compatibility and security contract", + }, + "LoggingBaseConfig": { + "source": ("nat.data_models.logging", "LoggingBaseConfig"), + "reason": "log sink surface; needs clearer trust guidance for sensitive logs", + }, + "OptimizerStrategyBaseConfig": { + "source": ("nat.data_models.optimizer", "OptimizerStrategyBaseConfig"), + "reason": "specialized optimizer subsystem API", + }, + "PromptOptimizationConfig": { + "source": ("nat.data_models.optimizer", "PromptOptimizationConfig"), + "reason": "specialized optimizer subsystem API", + }, + "RegistryHandlerBaseConfig": { + "source": ("nat.data_models.registry_handler", "RegistryHandlerBaseConfig"), + "reason": "registry resolution surface; needs extension-contract review", + }, + "TTCStrategyBaseConfig": { + "source": ("nat.data_models.ttc_strategy", "TTCStrategyBaseConfig"), + "reason": "advanced test-time compute subsystem API", + }, + "TTCStrategyRef": { + "source": ("nat.data_models.component_ref", "TTCStrategyRef"), + "reason": "advanced test-time compute subsystem API", + }, + "TrainerAdapterConfig": { + "source": ("nat.data_models.finetuning", "TrainerAdapterConfig"), + "reason": "specialized finetuning subsystem API", + }, + "TrainerAdapterRef": { + "source": ("nat.data_models.component_ref", "TrainerAdapterRef"), + "reason": "specialized finetuning subsystem API", + }, + "TrainerConfig": { + "source": ("nat.data_models.finetuning", "TrainerConfig"), + "reason": "specialized finetuning subsystem API", + }, + "TrainerRef": { + "source": ("nat.data_models.component_ref", "TrainerRef"), + "reason": "specialized finetuning subsystem API", + }, + "TrajectoryBuilderConfig": { + "source": ("nat.data_models.finetuning", "TrajectoryBuilderConfig"), + "reason": "specialized finetuning subsystem API", + }, + "TrajectoryBuilderRef": { + "source": ("nat.data_models.component_ref", "TrajectoryBuilderRef"), + "reason": "specialized finetuning subsystem API", + }, + "register_front_end": { + "source": ("nat.cli.register_workflow", "register_front_end"), + "reason": "runtime hosting surface; needs explicit compatibility and security contract", + }, + "register_logging_method": { + "source": ("nat.cli.register_workflow", "register_logging_method"), + "reason": "log sink surface; needs clearer trust guidance for sensitive logs", + }, + "register_optimizer": { + "source": ("nat.cli.register_workflow", "register_optimizer"), + "reason": "specialized optimizer subsystem API", + }, + "register_optimizer_callback": { + "source": ("nat.cli.register_workflow", "register_optimizer_callback"), + "reason": "specialized optimizer subsystem API", + }, + "register_registry_handler": { + "source": ("nat.cli.register_workflow", "register_registry_handler"), + "reason": "registry resolution surface; needs extension-contract review", + }, + "register_trainer": { + "source": ("nat.cli.register_workflow", "register_trainer"), + "reason": "specialized finetuning subsystem API", + }, + "register_trainer_adapter": { + "source": ("nat.cli.register_workflow", "register_trainer_adapter"), + "reason": "specialized finetuning subsystem API", + }, + "register_trajectory_builder": { + "source": ("nat.cli.register_workflow", "register_trajectory_builder"), + "reason": "specialized finetuning subsystem API", + }, + "register_ttc_strategy": { + "source": ("nat.cli.register_workflow", "register_ttc_strategy"), + "reason": "advanced test-time compute subsystem API", + }, +} + def test_plugin_api_exports_public_contract(): assert len(plugin_api.__all__) == len(set(plugin_api.__all__)) @@ -124,6 +194,17 @@ def test_plugin_api_exports_public_contract(): assert getattr(plugin_api, public_name) is getattr(source_module, source_name) +def test_deferred_plugin_api_candidates_remain_unpromoted(): + assert not (set(DEFERRED_PLUGIN_API_CANDIDATES) & set(EXPECTED_PLUGIN_API_EXPORTS)) + assert not (set(DEFERRED_PLUGIN_API_CANDIDATES) & set(plugin_api.__all__)) + + for candidate, metadata in DEFERRED_PLUGIN_API_CANDIDATES.items(): + module_name, source_name = metadata["source"] + source_module = importlib.import_module(module_name) + assert getattr(source_module, source_name) is not None, f"Deferred plugin API candidate {candidate} moved" + assert metadata["reason"] + + def test_plugin_authoring_docs_prefer_public_api_imports(): repo_root = Path(__file__).parents[4] paths = [ @@ -139,8 +220,25 @@ def test_plugin_authoring_docs_prefer_public_api_imports(): repo_root / "packages/nvidia_nat_core/src/nat/cli/commands/workflow/templates/workflow.py.j2", ] denied_patterns = [ - "nat.cli.register_workflow", - "nat.cli.register_llm_client", + "from nat.cli.register_workflow import register_auth_provider", + "from nat.cli.register_workflow import register_dataset_loader", + "from nat.cli.register_workflow import register_embedder_client", + "from nat.cli.register_workflow import register_embedder_provider", + "from nat.cli.register_workflow import register_eval_callback", + "from nat.cli.register_workflow import register_evaluator", + "from nat.cli.register_workflow import register_function", + "from nat.cli.register_workflow import register_function_group", + "from nat.cli.register_workflow import register_llm_client", + "from nat.cli.register_workflow import register_llm_provider", + "from nat.cli.register_workflow import register_memory", + "from nat.cli.register_workflow import register_middleware", + "from nat.cli.register_workflow import register_object_store", + "from nat.cli.register_workflow import register_per_user_function", + "from nat.cli.register_workflow import register_per_user_function_group", + "from nat.cli.register_workflow import register_retriever_client", + "from nat.cli.register_workflow import register_retriever_provider", + "from nat.cli.register_workflow import register_telemetry_exporter", + "from nat.cli.register_workflow import register_tool_wrapper", "from nat.builder.dataset_loader import DatasetLoaderInfo", "from nat.builder.embedder import EmbedderProviderInfo", "from nat.builder.evaluator import EvaluatorInfo", @@ -156,26 +254,17 @@ def test_plugin_authoring_docs_prefer_public_api_imports(): "from nat.data_models.dataset_handler import EvalDatasetBaseConfig", "from nat.data_models.embedder import EmbedderBaseConfig", "from nat.data_models.evaluator import EvaluatorBaseConfig", - "from nat.data_models.finetuning import TrainerAdapterConfig", - "from nat.data_models.finetuning import TrainerConfig", - "from nat.data_models.finetuning import TrajectoryBuilderConfig", - "from nat.data_models.front_end import FrontEndBaseConfig", "from nat.data_models.function import FunctionBaseConfig", "from nat.data_models.function import FunctionGroupBaseConfig", "from nat.data_models.llm import LLMBaseConfig", - "from nat.data_models.logging import LoggingBaseConfig", "from nat.data_models.memory import MemoryBaseConfig", "from nat.data_models.middleware import FunctionMiddlewareBaseConfig", "from nat.data_models.middleware import MiddlewareBaseConfig", "from nat.data_models.object_store import KeyAlreadyExistsError", "from nat.data_models.object_store import NoSuchKeyError", "from nat.data_models.object_store import ObjectStoreBaseConfig", - "from nat.data_models.optimizer import OptimizerStrategyBaseConfig", - "from nat.data_models.optimizer import PromptOptimizationConfig", - "from nat.data_models.registry_handler import RegistryHandlerBaseConfig", "from nat.data_models.retriever import RetrieverBaseConfig", "from nat.data_models.telemetry_exporter import TelemetryExporterBaseConfig", - "from nat.data_models.ttc_strategy import TTCStrategyBaseConfig", "from nat.memory.interfaces import MemoryEditor", "from nat.memory.interfaces import MemoryManager", "from nat.memory.interfaces import MemoryReader", @@ -198,10 +287,18 @@ def test_plugin_authoring_docs_prefer_public_api_imports(): files.append(path) violations = [] + public_imports = re.compile(r"from nat\.plugin_api import ([^\n]+)") for file_path in files: text = file_path.read_text(encoding="utf-8") for pattern in denied_patterns: if pattern in text: violations.append(f"{file_path.relative_to(repo_root)} contains {pattern!r}") + for match in public_imports.finditer(text): + imported_names = [name.strip().split(" as ")[0] for name in match.group(1).split(",")] + for imported_name in imported_names: + if imported_name and imported_name not in EXPECTED_PLUGIN_API_EXPORTS: + violations.append( + f"{file_path.relative_to(repo_root)} imports non-public nat.plugin_api symbol " + f"{imported_name!r}") assert not violations, "Plugin authoring docs should use nat.plugin_api:\n" + "\n".join(violations) From 30531470e2d64c436be9d0383d0289a21491842d Mon Sep 17 00:00:00 2001 From: Bryan Bednarski Date: Thu, 28 May 2026 11:46:28 -0700 Subject: [PATCH 05/13] docs: clarify plugin entry point compatibility policy Signed-off-by: Bryan Bednarski --- .../custom-functions/functions.md | 2 +- docs/source/extend/plugin-api.md | 4 +- docs/source/extend/plugins.md | 2 + .../nvidia_nat_core/src/nat/plugin_api.py | 14 +- .../tests/nat/test_plugin_api.py | 165 +++++++++++++++++- 5 files changed, 175 insertions(+), 12 deletions(-) diff --git a/docs/source/extend/custom-components/custom-functions/functions.md b/docs/source/extend/custom-components/custom-functions/functions.md index 8ae5fddd51..ed859a8245 100644 --- a/docs/source/extend/custom-components/custom-functions/functions.md +++ b/docs/source/extend/custom-components/custom-functions/functions.md @@ -68,7 +68,7 @@ Both of these methods will result in a function that can be used in the same way ### Function Configuration Object -To use a function from a configuration file, it must be registered with NeMo Agent Toolkit. Registering a function is done with the {py:deco}`nat.plugin_api.register_function` decorator. More information about registering components can be found in the [Plugin System](../../plugins.md) documentation. +To use a function from a configuration file, it must be registered with NVIDIA NeMo Agent Toolkit. Registering a function is done with the {py:deco}`nat.plugin_api.register_function` decorator. More information about registering components can be found in the [Plugin System](../../plugins.md) documentation. When registering a function, we first need to define the function configuration object. This object is used to configure the function and is passed to the function when it is invoked. Any options that are available to the function must be specified in the configuration object. diff --git a/docs/source/extend/plugin-api.md b/docs/source/extend/plugin-api.md index 6c7798dc75..6355aedb5f 100644 --- a/docs/source/extend/plugin-api.md +++ b/docs/source/extend/plugin-api.md @@ -15,9 +15,9 @@ See the License for the specific language governing permissions and limitations under the License. --> -# Public Plugin API +# NVIDIA NeMo Agent Toolkit — Public Plugin API -External plugin packages should import NeMo Agent Toolkit plugin-authoring APIs from `nat.plugin_api`. +NVIDIA NeMo Agent Toolkit external plugin packages should import plugin-authoring APIs from `nat.plugin_api`. This module is the stable public import surface for registering common plugin components and authoring functions or function groups. diff --git a/docs/source/extend/plugins.md b/docs/source/extend/plugins.md index 820bd275b9..fe4f097a7c 100644 --- a/docs/source/extend/plugins.md +++ b/docs/source/extend/plugins.md @@ -95,6 +95,8 @@ The `wrapper_type` argument can also be used with the library's `Builder` class Determining which plugins are available in a given environment is done through the use of [python entry points](https://packaging.python.org/en/latest/specifications/entry-points/). In NeMo Agent Toolkit, we scan the python environment for entry points which have the name `nat.plugins`. The value of the entry point is a python module that will be imported when the entry point is loaded. +New external plugin packages should use the `nat.plugins` entry point group. The runtime also continues to load `nat.components` entry points for backward compatibility with existing packages. + For example, the `nvidia-nat-langchain` distribution has the following entry point specified in the `pyproject.toml` file: ```toml diff --git a/packages/nvidia_nat_core/src/nat/plugin_api.py b/packages/nvidia_nat_core/src/nat/plugin_api.py index f617efad62..31aafa1fce 100644 --- a/packages/nvidia_nat_core/src/nat/plugin_api.py +++ b/packages/nvidia_nat_core/src/nat/plugin_api.py @@ -12,11 +12,23 @@ # 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. - """Stable public API for NeMo Agent Toolkit plugin authors. External plugin packages should prefer importing from this module instead of depending on implementation-oriented modules such as ``nat.cli.register_workflow`` or ``nat.builder.function`` directly. + +The ``Builder`` class is re-exported because plugin callables are typed around it, but only a subset of its methods +is part of the stable contract. The following ``Builder`` methods belong to subsystems that are intentionally +deferred from the public plugin API (see ``DEFERRED_PLUGIN_API_CANDIDATES`` in the plugin API tests) and may change +without notice; plugin authors must not depend on them: + +* Finetuning: ``add_trainer``, ``add_trainer_adapter``, ``add_trajectory_builder``, ``get_trainer``, + ``get_trainer_adapter``, ``get_trajectory_builder``, ``get_trainer_config``, ``get_trainer_adapter_config``, + ``get_trajectory_builder_config``. +* Test-time compute: ``add_ttc_strategy``, ``get_ttc_strategy``, ``get_ttc_strategy_config``. + +``test_builder_stable_surface_is_explicit`` pins the full method partition, so any new method added to ``Builder`` +must be explicitly categorized as stable or deferred. """ from nat.builder.builder import Builder diff --git a/packages/nvidia_nat_core/tests/nat/test_plugin_api.py b/packages/nvidia_nat_core/tests/nat/test_plugin_api.py index d52c168a46..fc4b7c0f8a 100644 --- a/packages/nvidia_nat_core/tests/nat/test_plugin_api.py +++ b/packages/nvidia_nat_core/tests/nat/test_plugin_api.py @@ -13,6 +13,7 @@ # See the License for the specific language governing permissions and # limitations under the License. +import ast import importlib import re from pathlib import Path @@ -184,6 +185,73 @@ }, } +# Stable subset of ``Builder``'s public method surface that plugin authors may rely on. +# Update this set together with ``Builder`` when promoting or deprecating plugin-authoring methods. +STABLE_BUILDER_METHODS = { + "add_auth_provider", + "add_embedder", + "add_function", + "add_function_group", + "add_llm", + "add_memory_client", + "add_middleware", + "add_object_store", + "add_retriever", + "current", + "get_auth_provider", + "get_auth_providers", + "get_embedder", + "get_embedder_config", + "get_embedders", + "get_function", + "get_function_config", + "get_function_dependencies", + "get_function_group", + "get_function_group_config", + "get_function_group_dependencies", + "get_function_groups", + "get_functions", + "get_llm", + "get_llm_config", + "get_llms", + "get_memory_client", + "get_memory_client_config", + "get_memory_clients", + "get_middleware", + "get_middleware_config", + "get_middleware_list", + "get_object_store_client", + "get_object_store_clients", + "get_object_store_config", + "get_retriever", + "get_retriever_config", + "get_retrievers", + "get_tool", + "get_tools", + "get_workflow", + "get_workflow_config", + "set_workflow", + "sync_builder", +} + +# ``Builder`` methods that belong to subsystems intentionally deferred from the public plugin API +# (mirrors the finetuning and test-time-compute entries in ``DEFERRED_PLUGIN_API_CANDIDATES``). +# Plugin authors must not depend on these even though they are reachable via the re-exported ``Builder``. +DEFERRED_BUILDER_METHODS = { + "add_trainer", + "add_trainer_adapter", + "add_trajectory_builder", + "add_ttc_strategy", + "get_trainer", + "get_trainer_adapter", + "get_trainer_adapter_config", + "get_trainer_config", + "get_trajectory_builder", + "get_trajectory_builder_config", + "get_ttc_strategy", + "get_ttc_strategy_config", +} + def test_plugin_api_exports_public_contract(): assert len(plugin_api.__all__) == len(set(plugin_api.__all__)) @@ -287,18 +355,99 @@ def test_plugin_authoring_docs_prefer_public_api_imports(): files.append(path) violations = [] - public_imports = re.compile(r"from nat\.plugin_api import ([^\n]+)") + md_code_block = re.compile(r"```(?:python|py)?\s*\n(.*?)```", re.DOTALL) + jinja_var = re.compile(r"\{\{\s*([\w.]+)\s*\}\}") for file_path in files: text = file_path.read_text(encoding="utf-8") for pattern in denied_patterns: if pattern in text: violations.append(f"{file_path.relative_to(repo_root)} contains {pattern!r}") - for match in public_imports.finditer(text): - imported_names = [name.strip().split(" as ")[0] for name in match.group(1).split(",")] - for imported_name in imported_names: - if imported_name and imported_name not in EXPECTED_PLUGIN_API_EXPORTS: - violations.append( - f"{file_path.relative_to(repo_root)} imports non-public nat.plugin_api symbol " - f"{imported_name!r}") + + if file_path.suffix == ".md": + snippets = md_code_block.findall(text) + elif file_path.name.endswith(".j2"): + snippets = [jinja_var.sub(r"\1", text)] + else: + snippets = [text] + + for snippet in snippets: + try: + tree = ast.parse(snippet) + except SyntaxError: + continue + for node in ast.walk(tree): + if isinstance(node, ast.ImportFrom) and node.module == "nat.plugin_api": + for alias in node.names: + if alias.name not in EXPECTED_PLUGIN_API_EXPORTS: + violations.append( + f"{file_path.relative_to(repo_root)} imports non-public nat.plugin_api symbol " + f"{alias.name!r}") assert not violations, "Plugin authoring docs should use nat.plugin_api:\n" + "\n".join(violations) + + +def test_builder_stable_surface_is_explicit(): + """Pin ``Builder``'s public method surface so deferred subsystems do not silently leak into the plugin contract.""" + from nat.builder.builder import Builder + + actual = {name for name in vars(Builder) if not name.startswith("_")} + expected = STABLE_BUILDER_METHODS | DEFERRED_BUILDER_METHODS + + new_methods = actual - expected + missing_methods = expected - actual + assert not new_methods, ( + f"New public methods on Builder: {sorted(new_methods)}. Add each to STABLE_BUILDER_METHODS " + "or DEFERRED_BUILDER_METHODS depending on whether the new surface is part of the stable plugin contract.") + assert not missing_methods, ( + f"Documented Builder methods missing from class: {sorted(missing_methods)}. " + "Update STABLE_BUILDER_METHODS / DEFERRED_BUILDER_METHODS.") + assert not (STABLE_BUILDER_METHODS & DEFERRED_BUILDER_METHODS), ( + "A Builder method appears in both STABLE and DEFERRED sets; pick one.") + + +def test_consumer_style_plugin_registration(): + """Exercise the external plugin authoring path using only ``nat.plugin_api`` imports. + + The snapshot tests above check that names are re-exported; this test verifies that the + public surface is actually sufficient to author and register a working plugin. + """ + # Imports scoped inside the test body to make the external authoring surface explicit. + from nat.plugin_api import Builder + from nat.plugin_api import FunctionBaseConfig + from nat.plugin_api import FunctionInfo + from nat.plugin_api import register_function + + class _ConsumerTestPluginConfig(FunctionBaseConfig, name="_consumer_test_plugin_api"): + prefix: str = "echo:" + + @register_function(config_type=_ConsumerTestPluginConfig) + async def _consumer_test_plugin_fn(config: _ConsumerTestPluginConfig, builder: Builder): + async def _run(text: str) -> str: + return f"{config.prefix} {text}" + + yield FunctionInfo.from_fn(_run, description="consumer-style plugin authoring test") + + from nat.cli.type_registry import GlobalTypeRegistry + registered = GlobalTypeRegistry.get().get_function(_ConsumerTestPluginConfig) + assert registered.config_type is _ConsumerTestPluginConfig + assert registered.build_fn is not None + + +def test_consumer_style_plugin_group_registration(): + """Same shape as ``test_consumer_style_plugin_registration`` for the function-group authoring path.""" + from nat.plugin_api import Builder + from nat.plugin_api import FunctionGroup + from nat.plugin_api import FunctionGroupBaseConfig + from nat.plugin_api import register_function_group + + class _ConsumerTestPluginGroupConfig(FunctionGroupBaseConfig, name="_consumer_test_plugin_api_group"): + prefix: str = "echo:" + + @register_function_group(config_type=_ConsumerTestPluginGroupConfig) + async def _consumer_test_plugin_group_fn(config: _ConsumerTestPluginGroupConfig, builder: Builder): + yield FunctionGroup(config=config) + + from nat.cli.type_registry import GlobalTypeRegistry + registered = GlobalTypeRegistry.get().get_function_group(_ConsumerTestPluginGroupConfig) + assert registered.config_type is _ConsumerTestPluginGroupConfig + assert registered.build_fn is not None From 2fd386c2e94846e74aaae3f3e25f1fdbca217bb6 Mon Sep 17 00:00:00 2001 From: Bryan Bednarski Date: Thu, 28 May 2026 12:11:46 -0700 Subject: [PATCH 06/13] update docs with 3P contributor guidelines and sync Signed-off-by: Bryan Bednarski --- docs/source/components/sharing-components.md | 5 +- .../adding-an-authentication-provider.md | 2 +- .../adding-an-llm-provider.md | 2 +- docs/source/extend/plugins.md | 3 + docs/source/extend/third-party-plugins.md | 302 ++++++++++++++++++ docs/source/index.md | 1 + .../tests/nat/test_plugin_api.py | 1 + 7 files changed, 313 insertions(+), 3 deletions(-) create mode 100644 docs/source/extend/third-party-plugins.md diff --git a/docs/source/components/sharing-components.md b/docs/source/components/sharing-components.md index 61b358fe7e..7d9e0f38e2 100644 --- a/docs/source/components/sharing-components.md +++ b/docs/source/components/sharing-components.md @@ -104,10 +104,13 @@ When building the `pyproject.toml` file, there are two critical sections: * Entrypoints: Provide the path to your plugins so they are registered with NeMo Agent Toolkit when installed. An example is provided below: ``` - [project.entry-points.'nat.components'] + [project.entry-points.'nat.plugins'] nat_notional_pkg_name = "nat_notional_pkg_name.register" ``` + The runtime continues to load `nat.components` entry points for backward compatibility, but new external packages + should use `nat.plugins`. + ### Building a Wheel Package After completing development and creating a `pyproject.toml` file that includes the necessary sections, the simplest diff --git a/docs/source/extend/custom-components/adding-an-authentication-provider.md b/docs/source/extend/custom-components/adding-an-authentication-provider.md index c5e764e16f..14d479915b 100644 --- a/docs/source/extend/custom-components/adding-an-authentication-provider.md +++ b/docs/source/extend/custom-components/adding-an-authentication-provider.md @@ -96,7 +96,7 @@ After implementing a new authentication provider, it’s important to verify tha ## Packaging the Provider The provider will need to be bundled into a Python package, which in turn will be registered with the toolkit as a [plugin](../plugins.md). In the `pyproject.toml` file of the package the -`project.entry-points.'nat.components'` section, defines a Python module as the entry point of the plugin. Details on how this is defined are found in the [Entry Point](../plugins.md#entry-point) section of the plugins document. By convention, the entry point module is named `register.py`, but this is not a requirement. +`project.entry-points.'nat.plugins'` section defines a Python module as the entry point of the plugin. Details on how this is defined are found in the [Entry Point](../plugins.md#entry-point) section of the plugins document. By convention, the entry point module is named `register.py`, but this is not a requirement. In the entry point module, the registration of provider, that is the function decorated with `register_auth_provider`, needs to be defined, either directly or imported from another module. A hypothetical `register.py` file could be defined as follows: diff --git a/docs/source/extend/custom-components/adding-an-llm-provider.md b/docs/source/extend/custom-components/adding-an-llm-provider.md index f1e0162477..a83dca6a02 100644 --- a/docs/source/extend/custom-components/adding-an-llm-provider.md +++ b/docs/source/extend/custom-components/adding-an-llm-provider.md @@ -221,7 +221,7 @@ Note: Since this test requires an API key, it's requesting the `nvidia_api_key` ## Packaging the Provider and Client -The provider and client will need to be bundled into a Python package, which in turn will be registered with NeMo Agent Toolkit as a [plugin](../plugins.md). In the `pyproject.toml` file of the package the `project.entry-points.'nat.components'` section, defines a Python module as the entry point of the plugin. Details on how this is defined are found in the [Entry Point](../plugins.md#entry-point) section of the plugins document. By convention, the entry point module is named `register.py`, but this is not a requirement. +The provider and client will need to be bundled into a Python package, which in turn will be registered with NeMo Agent Toolkit as a [plugin](../plugins.md). In the `pyproject.toml` file of the package the `project.entry-points.'nat.plugins'` section defines a Python module as the entry point of the plugin. Details on how this is defined are found in the [Entry Point](../plugins.md#entry-point) section of the plugins document. By convention, the entry point module is named `register.py`, but this is not a requirement. In the entry point module it is important that the provider is defined first followed by the client, this ensures that the provider is added to the NeMo Agent Toolkit registry before the client is registered. A hypothetical `register.py` file could be defined as follows: ```python diff --git a/docs/source/extend/plugins.md b/docs/source/extend/plugins.md index fe4f097a7c..71279a52b1 100644 --- a/docs/source/extend/plugins.md +++ b/docs/source/extend/plugins.md @@ -32,6 +32,9 @@ External plugin packages should import public plugin-authoring APIs from `nat.pl surface for decorators, function configuration bases, function groups, and common plugin helpers. See the [Public Plugin API](./plugin-api.md) documentation for the compatibility contract. +For guidance on partner-owned packages, repository layout, naming, testing, and documentation expectations, see +[Third-Party Plugin Packages](./third-party-plugins.md). + ## Supported Plugin Types diff --git a/docs/source/extend/third-party-plugins.md b/docs/source/extend/third-party-plugins.md new file mode 100644 index 0000000000..6dc0f4724c --- /dev/null +++ b/docs/source/extend/third-party-plugins.md @@ -0,0 +1,302 @@ + + +# Third-Party Plugin Packages + +NVIDIA NeMo Agent Toolkit supports plugin packages that are developed, released, and maintained outside of the main +NeMo Agent Toolkit repository. Third-party plugin packages use the same runtime discovery, configuration, and +observability paths as first-party packages. After a package is installed in the same Python environment as +`nvidia-nat-core`, the toolkit discovers it through Python entry points. + +This guide describes the recommended model for partner-owned plugin packages. For the stable Python import surface, see +the [Public Plugin API](./plugin-api.md). For general plugin discovery and supported plugin types, see the +[Plugin System](./plugins.md). + +## Ownership Model + +Third-party plugin packages are owned by the provider or partner that maintains the integration. The partner owns: + +- The GitHub repository, branch protection, contribution policy, and release process. +- Package publishing to PyPI or another package repository. +- Compatibility testing against supported NeMo Agent Toolkit versions. +- Issues and pull requests that affect only the partner integration. +- Documentation for installation, configuration, credentials, and bug routing. + +The NeMo Agent Toolkit project owns the stable public plugin API, first-party package behavior, and issues in +`nvidia-nat-core` that surface through third-party plugins. + +## When to Use a Third-Party Package + +Use a third-party-owned package for new partner integrations when the provider is best positioned to track its own API +roadmap, release cadence, and service behavior. This model is especially useful for: + +- Agentic tools and function groups. +- LLM, embedder, and retriever providers or framework clients. +- Telemetry exporters, memory backends, object stores, and authentication providers. +- Custom `nat` CLI subcommands. +- Specialized front-end integrations. + +New partner integrations should prefer this model over adding provider-specific code to the NeMo Agent Toolkit +monorepo. + +## Naming + +Use names that make the package discoverable and clearly associated with the toolkit: + +- GitHub repository: `/nemo-agent-toolkit-` +- PyPI package: `nemo-agent-toolkit-` +- Python import package: `nat.plugins.` +- Entry point name: `nat_` + +For example, a Tavily integration could use: + +- Repository: `tavily-ai/nemo-agent-toolkit-tavily` +- PyPI package: `nemo-agent-toolkit-tavily` +- Import package: `nat.plugins.tavily` +- Entry point: `nat_tavily` + +## Package Layout + +Third-party packages should use a PEP 420 namespace package layout compatible with other NeMo Agent Toolkit +distributions: + +```text +nemo-agent-toolkit-provider/ +|-- pyproject.toml +|-- README.md +|-- src/ +| `-- nat/ +| `-- plugins/ +| `-- provider/ +| |-- __init__.py +| |-- register.py +| `-- tools.py +`-- tests/ + `-- test_provider.py +``` + +Do not add `__init__.py` files at `src/nat/__init__.py` or `src/nat/plugins/__init__.py`. These directories are shared +namespace packages across NeMo Agent Toolkit distributions. The provider package itself, such as +`src/nat/plugins/provider`, should contain an `__init__.py`. + +The `register.py` module should import the modules that define plugin registration decorators so those decorators run +when the entry point is loaded. + +```python +from . import tools + +__all__ = ["tools"] +``` + +## Entry Points + +New external plugin packages should register component plugins under `nat.plugins`: + +```toml +[project.entry-points."nat.plugins"] +nat_provider = "nat.plugins.provider.register" +``` + +The runtime also continues to load `nat.components` entry points for backward compatibility with existing packages. +Do not use `nat.components` for new third-party plugin packages. + +Use these entry point groups for other extension points: + +| Entry point group | Use | +| --- | --- | +| `nat.plugins` | Functions, function groups, LLM clients, retrievers, embedders, telemetry exporters, memory backends, object stores, authentication providers, and other component plugins. | +| `nat.cli` | Custom `nat` CLI subcommands. | +| `nat.front_ends` | Specialized front-end implementations. Front-end registration is not part of the stable `nat.plugin_api` facade. | + +## Public API Imports + +Third-party plugin code should import stable plugin-authoring APIs from `nat.plugin_api`: + +```python +from nat.plugin_api import Builder +from nat.plugin_api import FunctionBaseConfig +from nat.plugin_api import FunctionInfo +from nat.plugin_api import register_function +``` + +Avoid depending on implementation modules such as `nat.cli.register_workflow`, `nat.builder.workflow_builder`, or +`nat.builder.function_info` unless a subsystem guide explicitly documents that module as the extension surface. Symbols +exported from `nat.plugin_api` are the stable public contract for external plugin packages. + +## Framework-Agnostic Tools + +Register tools with `register_function` or `register_function_group` unless the integration intentionally depends on a +specific framework. Framework-agnostic tools can be consumed through the toolkit's registered tool wrappers for +LangChain/LangGraph, CrewAI, LlamaIndex, AutoGen, Semantic Kernel, Google ADK, Agno, AWS Strands, and other supported +frameworks. + +```python +from pydantic import Field + +from nat.plugin_api import Builder +from nat.plugin_api import FunctionBaseConfig +from nat.plugin_api import FunctionInfo +from nat.plugin_api import register_function + + +class ProviderSearchConfig(FunctionBaseConfig, name="provider_search"): + """Search using the provider API.""" + + api_key: str = Field(description="Provider API key.") + + +@register_function(config_type=ProviderSearchConfig) +async def provider_search(config: ProviderSearchConfig, _builder: Builder): + async def search(query: str) -> str: + """Search for information using the provider API.""" + ... + + yield FunctionInfo.from_fn(search) +``` + +When possible, implement tools against the provider's own SDK or raw HTTP API rather than a framework-specific wrapper +package. Use framework-specific registration only when the integration cannot be expressed as a framework-agnostic +tool. + +## Dependencies + +Third-party packages should depend on `nvidia-nat-core` with a version range that allows compatible minor releases but +blocks unreviewed major releases: + +```toml +dependencies = [ + "nvidia-nat-core>=1.2,<2.0", +] +``` + +Choose the lower bound based on the first NeMo Agent Toolkit version that includes the public API symbols your package +uses. Update the upper bound when you have tested compatibility with a new major version. + +Use optional dependencies for provider SDK extras, development tools, and integration test dependencies when possible. + +## Installation and Discovery + +End users install third-party plugin packages directly into the same environment as `nvidia-nat-core`: + +```bash +uv add nemo-agent-toolkit-provider +``` + +or: + +```bash +pip install nemo-agent-toolkit-provider +``` + +After installation, the toolkit discovers the package through `importlib.metadata.entry_points()`. Users do not need to +edit a NeMo Agent Toolkit configuration file to load the package itself. They only need to reference the registered +component `_type` values in workflow configuration. + +Use `nat info components` to confirm that a package is installed and discoverable: + +```bash +uv run nat info components +``` + +## Development Workflow + +Use `uv` for local development when possible. This matches the primary NeMo Agent Toolkit development toolchain and +keeps lock files compatible with the toolkit's CI patterns. + +Common commands are: + +```bash +uv sync +uv lock +uv run pytest +``` + +## Required Documentation + +Each third-party plugin repository should include a README with: + +- Installation commands for `uv` and `pip`. +- A minimal workflow YAML example. +- Configuration schema and credential setup. +- Supported NeMo Agent Toolkit versions. +- Testing instructions. +- Bug routing guidance for provider-owned issues versus `nvidia-nat-core` issues. +- License information. + +## Testing + +At minimum, third-party plugin packages should include: + +- Unit tests for provider-specific logic. +- A loader smoke test that installs or imports the package, loads the entry point, and verifies that the registered + component can be discovered. +- At least one representative end-to-end test that invokes the plugin with a mock, stub, or local test service. +- Compatibility CI against supported NeMo Agent Toolkit versions. + +Provider integrations that require credentials should mark those tests as integration tests and skip them when the +required environment variables are not set. + +## Lifecycle + +Third-party plugin packages generally follow this lifecycle: + +1. Apply: Open an issue or pull request with the proposed package name, integration scope, license, and repository. +2. Develop: Build the package against the public plugin API and the package layout in this guide. +3. Review: Validate layout, license, naming, entry points, README, and smoke tests. +4. List: Add the approved package to NeMo Agent Toolkit documentation after the partner publishes it. +5. Verify or feature: Promote the package when it meets higher compatibility or announcement requirements. +6. Maintain: Keep compatibility CI passing and update the package as provider APIs or toolkit versions change. +7. Deprecate or archive: Remove inactive packages from toolkit documentation while leaving published package versions + in the partner's package repository. + +## Listing and Promotion + +The NeMo Agent Toolkit documentation may list third-party plugin packages that follow these guidelines. Plugin listings +can use the following lifecycle: + +| Tier | Requirements | Benefits | +| --- | --- | --- | +| Listed | Package layout, naming, license review, entry point registration, and one-time loader smoke check. | Listed in NeMo Agent Toolkit plugin documentation. | +| Verified | Listed requirements, partner-run compatibility CI for supported toolkit versions, and maintained README. | Eligible for coding assistant instructions and stronger discoverability. | +| Featured | Verified requirements plus a coordinated release or feature announcement. | Eligible for release-note placement and other promotion. | + +Inactive packages may be removed from NeMo Agent Toolkit documentation. Previously published package versions remain in +the partner's package repository. + +## Support Boundaries + +NVIDIA may provide design review, compatibility guidance, public API stability commitments, and documentation links for +approved third-party plugins. + +NVIDIA does not run partner CI, publish partner packages, accept liability for partner code, or guarantee feature parity +between third-party providers. Bugs in provider-owned integration code should be filed in the provider's repository. +Bugs in `nvidia-nat-core` should be filed in the NeMo Agent Toolkit repository. + +## Submission Checklist + +Before requesting inclusion in NeMo Agent Toolkit documentation, verify that the package has: + +- A repository and package name that follow the naming guidance. +- Source under `src/nat/plugins/`. +- No `__init__.py` files in shared namespace package directories. +- A `nat.plugins` entry point for component plugins. +- Imports from `nat.plugin_api` for public plugin-authoring APIs. +- A compatible `nvidia-nat-core` dependency range. +- An Apache-2.0 or approved permissive license. +- A README with install, configuration, workflow, testing, and bug routing instructions. +- Tests, including a loader smoke test. +- Compatibility CI for supported toolkit versions. diff --git a/docs/source/index.md b/docs/source/index.md index 3f7e005c49..d5d1bf3330 100644 --- a/docs/source/index.md +++ b/docs/source/index.md @@ -231,6 +231,7 @@ Sharing Components <./components/sharing-components.md> :caption: Extend Plugins <./extend/plugins.md> +Third-Party Plugin Packages <./extend/third-party-plugins.md> Plugin API <./extend/plugin-api.md> Custom Components <./extend/custom-components/index.md> ./extend/testing/index.md diff --git a/packages/nvidia_nat_core/tests/nat/test_plugin_api.py b/packages/nvidia_nat_core/tests/nat/test_plugin_api.py index fc4b7c0f8a..a746d95fa4 100644 --- a/packages/nvidia_nat_core/tests/nat/test_plugin_api.py +++ b/packages/nvidia_nat_core/tests/nat/test_plugin_api.py @@ -422,6 +422,7 @@ class _ConsumerTestPluginConfig(FunctionBaseConfig, name="_consumer_test_plugin_ @register_function(config_type=_ConsumerTestPluginConfig) async def _consumer_test_plugin_fn(config: _ConsumerTestPluginConfig, builder: Builder): + async def _run(text: str) -> str: return f"{config.prefix} {text}" From 434133c245fea45f5e07757d76e8bdc52a9a53d4 Mon Sep 17 00:00:00 2001 From: Bryan Bednarski Date: Thu, 28 May 2026 12:27:58 -0700 Subject: [PATCH 07/13] docs: resolve plugin API review feedback Signed-off-by: Bryan Bednarski --- .../custom-functions/functions.md | 2 +- docs/source/extend/plugin-api.md | 14 +++++++------- docs/source/extend/third-party-plugins.md | 12 +++++++----- 3 files changed, 15 insertions(+), 13 deletions(-) diff --git a/docs/source/extend/custom-components/custom-functions/functions.md b/docs/source/extend/custom-components/custom-functions/functions.md index ed859a8245..2733477674 100644 --- a/docs/source/extend/custom-components/custom-functions/functions.md +++ b/docs/source/extend/custom-components/custom-functions/functions.md @@ -194,7 +194,7 @@ With the configuration object defined, there are several options available to re ## Initialization and Cleanup -Its required to use an async context manager coroutine to register a function (it's not necessary to use `@asynccontextmanager`, since {py:deco}`nat.plugin_api.register_function` does this for you). This is because the function may need to execute some initialization before construction or cleanup after it is used. For example, if the function needs to load a model, connect to a resource, or download data, this can be done in the register function. +It's required to use an async context manager coroutine to register a function (it's not necessary to use `@asynccontextmanager`, since {py:deco}`nat.plugin_api.register_function` does this for you). This is because the function may need to execute some initialization before construction or cleanup after it is used. For example, if the function needs to load a model, connect to a resource, or download data, this can be done in the register function. ```python @register_function(config_type=MyFunctionConfig) diff --git a/docs/source/extend/plugin-api.md b/docs/source/extend/plugin-api.md index 6355aedb5f..0e62156a04 100644 --- a/docs/source/extend/plugin-api.md +++ b/docs/source/extend/plugin-api.md @@ -71,11 +71,11 @@ the current promotion decision for the major plugin-authoring surfaces. | Area | Public API status | Motivation | | --- | --- | --- | | Functions | Stable public | Core external plugin unit. Third-party tool and workflow packages need `register_function`, `FunctionBaseConfig`, `FunctionInfo`, and `Builder`. | -| Function groups | Stable public | Best fit for providers exposing multiple related tools. Supports external packages that share clients/resources and expose `group__function` names. | +| Function groups | Stable public | Best fit for providers exposing multiple related tools. Supports external packages that share clients and resources and expose `group__function` names. | | Builders | Stable public | Registered build functions receive a builder. Authors need a stable builder type without depending on `WorkflowBuilder`. | -| Config bases | Stable public | Public decorators require corresponding config base classes for typed YAML/discovery contracts. | +| Configuration bases | Stable public | Public decorators require corresponding configuration base classes for typed YAML and discovery contracts. | | Provider info objects | Stable public | LLM, embedder, retriever, dataset, and evaluator registrations yield these helper objects. | -| Component refs | Stable public | External configs need stable references to configured functions, LLMs, embedders, retrievers, memory, object stores, middleware, and auth providers. | +| Component refs | Stable public | External configurations need stable references to configured functions, LLMs, embedders, retrievers, memory, object stores, middleware, and auth providers. | | Secrets | Stable public | External providers commonly need API keys and environment-backed secrets. Public helpers reduce raw-string credential patterns. | | Registration decorators | Stable public | Decorators are the core plugin discovery and registration API. | | LLM | Stable public | External LLM providers and framework clients are primary integration points. | @@ -85,14 +85,14 @@ the current promotion decision for the major plugin-authoring surfaces. | Memory | Stable public, trusted plugin | External memory backends are documented integration points. They may handle user data, so plugins must be trusted. | | Object store | Stable public, trusted plugin | External storage backends need the config base, object-store interface, item model, and standard errors. Plugins must be trusted. | | Middleware | Stable public, trusted plugin | Middleware supports caching, policy, auth injection, redaction, and tracing. It can observe or alter calls, so plugins must be trusted. | -| Telemetry | Stable public, trusted plugin | External observability exporters are common integrations. They may receive traces or user data, so plugins must be trusted. | +| Telemetry registration | Stable public, trusted plugin | The facade covers telemetry exporter registration and configuration. Exporter runtime implementation APIs remain subsystem-specific until they are promoted deliberately. | | Auth provider | Stable public, trusted plugin | API integrations need auth providers. They handle credentials or tokens, so plugins must be trusted. | | Front end | Deferred | Runtime hosting surfaces need a more explicit compatibility and security contract before being promoted through `nat.plugin_api`. | -| Logging | Deferred | External log sinks may exfiltrate sensitive logs. Keep the existing implementation API until the stable contract and trust guidance are clearer. | -| Registry handler | Deferred | Registry handlers influence component discovery/resolution. Keep out of the stable facade until that extension contract is reviewed. | +| Logging | Deferred | External log sinks may export sensitive logs. Keep the existing implementation API until the stable contract and trust guidance are clearer. | +| Registry handler | Deferred | Registry handlers influence component discovery and resolution. Keep out of the stable facade until that extension contract is reviewed. | | Optimizer and optimizer callback | Deferred | Optimizer extension points are specialized subsystem APIs and are not required by common integration packages. | | Trainer, trainer adapter, and trajectory builder | Deferred | Finetuning extensions are broad subsystem APIs. Keep them in owning modules until the finetuning compatibility contract is promoted deliberately. | -| TTC strategy | Deferred | Test-time compute is an advanced/experimental subsystem. Do not imply stable public facade support until that API matures. | +| TTC strategy | Deferred | Test-time compute is an advanced and experimental subsystem. Do not imply stable public facade support until that API matures. | Deferred surfaces remain available through their existing modules where those subsystem guides document them, but they are not part of the stable `nat.plugin_api` facade. The deferred candidate list is also captured in diff --git a/docs/source/extend/third-party-plugins.md b/docs/source/extend/third-party-plugins.md index 6dc0f4724c..ab6ebb52c5 100644 --- a/docs/source/extend/third-party-plugins.md +++ b/docs/source/extend/third-party-plugins.md @@ -57,14 +57,16 @@ monorepo. Use names that make the package discoverable and clearly associated with the toolkit: -- GitHub repository: `/nemo-agent-toolkit-` +- GitHub repository owner: `` +- GitHub repository name: `nemo-agent-toolkit-` - PyPI package: `nemo-agent-toolkit-` - Python import package: `nat.plugins.` - Entry point name: `nat_` For example, a Tavily integration could use: -- Repository: `tavily-ai/nemo-agent-toolkit-tavily` +- Repository owner: `tavily-ai` +- Repository name: `nemo-agent-toolkit-tavily` - PyPI package: `nemo-agent-toolkit-tavily` - Import package: `nat.plugins.tavily` - Entry point: `nat_tavily` @@ -89,9 +91,9 @@ nemo-agent-toolkit-provider/ `-- test_provider.py ``` -Do not add `__init__.py` files at `src/nat/__init__.py` or `src/nat/plugins/__init__.py`. These directories are shared -namespace packages across NeMo Agent Toolkit distributions. The provider package itself, such as -`src/nat/plugins/provider`, should contain an `__init__.py`. +Do not add `__init__.py` files in the shared `nat` or `nat.plugins` namespace directories. These directories are shared +namespace packages across NeMo Agent Toolkit distributions. The provider package directory itself should contain an +`__init__.py`. The `register.py` module should import the modules that define plugin registration decorators so those decorators run when the entry point is loaded. From 1a8fa6d86fe65823a8f8bb47361627d1383159e4 Mon Sep 17 00:00:00 2001 From: Bryan Bednarski Date: Thu, 28 May 2026 12:44:56 -0700 Subject: [PATCH 08/13] (fix): stale references to legacy FunctionGroup.SEPARATOR Signed-off-by: Bryan Bednarski --- .../src/nat/middleware/defense/defense_middleware.py | 2 +- .../src/nat/middleware/red_teaming/red_teaming_middleware.py | 2 +- .../middleware/red_teaming/red_teaming_middleware_config.py | 4 ++-- .../tests/nat/builder/test_per_user_builder.py | 2 +- 4 files changed, 5 insertions(+), 5 deletions(-) diff --git a/packages/nvidia_nat_core/src/nat/middleware/defense/defense_middleware.py b/packages/nvidia_nat_core/src/nat/middleware/defense/defense_middleware.py index 4115939972..f6ca386b65 100644 --- a/packages/nvidia_nat_core/src/nat/middleware/defense/defense_middleware.py +++ b/packages/nvidia_nat_core/src/nat/middleware/defense/defense_middleware.py @@ -76,7 +76,7 @@ class DefenseMiddlewareConfig(FunctionMiddlewareBaseConfig): default=None, description="Optional function or function group to target. " "If None, defense applies to all functions. " - "Examples: 'my_calculator', 'my_calculator.divide', 'llm_agent.generate'") + "Examples: 'my_calculator', 'my_calculator__divide', 'llm_agent__generate'") target_location: TargetLocation = Field(default=TargetLocation.OUTPUT, description="Whether to analyze function input or output.") diff --git a/packages/nvidia_nat_core/src/nat/middleware/red_teaming/red_teaming_middleware.py b/packages/nvidia_nat_core/src/nat/middleware/red_teaming/red_teaming_middleware.py index e8ce6df7b8..5f72063425 100644 --- a/packages/nvidia_nat_core/src/nat/middleware/red_teaming/red_teaming_middleware.py +++ b/packages/nvidia_nat_core/src/nat/middleware/red_teaming/red_teaming_middleware.py @@ -67,7 +67,7 @@ class RedTeamingMiddleware(FunctionMiddleware): prompt_injection: _type: red_teaming attack_payload: "Ignore previous instructions" - target_function_or_group: my_llm.generate + target_function_or_group: my_llm__generate payload_placement: append_start target_location: input target_field: prompt diff --git a/packages/nvidia_nat_core/src/nat/middleware/red_teaming/red_teaming_middleware_config.py b/packages/nvidia_nat_core/src/nat/middleware/red_teaming/red_teaming_middleware_config.py index 8be146a040..c8064e40ca 100644 --- a/packages/nvidia_nat_core/src/nat/middleware/red_teaming/red_teaming_middleware_config.py +++ b/packages/nvidia_nat_core/src/nat/middleware/red_teaming/red_teaming_middleware_config.py @@ -43,7 +43,7 @@ class RedTeamingMiddlewareConfig(FunctionMiddlewareBaseConfig, name="red_teaming prompt_injection: _type: red_teaming attack_payload: "IGNORE ALL PREVIOUS INSTRUCTIONS" - target_function_or_group: my_llm.generate + target_function_or_group: my_llm__generate payload_placement: append_start target_location: input target_field: prompt @@ -67,7 +67,7 @@ class RedTeamingMiddlewareConfig(FunctionMiddlewareBaseConfig, name="red_teaming target_function_or_group: str | None = Field( default=None, description=("Optional function or group to target. " - "Format: 'group_name' for entire group, 'group_name.function_name' for specific function. " + "Format: 'group_name' for entire group, 'group_name__function_name' for specific function. " "If None, attacks all functions this middleware is applied to."), ) diff --git a/packages/nvidia_nat_core/tests/nat/builder/test_per_user_builder.py b/packages/nvidia_nat_core/tests/nat/builder/test_per_user_builder.py index 0d12df6260..9cfe8925c1 100644 --- a/packages/nvidia_nat_core/tests/nat/builder/test_per_user_builder.py +++ b/packages/nvidia_nat_core/tests/nat/builder/test_per_user_builder.py @@ -1088,7 +1088,7 @@ async def exposed_fn(inp: str) -> str: # Function group should be built assert "expose_fg" in per_user_builder._per_user_function_groups - # Exposed function should be accessible with prefixed name (group_name.function_name) + # Exposed function should be accessible with prefixed name (group_name__function_name) sep = FunctionGroup.SEPARATOR assert f"expose_fg{sep}exposed_tool" in per_user_builder._per_user_functions From 3a7851d4bb2cf490dd737578fe83efab16514eb4 Mon Sep 17 00:00:00 2001 From: Bryan Bednarski Date: Thu, 28 May 2026 13:48:31 -0700 Subject: [PATCH 09/13] docs: update 3P doc with mermaid diagram Signed-off-by: Bryan Bednarski --- docs/source/extend/plugins.md | 7 +- docs/source/extend/third-party-plugins.md | 412 ++++++++++++++-------- 2 files changed, 262 insertions(+), 157 deletions(-) diff --git a/docs/source/extend/plugins.md b/docs/source/extend/plugins.md index 71279a52b1..ef7d94c87f 100644 --- a/docs/source/extend/plugins.md +++ b/docs/source/extend/plugins.md @@ -96,9 +96,10 @@ The `wrapper_type` argument can also be used with the library's `Builder` class ### Entry Point -Determining which plugins are available in a given environment is done through the use of [python entry points](https://packaging.python.org/en/latest/specifications/entry-points/). In NeMo Agent Toolkit, we scan the python environment for entry points which have the name `nat.plugins`. The value of the entry point is a python module that will be imported when the entry point is loaded. - -New external plugin packages should use the `nat.plugins` entry point group. The runtime also continues to load `nat.components` entry points for backward compatibility with existing packages. +Determining which plugins are available in a given environment is done through the use of +[python entry points](https://packaging.python.org/en/latest/specifications/entry-points/). NeMo Agent Toolkit scans the +`nat.plugins` entry point group for plugin modules and also continues to load `nat.components` entry points for +backward compatibility with existing packages. New external plugin packages should use `nat.plugins`. For example, the `nvidia-nat-langchain` distribution has the following entry point specified in the `pyproject.toml` file: diff --git a/docs/source/extend/third-party-plugins.md b/docs/source/extend/third-party-plugins.md index ab6ebb52c5..97b110bef6 100644 --- a/docs/source/extend/third-party-plugins.md +++ b/docs/source/extend/third-party-plugins.md @@ -26,187 +26,274 @@ This guide describes the recommended model for partner-owned plugin packages. Fo the [Public Plugin API](./plugin-api.md). For general plugin discovery and supported plugin types, see the [Plugin System](./plugins.md). -## Ownership Model +## Ownership and Scope -Third-party plugin packages are owned by the provider or partner that maintains the integration. The partner owns: +Third-party plugin packages are provider-owned repositories. The provider owns the integration code, public package, +release process, compatibility testing, and user support for provider-specific behavior. The NeMo Agent Toolkit project +owns the stable plugin-authoring API, first-party package behavior, and issues in `nvidia-nat-core` that surface through +third-party packages. -- The GitHub repository, branch protection, contribution policy, and release process. -- Package publishing to PyPI or another package repository. -- Compatibility testing against supported NeMo Agent Toolkit versions. -- Issues and pull requests that affect only the partner integration. -- Documentation for installation, configuration, credentials, and bug routing. +This model is the default for new partner integrations where the provider is best positioned to track its own API +roadmap, service semantics, and release cadence. It works for function groups, tools, LLM clients, embedder clients, +retriever clients, telemetry exporters, memory backends, object stores, authentication providers, custom `nat` CLI +subcommands, and specialized front ends. -The NeMo Agent Toolkit project owns the stable public plugin API, first-party package behavior, and issues in -`nvidia-nat-core` that surface through third-party plugins. +Provider-specific behavior belongs in the provider repository. For example, a web search plugin can return the fields +and response shape exposed by the provider SDK. Do not introduce a shared web-search result schema unless the toolkit +defines that schema as a stable public API. -## When to Use a Third-Party Package +## Naming Convention -Use a third-party-owned package for new partner integrations when the provider is best positioned to track its own API -roadmap, release cadence, and service behavior. This model is especially useful for: +Use one provider token consistently across the repository, distribution, import package, entry point, registered +component `_type`, and function group namespace. The provider token should be short, lowercase, and stable. -- Agentic tools and function groups. -- LLM, embedder, and retriever providers or framework clients. -- Telemetry exporters, memory backends, object stores, and authentication providers. -- Custom `nat` CLI subcommands. -- Specialized front-end integrations. - -New partner integrations should prefer this model over adding provider-specific code to the NeMo Agent Toolkit -monorepo. - -## Naming - -Use names that make the package discoverable and clearly associated with the toolkit: - -- GitHub repository owner: `` -- GitHub repository name: `nemo-agent-toolkit-` -- PyPI package: `nemo-agent-toolkit-` -- Python import package: `nat.plugins.` -- Entry point name: `nat_` - -For example, a Tavily integration could use: +| Surface | Convention | Tavily example | +| --- | --- | --- | +| GitHub owner | Provider-owned account or organization | `tavily-ai` | +| GitHub repository | `NeMo-Agent-Toolkit-` | `NeMo-Agent-Toolkit-tavily` | +| Python distribution | `nemo-agent-toolkit-` | `nemo-agent-toolkit-tavily` | +| Python import package | `nat.plugins.` | `nat.plugins.tavily` | +| Component entry point name | `nat_` | `nat_tavily` | +| Registered function group `_type` | `` | `tavily` | +| Function group tool names | `__` | `tavily__search` | + +Function groups use a double underscore between the configured group instance name and each function name. This is the +current runtime convention in `FunctionGroup.SEPARATOR` and keeps tool names compatible with frameworks that reject or +reinterpret periods. For example, this configuration: + +```yaml +function_groups: + tavily: + _type: tavily +``` -- Repository owner: `tavily-ai` -- Repository name: `nemo-agent-toolkit-tavily` -- PyPI package: `nemo-agent-toolkit-tavily` -- Import package: `nat.plugins.tavily` -- Entry point: `nat_tavily` +exposes tools such as `tavily__search`, `tavily__extract`, and `tavily__research` when the plugin adds functions named +`search`, `extract`, and `research`. -## Package Layout +## Repository Layout -Third-party packages should use a PEP 420 namespace package layout compatible with other NeMo Agent Toolkit -distributions: +Use a PEP 420 namespace package layout so the provider package can share the `nat` namespace with other NeMo Agent +Toolkit distributions: ```text -nemo-agent-toolkit-provider/ +NeMo-Agent-Toolkit-tavily/ |-- pyproject.toml |-- README.md +|-- LICENSE |-- src/ | `-- nat/ | `-- plugins/ -| `-- provider/ +| `-- tavily/ | |-- __init__.py | |-- register.py | `-- tools.py `-- tests/ - `-- test_provider.py + `-- test_tools.py ``` -Do not add `__init__.py` files in the shared `nat` or `nat.plugins` namespace directories. These directories are shared -namespace packages across NeMo Agent Toolkit distributions. The provider package directory itself should contain an -`__init__.py`. +Do not add `__init__.py` files in the shared `nat` or `nat.plugins` namespace directories. The provider-owned package +directory, such as `tavily`, should contain an `__init__.py`. -The `register.py` module should import the modules that define plugin registration decorators so those decorators run -when the entry point is loaded. +The entry point target should import a registration module. That module should import the provider modules that define +registration decorators so the decorators run when the toolkit loads the entry point. ```python +# Registration module from . import tools __all__ = ["tools"] ``` -## Entry Points +## Package Metadata -New external plugin packages should register component plugins under `nat.plugins`: +The package should declare the shared namespace package, a bounded dependency on `nvidia-nat-core`, provider SDK +dependencies, optional test dependencies, repository metadata, and the component entry point. ```toml +[build-system] +requires = ["hatchling"] +build-backend = "hatchling.build" + +[tool.hatch.build.targets.wheel] +packages = ["src/nat"] + +[project] +name = "nemo-agent-toolkit-tavily" +version = "0.1.0" +requires-python = ">=3.11,<3.14" +description = "Tavily integration for NVIDIA NeMo Agent Toolkit" +readme = "README.md" +license = { text = "Apache-2.0" } +dependencies = [ + "nvidia-nat-core>=1.8,<2.0", + "tavily-python>=0.7.0,<1.0.0", +] + +[project.optional-dependencies] +test = [ + "pytest>=8.0", + "pytest-asyncio>=0.24", + "nvidia-nat-test>=1.8,<2.0", +] + +[project.urls] +documentation = "https://docs.nvidia.com/nemo/agent-toolkit/latest/" +source = "https://github.com/tavily-ai/NeMo-Agent-Toolkit-tavily" + [project.entry-points."nat.plugins"] -nat_provider = "nat.plugins.provider.register" +nat_tavily = "nat.plugins.tavily.register" ``` -The runtime also continues to load `nat.components` entry points for backward compatibility with existing packages. -Do not use `nat.components` for new third-party plugin packages. +New external component packages should use the `nat.plugins` entry point group. The runtime also loads +`nat.components` for backward compatibility with existing packages, but `nat.components` is compatibility-only for new +third-party packages. -Use these entry point groups for other extension points: +Other extension points use separate entry point groups: | Entry point group | Use | | --- | --- | -| `nat.plugins` | Functions, function groups, LLM clients, retrievers, embedders, telemetry exporters, memory backends, object stores, authentication providers, and other component plugins. | +| `nat.plugins` | Component plugins such as functions, function groups, model clients, retrievers, embedders, telemetry exporters, memory backends, object stores, middleware, and authentication providers. | | `nat.cli` | Custom `nat` CLI subcommands. | | `nat.front_ends` | Specialized front-end implementations. Front-end registration is not part of the stable `nat.plugin_api` facade. | -## Public API Imports +## Public API Surface -Third-party plugin code should import stable plugin-authoring APIs from `nat.plugin_api`: +Third-party packages should import stable plugin-authoring symbols from `nat.plugin_api`. ```python from nat.plugin_api import Builder -from nat.plugin_api import FunctionBaseConfig -from nat.plugin_api import FunctionInfo -from nat.plugin_api import register_function +from nat.plugin_api import FunctionGroup +from nat.plugin_api import FunctionGroupBaseConfig +from nat.plugin_api import SerializableSecretStr +from nat.plugin_api import register_function_group ``` -Avoid depending on implementation modules such as `nat.cli.register_workflow`, `nat.builder.workflow_builder`, or -`nat.builder.function_info` unless a subsystem guide explicitly documents that module as the extension surface. Symbols -exported from `nat.plugin_api` are the stable public contract for external plugin packages. +Avoid importing implementation modules such as `nat.cli.register_workflow`, `nat.builder.workflow_builder`, or +`nat.builder.function_info` unless another subsystem guide explicitly documents that module as the extension surface. +Symbols exported from `nat.plugin_api` are the public contract for external plugin packages. -## Framework-Agnostic Tools +## Function Group Implementation -Register tools with `register_function` or `register_function_group` unless the integration intentionally depends on a -specific framework. Framework-agnostic tools can be consumed through the toolkit's registered tool wrappers for -LangChain/LangGraph, CrewAI, LlamaIndex, AutoGen, Semantic Kernel, Google ADK, Agno, AWS Strands, and other supported -frameworks. +Use `register_function_group` when one provider exposes multiple related tools. A function group lets the integration +share configuration, credentials, clients, timeouts, and other resources while exposing individual tools through the +`instance_name__function_name` convention. ```python from pydantic import Field from nat.plugin_api import Builder -from nat.plugin_api import FunctionBaseConfig -from nat.plugin_api import FunctionInfo -from nat.plugin_api import register_function +from nat.plugin_api import FunctionGroup +from nat.plugin_api import FunctionGroupBaseConfig +from nat.plugin_api import SerializableSecretStr +from nat.plugin_api import register_function_group + +class TavilyToolsGroupConfig(FunctionGroupBaseConfig, name="tavily"): + """Tavily tools group.""" -class ProviderSearchConfig(FunctionBaseConfig, name="provider_search"): - """Search using the provider API.""" + api_key: SerializableSecretStr = Field( + default_factory=lambda: SerializableSecretStr(""), + description="Tavily API key. Falls back to the TAVILY_API_KEY environment variable.", + ) - api_key: str = Field(description="Provider API key.") +@register_function_group(config_type=TavilyToolsGroupConfig) +async def tavily_tools(config: TavilyToolsGroupConfig, _builder: Builder): + client = build_async_client(config.api_key) + group = FunctionGroup(config=config) -@register_function(config_type=ProviderSearchConfig) -async def provider_search(config: ProviderSearchConfig, _builder: Builder): - async def search(query: str) -> str: - """Search for information using the provider API.""" - ... + async def search(query: str) -> dict: + return await client.search(query=query) - yield FunctionInfo.from_fn(search) + async def extract(urls: list[str]) -> dict: + return await client.extract(urls=urls) + + group.add_function("search", search, description=search.__doc__) + group.add_function("extract", extract, description=extract.__doc__) + + yield group ``` -When possible, implement tools against the provider's own SDK or raw HTTP API rather than a framework-specific wrapper -package. Use framework-specific registration only when the integration cannot be expressed as a framework-agnostic -tool. +Use `register_function` instead when the integration exposes a single tool or workflow. Prefer provider SDKs or direct +HTTP clients over framework-specific wrappers when the tool can be expressed in a framework-agnostic way. Use +framework-specific registration only when the integration cannot be represented as a framework-agnostic toolkit tool. -## Dependencies +## README Requirements -Third-party packages should depend on `nvidia-nat-core` with a version range that allows compatible minor releases but -blocks unreviewed major releases: +Each provider repository should include a README that is complete enough for users and reviewers to install, configure, +test, and route bugs without reading the implementation. -```toml -dependencies = [ - "nvidia-nat-core>=1.2,<2.0", -] +At minimum, include: + +- Installation commands for `uv` and `pip`. +- A minimal workflow configuration. +- Configuration fields, defaults, and credential setup. +- The registered `_type` values and generated tool names. +- Supported NeMo Agent Toolkit versions. +- Local test commands. +- Bug routing for provider-owned integration bugs versus `nvidia-nat-core` bugs. +- License information. + +A minimal workflow example should be runnable from the repository root: + +```yaml +function_groups: + tavily: + _type: tavily + +llms: + my_llm: + _type: litellm + model_name: anthropic/claude-sonnet-4-6 + +workflow: + _type: react_agent + llm_name: my_llm + tool_names: + - tavily +``` + +```bash +export TAVILY_API_KEY=tvly-... +export ANTHROPIC_API_KEY=... + +uv run nat run --config_file config.yml --input "What changed in the latest release?" ``` -Choose the lower bound based on the first NeMo Agent Toolkit version that includes the public API symbols your package -uses. Update the upper bound when you have tested compatibility with a new major version. +## Testing and Compatibility + +At minimum, third-party plugin packages should include: + +- Unit tests for provider-specific logic. +- A loader smoke test that imports or installs the package, loads the entry point, and verifies that the registered + component can be discovered. +- Tool or function-group tests through `nvidia-nat-test` where possible. +- At least one representative end-to-end test that uses a mock, stub, or local test service. +- Compatibility CI against the supported NeMo Agent Toolkit versions for the plugin tier. + +Provider integrations that require live credentials should mark those tests as integration tests and skip them when the +required environment variables are not set. -Use optional dependencies for provider SDK extras, development tools, and integration test dependencies when possible. +Use a lower bound that matches the first NeMo Agent Toolkit release containing the `nat.plugin_api` symbols your package +uses. Keep the upper bound below the next major version until the package has been tested with that major version. ## Installation and Discovery -End users install third-party plugin packages directly into the same environment as `nvidia-nat-core`: +Users install a third-party plugin package into the same Python environment as `nvidia-nat-core`: ```bash -uv add nemo-agent-toolkit-provider +uv add nemo-agent-toolkit-tavily ``` or: ```bash -pip install nemo-agent-toolkit-provider +pip install nemo-agent-toolkit-tavily ``` -After installation, the toolkit discovers the package through `importlib.metadata.entry_points()`. Users do not need to -edit a NeMo Agent Toolkit configuration file to load the package itself. They only need to reference the registered -component `_type` values in workflow configuration. +After installation, the toolkit discovers the package with `importlib.metadata.entry_points()`. Users do not need to +edit a toolkit configuration file to load the package itself. They only reference the registered component `_type` +values in workflow configuration. Use `nat info components` to confirm that a package is installed and discoverable: @@ -219,65 +306,82 @@ uv run nat info components Use `uv` for local development when possible. This matches the primary NeMo Agent Toolkit development toolchain and keeps lock files compatible with the toolkit's CI patterns. -Common commands are: - ```bash -uv sync -uv lock -uv run pytest +git clone https://github.com/tavily-ai/NeMo-Agent-Toolkit-tavily.git +cd NeMo-Agent-Toolkit-tavily +uv sync --extra test +uv run pytest tests/ -v ``` -## Required Documentation - -Each third-party plugin repository should include a README with: - -- Installation commands for `uv` and `pip`. -- A minimal workflow YAML example. -- Configuration schema and credential setup. -- Supported NeMo Agent Toolkit versions. -- Testing instructions. -- Bug routing guidance for provider-owned issues versus `nvidia-nat-core` issues. -- License information. - -## Testing - -At minimum, third-party plugin packages should include: - -- Unit tests for provider-specific logic. -- A loader smoke test that installs or imports the package, loads the entry point, and verifies that the registered - component can be discovered. -- At least one representative end-to-end test that invokes the plugin with a mock, stub, or local test service. -- Compatibility CI against supported NeMo Agent Toolkit versions. - -Provider integrations that require credentials should mark those tests as integration tests and skip them when the -required environment variables are not set. - -## Lifecycle - -Third-party plugin packages generally follow this lifecycle: - -1. Apply: Open an issue or pull request with the proposed package name, integration scope, license, and repository. -2. Develop: Build the package against the public plugin API and the package layout in this guide. -3. Review: Validate layout, license, naming, entry points, README, and smoke tests. -4. List: Add the approved package to NeMo Agent Toolkit documentation after the partner publishes it. -5. Verify or feature: Promote the package when it meets higher compatibility or announcement requirements. -6. Maintain: Keep compatibility CI passing and update the package as provider APIs or toolkit versions change. -7. Deprecate or archive: Remove inactive packages from toolkit documentation while leaving published package versions - in the partner's package repository. +## Submission Review + +Open an issue or pull request in the NeMo Agent Toolkit repository before requesting documentation listing. Include the +provider name, package scope, license, repository URL, package name, entry point, registered `_type` values, and support +contacts. + +The review checks: + +- Repository and package names follow the naming convention. +- Package metadata uses `nat.plugins` for new component plugins. +- Source uses the shared `nat.plugins.` namespace package layout. +- Public imports come from `nat.plugin_api`. +- README covers installation, configuration, workflow examples, tests, and bug routing. +- License is Apache-2.0 or another approved permissive license. +- A smoke test proves the entry point can load and the registered component is discoverable. + +## Partner Plugin Lifecycle + +```mermaid +flowchart TB +subgraph Lifecycle["Partner Plugin Lifecycle"] +direction TB +Apply["1 Apply
Issue or PR in NeMo Agent Toolkit:
name, scope, license, repo"] +Develop["2 Develop
Build against template;
toolkit liaison available"] +Review["3 Submission review
Layout, license, naming,
entry points, smoke test"] +List["4 List
Added to NeMo Agent Toolkit plugin index;
partner publishes to PyPI"] +VerifyFeature["5 Verify or Feature
Upgrade tier
(see ladder below)"] +Maintain["6 Maintain
Keep compatibility CI green;
toolkit notifies of API changes"] +Deprecate["7 Deprecate or Archive
Inactive plugins archived;
removed from NeMo Agent Toolkit docs and instructions;
old versions stay on PyPI"] +Apply --> Develop --> Review --> List --> VerifyFeature --> Maintain +Maintain -->|"active"| Maintain +Maintain -->|"inactive"| Deprecate +end + +subgraph Ladder["Tier Ladder"] +direction TB +Listed["Listed
Requirements
• Package layout, naming, license
• Entry-point registration
• One-time loader smoke check
Benefits
• Documented in NeMo Agent Toolkit plugin docs"] +Verified["Verified
Requirements (Listed +)
• Partner-run compatibility CI
  (last 2 toolkit minor releases)
• Maintained README
Benefits
• Added to NeMo Agent Toolkit coding assistant instructions"] +Featured["Featured
Requirements (Verified +)
• Coordinated feature announcement
Benefits
• Featured plugin in release notes
• Eligible for additional promotion"] +Listed -->|"upgrade"| Verified -->|"upgrade"| Featured +end + +Lifecycle ~~~ Ladder +List -.->|"initial tier"| Listed +VerifyFeature -.->|"promote"| Verified +VerifyFeature -.->|"promote"| Featured + +classDef step fill:#e8f0fe,stroke:#1a73e8,stroke-width:1px,color:#000 +classDef warn fill:#fde7e9,stroke:#c5221f,stroke-width:1px,color:#000 +classDef tier fill:#fff,stroke:#444,stroke-width:1px,color:#000 +classDef wrapper fill:#f6f6f6,stroke:#888,stroke-width:1px,color:#000 +class Apply,Develop,Review,List,VerifyFeature,Maintain step +class Deprecate warn +class Listed,Verified,Featured tier +class Lifecycle,Ladder wrapper +``` -## Listing and Promotion +## Listing and Promotion Tiers -The NeMo Agent Toolkit documentation may list third-party plugin packages that follow these guidelines. Plugin listings -can use the following lifecycle: +The NeMo Agent Toolkit documentation may list third-party plugin packages that follow these guidelines. | Tier | Requirements | Benefits | | --- | --- | --- | | Listed | Package layout, naming, license review, entry point registration, and one-time loader smoke check. | Listed in NeMo Agent Toolkit plugin documentation. | -| Verified | Listed requirements, partner-run compatibility CI for supported toolkit versions, and maintained README. | Eligible for coding assistant instructions and stronger discoverability. | -| Featured | Verified requirements plus a coordinated release or feature announcement. | Eligible for release-note placement and other promotion. | +| Verified | Listed requirements plus partner-run compatibility CI for the last two toolkit minor releases and a maintained README. | Eligible for NeMo Agent Toolkit coding assistant instructions. | +| Featured | Verified requirements plus a coordinated feature announcement. | Eligible for release-note placement and additional promotion. | -Inactive packages may be removed from NeMo Agent Toolkit documentation. Previously published package versions remain in -the partner's package repository. +Inactive packages may be removed from NeMo Agent Toolkit documentation and coding assistant instructions. Previously +published package versions remain in the partner's package repository. ## Support Boundaries @@ -285,20 +389,20 @@ NVIDIA may provide design review, compatibility guidance, public API stability c approved third-party plugins. NVIDIA does not run partner CI, publish partner packages, accept liability for partner code, or guarantee feature parity -between third-party providers. Bugs in provider-owned integration code should be filed in the provider's repository. -Bugs in `nvidia-nat-core` should be filed in the NeMo Agent Toolkit repository. +between third-party providers. Bugs in provider-owned integration code should be filed in the provider repository. Bugs +in `nvidia-nat-core` should be filed in the NeMo Agent Toolkit repository. ## Submission Checklist Before requesting inclusion in NeMo Agent Toolkit documentation, verify that the package has: -- A repository and package name that follow the naming guidance. -- Source under `src/nat/plugins/`. +- A repository and package name that follow the naming convention. +- A PEP 420 namespace package under `nat.plugins.`. - No `__init__.py` files in shared namespace package directories. -- A `nat.plugins` entry point for component plugins. +- A `nat.plugins` entry point for new component plugins. - Imports from `nat.plugin_api` for public plugin-authoring APIs. - A compatible `nvidia-nat-core` dependency range. - An Apache-2.0 or approved permissive license. - A README with install, configuration, workflow, testing, and bug routing instructions. - Tests, including a loader smoke test. -- Compatibility CI for supported toolkit versions. +- Compatibility CI for the requested listing tier. From d95ee237ac984026388c7ae304ada3d392770cf3 Mon Sep 17 00:00:00 2001 From: Bryan Bednarski Date: Thu, 28 May 2026 14:50:13 -0700 Subject: [PATCH 10/13] Remove public API for auth for now, we are going to deffer this item. Add retriever, retoutput and document ro the plugin api Signed-off-by: Bryan Bednarski --- .../custom-components/adding-a-retriever.md | 41 ++++++++++++------- .../adding-an-authentication-provider.md | 8 ++-- .../custom-components/telemetry-exporters.md | 14 +++++++ docs/source/extend/plugin-api.md | 34 ++++++++------- docs/source/extend/plugins.md | 18 ++++---- docs/source/improve-workflows/evaluate.md | 6 +++ .../nvidia_nat_core/src/nat/plugin_api.py | 12 +++--- .../tests/nat/test_plugin_api.py | 35 ++++++++++++---- 8 files changed, 112 insertions(+), 56 deletions(-) diff --git a/docs/source/extend/custom-components/adding-a-retriever.md b/docs/source/extend/custom-components/adding-a-retriever.md index 02c2d57ab2..3885fd7e59 100644 --- a/docs/source/extend/custom-components/adding-a-retriever.md +++ b/docs/source/extend/custom-components/adding-a-retriever.md @@ -16,26 +16,28 @@ limitations under the License. --> # Adding a Retriever Provider -New [retrievers](../../build-workflows/retrievers.md) can be added to NeMo Agent Toolkit by creating a plugin. The general process is the same as for most plugins, but the retriever-specific steps are outlined here. +New [retrievers](../../build-workflows/retrievers.md) can be added to NVIDIA NeMo Agent Toolkit by creating a plugin. The general process is the same as for most plugins, but the retriever-specific steps are outlined here. First, create a retriever for the provider that implements the Retriever interface: ```python -class Retriever(ABC): - """ - Abstract interface for interacting with data stores. +from nat.plugin_api import Document +from nat.plugin_api import Retriever +from nat.plugin_api import RetrieverOutput - A Retriever is resposible for retrieving data from a configured data store. - Implemntations may integrate with vector stores or other indexing backends that allow for text-based search. - """ +class ExampleRetriever(Retriever): - @abstractmethod - async def search(self, query: str, **kwargs) -> RetrieverOutput: - """ - Retireve max(top_k) items from the data store based on vector similarity search (implementation dependent). + def __init__(self, client): + self._client = client - """ - raise NotImplementedError + async def search(self, query: str, **kwargs) -> RetrieverOutput: + result = await self._client.search(query=query, **kwargs) + return RetrieverOutput( + results=[ + Document(page_content=item.text, metadata=item.metadata, document_id=item.id) + for item in result.items + ] + ) ``` Next, create the config for the provider and register it with NeMo Agent Toolkit: @@ -45,12 +47,14 @@ from nat.plugin_api import Builder from nat.plugin_api import RetrieverBaseConfig from nat.plugin_api import RetrieverProviderInfo from nat.plugin_api import register_retriever_provider +from pydantic import Field +from pydantic import HttpUrl class ExampleRetrieverConfig(RetrieverBaseConfig, name="example_retriever"): """ Configuration for a Retriever provider. The parameters will depend on the particular provider. These are examples. """ - uri: HttpUrl = Field(description="The uri of the Nemo Retriever service.") + uri: HttpUrl = Field(description="The URI of the retriever service.") collection_name: str = Field(description="The name of the collection to search") top_k: int = Field(description="The number of results to return", gt=0, le=50, default=5) output_fields: list[str] | None = Field( @@ -71,9 +75,16 @@ from nat.plugin_api import register_retriever_client @register_retriever_client(config_type=ExampleRetrieverConfig, wrapper_type=None) async def nemo_retriever_client(config: ExampleRetrieverConfig, builder: Builder): + from example_plugin.client import ExampleClient from example_plugin.retriever import ExampleRetriever - retriever = ExampleRetriever(**config.model_dump()) + client = ExampleClient( + uri=str(config.uri), + collection_name=config.collection_name, + top_k=config.top_k, + output_fields=config.output_fields, + ) + retriever = ExampleRetriever(client=client) yield retriever ``` diff --git a/docs/source/extend/custom-components/adding-an-authentication-provider.md b/docs/source/extend/custom-components/adding-an-authentication-provider.md index 14d479915b..87f316890d 100644 --- a/docs/source/extend/custom-components/adding-an-authentication-provider.md +++ b/docs/source/extend/custom-components/adding-an-authentication-provider.md @@ -42,13 +42,13 @@ from the clients that facilitate the authentication process. Authentication prov ## Extending an API Authentication Provider The first step in adding an authentication provider is to create a configuration model that inherits from the -{py:class}`~nat.plugin_api.AuthProviderBaseConfig` class and define the credentials required to +{py:class}`~nat.data_models.authentication.AuthProviderBaseConfig` class and define the credentials required to authenticate with the target API resource. The following example shows how to define and register a custom evaluator and can be found here: {py:class}`~nat.authentication.oauth2.oauth2_auth_code_flow_provider_config.OAuth2AuthCodeFlowProviderConfig` class: ```python -from nat.plugin_api import AuthProviderBaseConfig +from nat.data_models.authentication import AuthProviderBaseConfig class OAuth2AuthCodeFlowProviderConfig(AuthProviderBaseConfig, name="oauth2_auth_code_flow"): @@ -72,13 +72,13 @@ class OAuth2AuthCodeFlowProviderConfig(AuthProviderBaseConfig, name="oauth2_auth ``` ### Registering the Provider -An asynchronous function decorated with {py:func}`~nat.plugin_api.register_auth_provider` is used to register the provider with NeMo Agent Toolkit by yielding an instance of +An asynchronous function decorated with {py:func}`~nat.cli.register_workflow.register_auth_provider` is used to register the provider with NeMo Agent Toolkit by yielding an instance of {py:class}`~nat.authentication.interfaces.AuthProviderBase`. The `OAuth2AuthCodeFlowProviderConfig` from the previous section is registered as follows: ```python from nat.plugin_api import Builder -from nat.plugin_api import register_auth_provider +from nat.cli.register_workflow import register_auth_provider @register_auth_provider(config_type=OAuth2AuthCodeFlowProviderConfig) async def oauth2_client(authentication_provider: OAuth2AuthCodeFlowProviderConfig, builder: Builder): diff --git a/docs/source/extend/custom-components/telemetry-exporters.md b/docs/source/extend/custom-components/telemetry-exporters.md index 59cdfca4b0..48cbe08af5 100644 --- a/docs/source/extend/custom-components/telemetry-exporters.md +++ b/docs/source/extend/custom-components/telemetry-exporters.md @@ -85,6 +85,14 @@ Examples of existing telemetry exporters include: Want to get started quickly? Here's a minimal working example that creates a console exporter to print traces to the terminal: +:::{important} +Telemetry exporter registration and configuration are available from the public `nat.plugin_api` facade. Exporter +implementation types such as `RawExporter`, `IntermediateStep`, span exporters, and processors are observability +subsystem APIs. They are documented here for telemetry exporter authors, but they are provisional and may evolve before +being promoted to the stable public plugin API. Telemetry plugins can observe workflow data and should only be installed +from trusted sources. +::: + ```python from pydantic import Field @@ -347,6 +355,12 @@ Start with the fields you need and add more as your integration becomes more sop Choose the appropriate base class based on your needs: +:::{note} +The exporter base classes and telemetry event models used in this section come from the observability subsystem, not +from `nat.plugin_api`. Treat them as subsystem-specific authoring APIs until the telemetry exporter implementation +contract is promoted deliberately. +::: + #### Raw Exporter (for simple trace exports) ```python diff --git a/docs/source/extend/plugin-api.md b/docs/source/extend/plugin-api.md index 0e62156a04..b1b97c7ae8 100644 --- a/docs/source/extend/plugin-api.md +++ b/docs/source/extend/plugin-api.md @@ -40,19 +40,22 @@ The following `nat.plugin_api` exports are intended for plugin authors: - Function authoring types, including `Builder`, `EvalBuilder`, `Function`, `FunctionInfo`, and `FunctionGroup`. - Component configuration bases, including `FunctionBaseConfig`, `FunctionGroupBaseConfig`, `LLMBaseConfig`, `EmbedderBaseConfig`, `RetrieverBaseConfig`, `MemoryBaseConfig`, `ObjectStoreBaseConfig`, - `MiddlewareBaseConfig`, `AuthProviderBaseConfig`, `EvaluatorBaseConfig`, `EvalDatasetBaseConfig`, and - `TelemetryExporterBaseConfig`. + `MiddlewareBaseConfig`, `FunctionMiddlewareBaseConfig`, `DynamicMiddlewareConfig`, `EvaluatorBaseConfig`, + `EvalDatasetBaseConfig`, and `TelemetryExporterBaseConfig`. - Registration return helpers, including `LLMProviderInfo`, `EmbedderProviderInfo`, `RetrieverProviderInfo`, `EvaluatorInfo`, and `DatasetLoaderInfo`. - Small implementation contracts needed by registered components, including `FunctionMiddleware`, - `DynamicFunctionMiddleware`, `MemoryEditor`, `ObjectStore`, and their associated context or value models. + `DynamicFunctionMiddleware`, `MemoryEditor`, `ObjectStore`, `Retriever`, `Document`, `RetrieverOutput`, and their + associated context or value models. - Component reference types, such as `FunctionRef`, `FunctionGroupRef`, `LLMRef`, `EmbedderRef`, `RetrieverRef`, - `MemoryRef`, `ObjectStoreRef`, `MiddlewareRef`, and `AuthenticationRef`. + `MemoryRef`, `ObjectStoreRef`, and `MiddlewareRef`. - Framework wrapper identifiers, including `LLMFrameworkEnum`. - Secret helpers, including `SerializableSecretStr`, `OptionalSecretStr`, `get_secret_value`, and `set_secret_from_env`. -When a symbol is exported from `nat.plugin_api`, external packages can depend on that symbol's documented behavior across -minor and patch releases. Breaking changes to this public surface require a major release. +When a symbol is exported from `nat.plugin_api` and marked stable in the surface review below, external packages can +depend on that symbol's documented behavior across minor and patch releases. Breaking changes to stable public surfaces +require a major release. Symbols marked provisional are available for plugin authors, but their compatibility contract is +still being refined and may evolve before being promoted to stable public status. Installed plugins execute as trusted Python code in the application environment. This public facade defines stable import paths and authoring contracts; it does not make untrusted plugin packages safe to install or execute. @@ -73,20 +76,21 @@ the current promotion decision for the major plugin-authoring surfaces. | Functions | Stable public | Core external plugin unit. Third-party tool and workflow packages need `register_function`, `FunctionBaseConfig`, `FunctionInfo`, and `Builder`. | | Function groups | Stable public | Best fit for providers exposing multiple related tools. Supports external packages that share clients and resources and expose `group__function` names. | | Builders | Stable public | Registered build functions receive a builder. Authors need a stable builder type without depending on `WorkflowBuilder`. | -| Configuration bases | Stable public | Public decorators require corresponding configuration base classes for typed YAML and discovery contracts. | +| Configuration bases | Stable public except where a component row below is provisional or deferred | Public decorators require corresponding configuration base classes for typed YAML and discovery contracts. A component's configuration base follows that component's support tier. | | Provider info objects | Stable public | LLM, embedder, retriever, dataset, and evaluator registrations yield these helper objects. | -| Component refs | Stable public | External configurations need stable references to configured functions, LLMs, embedders, retrievers, memory, object stores, middleware, and auth providers. | +| Component refs | Stable public | External configurations need stable references to configured functions, LLMs, embedders, retrievers, memory, object stores, and middleware. | | Secrets | Stable public | External providers commonly need API keys and environment-backed secrets. Public helpers reduce raw-string credential patterns. | -| Registration decorators | Stable public | Decorators are the core plugin discovery and registration API. | -| LLM | Stable public | External LLM providers and framework clients are primary integration points. | -| Embedder | Stable public | External embedding providers and framework clients are expected provider plugins. | -| Retriever | Stable public | External retrieval providers and framework clients are expected provider plugins. | -| Evaluator and dataset loader | Stable public | Evaluation integrations and dataset loaders are documented plugin types with direct external authoring use cases. | +| Registration decorators | Stable public except where a component row below is provisional or deferred | Decorators are the core plugin discovery and registration API. A component's registration decorator follows that component's support tier. | +| LLM registration and clients | Stable public | External LLM providers and framework clients are primary integration points. The stable facade covers registration, configuration, provider metadata, refs, and framework wrapper identifiers. Framework-native client runtime types and optional provider-specific config mixins remain framework or subsystem APIs unless exported from `nat.plugin_api`. | +| Embedder registration and clients | Stable public | External embedding providers and framework clients are expected provider plugins. The stable facade covers registration, configuration, provider metadata, refs, and wrapper selection. Framework-native client runtime types remain framework-specific. | +| Retriever | Stable public | External retrieval providers and framework clients are expected provider plugins. The stable facade includes retriever registration, configuration, provider metadata, refs, and the native retriever contract types `Retriever`, `RetrieverOutput`, and `Document`. | +| Evaluator and dataset loader registration | Stable public | Evaluation integrations and dataset loaders are documented plugin types with direct external authoring use cases. The stable facade covers registration, configuration, and info objects. Evaluator helper classes and ATIF-specific evaluator models remain subsystem-specific until promoted deliberately. | +| Evaluation callback registration | Provisional public | The facade exposes `register_eval_callback` for telemetry integrations that need evaluation lifecycle hooks. Callback protocol and result model types remain eval subsystem APIs and may evolve until they are promoted deliberately. | | Memory | Stable public, trusted plugin | External memory backends are documented integration points. They may handle user data, so plugins must be trusted. | | Object store | Stable public, trusted plugin | External storage backends need the config base, object-store interface, item model, and standard errors. Plugins must be trusted. | | Middleware | Stable public, trusted plugin | Middleware supports caching, policy, auth injection, redaction, and tracing. It can observe or alter calls, so plugins must be trusted. | -| Telemetry registration | Stable public, trusted plugin | The facade covers telemetry exporter registration and configuration. Exporter runtime implementation APIs remain subsystem-specific until they are promoted deliberately. | -| Auth provider | Stable public, trusted plugin | API integrations need auth providers. They handle credentials or tokens, so plugins must be trusted. | +| Telemetry registration | Provisional public, trusted plugin | The facade currently exposes telemetry exporter registration and configuration. Exporter implementation APIs, including raw exporters, span exporters, processors, and intermediate-step models, remain subsystem-specific and may evolve until they are promoted deliberately. Telemetry plugins can observe sensitive workflow data and must be trusted. | +| Auth provider | Deferred | Authentication provider authoring is still experimental and depends on subsystem APIs such as `AuthProviderBase`, `AuthProviderBaseConfig`, `AuthenticationRef`, and `register_auth_provider`. Keep it out of the stable facade until the auth compatibility and trust contract is promoted deliberately. | | Front end | Deferred | Runtime hosting surfaces need a more explicit compatibility and security contract before being promoted through `nat.plugin_api`. | | Logging | Deferred | External log sinks may export sensitive logs. Keep the existing implementation API until the stable contract and trust guidance are clearer. | | Registry handler | Deferred | Registry handlers influence component discovery and resolution. Keep out of the stable facade until that extension contract is reviewed. | diff --git a/docs/source/extend/plugins.md b/docs/source/extend/plugins.md index ef7d94c87f..9e8d9f23e1 100644 --- a/docs/source/extend/plugins.md +++ b/docs/source/extend/plugins.md @@ -56,7 +56,7 @@ NeMo Agent Toolkit currently supports the following plugin types: - **Retriever Providers**: Retriever providers are services that provide a way to retrieve information from a database. Examples of retriever providers include Chroma and Milvus. To register a retriever provider, you can use the {py:deco}`nat.plugin_api.register_retriever_provider` decorator. - **Telemetry Exporters**: [Telemetry exporters](../run-workflows/observe/observe.md) send telemetry data to a telemetry service. To register a telemetry exporter, you can use the {py:deco}`nat.plugin_api.register_telemetry_exporter` decorator. - **Tool Wrappers**: Tool wrappers are used to wrap functions in a way that is specific to a LLM framework. For example, when using the LangChain/LangGraph framework, NeMo Agent Toolkit functions need to be wrapped in `BaseTool` class to be compatible with LangChain/LangGraph. To register a tool wrapper, you can use the {py:deco}`nat.plugin_api.register_tool_wrapper` decorator. -- **API Authentication Providers**: [API authentication providers](../components/auth/api-authentication.md) are services that provide a way to authenticate requests to an API provider. Examples of authentication providers include OAuth 2.0 Authorization Code Grant and API Key. To register an API authentication provider, you can use the {py:deco}`nat.plugin_api.register_auth_provider` decorator. +- **API Authentication Providers**: [API authentication providers](../components/auth/api-authentication.md) are services that provide a way to authenticate requests to an API provider. Examples of authentication providers include OAuth 2.0 Authorization Code Grant and API Key. Authentication provider registration is experimental and remains a specialized extension point outside the stable `nat.plugin_api` facade. ## Anatomy of a Plugin @@ -101,14 +101,17 @@ Determining which plugins are available in a given environment is done through t `nat.plugins` entry point group for plugin modules and also continues to load `nat.components` entry points for backward compatibility with existing packages. New external plugin packages should use `nat.plugins`. -For example, the `nvidia-nat-langchain` distribution has the following entry point specified in the `pyproject.toml` file: +For example, a new external `nemo-agent-toolkit-my-provider` distribution could specify the following entry point in its +`pyproject.toml` file: ```toml [project.entry-points.'nat.plugins'] -nat_langchain = "nat.plugins.langchain.register" +nat_my_provider = "nat.plugins.my_provider.register" ``` -What this means is that when the `nvidia-nat-langchain` distribution is installed, the `nat.plugins.langchain.register` module will be imported when the entry point is loaded. This module must contain all the `@register_` decorators which need to be loaded when the library is initialized. +What this means is that when the `nemo-agent-toolkit-my-provider` distribution is installed, the +`nat.plugins.my_provider.register` module will be imported when the entry point is loaded. This module must contain all +the `@register_` decorators which need to be loaded when the library is initialized. :::{note} The above syntax in the `pyproject.toml` file is specific to [uv](https://docs.astral.sh/uv/concepts/projects/config/#plugin-entry-points). Other package managers may have a different syntax for specifying entry points. @@ -117,7 +120,8 @@ The above syntax in the `pyproject.toml` file is specific to [uv](https://docs.a #### Multiple Plugins in a Single Distribution -It is possible to have multiple plugins in a single distribution. For example, the `nvidia-nat-langchain` distribution contains both the LangChain/LangGraph LLM client and the LangChain/LangGraph embedder client. +It is possible to have multiple plugins in a single distribution. For example, a provider distribution could contain +both an LLM client and an embedder client. To register multiple plugins in a single distribution, there are two options: @@ -136,8 +140,8 @@ To register multiple plugins in a single distribution, there are two options: ```toml [project.entry-points.'nat.plugins'] - nat_langchain = "nat.plugins.langchain.register" - nat_langchain_tools = "nat.plugins.langchain.tools.register" + nat_my_provider = "nat.plugins.my_provider.register" + nat_my_provider_tools = "nat.plugins.my_provider.tools.register" ``` ### CLI Command Plugins diff --git a/docs/source/improve-workflows/evaluate.md b/docs/source/improve-workflows/evaluate.md index c432e4c912..2297bd7cbf 100644 --- a/docs/source/improve-workflows/evaluate.md +++ b/docs/source/improve-workflows/evaluate.md @@ -486,6 +486,12 @@ Note: Plotting metrics for individual dataset entries is only available across t The evaluation system provides a callback interface that allows observability providers to hook into the evaluation lifecycle. Callbacks enable providers to create structured experiments, link workflow runs to dataset examples, and attach evaluator scores in their respective platforms. +:::{note} +Evaluation callback registration is a provisional public plugin surface. The `register_eval_callback` decorator is +available from `nat.plugin_api`, but callback protocol and result model types remain eval subsystem APIs until that +runtime contract is promoted deliberately. +::: + ### `EvalCallback` Protocol Any class implementing the following methods can be registered as an evaluation callback: diff --git a/packages/nvidia_nat_core/src/nat/plugin_api.py b/packages/nvidia_nat_core/src/nat/plugin_api.py index 31aafa1fce..7c8da9b448 100644 --- a/packages/nvidia_nat_core/src/nat/plugin_api.py +++ b/packages/nvidia_nat_core/src/nat/plugin_api.py @@ -42,7 +42,6 @@ from nat.builder.function_info import FunctionInfo from nat.builder.llm import LLMProviderInfo from nat.builder.retriever import RetrieverProviderInfo -from nat.cli.register_workflow import register_auth_provider from nat.cli.register_workflow import register_dataset_loader from nat.cli.register_workflow import register_embedder_client from nat.cli.register_workflow import register_embedder_provider @@ -61,12 +60,10 @@ from nat.cli.register_workflow import register_retriever_provider from nat.cli.register_workflow import register_telemetry_exporter from nat.cli.register_workflow import register_tool_wrapper -from nat.data_models.authentication import AuthProviderBaseConfig from nat.data_models.common import OptionalSecretStr from nat.data_models.common import SerializableSecretStr from nat.data_models.common import get_secret_value from nat.data_models.common import set_secret_from_env -from nat.data_models.component_ref import AuthenticationRef from nat.data_models.component_ref import ComponentRef from nat.data_models.component_ref import EmbedderRef from nat.data_models.component_ref import FunctionGroupRef @@ -102,15 +99,17 @@ from nat.middleware.middleware import InvocationContext from nat.object_store.interfaces import ObjectStore from nat.object_store.models import ObjectStoreItem +from nat.retriever.interface import Retriever +from nat.retriever.models import Document +from nat.retriever.models import RetrieverOutput # Public contract: keep this list exact and update docs/source/extend/plugin-api.md plus # packages/nvidia_nat_core/tests/nat/test_plugin_api.py whenever symbols are added or removed. __all__ = [ - "AuthProviderBaseConfig", - "AuthenticationRef", "Builder", "ComponentRef", "DatasetLoaderInfo", + "Document", "DynamicFunctionMiddleware", "DynamicMiddlewareConfig", "EmbedderBaseConfig", @@ -151,13 +150,14 @@ "ObjectStoreItem", "ObjectStoreBaseConfig", "OptionalSecretStr", + "Retriever", "RetrieverBaseConfig", + "RetrieverOutput", "RetrieverProviderInfo", "RetrieverRef", "SerializableSecretStr", "TelemetryExporterBaseConfig", "get_secret_value", - "register_auth_provider", "register_dataset_loader", "register_embedder_client", "register_embedder_provider", diff --git a/packages/nvidia_nat_core/tests/nat/test_plugin_api.py b/packages/nvidia_nat_core/tests/nat/test_plugin_api.py index a746d95fa4..7f358ac901 100644 --- a/packages/nvidia_nat_core/tests/nat/test_plugin_api.py +++ b/packages/nvidia_nat_core/tests/nat/test_plugin_api.py @@ -21,11 +21,10 @@ from nat import plugin_api EXPECTED_PLUGIN_API_EXPORTS = { - "AuthProviderBaseConfig": ("nat.data_models.authentication", "AuthProviderBaseConfig"), - "AuthenticationRef": ("nat.data_models.component_ref", "AuthenticationRef"), "Builder": ("nat.builder.builder", "Builder"), "ComponentRef": ("nat.data_models.component_ref", "ComponentRef"), "DatasetLoaderInfo": ("nat.builder.dataset_loader", "DatasetLoaderInfo"), + "Document": ("nat.retriever.models", "Document"), "DynamicFunctionMiddleware": ("nat.middleware.dynamic.dynamic_function_middleware", "DynamicFunctionMiddleware"), "DynamicMiddlewareConfig": ("nat.middleware.dynamic.dynamic_middleware_config", "DynamicMiddlewareConfig"), "EmbedderBaseConfig": ("nat.data_models.embedder", "EmbedderBaseConfig"), @@ -66,13 +65,14 @@ "ObjectStoreItem": ("nat.object_store.models", "ObjectStoreItem"), "ObjectStoreBaseConfig": ("nat.data_models.object_store", "ObjectStoreBaseConfig"), "OptionalSecretStr": ("nat.data_models.common", "OptionalSecretStr"), + "Retriever": ("nat.retriever.interface", "Retriever"), "RetrieverBaseConfig": ("nat.data_models.retriever", "RetrieverBaseConfig"), + "RetrieverOutput": ("nat.retriever.models", "RetrieverOutput"), "RetrieverProviderInfo": ("nat.builder.retriever", "RetrieverProviderInfo"), "RetrieverRef": ("nat.data_models.component_ref", "RetrieverRef"), "SerializableSecretStr": ("nat.data_models.common", "SerializableSecretStr"), "TelemetryExporterBaseConfig": ("nat.data_models.telemetry_exporter", "TelemetryExporterBaseConfig"), "get_secret_value": ("nat.data_models.common", "get_secret_value"), - "register_auth_provider": ("nat.cli.register_workflow", "register_auth_provider"), "register_dataset_loader": ("nat.cli.register_workflow", "register_dataset_loader"), "register_embedder_client": ("nat.cli.register_workflow", "register_embedder_client"), "register_embedder_provider": ("nat.cli.register_workflow", "register_embedder_provider"), @@ -95,6 +95,18 @@ } DEFERRED_PLUGIN_API_CANDIDATES = { + "AuthenticationRef": { + "source": ("nat.data_models.component_ref", "AuthenticationRef"), + "reason": "authentication provider API is experimental and depends on subsystem interfaces", + }, + "AuthProviderBase": { + "source": ("nat.authentication.interfaces", "AuthProviderBase"), + "reason": "authentication provider API is experimental and depends on subsystem interfaces", + }, + "AuthProviderBaseConfig": { + "source": ("nat.data_models.authentication", "AuthProviderBaseConfig"), + "reason": "authentication provider API is experimental and depends on subsystem interfaces", + }, "FrontEndBaseConfig": { "source": ("nat.data_models.front_end", "FrontEndBaseConfig"), "reason": "runtime hosting surface; needs explicit compatibility and security contract", @@ -151,6 +163,10 @@ "source": ("nat.cli.register_workflow", "register_front_end"), "reason": "runtime hosting surface; needs explicit compatibility and security contract", }, + "register_auth_provider": { + "source": ("nat.cli.register_workflow", "register_auth_provider"), + "reason": "authentication provider API is experimental and depends on subsystem interfaces", + }, "register_logging_method": { "source": ("nat.cli.register_workflow", "register_logging_method"), "reason": "log sink surface; needs clearer trust guidance for sensitive logs", @@ -188,7 +204,6 @@ # Stable subset of ``Builder``'s public method surface that plugin authors may rely on. # Update this set together with ``Builder`` when promoting or deprecating plugin-authoring methods. STABLE_BUILDER_METHODS = { - "add_auth_provider", "add_embedder", "add_function", "add_function_group", @@ -198,8 +213,6 @@ "add_object_store", "add_retriever", "current", - "get_auth_provider", - "get_auth_providers", "get_embedder", "get_embedder_config", "get_embedders", @@ -235,13 +248,16 @@ } # ``Builder`` methods that belong to subsystems intentionally deferred from the public plugin API -# (mirrors the finetuning and test-time-compute entries in ``DEFERRED_PLUGIN_API_CANDIDATES``). +# (mirrors the auth, finetuning, and test-time-compute entries in ``DEFERRED_PLUGIN_API_CANDIDATES``). # Plugin authors must not depend on these even though they are reachable via the re-exported ``Builder``. DEFERRED_BUILDER_METHODS = { + "add_auth_provider", "add_trainer", "add_trainer_adapter", "add_trajectory_builder", "add_ttc_strategy", + "get_auth_provider", + "get_auth_providers", "get_trainer", "get_trainer_adapter", "get_trainer_adapter_config", @@ -288,7 +304,6 @@ def test_plugin_authoring_docs_prefer_public_api_imports(): repo_root / "packages/nvidia_nat_core/src/nat/cli/commands/workflow/templates/workflow.py.j2", ] denied_patterns = [ - "from nat.cli.register_workflow import register_auth_provider", "from nat.cli.register_workflow import register_dataset_loader", "from nat.cli.register_workflow import register_embedder_client", "from nat.cli.register_workflow import register_embedder_provider", @@ -317,7 +332,6 @@ def test_plugin_authoring_docs_prefer_public_api_imports(): "from nat.builder.function_info import FunctionInfo", "from nat.builder.llm import LLMProviderInfo", "from nat.builder.retriever import RetrieverProviderInfo", - "from nat.data_models.authentication import AuthProviderBaseConfig", "from nat.data_models.component_ref import", "from nat.data_models.dataset_handler import EvalDatasetBaseConfig", "from nat.data_models.embedder import EmbedderBaseConfig", @@ -345,6 +359,9 @@ def test_plugin_authoring_docs_prefer_public_api_imports(): "from nat.middleware.middleware import InvocationContext", "from nat.object_store.interfaces import ObjectStore", "from nat.object_store.models import ObjectStoreItem", + "from nat.retriever.interface import Retriever", + "from nat.retriever.models import Document", + "from nat.retriever.models import RetrieverOutput", ] files: list[Path] = [] From e27ea25aa2fc342209f5f0b4adcd10deff25759c Mon Sep 17 00:00:00 2001 From: Bryan Bednarski Date: Thu, 28 May 2026 14:56:35 -0700 Subject: [PATCH 11/13] - Added explicit Tool wrapper registration as Provisional public. - Split middleware into stable basic/function middleware vs provisional dynamic/runtime introspection. - Narrowed the builder row to the tested stable method subset. - Added evaluator-guide note that only registration/config/EvaluatorInfo are stable facade APIs. - Updated nat.plugin_api docstring to include deferred auth builder methods. Signed-off-by: Bryan Bednarski --- docs/source/extend/custom-components/custom-evaluator.md | 6 ++++++ docs/source/extend/plugin-api.md | 6 ++++-- docs/source/extend/plugins.md | 2 +- packages/nvidia_nat_core/src/nat/plugin_api.py | 1 + 4 files changed, 12 insertions(+), 3 deletions(-) diff --git a/docs/source/extend/custom-components/custom-evaluator.md b/docs/source/extend/custom-components/custom-evaluator.md index 99df98ee73..a59dc9b6ab 100644 --- a/docs/source/extend/custom-components/custom-evaluator.md +++ b/docs/source/extend/custom-components/custom-evaluator.md @@ -38,6 +38,12 @@ To extend NeMo Agent Toolkit with custom evaluators, you need to create an evalu This section provides a step-by-step guide to create and register a custom evaluator with NeMo Agent Toolkit. A similarity evaluator is used as an example to demonstrate the process. +:::{note} +Evaluator registration, configuration, and `EvaluatorInfo` are stable public plugin APIs available from +`nat.plugin_api`. Evaluator helper classes and ATIF-specific evaluator models remain eval subsystem APIs until they are +promoted deliberately. +::: + ### Evaluator Configuration The evaluator configuration defines the evaluator name and any evaluator-specific parameters. This configuration is paired with a registration function that yields an asynchronous evaluation method. diff --git a/docs/source/extend/plugin-api.md b/docs/source/extend/plugin-api.md index b1b97c7ae8..6ff340edda 100644 --- a/docs/source/extend/plugin-api.md +++ b/docs/source/extend/plugin-api.md @@ -75,7 +75,7 @@ the current promotion decision for the major plugin-authoring surfaces. | --- | --- | --- | | Functions | Stable public | Core external plugin unit. Third-party tool and workflow packages need `register_function`, `FunctionBaseConfig`, `FunctionInfo`, and `Builder`. | | Function groups | Stable public | Best fit for providers exposing multiple related tools. Supports external packages that share clients and resources and expose `group__function` names. | -| Builders | Stable public | Registered build functions receive a builder. Authors need a stable builder type without depending on `WorkflowBuilder`. | +| Builder type and common component access | Stable public | Registered build functions receive a builder. Authors need a stable builder type without depending on `WorkflowBuilder`. Only the builder methods categorized as stable in `packages/nvidia_nat_core/tests/nat/test_plugin_api.py` are part of the facade contract; deferred subsystem methods and concrete builders remain implementation details. | | Configuration bases | Stable public except where a component row below is provisional or deferred | Public decorators require corresponding configuration base classes for typed YAML and discovery contracts. A component's configuration base follows that component's support tier. | | Provider info objects | Stable public | LLM, embedder, retriever, dataset, and evaluator registrations yield these helper objects. | | Component refs | Stable public | External configurations need stable references to configured functions, LLMs, embedders, retrievers, memory, object stores, and middleware. | @@ -88,8 +88,10 @@ the current promotion decision for the major plugin-authoring surfaces. | Evaluation callback registration | Provisional public | The facade exposes `register_eval_callback` for telemetry integrations that need evaluation lifecycle hooks. Callback protocol and result model types remain eval subsystem APIs and may evolve until they are promoted deliberately. | | Memory | Stable public, trusted plugin | External memory backends are documented integration points. They may handle user data, so plugins must be trusted. | | Object store | Stable public, trusted plugin | External storage backends need the config base, object-store interface, item model, and standard errors. Plugins must be trusted. | -| Middleware | Stable public, trusted plugin | Middleware supports caching, policy, auth injection, redaction, and tracing. It can observe or alter calls, so plugins must be trusted. | +| Middleware registration and function middleware | Stable public, trusted plugin | Basic middleware registration and function middleware support caching, policy, auth injection, redaction, and tracing. Middleware can observe or alter calls, so plugins must be trusted. | +| Dynamic middleware and runtime introspection | Provisional public, trusted plugin | Dynamic middleware reaches deeper into runtime call interception and inventory metadata. The facade exposes the current dynamic middleware types, but inventory and unregister details remain subsystem-specific until the contract is promoted deliberately. | | Telemetry registration | Provisional public, trusted plugin | The facade currently exposes telemetry exporter registration and configuration. Exporter implementation APIs, including raw exporters, span exporters, processors, and intermediate-step models, remain subsystem-specific and may evolve until they are promoted deliberately. Telemetry plugins can observe sensitive workflow data and must be trusted. | +| Tool wrapper registration | Provisional public | The facade exposes `register_tool_wrapper` for framework integrations. Wrapper callables depend on framework-native tool types and currently return framework-specific objects, so keep this provisional until the wrapper callable contract is promoted deliberately. | | Auth provider | Deferred | Authentication provider authoring is still experimental and depends on subsystem APIs such as `AuthProviderBase`, `AuthProviderBaseConfig`, `AuthenticationRef`, and `register_auth_provider`. Keep it out of the stable facade until the auth compatibility and trust contract is promoted deliberately. | | Front end | Deferred | Runtime hosting surfaces need a more explicit compatibility and security contract before being promoted through `nat.plugin_api`. | | Logging | Deferred | External log sinks may export sensitive logs. Keep the existing implementation API until the stable contract and trust guidance are clearer. | diff --git a/docs/source/extend/plugins.md b/docs/source/extend/plugins.md index 9e8d9f23e1..c1e13ee646 100644 --- a/docs/source/extend/plugins.md +++ b/docs/source/extend/plugins.md @@ -55,7 +55,7 @@ NeMo Agent Toolkit currently supports the following plugin types: - **Retriever Clients**: [Retriever](../build-workflows/retrievers.md) clients are implementations of retriever providers, which are specific to a LLM framework. For example, when using the Milvus retriever provider with the LangChain/LangGraph framework, the LangChain/LangGraph Milvus retriever client needs to be registered. To register a retriever client, you can use the {py:deco}`nat.plugin_api.register_retriever_client` decorator. - **Retriever Providers**: Retriever providers are services that provide a way to retrieve information from a database. Examples of retriever providers include Chroma and Milvus. To register a retriever provider, you can use the {py:deco}`nat.plugin_api.register_retriever_provider` decorator. - **Telemetry Exporters**: [Telemetry exporters](../run-workflows/observe/observe.md) send telemetry data to a telemetry service. To register a telemetry exporter, you can use the {py:deco}`nat.plugin_api.register_telemetry_exporter` decorator. -- **Tool Wrappers**: Tool wrappers are used to wrap functions in a way that is specific to a LLM framework. For example, when using the LangChain/LangGraph framework, NeMo Agent Toolkit functions need to be wrapped in `BaseTool` class to be compatible with LangChain/LangGraph. To register a tool wrapper, you can use the {py:deco}`nat.plugin_api.register_tool_wrapper` decorator. +- **Tool Wrappers**: Tool wrappers are used to wrap functions in a way that is specific to a LLM framework. For example, when using the LangChain/LangGraph framework, NeMo Agent Toolkit functions need to be wrapped in `BaseTool` class to be compatible with LangChain/LangGraph. Tool wrapper registration is available through the provisional {py:deco}`nat.plugin_api.register_tool_wrapper` decorator while the wrapper callable contract is refined. - **API Authentication Providers**: [API authentication providers](../components/auth/api-authentication.md) are services that provide a way to authenticate requests to an API provider. Examples of authentication providers include OAuth 2.0 Authorization Code Grant and API Key. Authentication provider registration is experimental and remains a specialized extension point outside the stable `nat.plugin_api` facade. ## Anatomy of a Plugin diff --git a/packages/nvidia_nat_core/src/nat/plugin_api.py b/packages/nvidia_nat_core/src/nat/plugin_api.py index 7c8da9b448..d9fbf5bbca 100644 --- a/packages/nvidia_nat_core/src/nat/plugin_api.py +++ b/packages/nvidia_nat_core/src/nat/plugin_api.py @@ -22,6 +22,7 @@ deferred from the public plugin API (see ``DEFERRED_PLUGIN_API_CANDIDATES`` in the plugin API tests) and may change without notice; plugin authors must not depend on them: +* Auth providers: ``add_auth_provider``, ``get_auth_provider``, ``get_auth_providers``. * Finetuning: ``add_trainer``, ``add_trainer_adapter``, ``add_trajectory_builder``, ``get_trainer``, ``get_trainer_adapter``, ``get_trajectory_builder``, ``get_trainer_config``, ``get_trainer_adapter_config``, ``get_trajectory_builder_config``. From a9cf67f151dab6b959a85806a24e739c5c77b25e Mon Sep 17 00:00:00 2001 From: Bryan Bednarski Date: Thu, 28 May 2026 15:33:00 -0700 Subject: [PATCH 12/13] fix: pre-commit hook vale checks Signed-off-by: Bryan Bednarski --- docs/source/extend/third-party-plugins.md | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/docs/source/extend/third-party-plugins.md b/docs/source/extend/third-party-plugins.md index 97b110bef6..50555e12e1 100644 --- a/docs/source/extend/third-party-plugins.md +++ b/docs/source/extend/third-party-plugins.md @@ -36,7 +36,7 @@ third-party packages. This model is the default for new partner integrations where the provider is best positioned to track its own API roadmap, service semantics, and release cadence. It works for function groups, tools, LLM clients, embedder clients, retriever clients, telemetry exporters, memory backends, object stores, authentication providers, custom `nat` CLI -subcommands, and specialized front ends. +sub-commands, and specialized front ends. Provider-specific behavior belongs in the provider repository. For example, a web search plugin can return the fields and response shape exposed by the provider SDK. Do not introduce a shared web-search result schema unless the toolkit @@ -153,7 +153,7 @@ Other extension points use separate entry point groups: | Entry point group | Use | | --- | --- | | `nat.plugins` | Component plugins such as functions, function groups, model clients, retrievers, embedders, telemetry exporters, memory backends, object stores, middleware, and authentication providers. | -| `nat.cli` | Custom `nat` CLI subcommands. | +| `nat.cli` | Custom `nat` CLI sub-commands. | | `nat.front_ends` | Specialized front-end implementations. Front-end registration is not part of the stable `nat.plugin_api` facade. | ## Public API Surface @@ -304,7 +304,7 @@ uv run nat info components ## Development Workflow Use `uv` for local development when possible. This matches the primary NeMo Agent Toolkit development toolchain and -keeps lock files compatible with the toolkit's CI patterns. +keeps lock files compatible with the Toolkit CI patterns. ```bash git clone https://github.com/tavily-ai/NeMo-Agent-Toolkit-tavily.git From ab80bc9e899ae1fc05a5211a31cf47a3bf1dd9af Mon Sep 17 00:00:00 2001 From: Bryan Bednarski Date: Thu, 28 May 2026 16:00:46 -0700 Subject: [PATCH 13/13] path check exemption Signed-off-by: Bryan Bednarski --- docs/source/extend/third-party-plugins.md | 3 +++ 1 file changed, 3 insertions(+) diff --git a/docs/source/extend/third-party-plugins.md b/docs/source/extend/third-party-plugins.md index 50555e12e1..758499a4e4 100644 --- a/docs/source/extend/third-party-plugins.md +++ b/docs/source/extend/third-party-plugins.md @@ -109,6 +109,8 @@ __all__ = ["tools"] The package should declare the shared namespace package, a bounded dependency on `nvidia-nat-core`, provider SDK dependencies, optional test dependencies, repository metadata, and the component entry point. + + ```toml [build-system] requires = ["hatchling"] @@ -143,6 +145,7 @@ source = "https://github.com/tavily-ai/NeMo-Agent-Toolkit-tavily" [project.entry-points."nat.plugins"] nat_tavily = "nat.plugins.tavily.register" ``` + New external component packages should use the `nat.plugins` entry point group. The runtime also loads `nat.components` for backward compatibility with existing packages, but `nat.components` is compatibility-only for new