nvidia-nat 1.4.0a20251112__py3-none-any.whl → 1.4.0a20260113__py3-none-any.whl

nat/cli/commands/registry/remove.py

@@ -1,4 +1,4 @@
-# SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: Apache-2.0
 #
 # Licensed under the Apache License, Version 2.0 (the "License");

nat/cli/commands/registry/search.py

@@ -1,4 +1,4 @@
-# SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: Apache-2.0
 #
 # Licensed under the Apache License, Version 2.0 (the "License");

nat/cli/commands/sizing/__init__.py

@@ -1,4 +1,4 @@
-# SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: Apache-2.0
 #
 # Licensed under the Apache License, Version 2.0 (the "License");

nat/cli/commands/sizing/calc.py

@@ -1,4 +1,4 @@
-# SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: Apache-2.0
 #
 # Licensed under the Apache License, Version 2.0 (the "License");

nat/cli/commands/sizing/sizing.py

@@ -1,4 +1,4 @@
-# SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: Apache-2.0
 #
 # Licensed under the Apache License, Version 2.0 (the "License");

nat/cli/commands/start.py

@@ -1,4 +1,4 @@
-# SPDX-FileCopyrightText: Copyright (c) 2024-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: Apache-2.0
 #
 # Licensed under the Apache License, Version 2.0 (the "License");

nat/cli/commands/uninstall.py

@@ -1,4 +1,4 @@
-# SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: Apache-2.0
 #
 # Licensed under the Apache License, Version 2.0 (the "License");

nat/cli/commands/validate.py

@@ -1,4 +1,4 @@
-# SPDX-FileCopyrightText: Copyright (c) 2024-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: Apache-2.0
 #
 # Licensed under the Apache License, Version 2.0 (the "License");

nat/cli/commands/workflow/__init__.py

@@ -1,4 +1,4 @@
-# SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: Apache-2.0
 #
 # Licensed under the Apache License, Version 2.0 (the "License");

nat/cli/commands/workflow/workflow.py

@@ -1,4 +1,4 @@
-# SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: Apache-2.0
 #
 # Licensed under the Apache License, Version 2.0 (the "License");

nat/cli/commands/workflow/workflow_commands.py

@@ -1,4 +1,4 @@
-# SPDX-FileCopyrightText: Copyright (c) 2024-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: Apache-2.0
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
@@ -19,6 +19,7 @@ import shutil
 import subprocess
 from pathlib import Path
 from urllib.parse import urlparse
+from urllib.request import url2pathname
 
 import click
 from jinja2 import Environment
@@ -151,7 +152,7 @@ def find_package_root(package_name: str) -> Path | None:
             logger.error("Invalid URL scheme in direct_url.json: %s", url)
             return None
 
-        package_root = Path(parsed_url.path).resolve()
+        package_root = Path(url2pathname(parsed_url.path)).resolve()
 
         # Ensure the path exists and is within an allowed base directory
         if not package_root.exists() or not package_root.is_dir():

nat/cli/entrypoint.py

@@ -1,4 +1,4 @@
-# SPDX-FileCopyrightText: Copyright (c) 2024-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: Apache-2.0
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
@@ -13,7 +13,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-# SPDX-FileCopyrightText: Copyright (c) 2024-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: LicenseRef-NvidiaProprietary
 #
 # NVIDIA CORPORATION, its affiliates and licensors retain all intellectual
@@ -28,39 +28,25 @@ import sys
 import time
 
 import click
-import nest_asyncio
+import nest_asyncio2
 from dotenv import load_dotenv
 
 from nat.utils.log_levels import LOG_LEVELS
+from nat.utils.log_utils import setup_logging as log_utils_setup_logging
 
-from .commands.configure.configure import configure_command
-from .commands.evaluate import eval_command
-from .commands.info.info import info_command
-from .commands.mcp.mcp import mcp_command
-from .commands.object_store.object_store import object_store_command
-from .commands.optimize import optimizer_command
-from .commands.registry.registry import registry_command
-from .commands.sizing.sizing import sizing
-from .commands.start import start_command
-from .commands.uninstall import uninstall_command
-from .commands.validate import validate_command
-from .commands.workflow.workflow import workflow_command
+from .plugin_loader import discover_and_load_cli_plugins
 
 # Load environment variables from .env file, if it exists
 load_dotenv()
 
 # Apply at the beginning of the file to avoid issues with asyncio
-nest_asyncio.apply()
+nest_asyncio2.apply()
 
 
 def setup_logging(log_level: str):
     """Configure logging with the specified level"""
     numeric_level = LOG_LEVELS.get(log_level.upper(), logging.INFO)
