nvidia-nat 1.2.0__py3-none-any.whl

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (435) hide show
  1. aiq/__init__.py +66 -0
  2. nat/agent/__init__.py +0 -0
  3. nat/agent/base.py +256 -0
  4. nat/agent/dual_node.py +67 -0
  5. nat/agent/react_agent/__init__.py +0 -0
  6. nat/agent/react_agent/agent.py +363 -0
  7. nat/agent/react_agent/output_parser.py +104 -0
  8. nat/agent/react_agent/prompt.py +44 -0
  9. nat/agent/react_agent/register.py +149 -0
  10. nat/agent/reasoning_agent/__init__.py +0 -0
  11. nat/agent/reasoning_agent/reasoning_agent.py +225 -0
  12. nat/agent/register.py +23 -0
  13. nat/agent/rewoo_agent/__init__.py +0 -0
  14. nat/agent/rewoo_agent/agent.py +415 -0
  15. nat/agent/rewoo_agent/prompt.py +110 -0
  16. nat/agent/rewoo_agent/register.py +157 -0
  17. nat/agent/tool_calling_agent/__init__.py +0 -0
  18. nat/agent/tool_calling_agent/agent.py +119 -0
  19. nat/agent/tool_calling_agent/register.py +106 -0
  20. nat/authentication/__init__.py +14 -0
  21. nat/authentication/api_key/__init__.py +14 -0
  22. nat/authentication/api_key/api_key_auth_provider.py +96 -0
  23. nat/authentication/api_key/api_key_auth_provider_config.py +124 -0
  24. nat/authentication/api_key/register.py +26 -0
  25. nat/authentication/exceptions/__init__.py +14 -0
  26. nat/authentication/exceptions/api_key_exceptions.py +38 -0
  27. nat/authentication/http_basic_auth/__init__.py +0 -0
  28. nat/authentication/http_basic_auth/http_basic_auth_provider.py +81 -0
  29. nat/authentication/http_basic_auth/register.py +30 -0
  30. nat/authentication/interfaces.py +93 -0
  31. nat/authentication/oauth2/__init__.py +14 -0
  32. nat/authentication/oauth2/oauth2_auth_code_flow_provider.py +107 -0
  33. nat/authentication/oauth2/oauth2_auth_code_flow_provider_config.py +39 -0
  34. nat/authentication/oauth2/register.py +25 -0
  35. nat/authentication/register.py +21 -0
  36. nat/builder/__init__.py +0 -0
  37. nat/builder/builder.py +285 -0
  38. nat/builder/component_utils.py +316 -0
  39. nat/builder/context.py +270 -0
  40. nat/builder/embedder.py +24 -0
  41. nat/builder/eval_builder.py +161 -0
  42. nat/builder/evaluator.py +29 -0
  43. nat/builder/framework_enum.py +24 -0
  44. nat/builder/front_end.py +73 -0
  45. nat/builder/function.py +344 -0
  46. nat/builder/function_base.py +380 -0
  47. nat/builder/function_info.py +627 -0
  48. nat/builder/intermediate_step_manager.py +174 -0
  49. nat/builder/llm.py +25 -0
  50. nat/builder/retriever.py +25 -0
  51. nat/builder/user_interaction_manager.py +78 -0
  52. nat/builder/workflow.py +148 -0
  53. nat/builder/workflow_builder.py +1117 -0
  54. nat/cli/__init__.py +14 -0
  55. nat/cli/cli_utils/__init__.py +0 -0
  56. nat/cli/cli_utils/config_override.py +231 -0
  57. nat/cli/cli_utils/validation.py +37 -0
  58. nat/cli/commands/__init__.py +0 -0
  59. nat/cli/commands/configure/__init__.py +0 -0
  60. nat/cli/commands/configure/channel/__init__.py +0 -0
  61. nat/cli/commands/configure/channel/add.py +28 -0
  62. nat/cli/commands/configure/channel/channel.py +34 -0
  63. nat/cli/commands/configure/channel/remove.py +30 -0
  64. nat/cli/commands/configure/channel/update.py +30 -0
  65. nat/cli/commands/configure/configure.py +33 -0
  66. nat/cli/commands/evaluate.py +139 -0
  67. nat/cli/commands/info/__init__.py +14 -0
  68. nat/cli/commands/info/info.py +37 -0
  69. nat/cli/commands/info/list_channels.py +32 -0
  70. nat/cli/commands/info/list_components.py +129 -0
  71. nat/cli/commands/info/list_mcp.py +304 -0
  72. nat/cli/commands/registry/__init__.py +14 -0
  73. nat/cli/commands/registry/publish.py +88 -0
  74. nat/cli/commands/registry/pull.py +118 -0
  75. nat/cli/commands/registry/registry.py +36 -0
  76. nat/cli/commands/registry/remove.py +108 -0
  77. nat/cli/commands/registry/search.py +155 -0
  78. nat/cli/commands/sizing/__init__.py +14 -0
  79. nat/cli/commands/sizing/calc.py +297 -0
  80. nat/cli/commands/sizing/sizing.py +27 -0
  81. nat/cli/commands/start.py +246 -0
  82. nat/cli/commands/uninstall.py +81 -0
  83. nat/cli/commands/validate.py +47 -0
  84. nat/cli/commands/workflow/__init__.py +14 -0
  85. nat/cli/commands/workflow/templates/__init__.py.j2 +0 -0
  86. nat/cli/commands/workflow/templates/config.yml.j2 +16 -0
  87. nat/cli/commands/workflow/templates/pyproject.toml.j2 +22 -0
  88. nat/cli/commands/workflow/templates/register.py.j2 +5 -0
  89. nat/cli/commands/workflow/templates/workflow.py.j2 +36 -0
  90. nat/cli/commands/workflow/workflow.py +37 -0
  91. nat/cli/commands/workflow/workflow_commands.py +317 -0
  92. nat/cli/entrypoint.py +135 -0
  93. nat/cli/main.py +57 -0
  94. nat/cli/register_workflow.py +488 -0
  95. nat/cli/type_registry.py +1000 -0
  96. nat/data_models/__init__.py +14 -0
  97. nat/data_models/api_server.py +716 -0
  98. nat/data_models/authentication.py +231 -0
  99. nat/data_models/common.py +171 -0
  100. nat/data_models/component.py +58 -0
  101. nat/data_models/component_ref.py +168 -0
  102. nat/data_models/config.py +410 -0
  103. nat/data_models/dataset_handler.py +169 -0
  104. nat/data_models/discovery_metadata.py +305 -0
  105. nat/data_models/embedder.py +27 -0
  106. nat/data_models/evaluate.py +127 -0
  107. nat/data_models/evaluator.py +26 -0
  108. nat/data_models/front_end.py +26 -0
  109. nat/data_models/function.py +30 -0
  110. nat/data_models/function_dependencies.py +72 -0
  111. nat/data_models/interactive.py +246 -0
  112. nat/data_models/intermediate_step.py +302 -0
  113. nat/data_models/invocation_node.py +38 -0
  114. nat/data_models/llm.py +27 -0
  115. nat/data_models/logging.py +26 -0
  116. nat/data_models/memory.py +27 -0
  117. nat/data_models/object_store.py +44 -0
  118. nat/data_models/profiler.py +54 -0
  119. nat/data_models/registry_handler.py +26 -0
  120. nat/data_models/retriever.py +30 -0
  121. nat/data_models/retry_mixin.py +35 -0
  122. nat/data_models/span.py +190 -0
  123. nat/data_models/step_adaptor.py +64 -0
  124. nat/data_models/streaming.py +33 -0
  125. nat/data_models/swe_bench_model.py +54 -0
  126. nat/data_models/telemetry_exporter.py +26 -0
  127. nat/data_models/ttc_strategy.py +30 -0
  128. nat/embedder/__init__.py +0 -0
  129. nat/embedder/nim_embedder.py +59 -0
  130. nat/embedder/openai_embedder.py +43 -0
  131. nat/embedder/register.py +22 -0
  132. nat/eval/__init__.py +14 -0
  133. nat/eval/config.py +60 -0
  134. nat/eval/dataset_handler/__init__.py +0 -0
  135. nat/eval/dataset_handler/dataset_downloader.py +106 -0
  136. nat/eval/dataset_handler/dataset_filter.py +52 -0
  137. nat/eval/dataset_handler/dataset_handler.py +367 -0
  138. nat/eval/evaluate.py +510 -0
  139. nat/eval/evaluator/__init__.py +14 -0
  140. nat/eval/evaluator/base_evaluator.py +77 -0
  141. nat/eval/evaluator/evaluator_model.py +45 -0
  142. nat/eval/intermediate_step_adapter.py +99 -0
  143. nat/eval/rag_evaluator/__init__.py +0 -0
  144. nat/eval/rag_evaluator/evaluate.py +178 -0
  145. nat/eval/rag_evaluator/register.py +143 -0
  146. nat/eval/register.py +23 -0
  147. nat/eval/remote_workflow.py +133 -0
  148. nat/eval/runners/__init__.py +14 -0
  149. nat/eval/runners/config.py +39 -0
  150. nat/eval/runners/multi_eval_runner.py +54 -0
  151. nat/eval/runtime_event_subscriber.py +52 -0
  152. nat/eval/swe_bench_evaluator/__init__.py +0 -0
  153. nat/eval/swe_bench_evaluator/evaluate.py +215 -0
  154. nat/eval/swe_bench_evaluator/register.py +36 -0
  155. nat/eval/trajectory_evaluator/__init__.py +0 -0
  156. nat/eval/trajectory_evaluator/evaluate.py +75 -0
  157. nat/eval/trajectory_evaluator/register.py +40 -0
  158. nat/eval/tunable_rag_evaluator/__init__.py +0 -0
  159. nat/eval/tunable_rag_evaluator/evaluate.py +245 -0
  160. nat/eval/tunable_rag_evaluator/register.py +52 -0
  161. nat/eval/usage_stats.py +41 -0
  162. nat/eval/utils/__init__.py +0 -0
  163. nat/eval/utils/output_uploader.py +140 -0
  164. nat/eval/utils/tqdm_position_registry.py +40 -0
  165. nat/eval/utils/weave_eval.py +184 -0
  166. nat/experimental/__init__.py +0 -0
  167. nat/experimental/decorators/__init__.py +0 -0
  168. nat/experimental/decorators/experimental_warning_decorator.py +134 -0
  169. nat/experimental/test_time_compute/__init__.py +0 -0
  170. nat/experimental/test_time_compute/editing/__init__.py +0 -0
  171. nat/experimental/test_time_compute/editing/iterative_plan_refinement_editor.py +147 -0
  172. nat/experimental/test_time_compute/editing/llm_as_a_judge_editor.py +204 -0
  173. nat/experimental/test_time_compute/editing/motivation_aware_summarization.py +107 -0
  174. nat/experimental/test_time_compute/functions/__init__.py +0 -0
  175. nat/experimental/test_time_compute/functions/execute_score_select_function.py +105 -0
  176. nat/experimental/test_time_compute/functions/plan_select_execute_function.py +224 -0
  177. nat/experimental/test_time_compute/functions/ttc_tool_orchestration_function.py +205 -0
  178. nat/experimental/test_time_compute/functions/ttc_tool_wrapper_function.py +146 -0
  179. nat/experimental/test_time_compute/models/__init__.py +0 -0
  180. nat/experimental/test_time_compute/models/editor_config.py +132 -0
  181. nat/experimental/test_time_compute/models/scoring_config.py +112 -0
  182. nat/experimental/test_time_compute/models/search_config.py +120 -0
  183. nat/experimental/test_time_compute/models/selection_config.py +154 -0
  184. nat/experimental/test_time_compute/models/stage_enums.py +43 -0
  185. nat/experimental/test_time_compute/models/strategy_base.py +66 -0
  186. nat/experimental/test_time_compute/models/tool_use_config.py +41 -0
  187. nat/experimental/test_time_compute/models/ttc_item.py +48 -0
  188. nat/experimental/test_time_compute/register.py +36 -0
  189. nat/experimental/test_time_compute/scoring/__init__.py +0 -0
  190. nat/experimental/test_time_compute/scoring/llm_based_agent_scorer.py +168 -0
  191. nat/experimental/test_time_compute/scoring/llm_based_plan_scorer.py +168 -0
  192. nat/experimental/test_time_compute/scoring/motivation_aware_scorer.py +111 -0
  193. nat/experimental/test_time_compute/search/__init__.py +0 -0
  194. nat/experimental/test_time_compute/search/multi_llm_planner.py +128 -0
  195. nat/experimental/test_time_compute/search/multi_query_retrieval_search.py +122 -0
  196. nat/experimental/test_time_compute/search/single_shot_multi_plan_planner.py +128 -0
  197. nat/experimental/test_time_compute/selection/__init__.py +0 -0
  198. nat/experimental/test_time_compute/selection/best_of_n_selector.py +63 -0
  199. nat/experimental/test_time_compute/selection/llm_based_agent_output_selector.py +131 -0
  200. nat/experimental/test_time_compute/selection/llm_based_output_merging_selector.py +159 -0
  201. nat/experimental/test_time_compute/selection/llm_based_plan_selector.py +128 -0
  202. nat/experimental/test_time_compute/selection/threshold_selector.py +58 -0
  203. nat/front_ends/__init__.py +14 -0
  204. nat/front_ends/console/__init__.py +14 -0
  205. nat/front_ends/console/authentication_flow_handler.py +233 -0
  206. nat/front_ends/console/console_front_end_config.py +32 -0
  207. nat/front_ends/console/console_front_end_plugin.py +96 -0
  208. nat/front_ends/console/register.py +25 -0
  209. nat/front_ends/cron/__init__.py +14 -0
  210. nat/front_ends/fastapi/__init__.py +14 -0
  211. nat/front_ends/fastapi/auth_flow_handlers/__init__.py +0 -0
  212. nat/front_ends/fastapi/auth_flow_handlers/http_flow_handler.py +27 -0
  213. nat/front_ends/fastapi/auth_flow_handlers/websocket_flow_handler.py +107 -0
  214. nat/front_ends/fastapi/fastapi_front_end_config.py +241 -0
  215. nat/front_ends/fastapi/fastapi_front_end_controller.py +68 -0
  216. nat/front_ends/fastapi/fastapi_front_end_plugin.py +116 -0
  217. nat/front_ends/fastapi/fastapi_front_end_plugin_worker.py +1087 -0
  218. nat/front_ends/fastapi/html_snippets/__init__.py +14 -0
  219. nat/front_ends/fastapi/html_snippets/auth_code_grant_success.py +35 -0
  220. nat/front_ends/fastapi/intermediate_steps_subscriber.py +80 -0
  221. nat/front_ends/fastapi/job_store.py +183 -0
  222. nat/front_ends/fastapi/main.py +72 -0
  223. nat/front_ends/fastapi/message_handler.py +320 -0
  224. nat/front_ends/fastapi/message_validator.py +352 -0
  225. nat/front_ends/fastapi/register.py +25 -0
  226. nat/front_ends/fastapi/response_helpers.py +195 -0
  227. nat/front_ends/fastapi/step_adaptor.py +319 -0
  228. nat/front_ends/mcp/__init__.py +14 -0
  229. nat/front_ends/mcp/mcp_front_end_config.py +36 -0
  230. nat/front_ends/mcp/mcp_front_end_plugin.py +81 -0
  231. nat/front_ends/mcp/mcp_front_end_plugin_worker.py +143 -0
  232. nat/front_ends/mcp/register.py +27 -0
  233. nat/front_ends/mcp/tool_converter.py +241 -0
  234. nat/front_ends/register.py +22 -0
  235. nat/front_ends/simple_base/__init__.py +14 -0
  236. nat/front_ends/simple_base/simple_front_end_plugin_base.py +54 -0
  237. nat/llm/__init__.py +0 -0
  238. nat/llm/aws_bedrock_llm.py +57 -0
  239. nat/llm/nim_llm.py +46 -0
  240. nat/llm/openai_llm.py +46 -0
  241. nat/llm/register.py +23 -0
  242. nat/llm/utils/__init__.py +14 -0
  243. nat/llm/utils/env_config_value.py +94 -0
  244. nat/llm/utils/error.py +17 -0
  245. nat/memory/__init__.py +20 -0
  246. nat/memory/interfaces.py +183 -0
  247. nat/memory/models.py +112 -0
  248. nat/meta/pypi.md +58 -0
  249. nat/object_store/__init__.py +20 -0
  250. nat/object_store/in_memory_object_store.py +76 -0
  251. nat/object_store/interfaces.py +84 -0
  252. nat/object_store/models.py +38 -0
  253. nat/object_store/register.py +20 -0
  254. nat/observability/__init__.py +14 -0
  255. nat/observability/exporter/__init__.py +14 -0
  256. nat/observability/exporter/base_exporter.py +449 -0
  257. nat/observability/exporter/exporter.py +78 -0
  258. nat/observability/exporter/file_exporter.py +33 -0
  259. nat/observability/exporter/processing_exporter.py +322 -0
  260. nat/observability/exporter/raw_exporter.py +52 -0
  261. nat/observability/exporter/span_exporter.py +288 -0
  262. nat/observability/exporter_manager.py +335 -0
  263. nat/observability/mixin/__init__.py +14 -0
  264. nat/observability/mixin/batch_config_mixin.py +26 -0
  265. nat/observability/mixin/collector_config_mixin.py +23 -0
  266. nat/observability/mixin/file_mixin.py +288 -0
  267. nat/observability/mixin/file_mode.py +23 -0
  268. nat/observability/mixin/resource_conflict_mixin.py +134 -0
  269. nat/observability/mixin/serialize_mixin.py +61 -0
  270. nat/observability/mixin/type_introspection_mixin.py +183 -0
  271. nat/observability/processor/__init__.py +14 -0
  272. nat/observability/processor/batching_processor.py +310 -0
  273. nat/observability/processor/callback_processor.py +42 -0
  274. nat/observability/processor/intermediate_step_serializer.py +28 -0
  275. nat/observability/processor/processor.py +71 -0
  276. nat/observability/register.py +96 -0
  277. nat/observability/utils/__init__.py +14 -0
  278. nat/observability/utils/dict_utils.py +236 -0
  279. nat/observability/utils/time_utils.py +31 -0
  280. nat/plugins/.namespace +1 -0
  281. nat/profiler/__init__.py +0 -0
  282. nat/profiler/calc/__init__.py +14 -0
  283. nat/profiler/calc/calc_runner.py +627 -0
  284. nat/profiler/calc/calculations.py +288 -0
  285. nat/profiler/calc/data_models.py +188 -0
  286. nat/profiler/calc/plot.py +345 -0
  287. nat/profiler/callbacks/__init__.py +0 -0
  288. nat/profiler/callbacks/agno_callback_handler.py +295 -0
  289. nat/profiler/callbacks/base_callback_class.py +20 -0
  290. nat/profiler/callbacks/langchain_callback_handler.py +290 -0
  291. nat/profiler/callbacks/llama_index_callback_handler.py +205 -0
  292. nat/profiler/callbacks/semantic_kernel_callback_handler.py +238 -0
  293. nat/profiler/callbacks/token_usage_base_model.py +27 -0
  294. nat/profiler/data_frame_row.py +51 -0
  295. nat/profiler/data_models.py +24 -0
  296. nat/profiler/decorators/__init__.py +0 -0
  297. nat/profiler/decorators/framework_wrapper.py +131 -0
  298. nat/profiler/decorators/function_tracking.py +254 -0
  299. nat/profiler/forecasting/__init__.py +0 -0
  300. nat/profiler/forecasting/config.py +18 -0
  301. nat/profiler/forecasting/model_trainer.py +75 -0
  302. nat/profiler/forecasting/models/__init__.py +22 -0
  303. nat/profiler/forecasting/models/forecasting_base_model.py +40 -0
  304. nat/profiler/forecasting/models/linear_model.py +197 -0
  305. nat/profiler/forecasting/models/random_forest_regressor.py +269 -0
  306. nat/profiler/inference_metrics_model.py +28 -0
  307. nat/profiler/inference_optimization/__init__.py +0 -0
  308. nat/profiler/inference_optimization/bottleneck_analysis/__init__.py +0 -0
  309. nat/profiler/inference_optimization/bottleneck_analysis/nested_stack_analysis.py +460 -0
  310. nat/profiler/inference_optimization/bottleneck_analysis/simple_stack_analysis.py +258 -0
  311. nat/profiler/inference_optimization/data_models.py +386 -0
  312. nat/profiler/inference_optimization/experimental/__init__.py +0 -0
  313. nat/profiler/inference_optimization/experimental/concurrency_spike_analysis.py +468 -0
  314. nat/profiler/inference_optimization/experimental/prefix_span_analysis.py +405 -0
  315. nat/profiler/inference_optimization/llm_metrics.py +212 -0
  316. nat/profiler/inference_optimization/prompt_caching.py +163 -0
  317. nat/profiler/inference_optimization/token_uniqueness.py +107 -0
  318. nat/profiler/inference_optimization/workflow_runtimes.py +72 -0
  319. nat/profiler/intermediate_property_adapter.py +102 -0
  320. nat/profiler/profile_runner.py +473 -0
  321. nat/profiler/utils.py +184 -0
  322. nat/registry_handlers/__init__.py +0 -0
  323. nat/registry_handlers/local/__init__.py +0 -0
  324. nat/registry_handlers/local/local_handler.py +176 -0
  325. nat/registry_handlers/local/register_local.py +37 -0
  326. nat/registry_handlers/metadata_factory.py +60 -0
  327. nat/registry_handlers/package_utils.py +571 -0
  328. nat/registry_handlers/pypi/__init__.py +0 -0
  329. nat/registry_handlers/pypi/pypi_handler.py +251 -0
  330. nat/registry_handlers/pypi/register_pypi.py +40 -0
  331. nat/registry_handlers/register.py +21 -0
  332. nat/registry_handlers/registry_handler_base.py +157 -0
  333. nat/registry_handlers/rest/__init__.py +0 -0
  334. nat/registry_handlers/rest/register_rest.py +56 -0
  335. nat/registry_handlers/rest/rest_handler.py +237 -0
  336. nat/registry_handlers/schemas/__init__.py +0 -0
  337. nat/registry_handlers/schemas/headers.py +42 -0
  338. nat/registry_handlers/schemas/package.py +68 -0
  339. nat/registry_handlers/schemas/publish.py +68 -0
  340. nat/registry_handlers/schemas/pull.py +82 -0
  341. nat/registry_handlers/schemas/remove.py +36 -0
  342. nat/registry_handlers/schemas/search.py +91 -0
  343. nat/registry_handlers/schemas/status.py +47 -0
  344. nat/retriever/__init__.py +0 -0
  345. nat/retriever/interface.py +41 -0
  346. nat/retriever/milvus/__init__.py +14 -0
  347. nat/retriever/milvus/register.py +81 -0
  348. nat/retriever/milvus/retriever.py +228 -0
  349. nat/retriever/models.py +77 -0
  350. nat/retriever/nemo_retriever/__init__.py +14 -0
  351. nat/retriever/nemo_retriever/register.py +60 -0
  352. nat/retriever/nemo_retriever/retriever.py +190 -0
  353. nat/retriever/register.py +22 -0
  354. nat/runtime/__init__.py +14 -0
  355. nat/runtime/loader.py +220 -0
  356. nat/runtime/runner.py +195 -0
  357. nat/runtime/session.py +162 -0
  358. nat/runtime/user_metadata.py +130 -0
  359. nat/settings/__init__.py +0 -0
  360. nat/settings/global_settings.py +318 -0
  361. nat/test/.namespace +1 -0
  362. nat/tool/__init__.py +0 -0
  363. nat/tool/chat_completion.py +74 -0
  364. nat/tool/code_execution/README.md +151 -0
  365. nat/tool/code_execution/__init__.py +0 -0
  366. nat/tool/code_execution/code_sandbox.py +267 -0
  367. nat/tool/code_execution/local_sandbox/.gitignore +1 -0
  368. nat/tool/code_execution/local_sandbox/Dockerfile.sandbox +60 -0
  369. nat/tool/code_execution/local_sandbox/__init__.py +13 -0
  370. nat/tool/code_execution/local_sandbox/local_sandbox_server.py +198 -0
  371. nat/tool/code_execution/local_sandbox/sandbox.requirements.txt +6 -0
  372. nat/tool/code_execution/local_sandbox/start_local_sandbox.sh +50 -0
  373. nat/tool/code_execution/register.py +74 -0
  374. nat/tool/code_execution/test_code_execution_sandbox.py +414 -0
  375. nat/tool/code_execution/utils.py +100 -0
  376. nat/tool/datetime_tools.py +42 -0
  377. nat/tool/document_search.py +141 -0
  378. nat/tool/github_tools/__init__.py +0 -0
  379. nat/tool/github_tools/create_github_commit.py +133 -0
  380. nat/tool/github_tools/create_github_issue.py +87 -0
  381. nat/tool/github_tools/create_github_pr.py +106 -0
  382. nat/tool/github_tools/get_github_file.py +106 -0
  383. nat/tool/github_tools/get_github_issue.py +166 -0
  384. nat/tool/github_tools/get_github_pr.py +256 -0
  385. nat/tool/github_tools/update_github_issue.py +100 -0
  386. nat/tool/mcp/__init__.py +14 -0
  387. nat/tool/mcp/exceptions.py +142 -0
  388. nat/tool/mcp/mcp_client.py +255 -0
  389. nat/tool/mcp/mcp_tool.py +96 -0
  390. nat/tool/memory_tools/__init__.py +0 -0
  391. nat/tool/memory_tools/add_memory_tool.py +79 -0
  392. nat/tool/memory_tools/delete_memory_tool.py +67 -0
  393. nat/tool/memory_tools/get_memory_tool.py +72 -0
  394. nat/tool/nvidia_rag.py +95 -0
  395. nat/tool/register.py +38 -0
  396. nat/tool/retriever.py +94 -0
  397. nat/tool/server_tools.py +66 -0
  398. nat/utils/__init__.py +0 -0
  399. nat/utils/data_models/__init__.py +0 -0
  400. nat/utils/data_models/schema_validator.py +58 -0
  401. nat/utils/debugging_utils.py +43 -0
  402. nat/utils/dump_distro_mapping.py +32 -0
  403. nat/utils/exception_handlers/__init__.py +0 -0
  404. nat/utils/exception_handlers/automatic_retries.py +289 -0
  405. nat/utils/exception_handlers/mcp.py +211 -0
  406. nat/utils/exception_handlers/schemas.py +114 -0
  407. nat/utils/io/__init__.py +0 -0
  408. nat/utils/io/model_processing.py +28 -0
  409. nat/utils/io/yaml_tools.py +119 -0
  410. nat/utils/log_utils.py +37 -0
  411. nat/utils/metadata_utils.py +74 -0
  412. nat/utils/optional_imports.py +142 -0
  413. nat/utils/producer_consumer_queue.py +178 -0
  414. nat/utils/reactive/__init__.py +0 -0
  415. nat/utils/reactive/base/__init__.py +0 -0
  416. nat/utils/reactive/base/observable_base.py +65 -0
  417. nat/utils/reactive/base/observer_base.py +55 -0
  418. nat/utils/reactive/base/subject_base.py +79 -0
  419. nat/utils/reactive/observable.py +59 -0
  420. nat/utils/reactive/observer.py +76 -0
  421. nat/utils/reactive/subject.py +131 -0
  422. nat/utils/reactive/subscription.py +49 -0
  423. nat/utils/settings/__init__.py +0 -0
  424. nat/utils/settings/global_settings.py +197 -0
  425. nat/utils/string_utils.py +38 -0
  426. nat/utils/type_converter.py +290 -0
  427. nat/utils/type_utils.py +484 -0
  428. nat/utils/url_utils.py +27 -0
  429. nvidia_nat-1.2.0.dist-info/METADATA +365 -0
  430. nvidia_nat-1.2.0.dist-info/RECORD +435 -0
  431. nvidia_nat-1.2.0.dist-info/WHEEL +5 -0
  432. nvidia_nat-1.2.0.dist-info/entry_points.txt +21 -0
  433. nvidia_nat-1.2.0.dist-info/licenses/LICENSE-3rd-party.txt +5478 -0
  434. nvidia_nat-1.2.0.dist-info/licenses/LICENSE.md +201 -0
  435. nvidia_nat-1.2.0.dist-info/top_level.txt +2 -0
