nvidia-nat 1.1.0a20251020__py3-none-any.whl

nat/experimental/test_time_compute/editing/llm_as_a_judge_editor.py

@@ -0,0 +1,204 @@
+# 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 nat.builder.builder import Builder
+from nat.builder.framework_enum import LLMFrameworkEnum
+from nat.cli.register_workflow import register_ttc_strategy
+from nat.data_models.ttc_strategy import TTCStrategyBaseConfig
+from nat.experimental.test_time_compute.models.editor_config import LLMAsAJudgeEditorConfig
+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 LLMAsAJudgeEditor(StrategyBase):
+    """
+    Given a list of PlanningItems, uses a feedback LLM to generate feedback on each plan
+    Then edits the plan based on feedback.
+    """
+
+    def __init__(self, config: TTCStrategyBaseConfig) -> None:
+        super().__init__(config)
+        self.feedback_llm = None
+        self.editing_llm = None
+
+    async def build_components(self, builder: Builder) -> None:
+        """
+        Build the components required for the editor.
+        """
+        # Get the feedback LLM
+        self.feedback_llm = await builder.get_llm(self.config.feedback_llm, wrapper_type=LLMFrameworkEnum.LANGCHAIN)
+
+        self.editing_llm = await builder.get_llm(self.config.editing_llm, wrapper_type=LLMFrameworkEnum.LANGCHAIN)
+
+    def supported_pipeline_types(self) -> [PipelineTypeEnum]:
+        return [PipelineTypeEnum.PLANNING]
+
+    def stage_type(self) -> StageTypeEnum:
+        return StageTypeEnum.EDITING
+
+    async def generate_feedback(self, llm, template, context: str, prompt: str, item: TTCItem) -> TTCItem:
+        """
+        Helper function to generate feedback for a given planning item using the provided prompt.
+        """
+
+        prompt = await template.ainvoke(
+            input={
+                "context": context,
+                "original_prompt": prompt,  # Original prompt used to generate the plans
+                "plan": item.plan,
+                "num_feedback": self.config.num_feedback
+            })
+
+        feedback_result = await llm.ainvoke(prompt.to_string())
+        if not feedback_result:
+            logger.warning(f"No feedback generated for plan: {item.plan}.")
+            return item
+
+        # Update the planning item with the generated feedback
+        cleaned = remove_r1_think_tags(
+            feedback_result.content if hasattr(feedback_result, 'content') else str(feedback_result))
+
+        # Feedback is the string following 'FEEDBACK:'. Use Regex to extract
+        cleaned = re.sub(r'(?i)^\s*FEEDBACK:\s*', '', cleaned).strip()
+        if not cleaned:
+            logger.warning(f"Feedback was empty for plan: {item.plan}.")
+            return item
+
+        item.feedback = cleaned  # Set the feedback in the TTCItem
+
+        return item
+
+    async def edit_plan(self, llm, template, context: str, prompt: str, item: TTCItem) -> TTCItem:
+        """
+        Helper function to edit a plan based on feedback using the provided prompt.
+        """
+
+        if not item.feedback:
+            logger.warning(f"No feedback available for plan: {item.plan}. Cannot edit.")
+            return item
+
+        prompt = await template.ainvoke(
+            input={
+                "context": context,
+                "original_prompt": prompt,  # Original prompt used to generate the plans
+                "plan": item.plan,
+                "feedback": item.feedback
+            })
+
+        editing_result = await llm.ainvoke(prompt.to_string())
+        if not editing_result:
+            logger.warning(f"No editing result generated for plan: {item.plan}.")
+            return item
+
+        # Update the planning item with the edited plan
+        cleaned = remove_r1_think_tags(
+            editing_result.content if hasattr(editing_result, 'content') else str(editing_result))
+
+        # Plan is the string following 'EDITED PLAN:'. Use Regex to extract
+        cleaned = re.sub(r'(?i)^\s*EDITED PLAN:\s*', '', cleaned).strip()
+        if not cleaned:
+            logger.warning(f"Edited plan was empty for plan: {item.plan}. Returning original.")
+            return item
+
+        # Update the plan in the PlanningItem
+        item.plan = cleaned
+
+        return item
+
+    async def ainvoke(self,
+                      items: list[TTCItem],
+                      original_prompt: str | None = None,
+                      agent_context: str | None = None,
+                      **kwargs) -> list[TTCItem]:
+        """
+        Edit the provided planning items using a feedback LLM.
+        """
+        from langchain_core.language_models import BaseChatModel
+        from langchain_core.prompts import PromptTemplate
+
+        # assert self.config.feedback_llm is a BaseChatModel
+        if not isinstance(self.feedback_llm, BaseChatModel):
+            raise ValueError("The `feedback_llm` must be an instance of `BaseChatModel`.")
+
+        # assert self.config.editing_llm is a BaseChatModel
+        if not isinstance(self.editing_llm, BaseChatModel):
+            raise ValueError("The `editing_llm` must be an instance of `BaseChatModel`.")
+
+        feedback_model: BaseChatModel = self.feedback_llm
+        editing_model: BaseChatModel = self.editing_llm
+
+        feedback_template = PromptTemplate(template=self.config.feedback_template,
+                                           input_variables=["context", "original_prompt", "plan", "num_feedback"],
+                                           validate_template=True)
+
+        editing_template = PromptTemplate(template=self.config.editor_template,
+                                          input_variables=["context", "original_prompt", "plan", "feedback"],
+                                          validate_template=True)
+
+        # Generate feedback for each planning item concurrently
+        feedback_tasks = [
+            self.generate_feedback(
+                llm=feedback_model,
+                template=feedback_template,
+                context=agent_context,
+                prompt=original_prompt,  # Original prompt used to generate the plans
+                item=item) for item in items
+        ]
+        # Run the feedback tasks concurrently and gather results
+        planning_items_with_feedback = await asyncio.gather(*feedback_tasks)
+
+        if not planning_items_with_feedback:
+            raise ValueError("No feedback was generated for the planning items. Please check the LLM response.")
+
+        logger.info("Generated feedback for %d plans.", len(planning_items_with_feedback))
+
+        # Now edit each planning item based on the feedback concurrently
+        editing_tasks = [
+            self.edit_plan(
+                llm=editing_model,
+                template=editing_template,
+                context=agent_context,
+                prompt=original_prompt,  # Original prompt used to generate the plans
+                item=item) for item in planning_items_with_feedback
+        ]
+        # Run the editing tasks concurrently and gather results
+        edited_planning_items = await asyncio.gather(*editing_tasks)
+
+        if not edited_planning_items:
+            raise ValueError("No plans were edited. Please check the LLM response.")
+
+        logger.info("Edited %d plans based on feedback.", len(edited_planning_items))
+        return edited_planning_items
+
+
+@register_ttc_strategy(config_type=LLMAsAJudgeEditorConfig)
+async def register_llm_as_a_judge_editor(config: TTCStrategyBaseConfig, builder: Builder):
+    """
+    Register the LLMAsAJudgeEditor strategy with the provided configuration and builder.
+    """
+
+    editor = LLMAsAJudgeEditor(config)
+    await editor.build_components(builder)
+
+    yield editor

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 = await 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,228 @@
+# 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 = await 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_dependencies = builder.get_function_dependencies(config.augmented_fn)
+    function_used_tools = set(function_dependencies.functions)
+    for function_group in function_dependencies.function_groups:
+        function_used_tools.update(builder.get_function_group_dependencies(function_group).functions)
+
+    tool_list = "Tool: Description\n"
+
+    for tool in function_used_tools:
+        tool_impl = await 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)