nvidia-nat 1.3.dev0__py3-none-any.whl → 1.3.0rc1__py3-none-any.whl

nat/agent/react_agent/agent.py

@@ -14,20 +14,23 @@
 # limitations under the License.
 
 import json
-# pylint: disable=R0917
 import logging
+import re
+import typing
 from json import JSONDecodeError
 
 from langchain_core.agents import AgentAction
 from langchain_core.agents import AgentFinish
 from langchain_core.callbacks.base import AsyncCallbackHandler
 from langchain_core.language_models import BaseChatModel
+from langchain_core.language_models import LanguageModelInput
 from langchain_core.messages.ai import AIMessage
 from langchain_core.messages.base import BaseMessage
 from langchain_core.messages.human import HumanMessage
 from langchain_core.messages.tool import ToolMessage
 from langchain_core.prompts import ChatPromptTemplate
 from langchain_core.prompts import MessagesPlaceholder
+from langchain_core.runnables import Runnable
 from langchain_core.runnables.config import RunnableConfig
 from langchain_core.tools import BaseTool
 from pydantic import BaseModel
@@ -44,7 +47,9 @@ from nat.agent.react_agent.output_parser import ReActOutputParser
 from nat.agent.react_agent.output_parser import ReActOutputParserException
 from nat.agent.react_agent.prompt import SYSTEM_PROMPT
 from nat.agent.react_agent.prompt import USER_PROMPT
-from nat.agent.react_agent.register import ReActAgentWorkflowConfig
+
+if typing.TYPE_CHECKING:
+    from nat.agent.react_agent.register import ReActAgentWorkflowConfig
 
 logger = logging.getLogger(__name__)
 
@@ -54,6 +59,7 @@ class ReActGraphState(BaseModel):
     messages: list[BaseMessage] = Field(default_factory=list)  # input and output of the ReAct Agent
     agent_scratchpad: list[AgentAction] = Field(default_factory=list)  # agent thoughts / intermediate steps
     tool_responses: list[BaseMessage] = Field(default_factory=list)  # the responses from any tool calls
+    final_answer: str | None = Field(default=None)  # the final answer from the ReAct Agent
 
 
 class ReActAgentGraph(DualNodeAgent):
@@ -68,15 +74,22 @@ class ReActAgentGraph(DualNodeAgent):
                  use_tool_schema: bool = True,
                  callbacks: list[AsyncCallbackHandler] | None = None,
                  detailed_logs: bool = False,
+                 log_response_max_chars: int = 1000,
                  retry_agent_response_parsing_errors: bool = True,
                  parse_agent_response_max_retries: int = 1,
                  tool_call_max_retries: int = 1,
-                 pass_tool_call_errors_to_agent: bool = True):
-        super().__init__(llm=llm, tools=tools, callbacks=callbacks, detailed_logs=detailed_logs)
+                 pass_tool_call_errors_to_agent: bool = True,
+                 normalize_tool_input_quotes: bool = True):
+        super().__init__(llm=llm,
+                         tools=tools,
+                         callbacks=callbacks,
+                         detailed_logs=detailed_logs,
+                         log_response_max_chars=log_response_max_chars)
         self.parse_agent_response_max_retries = (parse_agent_response_max_retries
                                                  if retry_agent_response_parsing_errors else 1)
         self.tool_call_max_retries = tool_call_max_retries
         self.pass_tool_call_errors_to_agent = pass_tool_call_errors_to_agent
+        self.normalize_tool_input_quotes = normalize_tool_input_quotes
         logger.debug(
             "%s Filling the prompt variables 'tools' and 'tool_names', using the tools provided in the config.",
             AGENT_LOG_PREFIX)
@@ -94,21 +107,33 @@ class ReActAgentGraph(DualNodeAgent):
                          f"{INPUT_SCHEMA_MESSAGE.format(schema=tools[-1].input_schema.model_fields)}")
         prompt = prompt.partial(tools=tool_names_and_descriptions, tool_names=tool_names)
         # construct the ReAct Agent
