aiqtoolkit 1.1.0__py3-none-any.whl

aiq/agent/rewoo_agent/agent.py

@@ -0,0 +1,410 @@
+# 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 json
+# pylint: disable=R0917
+import logging
+from json import JSONDecodeError
+
+from langchain_core.callbacks.base import AsyncCallbackHandler
+from langchain_core.language_models import BaseChatModel
+from langchain_core.messages.ai import AIMessage
+from langchain_core.messages.human import HumanMessage
+from langchain_core.messages.tool import ToolMessage
+from langchain_core.prompts.chat import ChatPromptTemplate
+from langchain_core.runnables.config import RunnableConfig
+from langchain_core.tools import BaseTool
+from langgraph.graph import StateGraph
+from pydantic import BaseModel
+from pydantic import Field
+
+from aiq.agent.base import AGENT_LOG_PREFIX
+from aiq.agent.base import AGENT_RESPONSE_LOG_MESSAGE
+from aiq.agent.base import INPUT_SCHEMA_MESSAGE
+from aiq.agent.base import NO_INPUT_ERROR_MESSAGE
+from aiq.agent.base import TOOL_NOT_FOUND_ERROR_MESSAGE
+from aiq.agent.base import TOOL_RESPONSE_LOG_MESSAGE
+from aiq.agent.base import AgentDecision
+from aiq.agent.base import BaseAgent
+
+logger = logging.getLogger(__name__)
+
+
+class ReWOOGraphState(BaseModel):
+    """State schema for the ReWOO Agent Graph"""
+    task: HumanMessage = Field(default_factory=lambda: HumanMessage(content=""))  # the task provided by user
+    plan: AIMessage = Field(
+        default_factory=lambda: AIMessage(content=""))  # the plan generated by the planner to solve the task
+    steps: AIMessage = Field(
+        default_factory=lambda: AIMessage(content=""))  # the steps to solve the task, parsed from the plan
+    intermediate_results: dict[str, ToolMessage] = Field(default_factory=dict)  # the intermediate results of each step
+    result: AIMessage = Field(
+        default_factory=lambda: AIMessage(content=""))  # the final result of the task, generated by the solver
+
+
+class ReWOOAgentGraph(BaseAgent):
+    """Configurable LangGraph ReWOO Agent. A ReWOO Agent performs reasoning by interacting with other objects or tools
+    and utilizes their outputs to make decisions. Supports retrying on output parsing errors. Argument
+    "detailed_logs" toggles logging of inputs, outputs, and intermediate steps."""
+
+    def __init__(self,
+                 llm: BaseChatModel,
+                 planner_prompt: ChatPromptTemplate,
+                 solver_prompt: ChatPromptTemplate,
+                 tools: list[BaseTool],
+                 use_tool_schema: bool = True,
+                 callbacks: list[AsyncCallbackHandler] = None,
+                 detailed_logs: bool = False):
+        super().__init__(llm=llm, tools=tools, callbacks=callbacks, detailed_logs=detailed_logs)
+
+        logger.debug(
+            "%s Filling the prompt variables 'tools' and 'tool_names', using the tools provided in the config.",
+            AGENT_LOG_PREFIX)
+        tool_names = ",".join([tool.name for tool in tools[:-1]]) + ',' + tools[-1].name  # prevent trailing ","
+        if not use_tool_schema:
+            tool_names_and_descriptions = "\n".join(
+                [f"{tool.name}: {tool.description}"
+                 for tool in tools[:-1]]) + "\n" + f"{tools[-1].name}: {tools[-1].description}"  # prevent trailing "\n"
+        else:
+            logger.debug("%s Adding the tools' input schema to the tools' description", AGENT_LOG_PREFIX)
+            tool_names_and_descriptions = "\n".join([
+                f"{tool.name}: {tool.description}. {INPUT_SCHEMA_MESSAGE.format(schema=tool.input_schema.model_fields)}"
+                for tool in tools[:-1]
+            ]) + "\n" + (f"{tools[-1].name}: {tools[-1].description}. "
+                         f"{INPUT_SCHEMA_MESSAGE.format(schema=tools[-1].input_schema.model_fields)}")
+
+        self.planner_prompt = planner_prompt.partial(tools=tool_names_and_descriptions, tool_names=tool_names)
+        self.solver_prompt = solver_prompt
+        self.tools_dict = {tool.name: tool for tool in tools}
+
+        logger.debug("%s Initialized ReWOO Agent Graph", AGENT_LOG_PREFIX)
+
+    def _get_tool(self, tool_name):
+        try:
+            return self.tools_dict.get(tool_name)
+        except Exception as ex:
+            logger.exception("%s Unable to find tool with the name %s\n%s",
+                             AGENT_LOG_PREFIX,
+                             tool_name,
+                             ex,
+                             exc_info=True)
+            raise ex
+
+    @staticmethod
+    def _get_current_step(state: ReWOOGraphState) -> int:
+        steps = state.steps.content
+        if len(steps) == 0:
+            raise RuntimeError('No steps received in ReWOOGraphState')
+
+        if len(state.intermediate_results) == len(steps):
+            # all steps are done
+            return -1
+
+        return len(state.intermediate_results)
+
+    @staticmethod
+    def _parse_planner_output(planner_output: str) -> AIMessage:
+
+        try:
+            steps = json.loads(planner_output)
+        except json.JSONDecodeError as ex:
+            raise ValueError(f"The output of planner is invalid JSON format: {planner_output}") from ex
+
+        return AIMessage(content=steps)
+
+    @staticmethod
+    def _replace_placeholder(placeholder: str, tool_input: str | dict, tool_output: str | dict) -> str | dict:
+
+        # Replace the placeholders in the tool input with the previous tool output
+        if isinstance(tool_input, dict):
+            for key, value in tool_input.items():
+                if value is not None:
+                    if value == placeholder:
+                        tool_input[key] = tool_output
+                    elif placeholder in value:
+                        # If the placeholder is part of the value, replace it with the stringified output
+                        tool_input[key] = value.replace(placeholder, str(tool_output))
+
+        elif isinstance(tool_input, str):
+            tool_input = tool_input.replace(placeholder, str(tool_output))
+
+        else:
+            assert False, f"Unexpected type for tool_input: {type(tool_input)}"
+        return tool_input
+
+    @staticmethod
+    def _parse_tool_input(tool_input: str | dict):
+
+        # If the input is already a dictionary, return it as is
+        if isinstance(tool_input, dict):
+            logger.debug("%s Tool input is already a dictionary. Use the tool input as is.", AGENT_LOG_PREFIX)
+            return tool_input
+
+        # If the input is a string, attempt to parse it as JSON
+        try:
+            tool_input = tool_input.strip()
+            # If the input is already a valid JSON string, load it
+            tool_input_parsed = json.loads(tool_input)
+            logger.debug("%s Successfully parsed structured tool input", AGENT_LOG_PREFIX)
+
+        except JSONDecodeError:
+            try:
+                # Replace single quotes with double quotes and attempt parsing again
+                tool_input_fixed = tool_input.replace("'", '"')
+                tool_input_parsed = json.loads(tool_input_fixed)
+                logger.debug(
+                    "%s Successfully parsed structured tool input after replacing single quotes with double quotes",
+                    AGENT_LOG_PREFIX)
+
+            except JSONDecodeError:
+                # If it still fails, fall back to using the input as a raw string
+                tool_input_parsed = tool_input
+                logger.debug("%s Unable to parse structured tool input. Using raw tool input as is.", AGENT_LOG_PREFIX)
+
+        return tool_input_parsed
+
+    async def planner_node(self, state: ReWOOGraphState):
+        try:
+            logger.debug("%s Starting the ReWOO Planner Node", AGENT_LOG_PREFIX)
+
+            planner = self.planner_prompt | self.llm
+            task = state.task.content
+            if not task:
+                logger.error("%s No task provided to the ReWOO Agent. Please provide a valid task.", AGENT_LOG_PREFIX)
+                return {"result": NO_INPUT_ERROR_MESSAGE}
+
+            plan = ""
+            async for event in planner.astream({"task": task}, config=RunnableConfig(callbacks=self.callbacks)):
+                plan += event.content
+
+            steps = self._parse_planner_output(plan)
+
+            if self.detailed_logs:
+                agent_response_log_message = AGENT_RESPONSE_LOG_MESSAGE % (task, plan)
+                logger.info("ReWOO agent planner output: %s", agent_response_log_message)
+
+            return {"plan": AIMessage(content=plan), "steps": steps}
+
+        except Exception as ex:
+            logger.exception("%s Failed to call planner_node: %s", AGENT_LOG_PREFIX, ex, exc_info=True)
+            raise ex
+
+    async def executor_node(self, state: ReWOOGraphState):
+        try:
+            logger.debug("%s Starting the ReWOO Executor Node", AGENT_LOG_PREFIX)
+
+            current_step = self._get_current_step(state)
+            # The executor node should not be invoked after all steps are finished
+            if current_step < 0:
+                logger.error("%s ReWOO Executor is invoked with an invalid step number: %s",
+                             AGENT_LOG_PREFIX,
+                             current_step)
+                raise RuntimeError(f"ReWOO Executor is invoked with an invalid step number: {current_step}")
+
+            step_info = state.steps.content[current_step]["evidence"]
+            placeholder = step_info.get("placeholder", "")
+            tool = step_info.get("tool", "")
+            tool_input = step_info.get("tool_input", "")
+
+            intermediate_results = state.intermediate_results
+
+            # Replace the placeholder in the tool input with the previous tool output
+            for _placeholder, _tool_output in intermediate_results.items():
+                _tool_output = _tool_output.content
+                # If the content is a list, get the first element which should be a dict
+                if isinstance(_tool_output, list):
+                    _tool_output = _tool_output[0]
+                    assert isinstance(_tool_output, dict)
+
+                tool_input = self._replace_placeholder(_placeholder, tool_input, _tool_output)
+
+            requested_tool = self._get_tool(tool)
+            if not requested_tool:
+                configured_tool_names = list(self.tools_dict.keys())
+                logger.warning(
+                    "%s ReWOO Agent wants to call tool %s. In the ReWOO Agent's configuration within the config file,"
+                    "there is no tool with that name: %s",
+                    AGENT_LOG_PREFIX,
+                    tool,
+                    configured_tool_names)
+
+                intermediate_results[placeholder] = ToolMessage(content=TOOL_NOT_FOUND_ERROR_MESSAGE.format(
+                    tool_name=tool, tools=configured_tool_names),
+                                                                tool_call_id=tool)
+                return {"intermediate_results": intermediate_results}
+
+            if self.detailed_logs:
+                logger.debug("%s Calling tool %s with input: %s", AGENT_LOG_PREFIX, requested_tool.name, tool_input)
+
+            # Run the tool. Try to use structured input, if possible
+            tool_input_parsed = self._parse_tool_input(tool_input)
+            tool_response = await requested_tool.ainvoke(tool_input_parsed,
+                                                         config=RunnableConfig(callbacks=self.callbacks))
+
+            # some tools, such as Wikipedia, will return an empty response when no search results are found
+            if tool_response is None or tool_response == "":
+                tool_response = "The tool provided an empty response.\n"
+
+            # ToolMessage only accepts str or list[str | dict] as content.
+            # Convert into list if the response is a dict.
+            if isinstance(tool_response, dict):
+                tool_response = [tool_response]
+
+            tool_response_message = ToolMessage(name=tool, tool_call_id=tool, content=tool_response)
+
+            logger.debug("%s Successfully called the tool", AGENT_LOG_PREFIX)
+            if self.detailed_logs:
+                # The tool response can be very large, so we log only the first 1000 characters
+                tool_response_str = tool_response_message.content
+                tool_response_str = tool_response_str[:1000] + "..." if len(
+                    tool_response_str) > 1000 else tool_response_str
+                tool_response_log_message = TOOL_RESPONSE_LOG_MESSAGE % (
+                    requested_tool.name, tool_input_parsed, tool_response_str)
+                logger.info("ReWOO agent executor output: %s", tool_response_log_message)
+
+            intermediate_results[placeholder] = tool_response_message
+            return {"intermediate_results": intermediate_results}
+
+        except Exception as ex:
+            logger.exception("%s Failed to call executor_node: %s", AGENT_LOG_PREFIX, ex, exc_info=True)
+            raise ex
+
+    async def solver_node(self, state: ReWOOGraphState):
+        try:
+            logger.debug("%s Starting the ReWOO Solver Node", AGENT_LOG_PREFIX)
+
+            plan = ""
+            # Add the tool outputs of each step to the plan
+            for step in state.steps.content:
+                step_info = step["evidence"]
+                placeholder = step_info.get("placeholder", "")
+                tool_input = step_info.get("tool_input", "")
+
+                intermediate_results = state.intermediate_results
+                for _placeholder, _tool_output in intermediate_results.items():
+                    _tool_output = _tool_output.content
+                    # If the content is a list, get the first element which should be a dict
+                    if isinstance(_tool_output, list):
+                        _tool_output = _tool_output[0]
+                        assert isinstance(_tool_output, dict)
+
+                    tool_input = self._replace_placeholder(_placeholder, tool_input, _tool_output)
+
+                    placeholder = placeholder.replace(_placeholder, str(_tool_output))
+
+                _plan = step.get("plan")
+                tool = step_info.get("tool")
+                plan += f"Plan: {_plan}\n{placeholder} = {tool}[{tool_input}]"
+
+            task = state.task.content
+            solver_prompt = self.solver_prompt.partial(plan=plan)
+            solver = solver_prompt | self.llm
+            output_message = ""
+            async for event in solver.astream({"task": task}, config=RunnableConfig(callbacks=self.callbacks)):
+                output_message += event.content
+
+            output_message = AIMessage(content=output_message)
+            if self.detailed_logs:
+                solver_output_log_message = AGENT_RESPONSE_LOG_MESSAGE % (task, output_message.content)
+                logger.info("ReWOO agent solver output: %s", solver_output_log_message)
+
+            return {"result": output_message}
+
+        except Exception as ex:
+            logger.exception("%s Failed to call solver_node: %s", AGENT_LOG_PREFIX, ex, exc_info=True)
+            raise ex
+
+    async def conditional_edge(self, state: ReWOOGraphState):
+        try:
+            logger.debug("%s Starting the ReWOO Conditional Edge", AGENT_LOG_PREFIX)
+
+            current_step = self._get_current_step(state)
+            if current_step == -1:
+                logger.debug("%s The ReWOO Executor has finished its task", AGENT_LOG_PREFIX)
+                return AgentDecision.END
+
+            logger.debug("%s The ReWOO Executor is still working on the task", AGENT_LOG_PREFIX)
+            return AgentDecision.TOOL
+
+        except Exception as ex:
+            logger.exception("%s Failed to determine whether agent is calling a tool: %s",
+                             AGENT_LOG_PREFIX,
+                             ex,
+                             exc_info=True)
+            logger.warning("%s Ending graph traversal", AGENT_LOG_PREFIX)
+            return AgentDecision.END
+
+    async def _build_graph(self, state_schema):
+        try:
+            logger.debug("%s Building and compiling the ReWOO Graph", AGENT_LOG_PREFIX)
+
+            graph = StateGraph(state_schema)
+            graph.add_node("planner", self.planner_node)
+            graph.add_node("executor", self.executor_node)
+            graph.add_node("solver", self.solver_node)
+
+            graph.add_edge("planner", "executor")
+            conditional_edge_possible_outputs = {AgentDecision.TOOL: "executor", AgentDecision.END: "solver"}
+            graph.add_conditional_edges("executor", self.conditional_edge, conditional_edge_possible_outputs)
+
+            graph.set_entry_point("planner")
+            graph.set_finish_point("solver")
+
+            self.graph = graph.compile()
+            logger.debug("%s ReWOO Graph built and compiled successfully", AGENT_LOG_PREFIX)
+
+            return self.graph
+
+        except Exception as ex:
+            logger.exception("%s Failed to build ReWOO Graph: %s", AGENT_LOG_PREFIX, ex, exc_info=ex)
+            raise ex
+
+    async def build_graph(self):
+        try:
+            await self._build_graph(state_schema=ReWOOGraphState)
+            logger.debug("%s ReWOO Graph built and compiled successfully", AGENT_LOG_PREFIX)
+            return self.graph
+        except Exception as ex:
+            logger.exception("%s Failed to build ReWOO Graph: %s", AGENT_LOG_PREFIX, ex, exc_info=ex)
+            raise ex
+
+    @staticmethod
+    def validate_planner_prompt(planner_prompt: str) -> bool:
+        errors = []
+        if not planner_prompt:
+            errors.append("The planner prompt cannot be empty.")
+        required_prompt_variables = {
+            "{tools}": "The planner prompt must contain {tools} so the planner agent knows about configured tools.",
+            "{tool_names}": "The planner prompt must contain {tool_names} so the planner agent knows tool names."
+        }
+        for variable_name, error_message in required_prompt_variables.items():
+            if variable_name not in planner_prompt:
+                errors.append(error_message)
+        if errors:
+            error_text = "\n".join(errors)
+            logger.exception("%s %s", AGENT_LOG_PREFIX, error_text)
+            raise ValueError(error_text)
+        return True
+
+    @staticmethod
+    def validate_solver_prompt(solver_prompt: str) -> bool:
+        errors = []
+        if not solver_prompt:
+            errors.append("The solver prompt cannot be empty.")
+        if errors:
+            error_text = "\n".join(errors)
+            logger.exception("%s %s", AGENT_LOG_PREFIX, error_text)
+            raise ValueError(error_text)
+        return True

