aiqtoolkit 1.2.0a20250706__py3-none-any.whl → 1.2.0a20250730__py3-none-any.whl

aiq/experimental/inference_time_scaling/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 aiq.data_models.component_ref import LLMRef
+from aiq.data_models.its_strategy import ITSStrategyBaseConfig
+
+
+class SingleShotMultiPlanConfig(ITSStrategyBaseConfig, 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(ITSStrategyBaseConfig, 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(ITSStrategyBaseConfig, name="multi_query_retrieval_search"):
+    """
+    Configuration for the MultiQueryRetrievalSearch strategy.
+    This strategy generates multiple new 'ITSItem'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

aiq/experimental/inference_time_scaling/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.its_strategy import ITSStrategyBaseConfig
+
+
+class LLMBasedPlanSelectionConfig(ITSStrategyBaseConfig, 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(ITSStrategyBaseConfig, 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(ITSStrategyBaseConfig, 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(ITSStrategyBaseConfig, 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 ITSItems with score >= this value.")
+
+
+class BestOfNSelectionConfig(ITSStrategyBaseConfig, name="best_of_n_selection"):
+    """
+    Configuration for Best of N Selection
+    """
+    pass

aiq/experimental/inference_time_scaling/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/inference_time_scaling/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.inference_time_scaling.models.its_item import ITSItem
+from aiq.experimental.inference_time_scaling.models.stage_enums import StageTypeEnum, PipelineTypeEnum
+from aiq.data_models.its_strategy import ITSStrategyBaseConfig
+
+
+class StrategyBase(ABC):
+    """
+    Abstract base class for strategy implementations.
+
+    This class defines the interface for strategies that can be used in the
+    Inference Time Scaling (ITS) framework. Concrete strategy classes should
+    implement the methods defined in this class.
+    """
+
+    def __init__(self, config: ITSStrategyBaseConfig) -> None:
+        self.config: ITSStrategyBaseConfig = 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[ITSItem],
+                      original_prompt: str | None = None,
+                      agent_context: str | None = None,
+                      **kwargs) -> [ITSItem]:
+        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/inference_time_scaling/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/inference_time_scaling/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/inference_time_scaling/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_its_strategy
+from aiq.data_models.its_strategy import ITSStrategyBaseConfig
+from aiq.experimental.inference_time_scaling.models.its_item import ITSItem
+from aiq.experimental.inference_time_scaling.models.scoring_config import LLMBasedAgentScoringConfig
+from aiq.experimental.inference_time_scaling.models.stage_enums import PipelineTypeEnum
+from aiq.experimental.inference_time_scaling.models.stage_enums import StageTypeEnum
+from aiq.experimental.inference_time_scaling.models.strategy_base import StrategyBase
+from aiq.utils.io.model_processing import remove_r1_think_tags
+
+logger = logging.getLogger(__name__)
+
+
+class LLMBasedAgentScorer(StrategyBase):
+
+    def __init__(self, config: ITSStrategyBaseConfig) -> 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: ITSItem) -> float:
+        """
+        Score a single planning item using the LLM.
+
+        Args:
+            original_prompt (str): The original prompt.
+            agent_context (str): The agent context.
+            item (ITSItem): 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[ITSItem],
+                      original_prompt: str | None = None,
+                      agent_context: str | None = None,
+                      **kwargs) -> list[ITSItem]:
+        """
+        Score a list of planning items.
+
+        Args:
+            original_prompt (str): The original prompt.
+            agent_context (str): The agent context.
+            items (list[ITSItem]): 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_its_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