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

nat/cli/commands/info/list_mcp.py

@@ -23,14 +23,36 @@ import click
 from pydantic import BaseModel
 
 from nat.tool.mcp.exceptions import MCPError
-from nat.tool.mcp.mcp_client import MCPBuilder
 from nat.utils.exception_handlers.mcp import format_mcp_error
 
 # Suppress verbose logs from mcp.client.sse and httpx
 logging.getLogger("mcp.client.sse").setLevel(logging.WARNING)
 logging.getLogger("httpx").setLevel(logging.WARNING)
 
-logger = logging.getLogger(__name__)
+
+def validate_transport_cli_args(transport: str, command: str | None, args: str | None, env: str | None) -> bool:
+    """
+    Validate transport and parameter combinations, returning False if invalid.
+
+    Args:
+        transport: The transport type ('sse', 'stdio', or 'streamable-http')
+        command: Command for stdio transport
+        args: Arguments for stdio transport
+        env: Environment variables for stdio transport
+
+    Returns:
+        bool: True if valid, False if invalid (error message already displayed)
+    """
+    if transport == 'stdio':
+        if not command:
+            click.echo("--command is required when using stdio client type", err=True)
+            return False
+    elif transport in ['sse', 'streamable-http']:
+        if command or args or env:
+            click.echo("--command, --args, and --env are not allowed when using sse or streamable-http client type",
+                       err=True)
+            return False
+    return True
 
 
 class MCPPingResult(BaseModel):
@@ -64,12 +86,20 @@ def format_tool(tool: Any) -> dict[str, str | None]:
     description = getattr(tool, 'description', '')
     input_schema = getattr(tool, 'input_schema', None) or getattr(tool, 'inputSchema', None)
 
