nvidia-nat 1.2.0__py3-none-any.whl

nat/experimental/test_time_compute/selection/llm_based_agent_output_selector.py

@@ -0,0 +1,131 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import logging
+import re
+
+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.selection_config import LLMBasedAgentOutputSelectionConfig
+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__)
+
+
+class LLMBasedAgentOutputSelector(StrategyBase):
+
+    def __init__(self, config: TTCStrategyBaseConfig) -> None:
+        super().__init__(config)
+        self.llm_bound = None
+
+    async def build_components(self, builder: Builder) -> None:
+        """
+        Build the components required for the selector.
+        """
+        self.llm_bound = await builder.get_llm(self.config.selection_llm, wrapper_type=LLMFrameworkEnum.LANGCHAIN)
+
+    def supported_pipeline_types(self) -> [PipelineTypeEnum]:
+        return [PipelineTypeEnum.AGENT_EXECUTION]
+
+    def stage_type(self) -> StageTypeEnum:
+        return StageTypeEnum.SELECTION
+
+    async def ainvoke(self,
+                      items: list[TTCItem],
+                      original_prompt: str | None = None,
+                      agent_context: str | None = None,
+                      **kwargs) -> [TTCItem]:
+        """
+        Select the planning items based on the selection strategy.
+
+        Args:
+            original_prompt (str): The prompt the user provided the agent.
+            agent_context (str): The context of the agent, if applicable.
+            items (list[TTCItem]): The list of planning items to select from.
+
+        Returns:
+            TTCItem: The selected planning item.
+        """
+
+        try:
+            from langchain_core.language_models import BaseChatModel
+            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 nvidia-nat-langchain.")
+
+        from pydantic import BaseModel
+
+        if not isinstance(self.llm_bound, BaseChatModel):
+            raise ValueError("The `selection_llm` must be an instance of `BaseChatModel`.")
+
+        model: BaseChatModel = self.llm_bound
+
+        results = ""
+        for idx, item in enumerate(items):
+            item_str = str(item.output.model_dump()) if isinstance(item.output, BaseModel) else str(item.output)
+            results += f"{idx + 1}. {remove_r1_think_tags(item_str)}\n\n"
+
+        prompt_template = PromptTemplate(
+            template=self.config.selection_template,
+            input_variables=["objective", "input", "results"],
+            validate_template=True,
+        )
+
+        prompt = (await prompt_template.ainvoke(input={
+            "objective": agent_context, "input": original_prompt, "results": results
+        })).to_string()
+
+        selected_plan_index = remove_r1_think_tags((await model.ainvoke(prompt)).content)
+
+        # Model Response will be 'Plan {plan number}'
+        # Use RegEx to extrac Plan {idx} from response strong
+        if not isinstance(selected_plan_index, str):
+            logger.warning(f"Invalid response from LLM for selected plan index: {selected_plan_index}.")
+            raise ValueError("Unable to parse the selected plan index.")
+        selected_plan_index = selected_plan_index.strip()
+        match = re.match(r'^\s*SELECTED ITEM:\s+(\d+)', selected_plan_index)
+        if not match:
+            logger.warning(f"Could not parse the selected plan index from the response: {selected_plan_index}.")
+            raise ValueError("The response format for selecting the item is incorrect.")
+        index = match.group(1)
+
+        try:
+            selected_index = int(index) - 1
+            if selected_index < 0 or selected_index >= len(items):
+                raise ValueError("Selected index is out of range.")
+
+            # Return the selected planning item
+            return [items[selected_index]]
+        except ValueError as e:
+            logger.warning(f"Error parsing the selected plan index: {index}. Exception: {str(e)}")
+            raise ValueError(f"Failed to parse the selected plan index from the LLM response: {selected_plan_index}. "
+                             "Ensure the response follows the expected format.") from e
+
+
+@register_ttc_strategy(config_type=LLMBasedAgentOutputSelectionConfig)
+async def register_llm_based_agent_output_selector(config: LLMBasedAgentOutputSelectionConfig, builder: Builder):
+    """
+    Register the LLMBasedAgentOutputSelector with the builder.
+    """
+    selector = LLMBasedAgentOutputSelector(config)
+    await selector.build_components(builder)
+    yield selector

nat/experimental/test_time_compute/selection/llm_based_output_merging_selector.py

