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

nat/data_models/config.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");
@@ -16,6 +16,7 @@
 import logging
 import sys
 import typing
+from datetime import timedelta
 
 from pydantic import BaseModel
 from pydantic import ConfigDict
@@ -27,6 +28,10 @@ from pydantic import ValidatorFunctionWrapHandler
 from pydantic import field_validator
 
 from nat.data_models.evaluate import EvalConfig
+from nat.data_models.finetuning import FinetuneConfig
+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 EmptyFunctionConfig
 from nat.data_models.function import FunctionBaseConfig
@@ -89,6 +94,12 @@ def _process_validation_error(err: ValidationError, handler: ValidatorFunctionWr
                 registered_keys = GlobalTypeRegistry.get().get_registered_ttc_strategies()
             elif (info.field_name == "middleware"):
                 registered_keys = GlobalTypeRegistry.get().get_registered_middleware()
+            elif (info.field_name == "trainers"):
+                registered_keys = GlobalTypeRegistry.get().get_registered_trainers()
+            elif (info.field_name == "trainer_adapters"):
+                registered_keys = GlobalTypeRegistry.get().get_registered_trainer_adapters()
+            elif (info.field_name == "trajectory_builders"):
+                registered_keys = GlobalTypeRegistry.get().get_registered_trajectory_builders()
 
             else:
                 assert False, f"Unknown field name {info.field_name} in validator"
@@ -204,6 +215,19 @@ class GeneralConfig(BaseModel):
 
     telemetry: TelemetryConfig = TelemetryConfig()
 
+    per_user_workflow_timeout: timedelta = Field(
+        default=timedelta(minutes=30),
+        description="Time after which inactive per-user workflows are cleaned up. "
+        "Only applies when workflow is per-user. Defaults to 30 minutes.")
+    per_user_workflow_cleanup_interval: timedelta = Field(
+        default=timedelta(minutes=5),
+        description="Interval for running cleanup of inactive per-user workflows. "
+        "Only applies when workflow is per-user. Defaults to 5 minutes.")
+    enable_per_user_monitoring: bool = Field(
+        default=False,
+        description="Enable the /monitor/users endpoint for per-user workflow resource monitoring. "
+        "When enabled, exposes metrics like request counts, latency, LLM usage, and memory for each user.")
+
     # FrontEnd Configuration
     front_end: FrontEndBaseConfig = FastApiFrontEndConfig()
 
@@ -289,6 +313,12 @@ class Config(HashableBaseModel):
     # Evaluation Options
     eval: EvalConfig = EvalConfig()
 
+    # Finetuning Options
+    trainers: dict[str, TrainerConfig] = Field(default_factory=dict)
+    trainer_adapters: dict[str, TrainerAdapterConfig] = Field(default_factory=dict)
+    trajectory_builders: dict[str, TrajectoryBuilderConfig] = Field(default_factory=dict)
+    finetuning: FinetuneConfig = FinetuneConfig()
+
     def print_summary(self, stream: typing.TextIO = sys.stdout):
         """Print a summary of the configuration"""
 
@@ -317,6 +347,9 @@ class Config(HashableBaseModel):
                      "workflow",
                      "ttc_strategies",
                      "authentication",
+                     "trainers",
+                     "trainer_adapters",
+                     "trajectory_builders",
                      mode="wrap")
     @classmethod
     def validate_components(cls, value: typing.Any, handler: ValidatorFunctionWrapHandler, info: ValidationInfo):
@@ -377,6 +410,18 @@ class Config(HashableBaseModel):
         WorkflowAnnotation = typing.Annotated[(type_registry.compute_annotation(FunctionBaseConfig)),
                                               Discriminator(TypedBaseModel.discriminator)]
 
+        TrainersAnnotation = dict[str,
+                                  typing.Annotated[type_registry.compute_annotation(TrainerConfig),
+                                                   Discriminator(TypedBaseModel.discriminator)]]
+
+        TrainerAdaptersAnnotation = dict[str,
+                                         typing.Annotated[type_registry.compute_annotation(TrainerAdapterConfig),
+                                                          Discriminator(TypedBaseModel.discriminator)]]
+
+        TrajectoryBuildersAnnotation = dict[str,
+                                            typing.Annotated[type_registry.compute_annotation(TrajectoryBuilderConfig),
+                                                             Discriminator(TypedBaseModel.discriminator)]]
+
         should_rebuild = False
 
         auth_providers_field = cls.model_fields.get("authentication")
