nvidia-nat 1.1.0a20251020__py3-none-any.whl

nat/experimental/test_time_compute/search/multi_query_retrieval_search.py

@@ -0,0 +1,122 @@
+# 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 nat.builder.builder import Builder
+from nat.builder.framework_enum import LLMFrameworkEnum
+from nat.cli.register_workflow import register_ttc_strategy
+from nat.experimental.test_time_compute.models.search_config import MultiQueryRetrievalSearchConfig
+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 MultiQueryRetrievalSearch(StrategyBase):
+    """
+    A strategy that, for each incoming TTCItem, generates multiple new items by
+    re-writing the input 'task_description' from different perspectives.
+    Uses multiple LLMs to encourage diversity.
+    """
+
+    def __init__(self, config: MultiQueryRetrievalSearchConfig) -> None:
+        super().__init__(config)
+        self.config = config
+        self.llms_bound = []
+
+    async def build_components(self, builder: Builder) -> None:
+        """
+        Binds each LLMRef in self.config.llms to an actual LLM client.
+        """
+        self.llms_bound = []
+        for llm_ref in self.config.llms:
+            bound_llm = await builder.get_llm(llm_ref, wrapper_type=LLMFrameworkEnum.LANGCHAIN)
+            self.llms_bound.append(bound_llm)
+
+    def supported_pipeline_types(self) -> list[PipelineTypeEnum]:
+        return [PipelineTypeEnum.TOOL_USE]
+
+    def stage_type(self) -> StageTypeEnum:
+        return StageTypeEnum.SEARCH
+
+    async def ainvoke(self,
+                      items: list[TTCItem],
+                      original_prompt: str | None = None,
+                      agent_context: str | None = None,
+                      **kwargs) -> list[TTCItem]:
+        """
+        For each TTCItem, rewrite the 'input' using each LLM to create a new perspective.
+        The new TTCItems' 'output' field will store the newly generated query.
+        """
+        try:
+            from langchain_core.prompts import PromptTemplate
+        except ImportError:
+            raise ImportError("langchain-core is required for MultiQueryRetrievalSearch. "
+                              "Install nvidia-nat-langchain or similar.")
+
+        new_ttc_items: list[TTCItem] = []
+
+        # Create a single PromptTemplate object for rewriting the query
+        template_vars = ["task", "motivation"]
+        query_template = PromptTemplate(template=self.config.query_generation_template,
+                                        input_variables=template_vars,
+                                        validate_template=True)
+
+        for item in items:
+            original_task = str(item.input) or ""
+            motivation = str(item.metadata) if item.metadata else ""
+            new_ttc_items.append(
+                TTCItem(
+                    input=item.input,
+                    output=item.input,
+                    metadata=item.metadata,
+                    name=item.name,  # keep the original tool name
+                ))
+
+            for llm in self.llms_bound:
+                prompt_str = (await query_template.ainvoke({
+                    "task": original_task, "motivation": motivation
+                })).to_string()
+
+                # We'll call each LLM to produce a new query
+                response = await llm.ainvoke(prompt_str)
+                cleaned = remove_r1_think_tags(response.content if hasattr(response, 'content') else str(response))
+                cleaned = cleaned.strip()
+
+                # Create a new TTCItem for each newly generated query
+                new_item = TTCItem(
+                    input=item.input,  # keep the original input for reference
+                    output=cleaned,  # store the newly generated query in the output
+                    metadata=item.metadata,
+                    name=item.name,  # same tool name or optional new name
+                )
+                new_ttc_items.append(new_item)
+
+        logger.info("MultiQueryRetrievalSearch produced %d new items from %d original items.",
+                    len(new_ttc_items),
+                    len(items))
+
+        return new_ttc_items
+
+
+@register_ttc_strategy(config_type=MultiQueryRetrievalSearchConfig)
+async def register_multi_query_retrieval_search(config: MultiQueryRetrievalSearchConfig, builder: Builder):
+    strategy = MultiQueryRetrievalSearch(config)
+    await strategy.build_components(builder)
+    yield strategy

nat/experimental/test_time_compute/search/single_shot_multi_plan_planner.py

