nvidia-nat 1.3.0a20250826__py3-none-any.whl → 1.3.0a20250828__py3-none-any.whl

nat/agent/base.py

@@ -70,12 +70,14 @@ class BaseAgent(ABC):
                  llm: BaseChatModel,
                  tools: list[BaseTool],
                  callbacks: list[AsyncCallbackHandler] | None = None,
-                 detailed_logs: bool = False) -> None:
+                 detailed_logs: bool = False,
+                 log_response_max_chars: int = 1000) -> None:
         logger.debug("Initializing Agent Graph")
         self.llm = llm
         self.tools = tools
         self.callbacks = callbacks or []
         self.detailed_logs = detailed_logs
+        self.log_response_max_chars = log_response_max_chars
         self.graph = None
 
     async def _stream_llm(self,
@@ -184,7 +186,7 @@ class BaseAgent(ABC):
         logger.error("%s %s", AGENT_LOG_PREFIX, error_content)
         return ToolMessage(name=tool.name, tool_call_id=tool.name, content=error_content, status="error")
 
-    def _log_tool_response(self, tool_name: str, tool_input: Any, tool_response: str, max_chars: int = 1000) -> None:
+    def _log_tool_response(self, tool_name: str, tool_input: Any, tool_response: str) -> None:
         """
         Log tool response with consistent formatting and length limits.
 
@@ -196,13 +198,11 @@ class BaseAgent(ABC):
             The input that was passed to the tool
         tool_response : str
             The response from the tool
-        max_chars : int
-            Maximum number of characters to log (default: 1000)
         """
         if self.detailed_logs:
             # Truncate tool response if too long
-            display_response = tool_response[:max_chars] + "...(rest of response truncated)" if len(
-                tool_response) > max_chars else tool_response
+            display_response = tool_response[:self.log_response_max_chars] + "...(rest of response truncated)" if len(
+                tool_response) > self.log_response_max_chars else tool_response
 
             # Format the tool input for display
             tool_input_str = str(tool_input)

nat/agent/dual_node.py

@@ -35,8 +35,13 @@ class DualNodeAgent(BaseAgent):
                  llm: BaseChatModel,
                  tools: list[BaseTool],
                  callbacks: list[AsyncCallbackHandler] | None = None,
-                 detailed_logs: bool = False):
-        super().__init__(llm=llm, tools=tools, callbacks=callbacks, detailed_logs=detailed_logs)
+                 detailed_logs: bool = False,
+                 log_response_max_chars: int = 1000):
+        super().__init__(llm=llm,
+                         tools=tools,
+                         callbacks=callbacks,
+                         detailed_logs=detailed_logs,
+                         log_response_max_chars=log_response_max_chars)
 
     @abstractmethod
     async def agent_node(self, state: BaseModel) -> BaseModel:

nat/agent/react_agent/agent.py

@@ -73,11 +73,16 @@ class ReActAgentGraph(DualNodeAgent):
                  use_tool_schema: bool = True,
                  callbacks: list[AsyncCallbackHandler] | None = None,
                  detailed_logs: bool = False,
+                 log_response_max_chars: int = 1000,
                  retry_agent_response_parsing_errors: bool = True,
                  parse_agent_response_max_retries: int = 1,
                  tool_call_max_retries: int = 1,
                  pass_tool_call_errors_to_agent: bool = True):
-        super().__init__(llm=llm, tools=tools, callbacks=callbacks, detailed_logs=detailed_logs)
+        super().__init__(llm=llm,
+                         tools=tools,
+                         callbacks=callbacks,
+                         detailed_logs=detailed_logs,
+                         log_response_max_chars=log_response_max_chars)
         self.parse_agent_response_max_retries = (parse_agent_response_max_retries
                                                  if retry_agent_response_parsing_errors else 1)
         self.tool_call_max_retries = tool_call_max_retries

nat/agent/react_agent/register.py

@@ -17,6 +17,7 @@ import logging
 
 from pydantic import AliasChoices
 from pydantic import Field
+from pydantic import PositiveInt
 
 from nat.builder.builder import Builder
 from nat.builder.framework_enum import LLMFrameworkEnum