-    logging.basicConfig(
-        level=numeric_level,
-        format="%(asctime)s - %(levelname)-8s - %(name)s:%(lineno)d - %(message)s",
-        datefmt="%Y-%m-%d %H:%M:%S",
-    )
+    log_utils_setup_logging(numeric_level)
     return numeric_level
 
 
@@ -102,22 +88,14 @@ def cli(ctx: click.Context, log_level: str):
     ctx_dict["log_level"] = log_level
 
 
-cli.add_command(configure_command, name="configure")
-cli.add_command(eval_command, name="eval")
-cli.add_command(info_command, name="info")
-cli.add_command(registry_command, name="registry")
-cli.add_command(start_command, name="start")
-cli.add_command(uninstall_command, name="uninstall")
-cli.add_command(validate_command, name="validate")
-cli.add_command(workflow_command, name="workflow")
-cli.add_command(sizing, name="sizing")
-cli.add_command(optimizer_command, name="optimize")
-cli.add_command(object_store_command, name="object-store")
-cli.add_command(mcp_command, name="mcp")
-
-# Aliases
-cli.add_command(start_command.get_command(None, "console"), name="run")  # type: ignore
-cli.add_command(start_command.get_command(None, "fastapi"), name="serve")  # type: ignore
+# Discover and load ALL CLI commands (core + plugins) via entry points
+discover_and_load_cli_plugins(cli)
+
+# Aliases - need to get start_command from the loaded commands
+start_cmd = cli.commands.get("start")
+if start_cmd and hasattr(start_cmd, "get_command"):
+    cli.add_command(start_cmd.get_command(None, "console"), name="run")  # type: ignore
+    cli.add_command(start_cmd.get_command(None, "fastapi"), name="serve")  # type: ignore
 
 
 @cli.result_callback()

nat/cli/main.py

@@ -1,4 +1,4 @@
-# SPDX-FileCopyrightText: Copyright (c) 2024-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: Apache-2.0
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
@@ -13,7 +13,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-# SPDX-FileCopyrightText: Copyright (c) 2024-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: LicenseRef-NvidiaProprietary
 #
 # NVIDIA CORPORATION, its affiliates and licensors retain all intellectual

nat/cli/plugin_loader.py

@@ -0,0 +1,69 @@
+# 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.
+"""CLI plugin discovery system for loading plugin-specific commands."""
+
+import logging
+from importlib.metadata import entry_points
+
+import click
+
+logger = logging.getLogger(__name__)
+
+
+def discover_and_load_cli_plugins(cli_group: click.Group) -> None:
+    """Discover and load CLI command plugins from installed packages.
+
+    This function uses Python entry points to discover CLI commands provided by
+    plugin packages. Plugins register their commands under the 'nat.cli' entry
+    point group in their pyproject.toml.
+
+    The function handles import errors gracefully - if a plugin cannot be loaded
+    (e.g., due to missing dependencies), it logs a debug message but continues
+    loading other plugins.
+
+    Args:
+        cli_group: The Click group to add discovered commands to
+
+    Example plugin registration in pyproject.toml:
+        [project.entry-points.'nat.cli']
+        mcp = "nat.plugins.mcp.cli.commands:mcp_command"
+    """
+    discovered_eps = entry_points(group='nat.cli')
+
+    for ep in discovered_eps:
+        try:
+            # Load the command from the entry point
+            command = ep.load()
+
+            # Verify it's a Click command or group
+            if not isinstance(command, click.Command | click.Group):
+                logger.warning("CLI plugin '%s' from '%s' is not a Click command/group, skipping", ep.name, ep.value)
+                continue
+
+            # Add the command to the CLI group
+            cli_group.add_command(command, name=ep.name)
+            logger.debug("Loaded CLI plugin: %s from %s", ep.name, ep.value)
+
+        except ImportError as e:
+            # Plugin package not installed or missing dependencies - this is expected
+            logger.debug(
+                "Could not load CLI plugin '%s' from '%s': %s. "
+                "This is expected if the plugin package is not installed.",
+                ep.name,
+                ep.value,
+                e)
+        except Exception as e:  # noqa: BLE001
+            # Unexpected error - log as warning but continue
+            logger.warning("Error loading CLI plugin '%s' from '%s': %s", ep.name, ep.value, e, exc_info=True)

nat/cli/register_workflow.py

@@ -1,4 +1,4 @@
-# SPDX-FileCopyrightText: Copyright (c) 2024-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: Apache-2.0
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
@@ -15,6 +15,8 @@
 
 from contextlib import asynccontextmanager
 