-        bound_llm = llm.bind(stop=["Observation:"])  # type: ignore
-        self.agent = prompt | bound_llm
+        self.agent = prompt | self._maybe_bind_llm_and_yield()
         self.tools_dict = {tool.name: tool for tool in tools}
         logger.debug("%s Initialized ReAct Agent Graph", AGENT_LOG_PREFIX)
 
+    def _maybe_bind_llm_and_yield(self) -> Runnable[LanguageModelInput, BaseMessage]:
+        """
+        Bind additional parameters to the LLM if needed
+        - if the LLM is a smart model, no need to bind any additional parameters
+        - if the LLM is a non-smart model, bind a stop sequence to the LLM
+
+        Returns:
+            Runnable[LanguageModelInput, BaseMessage]: The LLM with any additional parameters bound.
+        """
+        # models that don't need (or don't support)a stop sequence
+        smart_models = re.compile(r"gpt-?5", re.IGNORECASE)
+        if any(smart_models.search(getattr(self.llm, model, "")) for model in ["model", "model_name"]):
+            # no need to bind any additional parameters to the LLM
+            return self.llm
+        # add a stop sequence to the LLM
+        return self.llm.bind(stop=["Observation:"])
+
     def _get_tool(self, tool_name: str):
         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
+            logger.error("%s Unable to find tool with the name %s\n%s", AGENT_LOG_PREFIX, tool_name, ex)
+            raise
 
     async def agent_node(self, state: ReActGraphState):
         try:
@@ -124,17 +149,19 @@ class ReActAgentGraph(DualNodeAgent):
                     if len(state.messages) == 0:
                         raise RuntimeError('No input received in state: "messages"')
                     # to check is any human input passed or not, if no input passed Agent will return the state
-                    content = str(state.messages[0].content)
+                    content = str(state.messages[-1].content)
                     if content.strip() == "":
                         logger.error("%s No human input passed to the agent.", AGENT_LOG_PREFIX)
                         state.messages += [AIMessage(content=NO_INPUT_ERROR_MESSAGE)]
                         return state
                     question = content
                     logger.debug("%s Querying agent, attempt: %s", AGENT_LOG_PREFIX, attempt)
-
+                    chat_history = self._get_chat_history(state.messages)
                     output_message = await self._stream_llm(
                         self.agent,
-                        {"question": question},
+                        {
+                            "question": question, "chat_history": chat_history
+                        },
                         RunnableConfig(callbacks=self.callbacks)  # type: ignore
                     )
 
@@ -152,13 +179,15 @@ class ReActAgentGraph(DualNodeAgent):
                         tool_response = HumanMessage(content=tool_response_content)
                         agent_scratchpad.append(tool_response)
                     agent_scratchpad += working_state
-                    question = str(state.messages[0].content)
+                    chat_history = self._get_chat_history(state.messages)
+                    question = str(state.messages[-1].content)
                     logger.debug("%s Querying agent, attempt: %s", AGENT_LOG_PREFIX, attempt)
 
-                    output_message = await self._stream_llm(self.agent, {
-                        "question": question, "agent_scratchpad": agent_scratchpad
-                    },
-                                                            RunnableConfig(callbacks=self.callbacks))
+                    output_message = await self._stream_llm(
+                        self.agent, {
+                            "question": question, "agent_scratchpad": agent_scratchpad, "chat_history": chat_history
+                        },
+                        RunnableConfig(callbacks=self.callbacks))
 
                     if self.detailed_logs:
                         logger.info(AGENT_CALL_LOG_MESSAGE, question, output_message.content)
@@ -176,6 +205,7 @@ class ReActAgentGraph(DualNodeAgent):
                         # this is where we handle the final output of the Agent, we can clean-up/format/postprocess here
                         # the final answer goes in the "messages" state channel
                         state.messages += [AIMessage(content=final_answer)]
+                        state.final_answer = final_answer
                     else:
                         # the agent wants to call a tool, ensure the thoughts are preserved for the next agentic cycle
                         agent_output.log = output_message.content
@@ -208,16 +238,15 @@ class ReActAgentGraph(DualNodeAgent):
                     working_state.append(output_message)
                     working_state.append(HumanMessage(content=str(ex.observation)))
         except Exception as ex:
