nvidia-nat 1.3.0a20250910__py3-none-any.whl → 1.3.0a20250917__py3-none-any.whl
This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
- nat/agent/base.py +9 -4
- nat/agent/prompt_optimizer/prompt.py +68 -0
- nat/agent/prompt_optimizer/register.py +149 -0
- nat/agent/react_agent/agent.py +1 -1
- nat/agent/react_agent/register.py +15 -5
- nat/agent/reasoning_agent/reasoning_agent.py +6 -1
- nat/agent/register.py +2 -0
- nat/agent/rewoo_agent/agent.py +4 -2
- nat/agent/rewoo_agent/register.py +8 -3
- nat/agent/router_agent/__init__.py +0 -0
- nat/agent/router_agent/agent.py +329 -0
- nat/agent/router_agent/prompt.py +48 -0
- nat/agent/router_agent/register.py +97 -0
- nat/agent/tool_calling_agent/agent.py +69 -7
- nat/agent/tool_calling_agent/register.py +11 -3
- nat/builder/builder.py +27 -4
- nat/builder/component_utils.py +7 -3
- nat/builder/function.py +167 -0
- nat/builder/function_info.py +1 -1
- nat/builder/workflow.py +5 -0
- nat/builder/workflow_builder.py +213 -16
- nat/cli/commands/optimize.py +90 -0
- nat/cli/commands/workflow/templates/config.yml.j2 +0 -1
- nat/cli/commands/workflow/workflow_commands.py +4 -7
- nat/cli/entrypoint.py +2 -0
- nat/cli/register_workflow.py +38 -4
- nat/cli/type_registry.py +71 -0
- nat/data_models/component.py +2 -0
- nat/data_models/component_ref.py +11 -0
- nat/data_models/config.py +40 -16
- nat/data_models/function.py +34 -0
- nat/data_models/function_dependencies.py +8 -0
- nat/data_models/optimizable.py +119 -0
- nat/data_models/optimizer.py +149 -0
- nat/data_models/temperature_mixin.py +4 -3
- nat/data_models/top_p_mixin.py +4 -3
- nat/embedder/nim_embedder.py +1 -1
- nat/embedder/openai_embedder.py +1 -1
- nat/eval/config.py +1 -1
- nat/eval/evaluate.py +5 -1
- nat/eval/register.py +4 -0
- nat/eval/runtime_evaluator/__init__.py +14 -0
- nat/eval/runtime_evaluator/evaluate.py +123 -0
- nat/eval/runtime_evaluator/register.py +100 -0
- nat/experimental/test_time_compute/functions/plan_select_execute_function.py +5 -1
- nat/front_ends/fastapi/dask_client_mixin.py +43 -0
- nat/front_ends/fastapi/fastapi_front_end_config.py +14 -3
- nat/front_ends/fastapi/fastapi_front_end_plugin.py +111 -3
- nat/front_ends/fastapi/fastapi_front_end_plugin_worker.py +243 -228
- nat/front_ends/fastapi/job_store.py +518 -99
- nat/front_ends/fastapi/main.py +11 -19
- nat/front_ends/fastapi/utils.py +57 -0
- nat/front_ends/mcp/mcp_front_end_plugin_worker.py +3 -2
- nat/llm/aws_bedrock_llm.py +14 -3
- nat/llm/nim_llm.py +14 -3
- nat/llm/openai_llm.py +8 -1
- nat/observability/exporter/processing_exporter.py +29 -55
- nat/observability/mixin/redaction_config_mixin.py +5 -4
- nat/observability/mixin/tagging_config_mixin.py +26 -14
- nat/observability/mixin/type_introspection_mixin.py +401 -107
- nat/observability/processor/processor.py +3 -0
- nat/observability/processor/redaction/__init__.py +24 -0
- nat/observability/processor/redaction/contextual_redaction_processor.py +125 -0
- nat/observability/processor/redaction/contextual_span_redaction_processor.py +66 -0
- nat/observability/processor/redaction/redaction_processor.py +177 -0
- nat/observability/processor/redaction/span_header_redaction_processor.py +92 -0
- nat/observability/processor/span_tagging_processor.py +21 -14
- nat/profiler/decorators/framework_wrapper.py +9 -6
- nat/profiler/parameter_optimization/__init__.py +0 -0
- nat/profiler/parameter_optimization/optimizable_utils.py +93 -0
- nat/profiler/parameter_optimization/optimizer_runtime.py +67 -0
- nat/profiler/parameter_optimization/parameter_optimizer.py +149 -0
- nat/profiler/parameter_optimization/parameter_selection.py +108 -0
- nat/profiler/parameter_optimization/pareto_visualizer.py +380 -0
- nat/profiler/parameter_optimization/prompt_optimizer.py +384 -0
- nat/profiler/parameter_optimization/update_helpers.py +66 -0
- nat/profiler/utils.py +3 -1
- nat/tool/chat_completion.py +4 -1
- nat/tool/github_tools.py +450 -0
- nat/tool/register.py +2 -7
- nat/utils/callable_utils.py +70 -0
- nat/utils/exception_handlers/automatic_retries.py +103 -48
- nat/utils/type_utils.py +4 -0
- {nvidia_nat-1.3.0a20250910.dist-info → nvidia_nat-1.3.0a20250917.dist-info}/METADATA +8 -1
- {nvidia_nat-1.3.0a20250910.dist-info → nvidia_nat-1.3.0a20250917.dist-info}/RECORD +91 -71
- nat/observability/processor/header_redaction_processor.py +0 -123
- nat/observability/processor/redaction_processor.py +0 -77
- nat/tool/github_tools/create_github_commit.py +0 -133
- nat/tool/github_tools/create_github_issue.py +0 -87
- nat/tool/github_tools/create_github_pr.py +0 -106
- nat/tool/github_tools/get_github_file.py +0 -106
- nat/tool/github_tools/get_github_issue.py +0 -166
- nat/tool/github_tools/get_github_pr.py +0 -256
- nat/tool/github_tools/update_github_issue.py +0 -100
- /nat/{tool/github_tools → agent/prompt_optimizer}/__init__.py +0 -0
- {nvidia_nat-1.3.0a20250910.dist-info → nvidia_nat-1.3.0a20250917.dist-info}/WHEEL +0 -0
- {nvidia_nat-1.3.0a20250910.dist-info → nvidia_nat-1.3.0a20250917.dist-info}/entry_points.txt +0 -0
- {nvidia_nat-1.3.0a20250910.dist-info → nvidia_nat-1.3.0a20250917.dist-info}/licenses/LICENSE-3rd-party.txt +0 -0
- {nvidia_nat-1.3.0a20250910.dist-info → nvidia_nat-1.3.0a20250917.dist-info}/licenses/LICENSE.md +0 -0
- {nvidia_nat-1.3.0a20250910.dist-info → nvidia_nat-1.3.0a20250917.dist-info}/top_level.txt +0 -0
|
@@ -0,0 +1,329 @@
|
|
|
1
|
+
# SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
|
|
2
|
+
# SPDX-License-Identifier: Apache-2.0
|
|
3
|
+
#
|
|
4
|
+
# Licensed under the Apache License, Version 2.0 (the "License");
|
|
5
|
+
# you may not use this file except in compliance with the License.
|
|
6
|
+
# You may obtain a copy of the License at
|
|
7
|
+
#
|
|
8
|
+
# http://www.apache.org/licenses/LICENSE-2.0
|
|
9
|
+
#
|
|
10
|
+
# Unless required by applicable law or agreed to in writing, software
|
|
11
|
+
# distributed under the License is distributed on an "AS IS" BASIS,
|
|
12
|
+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
13
|
+
# See the License for the specific language governing permissions and
|
|
14
|
+
# limitations under the License.
|
|
15
|
+
|
|
16
|
+
import logging
|
|
17
|
+
import typing
|
|
18
|
+
|
|
19
|
+
from langchain_core.callbacks.base import AsyncCallbackHandler
|
|
20
|
+
from langchain_core.language_models import BaseChatModel
|
|
21
|
+
from langchain_core.messages.base import BaseMessage
|
|
22
|
+
from langchain_core.messages.human import HumanMessage
|
|
23
|
+
from langchain_core.prompts.chat import ChatPromptTemplate
|
|
24
|
+
from langchain_core.tools import BaseTool
|
|
25
|
+
from langgraph.graph import StateGraph
|
|
26
|
+
from pydantic import BaseModel
|
|
27
|
+
from pydantic import Field
|
|
28
|
+
|
|
29
|
+
from nat.agent.base import AGENT_CALL_LOG_MESSAGE
|
|
30
|
+
from nat.agent.base import AGENT_LOG_PREFIX
|
|
31
|
+
from nat.agent.base import BaseAgent
|
|
32
|
+
|
|
33
|
+
if typing.TYPE_CHECKING:
|
|
34
|
+
from nat.agent.router_agent.register import RouterAgentWorkflowConfig
|
|
35
|
+
|
|
36
|
+
logger = logging.getLogger(__name__)
|
|
37
|
+
|
|
38
|
+
|
|
39
|
+
class RouterAgentGraphState(BaseModel):
|
|
40
|
+
"""State schema for the Router Agent Graph.
|
|
41
|
+
|
|
42
|
+
This class defines the state structure used throughout the Router Agent's
|
|
43
|
+
execution graph, containing messages, routing information, and branch selection.
|
|
44
|
+
|
|
45
|
+
Attributes:
|
|
46
|
+
messages: A list of messages representing the conversation history.
|
|
47
|
+
forward_message: The message to be forwarded to the chosen branch.
|
|
48
|
+
chosen_branch: The name of the branch selected by the router agent.
|
|
49
|
+
"""
|
|
50
|
+
messages: list[BaseMessage] = Field(default_factory=list)
|
|
51
|
+
forward_message: BaseMessage = Field(default_factory=lambda: HumanMessage(content=""))
|
|
52
|
+
chosen_branch: str = Field(default="")
|
|
53
|
+
|
|
54
|
+
|
|
55
|
+
class RouterAgentGraph(BaseAgent):
|
|
56
|
+
"""Configurable Router Agent for routing requests to different branches.
|
|
57
|
+
|
|
58
|
+
A Router Agent analyzes incoming requests and routes them to one of the
|
|
59
|
+
configured branches based on the conte nt and context. It makes a single
|
|
60
|
+
routing decision and executes only the selected branch before returning.
|
|
61
|
+
|
|
62
|
+
This agent is useful for creating multi-path workflows where different
|
|
63
|
+
types of requests need to be handled by specialized sub-agents or tools.
|
|
64
|
+
"""
|
|
65
|
+
|
|
66
|
+
def __init__(
|
|
67
|
+
self,
|
|
68
|
+
llm: BaseChatModel,
|
|
69
|
+
branches: list[BaseTool],
|
|
70
|
+
prompt: ChatPromptTemplate,
|
|
71
|
+
max_router_retries: int = 3,
|
|
72
|
+
callbacks: list[AsyncCallbackHandler] | None = None,
|
|
73
|
+
detailed_logs: bool = False,
|
|
74
|
+
log_response_max_chars: int = 1000,
|
|
75
|
+
):
|
|
76
|
+
"""Initialize the Router Agent.
|
|
77
|
+
|
|
78
|
+
Args:
|
|
79
|
+
llm: The language model to use for routing decisions.
|
|
80
|
+
branches: List of tools/branches that the agent can route to.
|
|
81
|
+
prompt: The chat prompt template for the routing agent.
|
|
82
|
+
max_router_retries: Maximum number of retries if branch selection fails.
|
|
83
|
+
callbacks: Optional list of async callback handlers.
|
|
84
|
+
detailed_logs: Whether to enable detailed logging.
|
|
85
|
+
log_response_max_chars: Maximum characters to log in responses.
|
|
86
|
+
"""
|
|
87
|
+
super().__init__(llm=llm,
|
|
88
|
+
tools=branches,
|
|
89
|
+
callbacks=callbacks,
|
|
90
|
+
detailed_logs=detailed_logs,
|
|
91
|
+
log_response_max_chars=log_response_max_chars)
|
|
92
|
+
|
|
93
|
+
self._branches = branches
|
|
94
|
+
self._branches_dict = {branch.name: branch for branch in branches}
|
|
95
|
+
branch_names = ",".join([branch.name for branch in branches])
|
|
96
|
+
branch_names_and_descriptions = "\n".join([f"{branch.name}: {branch.description}" for branch in branches])
|
|
97
|
+
|
|
98
|
+
prompt = prompt.partial(branches=branch_names_and_descriptions, branch_names=branch_names)
|
|
99
|
+
self.agent = prompt | self.llm
|
|
100
|
+
|
|
101
|
+
self.max_router_retries = max_router_retries
|
|
102
|
+
|
|
103
|
+
def _get_branch(self, branch_name: str) -> BaseTool | None:
|
|
104
|
+
return self._branches_dict.get(branch_name, None)
|
|
105
|
+
|
|
106
|
+
async def agent_node(self, state: RouterAgentGraphState):
|
|
107
|
+
"""Execute the agent node to select a branch for routing.
|
|
108
|
+
|
|
109
|
+
This method processes the incoming request and determines which branch
|
|
110
|
+
should handle it. It uses the configured LLM to analyze the request
|
|
111
|
+
and select the most appropriate branch.
|
|
112
|
+
|
|
113
|
+
Args:
|
|
114
|
+
state: The current state of the router agent graph.
|
|
115
|
+
|
|
116
|
+
Returns:
|
|
117
|
+
RouterAgentGraphState: Updated state with the chosen branch.
|
|
118
|
+
|
|
119
|
+
Raises:
|
|
120
|
+
RuntimeError: If the agent fails to choose a branch after max retries.
|
|
121
|
+
"""
|
|
122
|
+
logger.debug("%s Starting the Router Agent Node", AGENT_LOG_PREFIX)
|
|
123
|
+
chat_history = self._get_chat_history(state.messages)
|
|
124
|
+
request = state.forward_message.content
|
|
125
|
+
for attempt in range(1, self.max_router_retries + 1):
|
|
126
|
+
try:
|
|
127
|
+
agent_response = await self._call_llm(self.agent, {"request": request, "chat_history": chat_history})
|
|
128
|
+
if self.detailed_logs:
|
|
129
|
+
logger.info(AGENT_CALL_LOG_MESSAGE, request, agent_response)
|
|
130
|
+
|
|
131
|
+
state.messages += [agent_response]
|
|
132
|
+
|
|
133
|
+
# Determine chosen branch based on agent response
|
|
134
|
+
if state.chosen_branch == "":
|
|
135
|
+
for branch in self._branches:
|
|
136
|
+
if branch.name.lower() in str(agent_response.content).lower():
|
|
137
|
+
state.chosen_branch = branch.name
|
|
138
|
+
if self.detailed_logs:
|
|
139
|
+
logger.debug("%s Router Agent has chosen branch: %s", AGENT_LOG_PREFIX, branch.name)
|
|
140
|
+
return state
|
|
141
|
+
|
|
142
|
+
# The agent failed to choose a branch
|
|
143
|
+
if state.chosen_branch == "":
|
|
144
|
+
if attempt == self.max_router_retries:
|
|
145
|
+
logger.error("%s Router Agent has empty chosen branch", AGENT_LOG_PREFIX)
|
|
146
|
+
raise RuntimeError("Router Agent failed to choose a branch")
|
|
147
|
+
logger.warning("%s Router Agent failed to choose a branch, retrying %d out of %d",
|
|
148
|
+
AGENT_LOG_PREFIX,
|
|
149
|
+
attempt,
|
|
150
|
+
self.max_router_retries)
|
|
151
|
+
|
|
152
|
+
except Exception as ex:
|
|
153
|
+
logger.error("%s Router Agent failed to call agent_node: %s", AGENT_LOG_PREFIX, ex)
|
|
154
|
+
raise
|
|
155
|
+
|
|
156
|
+
return state
|
|
157
|
+
|
|
158
|
+
async def branch_node(self, state: RouterAgentGraphState):
|
|
159
|
+
"""Execute the selected branch with the forwarded message.
|
|
160
|
+
|
|
161
|
+
This method calls the tool/branch that was selected by the agent node
|
|
162
|
+
and processes the response.
|
|
163
|
+
|
|
164
|
+
Args:
|
|
165
|
+
state: The current state containing the chosen branch and message.
|
|
166
|
+
|
|
167
|
+
Returns:
|
|
168
|
+
RouterAgentGraphState: Updated state with the branch response.
|
|
169
|
+
|
|
170
|
+
Raises:
|
|
171
|
+
RuntimeError: If no branch was chosen or branch execution fails.
|
|
172
|
+
ValueError: If the requested tool is not found in the configuration.
|
|
173
|
+
"""
|
|
174
|
+
logger.debug("%s Starting Router Agent Tool Node", AGENT_LOG_PREFIX)
|
|
175
|
+
try:
|
|
176
|
+
if state.chosen_branch == "":
|
|
177
|
+
logger.error("%s Router Agent has empty chosen branch", AGENT_LOG_PREFIX)
|
|
178
|
+
raise RuntimeError("Router Agent failed to choose a branch")
|
|
179
|
+
requested_branch = self._get_branch(state.chosen_branch)
|
|
180
|
+
if not requested_branch:
|
|
181
|
+
logger.error("%s Router Agent wants to call tool %s but it is not in the config file",
|
|
182
|
+
AGENT_LOG_PREFIX,
|
|
183
|
+
state.chosen_branch)
|
|
184
|
+
raise ValueError("Tool not found in config file")
|
|
185
|
+
|
|
186
|
+
branch_input = state.forward_message.content
|
|
187
|
+
branch_response = await self._call_tool(requested_branch, branch_input)
|
|
188
|
+
state.messages += [branch_response]
|
|
189
|
+
if self.detailed_logs:
|
|
190
|
+
self._log_tool_response(requested_branch.name, branch_input, branch_response.content)
|
|
191
|
+
|
|
192
|
+
return state
|
|
193
|
+
|
|
194
|
+
except Exception as ex:
|
|
195
|
+
logger.error("%s Router Agent throws exception during branch node execution: %s", AGENT_LOG_PREFIX, ex)
|
|
196
|
+
raise
|
|
197
|
+
|
|
198
|
+
async def _build_graph(self, state_schema):
|
|
199
|
+
logger.debug("%s Building and compiling the Router Agent Graph", AGENT_LOG_PREFIX)
|
|
200
|
+
|
|
201
|
+
graph = StateGraph(state_schema)
|
|
202
|
+
graph.add_node("agent", self.agent_node)
|
|
203
|
+
graph.add_node("branch", self.branch_node)
|
|
204
|
+
graph.add_edge("agent", "branch")
|
|
205
|
+
graph.set_entry_point("agent")
|
|
206
|
+
|
|
207
|
+
self.graph = graph.compile()
|
|
208
|
+
logger.debug("%s Router Agent Graph built and compiled successfully", AGENT_LOG_PREFIX)
|
|
209
|
+
|
|
210
|
+
return self.graph
|
|
211
|
+
|
|
212
|
+
async def build_graph(self):
|
|
213
|
+
"""Build and compile the router agent execution graph.
|
|
214
|
+
|
|
215
|
+
Creates a state graph with agent and branch nodes, configures the
|
|
216
|
+
execution flow, and compiles the graph for execution.
|
|
217
|
+
|
|
218
|
+
Returns:
|
|
219
|
+
The compiled execution graph.
|
|
220
|
+
|
|
221
|
+
Raises:
|
|
222
|
+
Exception: If graph building or compilation fails.
|
|
223
|
+
"""
|
|
224
|
+
try:
|
|
225
|
+
await self._build_graph(state_schema=RouterAgentGraphState)
|
|
226
|
+
return self.graph
|
|
227
|
+
except Exception as ex:
|
|
228
|
+
logger.error("%s Router Agent failed to build graph: %s", AGENT_LOG_PREFIX, ex)
|
|
229
|
+
raise
|
|
230
|
+
|
|
231
|
+
@staticmethod
|
|
232
|
+
def validate_system_prompt(system_prompt: str) -> bool:
|
|
233
|
+
"""Validate that the system prompt contains required variables.
|
|
234
|
+
|
|
235
|
+
Checks that the system prompt includes necessary template variables
|
|
236
|
+
for branch information that the router agent needs.
|
|
237
|
+
|
|
238
|
+
Args:
|
|
239
|
+
system_prompt: The system prompt string to validate.
|
|
240
|
+
|
|
241
|
+
Returns:
|
|
242
|
+
True if the prompt is valid, False otherwise.
|
|
243
|
+
"""
|
|
244
|
+
errors = []
|
|
245
|
+
required_prompt_variables = {
|
|
246
|
+
"{branches}": "The system prompt must contain {branches} so the agent knows about configured branches.",
|
|
247
|
+
"{branch_names}": "The system prompt must contain {branch_names} so the agent knows branch names."
|
|
248
|
+
}
|
|
249
|
+
for variable_name, error_message in required_prompt_variables.items():
|
|
250
|
+
if variable_name not in system_prompt:
|
|
251
|
+
errors.append(error_message)
|
|
252
|
+
if errors:
|
|
253
|
+
error_text = "\n".join(errors)
|
|
254
|
+
logger.error("%s %s", AGENT_LOG_PREFIX, error_text)
|
|
255
|
+
return False
|
|
256
|
+
return True
|
|
257
|
+
|
|
258
|
+
@staticmethod
|
|
259
|
+
def validate_user_prompt(user_prompt: str) -> bool:
|
|
260
|
+
"""Validate that the user prompt contains required variables.
|
|
261
|
+
|
|
262
|
+
Checks that the user prompt includes necessary template variables
|
|
263
|
+
for chat history and other required information.
|
|
264
|
+
|
|
265
|
+
Args:
|
|
266
|
+
user_prompt: The user prompt string to validate.
|
|
267
|
+
|
|
268
|
+
Returns:
|
|
269
|
+
True if the prompt is valid, False otherwise.
|
|
270
|
+
"""
|
|
271
|
+
errors = []
|
|
272
|
+
if not user_prompt:
|
|
273
|
+
errors.append("The user prompt cannot be empty.")
|
|
274
|
+
else:
|
|
275
|
+
required_prompt_variables = {
|
|
276
|
+
"{chat_history}":
|
|
277
|
+
"The user prompt must contain {chat_history} so the agent knows about the conversation history.",
|
|
278
|
+
"{request}":
|
|
279
|
+
"The user prompt must contain {request} so the agent sees the current request.",
|
|
280
|
+
}
|
|
281
|
+
for variable_name, error_message in required_prompt_variables.items():
|
|
282
|
+
if variable_name not in user_prompt:
|
|
283
|
+
errors.append(error_message)
|
|
284
|
+
if errors:
|
|
285
|
+
error_text = "\n".join(errors)
|
|
286
|
+
logger.error("%s %s", AGENT_LOG_PREFIX, error_text)
|
|
287
|
+
return False
|
|
288
|
+
return True
|
|
289
|
+
|
|
290
|
+
|
|
291
|
+
def create_router_agent_prompt(config: "RouterAgentWorkflowConfig") -> ChatPromptTemplate:
|
|
292
|
+
"""Create a Router Agent prompt from the configuration.
|
|
293
|
+
|
|
294
|
+
Builds a ChatPromptTemplate using either custom prompts from the config
|
|
295
|
+
or default system and user prompts. Validates the prompts to ensure they
|
|
296
|
+
contain required template variables.
|
|
297
|
+
|
|
298
|
+
Args:
|
|
299
|
+
config: The router agent workflow configuration containing prompt settings.
|
|
300
|
+
|
|
301
|
+
Returns:
|
|
302
|
+
A configured ChatPromptTemplate for the router agent.
|
|
303
|
+
|
|
304
|
+
Raises:
|
|
305
|
+
ValueError: If the system_prompt or user_prompt validation fails.
|
|
306
|
+
"""
|
|
307
|
+
from nat.agent.router_agent.prompt import SYSTEM_PROMPT
|
|
308
|
+
from nat.agent.router_agent.prompt import USER_PROMPT
|
|
309
|
+
# the Router Agent prompt can be customized via config option system_prompt and user_prompt.
|
|
310
|
+
|
|
311
|
+
if config.system_prompt:
|
|
312
|
+
system_prompt = config.system_prompt
|
|
313
|
+
else:
|
|
314
|
+
system_prompt = SYSTEM_PROMPT
|
|
315
|
+
|
|
316
|
+
if config.user_prompt:
|
|
317
|
+
user_prompt = config.user_prompt
|
|
318
|
+
else:
|
|
319
|
+
user_prompt = USER_PROMPT
|
|
320
|
+
|
|
321
|
+
if not RouterAgentGraph.validate_system_prompt(system_prompt):
|
|
322
|
+
logger.error("%s Invalid system_prompt", AGENT_LOG_PREFIX)
|
|
323
|
+
raise ValueError("Invalid system_prompt")
|
|
324
|
+
|
|
325
|
+
if not RouterAgentGraph.validate_user_prompt(user_prompt):
|
|
326
|
+
logger.error("%s Invalid user_prompt", AGENT_LOG_PREFIX)
|
|
327
|
+
raise ValueError("Invalid user_prompt")
|
|
328
|
+
|
|
329
|
+
return ChatPromptTemplate([("system", system_prompt), ("user", user_prompt)])
|
|
@@ -0,0 +1,48 @@
|
|
|
1
|
+
# SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
|
|
2
|
+
# SPDX-License-Identifier: Apache-2.0
|
|
3
|
+
#
|
|
4
|
+
# Licensed under the Apache License, Version 2.0 (the "License");
|
|
5
|
+
# you may not use this file except in compliance with the License.
|
|
6
|
+
# You may obtain a copy of the License at
|
|
7
|
+
#
|
|
8
|
+
# http://www.apache.org/licenses/LICENSE-2.0
|
|
9
|
+
#
|
|
10
|
+
# Unless required by applicable law or agreed to in writing, software
|
|
11
|
+
# distributed under the License is distributed on an "AS IS" BASIS,
|
|
12
|
+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
13
|
+
# See the License for the specific language governing permissions and
|
|
14
|
+
# limitations under the License.
|
|
15
|
+
|
|
16
|
+
SYSTEM_PROMPT = """
|
|
17
|
+
You are a Router Agent responsible for analyzing incoming requests and routing them to the most appropriate branch.
|
|
18
|
+
|
|
19
|
+
Available branches:
|
|
20
|
+
{branches}
|
|
21
|
+
|
|
22
|
+
CRITICAL INSTRUCTIONS:
|
|
23
|
+
- Analyze the user's request carefully
|
|
24
|
+
- Select exactly ONE branch that best handles the request from: [{branch_names}]
|
|
25
|
+
- Respond with ONLY the exact branch name, nothing else
|
|
26
|
+
- Be decisive - choose the single best match, if the request could fit multiple branches,
|
|
27
|
+
choose the most specific/specialized one
|
|
28
|
+
- If no branch perfectly fits, choose the closest match
|
|
29
|
+
|
|
30
|
+
Your response MUST contain ONLY the branch name. Do not include any explanations, reasoning, or additional text.
|
|
31
|
+
|
|
32
|
+
Examples:
|
|
33
|
+
User: "How do I calculate 15 + 25?"
|
|
34
|
+
Response: calculator_tool
|
|
35
|
+
|
|
36
|
+
User: "What's the weather like today?"
|
|
37
|
+
Response: weather_service
|
|
38
|
+
|
|
39
|
+
User: "Send an email to John"
|
|
40
|
+
Response: email_tool"""
|
|
41
|
+
|
|
42
|
+
USER_PROMPT = """
|
|
43
|
+
Previous conversation history:
|
|
44
|
+
{chat_history}
|
|
45
|
+
|
|
46
|
+
To respond to the request: {request}, which branch should be chosen?
|
|
47
|
+
|
|
48
|
+
Respond with only the branch name."""
|
|
@@ -0,0 +1,97 @@
|
|
|
1
|
+
# SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
|
|
2
|
+
# SPDX-License-Identifier: Apache-2.0
|
|
3
|
+
#
|
|
4
|
+
# Licensed under the Apache License, Version 2.0 (the "License");
|
|
5
|
+
# you may not use this file except in compliance with the License.
|
|
6
|
+
# You may obtain a copy of the License at
|
|
7
|
+
#
|
|
8
|
+
# http://www.apache.org/licenses/LICENSE-2.0
|
|
9
|
+
#
|
|
10
|
+
# Unless required by applicable law or agreed to in writing, software
|
|
11
|
+
# distributed under the License is distributed on an "AS IS" BASIS,
|
|
12
|
+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
13
|
+
# See the License for the specific language governing permissions and
|
|
14
|
+
# limitations under the License.
|
|
15
|
+
|
|
16
|
+
import logging
|
|
17
|
+
|
|
18
|
+
from pydantic import Field
|
|
19
|
+
from pydantic import PositiveInt
|
|
20
|
+
|
|
21
|
+
from nat.builder.builder import Builder
|
|
22
|
+
from nat.builder.framework_enum import LLMFrameworkEnum
|
|
23
|
+
from nat.builder.function_info import FunctionInfo
|
|
24
|
+
from nat.cli.register_workflow import register_function
|
|
25
|
+
from nat.data_models.component_ref import FunctionRef
|
|
26
|
+
from nat.data_models.component_ref import LLMRef
|
|
27
|
+
from nat.data_models.function import FunctionBaseConfig
|
|
28
|
+
|
|
29
|
+
logger = logging.getLogger(__name__)
|
|
30
|
+
|
|
31
|
+
|
|
32
|
+
class RouterAgentWorkflowConfig(FunctionBaseConfig, name="router_agent"):
|
|
33
|
+
"""
|
|
34
|
+
A router agent takes in the incoming message, combines it with a prompt and the list of branches,
|
|
35
|
+
and ask a LLM about which branch to take.
|
|
36
|
+
"""
|
|
37
|
+
branches: list[FunctionRef] = Field(default_factory=list,
|
|
38
|
+
description="The list of branches to provide to the router agent.")
|
|
39
|
+
llm_name: LLMRef = Field(description="The LLM model to use with the routing agent.")
|
|
40
|
+
description: str = Field(default="Router Agent Workflow", description="Description of this functions use.")
|
|
41
|
+
system_prompt: str | None = Field(default=None, description="Provides the system prompt to use with the agent.")
|
|
42
|
+
user_prompt: str | None = Field(default=None, description="Provides the prompt to use with the agent.")
|
|
43
|
+
max_router_retries: int = Field(
|
|
44
|
+
default=3, description="Maximum number of retries if the router agent fails to choose a branch.")
|
|
45
|
+
detailed_logs: bool = Field(default=False, description="Set the verbosity of the router agent's logging.")
|
|
46
|
+
log_response_max_chars: PositiveInt = Field(
|
|
47
|
+
default=1000, description="Maximum number of characters to display in logs when logging branch responses.")
|
|
48
|
+
|
|
49
|
+
|
|
50
|
+
@register_function(config_type=RouterAgentWorkflowConfig, framework_wrappers=[LLMFrameworkEnum.LANGCHAIN])
|
|
51
|
+
async def router_agent_workflow(config: RouterAgentWorkflowConfig, builder: Builder):
|
|
52
|
+
from langchain_core.messages.human import HumanMessage
|
|
53
|
+
from langgraph.graph.state import CompiledStateGraph
|
|
54
|
+
|
|
55
|
+
from nat.agent.base import AGENT_LOG_PREFIX
|
|
56
|
+
from nat.agent.router_agent.agent import RouterAgentGraph
|
|
57
|
+
from nat.agent.router_agent.agent import RouterAgentGraphState
|
|
58
|
+
from nat.agent.router_agent.agent import create_router_agent_prompt
|
|
59
|
+
|
|
60
|
+
prompt = create_router_agent_prompt(config)
|
|
61
|
+
llm = await builder.get_llm(config.llm_name, wrapper_type=LLMFrameworkEnum.LANGCHAIN)
|
|
62
|
+
branches = builder.get_tools(tool_names=config.branches, wrapper_type=LLMFrameworkEnum.LANGCHAIN)
|
|
63
|
+
if not branches:
|
|
64
|
+
raise ValueError(f"No branches specified for Router Agent '{config.llm_name}'")
|
|
65
|
+
|
|
66
|
+
graph: CompiledStateGraph = await RouterAgentGraph(
|
|
67
|
+
llm=llm,
|
|
68
|
+
branches=branches,
|
|
69
|
+
prompt=prompt,
|
|
70
|
+
max_router_retries=config.max_router_retries,
|
|
71
|
+
detailed_logs=config.detailed_logs,
|
|
72
|
+
log_response_max_chars=config.log_response_max_chars,
|
|
73
|
+
).build_graph()
|
|
74
|
+
|
|
75
|
+
async def _response_fn(input_message: str) -> str:
|
|
76
|
+
try:
|
|
77
|
+
message = HumanMessage(content=input_message)
|
|
78
|
+
state = RouterAgentGraphState(forward_message=message)
|
|
79
|
+
|
|
80
|
+
result_dict = await graph.ainvoke(state)
|
|
81
|
+
result_state = RouterAgentGraphState(**result_dict)
|
|
82
|
+
|
|
83
|
+
output_message = result_state.messages[-1]
|
|
84
|
+
return str(output_message.content)
|
|
85
|
+
|
|
86
|
+
except Exception as ex:
|
|
87
|
+
logger.exception("%s Router Agent failed with exception: %s", AGENT_LOG_PREFIX, ex)
|
|
88
|
+
if config.detailed_logs:
|
|
89
|
+
return str(ex)
|
|
90
|
+
return "Router agent failed with exception: %s" % ex
|
|
91
|
+
|
|
92
|
+
try:
|
|
93
|
+
yield FunctionInfo.from_fn(_response_fn, description=config.description)
|
|
94
|
+
except GeneratorExit:
|
|
95
|
+
logger.exception("%s Workflow exited early!", AGENT_LOG_PREFIX)
|
|
96
|
+
finally:
|
|
97
|
+
logger.debug("%s Cleaning up router_agent workflow.", AGENT_LOG_PREFIX)
|
|
@@ -19,10 +19,13 @@ import typing
|
|
|
19
19
|
from langchain_core.callbacks.base import AsyncCallbackHandler
|
|
20
20
|
from langchain_core.language_models import BaseChatModel
|
|
21
21
|
from langchain_core.messages import SystemMessage
|
|
22
|
+
from langchain_core.messages import ToolMessage
|
|
22
23
|
from langchain_core.messages.base import BaseMessage
|
|
23
24
|
from langchain_core.runnables import RunnableLambda
|
|
24
25
|
from langchain_core.runnables.config import RunnableConfig
|
|
25
26
|
from langchain_core.tools import BaseTool
|
|
27
|
+
from langgraph.graph import StateGraph
|
|
28
|
+
from langgraph.graph.state import CompiledStateGraph
|
|
26
29
|
from langgraph.prebuilt import ToolNode
|
|
27
30
|
from pydantic import BaseModel
|
|
28
31
|
from pydantic import Field
|
|
@@ -57,12 +60,14 @@ class ToolCallAgentGraph(DualNodeAgent):
|
|
|
57
60
|
detailed_logs: bool = False,
|
|
58
61
|
log_response_max_chars: int = 1000,
|
|
59
62
|
handle_tool_errors: bool = True,
|
|
63
|
+
return_direct: list[BaseTool] | None = None,
|
|
60
64
|
):
|
|
61
65
|
super().__init__(llm=llm,
|
|
62
66
|
tools=tools,
|
|
63
67
|
callbacks=callbacks,
|
|
64
68
|
detailed_logs=detailed_logs,
|
|
65
69
|
log_response_max_chars=log_response_max_chars)
|
|
70
|
+
|
|
66
71
|
# some LLMs support tool calling
|
|
67
72
|
# these models accept the tool's input schema and decide when to use a tool based on the input's relevance
|
|
68
73
|
try:
|
|
@@ -85,8 +90,8 @@ class ToolCallAgentGraph(DualNodeAgent):
|
|
|
85
90
|
)
|
|
86
91
|
|
|
87
92
|
self.agent = prompt_runnable | self.bound_llm
|
|
88
|
-
|
|
89
93
|
self.tool_caller = ToolNode(tools, handle_tool_errors=handle_tool_errors)
|
|
94
|
+
self.return_direct = [tool.name for tool in return_direct] if return_direct else []
|
|
90
95
|
logger.debug("%s Initialized Tool Calling Agent Graph", AGENT_LOG_PREFIX)
|
|
91
96
|
|
|
92
97
|
async def agent_node(self, state: ToolCallAgentGraphState):
|
|
@@ -146,13 +151,70 @@ class ToolCallAgentGraph(DualNodeAgent):
|
|
|
146
151
|
logger.error("%s Failed to call tool_node: %s", AGENT_LOG_PREFIX, ex)
|
|
147
152
|
raise
|
|
148
153
|
|
|
149
|
-
async def
|
|
154
|
+
async def tool_conditional_edge(self, state: ToolCallAgentGraphState) -> AgentDecision:
|
|
155
|
+
"""
|
|
156
|
+
Determines whether to continue to the agent or end graph execution after a tool call.
|
|
157
|
+
|
|
158
|
+
Args:
|
|
159
|
+
state: The current state of the Tool Calling Agent graph containing messages and tool responses.
|
|
160
|
+
|
|
161
|
+
Returns:
|
|
162
|
+
AgentDecision: TOOL to continue to agent for processing, or END to terminate graph execution.
|
|
163
|
+
Returns END if the tool is in return_direct list, otherwise returns TOOL to continue processing.
|
|
164
|
+
"""
|
|
150
165
|
try:
|
|
151
|
-
|
|
152
|
-
|
|
153
|
-
"%s
|
|
154
|
-
|
|
155
|
-
|
|
166
|
+
logger.debug("%s Starting the Tool Conditional Edge", AGENT_LOG_PREFIX)
|
|
167
|
+
if not state.messages:
|
|
168
|
+
logger.debug("%s No messages in state; routing to agent", AGENT_LOG_PREFIX)
|
|
169
|
+
return AgentDecision.TOOL
|
|
170
|
+
|
|
171
|
+
last_message = state.messages[-1]
|
|
172
|
+
# Return directly if this tool is in the return_direct set
|
|
173
|
+
if (self.return_direct and isinstance(last_message, ToolMessage) and last_message.name
|
|
174
|
+
and last_message.name in self.return_direct):
|
|
175
|
+
# Return directly if this tool is in the return_direct list
|
|
176
|
+
logger.debug("%s Tool %s is set to return directly", AGENT_LOG_PREFIX, last_message.name)
|
|
177
|
+
return AgentDecision.END
|
|
178
|
+
else:
|
|
179
|
+
# Continue to agent for processing
|
|
180
|
+
logger.debug("%s Tool response will be processed by agent", AGENT_LOG_PREFIX)
|
|
181
|
+
return AgentDecision.TOOL
|
|
182
|
+
except Exception as ex:
|
|
183
|
+
logger.exception("%s Failed to determine tool conditional edge: %s", AGENT_LOG_PREFIX, ex)
|
|
184
|
+
logger.warning("%s Continuing to agent for processing", AGENT_LOG_PREFIX)
|
|
185
|
+
return AgentDecision.TOOL
|
|
186
|
+
|
|
187
|
+
async def _build_graph(self, state_schema: type) -> CompiledStateGraph:
|
|
188
|
+
try:
|
|
189
|
+
logger.debug("%s Building and compiling the Tool Calling Agent Graph", AGENT_LOG_PREFIX)
|
|
190
|
+
|
|
191
|
+
graph = StateGraph(state_schema)
|
|
192
|
+
graph.add_node("agent", self.agent_node)
|
|
193
|
+
graph.add_node("tool", self.tool_node)
|
|
194
|
+
|
|
195
|
+
if self.return_direct:
|
|
196
|
+
# go to end of graph if tool is set to return directly
|
|
197
|
+
tool_conditional_edge_possible_outputs = {AgentDecision.END: "__end__", AgentDecision.TOOL: "agent"}
|
|
198
|
+
graph.add_conditional_edges("tool", self.tool_conditional_edge, tool_conditional_edge_possible_outputs)
|
|
199
|
+
else:
|
|
200
|
+
# otherwise return to agent after tool call
|
|
201
|
+
graph.add_edge("tool", "agent")
|
|
202
|
+
|
|
203
|
+
conditional_edge_possible_outputs = {AgentDecision.TOOL: "tool", AgentDecision.END: "__end__"}
|
|
204
|
+
graph.add_conditional_edges("agent", self.conditional_edge, conditional_edge_possible_outputs)
|
|
205
|
+
|
|
206
|
+
graph.set_entry_point("agent")
|
|
207
|
+
self.graph = graph.compile()
|
|
208
|
+
|
|
209
|
+
return self.graph
|
|
210
|
+
except Exception as ex:
|
|
211
|
+
logger.error("%s Failed to build Tool Calling Agent Graph: %s", AGENT_LOG_PREFIX, ex)
|
|
212
|
+
raise
|
|
213
|
+
|
|
214
|
+
async def build_graph(self) -> CompiledStateGraph:
|
|
215
|
+
try:
|
|
216
|
+
await self._build_graph(state_schema=ToolCallAgentGraphState)
|
|
217
|
+
logger.debug("%s Tool Calling Agent Graph built and compiled successfully", AGENT_LOG_PREFIX)
|
|
156
218
|
return self.graph
|
|
157
219
|
except Exception as ex:
|
|
158
220
|
logger.error("%s Failed to build Tool Calling Agent Graph: %s", AGENT_LOG_PREFIX, ex)
|
|
@@ -22,6 +22,7 @@ from nat.builder.builder import Builder
|
|
|
22
22
|
from nat.builder.framework_enum import LLMFrameworkEnum
|
|
23
23
|
from nat.builder.function_info import FunctionInfo
|
|
24
24
|
from nat.cli.register_workflow import register_function
|
|
25
|
+
from nat.data_models.component_ref import FunctionGroupRef
|
|
25
26
|
from nat.data_models.component_ref import FunctionRef
|
|
26
27
|
from nat.data_models.component_ref import LLMRef
|
|
27
28
|
from nat.data_models.function import FunctionBaseConfig
|
|
@@ -35,8 +36,8 @@ class ToolCallAgentWorkflowConfig(FunctionBaseConfig, name="tool_calling_agent")
|
|
|
35
36
|
input parameters to select the optimal tool. Supports handling tool errors.
|
|
36
37
|
"""
|
|
37
38
|
|
|
38
|
-
tool_names: list[FunctionRef] = Field(
|
|
39
|
-
|
|
39
|
+
tool_names: list[FunctionRef | FunctionGroupRef] = Field(
|
|
40
|
+
default_factory=list, description="The list of tools to provide to the tool calling agent.")
|
|
40
41
|
llm_name: LLMRef = Field(description="The LLM model to use with the tool calling agent.")
|
|
41
42
|
verbose: bool = Field(default=False, description="Set the verbosity of the tool calling agent's logging.")
|
|
42
43
|
handle_tool_errors: bool = Field(default=True, description="Specify ability to handle tool calling errors.")
|
|
@@ -47,6 +48,8 @@ class ToolCallAgentWorkflowConfig(FunctionBaseConfig, name="tool_calling_agent")
|
|
|
47
48
|
system_prompt: str | None = Field(default=None, description="Provides the system prompt to use with the agent.")
|
|
48
49
|
additional_instructions: str | None = Field(default=None,
|
|
49
50
|
description="Additional instructions appended to the system prompt.")
|
|
51
|
+
return_direct: list[FunctionRef] | None = Field(
|
|
52
|
+
default=None, description="List of tool names that should return responses directly without LLM processing.")
|
|
50
53
|
|
|
51
54
|
|
|
52
55
|
@register_function(config_type=ToolCallAgentWorkflowConfig, framework_wrappers=[LLMFrameworkEnum.LANGCHAIN])
|
|
@@ -68,13 +71,18 @@ async def tool_calling_agent_workflow(config: ToolCallAgentWorkflowConfig, build
|
|
|
68
71
|
if not tools:
|
|
69
72
|
raise ValueError(f"No tools specified for Tool Calling Agent '{config.llm_name}'")
|
|
70
73
|
|
|
74
|
+
# convert return_direct FunctionRef objects to BaseTool objects
|
|
75
|
+
return_direct_tools = builder.get_tools(tool_names=config.return_direct,
|
|
76
|
+
wrapper_type=LLMFrameworkEnum.LANGCHAIN) if config.return_direct else None
|
|
77
|
+
|
|
71
78
|
# construct the Tool Calling Agent Graph from the configured llm, and tools
|
|
72
79
|
graph: CompiledStateGraph = await ToolCallAgentGraph(llm=llm,
|
|
73
80
|
tools=tools,
|
|
74
81
|
prompt=prompt,
|
|
75
82
|
detailed_logs=config.verbose,
|
|
76
83
|
log_response_max_chars=config.log_response_max_chars,
|
|
77
|
-
handle_tool_errors=config.handle_tool_errors
|
|
84
|
+
handle_tool_errors=config.handle_tool_errors,
|
|
85
|
+
return_direct=return_direct_tools).build_graph()
|
|
78
86
|
|
|
79
87
|
async def _response_fn(input_message: str) -> str:
|
|
80
88
|
try:
|