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

{aiq → nat}/profiler/callbacks/agno_callback_handler.py

@@ -23,15 +23,15 @@ from uuid import uuid4
 
 import litellm
 
-from aiq.builder.context import AIQContext
-from aiq.builder.framework_enum import LLMFrameworkEnum
-from aiq.data_models.intermediate_step import IntermediateStepPayload
-from aiq.data_models.intermediate_step import IntermediateStepType
-from aiq.data_models.intermediate_step import StreamEventData
-from aiq.data_models.intermediate_step import TraceMetadata
-from aiq.data_models.intermediate_step import UsageInfo
-from aiq.profiler.callbacks.base_callback_class import BaseProfilerCallback
-from aiq.profiler.callbacks.token_usage_base_model import TokenUsageBaseModel
+from nat.builder.context import Context
+from nat.builder.framework_enum import LLMFrameworkEnum
+from nat.data_models.intermediate_step import IntermediateStepPayload
+from nat.data_models.intermediate_step import IntermediateStepType
+from nat.data_models.intermediate_step import StreamEventData
+from nat.data_models.intermediate_step import TraceMetadata
+from nat.data_models.intermediate_step import UsageInfo
+from nat.profiler.callbacks.base_callback_class import BaseProfilerCallback
+from nat.profiler.callbacks.token_usage_base_model import TokenUsageBaseModel
 
 logger = logging.getLogger(__name__)
 
@@ -44,14 +44,14 @@ 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:
         super().__init__()
         self._lock = threading.Lock()
         self.last_call_ts = time.time()
-        self.step_manager = AIQContext.get().intermediate_step_manager
+        self.step_manager = Context.get().intermediate_step_manager
 
         # Original references to Agno methods (for uninstrumenting if needed)
         self._original_tool_execute = 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.
 

{aiq → nat}/profiler/callbacks/langchain_callback_handler.py

@@ -29,16 +29,16 @@ from langchain_core.messages import BaseMessage
 from langchain_core.outputs import ChatGeneration
 from langchain_core.outputs import LLMResult
 
-from aiq.builder.context import AIQContext
-from aiq.builder.framework_enum import LLMFrameworkEnum
-from aiq.data_models.intermediate_step import IntermediateStepPayload
-from aiq.data_models.intermediate_step import IntermediateStepType
-from aiq.data_models.intermediate_step import StreamEventData
-from aiq.data_models.intermediate_step import ToolSchema
-from aiq.data_models.intermediate_step import TraceMetadata
-from aiq.data_models.intermediate_step import UsageInfo
-from aiq.profiler.callbacks.base_callback_class import BaseProfilerCallback
-from aiq.profiler.callbacks.token_usage_base_model import TokenUsageBaseModel
+from nat.builder.context import Context
+from nat.builder.framework_enum import LLMFrameworkEnum
+from nat.data_models.intermediate_step import IntermediateStepPayload
+from nat.data_models.intermediate_step import IntermediateStepType
+from nat.data_models.intermediate_step import StreamEventData
+from nat.data_models.intermediate_step import ToolSchema
+from nat.data_models.intermediate_step import TraceMetadata
+from nat.data_models.intermediate_step import UsageInfo
+from nat.profiler.callbacks.base_callback_class import BaseProfilerCallback
+from nat.profiler.callbacks.token_usage_base_model import TokenUsageBaseModel
 
 logger = logging.getLogger(__name__)
 
@@ -68,7 +68,7 @@ class LangchainProfilerHandler(AsyncCallbackHandler, BaseProfilerCallback):  # p
         self._lock = threading.Lock()
         self.last_call_ts = time.time()
 
-        self.step_manager = AIQContext.get().intermediate_step_manager
+        self.step_manager = Context.get().intermediate_step_manager
         self._state = IntermediateStepType.LLM_END
 
         self._run_id_to_model_name = {}

{aiq → nat}/profiler/callbacks/llama_index_callback_handler.py

@@ -26,15 +26,15 @@ from llama_index.core.callbacks import EventPayload
 from llama_index.core.callbacks.base_handler import BaseCallbackHandler
 from llama_index.core.llms import ChatResponse
 