@@ -434,6 +479,22 @@ class Config(HashableBaseModel):
             workflow_field.annotation = WorkflowAnnotation
             should_rebuild = True
 
+        trainers_field = cls.model_fields.get("trainers")
+        if trainers_field is not None and trainers_field.annotation != TrainersAnnotation:
+            trainers_field.annotation = TrainersAnnotation
+            should_rebuild = True
+
+        trainer_adapters_field = cls.model_fields.get("trainer_adapters")
+        if trainer_adapters_field is not None and trainer_adapters_field.annotation != TrainerAdaptersAnnotation:
+            trainer_adapters_field.annotation = TrainerAdaptersAnnotation
+            should_rebuild = True
+
+        trajectory_builders_field = cls.model_fields.get("trajectory_builders")
+        if (trajectory_builders_field is not None
+                and trajectory_builders_field.annotation != TrajectoryBuildersAnnotation):
+            trajectory_builders_field.annotation = TrajectoryBuildersAnnotation
+            should_rebuild = True
+
         if (GeneralConfig.rebuild_annotations()):
             should_rebuild = True
 

nat/data_models/dataset_handler.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");
@@ -12,7 +12,6 @@
 # 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.
-
 import importlib
 import json
 import typing
@@ -21,6 +20,7 @@ from pathlib import Path
 
 import pandas as pd
 from pydantic import BaseModel
+from pydantic import ConfigDict
 from pydantic import Discriminator
 from pydantic import FilePath
 from pydantic import Tag
@@ -32,6 +32,8 @@ from nat.data_models.common import TypedBaseModel
 
 class EvalS3Config(BaseModel):
 
+    model_config = ConfigDict(extra="forbid")
+
     endpoint_url: str | None = None
     region_name: str | None = None
     bucket: str
@@ -40,16 +42,25 @@ class EvalS3Config(BaseModel):
 
 
 class EvalFilterEntryConfig(BaseModel):
+
+    model_config = ConfigDict(extra="forbid")
+
     # values are lists of allowed/blocked values
     field: dict[str, list[str | int | float]] = {}
 
 
 class EvalFilterConfig(BaseModel):
+
+    model_config = ConfigDict(extra="forbid")
+
     allowlist: EvalFilterEntryConfig | None = None
     denylist: EvalFilterEntryConfig | None = None
 
 
 class EvalDatasetStructureConfig(BaseModel):
+
+    model_config = ConfigDict(extra="forbid")
+
     disable: bool = False
     question_key: str = "question"
     answer_key: str = "answer"
@@ -61,6 +72,8 @@ class EvalDatasetStructureConfig(BaseModel):
 # Base model
 class EvalDatasetBaseConfig(TypedBaseModel, BaseModelRegistryTag):
 
+    model_config = ConfigDict(extra="forbid")
+
     id_key: str = "id"
     structure: EvalDatasetStructureConfig = EvalDatasetStructureConfig()
 

nat/data_models/discovery_metadata.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/embedder.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/evaluate.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");
@@ -91,6 +91,11 @@ class EvalGeneralConfig(BaseModel):
     # Inference profiler
     profiler: ProfilerConfig | None = None
 
+    # When enabled, validates that all LLM endpoints are accessible before starting evaluation.
+    # This catches deployment issues early (e.g., 404 errors from canceled training jobs).
+    # Recommended for production workflows. Opt-in for now, may become default in future.
+    validate_llm_endpoints: bool = False
+
     # overwrite the output_dir with the output config if present
     @model_validator(mode="before")
     @classmethod

nat/data_models/evaluator.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/finetuning.py

