nvidia-nat 1.1.0a20251020__py3-none-any.whl

nat/experimental/test_time_compute/models/search_config.py

@@ -0,0 +1,120 @@
+# 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 Field
+from pydantic import model_validator
+
+from nat.data_models.component_ref import LLMRef
+from nat.data_models.ttc_strategy import TTCStrategyBaseConfig
+
+
+class SingleShotMultiPlanConfig(TTCStrategyBaseConfig, name="single_shot_multi_plan"):
+    num_plans: int = Field(default=4, description="Number of plans to generate.")
+    max_temperature: float = Field(default=1.0,
+                                   description="Maximum temperature to use for sampling when generating plans. "
+                                   "This can help control the randomness of the generated plans.")
+    min_temperature: float = Field(default=0.5,
+                                   description="Minimum temperature to use for sampling when generating plans. "
+                                   "This can help control the randomness of the generated plans.")
+    # If strategy is provided, LLM must be
+    planning_llm: LLMRef | typing.Any | None = Field(
+        default=None,
+        description="The LLM to use for planning. This can be a callable or an "
+        "instance of an LLM client.")
+
+    planning_template: str = Field(
+        default=("You are an expert reasoning model task with creating a detailed execution plan"
+                 " for a system that has the following information to get the result of a given input:\n\n"
+                 "**System Information:**\n {context}"
+                 "**Input:** \n{prompt}\n\n"
+                 "An example plan could look like this:\n\n"
+                 "1. Call tool A with input X\n"
+                 "2. Call tool B with input Y\n"
+                 "3. Interpret the output of tool A and B\n"
+                 "4. Return the final result"
+                 "\n\nBegin the final plan with PLAN:\n"),
+        description="The template to use for generating plans.")
+
+    @model_validator(mode="before")
+    def validate_strategies(cls, values: dict[str, typing.Any]) -> dict[str, typing.Any]:
+        """
+        Ensure that the required LLMs are provided based on the selected strategies.
+        """
+        # Validate planning strategy: planning_llm must be provided if planning_strategy is set
+        if values.get('planning_llm') is None:
+            raise ValueError('planning_llm must be provided when planning_strategy is set.')
+
+        return values
+
+
+class MultiLLMPlanConfig(TTCStrategyBaseConfig, name="multi_llm_plan"):
+    """Configuration for a 'multi LLM plan generation' strategy."""
+    llms: list[LLMRef] = Field(
+        default_factory=list,
+        description="list of LLMs to use for plan generation. Each LLM can generate one or more plans.")
+    plans_per_llm: int = Field(default=2, description="Number of plans each LLM should generate.")
+    max_temperature: float = Field(default=1.0,
+                                   description="Maximum temperature to use for sampling when generating plans. "
+                                   "This can help control the randomness of the generated plans.")
+    min_temperature: float = Field(default=0.5,
+                                   description="Minimum temperature to use for sampling when generating plans. "
+                                   "This can help control the randomness of the generated plans.")
+    planning_template: str = Field(
+        default=("You are an expert reasoning model task with creating a detailed execution plan"
+                 " for a system that has the following information to get the result of a given input:\n\n"
+                 "**System Information:**\n {context}"
+                 "**Input:** \n{prompt}\n\n"
+                 "An example plan could look like this:\n\n"
+                 "1. Call tool A with input X\n"
+                 "2. Call tool B with input Y\n"
+                 "3. Interpret the output of tool A and B\n"
+                 "4. Return the final result"
+                 "\n\nBegin the final plan with PLAN:\n"),
+        description="The template to use for generating plans.")
+
+    @model_validator(mode="before")
+    def validate_multi_llm_strategies(cls, values: dict) -> dict:
+        if not values.get('llms'):
+            raise ValueError('Must provide at least one LLMRef in `llms` for multi-LLM strategy.')
+        return values
+
+
+class MultiQueryRetrievalSearchConfig(TTCStrategyBaseConfig, name="multi_query_retrieval_search"):
+    """
+    Configuration for the MultiQueryRetrievalSearch strategy.
+    This strategy generates multiple new 'TTCItem's per original item,
+    each containing a differently phrased or re-focused version of the original task.
+    """
+    llms: list[LLMRef] = Field(default_factory=list,
+                               description="list of LLM references to use for generating diverse queries.")
+
+    query_generation_template: str = Field(
+        default=("You are an expert at re-framing a user's query to encourage new solution paths. "
+                 "Given the task description and an optional motivation, produce a short alternative query "
+                 "that addresses the same task from a different angle. By generating multiple "
+                 "perspectives on the task, your goal is to help "
+                 "the user overcome some of the limitations of distance-based similarity search.\n\n"
+                 "Task: {task}\n"
+                 "Motivation: {motivation}\n\n"
+                 "Output a concise new query statement below. Only output the revised query and nothing else.\n"),
+        description="Prompt template for rewriting the task from a different perspective.")
+
+    @model_validator(mode="before")
+    def validate_llms(cls, values):
+        if not values.get('llms'):
+            raise ValueError("At least one LLMRef must be provided for multi_query_retrieval_search.")
+        return values

