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

{aiq → nat}/eval/tunable_rag_evaluator/evaluate.py

@@ -25,9 +25,9 @@ from langchain_core.language_models import BaseChatModel
 from langchain_core.runnables import RunnableLambda
 from tqdm import tqdm
 
-from aiq.eval.evaluator.base_evaluator import BaseEvaluator
-from aiq.eval.evaluator.evaluator_model import EvalInputItem
-from aiq.eval.evaluator.evaluator_model import EvalOutputItem
+from nat.eval.evaluator.base_evaluator import BaseEvaluator
+from nat.eval.evaluator.evaluator_model import EvalInputItem
+from nat.eval.evaluator.evaluator_model import EvalOutputItem
 
 logger = logging.getLogger(__name__)
 

{aiq → nat}/eval/tunable_rag_evaluator/register.py

@@ -15,12 +15,12 @@
 
 from pydantic import Field
 
-from aiq.builder.builder import EvalBuilder
-from aiq.builder.evaluator import EvaluatorInfo
-from aiq.builder.framework_enum import LLMFrameworkEnum
-from aiq.cli.register_workflow import register_evaluator
-from aiq.data_models.component_ref import LLMRef
-from aiq.data_models.evaluator import EvaluatorBaseConfig
+from nat.builder.builder import EvalBuilder
+from nat.builder.evaluator import EvaluatorInfo
+from nat.builder.framework_enum import LLMFrameworkEnum
+from nat.cli.register_workflow import register_evaluator
+from nat.data_models.component_ref import LLMRef
+from nat.data_models.evaluator import EvaluatorBaseConfig
 
 
 class TunableRagEvaluatorConfig(EvaluatorBaseConfig, name="tunable_rag_evaluator"):

{aiq → nat}/eval/utils/output_uploader.py

@@ -24,7 +24,7 @@ import aioboto3
 from botocore.exceptions import NoCredentialsError
 from tqdm import tqdm
 
-from aiq.data_models.evaluate import EvalOutputConfig
+from nat.data_models.evaluate import EvalOutputConfig
 
 logger = logging.getLogger(__name__)
 

{aiq → nat}/eval/utils/weave_eval.py

@@ -17,12 +17,12 @@ import asyncio
 import logging
 from typing import Any
 
-from aiq.eval.evaluator.evaluator_model import EvalInput
-from aiq.eval.evaluator.evaluator_model import EvalInputItem
-from aiq.eval.evaluator.evaluator_model import EvalOutput
-from aiq.eval.usage_stats import UsageStats
-from aiq.eval.usage_stats import UsageStatsItem
-from aiq.profiler.data_models import ProfilerResults
+from nat.eval.evaluator.evaluator_model import EvalInput
+from nat.eval.evaluator.evaluator_model import EvalInputItem
+from nat.eval.evaluator.evaluator_model import EvalOutput
+from nat.eval.usage_stats import UsageStats
+from nat.eval.usage_stats import UsageStatsItem
+from nat.profiler.data_models import ProfilerResults
 
 logger = logging.getLogger(__name__)
 

{aiq → nat}/experimental/decorators/experimental_warning_decorator.py