-    schema_str = None
-    if input_schema:
-        if hasattr(input_schema, "schema_json"):
-            schema_str = input_schema.schema_json(indent=2)
-        else:
-            schema_str = str(input_schema)
+    # Normalize schema to JSON string
+    if input_schema is None:
+        return {
+            "name": name,
+            "description": description,
+            "input_schema": None,
+        }
+    elif hasattr(input_schema, "schema_json"):
+        schema_str = input_schema.schema_json(indent=2)
+    elif isinstance(input_schema, dict):
+        schema_str = json.dumps(input_schema, indent=2)
+    else:
+        # Final fallback: attempt to dump stringified version wrapped as JSON string
+        schema_str = json.dumps({"raw": str(input_schema)}, indent=2)
 
     return {
         "name": name,
@@ -100,8 +130,8 @@ def print_tool(tool_dict: dict[str, str | None], detail: bool = False) -> None:
         click.echo("-" * 60)
 
 
-async def list_tools_and_schemas(url: str, tool_name: str | None = None) -> list[dict[str, str | None]]:
-    """List MCP tools using MCPBuilder with structured exception handling.
+async def list_tools_and_schemas(command, url, tool_name=None, transport='sse', args=None, env=None):
+    """List MCP tools using NAT MCPClient with structured exception handling.
 
     Args:
         url (str): MCP server URL to connect to
@@ -115,20 +145,35 @@ async def list_tools_and_schemas(url: str, tool_name: str | None = None) -> list
     Raises:
         MCPError: Caught internally and logged, returns empty list instead
     """
-    builder = MCPBuilder(url=url)
+    from nat.tool.mcp.mcp_client_base import MCPSSEClient
+    from nat.tool.mcp.mcp_client_base import MCPStdioClient
+    from nat.tool.mcp.mcp_client_base import MCPStreamableHTTPClient
+
+    if args is None:
+        args = []
+
     try:
-        if tool_name:
-            tool = await builder.get_tool(tool_name)
-            return [format_tool(tool)]
-        tools = await builder.get_tools()
-        return [format_tool(tool) for tool in tools.values()]
+        if transport == 'stdio':
+            client = MCPStdioClient(command=command, args=args, env=env)
+        elif transport == 'streamable-http':
+            client = MCPStreamableHTTPClient(url=url)
+        else:  # sse
+            client = MCPSSEClient(url=url)
+
+        async with client:
+            if tool_name:
+                tool = await client.get_tool(tool_name)
+                return [format_tool(tool)]
+            else:
+                tools = await client.get_tools()
+                return [format_tool(tool) for tool in tools.values()]
     except MCPError as e:
         format_mcp_error(e, include_traceback=False)
         return []
 
 
-async def list_tools_direct(url: str, tool_name: str | None = None) -> list[dict[str, str | None]]:
-    """List MCP tools using direct MCP protocol with exception conversion.
+async def list_tools_direct(command, url, tool_name=None, transport='sse', args=None, env=None):
+    """List MCP tools using direct MCP protocol with structured exception handling.
 
     Bypasses MCPBuilder and uses raw MCP ClientSession and SSE client directly.
     Converts raw exceptions to structured MCPErrors for consistent user experience.
@@ -147,25 +192,51 @@ async def list_tools_direct(url: str, tool_name: str | None = None) -> list[dict
         This function handles ExceptionGroup by extracting the most relevant exception
         and converting it to MCPError for consistent error reporting.
     """
+    if args is None:
+        args = []
     from mcp import ClientSession
     from mcp.client.sse import sse_client
+    from mcp.client.stdio import StdioServerParameters
+    from mcp.client.stdio import stdio_client
+    from mcp.client.streamable_http import streamablehttp_client
 
     try:
-        async with sse_client(url=url) as (read, write):
+        if transport == 'stdio':
+
+            def get_stdio_client():
+                return stdio_client(server=StdioServerParameters(command=command, args=args, env=env))
+
+            client = get_stdio_client
+        elif transport == 'streamable-http':
+
+            def get_streamable_http_client():
+                return streamablehttp_client(url=url)
+
+            client = get_streamable_http_client
+        else:
+
+            def get_sse_client():
+                return sse_client(url=url)
+
+            client = get_sse_client
+
+        async with client() as ctx:
+            read, write = (ctx[0], ctx[1]) if isinstance(ctx, tuple) else ctx
             async with ClientSession(read, write) as session:
                 await session.initialize()
                 response = await session.list_tools()
 
-                tools = []
-                for tool in response.tools:
-                    if tool_name:
-                        if tool.name == tool_name:
-                            return [format_tool(tool)]
-                    else:
-                        tools.append(format_tool(tool))
-                if tool_name and not tools:
-                    click.echo(f"[INFO] Tool '{tool_name}' not found.")
-                return tools
+        tools = []
+        for tool in response.tools:
+            if tool_name:
+                if tool.name == tool_name:
+                    tools.append(format_tool(tool))
+            else:
+                tools.append(format_tool(tool))
+
+        if tool_name and not tools:
+            click.echo(f"[INFO] Tool '{tool_name}' not found.")
+        return tools
     except Exception as e:
         # Convert raw exceptions to structured MCPError for consistency
         from nat.utils.exception_handlers.mcp import convert_to_mcp_error
@@ -181,7 +252,12 @@ async def list_tools_direct(url: str, tool_name: str | None = None) -> list[dict
         return []
 
 
-async def ping_mcp_server(url: str, timeout: int) -> MCPPingResult:
+async def ping_mcp_server(url: str,
+                          timeout: int,
+                          transport: str = 'streamable-http',
+                          command: str | None = None,
+                          args: list[str] | None = None,
+                          env: dict[str, str] | None = None) -> MCPPingResult:
     """Ping an MCP server to check if it's responsive.
 
     Args:
@@ -193,18 +269,29 @@ async def ping_mcp_server(url: str, timeout: int) -> MCPPingResult:
     """
     from mcp.client.session import ClientSession
     from mcp.client.sse import sse_client
+    from mcp.client.stdio import StdioServerParameters
+    from mcp.client.stdio import stdio_client
+    from mcp.client.streamable_http import streamablehttp_client
 
     async def _ping_operation():
-        async with sse_client(url) as (read, write):
+        # Select transport
+        if transport == 'stdio':
+            stdio_args_local: list[str] = args or []
+            if not command:
+                raise RuntimeError("--command is required for stdio transport")
+            client_ctx = stdio_client(server=StdioServerParameters(command=command, args=stdio_args_local, env=env))
+        elif transport == 'sse':
+            client_ctx = sse_client(url)
+        else:  # streamable-http
+            client_ctx = streamablehttp_client(url=url)
+
+        async with client_ctx as ctx:
+            read, write = (ctx[0], ctx[1]) if isinstance(ctx, tuple) else ctx
             async with ClientSession(read, write) as session:
-                # Initialize the session
                 await session.initialize()
 
-                # Record start time just before ping
                 start_time = time.time()
-                # Send ping request
                 await session.send_ping()
-
                 end_time = time.time()
                 response_time_ms = round((end_time - start_time) * 1000, 2)
 
@@ -226,12 +313,24 @@ async def ping_mcp_server(url: str, timeout: int) -> MCPPingResult:
 
 @click.group(invoke_without_command=True, help="List tool names (default), or show details with --detail or --tool.")
 @click.option('--direct', is_flag=True, help='Bypass MCPBuilder and use direct MCP protocol')
-@click.option('--url', default='http://localhost:9901/sse', show_default=True, help='MCP server URL')
+@click.option(
+    '--url',
+    default='http://localhost:9901/mcp',
+    show_default=True,
+    help='MCP server URL (e.g. http://localhost:8080/mcp for streamable-http, http://localhost:8080/sse for sse)')
+@click.option('--transport',
+              type=click.Choice(['sse', 'stdio', 'streamable-http']),
+              default='streamable-http',
+              show_default=True,
+              help='Type of client to use (default: streamable-http, backwards compatible with sse)')
+@click.option('--command', help='For stdio: The command to run (e.g. mcp-server)')
+@click.option('--args', help='For stdio: Additional arguments for the command (space-separated)')
+@click.option('--env', help='For stdio: Environment variables in KEY=VALUE format (space-separated)')
 @click.option('--tool', default=None, help='Get details for a specific tool by name')
 @click.option('--detail', is_flag=True, help='Show full details for all tools')
 @click.option('--json-output', is_flag=True, help='Output tool metadata in JSON format')
 @click.pass_context
-def list_mcp(ctx: click.Context, direct: bool, url: str, tool: str | None, detail: bool, json_output: bool) -> None:
+def list_mcp(ctx, direct, url, transport, command, args, env, tool, detail, json_output):
     """List MCP tool names (default) or show detailed tool information.
 
     Use --detail for full output including descriptions and input schemas.
@@ -242,7 +341,7 @@ def list_mcp(ctx: click.Context, direct: bool, url: str, tool: str | None, detai
     Args:
         ctx (click.Context): Click context object for command invocation
         direct (bool): Whether to bypass MCPBuilder and use direct MCP protocol
-        url (str): MCP server URL to connect to (default: http://localhost:9901/sse)
+        url (str): MCP server URL to connect to (default: http://localhost:9901/mcp)
         tool (str | None): Optional specific tool name to retrieve detailed info for
         detail (bool): Whether to show full details (description + schema) for all tools
         json_output (bool): Whether to output tool metadata in JSON format instead of text
@@ -256,44 +355,81 @@ def list_mcp(ctx: click.Context, direct: bool, url: str, tool: str | None, detai
     """
     if ctx.invoked_subcommand is not None:
         return
+
+    if not validate_transport_cli_args(transport, command, args, env):
+        return
+
+    if transport in ['sse', 'streamable-http']:
+        if not url:
+            click.echo("[ERROR] --url is required when using sse or streamable-http client type", err=True)
+            return
+
+    stdio_args = args.split() if args else []
+    stdio_env = dict(var.split('=', 1) for var in env.split()) if env else None
+
     fetcher = list_tools_direct if direct else list_tools_and_schemas
-    tools = asyncio.run(fetcher(url, tool))
+    tools = asyncio.run(fetcher(command, url, tool, transport, stdio_args, stdio_env))
 
     if json_output:
         click.echo(json.dumps(tools, indent=2))
     elif tool:
-        for tool_dict in tools:
+        for tool_dict in (tools or []):
             print_tool(tool_dict, detail=True)
     elif detail:
-        for tool_dict in tools:
+        for tool_dict in (tools or []):
             print_tool(tool_dict, detail=True)
     else:
-        for tool_dict in tools:
+        for tool_dict in (tools or []):
             click.echo(tool_dict.get('name', 'Unknown tool'))
 
 
 @list_mcp.command()
-@click.option('--url', default='http://localhost:9901/sse', show_default=True, help='MCP server URL')
+@click.option(
+    '--url',
+    default='http://localhost:9901/mcp',
+    show_default=True,
+    help='MCP server URL (e.g. http://localhost:8080/mcp for streamable-http, http://localhost:8080/sse for sse)')
+@click.option('--transport',
+              type=click.Choice(['sse', 'stdio', 'streamable-http']),
+              default='streamable-http',
+              show_default=True,
+              help='Type of client to use for ping')
+@click.option('--command', help='For stdio: The command to run (e.g. mcp-server)')
+@click.option('--args', help='For stdio: Additional arguments for the command (space-separated)')
+@click.option('--env', help='For stdio: Environment variables in KEY=VALUE format (space-separated)')
 @click.option('--timeout', default=60, show_default=True, help='Timeout in seconds for ping request')
 @click.option('--json-output', is_flag=True, help='Output ping result in JSON format')
-def ping(url: str, timeout: int, json_output: bool) -> None:
+def ping(url: str,
+         transport: str,
+         command: str | None,
+         args: str | None,
+         env: str | None,
+         timeout: int,
+         json_output: bool) -> None:
     """Ping an MCP server to check if it's responsive.
 
     This command sends a ping request to the MCP server and measures the response time.
     It's useful for health checks and monitoring server availability.
 
     Args:
-        url (str): MCP server URL to ping (default: http://localhost:9901/sse)
+        url (str): MCP server URL to ping (default: http://localhost:9901/mcp)
         timeout (int): Timeout in seconds for the ping request (default: 60)
         json_output (bool): Whether to output the result in JSON format
 
     Examples:
         nat info mcp ping                                    # Ping default server
-        nat info mcp ping --url http://custom-server:9901/sse # Ping custom server
+        nat info mcp ping --url http://custom-server:9901/mcp # Ping custom server
         nat info mcp ping --timeout 10                      # Use 10 second timeout
         nat info mcp ping --json-output                     # Get JSON format output
     """
-    result = asyncio.run(ping_mcp_server(url, timeout))
+    # Validate combinations similar to parent command
+    if not validate_transport_cli_args(transport, command, args, env):
+        return
+
+    stdio_args = args.split() if args else []
+    stdio_env = dict(var.split('=', 1) for var in env.split()) if env else None
+
+    result = asyncio.run(ping_mcp_server(url, timeout, transport, command, stdio_args, stdio_env))
 
     if json_output:
         click.echo(result.model_dump_json(indent=2))

nat/cli/commands/registry/publish.py

@@ -40,7 +40,7 @@ async def publish_artifact(registry_handler_config: RegistryHandlerBaseConfig, p
         try:
             artifact = build_artifact(package_root=package_root)
         except Exception as e:
-            logger.exception("Error building artifact: %s", e, exc_info=True)
+            logger.exception("Error building artifact: %s", e)
             return
         await stack.enter_async_context(registry_handler.publish(artifact=artifact))
 
@@ -82,7 +82,7 @@ def publish(channel: str, config_file: str, package_root: str) -> None:
             logger.error("Publish channel '%s' has not been configured.", channel)
             return
     except Exception as e:
-        logger.exception("Error loading user settings: %s", e, exc_info=True)
+        logger.exception("Error loading user settings: %s", e)
         return
 
     asyncio.run(publish_artifact(registry_handler_config=publish_channel_config, package_root=package_root))

nat/cli/commands/registry/pull.py

@@ -66,7 +66,7 @@ async def pull_artifact(registry_handler_config: RegistryHandlerBaseConfig, pack
             validated_packages = PullRequestPackages(packages=package_list)
 
         except Exception as e:
-            logger.exception("Error processing package names: %s", e, exc_info=True)
+            logger.exception("Error processing package names: %s", e)
             return
 
         await stack.enter_async_context(registry_handler.pull(packages=validated_packages))
@@ -112,7 +112,7 @@ def pull(channel: str, config_file: str, packages: str) -> None:
             logger.error("Pull channel '%s' has not been configured.", channel)
             return
     except Exception as e:
-        logger.exception("Error loading user settings: %s", e, exc_info=True)
+        logger.exception("Error loading user settings: %s", e)
         return
 
     asyncio.run(pull_artifact(pull_channel_config, packages))

nat/cli/commands/registry/remove.py

@@ -41,7 +41,7 @@ async def remove_artifact(registry_handler_config: RegistryHandlerBaseConfig, pa
         try:
             package_name_list = PackageNameVersionList(**{"packages": packages})
         except Exception as e:
-            logger.exception("Invalid package format: '%s'", e, exc_info=True)
+            logger.exception("Invalid package format: '%s'", e)
 
         await stack.enter_async_context(registry_handler.remove(packages=package_name_list))
 
@@ -102,7 +102,7 @@ def remove(channel: str, config_file: str, packages: str) -> None:
             logger.error("Remove channel '%s' has not been configured.", channel)
             return
     except Exception as e:
-        logger.exception("Error loading user settings: %s", e, exc_info=True)
+        logger.exception("Error loading user settings: %s", e)
         return
 
     asyncio.run(remove_artifact(registry_handler_config=remove_channel_config, packages=packages_versions))

nat/cli/commands/registry/search.py

@@ -140,7 +140,7 @@ def search(config_file: str,
             logger.error("Search channel '%s' has not been configured.", channel)
             return
     except Exception as e:
-        logger.exception("Error loading user settings: %s", e, exc_info=True)
+        logger.exception("Error loading user settings: %s", e)
         return
 
     asyncio.run(

nat/cli/commands/start.py

@@ -102,12 +102,24 @@ class StartCommandGroup(click.Group):
                 raise ValueError(f"Invalid field '{name}'.Unions are only supported for optional parameters.")
 
             # Handle the types
-            if (issubclass(decomposed_type.root, Path)):
+            # Literal[...] -> map to click.Choice([...])
+            if (decomposed_type.origin is typing.Literal):
+                # typing.get_args returns the literal values; ensure they are strings for Click
+                literal_values = [str(v) for v in decomposed_type.args]
+                param_type = click.Choice(literal_values)
+
+            elif (issubclass(decomposed_type.root, Path)):
                 param_type = click.Path(exists=True, file_okay=True, dir_okay=False, path_type=Path)
 
             elif (issubclass(decomposed_type.root, (list, tuple, set))):
                 if (len(decomposed_type.args) == 1):
-                    param_type = decomposed_type.args[0]
+                    inner = DecomposedType(decomposed_type.args[0])
+                    # Support containers of Literal values -> multiple Choice
+                    if (inner.origin is typing.Literal):
+                        literal_values = [str(v) for v in inner.args]
+                        param_type = click.Choice(literal_values)
+                    else:
+                        param_type = inner.root
                 else:
                     param_type = None
 
@@ -224,7 +236,7 @@ class StartCommandGroup(click.Group):
             return asyncio.run(run_plugin())
 
         except Exception as e:
-            logger.error("Failed to initialize workflow", exc_info=True)
+            logger.error("Failed to initialize workflow")
             raise click.ClickException(str(e)) from e
 
     def get_command(self, ctx: click.Context, cmd_name: str) -> click.Command | None:

nat/cli/commands/uninstall.py

@@ -44,7 +44,7 @@ async def uninstall_packages(packages: list[dict[str, str]]) -> None:
     try:
         package_name_list = PackageNameVersionList(**{"packages": packages})
     except Exception as e:
-        logger.exception("Error validating package format: %s", e, exc_info=True)
+        logger.exception("Error validating package format: %s", e)
         return
 
     async with AsyncExitStack() as stack:

nat/cli/commands/workflow/workflow_commands.py

@@ -97,7 +97,7 @@ def find_package_root(package_name: str) -> Path | None:
         try:
             info = json.loads(direct_url)
         except json.JSONDecodeError:
-            logger.error("Malformed direct_url.json for package: %s", package_name)
+            logger.exception("Malformed direct_url.json for package: %s", package_name)
             return None
 
         if not info.get("dir_info", {}).get("editable"):
@@ -271,7 +271,7 @@ def create_command(workflow_name: str, install: bool, workflow_dir: str, descrip
 
         click.echo(f"Workflow '{workflow_name}' created successfully in '{new_workflow_dir}'.")
     except Exception as e:
-        logger.exception("An error occurred while creating the workflow: %s", e, exc_info=True)
+        logger.exception("An error occurred while creating the workflow: %s", e)
         click.echo(f"An error occurred while creating the workflow: {e}")
 
 
@@ -307,7 +307,7 @@ def reinstall_command(workflow_name):
 
         click.echo(f"Workflow '{workflow_name}' reinstalled successfully.")
     except Exception as e:
-        logger.exception("An error occurred while reinstalling the workflow: %s", e, exc_info=True)
+        logger.exception("An error occurred while reinstalling the workflow: %s", e)
         click.echo(f"An error occurred while reinstalling the workflow: {e}")
 
 
@@ -354,7 +354,7 @@ def delete_command(workflow_name: str):
 
         click.echo(f"Workflow '{workflow_name}' deleted successfully.")
     except Exception as e:
-        logger.exception("An error occurred while deleting the workflow: %s", e, exc_info=True)
+        logger.exception("An error occurred while deleting the workflow: %s", e)
         click.echo(f"An error occurred while deleting the workflow: {e}")
 
 

nat/data_models/discovery_metadata.py

@@ -177,7 +177,7 @@ class DiscoveryMetadata(BaseModel):
                 logger.warning("Package metadata not found for %s", distro_name)
                 version = ""
         except Exception as e:
-            logger.exception("Encountered issue extracting module metadata for %s: %s", config_type, e, exc_info=True)
+            logger.exception("Encountered issue extracting module metadata for %s: %s", config_type, e)
             return DiscoveryMetadata(status=DiscoveryStatusEnum.FAILURE)
 
         description = generate_config_type_docs(config_type=config_type)
@@ -217,7 +217,7 @@ class DiscoveryMetadata(BaseModel):
                 logger.warning("Package metadata not found for %s", distro_name)
                 version = ""
         except Exception as e:
-            logger.exception("Encountered issue extracting module metadata for %s: %s", fn, e, exc_info=True)
+            logger.exception("Encountered issue extracting module metadata for %s: %s", fn, e)
             return DiscoveryMetadata(status=DiscoveryStatusEnum.FAILURE)
 
         if isinstance(wrapper_type, LLMFrameworkEnum):
@@ -252,7 +252,7 @@ class DiscoveryMetadata(BaseModel):
                 description = ""
                 package_version = package_version or ""
         except Exception as e:
-            logger.exception("Encountered issue extracting module metadata for %s: %s", package_name, e, exc_info=True)
+            logger.exception("Encountered issue extracting module metadata for %s: %s", package_name, e)
             return DiscoveryMetadata(status=DiscoveryStatusEnum.FAILURE)
 
         return DiscoveryMetadata(package=package_name,
@@ -290,7 +290,7 @@ class DiscoveryMetadata(BaseModel):
                 logger.warning("Package metadata not found for %s", distro_name)
                 version = ""
         except Exception as e:
-            logger.exception("Encountered issue extracting module metadata for %s: %s", config_type, e, exc_info=True)
+            logger.exception("Encountered issue extracting module metadata for %s: %s", config_type, e)
             return DiscoveryMetadata(status=DiscoveryStatusEnum.FAILURE)
 
         wrapper_type = wrapper_type.value if isinstance(wrapper_type, LLMFrameworkEnum) else wrapper_type

nat/data_models/thinking_mixin.py

@@ -22,8 +22,7 @@ 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)
+_NEMOTRON_REGEX = re.compile(r"^nvidia/(llama|nvidia).*nemotron", re.IGNORECASE)
 _MODEL_KEYS = ("model_name", "model", "azure_deployment")
 
 
@@ -33,7 +32,7 @@ class ThinkingMixin(
         field_name="thinking",
         default_if_supported=None,
         keys=_MODEL_KEYS,
-        supported=(_NVIDIA_NEMOTRON_REGEX, _LLAMA_NEMOTRON_REGEX),
+        supported=(_NEMOTRON_REGEX, ),
 ):
     """
     Mixin class for thinking configuration. Only supported on Nemotron models.
@@ -52,7 +51,8 @@ class ThinkingMixin(
         """
         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".
+        For Llama Nemotron v1.5, returns "/think" if enabled, else "/no_think".
+        For Llama Nemotron v1.0, returns "detailed thinking on" if enabled, else "detailed thinking off".
         If thinking is not supported on the model, returns None.
 
         Returns:
@@ -60,9 +60,28 @@ class ThinkingMixin(
         """
         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)):
+            model = getattr(self, key, None)
+            if not isinstance(model, str) or model is None:
+                continue
+
+            # Normalize name to reduce checks
+            model = model.lower().translate(str.maketrans("_.", "--"))
+
+            if model.startswith("nvidia/nvidia"):
+                return "/think" if self.thinking else "/no_think"
+
+            if model.startswith("nvidia/llama"):
+                if "v1-0" in model or "v1-1" in model:
                     return f"detailed thinking {'on' if self.thinking else 'off'}"
+
+                if "v1-5" in model:
+                    # v1.5 models are updated to use the /think and /no_think system prompts
+                    return "/think" if self.thinking else "/no_think"
+
+                # Assume any other model is a newer model that uses the /think and /no_think system prompts
+                return "/think" if self.thinking else "/no_think"
+
+        # Unknown model
+        return None

nat/eval/evaluate.py

@@ -168,17 +168,17 @@ class EvaluationRun:
                 intermediate_future = None
 
                 try:
-
                     # Start usage stats and intermediate steps collection in parallel
                     intermediate_future = pull_intermediate()
                     runner_result = runner.result()
                     base_output = await runner_result
                     intermediate_steps = await intermediate_future
                 except NotImplementedError as e:
+                    logger.error("Failed to run the workflow: %s", e)
                     # raise original error
-                    raise e
+                    raise
                 except Exception as e:
-                    logger.exception("Failed to run the workflow: %s", e, exc_info=True)
+                    logger.exception("Failed to run the workflow: %s", e)
                     # stop processing if a workflow error occurs
                     self.workflow_interrupted = True
 
@@ -317,7 +317,7 @@ class EvaluationRun:
                 logger.info("Deleting old job directory: %s", dir_to_delete)
                 shutil.rmtree(dir_to_delete)
             except Exception as e:
-                logger.exception("Failed to delete old job directory: %s: %s", dir_to_delete, e, exc_info=True)
+                logger.exception("Failed to delete old job directory: %s: %s", dir_to_delete, e)
 
     def write_output(self, dataset_handler: DatasetHandler, profiler_results: ProfilerResults):
         workflow_output_file = self.eval_config.general.output_dir / "workflow_output.json"
@@ -367,7 +367,7 @@ class EvaluationRun:
 
             await self.weave_eval.alog_score(eval_output, evaluator_name)
         except Exception as e:
-            logger.exception("An error occurred while running evaluator %s: %s", evaluator_name, e, exc_info=True)
+            logger.exception("An error occurred while running evaluator %s: %s", evaluator_name, e)
 
     async def run_evaluators(self, evaluators: dict[str, Any]):
         """Run all configured evaluators asynchronously."""
@@ -380,7 +380,7 @@ class EvaluationRun:
         try:
             await asyncio.gather(*tasks)
         except Exception as e:
-            logger.exception("An error occurred while running evaluators: %s", e, exc_info=True)
+            logger.error("An error occurred while running evaluators: %s", e)
             raise
         finally:
             # Finish prediction loggers in Weave