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

aiq/experimental/inference_time_scaling/editing/motivation_aware_summarization.py

@@ -1,107 +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.editor_config import MotivationAwareSummarizationConfig
-from aiq.experimental.inference_time_scaling.models.its_item import ITSItem
-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 MotivationAwareSummarization(StrategyBase):
-    """
-    A strategy that, for each incoming ITSItem, summarizes the output based on input
-    and motivation.
-    """
-
-    def __init__(self, config: MotivationAwareSummarizationConfig) -> None:
-        super().__init__(config)
-        self.config = config
-        self.llm_bound = None
-
-    async def build_components(self, builder: Builder) -> None:
-        """
-        Binds each LLMRef in self.config.llms to an actual LLM client.
-        """
-        bound_llm = await builder.get_llm(self.config.editor_llm, wrapper_type=LLMFrameworkEnum.LANGCHAIN)
-        self.llm_bound = bound_llm
-
-    def supported_pipeline_types(self) -> list[PipelineTypeEnum]:
-        return [PipelineTypeEnum.TOOL_USE]
-
-    def stage_type(self) -> StageTypeEnum:
-        return StageTypeEnum.EDITING
-
-    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", "output"]
-        query_template = PromptTemplate(template=self.config.editor_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 ""
-            output = str(item.output) if item.output else ""
-
-            prompt = await (query_template.ainvoke(input={
-                "task": original_task, "motivation": motivation, "output": output
-            }))
-
-            llm_response = await self.llm_bound.ainvoke(prompt.to_string())
-            llm_response = remove_r1_think_tags(llm_response.content)
-
-            logger.info("LLM response from summarization: %s", llm_response)
-
-            new_its_items.append(
-                ITSItem(
-                    input=item.input,
-                    output=remove_r1_think_tags(llm_response),
-                    metadata=item.metadata,
-                    name=item.name,  # keep the original tool name
-                ))
-
-        return new_its_items
-
-
-@register_its_strategy(config_type=MotivationAwareSummarizationConfig)
-async def register_multi_query_retrieval_search(config: MotivationAwareSummarizationConfig, builder: Builder):
-    strategy = MotivationAwareSummarization(config)
-    await strategy.build_components(builder)
-    yield strategy

aiq/experimental/inference_time_scaling/functions/execute_score_select_function.py

@@ -1,105 +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 pydantic import Field
-
-from aiq.builder.builder import Builder
-from aiq.builder.function import Function
-from aiq.builder.function_info import FunctionInfo
-from aiq.cli.register_workflow import register_function
-from aiq.data_models.component_ref import FunctionRef
-from aiq.data_models.component_ref import ITSStrategyRef
-from aiq.data_models.function import FunctionBaseConfig
-from aiq.experimental.inference_time_scaling.models.its_item import ITSItem
-from aiq.experimental.inference_time_scaling.models.stage_enums import PipelineTypeEnum
-from aiq.experimental.inference_time_scaling.models.stage_enums import StageTypeEnum
-
-logger = logging.getLogger(__name__)
-
-
-class ExecuteScoreSelectFunctionConfig(FunctionBaseConfig, name="execute_score_select_function"):
-    scorer: ITSStrategyRef | None = Field(description="Strategy to score the output of the function", default=None)
-    selector: ITSStrategyRef = Field(description="Strategy to select the best output of the function")
-    augmented_fn: FunctionRef = Field(description="Function that will be executed")
-
-    num_executions: int = Field(3, description="Number of times to execute the function")
-
-
-@register_function(config_type=ExecuteScoreSelectFunctionConfig)
-async def execute_score_select_function(config: ExecuteScoreSelectFunctionConfig, builder: Builder):
-    import asyncio
-    import warnings
-
-    from pydantic import BaseModel
-
-    executable_fn: Function = builder.get_function(name=config.augmented_fn)
-
-    if config.scorer:
-        scorer = await builder.get_its_strategy(strategy_name=config.scorer,
-                                                pipeline_type=PipelineTypeEnum.AGENT_EXECUTION,
-                                                stage_type=StageTypeEnum.SCORING)
-    else:
-        scorer = None
-
-    selector = await builder.get_its_strategy(strategy_name=config.selector,
-                                              pipeline_type=PipelineTypeEnum.AGENT_EXECUTION,
-                                              stage_type=StageTypeEnum.SELECTION)
-
-    if executable_fn.has_streaming_output:
-        warnings.warn("Streaming output is not supported for this function. "
-                      "The function will be executed in non-streaming mode.")
-
-    def convert_to_str(arg):
-        if isinstance(arg, BaseModel):
-            return str(arg.model_dump())
-        return str(arg)
-
-    async def execute_fn(input_msg: executable_fn.input_type) -> executable_fn.single_output_type:
-
-        logger.info("Executing function %d times", config.num_executions)
-        tasks = [executable_fn.ainvoke(input_msg) for _ in range(config.num_executions)]
-        results = await asyncio.gather(*tasks)
-
-        input_str = convert_to_str(input_msg)
-        function_outputs = [convert_to_str(out) for out in results]
-        its_items = [ITSItem(
-            input=input_str,
-            output=out,
-        ) for out in function_outputs]
-
-        if scorer:
-            logger.info("Beginning scoring")
-            its_items = await scorer.ainvoke(items=its_items)
-
-        logger.info("Beginning selection")
-        selected_item = (await selector.ainvoke(items=its_items, original_prompt=its_items[0].input))[0]
-
-        # Find the index of selected item in its_items by matching the output
-        selected_output = selected_item.output
-        selected_index = -1
-        for i, item in enumerate(its_items):
-            if item.output == selected_output:
-                selected_index = i
-                break
-
-        return results[selected_index] if selected_index != -1 else selected_output
-
-    yield FunctionInfo.from_fn(
-        fn=execute_fn,
-        description=("This function executes a given function multiple times, scores the outputs, "
-                     "and selects the best output based on the specified scoring and selection strategies."),
-    )

aiq/experimental/inference_time_scaling/functions/its_tool_orchestration_function.py

@@ -1,205 +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
-
-from pydantic import Field
-
-from aiq.builder.builder import Builder
-from aiq.builder.framework_enum import LLMFrameworkEnum
-from aiq.builder.function_info import FunctionInfo
-from aiq.cli.register_workflow import register_function
-from aiq.data_models.component_ref import FunctionRef
-from aiq.data_models.component_ref import ITSStrategyRef
-from aiq.data_models.function import FunctionBaseConfig
-from aiq.experimental.inference_time_scaling.models.its_item import ITSItem
-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.tool_use_config import ToolUseInputSchema
-from aiq.experimental.inference_time_scaling.models.tool_use_config import ToolUselist
-
-logger = logging.getLogger(__name__)
-
-
-class ITSToolOrchestrationFunctionConfig(FunctionBaseConfig, name="its_tool_orchestration"):
-    """
-    Configuration for the ITSToolOrchestrationFunction, which is used to orchestrate multiple functions.
-    """
-
-    augmented_fns: list[FunctionRef] = Field(
-        description="list of FunctionRefs for the functions to be orchestrated. Must be wrapped in `its_tool_wrapper`.")
-
-    search_strategy: ITSStrategyRef | None = Field(
-        description="The ITS search strategy to use for orchestrating invocation of the functions."
-        " If None, no search will be performed.",
-        default=None,
-    )
-
-    editing_strategy: ITSStrategyRef | None = Field(
-        default=None,
-        description="The ITS editing strategy to use for orchestrating invocation of the functions. "
-        "If None, no editing will be performed.",
-    )
-
-    scoring_strategy: ITSStrategyRef | None = Field(
-        default=None,
-        description="The ITS scoring strategy to use for orchestrating invocation of the functions. "
-        "If None, no scoring will be performed.",
-    )
-
-    selection_strategy: ITSStrategyRef = Field(
-        description="The ITS selection strategy to use for orchestrating invocation of the functions.")
-
-
-@register_function(config_type=ITSToolOrchestrationFunctionConfig, framework_wrappers=[LLMFrameworkEnum.LANGCHAIN])
-async def register_its_tool_orchestration_function(
-    config: ITSToolOrchestrationFunctionConfig,
-    builder: Builder,
-):
-    """
-    Registers an ITS-based orchestration function that:
-     1. Instantiates all relevant strategies (search, editing, scoring, selection).
-     2. Accepts a ToolUselist, converts each item to an ITSItem, optionally runs search/editing.
-     3. Calls the correct augmented_fn per item using name=tool name.
-     4. If configured, runs scoring and selection on the result.
-     5. Returns a new ToolUselist with each output set.
-    """
-
-    # 1) Gather references to all augmented (wrapped) functions
-    function_map = {}
-    for fn_ref in config.augmented_fns:
-        # Retrieve the actual function from the builder
-        fn_obj = builder.get_function(fn_ref)
-        function_map[fn_ref] = fn_obj
-
-    # 2) Instantiate search, editing, scoring, selection strategies (if any)
-    search = None
-    if config.search_strategy is not None:
-        search = await builder.get_its_strategy(
-            strategy_name=config.search_strategy,
-            pipeline_type=PipelineTypeEnum.TOOL_USE,
-            stage_type=StageTypeEnum.SEARCH,
-        )
-
-    editing = None
-    if config.editing_strategy is not None:
-        editing = await builder.get_its_strategy(
-            strategy_name=config.editing_strategy,
-            pipeline_type=PipelineTypeEnum.TOOL_USE,
-            stage_type=StageTypeEnum.EDITING,
-        )
-
-    scoring = None
-    if config.scoring_strategy is not None:
-        scoring = await builder.get_its_strategy(
-            strategy_name=config.scoring_strategy,
-            pipeline_type=PipelineTypeEnum.TOOL_USE,
-            stage_type=StageTypeEnum.SCORING,
-        )
-
-    selection = await builder.get_its_strategy(
-        strategy_name=config.selection_strategy,
-        pipeline_type=PipelineTypeEnum.TOOL_USE,
-        stage_type=StageTypeEnum.SELECTION,
-    )
-
-    fn_description = ("\n".join(f"- **{fn_ref}**: {function_map[fn_ref].description or 'No description provided.'}"
-                                for fn_ref in config.augmented_fns))
-
-    # 3) Create the inner function to handle single (non-streaming) calls.
-    async def single_inner(tool_list: ToolUselist) -> ToolUselist:
-        """
-        Orchestrates multiple tool usages, optionally using search/editing/scoring/selection steps.
-        """
-        # Convert each ToolUseInputSchema to ITSItem
-        its_items = []
-        for t in tool_list.tools:
-            item = ITSItem(
-                input=t.task_description,  # The user "task"
-                output=None,
-                name=t.tool_name,  # The "tool name"
-                metadata=t.motivation,  # The "justification"
-            )
-            its_items.append(item)
-
-        # Run search strategy if present
-        if search is not None:
-            its_items = await search.ainvoke(its_items)
-
-        logger.info("ITS orchestration function: %d items after search", len(its_items))
-
-        # Invoke the correct augmented function for each item concurrently
-        # Helper coroutine to invoke a tool function and capture result or error
-        async def _invoke_tool(item: ITSItem, fn):
-            try:
-                result = await fn.acall_invoke(item.output)
-                return item, result, None
-            except Exception as e:
-                logger.error(f"Error invoking function '{item.name}': {e}")
-                return item, None, str(e)
-
-        tasks = []
-        for item in its_items:
-            if item.name not in function_map:
-                logger.error(f"Function '{item.name}' not found in function map.")
-                item.output = f"Error: Function '{item.name}' not found in function map. Check your input"
-            else:
-                fn = function_map[item.name]
-                tasks.append(_invoke_tool(item, fn))
-
-        # Await all tasks and assign outputs
-        if tasks:
-            results = await asyncio.gather(*tasks)
-            for item, result, error in results:
-                if error:
-                    item.output = f"Error invoking function '{item.name}': {error}"
-                else:
-                    item.output = result
-
-        if editing:
-            its_items = await editing.ainvoke(its_items)
-
-        # Run scoring strategy if present
-        if scoring is not None:
-            its_items = await scoring.ainvoke(its_items)
-
-        # Run selection strategy
-        if selection is not None:
-            its_items = await selection.ainvoke(its_items)
-
-        logger.info("ITS orchestration function: %d items after selection", len(its_items))
-
-        # Convert final results from ITSItems back to a ToolUselist
-        final_list = ToolUselist(tools=[])
-        for item in its_items:
-            # Compose a new ToolUseInputSchema with final output
-            new_tool = ToolUseInputSchema(
-                tool_name=item.name,
-                task_description=str(item.input),
-                motivation=item.metadata if item.metadata else None,
-                output=str(item.output) if item.output is not None else None,
-            )
-            final_list.tools.append(new_tool)
-
-        return final_list
-
-    # 4) Return the function info (only a single_fn is needed; no streaming)
-    yield FunctionInfo.create(
-        single_fn=single_inner,
-        stream_fn=None,  # No streaming required
-        input_schema=ToolUselist,
-        single_output_schema=ToolUselist,
-        description=fn_description)

aiq/experimental/inference_time_scaling/functions/its_tool_wrapper_function.py

@@ -1,146 +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 pydantic import BaseModel
-from pydantic import Field
-
-from aiq.builder.builder import Builder
-from aiq.builder.framework_enum import LLMFrameworkEnum
-from aiq.builder.function import Function
-from aiq.builder.function_info import FunctionInfo
-from aiq.cli.register_workflow import register_function
-from aiq.data_models.component_ref import FunctionRef
-from aiq.data_models.component_ref import LLMRef
-from aiq.data_models.function import FunctionBaseConfig
-from aiq.utils.string_utils import convert_to_str
-
-logger = logging.getLogger(__name__)
-
-
-class ITSToolWrapperFunctionConfig(FunctionBaseConfig, name="its_tool_wrapper"):
-    """
-    Configuration for the ITSToolWrapperFunction, which is used to wrap a function that will be executed
-    in the inference time scaling pipeline.
-
-    This function is responsible for turning an 'objective' or description for the tool into tool input.
-
-    NOTE: Only supports LLMs with structured output.
-    """
-
-    augmented_fn: FunctionRef = Field(description="The name of the function to reason on.")
-
-    input_llm: LLMRef = Field(description="The LLM that will generate input to the function.")
-    verbose: bool = Field(default=False, description="Whether to log detailed information.")
-
-    downstream_template: str = Field(
-        description="The template for the input LLM to generate structured input to the function.",
-        default=("You are highly sophisticated generalist AI assistant. Your objective is to act as a"
-                 " conduit between a user's task for a function and the function itself. You will be given a general "
-                 "description of the task, or pseudo input for a function. You will also be provided with description "
-                 "of the function, its input schema, and the output schema. Your task is to generate structured input "
-                 "to the function based on the description of the task and the function's input schema. If you do not "
-                 "have enough information to generate structured input, you should respond with 'NOT ENOUGH "
-                 "INFORMATION'. \n\n The description of the function is: {function_description}\n\n"
-                 "The input schema of the function is: {input_schema}\n\n"
-                 "The output schema of the function is: {output_schema}\n\n"
-                 "The description of the task is: {task_description}\n\n"
-                 "The structured input to the function is: "))
-
-    tool_description: str | None = Field(description="The description of the tool to be used for the function.",
-                                         default=None)
-
-
-@register_function(config_type=ITSToolWrapperFunctionConfig, framework_wrappers=[LLMFrameworkEnum.LANGCHAIN])
-async def register_its_tool_wrapper_function(
-    config: ITSToolWrapperFunctionConfig,
-    builder: Builder,
-):
-    """
-    Register the ITSToolWrapperFunction with the provided builder and configuration.
-    """
-
-    try:
-        from langchain_core.language_models import BaseChatModel
-        from langchain_core.prompts import PromptTemplate
-    except ImportError:
-        raise ImportError("langchain-core is not installed. Please install it to use SingleShotMultiPlanPlanner.\n"
-                          "This error can be resolved by installing aiqtoolkit-langchain.")
-
-    augmented_function: Function = builder.get_function(config.augmented_fn)
-    input_llm: BaseChatModel = await builder.get_llm(config.input_llm, wrapper_type=LLMFrameworkEnum.LANGCHAIN)
-
-    if not augmented_function.has_single_output:
-        raise ValueError("ITSToolWrapperFunction only supports functions with a single output.")
-
-    if not augmented_function.has_single_output:
-        raise ValueError("ITSToolWrapperFunction only supports functions with a single output.")
-
-    if augmented_function.description and augmented_function.description != "":
-        augmented_function_desc = augmented_function.description
-    else:
-        if not config.tool_description:
-            raise ValueError(f"Function {config.augmented_fn} does not have a description. Cannot augment "
-                             f"function without a description and without a tool description.")
-
-        augmented_function_desc = config.tool_description
-
-    fn_input_schema: BaseModel = augmented_function.input_schema
-    fn_output_schema: BaseModel = augmented_function.single_output_schema
-
-    runnable_llm = input_llm.with_structured_output(schema=fn_input_schema)
-
-    template = PromptTemplate(
-        template=config.downstream_template,
-        input_variables=["function_description", "input_schema", "output_schema", "task_description"],
-        validate_template=True)
-
-    function_description = (f"\nDescription: {augmented_function_desc}\n" +
-                            "\n Input should be a thorough description with all relevant information on what "
-                            f"the tool should do.  The tool requires information about "
-                            f"{fn_input_schema.model_fields}")
-
-    async def single_inner(input_message: str) -> fn_output_schema:
-        """
-        Inner function to handle the streaming output of the ITSToolWrapperFunction.
-        It generates structured input for the augmented function based on the input message.
-        """
-
-        prompt = await template.ainvoke(
-            input={
-                "function_description": augmented_function_desc,
-                "input_schema": fn_input_schema,
-                "output_schema": fn_output_schema,
-                "task_description": input_message
-            })
-
-        prompt = prompt.to_string()
-
-        if config.verbose:
-            logger.info("ITSToolWrapperFunction: Generated prompt: %s", prompt)
-
-        llm_parsed = await runnable_llm.ainvoke(prompt)
-
-        if not llm_parsed:
-            logger.warning("ITSToolWrapperFunction: LLM parsing error")
-            return "Not enough information"
-
-        # Call the augmented function with the structured input
-        result = await augmented_function.acall_invoke(llm_parsed)
-
-        return result
-
-    yield FunctionInfo.from_fn(fn=single_inner, description=function_description, converters=[convert_to_str])