nvidia-nat 1.3.0.dev2__py3-none-any.whl → 1.3.0rc1__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.
Files changed (242) hide show
  1. aiq/__init__.py +2 -2
  2. nat/agent/base.py +24 -15
  3. nat/agent/dual_node.py +9 -4
  4. nat/agent/prompt_optimizer/prompt.py +68 -0
  5. nat/agent/prompt_optimizer/register.py +149 -0
  6. nat/agent/react_agent/agent.py +79 -47
  7. nat/agent/react_agent/register.py +41 -21
  8. nat/agent/reasoning_agent/reasoning_agent.py +11 -9
  9. nat/agent/register.py +1 -1
  10. nat/agent/rewoo_agent/agent.py +326 -148
  11. nat/agent/rewoo_agent/prompt.py +19 -22
  12. nat/agent/rewoo_agent/register.py +46 -26
  13. nat/agent/tool_calling_agent/agent.py +84 -28
  14. nat/agent/tool_calling_agent/register.py +51 -28
  15. nat/authentication/api_key/api_key_auth_provider.py +2 -2
  16. nat/authentication/credential_validator/bearer_token_validator.py +557 -0
  17. nat/authentication/http_basic_auth/http_basic_auth_provider.py +1 -1
  18. nat/authentication/interfaces.py +5 -2
  19. nat/authentication/oauth2/oauth2_auth_code_flow_provider.py +40 -20
  20. nat/authentication/oauth2/oauth2_resource_server_config.py +124 -0
  21. nat/authentication/register.py +0 -1
  22. nat/builder/builder.py +56 -24
  23. nat/builder/component_utils.py +9 -5
  24. nat/builder/context.py +46 -11
  25. nat/builder/eval_builder.py +16 -11
  26. nat/builder/framework_enum.py +1 -0
  27. nat/builder/front_end.py +1 -1
  28. nat/builder/function.py +378 -8
  29. nat/builder/function_base.py +3 -3
  30. nat/builder/function_info.py +6 -8
  31. nat/builder/user_interaction_manager.py +2 -2
  32. nat/builder/workflow.py +13 -1
  33. nat/builder/workflow_builder.py +281 -76
  34. nat/cli/cli_utils/config_override.py +2 -2
  35. nat/cli/commands/evaluate.py +1 -1
  36. nat/cli/commands/info/info.py +16 -6
  37. nat/cli/commands/info/list_channels.py +1 -1
  38. nat/cli/commands/info/list_components.py +7 -8
  39. nat/cli/commands/mcp/__init__.py +14 -0
  40. nat/cli/commands/mcp/mcp.py +986 -0
  41. nat/cli/commands/object_store/__init__.py +14 -0
  42. nat/cli/commands/object_store/object_store.py +227 -0
  43. nat/cli/commands/optimize.py +90 -0
  44. nat/cli/commands/registry/publish.py +2 -2
  45. nat/cli/commands/registry/pull.py +2 -2
  46. nat/cli/commands/registry/remove.py +2 -2
  47. nat/cli/commands/registry/search.py +15 -17
  48. nat/cli/commands/start.py +16 -5
  49. nat/cli/commands/uninstall.py +1 -1
  50. nat/cli/commands/workflow/templates/config.yml.j2 +0 -1
  51. nat/cli/commands/workflow/templates/pyproject.toml.j2 +4 -1
  52. nat/cli/commands/workflow/templates/register.py.j2 +0 -1
  53. nat/cli/commands/workflow/workflow_commands.py +9 -13
  54. nat/cli/entrypoint.py +8 -10
  55. nat/cli/register_workflow.py +38 -4
  56. nat/cli/type_registry.py +75 -6
  57. nat/control_flow/__init__.py +0 -0
  58. nat/control_flow/register.py +20 -0
  59. nat/control_flow/router_agent/__init__.py +0 -0
  60. nat/control_flow/router_agent/agent.py +329 -0
  61. nat/control_flow/router_agent/prompt.py +48 -0
  62. nat/control_flow/router_agent/register.py +91 -0
  63. nat/control_flow/sequential_executor.py +166 -0
  64. nat/data_models/agent.py +34 -0
  65. nat/data_models/api_server.py +10 -10
  66. nat/data_models/authentication.py +23 -9
  67. nat/data_models/common.py +1 -1
  68. nat/data_models/component.py +2 -0
  69. nat/data_models/component_ref.py +11 -0
  70. nat/data_models/config.py +41 -17
  71. nat/data_models/dataset_handler.py +1 -1
  72. nat/data_models/discovery_metadata.py +4 -4
  73. nat/data_models/evaluate.py +4 -1
  74. nat/data_models/function.py +34 -0
  75. nat/data_models/function_dependencies.py +14 -6
  76. nat/data_models/gated_field_mixin.py +242 -0
  77. nat/data_models/intermediate_step.py +3 -3
  78. nat/data_models/optimizable.py +119 -0
  79. nat/data_models/optimizer.py +149 -0
  80. nat/data_models/swe_bench_model.py +1 -1
  81. nat/data_models/temperature_mixin.py +44 -0
  82. nat/data_models/thinking_mixin.py +86 -0
  83. nat/data_models/top_p_mixin.py +44 -0
  84. nat/embedder/nim_embedder.py +1 -1
  85. nat/embedder/openai_embedder.py +1 -1
  86. nat/embedder/register.py +0 -1
  87. nat/eval/config.py +3 -1
  88. nat/eval/dataset_handler/dataset_handler.py +71 -7
  89. nat/eval/evaluate.py +86 -31
  90. nat/eval/evaluator/base_evaluator.py +1 -1
  91. nat/eval/evaluator/evaluator_model.py +13 -0
  92. nat/eval/intermediate_step_adapter.py +1 -1
  93. nat/eval/rag_evaluator/evaluate.py +2 -2
  94. nat/eval/rag_evaluator/register.py +3 -3
  95. nat/eval/register.py +4 -1
  96. nat/eval/remote_workflow.py +3 -3
  97. nat/eval/runtime_evaluator/__init__.py +14 -0
  98. nat/eval/runtime_evaluator/evaluate.py +123 -0
  99. nat/eval/runtime_evaluator/register.py +100 -0
  100. nat/eval/swe_bench_evaluator/evaluate.py +6 -6
  101. nat/eval/trajectory_evaluator/evaluate.py +1 -1
  102. nat/eval/trajectory_evaluator/register.py +1 -1
  103. nat/eval/tunable_rag_evaluator/evaluate.py +4 -7
  104. nat/eval/utils/eval_trace_ctx.py +89 -0
  105. nat/eval/utils/weave_eval.py +18 -9
  106. nat/experimental/decorators/experimental_warning_decorator.py +27 -7
  107. nat/experimental/test_time_compute/functions/plan_select_execute_function.py +7 -3
  108. nat/experimental/test_time_compute/functions/ttc_tool_orchestration_function.py +3 -3
  109. nat/experimental/test_time_compute/functions/ttc_tool_wrapper_function.py +1 -1
  110. nat/experimental/test_time_compute/models/strategy_base.py +5 -4
  111. nat/experimental/test_time_compute/register.py +0 -1
  112. nat/experimental/test_time_compute/selection/llm_based_output_merging_selector.py +1 -3
  113. nat/front_ends/console/authentication_flow_handler.py +82 -30
  114. nat/front_ends/console/console_front_end_plugin.py +8 -5
  115. nat/front_ends/fastapi/auth_flow_handlers/websocket_flow_handler.py +52 -17
  116. nat/front_ends/fastapi/dask_client_mixin.py +65 -0
  117. nat/front_ends/fastapi/fastapi_front_end_config.py +36 -5
  118. nat/front_ends/fastapi/fastapi_front_end_controller.py +4 -4
  119. nat/front_ends/fastapi/fastapi_front_end_plugin.py +135 -4
  120. nat/front_ends/fastapi/fastapi_front_end_plugin_worker.py +481 -281
  121. nat/front_ends/fastapi/job_store.py +518 -99
  122. nat/front_ends/fastapi/main.py +11 -19
  123. nat/front_ends/fastapi/message_handler.py +13 -14
  124. nat/front_ends/fastapi/message_validator.py +17 -19
  125. nat/front_ends/fastapi/response_helpers.py +4 -4
  126. nat/front_ends/fastapi/step_adaptor.py +2 -2
  127. nat/front_ends/fastapi/utils.py +57 -0
  128. nat/front_ends/mcp/introspection_token_verifier.py +73 -0
  129. nat/front_ends/mcp/mcp_front_end_config.py +10 -1
  130. nat/front_ends/mcp/mcp_front_end_plugin.py +45 -13
  131. nat/front_ends/mcp/mcp_front_end_plugin_worker.py +116 -8
  132. nat/front_ends/mcp/tool_converter.py +44 -14
  133. nat/front_ends/register.py +0 -1
  134. nat/front_ends/simple_base/simple_front_end_plugin_base.py +3 -1
  135. nat/llm/aws_bedrock_llm.py +24 -12
  136. nat/llm/azure_openai_llm.py +13 -6
  137. nat/llm/litellm_llm.py +69 -0
  138. nat/llm/nim_llm.py +20 -8
  139. nat/llm/openai_llm.py +14 -6
  140. nat/llm/register.py +4 -1
  141. nat/llm/utils/env_config_value.py +2 -3
  142. nat/llm/utils/thinking.py +215 -0
  143. nat/meta/pypi.md +9 -9
  144. nat/object_store/register.py +0 -1
  145. nat/observability/exporter/base_exporter.py +3 -3
  146. nat/observability/exporter/file_exporter.py +1 -1
  147. nat/observability/exporter/processing_exporter.py +309 -81
  148. nat/observability/exporter/span_exporter.py +1 -1
  149. nat/observability/exporter_manager.py +7 -7
  150. nat/observability/mixin/file_mixin.py +7 -7
  151. nat/observability/mixin/redaction_config_mixin.py +42 -0
  152. nat/observability/mixin/tagging_config_mixin.py +62 -0
  153. nat/observability/mixin/type_introspection_mixin.py +420 -107
  154. nat/observability/processor/batching_processor.py +5 -7
  155. nat/observability/processor/falsy_batch_filter_processor.py +55 -0
  156. nat/observability/processor/processor.py +3 -0
  157. nat/observability/processor/processor_factory.py +70 -0
  158. nat/observability/processor/redaction/__init__.py +24 -0
  159. nat/observability/processor/redaction/contextual_redaction_processor.py +125 -0
  160. nat/observability/processor/redaction/contextual_span_redaction_processor.py +66 -0
  161. nat/observability/processor/redaction/redaction_processor.py +177 -0
  162. nat/observability/processor/redaction/span_header_redaction_processor.py +92 -0
  163. nat/observability/processor/span_tagging_processor.py +68 -0
  164. nat/observability/register.py +6 -4
  165. nat/profiler/calc/calc_runner.py +3 -4
  166. nat/profiler/callbacks/agno_callback_handler.py +1 -1
  167. nat/profiler/callbacks/langchain_callback_handler.py +6 -6
  168. nat/profiler/callbacks/llama_index_callback_handler.py +3 -3
  169. nat/profiler/callbacks/semantic_kernel_callback_handler.py +3 -3
  170. nat/profiler/data_frame_row.py +1 -1
  171. nat/profiler/decorators/framework_wrapper.py +62 -13
  172. nat/profiler/decorators/function_tracking.py +160 -3
  173. nat/profiler/forecasting/models/forecasting_base_model.py +3 -1
  174. nat/profiler/inference_optimization/bottleneck_analysis/simple_stack_analysis.py +1 -1
  175. nat/profiler/inference_optimization/data_models.py +3 -3
  176. nat/profiler/inference_optimization/experimental/prefix_span_analysis.py +7 -8
  177. nat/profiler/inference_optimization/token_uniqueness.py +1 -1
  178. nat/profiler/parameter_optimization/__init__.py +0 -0
  179. nat/profiler/parameter_optimization/optimizable_utils.py +93 -0
  180. nat/profiler/parameter_optimization/optimizer_runtime.py +67 -0
  181. nat/profiler/parameter_optimization/parameter_optimizer.py +153 -0
  182. nat/profiler/parameter_optimization/parameter_selection.py +107 -0
  183. nat/profiler/parameter_optimization/pareto_visualizer.py +380 -0
  184. nat/profiler/parameter_optimization/prompt_optimizer.py +384 -0
  185. nat/profiler/parameter_optimization/update_helpers.py +66 -0
  186. nat/profiler/profile_runner.py +14 -9
  187. nat/profiler/utils.py +4 -2
  188. nat/registry_handlers/local/local_handler.py +2 -2
  189. nat/registry_handlers/package_utils.py +1 -2
  190. nat/registry_handlers/pypi/pypi_handler.py +23 -26
  191. nat/registry_handlers/register.py +3 -4
  192. nat/registry_handlers/rest/rest_handler.py +12 -13
  193. nat/retriever/milvus/retriever.py +2 -2
  194. nat/retriever/nemo_retriever/retriever.py +1 -1
  195. nat/retriever/register.py +0 -1
  196. nat/runtime/loader.py +2 -2
  197. nat/runtime/runner.py +3 -2
  198. nat/runtime/session.py +43 -8
  199. nat/settings/global_settings.py +16 -5
  200. nat/tool/chat_completion.py +5 -2
  201. nat/tool/code_execution/local_sandbox/local_sandbox_server.py +3 -3
  202. nat/tool/datetime_tools.py +49 -9
  203. nat/tool/document_search.py +2 -2
  204. nat/tool/github_tools.py +450 -0
  205. nat/tool/nvidia_rag.py +1 -1
  206. nat/tool/register.py +2 -9
  207. nat/tool/retriever.py +3 -2
  208. nat/utils/callable_utils.py +70 -0
  209. nat/utils/data_models/schema_validator.py +3 -3
  210. nat/utils/exception_handlers/automatic_retries.py +104 -51
  211. nat/utils/exception_handlers/schemas.py +1 -1
  212. nat/utils/io/yaml_tools.py +2 -2
  213. nat/utils/log_levels.py +25 -0
  214. nat/utils/reactive/base/observable_base.py +2 -2
  215. nat/utils/reactive/base/observer_base.py +1 -1
  216. nat/utils/reactive/observable.py +2 -2
  217. nat/utils/reactive/observer.py +4 -4
  218. nat/utils/reactive/subscription.py +1 -1
  219. nat/utils/settings/global_settings.py +6 -8
  220. nat/utils/type_converter.py +4 -3
  221. nat/utils/type_utils.py +9 -5
  222. {nvidia_nat-1.3.0.dev2.dist-info → nvidia_nat-1.3.0rc1.dist-info}/METADATA +42 -16
  223. {nvidia_nat-1.3.0.dev2.dist-info → nvidia_nat-1.3.0rc1.dist-info}/RECORD +230 -189
  224. {nvidia_nat-1.3.0.dev2.dist-info → nvidia_nat-1.3.0rc1.dist-info}/entry_points.txt +1 -0
  225. nat/cli/commands/info/list_mcp.py +0 -304
  226. nat/tool/github_tools/create_github_commit.py +0 -133
  227. nat/tool/github_tools/create_github_issue.py +0 -87
  228. nat/tool/github_tools/create_github_pr.py +0 -106
  229. nat/tool/github_tools/get_github_file.py +0 -106
  230. nat/tool/github_tools/get_github_issue.py +0 -166
  231. nat/tool/github_tools/get_github_pr.py +0 -256
  232. nat/tool/github_tools/update_github_issue.py +0 -100
  233. nat/tool/mcp/exceptions.py +0 -142
  234. nat/tool/mcp/mcp_client.py +0 -255
  235. nat/tool/mcp/mcp_tool.py +0 -96
  236. nat/utils/exception_handlers/mcp.py +0 -211
  237. /nat/{tool/github_tools → agent/prompt_optimizer}/__init__.py +0 -0
  238. /nat/{tool/mcp → authentication/credential_validator}/__init__.py +0 -0
  239. {nvidia_nat-1.3.0.dev2.dist-info → nvidia_nat-1.3.0rc1.dist-info}/WHEEL +0 -0
  240. {nvidia_nat-1.3.0.dev2.dist-info → nvidia_nat-1.3.0rc1.dist-info}/licenses/LICENSE-3rd-party.txt +0 -0
  241. {nvidia_nat-1.3.0.dev2.dist-info → nvidia_nat-1.3.0rc1.dist-info}/licenses/LICENSE.md +0 -0
  242. {nvidia_nat-1.3.0.dev2.dist-info → nvidia_nat-1.3.0rc1.dist-info}/top_level.txt +0 -0
