aiqtoolkit 1.2.0rc3__py3-none-any.whl → 1.2.0rc5__py3-none-any.whl

aiq/experimental/inference_time_scaling/scoring/motivation_aware_scorer.py

@@ -1,111 +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 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 MotivationAwareScoringConfig
-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 MotivationAwareScorer(StrategyBase):
-    """
-    A strategy that scores an ITSItem's output based on how well it
-    addresses both the original input (task) and the 'motivation' from metadata.
-    """
-
-    def __init__(self, config: ITSStrategyBaseConfig) -> 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.scoring_llm, wrapper_type=LLMFrameworkEnum.LANGCHAIN)
-
-    def supported_pipeline_types(self) -> list[PipelineTypeEnum]:
-        return [PipelineTypeEnum.TOOL_USE]
-
-    def stage_type(self) -> StageTypeEnum:
-        return StageTypeEnum.SCORING
-
-    async def ainvoke(self,
-                      items: list[ITSItem],
-                      original_prompt: str | None = None,
-                      agent_context: str | None = None,
-                      **kwargs) -> list[ITSItem]:
-        """
-        Scores each item by combining the original 'task_description' and 'motivation' with the 'output'.
-        The resulting score is stored in item.score.
-        """
-        from langchain_core.language_models import BaseChatModel
-        from langchain_core.prompts import PromptTemplate
-
-        if not isinstance(self.llm_bound, BaseChatModel):
-            raise ValueError("scoring_llm must be a BaseChatModel instance for MotivationAwareScorer.")
-
-        scoring_model: BaseChatModel = self.llm_bound
-
-        scoring_template = PromptTemplate(template=self.config.scoring_template,
-                                          input_variables=["task", "motivation", "output"],
-                                          validate_template=True)
-
-        async def score_item(item: ITSItem) -> float:
-            task_str = str(item.input) or ""
-            motivation_str = str(item.metadata) if item.metadata else ""
-            output_str = str(item.output) or ""
-
-            prompt = (await scoring_template.ainvoke({
-                "task": task_str, "motivation": motivation_str, "output": output_str
-            })).to_string()
-
-            response = (await scoring_model.ainvoke(prompt)).content
-            response = remove_r1_think_tags(response or "")
-
-            match = re.search(r'FINAL SCORE:\s*([\d.]+)', response)
-            if not match:
-                logger.warning(f"Could not parse score from response: {response}")
-                return 0.0
-
-            score_str = match.group(1)
-            try:
-                return float(score_str)
-            except ValueError:
-                logger.warning(f"Could not convert score '{score_str}' to float.")
-                return 0.0
-
-        tasks = [score_item(item) for item in items]
-        scores = await asyncio.gather(*tasks)
-
-        for i, s in enumerate(scores):
-            items[i].score = s
-
-        return items
-
-
-@register_its_strategy(config_type=MotivationAwareScoringConfig)
-async def register_motivation_aware_scorer(config: MotivationAwareScoringConfig, builder: Builder):
-    scorer = MotivationAwareScorer(config)
-    await scorer.build_components(builder)
-    yield scorer