@@ -65,6 +66,8 @@ class ReActAgentWorkflowConfig(FunctionBaseConfig, name="react_agent"):
         default=None,
         description="Provides the SYSTEM_PROMPT to use with the agent")  # defaults to SYSTEM_PROMPT in prompt.py
     max_history: int = Field(default=15, description="Maximum number of messages to keep in the conversation history.")
+    log_response_max_chars: PositiveInt = Field(
+        default=1000, description="Maximum number of characters to display in logs when logging tool responses.")
     use_openai_api: bool = Field(default=False,
                                  description=("Use OpenAI API for the input/output types to the function. "
                                               "If False, strings will be used."))
@@ -100,6 +103,7 @@ async def react_agent_workflow(config: ReActAgentWorkflowConfig, builder: Builde
         tools=tools,
         use_tool_schema=config.include_tool_input_schema_in_tool_description,
         detailed_logs=config.verbose,
+        log_response_max_chars=config.log_response_max_chars,
         retry_agent_response_parsing_errors=config.retry_agent_response_parsing_errors,
         parse_agent_response_max_retries=config.parse_agent_response_max_retries,
         tool_call_max_retries=config.tool_call_max_retries,

nat/agent/rewoo_agent/agent.py

@@ -66,8 +66,13 @@ class ReWOOAgentGraph(BaseAgent):
                  tools: list[BaseTool],
                  use_tool_schema: bool = True,
                  callbacks: list[AsyncCallbackHandler] | None = None,
-                 detailed_logs: bool = False):
-        super().__init__(llm=llm, tools=tools, callbacks=callbacks, detailed_logs=detailed_logs)
+                 detailed_logs: bool = False,
+                 log_response_max_chars: int = 1000):
+        super().__init__(llm=llm,
+                         tools=tools,
+                         callbacks=callbacks,
+                         detailed_logs=detailed_logs,
+                         log_response_max_chars=log_response_max_chars)
 
         logger.debug(
             "%s Filling the prompt variables 'tools' and 'tool_names', using the tools provided in the config.",

nat/agent/rewoo_agent/register.py

@@ -17,6 +17,7 @@ import logging
 
 from pydantic import AliasChoices
 from pydantic import Field
+from pydantic import PositiveInt
 
 from nat.builder.builder import Builder
 from nat.builder.framework_enum import LLMFrameworkEnum
@@ -52,6 +53,8 @@ class ReWOOAgentWorkflowConfig(FunctionBaseConfig, name="rewoo_agent"):
         default=None,
         description="Provides the SOLVER_PROMPT to use with the agent")  # defaults to SOLVER_PROMPT in prompt.py
     max_history: int = Field(default=15, description="Maximum number of messages to keep in the conversation history.")
+    log_response_max_chars: PositiveInt = Field(
+        default=1000, description="Maximum number of characters to display in logs when logging tool responses.")
     use_openai_api: bool = Field(default=False,
                                  description=("Use OpenAI API for the input/output types to the function. "
                                               "If False, strings will be used."))
@@ -113,7 +116,8 @@ async def rewoo_agent_workflow(config: ReWOOAgentWorkflowConfig, builder: Builde
                                                  solver_prompt=solver_prompt,
                                                  tools=tools,
                                                  use_tool_schema=config.include_tool_input_schema_in_tool_description,
-                                                 detailed_logs=config.verbose).build_graph()
+                                                 detailed_logs=config.verbose,
+                                                 log_response_max_chars=config.log_response_max_chars).build_graph()
 
     async def _response_fn(input_message: ChatRequest) -> ChatResponse:
         try:

nat/agent/tool_calling_agent/agent.py

@@ -55,9 +55,14 @@ class ToolCallAgentGraph(DualNodeAgent):
         prompt: str | None = None,
         callbacks: list[AsyncCallbackHandler] = None,
         detailed_logs: bool = False,
+        log_response_max_chars: int = 1000,
         handle_tool_errors: bool = True,
     ):
-        super().__init__(llm=llm, tools=tools, callbacks=callbacks, detailed_logs=detailed_logs)
+        super().__init__(llm=llm,
+                         tools=tools,
+                         callbacks=callbacks,
+                         detailed_logs=detailed_logs,
+                         log_response_max_chars=log_response_max_chars)
         # some LLMs support tool calling
         # these models accept the tool's input schema and decide when to use a tool based on the input's relevance
         try:

nat/agent/tool_calling_agent/register.py

@@ -16,6 +16,7 @@
 import logging
 
 from pydantic import Field