+from pydantic import BaseModel
+
 from nat.builder.framework_enum import LLMFrameworkEnum
 from nat.cli.type_registry import AuthProviderBuildCallableT
 from nat.cli.type_registry import AuthProviderRegisteredCallableT
@@ -38,6 +40,8 @@ from nat.cli.type_registry import LoggingMethodConfigT
 from nat.cli.type_registry import LoggingMethodRegisteredCallableT
 from nat.cli.type_registry import MemoryBuildCallableT
 from nat.cli.type_registry import MemoryRegisteredCallableT
+from nat.cli.type_registry import MiddlewareBuildCallableT
+from nat.cli.type_registry import MiddlewareRegisteredCallableT
 from nat.cli.type_registry import ObjectStoreBuildCallableT
 from nat.cli.type_registry import ObjectStoreRegisteredCallableT
 from nat.cli.type_registry import RegisteredLoggingMethod
@@ -53,21 +57,33 @@ from nat.cli.type_registry import TeleExporterRegisteredCallableT
 from nat.cli.type_registry import TelemetryExporterBuildCallableT
 from nat.cli.type_registry import TelemetryExporterConfigT
 from nat.cli.type_registry import ToolWrapperBuildCallableT
+from nat.cli.type_registry import TrainerAdapterBuildCallableT
+from nat.cli.type_registry import TrainerAdapterRegisteredCallableT
+from nat.cli.type_registry import TrainerBuildCallableT
+from nat.cli.type_registry import TrainerRegisteredCallableT
+from nat.cli.type_registry import TrajectoryBuilderBuildCallableT
+from nat.cli.type_registry import TrajectoryBuilderRegisteredCallableT
 from nat.cli.type_registry import TTCStrategyBuildCallableT
-from nat.cli.type_registry import TTCStrategyRegisterCallableT
+from nat.cli.type_registry import TTCStrategyRegisteredCallableT
 from nat.data_models.authentication import AuthProviderBaseConfigT
 from nat.data_models.component import ComponentEnum
 from nat.data_models.discovery_metadata import DiscoveryMetadata
 from nat.data_models.embedder import EmbedderBaseConfigT
 from nat.data_models.evaluator import EvaluatorBaseConfigT
+from nat.data_models.finetuning import TrainerAdapterConfigT
+from nat.data_models.finetuning import TrainerConfigT
+from nat.data_models.finetuning import TrajectoryBuilderConfigT
 from nat.data_models.front_end import FrontEndConfigT
 from nat.data_models.function import FunctionConfigT
 from nat.data_models.function import FunctionGroupConfigT
 from nat.data_models.llm import LLMBaseConfigT
 from nat.data_models.memory import MemoryBaseConfigT
+from nat.data_models.middleware import MiddlewareBaseConfigT
 from nat.data_models.object_store import ObjectStoreBaseConfigT
 from nat.data_models.registry_handler import RegistryHandlerBaseConfigT
 from nat.data_models.retriever import RetrieverBaseConfigT
+from nat.data_models.ttc_strategy import TTCStrategyBaseConfigT
+from nat.utils.type_utils import DecomposedType
 
 
 def register_telemetry_exporter(config_type: type[TelemetryExporterConfigT]):
