nvidia-nat 1.2.0__py3-none-any.whl

nat/experimental/test_time_compute/models/stage_enums.py

@@ -0,0 +1,43 @@
+# 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.
+
+from enum import Enum
+
+
+class PipelineTypeEnum(str, Enum):
+    """
+    Enum to represent the type of pipeline used in Inference Time Scaling.
+    """
+    PLANNING = "planning"
+    TOOL_USE = "tool_use"
+    AGENT_EXECUTION = "agent_execution"
+    CUSTOM = "custom"
+
+    def __str__(self) -> str:
+        return self.value
+
+
+class StageTypeEnum(str, Enum):
+    """
+    Enum to represent the type of stage in a pipeline.
+    """
+    SEARCH = "search"
+    EDITING = "editing"
+    SCORING = "scoring"
+    SELECTION = "selection"
+    CUSTOM = "custom"
+
+    def __str__(self) -> str:
+        return self.value

nat/experimental/test_time_compute/models/strategy_base.py

@@ -0,0 +1,66 @@
+# 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.
+
+from abc import ABC
+from abc import abstractmethod
+
+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):
+    """
+    Abstract base class for strategy implementations.
+
+    This class defines the interface for strategies that can be used in the
+    TTC framework. Concrete strategy classes should
+    implement the methods defined in this class.
+    """
+
+    def __init__(self, config: TTCStrategyBaseConfig) -> None:
+        self.config: TTCStrategyBaseConfig = config
+        self.pipeline_type: PipelineTypeEnum | None = None
+
+    @abstractmethod
+    async def build_components(self, builder: Builder) -> None:
+        """Build the components required for the selector."""
+        pass
+
+    @abstractmethod
+    async def ainvoke(self,
+                      items: list[TTCItem],
+                      original_prompt: str | None = None,
+                      agent_context: str | None = None,
+                      **kwargs) -> [TTCItem]:
+        pass
+
+    @abstractmethod
+    def supported_pipeline_types(self) -> [PipelineTypeEnum]:
+        """Return the stage types supported by this selector."""
+        pass
+
+    @abstractmethod
+    def stage_type(self) -> StageTypeEnum:
+        """Return the stage type of this strategy."""
+        pass
+
+    def set_pipeline_type(self, pipeline_type: PipelineTypeEnum) -> None:
+        """Set the pipeline type for this strategy."""
+        if pipeline_type in self.supported_pipeline_types():
+            self.pipeline_type = pipeline_type
+        else:
+            raise ValueError(f"Pipeline type {pipeline_type} is not supported by this strategy.")

nat/experimental/test_time_compute/models/tool_use_config.py

@@ -0,0 +1,41 @@
+# 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.
+
+from pydantic import BaseModel
+from pydantic import Field
+
+
+class ToolUseInputSchema(BaseModel):
+    """
+    Input schema for the tool use function.
+    """
+    tool_name: str = Field(description="The name of the tool to use. Must be registered in the system.", )
+    task_description: str = Field(description="The description of the task to perform with the tool.", )
+    motivation: str | None = Field(
+        default=None,
+        description="An optional motivation for the tool use, providing additional context or reasoning.",
+    )
+    output: str | None = Field(
+        default=None,
+        description="The output of the tool use. This can be used to store the result of the tool execution.",
+    )
+
+
+class ToolUselist(BaseModel):
+    """
+    A list of tools to use.
+    """
+    tools: list[ToolUseInputSchema] = Field(
+        description="A list of tool use inputs, each containing the tool name and task description.", )

nat/experimental/test_time_compute/models/ttc_item.py

@@ -0,0 +1,48 @@
+# 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 typing
+
+from pydantic import BaseModel
+from pydantic import ConfigDict
+from pydantic import Field
+
+
+class TTCItem(BaseModel):
+    """
+    Represents an item in the TTC functions and pipelines
+    """
+    model_config = ConfigDict(extra="allow")
+
+    input: typing.Any | None = Field(default=None,
+                                     description="Input to the function or pipeline. "
+                                     "This can be a structured tool call, or other info.")
+    output: typing.Any | None = Field(default=None,
+                                      description="Output from the function or pipeline. "
+                                      "This can be a structured tool call, or other info.")
+    plan: typing.Any | None = Field(default=None, description="Search plan for downstream agent(s).")
+    feedback: str | None = Field(default=None,
+                                 description="Feedback "
+                                 "provided by feedback steps to improve the plan.")
+    score: float | None = Field(default=None,
+                                description="Score of the plan based on feedback or other evaluation criteria. "
+                                "This can be used to rank plans.")
+    metadata: typing.Any | None = Field(default=None,
+                                        description="Additional information. This can be"
+                                        " a structured tool call, or other info not "
+                                        "in the plan.")
+    name: str | None = Field(default=None,
+                             description="Name of the item or function"
+                             ", used for identification in pipelines.")

nat/experimental/test_time_compute/register.py

@@ -0,0 +1,36 @@
+# 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.
+
+# pylint: disable=unused-import
+# flake8: noqa
+
+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 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
+from .search import multi_llm_planner
+from .search import multi_query_retrieval_search
+from .search import single_shot_multi_plan_planner
+from .selection import best_of_n_selector
+from .selection import llm_based_agent_output_selector
+from .selection import llm_based_output_merging_selector
+from .selection import llm_based_plan_selector
+from .selection import threshold_selector

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