aiq/agent/rewoo_agent/prompt.py

@@ -0,0 +1,108 @@
+# 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.
+
+# flake8: noqa
+from langchain_core.prompts.chat import ChatPromptTemplate
+
+PLANNER_SYSTEM_PROMPT = """
+For the following task, make plans that can solve the problem step by step. For each plan, indicate \
+which external tool together with tool input to retrieve evidence. You can store the evidence into a \
+placeholder #E that can be called by later tools. (Plan, #E1, Plan, #E2, Plan, ...)
+
+You may ask the human to the following tools:
+
+{tools}
+
+The tools should be one of the following: [{tool_names}]
+
+You are not required to use all the tools listed. Choose only the ones that best fit the needs of each plan step.
+
+Your output must be a JSON array where each element represents one planning step. Each step must be an object with exactly two keys:
+
+1. "plan": A string that describes in detail the action or reasoning for that step.
+
+2. "evidence": An object representing the external tool call associated with that plan step. This object must have the following keys:
+
+   -"placeholder": A string that identifies the evidence placeholder (e.g., "#E1", "#E2", etc.). The numbering should be sequential based on the order of steps.
+
+   -"tool": A string specifying the name of the external tool used.
+
+   -"tool_input": The input to the tool. This can be a string, array, or object, depending on the requirements of the tool.
+
+Do not include any additional keys or characters in your output, and do not wrap your response with markdown formatting. Your output must be strictly valid JSON.
+
+Important instructions:
+
+Do not output any additional text, comments, or markdown formatting.
+
+Do not include any explanation or reasoning text outside of the JSON array.
+
+The output must be a valid JSON array that can be parsed directly.
+
+Here is an example of how a valid JSON output should look:
+
+[
+  \'{{
+    "plan": "Calculate the result of 2023 minus 25.",
+    "evidence": \'{{
+      "placeholder": "#E1",
+      "tool": "calculator_subtract",
+      "tool_input": [2023, 25]
+    }}\'
+  }}\',
+  \'{{
+    "plan": "Retrieve the year represented by the result stored in #E1.",
+    "evidence": \'{{
+      "placeholder": "#E2",
+      "tool": "haystack_chitchat_agent",
+      "tool_input": "Response with the result number contained in #E1"
+    }}\'
+  }}\',
+  \'{{
+    "plan": "Search for the CEO of Golden State Warriors in the year stored in #E2.",
+    "evidence": \'{{
+      "placeholder": "#E3",
+      "tool": "internet_search",
+      "tool_input": "Who was the CEO of Golden State Warriors in the year #E2?"
+    }}\'
+  }}\'
+]
+
+Begin!
+"""
+
+PLANNER_USER_PROMPT = """
+task: {task}
+"""
+
+rewoo_planner_prompt = ChatPromptTemplate([("system", PLANNER_SYSTEM_PROMPT), ("user", PLANNER_USER_PROMPT)])
+
+SOLVER_SYSTEM_PROMPT = """
+Solve the following task or problem. To solve the problem, we have made step-by-step Plan and \
+retrieved corresponding Evidence to each Plan. Use them with caution since long evidence might \
+contain irrelevant information.
+
+Now solve the question or task according to provided Evidence above. Respond with the answer
+directly with no extra words.
+
+"""
+SOLVER_USER_PROMPT = """
+plan: {plan}
+task: {task}
+
+Response:
+"""
+
+rewoo_solver_prompt = ChatPromptTemplate([("system", SOLVER_SYSTEM_PROMPT), ("user", SOLVER_USER_PROMPT)])

