nvidia-nat 1.2.0__py3-none-any.whl

nat/experimental/test_time_compute/scoring/motivation_aware_scorer.py

@@ -0,0 +1,111 @@
+# 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 asyncio
+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.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__)
+
+
+class MotivationAwareScorer(StrategyBase):
+    """
+    A strategy that scores an TTCItem's output based on how well it
+    addresses both the original input (task) and the 'motivation' from metadata.
+    """
+
+    def __init__(self, config: TTCStrategyBaseConfig) -> None:
+        super().__init__(config)
+        self.llm_bound = None
+
+    async def build_components(self, builder: Builder) -> None:
+        self.llm_bound = await builder.get_llm(self.config.scoring_llm, wrapper_type=LLMFrameworkEnum.LANGCHAIN)
+
+    def supported_pipeline_types(self) -> list[PipelineTypeEnum]:
+        return [PipelineTypeEnum.TOOL_USE]
+
+    def stage_type(self) -> StageTypeEnum:
+        return StageTypeEnum.SCORING
+
+    async def ainvoke(self,
+                      items: list[TTCItem],
+                      original_prompt: str | None = None,
+                      agent_context: str | None = None,
+                      **kwargs) -> list[TTCItem]:
+        """
+        Scores each item by combining the original 'task_description' and 'motivation' with the 'output'.
+        The resulting score is stored in item.score.
+        """
+        from langchain_core.language_models import BaseChatModel
+        from langchain_core.prompts import PromptTemplate
+
+        if not isinstance(self.llm_bound, BaseChatModel):
+            raise ValueError("scoring_llm must be a BaseChatModel instance for MotivationAwareScorer.")
+
+        scoring_model: BaseChatModel = self.llm_bound
+
+        scoring_template = PromptTemplate(template=self.config.scoring_template,
+                                          input_variables=["task", "motivation", "output"],
+                                          validate_template=True)
+
+        async def score_item(item: TTCItem) -> float:
+            task_str = str(item.input) or ""
+            motivation_str = str(item.metadata) if item.metadata else ""
+            output_str = str(item.output) or ""
+
+            prompt = (await scoring_template.ainvoke({
+                "task": task_str, "motivation": motivation_str, "output": output_str
+            })).to_string()
+
+            response = (await scoring_model.ainvoke(prompt)).content
+            response = remove_r1_think_tags(response or "")
+
+            match = re.search(r'FINAL SCORE:\s*([\d.]+)', response)
+            if not match:
+                logger.warning(f"Could not parse score from response: {response}")
+                return 0.0
+
+            score_str = match.group(1)
+            try:
+                return float(score_str)
+            except ValueError:
+                logger.warning(f"Could not convert score '{score_str}' to float.")
+                return 0.0
+
+        tasks = [score_item(item) for item in items]
+        scores = await asyncio.gather(*tasks)
+
+        for i, s in enumerate(scores):
+            items[i].score = s
+
+        return items
+
+
+@register_ttc_strategy(config_type=MotivationAwareScoringConfig)
+async def register_motivation_aware_scorer(config: MotivationAwareScoringConfig, builder: Builder):
+    scorer = MotivationAwareScorer(config)
+    await scorer.build_components(builder)
+    yield scorer