-            logger.exception("%s Failed to call agent_node: %s", AGENT_LOG_PREFIX, ex, exc_info=True)
-            raise ex
+            logger.error("%s Failed to call agent_node: %s", AGENT_LOG_PREFIX, ex)
+            raise
 
     async def conditional_edge(self, state: ReActGraphState):
         try:
             logger.debug("%s Starting the ReAct Conditional Edge", AGENT_LOG_PREFIX)
-            if len(state.messages) > 1:
-                # the ReAct Agent has finished executing, the last agent output was AgentFinish
-                last_message_content = str(state.messages[-1].content)
-                logger.debug("%s Final answer:\n%s", AGENT_LOG_PREFIX, last_message_content)
+            if state.final_answer:
+                # the ReAct Agent has finished executing
+                logger.debug("%s Final answer:\n%s", AGENT_LOG_PREFIX, state.final_answer)
                 return AgentDecision.END
             # else the agent wants to call a tool
             agent_output = state.agent_scratchpad[-1]
@@ -227,7 +256,7 @@ class ReActAgentGraph(DualNodeAgent):
                          agent_output.tool_input)
             return AgentDecision.TOOL
         except Exception as ex:
-            logger.exception("Failed to determine whether agent is calling a tool: %s", ex, exc_info=True)
+            logger.exception("Failed to determine whether agent is calling a tool: %s", ex)
             logger.warning("%s Ending graph traversal", AGENT_LOG_PREFIX)
             return AgentDecision.END
 
@@ -260,35 +289,45 @@ class ReActAgentGraph(DualNodeAgent):
                      agent_thoughts.tool_input)
 
         # Run the tool. Try to use structured input, if possible.
+        tool_input_str = agent_thoughts.tool_input.strip()
+
         try:
-            tool_input_str = str(agent_thoughts.tool_input).strip().replace("'", '"')
-            tool_input_dict = json.loads(tool_input_str) if tool_input_str != 'None' else tool_input_str
+            tool_input = json.loads(tool_input_str) if tool_input_str != 'None' else tool_input_str
             logger.debug("%s Successfully parsed structured tool input from Action Input", AGENT_LOG_PREFIX)
 
-            tool_response = await self._call_tool(requested_tool,
-                                                  tool_input_dict,
-                                                  RunnableConfig(callbacks=self.callbacks),
-                                                  max_retries=self.tool_call_max_retries)
-
-            if self.detailed_logs:
-                self._log_tool_response(requested_tool.name, tool_input_dict, str(tool_response.content))
-
-        except JSONDecodeError as ex:
-            logger.debug(
-                "%s Unable to parse structured tool input from Action Input. Using Action Input as is."
-                "\nParsing error: %s",
-                AGENT_LOG_PREFIX,
-                ex,
-                exc_info=True)
-            tool_input_str = str(agent_thoughts.tool_input)
-
-            tool_response = await self._call_tool(requested_tool,
-                                                  tool_input_str,
-                                                  RunnableConfig(callbacks=self.callbacks),
-                                                  max_retries=self.tool_call_max_retries)
+        except JSONDecodeError as original_ex:
+            if self.normalize_tool_input_quotes:
+                # If initial JSON parsing fails, try with quote normalization as a fallback
+                normalized_str = tool_input_str.replace("'", '"')
+                try:
+                    tool_input = json.loads(normalized_str)
+                    logger.debug("%s Successfully parsed structured tool input after quote normalization",
+                                 AGENT_LOG_PREFIX)
+                except JSONDecodeError:
+                    # the quote normalization failed, use raw string input
+                    logger.debug(
+                        "%s Unable to parse structured tool input after quote normalization. Using Action Input as is."
+                        "\nParsing error: %s",
+                        AGENT_LOG_PREFIX,
+                        original_ex)
+                    tool_input = tool_input_str
+            else:
+                # use raw string input
+                logger.debug(
+                    "%s Unable to parse structured tool input from Action Input. Using Action Input as is."
+                    "\nParsing error: %s",
+                    AGENT_LOG_PREFIX,
+                    original_ex)
+                tool_input = tool_input_str
+
+        # Call tool once with the determined input (either parsed dict or raw string)
+        tool_response = await self._call_tool(requested_tool,
+                                              tool_input,
+                                              RunnableConfig(callbacks=self.callbacks),
+                                              max_retries=self.tool_call_max_retries)
 
         if self.detailed_logs:
