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,25 @@
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 nat.cli.register_workflow import register_front_end
17
+ from nat.data_models.config import Config
18
+ from nat.front_ends.fastapi.fastapi_front_end_config import FastApiFrontEndConfig
19
+
20
+
21
+ @register_front_end(config_type=FastApiFrontEndConfig)
22
+ async def register_fastapi_front_end(config: FastApiFrontEndConfig, full_config: Config):
23
+ from nat.front_ends.fastapi.fastapi_front_end_plugin import FastApiFrontEndPlugin
24
+
25
+ yield FastApiFrontEndPlugin(full_config=full_config)
@@ -0,0 +1,195 @@
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 typing
18
+ from collections.abc import AsyncGenerator
19
+
20
+ from nat.data_models.api_server import ResponseIntermediateStep
21
+ from nat.data_models.api_server import ResponsePayloadOutput
22
+ from nat.data_models.api_server import ResponseSerializable
23
+ from nat.data_models.step_adaptor import StepAdaptorConfig
24
+ from nat.front_ends.fastapi.intermediate_steps_subscriber import pull_intermediate
25
+ from nat.front_ends.fastapi.step_adaptor import StepAdaptor
26
+ from nat.runtime.session import SessionManager
27
+ from nat.utils.producer_consumer_queue import AsyncIOProducerConsumerQueue
28
+
29
+
30
+ async def generate_streaming_response_as_str(payload: typing.Any,
31
+ *,
32
+ session_manager: SessionManager,
33
+ streaming: bool,
34
+ step_adaptor: StepAdaptor = StepAdaptor(StepAdaptorConfig()),
35
+ result_type: type | None = None,
36
+ output_type: type | None = None) -> AsyncGenerator[str]:
37
+
38
+ async for item in generate_streaming_response(payload,
39
+ session_manager=session_manager,
40
+ streaming=streaming,
41
+ step_adaptor=step_adaptor,
42
+ result_type=result_type,
43
+ output_type=output_type):
44
+
45
+ if (isinstance(item, ResponseSerializable)):
46
+ yield item.get_stream_data()
47
+ else:
48
+ raise ValueError("Unexpected item type in stream. Expected ChatResponseSerializable, got: " +
49
+ str(type(item)))
50
+
51
+
52
+ async def generate_streaming_response(payload: typing.Any,
53
+ *,
54
+ session_manager: SessionManager,
55
+ streaming: bool,
56
+ step_adaptor: StepAdaptor = StepAdaptor(StepAdaptorConfig()),
57
+ result_type: type | None = None,
58
+ output_type: type | None = None) -> AsyncGenerator[ResponseSerializable]:
59
+
60
+ async with session_manager.run(payload) as runner:
61
+
62
+ q: AsyncIOProducerConsumerQueue[ResponseSerializable] = AsyncIOProducerConsumerQueue()
63
+
64
+ # Start the intermediate stream
65
+ intermediate_complete = await pull_intermediate(q, step_adaptor)
66
+
67
+ async def pull_result():
68
+ if session_manager.workflow.has_streaming_output and streaming:
69
+ async for chunk in runner.result_stream(to_type=output_type):
70
+ await q.put(chunk)
71
+ else:
72
+ result = await runner.result(to_type=result_type)
73
+ await q.put(runner.convert(result, output_type))
74
+
75
+ # Wait until the intermediate subscription is done before closing q
76
+ # But we have no direct "intermediate_done" reference here
77
+ # because it's encapsulated in pull_intermediate. So we can do:
78
+ # await some_event.wait()
79
+ # If needed. Alternatively, you can skip that if the intermediate
80
+ # subscriber won't block the main flow.
81
+ #
82
+ # For example, if you *need* to guarantee the subscriber is done before
83
+ # closing the queue, you can structure the code to store or return
84
+ # the 'intermediate_done' event from pull_intermediate.
85
+ #
86
+
87
+ await intermediate_complete.wait()
88
+
89
+ await q.close()
90
+
91
+ try:
92
+ # Start the result stream
93
+ asyncio.create_task(pull_result())
94
+
95
+ async for item in q:
96
+
97
+ if (isinstance(item, ResponseSerializable)):
98
+ yield item
99
+ else:
100
+ yield ResponsePayloadOutput(payload=item)
101
+ except Exception as e:
102
+ # Handle exceptions here
103
+ raise e
104
+ finally:
105
+ await q.close()
106
+
107
+
108
+ async def generate_single_response(
109
+ payload: typing.Any,
110
+ session_manager: SessionManager,
111
+ result_type: type | None = None,
112
+ ) -> typing.Any:
113
+ if (not session_manager.workflow.has_single_output):
114
+ raise ValueError("Cannot get a single output value for streaming workflows")
115
+
116
+ async with session_manager.run(payload) as runner:
117
+ return await runner.result(to_type=result_type)
118
+
119
+
120
+ async def generate_streaming_response_full(payload: typing.Any,
121
+ *,
122
+ session_manager: SessionManager,
123
+ streaming: bool,
124
+ result_type: type | None = None,
125
+ output_type: type | None = None,
126
+ filter_steps: str | None = None) -> AsyncGenerator[ResponseSerializable]:
127
+ """
128
+ Similar to generate_streaming_response but provides raw ResponseIntermediateStep objects
129
+ without any step adaptor translations.
130
+ """
131
+ # Parse filter_steps into a set of allowed types if provided
132
+ # Special case: if filter_steps is "none", suppress all steps
133
+ allowed_types = None
134
+ if filter_steps:
135
+ if filter_steps.lower() == "none":
136
+ allowed_types = set() # Empty set means no steps allowed
137
+ else:
138
+ allowed_types = set(filter_steps.split(','))
139
+
140
+ async with session_manager.run(payload) as runner:
141
+ q: AsyncIOProducerConsumerQueue[ResponseSerializable] = AsyncIOProducerConsumerQueue()
142
+
143
+ # Start the intermediate stream without step adaptor
144
+ intermediate_complete = await pull_intermediate(q, None)
145
+
146
+ async def pull_result():
147
+ if session_manager.workflow.has_streaming_output and streaming:
148
+ async for chunk in runner.result_stream(to_type=output_type):
149
+ await q.put(chunk)
150
+ else:
151
+ result = await runner.result(to_type=result_type)
152
+ await q.put(runner.convert(result, output_type))
153
+
154
+ await intermediate_complete.wait()
155
+ await q.close()
156
+
157
+ try:
158
+ # Start the result stream
159
+ asyncio.create_task(pull_result())
160
+
161
+ async for item in q:
162
+ if (isinstance(item, ResponseIntermediateStep)):
163
+ # Filter intermediate steps if filter_steps is provided
164
+ if allowed_types is None or item.type in allowed_types:
165
+ yield item
166
+ else:
167
+ yield ResponsePayloadOutput(payload=item)
168
+ except Exception as e:
169
+ # Handle exceptions here
170
+ raise e
171
+ finally:
172
+ await q.close()
173
+
174
+
175
+ async def generate_streaming_response_full_as_str(payload: typing.Any,
176
+ *,
177
+ session_manager: SessionManager,
178
+ streaming: bool,
179
+ result_type: type | None = None,
180
+ output_type: type | None = None,
181
+ filter_steps: str | None = None) -> AsyncGenerator[str]:
182
+ """
183
+ Similar to generate_streaming_response but converts the response to a string format.
184
+ """
185
+ async for item in generate_streaming_response_full(payload,
186
+ session_manager=session_manager,
187
+ streaming=streaming,
188
+ result_type=result_type,
189
+ output_type=output_type,
190
+ filter_steps=filter_steps):
191
+ if (isinstance(item, ResponseIntermediateStep) or isinstance(item, ResponsePayloadOutput)):
192
+ yield item.get_stream_data()
193
+ else:
194
+ raise ValueError("Unexpected item type in stream. Expected ChatResponseSerializable, got: " +
195
+ str(type(item)))
@@ -0,0 +1,319 @@
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 html
17
+ import logging
18
+ from functools import reduce
19
+ from textwrap import dedent
20
+
21
+ from nat.data_models.api_server import ResponseIntermediateStep
22
+ from nat.data_models.api_server import ResponseSerializable
23
+ from nat.data_models.intermediate_step import IntermediateStep
24
+ from nat.data_models.intermediate_step import IntermediateStepCategory
25
+ from nat.data_models.intermediate_step import IntermediateStepPayload
26
+ from nat.data_models.intermediate_step import IntermediateStepType
27
+ from nat.data_models.invocation_node import InvocationNode
28
+ from nat.data_models.step_adaptor import StepAdaptorConfig
29
+ from nat.data_models.step_adaptor import StepAdaptorMode
30
+ from nat.utils.type_utils import is_valid_json
31
+
32
+ logger = logging.getLogger(__name__)
33
+
34
+
35
+ class StepAdaptor:
36
+
37
+ def __init__(self, config: StepAdaptorConfig):
38
+
39
+ self._history: list[IntermediateStep] = []
40
+ self.config = config
41
+
42
+ def _step_matches_filter(self, step: IntermediateStep, config: StepAdaptorConfig) -> bool:
43
+ """
44
+ Returns True if this intermediate step should be included (based on the config.mode).
45
+ """
46
+
47
+ if config.mode == StepAdaptorMode.OFF:
48
+ return False
49
+
50
+ if config.mode == StepAdaptorMode.DEFAULT:
51
+ # default existing behavior: show LLM events + TOOL_END + FUNCTION events
52
+ if step.event_category == IntermediateStepCategory.LLM:
53
+ return True
54
+ if step.event_category == IntermediateStepCategory.TOOL:
55
+ return True
56
+ if step.event_category == IntermediateStepCategory.FUNCTION:
57
+ return True
58
+ return False
59
+
60
+ if config.mode == StepAdaptorMode.CUSTOM:
61
+ # pass only what the user explicitly listed
62
+ return step.event_type in config.custom_event_types
63
+
64
+ return False
65
+
66
+ def _handle_llm(self, step: IntermediateStepPayload, ancestry: InvocationNode) -> ResponseSerializable | None:
67
+ input_str: str | None = None
68
+ output_str: str | None = None
69
+
70
+ # Find the start in the history with matching run_id
71
+ start_step = next(
72
+ (x for x in self._history if x.event_type == IntermediateStepType.LLM_START and x.UUID == step.UUID), None)
73
+
74
+ if not start_step:
75
+ # If we don't have a start step, we can't do anything
76
+ return None
77
+
78
+ input_str = str(start_step.data.input)
79
+
80
+ if step.event_type == IntermediateStepType.LLM_NEW_TOKEN:
81
+
82
+ # Find all of the previous LLM chunks and concatenate them
83
+ output_str = reduce(
84
+ lambda x, y: x + y,
85
+ (str(x.data.chunk)
86
+ for x in self._history if x.event_type == IntermediateStepType.LLM_NEW_TOKEN and x.UUID == step.UUID),
87
+ "")
88
+
89
+ elif step.event_type == IntermediateStepType.LLM_END:
90
+ output_str = str(step.data.output)
91
+
92
+ if not input_str and not output_str:
93
+ return None
94
+
95
+ escaped_input = html.escape(input_str, quote=False)
96
+
97
+ # Dont use f-strings here because the payload is markdown and screws up the dedent
98
+ payload = dedent("""
99
+ **Input:**
100
+ ```python
101
+ {input_value}
102
+ ```
103
+ """).strip("\n").format(input_value=escaped_input)
104
+
105
+ if (output_str):
106
+ escaped_output = html.escape(output_str, quote=False) if output_str else ""
107
+
108
+ # Dont use f-strings here because the payload is markdown and screws up the dedent
109
+ payload = dedent("""
110
+ {payload}
111
+
112
+ **Output:**
113
+ {output_value}
114
+ """).strip("\n").format(payload=payload, output_value=escaped_output)
115
+
116
+ event = ResponseIntermediateStep(id=step.UUID,
117
+ name=step.name or "",
118
+ payload=payload,
119
+ parent_id=ancestry.function_id)
120
+
121
+ return event
122
+
123
+ def _handle_tool(self, step: IntermediateStepPayload, ancestry: InvocationNode) -> ResponseSerializable | None:
124
+ """
125
+ Handles both TOOL_START and TOOL_END events
126
+ """
127
+ input_str: str | None = None
128
+ output_str: str | None = None
129
+
130
+ # Find the start in the history with matching run_id
131
+ start_step = next(
132
+ (x for x in self._history if x.event_type == IntermediateStepType.TOOL_START and x.UUID == step.UUID), None)
133
+
134
+ if not start_step:
135
+ # If we don't have a start step, we can't do anything
136
+ return None
137
+
138
+ input_str = str(start_step.data.input)
139
+
140
+ if step.event_type == IntermediateStepType.TOOL_END:
141
+ output_str = str(step.data.output)
142
+
143
+ if not input_str and not output_str:
144
+ return None
145
+
146
+ escaped_input = html.escape(input_str, quote=False)
147
+ format_input_type = "json" if is_valid_json(escaped_input) else "python"
148
+
149
+ # Dont use f-strings here because the payload is markdown and screws up the dedent
150
+ payload = dedent("""
151
+ **Input:**
152
+ ```{format_input_type}
153
+ {input_value}
154
+ ```
155
+ """).strip("\n").format(input_value=escaped_input, format_input_type=format_input_type)
156
+
157
+ if output_str:
158
+ escaped_output = html.escape(output_str, quote=False)
159
+ format_output_type = "json" if is_valid_json(escaped_output) else "python"
160
+
161
+ # Dont use f-strings here because the payload is markdown and screws up the dedent
162
+ payload = dedent("""
163
+ {payload}
164
+
165
+ **Output:**
166
+ ```{format_output_type}
167
+ {output_value}
168
+ ```
169
+ """).strip("\n").format(payload=payload, output_value=escaped_output, format_output_type=format_output_type)
170
+
171
+ event = ResponseIntermediateStep(id=step.UUID,
172
+ name=f"Tool: {step.name}",
173
+ payload=payload,
174
+ parent_id=ancestry.function_id)
175
+
176
+ return event
177
+
178
+ def _handle_function(self, step: IntermediateStepPayload, ancestry: InvocationNode) -> ResponseSerializable | None:
179
+ """
180
+ Handles the FUNCTION_START and FUNCTION_END events
181
+ """
182
+ input_str: str | None = None
183
+ output_str: str | None = None
184
+
185
+ if step.event_type == IntermediateStepType.FUNCTION_START:
186
+ # For function start events, display input data
187
+ if step.data and hasattr(step.data, 'input'):
188
+ input_str = str(step.data.input)
189
+ elif step.data:
190
+ input_str = str(step.data)
191
+
192
+ if not input_str:
193
+ return None
194
+
195
+ escaped_input = html.escape(input_str, quote=False)
196
+ format_input_type = "json" if is_valid_json(escaped_input) else "python"
197
+
198
+ # Create payload for function start
199
+ payload_str = dedent("""
200
+ **Function Input:**
201
+ ```{format_input_type}
202
+ {input_value}
203
+ ```
204
+ """).strip("\n").format(input_value=escaped_input, format_input_type=format_input_type)
205
+
206
+ event = ResponseIntermediateStep(id=step.UUID,
207
+ name=f"Function Start: {step.name}",
208
+ payload=payload_str,
209
+ parent_id=ancestry.parent_id)
210
+ return event
211
+
212
+ if step.event_type == IntermediateStepType.FUNCTION_END:
213
+ # Find the start event with matching UUID
214
+ start_step = next(
215
+ (x
216
+ for x in self._history if x.event_type == IntermediateStepType.FUNCTION_START and x.UUID == step.UUID),
217
+ None)
218
+
219
+ # For function end events, display output data
220
+ if step.data and hasattr(step.data, 'output'):
221
+ output_str = str(step.data.output)
222
+ elif step.data:
223
+ output_str = str(step.data)
224
+
225
+ if not output_str:
226
+ return None
227
+
228
+ escaped_output = html.escape(output_str, quote=False)
229
+ format_output_type = "json" if is_valid_json(escaped_output) else "python"
230
+
231
+ # Get input from start step if available
232
+ input_payload = ""
233
+ if start_step and start_step.data:
234
+ if hasattr(start_step.data, 'input'):
235
+ input_str = str(start_step.data.input)
236
+ else:
237
+ input_str = str(start_step.data)
238
+
239
+ if input_str:
240
+ escaped_input = html.escape(input_str, quote=False)
241
+ format_input_type = "json" if is_valid_json(escaped_input) else "python"
242
+ input_payload = dedent("""
243
+ **Function Input:**
244
+ ```{format_input_type}
245
+ {input_value}
246
+ ```
247
+ """).strip("\n").format(input_value=escaped_input, format_input_type=format_input_type)
248
+
249
+ # Create payload for function end
250
+ payload_str = dedent("""
251
+ {input_payload}**Function Output:**
252
+ ```{format_output_type}
253
+ {output_value}
254
+ ```
255
+ """).strip("\n").format(input_payload=input_payload,
256
+ output_value=escaped_output,
257
+ format_output_type=format_output_type)
258
+
259
+ event = ResponseIntermediateStep(id=step.UUID,
260
+ name=f"Function Complete: {step.name}",
261
+ payload=payload_str,
262
+ parent_id=ancestry.parent_id)
263
+ return event
264
+
265
+ return None
266
+
267
+ def _handle_custom(self, payload: IntermediateStepPayload, ancestry: InvocationNode) -> ResponseSerializable | None:
268
+ """
269
+ Handles the CUSTOM event
270
+ """
271
+ escaped_payload = html.escape(str(payload), quote=False)
272
+ escaped_payload = escaped_payload.replace("\n", "")
273
+
274
+ # Attempt to determine type
275
+ format_type = "json" if is_valid_json(escaped_payload) else "python"
276
+
277
+ # Don't use f-strings here because the payload is markdown and screws up the dedent
278
+ payload_str = dedent("""
279
+ ```{format_type}
280
+ {payload}
281
+ ```
282
+ """).strip("\n").format(payload=escaped_payload, format_type=format_type)
283
+
284
+ # Return the event
285
+ event = ResponseIntermediateStep(id=payload.UUID,
286
+ name=f"{payload.event_type}",
287
+ payload=payload_str,
288
+ parent_id=ancestry.function_id)
289
+
290
+ return event
291
+
292
+ def process(self, step: IntermediateStep) -> ResponseSerializable | None: # pylint: disable=R1710
293
+
294
+ # Track the chunk
295
+ self._history.append(step)
296
+ payload = step.payload
297
+ ancestry = step.function_ancestry
298
+
299
+ if not self._step_matches_filter(step, self.config):
300
+ return None
301
+
302
+ try:
303
+
304
+ if step.event_category == IntermediateStepCategory.LLM:
305
+ return self._handle_llm(payload, ancestry)
306
+
307
+ if step.event_category == IntermediateStepCategory.TOOL:
308
+ return self._handle_tool(payload, ancestry)
309
+
310
+ if step.event_category == IntermediateStepCategory.FUNCTION:
311
+ return self._handle_function(payload, ancestry)
312
+
313
+ if step.event_category == IntermediateStepCategory.CUSTOM:
314
+ return self._handle_custom(payload, ancestry)
315
+
316
+ except Exception as e:
317
+ logger.error("Error processing intermediate step: %s", e, exc_info=True)
318
+
319
+ return None
@@ -0,0 +1,14 @@
1
+ # SPDX-FileCopyrightText: Copyright (c) 2024-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,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
+ from pydantic import Field
17
+
18
+ from nat.data_models.front_end import FrontEndBaseConfig
19
+
20
+
21
+ class MCPFrontEndConfig(FrontEndBaseConfig, name="mcp"):
22
+ """MCP front end configuration.
23
+
24
+ A simple MCP (Modular Communication Protocol) front end for NeMo Agent toolkit.
25
+ """
26
+
27
+ name: str = Field(default="NeMo Agent Toolkit MCP",
28
+ description="Name of the MCP server (default: NeMo Agent Toolkit MCP)")
29
+ host: str = Field(default="localhost", description="Host to bind the server to (default: localhost)")
30
+ port: int = Field(default=9901, description="Port to bind the server to (default: 9901)", ge=0, le=65535)
31
+ debug: bool = Field(default=False, description="Enable debug mode (default: False)")
32
+ log_level: str = Field(default="INFO", description="Log level for the MCP server (default: INFO)")
33
+ tool_names: list[str] = Field(default_factory=list,
34
+ description="The list of tools MCP server will expose (default: all tools)")
35
+ runner_class: str | None = Field(
36
+ default=None, description="Custom worker class for handling MCP routes (default: built-in worker)")
@@ -0,0 +1,81 @@
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
+ import typing
18
+
19
+ from nat.builder.front_end import FrontEndBase
20
+ from nat.builder.workflow_builder import WorkflowBuilder
21
+ from nat.front_ends.mcp.mcp_front_end_config import MCPFrontEndConfig
22
+ from nat.front_ends.mcp.mcp_front_end_plugin_worker import MCPFrontEndPluginWorkerBase
23
+
24
+ logger = logging.getLogger(__name__)
25
+
26
+
27
+ class MCPFrontEndPlugin(FrontEndBase[MCPFrontEndConfig]):
28
+ """MCP front end plugin implementation."""
29
+
30
+ def get_worker_class(self) -> type[MCPFrontEndPluginWorkerBase]:
31
+ """Get the worker class for handling MCP routes."""
32
+ from nat.front_ends.mcp.mcp_front_end_plugin_worker import MCPFrontEndPluginWorker
33
+
34
+ return MCPFrontEndPluginWorker
35
+
36
+ @typing.final
37
+ def get_worker_class_name(self) -> str:
38
+ """Get the worker class name from configuration or default."""
39
+ if self.front_end_config.runner_class:
40
+ return self.front_end_config.runner_class
41
+
42
+ worker_class = self.get_worker_class()
43
+ return f"{worker_class.__module__}.{worker_class.__qualname__}"
44
+
45
+ def _get_worker_instance(self) -> MCPFrontEndPluginWorkerBase:
46
+ """Get an instance of the worker class."""
47
+ # Import the worker class dynamically if specified in config
48
+ if self.front_end_config.runner_class:
49
+ module_name, class_name = self.front_end_config.runner_class.rsplit(".", 1)
50
+ import importlib
51
+ module = importlib.import_module(module_name)
52
+ worker_class = getattr(module, class_name)
53
+ else:
54
+ worker_class = self.get_worker_class()
55
+
56
+ return worker_class(self.full_config)
57
+
58
+ async def run(self) -> None:
59
+ """Run the MCP server."""
60
+ # Import FastMCP
61
+ from mcp.server.fastmcp import FastMCP
62
+
63
+ # Create an MCP server with the configured parameters
64
+ mcp = FastMCP(
65
+ self.front_end_config.name,
66
+ host=self.front_end_config.host,
67
+ port=self.front_end_config.port,
68
+ debug=self.front_end_config.debug,
69
+ log_level=self.front_end_config.log_level,
70
+ )
71
+
72
+ # Get the worker instance and set up routes
73
+ worker = self._get_worker_instance()
74
+
75
+ # Build the workflow and add routes using the worker
76
+ async with WorkflowBuilder.from_config(config=self.full_config) as builder:
77
+ # Add routes through the worker (includes health endpoint and function registration)
78
+ await worker.add_routes(mcp, builder)
79
+
80
+ # Start the MCP server
81
+ await mcp.run_sse_async()