nvidia-nat 1.3.0rc1__py3-none-any.whl → 1.3.0rc3__py3-none-any.whl

nat/agent/prompt_optimizer/register.py

@@ -51,7 +51,7 @@ async def prompt_optimizer_function(config: PromptOptimizerConfig, builder: Buil
         from .prompt import mutator_prompt
     except ImportError as exc:
         raise ImportError("langchain-core is not installed. Please install it to use MultiLLMPlanner.\n"
-                          "This error can be resolve by installing nvidia-nat[langchain]") from exc
+                          "This error can be resolve by installing \"nvidia-nat[langchain]\".") from exc
 
     llm = await builder.get_llm(config.optimizer_llm, wrapper_type=LLMFrameworkEnum.LANGCHAIN)
 
@@ -111,7 +111,7 @@ async def prompt_recombiner_function(config: PromptRecombinerConfig, builder: Bu
         from langchain_core.prompts import PromptTemplate
     except ImportError as exc:
         raise ImportError("langchain-core is not installed. Please install it to use MultiLLMPlanner.\n"
-                          "This error can be resolve by installing nvidia-nat[langchain].") from exc
+                          "This error can be resolve by installing \"nvidia-nat[langchain]\".") from exc
 
     llm = await builder.get_llm(config.optimizer_llm, wrapper_type=LLMFrameworkEnum.LANGCHAIN)
 

nat/agent/react_agent/register.py

@@ -24,7 +24,9 @@ from nat.builder.function_info import FunctionInfo
 from nat.cli.register_workflow import register_function
 from nat.data_models.agent import AgentBaseConfig
 from nat.data_models.api_server import ChatRequest
+from nat.data_models.api_server import ChatRequestOrMessage
 from nat.data_models.api_server import ChatResponse
+from nat.data_models.api_server import Usage
 from nat.data_models.component_ref import FunctionGroupRef
 from nat.data_models.component_ref import FunctionRef
 from nat.data_models.optimizable import OptimizableField
