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
@@ -5,6 +5,7 @@ nat = nat.cli.main:run_cli
5
5
  [nat.components]
6
6
  nat_agents = nat.agent.register
7
7
  nat_authentication = nat.authentication.register
8
+ nat_control_flow = nat.control_flow.register
8
9
  nat_embedders = nat.embedder.register
9
10
  nat_evaluators = nat.eval.register
10
11
  nat_llms = nat.llm.register
@@ -1,304 +0,0 @@
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
-
22
- import click
23
- from pydantic import BaseModel
24
-
25
- from nat.tool.mcp.exceptions import MCPError
26
- from nat.tool.mcp.mcp_client import MCPBuilder
27
- from nat.utils.exception_handlers.mcp import format_mcp_error
28
-
29
- # Suppress verbose logs from mcp.client.sse and httpx
30
- logging.getLogger("mcp.client.sse").setLevel(logging.WARNING)
31
- logging.getLogger("httpx").setLevel(logging.WARNING)
32
-
33
- logger = logging.getLogger(__name__)
34
-
35
-
36
- class MCPPingResult(BaseModel):
37
- """Result of an MCP server ping request.
38
-
39
- Attributes:
40
- url (str): The MCP server URL that was pinged
41
- status (str): Health status - 'healthy', 'unhealthy', or 'unknown'
42
- response_time_ms (float | None): Response time in milliseconds, None if request failed completely
43
- error (str | None): Error message if the ping failed, None if successful
44
- """
45
- url: str
46
- status: str
47
- response_time_ms: float | None
48
- error: str | None
49
-
50
-
51
- def format_tool(tool: Any) -> dict[str, str | None]:
52
- """Format an MCP tool into a dictionary for display.
53
-
54
- Extracts name, description, and input schema from various MCP tool object types
55
- and normalizes them into a consistent dictionary format for CLI display.
56
-
57
- Args:
58
- tool (Any): MCPToolClient or raw MCP Tool object (uses Any due to different types)
59
-
60
- Returns:
61
- dict[str, str | None]: Dictionary with name, description, and input_schema as keys
62
- """
63
- name = getattr(tool, 'name', None)
64
- description = getattr(tool, 'description', '')
65
- input_schema = getattr(tool, 'input_schema', None) or getattr(tool, 'inputSchema', None)
66
-
67
- schema_str = None
68
- if input_schema:
69
- if hasattr(input_schema, "schema_json"):
70
- schema_str = input_schema.schema_json(indent=2)
71
- else:
72
- schema_str = str(input_schema)
73
-
74
- return {
75
- "name": name,
76
- "description": description,
77
- "input_schema": schema_str,
78
- }
79
-
80
-
81
- def print_tool(tool_dict: dict[str, str | None], detail: bool = False) -> None:
82
- """Print a formatted tool to the console with optional detailed information.
83
-
84
- Outputs tool information in a user-friendly format to stdout. When detail=True
85
- or when description/schema are available, shows full information with separator.
86
-
87
- Args:
88
- tool_dict (dict[str, str | None]): Dictionary containing tool information with name, description, and
89
- input_schema as keys
90
- detail (bool, optional): Whether to force detailed output. Defaults to False.
91
- """
92
- click.echo(f"Tool: {tool_dict.get('name', 'Unknown')}")
93
- if detail or tool_dict.get('input_schema') or tool_dict.get('description'):
94
- click.echo(f"Description: {tool_dict.get('description', 'No description available')}")
95
- if tool_dict.get("input_schema"):
96
- click.echo("Input Schema:")
97
- click.echo(tool_dict.get("input_schema"))
98
- else:
99
- click.echo("Input Schema: None")
100
- click.echo("-" * 60)
101
-
102
-
103
- async def list_tools_and_schemas(url: str, tool_name: str | None = None) -> list[dict[str, str | None]]:
104
- """List MCP tools using MCPBuilder with structured exception handling.
105
-
106
- Args:
107
- url (str): MCP server URL to connect to
108
- tool_name (str | None, optional): Specific tool name to retrieve.
109
- If None, retrieves all available tools. Defaults to None.
110
-
111
- Returns:
112
- list[dict[str, str | None]]: List of formatted tool dictionaries, each containing name, description, and
113
- input_schema as keys
114
-
115
- Raises:
116
- MCPError: Caught internally and logged, returns empty list instead
117
- """
118
- builder = MCPBuilder(url=url)
119
- try:
120
- if tool_name:
121
- tool = await builder.get_tool(tool_name)
122
- return [format_tool(tool)]
123
- tools = await builder.get_tools()
124
- return [format_tool(tool) for tool in tools.values()]
125
- except MCPError as e:
126
- format_mcp_error(e, include_traceback=False)
127
- return []
128
-
129
-
130
- async def list_tools_direct(url: str, tool_name: str | None = None) -> list[dict[str, str | None]]:
131
- """List MCP tools using direct MCP protocol with exception conversion.
132
-
133
- Bypasses MCPBuilder and uses raw MCP ClientSession and SSE client directly.
134
- Converts raw exceptions to structured MCPErrors for consistent user experience.
135
- Used when --direct flag is specified in CLI.
136
-
137
- Args:
138
- url (str): MCP server URL to connect to
139
- tool_name (str | None, optional): Specific tool name to retrieve.
140
- If None, retrieves all available tools. Defaults to None.
141
-
142
- Returns:
143
- list[dict[str, str | None]]: List of formatted tool dictionaries, each containing name, description, and
144
- input_schema as keys
145
-
146
- Note:
147
- This function handles ExceptionGroup by extracting the most relevant exception
148
- and converting it to MCPError for consistent error reporting.
149
- """
150
- from mcp import ClientSession
151
- from mcp.client.sse import sse_client
152
-
153
- try:
154
- async with sse_client(url=url) as (read, write):
155
- async with ClientSession(read, write) as session:
156
- await session.initialize()
157
- response = await session.list_tools()
158
-
159
- tools = []
160
- for tool in response.tools:
161
- if tool_name:
162
- if tool.name == tool_name:
163
- return [format_tool(tool)]
164
- else:
165
- tools.append(format_tool(tool))
166
- if tool_name and not tools:
167
- click.echo(f"[INFO] Tool '{tool_name}' not found.")
168
- return tools
169
- except Exception as e:
170
- # Convert raw exceptions to structured MCPError for consistency
171
- from nat.utils.exception_handlers.mcp import convert_to_mcp_error
172
- from nat.utils.exception_handlers.mcp import extract_primary_exception
173
-
174
- if isinstance(e, ExceptionGroup): # noqa: F821
175
- primary_exception = extract_primary_exception(list(e.exceptions))
176
- mcp_error = convert_to_mcp_error(primary_exception, url)
177
- else:
178
- mcp_error = convert_to_mcp_error(e, url)
179
-
180
- format_mcp_error(mcp_error, include_traceback=False)
181
- return []
182
-
183
-
184
- async def ping_mcp_server(url: str, timeout: int) -> MCPPingResult:
185
- """Ping an MCP server to check if it's responsive.
186
-
187
- Args:
188
- url (str): MCP server URL to ping
189
- timeout (int): Timeout in seconds for the ping request
190
-
191
- Returns:
192
- MCPPingResult: Structured result with status, response_time, and any error info
193
- """
194
- from mcp.client.session import ClientSession
195
- from mcp.client.sse import sse_client
196
-
197
- async def _ping_operation():
198
- async with sse_client(url) as (read, write):
199
- async with ClientSession(read, write) as session:
200
- # Initialize the session
201
- await session.initialize()
202
-
203
- # Record start time just before ping
204
- start_time = time.time()
205
- # Send ping request
206
- await session.send_ping()
207
-
208
- end_time = time.time()
209
- response_time_ms = round((end_time - start_time) * 1000, 2)
210
-
211
- return MCPPingResult(url=url, status="healthy", response_time_ms=response_time_ms, error=None)
212
-
213
- try:
214
- # Apply timeout to the entire ping operation
215
- return await asyncio.wait_for(_ping_operation(), timeout=timeout)
216
-
217
- except asyncio.TimeoutError:
218
- return MCPPingResult(url=url,
219
- status="unreachable",
220
- response_time_ms=None,
221
- error=f"Timeout after {timeout} seconds")
222
-
223
- except Exception as e:
224
- return MCPPingResult(url=url, status="unhealthy", response_time_ms=None, error=str(e))
225
-
226
-
227
- @click.group(invoke_without_command=True, help="List tool names (default), or show details with --detail or --tool.")
228
- @click.option('--direct', is_flag=True, help='Bypass MCPBuilder and use direct MCP protocol')
229
- @click.option('--url', default='http://localhost:9901/sse', show_default=True, help='MCP server URL')
230
- @click.option('--tool', default=None, help='Get details for a specific tool by name')
231
- @click.option('--detail', is_flag=True, help='Show full details for all tools')
232
- @click.option('--json-output', is_flag=True, help='Output tool metadata in JSON format')
233
- @click.pass_context
234
- def list_mcp(ctx: click.Context, direct: bool, url: str, tool: str | None, detail: bool, json_output: bool) -> None:
235
- """List MCP tool names (default) or show detailed tool information.
236
-
237
- Use --detail for full output including descriptions and input schemas.
238
- If --tool is provided, always shows full output for that specific tool.
239
- Use --direct to bypass MCPBuilder and use raw MCP protocol.
240
- Use --json-output to get structured JSON data instead of formatted text.
241
-
242
- Args:
243
- ctx (click.Context): Click context object for command invocation
244
- direct (bool): Whether to bypass MCPBuilder and use direct MCP protocol
245
- url (str): MCP server URL to connect to (default: http://localhost:9901/sse)
246
- tool (str | None): Optional specific tool name to retrieve detailed info for
247
- detail (bool): Whether to show full details (description + schema) for all tools
248
- json_output (bool): Whether to output tool metadata in JSON format instead of text
249
-
250
- Examples:
251
- nat info mcp # List tool names only
252
- nat info mcp --detail # Show all tools with full details
253
- nat info mcp --tool my_tool # Show details for specific tool
254
- nat info mcp --json-output # Get JSON format output
255
- nat info mcp --direct --url http://... # Use direct protocol with custom URL
256
- """
257
- if ctx.invoked_subcommand is not None:
258
- return
259
- fetcher = list_tools_direct if direct else list_tools_and_schemas
260
- tools = asyncio.run(fetcher(url, tool))
261
-
262
- if json_output:
263
- click.echo(json.dumps(tools, indent=2))
264
- elif tool:
265
- for tool_dict in tools:
266
- print_tool(tool_dict, detail=True)
267
- elif detail:
268
- for tool_dict in tools:
269
- print_tool(tool_dict, detail=True)
270
- else:
271
- for tool_dict in tools:
272
- click.echo(tool_dict.get('name', 'Unknown tool'))
273
-
274
-
275
- @list_mcp.command()
276
- @click.option('--url', default='http://localhost:9901/sse', show_default=True, help='MCP server URL')
277
- @click.option('--timeout', default=60, show_default=True, help='Timeout in seconds for ping request')
278
- @click.option('--json-output', is_flag=True, help='Output ping result in JSON format')
279
- def ping(url: str, timeout: int, json_output: bool) -> None:
280
- """Ping an MCP server to check if it's responsive.
281
-
282
- This command sends a ping request to the MCP server and measures the response time.
283
- It's useful for health checks and monitoring server availability.
284
-
285
- Args:
286
- url (str): MCP server URL to ping (default: http://localhost:9901/sse)
287
- timeout (int): Timeout in seconds for the ping request (default: 60)
288
- json_output (bool): Whether to output the result in JSON format
289
-
290
- Examples:
291
- nat info mcp ping # Ping default server
292
- nat info mcp ping --url http://custom-server:9901/sse # Ping custom server
293
- nat info mcp ping --timeout 10 # Use 10 second timeout
294
- nat info mcp ping --json-output # Get JSON format output
295
- """
296
- result = asyncio.run(ping_mcp_server(url, timeout))
297
-
298
- if json_output:
299
- click.echo(result.model_dump_json(indent=2))
300
- else:
301
- if result.status == "healthy":
302
- click.echo(f"Server at {result.url} is healthy (response time: {result.response_time_ms}ms)")
303
- else:
304
- click.echo(f"Server at {result.url} {result.status}: {result.error}")
@@ -1,133 +0,0 @@
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
- from pydantic import BaseModel
17
- from pydantic import Field
18
-
19
- from nat.builder.builder import Builder
20
- from nat.builder.function_info import FunctionInfo
21
- from nat.cli.register_workflow import register_function
22
- from nat.data_models.function import FunctionBaseConfig
23
-
24
-
25
- class GithubCommitCodeModel(BaseModel):
26
- branch: str = Field(description="The branch of the remote repo to which the code will be committed")
27
- commit_msg: str = Field(description="Message with which the code will be committed to the remote repo")
28
- local_path: str = Field(description="Local filepath of the file that has been updated and "
29
- "needs to be committed to the remote repo")
30
- remote_path: str = Field(description="Remote filepath of the updated file in GitHub. Path is relative to "
31
- "root of current repository")
32
-
33
-
34
- class GithubCommitCodeModelList(BaseModel):
35
- updated_files: list[GithubCommitCodeModel] = Field(description=("A list of local filepaths and commit messages"))
36
-
37
-
38
- class GithubCommitCodeConfig(FunctionBaseConfig, name="github_commit_code_tool"):
39
- """
40
- Tool that commits and pushes modified code to a remote GitHub repository asynchronously.
41
- """
42
- repo_name: str = Field(description="The repository name in the format 'owner/repo'")
43
- local_repo_dir: str = Field(description="Absolute path to the root of the repo, cloned locally")
44
- timeout: int = Field(default=300, description="The timeout configuration to use when sending requests.")
45
-
46
-
47
- @register_function(config_type=GithubCommitCodeConfig)
48
- async def commit_code_async(config: GithubCommitCodeConfig, builder: Builder):
49
- """
50
- Commits and pushes modified code to a remote GitHub repository asynchronously.
51
-
52
- """
53
- import json
54
- import os
55
-
56
- import httpx
57
-
58
- github_pat = os.getenv("GITHUB_PAT")
59
- if not github_pat:
60
- raise ValueError("GITHUB_PAT environment variable must be set")
61
-
62
- # define the headers for the payload request
63
- headers = {"Authorization": f"Bearer {github_pat}", "Accept": "application/vnd.github+json"}
64
-
65
- async def _github_commit_code(updated_files) -> list:
66
- results = []
67
- async with httpx.AsyncClient(timeout=config.timeout) as client:
68
- for file_ in updated_files:
69
- branch = file_.branch
70
- commit_msg = file_.commit_msg
71
- local_path = file_.local_path
72
- remote_path = file_.remote_path
73
-
74
- # Read content from the local file
75
- local_path = os.path.join(config.local_repo_dir, local_path)
76
- with open(local_path, 'r', encoding='utf-8', errors='ignore') as f:
77
- content = f.read()
78
-
79
- # Step 1. Create a blob with the updated contents of the file
80
- blob_url = f'https://api.github.com/repos/{config.repo_name}/git/blobs'
81
- blob_data = {'content': content, 'encoding': 'utf-8'}
82
- blob_response = await client.request("POST", blob_url, json=blob_data, headers=headers)
83
- blob_response.raise_for_status()
84
- blob_sha = blob_response.json()['sha']
85
-
86
- # Step 2: Get the base tree SHA. The commit will be pushed to this ref node in the Git graph
87
- ref_url = f'https://api.github.com/repos/{config.repo_name}/git/refs/heads/{branch}'
88
- ref_response = await client.request("GET", ref_url, headers=headers)
89
- ref_response.raise_for_status()
90
- base_tree_sha = ref_response.json()['object']['sha']
91
-
92
- # Step 3. Create an updated tree (Git graph) with the new blob
93
- tree_url = f'https://api.github.com/repos/{config.repo_name}/git/trees'
94
- tree_data = {
95
- 'base_tree': base_tree_sha,
96
- 'tree': [{
97
- 'path': remote_path, 'mode': '100644', 'type': 'blob', 'sha': blob_sha
98
- }]
99
- }
100
- tree_response = await client.request("POST", tree_url, json=tree_data, headers=headers)
101
- tree_response.raise_for_status()
102
- tree_sha = tree_response.json()['sha']
103
-
104
- # Step 4: Create a commit
105
- commit_url = f'https://api.github.com/repos/{config.repo_name}/git/commits'
106
- commit_data = {'message': commit_msg, 'tree': tree_sha, 'parents': [base_tree_sha]}
107
- commit_response = await client.request("POST", commit_url, json=commit_data, headers=headers)
108
- commit_response.raise_for_status()
109
- commit_sha = commit_response.json()['sha']
110
-
111
- # Step 5: Update the reference in the Git graph
112
- update_ref_url = f'https://api.github.com/repos/{config.repo_name}/git/refs/heads/{branch}'
113
- update_ref_data = {'sha': commit_sha}
114
- update_ref_response = await client.request("PATCH",
115
- update_ref_url,
116
- json=update_ref_data,
117
- headers=headers)
118
- update_ref_response.raise_for_status()
119
-
120
- payload_responses = {
121
- 'blob_resp': blob_response.json(),
122
- 'original_tree_ref': tree_response.json(),
123
- 'commit_resp': commit_response.json(),
124
- 'updated_tree_ref_resp': update_ref_response.json()
125
- }
126
- results.append(payload_responses)
127
-
128
- return json.dumps(results)
129
-
130
- yield FunctionInfo.from_fn(_github_commit_code,
131
- description=(f"Commits and pushes modified code to a "
132
- f"GitHub repository in the repo named {config.repo_name}"),
133
- input_schema=GithubCommitCodeModelList)
@@ -1,87 +0,0 @@
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
- from pydantic import BaseModel
17
- from pydantic import Field
18
-
19
- from nat.builder.builder import Builder
20
- from nat.builder.function_info import FunctionInfo
21
- from nat.cli.register_workflow import register_function
22
- from nat.data_models.function import FunctionBaseConfig
23
-
24
-
25
- class GithubCreateIssueModel(BaseModel):
26
- title: str = Field(description="The title of the GitHub Issue")
27
- body: str = Field(description="The body of the GitHub Issue")
28
-
29
-
30
- class GithubCreateIssueModelList(BaseModel):
31
- issues: list[GithubCreateIssueModel] = Field(description=("A list of GitHub issues, "
32
- "each with a title and a body"))
33
-
34
-
35
- class GithubCreateIssueToolConfig(FunctionBaseConfig, name="github_create_issue_tool"):
36
- """
37
- Tool that creates an issue in a GitHub repository asynchronously.
38
- """
39
- repo_name: str = Field(description="The repository name in the format 'owner/repo'")
40
- timeout: int = Field(default=300, description="The timeout configuration to use when sending requests.")
41
-
42
-
43
- @register_function(config_type=GithubCreateIssueToolConfig)
44
- async def create_github_issue_async(config: GithubCreateIssueToolConfig, builder: Builder):
45
- """
46
- Creates an issue in a GitHub repository asynchronously.
47
- """
48
- import json
49
- import os
50
-
51
- import httpx
52
-
53
- github_pat = os.getenv("GITHUB_PAT")
54
- if not github_pat:
55
- raise ValueError("GITHUB_PAT environment variable must be set")
56
-
57
- url = f"https://api.github.com/repos/{config.repo_name}/issues"
58
-
59
- # define the headers for the payload request
60
- headers = {"Authorization": f"Bearer {github_pat}", "Accept": "application/vnd.github+json"}
61
-
62
- async def _github_post_issue(issues) -> list:
63
- results = []
64
- async with httpx.AsyncClient(timeout=config.timeout) as client:
65
- for issue in issues:
66
- # define the payload body
67
- payload = issue.dict(exclude_unset=True)
68
-
69
- response = await client.request("POST", url, json=payload, headers=headers)
70
-
71
- # Raise an exception for HTTP errors
72
- response.raise_for_status()
73
-
74
- # Parse and return the response JSON
75
- try:
76
- result = response.json()
77
- results.append(result)
78
-
79
- except ValueError as e:
80
- raise ValueError("The API response is not valid JSON.") from e
81
-
82
- return json.dumps(results)
83
-
84
- yield FunctionInfo.from_fn(_github_post_issue,
85
- description=(f"Creates a GitHub issue in the "
86
- f"repo named {config.repo_name}"),
87
- input_schema=GithubCreateIssueModelList)
@@ -1,106 +0,0 @@
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
- from pydantic import BaseModel
17
- from pydantic import Field
18
-
19
- from nat.builder.builder import Builder
20
- from nat.builder.function_info import FunctionInfo
21
- from nat.cli.register_workflow import register_function
22
- from nat.data_models.function import FunctionBaseConfig
23
-
24
-
25
- class GithubCreatePullModel(BaseModel):
26
- title: str = Field(description="Title of the pull request")
27
- body: str = Field(description="Description of the pull request")
28
- source_branch: str = Field(description="The name of the branch containing your changes")
29
- target_branch: str = Field(description="The name of the branch you want to merge into")
30
- assignees: list[str] | None = Field([],
31
- description="List of GitHub usernames to assign to the PR. "
32
- "Always the current user")
33
- reviewers: list[str] | None = Field([], description="List of GitHub usernames to request review from")
34
-
35
-
36
- class GithubCreatePullList(BaseModel):
37
- pull_details: GithubCreatePullModel = Field(description=("A list of params used for creating the PR in GitHub"))
38
-
39
-
40
- class GithubCreatePullConfig(FunctionBaseConfig, name="github_create_pull_tool"):
41
- """
42
- Tool that creates a pull request in a GitHub repository asynchronously with assignees and reviewers.
43
- """
44
- repo_name: str = Field(description="The repository name in the format 'owner/repo'")
45
- timeout: int = Field(default=300, description="The timeout configuration to use when sending requests.")
46
-
47
-
48
- @register_function(config_type=GithubCreatePullConfig)
49
- async def create_pull_request_async(config: GithubCreatePullConfig, builder: Builder):
50
- """
51
- Creates a pull request in a GitHub repository asynchronously with assignees and reviewers.
52
-
53
- """
54
- import json
55
- import os
56
-
57
- import httpx
58
-
59
- github_pat = os.getenv("GITHUB_PAT")
60
- if not github_pat:
61
- raise ValueError("GITHUB_PAT environment variable must be set")
62
-
63
- headers = {"Authorization": f"Bearer {github_pat}", "Accept": "application/vnd.github+json"}
64
-
65
- async def _github_create_pull(pull_details: GithubCreatePullList) -> str:
66
- results = []
67
- async with httpx.AsyncClient(timeout=config.timeout) as client:
68
- # Create pull request
69
- pr_url = f'https://api.github.com/repos/{config.repo_name}/pulls'
70
- pr_data = {
71
- 'title': pull_details.title,
72
- 'body': pull_details.body,
73
- 'head': pull_details.source_branch,
74
- 'base': pull_details.target_branch
75
- }
76
-
77
- pr_response = await client.request("POST", pr_url, json=pr_data, headers=headers)
78
- pr_response.raise_for_status()
79
- pr_number = pr_response.json()['number']
80
-
81
- # Add assignees if provided
82
- if pull_details.assignees:
83
- assignees_url = f'https://api.github.com/repos/{config.repo_name}/issues/{pr_number}/assignees'
84
- assignees_data = {'assignees': pull_details.assignees}
85
- assignees_response = await client.request("POST", assignees_url, json=assignees_data, headers=headers)
86
- assignees_response.raise_for_status()
87
-
88
- # Request reviewers if provided
89
- if pull_details.reviewers:
90
- reviewers_url = f'https://api.github.com/repos/{config.repo_name}/pulls/{pr_number}/requested_reviewers'
91
- reviewers_data = {'reviewers': pull_details.reviewers}
92
- reviewers_response = await client.request("POST", reviewers_url, json=reviewers_data, headers=headers)
93
- reviewers_response.raise_for_status()
94
-
95
- results.append({
96
- 'pull_request': pr_response.json(),
97
- 'assignees': assignees_response.json() if pull_details.assignees else None,
98
- 'reviewers': reviewers_response.json() if pull_details.reviewers else None
99
- })
100
-
101
- return json.dumps(results)
102
-
103
- yield FunctionInfo.from_fn(_github_create_pull,
104
- description=(f"Creates a pull request with assignees and reviewers in the "
105
- f"GitHub repository named {config.repo_name}"),
106
- input_schema=GithubCreatePullList)