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

{aiq → nat}/runtime/runner.py

@@ -17,12 +17,12 @@ import logging
 import typing
 from enum import Enum
 
-from aiq.builder.context import AIQContext
-from aiq.builder.context import AIQContextState
-from aiq.builder.function import Function
-from aiq.data_models.invocation_node import InvocationNode
-from aiq.observability.exporter_manager import ExporterManager
-from aiq.utils.reactive.subject import Subject
+from nat.builder.context import Context
+from nat.builder.context import ContextState
+from nat.builder.function import Function
+from nat.data_models.invocation_node import InvocationNode
+from nat.observability.exporter_manager import ExporterManager
+from nat.utils.reactive.subject import Subject
 
 logger = logging.getLogger(__name__)
 
@@ -31,7 +31,7 @@ class UserManagerBase:
     pass
 
 
-class AIQRunnerState(Enum):
+class RunnerState(Enum):
     UNINITIALIZED = 0
     INITIALIZED = 1
     RUNNING = 2
@@ -42,15 +42,15 @@ class AIQRunnerState(Enum):
 _T = typing.TypeVar("_T")
 
 
-class AIQRunner:
+class Runner:
 
     def __init__(self,
                  input_message: typing.Any,
                  entry_fn: Function,
-                 context_state: AIQContextState,
+                 context_state: ContextState,
                  exporter_manager: ExporterManager):
         """
-        The AIQRunner class is used to run a workflow. It handles converting input and output data types and running the
+        The Runner class is used to run a workflow. It handles converting input and output data types and running the
         workflow with the specified concurrency.
 
         Parameters
@@ -59,7 +59,7 @@ class AIQRunner:
             The input message to the workflow
         entry_fn : Function
             The entry function to the workflow
-        context_state : AIQContextState
+        context_state : ContextState
             The context state to use
         exporter_manager : ExporterManager
             The exporter manager to use
@@ -70,9 +70,9 @@ class AIQRunner:
 
         self._entry_fn = entry_fn
         self._context_state = context_state
-        self._context = AIQContext(self._context_state)
+        self._context = Context(self._context_state)
 
-        self._state = AIQRunnerState.UNINITIALIZED
+        self._state = RunnerState.UNINITIALIZED
 
         self._input_message_token = None
 
@@ -82,7 +82,7 @@ class AIQRunner:
         self._exporter_manager = exporter_manager
 
     @property
-    def context(self) -> AIQContext:
+    def context(self) -> Context:
         return self._context
 
     def convert(self, value: typing.Any, to_type: type[_T]) -> _T:
@@ -100,8 +100,8 @@ class AIQRunner:
             function_id="root",
         ))
 
-        if (self._state == AIQRunnerState.UNINITIALIZED):
-            self._state = AIQRunnerState.INITIALIZED
+        if (self._state == RunnerState.UNINITIALIZED):
+            self._state = RunnerState.INITIALIZED
         else:
             raise ValueError("Cannot enter the context more than once")
 
@@ -114,7 +114,7 @@ class AIQRunner:
 
         self._context_state.input_message.reset(self._input_message_token)
 
-        if (self._state not in (AIQRunnerState.COMPLETED, AIQRunnerState.FAILED)):
+        if (self._state not in (RunnerState.COMPLETED, RunnerState.FAILED)):
             raise ValueError("Cannot exit the context without completing the workflow")
 
     @typing.overload
@@ -127,11 +127,11 @@ class AIQRunner:
 
     async def result(self, to_type: type | None = None):
 
-        if (self._state != AIQRunnerState.INITIALIZED):
+        if (self._state != RunnerState.INITIALIZED):
             raise ValueError("Cannot run the workflow without entering the context")
 
         try:
-            self._state = AIQRunnerState.RUNNING
+            self._state = RunnerState.RUNNING
 
             if (not self._entry_fn.has_single_output):
                 raise ValueError("Workflow does not support single output")
@@ -145,7 +145,7 @@ class AIQRunner:
                 if event_stream:
                     event_stream.on_complete()
 
-            self._state = AIQRunnerState.COMPLETED
+            self._state = RunnerState.COMPLETED
 
             return result
         except Exception as e:
@@ -153,17 +153,17 @@ class AIQRunner:
             event_stream = self._context_state.event_stream.get()
             if event_stream:
                 event_stream.on_complete()
-            self._state = AIQRunnerState.FAILED
+            self._state = RunnerState.FAILED
 
             raise
 
     async def result_stream(self, to_type: type | None = None):
 
-        if (self._state != AIQRunnerState.INITIALIZED):
+        if (self._state != RunnerState.INITIALIZED):
             raise ValueError("Cannot run the workflow without entering the context")
 
         try:
-            self._state = AIQRunnerState.RUNNING
+            self._state = RunnerState.RUNNING
 
             if (not self._entry_fn.has_streaming_output):
                 raise ValueError("Workflow does not support streaming output")
@@ -173,7 +173,7 @@ class AIQRunner:
                 async for m in self._entry_fn.astream(self._input_message, to_type=to_type):
                     yield m
 
-                self._state = AIQRunnerState.COMPLETED
+                self._state = RunnerState.COMPLETED
 
                 # Close the intermediate stream
                 event_stream = self._context_state.event_stream.get()
@@ -185,6 +185,11 @@ class AIQRunner:
             event_stream = self._context_state.event_stream.get()
             if event_stream:
                 event_stream.on_complete()
-            self._state = AIQRunnerState.FAILED
+            self._state = RunnerState.FAILED
 
             raise
+
+
+# Compatibility aliases with previous releases
+AIQRunnerState = RunnerState
+AIQRunner = Runner

{aiq → nat}/runtime/session.py

@@ -23,15 +23,15 @@ from contextlib import nullcontext
 
 from starlette.requests import HTTPConnection
 
-from aiq.builder.context import AIQContext
-from aiq.builder.context import AIQContextState
-from aiq.builder.workflow import Workflow
-from aiq.data_models.authentication import AuthenticatedContext
-from aiq.data_models.authentication import AuthFlowType
-from aiq.data_models.authentication import AuthProviderBaseConfig
-from aiq.data_models.config import AIQConfig
-from aiq.data_models.interactive import HumanResponse
-from aiq.data_models.interactive import InteractionPrompt
+from nat.builder.context import Context
+from nat.builder.context import ContextState
+from nat.builder.workflow import Workflow
+from nat.data_models.authentication import AuthenticatedContext
+from nat.data_models.authentication import AuthFlowType
+from nat.data_models.authentication import AuthProviderBaseConfig
+from nat.data_models.config import Config
+from nat.data_models.interactive import HumanResponse
+from nat.data_models.interactive import InteractionPrompt
 
 _T = typing.TypeVar("_T")
 
@@ -40,11 +40,11 @@ class UserManagerBase:
     pass
 
 
-class AIQSessionManager:
+class SessionManager:
 
     def __init__(self, workflow: Workflow, max_concurrency: int = 8):
         """
