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

nat/cli/type_registry.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");
@@ -22,6 +22,7 @@ from contextlib import contextmanager
 from copy import deepcopy
 from functools import cached_property
 from logging import Handler
+from typing import Self
 
 from pydantic import BaseModel
 from pydantic import ConfigDict
@@ -29,6 +30,7 @@ from pydantic import Field
 from pydantic import Tag
 from pydantic import computed_field
 from pydantic import field_validator
+from pydantic import model_validator
 
 from nat.authentication.interfaces import AuthProviderBase
 from nat.builder.builder import Builder
@@ -52,6 +54,12 @@ from nat.data_models.embedder import EmbedderBaseConfig
 from nat.data_models.embedder import EmbedderBaseConfigT
 from nat.data_models.evaluator import EvaluatorBaseConfig
 from nat.data_models.evaluator import EvaluatorBaseConfigT
+from nat.data_models.finetuning import TrainerAdapterConfig
+from nat.data_models.finetuning import TrainerAdapterConfigT
+from nat.data_models.finetuning import TrainerConfig
+from nat.data_models.finetuning import TrainerConfigT
+from nat.data_models.finetuning import TrajectoryBuilderConfig
+from nat.data_models.finetuning import TrajectoryBuilderConfigT
 from nat.data_models.front_end import FrontEndBaseConfig
 from nat.data_models.front_end import FrontEndConfigT
 from nat.data_models.function import FunctionBaseConfig
@@ -77,6 +85,9 @@ from nat.data_models.telemetry_exporter import TelemetryExporterConfigT
 from nat.data_models.ttc_strategy import TTCStrategyBaseConfig
 from nat.data_models.ttc_strategy import TTCStrategyBaseConfigT
 from nat.experimental.test_time_compute.models.strategy_base import StrategyBase
+from nat.finetuning.interfaces.finetuning_runner import Trainer
+from nat.finetuning.interfaces.trainer_adapter import TrainerAdapter
+from nat.finetuning.interfaces.trajectory_builder import TrajectoryBuilder
 from nat.memory.interfaces import MemoryEditor
 from nat.middleware.middleware import Middleware
 from nat.object_store.interfaces import ObjectStore
