nvidia-nat 1.3.0a20250910__py3-none-any.whl → 1.3.0a20250922__py3-none-any.whl

nat/agent/base.py

@@ -27,6 +27,7 @@ from langchain_core.language_models import BaseChatModel
 from langchain_core.messages import AIMessage
 from langchain_core.messages import BaseMessage
 from langchain_core.messages import ToolMessage
+from langchain_core.runnables import Runnable
 from langchain_core.runnables import RunnableConfig
 from langchain_core.tools import BaseTool
 from langgraph.graph.state import CompiledStateGraph
@@ -107,21 +108,25 @@ class BaseAgent(ABC):
 
         return AIMessage(content=output_message)
 
-    async def _call_llm(self, messages: list[BaseMessage]) -> AIMessage:
+    async def _call_llm(self, llm: Runnable, inputs: dict[str, Any], config: RunnableConfig | None = None) -> AIMessage:
         """
         Call the LLM directly. Retry logic is handled automatically by the underlying LLM client.
 
         Parameters
         ----------
-        messages : list[BaseMessage]
-            The messages to send to the LLM
+        llm : Runnable
+            The LLM runnable (prompt | llm or similar)
+        inputs : dict[str, Any]
+            The inputs to pass to the runnable
+        config : RunnableConfig | None
+            The config to pass to the runnable (should include callbacks)
 
         Returns
         -------
         AIMessage
             The LLM response
         """