@@ -53,7 +53,7 @@ def issue_experimental_warning(function_name: str,
         _warning_issued.add(function_name)
 
 
-def aiq_experimental(func: Any = None, *, feature_name: str | None = None, metadata: dict[str, Any] | None = None):
+def experimental(func: Any = None, *, feature_name: str | None = None, metadata: dict[str, Any] | None = None):
     """
     Decorator that can wrap any type of function (sync, async, generator,
     async generator) and logs a warning that the function is experimental.
@@ -72,7 +72,7 @@ def aiq_experimental(func: Any = None, *, feature_name: str | None = None, metad
     if func is None:
 
         def decorator_wrapper(actual_func):
-            return aiq_experimental(actual_func, feature_name=feature_name, metadata=metadata)
+            return experimental(actual_func, feature_name=feature_name, metadata=metadata)
 
         return decorator_wrapper
 
@@ -128,3 +128,7 @@ def aiq_experimental(func: Any = None, *, feature_name: str | None = None, metad
         return result
 
     return sync_wrapper
+
+
+# Compatibility aliases with previous releases
+aiq_experimental = experimental

{aiq → nat}/experimental/test_time_compute/editing/iterative_plan_refinement_editor.py

@@ -17,16 +17,16 @@ import asyncio
 import logging
 import re
 
-from aiq.builder.builder import Builder
-from aiq.builder.framework_enum import LLMFrameworkEnum
-from aiq.cli.register_workflow import register_ttc_strategy
-from aiq.data_models.ttc_strategy import TTCStrategyBaseConfig
-from aiq.experimental.test_time_compute.models.editor_config import IterativePlanRefinementConfig
-from aiq.experimental.test_time_compute.models.stage_enums import PipelineTypeEnum
-from aiq.experimental.test_time_compute.models.stage_enums import StageTypeEnum
-from aiq.experimental.test_time_compute.models.strategy_base import StrategyBase
-from aiq.experimental.test_time_compute.models.ttc_item import TTCItem
-from aiq.utils.io.model_processing import remove_r1_think_tags
+from nat.builder.builder import Builder
+from nat.builder.framework_enum import LLMFrameworkEnum
+from nat.cli.register_workflow import register_ttc_strategy
+from nat.data_models.ttc_strategy import TTCStrategyBaseConfig
+from nat.experimental.test_time_compute.models.editor_config import IterativePlanRefinementConfig
+from nat.experimental.test_time_compute.models.stage_enums import PipelineTypeEnum
+from nat.experimental.test_time_compute.models.stage_enums import StageTypeEnum
+from nat.experimental.test_time_compute.models.strategy_base import StrategyBase
+from nat.experimental.test_time_compute.models.ttc_item import TTCItem
+from nat.utils.io.model_processing import remove_r1_think_tags
 
 logger = logging.getLogger(__name__)
 

{aiq → nat}/experimental/test_time_compute/editing/llm_as_a_judge_editor.py

@@ -17,16 +17,16 @@ import asyncio
 import logging
 import re
 
-from aiq.builder.builder import Builder
-from aiq.builder.framework_enum import LLMFrameworkEnum
-from aiq.cli.register_workflow import register_ttc_strategy
-from aiq.data_models.ttc_strategy import TTCStrategyBaseConfig
-from aiq.experimental.test_time_compute.models.editor_config import LLMAsAJudgeEditorConfig
-from aiq.experimental.test_time_compute.models.stage_enums import PipelineTypeEnum
-from aiq.experimental.test_time_compute.models.stage_enums import StageTypeEnum
-from aiq.experimental.test_time_compute.models.strategy_base import StrategyBase
-from aiq.experimental.test_time_compute.models.ttc_item import TTCItem
-from aiq.utils.io.model_processing import remove_r1_think_tags
+from nat.builder.builder import Builder
+from nat.builder.framework_enum import LLMFrameworkEnum
+from nat.cli.register_workflow import register_ttc_strategy
+from nat.data_models.ttc_strategy import TTCStrategyBaseConfig
+from nat.experimental.test_time_compute.models.editor_config import LLMAsAJudgeEditorConfig
+from nat.experimental.test_time_compute.models.stage_enums import PipelineTypeEnum
+from nat.experimental.test_time_compute.models.stage_enums import StageTypeEnum
+from nat.experimental.test_time_compute.models.strategy_base import StrategyBase
+from nat.experimental.test_time_compute.models.ttc_item import TTCItem
+from nat.utils.io.model_processing import remove_r1_think_tags
 
 logger = logging.getLogger(__name__)
 

{aiq → nat}/experimental/test_time_compute/editing/motivation_aware_summarization.py

@@ -15,15 +15,15 @@
 
 import logging
 
-from aiq.builder.builder import Builder
-from aiq.builder.framework_enum import LLMFrameworkEnum
-from aiq.cli.register_workflow import register_ttc_strategy
-from aiq.experimental.test_time_compute.models.editor_config import MotivationAwareSummarizationConfig
-from aiq.experimental.test_time_compute.models.stage_enums import PipelineTypeEnum
-from aiq.experimental.test_time_compute.models.stage_enums import StageTypeEnum
-from aiq.experimental.test_time_compute.models.strategy_base import StrategyBase
-from aiq.experimental.test_time_compute.models.ttc_item import TTCItem
-from aiq.utils.io.model_processing import remove_r1_think_tags
+from nat.builder.builder import Builder
+from nat.builder.framework_enum import LLMFrameworkEnum
+from nat.cli.register_workflow import register_ttc_strategy
+from nat.experimental.test_time_compute.models.editor_config import MotivationAwareSummarizationConfig
+from nat.experimental.test_time_compute.models.stage_enums import PipelineTypeEnum
+from nat.experimental.test_time_compute.models.stage_enums import StageTypeEnum
+from nat.experimental.test_time_compute.models.strategy_base import StrategyBase
+from nat.experimental.test_time_compute.models.ttc_item import TTCItem
+from nat.utils.io.model_processing import remove_r1_think_tags
 
 logger = logging.getLogger(__name__)
 
@@ -65,7 +65,7 @@ class MotivationAwareSummarization(StrategyBase):
             from langchain_core.prompts import PromptTemplate
         except ImportError:
             raise ImportError("langchain-core is required for MultiQueryRetrievalSearch. "
-                              "Install aiqtoolkit-langchain or similar.")
+                              "Install nvidia-nat-langchain or similar.")
 
         new_ttc_items: list[TTCItem] = []
 

{aiq → nat}/experimental/test_time_compute/functions/execute_score_select_function.py

@@ -17,16 +17,16 @@ import logging
 
 from pydantic import Field
 
-from aiq.builder.builder import Builder
-from aiq.builder.function import Function
-from aiq.builder.function_info import FunctionInfo
-from aiq.cli.register_workflow import register_function
-from aiq.data_models.component_ref import FunctionRef
-from aiq.data_models.component_ref import TTCStrategyRef
-from aiq.data_models.function import FunctionBaseConfig
-from aiq.experimental.test_time_compute.models.stage_enums import PipelineTypeEnum
-from aiq.experimental.test_time_compute.models.stage_enums import StageTypeEnum
-from aiq.experimental.test_time_compute.models.ttc_item import TTCItem
+from nat.builder.builder import Builder
+from nat.builder.function import Function
+from nat.builder.function_info import FunctionInfo
+from nat.cli.register_workflow import register_function
+from nat.data_models.component_ref import FunctionRef
+from nat.data_models.component_ref import TTCStrategyRef
+from nat.data_models.function import FunctionBaseConfig
+from nat.experimental.test_time_compute.models.stage_enums import PipelineTypeEnum
+from nat.experimental.test_time_compute.models.stage_enums import StageTypeEnum
+from nat.experimental.test_time_compute.models.ttc_item import TTCItem
 
 logger = logging.getLogger(__name__)
 

{aiq → nat}/experimental/test_time_compute/functions/plan_select_execute_function.py

@@ -18,24 +18,24 @@ from collections.abc import AsyncGenerator
 
 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.api_server import AIQChatRequest
-from aiq.data_models.component_ref import FunctionRef
-from aiq.data_models.component_ref import TTCStrategyRef
-from aiq.data_models.function import FunctionBaseConfig
-from aiq.experimental.test_time_compute.models.stage_enums import PipelineTypeEnum
-from aiq.experimental.test_time_compute.models.stage_enums import StageTypeEnum
-from aiq.experimental.test_time_compute.models.ttc_item import TTCItem
+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.api_server import ChatRequest
+from nat.data_models.component_ref import FunctionRef
+from nat.data_models.component_ref import TTCStrategyRef
+from nat.data_models.function import FunctionBaseConfig
+from nat.experimental.test_time_compute.models.stage_enums import PipelineTypeEnum
+from nat.experimental.test_time_compute.models.stage_enums import StageTypeEnum
+from nat.experimental.test_time_compute.models.ttc_item import TTCItem
 
 logger = logging.getLogger(__name__)
 
 
 class PlanSelectExecuteFunctionConfig(FunctionBaseConfig, name="plan_select_execute_function"):
     """
-    Defines an aiqtoolkit function that performs reasoning on the input data.
+    Defines a NAT function that performs reasoning on the input data.
     Output is passed to the next function in the workflow.
 
     Designed to be used with an InterceptingFunction.
@@ -83,7 +83,7 @@ async def plan_select_execute_function(config: PlanSelectExecuteFunctionConfig,
         from langchain_core.prompts import PromptTemplate
     except ImportError:
         raise ImportError("langchain-core is not installed. Please install it to use SingleShotMultiPlanPlanner.\n"
-                          "This error can be resolved by installing aiqtoolkit-langchain.")
+                          "This error can be resolved by installing nvidia-nat-langchain.")
 
     # Get the augmented function's description
     augmented_function = builder.get_function(config.augmented_fn)
@@ -154,12 +154,12 @@ async def plan_select_execute_function(config: PlanSelectExecuteFunctionConfig,
     if augmented_function.has_streaming_output:
 
         async def streaming_inner(
-                input_message: AIQChatRequest) -> AsyncGenerator[augmented_function.streaming_output_type]:
+                input_message: ChatRequest) -> AsyncGenerator[augmented_function.streaming_output_type]:
             """
             Perform reasoning on the input text.
 
             Args:
-                input_message (AIQChatRequest): The input text to reason on.
+                input_message (ChatRequest): The input text to reason on.
             """
 
             input_text = "".join([str(message.model_dump()) + "\n" for message in input_message.messages])
@@ -187,12 +187,12 @@ async def plan_select_execute_function(config: PlanSelectExecuteFunctionConfig,
 
     if augmented_function.has_single_output:
 
-        async def single_inner(input_message: AIQChatRequest) -> augmented_function.single_output_type:
+        async def single_inner(input_message: ChatRequest) -> augmented_function.single_output_type:
             """
             Perform reasoning on the input text.
 
             Args:
-                input_message (AIQChatRequest): The input text to reason on.
+                input_message (ChatRequest): The input text to reason on.
             """
 
             input_text = "".join([str(message.model_dump()) + "\n" for message in input_message.messages])

aiq/experimental/test_time_compute/functions/its_tool_orchestration_function.py → nat/experimental/test_time_compute/functions/ttc_tool_orchestration_function.py

@@ -18,18 +18,18 @@ 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 FunctionRef
-from aiq.data_models.component_ref import TTCStrategyRef
-from aiq.data_models.function import FunctionBaseConfig
-from aiq.experimental.test_time_compute.models.stage_enums import PipelineTypeEnum
-from aiq.experimental.test_time_compute.models.stage_enums import StageTypeEnum
-from aiq.experimental.test_time_compute.models.tool_use_config import ToolUseInputSchema
-from aiq.experimental.test_time_compute.models.tool_use_config import ToolUselist
-from aiq.experimental.test_time_compute.models.ttc_item import TTCItem
+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 FunctionRef
+from nat.data_models.component_ref import TTCStrategyRef
+from nat.data_models.function import FunctionBaseConfig
+from nat.experimental.test_time_compute.models.stage_enums import PipelineTypeEnum
+from nat.experimental.test_time_compute.models.stage_enums import StageTypeEnum
+from nat.experimental.test_time_compute.models.tool_use_config import ToolUseInputSchema
+from nat.experimental.test_time_compute.models.tool_use_config import ToolUselist
+from nat.experimental.test_time_compute.models.ttc_item import TTCItem
 
 logger = logging.getLogger(__name__)
 

aiq/experimental/test_time_compute/functions/its_tool_wrapper_function.py → nat/experimental/test_time_compute/functions/ttc_tool_wrapper_function.py

@@ -18,15 +18,15 @@ import logging
 from pydantic import BaseModel
 from pydantic import Field
 
-from aiq.builder.builder import Builder
-from aiq.builder.framework_enum import LLMFrameworkEnum
-from aiq.builder.function import Function
-from aiq.builder.function_info import FunctionInfo
-from aiq.cli.register_workflow import register_function
-from aiq.data_models.component_ref import FunctionRef
-from aiq.data_models.component_ref import LLMRef
-from aiq.data_models.function import FunctionBaseConfig
-from aiq.utils.string_utils import convert_to_str
+from nat.builder.builder import Builder
+from nat.builder.framework_enum import LLMFrameworkEnum
+from nat.builder.function import Function
+from nat.builder.function_info import FunctionInfo
+from nat.cli.register_workflow import register_function
+from nat.data_models.component_ref import FunctionRef
+from nat.data_models.component_ref import LLMRef
+from nat.data_models.function import FunctionBaseConfig
+from nat.utils.string_utils import convert_to_str
 
 logger = logging.getLogger(__name__)
 
@@ -78,7 +78,7 @@ async def register_ttc_tool_wrapper_function(
         from langchain_core.prompts import PromptTemplate
     except ImportError:
         raise ImportError("langchain-core is not installed. Please install it to use SingleShotMultiPlanPlanner.\n"
-                          "This error can be resolved by installing aiqtoolkit-langchain.")
+                          "This error can be resolved by installing nvidia-nat-langchain.")
 
     augmented_function: Function = builder.get_function(config.augmented_fn)
     input_llm: BaseChatModel = await builder.get_llm(config.input_llm, wrapper_type=LLMFrameworkEnum.LANGCHAIN)

{aiq → nat}/experimental/test_time_compute/models/editor_config.py

@@ -18,8 +18,8 @@ import typing
 from pydantic import Field
 from pydantic import model_validator
 
-from aiq.data_models.component_ref import LLMRef
-from aiq.data_models.ttc_strategy import TTCStrategyBaseConfig
+from nat.data_models.component_ref import LLMRef
+from nat.data_models.ttc_strategy import TTCStrategyBaseConfig
 
 
 class LLMAsAJudgeEditorConfig(TTCStrategyBaseConfig, name="llm_as_a_judge_editor"):

{aiq → nat}/experimental/test_time_compute/models/scoring_config.py

@@ -18,8 +18,8 @@ import typing
 from pydantic import Field
 from pydantic import model_validator
 
-from aiq.data_models.component_ref import LLMRef
-from aiq.data_models.ttc_strategy import TTCStrategyBaseConfig
+from nat.data_models.component_ref import LLMRef
+from nat.data_models.ttc_strategy import TTCStrategyBaseConfig
 
 
 class LLMBasedPlanScoringConfig(TTCStrategyBaseConfig, name="llm_based_plan_scoring"):

{aiq → nat}/experimental/test_time_compute/models/search_config.py

@@ -18,8 +18,8 @@ import typing
 from pydantic import Field
 from pydantic import model_validator
 
-from aiq.data_models.component_ref import LLMRef
-from aiq.data_models.ttc_strategy import TTCStrategyBaseConfig
+from nat.data_models.component_ref import LLMRef
+from nat.data_models.ttc_strategy import TTCStrategyBaseConfig
 
 
 class SingleShotMultiPlanConfig(TTCStrategyBaseConfig, name="single_shot_multi_plan"):

{aiq → nat}/experimental/test_time_compute/models/selection_config.py

@@ -18,8 +18,8 @@ import typing
 from pydantic import Field
 from pydantic import model_validator
 
-from aiq.data_models.component_ref import LLMRef
-from aiq.data_models.ttc_strategy import TTCStrategyBaseConfig
+from nat.data_models.component_ref import LLMRef
+from nat.data_models.ttc_strategy import TTCStrategyBaseConfig
 
 
 class LLMBasedPlanSelectionConfig(TTCStrategyBaseConfig, name="llm_based_plan_selection"):

{aiq → nat}/experimental/test_time_compute/models/strategy_base.py

@@ -16,10 +16,10 @@
 from abc import ABC
 from abc import abstractmethod
 
-from aiq.builder.builder import Builder
-from aiq.experimental.test_time_compute.models.ttc_item import TTCItem
-from aiq.experimental.test_time_compute.models.stage_enums import StageTypeEnum, PipelineTypeEnum
-from aiq.data_models.ttc_strategy import TTCStrategyBaseConfig
+from nat.builder.builder import Builder
+from nat.experimental.test_time_compute.models.ttc_item import TTCItem
+from nat.experimental.test_time_compute.models.stage_enums import StageTypeEnum, PipelineTypeEnum
+from nat.data_models.ttc_strategy import TTCStrategyBaseConfig
 
 
 class StrategyBase(ABC):

{aiq → nat}/experimental/test_time_compute/register.py

@@ -20,9 +20,9 @@ from .editing import iterative_plan_refinement_editor
 from .editing import llm_as_a_judge_editor
 from .editing import motivation_aware_summarization
 from .functions import execute_score_select_function
-from .functions import its_tool_orchestration_function
-from .functions import its_tool_wrapper_function
 from .functions import plan_select_execute_function
+from .functions import ttc_tool_orchestration_function
+from .functions import ttc_tool_wrapper_function
 from .scoring import llm_based_agent_scorer
 from .scoring import llm_based_plan_scorer
 from .scoring import motivation_aware_scorer

{aiq → nat}/experimental/test_time_compute/scoring/llm_based_agent_scorer.py

@@ -17,16 +17,16 @@ import asyncio
 import logging
 import re
 
-from aiq.builder.builder import Builder
-from aiq.builder.framework_enum import LLMFrameworkEnum
-from aiq.cli.register_workflow import register_ttc_strategy
-from aiq.data_models.ttc_strategy import TTCStrategyBaseConfig
-from aiq.experimental.test_time_compute.models.scoring_config import LLMBasedAgentScoringConfig
-from aiq.experimental.test_time_compute.models.stage_enums import PipelineTypeEnum
-from aiq.experimental.test_time_compute.models.stage_enums import StageTypeEnum
-from aiq.experimental.test_time_compute.models.strategy_base import StrategyBase
-from aiq.experimental.test_time_compute.models.ttc_item import TTCItem
-from aiq.utils.io.model_processing import remove_r1_think_tags
+from nat.builder.builder import Builder
+from nat.builder.framework_enum import LLMFrameworkEnum
+from nat.cli.register_workflow import register_ttc_strategy
+from nat.data_models.ttc_strategy import TTCStrategyBaseConfig
+from nat.experimental.test_time_compute.models.scoring_config import LLMBasedAgentScoringConfig
+from nat.experimental.test_time_compute.models.stage_enums import PipelineTypeEnum
+from nat.experimental.test_time_compute.models.stage_enums import StageTypeEnum
+from nat.experimental.test_time_compute.models.strategy_base import StrategyBase
+from nat.experimental.test_time_compute.models.ttc_item import TTCItem
+from nat.utils.io.model_processing import remove_r1_think_tags
 
 logger = logging.getLogger(__name__)
 
@@ -67,7 +67,7 @@ class LLMBasedAgentScorer(StrategyBase):
             from langchain_core.prompts import PromptTemplate
         except ImportError:
             raise ImportError("langchain-core is not installed. Please install it to use SingleShotMultiPlanPlanner.\n"
-                              "This error can be resolved by installing aiqtoolkit-langchain.")
+                              "This error can be resolved by installing nvidia-nat-langchain.")
 
         if not isinstance(self.llm_bound, BaseChatModel):
             raise ValueError("The `scoring_llm` must be an instance of `BaseChatModel`.")

{aiq → nat}/experimental/test_time_compute/scoring/llm_based_plan_scorer.py

@@ -17,16 +17,16 @@ import asyncio
 import logging
 import re
 
-from aiq.builder.builder import Builder
-from aiq.builder.framework_enum import LLMFrameworkEnum
-from aiq.cli.register_workflow import register_ttc_strategy
-from aiq.data_models.ttc_strategy import TTCStrategyBaseConfig
-from aiq.experimental.test_time_compute.models.scoring_config import LLMBasedPlanScoringConfig
-from aiq.experimental.test_time_compute.models.stage_enums import PipelineTypeEnum
-from aiq.experimental.test_time_compute.models.stage_enums import StageTypeEnum
-from aiq.experimental.test_time_compute.models.strategy_base import StrategyBase
-from aiq.experimental.test_time_compute.models.ttc_item import TTCItem
-from aiq.utils.io.model_processing import remove_r1_think_tags
+from nat.builder.builder import Builder
+from nat.builder.framework_enum import LLMFrameworkEnum
+from nat.cli.register_workflow import register_ttc_strategy
+from nat.data_models.ttc_strategy import TTCStrategyBaseConfig
+from nat.experimental.test_time_compute.models.scoring_config import LLMBasedPlanScoringConfig
+from nat.experimental.test_time_compute.models.stage_enums import PipelineTypeEnum
+from nat.experimental.test_time_compute.models.stage_enums import StageTypeEnum
+from nat.experimental.test_time_compute.models.strategy_base import StrategyBase
+from nat.experimental.test_time_compute.models.ttc_item import TTCItem
+from nat.utils.io.model_processing import remove_r1_think_tags
 
 logger = logging.getLogger(__name__)
 
@@ -67,7 +67,7 @@ class LLMBasedPlanScorer(StrategyBase):
             from langchain_core.prompts import PromptTemplate
         except ImportError:
             raise ImportError("langchain-core is not installed. Please install it to use SingleShotMultiPlanPlanner.\n"
-                              "This error can be resolved by installing aiqtoolkit-langchain.")
+                              "This error can be resolved by installing nvidia-nat-langchain.")
 
         if not isinstance(self.llm_bound, BaseChatModel):
             raise ValueError("The `scoring_llm` must be an instance of `BaseChatModel`.")