@@ -0,0 +1,128 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import asyncio
+import logging
+import re
+
+from nat.builder.builder import Builder
+from nat.builder.framework_enum import LLMFrameworkEnum
+from nat.cli.register_workflow import register_ttc_strategy
+from nat.data_models.ttc_strategy import TTCStrategyBaseConfig
+from nat.experimental.test_time_compute.models.search_config import SingleShotMultiPlanConfig
+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 SingleShotMultiPlanPlanner(StrategyBase):
+    """
+    Implementation of the Single Shot Multi Plan Planner.
+    This planner generates multiple plans in a single shot.
+    """
+
+    def __init__(self, config: TTCStrategyBaseConfig) -> None:
+        super().__init__(config)
+        self.llm_bound = None
+
+    async def build_components(self, builder: Builder) -> None:
+        self.llm_bound = await builder.get_llm(self.config.planning_llm, wrapper_type=LLMFrameworkEnum.LANGCHAIN)
+
+    def supported_pipeline_types(self) -> [PipelineTypeEnum]:
+        return [PipelineTypeEnum.PLANNING]
+
+    def stage_type(self) -> StageTypeEnum:
+        return StageTypeEnum.SEARCH
+
+    async def ainvoke(self,
+                      items: list[TTCItem],
+                      original_prompt: str | None = None,
+                      agent_context: str | None = None,
+                      **kwargs) -> list[TTCItem]:
+        """
+        Generate a TTCItem based on the provided prompt.
+        """
+
+        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 resolve by installing nvidia-nat-langchain.")
+
+        planning_template = PromptTemplate(template=self.config.planning_template,
+                                           input_variables=["context", "prompt"],
+                                           validate_template=True)
+        prompt = (await planning_template.ainvoke(input={
+            "context": agent_context, "prompt": original_prompt
+        })).to_string()
+
+        # assert self.config.planning llm is a BaseChatModel
+        if not isinstance(self.llm_bound, BaseChatModel):
+            raise ValueError("The `planning_llm` must be an instance of `BaseChatModel`.")
+
+        model: BaseChatModel = self.llm_bound
+
+        async def generate_plan(llm: BaseChatModel, plan_prompt: str, temperature: float) -> TTCItem:
+            """
+            Helper function to generate a plan using the provided prompt and temperature.
+            """
+            llm_bound = llm.bind(temperature=temperature)
+            response = await llm_bound.ainvoke(plan_prompt)
+            cleaned = remove_r1_think_tags(response.content if hasattr(response, 'content') else str(response))
+
+            # Plan will be the string following 'PLAN:'. Use Regex tpo extract
+            cleaned = re.sub(r'(?i)^\s*PLAN:\s*', '', cleaned).strip()
+
+            if not cleaned:
+                logger.warning(f"No plan generated for the prompt: {plan_prompt}.")
+                # Return an empty PlanningItem to avoid breaking the generation loop
+                return TTCItem(plan="Plan was not generated")
+
+            return TTCItem(plan=cleaned)
+
+        # Define a list of temperatures based on min and max temperature in the config and number of plans to generate
+        temperatures = [
+            self.config.min_temperature + (i / (self.config.num_plans - 1)) *
+            (self.config.max_temperature - self.config.min_temperature) for i in range(self.config.num_plans)
+        ]
+
+        # Generate plans using the defined temperatures in parallel using asyncio
+        tasks = [generate_plan(model, prompt, temperature) for temperature in temperatures]
+        # Run the tasks concurrently and gather results
+        plans = await asyncio.gather(*tasks)
+
+        if not plans:
+            raise ValueError("No plans were generated. Please check the LLM response.")
+
+        logger.info("Generated %d plans from the SingleShotMultiPlanPlanner", self.config.num_plans)
+
+        logger.debug("Generated plans: %s", [plan.dict() for plan in plans])
+
+        return plans
+
+
+@register_ttc_strategy(config_type=SingleShotMultiPlanConfig)
+async def register_single_shot_multi_plan_planner(config: SingleShotMultiPlanConfig, builder: Builder):
+    """
+    Register the SingleShotMultiPlanPlanner strategy with the provided configuration.
+    """
+    planner = SingleShotMultiPlanPlanner(config)
+    await planner.build_components(builder)
+    yield planner

nat/experimental/test_time_compute/selection/best_of_n_selector.py