@@ -0,0 +1,129 @@
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 logging
18
+ from contextlib import AsyncExitStack
19
+
20
+ import click
21
+
22
+ from nat.data_models.component import ComponentEnum
23
+ from nat.data_models.registry_handler import RegistryHandlerBaseConfig
24
+ from nat.registry_handlers.schemas.search import SearchFields
25
+
26
+ logger = logging.getLogger(__name__)
27
+
28
+
29
+ async def search_artifacts( # pylint: disable=R0917
30
+ registry_handler_config: RegistryHandlerBaseConfig,
31
+ component_types: list[ComponentEnum],
32
+ visualize: bool,
33
+ query: str,
34
+ num_results: int,
35
+ query_fields: list[SearchFields],
36
+ save_path: str | None) -> None:
37
+
38
+ from nat.cli.type_registry import GlobalTypeRegistry
39
+ from nat.registry_handlers.schemas.search import SearchQuery
40
+
41
+ registry = GlobalTypeRegistry.get()
42
+
43
+ async with AsyncExitStack() as stack:
44
+
45
+ registry_handler_info = registry.get_registry_handler(type(registry_handler_config))
46
+ registry_handler = await stack.enter_async_context(registry_handler_info.build_fn(registry_handler_config))
47
+
48
+ if (len(component_types) == 0):
49
+ component_types = [t.value for t in ComponentEnum]
50
+
51
+ if (len(query_fields) == 0):
52
+ query_fields = (SearchFields.ALL, )
53
+
54
+ query = SearchQuery(query=query, fields=query_fields, top_k=num_results, component_types=component_types)
55
+ search_response = await stack.enter_async_context(registry_handler.search(query=query))
56
+
57
+ if (visualize):
58
+ registry_handler.visualize_search_results(search_response=search_response)
59
+ if (save_path is not None):
60
+ registry_handler.save_search_results(search_response=search_response, save_path=save_path)
61
+
62
+
63
+ @click.group(name=__name__, invoke_without_command=True, help="List the locally registered NAT components.")
64
+ @click.option(
65
+ "-t",
66
+ "--types",
67
+ "component_types",
68
+ multiple=True,
69
+ type=click.Choice([e.value for e in ComponentEnum], case_sensitive=False),
70
+ required=False,
71
+ help=("Filter the search by NAT component type."),
72
+ )
73
+ @click.option(
74
+ "-o",
75
+ "--output_path",
76
+ type=str,
77
+ required=False,
78
+ help=("Path to save search results."),
79
+ )
80
+ @click.option(
81
+ "-q",
82
+ "--query",
83
+ type=str,
84
+ default="",
85
+ required=False,
86
+ help=("The query string."),
87
+ )
88
+ @click.option(
89
+ "-n",
90
+ "--num_results",
91
+ type=int,
92
+ default=-1,
93
+ required=False,
94
+ help=("Number of results to return."),
95
+ )
96
+ @click.option(
97
+ "-f",
98
+ "--fields",
99
+ multiple=True,
100
+ type=click.Choice([e.value for e in SearchFields], case_sensitive=False),
101
+ required=False,
102
+ help=("Fields used when applying query."),
103
+ )
104
+ def list_components(fields: list[SearchFields],
105
+ query: str,
106
+ num_results: int,
107
+ component_types: list[ComponentEnum],
108
+ output_path: str | None = None) -> None:
109
+
110
+ from nat.runtime.loader import PluginTypes
111
+ from nat.runtime.loader import discover_and_register_plugins
112
+ from nat.settings.global_settings import GlobalSettings
113
+
114
+ discover_and_register_plugins(PluginTypes.ALL)
115
+
116
+ config_dict = {"channels": {"list_components": {"_type": "local"}}}
117
+ registry_config = GlobalSettings.get().model_validate(config_dict)
118
+ local_registry_config = registry_config.channels.get("list_components", None)
119
+ if (local_registry_config is None):
120
+ logger.error("Channel runtime instance not found.")
121
+
122
+ asyncio.run(
123
+ search_artifacts(local_registry_config,
124
+ query=query,
125
+ num_results=num_results,
126
+ query_fields=fields,
127
+ component_types=component_types,
128
+ visualize=True,
129
+ save_path=output_path))
@@ -0,0 +1,304 @@
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}")
@@ -0,0 +1,14 @@
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.
@@ -0,0 +1,88 @@
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 logging
18
+ from contextlib import AsyncExitStack
19
+ from pathlib import Path
20
+
21
+ import click
22
+
23
+ from nat.data_models.registry_handler import RegistryHandlerBaseConfig
24
+ from nat.utils.data_models.schema_validator import validate_yaml
25
+
26
+ logger = logging.getLogger(__name__)
27
+
28
+
29
+ async def publish_artifact(registry_handler_config: RegistryHandlerBaseConfig, package_root: str) -> None:
30
+
31
+ from nat.cli.type_registry import GlobalTypeRegistry
32
+ from nat.registry_handlers.package_utils import build_artifact
33
+
34
+ registry = GlobalTypeRegistry.get()
35
+
36
+ async with AsyncExitStack() as stack:
37
+
38
+ registry_handler_info = registry.get_registry_handler(type(registry_handler_config))
39
+ registry_handler = await stack.enter_async_context(registry_handler_info.build_fn(registry_handler_config))
40
+ try:
41
+ artifact = build_artifact(package_root=package_root)
42
+ except Exception as e:
43
+ logger.exception("Error building artifact: %s", e, exc_info=True)
44
+ return
45
+ await stack.enter_async_context(registry_handler.publish(artifact=artifact))
46
+
47
+
48
+ @click.group(name=__name__,
49
+ invoke_without_command=True,
50
+ help=("Publish local NAT artifacts to a remote "
51
+ "registry from package repository."))
52
+ @click.option(
53
+ "--config_file",
54
+ type=click.Path(exists=True, file_okay=True, dir_okay=False, path_type=Path),
55
+ callback=validate_yaml,
56
+ required=False,
57
+ help=("A YAML file to override configured channel settings."),
58
+ )
59
+ @click.option(
60
+ "-c",
61
+ "--channel",
62
+ type=str,
63
+ required=True,
64
+ help=("The remote registry channel to use when publishing the NAT artifact."),
65
+ )
66
+ @click.argument("package_root", type=str)
67
+ def publish(channel: str, config_file: str, package_root: str) -> None:
68
+ """
69
+ Publish NAT artifacts with the specified configuration
70
+ """
71
+ from nat.settings.global_settings import GlobalSettings
72
+
73
+ settings = GlobalSettings().get()
74
+
75
+ if (config_file is not None):
76
+ settings = settings.override_settings(config_file)
77
+
78
+ try:
79
+ publish_channel_config = settings.channels.get(channel)
80
+
81
+ if (publish_channel_config is None):
82
+ logger.error("Publish channel '%s' has not been configured.", channel)
83
+ return
84
+ except Exception as e:
85
+ logger.exception("Error loading user settings: %s", e, exc_info=True)
86
+ return
87
+
88
+ asyncio.run(publish_artifact(registry_handler_config=publish_channel_config, package_root=package_root))
@@ -0,0 +1,118 @@
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 logging
18
+ from contextlib import AsyncExitStack
19
+ from pathlib import Path
20
+
21
+ import click
22
+
23
+ from nat.data_models.registry_handler import RegistryHandlerBaseConfig
24
+ from nat.utils.data_models.schema_validator import validate_yaml
25
+
26
+ logger = logging.getLogger(__name__)
27
+
28
+
29
+ async def pull_artifact(registry_handler_config: RegistryHandlerBaseConfig, packages: list[str]) -> None:
30
+
31
+ from nat.cli.type_registry import GlobalTypeRegistry
32
+ from nat.registry_handlers.schemas.package import PackageNameVersion
33
+ from nat.registry_handlers.schemas.pull import PullPackageWhl
34
+ from nat.registry_handlers.schemas.pull import PullRequestPackages
35
+
36
+ registry = GlobalTypeRegistry.get()
37
+
38
+ async with AsyncExitStack() as stack:
39
+
40
+ registry_handler_info = registry.get_registry_handler(type(registry_handler_config))
41
+ registry_handler = await stack.enter_async_context(registry_handler_info.build_fn(registry_handler_config))
42
+
43
+ try:
44
+ package_list = []
45
+ for package in packages:
46
+
47
+ package_data = {}
48
+
49
+ assert len(package) > 0, f"Supplied invalid package '{package}'."
50
+
51
+ if package[:-4] == ".whl":
52
+ package_data["whl_path"] = package
53
+ package_list.append(PullPackageWhl(**package_data))
54
+ else:
55
+ package_split = package.split("==")
56
+
57
+ assert len(package_split) in (1, 2), f"Supplied invalid package '{package}'."
58
+
59
+ package_data["name"] = package_split[0]
60
+
61
+ if (package_split == 2):
62
+ package_data["version"] = package_split[1]
63
+
64
+ package_list.append(PackageNameVersion(**package_data))
65
+
66
+ validated_packages = PullRequestPackages(packages=package_list)
67
+
68
+ except Exception as e:
69
+ logger.exception("Error processing package names: %s", e, exc_info=True)
70
+ return
71
+
72
+ await stack.enter_async_context(registry_handler.pull(packages=validated_packages))
73
+
74
+
75
+ @click.group(name=__name__,
76
+ invoke_without_command=True,
77
+ help=("Pull NAT artifacts from a remote registry "
78
+ "by package name."))
79
+ @click.option(
80
+ "--config_file",
81
+ type=click.Path(exists=True, file_okay=True, dir_okay=False, path_type=Path),
82
+ callback=validate_yaml,
83
+ required=False,
84
+ help=("A YAML file to override the channel settings."),
85
+ )
86
+ @click.option(
87
+ "-c",
88
+ "--channel",
89
+ type=str,
90
+ required=True,
91
+ help=("The remote registry channel to use when pulling the NAT artifact."),
92
+ )
93
+ @click.argument("packages", type=str)
94
+ def pull(channel: str, config_file: str, packages: str) -> None:
95
+ """
96
+ Pull NAT artifacts from a remote registry channel.
97
+ """
98
+
99
+ from nat.settings.global_settings import GlobalSettings
100
+
101
+ packages = packages.split()
102
+
103
+ settings = GlobalSettings().get()
104
+
105
+ if (config_file is not None):
106
+ settings = settings.override_settings(config_file)
107
+
108
+ try:
109
+ pull_channel_config = settings.channels.get(channel)
110
+
111
+ if (pull_channel_config is None):
112
+ logger.error("Pull channel '%s' has not been configured.", channel)
113
+ return
114
+ except Exception as e:
115
+ logger.exception("Error loading user settings: %s", e, exc_info=True)
116
+ return
117
+
118
+ asyncio.run(pull_artifact(pull_channel_config, packages))
@@ -0,0 +1,36 @@
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 logging
17
+
18
+ import click
19
+
20
+ from nat.cli.commands.registry.publish import publish
21
+ from nat.cli.commands.registry.pull import pull
22
+ from nat.cli.commands.registry.remove import remove
23
+ from nat.cli.commands.registry.search import search
24
+
25
+ logger = logging.getLogger(__name__)
26
+
27
+
28
+ @click.group(name=__name__, invoke_without_command=False, help="Utility to configure NAT remote registry channels.")
29
+ def registry_command(**kwargs):
30
+ pass
31
+
32
+
33
+ registry_command.add_command(publish, "publish")
34
+ registry_command.add_command(pull, "pull")
35
+ registry_command.add_command(remove, "remove")
36
+ registry_command.add_command(search, "search")