+from pydantic import PositiveInt
 
 from nat.builder.builder import Builder
 from nat.builder.framework_enum import LLMFrameworkEnum
@@ -41,6 +42,8 @@ class ToolCallAgentWorkflowConfig(FunctionBaseConfig, name="tool_calling_agent")
     handle_tool_errors: bool = Field(default=True, description="Specify ability to handle tool calling errors.")
     description: str = Field(default="Tool Calling Agent Workflow", description="Description of this functions use.")
     max_iterations: int = Field(default=15, description="Number of tool calls before stoping the tool calling agent.")
+    log_response_max_chars: PositiveInt = Field(
+        default=1000, description="Maximum number of characters to display in logs when logging tool responses.")
     system_prompt: str | None = Field(default=None, description="Provides the system prompt to use with the agent.")
     additional_instructions: str | None = Field(default=None,
                                                 description="Additional instructions appended to the system prompt.")
@@ -70,6 +73,7 @@ async def tool_calling_agent_workflow(config: ToolCallAgentWorkflowConfig, build
                                                     tools=tools,
                                                     prompt=prompt,
                                                     detailed_logs=config.verbose,
+                                                    log_response_max_chars=config.log_response_max_chars,
                                                     handle_tool_errors=config.handle_tool_errors).build_graph()
 
     async def _response_fn(input_message: str) -> str:

nat/builder/context.py

@@ -31,6 +31,7 @@ from nat.data_models.intermediate_step import IntermediateStep
 from nat.data_models.intermediate_step import IntermediateStepPayload
 from nat.data_models.intermediate_step import IntermediateStepType
 from nat.data_models.intermediate_step import StreamEventData
+from nat.data_models.intermediate_step import TraceMetadata
 from nat.data_models.invocation_node import InvocationNode
 from nat.runtime.user_metadata import RequestAttributes
 from nat.utils.reactive.subject import Subject
@@ -174,7 +175,10 @@ class Context:
         return self._context_state.user_message_id.get()
 
     @contextmanager