-        The AIQSessionManager class is used to run and manage a user workflow session. It runs and manages the context,
+        The SessionManager class is used to run and manage a user workflow session. It runs and manages the context,
         and configuration of a workflow with the specified concurrency.
 
         Parameters
@@ -61,8 +61,8 @@ class AIQSessionManager:
         self._workflow: Workflow = workflow
 
         self._max_concurrency = max_concurrency
-        self._context_state = AIQContextState.get()
-        self._context = AIQContext(self._context_state)
+        self._context_state = ContextState.get()
+        self._context = Context(self._context_state)
 
         # We save the context because Uvicorn spawns a new process
         # for each request, and we need to restore the context vars
@@ -75,7 +75,7 @@ class AIQSessionManager:
             self._semaphore = nullcontext()
 
     @property
-    def config(self) -> AIQConfig:
+    def config(self) -> Config:
         return self._workflow.config
 
     @property
@@ -83,7 +83,7 @@ class AIQSessionManager:
         return self._workflow
 
     @property
-    def context(self) -> AIQContext:
+    def context(self) -> Context:
         return self._context
 
     @asynccontextmanager
@@ -156,3 +156,7 @@ class AIQSessionManager:
 
         if request.headers.get("conversation-id"):
             self._context_state.conversation_id.set(request.headers["conversation-id"])
+
+
+# Compatibility aliases with previous releases
+AIQSessionManager = SessionManager

{aiq → nat}/runtime/user_metadata.py

@@ -16,7 +16,7 @@
 from starlette.datastructures import Headers
 from starlette.datastructures import QueryParams
 