@@ -0,0 +1,159 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import logging
+
+from 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.selection_config import LLMBasedOutputMergingConfig
+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__)
+
+
+class LLMBasedOutputMergingSelector(StrategyBase):
+
+    def __init__(self, config: TTCStrategyBaseConfig) -> None:
+        super().__init__(config)
+        self.llm_bound = None
+
+    async def build_components(self, builder: Builder) -> None:
+        """
+        Build the components required for the selector.
+        """
+        self.llm_bound = await builder.get_llm(self.config.selection_llm, wrapper_type=LLMFrameworkEnum.LANGCHAIN)
+
+    def supported_pipeline_types(self) -> [PipelineTypeEnum]:
+        return [PipelineTypeEnum.AGENT_EXECUTION, PipelineTypeEnum.PLANNING]
+
+    def stage_type(self) -> StageTypeEnum:
+        return StageTypeEnum.SELECTION
+
+    async def ainvoke(self,
+                      items: list[TTCItem],
+                      original_prompt: str | None = None,
+                      agent_context: str | None = None,
+                      **kwargs) -> [TTCItem]:
+        """
+        Merge the outputs of multiple planning items into a single output
+
+        Args:
+            original_prompt (str): The prompt the user provided the agent.
+            agent_context (str): The context of the agent, if applicable.
+            items (list[TTCItem]): The list of planning items to select from.
+
+        Returns:
+            TTCItem: The selected planning item.
+        """
+
+        try:
+            from langchain_core.language_models import BaseChatModel
+            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 nvidia-nat-langchain.")
+
+        from typing import Callable
+
+        from pydantic import BaseModel
+
+        if not isinstance(self.llm_bound, BaseChatModel):
+            raise ValueError("The `selection_llm` must be an instance of `BaseChatModel`.")
+
+        if not self.pipeline_type:
+            raise RuntimeError("Pipeline type is not set. Ensure that the pipeline "
+                               "type is set before invoking the selector.")
+
+        model: BaseChatModel = self.llm_bound
+
+        results = ""
+        if self.pipeline_type == PipelineTypeEnum.AGENT_EXECUTION:
+            for idx, item in enumerate(items):
+                item_str = str(item.output.model_dump()) if isinstance(item.output, BaseModel) else str(item.output)
+                results += f"{idx + 1}. {remove_r1_think_tags(item_str)}\n\n"
+        else:
+            for idx, item in enumerate(items):
+                item_str = str(item.plan)
+                results += f"{idx + 1}. {remove_r1_think_tags(item_str)}\n\n"
+
+        prompt_template = PromptTemplate(
+            template=self.config.selection_template,
+            input_variables=["pipeline_type", "objective", "input", "results"],
+            validate_template=True,
+        )
+
+        if self.pipeline_type == PipelineTypeEnum.PLANNING:
+            pipeline_objective = "execution plans for a given objective and input."
+        else:
+            pipeline_objective = "outputs from an agent system based on the provided objective and input."
+
+        prompt = (await prompt_template.ainvoke(
+            input={
+                "objective": agent_context,
+                "input": original_prompt,
+                "results": results,
+                "pipeline_type": pipeline_objective
+            })).to_string()
+
+        merged_output = remove_r1_think_tags((await model.ainvoke(prompt)).content)
+
+        if not isinstance(merged_output, str):
+            logger.warning(f"Invalid response from LLM for merged_plan: {merged_output}.")
+            raise ValueError("Unable to parse merged plan.")
+        merged_output = merged_output.strip()
+
+        # match = split the string after 'MERGED OUTPUT:'
+        matches = merged_output.split("MERGED OUTPUT:")
+        if len(matches) > 1:
+            merged_output = matches[-1].strip()
+        else:
+            raise ValueError("Merged output does not contain 'MERGED OUTPUT:' prefix.")
+
+        # Check if a callable argument is provided in kwargs called output_parser
+        output_parser: Callable | None = kwargs.get('output_parser', None)
+        if output_parser:
+            try:
+                merged_output = output_parser(merged_output)
+            except Exception as e:
+                logger.error(f"Error parsing merged output: {e}")
+                raise ValueError("Failed to parse merged output.")
+        else:
+            merged_output = merged_output
+
+        logger.info("Merged output: %s", str(merged_output))
+
+        # Create a new TTCItem with the merged plan or output
+        if self.pipeline_type == PipelineTypeEnum.PLANNING:
+            merged_item = TTCItem(input=items[0].input, output=merged_output, plan=merged_output)
+        else:
+            merged_item = TTCItem(input=items[0].input, output=merged_output)
+
+        return [merged_item]
+
+
+@register_ttc_strategy(config_type=LLMBasedOutputMergingConfig)
+async def register_llm_based_output_merging_selector(config: LLMBasedOutputMergingConfig, builder: Builder):
+    """
+    Register the LLMBasedOutputMergingSelector with the builder.
+    """
+    selector = LLMBasedOutputMergingSelector(config)
+    await selector.build_components(builder)
+    yield selector

nat/experimental/test_time_compute/selection/llm_based_plan_selector.py

