nvidia-nat 1.2.0rc5__py3-none-any.whl

aiq/experimental/test_time_compute/editing/motivation_aware_summarization.py

@@ -0,0 +1,107 @@
+# 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_ttc_strategy
+from aiq.experimental.test_time_compute.models.editor_config import MotivationAwareSummarizationConfig
+from aiq.experimental.test_time_compute.models.stage_enums import PipelineTypeEnum
+from aiq.experimental.test_time_compute.models.stage_enums import StageTypeEnum
+from aiq.experimental.test_time_compute.models.strategy_base import StrategyBase
+from aiq.experimental.test_time_compute.models.ttc_item import TTCItem
+from aiq.utils.io.model_processing import remove_r1_think_tags
+
+logger = logging.getLogger(__name__)
+
+
+class MotivationAwareSummarization(StrategyBase):
+    """
+    A strategy that, for each incoming TTCItem, 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[TTCItem],
+                      original_prompt: str | None = None,
+                      agent_context: str | None = None,
+                      **kwargs) -> list[TTCItem]:
+        """
+        For each TTCItem, rewrite the 'input' using each LLM to create a new perspective.
+        The new TTCItems' 'output' field will store the newly generated query.
+        """
+        try:
+            from langchain_core.prompts import PromptTemplate
+        except ImportError:
+            raise ImportError("langchain-core is required for MultiQueryRetrievalSearch. "
+                              "Install aiqtoolkit-langchain or similar.")
+
+        new_ttc_items: list[TTCItem] = []
+
+        # 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_ttc_items.append(
+                TTCItem(
+                    input=item.input,
+                    output=remove_r1_think_tags(llm_response),
+                    metadata=item.metadata,
+                    name=item.name,  # keep the original tool name
+                ))
+
+        return new_ttc_items
+
+
+@register_ttc_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/test_time_compute/functions/execute_score_select_function.py

@@ -0,0 +1,105 @@
+# 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 TTCStrategyRef
+from aiq.data_models.function import FunctionBaseConfig
+from aiq.experimental.test_time_compute.models.stage_enums import PipelineTypeEnum
+from aiq.experimental.test_time_compute.models.stage_enums import StageTypeEnum
+from aiq.experimental.test_time_compute.models.ttc_item import TTCItem
+
+logger = logging.getLogger(__name__)
+
+
+class ExecuteScoreSelectFunctionConfig(FunctionBaseConfig, name="execute_score_select_function"):
+    scorer: TTCStrategyRef | None = Field(description="Strategy to score the output of the function", default=None)
+    selector: TTCStrategyRef = 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_ttc_strategy(strategy_name=config.scorer,
+                                                pipeline_type=PipelineTypeEnum.AGENT_EXECUTION,
+                                                stage_type=StageTypeEnum.SCORING)
+    else:
+        scorer = None
+
+    selector = await builder.get_ttc_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 = [TTCItem(
+            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/test_time_compute/functions/its_tool_orchestration_function.py

@@ -0,0 +1,205 @@
+# 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 TTCStrategyRef
+from aiq.data_models.function import FunctionBaseConfig
+from aiq.experimental.test_time_compute.models.stage_enums import PipelineTypeEnum
+from aiq.experimental.test_time_compute.models.stage_enums import StageTypeEnum
+from aiq.experimental.test_time_compute.models.tool_use_config import ToolUseInputSchema
+from aiq.experimental.test_time_compute.models.tool_use_config import ToolUselist
+from aiq.experimental.test_time_compute.models.ttc_item import TTCItem
+
+logger = logging.getLogger(__name__)
+
+
+class TTCToolOrchestrationFunctionConfig(FunctionBaseConfig, name="ttc_tool_orchestration"):
+    """
+    Configuration for the TTCToolOrchestrationFunction, 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 `ttc_tool_wrapper`.")
+
+    search_strategy: TTCStrategyRef | None = Field(
+        description="The TTC search strategy to use for orchestrating invocation of the functions."
+        " If None, no search will be performed.",
+        default=None,
+    )
+
+    editing_strategy: TTCStrategyRef | None = Field(
+        default=None,
+        description="The TTC editing strategy to use for orchestrating invocation of the functions. "
+        "If None, no editing will be performed.",
+    )
+
+    scoring_strategy: TTCStrategyRef | None = Field(
+        default=None,
+        description="The TTC scoring strategy to use for orchestrating invocation of the functions. "
+        "If None, no scoring will be performed.",
+    )
+
+    selection_strategy: TTCStrategyRef = Field(
+        description="The TTC selection strategy to use for orchestrating invocation of the functions.")
+
+
+@register_function(config_type=TTCToolOrchestrationFunctionConfig, framework_wrappers=[LLMFrameworkEnum.LANGCHAIN])
+async def register_ttc_tool_orchestration_function(
+    config: TTCToolOrchestrationFunctionConfig,
+    builder: Builder,
+):
+    """
+    Registers an TTC-based orchestration function that:
+     1. Instantiates all relevant strategies (search, editing, scoring, selection).
+     2. Accepts a ToolUselist, converts each item to an TTCItem, 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_ttc_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_ttc_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_ttc_strategy(
+            strategy_name=config.scoring_strategy,
+            pipeline_type=PipelineTypeEnum.TOOL_USE,
+            stage_type=StageTypeEnum.SCORING,
+        )
+
+    selection = await builder.get_ttc_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 TTCItem
+        ttc_items = []
+        for t in tool_list.tools:
+            item = TTCItem(
+                input=t.task_description,  # The user "task"
+                output=None,
+                name=t.tool_name,  # The "tool name"
+                metadata=t.motivation,  # The "justification"
+            )
+            ttc_items.append(item)
+
+        # Run search strategy if present
+        if search is not None:
+            ttc_items = await search.ainvoke(ttc_items)
+
+        logger.info("TTC orchestration function: %d items after search", len(ttc_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: TTCItem, 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 ttc_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:
+            ttc_items = await editing.ainvoke(ttc_items)
+
+        # Run scoring strategy if present
+        if scoring is not None:
+            ttc_items = await scoring.ainvoke(ttc_items)
+
+        # Run selection strategy
+        if selection is not None:
+            ttc_items = await selection.ainvoke(ttc_items)
+
+        logger.info("TTC orchestration function: %d items after selection", len(ttc_items))
+
+        # Convert final results from TTCItems back to a ToolUselist
+        final_list = ToolUselist(tools=[])
+        for item in ttc_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/test_time_compute/functions/its_tool_wrapper_function.py

@@ -0,0 +1,146 @@
+# 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 TTCToolWrapperFunctionConfig(FunctionBaseConfig, name="ttc_tool_wrapper"):
+    """
+    Configuration for the TTCToolWrapperFunction, 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=TTCToolWrapperFunctionConfig, framework_wrappers=[LLMFrameworkEnum.LANGCHAIN])
+async def register_ttc_tool_wrapper_function(
+    config: TTCToolWrapperFunctionConfig,
+    builder: Builder,
+):
+    """
+    Register the TTCToolWrapperFunction 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("TTCToolWrapperFunction only supports functions with a single output.")
+
+    if not augmented_function.has_single_output:
+        raise ValueError("TTCToolWrapperFunction 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 TTCToolWrapperFunction.
+        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("TTCToolWrapperFunction: Generated prompt: %s", prompt)
+
+        llm_parsed = await runnable_llm.ainvoke(prompt)
+
+        if not llm_parsed:
+            logger.warning("TTCToolWrapperFunction: 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])