@@ -89,6 +100,9 @@ AuthProviderBuildCallableT = Callable[[AuthProviderBaseConfigT, Builder], AsyncI
 EmbedderClientBuildCallableT = Callable[[EmbedderBaseConfigT, Builder], AsyncIterator[typing.Any]]
 EmbedderProviderBuildCallableT = Callable[[EmbedderBaseConfigT, Builder], AsyncIterator[EmbedderProviderInfo]]
 EvaluatorBuildCallableT = Callable[[EvaluatorBaseConfigT, EvalBuilder], AsyncIterator[EvaluatorInfo]]
+TrainerBuildCallableT = Callable[[TrainerConfigT, Builder], AsyncIterator[Trainer]]
+TrainerAdapterBuildCallableT = Callable[[TrainerAdapterConfigT, Builder], AsyncIterator[TrainerAdapter]]
+TrajectoryBuilderBuildCallableT = Callable[[TrajectoryBuilderConfigT, Builder], AsyncIterator[TrajectoryBuilder]]
 FrontEndBuildCallableT = Callable[[FrontEndConfigT, Config], AsyncIterator[FrontEndBase]]
 FunctionBuildCallableT = Callable[[FunctionConfigT, Builder], AsyncIterator[FunctionInfo | Callable | FunctionBase]]
 FunctionGroupBuildCallableT = Callable[[FunctionGroupConfigT, Builder], AsyncIterator[FunctionGroup]]
@@ -111,12 +125,17 @@ EmbedderClientRegisteredCallableT = Callable[[EmbedderBaseConfigT, Builder], Abs
 EmbedderProviderRegisteredCallableT = Callable[[EmbedderBaseConfigT, Builder],
                                                AbstractAsyncContextManager[EmbedderProviderInfo]]
 EvaluatorRegisteredCallableT = Callable[[EvaluatorBaseConfigT, EvalBuilder], AbstractAsyncContextManager[EvaluatorInfo]]
+TrainerRegisteredCallableT = Callable[[TrainerConfigT, Builder], AbstractAsyncContextManager[Trainer]]
+TrainerAdapterRegisteredCallableT = Callable[[TrainerAdapterConfigT, Builder],
+                                             AbstractAsyncContextManager[TrainerAdapter]]
+TrajectoryBuilderRegisteredCallableT = Callable[[TrajectoryBuilderConfigT, Builder],
+                                                AbstractAsyncContextManager[TrajectoryBuilder]]
 FrontEndRegisteredCallableT = Callable[[FrontEndConfigT, Config], AbstractAsyncContextManager[FrontEndBase]]
 FunctionRegisteredCallableT = Callable[[FunctionConfigT, Builder],
                                        AbstractAsyncContextManager[FunctionInfo | Callable | FunctionBase]]
 FunctionGroupRegisteredCallableT = Callable[[FunctionGroupConfigT, Builder], AbstractAsyncContextManager[FunctionGroup]]
 MiddlewareRegisteredCallableT = Callable[[MiddlewareBaseConfigT, Builder], AbstractAsyncContextManager[Middleware]]
-TTCStrategyRegisterCallableT = Callable[[TTCStrategyBaseConfigT, Builder], AbstractAsyncContextManager[StrategyBase]]
+TTCStrategyRegisteredCallableT = Callable[[TTCStrategyBaseConfigT, Builder], AbstractAsyncContextManager[StrategyBase]]
 LLMClientRegisteredCallableT = Callable[[LLMBaseConfigT, Builder], AbstractAsyncContextManager[typing.Any]]
 LLMProviderRegisteredCallableT = Callable[[LLMBaseConfigT, Builder], AbstractAsyncContextManager[LLMProviderInfo]]
 LoggingMethodRegisteredCallableT = Callable[[LoggingMethodConfigT, Builder], AbstractAsyncContextManager[typing.Any]]
@@ -138,6 +157,12 @@ class RegisteredInfo(BaseModel, typing.Generic[TypedBaseModelT]):
     config_type: type[TypedBaseModelT]
     discovery_metadata: DiscoveryMetadata = DiscoveryMetadata()
 
+    is_per_user: bool = Field(default=False,
+                              description="Whether the component is per-user. Default to False. If True, "
+                              "each user will have their own separate instance of the component. The per-user "
+                              "component instance will be lazily built on user's first invocation. If False, the "
+                              "component instance will be shared across all users.")
+
     @computed_field
     @cached_property
     def module_name(self) -> str:
@@ -169,6 +194,32 @@ class RegisteredLoggingMethod(RegisteredInfo[LoggingBaseConfig]):
     build_fn: LoggingMethodRegisteredCallableT = Field(repr=False)
 
 
+class RegisteredTrainerInfo(RegisteredInfo[TrainerConfig]):
+    """
+    Represents a registered Trainer. Trainers are responsible for fine-tuning LLMs.
+    """
+
+    build_fn: TrainerRegisteredCallableT = Field(repr=False)
+
+
+class RegisteredTrainerAdapterInfo(RegisteredInfo[TrainerAdapterConfig]):
+    """
+    Represents a registered Trainer Adapter. Trainer Adapters are responsible for adapting the training process to
+    different frameworks.
+    """
+
+    build_fn: TrainerAdapterRegisteredCallableT = Field(repr=False)
+
+
+class RegisteredTrajectoryBuilderInfo(RegisteredInfo[TrajectoryBuilderConfig]):
+    """
+    Represents a registered Trajectory Builder. Trajectory Builders are responsible for building trajectories for
+    fine-tuning.
+    """
+
+    build_fn: TrajectoryBuilderRegisteredCallableT = Field(repr=False)
+
+
 class RegisteredFrontEndInfo(RegisteredInfo[FrontEndBaseConfig]):
     """
     Represents a registered front end. Front ends are the entry points to the workflow and are responsible for
@@ -189,6 +240,38 @@ class RegisteredFunctionInfo(RegisteredInfo[FunctionBaseConfig]):
     build_fn: FunctionRegisteredCallableT = Field(repr=False)
     framework_wrappers: list[str] = Field(default_factory=list)
 
+    # Declared schemas for per-user functions which are lazy-loaded. Must be provided if is_per_user is True.
+    per_user_function_input_schema: type[BaseModel] | type[None] | None = Field(
+        default=None,
+        description="Declared input schema for per-user functions. Must be provided if is_per_user "
+        "is True. This is for enabling OpenAPI documentation generation without a concrete function instance.")
+    per_user_function_single_output_schema: type[BaseModel] | type[None] | None = Field(
+        default=None,
+        description="Declared single output schema for per-user functions. Must be provided if is_per_user "
+        "is True. This is for enabling OpenAPI documentation generation without a concrete function instance.")
+    per_user_function_streaming_output_schema: type[BaseModel] | type[None] | None = Field(
+        default=None,
+        description="Declared streaming output schema for per-user functions. Must be provided if is_per_user "
+        "is True. This is for enabling OpenAPI documentation generation without a concrete function instance.")
+
+    @model_validator(mode="after")
+    def validate_per_user_function_schema_declaration(self) -> Self:
+        """
+        Validate if the schemas are explicitly declared when is_per_user is True
+        """
+        if self.is_per_user:
+
+            if self.per_user_function_input_schema is None:
+                raise ValueError("per_user_function_input_schema must be provided if is_per_user is True")
+
+            if self.per_user_function_single_output_schema is None and \
+                self.per_user_function_streaming_output_schema is None:
+                raise ValueError(
+                    "per_user_function_single_output_schema or per_user_function_streaming_output_schema must be "
+                    "provided if is_per_user is True")
+
+        return self
+
 
 class RegisteredFunctionGroupInfo(RegisteredInfo[FunctionGroupBaseConfig]):
     """
@@ -284,7 +367,7 @@ class RegisteredTTCStrategyInfo(RegisteredInfo[TTCStrategyBaseConfig]):
     Represents a registered TTC strategy.
     """
 
-    build_fn: TTCStrategyRegisterCallableT = Field(repr=False)
+    build_fn: TTCStrategyRegisteredCallableT = Field(repr=False)
 
 
 class RegisteredToolWrapper(BaseModel):
@@ -393,6 +476,12 @@ class TypeRegistry:
         # TTC Strategies
         self._registered_ttc_strategies: dict[type[TTCStrategyBaseConfig], RegisteredTTCStrategyInfo] = {}
 
+        # Registered training things
+        self._registered_trainer_infos: dict[type[TrainerConfig], RegisteredTrainerInfo] = {}
+        self._registered_trainer_adapter_infos: dict[type[TrainerAdapterConfig], RegisteredTrainerAdapterInfo] = {}
+        self._registered_trajectory_builder_infos: dict[type[TrajectoryBuilderConfig],
+                                                        RegisteredTrajectoryBuilderInfo] = {}
+
         # Packages
         self._registered_packages: dict[str, RegisteredPackage] = {}
 
@@ -471,6 +560,65 @@ class TypeRegistry:
 
         return list(self._registered_logging_methods.values())
 
+    def register_trainer(self, registration: RegisteredTrainerInfo):
+
+        if (registration.config_type in self._registered_trainer_infos):
+            raise ValueError(f"A trainer with the same config type `{registration.config_type}` has already "
+                             "been registered.")
+
+        self._registered_trainer_infos[registration.config_type] = registration
+
+        self._registration_changed()
+
+    def register_trainer_adapter(self, registration: RegisteredTrainerAdapterInfo):
+        if (registration.config_type in self._registered_trainer_adapter_infos):
+            raise ValueError(f"A trainer adapter with the same config type `{registration.config_type}` has already "
+                             "been registered.")
+
+        self._registered_trainer_adapter_infos[registration.config_type] = registration
+
+        self._registration_changed()
+
+    def register_trajectory_builder(self, registration: RegisteredTrajectoryBuilderInfo):
+        if (registration.config_type in self._registered_trajectory_builder_infos):
+            raise ValueError(f"A trajectory builder with the same config type `{registration.config_type}` has already "
+                             "been registered.")
+
+        self._registered_trajectory_builder_infos[registration.config_type] = registration
+
+        self._registration_changed()
+
+    def get_trainer(self, config_type: type[TrainerConfig]) -> RegisteredTrainerInfo:
+
+        try:
+            return self._registered_trainer_infos[config_type]
+        except KeyError as err:
+            raise KeyError(f"Could not find a registered trainer for config `{config_type}`. "
+                           f"Registered configs: {set(self._registered_trainer_infos.keys())}") from err
+
+    def get_trainer_adapter(self, config_type: type[TrainerAdapterConfig]) -> RegisteredTrainerAdapterInfo:
+        try:
+            return self._registered_trainer_adapter_infos[config_type]
+        except KeyError as err:
+            raise KeyError(f"Could not find a registered trainer adapter for config `{config_type}`. "
+                           f"Registered configs: {set(self._registered_trainer_adapter_infos.keys())}") from err
+
+    def get_trajectory_builder(self, config_type: type[TrajectoryBuilderConfig]) -> RegisteredTrajectoryBuilderInfo:
+        try:
+            return self._registered_trajectory_builder_infos[config_type]
+        except KeyError as err:
+            raise KeyError(f"Could not find a registered trajectory builder for config `{config_type}`. "
+                           f"Registered configs: {set(self._registered_trajectory_builder_infos.keys())}") from err
+
+    def get_registered_trainers(self) -> list[RegisteredInfo[TrainerConfig]]:
+        return list(self._registered_trainer_infos.values())
+
+    def get_registered_trainer_adapters(self) -> list[RegisteredInfo[TrainerAdapterConfig]]:
+        return list(self._registered_trainer_adapter_infos.values())
+
+    def get_registered_trajectory_builders(self) -> list[RegisteredInfo[TrajectoryBuilderConfig]]:
+        return list(self._registered_trajectory_builder_infos.values())
+
     def register_front_end(self, registration: RegisteredFrontEndInfo):
 
         if (registration.config_type in self._registered_front_end_infos):
@@ -977,6 +1125,15 @@ class TypeRegistry:
         if component_type == ComponentEnum.MIDDLEWARE:
             return self._registered_middleware
 
+        if component_type == ComponentEnum.TRAINER:
+            return self._registered_trainer_infos
+
+        if component_type == ComponentEnum.TRAJECTORY_BUILDER:
+            return self._registered_trajectory_builder_infos
+
+        if component_type == ComponentEnum.TRAINER_ADAPTER:
+            return self._registered_trainer_adapter_infos
+
         raise ValueError(f"Supplied an unsupported component type {component_type}")
 
     def get_registered_types_by_component_type(self, component_type: ComponentEnum) -> list[str]:
@@ -1106,6 +1263,15 @@ class TypeRegistry:
         if issubclass(cls, MiddlewareBaseConfig):
             return self._do_compute_annotation(cls, self.get_registered_middleware())
 
+        if issubclass(cls, TrainerConfig):
+            return self._do_compute_annotation(cls, self.get_registered_trainers())
+
+        if issubclass(cls, TrainerAdapterConfig):
+            return self._do_compute_annotation(cls, self.get_registered_trainer_adapters())
+
+        if issubclass(cls, TrajectoryBuilderConfig):
+            return self._do_compute_annotation(cls, self.get_registered_trajectory_builders())
+
         raise ValueError(f"Supplied an unsupported component type {cls}")
 
 

nat/control_flow/register.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/control_flow/router_agent/agent.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/control_flow/router_agent/prompt.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/control_flow/router_agent/register.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/control_flow/sequential_executor.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");
@@ -32,6 +32,14 @@ from nat.utils.type_utils import DecomposedType
 logger = logging.getLogger(__name__)
 
 
+class SequentialExecutorExit(Exception):
+    """Raised when a tool wants to exit the sequential executor chain early with a custom message."""
+
+    def __init__(self, message: str):
+        self.message = message
+        super().__init__(message)
+
+
 class ToolExecutionConfig(BaseModel):
     """Configuration for individual tool execution within sequential execution."""
 
@@ -41,6 +49,7 @@ class ToolExecutionConfig(BaseModel):
 class SequentialExecutorConfig(FunctionBaseConfig, name="sequential_executor"):
     """Configuration for sequential execution of a list of functions."""
 
+    description: str = Field(default="Sequential Executor Workflow", description="Description of this functions use.")
     tool_list: list[FunctionRef] = Field(default_factory=list,
                                          description="A list of functions to execute sequentially.")
     tool_execution_config: dict[str, ToolExecutionConfig] = Field(default_factory=dict,
@@ -54,6 +63,11 @@ class SequentialExecutorConfig(FunctionBaseConfig, name="sequential_executor"):
         "which means the output type of the previous function is compatible with the input type of the next function."
         "If set to True, any incompatibility will raise an exception. If set to false, the incompatibility will only"
         "generate a warning message and the sequential execution will continue.")
+    return_error_on_exception: bool = Field(
+        default=False,
+        description="If set to True, when an uncaught exception occurs during tool execution, the sequential executor "
+        "will exit early and return an error message as the workflow output instead of raising the exception. "
+        "If set to False (default), exceptions are re-raised.")
 
 
 def _get_function_output_type(function: Function, tool_execution_config: dict[str, ToolExecutionConfig]) -> type:
@@ -126,11 +140,11 @@ async def sequential_execution(config: SequentialExecutorConfig, builder: Builde
         raise ValueError(f"Error with the sequential executor tool list: {e}")
 
     # The type annotation of _sequential_function_execution is dynamically set according to the tool list
-    async def _sequential_function_execution(initial_tool_input):
+    async def _sequential_function_execution(input_message):
         logger.debug(f"Executing sequential executor with tool list: {config.tool_list}")
 
         tool_list: list[FunctionRef] = config.tool_list
-        tool_input = initial_tool_input
+        tool_input = input_message
         tool_response = None
 
         for tool_name in tool_list:
@@ -148,7 +162,16 @@ async def sequential_execution(config: SequentialExecutorConfig, builder: Builde
                         tool_response = await tool.ainvoke(tool_input)
                 else:
                     tool_response = await tool.ainvoke(tool_input)
+            except SequentialExecutorExit as e:
+                # Tool explicitly requested early exit - always return the message
+                logger.info(f"Tool {tool_name} requested early exit: {e.message}")
+                return e.message
             except Exception as e:
+                if config.return_error_on_exception:
+                    # Return error message as workflow output instead of raising exception
+                    logger.exception(f"Error with tool {tool_name}, returning error message")
+                    error_message = f"Error in {tool_name}: {type(e).__name__}: {str(e)}"
+                    return error_message
                 logger.error(f"Error with tool {tool_name}: {e}")
                 raise
 
@@ -158,9 +181,7 @@ async def sequential_execution(config: SequentialExecutorConfig, builder: Builde
         return tool_response
 
     # Dynamically set the annotations for the function
-    _sequential_function_execution.__annotations__ = {"initial_tool_input": input_type, "return": output_type}
+    _sequential_function_execution.__annotations__ = {"input_message": input_type, "return": output_type}
     logger.debug(f"Sequential executor function annotations: {_sequential_function_execution.__annotations__}")
 
-    yield FunctionInfo.from_fn(_sequential_function_execution,
-                               description="Executes a list of functions sequentially."
-                               "The input of the next tool is the response of the previous tool.")
+    yield FunctionInfo.from_fn(_sequential_function_execution, description=config.description)

nat/data_models/__init__.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/data_models/agent.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/data_models/api_server.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");
@@ -24,7 +24,6 @@ from pydantic import BaseModel
 from pydantic import ConfigDict
 from pydantic import Discriminator
 from pydantic import Field
-from pydantic import HttpUrl
 from pydantic import conlist
 from pydantic import field_serializer
 from pydantic import field_validator
@@ -90,7 +89,8 @@ class AudioContent(BaseModel):
 
 
 class ImageUrl(BaseModel):
-    url: HttpUrl = HttpUrl(url="http://default.com")
+    url: str = Field(default="http://default.com",
+                     description="Either a URL of the image or the base64 encoded image data.")
 
 
 class ImageContent(BaseModel):
@@ -494,6 +494,18 @@ class ResponseIntermediateStep(ResponseBaseModelIntermediate):
     payload: str
 
 
+class ResponseObservabilityTrace(BaseModel, ResponseSerializable):
+    """
+    ResponseObservabilityTrace is a data model that represents an observability trace event
+    sent once when the trace ID becomes available
+    """
+
+    observability_trace_id: str
+
+    def get_stream_data(self) -> str:
+        return f"observability_trace: {self.model_dump_json()}\n\n"
+
+
 class ResponsePayloadOutput(BaseModel, ResponseSerializable):
 
     payload: typing.Any
@@ -525,6 +537,7 @@ class WebSocketMessageType(str, Enum):
     INTERMEDIATE_STEP_MESSAGE = "system_intermediate_message"
     SYSTEM_INTERACTION_MESSAGE = "system_interaction_message"
     USER_INTERACTION_MESSAGE = "user_interaction_message"
+    OBSERVABILITY_TRACE_MESSAGE = "observability_trace_message"
     ERROR_MESSAGE = "error_message"
 
 
@@ -700,6 +713,28 @@ class WebSocketSystemInteractionMessage(BaseModel):
     timestamp: str = str(datetime.datetime.now(datetime.UTC))
 
 
+class ObservabilityTraceContent(BaseModel):
+    model_config = ConfigDict(extra="forbid")
+    observability_trace_id: str
+
+
+class WebSocketObservabilityTraceMessage(BaseModel):
+    """
+    WebSocket message for observability trace ID.
+    Sent once after the workflow completes to correlate the request with observability traces.
+    """
+    # Allow extra fields in the model_config to support derived models
+    model_config = ConfigDict(extra="allow")
+
+    type: typing.Literal[
+        WebSocketMessageType.OBSERVABILITY_TRACE_MESSAGE] = WebSocketMessageType.OBSERVABILITY_TRACE_MESSAGE
+    id: str = "default"
+    parent_id: str = "default"
+    conversation_id: str | None = None
+    content: ObservabilityTraceContent
+    timestamp: str = str(datetime.datetime.now(datetime.UTC))
+
+
 # ======== GenerateResponse Converters ========
 
 

nat/data_models/authentication.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/data_models/common.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/data_models/component.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");
@@ -41,6 +41,9 @@ class ComponentEnum(StrEnum):
     RETRIEVER_PROVIDER = "retriever_provider"
     TOOL_WRAPPER = "tool_wrapper"
     TRACING = "tracing"
+    TRAINER = "trainer"
+    TRAJECTORY_BUILDER = "trajectory_builder"
+    TRAINER_ADAPTER = "trainer_adapter"
     UNDEFINED = "undefined"
 
 
@@ -56,6 +59,9 @@ class ComponentGroup(StrEnum):
     MEMORY = "memory"
     OBJECT_STORES = "object_stores"
     RETRIEVERS = "retrievers"
+    TRAINERS = "trainers"
+    TRAINER_ADAPTERS = "trainer_adapters"
+    TRAJECTORY_BUILDERS = "trajectory_builders"
 
 
 # Compatibility aliases with previous releases

nat/data_models/component_ref.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");
@@ -188,3 +188,36 @@ class MiddlewareRef(ComponentRef):
     @override
     def component_group(self):
         return ComponentGroup.MIDDLEWARE
+
+
+class TrainerRef(ComponentRef):
+    """
+    A reference to a trainer in a NAT configuration object.
+    """
+
+    @property
+    @override
+    def component_group(self):
+        return ComponentGroup.TRAINERS
+
+
+class TrajectoryBuilderRef(ComponentRef):
+    """
+    A reference to a trajectory builder in a NAT configuration object.
+    """
+
+    @property
+    @override
+    def component_group(self):
+        return ComponentGroup.TRAJECTORY_BUILDERS
+
+
+class TrainerAdapterRef(ComponentRef):
+    """
+    A reference to a trainer adapter in a NAT configuration object.
+    """
+
+    @property
+    @override
+    def component_group(self):
+        return ComponentGroup.TRAINER_ADAPTERS