-        response = await self.llm.ainvoke(messages)
+        response = await llm.ainvoke(inputs, config=config)
         return AIMessage(content=str(response.content))
 
     async def _call_tool(self,

nat/agent/prompt_optimizer/prompt.py

@@ -0,0 +1,68 @@
+# SPDX-FileCopyrightText: Copyright (c) 2024-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 W291
+
+mutator_prompt = """
+
+## CORE DIRECTIVES
+- **Preserve the original objective and task.** Do not change what the prompt is meant to accomplish.  
+- **Keep the intent intact.** The improved prompt must solve the same problem as the original.  
+- **Do not invent new goals.** Only improve clarity, structure, constraints, and usability.  
+- **Do not drop critical instructions.** Everything essential from the original prompt must remain.  
+- **Return only the mutated prompt text.** No rationale, no diffs, no explanations.  
+- **Be Creative within bounds.** You may rephrase, reorganize, and enhance, but not alter meaning.
+- **DO NOT use curly braces in your prompt** for anything other than existing variables in the prompt as the string
+will be treated as an f-string.
+- **Examples are a good idea** if the original prompt lacks them. They help clarify expected output.
+
+---
+
+## IMPROVEMENT HINTS
+When modifying, apply these principles:
+1. **Clarity & Precision** – remove vague language, strengthen directives.  
+2. **Structure & Flow** – order sections as: *Objective → Constraints → Tools → Steps → Output Schema → Examples*.  
+3. **Schema Adherence** – enforce a single canonical output schema (JSON/XML) with `schema_version`.  
+4. **Tool Governance** – clarify when/how tools are used, their inputs/outputs, and fallback behavior.  
+5. **Error Handling** – specify behavior if tools fail or inputs are insufficient.  
+6. **Budget Awareness** – minimize verbosity, respect token/latency limits.  
+7. **Safety** – include refusals for unsafe requests, enforce compliance with rules.  
+8. **Consistency** – avoid format drift; always maintain the same schema.  
+9. **Integrity** – confirm the task, objective, and intent are preserved.  
+
+---
+
+## MUTATION OPERATORS
+You may:
+- **Tighten** (remove fluff, redundancies)  
+- **Reorder** (improve logical flow)  
+- **Constrain** (add explicit rules/limits)  
+- **Harden** (improve error handling/fallbacks)  
+- **Defuse** (replace ambiguous verbs with measurable actions)  
+- **Format-lock** (wrap outputs in JSON/XML fenced blocks)
+- **Example-ify** (add examples if missing or weak)  
+
+---
+
+## INPUT
+Here is the prompt to mutate:
+{original_prompt}
+
+## OBJECTIVE
+The prompt must acheive the following objective:
+{objective}
+
+The modified prompt is: \n
+
+"""

nat/agent/prompt_optimizer/register.py

@@ -0,0 +1,149 @@
+# SPDX-FileCopyrightText: Copyright (c) 2021-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.
+
+from pydantic import Field
+
+from nat.builder.builder import Builder
+from nat.builder.framework_enum import LLMFrameworkEnum
+from nat.builder.function_info import FunctionInfo
+from nat.cli.register_workflow import register_function
+from nat.data_models.component_ref import LLMRef
+from nat.data_models.function import FunctionBaseConfig
+from nat.profiler.parameter_optimization.prompt_optimizer import PromptOptimizerInputSchema
+
+
+class PromptOptimizerConfig(FunctionBaseConfig, name="prompt_init"):
+
+    optimizer_llm: LLMRef = Field(description="LLM to use for prompt optimization")
+    optimizer_prompt: str = Field(
+        description="Prompt template for the optimizer",
+        default=(
+            "You are an expert at optimizing prompts for LLMs. "
+            "Your task is to take a given prompt and suggest an optimized version of it. "
+            "Note that the prompt might be a template with variables and curly braces. Remember to always keep the "
+            "variables and curly braces in the prompt the same. Only modify the instructions in the prompt that are"
+            "not variables. The system is meant to achieve the following objective\n"
+            "{system_objective}\n Of which, the prompt is one part. The details of the prompt and context as below.\n"))
+    system_objective: str = Field(description="Objective of the workflow")
+
+
+@register_function(config_type=PromptOptimizerConfig, framework_wrappers=[LLMFrameworkEnum.LANGCHAIN])
+async def prompt_optimizer_function(config: PromptOptimizerConfig, builder: Builder):
+    """
+    Function to optimize prompts for LLMs.
+    """
+
+    try:
+        from langchain_core.prompts import PromptTemplate
+
+        from .prompt import mutator_prompt
+    except ImportError as exc:
+        raise ImportError("langchain-core is not installed. Please install it to use MultiLLMPlanner.\n"
+                          "This error can be resolve by installing nvidia-nat[langchain]") from exc
+
+    llm = await builder.get_llm(config.optimizer_llm, wrapper_type=LLMFrameworkEnum.LANGCHAIN)
+
+    template = PromptTemplate(template=config.optimizer_prompt,
+                              input_variables=["system_objective"],
+                              validate_template=True)
+
+    base_prompt: str = (await template.ainvoke(input={"system_objective": config.system_objective})).to_string()
+    prompt_extension_template = PromptTemplate(template=mutator_prompt,
+                                               input_variables=["original_prompt", "objective"],
+                                               validate_template=True)
+
+    async def _inner(input_message: PromptOptimizerInputSchema) -> str:
+        """
+        Optimize the prompt using the provided LLM.
+        """
+
+        original_prompt = input_message.original_prompt
+        prompt_objective = input_message.objective
+
+        prompt_extension = (await prompt_extension_template.ainvoke(input={
+            "original_prompt": original_prompt,
+            "objective": prompt_objective,
+        })).to_string()
+
+        prompt = f"{base_prompt}\n\n{prompt_extension}"
+
+        optimized_prompt = await llm.ainvoke(prompt)
+        return optimized_prompt.content
+
+    yield FunctionInfo.from_fn(
+        fn=_inner,
+        description="Optimize prompts for LLMs using a feedback LLM.",
+    )
+
+
+class PromptRecombinerConfig(FunctionBaseConfig, name="prompt_recombiner"):
+
+    optimizer_llm: LLMRef = Field(description="LLM to use for prompt recombination")
+    optimizer_prompt: str = Field(
+        description="Prompt template for the recombiner",
+        default=("You are an expert at combining prompt instructions for LLMs. "
+                 "Your task is to merge two prompts for the same objective into a single, stronger prompt. "
+                 "Do not introduce new variables or modify existing placeholders."),
+    )
+    system_objective: str = Field(description="Objective of the workflow")
+
+
+@register_function(config_type=PromptRecombinerConfig, framework_wrappers=[LLMFrameworkEnum.LANGCHAIN])
+async def prompt_recombiner_function(config: PromptRecombinerConfig, builder: Builder):
+    """
+    Function to recombine two parent prompts into a child prompt using the optimizer LLM.
+    Uses the same base template and objective instructions.
+    """
+
+    try:
+        from langchain_core.prompts import PromptTemplate
+    except ImportError as exc:
+        raise ImportError("langchain-core is not installed. Please install it to use MultiLLMPlanner.\n"
+                          "This error can be resolve by installing nvidia-nat[langchain].") from exc
+
+    llm = await builder.get_llm(config.optimizer_llm, wrapper_type=LLMFrameworkEnum.LANGCHAIN)
+
+    template = PromptTemplate(template=config.optimizer_prompt,
+                              input_variables=["system_objective"],
+                              validate_template=True)
+
+    base_prompt: str = (await template.ainvoke(input={"system_objective": config.system_objective})).to_string()
+
+    class RecombineSchema(PromptOptimizerInputSchema):
+        parent_b: str | None = None
+
+    async def _inner(input_message: RecombineSchema) -> str:
+        parent_a = input_message.original_prompt
+        parent_b = input_message.parent_b or ""
+        prompt_objective = input_message.objective
+
+        prompt = (
+            f"{base_prompt}\n\n"
+            "We are performing genetic recombination between two prompts that satisfy the same objective.\n"
+            f"Objective: {prompt_objective}\n\n"
+            f"Parent A:\n{parent_a}\n\n"
+            f"Parent B:\n{parent_b}\n\n"
+            "Combine the strongest instructions and phrasing from both parents to produce a single, coherent child "
+            "prompt.\n"
+            "Maintain variables and placeholders unchanged.\n"
+            "Return only the child prompt text, with no additional commentary.")
+
+        child_prompt = await llm.ainvoke(prompt)
+        return child_prompt.content
+
+    yield FunctionInfo.from_fn(
+        fn=_inner,
+        description="Recombine two prompts into a stronger child prompt.",
+    )

nat/agent/react_agent/agent.py

@@ -360,7 +360,7 @@ class ReActAgentGraph(DualNodeAgent):
         if errors:
             error_text = "\n".join(errors)
             logger.error("%s %s", AGENT_LOG_PREFIX, error_text)
-            raise ValueError(error_text)
+            return False
         return True
 
 

nat/agent/react_agent/register.py

@@ -17,32 +17,32 @@ 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.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"),
@@ -61,7 +61,6 @@ 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. "
@@ -70,13 +69,17 @@ class ReActAgentWorkflowConfig(FunctionBaseConfig, name="react_agent"):
         default=None,
         description="Provides the SYSTEM_PROMPT to use with the agent")  # defaults to SYSTEM_PROMPT in prompt.py
     max_history: int = Field(default=15, description="Maximum number of messages to keep in the conversation history.")