@@ -0,0 +1,63 @@
+# 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 nat.builder.builder import Builder
+from nat.cli.register_workflow import register_ttc_strategy
+from nat.experimental.test_time_compute.models.selection_config import BestOfNSelectionConfig
+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
+
+logger = logging.getLogger(__name__)
+
+
+class BestOfNSelector(StrategyBase):
+
+    async def build_components(self, builder: Builder) -> None:
+        pass
+
+    def supported_pipeline_types(self) -> [PipelineTypeEnum]:
+        return [PipelineTypeEnum.PLANNING, PipelineTypeEnum.AGENT_EXECUTION, PipelineTypeEnum.TOOL_USE]
+
+    def stage_type(self) -> StageTypeEnum:
+        return StageTypeEnum.SELECTION
+
+    async def ainvoke(self,
+                      items: list[TTCItem],
+                      original_prompt: str | None = None,
+                      agent_context: str | None = None,
+                      **kwargs) -> [TTCItem]:
+
+        # Assert that every planning item has a non NoneType score
+        for item in items:
+            if item.score is None:
+                raise ValueError("Every planning item must have a score. Did you use a scorer before this?")
+
+        # Pick the planning item with the highest score
+        best_item = max(items, key=lambda x: x.score)
+
+        return [best_item]
+
+
+@register_ttc_strategy(config_type=BestOfNSelectionConfig)
+async def register_best_of_n_selector(config: BestOfNSelectionConfig, builder: Builder):
+    """
+    Register the BestOfNSelector strategy.
+    """
+    selector = BestOfNSelector(config)
+    yield selector

nat/experimental/test_time_compute/selection/llm_based_agent_output_selector.py

