nvidia-nat 1.2.0rc5__py3-none-any.whl → 1.2.0rc6__py3-none-any.whl

{aiq → nat}/front_ends/fastapi/step_adaptor.py

@@ -18,16 +18,16 @@ import logging
 from functools import reduce
 from textwrap import dedent
 
-from aiq.data_models.api_server import AIQResponseIntermediateStep
-from aiq.data_models.api_server import AIQResponseSerializable
-from aiq.data_models.intermediate_step import IntermediateStep
-from aiq.data_models.intermediate_step import IntermediateStepCategory
-from aiq.data_models.intermediate_step import IntermediateStepPayload
-from aiq.data_models.intermediate_step import IntermediateStepType
-from aiq.data_models.invocation_node import InvocationNode
-from aiq.data_models.step_adaptor import StepAdaptorConfig
-from aiq.data_models.step_adaptor import StepAdaptorMode
-from aiq.utils.type_utils import is_valid_json
+from nat.data_models.api_server import ResponseIntermediateStep
+from nat.data_models.api_server import ResponseSerializable
+from nat.data_models.intermediate_step import IntermediateStep
+from nat.data_models.intermediate_step import IntermediateStepCategory
+from nat.data_models.intermediate_step import IntermediateStepPayload
+from nat.data_models.intermediate_step import IntermediateStepType
+from nat.data_models.invocation_node import InvocationNode
+from nat.data_models.step_adaptor import StepAdaptorConfig
+from nat.data_models.step_adaptor import StepAdaptorMode
+from nat.utils.type_utils import is_valid_json
 
 logger = logging.getLogger(__name__)
 
@@ -63,7 +63,7 @@ class StepAdaptor:
 
         return False
 
-    def _handle_llm(self, step: IntermediateStepPayload, ancestry: InvocationNode) -> AIQResponseSerializable | None:
+    def _handle_llm(self, step: IntermediateStepPayload, ancestry: InvocationNode) -> ResponseSerializable | None:
         input_str: str | None = None
         output_str: str | None = None
 
@@ -113,14 +113,14 @@ class StepAdaptor:
             {output_value}
             """).strip("\n").format(payload=payload, output_value=escaped_output)
 
-        event = AIQResponseIntermediateStep(id=step.UUID,
-                                            name=step.name or "",
-                                            payload=payload,
-                                            parent_id=ancestry.function_id)
+        event = ResponseIntermediateStep(id=step.UUID,
+                                         name=step.name or "",
+                                         payload=payload,
+                                         parent_id=ancestry.function_id)
 
         return event
 
-    def _handle_tool(self, step: IntermediateStepPayload, ancestry: InvocationNode) -> AIQResponseSerializable | None:
+    def _handle_tool(self, step: IntermediateStepPayload, ancestry: InvocationNode) -> ResponseSerializable | None:
         """
         Handles both TOOL_START and TOOL_END events
         """
@@ -168,15 +168,14 @@ class StepAdaptor:
             ```
             """).strip("\n").format(payload=payload, output_value=escaped_output, format_output_type=format_output_type)
 
-        event = AIQResponseIntermediateStep(id=step.UUID,
-                                            name=f"Tool: {step.name}",
-                                            payload=payload,
-                                            parent_id=ancestry.function_id)
+        event = ResponseIntermediateStep(id=step.UUID,
+                                         name=f"Tool: {step.name}",
+                                         payload=payload,
+                                         parent_id=ancestry.function_id)
 
         return event
 
-    def _handle_function(self, step: IntermediateStepPayload,
-                         ancestry: InvocationNode) -> AIQResponseSerializable | None:
+    def _handle_function(self, step: IntermediateStepPayload, ancestry: InvocationNode) -> ResponseSerializable | None:
         """
         Handles the FUNCTION_START and FUNCTION_END events
         """