{aiq → nat}/experimental/test_time_compute/scoring/motivation_aware_scorer.py

@@ -17,16 +17,16 @@ import asyncio
 import logging
 import re
 
-from aiq.builder.builder import Builder
-from aiq.builder.framework_enum import LLMFrameworkEnum
-from aiq.cli.register_workflow import register_ttc_strategy
-from aiq.data_models.ttc_strategy import TTCStrategyBaseConfig
-from aiq.experimental.test_time_compute.models.scoring_config import MotivationAwareScoringConfig
-from aiq.experimental.test_time_compute.models.stage_enums import PipelineTypeEnum
-from aiq.experimental.test_time_compute.models.stage_enums import StageTypeEnum
-from aiq.experimental.test_time_compute.models.strategy_base import StrategyBase
-from aiq.experimental.test_time_compute.models.ttc_item import TTCItem
-from aiq.utils.io.model_processing import remove_r1_think_tags
+from nat.builder.builder import Builder
+from nat.builder.framework_enum import LLMFrameworkEnum
+from nat.cli.register_workflow import register_ttc_strategy
+from nat.data_models.ttc_strategy import TTCStrategyBaseConfig
+from nat.experimental.test_time_compute.models.scoring_config import MotivationAwareScoringConfig
+from nat.experimental.test_time_compute.models.stage_enums import PipelineTypeEnum
+from nat.experimental.test_time_compute.models.stage_enums import StageTypeEnum
+from nat.experimental.test_time_compute.models.strategy_base import StrategyBase
+from nat.experimental.test_time_compute.models.ttc_item import TTCItem
+from nat.utils.io.model_processing import remove_r1_think_tags
 
 logger = logging.getLogger(__name__)
 