@@ -69,9 +71,6 @@ class ReActAgentWorkflowConfig(AgentBaseConfig, OptimizableMixin, name="react_ag
         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.")
-    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."))
     additional_instructions: str | None = OptimizableField(
         default=None,
         description="Additional instructions to provide to the agent in addition to the base prompt.",
@@ -117,21 +116,23 @@ async def react_agent_workflow(config: ReActAgentWorkflowConfig, builder: Builde
         pass_tool_call_errors_to_agent=config.pass_tool_call_errors_to_agent,
         normalize_tool_input_quotes=config.normalize_tool_input_quotes).build_graph()
 
-    async def _response_fn(input_message: ChatRequest) -> ChatResponse:
+    async def _response_fn(chat_request_or_message: ChatRequestOrMessage) -> ChatResponse | str:
         """
         Main workflow entry function for the ReAct Agent.
 
         This function invokes the ReAct Agent Graph and returns the response.
 
         Args:
-            input_message (ChatRequest): The input message to process
+            chat_request_or_message (ChatRequestOrMessage): The input message to process
 
         Returns:
-            ChatResponse: The response from the agent or error message
+            ChatResponse | str: The response from the agent or error message
         """
         try:
+            message = GlobalTypeConverter.get().convert(chat_request_or_message, to_type=ChatRequest)
+
             # initialize the starting state with the user query
-            messages: list[BaseMessage] = trim_messages(messages=[m.model_dump() for m in input_message.messages],
+            messages: list[BaseMessage] = trim_messages(messages=[m.model_dump() for m in message.messages],
                                                         max_tokens=config.max_history,
                                                         strategy="last",
                                                         token_counter=len,
@@ -149,21 +150,19 @@ async def react_agent_workflow(config: ReActAgentWorkflowConfig, builder: Builde
             # get and return the output from the state
             state = ReActGraphState(**state)
             output_message = state.messages[-1]
-            return ChatResponse.from_string(str(output_message.content))
-
+            content = str(output_message.content)
+
+            # Create usage statistics for the response
+            prompt_tokens = sum(len(str(msg.content).split()) for msg in message.messages)
+            completion_tokens = len(content.split()) if content else 0
+            total_tokens = prompt_tokens + completion_tokens
+            usage = Usage(prompt_tokens=prompt_tokens, completion_tokens=completion_tokens, total_tokens=total_tokens)
+            response = ChatResponse.from_string(content, usage=usage)
+            if chat_request_or_message.is_string:
+                return GlobalTypeConverter.get().convert(response, to_type=str)
+            return response
         except Exception as ex:
             logger.exception("%s ReAct Agent failed with exception: %s", AGENT_LOG_PREFIX, str(ex))
             raise RuntimeError
 
-    if (config.use_openai_api):
-        yield FunctionInfo.from_fn(_response_fn, description=config.description)
-    else:
-
-        async def _str_api_fn(input_message: str) -> str:
-            oai_input = GlobalTypeConverter.get().try_convert(input_message, to_type=ChatRequest)
-
-            oai_output = await _response_fn(oai_input)
-
-            return GlobalTypeConverter.get().try_convert(oai_output, to_type=str)
-
-        yield FunctionInfo.from_fn(_str_api_fn, description=config.description)
+    yield FunctionInfo.from_fn(_response_fn, description=config.description)

nat/agent/rewoo_agent/register.py

@@ -25,7 +25,9 @@ from nat.builder.function_info import FunctionInfo
 from nat.cli.register_workflow import register_function
 from nat.data_models.agent import AgentBaseConfig
 from nat.data_models.api_server import ChatRequest
+from nat.data_models.api_server import ChatRequestOrMessage
 from nat.data_models.api_server import ChatResponse
+from nat.data_models.api_server import Usage
 from nat.data_models.component_ref import FunctionGroupRef
 from nat.data_models.component_ref import FunctionRef
 from nat.utils.type_converter import GlobalTypeConverter
@@ -53,9 +55,6 @@ class ReWOOAgentWorkflowConfig(AgentBaseConfig, name="rewoo_agent"):
                                                description="The number of retries before raising a tool call error.",
                                                ge=1)
     max_history: int = Field(default=15, description="Maximum number of messages to keep in the conversation history.")
-    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."))
     additional_planner_instructions: str | None = Field(
         default=None,
         validation_alias=AliasChoices("additional_planner_instructions", "additional_instructions"),
@@ -124,21 +123,23 @@ async def rewoo_agent_workflow(config: ReWOOAgentWorkflowConfig, builder: Builde
         tool_call_max_retries=config.tool_call_max_retries,
         raise_tool_call_error=config.raise_tool_call_error).build_graph()
 
-    async def _response_fn(input_message: ChatRequest) -> ChatResponse:
+    async def _response_fn(chat_request_or_message: ChatRequestOrMessage) -> ChatResponse | str:
         """
         Main workflow entry function for the ReWOO Agent.
 
         This function invokes the ReWOO Agent Graph and returns the response.
 
         Args:
-            input_message (ChatRequest): The input message to process
+            chat_request_or_message (ChatRequestOrMessage): The input message to process
 
         Returns:
-            ChatResponse: The response from the agent or error message
+            ChatResponse | str: The response from the agent or error message
         """
         try:
+            message = GlobalTypeConverter.get().convert(chat_request_or_message, to_type=ChatRequest)
+
             # initialize the starting state with the user query
-            messages: list[BaseMessage] = trim_messages(messages=[m.model_dump() for m in input_message.messages],
+            messages: list[BaseMessage] = trim_messages(messages=[m.model_dump() for m in message.messages],
                                                         max_tokens=config.max_history,
                                                         strategy="last",
                                                         token_counter=len,
@@ -157,21 +158,18 @@ async def rewoo_agent_workflow(config: ReWOOAgentWorkflowConfig, builder: Builde
             # Ensure output_message is a string
             if isinstance(output_message, list | dict):
                 output_message = str(output_message)
-            return ChatResponse.from_string(output_message)
 
+            # Create usage statistics for the response
+            prompt_tokens = sum(len(str(msg.content).split()) for msg in message.messages)
+            completion_tokens = len(output_message.split()) if output_message else 0
+            total_tokens = prompt_tokens + completion_tokens
+            usage = Usage(prompt_tokens=prompt_tokens, completion_tokens=completion_tokens, total_tokens=total_tokens)
+            response = ChatResponse.from_string(output_message, usage=usage)
+            if chat_request_or_message.is_string:
+                return GlobalTypeConverter.get().convert(response, to_type=str)
+            return response
         except Exception as ex:
             logger.exception("ReWOO Agent failed with exception: %s", ex)
             raise RuntimeError
 
-    if (config.use_openai_api):
-        yield FunctionInfo.from_fn(_response_fn, description=config.description)
-
-    else:
-
-        async def _str_api_fn(input_message: str) -> str:
-            oai_input = GlobalTypeConverter.get().try_convert(input_message, to_type=ChatRequest)
-            oai_output = await _response_fn(oai_input)
-
-            return GlobalTypeConverter.get().try_convert(oai_output, to_type=str)
-
-        yield FunctionInfo.from_fn(_str_api_fn, description=config.description)
+    yield FunctionInfo.from_fn(_response_fn, description=config.description)

nat/agent/tool_calling_agent/register.py

@@ -23,8 +23,10 @@ from nat.builder.function_info import FunctionInfo
 from nat.cli.register_workflow import register_function
 from nat.data_models.agent import AgentBaseConfig
 from nat.data_models.api_server import ChatRequest
+from nat.data_models.api_server import ChatRequestOrMessage
 from nat.data_models.component_ref import FunctionGroupRef
 from nat.data_models.component_ref import FunctionRef
+from nat.utils.type_converter import GlobalTypeConverter
 
 logger = logging.getLogger(__name__)
 
@@ -81,21 +83,23 @@ async def tool_calling_agent_workflow(config: ToolCallAgentWorkflowConfig, build
                                                          handle_tool_errors=config.handle_tool_errors,
                                                          return_direct=return_direct_tools).build_graph()
 
-    async def _response_fn(input_message: ChatRequest) -> str:
+    async def _response_fn(chat_request_or_message: ChatRequestOrMessage) -> str:
         """
         Main workflow entry function for the Tool Calling Agent.
 
         This function invokes the Tool Calling Agent Graph and returns the response.
 
         Args:
-            input_message (ChatRequest): The input message to process
+            chat_request_or_message (ChatRequestOrMessage): The input message to process
 
         Returns:
             str: The response from the agent or error message
         """
         try:
+            message = GlobalTypeConverter.get().convert(chat_request_or_message, to_type=ChatRequest)
+
             # initialize the starting state with the user query
-            messages: list[BaseMessage] = trim_messages(messages=[m.model_dump() for m in input_message.messages],
+            messages: list[BaseMessage] = trim_messages(messages=[m.model_dump() for m in message.messages],
                                                         max_tokens=config.max_history,
                                                         strategy="last",
                                                         token_counter=len,

nat/authentication/oauth2/oauth2_auth_code_flow_provider.py

@@ -14,6 +14,7 @@
 # limitations under the License.
 
 import logging
+from collections.abc import Awaitable
 from collections.abc import Callable
 from datetime import UTC
 from datetime import datetime
@@ -35,10 +36,15 @@ logger = logging.getLogger(__name__)
 
 class OAuth2AuthCodeFlowProvider(AuthProviderBase[OAuth2AuthCodeFlowProviderConfig]):
 
-    def __init__(self, config: OAuth2AuthCodeFlowProviderConfig):
+    def __init__(self, config: OAuth2AuthCodeFlowProviderConfig, token_storage=None):
         super().__init__(config)
-        self._authenticated_tokens: dict[str, AuthResult] = {}
         self._auth_callback = None
+        # Always use token storage - defaults to in-memory if not provided
+        if token_storage is None:
+            from nat.plugins.mcp.auth.token_storage import InMemoryTokenStorage
+            self._token_storage = InMemoryTokenStorage()
+        else:
+            self._token_storage = token_storage
 
     async def _attempt_token_refresh(self, user_id: str, auth_result: AuthResult) -> AuthResult | None:
         refresh_token = auth_result.raw.get("refresh_token")
@@ -61,7 +67,7 @@ class OAuth2AuthCodeFlowProvider(AuthProviderBase[OAuth2AuthCodeFlowProviderConf
                 raw=new_token_data,
             )
 
-            self._authenticated_tokens[user_id] = new_auth_result
+            await self._token_storage.store(user_id, new_auth_result)
         except httpx.HTTPStatusError:
             return None
         except httpx.RequestError:
@@ -74,26 +80,30 @@ class OAuth2AuthCodeFlowProvider(AuthProviderBase[OAuth2AuthCodeFlowProviderConf
 
     def _set_custom_auth_callback(self,
                                   auth_callback: Callable[[OAuth2AuthCodeFlowProviderConfig, AuthFlowType],
-                                                          AuthenticatedContext]):
+                                                          Awaitable[AuthenticatedContext]]):
         self._auth_callback = auth_callback
 
     async def authenticate(self, user_id: str | None = None, **kwargs) -> AuthResult:
-        if user_id is None and hasattr(Context.get(), "metadata") and hasattr(
-                Context.get().metadata, "cookies") and Context.get().metadata.cookies is not None:
-            session_id = Context.get().metadata.cookies.get("nat-session", None)
+        context = Context.get()
+        if user_id is None and hasattr(context, "metadata") and hasattr(
+                context.metadata, "cookies") and context.metadata.cookies is not None:
+            session_id = context.metadata.cookies.get("nat-session", None)
             if not session_id:
                 raise RuntimeError("Authentication failed. No session ID found. Cannot identify user.")
 
             user_id = session_id
 
-        if user_id and user_id in self._authenticated_tokens:
-            auth_result = self._authenticated_tokens[user_id]
-            if not auth_result.is_expired():
-                return auth_result
+        if user_id:
+            # Try to retrieve from token storage
+            auth_result = await self._token_storage.retrieve(user_id)
+
+            if auth_result:
+                if not auth_result.is_expired():
+                    return auth_result
 
-            refreshed_auth_result = await self._attempt_token_refresh(user_id, auth_result)
-            if refreshed_auth_result:
-                return refreshed_auth_result
+                refreshed_auth_result = await self._attempt_token_refresh(user_id, auth_result)
+                if refreshed_auth_result:
+                    return refreshed_auth_result
 
         # Try getting callback from the context if that's not set, use the default callback
         try:
@@ -109,19 +119,22 @@ class OAuth2AuthCodeFlowProvider(AuthProviderBase[OAuth2AuthCodeFlowProviderConf
         except Exception as e:
             raise RuntimeError(f"Authentication callback failed: {e}") from e
 
-        auth_header = authenticated_context.headers.get("Authorization", "")
+        headers = authenticated_context.headers or {}
+        auth_header = headers.get("Authorization", "")
         if not auth_header.startswith("Bearer "):
             raise RuntimeError("Invalid Authorization header")
 
         token = auth_header.split(" ")[1]
 
+        # Safely access metadata
+        metadata = authenticated_context.metadata or {}
         auth_result = AuthResult(
             credentials=[BearerTokenCred(token=SecretStr(token))],
-            token_expires_at=authenticated_context.metadata.get("expires_at"),
-            raw=authenticated_context.metadata.get("raw_token"),
+            token_expires_at=metadata.get("expires_at"),
+            raw=metadata.get("raw_token") or {},
         )
 
         if user_id:
-            self._authenticated_tokens[user_id] = auth_result
+            await self._token_storage.store(user_id, auth_result)
 
         return auth_result

nat/builder/component_utils.py

@@ -153,7 +153,7 @@ def recursive_componentref_discovery(cls: TypedBaseModel, value: typing.Any,
         for v in value.values():
             yield from recursive_componentref_discovery(cls, v, decomposed_type.args[1])
     elif (issubclass(type(value), BaseModel)):
-        for field, field_info in value.model_fields.items():
+        for field, field_info in type(value).model_fields.items():
             field_data = getattr(value, field)
             yield from recursive_componentref_discovery(cls, field_data, field_info.annotation)
     if (decomposed_type.is_union):

nat/builder/context.py

@@ -67,6 +67,8 @@ class ContextState(metaclass=Singleton):
     def __init__(self):
         self.conversation_id: ContextVar[str | None] = ContextVar("conversation_id", default=None)
         self.user_message_id: ContextVar[str | None] = ContextVar("user_message_id", default=None)
+        self.workflow_run_id: ContextVar[str | None] = ContextVar("workflow_run_id", default=None)
+        self.workflow_trace_id: ContextVar[int | None] = ContextVar("workflow_trace_id", default=None)
         self.input_message: ContextVar[typing.Any] = ContextVar("input_message", default=None)
         self.user_manager: ContextVar[typing.Any] = ContextVar("user_manager", default=None)
         self._metadata: ContextVar[RequestAttributes | None] = ContextVar("request_attributes", default=None)
@@ -120,14 +122,14 @@ class Context:
     @property
     def input_message(self):
         """
-            Retrieves the input message from the context state.
+        Retrieves the input message from the context state.
 
-            The input_message property is used to access the message stored in the
-            context state. This property returns the message as it is currently
-            maintained in the context.
+        The input_message property is used to access the message stored in the
+        context state. This property returns the message as it is currently
+        maintained in the context.
 
-            Returns:
-                str: The input message retrieved from the context state.
+        Returns:
+            str: The input message retrieved from the context state.
         """
         return self._context_state.input_message.get()
 
@@ -196,6 +198,20 @@ class Context:
         """
         return self._context_state.user_message_id.get()
 
+    @property
+    def workflow_run_id(self) -> str | None:
+        """
+        Returns a stable identifier for the current workflow/agent invocation (UUID string).
+        """
+        return self._context_state.workflow_run_id.get()
+
+    @property
+    def workflow_trace_id(self) -> int | None:
+        """
+        Returns the 128-bit trace identifier for the current run, used as the OpenTelemetry trace_id.
+        """
+        return self._context_state.workflow_trace_id.get()
+
     @contextmanager
     def push_active_function(self,
                              function_name: str,

nat/builder/function.py

@@ -416,8 +416,9 @@ class FunctionGroup:
         """
         if not name.strip():
             raise ValueError("Function name cannot be empty or blank")
-        if not re.match(r"^[a-zA-Z0-9_-]+$", name):
-            raise ValueError(f"Function name can only contain letters, numbers, underscores, and hyphens: {name}")
+        if not re.match(r"^[a-zA-Z0-9_.-]+$", name):
+            raise ValueError(
+                f"Function name can only contain letters, numbers, underscores, periods, and hyphens: {name}")
         if name in self._functions:
             raise ValueError(f"Function {name} already exists in function group {self._instance_name}")
 

nat/builder/workflow_builder.py

@@ -156,6 +156,7 @@ class WorkflowBuilder(Builder, AbstractAsyncContextManager):
         self._registry = registry
 
         self._logging_handlers: dict[str, logging.Handler] = {}
+        self._removed_root_handlers: list[tuple[logging.Handler, int]] = []
         self._telemetry_exporters: dict[str, ConfiguredTelemetryExporter] = {}
 
         self._functions: dict[str, ConfiguredFunction] = {}
@@ -187,6 +188,15 @@ class WorkflowBuilder(Builder, AbstractAsyncContextManager):
         # Get the telemetry info from the config
         telemetry_config = self.general_config.telemetry
 
+        # If we have logging configuration, we need to manage the root logger properly
+        root_logger = logging.getLogger()
+
+        # Collect configured handler types to determine if we need to adjust existing handlers
+        # This is somewhat of a hack by inspecting the class name of the config object
+        has_console_handler = any(
+            hasattr(config, "__class__") and "console" in config.__class__.__name__.lower()
+            for config in telemetry_config.logging.values())
+
         for key, logging_config in telemetry_config.logging.items():
             # Use the same pattern as tracing, but for logging
             logging_info = self._registry.get_logging_method(type(logging_config))
@@ -200,7 +210,31 @@ class WorkflowBuilder(Builder, AbstractAsyncContextManager):
             self._logging_handlers[key] = handler
 
             # Now attach to NAT's root logger
-            logging.getLogger().addHandler(handler)
+            root_logger.addHandler(handler)
+
+        # If we added logging handlers, manage existing handlers appropriately
+        if self._logging_handlers:
+            min_handler_level = min((handler.level for handler in root_logger.handlers), default=logging.CRITICAL)
+
+            # Ensure the root logger level allows messages through
+            root_logger.level = max(root_logger.level, min_handler_level)
+
+            # If a console handler is configured, adjust or remove default CLI handlers
+            # to avoid duplicate output while preserving workflow visibility
+            if has_console_handler:
+                # Remove existing StreamHandlers that are not the newly configured ones
+                for handler in root_logger.handlers[:]:
+                    if type(handler) is logging.StreamHandler and handler not in self._logging_handlers.values():
+                        self._removed_root_handlers.append((handler, handler.level))
+                        root_logger.removeHandler(handler)
+            else:
+                # No console handler configured, but adjust existing handler levels
+                # to respect the minimum configured level for file/other handlers
+                for handler in root_logger.handlers[:]:
+                    if type(handler) is logging.StreamHandler:
+                        old_level = handler.level
+                        handler.setLevel(min_handler_level)
+                        self._removed_root_handlers.append((handler, old_level))
 
         # Add the telemetry exporters
         for key, telemetry_exporter_config in telemetry_config.tracing.items():
@@ -212,8 +246,17 @@ class WorkflowBuilder(Builder, AbstractAsyncContextManager):
 
         assert self._exit_stack is not None, "Exit stack not initialized"
 
-        for _, handler in self._logging_handlers.items():
-            logging.getLogger().removeHandler(handler)
+        root_logger = logging.getLogger()
+
+        # Remove custom logging handlers
+        for handler in self._logging_handlers.values():
+            root_logger.removeHandler(handler)
+
+        # Restore original handlers and their levels
+        for handler, old_level in self._removed_root_handlers:
+            if handler not in root_logger.handlers:
+                root_logger.addHandler(handler)
+            handler.setLevel(old_level)
 
         await self._exit_stack.__aexit__(*exc_details)
 

nat/cli/commands/mcp/mcp.py

@@ -194,7 +194,7 @@ async def _create_mcp_client_config(
     auth_user_id: str | None,
     auth_scopes: list[str] | None,
 ):
-    from nat.plugins.mcp.client_impl import MCPClientConfig
+    from nat.plugins.mcp.client_config import MCPClientConfig
 
     if url and transport == "streamable-http" and auth_redirect_uri:
         try:
@@ -236,8 +236,8 @@ async def list_tools_via_function_group(
     try:
         # Ensure the registration side-effects are loaded
         from nat.builder.workflow_builder import WorkflowBuilder
-        from nat.plugins.mcp.client_impl import MCPClientConfig
-        from nat.plugins.mcp.client_impl import MCPServerConfig
+        from nat.plugins.mcp.client_config import MCPClientConfig
+        from nat.plugins.mcp.client_config import MCPServerConfig
     except ImportError:
         click.echo(
             "MCP client functionality requires nvidia-nat-mcp package. Install with: uv pip install nvidia-nat-mcp",
@@ -297,7 +297,7 @@ async def list_tools_via_function_group(
             if fn is not None:
                 tools.append(to_tool_entry(full, fn))
         else:
-            for full, fn in (await fns).items():
+            for full, fn in fns.items():
                 tools.append(to_tool_entry(full, fn))
 
     return tools
@@ -826,8 +826,8 @@ async def call_tool_and_print(command: str | None,
 
     try:
         from nat.builder.workflow_builder import WorkflowBuilder
-        from nat.plugins.mcp.client_impl import MCPClientConfig
-        from nat.plugins.mcp.client_impl import MCPServerConfig
+        from nat.plugins.mcp.client_config import MCPClientConfig
+        from nat.plugins.mcp.client_config import MCPServerConfig
     except ImportError:
         click.echo(
             "MCP client functionality requires nvidia-nat-mcp package. Install with: uv pip install nvidia-nat-mcp",

nat/cli/commands/workflow/templates/config.yml.j2

@@ -1,15 +1,17 @@
-general:
-  logging:
-    console:
-      _type: console
-      level: WARN
+functions:
+  current_datetime:
+    _type: current_datetime
+  {{python_safe_workflow_name}}:
+    _type: {{python_safe_workflow_name}}
+    prefix: "Hello:"
 
-  front_end:
-    _type: fastapi
-
-  front_end:
-    _type: console
+llms:
+  nim_llm:
+    _type: nim
+    model_name: meta/llama-3.1-70b-instruct
+    temperature: 0.0
 
 workflow:
-  _type: {{workflow_name}}
-  parameter: default_value
+  _type: react_agent
+  llm_name: nim_llm
+  tool_names: [current_datetime, {{python_safe_workflow_name}}]

nat/cli/commands/workflow/templates/register.py.j2

@@ -1,4 +1,4 @@
 # flake8: noqa
 
-# Import any tools which need to be automatically registered here
-from {{package_name}} import {{workflow_name}}_function
+# Import the generated workflow function to trigger registration
+from .{{package_name}} import {{ python_safe_workflow_name }}_function

nat/cli/commands/workflow/templates/workflow.py.j2

@@ -3,6 +3,7 @@ import logging
 from pydantic import Field
 
 from nat.builder.builder import Builder
+from nat.builder.framework_enum import LLMFrameworkEnum
 from nat.builder.function_info import FunctionInfo
 from nat.cli.register_workflow import register_function
 from nat.data_models.function import FunctionBaseConfig
@@ -12,25 +13,38 @@ logger = logging.getLogger(__name__)
 
 class {{ workflow_class_name }}(FunctionBaseConfig, name="{{ workflow_name }}"):
     """
-    {{workflow_description}}
+    {{ workflow_description }}
     """
-    # Add your custom configuration parameters here
-    parameter: str = Field(default="default_value", description="Notional description for this parameter")
-
-
-@register_function(config_type={{ workflow_class_name }})
-async def {{ python_safe_workflow_name }}_function(
-    config: {{ workflow_class_name }}, builder: Builder
-):
-    # Implement your function logic here
-    async def _response_fn(input_message: str) -> str:
-        # Process the input_message and generate output
-        output_message = f"Hello from {{ workflow_name }} workflow! You said: {input_message}"
-        return output_message
-
-    try:
-        yield FunctionInfo.create(single_fn=_response_fn)
-    except GeneratorExit:
-        logger.warning("Function exited early!")
-    finally:
-        logger.info("Cleaning up {{ workflow_name }} workflow.")
+    prefix: str = Field(default="Echo:", description="Prefix to add before the echoed text.")
+
+
+@register_function(config_type={{ workflow_class_name }}, framework_wrappers=[LLMFrameworkEnum.LANGCHAIN])
+async def {{ python_safe_workflow_name }}_function(config: {{ workflow_class_name }}, builder: Builder):
+    """
+    Registers a function (addressable via `{{ workflow_name }}` in the configuration).
+    This registration ensures a static mapping of the function type, `{{ workflow_name }}`, to the `{{ workflow_class_name }}` configuration object.
+
+    Args:
+        config ({{ workflow_class_name }}): The configuration for the function.
+        builder (Builder): The builder object.
+
+    Returns:
+        FunctionInfo: The function info object for the function.
+    """
+
+    # Define the function that will be registered.
+    async def _echo(text: str) -> str:
+        """
+        Takes a text input and echoes back with a pre-defined prefix.
+
+        Args:
+            text (str): The text to echo back.
+
+        Returns:
+            str: The text with the prefix.
+        """
+        return f"{config.prefix} {text}"
+
+    # The callable is wrapped in a FunctionInfo object.
+    # The description parameter is used to describe the function.
+    yield FunctionInfo.from_fn(_echo, description=_echo.__doc__)