-from aiq.data_models.api_server import Request
+from nat.data_models.api_server import Request
 
 
 class RequestAttributes:

{aiq → nat}/settings/global_settings.py

@@ -30,12 +30,12 @@ from pydantic import ValidationInfo
 from pydantic import ValidatorFunctionWrapHandler
 from pydantic import field_validator
 
-from aiq.cli.type_registry import GlobalTypeRegistry
-from aiq.cli.type_registry import RegisteredInfo
-from aiq.data_models.common import HashableBaseModel
-from aiq.data_models.common import TypedBaseModel
-from aiq.data_models.common import TypedBaseModelT
-from aiq.data_models.registry_handler import RegistryHandlerBaseConfig
+from nat.cli.type_registry import GlobalTypeRegistry
+from nat.cli.type_registry import RegisteredInfo
+from nat.data_models.common import HashableBaseModel
+from nat.data_models.common import TypedBaseModel
+from nat.data_models.common import TypedBaseModelT
+from nat.data_models.registry_handler import RegistryHandlerBaseConfig
 
 logger = logging.getLogger(__name__)
 
@@ -152,7 +152,7 @@ class Settings(HashableBaseModel):
     @staticmethod
     def from_file():
 
-        configuration_directory = os.getenv("AIQ_CONFIG_DIR", user_config_dir(appname="aiq"))
+        configuration_directory = os.getenv("NAT_CONFIG_DIR", user_config_dir(appname="nat"))
 
         if not os.path.exists(configuration_directory):
             os.makedirs(configuration_directory, exist_ok=True)
@@ -181,7 +181,7 @@ class Settings(HashableBaseModel):
         if (remove):
             if os.path.exists(self.configuration_directory):
                 os.rmdir(self.configuration_directory)
-        self._configuration_directory = os.getenv("AIQ_CONFIG_DIR", user_config_dir(appname="aiq"))
+        self._configuration_directory = os.getenv("NAT_CONFIG_DIR", user_config_dir(appname="nat"))
 
     def _save_settings(self) -> None:
 
@@ -246,7 +246,7 @@ class Settings(HashableBaseModel):
 
     def override_settings(self, config_file: str) -> "Settings":
 
-        from aiq.utils.io.yaml_tools import yaml_load
+        from nat.utils.io.yaml_tools import yaml_load
 
         override_settings_dict = yaml_load(config_file)
 
@@ -293,8 +293,8 @@ class GlobalSettings:
     def get() -> Settings:
 
         if (GlobalSettings._global_settings is None):
-            from aiq.runtime.loader import PluginTypes
-            from aiq.runtime.loader import discover_and_register_plugins
+            from nat.runtime.loader import PluginTypes
+            from nat.runtime.loader import discover_and_register_plugins
 
             discover_and_register_plugins(PluginTypes.REGISTRY_HANDLER)
 

{aiq → nat}/tool/chat_completion.py