aiq/experimental/inference_time_scaling/search/multi_llm_planner.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 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.experimental.inference_time_scaling.models.its_item import ITSItem
-from aiq.experimental.inference_time_scaling.models.search_config import MultiLLMPlanConfig
-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 MultiLLMPlanner(StrategyBase):
-    """
-    A planner that uses multiple LLMs to generate plans. Each LLM can generate
-    a specified number of plans, and all plans are combined.
-    """
-
-    def __init__(self, config: MultiLLMPlanConfig) -> None:
-        super().__init__(config)
-        self.config = config
-        self.llms_bound = []  # Will hold the "bound" LLMs after build_components
-
-    async def build_components(self, builder: Builder) -> None:
-        """
-        Build the components required for this multi-LLM planner.
-        Binds each LLMRef from the config with the selected framework wrapper (LANGCHAIN).
-        """
-        logger.debug("Building components for MultiLLMPlanner")
-        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) -> [PipelineTypeEnum]:
-        return [PipelineTypeEnum.PLANNING]
-
-    def stage_type(self) -> StageTypeEnum:
-        return StageTypeEnum.SEARCH
-
-    async def _generate_plan_for_temperature(self, llm, base_prompt: str, temperature: float) -> ITSItem:
-        bound_llm = llm.bind(temperature=temperature)
-        response = await bound_llm.ainvoke(base_prompt)
-        cleaned = remove_r1_think_tags(response.content if hasattr(response, 'content') else str(response))
-        # The plan is expected to start with "PLAN:" and all the text after it is the plan
-        cleaned = re.sub(r'(?i)^\s*PLAN:\s*', '', cleaned).strip()
-
-        if not cleaned:
-            logger.warning(f"No plan generated for the prompt: {base_prompt}.")
-            # Return an empty PlanningItem to avoid breaking the generation loop
-            return ITSItem(plan="Plan was not generated")
-
-        return ITSItem(plan=cleaned)
-
-    async def _generate_plans_for_llm(self, llm, base_prompt: str) -> list[ITSItem]:
-        if self.config.plans_per_llm == 1:
-            temps = [self.config.min_temperature]
-        else:
-            temps = [
-                self.config.min_temperature + (i / (self.config.plans_per_llm - 1)) *
-                (self.config.max_temperature - self.config.min_temperature) for i in range(self.config.plans_per_llm)
-            ]
-        tasks = [self._generate_plan_for_temperature(llm, base_prompt, temp) for temp in temps]
-        return await asyncio.gather(*tasks)
-
-    async def ainvoke(self,
-                      items: list[ITSItem],
-                      original_prompt: str | None = None,
-                      agent_context: str | None = None,
-                      **kwargs) -> list[ITSItem]:
-        """
-        Generate a list of PlanningItems by querying each LLM in self.llms_bound.
-        Each LLM produces 'plans_per_llm' plans.
-        """
-        try:
-            from langchain_core.prompts import PromptTemplate
-        except ImportError:
-            raise ImportError("langchain-core is not installed. Please install it to use MultiLLMPlanner.\n"
-                              "This error can be resolve by installing aiqtoolkit-langchain.")
-
-        # Create a single PromptTemplate
-        planning_template = PromptTemplate(template=self.config.planning_template,
-                                           input_variables=["context", "prompt"],
-                                           validate_template=True)
-
-        # Format the prompt once
-        base_prompt = (await planning_template.ainvoke({
-            "context": agent_context, "prompt": original_prompt
-        })).to_string()
-
-        # Launch generation for each llm concurrently using the new helper method
-        tasks = [self._generate_plans_for_llm(llm, base_prompt) for llm in self.llms_bound]
-        results_nested = await asyncio.gather(*tasks)
-
-        # Flatten the nested lists of ITSItem
-        all_plans: list[ITSItem] = [p for sub in results_nested for p in sub]
-        logger.info("MultiLLMPlanner generated %d plans total.", len(all_plans))
-        return all_plans
-
-
-@register_its_strategy(config_type=MultiLLMPlanConfig)
-async def register_multi_llm_planner(config: MultiLLMPlanConfig, builder: Builder):
-    """
-    Register the MultiLLMPlanner strategy with the provided configuration.
-    """
-    planner = MultiLLMPlanner(config)
-    await planner.build_components(builder)
-    yield planner

aiq/experimental/inference_time_scaling/search/multi_query_retrieval_search.py

@@ -1,122 +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.experimental.inference_time_scaling.models.its_item import ITSItem
-from aiq.experimental.inference_time_scaling.models.search_config import MultiQueryRetrievalSearchConfig
-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 MultiQueryRetrievalSearch(StrategyBase):
-    """
-    A strategy that, for each incoming ITSItem, 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[ITSItem],
-                      original_prompt: str | None = None,
-                      agent_context: str | None = None,
-                      **kwargs) -> list[ITSItem]:
-        """
-        For each ITSItem, rewrite the 'input' using each LLM to create a new perspective.
-        The new ITSItems' '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 aiqtoolkit-langchain or similar.")
-
-        new_its_items: list[ITSItem] = []
-
-        # 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_its_items.append(
-                ITSItem(
-                    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 ITSItem for each newly generated query
-                new_item = ITSItem(
-                    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_its_items.append(new_item)
-
-        logger.info("MultiQueryRetrievalSearch produced %d new items from %d original items.",
-                    len(new_its_items),
-                    len(items))
-
-        return new_its_items
-
-
-@register_its_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

aiq/experimental/inference_time_scaling/search/single_shot_multi_plan_planner.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 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.search_config import SingleShotMultiPlanConfig
-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 SingleShotMultiPlanPlanner(StrategyBase):
-    """
-    Implementation of the Single Shot Multi Plan Planner.
-    This planner generates multiple plans in a single shot.
-    """
-
-    def __init__(self, config: ITSStrategyBaseConfig) -> 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[ITSItem],
-                      original_prompt: str | None = None,
-                      agent_context: str | None = None,
-                      **kwargs) -> list[ITSItem]:
-        """
-        Generate a ITSItem 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 aiqtoolkit-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) -> ITSItem:
-            """
-            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 ITSItem(plan="Plan was not generated")
-
-            return ITSItem(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_its_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

aiq/experimental/inference_time_scaling/selection/best_of_n_selector.py

@@ -1,63 +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 BestOfNSelectionConfig
-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 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[ITSItem],
-                      original_prompt: str | None = None,
-                      agent_context: str | None = None,
-                      **kwargs) -> [ITSItem]:
-
-        # 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_its_strategy(config_type=BestOfNSelectionConfig)
-async def register_best_of_n_selector(config: BestOfNSelectionConfig, builder: Builder):
-    """
-    Register the BestOfNSelector strategy.
-    """
-    selector = BestOfNSelector(config)
-    yield selector