nat/experimental/test_time_compute/search/multi_llm_planner.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 asyncio
+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.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__)
+
+
+class MultiLLMPlanner(StrategyBase):
+    """
+    A planner that uses multiple LLMs to generate plans. Each LLM can generate
+    a specified number of plans, and all plans are combined.
+    """
+
+    def __init__(self, config: MultiLLMPlanConfig) -> None:
+        super().__init__(config)
+        self.config = config
+        self.llms_bound = []  # Will hold the "bound" LLMs after build_components
+
+    async def build_components(self, builder: Builder) -> None:
+        """
+        Build the components required for this multi-LLM planner.
+        Binds each LLMRef from the config with the selected framework wrapper (LANGCHAIN).
+        """
+        logger.debug("Building components for MultiLLMPlanner")
+        self.llms_bound = []
+        for llm_ref in self.config.llms:
+            bound_llm = await builder.get_llm(llm_ref, wrapper_type=LLMFrameworkEnum.LANGCHAIN)
+            self.llms_bound.append(bound_llm)
+
+    def supported_pipeline_types(self) -> [PipelineTypeEnum]:
+        return [PipelineTypeEnum.PLANNING]
+
+    def stage_type(self) -> StageTypeEnum:
+        return StageTypeEnum.SEARCH
+
+    async def _generate_plan_for_temperature(self, llm, base_prompt: str, temperature: float) -> TTCItem:
+        bound_llm = llm.bind(temperature=temperature)
+        response = await bound_llm.ainvoke(base_prompt)
+        cleaned = remove_r1_think_tags(response.content if hasattr(response, 'content') else str(response))
+        # The plan is expected to start with "PLAN:" and all the text after it is the plan
+        cleaned = re.sub(r'(?i)^\s*PLAN:\s*', '', cleaned).strip()
+
+        if not cleaned:
+            logger.warning(f"No plan generated for the prompt: {base_prompt}.")
+            # Return an empty PlanningItem to avoid breaking the generation loop
+            return TTCItem(plan="Plan was not generated")
+
+        return TTCItem(plan=cleaned)
+
+    async def _generate_plans_for_llm(self, llm, base_prompt: str) -> list[TTCItem]:
+        if self.config.plans_per_llm == 1:
+            temps = [self.config.min_temperature]
+        else:
+            temps = [
+                self.config.min_temperature + (i / (self.config.plans_per_llm - 1)) *
+                (self.config.max_temperature - self.config.min_temperature) for i in range(self.config.plans_per_llm)
+            ]
+        tasks = [self._generate_plan_for_temperature(llm, base_prompt, temp) for temp in temps]
+        return await asyncio.gather(*tasks)
+
+    async def ainvoke(self,
+                      items: list[TTCItem],
+                      original_prompt: str | None = None,
+                      agent_context: str | None = None,
+                      **kwargs) -> list[TTCItem]:
+        """
+        Generate a list of PlanningItems by querying each LLM in self.llms_bound.
+        Each LLM produces 'plans_per_llm' plans.
+        """
+        try:
+            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 nvidia-nat-langchain.")
+
+        # Create a single PromptTemplate
+        planning_template = PromptTemplate(template=self.config.planning_template,
+                                           input_variables=["context", "prompt"],
+                                           validate_template=True)
+
+        # Format the prompt once
+        base_prompt = (await planning_template.ainvoke({
+            "context": agent_context, "prompt": original_prompt
+        })).to_string()
+
+        # Launch generation for each llm concurrently using the new helper method
+        tasks = [self._generate_plans_for_llm(llm, base_prompt) for llm in self.llms_bound]
+        results_nested = await asyncio.gather(*tasks)
+
+        # Flatten the nested lists of TTCItem
+        all_plans: list[TTCItem] = [p for sub in results_nested for p in sub]
+        logger.info("MultiLLMPlanner generated %d plans total.", len(all_plans))
+        return all_plans
+
+
+@register_ttc_strategy(config_type=MultiLLMPlanConfig)
+async def register_multi_llm_planner(config: MultiLLMPlanConfig, builder: Builder):
+    """
+    Register the MultiLLMPlanner strategy with the provided configuration.
+    """
+    planner = MultiLLMPlanner(config)
+    await planner.build_components(builder)
+    yield planner

nat/experimental/test_time_compute/search/multi_query_retrieval_search.py