{aiq → nat}/experimental/test_time_compute/search/multi_llm_planner.py

@@ -17,15 +17,15 @@ import asyncio
 import logging
 import re
 
-from aiq.builder.builder import Builder
-from aiq.builder.framework_enum import LLMFrameworkEnum
-from aiq.cli.register_workflow import register_ttc_strategy
-from aiq.experimental.test_time_compute.models.search_config import MultiLLMPlanConfig
-from aiq.experimental.test_time_compute.models.stage_enums import PipelineTypeEnum
-from aiq.experimental.test_time_compute.models.stage_enums import StageTypeEnum
-from aiq.experimental.test_time_compute.models.strategy_base import StrategyBase
-from aiq.experimental.test_time_compute.models.ttc_item import TTCItem
-from aiq.utils.io.model_processing import remove_r1_think_tags
+from nat.builder.builder import Builder
+from nat.builder.framework_enum import LLMFrameworkEnum
+from nat.cli.register_workflow import register_ttc_strategy
+from nat.experimental.test_time_compute.models.search_config import MultiLLMPlanConfig
+from nat.experimental.test_time_compute.models.stage_enums import PipelineTypeEnum
+from nat.experimental.test_time_compute.models.stage_enums import StageTypeEnum
+from nat.experimental.test_time_compute.models.strategy_base import StrategyBase
+from nat.experimental.test_time_compute.models.ttc_item import TTCItem
+from nat.utils.io.model_processing import remove_r1_think_tags
 
 logger = logging.getLogger(__name__)
 