@@ -0,0 +1,128 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import logging
+import re
+
+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.selection_config import LLMBasedPlanSelectionConfig
+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__)
+
+
+class LLMBasedPlanSelector(StrategyBase):
+
+    def __init__(self, config: TTCStrategyBaseConfig) -> None:
+        super().__init__(config)
+        self.llm_bound = None
+
+    async def build_components(self, builder: Builder) -> None:
+        """
+        Build the components required for the selector.
+        """
+        self.llm_bound = await builder.get_llm(self.config.selection_llm, wrapper_type=LLMFrameworkEnum.LANGCHAIN)
+
+    def supported_pipeline_types(self) -> [PipelineTypeEnum]:
+        return [PipelineTypeEnum.PLANNING]
+
+    def stage_type(self) -> StageTypeEnum:
+        return StageTypeEnum.SELECTION
+
+    async def ainvoke(self,
+                      items: list[TTCItem],
+                      original_prompt: str | None = None,
+                      agent_context: str | None = None,
+                      **kwargs) -> [TTCItem]:
+        """
+        Select the planning items based on the selection strategy.
+
+        Args:
+            original_prompt (str): The prompt the user provided the agent.
+            agent_context (str): The context of the agent, if applicable.
+            items (list[TTCItem]): The list of planning items to select from.
+
+        Returns:
+            TTCItem: The selected planning item.
+        """
+
+        try:
+            from langchain_core.language_models import BaseChatModel
+            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 nvidia-nat-langchain.")
+
+        if not isinstance(self.llm_bound, BaseChatModel):
+            raise ValueError("The `selection_llm` must be an instance of `BaseChatModel`.")
+
+        model: BaseChatModel = self.llm_bound
+
+        plans = ""
+        for idx, item in enumerate(items):
+            plans += f"{idx + 1}. {remove_r1_think_tags(item.plan)}\n"
+
+        prompt_template = PromptTemplate(
+            template=self.config.selection_template,
+            input_variables=["original_prompt", "context", "plans"],
+            validate_template=True,
+        )
+
+        prompt = (await prompt_template.ainvoke(input={
+            "original_prompt": original_prompt, "context": agent_context, "plans": plans
+        })).to_string()
+
+        selected_plan_index = remove_r1_think_tags((await model.ainvoke(prompt)).content)
+
+        # Model Response will be 'Plan {plan number}'
+        # Use RegEx to extrac Plan {idx} from response strong
+        if not isinstance(selected_plan_index, str):
+            logger.warning(f"Invalid response from LLM for selected plan index: {selected_plan_index}.")
+            raise ValueError("Unable to parse the selected plan index.")
+        selected_plan_index = selected_plan_index.strip()
+        match = re.match(r'^\s*SELECTED PLAN:\s+(\d+)', selected_plan_index)
+        if not match:
+            logger.warning(f"Could not parse the selected plan index from the response: {selected_plan_index}.")
+            raise ValueError("The response format for selecting the plan is incorrect.")
+        index = match.group(1)
+
+        try:
+            selected_index = int(index) - 1
+            if selected_index < 0 or selected_index >= len(items):
+                raise ValueError("Selected index is out of range.")
+
+            # Return the selected planning item
+            return [items[selected_index]]
+        except ValueError as e:
+            logger.warning(f"Error parsing the selected plan index: {index}. Exception: {str(e)}")
+            raise ValueError(f"Failed to parse the selected plan index from the LLM response: {selected_plan_index}. "
+                             "Ensure the response follows the expected format.") from e
+
+
+@register_ttc_strategy(config_type=LLMBasedPlanSelectionConfig)
+async def register_llm_based_plan_selection(config: LLMBasedPlanSelectionConfig, builder: Builder):
+    """
+    Register the LLMBasedPlanSelector with the provided configuration.
+    """
+    selector = LLMBasedPlanSelector(config)
+    await selector.build_components(Builder())
+    yield selector

nat/experimental/test_time_compute/selection/threshold_selector.py

@@ -0,0 +1,58 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import logging
+
+from nat.builder.builder import Builder
+from nat.cli.register_workflow import register_ttc_strategy
+from nat.experimental.test_time_compute.models.selection_config import ThresholdSelectionConfig
+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
+
+logger = logging.getLogger(__name__)
+
+
+class ThresholdSelector(StrategyBase):
+    """
+    Downselects only those TTCItems whose 'score' >= config.threshold.
+    """
+
+    async def build_components(self, builder: Builder) -> None:
+        # No special components needed
+        pass
+
+    def supported_pipeline_types(self) -> list[PipelineTypeEnum]:
+        return [PipelineTypeEnum.TOOL_USE]
+
+    def stage_type(self) -> StageTypeEnum:
+        return StageTypeEnum.SELECTION
+
+    async def ainvoke(self,
+                      items: list[TTCItem],
+                      original_prompt: str | None = None,
+                      agent_context: str | None = None,
+                      **kwargs) -> list[TTCItem]:
+        threshold = self.config.threshold
+        selected = [itm for itm in items if (itm.score is not None and itm.score >= threshold)]
+        logger.info("ThresholdSelector: %d items => %d items (threshold=%.1f)", len(items), len(selected), threshold)
+        return selected
+
+
+@register_ttc_strategy(config_type=ThresholdSelectionConfig)
+async def register_threshold_selector(config: ThresholdSelectionConfig, builder: Builder):
+    selector = ThresholdSelector(config)
+    yield selector

nat/front_ends/__init__.py

@@ -0,0 +1,14 @@
+# SPDX-FileCopyrightText: Copyright (c) 2024-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.

nat/front_ends/console/__init__.py

@@ -0,0 +1,14 @@
+# SPDX-FileCopyrightText: Copyright (c) 2024-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.