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

nat/agent/react_agent/register.py

@@ -24,6 +24,7 @@ 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
@@ -70,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.",
@@ -118,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,
@@ -153,25 +153,16 @@ async def react_agent_workflow(config: ReActAgentWorkflowConfig, builder: Builde
             content = str(output_message.content)
 
             # Create usage statistics for the response
-            prompt_tokens = sum(len(str(msg.content).split()) for msg in input_message.messages)
+            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)
-            return ChatResponse.from_string(content, usage=usage)
-
+            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,6 +25,7 @@ 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
@@ -54,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"),
@@ -125,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,
@@ -160,25 +160,16 @@ async def rewoo_agent_workflow(config: ReWOOAgentWorkflowConfig, builder: Builde
                 output_message = str(output_message)
 
             # Create usage statistics for the response
-            prompt_tokens = sum(len(str(msg.content).split()) for msg in input_message.messages)
+            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)
-            return ChatResponse.from_string(output_message, usage=usage)
-
+            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/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/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/entrypoint.py

@@ -29,6 +29,7 @@ import time
 
 import click
 import nest_asyncio
+from dotenv import load_dotenv
 
 from nat.utils.log_levels import LOG_LEVELS
 
@@ -45,6 +46,9 @@ from .commands.uninstall import uninstall_command
 from .commands.validate import validate_command
 from .commands.workflow.workflow import workflow_command
 
+# Load environment variables from .env file, if it exists
+load_dotenv()
+
 # Apply at the beginning of the file to avoid issues with asyncio
 nest_asyncio.apply()
 
@@ -52,7 +56,11 @@ nest_asyncio.apply()
 def setup_logging(log_level: str):
     """Configure logging with the specified level"""
     numeric_level = LOG_LEVELS.get(log_level.upper(), logging.INFO)
-    logging.basicConfig(level=numeric_level, format='%(asctime)s - %(name)s - %(levelname)s - %(message)s')
+    logging.basicConfig(
+        level=numeric_level,
+        format="%(asctime)s - %(levelname)-8s - %(name)s:%(lineno)d - %(message)s",
+        datefmt="%Y-%m-%d %H:%M:%S",
+    )
     return numeric_level
 
 

nat/data_models/api_server.py

@@ -28,6 +28,7 @@ from pydantic import HttpUrl
 from pydantic import conlist
 from pydantic import field_serializer
 from pydantic import field_validator
+from pydantic import model_validator
 from pydantic_core.core_schema import ValidationInfo
 
 from nat.data_models.interactive import HumanPrompt
@@ -120,15 +121,7 @@ class Message(BaseModel):
     role: UserMessageContentRoleType
 
 