@@ -96,7 +96,7 @@ class MultiLLMPlanner(StrategyBase):
             from langchain_core.prompts import PromptTemplate
         except ImportError:
             raise ImportError("langchain-core is not installed. Please install it to use MultiLLMPlanner.\n"
-                              "This error can be resolve by installing aiqtoolkit-langchain.")
+                              "This error can be resolve by installing nvidia-nat-langchain.")
 
         # Create a single PromptTemplate
         planning_template = PromptTemplate(template=self.config.planning_template,

{aiq → nat}/experimental/test_time_compute/search/multi_query_retrieval_search.py

@@ -15,15 +15,15 @@
 
 import logging
 
-from aiq.builder.builder import Builder
-from aiq.builder.framework_enum import LLMFrameworkEnum
-from aiq.cli.register_workflow import register_ttc_strategy
-from aiq.experimental.test_time_compute.models.search_config import MultiQueryRetrievalSearchConfig
-from aiq.experimental.test_time_compute.models.stage_enums import PipelineTypeEnum
-from aiq.experimental.test_time_compute.models.stage_enums import StageTypeEnum
-from aiq.experimental.test_time_compute.models.strategy_base import StrategyBase
-from aiq.experimental.test_time_compute.models.ttc_item import TTCItem
-from aiq.utils.io.model_processing import remove_r1_think_tags
+from nat.builder.builder import Builder
+from nat.builder.framework_enum import LLMFrameworkEnum
+from nat.cli.register_workflow import register_ttc_strategy
+from nat.experimental.test_time_compute.models.search_config import MultiQueryRetrievalSearchConfig
+from nat.experimental.test_time_compute.models.stage_enums import PipelineTypeEnum
+from nat.experimental.test_time_compute.models.stage_enums import StageTypeEnum
+from nat.experimental.test_time_compute.models.strategy_base import StrategyBase
+from nat.experimental.test_time_compute.models.ttc_item import TTCItem
+from nat.utils.io.model_processing import remove_r1_think_tags
 
 logger = logging.getLogger(__name__)
 
@@ -68,7 +68,7 @@ class MultiQueryRetrievalSearch(StrategyBase):
             from langchain_core.prompts import PromptTemplate
         except ImportError:
             raise ImportError("langchain-core is required for MultiQueryRetrievalSearch. "
-                              "Install aiqtoolkit-langchain or similar.")
+                              "Install nvidia-nat-langchain or similar.")
 
         new_ttc_items: list[TTCItem] = []