@@ -0,0 +1,122 @@
+# 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.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__)
+
+
+class MultiQueryRetrievalSearch(StrategyBase):
+    """
+    A strategy that, for each incoming TTCItem, generates multiple new items by
+    re-writing the input 'task_description' from different perspectives.
+    Uses multiple LLMs to encourage diversity.
+    """
+
+    def __init__(self, config: MultiQueryRetrievalSearchConfig) -> None:
+        super().__init__(config)
+        self.config = config
+        self.llms_bound = []
+
+    async def build_components(self, builder: Builder) -> None:
+        """
+        Binds each LLMRef in self.config.llms to an actual LLM client.
+        """
+        self.llms_bound = []
+        for llm_ref in self.config.llms:
+            bound_llm = await builder.get_llm(llm_ref, wrapper_type=LLMFrameworkEnum.LANGCHAIN)
+            self.llms_bound.append(bound_llm)
+
+    def supported_pipeline_types(self) -> list[PipelineTypeEnum]:
+        return [PipelineTypeEnum.TOOL_USE]
+
+    def stage_type(self) -> StageTypeEnum:
+        return StageTypeEnum.SEARCH
+
+    async def ainvoke(self,
+                      items: list[TTCItem],
+                      original_prompt: str | None = None,
+                      agent_context: str | None = None,
+                      **kwargs) -> list[TTCItem]:
+        """
+        For each TTCItem, rewrite the 'input' using each LLM to create a new perspective.
+        The new TTCItems' 'output' field will store the newly generated query.
+        """
+        try:
+            from langchain_core.prompts import PromptTemplate
+        except ImportError:
+            raise ImportError("langchain-core is required for MultiQueryRetrievalSearch. "
+                              "Install nvidia-nat-langchain or similar.")
+
+        new_ttc_items: list[TTCItem] = []
+
+        # Create a single PromptTemplate object for rewriting the query
+        template_vars = ["task", "motivation"]
+        query_template = PromptTemplate(template=self.config.query_generation_template,
+                                        input_variables=template_vars,
+                                        validate_template=True)
+
+        for item in items:
+            original_task = str(item.input) or ""
+            motivation = str(item.metadata) if item.metadata else ""
+            new_ttc_items.append(
+                TTCItem(
+                    input=item.input,
+                    output=item.input,
+                    metadata=item.metadata,
+                    name=item.name,  # keep the original tool name
+                ))
+
+            for llm in self.llms_bound:
+                prompt_str = (await query_template.ainvoke({
+                    "task": original_task, "motivation": motivation
+                })).to_string()
+
+                # We'll call each LLM to produce a new query
+                response = await llm.ainvoke(prompt_str)
+                cleaned = remove_r1_think_tags(response.content if hasattr(response, 'content') else str(response))
+                cleaned = cleaned.strip()
+
+                # Create a new TTCItem for each newly generated query
+                new_item = TTCItem(
+                    input=item.input,  # keep the original input for reference
+                    output=cleaned,  # store the newly generated query in the output
+                    metadata=item.metadata,
+                    name=item.name,  # same tool name or optional new name
+                )
+                new_ttc_items.append(new_item)
+
+        logger.info("MultiQueryRetrievalSearch produced %d new items from %d original items.",
+                    len(new_ttc_items),
+                    len(items))
+
+        return new_ttc_items
+
+
+@register_ttc_strategy(config_type=MultiQueryRetrievalSearchConfig)
+async def register_multi_query_retrieval_search(config: MultiQueryRetrievalSearchConfig, builder: Builder):
+    strategy = MultiQueryRetrievalSearch(config)
+    await strategy.build_components(builder)
+    yield strategy