-from aiq.builder.context import AIQContext
-from aiq.builder.framework_enum import LLMFrameworkEnum
-from aiq.data_models.intermediate_step import IntermediateStepPayload
-from aiq.data_models.intermediate_step import IntermediateStepType
-from aiq.data_models.intermediate_step import StreamEventData
-from aiq.data_models.intermediate_step import TraceMetadata
-from aiq.data_models.intermediate_step import UsageInfo
-from aiq.profiler.callbacks.base_callback_class import BaseProfilerCallback
-from aiq.profiler.callbacks.token_usage_base_model import TokenUsageBaseModel
+from nat.builder.context import Context
+from nat.builder.framework_enum import LLMFrameworkEnum
+from nat.data_models.intermediate_step import IntermediateStepPayload
+from nat.data_models.intermediate_step import IntermediateStepType
+from nat.data_models.intermediate_step import StreamEventData
+from nat.data_models.intermediate_step import TraceMetadata
+from nat.data_models.intermediate_step import UsageInfo
+from nat.profiler.callbacks.base_callback_class import BaseProfilerCallback
+from nat.profiler.callbacks.token_usage_base_model import TokenUsageBaseModel
 
 logger = logging.getLogger(__name__)
 
@@ -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:
@@ -58,7 +58,7 @@ class LlamaIndexProfilerHandler(BaseCallbackHandler, BaseProfilerCallback):
         self._lock = threading.Lock()
         self.last_call_ts = time.time()
         self._last_tool_map: dict[str, str] = {}
-        self.step_manager = AIQContext.get().intermediate_step_manager
+        self.step_manager = Context.get().intermediate_step_manager
 
         self._run_id_to_llm_input = {}
         self._run_id_to_tool_input = {}
@@ -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,

{aiq → nat}/profiler/callbacks/semantic_kernel_callback_handler.py

@@ -26,15 +26,15 @@ from pydantic import BaseModel
 from pydantic import Field
 from semantic_kernel.connectors.ai.open_ai.services.open_ai_chat_completion_base import OpenAIChatCompletionBase
 
-from aiq.builder.context import AIQContext
-from aiq.builder.framework_enum import LLMFrameworkEnum
-from aiq.data_models.intermediate_step import IntermediateStepPayload
-from aiq.data_models.intermediate_step import IntermediateStepType
-from aiq.data_models.intermediate_step import StreamEventData
-from aiq.data_models.intermediate_step import TraceMetadata
-from aiq.data_models.intermediate_step import UsageInfo
-from aiq.profiler.callbacks.base_callback_class import BaseProfilerCallback
-from aiq.profiler.callbacks.token_usage_base_model import TokenUsageBaseModel
+from nat.builder.context import Context
+from nat.builder.framework_enum import LLMFrameworkEnum
+from nat.data_models.intermediate_step import IntermediateStepPayload
+from nat.data_models.intermediate_step import IntermediateStepType
+from nat.data_models.intermediate_step import StreamEventData
+from nat.data_models.intermediate_step import TraceMetadata
+from nat.data_models.intermediate_step import UsageInfo
+from nat.profiler.callbacks.base_callback_class import BaseProfilerCallback
+from nat.profiler.callbacks.token_usage_base_model import TokenUsageBaseModel
 
 logger = logging.getLogger(__name__)
 
@@ -55,14 +55,14 @@ 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:
         super().__init__()
         self._lock = threading.Lock()
         self.last_call_ts = time.time()
-        self.step_manager = AIQContext.get().intermediate_step_manager
+        self.step_manager = Context.get().intermediate_step_manager
         self._builder_llms = workflow_llms
 
         # Original references to SK methods

{aiq → nat}/profiler/data_models.py

@@ -15,8 +15,8 @@
 
 from pydantic import BaseModel
 
-from aiq.profiler.inference_metrics_model import InferenceMetricsModel
-from aiq.profiler.inference_optimization.data_models import WorkflowRuntimeMetrics
+from nat.profiler.inference_metrics_model import InferenceMetricsModel
+from nat.profiler.inference_optimization.data_models import WorkflowRuntimeMetrics
 
 
 class ProfilerResults(BaseModel):

