aiqtoolkit 1.2.0rc4__py3-none-any.whl → 1.2rc9__py3-none-any.whl

aiq/experimental/inference_time_scaling/selection/llm_based_agent_output_selector.py

@@ -1,131 +0,0 @@
-# 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 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.selection_config import LLMBasedAgentOutputSelectionConfig
-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 LLMBasedAgentOutputSelector(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 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[ITSItem],
-                      original_prompt: str | None = None,
-                      agent_context: str | None = None,
-                      **kwargs) -> [ITSItem]:
-        """
-        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[ITSItem]): The list of planning items to select from.
-
-        Returns:
-            ITSItem: 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 aiqtoolkit-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_its_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

aiq/experimental/inference_time_scaling/selection/llm_based_output_merging_selector.py

@@ -1,159 +0,0 @@
-# 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 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.selection_config import LLMBasedOutputMergingConfig
-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 LLMBasedOutputMergingSelector(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 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[ITSItem],
-                      original_prompt: str | None = None,
-                      agent_context: str | None = None,
-                      **kwargs) -> [ITSItem]:
-        """
-        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[ITSItem]): The list of planning items to select from.
-
-        Returns:
-            ITSItem: 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 aiqtoolkit-langchain.")
-
-        from typing 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.")
-        else:
-            merged_output = merged_output
-
-        logger.info("Merged output: %s", str(merged_output))
-
-        # Create a new ITSItem with the merged plan or output
-        if self.pipeline_type == PipelineTypeEnum.PLANNING:
-            merged_item = ITSItem(input=items[0].input, output=merged_output, plan=merged_output)
-        else:
-            merged_item = ITSItem(input=items[0].input, output=merged_output)
-
-        return [merged_item]
-
-
-@register_its_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

aiq/experimental/inference_time_scaling/selection/llm_based_plan_selector.py

@@ -1,128 +0,0 @@
-# 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 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.selection_config import LLMBasedPlanSelectionConfig
-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 LLMBasedPlanSelector(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 selector.
-        """
-        self.llm_bound = await builder.get_llm(self.config.selection_llm, wrapper_type=LLMFrameworkEnum.LANGCHAIN)
-
-    def supported_pipeline_types(self) -> [PipelineTypeEnum]:
-        return [PipelineTypeEnum.PLANNING]
-
-    def stage_type(self) -> StageTypeEnum:
-        return StageTypeEnum.SELECTION
-
-    async def ainvoke(self,
-                      items: list[ITSItem],
-                      original_prompt: str | None = None,
-                      agent_context: str | None = None,
-                      **kwargs) -> [ITSItem]:
-        """
-        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[ITSItem]): The list of planning items to select from.
-
-        Returns:
-            ITSItem: 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 aiqtoolkit-langchain.")
-
-        if not isinstance(self.llm_bound, BaseChatModel):
-            raise ValueError("The `selection_llm` must be an instance of `BaseChatModel`.")
-
-        model: BaseChatModel = self.llm_bound
-
-        plans = ""
-        for idx, item in enumerate(items):
-            plans += f"{idx + 1}. {remove_r1_think_tags(item.plan)}\n"
-
-        prompt_template = PromptTemplate(
-            template=self.config.selection_template,
-            input_variables=["original_prompt", "context", "plans"],
-            validate_template=True,
-        )
-
-        prompt = (await prompt_template.ainvoke(input={
-            "original_prompt": original_prompt, "context": agent_context, "plans": plans
-        })).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 PLAN:\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 plan 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_its_strategy(config_type=LLMBasedPlanSelectionConfig)
-async def register_llm_based_plan_selection(config: LLMBasedPlanSelectionConfig, builder: Builder):
-    """
-    Register the LLMBasedPlanSelector with the provided configuration.
-    """
-    selector = LLMBasedPlanSelector(config)
-    await selector.build_components(Builder())
-    yield selector

aiq/experimental/inference_time_scaling/selection/threshold_selector.py

@@ -1,58 +0,0 @@
-# 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 aiq.builder.builder import Builder
-from aiq.cli.register_workflow import register_its_strategy
-from aiq.experimental.inference_time_scaling.models.its_item import ITSItem
-from aiq.experimental.inference_time_scaling.models.selection_config import ThresholdSelectionConfig
-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
-
-logger = logging.getLogger(__name__)
-
-
-class ThresholdSelector(StrategyBase):
-    """
-    Downselects only those ITSItems whose 'score' >= config.threshold.
-    """
-
-    async def build_components(self, builder: Builder) -> None:
-        # No special components needed
-        pass
-
-    def supported_pipeline_types(self) -> list[PipelineTypeEnum]:
-        return [PipelineTypeEnum.TOOL_USE]
-
-    def stage_type(self) -> StageTypeEnum:
-        return StageTypeEnum.SELECTION
-
-    async def ainvoke(self,
-                      items: list[ITSItem],
-                      original_prompt: str | None = None,
-                      agent_context: str | None = None,
-                      **kwargs) -> list[ITSItem]:
-        threshold = self.config.threshold
-        selected = [itm for itm in items if (itm.score is not None and itm.score >= threshold)]
-        logger.info("ThresholdSelector: %d items => %d items (threshold=%.1f)", len(items), len(selected), threshold)
-        return selected
-
-
-@register_its_strategy(config_type=ThresholdSelectionConfig)
-async def register_threshold_selector(config: ThresholdSelectionConfig, builder: Builder):
-    selector = ThresholdSelector(config)
-    yield selector

aiq/front_ends/__init__.py

@@ -1,14 +0,0 @@
-# SPDX-FileCopyrightText: Copyright (c) 2024-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.

aiq/front_ends/console/__init__.py

@@ -1,14 +0,0 @@
-# SPDX-FileCopyrightText: Copyright (c) 2024-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.