-            self._log_tool_response(requested_tool.name, tool_input_str, str(tool_response.content))
+            self._log_tool_response(requested_tool.name, tool_input, str(tool_response.content))
 
         if not self.pass_tool_call_errors_to_agent:
             if tool_response.status == "error":
@@ -304,8 +343,8 @@ class ReActAgentGraph(DualNodeAgent):
             logger.debug("%s ReAct Graph built and compiled successfully", AGENT_LOG_PREFIX)
             return self.graph
         except Exception as ex:
-            logger.exception("%s Failed to build ReAct Graph: %s", AGENT_LOG_PREFIX, ex, exc_info=ex)
-            raise ex
+            logger.error("%s Failed to build ReAct Graph: %s", AGENT_LOG_PREFIX, ex)
+            raise
 
     @staticmethod
     def validate_system_prompt(system_prompt: str) -> bool:
@@ -321,12 +360,12 @@ class ReActAgentGraph(DualNodeAgent):
                 errors.append(error_message)
         if errors:
             error_text = "\n".join(errors)
-            logger.exception("%s %s", AGENT_LOG_PREFIX, error_text)
-            raise ValueError(error_text)
+            logger.error("%s %s", AGENT_LOG_PREFIX, error_text)
+            return False
         return True
 
 
-def create_react_agent_prompt(config: ReActAgentWorkflowConfig) -> ChatPromptTemplate:
+def create_react_agent_prompt(config: "ReActAgentWorkflowConfig") -> ChatPromptTemplate:
     """
     Create a ReAct Agent prompt from the config.
 
@@ -348,7 +387,7 @@ def create_react_agent_prompt(config: ReActAgentWorkflowConfig) -> ChatPromptTem
 
     valid_prompt = ReActAgentGraph.validate_system_prompt(prompt_str)
     if not valid_prompt:
-        logger.exception("%s Invalid system_prompt", AGENT_LOG_PREFIX)
+        logger.error("%s Invalid system_prompt", AGENT_LOG_PREFIX)
         raise ValueError("Invalid system_prompt")
     prompt = ChatPromptTemplate([("system", prompt_str), ("user", USER_PROMPT),
                                  MessagesPlaceholder(variable_name='agent_scratchpad', optional=True)])

nat/agent/react_agent/prompt.py

@@ -26,7 +26,7 @@ Use the following format exactly to ask the human to use a tool:
 Question: the input question you must answer
 Thought: you should always think about what to do
 Action: the action to take, should be one of [{tool_names}]
-Action Input: the input to the action (if there is no required input, include "Action Input: None")  
+Action Input: the input to the action (if there is no required input, include "Action Input: None")
 Observation: wait for the human to respond with the result from the tool, do not assume the response
 
 ... (this Thought/Action/Action Input/Observation can repeat N times. If you do not need to use a tool, or after asking the human to use any tools and waiting for the human to respond, you might know the final answer.)
@@ -37,5 +37,8 @@ Final Answer: the final answer to the original input question
 """
 
 USER_PROMPT = """
+Previous conversation history:
+{chat_history}
+
 Question: {question}
 """

nat/agent/react_agent/register.py

@@ -22,26 +22,27 @@ 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.agent import AgentBaseConfig
 from nat.data_models.api_server import ChatRequest
 from nat.data_models.api_server import ChatResponse
+from nat.data_models.component_ref import FunctionGroupRef
 from nat.data_models.component_ref import FunctionRef
-from nat.data_models.component_ref import LLMRef
-from nat.data_models.function import FunctionBaseConfig
+from nat.data_models.optimizable import OptimizableField
+from nat.data_models.optimizable import OptimizableMixin
+from nat.data_models.optimizable import SearchSpace
 from nat.utils.type_converter import GlobalTypeConverter
 
 logger = logging.getLogger(__name__)
 
 
