nvidia-nat 1.3.0.dev2__py3-none-any.whl → 1.3.0rc2__py3-none-any.whl

nat/agent/rewoo_agent/prompt.py

@@ -18,33 +18,29 @@ For the following task, make plans that can solve the problem step by step. For
 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:
+The following tools and respective requirements are available to you:
 
 {tools}
 
-The tools should be one of the following: [{tool_names}]
+The tool calls you make 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
-
+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
+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.
+   -"placeholder": A string that identifies the evidence placeholder ("#E1", "#E2", ...). 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.
+   -"tool_input": The input to the tool. This can be a string, array, or object, depending on the requirements of the \
+tool. Be careful about type assumptions because the output of former tools might contain noise.
 
 Important instructions:
 
@@ -58,27 +54,28 @@ Here is an example of how a valid JSON output should look:
 
 [
   \'{{
-    "plan": "Calculate the result of 2023 minus 25.",
+    "plan": "Find Alex's schedule on Sep 25, 2025",
     "evidence": \'{{
       "placeholder": "#E1",
-      "tool": "calculator_subtract",
-      "tool_input": [2023, 25]
+      "tool": "search_calendar",
+      "tool_input": ("Alex", "09/25/2025")
     }}\'
   }}\',
   \'{{
-    "plan": "Retrieve the year represented by the result stored in #E1.",
+    "plan": "Find Bill's schedule on sep 25, 2025",
     "evidence": \'{{
       "placeholder": "#E2",
-      "tool": "haystack_chitchat_agent",
-      "tool_input": "Response with the result number contained in #E1"
+      "tool": "search_calendar",
+      "tool_input": ("Bill", "09/25/2025")
     }}\'
   }}\',
   \'{{
-    "plan": "Search for the CEO of Golden State Warriors in the year stored in #E2.",
+    "plan": "Suggest a time for 1-hour meeting given Alex's and Bill's schedule.",
     "evidence": \'{{
       "placeholder": "#E3",
-      "tool": "internet_search",
-      "tool_input": "Who was the CEO of Golden State Warriors in the year #E2?"
+      "tool": "llm_chat",
+      "tool_input": "Find a common 1-hour time slot for Alex and Bill given their schedules. \
+Alex's schedule: #E1; Bill's schedule: #E2?"
     }}\'
   }}\'
 ]
@@ -94,7 +91,7 @@ task: {task}
 """
 
 SOLVER_SYSTEM_PROMPT = """
-Solve the following task or problem. To solve the problem, we have made step-by-step Plan and \
+Solve the following task or problem. To solve the problem, we have made some Plans ahead and \
 retrieved corresponding Evidence to each Plan. Use them with caution since long evidence might \
 contain irrelevant information.
 

nat/agent/rewoo_agent/register.py

@@ -17,40 +17,42 @@ import logging
 
 from pydantic import AliasChoices
 from pydantic import Field
+from pydantic import PositiveInt
 
 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.api_server import Usage
+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.utils.type_converter import GlobalTypeConverter
 
 logger = logging.getLogger(__name__)
 
 