nat/experimental/test_time_compute/models/selection_config.py

@@ -0,0 +1,154 @@
+# 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 Field
+from pydantic import model_validator
+
+from nat.data_models.component_ref import LLMRef
+from nat.data_models.ttc_strategy import TTCStrategyBaseConfig
+
+
+class LLMBasedPlanSelectionConfig(TTCStrategyBaseConfig, name="llm_based_plan_selection"):
+    """
+    Configuration for LLMBasedSelection.
+    """
+    selection_llm: LLMRef | typing.Any | None = Field(
+        default=None,
+        description="The LLM to use for selecting the best plan. This can be an instance of an LLM client.")
+
+    selection_template: str = Field(
+        default=("You are tasked with selecting the best plan from several alternative plans."
+                 " Review the following plans and their feedback carefully to select the most "
+                 "comprehensive, efficient, and effective one."
+                 "The plan is for an agent system with the following objective and context:\n\n"
+                 "{context}\n\n"
+                 "The system is asked to achieve the following goal:\n\n"
+                 "{original_prompt}\n\n"
+                 "The generated plans are as follows."
+                 "\n\n{plans}"
+                 "\n\nBased on your analysis, which plan (numbered 1 and onwards) is the best? "
+                 "Provide a thorough explanation of your choice,"
+                 " referencing specific strengths from the feedback and how they outweigh any weaknesses."
+                 "Make sure you begin your choice of selected plan with the words 'SELECTED PLAN:' "
+                 "followed by the plan number."),
+        description="The template to use for selecting the best plan. This should guide the LLM on how to evaluate "
+        "the plans and select the best one. Ensure it is clear and concise.")
+
+    @model_validator(mode="before")
+    def validate_strategies(cls, values: dict[str, typing.Any]) -> dict[str, typing.Any]:
+        """
+        Ensure that the selection_llm is provided when using LLMBasedSelection.
+        """
+        if values.get('selection_llm') is None:
+            raise ValueError('selection_llm must be provided when'
+                             ' selection_strategy is set to LLM_BASED_PLAN_SELECTION.')
+
+        return values
+
+
+class LLMBasedAgentOutputSelectionConfig(TTCStrategyBaseConfig, name="llm_based_agent_output_selection"):
+    """
+    Configuration for LLMBasedSelection.
+    """
+    selection_llm: LLMRef | typing.Any | None = Field(
+        default=None,
+        description="The LLM to use for selecting the best plan. This can be an instance of an LLM client.")
+
+    selection_template: str = Field(
+        default=("You are tasked with selecting the best output from several output."
+                 "The outputs are from an agent system whose object and input will be provided below.\n "
+                 "Review all the outputs and select one that fits the best. You will do this by "
+                 "looking at how many outputs have the same classification. Chose the one that has the most. "
+                 "Of the ones that have the same classification, choose the one that is the most complete, "
+                 "clear, and comprehensive. The objective of the agent is: \n"
+                 "{objective}\n\n"
+                 "\n\nThe agent is asked to achieve the following goal:\n\n"
+                 "{input}\n\n"
+                 "The generated outputs are as follows."
+                 "\n\n{results}"
+                 "\n\nBased on your analysis, which plan (numbered 1 and onwards) is the best? "
+                 "Provide a thorough explanation of your choice,"
+                 " referencing specific strengths from the feedback and how they outweigh any weaknesses."
+                 "You must ALWAYS select an option, even if the options are identical or similar. "
+                 "Make sure you begin your choice of selected plan with the words 'SELECTED ITEM:' "
+                 "followed by the plan number."),
+        description="The template to use for selecting the best output. This should guide the LLM on how to evaluate "
+        "the outputs and select the best one. Ensure it is clear and concise. Must contain {objective}, "
+        "{input}, and {results} ")
+
+    @model_validator(mode="before")
+    def validate_strategies(cls, values: dict[str, typing.Any]) -> dict[str, typing.Any]:
+        """
+        Ensure that the selection_llm is provided when using LLMBasedSelection.
+        """
+        if values.get('selection_llm') is None:
+            raise ValueError('selection_llm must be provided when '
+                             'selection_strategy is set to LLM_BASED_AGENT_OUTPUT_SELECTION.')
+
+        return values
+
+
+class LLMBasedOutputMergingConfig(TTCStrategyBaseConfig, name="llm_based_agent_output_merging"):
+    """
+    Configuration for LLMBasedSelection.
+    """
+    selection_llm: LLMRef | typing.Any | None = Field(
+        default=None,
+        description="The LLM to use for selecting the best plan. This can be an instance of an LLM client.")
+
+    selection_template: str = Field(
+        default=("You are tasked with merging the output of an agent systems that produces {pipeline_type}."
+                 "The outputs are from an agent system whose objective and input will be provided below.\n "
+                 "Review all the outputs, please combine them all into one output, keeping with the intended structure "
+                 "generated by the outputs and general tone. Capture the important pieces of each of the outputs "
+                 "to create comprehensive output that achieves the input and objective. "
+                 "The objective of the agent is: \n"
+                 "{objective}\n\n"
+                 "\n\nThe agent is asked to achieve the following goal:\n\n"
+                 "{input}\n\n"
+                 "The generated outputs are as follows."
+                 "\n\n{results}"
+                 "\n\n Make sure you begin your updated output with the words 'MERGED OUTPUT:' "),
+        description="The template to use for selecting the best output. This should guide the LLM on how to evaluate "
+        "the outputs and select the best one. Ensure it is clear and concise. Must contain {objective}, "
+        "{input}, and {results} ")
+
+    @model_validator(mode="before")
+    def validate_strategies(cls, values: dict[str, typing.Any]) -> dict[str, typing.Any]:
+        """
+        Ensure that the selection_llm is provided when using LLMBasedSelection.
+        """
+        if values.get('selection_llm') is None:
+            raise ValueError('selection_llm must be provided when '
+                             'selection_strategy is set to LLM_BASED_AGENT_OUTPUT_SELECTION.')
+
+        return values
+
+
+class ThresholdSelectionConfig(TTCStrategyBaseConfig, name="threshold_selection"):
+    """
+    Configuration for a selection strategy that keeps only the items
+    whose scores exceed a specified threshold.
+    """
+    threshold: float = Field(default=5.0, description="Only keep TTCItems with score >= this value.")
+
+
+class BestOfNSelectionConfig(TTCStrategyBaseConfig, name="best_of_n_selection"):
+    """
+    Configuration for Best of N Selection
+    """
+    pass

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,67 @@
+# 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.data_models.ttc_strategy import TTCStrategyBaseConfig
+from nat.experimental.test_time_compute.models.stage_enums import PipelineTypeEnum
+from nat.experimental.test_time_compute.models.stage_enums import StageTypeEnum
+from nat.experimental.test_time_compute.models.ttc_item import TTCItem
+
+
+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) -> list[TTCItem]:
+        pass
+
+    @abstractmethod
+    def supported_pipeline_types(self) -> list[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,35 @@
+# 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.
+
+# 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