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

nat/agent/base.py

@@ -27,9 +27,10 @@ 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.graph import CompiledGraph
+from langgraph.graph.state import CompiledStateGraph
 
 logger = logging.getLogger(__name__)
 
@@ -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,
@@ -256,5 +261,5 @@ class BaseAgent(ABC):
         return "\n".join([f"{message.type}: {message.content}" for message in messages[:-1]])
 
     @abstractmethod
-    async def _build_graph(self, state_schema: type) -> CompiledGraph:
+    async def _build_graph(self, state_schema: type) -> CompiledStateGraph:
         pass

nat/agent/dual_node.py

@@ -20,7 +20,7 @@ from langchain_core.callbacks import AsyncCallbackHandler
 from langchain_core.language_models import BaseChatModel
 from langchain_core.tools import BaseTool
 from langgraph.graph import StateGraph
-from langgraph.graph.graph import CompiledGraph
+from langgraph.graph.state import CompiledStateGraph
 from pydantic import BaseModel
 
 from .base import AgentDecision
@@ -55,7 +55,7 @@ class DualNodeAgent(BaseAgent):
     async def conditional_edge(self, state: BaseModel) -> str:
         pass
 
-    async def _build_graph(self, state_schema) -> CompiledGraph:
+    async def _build_graph(self, state_schema: type) -> CompiledStateGraph:
         log.debug("Building and compiling the Agent Graph")
 
         graph = StateGraph(state_schema)

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

@@ -25,22 +25,26 @@ from nat.builder.function_info import FunctionInfo
 from nat.cli.register_workflow import register_function
 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(FunctionBaseConfig, 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.")
+    tool_names: list[FunctionRef | FunctionGroupRef] = 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.")
     retry_agent_response_parsing_errors: bool = Field(
@@ -75,15 +79,21 @@ class ReActAgentWorkflowConfig(FunctionBaseConfig, name="react_agent"):
     use_openai_api: bool = Field(default=False,
                                  description=("Use OpenAI API for the input/output types to the function. "
                                               "If False, strings will be used."))
-    additional_instructions: str | None = Field(
-        default=None, description="Additional instructions to provide to the agent in addition to the base prompt.")
+    additional_instructions: str | None = OptimizableField(
+        default=None,
+        description="Additional instructions to provide to the agent in addition to the base prompt.",
+        space=SearchSpace(
+            is_prompt=True,
+            prompt="No additional instructions.",
+            prompt_purpose="Additional instructions to provide to the agent in addition to the base prompt.",
+        ))
 
 
 @register_function(config_type=ReActAgentWorkflowConfig, framework_wrappers=[LLMFrameworkEnum.LANGCHAIN])
 async def react_agent_workflow(config: ReActAgentWorkflowConfig, builder: Builder):
     from langchain.schema import BaseMessage
     from langchain_core.messages import trim_messages
-    from langgraph.graph.graph import CompiledGraph
+    from langgraph.graph.state import CompiledStateGraph
 
     from nat.agent.base import AGENT_LOG_PREFIX
     from nat.agent.react_agent.agent import ReActAgentGraph
@@ -101,7 +111,7 @@ async def react_agent_workflow(config: ReActAgentWorkflowConfig, builder: Builde
         raise ValueError(f"No tools specified for ReAct Agent '{config.llm_name}'")
     # configure callbacks, for sending intermediate steps
     # construct the ReAct Agent Graph from the configured llm, prompt, and tools
-    graph: CompiledGraph = await ReActAgentGraph(
+    graph: CompiledStateGraph = await ReActAgentGraph(
         llm=llm,
         prompt=prompt,
         tools=tools,

nat/agent/reasoning_agent/reasoning_agent.py

@@ -113,7 +113,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,7 +16,9 @@
 # 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
+from .router_agent import register as router_agent
 from .tool_calling_agent import register as tool_calling_agent

nat/agent/rewoo_agent/agent.py

@@ -27,6 +27,7 @@ from langchain_core.prompts.chat import ChatPromptTemplate
 from langchain_core.runnables.config import RunnableConfig
 from langchain_core.tools import BaseTool
 from langgraph.graph import StateGraph
+from langgraph.graph.state import CompiledStateGraph
 from pydantic import BaseModel
 from pydantic import Field
 
@@ -67,7 +68,8 @@ 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):
         super().__init__(llm=llm,
                          tools=tools,
                          callbacks=callbacks,
@@ -93,6 +95,7 @@ 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
 
         logger.debug("%s Initialized ReWOO Agent Graph", AGENT_LOG_PREFIX)
 
@@ -268,7 +271,7 @@ 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))
@@ -341,7 +344,7 @@ class ReWOOAgentGraph(BaseAgent):
             logger.warning("%s Ending graph traversal", AGENT_LOG_PREFIX)
             return AgentDecision.END
 
-    async def _build_graph(self, state_schema):
+    async def _build_graph(self, state_schema: type) -> CompiledStateGraph:
         try:
             logger.debug("%s Building and compiling the ReWOO Graph", AGENT_LOG_PREFIX)
 

nat/agent/rewoo_agent/register.py

@@ -25,6 +25,7 @@ from nat.builder.function_info import FunctionInfo
 from nat.cli.register_workflow import register_function
 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
@@ -39,8 +40,8 @@ class ReWOOAgentWorkflowConfig(FunctionBaseConfig, name="rewoo_agent"):
     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.")
+    tool_names: list[FunctionRef | FunctionGroupRef] = Field(
+        default_factory=list, description="The list of tools to provide to the rewoo agent.")
     llm_name: LLMRef = Field(description="The LLM model to use with the rewoo agent.")
     verbose: bool = Field(default=False, description="Set the verbosity of the rewoo agent's logging.")
     include_tool_input_schema_in_tool_description: bool = Field(
@@ -52,6 +53,9 @@ class ReWOOAgentWorkflowConfig(FunctionBaseConfig, name="rewoo_agent"):
     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.")
@@ -73,7 +77,7 @@ async def rewoo_agent_workflow(config: ReWOOAgentWorkflowConfig, builder: Builde
     from langchain_core.messages import trim_messages
     from langchain_core.messages.human import HumanMessage
     from langchain_core.prompts import ChatPromptTemplate
-    from langgraph.graph.graph import CompiledGraph
+    from 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
@@ -111,13 +115,15 @@ async def rewoo_agent_workflow(config: ReWOOAgentWorkflowConfig, builder: Builde
         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,
-                                                 log_response_max_chars=config.log_response_max_chars).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).build_graph()
 
     async def _response_fn(input_message: ChatRequest) -> ChatResponse:
         try: