nvidia-nat 1.3.0a20250909__py3-none-any.whl → 1.3.0a20250917__py3-none-any.whl

nat/agent/router_agent/agent.py

@@ -0,0 +1,329 @@
+# 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
+import typing
+
+from langchain_core.callbacks.base import AsyncCallbackHandler
+from langchain_core.language_models import BaseChatModel
+from langchain_core.messages.base import BaseMessage
+from langchain_core.messages.human import HumanMessage
+from langchain_core.prompts.chat import ChatPromptTemplate
+from langchain_core.tools import BaseTool
+from langgraph.graph import StateGraph
+from pydantic import BaseModel
+from pydantic import Field
+
+from nat.agent.base import AGENT_CALL_LOG_MESSAGE
+from nat.agent.base import AGENT_LOG_PREFIX
+from nat.agent.base import BaseAgent
+
+if typing.TYPE_CHECKING:
+    from nat.agent.router_agent.register import RouterAgentWorkflowConfig
+
+logger = logging.getLogger(__name__)
+
+
+class RouterAgentGraphState(BaseModel):
+    """State schema for the Router Agent Graph.
+
+    This class defines the state structure used throughout the Router Agent's
+    execution graph, containing messages, routing information, and branch selection.
+
+    Attributes:
+        messages: A list of messages representing the conversation history.
+        forward_message: The message to be forwarded to the chosen branch.
+        chosen_branch: The name of the branch selected by the router agent.
+    """
+    messages: list[BaseMessage] = Field(default_factory=list)
+    forward_message: BaseMessage = Field(default_factory=lambda: HumanMessage(content=""))
+    chosen_branch: str = Field(default="")
+
+
+class RouterAgentGraph(BaseAgent):
+    """Configurable Router Agent for routing requests to different branches.
+
+    A Router Agent analyzes incoming requests and routes them to one of the
+    configured branches based on the conte nt and context. It makes a single
+    routing decision and executes only the selected branch before returning.
+
+    This agent is useful for creating multi-path workflows where different
+    types of requests need to be handled by specialized sub-agents or tools.
+    """
+
+    def __init__(
+        self,
+        llm: BaseChatModel,
+        branches: list[BaseTool],
+        prompt: ChatPromptTemplate,
+        max_router_retries: int = 3,
+        callbacks: list[AsyncCallbackHandler] | None = None,
+        detailed_logs: bool = False,
+        log_response_max_chars: int = 1000,
+    ):
+        """Initialize the Router Agent.
+
+        Args:
+            llm: The language model to use for routing decisions.
+            branches: List of tools/branches that the agent can route to.
+            prompt: The chat prompt template for the routing agent.
+            max_router_retries: Maximum number of retries if branch selection fails.
+            callbacks: Optional list of async callback handlers.
+            detailed_logs: Whether to enable detailed logging.
+            log_response_max_chars: Maximum characters to log in responses.
+        """
+        super().__init__(llm=llm,
+                         tools=branches,
+                         callbacks=callbacks,
+                         detailed_logs=detailed_logs,
+                         log_response_max_chars=log_response_max_chars)
+
+        self._branches = branches
+        self._branches_dict = {branch.name: branch for branch in branches}
+        branch_names = ",".join([branch.name for branch in branches])
+        branch_names_and_descriptions = "\n".join([f"{branch.name}: {branch.description}" for branch in branches])
+
+        prompt = prompt.partial(branches=branch_names_and_descriptions, branch_names=branch_names)
+        self.agent = prompt | self.llm
+
+        self.max_router_retries = max_router_retries
+
+    def _get_branch(self, branch_name: str) -> BaseTool | None:
+        return self._branches_dict.get(branch_name, None)
+
+    async def agent_node(self, state: RouterAgentGraphState):
+        """Execute the agent node to select a branch for routing.
+
+        This method processes the incoming request and determines which branch
+        should handle it. It uses the configured LLM to analyze the request
+        and select the most appropriate branch.
+
+        Args:
+            state: The current state of the router agent graph.
+
+        Returns:
+            RouterAgentGraphState: Updated state with the chosen branch.
+
+        Raises:
+            RuntimeError: If the agent fails to choose a branch after max retries.
+        """
+        logger.debug("%s Starting the Router Agent Node", AGENT_LOG_PREFIX)
+        chat_history = self._get_chat_history(state.messages)
+        request = state.forward_message.content
+        for attempt in range(1, self.max_router_retries + 1):
+            try:
+                agent_response = await self._call_llm(self.agent, {"request": request, "chat_history": chat_history})
+                if self.detailed_logs:
+                    logger.info(AGENT_CALL_LOG_MESSAGE, request, agent_response)
+
+                state.messages += [agent_response]
+
+                # Determine chosen branch based on agent response
+                if state.chosen_branch == "":
+                    for branch in self._branches:
+                        if branch.name.lower() in str(agent_response.content).lower():
+                            state.chosen_branch = branch.name
+                            if self.detailed_logs:
+                                logger.debug("%s Router Agent has chosen branch: %s", AGENT_LOG_PREFIX, branch.name)
+                            return state
+
+                # The agent failed to choose a branch
+                if state.chosen_branch == "":
+                    if attempt == self.max_router_retries:
+                        logger.error("%s Router Agent has empty chosen branch", AGENT_LOG_PREFIX)
+                        raise RuntimeError("Router Agent failed to choose a branch")
+                    logger.warning("%s Router Agent failed to choose a branch, retrying %d out of %d",
+                                   AGENT_LOG_PREFIX,
+                                   attempt,
+                                   self.max_router_retries)
+
+            except Exception as ex:
+                logger.error("%s Router Agent failed to call agent_node: %s", AGENT_LOG_PREFIX, ex)
+                raise
+
+        return state
+
+    async def branch_node(self, state: RouterAgentGraphState):
+        """Execute the selected branch with the forwarded message.
+
+        This method calls the tool/branch that was selected by the agent node
+        and processes the response.
+
+        Args:
+            state: The current state containing the chosen branch and message.
+
+        Returns:
+            RouterAgentGraphState: Updated state with the branch response.
+
+        Raises:
+            RuntimeError: If no branch was chosen or branch execution fails.
+            ValueError: If the requested tool is not found in the configuration.
+        """
+        logger.debug("%s Starting Router Agent Tool Node", AGENT_LOG_PREFIX)
+        try:
+            if state.chosen_branch == "":
+                logger.error("%s Router Agent has empty chosen branch", AGENT_LOG_PREFIX)
+                raise RuntimeError("Router Agent failed to choose a branch")
+            requested_branch = self._get_branch(state.chosen_branch)
+            if not requested_branch:
+                logger.error("%s Router Agent wants to call tool %s but it is not in the config file",
+                             AGENT_LOG_PREFIX,
+                             state.chosen_branch)
+                raise ValueError("Tool not found in config file")
+
+            branch_input = state.forward_message.content
+            branch_response = await self._call_tool(requested_branch, branch_input)
+            state.messages += [branch_response]
+            if self.detailed_logs:
+                self._log_tool_response(requested_branch.name, branch_input, branch_response.content)
+
+            return state
+
+        except Exception as ex:
+            logger.error("%s Router Agent throws exception during branch node execution: %s", AGENT_LOG_PREFIX, ex)
+            raise
+
+    async def _build_graph(self, state_schema):
+        logger.debug("%s Building and compiling the Router Agent Graph", AGENT_LOG_PREFIX)
+
+        graph = StateGraph(state_schema)
+        graph.add_node("agent", self.agent_node)
+        graph.add_node("branch", self.branch_node)
+        graph.add_edge("agent", "branch")
+        graph.set_entry_point("agent")
+
+        self.graph = graph.compile()
+        logger.debug("%s Router Agent Graph built and compiled successfully", AGENT_LOG_PREFIX)
+
+        return self.graph
+
+    async def build_graph(self):
+        """Build and compile the router agent execution graph.
+
+        Creates a state graph with agent and branch nodes, configures the
+        execution flow, and compiles the graph for execution.
+
+        Returns:
+            The compiled execution graph.
+
+        Raises:
+            Exception: If graph building or compilation fails.
+        """
+        try:
+            await self._build_graph(state_schema=RouterAgentGraphState)
+            return self.graph
+        except Exception as ex:
+            logger.error("%s Router Agent failed to build graph: %s", AGENT_LOG_PREFIX, ex)
+            raise
+
+    @staticmethod
+    def validate_system_prompt(system_prompt: str) -> bool:
+        """Validate that the system prompt contains required variables.
+
+        Checks that the system prompt includes necessary template variables
+        for branch information that the router agent needs.
+
+        Args:
+            system_prompt: The system prompt string to validate.
+
+        Returns:
+            True if the prompt is valid, False otherwise.
+        """
+        errors = []
+        required_prompt_variables = {
+            "{branches}": "The system prompt must contain {branches} so the agent knows about configured branches.",
+            "{branch_names}": "The system prompt must contain {branch_names} so the agent knows branch names."
+        }
+        for variable_name, error_message in required_prompt_variables.items():
+            if variable_name not in system_prompt:
+                errors.append(error_message)
+        if errors:
+            error_text = "\n".join(errors)
+            logger.error("%s %s", AGENT_LOG_PREFIX, error_text)
+            return False
+        return True
+
+    @staticmethod
+    def validate_user_prompt(user_prompt: str) -> bool:
+        """Validate that the user prompt contains required variables.
+
+        Checks that the user prompt includes necessary template variables
+        for chat history and other required information.
+
+        Args:
+            user_prompt: The user prompt string to validate.
+
+        Returns:
+            True if the prompt is valid, False otherwise.
+        """
+        errors = []
+        if not user_prompt:
+            errors.append("The user prompt cannot be empty.")
+        else:
+            required_prompt_variables = {
+                "{chat_history}":
+                    "The user prompt must contain {chat_history} so the agent knows about the conversation history.",
+                "{request}":
+                    "The user prompt must contain {request} so the agent sees the current request.",
+            }
+            for variable_name, error_message in required_prompt_variables.items():
+                if variable_name not in user_prompt:
+                    errors.append(error_message)
+        if errors:
+            error_text = "\n".join(errors)
+            logger.error("%s %s", AGENT_LOG_PREFIX, error_text)
+            return False
+        return True
+
+
+def create_router_agent_prompt(config: "RouterAgentWorkflowConfig") -> ChatPromptTemplate:
+    """Create a Router Agent prompt from the configuration.
+
+    Builds a ChatPromptTemplate using either custom prompts from the config
+    or default system and user prompts. Validates the prompts to ensure they
+    contain required template variables.
+
+    Args:
+        config: The router agent workflow configuration containing prompt settings.
+
+    Returns:
+        A configured ChatPromptTemplate for the router agent.
+
+    Raises:
+        ValueError: If the system_prompt or user_prompt validation fails.
+    """
+    from nat.agent.router_agent.prompt import SYSTEM_PROMPT
+    from nat.agent.router_agent.prompt import USER_PROMPT
+    # the Router Agent prompt can be customized via config option system_prompt and user_prompt.
+
+    if config.system_prompt:
+        system_prompt = config.system_prompt
+    else:
+        system_prompt = SYSTEM_PROMPT
+
+    if config.user_prompt:
+        user_prompt = config.user_prompt
+    else:
+        user_prompt = USER_PROMPT
+
+    if not RouterAgentGraph.validate_system_prompt(system_prompt):
+        logger.error("%s Invalid system_prompt", AGENT_LOG_PREFIX)
+        raise ValueError("Invalid system_prompt")
+
+    if not RouterAgentGraph.validate_user_prompt(user_prompt):
+        logger.error("%s Invalid user_prompt", AGENT_LOG_PREFIX)
+        raise ValueError("Invalid user_prompt")
+
+    return ChatPromptTemplate([("system", system_prompt), ("user", user_prompt)])