-class ReWOOAgentWorkflowConfig(FunctionBaseConfig, name="rewoo_agent"):
+class ReWOOAgentWorkflowConfig(AgentBaseConfig, name="rewoo_agent"):
     """
     Defines a NAT 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.")
+    description: str = Field(default="ReWOO 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 rewoo agent.")
     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
+    tool_call_max_retries: PositiveInt = Field(default=3,
+                                               description="The number of retries before raising a tool call error.",
+                                               ge=1)
     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. "
@@ -62,15 +64,19 @@ class ReWOOAgentWorkflowConfig(FunctionBaseConfig, name="rewoo_agent"):
     additional_solver_instructions: str | None = Field(
         default=None,
         description="Additional instructions to provide to the agent in addition to the base solver prompt.")
+    raise_tool_call_error: bool = Field(default=True,
+                                        description="Whether to raise a exception immediately if a tool"
+                                        "call fails. If set to False, the tool call error message will be included in"
+                                        "the tool response and passed to the next tool.")
 
 
 @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.base import BaseMessage
     from langchain_core.messages.human import HumanMessage
     from langchain_core.prompts import ChatPromptTemplate
-    from langgraph.graph.graph import CompiledGraph
+    from langgraph.graph.state import CompiledStateGraph
 
     from nat.agent.rewoo_agent.prompt import PLANNER_SYSTEM_PROMPT
     from nat.agent.rewoo_agent.prompt import PLANNER_USER_PROMPT
@@ -86,7 +92,7 @@ async def rewoo_agent_workflow(config: ReWOOAgentWorkflowConfig, builder: Builde
     if config.additional_planner_instructions:
         planner_system_prompt += f"{config.additional_planner_instructions}"
     if not ReWOOAgentGraph.validate_planner_prompt(planner_system_prompt):
-        logger.exception("Invalid planner prompt")
+        logger.error("Invalid planner prompt")
         raise ValueError("Invalid planner prompt")
     planner_prompt = ChatPromptTemplate([("system", planner_system_prompt), ("user", PLANNER_USER_PROMPT)])
 
@@ -94,7 +100,7 @@ async def rewoo_agent_workflow(config: ReWOOAgentWorkflowConfig, builder: Builde
     if config.additional_solver_instructions:
         solver_system_prompt += f"{config.additional_solver_instructions}"
     if not ReWOOAgentGraph.validate_solver_prompt(solver_system_prompt):
-        logger.exception("Invalid solver prompt")
+        logger.error("Invalid solver prompt")
         raise ValueError("Invalid solver prompt")
     solver_prompt = ChatPromptTemplate([("system", solver_system_prompt), ("user", SOLVER_USER_PROMPT)])
 
@@ -103,19 +109,34 @@ async def rewoo_agent_workflow(config: ReWOOAgentWorkflowConfig, builder: Builde
 
     # 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 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()
+    graph: CompiledStateGraph = 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,
+        log_response_max_chars=config.log_response_max_chars,
+        tool_call_max_retries=config.tool_call_max_retries,
+        raise_tool_call_error=config.raise_tool_call_error).build_graph()
 
     async def _response_fn(input_message: ChatRequest) -> ChatResponse:
+        """
+        Main workflow entry function for the ReWOO Agent.
+
+        This function invokes the ReWOO 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],
@@ -133,15 +154,21 @@ async def rewoo_agent_workflow(config: ReWOOAgentWorkflowConfig, builder: Builde
 
             # get and return the output from the state
             state = ReWOOGraphState(**state)
-            output_message = state.result.content  # pylint: disable=E1101
-            return ChatResponse.from_string(output_message)
+            output_message = state.result.content
+            # Ensure output_message is a string
+            if isinstance(output_message, list | dict):
+                output_message = str(output_message)
+
+            # Create usage statistics for the response
+            prompt_tokens = sum(len(str(msg.content).split()) for msg in input_message.messages)
+            completion_tokens = len(output_message.split()) if output_message else 0
+            total_tokens = prompt_tokens + completion_tokens
+            usage = Usage(prompt_tokens=prompt_tokens, completion_tokens=completion_tokens, total_tokens=total_tokens)
+            return ChatResponse.from_string(output_message, usage=usage)
 
         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 ChatResponse.from_string(str(ex))
-            return ChatResponse.from_string("I seem to be having a problem.")
+            logger.exception("ReWOO Agent failed with exception: %s", ex)
+            raise RuntimeError
 
     if (config.use_openai_api):
         yield FunctionInfo.from_fn(_response_fn, description=config.description)

nat/agent/tool_calling_agent/agent.py

@@ -13,17 +13,19 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-# pylint: disable=R0917
 import logging
 import typing
 
 from langchain_core.callbacks.base import AsyncCallbackHandler
 from langchain_core.language_models import BaseChatModel
 from langchain_core.messages import SystemMessage
+from langchain_core.messages import ToolMessage
 from langchain_core.messages.base import BaseMessage
 from langchain_core.runnables import RunnableLambda
 from langchain_core.runnables.config import RunnableConfig
 from langchain_core.tools import BaseTool
+from langgraph.graph import StateGraph
+from langgraph.graph.state import CompiledStateGraph
 from langgraph.prebuilt import ToolNode
 from pydantic import BaseModel
 from pydantic import Field
@@ -56,17 +58,24 @@ class ToolCallAgentGraph(DualNodeAgent):
         prompt: str | None = None,
         callbacks: list[AsyncCallbackHandler] = None,
         detailed_logs: bool = False,
+        log_response_max_chars: int = 1000,
         handle_tool_errors: bool = True,
+        return_direct: list[BaseTool] | None = None,
     ):