-    def push_active_function(self, function_name: str, input_data: typing.Any | None):
+    def push_active_function(self,
+                             function_name: str,
+                             input_data: typing.Any | None,
+                             metadata: dict[str, typing.Any] | TraceMetadata | None = None):
         """
         Set the 'active_function' in context, push an invocation node,
         AND create an OTel child span for that function call.
@@ -195,7 +199,8 @@ class Context:
             IntermediateStepPayload(UUID=current_function_id,
                                     event_type=IntermediateStepType.FUNCTION_START,
                                     name=function_name,
-                                    data=StreamEventData(input=input_data)))
+                                    data=StreamEventData(input=input_data),
+                                    metadata=metadata))
 
         manager = ActiveFunctionContextManager()
 

nat/cli/commands/object_store/__init__.py

@@ -0,0 +1,14 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025, 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.

nat/cli/commands/object_store/object_store.py

@@ -0,0 +1,227 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025, 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 asyncio
+import importlib
+import logging
+import mimetypes
+import time
+from pathlib import Path
+
+import click
+
+from nat.builder.workflow_builder import WorkflowBuilder
+from nat.data_models.object_store import ObjectStoreBaseConfig
+from nat.object_store.interfaces import ObjectStore
+from nat.object_store.models import ObjectStoreItem
+
+logger = logging.getLogger(__name__)
+
+STORE_CONFIGS = {
+    "s3": {
+        "module": "nat.plugins.s3.object_store", "config_class": "S3ObjectStoreClientConfig"
+    },
+    "mysql": {
+        "module": "nat.plugins.mysql.object_store", "config_class": "MySQLObjectStoreClientConfig"
+    },
+    "redis": {
+        "module": "nat.plugins.redis.object_store", "config_class": "RedisObjectStoreClientConfig"
+    }
+}
+
+
+def get_object_store_config(**kwargs) -> ObjectStoreBaseConfig:
+    """Process common object store arguments and return the config class"""
+    store_type = kwargs.pop("store_type")
+    config = STORE_CONFIGS[store_type]
+    module = importlib.import_module(config["module"])
+    config_class = getattr(module, config["config_class"])
+    return config_class(**kwargs)
+
+
+async def upload_file(object_store: ObjectStore, file_path: Path, key: str):
+    """
+    Upload a single file to object store.
+
+    Args:
+        object_store: The object store instance to use.
+        file_path: The path to the file to upload.
+        key: The key to upload the file to.
+    """
+    try:
+        data = await asyncio.to_thread(file_path.read_bytes)
+
+        item = ObjectStoreItem(data=data,
+                               content_type=mimetypes.guess_type(str(file_path))[0],
+                               metadata={
+                                   "original_filename": file_path.name,
+                                   "file_size": str(len(data)),
+                                   "file_extension": file_path.suffix,
+                                   "upload_timestamp": str(int(time.time()))
+                               })
+
+        # Upload using upsert to allow overwriting
+        await object_store.upsert_object(key, item)
+        click.echo(f"✅ Uploaded: {file_path.name} -> {key}")
+
+    except Exception as e:
+        raise RuntimeError(f"Failed to upload {file_path.name}:\n{e}") from e
+
+
+def object_store_command_decorator(async_func):
+    """
+    Decorator that handles the common object store command pattern.
+
+    The decorated function should take (store: ObjectStore, kwargs) as parameters
+    and return an exit code (0 for success).
+    """
+
+    @click.pass_context
+    def wrapper(ctx: click.Context, **kwargs):
+        config = ctx.obj["store_config"]
+
+        async def work():
+            async with WorkflowBuilder() as builder:
+                await builder.add_object_store(name="store", config=config)
+                store = await builder.get_object_store_client("store")
+                return await async_func(store, **kwargs)
+
+        try:
+            exit_code = asyncio.run(work())
+        except Exception as e:
+            raise click.ClickException(f"Command failed: {e}") from e
+        if exit_code != 0:
+            raise click.ClickException(f"Command failed with exit code {exit_code}")
+        return exit_code
+
+    return wrapper
+
+
+@click.command(name="upload", help="Upload a directory to an object store.")
+@click.argument("local_dir",
+                type=click.Path(exists=True, file_okay=False, dir_okay=True, path_type=Path),
+                required=True)
+@click.help_option("--help", "-h")
+@object_store_command_decorator
+async def upload_command(store: ObjectStore, local_dir: Path, **_kwargs):
+    """
+    Upload a directory to an object store.
+
+    Args:
+        local_dir: The local directory to upload.
+        store: The object store to use.
+        _kwargs: Additional keyword arguments.
+    """
+    try:
+        click.echo(f"📁 Processing directory: {local_dir}")
+        file_count = 0
+
+        # Process each file recursively
+        for file_path in local_dir.rglob("*"):
+            if file_path.is_file():
+                key = file_path.relative_to(local_dir).as_posix()
+                await upload_file(store, file_path, key)
+                file_count += 1
+
+        click.echo(f"✅ Directory uploaded successfully! {file_count} files uploaded.")
+        return 0
+
+    except Exception as e:
+        raise click.ClickException(f"❌ Failed to upload directory {local_dir}:\n  {e}") from e
+
+
+@click.command(name="delete", help="Delete files from an object store.")
+@click.argument("keys", type=str, required=True, nargs=-1)
+@click.help_option("--help", "-h")
+@object_store_command_decorator
+async def delete_command(store: ObjectStore, keys: list[str], **_kwargs):
+    """
+    Delete files from an object store.
+
+    Args:
+        store: The object store to use.
+        keys: The keys to delete.
+        _kwargs: Additional keyword arguments.
+    """
+    deleted_count = 0
+    failed_count = 0
+    for key in keys:
+        try:
+            await store.delete_object(key)
+            click.echo(f"✅ Deleted: {key}")
+            deleted_count += 1
+        except Exception as e:
+            click.echo(f"❌ Failed to delete {key}: {e}")
+            failed_count += 1
+
+    click.echo(f"✅ Deletion completed! {deleted_count} keys deleted. {failed_count} keys failed to delete.")
+    return 0 if failed_count == 0 else 1
+
+
+@click.group(name="object-store", invoke_without_command=False, help="Manage object store operations.")
+def object_store_command(**_kwargs):
+    """Manage object store operations including uploading files and directories."""
+    pass
+
+
+def register_object_store_commands():
+
+    @click.group(name="s3", invoke_without_command=False, help="S3 object store operations.")
+    @click.argument("bucket_name", type=str, required=True)
+    @click.option("--endpoint-url", type=str, help="S3 endpoint URL")
+    @click.option("--access-key", type=str, help="S3 access key")
+    @click.option("--secret-key", type=str, help="S3 secret key")
+    @click.option("--region", type=str, help="S3 region")
+    @click.pass_context
+    def s3(ctx: click.Context, **kwargs):
+        ctx.ensure_object(dict)
+        ctx.obj["store_config"] = get_object_store_config(store_type="s3", **kwargs)
+
+    @click.group(name="mysql", invoke_without_command=False, help="MySQL object store operations.")
+    @click.argument("bucket_name", type=str, required=True)
+    @click.option("--host", type=str, help="MySQL host")
+    @click.option("--port", type=int, help="MySQL port")
+    @click.option("--db", type=str, help="MySQL database name")
+    @click.option("--username", type=str, help="MySQL username")
+    @click.option("--password", type=str, help="MySQL password")
+    @click.pass_context
+    def mysql(ctx: click.Context, **kwargs):
+        ctx.ensure_object(dict)
+        ctx.obj["store_config"] = get_object_store_config(store_type="mysql", **kwargs)
+
+    @click.group(name="redis", invoke_without_command=False, help="Redis object store operations.")
+    @click.argument("bucket_name", type=str, required=True)
+    @click.option("--host", type=str, help="Redis host")
+    @click.option("--port", type=int, help="Redis port")
+    @click.option("--db", type=int, help="Redis db")
+    @click.pass_context
+    def redis(ctx: click.Context, **kwargs):
+        ctx.ensure_object(dict)
+        ctx.obj["store_config"] = get_object_store_config(store_type="redis", **kwargs)
+
+    commands = {"s3": s3, "mysql": mysql, "redis": redis}
+
+    for store_type, config in STORE_CONFIGS.items():
+        try:
+            importlib.import_module(config["module"])
+            command = commands[store_type]
+            object_store_command.add_command(command, name=store_type)
+            command.add_command(upload_command, name="upload")
+            command.add_command(delete_command, name="delete")
+        except ImportError:
+            pass
+
+
+register_object_store_commands()

nat/cli/entrypoint.py

@@ -33,6 +33,7 @@ import nest_asyncio
 from .commands.configure.configure import configure_command
 from .commands.evaluate import eval_command
 from .commands.info.info import info_command
+from .commands.object_store.object_store import object_store_command
 from .commands.registry.registry import registry_command
 from .commands.sizing.sizing import sizing
 from .commands.start import start_command
@@ -107,11 +108,12 @@ 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(object_store_command, name="object-store")
 
 # 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
-cli.add_command(start_command.get_command(None, "mcp"), name="mcp")
+cli.add_command(start_command.get_command(None, "mcp"), name="mcp")  # type: ignore
 
 
 @cli.result_callback()

nat/data_models/gated_field_mixin.py

@@ -16,8 +16,6 @@
 from collections.abc import Sequence
 from dataclasses import dataclass
 from re import Pattern
-from typing import Generic
-from typing import TypeVar
 
 from pydantic import model_validator
 
@@ -33,10 +31,7 @@ class GatedFieldMixinConfig:
     keys: Sequence[str]
 
 
-T = TypeVar("T")
-
-
-class GatedFieldMixin(Generic[T]):
+class GatedFieldMixin:
     """
     A mixin that gates a field based on specified keys.
 
