nvidia-nat 1.2.0__py3-none-any.whl

nat/experimental/test_time_compute/functions/ttc_tool_wrapper_function.py

@@ -0,0 +1,146 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import logging
+
+from pydantic import BaseModel
+from pydantic import Field
+
+from nat.builder.builder import Builder
+from nat.builder.framework_enum import LLMFrameworkEnum
+from nat.builder.function import Function
+from nat.builder.function_info import FunctionInfo
+from nat.cli.register_workflow import register_function
+from nat.data_models.component_ref import FunctionRef
+from nat.data_models.component_ref import LLMRef
+from nat.data_models.function import FunctionBaseConfig
+from nat.utils.string_utils import convert_to_str
+
+logger = logging.getLogger(__name__)
+
+
+class TTCToolWrapperFunctionConfig(FunctionBaseConfig, name="ttc_tool_wrapper"):
+    """
+    Configuration for the TTCToolWrapperFunction, which is used to wrap a function that will be executed
+    in the inference time scaling pipeline.
+
+    This function is responsible for turning an 'objective' or description for the tool into tool input.
+
+    NOTE: Only supports LLMs with structured output.
+    """
+
+    augmented_fn: FunctionRef = Field(description="The name of the function to reason on.")
+
+    input_llm: LLMRef = Field(description="The LLM that will generate input to the function.")
+    verbose: bool = Field(default=False, description="Whether to log detailed information.")
+
+    downstream_template: str = Field(
+        description="The template for the input LLM to generate structured input to the function.",
+        default=("You are highly sophisticated generalist AI assistant. Your objective is to act as a"
+                 " conduit between a user's task for a function and the function itself. You will be given a general "
+                 "description of the task, or pseudo input for a function. You will also be provided with description "
+                 "of the function, its input schema, and the output schema. Your task is to generate structured input "
+                 "to the function based on the description of the task and the function's input schema. If you do not "
+                 "have enough information to generate structured input, you should respond with 'NOT ENOUGH "
+                 "INFORMATION'. \n\n The description of the function is: {function_description}\n\n"
+                 "The input schema of the function is: {input_schema}\n\n"
+                 "The output schema of the function is: {output_schema}\n\n"
+                 "The description of the task is: {task_description}\n\n"
+                 "The structured input to the function is: "))
+
+    tool_description: str | None = Field(description="The description of the tool to be used for the function.",
+                                         default=None)
+
+
+@register_function(config_type=TTCToolWrapperFunctionConfig, framework_wrappers=[LLMFrameworkEnum.LANGCHAIN])
+async def register_ttc_tool_wrapper_function(
+    config: TTCToolWrapperFunctionConfig,
+    builder: Builder,
+):
+    """
+    Register the TTCToolWrapperFunction with the provided builder and configuration.
+    """
+
+    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.")
+
+    augmented_function: Function = builder.get_function(config.augmented_fn)
+    input_llm: BaseChatModel = await builder.get_llm(config.input_llm, wrapper_type=LLMFrameworkEnum.LANGCHAIN)
+
+    if not augmented_function.has_single_output:
+        raise ValueError("TTCToolWrapperFunction only supports functions with a single output.")
+
+    if not augmented_function.has_single_output:
+        raise ValueError("TTCToolWrapperFunction only supports functions with a single output.")
+
+    if augmented_function.description and augmented_function.description != "":
+        augmented_function_desc = augmented_function.description
+    else:
+        if not config.tool_description:
+            raise ValueError(f"Function {config.augmented_fn} does not have a description. Cannot augment "
+                             f"function without a description and without a tool description.")
+
+        augmented_function_desc = config.tool_description
+
+    fn_input_schema: BaseModel = augmented_function.input_schema
+    fn_output_schema: BaseModel = augmented_function.single_output_schema
+
+    runnable_llm = input_llm.with_structured_output(schema=fn_input_schema)
+
+    template = PromptTemplate(
+        template=config.downstream_template,
+        input_variables=["function_description", "input_schema", "output_schema", "task_description"],
+        validate_template=True)
+
+    function_description = (f"\nDescription: {augmented_function_desc}\n" +
+                            "\n Input should be a thorough description with all relevant information on what "
+                            f"the tool should do.  The tool requires information about "
+                            f"{fn_input_schema.model_fields}")
+
+    async def single_inner(input_message: str) -> fn_output_schema:
+        """
+        Inner function to handle the streaming output of the TTCToolWrapperFunction.
+        It generates structured input for the augmented function based on the input message.
+        """
+
+        prompt = await template.ainvoke(
+            input={
+                "function_description": augmented_function_desc,
+                "input_schema": fn_input_schema,
+                "output_schema": fn_output_schema,
+                "task_description": input_message
+            })
+
+        prompt = prompt.to_string()
+
+        if config.verbose:
+            logger.info("TTCToolWrapperFunction: Generated prompt: %s", prompt)
+
+        llm_parsed = await runnable_llm.ainvoke(prompt)
+
+        if not llm_parsed:
+            logger.warning("TTCToolWrapperFunction: LLM parsing error")
+            return "Not enough information"
+
+        # Call the augmented function with the structured input
+        result = await augmented_function.acall_invoke(llm_parsed)
+
+        return result
+
+    yield FunctionInfo.from_fn(fn=single_inner, description=function_description, converters=[convert_to_str])