-class ReActAgentWorkflowConfig(FunctionBaseConfig, name="react_agent"):
+class ReActAgentWorkflowConfig(AgentBaseConfig, OptimizableMixin, name="react_agent"):
     """
     Defines a NAT function that uses a ReAct 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 react agent.")
-    llm_name: LLMRef = Field(description="The LLM model to use with the react agent.")
-    verbose: bool = Field(default=False, description="Set the verbosity of the react agent's logging.")
+    description: str = Field(default="ReAct Agent Workflow", description="The description of this functions use.")
+    tool_names: list[FunctionRef | FunctionGroupRef] = Field(
+        default_factory=list, description="The list of tools to provide to the react agent.")
     retry_agent_response_parsing_errors: bool = Field(
         default=True,
         validation_alias=AliasChoices("retry_agent_response_parsing_errors", "retry_parsing_errors"),
@@ -60,7 +61,10 @@ class ReActAgentWorkflowConfig(FunctionBaseConfig, name="react_agent"):
         description="Whether to pass tool call errors to agent. If False, failed tool calls will raise an exception.")
     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="ReAct Agent Workflow", description="The description of this functions use.")
+    normalize_tool_input_quotes: bool = Field(
+        default=True,
+        description="Whether to replace single quotes with double quotes in the tool input. "
+        "This is useful for tools that expect structured json input.")
     system_prompt: str | None = Field(
         default=None,
         description="Provides the SYSTEM_PROMPT to use with the agent")  # defaults to SYSTEM_PROMPT in prompt.py
@@ -68,15 +72,21 @@ class ReActAgentWorkflowConfig(FunctionBaseConfig, name="react_agent"):
     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.")
+    additional_instructions: str | None = OptimizableField(
+        default=None,
+        description="Additional instructions to provide to the agent in addition to the base prompt.",
+        space=SearchSpace(
+            is_prompt=True,
+            prompt="No additional instructions.",
+            prompt_purpose="Additional instructions to provide to the agent in addition to the base prompt.",
+        ))
 
 
 @register_function(config_type=ReActAgentWorkflowConfig, framework_wrappers=[LLMFrameworkEnum.LANGCHAIN])
 async def react_agent_workflow(config: ReActAgentWorkflowConfig, builder: Builder):
     from langchain.schema import BaseMessage
     from langchain_core.messages import trim_messages
-    from langgraph.graph.graph import CompiledGraph
+    from langgraph.graph.state import CompiledStateGraph
 
     from nat.agent.base import AGENT_LOG_PREFIX
     from nat.agent.react_agent.agent import ReActAgentGraph
@@ -89,23 +99,36 @@ async def react_agent_workflow(config: ReActAgentWorkflowConfig, builder: Builde
     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)
+    tools = await builder.get_tools(tool_names=config.tool_names, wrapper_type=LLMFrameworkEnum.LANGCHAIN)
     if not tools:
         raise ValueError(f"No tools specified for ReAct Agent '{config.llm_name}'")
     # configure callbacks, for sending intermediate steps
     # construct the ReAct Agent Graph from the configured llm, prompt, and tools
-    graph: CompiledGraph = await ReActAgentGraph(
+    graph: CompiledStateGraph = await ReActAgentGraph(
         llm=llm,
         prompt=prompt,
         tools=tools,
         use_tool_schema=config.include_tool_input_schema_in_tool_description,
         detailed_logs=config.verbose,
+        log_response_max_chars=config.log_response_max_chars,
         retry_agent_response_parsing_errors=config.retry_agent_response_parsing_errors,
         parse_agent_response_max_retries=config.parse_agent_response_max_retries,
         tool_call_max_retries=config.tool_call_max_retries,
-        pass_tool_call_errors_to_agent=config.pass_tool_call_errors_to_agent).build_graph()
+        pass_tool_call_errors_to_agent=config.pass_tool_call_errors_to_agent,
+        normalize_tool_input_quotes=config.normalize_tool_input_quotes).build_graph()
 
     async def _response_fn(input_message: ChatRequest) -> ChatResponse:
+        """
+        Main workflow entry function for the ReAct Agent.
+
+        This function invokes the ReAct Agent Graph and returns the response.
+
+        Args:
+            input_message (ChatRequest): The input message to process
+
+        Returns:
+            ChatResponse: The response from the agent or error message
+        """
         try:
             # initialize the starting state with the user query
             messages: list[BaseMessage] = trim_messages(messages=[m.model_dump() for m in input_message.messages],
@@ -125,15 +148,12 @@ async def react_agent_workflow(config: ReActAgentWorkflowConfig, builder: Builde
 
             # get and return the output from the state
             state = ReActGraphState(**state)
-            output_message = state.messages[-1]  # pylint: disable=E1136
+            output_message = state.messages[-1]
             return ChatResponse.from_string(str(output_message.content))
 
         except Exception as ex:
-            logger.exception("%s ReAct Agent failed with exception: %s", AGENT_LOG_PREFIX, ex, exc_info=ex)
-            # here, we can implement custom error messages
-            if config.verbose:
-                return ChatResponse.from_string(str(ex))
-            return ChatResponse.from_string("I seem to be having a problem.")
+            logger.exception("%s ReAct Agent failed with exception: %s", AGENT_LOG_PREFIX, str(ex))
+            raise RuntimeError
 
     if (config.use_openai_api):
         yield FunctionInfo.from_fn(_response_fn, description=config.description)

nat/agent/reasoning_agent/reasoning_agent.py

@@ -23,25 +23,22 @@ 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.agent import AgentBaseConfig
 from nat.data_models.api_server import ChatRequest
 from nat.data_models.component_ref import FunctionRef
-from nat.data_models.component_ref import LLMRef
-from nat.data_models.function import FunctionBaseConfig
 
 logger = logging.getLogger(__name__)
 
 
-class ReasoningFunctionConfig(FunctionBaseConfig, name="reasoning_agent"):
+class ReasoningFunctionConfig(AgentBaseConfig, name="reasoning_agent"):
     """
     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.
     """
-
-    llm_name: LLMRef = Field(description="The name of the LLM to use for reasoning.")
+    description: str = Field(default="Reasoning Agent", description="The description of this function's use.")
     augmented_fn: FunctionRef = Field(description="The name of the function to reason on.")
-    verbose: bool = Field(default=False, description="Whether to log detailed information.")
     reasoning_prompt_template: str = Field(
         default=("You are an expert reasoning model task with creating a detailed execution plan"
                  " for a system that has the following description:\n\n"
@@ -102,7 +99,7 @@ async def build_reasoning_function(config: ReasoningFunctionConfig, builder: Bui
     llm: BaseChatModel = await builder.get_llm(config.llm_name, wrapper_type=LLMFrameworkEnum.LANGCHAIN)
 
     # Get the augmented function's description
-    augmented_function = builder.get_function(config.augmented_fn)
+    augmented_function = await builder.get_function(config.augmented_fn)
 
     # For now, we rely on runtime checking for type conversion
 
@@ -113,11 +110,16 @@ async def build_reasoning_function(config: ReasoningFunctionConfig, builder: Bui
                          f"function without a description.")
 
     # Get the function dependencies of the augmented function
-    function_used_tools = builder.get_function_dependencies(config.augmented_fn).functions
+    function_dependencies = builder.get_function_dependencies(config.augmented_fn)
+    function_used_tools = set()
+    function_used_tools.update(function_dependencies.functions)
+    for function_group in function_dependencies.function_groups:
+        function_used_tools.update(builder.get_function_group_dependencies(function_group).functions)
+
     tool_names_with_desc: list[tuple[str, str]] = []
 
     for tool in function_used_tools:
-        tool_impl = builder.get_function(tool)
+        tool_impl = await builder.get_function(tool)
         tool_names_with_desc.append((tool, tool_impl.description if hasattr(tool_impl, "description") else ""))
 
     # Draft the reasoning prompt for the augmented function

nat/agent/register.py

@@ -13,10 +13,10 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-# pylint: disable=unused-import
 # flake8: noqa
 
 # Import any workflows which need to be automatically registered here
+from .prompt_optimizer import register as prompt_optimizer
 from .react_agent import register as react_agent
 from .reasoning_agent import reasoning_agent
 from .rewoo_agent import register as rewoo_agent