@@ -46,7 +41,7 @@ class GatedFieldMixin(Generic[T]):
     ----------
     field_name: `str`
                 The name of the field.
-    default_if_supported: `T | None`
+    default_if_supported: `object | None`
                           The default value of the field if it is supported for the key.
     keys: `Sequence[str]`
           A sequence of keys that are used to validate the field.
@@ -61,7 +56,7 @@ class GatedFieldMixin(Generic[T]):
     def __init_subclass__(
         cls,
         field_name: str | None = None,
-        default_if_supported: T | None = None,
+        default_if_supported: object | None = None,
         keys: Sequence[str] | None = None,
         unsupported: Sequence[Pattern[str]] | None = None,
         supported: Sequence[Pattern[str]] | None = None,
@@ -90,7 +85,7 @@ class GatedFieldMixin(Generic[T]):
     def _setup_direct_mixin(
         cls,
         field_name: str,
-        default_if_supported: T | None,
+        default_if_supported: object | None,
         unsupported: Sequence[Pattern[str]] | None,
         supported: Sequence[Pattern[str]] | None,
         keys: Sequence[str],
@@ -135,7 +130,7 @@ class GatedFieldMixin(Generic[T]):
     def _create_gated_field_validator(
         cls,
         field_name: str,
-        default_if_supported: T | None,
+        default_if_supported: object | None,
         unsupported: Sequence[Pattern[str]] | None,
         supported: Sequence[Pattern[str]] | None,
         keys: Sequence[str],
@@ -167,16 +162,19 @@ class GatedFieldMixin(Generic[T]):
         keys: Sequence[str],
     ) -> bool:
         """Check if a specific field is supported based on its configuration and keys."""
+        seen = False
         for key in keys:
             if not hasattr(instance, key):
                 continue
+            seen = True
             value = str(getattr(instance, key))
             if supported is not None:
-                return any(p.search(value) for p in supported)
+                if any(p.search(value) for p in supported):
+                    return True
             elif unsupported is not None:
-                return not any(p.search(value) for p in unsupported)
-        # Default to supported if no model keys found
-        return True
+                if any(p.search(value) for p in unsupported):
+                    return False
+        return True if not seen else (unsupported is not None)
 
     @classmethod
     def _find_blocking_key(

nat/data_models/temperature_mixin.py

@@ -23,7 +23,7 @@ from nat.data_models.gated_field_mixin import GatedFieldMixin
 
 class TemperatureMixin(
         BaseModel,
-        GatedFieldMixin[float],
+        GatedFieldMixin,
         field_name="temperature",
         default_if_supported=0.0,
         keys=("model_name", "model", "azure_deployment"),

nat/data_models/thinking_mixin.py

@@ -0,0 +1,68 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025, 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 re
+
+from pydantic import BaseModel
+from pydantic import Field
+
+from nat.data_models.gated_field_mixin import GatedFieldMixin
+
+# The system prompt format for thinking is different for these, so we need to distinguish them here with two separate
+# regex patterns
+_NVIDIA_NEMOTRON_REGEX = re.compile(r"^nvidia/nvidia.*nemotron", re.IGNORECASE)
+_LLAMA_NEMOTRON_REGEX = re.compile(r"^nvidia/llama.*nemotron", re.IGNORECASE)
+_MODEL_KEYS = ("model_name", "model", "azure_deployment")
+
+
+class ThinkingMixin(
+        BaseModel,
+        GatedFieldMixin,
+        field_name="thinking",
+        default_if_supported=None,
+        keys=_MODEL_KEYS,
+        supported=(_NVIDIA_NEMOTRON_REGEX, _LLAMA_NEMOTRON_REGEX),
+):
+    """
+    Mixin class for thinking configuration. Only supported on Nemotron models.
+
+    Attributes:
+        thinking: Whether to enable thinking. Defaults to None when supported on the model.
+    """
+    thinking: bool | None = Field(
+        default=None,
+        description="Whether to enable thinking. Defaults to None when supported on the model.",
+        exclude=True,
+    )
+
+    @property
+    def thinking_system_prompt(self) -> str | None:
+        """
+        Returns the system prompt to use for thinking.
+        For NVIDIA Nemotron, returns "/think" if enabled, else "/no_think".
+        For Llama Nemotron, returns "detailed thinking on" if enabled, else "detailed thinking off".
+        If thinking is not supported on the model, returns None.
+
+        Returns:
+            str | None: The system prompt to use for thinking.
+        """
+        if self.thinking is None:
+            return None
+        for key in _MODEL_KEYS:
+            if hasattr(self, key):
+                if _NVIDIA_NEMOTRON_REGEX.match(getattr(self, key)):
+                    return "/think" if self.thinking else "/no_think"
+                elif _LLAMA_NEMOTRON_REGEX.match(getattr(self, key)):
+                    return f"detailed thinking {'on' if self.thinking else 'off'}"

nat/data_models/top_p_mixin.py

@@ -23,7 +23,7 @@ from nat.data_models.gated_field_mixin import GatedFieldMixin
 
 class TopPMixin(
         BaseModel,
-        GatedFieldMixin[float],
+        GatedFieldMixin,
         field_name="top_p",
         default_if_supported=1.0,
         keys=("model_name", "model", "azure_deployment"),