@@ -0,0 +1,131 @@
+# 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
+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.selection_config import LLMBasedAgentOutputSelectionConfig
+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 LLMBasedAgentOutputSelector(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 selector.
+        """
+        self.llm_bound = await builder.get_llm(self.config.selection_llm, wrapper_type=LLMFrameworkEnum.LANGCHAIN)
+
+    def supported_pipeline_types(self) -> [PipelineTypeEnum]:
+        return [PipelineTypeEnum.AGENT_EXECUTION]
+
+    def stage_type(self) -> StageTypeEnum:
+        return StageTypeEnum.SELECTION
+
+    async def ainvoke(self,
+                      items: list[TTCItem],
+                      original_prompt: str | None = None,
+                      agent_context: str | None = None,
+                      **kwargs) -> [TTCItem]:
+        """
+        Select the planning items based on the selection strategy.
+
+        Args:
+            original_prompt (str): The prompt the user provided the agent.
+            agent_context (str): The context of the agent, if applicable.
+            items (list[TTCItem]): The list of planning items to select from.
+
+        Returns:
+            TTCItem: The selected planning 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.")
+
+        from pydantic import BaseModel
+
+        if not isinstance(self.llm_bound, BaseChatModel):
+            raise ValueError("The `selection_llm` must be an instance of `BaseChatModel`.")
+
+        model: BaseChatModel = self.llm_bound
+
+        results = ""
+        for idx, item in enumerate(items):
+            item_str = str(item.output.model_dump()) if isinstance(item.output, BaseModel) else str(item.output)
+            results += f"{idx + 1}. {remove_r1_think_tags(item_str)}\n\n"
+
+        prompt_template = PromptTemplate(
+            template=self.config.selection_template,
+            input_variables=["objective", "input", "results"],
+            validate_template=True,
+        )
+
+        prompt = (await prompt_template.ainvoke(input={
+            "objective": agent_context, "input": original_prompt, "results": results
+        })).to_string()
+
+        selected_plan_index = remove_r1_think_tags((await model.ainvoke(prompt)).content)
+
+        # Model Response will be 'Plan {plan number}'
+        # Use RegEx to extrac Plan {idx} from response strong
+        if not isinstance(selected_plan_index, str):
+            logger.warning(f"Invalid response from LLM for selected plan index: {selected_plan_index}.")
+            raise ValueError("Unable to parse the selected plan index.")
+        selected_plan_index = selected_plan_index.strip()
+        match = re.match(r'^\s*SELECTED ITEM:\s+(\d+)', selected_plan_index)
+        if not match:
+            logger.warning(f"Could not parse the selected plan index from the response: {selected_plan_index}.")
+            raise ValueError("The response format for selecting the item is incorrect.")
+        index = match.group(1)
+
+        try:
+            selected_index = int(index) - 1
+            if selected_index < 0 or selected_index >= len(items):
+                raise ValueError("Selected index is out of range.")
+
+            # Return the selected planning item
+            return [items[selected_index]]
+        except ValueError as e:
+            logger.warning(f"Error parsing the selected plan index: {index}. Exception: {str(e)}")
+            raise ValueError(f"Failed to parse the selected plan index from the LLM response: {selected_plan_index}. "
+                             "Ensure the response follows the expected format.") from e
+
+
+@register_ttc_strategy(config_type=LLMBasedAgentOutputSelectionConfig)
+async def register_llm_based_agent_output_selector(config: LLMBasedAgentOutputSelectionConfig, builder: Builder):
+    """
+    Register the LLMBasedAgentOutputSelector with the builder.
+    """
+    selector = LLMBasedAgentOutputSelector(config)
+    await selector.build_components(builder)
+    yield selector

nat/experimental/test_time_compute/selection/llm_based_output_merging_selector.py

@@ -0,0 +1,157 @@
+# 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 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.selection_config import LLMBasedOutputMergingConfig
+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 LLMBasedOutputMergingSelector(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 selector.
+        """
+        self.llm_bound = await builder.get_llm(self.config.selection_llm, wrapper_type=LLMFrameworkEnum.LANGCHAIN)
+
+    def supported_pipeline_types(self) -> [PipelineTypeEnum]:
+        return [PipelineTypeEnum.AGENT_EXECUTION, PipelineTypeEnum.PLANNING]
+
+    def stage_type(self) -> StageTypeEnum:
+        return StageTypeEnum.SELECTION
+
+    async def ainvoke(self,
+                      items: list[TTCItem],
+                      original_prompt: str | None = None,
+                      agent_context: str | None = None,
+                      **kwargs) -> [TTCItem]:
+        """
+        Merge the outputs of multiple planning items into a single output
+
+        Args:
+            original_prompt (str): The prompt the user provided the agent.
+            agent_context (str): The context of the agent, if applicable.
+            items (list[TTCItem]): The list of planning items to select from.
+
+        Returns:
+            TTCItem: The selected planning 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.")
+
+        from collections.abc import Callable
+
+        from pydantic import BaseModel
+
+        if not isinstance(self.llm_bound, BaseChatModel):
+            raise ValueError("The `selection_llm` must be an instance of `BaseChatModel`.")
+
+        if not self.pipeline_type:
+            raise RuntimeError("Pipeline type is not set. Ensure that the pipeline "
+                               "type is set before invoking the selector.")
+
+        model: BaseChatModel = self.llm_bound
+
+        results = ""
+        if self.pipeline_type == PipelineTypeEnum.AGENT_EXECUTION:
+            for idx, item in enumerate(items):
+                item_str = str(item.output.model_dump()) if isinstance(item.output, BaseModel) else str(item.output)
+                results += f"{idx + 1}. {remove_r1_think_tags(item_str)}\n\n"
+        else:
+            for idx, item in enumerate(items):
+                item_str = str(item.plan)
+                results += f"{idx + 1}. {remove_r1_think_tags(item_str)}\n\n"
+
+        prompt_template = PromptTemplate(
+            template=self.config.selection_template,
+            input_variables=["pipeline_type", "objective", "input", "results"],
+            validate_template=True,
+        )
+
+        if self.pipeline_type == PipelineTypeEnum.PLANNING:
+            pipeline_objective = "execution plans for a given objective and input."
+        else:
+            pipeline_objective = "outputs from an agent system based on the provided objective and input."
+
+        prompt = (await prompt_template.ainvoke(
+            input={
+                "objective": agent_context,
+                "input": original_prompt,
+                "results": results,
+                "pipeline_type": pipeline_objective
+            })).to_string()
+
+        merged_output = remove_r1_think_tags((await model.ainvoke(prompt)).content)
+
+        if not isinstance(merged_output, str):
+            logger.warning(f"Invalid response from LLM for merged_plan: {merged_output}.")
+            raise ValueError("Unable to parse merged plan.")
+        merged_output = merged_output.strip()
+
+        # match = split the string after 'MERGED OUTPUT:'
+        matches = merged_output.split("MERGED OUTPUT:")
+        if len(matches) > 1:
+            merged_output = matches[-1].strip()
+        else:
+            raise ValueError("Merged output does not contain 'MERGED OUTPUT:' prefix.")
+
+        # Check if a callable argument is provided in kwargs called output_parser
+        output_parser: Callable | None = kwargs.get('output_parser', None)
+        if output_parser:
+            try:
+                merged_output = output_parser(merged_output)
+            except Exception as e:
+                logger.error(f"Error parsing merged output: {e}")
+                raise ValueError("Failed to parse merged output.")
+
+        logger.info("Merged output: %s", str(merged_output))
+
+        # Create a new TTCItem with the merged plan or output
+        if self.pipeline_type == PipelineTypeEnum.PLANNING:
+            merged_item = TTCItem(input=items[0].input, output=merged_output, plan=merged_output)
+        else:
+            merged_item = TTCItem(input=items[0].input, output=merged_output)
+
+        return [merged_item]
+
+
+@register_ttc_strategy(config_type=LLMBasedOutputMergingConfig)
+async def register_llm_based_output_merging_selector(config: LLMBasedOutputMergingConfig, builder: Builder):
+    """
+    Register the LLMBasedOutputMergingSelector with the builder.
+    """
+    selector = LLMBasedOutputMergingSelector(config)
+    await selector.build_components(builder)
+    yield selector