@@ -13,7 +13,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 """
-Simple Completion Function for AIQ Toolkit
+Simple Completion Function for NAT
 
 This module provides a simple completion function that can handle
 natural language queries and perform basic text completion tasks.
@@ -21,11 +21,11 @@ natural language queries and perform basic text completion tasks.
 
 from pydantic import Field
 
-from aiq.builder.builder import Builder
-from aiq.builder.framework_enum import LLMFrameworkEnum
-from aiq.cli.register_workflow import register_function
-from aiq.data_models.component_ref import LLMRef
-from aiq.data_models.function import FunctionBaseConfig
+from nat.builder.builder import Builder
+from nat.builder.framework_enum import LLMFrameworkEnum
+from nat.cli.register_workflow import register_function
+from nat.data_models.component_ref import LLMRef
+from nat.data_models.function import FunctionBaseConfig
 
 
 class ChatCompletionConfig(FunctionBaseConfig, name="chat_completion"):

{aiq → nat}/tool/code_execution/README.md

@@ -35,7 +35,7 @@ The Code Execution Sandbox provides:
 Navigate to the local sandbox directory and start the server:
 
 ```bash
-cd src/aiq/tool/code_execution/local_sandbox
+cd src/nat/tool/code_execution/local_sandbox
 ./start_local_sandbox.sh
 ```
 
@@ -62,7 +62,7 @@ export OUTPUT_DATA_PATH=/path/to/output
 Run the comprehensive test suite to verify everything is working:
 
 ```bash
-cd src/aiq/tool/code_execution
+cd src/nat/tool/code_execution
 pytest test_code_execution_sandbox.py
 ```
 

{aiq → nat}/tool/code_execution/code_sandbox.py

@@ -23,7 +23,7 @@ import requests
 import requests.adapters
 from pydantic import HttpUrl
 
-from aiq.utils.type_utils import override
+from nat.utils.type_utils import override
 
 logger = logging.getLogger(__file__)
 

{aiq → nat}/tool/code_execution/register.py

@@ -20,10 +20,10 @@ from pydantic import BaseModel
 from pydantic import Field
 from pydantic import HttpUrl
 
-from aiq.builder.builder import Builder
-from aiq.builder.function_info import FunctionInfo
-from aiq.cli.register_workflow import register_function
-from aiq.data_models.function import FunctionBaseConfig
+from nat.builder.builder import Builder
+from nat.builder.function_info import FunctionInfo
+from nat.cli.register_workflow import register_function
+from nat.data_models.function import FunctionBaseConfig
 
 logger = logging.getLogger(__name__)
 
@@ -41,7 +41,7 @@ class CodeExecutionToolConfig(FunctionBaseConfig, name="code_execution"):
 
 @register_function(config_type=CodeExecutionToolConfig)
 async def code_execution_tool(config: CodeExecutionToolConfig, builder: Builder):
-    from aiq.tool.code_execution.code_sandbox import get_sandbox
+    from nat.tool.code_execution.code_sandbox import get_sandbox
 
     class CodeExecutionInputSchema(BaseModel):
         generated_code: str = Field(description="String containing the code to be executed")

{aiq → nat}/tool/code_execution/test_code_execution_sandbox.py

@@ -50,7 +50,7 @@ class TestCodeExecutionSandbox:
         except (ConnectionError, Timeout, RequestException):
             pytest.skip(
                 f"Sandbox server is not running at {sandbox_config['url']}. "
-                "Please start it with: cd src/aiq/tool/code_execution/local_sandbox && ./start_local_sandbox.sh")
+                "Please start it with: cd src/nat/tool/code_execution/local_sandbox && ./start_local_sandbox.sh")
 
     def execute_code(self, sandbox_config: dict[str, Any], code: str, language: str = "python") -> dict[str, Any]:
         """

{aiq → nat}/tool/datetime_tools.py

@@ -13,10 +13,10 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from aiq.builder.builder import Builder
-from aiq.builder.function_info import FunctionInfo
-from aiq.cli.register_workflow import register_function
-from aiq.data_models.function import FunctionBaseConfig
+from nat.builder.builder import Builder
+from nat.builder.function_info import FunctionInfo
+from nat.cli.register_workflow import register_function
+from nat.data_models.function import FunctionBaseConfig
 
 
 class CurrentTimeToolConfig(FunctionBaseConfig, name="current_datetime"):

{aiq → nat}/tool/document_search.py

@@ -18,12 +18,12 @@ import logging
 
 from pydantic import Field
 
-from aiq.builder.builder import Builder
-from aiq.builder.framework_enum import LLMFrameworkEnum
-from aiq.builder.function_info import FunctionInfo
-from aiq.cli.register_workflow import register_function
-from aiq.data_models.component_ref import LLMRef
-from aiq.data_models.function import FunctionBaseConfig
+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.component_ref import LLMRef
+from nat.data_models.function import FunctionBaseConfig
 
 logger = logging.getLogger(__name__)
 

{aiq → nat}/tool/github_tools/create_github_commit.py

@@ -16,10 +16,10 @@
 from pydantic import BaseModel
 from pydantic import Field
 
-from aiq.builder.builder import Builder
-from aiq.builder.function_info import FunctionInfo
-from aiq.cli.register_workflow import register_function
-from aiq.data_models.function import FunctionBaseConfig
+from nat.builder.builder import Builder
+from nat.builder.function_info import FunctionInfo
+from nat.cli.register_workflow import register_function
+from nat.data_models.function import FunctionBaseConfig
 
 
 class GithubCommitCodeModel(BaseModel):

{aiq → nat}/tool/github_tools/create_github_issue.py