@@ -149,6 +165,10 @@ def register_function(config_type: type[FunctionConfigT],
                       framework_wrappers: list[LLMFrameworkEnum | str] | None = None):
     """
     Register a workflow with optional framework_wrappers for automatic profiler hooking.
+
+    Args:
+        config_type: The function configuration type
+        framework_wrappers: Optional list of framework wrappers for automatic profiler hooking
     """
 
     def register_function_inner(
@@ -177,6 +197,64 @@ def register_function(config_type: type[FunctionConfigT],
     return register_function_inner
 
 
+def register_per_user_function(config_type: type[FunctionConfigT],
+                               input_type: type | type[BaseModel],
+                               single_output_type: type | type[BaseModel] | None = None,
+                               streaming_output_type: type | type[BaseModel] | None = None,
+                               framework_wrappers: list[LLMFrameworkEnum | str] | None = None):
+    """
+    Register a per-user function with optional framework_wrappers for automatic profiler hooking.
+
+    The per-user function is instantiated lazily on user's first invocation. Each user will have a separate instance of
+    the function. Schemas must be provided to enable OpenAPI documentation generation without a concrete instance.
+
+    Args:
+        config_type: The function configuration type
+        input_type: The input type for the function (can be a type or a Pydantic model)
+        single_output_type: The single output type for the function (can be a type or a Pydantic model)
+        streaming_output_type: The streaming output type for the function (can be a type or a Pydantic model)
+        framework_wrappers: Optional list of framework wrappers for automatic profiler hooking
+    """
+
+    def register_per_user_function_inner(
+            fn: FunctionBuildCallableT[FunctionConfigT]) -> FunctionRegisteredCallableT[FunctionConfigT]:
+        from .type_registry import GlobalTypeRegistry
+        from .type_registry import RegisteredFunctionInfo
+
+        if not input_type:
+            raise ValueError("input_type must be provided to register a per-user function")
+
+        context_manager_fn = asynccontextmanager(fn)
+
+        framework_wrappers_list = list(framework_wrappers or [])
+
+        discovery_metadata = DiscoveryMetadata.from_config_type(config_type=config_type,
+                                                                component_type=ComponentEnum.FUNCTION)
+
+        # Convert types to Pydantic models if they are not already
+        input_schema = DecomposedType(input_type).get_pydantic_schema()
+        single_output_schema = DecomposedType(single_output_type).get_pydantic_schema() if single_output_type else None
+        streaming_output_schema = DecomposedType(
+            streaming_output_type).get_pydantic_schema() if streaming_output_type else None
+
+        GlobalTypeRegistry.get().register_function(
+            RegisteredFunctionInfo(
+                full_type=config_type.full_type,
+                config_type=config_type,
+                build_fn=context_manager_fn,
+                framework_wrappers=framework_wrappers_list,
+                discovery_metadata=discovery_metadata,
+                is_per_user=True,
+                per_user_function_input_schema=input_schema,
+                per_user_function_single_output_schema=single_output_schema,
+                per_user_function_streaming_output_schema=streaming_output_schema,
+            ))
+
+        return context_manager_fn
+
+    return register_per_user_function_inner
+
+
 def register_function_group(config_type: type[FunctionGroupConfigT],
                             framework_wrappers: list[LLMFrameworkEnum | str] | None = None):
     """
@@ -211,6 +289,86 @@ def register_function_group(config_type: type[FunctionGroupConfigT],
     return register_function_group_inner
 
 
+def register_per_user_function_group(config_type: type[FunctionGroupConfigT],
+                                     framework_wrappers: list[LLMFrameworkEnum | str] | None = None):
+    """
+    Register a per-user function group with optional framework_wrappers for automatic profiler hooking.
+
+    Per-user function groups are instantiated separately for each user, allowing for user-specific
+    shared state across multiple functions within the group.
+    """
+
+    def register_per_user_function_group_inner(
+        fn: FunctionGroupBuildCallableT[FunctionGroupConfigT]
+    ) -> FunctionGroupRegisteredCallableT[FunctionGroupConfigT]:
+        from .type_registry import GlobalTypeRegistry
+        from .type_registry import RegisteredFunctionGroupInfo
+
+        context_manager_fn = asynccontextmanager(fn)
+
+        framework_wrappers_list = list(framework_wrappers or [])
+
+        discovery_metadata = DiscoveryMetadata.from_config_type(config_type=config_type,
+                                                                component_type=ComponentEnum.FUNCTION_GROUP)
+
+        GlobalTypeRegistry.get().register_function_group(
+            RegisteredFunctionGroupInfo(
+                full_type=config_type.full_type,
+                config_type=config_type,
+                build_fn=context_manager_fn,
+                framework_wrappers=framework_wrappers_list,
+                discovery_metadata=discovery_metadata,
+                is_per_user=True,
+            ))
+
+        return context_manager_fn
+
+    return register_per_user_function_group_inner
+
+
+def register_middleware(config_type: type[MiddlewareBaseConfigT]):
+    """
+    Register a middleware component.
+
+    Middleware provides middleware-style wrapping of calls with
+    preprocessing and postprocessing logic. They are built as components that can
+    be configured in YAML and referenced by name in configurations.
+
+    Args:
+        config_type: The middleware configuration type to register
+
+    Returns:
+        A decorator that wraps the build function as an async context manager
+    """
+
+    def register_middleware_inner(
+            fn: MiddlewareBuildCallableT[MiddlewareBaseConfigT]
+    ) -> MiddlewareRegisteredCallableT[MiddlewareBaseConfigT]:
+        from .type_registry import GlobalTypeRegistry
+        from .type_registry import RegisteredMiddlewareInfo
+
+        context_manager_fn = asynccontextmanager(fn)
+
+        discovery_metadata = DiscoveryMetadata.from_config_type(config_type=config_type,
+                                                                component_type=ComponentEnum.MIDDLEWARE)
+
+        GlobalTypeRegistry.get().register_middleware(
+            RegisteredMiddlewareInfo(
+                full_type=config_type.full_type,
+                config_type=config_type,
+                build_fn=context_manager_fn,
+                discovery_metadata=discovery_metadata,
+            ))
+
+        return context_manager_fn
+
+    return register_middleware_inner
+
+
+# Compatibility alias for backwards compatibility
+register_function_middleware = register_middleware
+
+
 def register_llm_provider(config_type: type[LLMBaseConfigT]):
 
     def register_llm_provider_inner(
@@ -405,11 +563,11 @@ def register_object_store(config_type: type[ObjectStoreBaseConfigT]):
     return register_kv_store_inner
 
 
-def register_ttc_strategy(config_type: type[TTCStrategyRegisterCallableT]):
+def register_ttc_strategy(config_type: type[TTCStrategyBaseConfigT]):
 
     def register_ttc_strategy_inner(
-        fn: TTCStrategyBuildCallableT[TTCStrategyRegisterCallableT]
-    ) -> TTCStrategyRegisterCallableT[TTCStrategyRegisterCallableT]:
+        fn: TTCStrategyBuildCallableT[TTCStrategyBaseConfigT]
+    ) -> TTCStrategyRegisteredCallableT[TTCStrategyBaseConfigT]:
         from .type_registry import GlobalTypeRegistry
         from .type_registry import RegisteredTTCStrategyInfo
 
@@ -429,6 +587,76 @@ def register_ttc_strategy(config_type: type[TTCStrategyRegisterCallableT]):
     return register_ttc_strategy_inner
 
 
+def register_trainer(config_type: type[TrainerConfigT]):
+
+    def register_trainer_inner(fn: TrainerBuildCallableT[TrainerConfigT]) -> TrainerRegisteredCallableT[TrainerConfigT]:
+        from .type_registry import GlobalTypeRegistry
+        from .type_registry import RegisteredTrainerInfo
+
+        context_manager_fn = asynccontextmanager(fn)
+
+        discovery_metadata = DiscoveryMetadata.from_config_type(config_type=config_type,
+                                                                component_type=ComponentEnum.TRAINER)
+
+        GlobalTypeRegistry.get().register_trainer(
+            RegisteredTrainerInfo(full_type=config_type.full_type,
+                                  config_type=config_type,
+                                  build_fn=context_manager_fn,
+                                  discovery_metadata=discovery_metadata))
+
+        return context_manager_fn
+
+    return register_trainer_inner
+
+
+def register_trainer_adapter(config_type: type[TrainerAdapterConfigT]):
+
+    def register_trainer_adapter_inner(
+        fn: TrainerAdapterBuildCallableT[TrainerAdapterConfigT]
+    ) -> TrainerAdapterRegisteredCallableT[TrainerAdapterConfigT]:
+        from .type_registry import GlobalTypeRegistry
+        from .type_registry import RegisteredTrainerAdapterInfo
+
+        context_manager_fn = asynccontextmanager(fn)
+
+        discovery_metadata = DiscoveryMetadata.from_config_type(config_type=config_type,
+                                                                component_type=ComponentEnum.TRAINER_ADAPTER)
+
+        GlobalTypeRegistry.get().register_trainer_adapter(
+            RegisteredTrainerAdapterInfo(full_type=config_type.full_type,
+                                         config_type=config_type,
+                                         build_fn=context_manager_fn,
+                                         discovery_metadata=discovery_metadata))
+
+        return context_manager_fn
+
+    return register_trainer_adapter_inner
+
+
+def register_trajectory_builder(config_type: type[TrajectoryBuilderConfigT]):
+
+    def register_trajectory_builder_inner(
+        fn: TrajectoryBuilderBuildCallableT[TrajectoryBuilderConfigT]
+    ) -> TrajectoryBuilderRegisteredCallableT[TrajectoryBuilderConfigT]:
+        from .type_registry import GlobalTypeRegistry
+        from .type_registry import RegisteredTrajectoryBuilderInfo
+
+        context_manager_fn = asynccontextmanager(fn)
+
+        discovery_metadata = DiscoveryMetadata.from_config_type(config_type=config_type,
+                                                                component_type=ComponentEnum.TRAJECTORY_BUILDER)
+
+        GlobalTypeRegistry.get().register_trajectory_builder(
+            RegisteredTrajectoryBuilderInfo(full_type=config_type.full_type,
+                                            config_type=config_type,
+                                            build_fn=context_manager_fn,
+                                            discovery_metadata=discovery_metadata))
+
+        return context_manager_fn
+
+    return register_trajectory_builder_inner
+
+
 def register_retriever_provider(config_type: type[RetrieverBaseConfigT]):
 
     def register_retriever_provider_inner(