@@ -0,0 +1,986 @@
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 asyncio
17
+ import json
18
+ import logging
19
+ import time
20
+ from typing import Any
21
+ from typing import Literal
22
+ from typing import cast
23
+
24
+ import click
25
+ from pydantic import BaseModel
26
+
27
+ from nat.cli.commands.start import start_command
28
+
29
+ logger = logging.getLogger(__name__)
30
+
31
+
32
+ @click.group(name=__name__, invoke_without_command=False, help="MCP-related commands.")
33
+ def mcp_command():
34
+ """
35
+ MCP-related commands.
36
+ """
37
+ return None
38
+
39
+
40
+ # nat mcp serve: reuses the start/mcp frontend command
41
+ mcp_command.add_command(start_command.get_command(None, "mcp"), name="serve") # type: ignore
42
+
43
+ # Suppress verbose logs from mcp.client.sse and httpx
44
+ logging.getLogger("mcp.client.sse").setLevel(logging.WARNING)
45
+ logging.getLogger("httpx").setLevel(logging.WARNING)
46
+
47
+ try:
48
+ from nat.plugins.mcp.exception_handler import format_mcp_error
49
+ from nat.plugins.mcp.exceptions import MCPError
50
+ except ImportError:
51
+ # Fallback for when MCP client package is not installed
52
+ MCPError = Exception
53
+
54
+ def format_mcp_error(error, include_traceback=False):
55
+ click.echo(f"Error: {error}", err=True)
56
+
57
+
58
+ def validate_transport_cli_args(transport: str, command: str | None, args: str | None, env: str | None) -> bool:
59
+ """
60
+ Validate transport and parameter combinations, returning False if invalid.
61
+
62
+ Args:
63
+ transport: The transport type ('sse', 'stdio', or 'streamable-http')
64
+ command: Command for stdio transport
65
+ args: Arguments for stdio transport
66
+ env: Environment variables for stdio transport
67
+
68
+ Returns:
69
+ bool: True if valid, False if invalid (error message already displayed)
70
+ """
71
+ if transport == 'stdio':
72
+ if not command:
73
+ click.echo("--command is required when using stdio client type", err=True)
74
+ return False
75
+ elif transport in ['sse', 'streamable-http']:
76
+ if command or args or env:
77
+ click.echo("--command, --args, and --env are not allowed when using sse or streamable-http client type",
78
+ err=True)
79
+ return False
80
+ return True
81
+
82
+
83
+ class MCPPingResult(BaseModel):
84
+ """Result of an MCP server ping request.
85
+
86
+ Attributes:
87
+ url (str): The MCP server URL that was pinged
88
+ status (str): Health status - 'healthy', 'unhealthy', or 'unknown'
89
+ response_time_ms (float | None): Response time in milliseconds, None if request failed completely
90
+ error (str | None): Error message if the ping failed, None if successful
91
+ """
92
+ url: str
93
+ status: str
94
+ response_time_ms: float | None
95
+ error: str | None
96
+
97
+
98
+ def format_tool(tool: Any) -> dict[str, str | None]:
99
+ """Format an MCP tool into a dictionary for display.
100
+
101
+ Extracts name, description, and input schema from various MCP tool object types
102
+ and normalizes them into a consistent dictionary format for CLI display.
103
+
104
+ Args:
105
+ tool (Any): MCPToolClient or raw MCP Tool object (uses Any due to different types)
106
+
107
+ Returns:
108
+ dict[str, str | None]: Dictionary with name, description, and input_schema as keys
109
+ """
110
+ name = getattr(tool, 'name', None)
111
+ description = getattr(tool, 'description', '')
112
+ input_schema = getattr(tool, 'input_schema', None) or getattr(tool, 'inputSchema', None)
113
+
114
+ # Normalize schema to JSON string
115
+ if input_schema is None:
116
+ return {
117
+ "name": name,
118
+ "description": description,
119
+ "input_schema": None,
120
+ }
121
+ elif hasattr(input_schema, "schema_json"):
122
+ schema_str = input_schema.schema_json(indent=2)
123
+ elif hasattr(input_schema, "model_json_schema"):
124
+ schema_str = json.dumps(input_schema.model_json_schema(), indent=2)
125
+ elif isinstance(input_schema, dict):
126
+ schema_str = json.dumps(input_schema, indent=2)
127
+ else:
128
+ # Final fallback: attempt to dump stringified version wrapped as JSON string
129
+ schema_str = json.dumps({"raw": str(input_schema)}, indent=2)
130
+
131
+ return {
132
+ "name": name,
133
+ "description": description,
134
+ "input_schema": schema_str,
135
+ }
136
+
137
+
138
+ def print_tool(tool_dict: dict[str, str | None], detail: bool = False) -> None:
139
+ """Print a formatted tool to the console with optional detailed information.
140
+
141
+ Outputs tool information in a user-friendly format to stdout. When detail=True
142
+ or when description/schema are available, shows full information with separator.
143
+
144
+ Args:
145
+ tool_dict (dict[str, str | None]): Dictionary containing tool information with name, description, and
146
+ input_schema as keys
147
+ detail (bool, optional): Whether to force detailed output. Defaults to False.
148
+ """
149
+ click.echo(f"Tool: {tool_dict.get('name', 'Unknown')}")
150
+ if detail or tool_dict.get('input_schema') or tool_dict.get('description'):
151
+ click.echo(f"Description: {tool_dict.get('description', 'No description available')}")
152
+ if tool_dict.get("input_schema"):
153
+ click.echo("Input Schema:")
154
+ click.echo(tool_dict.get("input_schema"))
155
+ else:
156
+ click.echo("Input Schema: None")
157
+ click.echo("-" * 60)
158
+
159
+
160
+ def _set_auth_defaults(auth: bool,
161
+ url: str | None,
162
+ auth_redirect_uri: str | None,
163
+ auth_user_id: str | None,
164
+ auth_scopes: str | None) -> tuple[str | None, str | None, list[str] | None]:
165
+ """Set default auth values when --auth flag is used.
166
+
167
+ Args:
168
+ auth: Whether --auth flag was used
169
+ url: MCP server URL
170
+ auth_redirect_uri: OAuth2 redirect URI
171
+ auth_user_id: User ID for authentication
172
+ auth_scopes: OAuth2 scopes (comma-separated string)
173
+
174
+ Returns:
175
+ Tuple of (auth_redirect_uri, auth_user_id, auth_scopes_list) with defaults applied
176
+ """
177
+ if auth:
178
+ auth_redirect_uri = auth_redirect_uri or "http://localhost:8000/auth/redirect"
179
+ auth_user_id = auth_user_id or url
180
+ auth_scopes = auth_scopes or ""
181
+
182
+ # Convert comma-separated string to list, stripping whitespace
183
+ auth_scopes_list = [scope.strip() for scope in auth_scopes.split(',')] if auth_scopes else None
184
+
185
+ return auth_redirect_uri, auth_user_id, auth_scopes_list
186
+
187
+
188
+ async def _create_mcp_client_config(
189
+ builder,
190
+ server_cfg,
191
+ url: str | None,
192
+ transport: str,
193
+ auth_redirect_uri: str | None,
194
+ auth_user_id: str | None,
195
+ auth_scopes: list[str] | None,
196
+ ):
197
+ from nat.plugins.mcp.client_impl import MCPClientConfig
198
+
199
+ if url and transport == "streamable-http" and auth_redirect_uri:
200
+ try:
201
+ from nat.plugins.mcp.auth.auth_provider_config import MCPOAuth2ProviderConfig
202
+ auth_config = MCPOAuth2ProviderConfig(
203
+ server_url=url,
204
+ redirect_uri=auth_redirect_uri,
205
+ default_user_id=auth_user_id or url,
206
+ scopes=auth_scopes or [],
207
+ )
208
+ auth_provider_name = "mcp_oauth2_cli"
209
+ await builder.add_auth_provider(auth_provider_name, auth_config)
210
+ server_cfg.auth_provider = auth_provider_name
211
+ except ImportError:
212
+ click.echo(
213
+ "[WARNING] MCP OAuth2 authentication requires nvidia-nat-mcp package.",
214
+ err=True,
215
+ )
216
+
217
+ return MCPClientConfig(server=server_cfg)
218
+
219
+
220
+ async def list_tools_via_function_group(
221
+ command: str | None,
222
+ url: str | None,
223
+ tool_name: str | None = None,
224
+ transport: str = 'sse',
225
+ args: list[str] | None = None,
226
+ env: dict[str, str] | None = None,
227
+ auth_redirect_uri: str | None = None,
228
+ auth_user_id: str | None = None,
229
+ auth_scopes: list[str] | None = None,
230
+ ) -> list[dict[str, str | None]]:
231
+ """List tools by constructing the mcp_client function group and introspecting functions.
232
+
233
+ Mirrors the behavior of list_mcp.py but routes through the registered function group to ensure
234
+ parity with workflow configuration.
235
+ """
236
+ try:
237
+ # Ensure the registration side-effects are loaded
238
+ from nat.builder.workflow_builder import WorkflowBuilder
239
+ from nat.plugins.mcp.client_impl import MCPClientConfig
240
+ from nat.plugins.mcp.client_impl import MCPServerConfig
241
+ except ImportError:
242
+ click.echo(
243
+ "MCP client functionality requires nvidia-nat-mcp package. Install with: uv pip install nvidia-nat-mcp",
244
+ err=True)
245
+ return []
246
+
247
+ if args is None:
248
+ args = []
249
+
250
+ # Build server config according to transport
251
+ server_cfg = MCPServerConfig(
252
+ transport=cast(Literal["stdio", "sse", "streamable-http"], transport),
253
+ url=cast(Any, url) if transport in ('sse', 'streamable-http') else None,
254
+ command=command if transport == 'stdio' else None,
255
+ args=args if transport == 'stdio' else None,
256
+ env=env if transport == 'stdio' else None,
257
+ )
258
+
259
+ group_cfg = MCPClientConfig(server=server_cfg)
260
+
261
+ tools: list[dict[str, str | None]] = []
262
+
263
+ async with WorkflowBuilder() as builder: # type: ignore
264
+ # Add auth provider if url is provided and auth_redirect_uri is given (only for streamable-http)
265
+ group_cfg = await _create_mcp_client_config(builder,
266
+ server_cfg,
267
+ url,
268
+ transport,
269
+ auth_redirect_uri,
270
+ auth_user_id,
271
+ auth_scopes)
272
+ group = await builder.add_function_group("mcp_client", group_cfg)
273
+
274
+ # Access functions exposed by the group
275
+ fns = await group.get_accessible_functions()
276
+
277
+ def to_tool_entry(full_name: str, fn_obj) -> dict[str, str | None]:
278
+ # full_name like "mcp_client.<tool>"
279
+ name = full_name.split(".", 1)[1] if "." in full_name else full_name
280
+ schema = getattr(fn_obj, 'input_schema', None)
281
+ if schema is None:
282
+ schema_str = None
283
+ elif hasattr(schema, "schema_json"):
284
+ schema_str = schema.schema_json(indent=2)
285
+ elif hasattr(schema, "model_json_schema"):
286
+ try:
287
+ schema_str = json.dumps(schema.model_json_schema(), indent=2)
288
+ except Exception:
289
+ schema_str = None
290
+ else:
291
+ schema_str = None
292
+ return {"name": name, "description": getattr(fn_obj, 'description', ''), "input_schema": schema_str}
293
+
294
+ if tool_name:
295
+ full = f"mcp_client.{tool_name}"
296
+ fn = fns.get(full)
297
+ if fn is not None:
298
+ tools.append(to_tool_entry(full, fn))
299
+ else:
300
+ for full, fn in (await fns).items():
301
+ tools.append(to_tool_entry(full, fn))
302
+
303
+ return tools
304
+
305
+
306
+ async def list_tools_direct(command, url, tool_name=None, transport='sse', args=None, env=None):
307
+ """List MCP tools using direct MCP protocol with structured exception handling.
308
+
309
+ Bypasses MCPBuilder and uses raw MCP ClientSession and SSE client directly.
310
+ Converts raw exceptions to structured MCPErrors for consistent user experience.
311
+ Used when --direct flag is specified in CLI.
312
+
313
+ Args:
314
+ url (str): MCP server URL to connect to
315
+ tool_name (str | None, optional): Specific tool name to retrieve.
316
+ If None, retrieves all available tools. Defaults to None.
317
+
318
+ Returns:
319
+ list[dict[str, str | None]]: List of formatted tool dictionaries, each containing name, description, and
320
+ input_schema as keys
321
+
322
+ Note:
323
+ This function handles ExceptionGroup by extracting the most relevant exception
324
+ and converting it to MCPError for consistent error reporting.
325
+ """
326
+ if args is None:
327
+ args = []
328
+ from mcp import ClientSession
329
+ from mcp.client.sse import sse_client
330
+ from mcp.client.stdio import StdioServerParameters
331
+ from mcp.client.stdio import stdio_client
332
+ from mcp.client.streamable_http import streamablehttp_client
333
+
334
+ try:
335
+ if transport == 'stdio':
336
+
337
+ def get_stdio_client():
338
+ return stdio_client(server=StdioServerParameters(command=command, args=args, env=env))
339
+
340
+ client = get_stdio_client
341
+ elif transport == 'streamable-http':
342
+
343
+ def get_streamable_http_client():
344
+ return streamablehttp_client(url=url)
345
+
346
+ client = get_streamable_http_client
347
+ else:
348
+
349
+ def get_sse_client():
350
+ return sse_client(url=url)
351
+
352
+ client = get_sse_client
353
+
354
+ async with client() as ctx:
355
+ read, write = (ctx[0], ctx[1]) if isinstance(ctx, tuple) else ctx
356
+ async with ClientSession(read, write) as session:
357
+ await session.initialize()
358
+ response = await session.list_tools()
359
+
360
+ tools = []
361
+ for tool in response.tools:
362
+ if tool_name:
363
+ if tool.name == tool_name:
364
+ tools.append(format_tool(tool))
365
+ else:
366
+ tools.append(format_tool(tool))
367
+
368
+ if tool_name and not tools:
369
+ click.echo(f"[INFO] Tool '{tool_name}' not found.")
370
+ return tools
371
+ except Exception as e:
372
+ # Convert raw exceptions to structured MCPError for consistency
373
+ try:
374
+ from nat.plugins.mcp.exception_handler import convert_to_mcp_error
375
+ from nat.plugins.mcp.exception_handler import extract_primary_exception
376
+ except ImportError:
377
+ # Fallback when MCP client package is not installed
378
+ def convert_to_mcp_error(exception, url):
379
+ return Exception(f"Error connecting to {url}: {exception}")
380
+
381
+ def extract_primary_exception(exceptions):
382
+ return exceptions[0] if exceptions else Exception("Unknown error")
383
+
384
+ if isinstance(e, ExceptionGroup):
385
+ primary_exception = extract_primary_exception(list(e.exceptions))
386
+ mcp_error = convert_to_mcp_error(primary_exception, url)
387
+ else:
388
+ mcp_error = convert_to_mcp_error(e, url)
389
+
390
+ format_mcp_error(mcp_error, include_traceback=False)
391
+ return []
392
+
393
+
394
+ async def ping_mcp_server(url: str,
395
+ timeout: int,
396
+ transport: str = 'streamable-http',
397
+ command: str | None = None,
398
+ args: list[str] | None = None,
399
+ env: dict[str, str] | None = None,
400
+ auth_redirect_uri: str | None = None,
401
+ auth_user_id: str | None = None,
402
+ auth_scopes: list[str] | None = None) -> MCPPingResult:
403
+ """Ping an MCP server to check if it's responsive.
404
+
405
+ Args:
406
+ url (str): MCP server URL to ping
407
+ timeout (int): Timeout in seconds for the ping request
408
+
409
+ Returns:
410
+ MCPPingResult: Structured result with status, response_time, and any error info
411
+ """
412
+ from mcp.client.session import ClientSession
413
+ from mcp.client.sse import sse_client
414
+ from mcp.client.stdio import StdioServerParameters
415
+ from mcp.client.stdio import stdio_client
416
+ from mcp.client.streamable_http import streamablehttp_client
417
+
418
+ async def _ping_operation():
419
+ # Select transport
420
+ if transport == 'stdio':
421
+ stdio_args_local: list[str] = args or []
422
+ if not command:
423
+ raise RuntimeError("--command is required for stdio transport")
424
+ client_ctx = stdio_client(server=StdioServerParameters(command=command, args=stdio_args_local, env=env))
425
+ elif transport == 'sse':
426
+ client_ctx = sse_client(url)
427
+ else: # streamable-http
428
+ client_ctx = streamablehttp_client(url=url)
429
+
430
+ async with client_ctx as ctx:
431
+ read, write = (ctx[0], ctx[1]) if isinstance(ctx, tuple) else ctx
432
+ async with ClientSession(read, write) as session:
433
+ await session.initialize()
434
+
435
+ start_time = time.time()
436
+ await session.send_ping()
437
+ end_time = time.time()
438
+ response_time_ms = round((end_time - start_time) * 1000, 2)
439
+
440
+ return MCPPingResult(url=url, status="healthy", response_time_ms=response_time_ms, error=None)
441
+
442
+ try:
443
+ # Apply timeout to the entire ping operation
444
+ return await asyncio.wait_for(_ping_operation(), timeout=timeout)
445
+
446
+ except TimeoutError:
447
+ return MCPPingResult(url=url,
448
+ status="unhealthy",
449
+ response_time_ms=None,
450
+ error=f"Timeout after {timeout} seconds")
451
+
452
+ except Exception as e:
453
+ return MCPPingResult(url=url, status="unhealthy", response_time_ms=None, error=str(e))
454
+
455
+
456
+ @mcp_command.group(name="client", invoke_without_command=False, help="MCP client commands.")
457
+ def mcp_client_command():
458
+ """
459
+ MCP client commands.
460
+ """
461
+ try:
462
+ from nat.runtime.loader import PluginTypes
463
+ from nat.runtime.loader import discover_and_register_plugins
464
+ discover_and_register_plugins(PluginTypes.CONFIG_OBJECT)
465
+ except ImportError:
466
+ click.echo("[WARNING] MCP client functionality requires nvidia-nat-mcp package.", err=True)
467
+ pass
468
+
469
+
470
+ @mcp_client_command.group(name="tool", invoke_without_command=False, help="Inspect and call MCP tools.")
471
+ def mcp_client_tool_group():
472
+ """
473
+ MCP client tool commands.
474
+ """
475
+ return None
476
+
477
+
478
+ @mcp_client_tool_group.command(name="list", help="List tool names (default), or show details with --detail or --tool.")
479
+ @click.option('--direct', is_flag=True, help='Bypass MCPBuilder and use direct MCP protocol')
480
+ @click.option(
481
+ '--url',
482
+ default='http://localhost:9901/mcp',
483
+ show_default=True,
484
+ help='MCP server URL (e.g. http://localhost:8080/mcp for streamable-http, http://localhost:8080/sse for sse)')
485
+ @click.option('--transport',
486
+ type=click.Choice(['sse', 'stdio', 'streamable-http']),
487
+ default='streamable-http',
488
+ show_default=True,
489
+ help='Type of client to use (default: streamable-http, backwards compatible with sse)')
490
+ @click.option('--command', help='For stdio: The command to run (e.g. mcp-server)')
491
+ @click.option('--args', help='For stdio: Additional arguments for the command (space-separated)')
492
+ @click.option('--env', help='For stdio: Environment variables in KEY=VALUE format (space-separated)')
493
+ @click.option('--tool', default=None, help='Get details for a specific tool by name')
494
+ @click.option('--detail', is_flag=True, help='Show full details for all tools')
495
+ @click.option('--json-output', is_flag=True, help='Output tool metadata in JSON format')
496
+ @click.option('--auth',
497
+ is_flag=True,
498
+ help='Enable OAuth2 authentication with default settings (streamable-http only, not with --direct)')
499
+ @click.option('--auth-redirect-uri',
500
+ help='OAuth2 redirect URI for authentication (streamable-http only, not with --direct)')
501
+ @click.option('--auth-user-id', help='User ID for authentication (streamable-http only, not with --direct)')
502
+ @click.option('--auth-scopes', help='OAuth2 scopes (comma-separated, streamable-http only, not with --direct)')
503
+ @click.pass_context
504
+ def mcp_client_tool_list(ctx,
505
+ direct,
506
+ url,
507
+ transport,
508
+ command,
509
+ args,
510
+ env,
511
+ tool,
512
+ detail,
513
+ json_output,
514
+ auth,
515
+ auth_redirect_uri,
516
+ auth_user_id,
517
+ auth_scopes):
518
+ """List MCP tool names (default) or show detailed tool information.
519
+
520
+ Use --detail for full output including descriptions and input schemas.
521
+ If --tool is provided, always shows full output for that specific tool.
522
+ Use --direct to bypass MCPBuilder and use raw MCP protocol.
523
+ Use --json-output to get structured JSON data instead of formatted text.
524
+ Use --auth to enable auth with default settings (streamable-http only, not with --direct).
525
+ Use --auth-redirect-uri to enable auth for protected MCP servers (streamable-http only, not with --direct).
526
+
527
+ Args:
528
+ ctx (click.Context): Click context object for command invocation
529
+ direct (bool): Whether to bypass MCPBuilder and use direct MCP protocol
530
+ url (str): MCP server URL to connect to (default: http://localhost:9901/mcp)
531
+ tool (str | None): Optional specific tool name to retrieve detailed info for
532
+ detail (bool): Whether to show full details (description + schema) for all tools
533
+ json_output (bool): Whether to output tool metadata in JSON format instead of text
534
+ auth (bool): Whether to enable OAuth2 authentication (streamable-http only, not with --direct)
535
+ auth_redirect_uri (str | None): redirect URI for auth (streamable-http only, not with --direct)
536
+ auth_user_id (str | None): User ID for authentication (streamable-http only, not with --direct)
537
+ auth_scopes (str | None): OAuth2 scopes (comma-separated, streamable-http only, not with --direct)
538
+
539
+ Examples:
540
+ nat mcp client tool list # List tool names only
541
+ nat mcp client tool list --detail # Show all tools with full details
542
+ nat mcp client tool list --tool my_tool # Show details for specific tool
543
+ nat mcp client tool list --json-output # Get JSON format output
544
+ nat mcp client tool list --direct --url http://... # Use direct protocol with custom URL (no auth)
545
+ nat mcp client tool list --url https://example.com/mcp/ --auth # With auth using defaults
546
+ nat mcp client tool list --url https://example.com/mcp/ --transport streamable-http \
547
+ --auth-redirect-uri http://localhost:8000/auth/redirect # With custom auth settings
548
+ nat mcp client tool list --url https://example.com/mcp/ --transport streamable-http \
549
+ --auth-redirect-uri http://localhost:8000/auth/redirect --auth-user-id myuser # With auth and user ID
550
+ """
551
+ if ctx.invoked_subcommand is not None:
552
+ return
553
+
554
+ if not validate_transport_cli_args(transport, command, args, env):
555
+ return
556
+
557
+ if transport in ['sse', 'streamable-http']:
558
+ if not url:
559
+ click.echo("[ERROR] --url is required when using sse or streamable-http client type", err=True)
560
+ return
561
+
562
+ # Set auth defaults if --auth flag is used
563
+ auth_redirect_uri, auth_user_id, auth_scopes_list = _set_auth_defaults(
564
+ auth, url, auth_redirect_uri, auth_user_id, auth_scopes
565
+ )
566
+
567
+ stdio_args = args.split() if args else []
568
+ stdio_env = dict(var.split('=', 1) for var in env.split()) if env else None
569
+
570
+ if direct:
571
+ tools = asyncio.run(
572
+ list_tools_direct(command, url, tool_name=tool, transport=transport, args=stdio_args, env=stdio_env))
573
+ else:
574
+ tools = asyncio.run(
575
+ list_tools_via_function_group(command,
576
+ url,
577
+ tool_name=tool,
578
+ transport=transport,
579
+ args=stdio_args,
580
+ env=stdio_env,
581
+ auth_redirect_uri=auth_redirect_uri,
582
+ auth_user_id=auth_user_id,
583
+ auth_scopes=auth_scopes_list))
584
+
585
+ if json_output:
586
+ click.echo(json.dumps(tools, indent=2))
587
+ elif tool:
588
+ for tool_dict in (tools or []):
589
+ print_tool(tool_dict, detail=True)
590
+ elif detail:
591
+ for tool_dict in (tools or []):
592
+ print_tool(tool_dict, detail=True)
593
+ else:
594
+ for tool_dict in (tools or []):
595
+ click.echo(tool_dict.get('name', 'Unknown tool'))
596
+
597
+
598
+ @mcp_client_command.command(name="ping", help="Ping an MCP server to check if it's responsive.")
599
+ @click.option(
600
+ '--url',
601
+ default='http://localhost:9901/mcp',
602
+ show_default=True,
603
+ help='MCP server URL (e.g. http://localhost:8080/mcp for streamable-http, http://localhost:8080/sse for sse)')
604
+ @click.option('--transport',
605
+ type=click.Choice(['sse', 'stdio', 'streamable-http']),
606
+ default='streamable-http',
607
+ show_default=True,
608
+ help='Type of client to use for ping')
609
+ @click.option('--command', help='For stdio: The command to run (e.g. mcp-server)')
610
+ @click.option('--args', help='For stdio: Additional arguments for the command (space-separated)')
611
+ @click.option('--env', help='For stdio: Environment variables in KEY=VALUE format (space-separated)')
612
+ @click.option('--timeout', default=60, show_default=True, help='Timeout in seconds for ping request')
613
+ @click.option('--json-output', is_flag=True, help='Output ping result in JSON format')
614
+ @click.option('--auth-redirect-uri',
615
+ help='OAuth2 redirect URI for authentication (streamable-http only, not with --direct)')
616
+ @click.option('--auth-user-id', help='User ID for authentication (streamable-http only, not with --direct)')
617
+ @click.option('--auth-scopes', help='OAuth2 scopes (comma-separated, streamable-http only, not with --direct)')
618
+ def mcp_client_ping(url: str,
619
+ transport: str,
620
+ command: str | None,
621
+ args: str | None,
622
+ env: str | None,
623
+ timeout: int,
624
+ json_output: bool,
625
+ auth_redirect_uri: str | None,
626
+ auth_user_id: str | None,
627
+ auth_scopes: str | None) -> None:
628
+ """Ping an MCP server to check if it's responsive.
629
+
630
+ This command sends a ping request to the MCP server and measures the response time.
631
+ It's useful for health checks and monitoring server availability.
632
+
633
+ Args:
634
+ url (str): MCP server URL to ping (default: http://localhost:9901/mcp)
635
+ timeout (int): Timeout in seconds for the ping request (default: 60)
636
+ json_output (bool): Whether to output the result in JSON format
637
+ auth_redirect_uri (str | None): redirect URI for auth (streamable-http only, not with --direct)
638
+ auth_user_id (str | None): User ID for auth (streamable-http only, not with --direct)
639
+ auth_scopes (str | None): OAuth2 scopes (comma-separated, streamable-http only, not with --direct)
640
+
641
+ Examples:
642
+ nat mcp client ping # Ping default server
643
+ nat mcp client ping --url http://custom-server:9901/mcp # Ping custom server
644
+ nat mcp client ping --timeout 10 # Use 10 second timeout
645
+ nat mcp client ping --json-output # Get JSON format output
646
+ nat mcp client ping --url https://example.com/mcp/ --transport streamable-http --auth # With auth
647
+ """
648
+ # Validate combinations similar to list command
649
+ if not validate_transport_cli_args(transport, command, args, env):
650
+ return
651
+
652
+ stdio_args = args.split() if args else []
653
+ stdio_env = dict(var.split('=', 1) for var in env.split()) if env else None
654
+
655
+ # Auth validation: if user_id or scopes provided, require redirect_uri
656
+ if (auth_user_id or auth_scopes) and not auth_redirect_uri:
657
+ click.echo("[ERROR] --auth-redirect-uri is required when using --auth-user-id or --auth-scopes", err=True)
658
+ return
659
+
660
+ # Parse auth scopes, stripping whitespace
661
+ auth_scopes_list = [scope.strip() for scope in auth_scopes.split(',')] if auth_scopes else None
662
+
663
+ result = asyncio.run(
664
+ ping_mcp_server(url,
665
+ timeout,
666
+ transport,
667
+ command,
668
+ stdio_args,
669
+ stdio_env,
670
+ auth_redirect_uri,
671
+ auth_user_id,
672
+ auth_scopes_list))
673
+
674
+ if json_output:
675
+ click.echo(result.model_dump_json(indent=2))
676
+ elif result.status == "healthy":
677
+ click.echo(f"Server at {result.url} is healthy (response time: {result.response_time_ms}ms)")
678
+ else:
679
+ click.echo(f"Server at {result.url} {result.status}: {result.error}")
680
+
681
+
682
+ async def call_tool_direct(command: str | None,
683
+ url: str | None,
684
+ tool_name: str,
685
+ transport: str,
686
+ args: list[str] | None,
687
+ env: dict[str, str] | None,
688
+ tool_args: dict[str, Any] | None) -> str:
689
+ """Call an MCP tool directly via the selected transport.
690
+
691
+ Bypasses the WorkflowBuilder and talks to the MCP server using the raw
692
+ protocol client for the given transport. Aggregates tool outputs into a
693
+ plain string suitable for terminal display. Converts transport/protocol
694
+ exceptions into a structured MCPError for consistency.
695
+
696
+ Args:
697
+ command (str | None): For ``stdio`` transport, the command to execute.
698
+ url (str | None): For ``sse`` or ``streamable-http`` transports, the server URL.
699
+ tool_name (str): Name of the tool to call.
700
+ transport (str): One of ``'stdio'``, ``'sse'``, or ``'streamable-http'``.
701
+ args (list[str] | None): For ``stdio`` transport, additional command arguments.
702
+ env (dict[str, str] | None): For ``stdio`` transport, environment variables.
703
+ tool_args (dict[str, Any] | None): JSON-serializable arguments passed to the tool.
704
+
705
+ Returns:
706
+ str: Concatenated textual output from the tool invocation.
707
+
708
+ Raises:
709
+ MCPError: If the connection, initialization, or tool call fails. When the
710
+ MCP client package is not installed, a generic ``Exception`` is raised
711
+ with an MCP-like error message.
712
+ RuntimeError: If required parameters for the chosen transport are missing
713
+ or if the tool returns an error response.
714
+ """
715
+ from mcp import ClientSession
716
+ from mcp.client.sse import sse_client
717
+ from mcp.client.stdio import StdioServerParameters
718
+ from mcp.client.stdio import stdio_client
719
+ from mcp.client.streamable_http import streamablehttp_client
720
+ from mcp.types import TextContent
721
+
722
+ try:
723
+ if transport == 'stdio':
724
+ if not command:
725
+ raise RuntimeError("--command is required for stdio transport")
726
+
727
+ def get_stdio_client():
728
+ return stdio_client(server=StdioServerParameters(command=command, args=args or [], env=env))
729
+
730
+ client = get_stdio_client
731
+ elif transport == 'streamable-http':
732
+
733
+ def get_streamable_http_client():
734
+ if not url:
735
+ raise RuntimeError("--url is required for streamable-http transport")
736
+ return streamablehttp_client(url=url)
737
+
738
+ client = get_streamable_http_client
739
+ else:
740
+
741
+ def get_sse_client():
742
+ if not url:
743
+ raise RuntimeError("--url is required for sse transport")
744
+ return sse_client(url=url)
745
+
746
+ client = get_sse_client
747
+
748
+ async with client() as ctx:
749
+ read, write = (ctx[0], ctx[1]) if isinstance(ctx, tuple) else ctx
750
+ async with ClientSession(read, write) as session:
751
+ await session.initialize()
752
+ result = await session.call_tool(tool_name, tool_args or {})
753
+
754
+ outputs: list[str] = []
755
+ for content in result.content:
756
+ if isinstance(content, TextContent):
757
+ outputs.append(content.text)
758
+ else:
759
+ outputs.append(str(content))
760
+
761
+ # If the result indicates an error, raise to surface in CLI
762
+ if getattr(result, "isError", False):
763
+ raise RuntimeError("\n".join(outputs) or f"Tool call '{tool_name}' returned an error")
764
+
765
+ return "\n".join(outputs)
766
+ except Exception as e:
767
+ # Convert raw exceptions to structured MCPError for consistency
768
+ try:
769
+ from nat.plugins.mcp.exception_handler import convert_to_mcp_error
770
+ from nat.plugins.mcp.exception_handler import extract_primary_exception
771
+ except ImportError:
772
+ # Fallback when MCP client package is not installed
773
+ def convert_to_mcp_error(exception: Exception, url: str):
774
+ return Exception(f"Error connecting to {url}: {exception}")
775
+
776
+ def extract_primary_exception(exceptions):
777
+ return exceptions[0] if exceptions else Exception("Unknown error")
778
+
779
+ endpoint = url or (f"stdio:{command}" if transport == 'stdio' else "unknown")
780
+ if isinstance(e, ExceptionGroup):
781
+ primary_exception = extract_primary_exception(list(e.exceptions))
782
+ mcp_error = convert_to_mcp_error(primary_exception, endpoint)
783
+ else:
784
+ mcp_error = convert_to_mcp_error(e, endpoint)
785
+ raise mcp_error from e
786
+
787
+
788
+ async def call_tool_and_print(command: str | None,
789
+ url: str | None,
790
+ tool_name: str,
791
+ transport: str,
792
+ args: list[str] | None,
793
+ env: dict[str, str] | None,
794
+ tool_args: dict[str, Any] | None,
795
+ direct: bool,
796
+ auth_redirect_uri: str | None = None,
797
+ auth_user_id: str | None = None,
798
+ auth_scopes: list[str] | None = None) -> str:
799
+ """Call an MCP tool either directly or via the function group and return output.
800
+
801
+ When ``direct`` is True, uses the raw MCP protocol client (bypassing the
802
+ builder). Otherwise, constructs the ``mcp_client`` function group and
803
+ invokes the corresponding function, mirroring workflow configuration.
804
+
805
+ Args:
806
+ command (str | None): For ``stdio`` transport, the command to execute.
807
+ url (str | None): For ``sse`` or ``streamable-http`` transports, the server URL.
808
+ tool_name (str): Name of the tool to call.
809
+ transport (str): One of ``'stdio'``, ``'sse'``, or ``'streamable-http'``.
810
+ args (list[str] | None): For ``stdio`` transport, additional command arguments.
811
+ env (dict[str, str] | None): For ``stdio`` transport, environment variables.
812
+ tool_args (dict[str, Any] | None): JSON-serializable arguments passed to the tool.
813
+ direct (bool): If True, bypass WorkflowBuilder and use direct MCP client.
814
+
815
+ Returns:
816
+ str: Stringified tool output suitable for terminal display. May be an
817
+ empty string when the MCP client package is not installed and ``direct``
818
+ is False.
819
+
820
+ Raises:
821
+ RuntimeError: If the tool is not found when using the function group.
822
+ MCPError: Propagated from ``call_tool_direct`` when direct mode fails.
823
+ """
824
+ if direct:
825
+ return await call_tool_direct(command, url, tool_name, transport, args, env, tool_args)
826
+
827
+ try:
828
+ from nat.builder.workflow_builder import WorkflowBuilder
829
+ from nat.plugins.mcp.client_impl import MCPClientConfig
830
+ from nat.plugins.mcp.client_impl import MCPServerConfig
831
+ except ImportError:
832
+ click.echo(
833
+ "MCP client functionality requires nvidia-nat-mcp package. Install with: uv pip install nvidia-nat-mcp",
834
+ err=True)
835
+ return ""
836
+
837
+ server_cfg = MCPServerConfig(
838
+ transport=cast(Literal["stdio", "sse", "streamable-http"], transport),
839
+ url=cast(Any, url) if transport in ('sse', 'streamable-http') else None,
840
+ command=command if transport == 'stdio' else None,
841
+ args=args if transport == 'stdio' else None,
842
+ env=env if transport == 'stdio' else None,
843
+ )
844
+
845
+ group_cfg = MCPClientConfig(server=server_cfg)
846
+
847
+ async with WorkflowBuilder() as builder: # type: ignore
848
+ # Add auth provider if url is provided and auth_redirect_uri is given (only for streamable-http)
849
+ if url and transport == 'streamable-http' and auth_redirect_uri:
850
+ try:
851
+ group_cfg = await _create_mcp_client_config(builder,
852
+ server_cfg,
853
+ url,
854
+ transport,
855
+ auth_redirect_uri,
856
+ auth_user_id,
857
+ auth_scopes)
858
+ except ImportError:
859
+ click.echo("[WARNING] MCP OAuth2 authentication requires nvidia-nat-mcp package.", err=True)
860
+
861
+ group = await builder.add_function_group("mcp_client", group_cfg)
862
+ fns = await group.get_accessible_functions()
863
+ full = f"mcp_client.{tool_name}"
864
+ fn = fns.get(full)
865
+ if fn is None:
866
+ raise RuntimeError(f"Tool '{tool_name}' not found")
867
+ # The group exposes a Function that we can invoke with kwargs
868
+ result = await fn.acall_invoke(**(tool_args or {}))
869
+ # Ensure string output for terminal
870
+ return str(result)
871
+
872
+
873
+ @mcp_client_tool_group.command(name="call", help="Call a tool by name with optional arguments.")
874
+ @click.argument('tool_name', nargs=1, required=True)
875
+ @click.option('--direct', is_flag=True, help='Bypass MCPBuilder and use direct MCP protocol')
876
+ @click.option(
877
+ '--url',
878
+ default='http://localhost:9901/mcp',
879
+ show_default=True,
880
+ help='MCP server URL (e.g. http://localhost:8080/mcp for streamable-http, http://localhost:8080/sse for sse)')
881
+ @click.option('--transport',
882
+ type=click.Choice(['sse', 'stdio', 'streamable-http']),
883
+ default='streamable-http',
884
+ show_default=True,
885
+ help='Type of client to use (default: streamable-http, backwards compatible with sse)')
886
+ @click.option('--command', help='For stdio: The command to run (e.g. mcp-server)')
887
+ @click.option('--args', help='For stdio: Additional arguments for the command (space-separated)')
888
+ @click.option('--env', help='For stdio: Environment variables in KEY=VALUE format (space-separated)')
889
+ @click.option('--json-args', default=None, help='Pass tool args as a JSON object string')
890
+ @click.option('--auth',
891
+ is_flag=True,
892
+ help='Enable OAuth2 authentication with default settings (streamable-http only, not with --direct)')
893
+ @click.option('--auth-redirect-uri',
894
+ help='OAuth2 redirect URI for authentication (streamable-http only, not with --direct)')
895
+ @click.option('--auth-user-id', help='User ID for authentication (streamable-http only, not with --direct)')
896
+ @click.option('--auth-scopes', help='OAuth2 scopes (comma-separated, streamable-http only, not with --direct)')
897
+ def mcp_client_tool_call(tool_name: str,
898
+ direct: bool,
899
+ url: str | None,
900
+ transport: str,
901
+ command: str | None,
902
+ args: str | None,
903
+ env: str | None,
904
+ json_args: str | None,
905
+ auth: bool,
906
+ auth_redirect_uri: str | None,
907
+ auth_user_id: str | None,
908
+ auth_scopes: str | None) -> None:
909
+ """Call an MCP tool by name with optional JSON arguments.
910
+
911
+ Validates transport parameters, parses ``--json-args`` into a dictionary,
912
+ invokes the tool (either directly or via the function group), and prints
913
+ the resulting output to stdout. Errors are formatted consistently with
914
+ other MCP CLI commands.
915
+
916
+ Args:
917
+ tool_name (str): Name of the tool to call.
918
+ direct (bool): If True, bypass WorkflowBuilder and use the direct MCP client.
919
+ url (str | None): For ``sse`` or ``streamable-http`` transports, the server URL.
920
+ transport (str): One of ``'stdio'``, ``'sse'``, or ``'streamable-http'``.
921
+ command (str | None): For ``stdio`` transport, the command to execute.
922
+ args (str | None): For ``stdio`` transport, space-separated command arguments.
923
+ env (str | None): For ``stdio`` transport, space-separated ``KEY=VALUE`` pairs.
924
+ json_args (str | None): JSON object string with tool arguments (e.g. '{"q": "hello"}').
925
+ auth_redirect_uri (str | None): redirect URI for auth (streamable-http only, not with --direct)
926
+ auth_user_id (str | None): User ID for authentication (streamable-http only, not with --direct)
927
+ auth_scopes (str | None): OAuth2 scopes (comma-separated, streamable-http only, not with --direct)
928
+
929
+ Examples:
930
+ nat mcp client tool call echo --json-args '{"text": "Hello"}'
931
+ nat mcp client tool call search --direct --url http://localhost:9901/mcp \
932
+ --json-args '{"query": "NVIDIA"}' # Direct mode (no auth)
933
+ nat mcp client tool call run --transport stdio --command mcp-server \
934
+ --args "--flag1 --flag2" --env "ENV1=V1 ENV2=V2" --json-args '{}'
935
+ nat mcp client tool call search --url https://example.com/mcp/ --auth \
936
+ --json-args '{"query": "test"}' # With auth using defaults
937
+ nat mcp client tool call search --url https://example.com/mcp/ \
938
+ --transport streamable-http --json-args '{"query": "test"}' --auth
939
+ """
940
+ # Validate transport args
941
+ if not validate_transport_cli_args(transport, command, args, env):
942
+ return
943
+
944
+ # Parse stdio params
945
+ stdio_args = args.split() if args else []
946
+ stdio_env = dict(var.split('=', 1) for var in env.split()) if env else None
947
+
948
+ # Set auth defaults if --auth flag is used
949
+ auth_redirect_uri, auth_user_id, auth_scopes_list = _set_auth_defaults(
950
+ auth, url, auth_redirect_uri, auth_user_id, auth_scopes
951
+ )
952
+
953
+ # Parse tool args
954
+ arg_obj: dict[str, Any] = {}
955
+ if json_args:
956
+ try:
957
+ parsed = json.loads(json_args)
958
+ if not isinstance(parsed, dict):
959
+ click.echo("[ERROR] --json-args must be a JSON object", err=True)
960
+ return
961
+ arg_obj.update(parsed)
962
+ except json.JSONDecodeError as e:
963
+ click.echo(f"[ERROR] Failed to parse --json-args: {e}", err=True)
964
+ return
965
+
966
+ try:
967
+ output = asyncio.run(
968
+ call_tool_and_print(
969
+ command=command,
970
+ url=url,
971
+ tool_name=tool_name,
972
+ transport=transport,
973
+ args=stdio_args,
974
+ env=stdio_env,
975
+ tool_args=arg_obj,
976
+ direct=direct,
977
+ auth_redirect_uri=auth_redirect_uri,
978
+ auth_user_id=auth_user_id,
979
+ auth_scopes=auth_scopes_list,
980
+ ))
981
+ if output:
982
+ click.echo(output)
983
+ except MCPError as e:
984
+ format_mcp_error(e, include_traceback=False)
985
+ except Exception as e:
986
+ click.echo(f"[ERROR] {e}", err=True)