@@ -0,0 +1,260 @@
+# 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");
+# 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.
+
+import logging
+import typing
+from enum import Enum
+from pathlib import Path
+from typing import Any
+
+from pydantic import BaseModel
+from pydantic import Field
+from pydantic import model_validator
+
+from .common import BaseModelRegistryTag
+from .common import TypedBaseModel
+
+logger = logging.getLogger(__name__)
+
+
+class RewardFunctionConfig(BaseModel):
+    """
+    Configuration for the reward function
+    """
+    name: str = Field(description="Name of the reward function.")
+
+
+class TrainerConfig(TypedBaseModel, BaseModelRegistryTag):
+    """
+    Base configuration for the Trainer
+    """
+    reward: RewardFunctionConfig | None = Field(
+        description="Configuration for the reward function used during training.", default=None)
+
+
+class TrajectoryBuilderConfig(TypedBaseModel, BaseModelRegistryTag):
+    """
+    Configuration for the trajectory collector
+    """
+    reward: RewardFunctionConfig | None = Field(
+        description="Configuration for the reward function used during trajectory building.", default=None)
+
+
+class TrainerAdapterConfig(TypedBaseModel, BaseModelRegistryTag):
+    """
+    Configuration for the trainer adapter
+    """
+    reward: RewardFunctionConfig | None = Field(
+        description="Configuration for the reward function used during training.", default=None)
+
+
+TrainerConfigT = typing.TypeVar("TrainerConfigT", bound=TrainerConfig)
+TrajectoryBuilderConfigT = typing.TypeVar("TrajectoryBuilderConfigT", bound=TrajectoryBuilderConfig)
+TrainerAdapterConfigT = typing.TypeVar("TrainerAdapterConfigT", bound=TrainerAdapterConfig)
+
+
+class TrainingJobRef(BaseModel):
+    """
+    A reference to a training job.
+    """
+    run_id: str = Field(description="The ID of the run this job belongs to.")
+    backend: str = Field(description="The backend used for the training job.")
+    metadata: dict | None = Field(description="Any additional metadata for the training job.", default=None)
+
+
+class TrainingStatusEnum(str, Enum):
+    PENDING = "pending"
+    RUNNING = "running"
+    COMPLETED = "completed"
+    FAILED = "failed"
+    CANCELED = "canceled"
+
+
+class TrainingJobStatus(BaseModel):
+    """
+    The status of a training job.
+    """
+    run_id: str = Field(description="The ID of the run this job belongs to.")
+    backend: str = Field(description="The backend used for the training job.")
+    status: TrainingStatusEnum = Field(description="The current status of the training job.")
+    progress: float | None = Field(description="The progress of the training job as a percentage (0.0 to 100.0).",
+                                   default=None)
+    message: str | None = Field(description="Any additional message or information about the training job.",
+                                default=None)
+    metadata: dict | None = Field(description="Any additional metadata for the training job.", default=None)
+
+
+class EpisodeItemRole(str, Enum):
+    USER = "user"
+    ASSISTANT = "assistant"
+    SYSTEM = "system"
+    FUNCTION = "function"
+    TOOL = "tool"
+    ENVIRONMENT = "environment"
+    OTHER = "other"
+
+
+class EpisodeItem(BaseModel):
+    """
+    A single step in an episode.
+    """
+    role: EpisodeItemRole = Field(description="The role of the agent (e.g., 'user', 'assistant').")
+    content: str = Field(description="The content of the message.")
+    logprobs: Any | None = Field(description="The log probabilities of the tokens in the message.", default=None)
+    metadata: dict | None = Field(description="Any additional metadata for the step.", default=None)
+
+    # Add model validator after construction that checks that logprobs can't be none of role is assistant
+    @model_validator(mode="after")
+    def check_logprobs(self) -> "EpisodeItem":
+        if self.role == EpisodeItemRole.ASSISTANT and self.logprobs is None:
+            raise ValueError("logprobs must be provided for assistant role.")
+        return self
+
+
+class OpenAIMessage(BaseModel):
+    """
+    A message in the OpenAI chat format.
+    """
+    role: str = Field(description="The role of the message (e.g., 'user', 'assistant').")
+    content: str = Field(description="The content of the message.")
+
+
+class DPOItem(BaseModel):
+    """
+    A single step in an episode for DPO training.
+    """
+    prompt: list[OpenAIMessage] | str = Field(description="The prompt messages leading to the response.")
+    chosen_response: str = Field(description="The response chosen as better by the reward model.")
+    rejected_response: str = Field(description="The response rejected as worse by the reward model.")
+
+
+class Trajectory(BaseModel):
+    """
+    A trajectory is a sequence of states, actions, and rewards.
+    """
+    episode: list[EpisodeItem] | list[DPOItem] = Field(description="A list of steps in the episode.")
+    reward: float = Field(description="The total reward for the episode.")
+    shaped_rewards: list[float] | None = Field(description="The shaped rewards for each step in the episode.",
+                                               default=None)
+    metadata: dict | None = Field(description="Any additional metadata for the trajectory.", default=None)
+
+
+class TrajectoryCollection(BaseModel):
+    """
+    A collection of trajectories.
+    """
+    trajectories: list[list[Trajectory]] = Field(
+        description="A list of trajectory lists, each inner list contains trajectories for one example.")
+    run_id: str = Field(description="The ID of the run this collection belongs to.")
+
+
+class CurriculumLearningConfig(BaseModel):
+    """
+    Configuration for curriculum learning in fine-tuning.
+
+    Curriculum learning progressively introduces harder training examples
+    to improve model learning and convergence.
+    """
+    enabled: bool = Field(default=False, description="Whether to enable curriculum learning")
+    initial_percentile: float = Field(default=0.3,
+                                      description="Initial percentile of trajectory groups to include (0.0-1.0). "
+                                      "E.g., 0.3 means start with top 30% easiest groups")
+    increment_percentile: float = Field(default=0.2,
+                                        description="Percentile increment when expanding curriculum. "
+                                        "E.g., 0.2 means add 20% more groups each expansion")
+    expansion_interval: int = Field(default=5, description="Number of epochs between curriculum expansions", ge=1)
+    min_reward_diff: float = Field(default=0.1,
+                                   description="Minimum reward difference within a group to be included. "
+                                   "Groups with all same rewards provide no learning signal")
+    sort_ascending: bool = Field(default=False,
+                                 description="If True, sort groups from low to high reward (hard to easy). "
+                                 "If False, sort from high to low reward (easy to hard)")
+
+    random_subsample: float | None = Field(
+        default=None, description="If set, randomly subsample this fraction of trajectories from each group.")
+
+    @model_validator(mode="after")
+    def validate_percentiles(self) -> "CurriculumLearningConfig":
+        """Validate that percentile values are in valid range."""
+        if not 0.0 < self.initial_percentile <= 1.0:
+            raise ValueError("initial_percentile must be between 0 and 1")
+        if not 0.0 < self.increment_percentile <= 1.0:
+            raise ValueError("increment_percentile must be between 0 and 1")
+        return self
+
+
+class FinetuneRunConfig(BaseModel):
+    """
+    CLI Args for running finetuning and configuring
+    """
+    config_file: Path | BaseModel = Field(description="Config file for NAT", default=None)
+    dataset: str | Path | None = None  # dataset file path can be specified in the config file
+    result_json_path: str = "$"
+    endpoint: str | None = None  # only used when running the workflow remotely
+    endpoint_timeout: int = 300
+    override: tuple[tuple[str, str], ...] = ()
+    validation_dataset: str | Path | None = Field(default=None,
+                                                  description="Validation dataset file path for periodic validation")
+
+    validation_interval: int = Field(default=5, description="Run validation every N epochs", ge=1)
+
+    validation_config_file: str | Path | None = Field(default=None,
+                                                      description="Optional separate config file for validation runs")
+
+
+class FinetuneConfig(BaseModel):
+    """
+    Parameters used for a Trainer run
+    """
+
+    enabled: bool = Field(description="Whether fine-tuning is enabled.", default=False)
+    trainer: str | None = Field(description="The trainer to use for fine-tuning.", default=None)
+    trajectory_builder: str | None = Field(description="The trajectory builder to use for fine-tuning.", default=None)
+
+    trainer_adapter: str | None = Field(description="The trainer adapter to use for fine-tuning.", default=None)
+    reward_function: RewardFunctionConfig | None = Field(description="Configuration for the reward function.",
+                                                         default=None)
+    target_functions: list[str] = ["<workflow>"]
+    target_model: str | None = Field(
+        description="Target model name to fine-tune. If None, all intermediate steps will be used without "
+        "filtering. This can lead to issues if multiple models are used in the workflow.",
+        default=None)
+    curriculum_learning: CurriculumLearningConfig = Field(
+        default=CurriculumLearningConfig(), description="Configuration for curriculum learning during fine-tuning")
+
+    num_epochs: int = Field(default=1, description="Number of epochs to run", ge=1)
+    output_dir: Path = Field(default=Path("./.tmp/nat/finetuning/"),
+                             description="Directory for outputs and checkpoints")
+
+    # Overridden by command line args
+    run_configuration: FinetuneRunConfig | None = Field(
+        description="Run-time configuration for fine-tuning (overrides CLI arguments).", default=None)
+
+    # Before validator: if enabled, config file, trainer, trajectory builder, trainer adapter and reward
+    # function must be set
+    @model_validator(mode="before")
+    def validate_finetuning_enabled(cls, values: dict[str, Any]) -> dict[str, Any]:
+        if values.get("enabled", False):
+            required_fields = ["trainer", "trajectory_builder", "trainer_adapter"]
+            missing_fields = [field for field in required_fields if values.get(field) is None]
+            if missing_fields:
+                raise ValueError(f"When fine-tuning is enabled, the following fields must be set: "
+                                 f"{', '.join(missing_fields)}")
+
+        # Warn user their config will be overridden by CLI args
+        if "run_configuration" in values and values["run_configuration"] is not None:
+            logger.warning("run_configuration will be overridden by CLI arguments during finetuning run.")
+
+        return values

