nvidia-nat 1.1.0a20251020__py3-none-any.whl

nat/experimental/test_time_compute/scoring/llm_based_agent_scorer.py

@@ -0,0 +1,168 @@
+# 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 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__)
+
+
+class LLMBasedAgentScorer(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 planner.
+        """
+        self.llm_bound = await builder.get_llm(self.config.scoring_llm, wrapper_type=LLMFrameworkEnum.LANGCHAIN)
+
+    def supported_pipeline_types(self) -> [PipelineTypeEnum]:
+        return [PipelineTypeEnum.AGENT_EXECUTION]
+
+    def stage_type(self) -> StageTypeEnum:
+        return StageTypeEnum.SCORING
+
+    async def score_single(self, original_prompt: str, agent_context: str, item: TTCItem) -> float:
+        """
+        Score a single planning item using the LLM.
+
+        Args:
+            original_prompt (str): The original prompt.
+            agent_context (str): The agent context.
+            item (TTCItem): The item to score.
+
+        Returns:
+            float: The score of the 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 `scoring_llm` must be an instance of `BaseChatModel`.")
+
+        model: BaseChatModel = self.llm_bound
+
+        prompt_template = PromptTemplate(
+            template=self.config.scoring_template,
+            input_variables=["objective", "input", "output"],
+            validate_template=True,
+        )
+
+        prompt = (await prompt_template.ainvoke(
+            input={
+                "objective": agent_context,
+                "input": str(item.input) if not original_prompt else original_prompt,
+                "output": str(item.output)
+            }))
+
+        response = (await model.ainvoke(prompt)).content
+        response = remove_r1_think_tags(response)
+
+        # Score will following the format of `FINAL SCORE: <float>` in the response from the LLM
+        if not isinstance(response, str):
+            logger.warning(f"Invalid response from LLM for scoring: {response}.")
+            raise ValueError("Unable to parse the score from the LLM response.")
+
+        response = response.strip()
+        match = re.search(r'FINAL SCORE:\s*([\d.]+)', response)
+        if not match:
+            logger.warning(f"Could not parse the score from the response: {response}.")
+            score_str = '0.0'
+        else:
+            score_str = match.group(1)
+
+        try:
+            score = float(score_str)
+        except ValueError:
+            logger.warning(f"Could not convert the score string '{score_str}' to float.")
+            raise ValueError(f"Unable to convert the extracted score '{score_str}' to a float.")
+
+        return score
+
+    async def ainvoke(self,
+                      items: list[TTCItem],
+                      original_prompt: str | None = None,
+                      agent_context: str | None = None,
+                      **kwargs) -> list[TTCItem]:
+        """
+        Score a list of planning items.
+
+        Args:
+            original_prompt (str): The original prompt.
+            agent_context (str): The agent context.
+            items (list[TTCItem]): The list of planning items to score.
+
+        Returns:
+            list[float]: A list of scores corresponding to each planning item.
+        """
+        # Run score single concurrently for all planning items
+        # Then set the score attribute on each planning item
+        if not items:
+            return []
+        tasks = [
+            self.score_single(original_prompt=original_prompt, agent_context=agent_context, item=item) for item in items
+        ]
+
+        # Gather all scores concurrently
+        scores = await asyncio.gather(*tasks)
+
+        if len(scores) != len(items):
+            logger.warning(f"Number of scores {len(scores)} does not match the number of items {len(items)}.")
+            raise ValueError("Mismatch in number of scores and planning items.")
+
+        logger.debug("Scores for planning items: %s", scores)
+
+        # Set the score on each planning item for reference
+        for idx, score in enumerate(scores):
+            items[idx].score = score
+
+        return items
+
+
+@register_ttc_strategy(config_type=LLMBasedAgentScoringConfig)
+async def register_llm_based_agent_scorer(config: LLMBasedAgentScoringConfig, builder: Builder):
+    """
+    Register the LLM-based agent scorer with the provided configuration and builder.
+
+    Args:
+        config (LLMBasedAgentScoringConfig): The configuration for the LLM-based agent scorer.
+        builder (Builder): The builder instance to use for building components.
+
+    Returns:
+        LLMBasedAgentScorer: The registered LLM-based agent scorer.
+    """
+    scorer = LLMBasedAgentScorer(config)
+    await scorer.build_components(builder)
+    yield scorer

nat/experimental/test_time_compute/scoring/llm_based_plan_scorer.py

@@ -0,0 +1,168 @@
+# 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 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__)
+
+
+class LLMBasedPlanScorer(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 planner.
+        """
+        self.llm_bound = await builder.get_llm(self.config.scoring_llm, wrapper_type=LLMFrameworkEnum.LANGCHAIN)
+
+    def supported_pipeline_types(self) -> [PipelineTypeEnum]:
+        return [PipelineTypeEnum.PLANNING]
+
+    def stage_type(self) -> StageTypeEnum:
+        return StageTypeEnum.SCORING
+
+    async def score_single(self, original_prompt: str, agent_context: str, planning_item: TTCItem) -> float:
+        """
+        Score a single planning item using the LLM.
+
+        Args:
+            original_prompt (str): The original prompt.
+            agent_context (str): The agent context.
+            planning_item (TTCItem): The item to score.
+
+        Returns:
+            float: The score of the 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 `scoring_llm` must be an instance of `BaseChatModel`.")
+
+        model: BaseChatModel = self.llm_bound
+
+        prompt_template = PromptTemplate(
+            template=self.config.scoring_template,
+            input_variables=["original_prompt", "context", "plan"],
+            validate_template=True,
+        )
+
+        prompt = (await prompt_template.ainvoke(
+            input={
+                "original_prompt": original_prompt,
+                "context": agent_context,
+                "plan": remove_r1_think_tags(planning_item.plan)
+            }))
+
+        response = (await model.ainvoke(prompt)).content
+
+        # Score will following the format of `FINAL SCORE: <float>` in the response from the LLM
+        if not isinstance(response, str):
+            logger.warning(f"Invalid response from LLM for scoring: {response}.")
+            raise ValueError("Unable to parse the score from the LLM response.")
+
+        response = response.strip()
+        match = re.search(r'FINAL SCORE:\s*([\d.]+)', response)
+        if not match:
+            logger.warning(f"Could not parse the score from the response: {response}.")
+            score_str = '0.0'
+        else:
+            score_str = match.group(1)
+
+        try:
+            score = float(score_str)
+        except ValueError:
+            logger.warning(f"Could not convert the score string '{score_str}' to float.")
+            raise ValueError(f"Unable to convert the extracted score '{score_str}' to a float.")
+
+        return score
+
+    async def ainvoke(self,
+                      items: list[TTCItem],
+                      original_prompt: str | None = None,
+                      agent_context: str | None = None,
+                      **kwargs) -> list[TTCItem]:
+        """
+        Score a list of planning items.
+
+        Args:
+            original_prompt (str): The original prompt.
+            agent_context (str): The agent context.
+            items (list[TTCItem]): The list of planning items to score.
+
+        Returns:
+            list[float]: A list of scores corresponding to each planning item.
+        """
+        # Run score single concurrently for all planning items
+        # Then set the score attribute on each planning item
+        if not items:
+            return []
+        tasks = [
+            self.score_single(original_prompt=original_prompt, agent_context=agent_context, planning_item=item)
+            for item in items
+        ]
+
+        # Gather all scores concurrently
+        scores = await asyncio.gather(*tasks)
+
+        if len(scores) != len(items):
+            logger.warning(f"Number of scores {len(scores)} does not match the number of planning items {len(items)}.")
+            raise ValueError("Mismatch in number of scores and planning items.")
+
+        logger.debug("Scores for planning items: %s", scores)
+
+        # Set the score on each planning item for reference
+        for idx, score in enumerate(scores):
+            items[idx].score = score
+
+        return items
+
+
+@register_ttc_strategy(config_type=LLMBasedPlanScoringConfig)
+async def register_llm_based_plan_scorer(config: LLMBasedPlanScoringConfig, builder: Builder):
+    """
+    Register the LLM-based plan scorer strategy.
+
+    Args:
+        config (LLMBasedPlanScoringConfig): The configuration for the strategy.
+        builder (Builder): The builder instance.
+
+    Returns:
+        LLMBasedPlanScorer: The registered LLM-based plan scorer.
+    """
+    scorer = LLMBasedPlanScorer(config)
+    await scorer.build_components(builder)
+    yield scorer

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