nvidia-nat 1.2.0__py3-none-any.whl

nat/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 nat.builder.builder import Builder
+from nat.builder.framework_enum import LLMFrameworkEnum
+from nat.cli.register_workflow import register_ttc_strategy
+from nat.experimental.test_time_compute.models.editor_config import MotivationAwareSummarizationConfig
+from nat.experimental.test_time_compute.models.stage_enums import PipelineTypeEnum
+from nat.experimental.test_time_compute.models.stage_enums import StageTypeEnum
+from nat.experimental.test_time_compute.models.strategy_base import StrategyBase
+from nat.experimental.test_time_compute.models.ttc_item import TTCItem
+from nat.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 nvidia-nat-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

nat/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 nat.builder.builder import Builder
+from nat.builder.function import Function
+from nat.builder.function_info import FunctionInfo
+from nat.cli.register_workflow import register_function
+from nat.data_models.component_ref import FunctionRef
+from nat.data_models.component_ref import TTCStrategyRef
+from nat.data_models.function import FunctionBaseConfig
+from nat.experimental.test_time_compute.models.stage_enums import PipelineTypeEnum
+from nat.experimental.test_time_compute.models.stage_enums import StageTypeEnum
+from nat.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."),
+    )

nat/experimental/test_time_compute/functions/plan_select_execute_function.py