nat/data_models/front_end.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/function.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/function_dependencies.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/gated_field_mixin.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/interactive.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/intermediate_step.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");
@@ -22,6 +22,7 @@ from typing import Literal
 from pydantic import BaseModel
 from pydantic import ConfigDict
 from pydantic import Field
+from pydantic import SerializeAsAny
 from pydantic import model_validator
 
 from nat.builder.framework_enum import LLMFrameworkEnum
@@ -37,6 +38,7 @@ class IntermediateStepCategory(str, Enum):
     FUNCTION = "FUNCTION"
     CUSTOM = "CUSTOM"
     SPAN = "SPAN"
+    TTC = "TTC"
 
 
 class IntermediateStepType(str, Enum):
@@ -49,6 +51,8 @@ class IntermediateStepType(str, Enum):
     WORKFLOW_END = "WORKFLOW_END"
     TASK_START = "TASK_START"
     TASK_END = "TASK_END"
+    TTC_START = "TTC_START"
+    TTC_END = "TTC_END"
     FUNCTION_START = "FUNCTION_START"
     FUNCTION_END = "FUNCTION_END"
     CUSTOM_START = "CUSTOM_START"
@@ -75,6 +79,21 @@ class StreamEventData(BaseModel):
     input: typing.Any | None = None
     output: typing.Any | None = None
     chunk: typing.Any | None = None