-        super().__init__(llm=llm, tools=tools, callbacks=callbacks, detailed_logs=detailed_logs)
+        super().__init__(llm=llm,
+                         tools=tools,
+                         callbacks=callbacks,
+                         detailed_logs=detailed_logs,
+                         log_response_max_chars=log_response_max_chars)
+
         # some LLMs support tool calling
         # these models accept the tool's input schema and decide when to use a tool based on the input's relevance
         try:
             # in tool calling agents, we bind the tools to the LLM, to pass the tools' input schemas at runtime
             self.bound_llm = llm.bind_tools(tools)
         except NotImplementedError as ex:
-            logger.error("%s Failed to bind tools: %s", AGENT_LOG_PREFIX, ex, exc_info=True)
-            raise ex
+            logger.error("%s Failed to bind tools: %s", AGENT_LOG_PREFIX, ex)
+            raise
 
         if prompt is not None:
             system_prompt = SystemMessage(content=prompt)
@@ -81,8 +90,8 @@ class ToolCallAgentGraph(DualNodeAgent):
             )
 
         self.agent = prompt_runnable | self.bound_llm
-
         self.tool_caller = ToolNode(tools, handle_tool_errors=handle_tool_errors)
+        self.return_direct = [tool.name for tool in return_direct] if return_direct else []
         logger.debug("%s Initialized Tool Calling Agent Graph", AGENT_LOG_PREFIX)
 
     async def agent_node(self, state: ToolCallAgentGraphState):
@@ -101,8 +110,8 @@ class ToolCallAgentGraph(DualNodeAgent):
             state.messages += [response]
             return state
         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: ToolCallAgentGraphState):
         try:
@@ -116,12 +125,7 @@ class ToolCallAgentGraph(DualNodeAgent):
                 logger.debug("%s Final answer:\n%s", AGENT_LOG_PREFIX, state.messages[-1].content)
             return AgentDecision.END
         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.exception("%s Failed to determine whether agent is calling a tool: %s", AGENT_LOG_PREFIX, ex)
             logger.warning("%s Ending graph traversal", AGENT_LOG_PREFIX)
             return AgentDecision.END
 
@@ -144,25 +148,77 @@ class ToolCallAgentGraph(DualNodeAgent):
 
             return state
         except Exception as ex:
-            logger.exception("%s Failed to call tool_node: %s", AGENT_LOG_PREFIX, ex, exc_info=ex)
-            raise ex
+            logger.error("%s Failed to call tool_node: %s", AGENT_LOG_PREFIX, ex)
+            raise
 
-    async def build_graph(self):
+    async def tool_conditional_edge(self, state: ToolCallAgentGraphState) -> AgentDecision:
+        """
+        Determines whether to continue to the agent or end graph execution after a tool call.
+
+        Args:
+            state: The current state of the Tool Calling Agent graph containing messages and tool responses.
+
+        Returns:
+            AgentDecision: TOOL to continue to agent for processing, or END to terminate graph execution.
+            Returns END if the tool is in return_direct list, otherwise returns TOOL to continue processing.
+        """
         try:
-            await super()._build_graph(state_schema=ToolCallAgentGraphState)
-            logger.debug(
-                "%s Tool Calling Agent Graph built and compiled successfully",
-                AGENT_LOG_PREFIX,
-            )
+            logger.debug("%s Starting the Tool Conditional Edge", AGENT_LOG_PREFIX)
+            if not state.messages:
+                logger.debug("%s No messages in state; routing to agent", AGENT_LOG_PREFIX)
+                return AgentDecision.TOOL
+
+            last_message = state.messages[-1]
+            # Return directly if this tool is in the return_direct set
+            if (self.return_direct and isinstance(last_message, ToolMessage) and last_message.name
+                    and last_message.name in self.return_direct):
+                # Return directly if this tool is in the return_direct list
+                logger.debug("%s Tool %s is set to return directly", AGENT_LOG_PREFIX, last_message.name)
+                return AgentDecision.END
+            else:
+                # Continue to agent for processing
+                logger.debug("%s Tool response will be processed by agent", AGENT_LOG_PREFIX)
+                return AgentDecision.TOOL
+        except Exception as ex:
+            logger.exception("%s Failed to determine tool conditional edge: %s", AGENT_LOG_PREFIX, ex)
+            logger.warning("%s Continuing to agent for processing", AGENT_LOG_PREFIX)
+            return AgentDecision.TOOL
+
+    async def _build_graph(self, state_schema: type) -> CompiledStateGraph:
+        try:
+            logger.debug("%s Building and compiling the Tool Calling Agent Graph", AGENT_LOG_PREFIX)
+
+            graph = StateGraph(state_schema)
+            graph.add_node("agent", self.agent_node)
+            graph.add_node("tool", self.tool_node)
+
+            if self.return_direct:
+                # go to end of graph if tool is set to return directly
+                tool_conditional_edge_possible_outputs = {AgentDecision.END: "__end__", AgentDecision.TOOL: "agent"}
+                graph.add_conditional_edges("tool", self.tool_conditional_edge, tool_conditional_edge_possible_outputs)
+            else:
+                # otherwise return to agent after tool call
+                graph.add_edge("tool", "agent")
+
+            conditional_edge_possible_outputs = {AgentDecision.TOOL: "tool", AgentDecision.END: "__end__"}
+            graph.add_conditional_edges("agent", self.conditional_edge, conditional_edge_possible_outputs)
+
+            graph.set_entry_point("agent")
+            self.graph = graph.compile()
+
             return self.graph
         except Exception as ex:
-            logger.exception(
-                "%s Failed to build Tool Calling Agent Graph: %s",
-                AGENT_LOG_PREFIX,
-                ex,
-                exc_info=ex,
-            )
-            raise ex
+            logger.error("%s Failed to build Tool Calling Agent Graph: %s", AGENT_LOG_PREFIX, ex)
+            raise
+
+    async def build_graph(self) -> CompiledStateGraph:
+        try:
+            await self._build_graph(state_schema=ToolCallAgentGraphState)
+            logger.debug("%s Tool Calling Agent Graph built and compiled successfully", AGENT_LOG_PREFIX)
+            return self.graph
+        except Exception as ex:
+            logger.error("%s Failed to build Tool Calling Agent Graph: %s", AGENT_LOG_PREFIX, ex)
+            raise
 
 
 def create_tool_calling_agent_prompt(config: "ToolCallAgentWorkflowConfig") -> str | None:

nat/agent/tool_calling_agent/register.py

@@ -21,35 +21,38 @@ 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 FunctionGroupRef
 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 ToolCallAgentWorkflowConfig(FunctionBaseConfig, name="tool_calling_agent"):
+class ToolCallAgentWorkflowConfig(AgentBaseConfig, name="tool_calling_agent"):
     """
     A Tool Calling Agent requires an LLM which supports tool calling. A tool Calling Agent utilizes the tool
     input parameters to select the optimal tool.  Supports handling tool errors.
     """
-
-    tool_names: list[FunctionRef] = Field(default_factory=list,
-                                          description="The list of tools to provide to the tool calling agent.")
-    llm_name: LLMRef = Field(description="The LLM model to use with the tool calling agent.")
-    verbose: bool = Field(default=False, description="Set the verbosity of the tool calling agent's logging.")
-    handle_tool_errors: bool = Field(default=True, description="Specify ability to handle tool calling errors.")
     description: str = Field(default="Tool Calling Agent Workflow", description="Description of this functions use.")
+    tool_names: list[FunctionRef | FunctionGroupRef] = Field(
+        default_factory=list, description="The list of tools to provide to the tool calling agent.")
+    handle_tool_errors: bool = Field(default=True, description="Specify ability to handle tool calling errors.")
     max_iterations: int = Field(default=15, description="Number of tool calls before stoping the tool calling agent.")
+    max_history: int = Field(default=15, description="Maximum number of messages to keep in the conversation history.")
+
     system_prompt: str | None = Field(default=None, description="Provides the system prompt to use with the agent.")
     additional_instructions: str | None = Field(default=None,
                                                 description="Additional instructions appended to the system prompt.")
+    return_direct: list[FunctionRef] | None = Field(
+        default=None, description="List of tool names that should return responses directly without LLM processing.")
 
 
 @register_function(config_type=ToolCallAgentWorkflowConfig, framework_wrappers=[LLMFrameworkEnum.LANGCHAIN])
 async def tool_calling_agent_workflow(config: ToolCallAgentWorkflowConfig, builder: Builder):
-    from langchain_core.messages.human import HumanMessage
-    from langgraph.graph.graph import CompiledGraph
+    from langchain_core.messages import trim_messages
+    from langchain_core.messages.base import BaseMessage
+    from langgraph.graph.state import CompiledStateGraph
 
     from nat.agent.base import AGENT_LOG_PREFIX
     from nat.agent.tool_calling_agent.agent import ToolCallAgentGraph