@@ -16,10 +16,10 @@
 from pydantic import BaseModel
 from pydantic import Field
 
-from aiq.builder.builder import Builder
-from aiq.builder.function_info import FunctionInfo
-from aiq.cli.register_workflow import register_function
-from aiq.data_models.function import FunctionBaseConfig
+from nat.builder.builder import Builder
+from nat.builder.function_info import FunctionInfo
+from nat.cli.register_workflow import register_function
+from nat.data_models.function import FunctionBaseConfig
 
 
 class GithubCreateIssueModel(BaseModel):

{aiq → nat}/tool/github_tools/create_github_pr.py

@@ -16,10 +16,10 @@
 from pydantic import BaseModel
 from pydantic import Field
 
-from aiq.builder.builder import Builder
-from aiq.builder.function_info import FunctionInfo
-from aiq.cli.register_workflow import register_function
-from aiq.data_models.function import FunctionBaseConfig
+from nat.builder.builder import Builder
+from nat.builder.function_info import FunctionInfo
+from nat.cli.register_workflow import register_function
+from nat.data_models.function import FunctionBaseConfig
 
 
 class GithubCreatePullModel(BaseModel):

{aiq → nat}/tool/github_tools/get_github_file.py

@@ -13,10 +13,10 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from aiq.builder.builder import Builder
-from aiq.builder.function_info import FunctionInfo
-from aiq.cli.register_workflow import register_function
-from aiq.data_models.function import FunctionBaseConfig
+from nat.builder.builder import Builder
+from nat.builder.function_info import FunctionInfo
+from nat.cli.register_workflow import register_function
+from nat.data_models.function import FunctionBaseConfig
 
 
 class GithubGetFileToolConfig(FunctionBaseConfig, name="github_getfile"):

{aiq → nat}/tool/github_tools/get_github_issue.py

@@ -20,10 +20,10 @@ from pydantic import BaseModel
 from pydantic import Field
 from pydantic import field_validator
 
-from aiq.builder.builder import Builder
-from aiq.builder.function_info import FunctionInfo
-from aiq.cli.register_workflow import register_function
-from aiq.data_models.function import FunctionBaseConfig
+from nat.builder.builder import Builder
+from nat.builder.function_info import FunctionInfo
+from nat.cli.register_workflow import register_function
+from nat.data_models.function import FunctionBaseConfig
 
 
 class GithubListIssueModel(BaseModel):

{aiq → nat}/tool/github_tools/get_github_pr.py

@@ -18,10 +18,10 @@ from typing import Literal
 from pydantic import BaseModel
 from pydantic import Field
 
-from aiq.builder.builder import Builder
-from aiq.builder.function_info import FunctionInfo
-from aiq.cli.register_workflow import register_function
-from aiq.data_models.function import FunctionBaseConfig
+from nat.builder.builder import Builder
+from nat.builder.function_info import FunctionInfo
+from nat.cli.register_workflow import register_function
+from nat.data_models.function import FunctionBaseConfig
 
 
 class GithubListPullsModel(BaseModel):

{aiq → nat}/tool/github_tools/update_github_issue.py

@@ -18,10 +18,10 @@ from typing import Literal
 from pydantic import BaseModel
 from pydantic import Field
 
-from aiq.builder.builder import Builder
-from aiq.builder.function_info import FunctionInfo
-from aiq.cli.register_workflow import register_function
-from aiq.data_models.function import FunctionBaseConfig
+from nat.builder.builder import Builder
+from nat.builder.function_info import FunctionInfo
+from nat.cli.register_workflow import register_function
+from nat.data_models.function import FunctionBaseConfig
 
 
 class GithubUpdateIssueModel(BaseModel):

{aiq → nat}/tool/mcp/exceptions.py

@@ -108,7 +108,7 @@ class MCPToolNotFoundError(MCPError):
                          url=url,
                          category=MCPErrorCategory.TOOL_NOT_FOUND,
                          suggestions=[
-                             "Use 'aiq info mcp --detail' to see available tools",
+                             "Use 'nat info mcp --detail' to see available tools",
                              "Check that the tool name is spelled correctly"
                          ],
                          original_exception=original_exception)

{aiq → nat}/tool/mcp/mcp_client.py

