nvidia-nat 1.2.0rc5__py3-none-any.whl

aiq/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 aiq.data_models.component_ref import LLMRef
+from aiq.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

aiq/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

aiq/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 aiq.builder.builder import Builder
+from aiq.experimental.test_time_compute.models.ttc_item import TTCItem
+from aiq.experimental.test_time_compute.models.stage_enums import StageTypeEnum, PipelineTypeEnum
+from aiq.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.")

aiq/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.", )

aiq/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.")

aiq/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 its_tool_orchestration_function
+from .functions import its_tool_wrapper_function
+from .functions import plan_select_execute_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

aiq/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 aiq.builder.builder import Builder
+from aiq.builder.framework_enum import LLMFrameworkEnum
+from aiq.cli.register_workflow import register_ttc_strategy
+from aiq.data_models.ttc_strategy import TTCStrategyBaseConfig
+from aiq.experimental.test_time_compute.models.scoring_config import LLMBasedAgentScoringConfig
+from aiq.experimental.test_time_compute.models.stage_enums import PipelineTypeEnum
+from aiq.experimental.test_time_compute.models.stage_enums import StageTypeEnum
+from aiq.experimental.test_time_compute.models.strategy_base import StrategyBase
+from aiq.experimental.test_time_compute.models.ttc_item import TTCItem
+from aiq.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 aiqtoolkit-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