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

nat/front_ends/fastapi/fastapi_front_end_plugin.py

@@ -47,7 +47,7 @@ class FastApiFrontEndPlugin(FrontEndBase[FastApiFrontEndConfig]):
     async def run(self):
 
         # Write the entire config to a temporary file
-        with tempfile.NamedTemporaryFile(mode="w", prefix="aiq_config", suffix=".yml", delete=False) as config_file:
+        with tempfile.NamedTemporaryFile(mode="w", prefix="nat_config", suffix=".yml", delete=False) as config_file:
 
             # Get as dict
             config_dict = self.full_config.model_dump(mode="json", by_alias=True, round_trip=True)
@@ -59,10 +59,10 @@ class FastApiFrontEndPlugin(FrontEndBase[FastApiFrontEndConfig]):
             config_file_name = config_file.name
 
             # Set the config file in the environment
-            os.environ["AIQ_CONFIG_FILE"] = str(config_file.name)
+            os.environ["NAT_CONFIG_FILE"] = str(config_file.name)
 
             # Set the worker class in the environment
-            os.environ["AIQ_FRONT_END_WORKER"] = self.get_worker_class_name()
+            os.environ["NAT_FRONT_END_WORKER"] = self.get_worker_class_name()
 
         try:
             if not self.front_end_config.use_gunicorn:
@@ -70,7 +70,7 @@ class FastApiFrontEndPlugin(FrontEndBase[FastApiFrontEndConfig]):
 
                 reload_excludes = ["./.*"]
 