@@ -61,22 +64,44 @@ async def tool_calling_agent_workflow(config: ToolCallAgentWorkflowConfig, build
     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 tools 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 Tool Calling Agent '{config.llm_name}'")
 
-    # construct the Tool Calling Agent Graph from the configured llm, and tools
-    graph: CompiledGraph = await ToolCallAgentGraph(llm=llm,
-                                                    tools=tools,
-                                                    prompt=prompt,
-                                                    detailed_logs=config.verbose,
-                                                    handle_tool_errors=config.handle_tool_errors).build_graph()
+    # convert return_direct FunctionRef objects to BaseTool objects
+    return_direct_tools = await builder.get_tools(
+        tool_names=config.return_direct, wrapper_type=LLMFrameworkEnum.LANGCHAIN) if config.return_direct else None
 
-    async def _response_fn(input_message: str) -> str:
+    # construct the Tool Calling Agent Graph from the configured llm, and tools
+    graph: CompiledStateGraph = await ToolCallAgentGraph(llm=llm,
+                                                         tools=tools,
+                                                         prompt=prompt,
+                                                         detailed_logs=config.verbose,
+                                                         log_response_max_chars=config.log_response_max_chars,
+                                                         handle_tool_errors=config.handle_tool_errors,
+                                                         return_direct=return_direct_tools).build_graph()
+
+    async def _response_fn(input_message: ChatRequest) -> str:
+        """
+        Main workflow entry function for the Tool Calling Agent.
+
+        This function invokes the Tool Calling Agent Graph and returns the response.
+
+        Args:
+            input_message (ChatRequest): The input message to process
+
+        Returns:
+            str: The response from the agent or error message
+        """
         try:
             # initialize the starting state with the user query
-            input_message = HumanMessage(content=input_message)
-            state = ToolCallAgentGraphState(messages=[input_message])
+            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)
+            state = ToolCallAgentGraphState(messages=messages)
 
             # run the Tool Calling Agent Graph
             state = await graph.ainvoke(state, config={'recursion_limit': (config.max_iterations + 1) * 2})
@@ -86,17 +111,15 @@ async def tool_calling_agent_workflow(config: ToolCallAgentWorkflowConfig, build
 
             # get and return the output from the state
             state = ToolCallAgentGraphState(**state)
-            output_message = state.messages[-1]  # pylint: disable=E1136
-            return output_message.content
+            output_message = state.messages[-1]
+            return str(output_message.content)
         except Exception as ex:
-            logger.exception("%s Tool Calling Agent failed with exception: %s", AGENT_LOG_PREFIX, ex, exc_info=ex)
-            if config.verbose:
-                return str(ex)
-            return "I seem to be having a problem."
+            logger.exception("%s Tool Calling Agent failed with exception: %s", AGENT_LOG_PREFIX, ex)
+            raise RuntimeError
 
     try:
         yield FunctionInfo.from_fn(_response_fn, description=config.description)
     except GeneratorExit:
-        logger.exception("%s Workflow exited early!", AGENT_LOG_PREFIX, exc_info=True)
+        logger.exception("%s Workflow exited early!", AGENT_LOG_PREFIX)
     finally:
         logger.debug("%s Cleaning up react_agent workflow.", AGENT_LOG_PREFIX)

nat/authentication/api_key/api_key_auth_provider.py

@@ -31,7 +31,7 @@ class APIKeyAuthProvider(AuthProviderBase[APIKeyAuthProviderConfig]):
     # fmt: off
     def __init__(self,
                  config: APIKeyAuthProviderConfig,
-                 config_name: str | None = None) -> None:  # pylint: disable=unused-argument
+                 config_name: str | None = None) -> None:
         assert isinstance(config, APIKeyAuthProviderConfig), ("Config is not APIKeyAuthProviderConfig")
         super().__init__(config)
     # fmt: on
@@ -80,7 +80,7 @@ class APIKeyAuthProvider(AuthProviderBase[APIKeyAuthProviderConfig]):
 
         raise ValueError(f"Unsupported header auth scheme: {header_auth_scheme}")
 
-    async def authenticate(self, user_id: str | None = None) -> AuthResult | None:
+    async def authenticate(self, user_id: str | None = None, **kwargs) -> AuthResult | None:
         """
         Authenticate the user using the API key credentials.