nvidia-nat 1.3.0a20250910__py3-none-any.whl → 1.4.0a20251112__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 +13 -8
- nat/agent/prompt_optimizer/prompt.py +68 -0
- nat/agent/prompt_optimizer/register.py +149 -0
- nat/agent/react_agent/agent.py +6 -5
- nat/agent/react_agent/register.py +49 -39
- nat/agent/reasoning_agent/reasoning_agent.py +17 -15
- nat/agent/register.py +2 -0
- nat/agent/responses_api_agent/__init__.py +14 -0
- nat/agent/responses_api_agent/register.py +126 -0
- nat/agent/rewoo_agent/agent.py +304 -117
- nat/agent/rewoo_agent/prompt.py +19 -22
- nat/agent/rewoo_agent/register.py +51 -38
- nat/agent/tool_calling_agent/agent.py +75 -17
- nat/agent/tool_calling_agent/register.py +46 -23
- nat/authentication/api_key/api_key_auth_provider.py +6 -11
- nat/authentication/api_key/api_key_auth_provider_config.py +8 -5
- nat/authentication/credential_validator/__init__.py +14 -0
- nat/authentication/credential_validator/bearer_token_validator.py +557 -0
- nat/authentication/http_basic_auth/http_basic_auth_provider.py +1 -1
- nat/authentication/interfaces.py +5 -2
- nat/authentication/oauth2/oauth2_auth_code_flow_provider.py +69 -36
- nat/authentication/oauth2/oauth2_auth_code_flow_provider_config.py +2 -1
- nat/authentication/oauth2/oauth2_resource_server_config.py +125 -0
- nat/builder/builder.py +55 -23
- nat/builder/component_utils.py +9 -5
- nat/builder/context.py +54 -15
- nat/builder/eval_builder.py +14 -9
- nat/builder/framework_enum.py +1 -0
- nat/builder/front_end.py +1 -1
- nat/builder/function.py +370 -0
- nat/builder/function_info.py +1 -1
- nat/builder/intermediate_step_manager.py +38 -2
- nat/builder/workflow.py +5 -0
- nat/builder/workflow_builder.py +306 -54
- nat/cli/cli_utils/config_override.py +1 -1
- nat/cli/commands/info/info.py +16 -6
- nat/cli/commands/mcp/__init__.py +14 -0
- nat/cli/commands/mcp/mcp.py +986 -0
- nat/cli/commands/optimize.py +90 -0
- nat/cli/commands/start.py +1 -1
- nat/cli/commands/workflow/templates/config.yml.j2 +14 -13
- nat/cli/commands/workflow/templates/register.py.j2 +2 -2
- nat/cli/commands/workflow/templates/workflow.py.j2 +35 -21
- nat/cli/commands/workflow/workflow_commands.py +60 -18
- nat/cli/entrypoint.py +15 -11
- nat/cli/main.py +3 -0
- nat/cli/register_workflow.py +38 -4
- nat/cli/type_registry.py +72 -1
- nat/control_flow/__init__.py +0 -0
- nat/control_flow/register.py +20 -0
- nat/control_flow/router_agent/__init__.py +0 -0
- nat/control_flow/router_agent/agent.py +329 -0
- nat/control_flow/router_agent/prompt.py +48 -0
- nat/control_flow/router_agent/register.py +91 -0
- nat/control_flow/sequential_executor.py +166 -0
- nat/data_models/agent.py +34 -0
- nat/data_models/api_server.py +199 -69
- nat/data_models/authentication.py +23 -9
- nat/data_models/common.py +47 -0
- nat/data_models/component.py +2 -0
- nat/data_models/component_ref.py +11 -0
- nat/data_models/config.py +41 -17
- nat/data_models/dataset_handler.py +4 -3
- nat/data_models/function.py +34 -0
- nat/data_models/function_dependencies.py +8 -0
- nat/data_models/intermediate_step.py +9 -1
- nat/data_models/llm.py +15 -1
- nat/data_models/openai_mcp.py +46 -0
- nat/data_models/optimizable.py +208 -0
- nat/data_models/optimizer.py +161 -0
- nat/data_models/span.py +41 -3
- nat/data_models/thinking_mixin.py +2 -2
- nat/embedder/azure_openai_embedder.py +2 -1
- nat/embedder/nim_embedder.py +3 -2
- nat/embedder/openai_embedder.py +3 -2
- nat/eval/config.py +1 -1
- nat/eval/dataset_handler/dataset_downloader.py +3 -2
- nat/eval/dataset_handler/dataset_filter.py +34 -2
- nat/eval/evaluate.py +10 -3
- nat/eval/evaluator/base_evaluator.py +1 -1
- nat/eval/rag_evaluator/evaluate.py +7 -4
- 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/eval/swe_bench_evaluator/evaluate.py +1 -1
- nat/eval/trajectory_evaluator/register.py +1 -1
- nat/eval/tunable_rag_evaluator/evaluate.py +1 -1
- nat/eval/usage_stats.py +2 -0
- nat/eval/utils/output_uploader.py +3 -2
- nat/eval/utils/weave_eval.py +17 -3
- nat/experimental/decorators/experimental_warning_decorator.py +27 -7
- nat/experimental/test_time_compute/functions/execute_score_select_function.py +1 -1
- nat/experimental/test_time_compute/functions/plan_select_execute_function.py +7 -3
- nat/experimental/test_time_compute/functions/ttc_tool_orchestration_function.py +1 -1
- nat/experimental/test_time_compute/functions/ttc_tool_wrapper_function.py +3 -3
- nat/experimental/test_time_compute/models/strategy_base.py +2 -2
- nat/experimental/test_time_compute/selection/llm_based_output_merging_selector.py +1 -1
- nat/front_ends/console/authentication_flow_handler.py +82 -30
- nat/front_ends/console/console_front_end_plugin.py +19 -7
- nat/front_ends/fastapi/auth_flow_handlers/http_flow_handler.py +1 -1
- nat/front_ends/fastapi/auth_flow_handlers/websocket_flow_handler.py +52 -17
- nat/front_ends/fastapi/dask_client_mixin.py +65 -0
- nat/front_ends/fastapi/fastapi_front_end_config.py +25 -3
- nat/front_ends/fastapi/fastapi_front_end_plugin.py +140 -3
- nat/front_ends/fastapi/fastapi_front_end_plugin_worker.py +445 -265
- nat/front_ends/fastapi/job_store.py +518 -99
- nat/front_ends/fastapi/main.py +11 -19
- nat/front_ends/fastapi/message_handler.py +69 -44
- nat/front_ends/fastapi/message_validator.py +8 -7
- nat/front_ends/fastapi/utils.py +57 -0
- nat/front_ends/mcp/introspection_token_verifier.py +73 -0
- nat/front_ends/mcp/mcp_front_end_config.py +71 -3
- nat/front_ends/mcp/mcp_front_end_plugin.py +85 -21
- nat/front_ends/mcp/mcp_front_end_plugin_worker.py +248 -29
- nat/front_ends/mcp/memory_profiler.py +320 -0
- nat/front_ends/mcp/tool_converter.py +78 -25
- nat/front_ends/simple_base/simple_front_end_plugin_base.py +3 -1
- nat/llm/aws_bedrock_llm.py +21 -8
- nat/llm/azure_openai_llm.py +14 -5
- nat/llm/litellm_llm.py +80 -0
- nat/llm/nim_llm.py +23 -9
- nat/llm/openai_llm.py +19 -7
- nat/llm/register.py +4 -0
- nat/llm/utils/thinking.py +1 -1
- nat/observability/exporter/base_exporter.py +1 -1
- nat/observability/exporter/processing_exporter.py +29 -55
- nat/observability/exporter/span_exporter.py +43 -15
- nat/observability/exporter_manager.py +2 -2
- 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 +420 -107
- nat/observability/processor/batching_processor.py +1 -1
- 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/observability/register.py +16 -0
- nat/profiler/callbacks/langchain_callback_handler.py +32 -7
- nat/profiler/callbacks/llama_index_callback_handler.py +36 -2
- nat/profiler/callbacks/token_usage_base_model.py +2 -0
- nat/profiler/decorators/framework_wrapper.py +61 -9
- nat/profiler/decorators/function_tracking.py +35 -3
- nat/profiler/forecasting/models/linear_model.py +1 -1
- nat/profiler/forecasting/models/random_forest_regressor.py +1 -1
- nat/profiler/inference_optimization/bottleneck_analysis/nested_stack_analysis.py +1 -1
- nat/profiler/inference_optimization/experimental/prefix_span_analysis.py +1 -1
- 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 +189 -0
- nat/profiler/parameter_optimization/parameter_selection.py +107 -0
- nat/profiler/parameter_optimization/pareto_visualizer.py +460 -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/registry_handlers/pypi/register_pypi.py +5 -3
- nat/registry_handlers/rest/register_rest.py +5 -3
- nat/retriever/milvus/retriever.py +1 -1
- nat/retriever/nemo_retriever/register.py +2 -1
- nat/runtime/loader.py +1 -1
- nat/runtime/runner.py +111 -6
- nat/runtime/session.py +49 -3
- nat/settings/global_settings.py +2 -2
- nat/tool/chat_completion.py +4 -1
- nat/tool/code_execution/code_sandbox.py +3 -6
- nat/tool/code_execution/local_sandbox/Dockerfile.sandbox +19 -32
- nat/tool/code_execution/local_sandbox/local_sandbox_server.py +6 -1
- nat/tool/code_execution/local_sandbox/sandbox.requirements.txt +2 -0
- nat/tool/code_execution/local_sandbox/start_local_sandbox.sh +10 -4
- nat/tool/datetime_tools.py +1 -1
- nat/tool/github_tools.py +450 -0
- nat/tool/memory_tools/add_memory_tool.py +3 -3
- nat/tool/memory_tools/delete_memory_tool.py +3 -4
- nat/tool/memory_tools/get_memory_tool.py +4 -4
- nat/tool/register.py +2 -7
- nat/tool/server_tools.py +15 -2
- nat/utils/__init__.py +76 -0
- nat/utils/callable_utils.py +70 -0
- nat/utils/data_models/schema_validator.py +1 -1
- nat/utils/decorators.py +210 -0
- nat/utils/exception_handlers/automatic_retries.py +278 -72
- nat/utils/io/yaml_tools.py +73 -3
- nat/utils/log_levels.py +25 -0
- nat/utils/responses_api.py +26 -0
- nat/utils/string_utils.py +16 -0
- nat/utils/type_converter.py +12 -3
- nat/utils/type_utils.py +6 -2
- nvidia_nat-1.4.0a20251112.dist-info/METADATA +197 -0
- {nvidia_nat-1.3.0a20250910.dist-info → nvidia_nat-1.4.0a20251112.dist-info}/RECORD +199 -165
- {nvidia_nat-1.3.0a20250910.dist-info → nvidia_nat-1.4.0a20251112.dist-info}/entry_points.txt +1 -0
- nat/cli/commands/info/list_mcp.py +0 -461
- nat/data_models/temperature_mixin.py +0 -43
- nat/data_models/top_p_mixin.py +0 -43
- nat/observability/processor/header_redaction_processor.py +0 -123
- nat/observability/processor/redaction_processor.py +0 -77
- nat/tool/code_execution/test_code_execution_sandbox.py +0 -414
- 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
- nvidia_nat-1.3.0a20250910.dist-info/METADATA +0 -373
- /nat/{tool/github_tools → agent/prompt_optimizer}/__init__.py +0 -0
- {nvidia_nat-1.3.0a20250910.dist-info → nvidia_nat-1.4.0a20251112.dist-info}/WHEEL +0 -0
- {nvidia_nat-1.3.0a20250910.dist-info → nvidia_nat-1.4.0a20251112.dist-info}/licenses/LICENSE-3rd-party.txt +0 -0
- {nvidia_nat-1.3.0a20250910.dist-info → nvidia_nat-1.4.0a20251112.dist-info}/licenses/LICENSE.md +0 -0
- {nvidia_nat-1.3.0a20250910.dist-info → nvidia_nat-1.4.0a20251112.dist-info}/top_level.txt +0 -0
nat/agent/rewoo_agent/agent.py
CHANGED
|
@@ -13,9 +13,12 @@
|
|
|
13
13
|
# See the License for the specific language governing permissions and
|
|
14
14
|
# limitations under the License.
|
|
15
15
|
|
|
16
|
+
import asyncio
|
|
16
17
|
import json
|
|
17
18
|
import logging
|
|
19
|
+
import re
|
|
18
20
|
from json import JSONDecodeError
|
|
21
|
+
from typing import Any
|
|
19
22
|
|
|
20
23
|
from langchain_core.callbacks.base import AsyncCallbackHandler
|
|
21
24
|
from langchain_core.language_models import BaseChatModel
|
|
@@ -42,6 +45,17 @@ from nat.agent.base import BaseAgent
|
|
|
42
45
|
logger = logging.getLogger(__name__)
|
|
43
46
|
|
|
44
47
|
|
|
48
|
+
class ReWOOEvidence(BaseModel):
|
|
49
|
+
placeholder: str
|
|
50
|
+
tool: str
|
|
51
|
+
tool_input: Any
|
|
52
|
+
|
|
53
|
+
|
|
54
|
+
class ReWOOPlanStep(BaseModel):
|
|
55
|
+
plan: str
|
|
56
|
+
evidence: ReWOOEvidence
|
|
57
|
+
|
|
58
|
+
|
|
45
59
|
class ReWOOGraphState(BaseModel):
|
|
46
60
|
"""State schema for the ReWOO Agent Graph"""
|
|
47
61
|
messages: list[BaseMessage] = Field(default_factory=list) # input and output of the ReWOO Agent
|
|
@@ -50,15 +64,21 @@ class ReWOOGraphState(BaseModel):
|
|
|
50
64
|
default_factory=lambda: AIMessage(content="")) # the plan generated by the planner to solve the task
|
|
51
65
|
steps: AIMessage = Field(
|
|
52
66
|
default_factory=lambda: AIMessage(content="")) # the steps to solve the task, parsed from the plan
|
|
67
|
+
# New fields for parallel execution support
|
|
68
|
+
evidence_map: dict[str, ReWOOPlanStep] = Field(default_factory=dict) # mapping from placeholders to step info
|
|
69
|
+
execution_levels: list[list[str]] = Field(default_factory=list) # levels for parallel execution
|
|
70
|
+
current_level: int = Field(default=0) # current execution level
|
|
53
71
|
intermediate_results: dict[str, ToolMessage] = Field(default_factory=dict) # the intermediate results of each step
|
|
54
72
|
result: AIMessage = Field(
|
|
55
73
|
default_factory=lambda: AIMessage(content="")) # the final result of the task, generated by the solver
|
|
56
74
|
|
|
57
75
|
|
|
58
76
|
class ReWOOAgentGraph(BaseAgent):
|
|
59
|
-
"""Configurable
|
|
60
|
-
|
|
61
|
-
|
|
77
|
+
"""Configurable ReWOO Agent.
|
|
78
|
+
|
|
79
|
+
Args:
|
|
80
|
+
detailed_logs: Toggles logging of inputs, outputs, and intermediate steps.
|
|
81
|
+
"""
|
|
62
82
|
|
|
63
83
|
def __init__(self,
|
|
64
84
|
llm: BaseChatModel,
|
|
@@ -68,7 +88,9 @@ class ReWOOAgentGraph(BaseAgent):
|
|
|
68
88
|
use_tool_schema: bool = True,
|
|
69
89
|
callbacks: list[AsyncCallbackHandler] | None = None,
|
|
70
90
|
detailed_logs: bool = False,
|
|
71
|
-
log_response_max_chars: int = 1000
|
|
91
|
+
log_response_max_chars: int = 1000,
|
|
92
|
+
tool_call_max_retries: int = 3,
|
|
93
|
+
raise_tool_call_error: bool = True):
|
|
72
94
|
super().__init__(llm=llm,
|
|
73
95
|
tools=tools,
|
|
74
96
|
callbacks=callbacks,
|
|
@@ -78,22 +100,21 @@ class ReWOOAgentGraph(BaseAgent):
|
|
|
78
100
|
logger.debug(
|
|
79
101
|
"%s Filling the prompt variables 'tools' and 'tool_names', using the tools provided in the config.",
|
|
80
102
|
AGENT_LOG_PREFIX)
|
|
81
|
-
|
|
82
|
-
|
|
83
|
-
|
|
84
|
-
|
|
85
|
-
|
|
86
|
-
|
|
87
|
-
|
|
88
|
-
|
|
89
|
-
|
|
90
|
-
for tool in tools[:-1]
|
|
91
|
-
]) + "\n" + (f"{tools[-1].name}: {tools[-1].description}. "
|
|
92
|
-
f"{INPUT_SCHEMA_MESSAGE.format(schema=tools[-1].input_schema.model_fields)}")
|
|
103
|
+
|
|
104
|
+
def describe_tool(tool: BaseTool) -> str:
|
|
105
|
+
description = f"{tool.name}: {tool.description}"
|
|
106
|
+
if use_tool_schema:
|
|
107
|
+
description += f". {INPUT_SCHEMA_MESSAGE.format(schema=tool.input_schema.model_fields)}"
|
|
108
|
+
return description
|
|
109
|
+
|
|
110
|
+
tool_names = ",".join(tool.name for tool in tools)
|
|
111
|
+
tool_names_and_descriptions = "\n".join(describe_tool(tool) for tool in tools)
|
|
93
112
|
|
|
94
113
|
self.planner_prompt = planner_prompt.partial(tools=tool_names_and_descriptions, tool_names=tool_names)
|
|
95
114
|
self.solver_prompt = solver_prompt
|
|
96
115
|
self.tools_dict = {tool.name: tool for tool in tools}
|
|
116
|
+
self.tool_call_max_retries = tool_call_max_retries
|
|
117
|
+
self.raise_tool_call_error = raise_tool_call_error
|
|
97
118
|
|
|
98
119
|
logger.debug("%s Initialized ReWOO Agent Graph", AGENT_LOG_PREFIX)
|
|
99
120
|
|
|
@@ -105,26 +126,87 @@ class ReWOOAgentGraph(BaseAgent):
|
|
|
105
126
|
raise
|
|
106
127
|
|
|
107
128
|
@staticmethod
|
|
108
|
-
def
|
|
109
|
-
|
|
110
|
-
|
|
111
|
-
|
|
129
|
+
def _get_current_level_status(state: ReWOOGraphState) -> tuple[int, bool]:
|
|
130
|
+
"""
|
|
131
|
+
Get the current execution level and whether it's complete.
|
|
132
|
+
|
|
133
|
+
Args:
|
|
134
|
+
state: The ReWOO graph state.
|
|
112
135
|
|
|
113
|
-
|
|
114
|
-
|
|
115
|
-
|
|
136
|
+
Returns:
|
|
137
|
+
tuple of (current_level, is_complete). Level -1 means all execution is complete.
|
|
138
|
+
"""
|
|
139
|
+
if not state.execution_levels:
|
|
140
|
+
return -1, True
|
|
116
141
|
|
|
117
|
-
|
|
142
|
+
current_level = state.current_level
|
|
143
|
+
|
|
144
|
+
# Check if we've completed all levels
|
|
145
|
+
if current_level >= len(state.execution_levels):
|
|
146
|
+
return -1, True
|
|
147
|
+
|
|
148
|
+
# Check if current level is complete
|
|
149
|
+
current_level_placeholders = state.execution_levels[current_level]
|
|
150
|
+
level_complete = all(placeholder in state.intermediate_results for placeholder in current_level_placeholders)
|
|
151
|
+
|
|
152
|
+
return current_level, level_complete
|
|
118
153
|
|
|
119
154
|
@staticmethod
|
|
120
|
-
def _parse_planner_output(planner_output: str) ->
|
|
155
|
+
def _parse_planner_output(planner_output: str) -> list[ReWOOPlanStep]:
|
|
121
156
|
|
|
122
157
|
try:
|
|
123
|
-
|
|
124
|
-
except
|
|
158
|
+
return [ReWOOPlanStep(**step) for step in json.loads(planner_output)]
|
|
159
|
+
except Exception as ex:
|
|
125
160
|
raise ValueError(f"The output of planner is invalid JSON format: {planner_output}") from ex
|
|
126
161
|
|
|
127
|
-
|
|
162
|
+
@staticmethod
|
|
163
|
+
def _parse_planner_dependencies(steps: list[ReWOOPlanStep]) -> tuple[dict[str, ReWOOPlanStep], list[list[str]]]:
|
|
164
|
+
"""
|
|
165
|
+
Parse planner steps to identify dependencies and create execution levels for parallel processing.
|
|
166
|
+
This creates a dependency map and identifies which evidence placeholders can be executed in parallel.
|
|
167
|
+
|
|
168
|
+
Args:
|
|
169
|
+
steps: list of plan steps from the planner.
|
|
170
|
+
|
|
171
|
+
Returns:
|
|
172
|
+
A mapping from evidence placeholders to step info and execution levels for parallel processing.
|
|
173
|
+
"""
|
|
174
|
+
# First pass: collect all evidence placeholders and their info
|
|
175
|
+
evidences: dict[str, ReWOOPlanStep] = {
|
|
176
|
+
step.evidence.placeholder: step
|
|
177
|
+
for step in steps if step.evidence and step.evidence.placeholder
|
|
178
|
+
}
|
|
179
|
+
|
|
180
|
+
# Second pass: find dependencies now that we have all placeholders
|
|
181
|
+
dependencies = {
|
|
182
|
+
step.evidence.placeholder: [
|
|
183
|
+
var for var in re.findall(r"#E\d+", str(step.evidence.tool_input))
|
|
184
|
+
if var in evidences and var != step.evidence.placeholder
|
|
185
|
+
]
|
|
186
|
+
for step in steps if step.evidence and step.evidence.placeholder
|
|
187
|
+
}
|
|
188
|
+
|
|
189
|
+
# Create execution levels using topological sort
|
|
190
|
+
levels: list[list[str]] = []
|
|
191
|
+
remaining = dict(dependencies)
|
|
192
|
+
|
|
193
|
+
while remaining:
|
|
194
|
+
# Find items with no dependencies (can be executed in parallel)
|
|
195
|
+
ready = [placeholder for placeholder, deps in remaining.items() if not deps]
|
|
196
|
+
|
|
197
|
+
if not ready:
|
|
198
|
+
raise ValueError("Circular dependency detected in planner output")
|
|
199
|
+
|
|
200
|
+
levels.append(ready)
|
|
201
|
+
|
|
202
|
+
# Remove completed items from remaining
|
|
203
|
+
for placeholder in ready:
|
|
204
|
+
remaining.pop(placeholder)
|
|
205
|
+
|
|
206
|
+
# Remove completed items from other dependencies
|
|
207
|
+
for ph, deps in list(remaining.items()):
|
|
208
|
+
remaining[ph] = list(set(deps) - set(ready))
|
|
209
|
+
return evidences, levels
|
|
128
210
|
|
|
129
211
|
@staticmethod
|
|
130
212
|
def _replace_placeholder(placeholder: str, tool_input: str | dict, tool_output: str | dict) -> str | dict:
|
|
@@ -135,7 +217,7 @@ class ReWOOAgentGraph(BaseAgent):
|
|
|
135
217
|
if value is not None:
|
|
136
218
|
if value == placeholder:
|
|
137
219
|
tool_input[key] = tool_output
|
|
138
|
-
elif placeholder in value:
|
|
220
|
+
elif isinstance(value, str) and placeholder in value:
|
|
139
221
|
# If the placeholder is part of the value, replace it with the stringified output
|
|
140
222
|
tool_input[key] = value.replace(placeholder, str(tool_output))
|
|
141
223
|
|
|
@@ -144,6 +226,7 @@ class ReWOOAgentGraph(BaseAgent):
|
|
|
144
226
|
|
|
145
227
|
else:
|
|
146
228
|
assert False, f"Unexpected type for tool_input: {type(tool_input)}"
|
|
229
|
+
|
|
147
230
|
return tool_input
|
|
148
231
|
|
|
149
232
|
@staticmethod
|
|
@@ -197,116 +280,206 @@ class ReWOOAgentGraph(BaseAgent):
|
|
|
197
280
|
|
|
198
281
|
steps = self._parse_planner_output(str(plan.content))
|
|
199
282
|
|
|
283
|
+
# Parse dependencies and create execution levels for parallel processing
|
|
284
|
+
evidence_map, execution_levels = self._parse_planner_dependencies(steps)
|
|
285
|
+
|
|
200
286
|
if self.detailed_logs:
|
|
201
287
|
agent_response_log_message = AGENT_CALL_LOG_MESSAGE % (task, str(plan.content))
|
|
202
288
|
logger.info("ReWOO agent planner output: %s", agent_response_log_message)
|
|
289
|
+
logger.info("ReWOO agent execution levels: %s", execution_levels)
|
|
203
290
|
|
|
204
|
-
return {
|
|
291
|
+
return {
|
|
292
|
+
"plan": plan,
|
|
293
|
+
"evidence_map": evidence_map,
|
|
294
|
+
"execution_levels": execution_levels,
|
|
295
|
+
"current_level": 0,
|
|
296
|
+
}
|
|
205
297
|
|
|
206
298
|
except Exception as ex:
|
|
207
299
|
logger.error("%s Failed to call planner_node: %s", AGENT_LOG_PREFIX, ex)
|
|
208
300
|
raise
|
|
209
301
|
|
|
210
302
|
async def executor_node(self, state: ReWOOGraphState):
|
|
303
|
+
"""
|
|
304
|
+
Execute tools in parallel for the current dependency level.
|
|
305
|
+
|
|
306
|
+
This replaces the sequential execution with parallel execution of tools
|
|
307
|
+
that have no dependencies between them.
|
|
308
|
+
"""
|
|
211
309
|
try:
|
|
212
310
|
logger.debug("%s Starting the ReWOO Executor Node", AGENT_LOG_PREFIX)
|
|
213
311
|
|
|
214
|
-
|
|
215
|
-
|
|
216
|
-
if
|
|
217
|
-
|
|
218
|
-
|
|
219
|
-
|
|
220
|
-
|
|
221
|
-
|
|
222
|
-
|
|
223
|
-
|
|
224
|
-
|
|
225
|
-
|
|
226
|
-
|
|
227
|
-
|
|
228
|
-
|
|
229
|
-
|
|
312
|
+
current_level, level_complete = self._get_current_level_status(state)
|
|
313
|
+
|
|
314
|
+
# Should not be invoked if all levels are complete
|
|
315
|
+
if current_level < 0:
|
|
316
|
+
logger.error("%s ReWOO Executor invoked after all levels complete", AGENT_LOG_PREFIX)
|
|
317
|
+
raise RuntimeError("ReWOO Executor invoked after all levels complete")
|
|
318
|
+
|
|
319
|
+
# If current level is already complete, move to next level
|
|
320
|
+
if level_complete:
|
|
321
|
+
new_level = current_level + 1
|
|
322
|
+
logger.debug("%s Level %s complete, moving to level %s", AGENT_LOG_PREFIX, current_level, new_level)
|
|
323
|
+
return {"current_level": new_level}
|
|
324
|
+
|
|
325
|
+
# Get placeholders for current level
|
|
326
|
+
current_level_placeholders = state.execution_levels[current_level]
|
|
327
|
+
|
|
328
|
+
# Filter to only placeholders not yet completed
|
|
329
|
+
pending_placeholders = list(set(current_level_placeholders) - set(state.intermediate_results.keys()))
|
|
330
|
+
|
|
331
|
+
if not pending_placeholders:
|
|
332
|
+
# All placeholders in this level are done, move to next level
|
|
333
|
+
new_level = current_level + 1
|
|
334
|
+
return {"current_level": new_level}
|
|
335
|
+
|
|
336
|
+
logger.debug("%s Executing level %s with %s tools in parallel: %s",
|
|
337
|
+
AGENT_LOG_PREFIX,
|
|
338
|
+
current_level,
|
|
339
|
+
len(pending_placeholders),
|
|
340
|
+
pending_placeholders)
|
|
341
|
+
|
|
342
|
+
# Execute all tools in current level in parallel
|
|
343
|
+
tasks = []
|
|
344
|
+
for placeholder in pending_placeholders:
|
|
345
|
+
step_info = state.evidence_map[placeholder]
|
|
346
|
+
task = self._execute_single_tool(placeholder, step_info, state.intermediate_results)
|
|
347
|
+
tasks.append(task)
|
|
348
|
+
|
|
349
|
+
# Wait for all tasks in current level to complete
|
|
350
|
+
results = await asyncio.gather(*tasks, return_exceptions=True)
|
|
351
|
+
|
|
352
|
+
# Process results and update intermediate_results
|
|
353
|
+
updated_intermediate_results = dict(state.intermediate_results)
|
|
354
|
+
|
|
355
|
+
for placeholder, result in zip(pending_placeholders, results):
|
|
356
|
+
if isinstance(result, BaseException):
|
|
357
|
+
logger.error("%s Tool execution failed for %s: %s", AGENT_LOG_PREFIX, placeholder, result)
|
|
358
|
+
# Create error tool message
|
|
359
|
+
error_message = f"Tool execution failed: {str(result)}"
|
|
360
|
+
updated_intermediate_results[placeholder] = ToolMessage(content=error_message,
|
|
361
|
+
tool_call_id=placeholder)
|
|
362
|
+
if self.raise_tool_call_error:
|
|
363
|
+
raise result
|
|
230
364
|
else:
|
|
231
|
-
|
|
232
|
-
|
|
233
|
-
|
|
234
|
-
|
|
235
|
-
|
|
236
|
-
|
|
237
|
-
intermediate_results = state.intermediate_results
|
|
238
|
-
|
|
239
|
-
# Replace the placeholder in the tool input with the previous tool output
|
|
240
|
-
for _placeholder, _tool_output in intermediate_results.items():
|
|
241
|
-
_tool_output = _tool_output.content
|
|
242
|
-
# If the content is a list, get the first element which should be a dict
|
|
243
|
-
if isinstance(_tool_output, list):
|
|
244
|
-
_tool_output = _tool_output[0]
|
|
245
|
-
assert isinstance(_tool_output, dict)
|
|
246
|
-
|
|
247
|
-
tool_input = self._replace_placeholder(_placeholder, tool_input, _tool_output)
|
|
248
|
-
|
|
249
|
-
requested_tool = self._get_tool(tool)
|
|
250
|
-
if not requested_tool:
|
|
251
|
-
configured_tool_names = list(self.tools_dict.keys())
|
|
252
|
-
logger.warning(
|
|
253
|
-
"%s ReWOO Agent wants to call tool %s. In the ReWOO Agent's configuration within the config file,"
|
|
254
|
-
"there is no tool with that name: %s",
|
|
255
|
-
AGENT_LOG_PREFIX,
|
|
256
|
-
tool,
|
|
257
|
-
configured_tool_names)
|
|
258
|
-
|
|
259
|
-
intermediate_results[placeholder] = ToolMessage(content=TOOL_NOT_FOUND_ERROR_MESSAGE.format(
|
|
260
|
-
tool_name=tool, tools=configured_tool_names),
|
|
261
|
-
tool_call_id=tool)
|
|
262
|
-
return {"intermediate_results": intermediate_results}
|
|
365
|
+
updated_intermediate_results[placeholder] = result
|
|
366
|
+
# Check if the ToolMessage has error status and raise_tool_call_error is True
|
|
367
|
+
if (isinstance(result, ToolMessage) and hasattr(result, 'status') and result.status == "error"
|
|
368
|
+
and self.raise_tool_call_error):
|
|
369
|
+
logger.error("%s Tool call failed for %s: %s", AGENT_LOG_PREFIX, placeholder, result.content)
|
|
370
|
+
raise RuntimeError(f"Tool call failed: {result.content}")
|
|
263
371
|
|
|
264
372
|
if self.detailed_logs:
|
|
265
|
-
logger.
|
|
373
|
+
logger.info("%s Completed level %s with %s tools",
|
|
374
|
+
AGENT_LOG_PREFIX,
|
|
375
|
+
current_level,
|
|
376
|
+
len(pending_placeholders))
|
|
266
377
|
|
|
267
|
-
|
|
268
|
-
tool_input_parsed = self._parse_tool_input(tool_input)
|
|
269
|
-
tool_response = await self._call_tool(requested_tool,
|
|
270
|
-
tool_input_parsed,
|
|
271
|
-
RunnableConfig(callbacks=self.callbacks),
|
|
272
|
-
max_retries=3)
|
|
273
|
-
|
|
274
|
-
if self.detailed_logs:
|
|
275
|
-
self._log_tool_response(requested_tool.name, tool_input_parsed, str(tool_response))
|
|
276
|
-
|
|
277
|
-
intermediate_results[placeholder] = tool_response
|
|
278
|
-
return {"intermediate_results": intermediate_results}
|
|
378
|
+
return {"intermediate_results": updated_intermediate_results}
|
|
279
379
|
|
|
280
380
|
except Exception as ex:
|
|
281
381
|
logger.error("%s Failed to call executor_node: %s", AGENT_LOG_PREFIX, ex)
|
|
282
382
|
raise
|
|
283
383
|
|
|
384
|
+
async def _execute_single_tool(self,
|
|
385
|
+
placeholder: str,
|
|
386
|
+
step_info: ReWOOPlanStep,
|
|
387
|
+
intermediate_results: dict[str, ToolMessage]) -> ToolMessage:
|
|
388
|
+
"""
|
|
389
|
+
Execute a single tool with proper placeholder replacement.
|
|
390
|
+
|
|
391
|
+
Args:
|
|
392
|
+
placeholder: The evidence placeholder (e.g., "#E1").
|
|
393
|
+
step_info: Step information containing tool and tool_input.
|
|
394
|
+
intermediate_results: Current intermediate results for placeholder replacement.
|
|
395
|
+
|
|
396
|
+
Returns:
|
|
397
|
+
ToolMessage with the tool execution result.
|
|
398
|
+
"""
|
|
399
|
+
evidence_info = step_info.evidence
|
|
400
|
+
tool_name = evidence_info.tool
|
|
401
|
+
tool_input = evidence_info.tool_input
|
|
402
|
+
|
|
403
|
+
# Replace placeholders in tool input with previous results
|
|
404
|
+
for ph_key, tool_output in intermediate_results.items():
|
|
405
|
+
tool_output_content = tool_output.content
|
|
406
|
+
# If the content is a list, get the first element which should be a dict
|
|
407
|
+
if isinstance(tool_output_content, list):
|
|
408
|
+
tool_output_content = tool_output_content[0]
|
|
409
|
+
assert isinstance(tool_output_content, dict)
|
|
410
|
+
|
|
411
|
+
tool_input = self._replace_placeholder(ph_key, tool_input, tool_output_content)
|
|
412
|
+
|
|
413
|
+
# Get the requested tool
|
|
414
|
+
requested_tool = self._get_tool(tool_name)
|
|
415
|
+
if not requested_tool:
|
|
416
|
+
configured_tool_names = list(self.tools_dict.keys())
|
|
417
|
+
logger.warning(
|
|
418
|
+
"%s ReWOO Agent wants to call tool %s. In the ReWOO Agent's configuration within the config file,"
|
|
419
|
+
"there is no tool with that name: %s",
|
|
420
|
+
AGENT_LOG_PREFIX,
|
|
421
|
+
tool_name,
|
|
422
|
+
configured_tool_names)
|
|
423
|
+
|
|
424
|
+
return ToolMessage(content=TOOL_NOT_FOUND_ERROR_MESSAGE.format(tool_name=tool_name,
|
|
425
|
+
tools=configured_tool_names),
|
|
426
|
+
tool_call_id=placeholder)
|
|
427
|
+
|
|
428
|
+
if self.detailed_logs:
|
|
429
|
+
logger.debug("%s Calling tool %s with input: %s", AGENT_LOG_PREFIX, requested_tool.name, tool_input)
|
|
430
|
+
|
|
431
|
+
# Parse and execute the tool
|
|
432
|
+
tool_input_parsed = self._parse_tool_input(tool_input)
|
|
433
|
+
tool_response = await self._call_tool(
|
|
434
|
+
requested_tool,
|
|
435
|
+
tool_input_parsed,
|
|
436
|
+
RunnableConfig(callbacks=self.callbacks), # type: ignore
|
|
437
|
+
max_retries=self.tool_call_max_retries)
|
|
438
|
+
|
|
439
|
+
if self.detailed_logs:
|
|
440
|
+
self._log_tool_response(requested_tool.name, tool_input_parsed, str(tool_response))
|
|
441
|
+
|
|
442
|
+
return tool_response
|
|
443
|
+
|
|
284
444
|
async def solver_node(self, state: ReWOOGraphState):
|
|
285
445
|
try:
|
|
286
446
|
logger.debug("%s Starting the ReWOO Solver Node", AGENT_LOG_PREFIX)
|
|
287
447
|
|
|
288
448
|
plan = ""
|
|
289
|
-
# Add the tool outputs of each step to the plan
|
|
290
|
-
for
|
|
291
|
-
|
|
292
|
-
|
|
293
|
-
|
|
294
|
-
|
|
295
|
-
|
|
296
|
-
|
|
297
|
-
|
|
449
|
+
# Add the tool outputs of each step to the plan using evidence_map
|
|
450
|
+
for placeholder, step_info in state.evidence_map.items():
|
|
451
|
+
evidence_info = step_info.evidence
|
|
452
|
+
original_tool_input = evidence_info.tool_input
|
|
453
|
+
tool_name = evidence_info.tool
|
|
454
|
+
|
|
455
|
+
# Replace placeholders in tool input with actual results
|
|
456
|
+
final_tool_input = original_tool_input
|
|
457
|
+
for ph_key, tool_output in state.intermediate_results.items():
|
|
458
|
+
tool_output_content = tool_output.content
|
|
298
459
|
# If the content is a list, get the first element which should be a dict
|
|
299
|
-
if isinstance(
|
|
300
|
-
|
|
301
|
-
assert isinstance(
|
|
302
|
-
|
|
303
|
-
|
|
304
|
-
|
|
305
|
-
|
|
306
|
-
|
|
307
|
-
|
|
308
|
-
|
|
309
|
-
|
|
460
|
+
if isinstance(tool_output_content, list):
|
|
461
|
+
tool_output_content = tool_output_content[0]
|
|
462
|
+
assert isinstance(tool_output_content, dict)
|
|
463
|
+
|
|
464
|
+
final_tool_input = self._replace_placeholder(ph_key, final_tool_input, tool_output_content)
|
|
465
|
+
|
|
466
|
+
# Get the final result for this placeholder
|
|
467
|
+
final_result = ""
|
|
468
|
+
if placeholder in state.intermediate_results:
|
|
469
|
+
result_content = state.intermediate_results[placeholder].content
|
|
470
|
+
if isinstance(result_content, list):
|
|
471
|
+
result_content = result_content[0]
|
|
472
|
+
if isinstance(result_content, dict):
|
|
473
|
+
final_result = str(result_content)
|
|
474
|
+
else:
|
|
475
|
+
final_result = str(result_content)
|
|
476
|
+
|
|
477
|
+
step_plan = step_info.plan
|
|
478
|
+
plan += '\n'.join([
|
|
479
|
+
f"Plan: {step_plan}",
|
|
480
|
+
f"{placeholder} = {tool_name}[{final_tool_input}",
|
|
481
|
+
f"Result: {final_result}\n\n"
|
|
482
|
+
])
|
|
310
483
|
|
|
311
484
|
task = str(state.task.content)
|
|
312
485
|
solver_prompt = self.solver_prompt.partial(plan=plan)
|
|
@@ -329,12 +502,24 @@ class ReWOOAgentGraph(BaseAgent):
|
|
|
329
502
|
try:
|
|
330
503
|
logger.debug("%s Starting the ReWOO Conditional Edge", AGENT_LOG_PREFIX)
|
|
331
504
|
|
|
332
|
-
|
|
333
|
-
|
|
334
|
-
|
|
505
|
+
current_level, level_complete = self._get_current_level_status(state)
|
|
506
|
+
|
|
507
|
+
# If all levels are complete, move to solver
|
|
508
|
+
if current_level == -1:
|
|
509
|
+
logger.debug("%s All execution levels complete, moving to solver", AGENT_LOG_PREFIX)
|
|
335
510
|
return AgentDecision.END
|
|
336
511
|
|
|
337
|
-
|
|
512
|
+
# If current level is complete, check if there are more levels
|
|
513
|
+
if level_complete:
|
|
514
|
+
next_level = current_level + 1
|
|
515
|
+
if next_level >= len(state.execution_levels):
|
|
516
|
+
logger.debug("%s All execution levels complete, moving to solver", AGENT_LOG_PREFIX)
|
|
517
|
+
return AgentDecision.END
|
|
518
|
+
|
|
519
|
+
logger.debug("%s Continuing with executor (level %s, complete: %s)",
|
|
520
|
+
AGENT_LOG_PREFIX,
|
|
521
|
+
current_level,
|
|
522
|
+
level_complete)
|
|
338
523
|
return AgentDecision.TOOL
|
|
339
524
|
|
|
340
525
|
except Exception as ex:
|
|
@@ -352,8 +537,10 @@ class ReWOOAgentGraph(BaseAgent):
|
|
|
352
537
|
graph.add_node("solver", self.solver_node)
|
|
353
538
|
|
|
354
539
|
graph.add_edge("planner", "executor")
|
|
355
|
-
|
|
356
|
-
|
|
540
|
+
graph.add_conditional_edges("executor",
|
|
541
|
+
self.conditional_edge, {
|
|
542
|
+
AgentDecision.TOOL: "executor", AgentDecision.END: "solver"
|
|
543
|
+
})
|
|
357
544
|
|
|
358
545
|
graph.set_entry_point("planner")
|
|
359
546
|
graph.set_finish_point("solver")
|
nat/agent/rewoo_agent/prompt.py
CHANGED
|
@@ -18,33 +18,29 @@ For the following task, make plans that can solve the problem step by step. For
|
|
|
18
18
|
which external tool together with tool input to retrieve evidence. You can store the evidence into a \
|
|
19
19
|
placeholder #E that can be called by later tools. (Plan, #E1, Plan, #E2, Plan, ...)
|
|
20
20
|
|
|
21
|
-
|
|
21
|
+
The following tools and respective requirements are available to you:
|
|
22
22
|
|
|
23
23
|
{tools}
|
|
24
24
|
|
|
25
|
-
The
|
|
25
|
+
The tool calls you make should be one of the following: [{tool_names}]
|
|
26
26
|
|
|
27
27
|
You are not required to use all the tools listed. Choose only the ones that best fit the needs of each plan step.
|
|
28
28
|
|
|
29
|
-
Your output must be a JSON array where each element represents one planning step. Each step must be an object with
|
|
30
|
-
|
|
29
|
+
Your output must be a JSON array where each element represents one planning step. Each step must be an object with \
|
|
31
30
|
exactly two keys:
|
|
32
31
|
|
|
33
32
|
1. "plan": A string that describes in detail the action or reasoning for that step.
|
|
34
33
|
|
|
35
|
-
2. "evidence": An object representing the external tool call associated with that plan step. This object must have the
|
|
34
|
+
2. "evidence": An object representing the external tool call associated with that plan step. This object must have the \
|
|
36
35
|
following keys:
|
|
37
36
|
|
|
38
|
-
-"placeholder": A string that identifies the evidence placeholder (
|
|
39
|
-
|
|
37
|
+
-"placeholder": A string that identifies the evidence placeholder ("#E1", "#E2", ...). The numbering should \
|
|
38
|
+
be sequential based on the order of steps.
|
|
40
39
|
|
|
41
40
|
-"tool": A string specifying the name of the external tool used.
|
|
42
41
|
|
|
43
|
-
-"tool_input": The input to the tool. This can be a string, array, or object, depending on the requirements of the
|
|
44
|
-
|
|
45
|
-
|
|
46
|
-
Do not include any additional keys or characters in your output, and do not wrap your response with markdown formatting.
|
|
47
|
-
Your output must be strictly valid JSON.
|
|
42
|
+
-"tool_input": The input to the tool. This can be a string, array, or object, depending on the requirements of the \
|
|
43
|
+
tool. Be careful about type assumptions because the output of former tools might contain noise.
|
|
48
44
|
|
|
49
45
|
Important instructions:
|
|
50
46
|
|
|
@@ -58,27 +54,28 @@ Here is an example of how a valid JSON output should look:
|
|
|
58
54
|
|
|
59
55
|
[
|
|
60
56
|
\'{{
|
|
61
|
-
"plan": "
|
|
57
|
+
"plan": "Find Alex's schedule on Sep 25, 2025",
|
|
62
58
|
"evidence": \'{{
|
|
63
59
|
"placeholder": "#E1",
|
|
64
|
-
"tool": "
|
|
65
|
-
"tool_input":
|
|
60
|
+
"tool": "search_calendar",
|
|
61
|
+
"tool_input": ("Alex", "09/25/2025")
|
|
66
62
|
}}\'
|
|
67
63
|
}}\',
|
|
68
64
|
\'{{
|
|
69
|
-
"plan": "
|
|
65
|
+
"plan": "Find Bill's schedule on sep 25, 2025",
|
|
70
66
|
"evidence": \'{{
|
|
71
67
|
"placeholder": "#E2",
|
|
72
|
-
"tool": "
|
|
73
|
-
"tool_input": "
|
|
68
|
+
"tool": "search_calendar",
|
|
69
|
+
"tool_input": ("Bill", "09/25/2025")
|
|
74
70
|
}}\'
|
|
75
71
|
}}\',
|
|
76
72
|
\'{{
|
|
77
|
-
"plan": "
|
|
73
|
+
"plan": "Suggest a time for 1-hour meeting given Alex's and Bill's schedule.",
|
|
78
74
|
"evidence": \'{{
|
|
79
75
|
"placeholder": "#E3",
|
|
80
|
-
"tool": "
|
|
81
|
-
"tool_input": "
|
|
76
|
+
"tool": "llm_chat",
|
|
77
|
+
"tool_input": "Find a common 1-hour time slot for Alex and Bill given their schedules. \
|
|
78
|
+
Alex's schedule: #E1; Bill's schedule: #E2?"
|
|
82
79
|
}}\'
|
|
83
80
|
}}\'
|
|
84
81
|
]
|
|
@@ -94,7 +91,7 @@ task: {task}
|
|
|
94
91
|
"""
|
|
95
92
|
|
|
96
93
|
SOLVER_SYSTEM_PROMPT = """
|
|
97
|
-
Solve the following task or problem. To solve the problem, we have made
|
|
94
|
+
Solve the following task or problem. To solve the problem, we have made some Plans ahead and \
|
|
98
95
|
retrieved corresponding Evidence to each Plan. Use them with caution since long evidence might \
|
|
99
96
|
contain irrelevant information.
|
|
100
97
|
|