{aiq → nat}/profiler/decorators/framework_wrapper.py

@@ -25,7 +25,7 @@ from contextlib import asynccontextmanager
 from contextvars import ContextVar
 from typing import Any
 
-from aiq.builder.framework_enum import LLMFrameworkEnum
+from nat.builder.framework_enum import LLMFrameworkEnum
 
 logger = logging.getLogger(__name__)
 
@@ -56,7 +56,7 @@ def set_framework_profiler_handler(
             if LLMFrameworkEnum.LANGCHAIN in frameworks and not _library_instrumented["langchain"]:
                 from langchain_core.tracers.context import register_configure_hook
 
-                from aiq.profiler.callbacks.langchain_callback_handler import LangchainProfilerHandler
+                from nat.profiler.callbacks.langchain_callback_handler import LangchainProfilerHandler
 
                 handler = LangchainProfilerHandler()
                 callback_handler_var.set(handler)
@@ -68,14 +68,14 @@ def set_framework_profiler_handler(
                 from llama_index.core import Settings
                 from llama_index.core.callbacks import CallbackManager
 
-                from aiq.profiler.callbacks.llama_index_callback_handler import LlamaIndexProfilerHandler
+                from nat.profiler.callbacks.llama_index_callback_handler import LlamaIndexProfilerHandler
 
                 handler = LlamaIndexProfilerHandler()
                 Settings.callback_manager = CallbackManager([handler])
                 logger.debug("LlamaIndex callback handler registered")
 
             if LLMFrameworkEnum.CREWAI in frameworks and not _library_instrumented["crewai"]:
-                from aiq.plugins.crewai.crewai_callback_handler import \
+                from nat.plugins.crewai.crewai_callback_handler import \
                     CrewAIProfilerHandler  # pylint: disable=ungrouped-imports,line-too-long  # noqa E501
 
                 handler = CrewAIProfilerHandler()
@@ -84,7 +84,7 @@ def set_framework_profiler_handler(
                 logger.debug("CrewAI callback handler registered")
 
             if LLMFrameworkEnum.SEMANTIC_KERNEL in frameworks and not _library_instrumented["semantic_kernel"]:
-                from aiq.profiler.callbacks.semantic_kernel_callback_handler import SemanticKernelProfilerHandler
+                from nat.profiler.callbacks.semantic_kernel_callback_handler import SemanticKernelProfilerHandler
 
                 handler = SemanticKernelProfilerHandler(workflow_llms=workflow_llms)
                 handler.instrument()
@@ -92,7 +92,7 @@ def set_framework_profiler_handler(
                 logger.debug("SemanticKernel callback handler registered")
 
             if LLMFrameworkEnum.AGNO in frameworks and not _library_instrumented["agno"]:
-                from aiq.profiler.callbacks.agno_callback_handler import AgnoProfilerHandler
+                from nat.profiler.callbacks.agno_callback_handler import AgnoProfilerHandler
 
                 handler = AgnoProfilerHandler()
                 handler.instrument()

{aiq → nat}/profiler/decorators/function_tracking.py

@@ -20,11 +20,11 @@ from typing import Any
 
 from pydantic import BaseModel
 
-from aiq.builder.context import AIQContext
-from aiq.builder.intermediate_step_manager import IntermediateStepManager
-from aiq.data_models.intermediate_step import IntermediateStepPayload
-from aiq.data_models.intermediate_step import IntermediateStepType
-from aiq.data_models.intermediate_step import TraceMetadata
+from nat.builder.context import Context
+from nat.builder.intermediate_step_manager import IntermediateStepManager
+from nat.data_models.intermediate_step import IntermediateStepPayload
+from nat.data_models.intermediate_step import IntermediateStepType
+from nat.data_models.intermediate_step import TraceMetadata
 
 
 # --- Helper function to recursively serialize any object into JSON-friendly data ---
@@ -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,
@@ -110,7 +110,7 @@ def track_function(func: Any = None, *, metadata: dict[str, Any] | None = None):
 
         @functools.wraps(func)
         async def async_gen_wrapper(*args, **kwargs):
-            step_manager: IntermediateStepManager = AIQContext.get().intermediate_step_manager
+            step_manager: IntermediateStepManager = Context.get().intermediate_step_manager
             # 1) Serialize input
             serialized_args, serialized_kwargs = _prepare_serialized_args_kwargs(*args, **kwargs)
 
@@ -156,7 +156,7 @@ def track_function(func: Any = None, *, metadata: dict[str, Any] | None = None):
         # ---------------------
         @functools.wraps(func)
         async def async_wrapper(*args, **kwargs):
-            step_manager: IntermediateStepManager = AIQContext.get().intermediate_step_manager
+            step_manager: IntermediateStepManager = Context.get().intermediate_step_manager
             serialized_args, serialized_kwargs = _prepare_serialized_args_kwargs(*args, **kwargs)
             invocation_id = str(uuid.uuid4())
             push_intermediate_step(step_manager,
@@ -189,7 +189,7 @@ def track_function(func: Any = None, *, metadata: dict[str, Any] | None = None):
         # ---------------------
         @functools.wraps(func)
         def sync_gen_wrapper(*args, **kwargs):
-            step_manager: IntermediateStepManager = AIQContext.get().intermediate_step_manager
+            step_manager: IntermediateStepManager = Context.get().intermediate_step_manager
             serialized_args, serialized_kwargs = _prepare_serialized_args_kwargs(*args, **kwargs)
             invocation_id = str(uuid.uuid4())
             push_intermediate_step(step_manager,
@@ -226,7 +226,7 @@ def track_function(func: Any = None, *, metadata: dict[str, Any] | None = None):
 
     @functools.wraps(func)
     def sync_wrapper(*args, **kwargs):
-        step_manager: IntermediateStepManager = AIQContext.get().intermediate_step_manager
+        step_manager: IntermediateStepManager = Context.get().intermediate_step_manager
         serialized_args, serialized_kwargs = _prepare_serialized_args_kwargs(*args, **kwargs)
         invocation_id = str(uuid.uuid4())
         push_intermediate_step(step_manager,

{aiq → nat}/profiler/forecasting/model_trainer.py

@@ -17,11 +17,11 @@
 
 import logging
 
-from aiq.profiler.forecasting.config import DEFAULT_MODEL_TYPE
-from aiq.profiler.forecasting.models import ForecastingBaseModel
-from aiq.profiler.forecasting.models import LinearModel
-from aiq.profiler.forecasting.models import RandomForestModel
-from aiq.profiler.intermediate_property_adapter import IntermediatePropertyAdaptor
+from nat.profiler.forecasting.config import DEFAULT_MODEL_TYPE
+from nat.profiler.forecasting.models import ForecastingBaseModel
+from nat.profiler.forecasting.models import LinearModel
+from nat.profiler.forecasting.models import RandomForestModel
+from nat.profiler.intermediate_property_adapter import IntermediatePropertyAdaptor
 
 logger = logging.getLogger(__name__)
 

{aiq → nat}/profiler/forecasting/models/linear_model.py

@@ -17,8 +17,8 @@ import logging
 
 import numpy as np
 
-from aiq.profiler.forecasting.models.forecasting_base_model import ForecastingBaseModel
-from aiq.profiler.intermediate_property_adapter import IntermediatePropertyAdaptor
+from nat.profiler.forecasting.models.forecasting_base_model import ForecastingBaseModel
+from nat.profiler.intermediate_property_adapter import IntermediatePropertyAdaptor
 
 logger = logging.getLogger(__name__)
 
@@ -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
 

{aiq → nat}/profiler/forecasting/models/random_forest_regressor.py

@@ -17,8 +17,8 @@ import logging
 
 import numpy as np
 
-from aiq.profiler.forecasting.models.forecasting_base_model import ForecastingBaseModel
-from aiq.profiler.intermediate_property_adapter import IntermediatePropertyAdaptor
+from nat.profiler.forecasting.models.forecasting_base_model import ForecastingBaseModel
+from nat.profiler.intermediate_property_adapter import IntermediatePropertyAdaptor
 
 logger = logging.getLogger(__name__)
 
@@ -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
 

{aiq → nat}/profiler/inference_optimization/bottleneck_analysis/nested_stack_analysis.py

@@ -34,12 +34,12 @@ import os
 
 import pandas as pd
 
-from aiq.data_models.intermediate_step import IntermediateStep
-from aiq.profiler.inference_optimization.data_models import CallNode
-from aiq.profiler.inference_optimization.data_models import ConcurrencyDistribution
-from aiq.profiler.inference_optimization.data_models import NestedCallProfilingResult
-from aiq.profiler.inference_optimization.data_models import NodeMetrics
-from aiq.profiler.utils import create_standardized_dataframe
+from nat.data_models.intermediate_step import IntermediateStep
+from nat.profiler.inference_optimization.data_models import CallNode
+from nat.profiler.inference_optimization.data_models import ConcurrencyDistribution
+from nat.profiler.inference_optimization.data_models import NestedCallProfilingResult
+from nat.profiler.inference_optimization.data_models import NodeMetrics
+from nat.profiler.utils import create_standardized_dataframe
 
 logger = logging.getLogger(__name__)
 
@@ -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
 

{aiq → nat}/profiler/inference_optimization/bottleneck_analysis/simple_stack_analysis.py

@@ -28,10 +28,10 @@ then analyze concurrency and produce a summary report.
 import numpy as np
 import pandas as pd
 
-from aiq.data_models.intermediate_step import IntermediateStep
-from aiq.profiler.inference_optimization.data_models import SimpleBottleneckReport
-from aiq.profiler.inference_optimization.data_models import SimpleOperationStats
-from aiq.profiler.utils import create_standardized_dataframe
+from nat.data_models.intermediate_step import IntermediateStep
+from nat.profiler.inference_optimization.data_models import SimpleBottleneckReport
+from nat.profiler.inference_optimization.data_models import SimpleOperationStats
+from nat.profiler.utils import create_standardized_dataframe
 
 
 # ----------------------------------------------------------------------

{aiq → nat}/profiler/inference_optimization/experimental/concurrency_spike_analysis.py

@@ -34,12 +34,12 @@ Changes from previous version:
 import numpy as np
 import pandas as pd
 
-from aiq.data_models.intermediate_step import IntermediateStep
-from aiq.profiler.inference_optimization.data_models import ConcurrencyAnalysisResult
-from aiq.profiler.inference_optimization.data_models import ConcurrencyCallNode
-from aiq.profiler.inference_optimization.data_models import ConcurrencyCorrelationStats
-from aiq.profiler.inference_optimization.data_models import ConcurrencySpikeInfo
-from aiq.profiler.utils import create_standardized_dataframe
+from nat.data_models.intermediate_step import IntermediateStep
+from nat.profiler.inference_optimization.data_models import ConcurrencyAnalysisResult
+from nat.profiler.inference_optimization.data_models import ConcurrencyCallNode
+from nat.profiler.inference_optimization.data_models import ConcurrencyCorrelationStats
+from nat.profiler.inference_optimization.data_models import ConcurrencySpikeInfo
+from nat.profiler.utils import create_standardized_dataframe
 
 # --------------------------------------------------------------------------------
 # 1) Building the Per-Example Call Trees

{aiq → nat}/profiler/inference_optimization/experimental/prefix_span_analysis.py

@@ -32,11 +32,11 @@ import logging
 import numpy as np
 import pandas as pd
 
-from aiq.data_models.intermediate_step import IntermediateStep
-from aiq.profiler.inference_optimization.data_models import FrequentPattern
-from aiq.profiler.inference_optimization.data_models import PrefixCallNode
-from aiq.profiler.inference_optimization.data_models import PrefixSpanSubworkflowResult
-from aiq.profiler.utils import create_standardized_dataframe
+from nat.data_models.intermediate_step import IntermediateStep
+from nat.profiler.inference_optimization.data_models import FrequentPattern
+from nat.profiler.inference_optimization.data_models import PrefixCallNode
+from nat.profiler.inference_optimization.data_models import PrefixSpanSubworkflowResult
+from nat.profiler.utils import create_standardized_dataframe
 
 logger = logging.getLogger(__name__)
 
@@ -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
 

{aiq → nat}/profiler/inference_optimization/llm_metrics.py

@@ -16,8 +16,8 @@
 import numpy as np
 import pandas as pd
 
-from aiq.data_models.intermediate_step import IntermediateStep
-from aiq.profiler.utils import create_standardized_dataframe
+from nat.data_models.intermediate_step import IntermediateStep
+from nat.profiler.utils import create_standardized_dataframe
 
 
 class LLMMetrics:

{aiq → nat}/profiler/inference_optimization/prompt_caching.py

@@ -13,11 +13,11 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from aiq.data_models.intermediate_step import IntermediateStep
-from aiq.profiler.inference_optimization.data_models import CommonPrefixesOutput
-from aiq.profiler.inference_optimization.data_models import FrameworkLLMPrefixData
-from aiq.profiler.inference_optimization.data_models import PrefixInfo
-from aiq.profiler.utils import create_standardized_dataframe
+from nat.data_models.intermediate_step import IntermediateStep
+from nat.profiler.inference_optimization.data_models import CommonPrefixesOutput
+from nat.profiler.inference_optimization.data_models import FrameworkLLMPrefixData
+from nat.profiler.inference_optimization.data_models import PrefixInfo
+from nat.profiler.utils import create_standardized_dataframe
 
 
 # -----------------------------------------------------------

{aiq → nat}/profiler/inference_optimization/token_uniqueness.py

@@ -17,10 +17,10 @@ import re
 
 import numpy as np
 
-from aiq.data_models.intermediate_step import IntermediateStep
-from aiq.profiler.inference_optimization.data_models import LLMUniquenessMetrics
-from aiq.profiler.inference_optimization.data_models import LLMUniquenessMetricsByLLM
-from aiq.profiler.utils import create_standardized_dataframe
+from nat.data_models.intermediate_step import IntermediateStep
+from nat.profiler.inference_optimization.data_models import LLMUniquenessMetrics
+from nat.profiler.inference_optimization.data_models import LLMUniquenessMetricsByLLM
+from nat.profiler.utils import create_standardized_dataframe
 
 
 # ----------------------------------------------------------------

{aiq → nat}/profiler/inference_optimization/workflow_runtimes.py

@@ -15,9 +15,9 @@
 
 import numpy as np
 
-from aiq.data_models.intermediate_step import IntermediateStep
-from aiq.profiler.inference_optimization.data_models import WorkflowRuntimeMetrics
-from aiq.profiler.utils import create_standardized_dataframe
+from nat.data_models.intermediate_step import IntermediateStep
+from nat.profiler.inference_optimization.data_models import WorkflowRuntimeMetrics
+from nat.profiler.utils import create_standardized_dataframe
 
 
 def compute_workflow_runtime_metrics(all_steps: list[list[IntermediateStep]]) -> WorkflowRuntimeMetrics:

{aiq → nat}/profiler/intermediate_property_adapter.py

@@ -13,9 +13,9 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from aiq.data_models.intermediate_step import IntermediateStep
-from aiq.data_models.intermediate_step import IntermediateStepType
-from aiq.data_models.intermediate_step import TokenUsageBaseModel
+from nat.data_models.intermediate_step import IntermediateStep
+from nat.data_models.intermediate_step import IntermediateStepType
+from nat.data_models.intermediate_step import TokenUsageBaseModel
 
 
 class IntermediatePropertyAdaptor(IntermediateStep):

{aiq → nat}/profiler/profile_runner.py

@@ -23,13 +23,13 @@ from typing import Any
 
 from pydantic import BaseModel
 
-from aiq.data_models.evaluate import ProfilerConfig
-from aiq.data_models.intermediate_step import IntermediateStep
-from aiq.profiler.data_models import ProfilerResults
-from aiq.profiler.forecasting.model_trainer import ModelTrainer
-from aiq.profiler.inference_metrics_model import InferenceMetricsModel
-from aiq.profiler.utils import create_standardized_dataframe
-from aiq.utils.type_converter import TypeConverter
+from nat.data_models.evaluate import ProfilerConfig
+from nat.data_models.intermediate_step import IntermediateStep
+from nat.profiler.data_models import ProfilerResults
+from nat.profiler.forecasting.model_trainer import ModelTrainer
+from nat.profiler.inference_metrics_model import InferenceMetricsModel
+from nat.profiler.utils import create_standardized_dataframe
+from nat.utils.type_converter import TypeConverter
 
 logger = logging.getLogger(__name__)
 
@@ -49,7 +49,7 @@ class InferenceOptimizationHolder(BaseModel):
 
 class ProfilerRunner:
     """
-    A utility to run a series of prompts through an AIQ Toolkit workflow for profiling:
+    A utility to run a series of prompts through a NAT workflow for profiling:
 
     - can load prompts from a file
     - or generate them via an LLM
@@ -88,19 +88,19 @@ class ProfilerRunner:
         writes out combined requests JSON, then computes and saves additional metrics,
         and optionally fits a forecasting model.
         """
-        from aiq.profiler.inference_optimization.bottleneck_analysis.nested_stack_analysis import \
+        from nat.profiler.inference_optimization.bottleneck_analysis.nested_stack_analysis import \
             multi_example_call_profiling
-        from aiq.profiler.inference_optimization.bottleneck_analysis.simple_stack_analysis import \
+        from nat.profiler.inference_optimization.bottleneck_analysis.simple_stack_analysis import \
             profile_workflow_bottlenecks
-        from aiq.profiler.inference_optimization.experimental.concurrency_spike_analysis import \
+        from nat.profiler.inference_optimization.experimental.concurrency_spike_analysis import \
             concurrency_spike_analysis
-        from aiq.profiler.inference_optimization.experimental.prefix_span_analysis import \
+        from nat.profiler.inference_optimization.experimental.prefix_span_analysis import \
             prefixspan_subworkflow_with_text
-        from aiq.profiler.inference_optimization.llm_metrics import LLMMetrics
-        from aiq.profiler.inference_optimization.prompt_caching import get_common_prefixes
-        from aiq.profiler.inference_optimization.token_uniqueness import compute_inter_query_token_uniqueness_by_llm
-        from aiq.profiler.inference_optimization.workflow_runtimes import compute_workflow_runtime_metrics
-        from aiq.profiler.intermediate_property_adapter import IntermediatePropertyAdaptor
+        from nat.profiler.inference_optimization.llm_metrics import LLMMetrics
+        from nat.profiler.inference_optimization.prompt_caching import get_common_prefixes
+        from nat.profiler.inference_optimization.token_uniqueness import compute_inter_query_token_uniqueness_by_llm
+        from nat.profiler.inference_optimization.workflow_runtimes import compute_workflow_runtime_metrics
+        from nat.profiler.intermediate_property_adapter import IntermediatePropertyAdaptor
 
         # Convert the incoming DataFrame to a list of dicts and store
         all_steps = [[IntermediatePropertyAdaptor.from_intermediate_step(step) for step in steps]

{aiq → nat}/profiler/utils.py

@@ -21,10 +21,10 @@ from typing import Any
 
 import pandas as pd
 
-from aiq.builder.framework_enum import LLMFrameworkEnum
-from aiq.cli.type_registry import RegisteredFunctionInfo
-from aiq.data_models.intermediate_step import IntermediateStep
-from aiq.profiler.data_frame_row import DataFrameRow
+from nat.builder.framework_enum import LLMFrameworkEnum
+from nat.cli.type_registry import RegisteredFunctionInfo
+from nat.data_models.intermediate_step import IntermediateStep
+from nat.profiler.data_frame_row import DataFrameRow
 
 # A simple set of regex patterns to scan for direct references to LLMFrameworkEnum
 _FRAMEWORK_REGEX_MAP = {t: fr'\b{t._name_}\b' for t in LLMFrameworkEnum}