-    log_response_max_chars: PositiveInt = Field(
-        default=1000, description="Maximum number of characters to display in logs when logging tool responses.")
     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])

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"
@@ -113,7 +110,12 @@ 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:

nat/agent/register.py

@@ -16,6 +16,7 @@
 # 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

nat/agent/rewoo_agent/agent.py

@@ -68,7 +68,9 @@ class ReWOOAgentGraph(BaseAgent):
                  use_tool_schema: bool = True,
                  callbacks: list[AsyncCallbackHandler] | None = None,
                  detailed_logs: bool = False,
-                 log_response_max_chars: int = 1000):
+                 log_response_max_chars: int = 1000,
+                 tool_call_max_retries: int = 3,
+                 raise_tool_call_error: bool = True):
         super().__init__(llm=llm,
                          tools=tools,
                          callbacks=callbacks,
@@ -94,6 +96,8 @@ class ReWOOAgentGraph(BaseAgent):
         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}
+        self.tool_call_max_retries = tool_call_max_retries
+        self.raise_tool_call_error = raise_tool_call_error
 
         logger.debug("%s Initialized ReWOO Agent Graph", AGENT_LOG_PREFIX)
 
@@ -269,11 +273,14 @@ class ReWOOAgentGraph(BaseAgent):
             tool_response = await self._call_tool(requested_tool,
                                                   tool_input_parsed,
                                                   RunnableConfig(callbacks=self.callbacks),
-                                                  max_retries=3)
+                                                  max_retries=self.tool_call_max_retries)
 
             if self.detailed_logs:
                 self._log_tool_response(requested_tool.name, tool_input_parsed, str(tool_response))
 
+            if self.raise_tool_call_error and tool_response.status == "error":
+                raise RuntimeError(f"Tool call failed: {tool_response.content}")
+
             intermediate_results[placeholder] = tool_response
             return {"intermediate_results": intermediate_results}
 

nat/agent/rewoo_agent/register.py

@@ -23,38 +23,36 @@ 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.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.")
-    log_response_max_chars: PositiveInt = Field(
-        default=1000, description="Maximum number of characters to display in logs when logging tool responses.")
     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."))
@@ -65,6 +63,10 @@ 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])
@@ -118,7 +120,9 @@ async def rewoo_agent_workflow(config: ReWOOAgentWorkflowConfig, builder: Builde
         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).build_graph()
+        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:
         try:

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)