nat/agent/router_agent/prompt.py

@@ -0,0 +1,48 @@
+# 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.
+
+SYSTEM_PROMPT = """
+You are a Router Agent responsible for analyzing incoming requests and routing them to the most appropriate branch.
+
+Available branches:
+{branches}
+
+CRITICAL INSTRUCTIONS:
+- Analyze the user's request carefully
+- Select exactly ONE branch that best handles the request from: [{branch_names}]
+- Respond with ONLY the exact branch name, nothing else
+- Be decisive - choose the single best match, if the request could fit multiple branches,
+  choose the most specific/specialized one
+- If no branch perfectly fits, choose the closest match
+
+Your response MUST contain ONLY the branch name. Do not include any explanations, reasoning, or additional text.
+
+Examples:
+User: "How do I calculate 15 + 25?"
+Response: calculator_tool
+
+User: "What's the weather like today?"
+Response: weather_service
+
+User: "Send an email to John"
+Response: email_tool"""
+
+USER_PROMPT = """
+Previous conversation history:
+{chat_history}
+
+To respond to the request: {request}, which branch should be chosen?
+
+Respond with only the branch name."""

nat/agent/router_agent/register.py

@@ -0,0 +1,97 @@
+# 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 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.component_ref import FunctionRef
+from nat.data_models.component_ref import LLMRef
+from nat.data_models.function import FunctionBaseConfig
+
+logger = logging.getLogger(__name__)
+
+
+class RouterAgentWorkflowConfig(FunctionBaseConfig, name="router_agent"):
+    """
+    A router agent takes in the incoming message, combines it with a prompt and the list of branches,
+    and ask a LLM about which branch to take.
+    """
+    branches: list[FunctionRef] = Field(default_factory=list,
+                                        description="The list of branches to provide to the router agent.")
+    llm_name: LLMRef = Field(description="The LLM model to use with the routing agent.")
+    description: str = Field(default="Router Agent Workflow", description="Description of this functions use.")
+    system_prompt: str | None = Field(default=None, description="Provides the system prompt to use with the agent.")
+    user_prompt: str | None = Field(default=None, description="Provides the prompt to use with the agent.")
+    max_router_retries: int = Field(
+        default=3, description="Maximum number of retries if the router agent fails to choose a branch.")
+    detailed_logs: bool = Field(default=False, description="Set the verbosity of the router agent's logging.")
+    log_response_max_chars: PositiveInt = Field(
+        default=1000, description="Maximum number of characters to display in logs when logging branch responses.")
+
+
+@register_function(config_type=RouterAgentWorkflowConfig, framework_wrappers=[LLMFrameworkEnum.LANGCHAIN])
+async def router_agent_workflow(config: RouterAgentWorkflowConfig, builder: Builder):
+    from langchain_core.messages.human import HumanMessage
+    from langgraph.graph.state import CompiledStateGraph
+
+    from nat.agent.base import AGENT_LOG_PREFIX
+    from nat.agent.router_agent.agent import RouterAgentGraph
+    from nat.agent.router_agent.agent import RouterAgentGraphState
+    from nat.agent.router_agent.agent import create_router_agent_prompt
+
+    prompt = create_router_agent_prompt(config)
+    llm = await builder.get_llm(config.llm_name, wrapper_type=LLMFrameworkEnum.LANGCHAIN)
+    branches = builder.get_tools(tool_names=config.branches, wrapper_type=LLMFrameworkEnum.LANGCHAIN)
+    if not branches:
+        raise ValueError(f"No branches specified for Router Agent '{config.llm_name}'")
+
+    graph: CompiledStateGraph = await RouterAgentGraph(
+        llm=llm,
+        branches=branches,
+        prompt=prompt,
+        max_router_retries=config.max_router_retries,
+        detailed_logs=config.detailed_logs,
+        log_response_max_chars=config.log_response_max_chars,
+    ).build_graph()
+
+    async def _response_fn(input_message: str) -> str:
+        try:
+            message = HumanMessage(content=input_message)
+            state = RouterAgentGraphState(forward_message=message)
+
+            result_dict = await graph.ainvoke(state)
+            result_state = RouterAgentGraphState(**result_dict)
+
+            output_message = result_state.messages[-1]
+            return str(output_message.content)
+
+        except Exception as ex:
+            logger.exception("%s Router Agent failed with exception: %s", AGENT_LOG_PREFIX, ex)
+            if config.detailed_logs:
+                return str(ex)
+            return "Router agent failed with exception: %s" % ex
+
+    try:
+        yield FunctionInfo.from_fn(_response_fn, description=config.description)
+    except GeneratorExit:
+        logger.exception("%s Workflow exited early!", AGENT_LOG_PREFIX)
+    finally:
+        logger.debug("%s Cleaning up router_agent workflow.", AGENT_LOG_PREFIX)