@@ -0,0 +1,224 @@
+# 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 collections.abc import AsyncGenerator
+
+from pydantic import Field
+
+from nat.builder.builder import Builder
+from nat.builder.framework_enum import LLMFrameworkEnum
+from nat.builder.function_info import FunctionInfo
+from nat.cli.register_workflow import register_function
+from nat.data_models.api_server import ChatRequest
+from nat.data_models.component_ref import FunctionRef
+from nat.data_models.component_ref import TTCStrategyRef
+from nat.data_models.function import FunctionBaseConfig
+from nat.experimental.test_time_compute.models.stage_enums import PipelineTypeEnum
+from nat.experimental.test_time_compute.models.stage_enums import StageTypeEnum
+from nat.experimental.test_time_compute.models.ttc_item import TTCItem
+
+logger = logging.getLogger(__name__)
+
+
+class PlanSelectExecuteFunctionConfig(FunctionBaseConfig, name="plan_select_execute_function"):
+    """
+    Defines a NAT function that performs reasoning on the input data.
+    Output is passed to the next function in the workflow.
+
+    Designed to be used with an InterceptingFunction.
+    """
+
+    augmented_fn: FunctionRef = Field(description="The name of the function to reason on.")
+
+    planner: TTCStrategyRef = Field(description="The configuration for the planner.")
+    editor: TTCStrategyRef | None = Field(description="The configuration for the editor.", default=None)
+    scorer: TTCStrategyRef | None = Field(description="The configuration for the scorer.", default=None)
+    selector: TTCStrategyRef = Field(description="The configuration for the selector.")
+
+    verbose: bool = Field(default=False, description="Whether to log detailed information.")
+    agent_context_prompt_template: str = Field(
+        description="The template for the agent context prompt. This prompt is used to provide context about the agent",
+        default=("\nThe agent system has the following description:\n"
+                 "{description}\n"
+                 "And has access to the following tools with functionality:\n"
+                 "{tools}\n\n"))
+
+    downstream_template: str = Field(
+        description=("The template for the downstream prompt. This prompt is used to provide the reasoning output to"
+                     " the executing agent"),
+        default=("Answer the following question based on message history: {input_text}"
+                 "\n\nHere is a plan for execution that you could use to guide you if you wanted to:"
+                 "\n\n{reasoning_output}"
+                 "\n\nNOTE: Remember to follow your guidance on how to format output, etc."
+                 "\n\n You must respond with the answer to the original question directly to the user."))
+
+
+@register_function(config_type=PlanSelectExecuteFunctionConfig, framework_wrappers=[LLMFrameworkEnum.LANGCHAIN])
+async def plan_select_execute_function(config: PlanSelectExecuteFunctionConfig, builder: Builder):
+    """
+    Build a ExecutionPlanningFunction from the provided config.
+
+    Args:
+        config (ExecutionPlanningFunctionConfig): The config for the ExecutionPlanningFunction.
+        builder (Builder): The Builder instance to use for building the function.
+
+    Returns:
+        ExecutionPlanningFunction: The built ExecutionPlanningFunction.
+    """
+
+    try:
+        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 nvidia-nat-langchain.")
+
+    # Get the augmented function's description
+    augmented_function = builder.get_function(config.augmented_fn)
+
+    # For now, we rely on runtime checking for type conversion
+
+    if augmented_function.description and augmented_function.description != "":
+        augmented_function_desc = augmented_function.description
+    else:
+        raise ValueError(f"Function {config.augmented_fn} does not have a description. Cannot augment "
+                         f"function without a description.")
+
+    # Get the function dependencies of the augmented function
+    function_used_tools = builder.get_function_dependencies(config.augmented_fn).functions
+    tool_list = "Tool: Description\n"
+
+    for tool in function_used_tools:
+        tool_impl = builder.get_function(tool)
+        tool_list += f"- {tool}: {tool_impl.description if hasattr(tool_impl, 'description') else ''}\n"
+
+    # Draft the reasoning prompt for the augmented function
+    template = PromptTemplate(template=config.agent_context_prompt_template,
+                              input_variables=["description", "tools"],
+                              validate_template=True)
+
+    downstream_template = PromptTemplate(template=config.downstream_template,
+                                         input_variables=["input_text", "reasoning_output"],
+                                         validate_template=True)
+
+    planner = await builder.get_ttc_strategy(strategy_name=config.planner,
+                                             pipeline_type=PipelineTypeEnum.PLANNING,
+                                             stage_type=StageTypeEnum.SEARCH)
+
+    selector = await builder.get_ttc_strategy(strategy_name=config.selector,
+                                              pipeline_type=PipelineTypeEnum.PLANNING,
+                                              stage_type=StageTypeEnum.SELECTION)
+
+    if config.editor:
+        editor = await builder.get_ttc_strategy(strategy_name=config.editor,
+                                                pipeline_type=PipelineTypeEnum.PLANNING,
+                                                stage_type=StageTypeEnum.EDITING)
+    else:
+        editor = None
+
+    if config.scorer:
+        scorer = await builder.get_ttc_strategy(strategy_name=config.scorer,
+                                                pipeline_type=PipelineTypeEnum.PLANNING,
+                                                stage_type=StageTypeEnum.SCORING)
+    else:
+        scorer = None
+
+    async def planning_pipeline(prompt, context):
+
+        plans = await planner.ainvoke([TTCItem()], prompt, context)
+
+        if editor:
+            plans = await editor.ainvoke(plans, prompt, context)
+        if scorer:
+            plans = await scorer.ainvoke(plans, prompt, context)
+
+        selected_plan = (await selector.ainvoke(plans, prompt, context))[0]
+
+        return selected_plan
+
+    streaming_inner_fn = None
+    single_inner_fn = None
+
+    if augmented_function.has_streaming_output:
+
+        async def streaming_inner(
+                input_message: ChatRequest) -> AsyncGenerator[augmented_function.streaming_output_type]:
+            """
+            Perform reasoning on the input text.
+
+            Args:
+                input_message (ChatRequest): The input text to reason on.
+            """
+
+            input_text = "".join([str(message.model_dump()) + "\n" for message in input_message.messages])
+
+            context_prompt = await template.ainvoke(input={"description": augmented_function_desc, "tools": tool_list})
+
+            context_prompt = context_prompt.to_string()
+
+            # Run the TTC pipeline
+            planning_item: TTCItem = await planning_pipeline(prompt=input_text, context=context_prompt)
+
+            output = await downstream_template.ainvoke(input={
+                "input_text": input_text, "reasoning_output": planning_item.plan
+            })
+
+            output = output.to_string()
+
+            if config.verbose:
+                logger.info("Reasoning plan and input to agent: \n\n%s", output)
+
+            async for chunk in augmented_function.acall_stream(output):
+                yield chunk
+
+        streaming_inner_fn = streaming_inner
+
+    if augmented_function.has_single_output:
+
+        async def single_inner(input_message: ChatRequest) -> augmented_function.single_output_type:
+            """
+            Perform reasoning on the input text.
+
+            Args:
+                input_message (ChatRequest): The input text to reason on.
+            """
+
+            input_text = "".join([str(message.model_dump()) + "\n" for message in input_message.messages])
+
+            context_prompt = await template.ainvoke(input={"description": augmented_function_desc, "tools": tool_list})
+
+            context_prompt = context_prompt.to_string()
+
+            # Run the TTC pipeline
+            planning_item: TTCItem = await planning_pipeline(prompt=input_text, context=context_prompt)
+
+            output = await downstream_template.ainvoke(input={
+                "input_text": input_text, "reasoning_output": planning_item.plan
+            })
+
+            output = output.to_string()
+
+            if config.verbose:
+                logger.info("Reasoning plan and input to agent: \n\n%s", output)
+
+            return await augmented_function.acall_invoke(output)
+
+        single_inner_fn = single_inner
+
+    yield FunctionInfo.create(
+        single_fn=single_inner_fn,
+        stream_fn=streaming_inner_fn,
+        description=("Function that runs an TTC execution planner on input and sends plan downstream"),
+        converters=augmented_function.converter_list)

nat/experimental/test_time_compute/functions/ttc_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 nat.builder.builder import Builder
+from nat.builder.framework_enum import LLMFrameworkEnum
+from nat.builder.function_info import FunctionInfo
+from nat.cli.register_workflow import register_function
+from nat.data_models.component_ref import FunctionRef
+from nat.data_models.component_ref import TTCStrategyRef
+from nat.data_models.function import FunctionBaseConfig
+from nat.experimental.test_time_compute.models.stage_enums import PipelineTypeEnum
+from nat.experimental.test_time_compute.models.stage_enums import StageTypeEnum
+from nat.experimental.test_time_compute.models.tool_use_config import ToolUseInputSchema
+from nat.experimental.test_time_compute.models.tool_use_config import ToolUselist
+from nat.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)