-class ChatRequest(BaseModel):
-    """
-    ChatRequest is a data model that represents a request to the NAT chat API.
-    Fully compatible with OpenAI Chat Completions API specification.
-    """
-
-    # Required fields
-    messages: typing.Annotated[list[Message], conlist(Message, min_length=1)]
-
+class ChatRequestOptionals(BaseModel):
     # Optional fields (OpenAI Chat Completions API compatible)
     model: str | None = Field(default=None, description="name of the model to use")
     frequency_penalty: float | None = Field(default=0.0,
@@ -153,6 +146,16 @@ class ChatRequest(BaseModel):
     parallel_tool_calls: bool | None = Field(default=True, description="Whether to enable parallel function calling")
     user: str | None = Field(default=None, description="Unique identifier representing end-user")
 
+
+class ChatRequest(ChatRequestOptionals):
+    """
+    ChatRequest is a data model that represents a request to the NAT chat API.
+    Fully compatible with OpenAI Chat Completions API specification.
+    """
+
+    # Required fields
+    messages: typing.Annotated[list[Message], conlist(Message, min_length=1)]
+
     model_config = ConfigDict(extra="allow",
                               json_schema_extra={
                                   "example": {
@@ -194,6 +197,42 @@ class ChatRequest(BaseModel):
                            top_p=top_p)
 
 
+class ChatRequestOrMessage(ChatRequestOptionals):
+    """
+    ChatRequestOrMessage is a data model that represents either a conversation or a string input.
+    This is useful for functions that can handle either type of input.
+
+    `messages` is compatible with the OpenAI Chat Completions API specification.
+
+    `input_string` is a string input that can be used for functions that do not require a conversation.
+    """
+
+    messages: typing.Annotated[list[Message] | None, conlist(Message, min_length=1)] = Field(
+        default=None, description="The conversation messages to process.")
+
+    input_string: str | None = Field(default=None, alias="input_message", description="The input message to process.")
+
+    @property
+    def is_string(self) -> bool:
+        return self.input_string is not None
+
+    @property
+    def is_conversation(self) -> bool:
+        return self.messages is not None
+
+    @model_validator(mode="after")
+    def validate_messages_or_input_string(self):
+        if self.messages is not None and self.input_string is not None:
+            raise ValueError("Either messages or input_message/input_string must be provided, not both")
+        if self.messages is None and self.input_string is None:
+            raise ValueError("Either messages or input_message/input_string must be provided")
+        if self.input_string is not None:
+            extra_fields = self.model_dump(exclude={"input_string"}, exclude_none=True, exclude_unset=True)
+            if len(extra_fields) > 0:
+                raise ValueError("no extra fields are permitted when input_message/input_string is provided")
+        return self
+
+
 class ChoiceMessage(BaseModel):
     content: str | None = None
     role: UserMessageContentRoleType | None = None
@@ -661,6 +700,36 @@ def _string_to_nat_chat_request(data: str) -> ChatRequest:
 GlobalTypeConverter.register_converter(_string_to_nat_chat_request)
 
 
+def _chat_request_or_message_to_chat_request(data: ChatRequestOrMessage) -> ChatRequest:
+    if data.input_string is not None:
+        return _string_to_nat_chat_request(data.input_string)
+    return ChatRequest(**data.model_dump(exclude={"input_string"}))
+
+
+GlobalTypeConverter.register_converter(_chat_request_or_message_to_chat_request)
+
+
+def _chat_request_to_chat_request_or_message(data: ChatRequest) -> ChatRequestOrMessage:
+    return ChatRequestOrMessage(**data.model_dump(by_alias=True))
+
+
+GlobalTypeConverter.register_converter(_chat_request_to_chat_request_or_message)
+
+
+def _chat_request_or_message_to_string(data: ChatRequestOrMessage) -> str:
+    return data.input_string or ""
+
+
+GlobalTypeConverter.register_converter(_chat_request_or_message_to_string)
+
+
+def _string_to_chat_request_or_message(data: str) -> ChatRequestOrMessage:
+    return ChatRequestOrMessage(input_message=data)
+
+
+GlobalTypeConverter.register_converter(_string_to_chat_request_or_message)
+
+
 # ======== ChatResponse Converters ========
 def _nat_chat_response_to_string(data: ChatResponse) -> str:
     if data.choices and data.choices[0].message:

nat/data_models/config.py

@@ -187,7 +187,7 @@ class TelemetryConfig(BaseModel):
 
 class GeneralConfig(BaseModel):
 
-    model_config = ConfigDict(protected_namespaces=())
+    model_config = ConfigDict(protected_namespaces=(), extra="forbid")
 
     use_uvloop: bool | None = Field(
         default=None,

nat/experimental/test_time_compute/functions/execute_score_select_function.py

@@ -46,7 +46,7 @@ async def execute_score_select_function(config: ExecuteScoreSelectFunctionConfig
 
     from pydantic import BaseModel
 
-    executable_fn: Function = builder.get_function(name=config.augmented_fn)
+    executable_fn: Function = await builder.get_function(name=config.augmented_fn)
 
     if config.scorer:
         scorer = await builder.get_ttc_strategy(strategy_name=config.scorer,

nat/experimental/test_time_compute/functions/ttc_tool_wrapper_function.py

@@ -98,8 +98,8 @@ async def register_ttc_tool_wrapper_function(
 
         augmented_function_desc = config.tool_description
 
-    fn_input_schema: BaseModel = augmented_function.input_schema
-    fn_output_schema: BaseModel = augmented_function.single_output_schema
+    fn_input_schema: type[BaseModel] = augmented_function.input_schema
+    fn_output_schema: type[BaseModel] | type[None] = augmented_function.single_output_schema
 
     runnable_llm = input_llm.with_structured_output(schema=fn_input_schema)
 

nat/front_ends/console/console_front_end_plugin.py

@@ -95,5 +95,14 @@ class ConsoleFrontEndPlugin(SimpleFrontEndPluginBase[ConsoleFrontEndConfig]):
         else:
             assert False, "Should not reach here. Should have been caught by pre_run"
 
-        # Print result
-        logger.info(f"\n{'-' * 50}\n{Fore.GREEN}Workflow Result:\n%s{Fore.RESET}\n{'-' * 50}", runner_outputs)
+        line = f"{'-' * 50}"
+        prefix = f"{line}\n{Fore.GREEN}Workflow Result:\n"
+        suffix = f"{Fore.RESET}\n{line}"
+
+        logger.info(f"{prefix}%s{suffix}", runner_outputs)
+
+        # (handler is a stream handler) => (level > INFO)
+        effective_level_too_high = all(
+            type(h) is not logging.StreamHandler or h.level > logging.INFO for h in logging.getLogger().handlers)
+        if effective_level_too_high:
+            print(f"{prefix}{runner_outputs}{suffix}")

nat/front_ends/fastapi/auth_flow_handlers/http_flow_handler.py

@@ -24,4 +24,4 @@ class HTTPAuthenticationFlowHandler(FlowHandlerBase):
     async def authenticate(self, config: AuthProviderBaseConfig, method: AuthFlowType) -> AuthenticatedContext:
 
         raise NotImplementedError(f"Authentication method '{method}' is not supported by the HTTP frontend."
-                                  f" Do you have Websockets enabled?")
+                                  f" Do you have WebSockets enabled?")

nat/observability/register.py

@@ -77,6 +77,14 @@ async def console_logging_method(config: ConsoleLoggingMethodConfig, builder: Bu
     level = getattr(logging, config.level.upper(), logging.INFO)
     handler = logging.StreamHandler(stream=sys.stdout)
     handler.setLevel(level)
+
+    # Set formatter to match the default CLI format
+    formatter = logging.Formatter(
+        fmt="%(asctime)s - %(levelname)-8s - %(name)s:%(lineno)d - %(message)s",
+        datefmt="%Y-%m-%d %H:%M:%S",
+    )
+    handler.setFormatter(formatter)
+
     yield handler
 
 
@@ -95,4 +103,12 @@ async def file_logging_method(config: FileLoggingMethod, builder: Builder):
     level = getattr(logging, config.level.upper(), logging.INFO)
     handler = logging.FileHandler(filename=config.path, mode="a", encoding="utf-8")
     handler.setLevel(level)
+
+    # Set formatter to match the default CLI format
+    formatter = logging.Formatter(
+        fmt="%(asctime)s - %(levelname)-8s - %(name)s:%(lineno)d - %(message)s",
+        datefmt="%Y-%m-%d %H:%M:%S",
+    )
+    handler.setFormatter(formatter)
+
     yield handler

nat/runtime/session.py

@@ -192,7 +192,7 @@ class SessionManager:
                                     user_message_id: str | None,
                                     conversation_id: str | None) -> None:
         """
-        Extracts and sets user metadata for Websocket connections.
+        Extracts and sets user metadata for WebSocket connections.
         """
 
         # Extract cookies from WebSocket headers (similar to HTTP request)

nat/tool/memory_tools/add_memory_tool.py

@@ -30,10 +30,10 @@ logger = logging.getLogger(__name__)
 class AddToolConfig(FunctionBaseConfig, name="add_memory"):
     """Function to add memory to a hosted memory platform."""
 
-    description: str = Field(default=("Tool to add memory about a user's interactions to a system "
+    description: str = Field(default=("Tool to add a memory about a user's interactions to a system "
                                       "for retrieval later."),
                              description="The description of this function's use for tool calling agents.")
-    memory: MemoryRef = Field(default="saas_memory",
+    memory: MemoryRef = Field(default=MemoryRef("saas_memory"),
                               description=("Instance name of the memory client instance from the workflow "
                                            "configuration object."))
 
@@ -46,7 +46,7 @@ async def add_memory_tool(config: AddToolConfig, builder: Builder):
     from langchain_core.tools import ToolException
 
     # First, retrieve the memory client
-    memory_editor = builder.get_memory_client(config.memory)
+    memory_editor = await builder.get_memory_client(config.memory)
 
     async def _arun(item: MemoryItem) -> str:
         """

nat/tool/memory_tools/delete_memory_tool.py

@@ -30,10 +30,9 @@ logger = logging.getLogger(__name__)
 class DeleteToolConfig(FunctionBaseConfig, name="delete_memory"):
     """Function to delete memory from a hosted memory platform."""
 
-    description: str = Field(default=("Tool to retrieve memory about a user's "
-                                      "interactions to help answer questions in a personalized way."),
+    description: str = Field(default="Tool to delete a memory from a hosted memory platform.",
                              description="The description of this function's use for tool calling agents.")
-    memory: MemoryRef = Field(default="saas_memory",
+    memory: MemoryRef = Field(default=MemoryRef("saas_memory"),
                               description=("Instance name of the memory client instance from the workflow "
                                            "configuration object."))
 
@@ -47,7 +46,7 @@ async def delete_memory_tool(config: DeleteToolConfig, builder: Builder):
     from langchain_core.tools import ToolException
 
     # First, retrieve the memory client
-    memory_editor = builder.get_memory_client(config.memory)
+    memory_editor = await builder.get_memory_client(config.memory)
 
     async def _arun(user_id: str) -> str:
         """

nat/tool/memory_tools/get_memory_tool.py

@@ -30,10 +30,10 @@ logger = logging.getLogger(__name__)
 class GetToolConfig(FunctionBaseConfig, name="get_memory"):
     """Function to get memory to a hosted memory platform."""
 
-    description: str = Field(default=("Tool to retrieve memory about a user's "
+    description: str = Field(default=("Tool to retrieve a memory about a user's "
                                       "interactions to help answer questions in a personalized way."),
                              description="The description of this function's use for tool calling agents.")
-    memory: MemoryRef = Field(default="saas_memory",
+    memory: MemoryRef = Field(default=MemoryRef("saas_memory"),
                               description=("Instance name of the memory client instance from the workflow "
                                            "configuration object."))
 
@@ -49,7 +49,7 @@ async def get_memory_tool(config: GetToolConfig, builder: Builder):
     from langchain_core.tools import ToolException
 
     # First, retrieve the memory client
-    memory_editor = builder.get_memory_client(config.memory)
+    memory_editor = await builder.get_memory_client(config.memory)
 
     async def _arun(search_input: SearchMemoryInput) -> str:
         """

nat/utils/type_converter.py

@@ -93,6 +93,14 @@ class TypeConverter:
         if to_type is None or decomposed.is_instance(data):
             return data
 
+        # 2) If data is a union type, try to convert to each type in the union
+        if decomposed.is_union:
+            for union_type in decomposed.args:
+                result = self._convert(data, union_type)
+                if result is not None:
+                    return result
+            return None
+
         root = decomposed.root
 
         # 2) Attempt direct in *this* converter