nvidia-nat 1.2.1__py3-none-any.whl → 1.3.0__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.
- aiq/__init__.py +2 -2
- nat/agent/base.py +27 -18
- nat/agent/dual_node.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 +81 -50
- nat/agent/react_agent/register.py +59 -40
- nat/agent/reasoning_agent/reasoning_agent.py +17 -15
- nat/agent/register.py +1 -1
- nat/agent/rewoo_agent/agent.py +327 -149
- nat/agent/rewoo_agent/prompt.py +19 -22
- nat/agent/rewoo_agent/register.py +64 -46
- nat/agent/tool_calling_agent/agent.py +152 -29
- nat/agent/tool_calling_agent/register.py +61 -38
- nat/authentication/api_key/api_key_auth_provider.py +2 -2
- 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_resource_server_config.py +124 -0
- nat/authentication/register.py +0 -1
- nat/builder/builder.py +56 -24
- nat/builder/component_utils.py +10 -6
- nat/builder/context.py +70 -18
- nat/builder/eval_builder.py +16 -11
- nat/builder/framework_enum.py +1 -0
- nat/builder/front_end.py +1 -1
- nat/builder/function.py +378 -8
- nat/builder/function_base.py +3 -3
- nat/builder/function_info.py +6 -8
- nat/builder/intermediate_step_manager.py +6 -2
- nat/builder/user_interaction_manager.py +2 -2
- nat/builder/workflow.py +13 -1
- nat/builder/workflow_builder.py +327 -79
- nat/cli/cli_utils/config_override.py +2 -2
- nat/cli/commands/evaluate.py +1 -1
- nat/cli/commands/info/info.py +16 -6
- nat/cli/commands/info/list_channels.py +1 -1
- nat/cli/commands/info/list_components.py +7 -8
- nat/cli/commands/mcp/__init__.py +14 -0
- nat/cli/commands/mcp/mcp.py +986 -0
- nat/cli/commands/object_store/__init__.py +14 -0
- nat/cli/commands/object_store/object_store.py +227 -0
- nat/cli/commands/optimize.py +90 -0
- nat/cli/commands/registry/publish.py +2 -2
- nat/cli/commands/registry/pull.py +2 -2
- nat/cli/commands/registry/remove.py +2 -2
- nat/cli/commands/registry/search.py +15 -17
- nat/cli/commands/start.py +16 -5
- nat/cli/commands/uninstall.py +1 -1
- nat/cli/commands/workflow/templates/config.yml.j2 +14 -13
- nat/cli/commands/workflow/templates/pyproject.toml.j2 +5 -2
- nat/cli/commands/workflow/templates/register.py.j2 +2 -3
- nat/cli/commands/workflow/templates/workflow.py.j2 +35 -21
- nat/cli/commands/workflow/workflow_commands.py +105 -19
- nat/cli/entrypoint.py +17 -11
- nat/cli/main.py +3 -0
- nat/cli/register_workflow.py +38 -4
- nat/cli/type_registry.py +79 -10
- 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 +196 -67
- nat/data_models/authentication.py +23 -9
- nat/data_models/common.py +1 -1
- nat/data_models/component.py +2 -0
- nat/data_models/component_ref.py +11 -0
- nat/data_models/config.py +42 -18
- nat/data_models/dataset_handler.py +1 -1
- nat/data_models/discovery_metadata.py +4 -4
- nat/data_models/evaluate.py +4 -1
- nat/data_models/function.py +34 -0
- nat/data_models/function_dependencies.py +14 -6
- nat/data_models/gated_field_mixin.py +242 -0
- nat/data_models/intermediate_step.py +3 -3
- nat/data_models/optimizable.py +119 -0
- nat/data_models/optimizer.py +149 -0
- nat/data_models/span.py +41 -3
- nat/data_models/swe_bench_model.py +1 -1
- nat/data_models/temperature_mixin.py +44 -0
- nat/data_models/thinking_mixin.py +86 -0
- nat/data_models/top_p_mixin.py +44 -0
- nat/embedder/azure_openai_embedder.py +46 -0
- nat/embedder/nim_embedder.py +1 -1
- nat/embedder/openai_embedder.py +2 -3
- nat/embedder/register.py +1 -1
- nat/eval/config.py +3 -1
- nat/eval/dataset_handler/dataset_handler.py +71 -7
- nat/eval/evaluate.py +86 -31
- nat/eval/evaluator/base_evaluator.py +1 -1
- nat/eval/evaluator/evaluator_model.py +13 -0
- nat/eval/intermediate_step_adapter.py +1 -1
- nat/eval/rag_evaluator/evaluate.py +9 -6
- nat/eval/rag_evaluator/register.py +3 -3
- nat/eval/register.py +4 -1
- nat/eval/remote_workflow.py +3 -3
- 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 +6 -6
- nat/eval/trajectory_evaluator/evaluate.py +1 -1
- nat/eval/trajectory_evaluator/register.py +1 -1
- nat/eval/tunable_rag_evaluator/evaluate.py +4 -7
- nat/eval/utils/eval_trace_ctx.py +89 -0
- nat/eval/utils/weave_eval.py +18 -9
- 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 +3 -3
- nat/experimental/test_time_compute/functions/ttc_tool_wrapper_function.py +3 -3
- nat/experimental/test_time_compute/models/strategy_base.py +5 -4
- nat/experimental/test_time_compute/register.py +0 -1
- nat/experimental/test_time_compute/selection/llm_based_output_merging_selector.py +1 -3
- 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 +36 -5
- nat/front_ends/fastapi/fastapi_front_end_controller.py +4 -4
- nat/front_ends/fastapi/fastapi_front_end_plugin.py +135 -4
- nat/front_ends/fastapi/fastapi_front_end_plugin_worker.py +455 -282
- nat/front_ends/fastapi/job_store.py +518 -99
- nat/front_ends/fastapi/main.py +11 -19
- nat/front_ends/fastapi/message_handler.py +74 -50
- nat/front_ends/fastapi/message_validator.py +20 -21
- nat/front_ends/fastapi/response_helpers.py +4 -4
- nat/front_ends/fastapi/step_adaptor.py +2 -2
- 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 +47 -3
- nat/front_ends/mcp/mcp_front_end_plugin.py +48 -13
- nat/front_ends/mcp/mcp_front_end_plugin_worker.py +120 -8
- nat/front_ends/mcp/tool_converter.py +44 -14
- nat/front_ends/register.py +0 -1
- nat/front_ends/simple_base/simple_front_end_plugin_base.py +3 -1
- nat/llm/aws_bedrock_llm.py +24 -12
- nat/llm/azure_openai_llm.py +57 -0
- nat/llm/litellm_llm.py +69 -0
- nat/llm/nim_llm.py +20 -8
- nat/llm/openai_llm.py +14 -6
- nat/llm/register.py +5 -1
- nat/llm/utils/env_config_value.py +2 -3
- nat/llm/utils/thinking.py +215 -0
- nat/meta/pypi.md +9 -9
- nat/object_store/register.py +0 -1
- nat/observability/exporter/base_exporter.py +3 -3
- nat/observability/exporter/file_exporter.py +1 -1
- nat/observability/exporter/processing_exporter.py +309 -81
- nat/observability/exporter/span_exporter.py +35 -15
- nat/observability/exporter_manager.py +7 -7
- nat/observability/mixin/file_mixin.py +7 -7
- nat/observability/mixin/redaction_config_mixin.py +42 -0
- nat/observability/mixin/tagging_config_mixin.py +62 -0
- nat/observability/mixin/type_introspection_mixin.py +420 -107
- nat/observability/processor/batching_processor.py +5 -7
- nat/observability/processor/falsy_batch_filter_processor.py +55 -0
- nat/observability/processor/processor.py +3 -0
- nat/observability/processor/processor_factory.py +70 -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 +68 -0
- nat/observability/register.py +22 -4
- nat/profiler/calc/calc_runner.py +3 -4
- nat/profiler/callbacks/agno_callback_handler.py +1 -1
- nat/profiler/callbacks/langchain_callback_handler.py +14 -7
- nat/profiler/callbacks/llama_index_callback_handler.py +3 -3
- nat/profiler/callbacks/semantic_kernel_callback_handler.py +3 -3
- nat/profiler/data_frame_row.py +1 -1
- nat/profiler/decorators/framework_wrapper.py +62 -13
- nat/profiler/decorators/function_tracking.py +160 -3
- nat/profiler/forecasting/models/forecasting_base_model.py +3 -1
- 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/bottleneck_analysis/simple_stack_analysis.py +1 -1
- nat/profiler/inference_optimization/data_models.py +3 -3
- nat/profiler/inference_optimization/experimental/prefix_span_analysis.py +8 -9
- nat/profiler/inference_optimization/token_uniqueness.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 +164 -0
- nat/profiler/parameter_optimization/parameter_selection.py +107 -0
- nat/profiler/parameter_optimization/pareto_visualizer.py +395 -0
- nat/profiler/parameter_optimization/prompt_optimizer.py +384 -0
- nat/profiler/parameter_optimization/update_helpers.py +66 -0
- nat/profiler/profile_runner.py +14 -9
- nat/profiler/utils.py +4 -2
- nat/registry_handlers/local/local_handler.py +2 -2
- nat/registry_handlers/package_utils.py +1 -2
- nat/registry_handlers/pypi/pypi_handler.py +23 -26
- nat/registry_handlers/register.py +3 -4
- nat/registry_handlers/rest/rest_handler.py +12 -13
- nat/retriever/milvus/retriever.py +2 -2
- nat/retriever/nemo_retriever/retriever.py +1 -1
- nat/retriever/register.py +0 -1
- nat/runtime/loader.py +2 -2
- nat/runtime/runner.py +105 -8
- nat/runtime/session.py +69 -8
- nat/settings/global_settings.py +16 -5
- nat/tool/chat_completion.py +5 -2
- nat/tool/code_execution/local_sandbox/local_sandbox_server.py +3 -3
- nat/tool/datetime_tools.py +49 -9
- nat/tool/document_search.py +2 -2
- 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/nvidia_rag.py +1 -1
- nat/tool/register.py +2 -9
- nat/tool/retriever.py +3 -2
- nat/utils/callable_utils.py +70 -0
- nat/utils/data_models/schema_validator.py +3 -3
- nat/utils/decorators.py +210 -0
- nat/utils/exception_handlers/automatic_retries.py +104 -51
- nat/utils/exception_handlers/schemas.py +1 -1
- nat/utils/io/yaml_tools.py +2 -2
- nat/utils/log_levels.py +25 -0
- nat/utils/reactive/base/observable_base.py +2 -2
- nat/utils/reactive/base/observer_base.py +1 -1
- nat/utils/reactive/observable.py +2 -2
- nat/utils/reactive/observer.py +4 -4
- nat/utils/reactive/subscription.py +1 -1
- nat/utils/settings/global_settings.py +6 -8
- nat/utils/type_converter.py +12 -3
- nat/utils/type_utils.py +9 -5
- nvidia_nat-1.3.0.dist-info/METADATA +195 -0
- {nvidia_nat-1.2.1.dist-info → nvidia_nat-1.3.0.dist-info}/RECORD +244 -200
- {nvidia_nat-1.2.1.dist-info → nvidia_nat-1.3.0.dist-info}/entry_points.txt +1 -0
- nat/cli/commands/info/list_mcp.py +0 -304
- 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/mcp/exceptions.py +0 -142
- nat/tool/mcp/mcp_client.py +0 -255
- nat/tool/mcp/mcp_tool.py +0 -96
- nat/utils/exception_handlers/mcp.py +0 -211
- nvidia_nat-1.2.1.dist-info/METADATA +0 -365
- /nat/{tool/github_tools → agent/prompt_optimizer}/__init__.py +0 -0
- /nat/{tool/mcp → authentication/credential_validator}/__init__.py +0 -0
- {nvidia_nat-1.2.1.dist-info → nvidia_nat-1.3.0.dist-info}/WHEEL +0 -0
- {nvidia_nat-1.2.1.dist-info → nvidia_nat-1.3.0.dist-info}/licenses/LICENSE-3rd-party.txt +0 -0
- {nvidia_nat-1.2.1.dist-info → nvidia_nat-1.3.0.dist-info}/licenses/LICENSE.md +0 -0
- {nvidia_nat-1.2.1.dist-info → nvidia_nat-1.3.0.dist-info}/top_level.txt +0 -0
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
|
|
|
@@ -17,44 +17,44 @@ import logging
|
|
|
17
17
|
|
|
18
18
|
from pydantic import AliasChoices
|
|
19
19
|
from pydantic import Field
|
|
20
|
+
from pydantic import PositiveInt
|
|
20
21
|
|
|
21
22
|
from nat.builder.builder import Builder
|
|
22
23
|
from nat.builder.framework_enum import LLMFrameworkEnum
|
|
23
24
|
from nat.builder.function_info import FunctionInfo
|
|
24
25
|
from nat.cli.register_workflow import register_function
|
|
26
|
+
from nat.data_models.agent import AgentBaseConfig
|
|
25
27
|
from nat.data_models.api_server import ChatRequest
|
|
28
|
+
from nat.data_models.api_server import ChatRequestOrMessage
|
|
26
29
|
from nat.data_models.api_server import ChatResponse
|
|
30
|
+
from nat.data_models.api_server import Usage
|
|
31
|
+
from nat.data_models.component_ref import FunctionGroupRef
|
|
27
32
|
from nat.data_models.component_ref import FunctionRef
|
|
28
|
-
from nat.data_models.component_ref import LLMRef
|
|
29
|
-
from nat.data_models.function import FunctionBaseConfig
|
|
30
33
|
from nat.utils.type_converter import GlobalTypeConverter
|
|
31
34
|
|
|
32
35
|
logger = logging.getLogger(__name__)
|
|
33
36
|
|
|
34
37
|
|
|
35
|
-
class ReWOOAgentWorkflowConfig(
|
|
38
|
+
class ReWOOAgentWorkflowConfig(AgentBaseConfig, name="rewoo_agent"):
|
|
36
39
|
"""
|
|
37
40
|
Defines a NAT function that uses a ReWOO Agent performs reasoning inbetween tool calls, and utilizes the
|
|
38
41
|
tool names and descriptions to select the optimal tool.
|
|
39
42
|
"""
|
|
40
|
-
|
|
41
|
-
tool_names: list[FunctionRef] = Field(
|
|
42
|
-
|
|
43
|
-
llm_name: LLMRef = Field(description="The LLM model to use with the rewoo agent.")
|
|
44
|
-
verbose: bool = Field(default=False, description="Set the verbosity of the rewoo agent's logging.")
|
|
43
|
+
description: str = Field(default="ReWOO Agent Workflow", description="The description of this functions use.")
|
|
44
|
+
tool_names: list[FunctionRef | FunctionGroupRef] = Field(
|
|
45
|
+
default_factory=list, description="The list of tools to provide to the rewoo agent.")
|
|
45
46
|
include_tool_input_schema_in_tool_description: bool = Field(
|
|
46
47
|
default=True, description="Specify inclusion of tool input schemas in the prompt.")
|
|
47
|
-
description: str = Field(default="ReWOO Agent Workflow", description="The description of this functions use.")
|
|
48
48
|
planner_prompt: str | None = Field(
|
|
49
49
|
default=None,
|
|
50
50
|
description="Provides the PLANNER_PROMPT to use with the agent") # defaults to PLANNER_PROMPT in prompt.py
|
|
51
51
|
solver_prompt: str | None = Field(
|
|
52
52
|
default=None,
|
|
53
53
|
description="Provides the SOLVER_PROMPT to use with the agent") # defaults to SOLVER_PROMPT in prompt.py
|
|
54
|
+
tool_call_max_retries: PositiveInt = Field(default=3,
|
|
55
|
+
description="The number of retries before raising a tool call error.",
|
|
56
|
+
ge=1)
|
|
54
57
|
max_history: int = Field(default=15, description="Maximum number of messages to keep in the conversation history.")
|
|
55
|
-
use_openai_api: bool = Field(default=False,
|
|
56
|
-
description=("Use OpenAI API for the input/output types to the function. "
|
|
57
|
-
"If False, strings will be used."))
|
|
58
58
|
additional_planner_instructions: str | None = Field(
|
|
59
59
|
default=None,
|
|
60
60
|
validation_alias=AliasChoices("additional_planner_instructions", "additional_instructions"),
|
|
@@ -62,15 +62,19 @@ class ReWOOAgentWorkflowConfig(FunctionBaseConfig, name="rewoo_agent"):
|
|
|
62
62
|
additional_solver_instructions: str | None = Field(
|
|
63
63
|
default=None,
|
|
64
64
|
description="Additional instructions to provide to the agent in addition to the base solver prompt.")
|
|
65
|
+
raise_tool_call_error: bool = Field(default=True,
|
|
66
|
+
description="Whether to raise a exception immediately if a tool"
|
|
67
|
+
"call fails. If set to False, the tool call error message will be included in"
|
|
68
|
+
"the tool response and passed to the next tool.")
|
|
65
69
|
|
|
66
70
|
|
|
67
71
|
@register_function(config_type=ReWOOAgentWorkflowConfig, framework_wrappers=[LLMFrameworkEnum.LANGCHAIN])
|
|
68
72
|
async def rewoo_agent_workflow(config: ReWOOAgentWorkflowConfig, builder: Builder):
|
|
69
|
-
from langchain.schema import BaseMessage
|
|
70
73
|
from langchain_core.messages import trim_messages
|
|
74
|
+
from langchain_core.messages.base import BaseMessage
|
|
71
75
|
from langchain_core.messages.human import HumanMessage
|
|
72
76
|
from langchain_core.prompts import ChatPromptTemplate
|
|
73
|
-
from langgraph.graph.
|
|
77
|
+
from langgraph.graph.state import CompiledStateGraph
|
|
74
78
|
|
|
75
79
|
from nat.agent.rewoo_agent.prompt import PLANNER_SYSTEM_PROMPT
|
|
76
80
|
from nat.agent.rewoo_agent.prompt import PLANNER_USER_PROMPT
|
|
@@ -86,7 +90,7 @@ async def rewoo_agent_workflow(config: ReWOOAgentWorkflowConfig, builder: Builde
|
|
|
86
90
|
if config.additional_planner_instructions:
|
|
87
91
|
planner_system_prompt += f"{config.additional_planner_instructions}"
|
|
88
92
|
if not ReWOOAgentGraph.validate_planner_prompt(planner_system_prompt):
|
|
89
|
-
logger.
|
|
93
|
+
logger.error("Invalid planner prompt")
|
|
90
94
|
raise ValueError("Invalid planner prompt")
|
|
91
95
|
planner_prompt = ChatPromptTemplate([("system", planner_system_prompt), ("user", PLANNER_USER_PROMPT)])
|
|
92
96
|
|
|
@@ -94,7 +98,7 @@ async def rewoo_agent_workflow(config: ReWOOAgentWorkflowConfig, builder: Builde
|
|
|
94
98
|
if config.additional_solver_instructions:
|
|
95
99
|
solver_system_prompt += f"{config.additional_solver_instructions}"
|
|
96
100
|
if not ReWOOAgentGraph.validate_solver_prompt(solver_system_prompt):
|
|
97
|
-
logger.
|
|
101
|
+
logger.error("Invalid solver prompt")
|
|
98
102
|
raise ValueError("Invalid solver prompt")
|
|
99
103
|
solver_prompt = ChatPromptTemplate([("system", solver_system_prompt), ("user", SOLVER_USER_PROMPT)])
|
|
100
104
|
|
|
@@ -103,22 +107,39 @@ async def rewoo_agent_workflow(config: ReWOOAgentWorkflowConfig, builder: Builde
|
|
|
103
107
|
|
|
104
108
|
# the agent can run any installed tool, simply install the tool and add it to the config file
|
|
105
109
|
# the sample tool provided can easily be copied or changed
|
|
106
|
-
tools = builder.get_tools(tool_names=config.tool_names, wrapper_type=LLMFrameworkEnum.LANGCHAIN)
|
|
110
|
+
tools = await builder.get_tools(tool_names=config.tool_names, wrapper_type=LLMFrameworkEnum.LANGCHAIN)
|
|
107
111
|
if not tools:
|
|
108
112
|
raise ValueError(f"No tools specified for ReWOO Agent '{config.llm_name}'")
|
|
109
113
|
|
|
110
114
|
# construct the ReWOO Agent Graph from the configured llm, prompt, and tools
|
|
111
|
-
graph:
|
|
112
|
-
|
|
113
|
-
|
|
114
|
-
|
|
115
|
-
|
|
116
|
-
|
|
117
|
-
|
|
118
|
-
|
|
115
|
+
graph: CompiledStateGraph = await ReWOOAgentGraph(
|
|
116
|
+
llm=llm,
|
|
117
|
+
planner_prompt=planner_prompt,
|
|
118
|
+
solver_prompt=solver_prompt,
|
|
119
|
+
tools=tools,
|
|
120
|
+
use_tool_schema=config.include_tool_input_schema_in_tool_description,
|
|
121
|
+
detailed_logs=config.verbose,
|
|
122
|
+
log_response_max_chars=config.log_response_max_chars,
|
|
123
|
+
tool_call_max_retries=config.tool_call_max_retries,
|
|
124
|
+
raise_tool_call_error=config.raise_tool_call_error).build_graph()
|
|
125
|
+
|
|
126
|
+
async def _response_fn(chat_request_or_message: ChatRequestOrMessage) -> ChatResponse | str:
|
|
127
|
+
"""
|
|
128
|
+
Main workflow entry function for the ReWOO Agent.
|
|
129
|
+
|
|
130
|
+
This function invokes the ReWOO Agent Graph and returns the response.
|
|
131
|
+
|
|
132
|
+
Args:
|
|
133
|
+
chat_request_or_message (ChatRequestOrMessage): The input message to process
|
|
134
|
+
|
|
135
|
+
Returns:
|
|
136
|
+
ChatResponse | str: The response from the agent or error message
|
|
137
|
+
"""
|
|
119
138
|
try:
|
|
139
|
+
message = GlobalTypeConverter.get().convert(chat_request_or_message, to_type=ChatRequest)
|
|
140
|
+
|
|
120
141
|
# initialize the starting state with the user query
|
|
121
|
-
messages: list[BaseMessage] = trim_messages(messages=[m.model_dump() for m in
|
|
142
|
+
messages: list[BaseMessage] = trim_messages(messages=[m.model_dump() for m in message.messages],
|
|
122
143
|
max_tokens=config.max_history,
|
|
123
144
|
strategy="last",
|
|
124
145
|
token_counter=len,
|
|
@@ -133,25 +154,22 @@ async def rewoo_agent_workflow(config: ReWOOAgentWorkflowConfig, builder: Builde
|
|
|
133
154
|
|
|
134
155
|
# get and return the output from the state
|
|
135
156
|
state = ReWOOGraphState(**state)
|
|
136
|
-
output_message = state.result.content
|
|
137
|
-
|
|
138
|
-
|
|
157
|
+
output_message = state.result.content
|
|
158
|
+
# Ensure output_message is a string
|
|
159
|
+
if isinstance(output_message, list | dict):
|
|
160
|
+
output_message = str(output_message)
|
|
161
|
+
|
|
162
|
+
# Create usage statistics for the response
|
|
163
|
+
prompt_tokens = sum(len(str(msg.content).split()) for msg in message.messages)
|
|
164
|
+
completion_tokens = len(output_message.split()) if output_message else 0
|
|
165
|
+
total_tokens = prompt_tokens + completion_tokens
|
|
166
|
+
usage = Usage(prompt_tokens=prompt_tokens, completion_tokens=completion_tokens, total_tokens=total_tokens)
|
|
167
|
+
response = ChatResponse.from_string(output_message, usage=usage)
|
|
168
|
+
if chat_request_or_message.is_string:
|
|
169
|
+
return GlobalTypeConverter.get().convert(response, to_type=str)
|
|
170
|
+
return response
|
|
139
171
|
except Exception as ex:
|
|
140
|
-
logger.
|
|
141
|
-
|
|
142
|
-
if config.verbose:
|
|
143
|
-
return ChatResponse.from_string(str(ex))
|
|
144
|
-
return ChatResponse.from_string("I seem to be having a problem.")
|
|
145
|
-
|
|
146
|
-
if (config.use_openai_api):
|
|
147
|
-
yield FunctionInfo.from_fn(_response_fn, description=config.description)
|
|
148
|
-
|
|
149
|
-
else:
|
|
150
|
-
|
|
151
|
-
async def _str_api_fn(input_message: str) -> str:
|
|
152
|
-
oai_input = GlobalTypeConverter.get().try_convert(input_message, to_type=ChatRequest)
|
|
153
|
-
oai_output = await _response_fn(oai_input)
|
|
154
|
-
|
|
155
|
-
return GlobalTypeConverter.get().try_convert(oai_output, to_type=str)
|
|
172
|
+
logger.error("ReWOO Agent failed with exception: %s", ex)
|
|
173
|
+
raise
|
|
156
174
|
|
|
157
|
-
|
|
175
|
+
yield FunctionInfo.from_fn(_response_fn, description=config.description)
|
|
@@ -13,14 +13,19 @@
|
|
|
13
13
|
# See the License for the specific language governing permissions and
|
|
14
14
|
# limitations under the License.
|
|
15
15
|
|
|
16
|
-
# pylint: disable=R0917
|
|
17
16
|
import logging
|
|
17
|
+
import typing
|
|
18
18
|
|
|
19
19
|
from langchain_core.callbacks.base import AsyncCallbackHandler
|
|
20
20
|
from langchain_core.language_models import BaseChatModel
|
|
21
|
+
from langchain_core.messages import SystemMessage
|
|
22
|
+
from langchain_core.messages import ToolMessage
|
|
21
23
|
from langchain_core.messages.base import BaseMessage
|
|
22
|
-
from langchain_core.runnables import
|
|
24
|
+
from langchain_core.runnables import RunnableLambda
|
|
25
|
+
from langchain_core.runnables.config import RunnableConfig
|
|
23
26
|
from langchain_core.tools import BaseTool
|
|
27
|
+
from langgraph.graph import StateGraph
|
|
28
|
+
from langgraph.graph.state import CompiledStateGraph
|
|
24
29
|
from langgraph.prebuilt import ToolNode
|
|
25
30
|
from pydantic import BaseModel
|
|
26
31
|
from pydantic import Field
|
|
@@ -30,6 +35,9 @@ from nat.agent.base import AGENT_LOG_PREFIX
|
|
|
30
35
|
from nat.agent.base import AgentDecision
|
|
31
36
|
from nat.agent.dual_node import DualNodeAgent
|
|
32
37
|
|
|
38
|
+
if typing.TYPE_CHECKING:
|
|
39
|
+
from nat.agent.tool_calling_agent.register import ToolCallAgentWorkflowConfig
|
|
40
|
+
|
|
33
41
|
logger = logging.getLogger(__name__)
|
|
34
42
|
|
|
35
43
|
|
|
@@ -43,22 +51,58 @@ class ToolCallAgentGraph(DualNodeAgent):
|
|
|
43
51
|
A tool Calling Agent utilizes the tool input parameters to select the optimal tool. Supports handling tool errors.
|
|
44
52
|
Argument "detailed_logs" toggles logging of inputs, outputs, and intermediate steps."""
|
|
45
53
|
|
|
46
|
-
def __init__(
|
|
47
|
-
|
|
48
|
-
|
|
49
|
-
|
|
50
|
-
|
|
51
|
-
|
|
52
|
-
|
|
54
|
+
def __init__(
|
|
55
|
+
self,
|
|
56
|
+
llm: BaseChatModel,
|
|
57
|
+
tools: list[BaseTool],
|
|
58
|
+
prompt: str | None = None,
|
|
59
|
+
callbacks: list[AsyncCallbackHandler] = None,
|
|
60
|
+
detailed_logs: bool = False,
|
|
61
|
+
log_response_max_chars: int = 1000,
|
|
62
|
+
handle_tool_errors: bool = True,
|
|
63
|
+
return_direct: list[BaseTool] | None = None,
|
|
64
|
+
):
|
|
65
|
+
super().__init__(llm=llm,
|
|
66
|
+
tools=tools,
|
|
67
|
+
callbacks=callbacks,
|
|
68
|
+
detailed_logs=detailed_logs,
|
|
69
|
+
log_response_max_chars=log_response_max_chars)
|
|
70
|
+
|
|
71
|
+
# some LLMs support tool calling
|
|
72
|
+
# these models accept the tool's input schema and decide when to use a tool based on the input's relevance
|
|
73
|
+
try:
|
|
74
|
+
# in tool calling agents, we bind the tools to the LLM, to pass the tools' input schemas at runtime
|
|
75
|
+
self.bound_llm = llm.bind_tools(tools)
|
|
76
|
+
except NotImplementedError as ex:
|
|
77
|
+
logger.error("%s Failed to bind tools: %s", AGENT_LOG_PREFIX, ex)
|
|
78
|
+
raise
|
|
79
|
+
|
|
80
|
+
if prompt is not None:
|
|
81
|
+
system_prompt = SystemMessage(content=prompt)
|
|
82
|
+
prompt_runnable = RunnableLambda(
|
|
83
|
+
lambda state: [system_prompt] + state.get("messages", []),
|
|
84
|
+
name="SystemPrompt",
|
|
85
|
+
)
|
|
86
|
+
else:
|
|
87
|
+
prompt_runnable = RunnableLambda(
|
|
88
|
+
lambda state: state.get("messages", []),
|
|
89
|
+
name="PromptPassthrough",
|
|
90
|
+
)
|
|
91
|
+
|
|
92
|
+
self.agent = prompt_runnable | self.bound_llm
|
|
53
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 []
|
|
54
95
|
logger.debug("%s Initialized Tool Calling Agent Graph", AGENT_LOG_PREFIX)
|
|
55
96
|
|
|
56
97
|
async def agent_node(self, state: ToolCallAgentGraphState):
|
|
57
98
|
try:
|
|
58
|
-
logger.debug(
|
|
99
|
+
logger.debug("%s Starting the Tool Calling Agent Node", AGENT_LOG_PREFIX)
|
|
59
100
|
if len(state.messages) == 0:
|
|
60
101
|
raise RuntimeError('No input received in state: "messages"')
|
|
61
|
-
response = await self.
|
|
102
|
+
response = await self.agent.ainvoke(
|
|
103
|
+
{"messages": state.messages},
|
|
104
|
+
config=RunnableConfig(callbacks=self.callbacks),
|
|
105
|
+
)
|
|
62
106
|
if self.detailed_logs:
|
|
63
107
|
agent_input = "\n".join(str(message.content) for message in state.messages)
|
|
64
108
|
logger.info(AGENT_CALL_LOG_MESSAGE, agent_input, response)
|
|
@@ -66,8 +110,8 @@ class ToolCallAgentGraph(DualNodeAgent):
|
|
|
66
110
|
state.messages += [response]
|
|
67
111
|
return state
|
|
68
112
|
except Exception as ex:
|
|
69
|
-
logger.
|
|
70
|
-
raise
|
|
113
|
+
logger.error("%s Failed to call agent_node: %s", AGENT_LOG_PREFIX, ex)
|
|
114
|
+
raise
|
|
71
115
|
|
|
72
116
|
async def conditional_edge(self, state: ToolCallAgentGraphState):
|
|
73
117
|
try:
|
|
@@ -75,16 +119,13 @@ class ToolCallAgentGraph(DualNodeAgent):
|
|
|
75
119
|
last_message = state.messages[-1]
|
|
76
120
|
if last_message.tool_calls:
|
|
77
121
|
# the agent wants to call a tool
|
|
78
|
-
logger.debug(
|
|
122
|
+
logger.debug("%s Agent is calling a tool", AGENT_LOG_PREFIX)
|
|
79
123
|
return AgentDecision.TOOL
|
|
80
124
|
if self.detailed_logs:
|
|
81
125
|
logger.debug("%s Final answer:\n%s", AGENT_LOG_PREFIX, state.messages[-1].content)
|
|
82
126
|
return AgentDecision.END
|
|
83
127
|
except Exception as ex:
|
|
84
|
-
logger.exception("%s Failed to determine whether agent is calling a tool: %s",
|
|
85
|
-
AGENT_LOG_PREFIX,
|
|
86
|
-
ex,
|
|
87
|
-
exc_info=True)
|
|
128
|
+
logger.exception("%s Failed to determine whether agent is calling a tool: %s", AGENT_LOG_PREFIX, ex)
|
|
88
129
|
logger.warning("%s Ending graph traversal", AGENT_LOG_PREFIX)
|
|
89
130
|
return AgentDecision.END
|
|
90
131
|
|
|
@@ -92,28 +133,110 @@ class ToolCallAgentGraph(DualNodeAgent):
|
|
|
92
133
|
try:
|
|
93
134
|
logger.debug("%s Starting Tool Node", AGENT_LOG_PREFIX)
|
|
94
135
|
tool_calls = state.messages[-1].tool_calls
|
|
95
|
-
tools = [tool.get(
|
|
136
|
+
tools = [tool.get("name") for tool in tool_calls]
|
|
96
137
|
tool_input = state.messages[-1]
|
|
97
|
-
tool_response = await self.tool_caller.ainvoke(
|
|
98
|
-
|
|
99
|
-
|
|
138
|
+
tool_response = await self.tool_caller.ainvoke(
|
|
139
|
+
input={"messages": [tool_input]},
|
|
140
|
+
config=RunnableConfig(callbacks=self.callbacks, configurable={}),
|
|
141
|
+
)
|
|
100
142
|
# this configurable = {} argument is needed due to a bug in LangGraph PreBuilt ToolNode ^
|
|
101
143
|
|
|
102
|
-
for response in tool_response.get(
|
|
144
|
+
for response in tool_response.get("messages"):
|
|
103
145
|
if self.detailed_logs:
|
|
104
146
|
self._log_tool_response(str(tools), str(tool_input), response.content)
|
|
105
147
|
state.messages += [response]
|
|
106
148
|
|
|
107
149
|
return state
|
|
108
150
|
except Exception as ex:
|
|
109
|
-
logger.
|
|
110
|
-
raise
|
|
151
|
+
logger.error("%s Failed to call tool_node: %s", AGENT_LOG_PREFIX, ex)
|
|
152
|
+
raise
|
|
153
|
+
|
|
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.
|
|
111
160
|
|
|
112
|
-
|
|
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
|
+
"""
|
|
113
165
|
try:
|
|
114
|
-
|
|
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)
|
|
115
217
|
logger.debug("%s Tool Calling Agent Graph built and compiled successfully", AGENT_LOG_PREFIX)
|
|
116
218
|
return self.graph
|
|
117
219
|
except Exception as ex:
|
|
118
|
-
logger.
|
|
119
|
-
raise
|
|
220
|
+
logger.error("%s Failed to build Tool Calling Agent Graph: %s", AGENT_LOG_PREFIX, ex)
|
|
221
|
+
raise
|
|
222
|
+
|
|
223
|
+
|
|
224
|
+
def create_tool_calling_agent_prompt(config: "ToolCallAgentWorkflowConfig") -> str | None:
|
|
225
|
+
"""
|
|
226
|
+
Create a Tool Calling Agent prompt from the config.
|
|
227
|
+
|
|
228
|
+
Args:
|
|
229
|
+
config (ToolCallAgentWorkflowConfig): The config to use for the prompt.
|
|
230
|
+
|
|
231
|
+
Returns:
|
|
232
|
+
ChatPromptTemplate: The Tool Calling Agent prompt.
|
|
233
|
+
"""
|
|
234
|
+
# the Tool Calling Agent prompt can be customized via config option system_prompt and additional_instructions.
|
|
235
|
+
|
|
236
|
+
prompt_strs = []
|
|
237
|
+
for msg in [config.system_prompt, config.additional_instructions]:
|
|
238
|
+
if msg is not None:
|
|
239
|
+
prompt_strs.append(msg)
|
|
240
|
+
if prompt_strs:
|
|
241
|
+
return " ".join(prompt_strs)
|
|
242
|
+
return None
|