nat/experimental/test_time_compute/models/editor_config.py

@@ -0,0 +1,132 @@
+# 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 LLMAsAJudgeEditorConfig(TTCStrategyBaseConfig, name="llm_as_a_judge_editor"):
+    """
+    Configuration for the LLMAsAJudgeEditor.
+    """
+    num_feedback: int = Field(default=10,
+                              description="Number of feedback items to generate for each plan during editing. "
+                              "This can help in refining the plans based on feedback.")
+
+    # If strategy is provided, LLM must be
+    editing_llm: LLMRef | typing.Any | None = Field(
+        default=None,
+        description="The LLM to use for editing the plans. This can be a callable or an instance of an LLM client.")
+
+    # If strategy is LLM_AS_A_JUDGE, ensure that the feedback_llm is provided.
+    feedback_llm: LLMRef | typing.Any | None = Field(default=None,
+                                                     description="The LLM to use for generating feedback on the plans."
+                                                     " This can be a callable or an instance of an LLM client.")
+
+    editor_template: str = Field(default=(
+        "You are an expert at improving execution plans. You will be given a plan and feedback on that plan."
+        " Your task is to create an improved version of the plan that addresses the feedback "
+        "while maintaining its strengths.\n\n"
+        "Here is the context:\n\n"
+        "{context}\n\n"
+        "**Input:** \n{original_prompt}\n\n"
+        "**Original Plan:**\n{plan}\n\n"
+        "**Feedback on the Plan:**\n{feedback}\n\n"
+        "Please provide an improved version of the plan that addresses"
+        " the feedback points. Maintain the same structure and "
+        "step-by-step format, but enhance the content. Do not include explanations of your changes, just provide the "
+        "improved plan directly:\n\n"
+        "Begin the final improve plan with 'EDITED PLAN:'"),
+                                 description="The template to use for editing the planning items based on feedback.")
+
+    feedback_template: str = Field(
+        default=("You are an expert at evaluating execution plans. You will be given a plan and "
+                 "need to provide {num_feedback} "
+                 "specific points of feedback about its strengths and weaknesses.\n\n"
+                 "Your feedback should cover aspects like:\n"
+                 "- Comprehensiveness of the plan\n"
+                 "- Logical flow and sequencing\n"
+                 "- Appropriate use of available tools\n"
+                 "- Potential edge cases or failure points\n"
+                 "- Efficiency and optimization opportunities\n\n"
+                 "Here is the context and plan to evaluate:\n\n"
+                 "{context}\n\n"
+                 "**Objective:** \n{original_prompt}\n\n"
+                 "**Plan to Evaluate:**\n{plan}\n\n"
+                 "Please provide exactly {num_feedback} numbered points of feedback, including "
+                 "both strengths and areas for improvement. Begin the feedback with 'FEEDBACK:' and provide"
+                 "{num_feedback} specific feedback points."),
+        description="The template to use for generating feedback for each planning item.")
+
+    @model_validator(mode="before")
+    def validate_strategies(cls, values: dict[str, typing.Any]) -> dict[str, typing.Any]:
+
+        if values.get('editing_llm') is None:
+            raise ValueError('editing_llm must be provided when editing_strategy is set.')
+        # If editing strategy is LLM_AS_A_JUDGE, feedback_llm must also be provided
+        if (values.get('feedback_llm') is None):
+            raise ValueError('feedback_llm must be provided when editing_strategy is LLM_AS_A_JUDGE.')
+
+        return values
+
+
+class IterativePlanRefinementConfig(TTCStrategyBaseConfig, name="iterative_plan_refinement"):
+    """Configuration for an 'iterative plan refinement' strategy."""
+    editor_llm: LLMRef | typing.Any | None = Field(
+        default=None, description="The LLM to use for generating and refining the plan across multiple iterations.")
+    num_iterations: int = Field(default=3, description="How many refinement steps to perform.")
+    refinement_template: str = Field(
+        default=("You have the current plan:\n{current_plan}\n\n"
+                 "The plan was generated to achieve the following objective:\n{original_prompt}\n\n"
+                 "Using an agent system with the following description:\n{context}\n\n"
+                 "Refine or improve it to achieve the objective better."
+                 "Output the updated plan, beginning with:\nEDITED PLAN:\n"),
+        description="Prompt used in each iteration to refine the plan.")
+
+    @model_validator(mode="before")
+    def validate_iterative_strategies(cls, values: dict) -> dict:
+        if not values.get('editor_llm'):
+            raise ValueError('planning_llm must be provided for iterative plan refinement.')
+        if values.get('num_iterations', 0) < 1:
+            raise ValueError('num_iterations must be >= 1 for iterative plan refinement.')
+        return values
+
+
+class MotivationAwareSummarizationConfig(TTCStrategyBaseConfig, name="motivation_aware_editing"):
+    """
+    Configuration for the MotivationAwareSummarization strategy.
+    """
+    editor_llm: LLMRef | typing.Any | None = Field(
+        default=None,
+        description="The LLM to use for editing the plans. This can be a callable or an instance of an LLM client.")
+
+    editor_template: str = Field(
+        default=("You are an expert at summarizing key information from relevant documents based on an input task"
+                 "and motivation. Given a task and motivation, and documents, your task is to create a concise "
+                 "a summarized response to the task and motivation grounded in the documents .\n\n"
+                 "Here is the task:\n\n"
+                 "{task}\n\n"
+                 "Here is the motivation:\n\n"
+                 "{motivation}\n\n"
+                 "and here are the documents:\n\n"
+                 "{output}\n\n"
+                 "Please respond with a concise summary that addresses the task and motivation, in at most one"
+                 "or two sentences. Do not include any other output except the summary. "),
+        description="The template to use for summarizing documents.")