@@ -27,8 +27,8 @@ from pydantic import BaseModel
 from pydantic import Field
 from pydantic import create_model
 
-from aiq.tool.mcp.exceptions import MCPToolNotFoundError
-from aiq.utils.exception_handlers.mcp import mcp_exception_handler
+from nat.tool.mcp.exceptions import MCPToolNotFoundError
+from nat.utils.exception_handlers.mcp import mcp_exception_handler
 
 logger = logging.getLogger(__name__)
 

{aiq → nat}/tool/mcp/mcp_tool.py

@@ -19,10 +19,10 @@ from pydantic import BaseModel
 from pydantic import Field
 from pydantic import HttpUrl
 
-from aiq.builder.builder import Builder
-from aiq.builder.function_info import FunctionInfo
-from aiq.cli.register_workflow import register_function
-from aiq.data_models.function import FunctionBaseConfig
+from nat.builder.builder import Builder
+from nat.builder.function_info import FunctionInfo
+from nat.cli.register_workflow import register_function
+from nat.data_models.function import FunctionBaseConfig
 
 logger = logging.getLogger(__name__)
 
@@ -50,11 +50,11 @@ class MCPToolConfig(FunctionBaseConfig, name="mcp_tool_wrapper"):
 @register_function(config_type=MCPToolConfig)
 async def mcp_tool(config: MCPToolConfig, builder: Builder):  # pylint: disable=unused-argument
     """
-    Generate an AIQ Toolkit Function that wraps a tool provided by the MCP server.
+    Generate a NAT Function that wraps a tool provided by the MCP server.
     """
 
-    from aiq.tool.mcp.mcp_client import MCPBuilder
-    from aiq.tool.mcp.mcp_client import MCPToolClient
+    from nat.tool.mcp.mcp_client import MCPBuilder
+    from nat.tool.mcp.mcp_client import MCPToolClient
 
     client = MCPBuilder(url=str(config.url))
 

{aiq → nat}/tool/memory_tools/add_memory_tool.py

@@ -17,12 +17,12 @@ import logging
 
 from pydantic import Field
 
-from aiq.builder.builder import Builder
-from aiq.builder.function_info import FunctionInfo
-from aiq.cli.register_workflow import register_function
-from aiq.data_models.component_ref import MemoryRef
-from aiq.data_models.function import FunctionBaseConfig
-from aiq.memory.models import MemoryItem
+from nat.builder.builder import Builder
+from nat.builder.function_info import FunctionInfo
+from nat.cli.register_workflow import register_function
+from nat.data_models.component_ref import MemoryRef
+from nat.data_models.function import FunctionBaseConfig
+from nat.memory.models import MemoryItem
 
 logger = logging.getLogger(__name__)
 

{aiq → nat}/tool/memory_tools/delete_memory_tool.py

@@ -17,12 +17,12 @@ import logging
 
 from pydantic import Field
 
-from aiq.builder.builder import Builder
-from aiq.builder.function_info import FunctionInfo
-from aiq.cli.register_workflow import register_function
-from aiq.data_models.component_ref import MemoryRef
-from aiq.data_models.function import FunctionBaseConfig
-from aiq.memory.models import DeleteMemoryInput
+from nat.builder.builder import Builder
+from nat.builder.function_info import FunctionInfo
+from nat.cli.register_workflow import register_function
+from nat.data_models.component_ref import MemoryRef
+from nat.data_models.function import FunctionBaseConfig
+from nat.memory.models import DeleteMemoryInput
 
 logger = logging.getLogger(__name__)
 

{aiq → nat}/tool/memory_tools/get_memory_tool.py

@@ -17,12 +17,12 @@ import logging
 
 from pydantic import Field
 
-from aiq.builder.builder import Builder
-from aiq.builder.function_info import FunctionInfo
-from aiq.cli.register_workflow import register_function
-from aiq.data_models.component_ref import MemoryRef
-from aiq.data_models.function import FunctionBaseConfig
-from aiq.memory.models import SearchMemoryInput
+from nat.builder.builder import Builder
+from nat.builder.function_info import FunctionInfo
+from nat.cli.register_workflow import register_function
+from nat.data_models.component_ref import MemoryRef
+from nat.data_models.function import FunctionBaseConfig
+from nat.memory.models import SearchMemoryInput
 
 logger = logging.getLogger(__name__)