nat/experimental/test_time_compute/search/single_shot_multi_plan_planner.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 asyncio
+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.search_config import SingleShotMultiPlanConfig
+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 SingleShotMultiPlanPlanner(StrategyBase):
+    """
+    Implementation of the Single Shot Multi Plan Planner.
+    This planner generates multiple plans in a single shot.
+    """
+
+    def __init__(self, config: TTCStrategyBaseConfig) -> None:
+        super().__init__(config)
+        self.llm_bound = None
+
+    async def build_components(self, builder: Builder) -> None:
+        self.llm_bound = await builder.get_llm(self.config.planning_llm, wrapper_type=LLMFrameworkEnum.LANGCHAIN)
+
+    def supported_pipeline_types(self) -> [PipelineTypeEnum]:
+        return [PipelineTypeEnum.PLANNING]
+
+    def stage_type(self) -> StageTypeEnum:
+        return StageTypeEnum.SEARCH
+
+    async def ainvoke(self,
+                      items: list[TTCItem],
+                      original_prompt: str | None = None,
+                      agent_context: str | None = None,
+                      **kwargs) -> list[TTCItem]:
+        """
+        Generate a TTCItem based on the provided prompt.
+        """
+
+        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 resolve by installing nvidia-nat-langchain.")
+
+        planning_template = PromptTemplate(template=self.config.planning_template,
+                                           input_variables=["context", "prompt"],
+                                           validate_template=True)
+        prompt = (await planning_template.ainvoke(input={
+            "context": agent_context, "prompt": original_prompt
+        })).to_string()
+
+        # assert self.config.planning llm is a BaseChatModel
+        if not isinstance(self.llm_bound, BaseChatModel):
+            raise ValueError("The `planning_llm` must be an instance of `BaseChatModel`.")
+
+        model: BaseChatModel = self.llm_bound
+
+        async def generate_plan(llm: BaseChatModel, plan_prompt: str, temperature: float) -> TTCItem:
+            """
+            Helper function to generate a plan using the provided prompt and temperature.
+            """
+            llm_bound = llm.bind(temperature=temperature)
+            response = await llm_bound.ainvoke(plan_prompt)
+            cleaned = remove_r1_think_tags(response.content if hasattr(response, 'content') else str(response))
+
+            # Plan will be the string following 'PLAN:'. Use Regex tpo extract
+            cleaned = re.sub(r'(?i)^\s*PLAN:\s*', '', cleaned).strip()
+
+            if not cleaned:
+                logger.warning(f"No plan generated for the prompt: {plan_prompt}.")
+                # Return an empty PlanningItem to avoid breaking the generation loop
+                return TTCItem(plan="Plan was not generated")
+
+            return TTCItem(plan=cleaned)
+
+        # Define a list of temperatures based on min and max temperature in the config and number of plans to generate
+        temperatures = [
+            self.config.min_temperature + (i / (self.config.num_plans - 1)) *
+            (self.config.max_temperature - self.config.min_temperature) for i in range(self.config.num_plans)
+        ]
+
+        # Generate plans using the defined temperatures in parallel using asyncio
+        tasks = [generate_plan(model, prompt, temperature) for temperature in temperatures]
+        # Run the tasks concurrently and gather results
+        plans = await asyncio.gather(*tasks)
+
+        if not plans:
+            raise ValueError("No plans were generated. Please check the LLM response.")
+
+        logger.info("Generated %d plans from the SingleShotMultiPlanPlanner", self.config.num_plans)
+
+        logger.debug("Generated plans: %s", [plan.dict() for plan in plans])
+
+        return plans
+
+
+@register_ttc_strategy(config_type=SingleShotMultiPlanConfig)
+async def register_single_shot_multi_plan_planner(config: SingleShotMultiPlanConfig, builder: Builder):
+    """
+    Register the SingleShotMultiPlanPlanner strategy with the provided configuration.
+    """
+    planner = SingleShotMultiPlanPlanner(config)
+    await planner.build_components(builder)
+    yield planner

nat/experimental/test_time_compute/selection/best_of_n_selector.py

@@ -0,0 +1,63 @@
+# 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 BestOfNSelectionConfig
+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 BestOfNSelector(StrategyBase):
+
+    async def build_components(self, builder: Builder) -> None:
+        pass
+
+    def supported_pipeline_types(self) -> [PipelineTypeEnum]:
+        return [PipelineTypeEnum.PLANNING, PipelineTypeEnum.AGENT_EXECUTION, 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) -> [TTCItem]:
+
+        # Assert that every planning item has a non NoneType score
+        for item in items:
+            if item.score is None:
+                raise ValueError("Every planning item must have a score. Did you use a scorer before this?")
+
+        # Pick the planning item with the highest score
+        best_item = max(items, key=lambda x: x.score)
+
+        return [best_item]
+
+
+@register_ttc_strategy(config_type=BestOfNSelectionConfig)
+async def register_best_of_n_selector(config: BestOfNSelectionConfig, builder: Builder):
+    """
+    Register the BestOfNSelector strategy.
+    """
+    selector = BestOfNSelector(config)
+    yield selector