+    payload: typing.Any | None = None  # Raw generation payload for downstream processing
+
+
+class TTCEventData(StreamEventData):
+    """
+    TTCEventData is a data model that represents the data field in a TTC event.
+    """
+
+    # Allow extra fields in the model_config to support derived models
+    model_config = ConfigDict(extra="allow")
+
+    turn_id: str | None = None
+    turn_index: int | None = None
+    candidate_index: int | None = None
+    score: float | None = None
 
 
 class UsageInfo(BaseModel):
@@ -145,7 +164,7 @@ class IntermediateStepPayload(BaseModel):
     name: str | None = None
     tags: list[str] | None = None
     metadata: dict[str, typing.Any] | TraceMetadata | None = None
-    data: StreamEventData | None = None
+    data: SerializeAsAny[StreamEventData] | None = None
     usage_info: UsageInfo | None = None
     UUID: str = Field(default_factory=lambda: str(uuid.uuid4()))
 
@@ -170,6 +189,10 @@ class IntermediateStepPayload(BaseModel):
                 return IntermediateStepCategory.TASK
             case IntermediateStepType.TASK_END:
                 return IntermediateStepCategory.TASK
+            case IntermediateStepType.TTC_START:
+                return IntermediateStepCategory.TTC
+            case IntermediateStepType.TTC_END:
+                return IntermediateStepCategory.TTC
             case IntermediateStepType.FUNCTION_START:
                 return IntermediateStepCategory.FUNCTION
             case IntermediateStepType.FUNCTION_END:
@@ -196,6 +219,10 @@ class IntermediateStepPayload(BaseModel):
                 return IntermediateStepState.END
             case IntermediateStepType.LLM_NEW_TOKEN:
                 return IntermediateStepState.CHUNK
+            case IntermediateStepType.TTC_START:
+                return IntermediateStepState.START
+            case IntermediateStepType.TTC_END:
+                return IntermediateStepState.END
             case IntermediateStepType.TOOL_START:
                 return IntermediateStepState.START
             case IntermediateStepType.TOOL_END:

nat/data_models/invocation_node.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/llm.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/logging.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/memory.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/middleware.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");
@@ -13,6 +13,8 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+from __future__ import annotations
+
 import typing
 
 from .common import BaseModelRegistryTag

nat/data_models/object_store.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/openai_mcp.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");