@@ -204,10 +203,10 @@ class StepAdaptor:
             ```
             """).strip("\n").format(input_value=escaped_input, format_input_type=format_input_type)
 
-            event = AIQResponseIntermediateStep(id=step.UUID,
-                                                name=f"Function Start: {step.name}",
-                                                payload=payload_str,
-                                                parent_id=ancestry.parent_id)
+            event = ResponseIntermediateStep(id=step.UUID,
+                                             name=f"Function Start: {step.name}",
+                                             payload=payload_str,
+                                             parent_id=ancestry.parent_id)
             return event
 
         if step.event_type == IntermediateStepType.FUNCTION_END:
@@ -257,16 +256,15 @@ class StepAdaptor:
                                     output_value=escaped_output,
                                     format_output_type=format_output_type)
 
-            event = AIQResponseIntermediateStep(id=step.UUID,
-                                                name=f"Function Complete: {step.name}",
-                                                payload=payload_str,
-                                                parent_id=ancestry.parent_id)
+            event = ResponseIntermediateStep(id=step.UUID,
+                                             name=f"Function Complete: {step.name}",
+                                             payload=payload_str,
+                                             parent_id=ancestry.parent_id)
             return event
 
         return None
 
-    def _handle_custom(self, payload: IntermediateStepPayload,
-                       ancestry: InvocationNode) -> AIQResponseSerializable | None:
+    def _handle_custom(self, payload: IntermediateStepPayload, ancestry: InvocationNode) -> ResponseSerializable | None:
         """
         Handles the CUSTOM event
         """
@@ -284,14 +282,14 @@ class StepAdaptor:
         """).strip("\n").format(payload=escaped_payload, format_type=format_type)
 
         # Return the event
-        event = AIQResponseIntermediateStep(id=payload.UUID,
-                                            name=f"{payload.event_type}",
-                                            payload=payload_str,
-                                            parent_id=ancestry.function_id)
+        event = ResponseIntermediateStep(id=payload.UUID,
+                                         name=f"{payload.event_type}",
+                                         payload=payload_str,
+                                         parent_id=ancestry.function_id)
 
         return event
 
-    def process(self, step: IntermediateStep) -> AIQResponseSerializable | None:  # pylint: disable=R1710
+    def process(self, step: IntermediateStep) -> ResponseSerializable | None:  # pylint: disable=R1710
 
         # Track the chunk
         self._history.append(step)

{aiq → nat}/front_ends/mcp/mcp_front_end_config.py

@@ -15,18 +15,22 @@
 
 from pydantic import Field
 
-from aiq.data_models.front_end import FrontEndBaseConfig
+from nat.data_models.front_end import FrontEndBaseConfig
 
 
 class MCPFrontEndConfig(FrontEndBaseConfig, name="mcp"):
     """MCP front end configuration.
 
-    A simple MCP (Modular Communication Protocol) front end for AIQ.
+    A simple MCP (Modular Communication Protocol) front end for NeMo Agent toolkit.
     """
 
-    name: str = Field(default="AIQ MCP", description="Name of the MCP server")
-    host: str = Field(default="localhost", description="Host to bind the server to")
-    port: int = Field(default=9901, description="Port to bind the server to", ge=0, le=65535)
-    debug: bool = Field(default=False, description="Enable debug mode")
-    log_level: str = Field(default="INFO", description="Log level for the MCP server")
-    tool_names: list[str] = Field(default_factory=list, description="The list of tools MCP server will expose.")
+    name: str = Field(default="NeMo Agent Toolkit MCP",
+                      description="Name of the MCP server (default: NeMo Agent Toolkit MCP)")
+    host: str = Field(default="localhost", description="Host to bind the server to (default: localhost)")
+    port: int = Field(default=9901, description="Port to bind the server to (default: 9901)", ge=0, le=65535)
+    debug: bool = Field(default=False, description="Enable debug mode (default: False)")
+    log_level: str = Field(default="INFO", description="Log level for the MCP server (default: INFO)")
+    tool_names: list[str] = Field(default_factory=list,
+                                  description="The list of tools MCP server will expose (default: all tools)")
+    runner_class: str | None = Field(
+        default=None, description="Custom worker class for handling MCP routes (default: built-in worker)")

nat/front_ends/mcp/mcp_front_end_plugin.py

@@ -0,0 +1,81 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import logging
+import typing
+
+from nat.builder.front_end import FrontEndBase
+from nat.builder.workflow_builder import WorkflowBuilder
+from nat.front_ends.mcp.mcp_front_end_config import MCPFrontEndConfig
+from nat.front_ends.mcp.mcp_front_end_plugin_worker import MCPFrontEndPluginWorkerBase
+
+logger = logging.getLogger(__name__)
+
+
+class MCPFrontEndPlugin(FrontEndBase[MCPFrontEndConfig]):
+    """MCP front end plugin implementation."""
+
+    def get_worker_class(self) -> type[MCPFrontEndPluginWorkerBase]:
+        """Get the worker class for handling MCP routes."""
+        from nat.front_ends.mcp.mcp_front_end_plugin_worker import MCPFrontEndPluginWorker
+
+        return MCPFrontEndPluginWorker
+
+    @typing.final
+    def get_worker_class_name(self) -> str:
+        """Get the worker class name from configuration or default."""
+        if self.front_end_config.runner_class:
+            return self.front_end_config.runner_class
+
+        worker_class = self.get_worker_class()
+        return f"{worker_class.__module__}.{worker_class.__qualname__}"
+
+    def _get_worker_instance(self) -> MCPFrontEndPluginWorkerBase:
+        """Get an instance of the worker class."""
+        # Import the worker class dynamically if specified in config
+        if self.front_end_config.runner_class:
+            module_name, class_name = self.front_end_config.runner_class.rsplit(".", 1)
+            import importlib
+            module = importlib.import_module(module_name)
+            worker_class = getattr(module, class_name)
+        else:
+            worker_class = self.get_worker_class()
+
+        return worker_class(self.full_config)
+
+    async def run(self) -> None:
+        """Run the MCP server."""
+        # Import FastMCP
+        from mcp.server.fastmcp import FastMCP
+
+        # Create an MCP server with the configured parameters
+        mcp = FastMCP(
+            self.front_end_config.name,
+            host=self.front_end_config.host,
+            port=self.front_end_config.port,
+            debug=self.front_end_config.debug,
+            log_level=self.front_end_config.log_level,
+        )
+
+        # Get the worker instance and set up routes
+        worker = self._get_worker_instance()
+
+        # Build the workflow and add routes using the worker
+        async with WorkflowBuilder.from_config(config=self.full_config) as builder:
+            # Add routes through the worker (includes health endpoint and function registration)
+            await worker.add_routes(mcp, builder)
+
+            # Start the MCP server
+            await mcp.run_sse_async()

nat/front_ends/mcp/mcp_front_end_plugin_worker.py

@@ -0,0 +1,143 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import logging
+from abc import ABC
+from abc import abstractmethod
+
+from mcp.server.fastmcp import FastMCP
+from starlette.requests import Request
+
+from nat.builder.function import Function
+from nat.builder.workflow import Workflow
+from nat.builder.workflow_builder import WorkflowBuilder
+from nat.data_models.config import Config
+from nat.front_ends.mcp.mcp_front_end_config import MCPFrontEndConfig
+
+logger = logging.getLogger(__name__)
+
+
+class MCPFrontEndPluginWorkerBase(ABC):
+    """Base class for MCP front end plugin workers."""
+
+    def __init__(self, config: Config):
+        """Initialize the MCP worker with configuration.
+
+        Args:
+            config: The full NAT configuration
+        """
+        self.full_config = config
+        self.front_end_config: MCPFrontEndConfig = config.general.front_end
+
+    def _setup_health_endpoint(self, mcp: FastMCP):
+        """Set up the HTTP health endpoint that exercises MCP ping handler."""
+
+        @mcp.custom_route("/health", methods=["GET"])
+        async def health_check(_request: Request):
+            """HTTP health check using server's internal ping handler"""
+            from starlette.responses import JSONResponse
+
+            try:
+                from mcp.types import PingRequest
+
+                # Create a ping request
+                ping_request = PingRequest(method="ping")
+
+                # Call the ping handler directly (same one that responds to MCP pings)
+                await mcp._mcp_server.request_handlers[PingRequest](ping_request)
+
+                return JSONResponse({
+                    "status": "healthy",
+                    "error": None,
+                    "server_name": mcp.name,
+                })
+
+            except Exception as e:
+                return JSONResponse({
+                    "status": "unhealthy",
+                    "error": str(e),
+                    "server_name": mcp.name,
+                },
+                                    status_code=503)
+
+    @abstractmethod
+    async def add_routes(self, mcp: FastMCP, builder: WorkflowBuilder):
+        """Add routes to the MCP server.
+
+        Args:
+            mcp: The FastMCP server instance
+            builder (WorkflowBuilder): The workflow builder instance
+        """
+        pass
+
+    def _get_all_functions(self, workflow: Workflow) -> dict[str, Function]:
+        """Get all functions from the workflow.
+
+        Args:
+            workflow: The NAT workflow.
+
+        Returns:
+            Dict mapping function names to Function objects.
+        """
+        functions: dict[str, Function] = {}
+
+        # Extract all functions from the workflow
+        for function_name, function in workflow.functions.items():
+            functions[function_name] = function
+
+        functions[workflow.config.workflow.type] = workflow
+
+        return functions
+
+
+class MCPFrontEndPluginWorker(MCPFrontEndPluginWorkerBase):
+    """Default MCP front end plugin worker implementation."""
+
+    async def add_routes(self, mcp: FastMCP, builder: WorkflowBuilder):
+        """Add default routes to the MCP server.
+
+        Args:
+            mcp: The FastMCP server instance
+            builder (WorkflowBuilder): The workflow builder instance
+        """
+        from nat.front_ends.mcp.tool_converter import register_function_with_mcp
+
+        # Set up the health endpoint
+        self._setup_health_endpoint(mcp)
+
+        # Build the workflow and register all functions with MCP
+        workflow = builder.build()
+
+        # Get all functions from the workflow
+        functions = self._get_all_functions(workflow)
+
+        # Filter functions based on tool_names if provided
+        if self.front_end_config.tool_names:
+            logger.info("Filtering functions based on tool_names: %s", self.front_end_config.tool_names)
+            filtered_functions: dict[str, Function] = {}
+            for function_name, function in functions.items():
+                if function_name in self.front_end_config.tool_names:
+                    filtered_functions[function_name] = function
+                else:
+                    logger.debug("Skipping function %s as it's not in tool_names", function_name)
+            functions = filtered_functions
+
+        # Register each function with MCP
+        for function_name, function in functions.items():
+            register_function_with_mcp(mcp, function_name, function)
+
+        # Add a simple fallback function if no functions were found
+        if not functions:
+            raise RuntimeError("No functions found in workflow. Please check your configuration.")

{aiq → nat}/front_ends/mcp/register.py

@@ -15,13 +15,13 @@
 
 from collections.abc import AsyncIterator
 
-from aiq.cli.register_workflow import register_front_end
-from aiq.data_models.config import AIQConfig
-from aiq.front_ends.mcp.mcp_front_end_config import MCPFrontEndConfig
+from nat.cli.register_workflow import register_front_end
+from nat.data_models.config import Config
+from nat.front_ends.mcp.mcp_front_end_config import MCPFrontEndConfig
 
 
 @register_front_end(config_type=MCPFrontEndConfig)
-async def register_mcp_front_end(config: MCPFrontEndConfig, full_config: AIQConfig) -> AsyncIterator:
-    from aiq.front_ends.mcp.mcp_front_end_plugin import MCPFrontEndPlugin
+async def register_mcp_front_end(config: MCPFrontEndConfig, full_config: Config) -> AsyncIterator:
+    from nat.front_ends.mcp.mcp_front_end_plugin import MCPFrontEndPlugin
 
     yield MCPFrontEndPlugin(full_config=full_config)

{aiq → nat}/front_ends/mcp/tool_converter.py

@@ -21,9 +21,9 @@ from inspect import Signature
 from mcp.server.fastmcp import FastMCP
 from pydantic import BaseModel
 
-from aiq.builder.function import Function
-from aiq.builder.function_base import FunctionBase
-from aiq.builder.workflow import Workflow
+from nat.builder.function import Function
+from nat.builder.function_base import FunctionBase
+from nat.builder.workflow import Workflow
 
 logger = logging.getLogger(__name__)
 
@@ -34,27 +34,26 @@ def create_function_wrapper(
     schema: type[BaseModel],
     is_workflow: bool = False,
 ):
-    """Create a wrapper function that exposes the actual parameters of an AIQ Function as an MCP tool.
+    """Create a wrapper function that exposes the actual parameters of a NAT Function as an MCP tool.
 
     Args:
         function_name: The name of the function/tool
-        function: The AIQ Function object
+        function: The NAT Function object
         schema: The input schema of the function
         is_workflow: Whether the function is a Workflow
 
     Returns:
         A wrapper function suitable for registration with MCP
     """
-    # Check if we're dealing with AIQChatRequest - special case
+    # Check if we're dealing with ChatRequest - special case
     is_chat_request = False
 
-    # Check if the schema name is AIQChatRequest
-    if schema.__name__ == "AIQChatRequest" or (hasattr(schema, "__qualname__")
-                                               and "AIQChatRequest" in schema.__qualname__):
+    # Check if the schema name is ChatRequest
+    if schema.__name__ == "ChatRequest" or (hasattr(schema, "__qualname__") and "ChatRequest" in schema.__qualname__):
         is_chat_request = True
-        logger.info("Function %s uses AIQChatRequest - creating simplified interface", function_name)
+        logger.info("Function %s uses ChatRequest - creating simplified interface", function_name)
 
-        # For AIQChatRequest, we'll create a simple wrapper with just a query parameter
+        # For ChatRequest, we'll create a simple wrapper with just a query parameter
         parameters = [Parameter(
             name="query",
             kind=Parameter.KEYWORD_ONLY,
@@ -102,18 +101,18 @@ def create_function_wrapper(
                 await ctx.report_progress(0, 100)
 
             try:
-                # Special handling for AIQChatRequest
+                # Special handling for ChatRequest
                 if is_chat_request:
-                    from aiq.data_models.api_server import AIQChatRequest
+                    from nat.data_models.api_server import ChatRequest
 
                     # Create a chat request from the query string
                     query = kwargs.get("query", "")
-                    chat_request = AIQChatRequest.from_string(query)
+                    chat_request = ChatRequest.from_string(query)
 
                     # Special handling for Workflow objects
                     if is_workflow:
                         # Workflows have a run method that is an async context manager
-                        # that returns an AIQRunner
+                        # that returns a Runner
                         async with function.run(chat_request) as runner:
                             # Get the result from the runner
                             result = await runner.result(to_type=str)
@@ -136,13 +135,13 @@ def create_function_wrapper(
                             # Call with the nested object
                             kwargs = {field_name: nested_obj}
 
-                    # Call the AIQ function with the parameters - special handling for Workflow
+                    # Call the NAT function with the parameters - special handling for Workflow
                     if is_workflow:
                         # For workflow with regular input, we'll assume the first parameter is the input
                         input_value = list(kwargs.values())[0] if kwargs else ""
 
                         # Workflows have a run method that is an async context manager
-                        # that returns an AIQRunner
+                        # that returns a Runner
                         async with function.run(input_value) as runner:
                             # Get the result from the runner
                             result = await runner.result(to_type=str)
@@ -180,7 +179,7 @@ def create_function_wrapper(
 
 def get_function_description(function: FunctionBase) -> str:
     """
-    Retrieve a human-readable description for an AIQ function or workflow.
+    Retrieve a human-readable description for a NAT function or workflow.
 
     The description is determined using the following precedence:
        1. If the function is a Workflow and has a 'description' attribute, use it.
@@ -189,7 +188,7 @@ def get_function_description(function: FunctionBase) -> str:
        4. If the function is a regular Function, use its 'description' attribute.
 
     Args:
-        function: The AIQ FunctionBase instance (Function or Workflow).
+        function: The NAT FunctionBase instance (Function or Workflow).
 
     Returns:
         The best available description string for the function.
@@ -216,12 +215,12 @@ def get_function_description(function: FunctionBase) -> str:
 
 
 def register_function_with_mcp(mcp: FastMCP, function_name: str, function: FunctionBase) -> None:
-    """Register an AIQ Function as an MCP tool.
+    """Register a NAT Function as an MCP tool.
 
     Args:
         mcp: The FastMCP instance
         function_name: The name to register the function under
-        function: The AIQ Function to register
+        function: The NAT Function to register
     """
     logger.info("Registering function %s with MCP", function_name)
 

{aiq → nat}/front_ends/simple_base/simple_front_end_plugin_base.py

@@ -20,10 +20,10 @@ from io import StringIO
 
 import click
 
-from aiq.builder.front_end import FrontEndBase
-from aiq.builder.workflow_builder import WorkflowBuilder
-from aiq.data_models.front_end import FrontEndConfigT
-from aiq.runtime.session import AIQSessionManager
+from nat.builder.front_end import FrontEndBase
+from nat.builder.workflow_builder import WorkflowBuilder
+from nat.data_models.front_end import FrontEndConfigT
+from nat.runtime.session import SessionManager
 
 logger = logging.getLogger(__name__)
 
@@ -46,9 +46,9 @@ class SimpleFrontEndPluginBase(FrontEndBase[FrontEndConfigT], ABC):
                 click.echo(stream.getvalue())
 
             workflow = builder.build()
-            session_manager = AIQSessionManager(workflow)
+            session_manager = SessionManager(workflow)
             await self.run_workflow(session_manager)
 
     @abstractmethod
-    async def run_workflow(self, session_manager: AIQSessionManager):
+    async def run_workflow(self, session_manager: SessionManager):
         pass

{aiq → nat}/llm/aws_bedrock_llm.py

@@ -17,11 +17,11 @@ from pydantic import AliasChoices
 from pydantic import ConfigDict
 from pydantic import Field
 
-from aiq.builder.builder import Builder
-from aiq.builder.llm import LLMProviderInfo
-from aiq.cli.register_workflow import register_llm_provider
-from aiq.data_models.llm import LLMBaseConfig
-from aiq.data_models.retry_mixin import RetryMixin
+from nat.builder.builder import Builder
+from nat.builder.llm import LLMProviderInfo
+from nat.cli.register_workflow import register_llm_provider
+from nat.data_models.llm import LLMBaseConfig
+from nat.data_models.retry_mixin import RetryMixin
 
 
 class AWSBedrockModelConfig(LLMBaseConfig, RetryMixin, name="aws_bedrock"):

{aiq → nat}/llm/nim_llm.py

@@ -18,11 +18,11 @@ from pydantic import ConfigDict
 from pydantic import Field
 from pydantic import PositiveInt
 
-from aiq.builder.builder import Builder
-from aiq.builder.llm import LLMProviderInfo
-from aiq.cli.register_workflow import register_llm_provider
-from aiq.data_models.llm import LLMBaseConfig
-from aiq.data_models.retry_mixin import RetryMixin
+from nat.builder.builder import Builder
+from nat.builder.llm import LLMProviderInfo
+from nat.cli.register_workflow import register_llm_provider
+from nat.data_models.llm import LLMBaseConfig
+from nat.data_models.retry_mixin import RetryMixin
 
 
 class NIMModelConfig(LLMBaseConfig, RetryMixin, name="nim"):

{aiq → nat}/llm/openai_llm.py

@@ -17,11 +17,11 @@ from pydantic import AliasChoices
 from pydantic import ConfigDict
 from pydantic import Field
 
-from aiq.builder.builder import Builder
-from aiq.builder.llm import LLMProviderInfo
-from aiq.cli.register_workflow import register_llm_provider
-from aiq.data_models.llm import LLMBaseConfig
-from aiq.data_models.retry_mixin import RetryMixin
+from nat.builder.builder import Builder
+from nat.builder.llm import LLMProviderInfo
+from nat.cli.register_workflow import register_llm_provider
+from nat.data_models.llm import LLMBaseConfig
+from nat.data_models.retry_mixin import RetryMixin
 
 
 class OpenAIModelConfig(LLMBaseConfig, RetryMixin, name="openai"):

{aiq → nat}/memory/__init__.py

@@ -13,8 +13,8 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 """
-AIQ Toolkit Memory Module
+NAT Memory Module
 
 This package provides foundational classes and interfaces
-for managing text-based memory in AIQ Toolkit's LLM-based agents.
+for managing text-based memory in NAT's LLM-based agents.
 """