nat/agent/tool_calling_agent/agent.py

@@ -19,10 +19,13 @@ 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
@@ -57,12 +60,14 @@ class ToolCallAgentGraph(DualNodeAgent):
         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,
                          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:
@@ -85,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):
@@ -146,13 +151,70 @@ class ToolCallAgentGraph(DualNodeAgent):
             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.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)

nat/agent/tool_calling_agent/register.py

@@ -22,6 +22,7 @@ from nat.builder.builder import Builder
 from nat.builder.framework_enum import LLMFrameworkEnum
 from nat.builder.function_info import FunctionInfo
 from nat.cli.register_workflow import register_function
+from nat.data_models.component_ref import FunctionGroupRef
 from nat.data_models.component_ref import FunctionRef
 from nat.data_models.component_ref import LLMRef
 from nat.data_models.function import FunctionBaseConfig
@@ -35,8 +36,8 @@ class ToolCallAgentWorkflowConfig(FunctionBaseConfig, name="tool_calling_agent")
     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.")
+    tool_names: list[FunctionRef | FunctionGroupRef] = 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.")
@@ -47,12 +48,14 @@ class ToolCallAgentWorkflowConfig(FunctionBaseConfig, name="tool_calling_agent")
     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 langgraph.graph.state import CompiledStateGraph
 
     from nat.agent.base import AGENT_LOG_PREFIX
     from nat.agent.tool_calling_agent.agent import ToolCallAgentGraph
@@ -68,13 +71,18 @@ async def tool_calling_agent_workflow(config: ToolCallAgentWorkflowConfig, build
     if not tools:
         raise ValueError(f"No tools specified for Tool Calling Agent '{config.llm_name}'")
 
+    # convert return_direct FunctionRef objects to BaseTool objects
+    return_direct_tools = builder.get_tools(tool_names=config.return_direct,
+                                            wrapper_type=LLMFrameworkEnum.LANGCHAIN) if config.return_direct else None
+
     # 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,
-                                                    log_response_max_chars=config.log_response_max_chars,
-                                                    handle_tool_errors=config.handle_tool_errors).build_graph()
+    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: str) -> str:
         try: