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.
- aiq/__init__.py +2 -2
- nat/agent/base.py +24 -15
- 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 +79 -47
- nat/agent/react_agent/register.py +41 -21
- nat/agent/reasoning_agent/reasoning_agent.py +11 -9
- nat/agent/register.py +1 -1
- nat/agent/rewoo_agent/agent.py +326 -148
- nat/agent/rewoo_agent/prompt.py +19 -22
- nat/agent/rewoo_agent/register.py +46 -26
- nat/agent/tool_calling_agent/agent.py +84 -28
- nat/agent/tool_calling_agent/register.py +51 -28
- 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 +40 -20
- 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 +9 -5
- nat/builder/context.py +46 -11
- 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/user_interaction_manager.py +2 -2
- nat/builder/workflow.py +13 -1
- nat/builder/workflow_builder.py +281 -76
- 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 +0 -1
- nat/cli/commands/workflow/templates/pyproject.toml.j2 +4 -1
- nat/cli/commands/workflow/templates/register.py.j2 +0 -1
- nat/cli/commands/workflow/workflow_commands.py +9 -13
- nat/cli/entrypoint.py +8 -10
- nat/cli/register_workflow.py +38 -4
- nat/cli/type_registry.py +75 -6
- 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 +10 -10
- 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 +41 -17
- 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/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/nim_embedder.py +1 -1
- nat/embedder/openai_embedder.py +1 -1
- nat/embedder/register.py +0 -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 +2 -2
- 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/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 +1 -1
- 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 +8 -5
- 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 +481 -281
- nat/front_ends/fastapi/job_store.py +518 -99
- nat/front_ends/fastapi/main.py +11 -19
- nat/front_ends/fastapi/message_handler.py +13 -14
- nat/front_ends/fastapi/message_validator.py +17 -19
- 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 +10 -1
- nat/front_ends/mcp/mcp_front_end_plugin.py +45 -13
- nat/front_ends/mcp/mcp_front_end_plugin_worker.py +116 -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 +13 -6
- 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 +4 -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 +1 -1
- 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 +6 -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 +6 -6
- 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/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 +7 -8
- 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 +153 -0
- nat/profiler/parameter_optimization/parameter_selection.py +107 -0
- nat/profiler/parameter_optimization/pareto_visualizer.py +380 -0
- nat/profiler/parameter_optimization/prompt_optimizer.py +384 -0
- nat/profiler/parameter_optimization/update_helpers.py +66 -0
- nat/profiler/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 +3 -2
- nat/runtime/session.py +43 -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/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/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 +4 -3
- nat/utils/type_utils.py +9 -5
- {nvidia_nat-1.3.0.dev2.dist-info → nvidia_nat-1.3.0rc1.dist-info}/METADATA +42 -16
- {nvidia_nat-1.3.0.dev2.dist-info → nvidia_nat-1.3.0rc1.dist-info}/RECORD +230 -189
- {nvidia_nat-1.3.0.dev2.dist-info → nvidia_nat-1.3.0rc1.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
- /nat/{tool/github_tools → agent/prompt_optimizer}/__init__.py +0 -0
- /nat/{tool/mcp → authentication/credential_validator}/__init__.py +0 -0
- {nvidia_nat-1.3.0.dev2.dist-info → nvidia_nat-1.3.0rc1.dist-info}/WHEEL +0 -0
- {nvidia_nat-1.3.0.dev2.dist-info → nvidia_nat-1.3.0rc1.dist-info}/licenses/LICENSE-3rd-party.txt +0 -0
- {nvidia_nat-1.3.0.dev2.dist-info → nvidia_nat-1.3.0rc1.dist-info}/licenses/LICENSE.md +0 -0
- {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)
|