aiq/agent/rewoo_agent/register.py

@@ -0,0 +1,158 @@
+# 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.framework_enum import LLMFrameworkEnum
+from aiq.builder.function_info import FunctionInfo
+from aiq.cli.register_workflow import register_function
+from aiq.data_models.api_server import AIQChatRequest
+from aiq.data_models.api_server import AIQChatResponse
+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.type_converter import GlobalTypeConverter
+
+logger = logging.getLogger(__name__)
+
+
+class ReWOOAgentWorkflowConfig(FunctionBaseConfig, name="rewoo_agent"):
+    """
+    Defines an AIQ Toolkit function that uses a ReWOO Agent performs reasoning inbetween tool calls, and utilizes the
+    tool names and descriptions to select the optimal tool.
+    """
+
+    tool_names: list[FunctionRef] = Field(default_factory=list,
+                                          description="The list of tools to provide to the rewoo agent.")
+    llm_name: LLMRef = Field(description="The LLM model to use with the rewoo agent.")
+    verbose: bool = Field(default=False, description="Set the verbosity of the rewoo agent's logging.")
+    include_tool_input_schema_in_tool_description: bool = Field(
+        default=True, description="Specify inclusion of tool input schemas in the prompt.")
+    description: str = Field(default="ReWOO Agent Workflow", description="The description of this functions use.")
+    planner_prompt: str | None = Field(
+        default=None,
+        description="Provides the PLANNER_PROMPT to use with the agent")  # defaults to PLANNER_PROMPT in prompt.py
+    solver_prompt: str | None = Field(
+        default=None,
+        description="Provides the SOLVER_PROMPT to use with the agent")  # defaults to SOLVER_PROMPT in prompt.py
+    max_history: int = Field(default=15, description="Maximum number of messages to keep in the conversation history.")
+    use_openai_api: bool = Field(default=False,
+                                 description=("Use OpenAI API for the input/output types to the function. "
+                                              "If False, strings will be used."))
+    additional_instructions: str | None = Field(
+        default=None, description="Additional instructions to provide to the agent in addition to the base prompt.")
+
+
+@register_function(config_type=ReWOOAgentWorkflowConfig, framework_wrappers=[LLMFrameworkEnum.LANGCHAIN])
+async def ReWOO_agent_workflow(config: ReWOOAgentWorkflowConfig, builder: Builder):
+    from langchain.schema import BaseMessage
+    from langchain_core.messages import trim_messages
+    from langchain_core.messages.human import HumanMessage
+    from langchain_core.prompts import ChatPromptTemplate
+    from langgraph.graph.graph import CompiledGraph
+
+    from aiq.agent.rewoo_agent.prompt import PLANNER_USER_PROMPT
+    from aiq.agent.rewoo_agent.prompt import SOLVER_USER_PROMPT
+
+    from .agent import ReWOOAgentGraph
+    from .agent import ReWOOGraphState
+    from .prompt import rewoo_planner_prompt
+    from .prompt import rewoo_solver_prompt
+
+    # the ReWOO Agent prompt comes from prompt.py, and can be customized there or via config option planner_prompt and
+    # solver_prompt.
+    if config.planner_prompt:
+        planner_prompt = config.planner_prompt
+        if config.additional_instructions:
+            planner_prompt += f"{config.additional_instructions}"
+        valid = ReWOOAgentGraph.validate_planner_prompt(config.planner_prompt)
+        if not valid:
+            logger.exception("Invalid planner_prompt")
+            raise ValueError("Invalid planner_prompt")
+        planner_prompt = ChatPromptTemplate([("system", config.planner_prompt), ("user", PLANNER_USER_PROMPT)])
+    else:
+        planner_prompt = rewoo_planner_prompt
+
+    if config.solver_prompt:
+        solver_prompt = config.solver_prompt
+        if config.additional_instructions:
+            solver_prompt += f"{config.additional_instructions}"
+        valid = ReWOOAgentGraph.validate_solver_prompt(config.solver_prompt)
+        if not valid:
+            logger.exception("Invalid solver_prompt")
+            raise ValueError("Invalid solver_prompt")
+        solver_prompt = ChatPromptTemplate([("system", config.solver_prompt), ("user", SOLVER_USER_PROMPT)])
+    else:
+        solver_prompt = rewoo_solver_prompt
+
+    # we can choose an LLM for the ReWOO agent in the config file
+    llm = await builder.get_llm(config.llm_name, wrapper_type=LLMFrameworkEnum.LANGCHAIN)
+
+    # the agent can run any installed tool, simply install the tool and add it to the config file
+    # the sample tool provided can easily be copied or changed
+    tools = builder.get_tools(tool_names=config.tool_names, wrapper_type=LLMFrameworkEnum.LANGCHAIN)
+    if not tools:
+        raise ValueError(f"No tools specified for ReWOO Agent '{config.llm_name}'")
+
+    # construct the ReWOO Agent Graph from the configured llm, prompt, and tools
+    graph: CompiledGraph = await ReWOOAgentGraph(llm=llm,
+                                                 planner_prompt=planner_prompt,
+                                                 solver_prompt=solver_prompt,
+                                                 tools=tools,
+                                                 use_tool_schema=config.include_tool_input_schema_in_tool_description,
+                                                 detailed_logs=config.verbose).build_graph()
+
+    async def _response_fn(input_message: AIQChatRequest) -> AIQChatResponse:
+        try:
+            # initialize the starting state with the user query
+            messages: list[BaseMessage] = trim_messages(messages=[m.model_dump() for m in input_message.messages],
+                                                        max_tokens=config.max_history,
+                                                        strategy="last",
+                                                        token_counter=len,
+                                                        start_on="human",
+                                                        include_system=True)
+            task = HumanMessage(content=messages[0].content)
+            state = ReWOOGraphState(task=task)
+
+            # run the ReWOO Agent Graph
+            state = await graph.ainvoke(state)
+
+            # get and return the output from the state
+            state = ReWOOGraphState(**state)
+            output_message = state.result.content  # pylint: disable=E1101
+            return AIQChatResponse.from_string(output_message)
+
+        except Exception as ex:
+            logger.exception("ReWOO Agent failed with exception: %s", ex, exc_info=ex)
+            # here, we can implement custom error messages
+            if config.verbose:
+                return AIQChatResponse.from_string(str(ex))
+            return AIQChatResponse.from_string("I seem to be having a problem.")
+
+    if (config.use_openai_api):
+        yield FunctionInfo.from_fn(_response_fn, description=config.description)
+
+    else:
+
+        async def _str_api_fn(input_message: str) -> str:
+            oai_input = GlobalTypeConverter.get().convert(input_message, to_type=AIQChatRequest)
+            oai_output = await _response_fn(oai_input)
+
+            return GlobalTypeConverter.get().convert(oai_output, to_type=str)
+
+        yield FunctionInfo.from_fn(_str_api_fn, description=config.description)