nat/experimental/test_time_compute/models/scoring_config.py

@@ -0,0 +1,112 @@
+# 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 LLMBasedPlanScoringConfig(TTCStrategyBaseConfig, name="llm_based_plan_scoring"):
+    """
+    Configuration for LLMBasedScoring.
+    """
+    scoring_llm: LLMRef | typing.Any | None = Field(
+        default=None,
+        description="The LLM to use for scoring the plans. This can be a callable or an instance of an LLM client.")
+
+    scoring_template: str = Field(
+        default=("You are an expert reasoning model tasked with scoring the following execution plan based on its"
+                 "quality and relevance to the provided input to an agent system.\n\n"
+                 "The agent system's role is:\n{context}\n\n"
+                 "It has been tasked with achieving the following goal: \n{original_prompt}\n\n"
+                 "The following plan has been generated to achieve this goal:\n\n{plan}\n\n"
+                 "Score the plan on a scale from 1 to 10, where 10 is the best. "
+                 "Return the final score as a floating point number preceded by `FINAL SCORE:` without any "
+                 "other text before or after it\n"),
+        description="The template to use for scoring the plans.")
+
+    @model_validator(mode="before")
+    def validate_strategies(cls, values: dict[str, typing.Any]) -> dict[str, typing.Any]:
+        """
+        Ensure that the scoring_llm is provided when using LLMBasedScoring.
+        """
+        if values.get('scoring_llm') is None:
+            raise ValueError('scoring_llm must be provided when scorer_type is set to LLM_BASED_SCORING.')
+
+        return values
+
+
+class LLMBasedAgentScoringConfig(TTCStrategyBaseConfig, name="llm_based_agent_scoring"):
+    """
+    Configuration for LLMBasedScoring.
+    """
+    scoring_llm: LLMRef | typing.Any | None = Field(
+        default=None,
+        description="The LLM to use for scoring the plans. This can be a callable or an instance of an LLM client.")
+
+    scoring_template: str = Field(
+        description="Prompt template to use for scoring the function output",
+        default=("You are an expert reasoning model tasked with scoring the following "
+                 "result of an agent system based on its input and objective. Judge"
+                 " the quality and relevance of the answer to score it.\n\n"
+                 "The agent system's objective is:\n{objective}\n\n"
+                 "It has been tasked with achieving the following goal: \n{input}\n\n"
+                 "The following output has been generated by the agent:\n\n{output}\n\n"
+                 "Score the result on a scale from 1 to 10, where 10 is the best. "
+                 "Return the final score as a floating point number preceded by `FINAL SCORE:` without any "
+                 "other text before or after it\n"),
+    )
+
+    @model_validator(mode="before")
+    def validate_strategies(cls, values: dict[str, typing.Any]) -> dict[str, typing.Any]:
+        """
+        Ensure that the scoring_llm is provided when using LLMBasedScoring.
+        """
+        if values.get('scoring_llm') is None:
+            raise ValueError('scoring_llm must be provided when scorer_type is set to LLM_BASED_SCORING.')
+
+        return values
+
+
+class MotivationAwareScoringConfig(TTCStrategyBaseConfig, name="motivation_aware_scoring"):
+    """
+    Configuration for a scoring strategy that considers both the original input (task)
+    and the motivation (from metadata) along with the current output.
+    """
+
+    scoring_llm: LLMRef | None = Field(
+        default=None, description="The LLM used to evaluate how well the output addresses the task plus motivation.")
+
+    scoring_template: str = Field(
+        default=("You are an expert at assessing the quality of an output in relation to its task and motivation.\n"
+                 "Task: {task}\n"
+                 "Motivation: {motivation}\n"
+                 "Output: {output}\n"
+                 "On a scale from 1 to 10 (10 being the best), how well does this output fulfill "
+                 "the original task in the context "
+                 "of the provided motivation? Note that the task might answer one part of a bigger question "
+                 "which should count as a satisfactory response and should not receive a lower score.\n"
+                 "Return the final score as a floating point number preceded by 'FINAL SCORE:'."),
+        description="The prompt template used to evaluate and score the output.")
+
+    @model_validator(mode="before")
+    def validate_scoring_llm(cls, values):
+        if values.get('scoring_llm') is None:
+            raise ValueError("A scoring_llm must be provided for motivation_aware_scoring.")
+        return values

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