-                uvicorn.run("aiq.front_ends.fastapi.main:get_app",
+                uvicorn.run("nat.front_ends.fastapi.main:get_app",
                             host=self.front_end_config.host,
                             port=self.front_end_config.port,
                             workers=self.front_end_config.workers,

nat/front_ends/fastapi/fastapi_front_end_plugin_worker.py

@@ -99,7 +99,7 @@ class FastApiFrontEndPluginWorkerBase(ABC):
         @asynccontextmanager
         async def lifespan(starting_app: FastAPI):
 
-            logger.debug("Starting AIQ Toolkit server from process %s", os.getpid())
+            logger.debug("Starting NAT server from process %s", os.getpid())
 
             async with WorkflowBuilder.from_config(self.config) as builder:
 
@@ -121,20 +121,20 @@ class FastApiFrontEndPluginWorkerBase(ABC):
 
                     self._cleanup_tasks.clear()
 
-            logger.debug("Closing AIQ Toolkit server from process %s", os.getpid())
+            logger.debug("Closing NAT server from process %s", os.getpid())
 
-        aiq_app = FastAPI(lifespan=lifespan)
+        nat_app = FastAPI(lifespan=lifespan)
 
         # Configure app CORS.
-        self.set_cors_config(aiq_app)
+        self.set_cors_config(nat_app)
 
-        @aiq_app.middleware("http")
+        @nat_app.middleware("http")
         async def authentication_log_filter(request: Request, call_next: Callable[[Request], Awaitable[Response]]):
             return await self._suppress_authentication_logs(request, call_next)
 
-        return aiq_app
+        return nat_app
 
-    def set_cors_config(self, aiq_app: FastAPI) -> None:
+    def set_cors_config(self, nat_app: FastAPI) -> None:
         """
         Set the cross origin resource sharing configuration.
         """
@@ -161,7 +161,7 @@ class FastApiFrontEndPluginWorkerBase(ABC):
         if self.front_end_config.cors.max_age is not None:
             cors_kwargs["max_age"] = self.front_end_config.cors.max_age
 
-        aiq_app.add_middleware(
+        nat_app.add_middleware(
             CORSMiddleware,
             **cors_kwargs,
         )
@@ -322,7 +322,7 @@ class FastApiFrontEndPluginWorker(FastApiFrontEndPluginWorkerBase):
                 return EvaluateResponse(job_id=job_id, status="submitted")
 
         def translate_job_to_response(job: JobInfo) -> EvaluateStatusResponse:
-            """Translate a JobInfo object to an AIQEvaluateStatusResponse."""
+            """Translate a JobInfo object to an EvaluateStatusResponse."""
             return EvaluateStatusResponse(job_id=job.job_id,
                                           status=job.status,
                                           config_file=str(job.config_file),
@@ -527,7 +527,7 @@ class FastApiFrontEndPluginWorker(FastApiFrontEndPluginWorkerBase):
         GenerateSingleResponseType = workflow.single_output_schema  # pylint: disable=invalid-name
 
         # Append job_id and expiry_seconds to the input schema, this effectively makes these reserved keywords
-        # Consider prefixing these with "aiq_" to avoid conflicts
+        # Consider prefixing these with "nat_" to avoid conflicts
         class AsyncGenerateRequest(GenerateBodyType):
             job_id: str | None = Field(default=None, description="Unique identifier for the evaluation job")
             sync_timeout: int = Field(
@@ -820,7 +820,7 @@ class FastApiFrontEndPluginWorker(FastApiFrontEndPluginWorkerBase):
             session_id = websocket.query_params.get("session")
             if session_id:
                 headers = list(websocket.scope.get("headers", []))
-                cookie_header = f"aiqtoolkit-session={session_id}"
+                cookie_header = f"nat-session={session_id}"
 
                 # Check if the session cookie already exists to avoid duplicates
                 cookie_exists = False
@@ -831,8 +831,8 @@ class FastApiFrontEndPluginWorker(FastApiFrontEndPluginWorkerBase):
                         cookie_exists = True
                         cookie_str = value.decode()
 
-                        # Check if aiqtoolkit-session already exists in cookies
-                        if "aiqtoolkit-session=" in cookie_str:
+                        # Check if nat-session already exists in cookies
+                        if "nat-session=" in cookie_str:
                             existing_session_cookie = True
                             logger.info("WebSocket: Session cookie already present in headers (same-origin)")
                         else:

nat/front_ends/fastapi/intermediate_steps_subscriber.py

@@ -41,7 +41,7 @@ async def pull_intermediate(_q, adapter):
         Synchronously called whenever the runner publishes an event.
         We process it, then place it into the async queue (via a small async task).
         If adapter is None, convert the raw IntermediateStep into the complete
-        AIQResponseIntermediateStep and place it into the queue.
+        ResponseIntermediateStep and place it into the queue.
         """
         if adapter is None:
             adapted = ResponseIntermediateStep(id=item.UUID,

nat/front_ends/fastapi/main.py

@@ -25,14 +25,14 @@ logger = logging.getLogger(__name__)
 
 def get_app():
 
-    config_file_path = os.getenv("AIQ_CONFIG_FILE")
-    front_end_worker_full_name = os.getenv("AIQ_FRONT_END_WORKER")
+    config_file_path = os.getenv("NAT_CONFIG_FILE")
+    front_end_worker_full_name = os.getenv("NAT_FRONT_END_WORKER")
 
     if (not config_file_path):
-        raise ValueError("Config file not found in environment variable AIQ_CONFIG_FILE.")
+        raise ValueError("Config file not found in environment variable NAT_CONFIG_FILE.")
 
     if (not front_end_worker_full_name):
-        raise ValueError("Front end worker not found in environment variable AIQ_FRONT_END_WORKER.")
+        raise ValueError("Front end worker not found in environment variable NAT_FRONT_END_WORKER.")
 
     # Try to import the front end worker class
     try:
@@ -62,9 +62,9 @@ def get_app():
         # Create an instance of the front end worker class
         front_end_worker = front_end_worker_class(config)
 
-        aiq_app = front_end_worker.build_app()
+        nat_app = front_end_worker.build_app()
 
-        return aiq_app
+        return nat_app
 
     except ImportError as e:
         raise ValueError(f"Front end worker {front_end_worker_full_name} not found.") from e

nat/front_ends/fastapi/message_handler.py

@@ -26,6 +26,7 @@ from starlette.websockets import WebSocketDisconnect
 
 from nat.authentication.interfaces import FlowHandlerBase
 from nat.data_models.api_server import ChatResponse
+from nat.data_models.api_server import ChatResponseChunk
 from nat.data_models.api_server import Error
 from nat.data_models.api_server import ErrorTypes
 from nat.data_models.api_server import ResponsePayloadOutput
@@ -39,6 +40,7 @@ from nat.data_models.api_server import WebSocketSystemIntermediateStepMessage
 from nat.data_models.api_server import WebSocketSystemResponseTokenMessage
 from nat.data_models.api_server import WebSocketUserInteractionResponseMessage
 from nat.data_models.api_server import WebSocketUserMessage
+from nat.data_models.api_server import WorkflowSchemaType
 from nat.data_models.interactive import HumanPromptNotification
 from nat.data_models.interactive import HumanResponse
 from nat.data_models.interactive import HumanResponseNotification
@@ -67,6 +69,13 @@ class WebSocketMessageHandler:
 
         self._flow_handler: FlowHandlerBase | None = None
 
+        self._schema_output_mapping: dict[str, type[BaseModel] | None] = {
+            WorkflowSchemaType.GENERATE: self._session_manager.workflow.single_output_schema,
+            WorkflowSchemaType.CHAT: ChatResponse,
+            WorkflowSchemaType.CHAT_STREAM: ChatResponseChunk,
+            WorkflowSchemaType.GENERATE_STREAM: self._session_manager.workflow.streaming_output_schema,
+        }
+
     def set_flow_handler(self, flow_handler: FlowHandlerBase) -> None:
         self._flow_handler = flow_handler
 
@@ -156,10 +165,12 @@ class WebSocketMessageHandler:
                 def _done_callback(task: asyncio.Task):  # pylint: disable=unused-argument
                     self._running_workflow_task = None
 
-                # await self._process_response()
                 self._running_workflow_task = asyncio.create_task(
-                    self._run_workflow(content.text, self._conversation_id,
-                                       result_type=ChatResponse)).add_done_callback(_done_callback)
+                    self._run_workflow(content.text,
+                                       self._conversation_id,
+                                       result_type=self._schema_output_mapping[self._workflow_schema_type],
+                                       output_type=self._schema_output_mapping[
+                                           self._workflow_schema_type])).add_done_callback(_done_callback)
 
         except ValueError as e:
             logger.error("User message content not found: %s", str(e), exc_info=True)

nat/front_ends/fastapi/message_validator.py

@@ -39,7 +39,6 @@ from nat.data_models.api_server import WebSocketSystemIntermediateStepMessage
 from nat.data_models.api_server import WebSocketSystemResponseTokenMessage
 from nat.data_models.api_server import WebSocketUserInteractionResponseMessage
 from nat.data_models.api_server import WebSocketUserMessage
-from nat.data_models.api_server import WorkflowSchemaType
 from nat.data_models.interactive import BinaryHumanPromptOption
 from nat.data_models.interactive import HumanPrompt
 from nat.data_models.interactive import HumanPromptBase
@@ -70,12 +69,7 @@ class MessageValidator:
             WebSocketMessageType.USER_INTERACTION_MESSAGE: WebSocketUserInteractionResponseMessage,
             WebSocketMessageType.ERROR_MESSAGE: Error
         }
-        self._data_type_schema_mapping: dict[str, type[BaseModel]] = {
-            WorkflowSchemaType.GENERATE: ResponsePayloadOutput,
-            WorkflowSchemaType.CHAT: ChatResponse,
-            WorkflowSchemaType.CHAT_STREAM: ChatResponseChunk,
-            WorkflowSchemaType.GENERATE_STREAM: ResponseIntermediateStep,
-        }
+
         self._message_parent_id: str = "default_id"
 
     async def validate_message(self, message: dict[str, Any]) -> BaseModel:
@@ -139,7 +133,11 @@ class MessageValidator:
         validated_message_content: BaseModel = None
         try:
             if (isinstance(data_model, ResponsePayloadOutput)):
-                validated_message_content = SystemResponseContent(text=data_model.payload)
+                if hasattr(data_model.payload, 'model_dump_json'):
+                    text_content: str = data_model.payload.model_dump_json()
+                else:
+                    text_content: str = str(data_model.payload)
+                validated_message_content = SystemResponseContent(text=text_content)
 
             elif (isinstance(data_model, (ChatResponse, ChatResponseChunk))):
                 validated_message_content = SystemResponseContent(text=data_model.choices[0].message.content)

nat/front_ends/fastapi/response_helpers.py

@@ -45,7 +45,7 @@ async def generate_streaming_response_as_str(payload: typing.Any,
         if (isinstance(item, ResponseSerializable)):
             yield item.get_stream_data()
         else:
-            raise ValueError("Unexpected item type in stream. Expected AIQChatResponseSerializable, got: " +
+            raise ValueError("Unexpected item type in stream. Expected ChatResponseSerializable, got: " +
                              str(type(item)))
 
 
@@ -125,7 +125,7 @@ async def generate_streaming_response_full(payload: typing.Any,
                                            output_type: type | None = None,
                                            filter_steps: str | None = None) -> AsyncGenerator[ResponseSerializable]:
     """
-    Similar to generate_streaming_response but provides raw AIQResponseIntermediateStep objects
+    Similar to generate_streaming_response but provides raw ResponseIntermediateStep objects
     without any step adaptor translations.
     """
     # Parse filter_steps into a set of allowed types if provided
@@ -191,5 +191,5 @@ async def generate_streaming_response_full_as_str(payload: typing.Any,
         if (isinstance(item, ResponseIntermediateStep) or isinstance(item, ResponsePayloadOutput)):
             yield item.get_stream_data()
         else:
-            raise ValueError("Unexpected item type in stream. Expected AIQChatResponseSerializable, got: " +
+            raise ValueError("Unexpected item type in stream. Expected ChatResponseSerializable, got: " +
                              str(type(item)))

nat/front_ends/mcp/mcp_front_end_config.py

@@ -21,10 +21,11 @@ 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 (default: AIQ MCP)")
+    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)")

nat/front_ends/mcp/mcp_front_end_plugin_worker.py

@@ -23,7 +23,7 @@ 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 AIQConfig
+from nat.data_models.config import Config
 from nat.front_ends.mcp.mcp_front_end_config import MCPFrontEndConfig
 
 logger = logging.getLogger(__name__)
@@ -32,11 +32,11 @@ logger = logging.getLogger(__name__)
 class MCPFrontEndPluginWorkerBase(ABC):
     """Base class for MCP front end plugin workers."""
 
-    def __init__(self, config: AIQConfig):
+    def __init__(self, config: Config):
         """Initialize the MCP worker with configuration.
 
         Args:
-            config: The full AIQ configuration
+            config: The full NAT configuration
         """
         self.full_config = config
         self.front_end_config: MCPFrontEndConfig = config.general.front_end
@@ -86,7 +86,7 @@ class MCPFrontEndPluginWorkerBase(ABC):
         """Get all functions from the workflow.
 
         Args:
-            workflow: The AIQ workflow.
+            workflow: The NAT workflow.
 
         Returns:
             Dict mapping function names to Function objects.

nat/front_ends/mcp/tool_converter.py

@@ -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,7 +101,7 @@ def create_function_wrapper(
                 await ctx.report_progress(0, 100)
 
             try:
-                # Special handling for AIQChatRequest
+                # Special handling for ChatRequest
                 if is_chat_request:
                     from nat.data_models.api_server import ChatRequest
 
@@ -113,7 +112,7 @@ def create_function_wrapper(
                     # 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)
 

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.
 """

nat/meta/pypi.md

@@ -15,11 +15,11 @@ See the License for the specific language governing permissions and
 limitations under the License.
 -->
 
-![NVIDIA NeMo Agent Toolkit](https://media.githubusercontent.com/media/NVIDIA/NeMo-Agent-Toolkit/refs/heads/main/docs/source/_static/aiqtoolkit_banner.png "NeMo Agent toolkit banner image")
+![NVIDIA NeMo Agent Toolkit](https://media.githubusercontent.com/media/NVIDIA/NeMo-Agent-Toolkit/refs/heads/main/docs/source/_static/banner.png "NeMo Agent toolkit banner image")
 
 # NVIDIA NeMo Agent Toolkit
 
-AIQ toolkit is a flexible library designed to seamlessly integrate your enterprise agents—regardless of framework—with various data sources and tools. By treating agents, tools, and agentic workflows as simple function calls, AIQ toolkit enables true composability: build once and reuse anywhere.
+NeMo Agent toolkit is a flexible library designed to seamlessly integrate your enterprise agents—regardless of framework—with various data sources and tools. By treating agents, tools, and agentic workflows as simple function calls, NeMo Agent toolkit enables true composability: build once and reuse anywhere.
 
 ## Key Features
 
@@ -29,16 +29,16 @@ AIQ toolkit is a flexible library designed to seamlessly integrate your enterpri
 - [**Profiling:**](https://docs.nvidia.com/nemo/agent-toolkit/1.2.0/workflows/profiler.html) Profile entire workflows down to the tool and agent level, track input/output tokens and timings, and identify bottlenecks.
 - [**Observability:**](https://docs.nvidia.com/nemo/agent-toolkit/1.2.0/workflows/observe/observe-workflow-with-phoenix.html) Monitor and debug your workflows with any OpenTelemetry-compatible observability tool, with examples using [Phoenix](https://docs.nvidia.com/nemo/agent-toolkit/1.2.0/workflows/observe/observe-workflow-with-phoenix.html) and [W&B Weave](https://docs.nvidia.com/nemo/agent-toolkit/1.2.0/workflows/observe/observe-workflow-with-weave.html).
 - [**Evaluation System:**](https://docs.nvidia.com/nemo/agent-toolkit/1.2.0/workflows/evaluate.html) Validate and maintain accuracy of agentic workflows with built-in evaluation tools.
-- [**User Interface:**](https://docs.nvidia.com/nemo/agent-toolkit/1.2.0/quick-start/launching-ui.html) Use the AIQ toolkit UI chat interface to interact with your agents, visualize output, and debug workflows.
-- [**MCP Compatibility**](https://docs.nvidia.com/nemo/agent-toolkit/1.2.0/workflows/mcp/mcp-client.html) Compatible with Model Context Protocol (MCP), allowing tools served by MCP Servers to be used as AIQ toolkit functions.
+- [**User Interface:**](https://docs.nvidia.com/nemo/agent-toolkit/1.2.0/quick-start/launching-ui.html) Use the NeMo Agent toolkit UI chat interface to interact with your agents, visualize output, and debug workflows.
+- [**MCP Compatibility**](https://docs.nvidia.com/nemo/agent-toolkit/1.2.0/workflows/mcp/mcp-client.html) Compatible with Model Context Protocol (MCP), allowing tools served by MCP Servers to be used as NeMo Agent toolkit functions.
 
-With AIQ toolkit, you can move quickly, experiment freely, and ensure reliability across all your agent-driven projects.
+With NeMo Agent toolkit, you can move quickly, experiment freely, and ensure reliability across all your agent-driven projects.
 
 ## Links
- * [Documentation](https://docs.nvidia.com/nemo/agent-toolkit/1.2.0/index.html): Explore the full documentation for AIQ toolkit.
+ * [Documentation](https://docs.nvidia.com/nemo/agent-toolkit/1.2.0/index.html): Explore the full documentation for NeMo Agent toolkit.
 
 ## First time user?
- If this is your first time using AIQ toolkit, it is recommended to install the latest version from the [source repository](https://github.com/NVIDIA/NeMo-Agent-Toolkit?tab=readme-ov-file#quick-start) on GitHub. This package is intended for users who are familiar with AIQ toolkit applications and need to add AIQ toolkit as a dependency to their project.
+ If this is your first time using NeMo Agent toolkit, it is recommended to install the latest version from the [source repository](https://github.com/NVIDIA/NeMo-Agent-Toolkit?tab=readme-ov-file#quick-start) on GitHub. This package is intended for users who are familiar with NeMo Agent toolkit applications and need to add NeMo Agent toolkit as a dependency to their project.
 
 ## Feedback
 
@@ -46,7 +46,7 @@ We would love to hear from you! Please file an issue on [GitHub](https://github.
 
 ## Acknowledgements
 
-We would like to thank the following open source projects that made AIQ toolkit possible:
+We would like to thank the following open source projects that made NeMo Agent toolkit possible:
 
 - [CrewAI](https://github.com/crewAIInc/crewAI)
 - [FastAPI](https://github.com/tiangolo/fastapi)

nat/object_store/__init__.py

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

nat/observability/exporter/base_exporter.py

@@ -117,7 +117,7 @@ class BaseExporter(Exporter):
         - This keeps observability "off the hot path" for optimal performance
 
     Args:
-        context_state (AIQContextState, optional): The context state to use for the exporter. Defaults to None.
+        context_state (ContextState, optional): The context state to use for the exporter. Defaults to None.
     """
 
     # Class-level tracking for debugging and monitoring

nat/observability/exporter/raw_exporter.py

@@ -37,7 +37,7 @@ class RawExporter(ProcessingExporter[InputT, OutputT]):
     The flow is: IntermediateStep -> [Processing Pipeline] -> OutputT -> Export
 
     Args:
-        context_state (AIQContextState, optional): The context state to use for the exporter. Defaults to None.
+        context_state (ContextState, optional): The context state to use for the exporter. Defaults to None.
     """
 
     @abstractmethod

nat/observability/exporter_manager.py

@@ -141,7 +141,7 @@ class ExporterManager:
         expensive resources but have separate mutable state.
 
         Args:
-            context_state (AIQContextState | None, optional): The isolated context state for the new exporter instances.
+            context_state (ContextState | None, optional): The isolated context state for the new exporter instances.
                 If not provided, a new context state will be created.
 
         Returns:

nat/profiler/callbacks/agno_callback_handler.py

@@ -44,7 +44,7 @@ class AgnoProfilerHandler(BaseProfilerCallback):
       - LLM Calls
 
     to collect usage statistics (tokens, inputs, outputs, time intervals, etc.)
-    and store them in AIQ Toolkit's usage_stats queue for subsequent analysis.
+    and store them in NAT's usage_stats queue for subsequent analysis.
     """
 
     def __init__(self) -> None:
@@ -73,8 +73,8 @@ class AgnoProfilerHandler(BaseProfilerCallback):
 
         # Note: Agno doesn't have a class-based tool structure to patch directly.
         # Instead, it uses decorators to convert functions to tools.
-        # In AIQ Toolkit, tool executions are captured at the execute_agno_tool level
-        # in packages/aiqtoolkit_agno/src/aiq/plugins/agno/tool_wrapper.py
+        # In NAT, tool executions are captured at the execute_agno_tool level
+        # in packages/nvidia_nat_agno/src/nat/plugins/agno/tool_wrapper.py
 
         # To properly monitor Agno tool executions, we would need to either:
         # 1. Patch the execute_agno_tool function in tool_wrapper.py
@@ -83,7 +83,7 @@ class AgnoProfilerHandler(BaseProfilerCallback):
         #    to patch those classes
 
         # Recommended future enhancement:
-        # The execute_agno_tool function in packages/aiqtoolkit_agno/src/aiq/plugins/agno/tool_wrapper.py
+        # The execute_agno_tool function in packages/nvidia_nat_agno/src/nat/plugins/agno/tool_wrapper.py
         # should be updated to directly push IntermediateStepPayload events to the step manager
         # at the beginning and end of tool execution, similar to what this handler does for LLM calls.
 

nat/profiler/callbacks/llama_index_callback_handler.py

@@ -49,7 +49,7 @@ class LlamaIndexProfilerHandler(BaseCallbackHandler, BaseProfilerCallback):
     - Response data
     - Time intervals between calls
 
-    and appends them to AIQContextState.usage_stats.
+    and appends them to ContextState.usage_stats.
     """
 
     def __init__(self) -> None:
@@ -167,7 +167,7 @@ class LlamaIndexProfilerHandler(BaseCallbackHandler, BaseProfilerCallback):
                 except Exception as e:
                     logger.exception("Error getting model name: %s", e, exc_info=True)
 
-                # Append usage data to AIQ Toolkit usage stats
+                # Append usage data to NAT usage stats
                 with self._lock:
                     stats = IntermediateStepPayload(
                         event_type=IntermediateStepType.LLM_END,

nat/profiler/callbacks/semantic_kernel_callback_handler.py

@@ -55,7 +55,7 @@ class SemanticKernelProfilerHandler(BaseProfilerCallback):
     - Tool calls
 
     to collect usage statistics (tokens, inputs, outputs, time intervals, etc.)
-    and store them in AIQ Toolkit's usage_stats queue for subsequent analysis.
+    and store them in NAT's usage_stats queue for subsequent analysis.
     """
 
     def __init__(self, workflow_llms: dict) -> None:

nat/profiler/decorators/function_tracking.py

@@ -61,7 +61,7 @@ def push_intermediate_step(step_manager: IntermediateStepManager,
                            kwargs: Any = None,
                            output: Any = None,
                            metadata: dict[str, Any] | None = None) -> None:
-    """Push an intermediate step to the AIQ Toolkit Event Stream."""
+    """Push an intermediate step to the NAT Event Stream."""
 
     payload = IntermediateStepPayload(UUID=identifier,
                                       event_type=event_type,

nat/profiler/forecasting/models/linear_model.py

@@ -34,8 +34,9 @@ class LinearModel(ForecastingBaseModel):
         try:
             from sklearn.linear_model import LinearRegression
         except ImportError:
-            logger.error("scikit-learn is not installed. Please install scikit-learn to use the LinearModel "
-                         "profiling model or install `aiq[profiler]` to install all necessary profiling packages.")
+            logger.error(
+                "scikit-learn is not installed. Please install scikit-learn to use the LinearModel "
+                "profiling model or install `nvidia-nat[profiler]` to install all necessary profiling packages.")
 
             raise
 

nat/profiler/forecasting/models/random_forest_regressor.py

@@ -34,8 +34,9 @@ class RandomForestModel(ForecastingBaseModel):
         try:
             from sklearn.ensemble import RandomForestRegressor
         except ImportError:
-            logger.error("scikit-learn is not installed. Please install scikit-learn to use the RandomForest "
-                         "profiling model or install `aiq[profiler]` to install all necessary profiling packages.")
+            logger.error(
+                "scikit-learn is not installed. Please install scikit-learn to use the RandomForest "
+                "profiling model or install `nvidia-nat[profiler]` to install all necessary profiling packages.")
 
             raise
 

nat/profiler/inference_optimization/bottleneck_analysis/nested_stack_analysis.py

@@ -304,7 +304,7 @@ def save_gantt_chart(all_nodes: list[CallNode], output_path: str) -> None:
         import matplotlib.pyplot as plt
     except ImportError:
         logger.error("matplotlib is not installed. Please install matplotlib to use generate plots for the profiler "
-                     "or install `aiq[profiler]` to install all necessary profiling packages.")
+                     "or install `nvidia-nat[profiler]` to install all necessary profiling packages.")
 
         raise
 

nat/profiler/inference_optimization/experimental/prefix_span_analysis.py

@@ -212,7 +212,7 @@ def run_prefixspan(sequences_map: dict[int, list[PrefixCallNode]],
         from prefixspan import PrefixSpan
     except ImportError:
         logger.error("prefixspan is not installed. Please install prefixspan to run the prefix analysis in the "
-                     "profiler or install `aiq[profiler]` to install all necessary profiling packages.")
+                     "profiler or install `nvidia-nat[profiler]` to install all necessary profiling packages.")
 
         raise