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,124 @@
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.
15
+
16
+ import logging
17
+ import re
18
+ import string
19
+
20
+ from pydantic import Field
21
+ from pydantic import field_validator
22
+
23
+ from nat.authentication.exceptions.api_key_exceptions import APIKeyFieldError
24
+ from nat.authentication.exceptions.api_key_exceptions import HeaderNameFieldError
25
+ from nat.authentication.exceptions.api_key_exceptions import HeaderPrefixFieldError
26
+ from nat.data_models.authentication import AuthProviderBaseConfig
27
+ from nat.data_models.authentication import HeaderAuthScheme
28
+
29
+ logger = logging.getLogger(__name__)
30
+
31
+ # Strict RFC 7230 compliant header name regex
32
+ HEADER_NAME_REGEX = re.compile(r"^[!#$%&'*+\-.^_`|~0-9a-zA-Z]+$")
33
+
34
+
35
+ class APIKeyAuthProviderConfig(AuthProviderBaseConfig, name="api_key"):
36
+ """
37
+ API Key authentication configuration model.
38
+ """
39
+
40
+ raw_key: str = Field(description=("Raw API token or credential to be injected into the request parameter. "
41
+ "Used for 'bearer','x-api-key','custom', and other schemes. "))
42
+
43
+ auth_scheme: HeaderAuthScheme = Field(default=HeaderAuthScheme.BEARER,
44
+ description=("The HTTP authentication scheme to use. "
45
+ "Supported schemes: BEARER, X_API_KEY, BASIC, CUSTOM."))
46
+
47
+ custom_header_name: str | None = Field(description="The HTTP header name that MUST be used in conjunction "
48
+ "with the custom_header_prefix when HeaderAuthScheme is CUSTOM.",
49
+ default=None)
50
+ custom_header_prefix: str | None = Field(description="The HTTP header prefix that MUST be used in conjunction "
51
+ "with the custom_header_name when HeaderAuthScheme is CUSTOM.",
52
+ default=None)
53
+
54
+ @field_validator('raw_key')
55
+ @classmethod
56
+ def validate_raw_key(cls, value: str) -> str:
57
+ if not value:
58
+ raise APIKeyFieldError('value_missing', 'raw_key field value is required.')
59
+
60
+ if len(value) < 8:
61
+ raise APIKeyFieldError(
62
+ 'value_too_short',
63
+ 'raw_key field value must be at least 8 characters long for security. '
64
+ f'Got: {len(value)} characters.')
65
+
66
+ if len(value.strip()) != len(value):
67
+ raise APIKeyFieldError('whitespace_found',
68
+ 'raw_key field value cannot have leading or trailing whitespace.')
69
+
70
+ if any(c in string.whitespace for c in value):
71
+ raise APIKeyFieldError('contains_whitespace', 'raw_key must not contain any '
72
+ 'whitespace characters.')
73
+
74
+ return value
75
+
76
+ @field_validator('custom_header_name')
77
+ @classmethod
78
+ def validate_custom_header_name(cls, value: str) -> str:
79
+ if not value:
80
+ raise HeaderNameFieldError('value_missing', 'custom_header_name is required.')
81
+
82
+ if value != value.strip():
83
+ raise HeaderNameFieldError('whitespace_found',
84
+ 'custom_header_name field value cannot have leading or trailing whitespace.')
85
+
86
+ if any(c in string.whitespace for c in value):
87
+ raise HeaderNameFieldError('contains_whitespace',
88
+ 'custom_header_name must not contain any whitespace characters.')
89
+
90
+ if not HEADER_NAME_REGEX.fullmatch(value):
91
+ raise HeaderNameFieldError(
92
+ 'invalid_format',
93
+ 'custom_header_name must match the HTTP token syntax: ASCII letters, digits, or allowed symbols.')
94
+
95
+ return value
96
+
97
+ @field_validator('custom_header_prefix')
98
+ @classmethod
99
+ def validate_custom_header_prefix(cls, value: str) -> str:
100
+ if not value:
101
+ raise HeaderPrefixFieldError('value_missing', 'custom_header_prefix is required.')
102
+
103
+ if value != value.strip():
104
+ raise HeaderPrefixFieldError(
105
+ 'whitespace_found', 'custom_header_prefix field value cannot have '
106
+ 'leading or trailing whitespace.')
107
+
108
+ if any(c in string.whitespace for c in value):
109
+ raise HeaderPrefixFieldError('contains_whitespace',
110
+ 'custom_header_prefix must not contain any whitespace characters.')
111
+
112
+ if not value.isascii():
113
+ raise HeaderPrefixFieldError('invalid_format', 'custom_header_prefix must be ASCII.')
114
+
115
+ return value
116
+
117
+ @field_validator('raw_key', mode='after')
118
+ @classmethod
119
+ def validate_raw_key_after(cls, value: str) -> str:
120
+ if not value:
121
+ raise APIKeyFieldError('value_missing', 'raw_key field value is '
122
+ 'required after construction.')
123
+
124
+ return value
@@ -0,0 +1,26 @@
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.authentication.api_key.api_key_auth_provider_config import APIKeyAuthProviderConfig
17
+ from nat.builder.builder import Builder
18
+ from nat.cli.register_workflow import register_auth_provider
19
+
20
+
21
+ @register_auth_provider(config_type=APIKeyAuthProviderConfig)
22
+ async def api_key_client(config: APIKeyAuthProviderConfig, builder: Builder):
23
+
24
+ from nat.authentication.api_key.api_key_auth_provider import APIKeyAuthProvider
25
+
26
+ yield APIKeyAuthProvider(config=config)
@@ -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,38 @@
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.
15
+
16
+
17
+ class APIKeyFieldError(Exception):
18
+ """Raised when API Key Config api_key field validation fails unexpectedly."""
19
+
20
+ def __init__(self, error_code: str, message: str, *args):
21
+ self.error_code = error_code
22
+ super().__init__(f"[{error_code}] {message}", *args)
23
+
24
+
25
+ class HeaderNameFieldError(Exception):
26
+ """Raised when API Key Config header_name field validation fails unexpectedly."""
27
+
28
+ def __init__(self, error_code: str, message: str, *args):
29
+ self.error_code = error_code
30
+ super().__init__(f"[{error_code}] {message}", *args)
31
+
32
+
33
+ class HeaderPrefixFieldError(Exception):
34
+ """Raised when API Key Config header_prefix field validation fails unexpectedly."""
35
+
36
+ def __init__(self, error_code: str, message: str, *args):
37
+ self.error_code = error_code
38
+ super().__init__(f"[{error_code}] {message}", *args)
File without changes
@@ -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
+ from pydantic import SecretStr
17
+
18
+ from nat.authentication.interfaces import AuthProviderBase
19
+ from nat.builder.context import Context
20
+ from nat.data_models.authentication import AuthenticatedContext
21
+ from nat.data_models.authentication import AuthFlowType
22
+ from nat.data_models.authentication import AuthProviderBaseConfig
23
+ from nat.data_models.authentication import AuthResult
24
+ from nat.data_models.authentication import BasicAuthCred
25
+ from nat.data_models.authentication import BearerTokenCred
26
+
27
+
28
+ class HTTPBasicAuthProvider(AuthProviderBase):
29
+ """
30
+ Abstract base class for HTTP Basic Authentication exchangers.
31
+ """
32
+
33
+ def __init__(self, config: AuthProviderBaseConfig):
34
+ """
35
+ Initialize the HTTP Basic Auth Exchanger with the given configuration.
36
+ """
37
+ super().__init__(config)
38
+
39
+ self._authenticated_tokens: dict[str, AuthResult] = {}
40
+
41
+ async def authenticate(self, user_id: str | None = None) -> AuthResult:
42
+ """
43
+ Performs simple HTTP Authentication using the provided user ID.
44
+ """
45
+
46
+ context = Context.get()
47
+
48
+ if user_id is None and hasattr(context, "metadata") and hasattr(
49
+ context.metadata, "cookies") and context.metadata.cookies is not None:
50
+ session_id = context.metadata.cookies.get("nat-session", None)
51
+ if not session_id:
52
+ raise RuntimeError("Authentication failed. No session ID found. Cannot identify user.")
53
+
54
+ user_id = session_id
55
+
56
+ if user_id and user_id in self._authenticated_tokens:
57
+ return self._authenticated_tokens[user_id]
58
+
59
+ auth_callback = context.user_auth_callback
60
+
61
+ try:
62
+ auth_context: AuthenticatedContext = await auth_callback(self.config, AuthFlowType.HTTP_BASIC)
63
+ except RuntimeError as e:
64
+ raise RuntimeError(f"Authentication callback failed: {str(e)}. Did you forget to set a "
65
+ f"callback handler for your frontend?") from e
66
+
67
+ basic_auth_credentials = BasicAuthCred(username=SecretStr(auth_context.metadata.get("username", "")),
68
+ password=SecretStr(auth_context.metadata.get("password", "")))
69
+
70
+ # Get the auth token from the headers of auth context
71
+ bearer_token = auth_context.headers.get("Authorization", "").split(" ")[-1]
72
+ if not bearer_token:
73
+ raise RuntimeError("Authentication failed: No Authorization header found in the response.")
74
+
75
+ bearer_token_cred = BearerTokenCred(token=SecretStr(bearer_token), scheme="Basic")
76
+
77
+ auth_result = AuthResult(credentials=[basic_auth_credentials, bearer_token_cred])
78
+
79
+ self._authenticated_tokens[user_id] = auth_result
80
+
81
+ return auth_result
@@ -0,0 +1,30 @@
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.builder.builder import Builder
17
+ from nat.cli.register_workflow import register_auth_provider
18
+ from nat.data_models.authentication import AuthProviderBaseConfig
19
+
20
+
21
+ class HTTPBasicAuthProviderConfig(AuthProviderBaseConfig, name="http_basic"):
22
+ pass
23
+
24
+
25
+ @register_auth_provider(config_type=HTTPBasicAuthProviderConfig)
26
+ async def http_basic_auth_provider(config: HTTPBasicAuthProviderConfig, builder: Builder):
27
+
28
+ from nat.authentication.http_basic_auth.http_basic_auth_provider import HTTPBasicAuthProvider
29
+
30
+ yield HTTPBasicAuthProvider(config)
@@ -0,0 +1,93 @@
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.
15
+
16
+ import typing
17
+ from abc import ABC
18
+ from abc import abstractmethod
19
+
20
+ from nat.data_models.authentication import AuthenticatedContext
21
+ from nat.data_models.authentication import AuthFlowType
22
+ from nat.data_models.authentication import AuthProviderBaseConfig
23
+ from nat.data_models.authentication import AuthProviderBaseConfigT
24
+ from nat.data_models.authentication import AuthResult
25
+
26
+ AUTHORIZATION_HEADER = "Authorization"
27
+
28
+
29
+ class AuthProviderBase(typing.Generic[AuthProviderBaseConfigT], ABC):
30
+ """
31
+ Base class for authenticating to API services.
32
+ This class provides an interface for authenticating to API services.
33
+ """
34
+
35
+ def __init__(self, config: AuthProviderBaseConfigT):
36
+ """
37
+ Initialize the AuthProviderBase with the given configuration.
38
+
39
+ Args:
40
+ config (AuthProviderBaseConfig): Configuration items for authentication.
41
+ """
42
+ self._config = config
43
+
44
+ @property
45
+ def config(self) -> AuthProviderBaseConfigT:
46
+ """
47
+ Returns the auth provider configuration object.
48
+
49
+ Returns
50
+ -------
51
+ AuthProviderBaseConfigT
52
+ The auth provider configuration object.
53
+ """
54
+ return self._config
55
+
56
+ @abstractmethod
57
+ async def authenticate(self, user_id: str | None = None) -> AuthResult:
58
+ """
59
+ Perform the authentication process for the client.
60
+
61
+ This method handles the necessary steps to authenticate the client with the
62
+ target API service, which may include obtaining tokens, refreshing credentials,
63
+ or completing multi-step authentication flows.
64
+
65
+ Raises:
66
+ NotImplementedError: Must be implemented by subclasses.
67
+ """
68
+ # This method will call the frontend FlowHandlerBase `authenticate` method
69
+ pass
70
+
71
+
72
+ class FlowHandlerBase(ABC):
73
+ """
74
+ Handles front-end specifc flows for authentication clients.
75
+
76
+ Each front end will define a FlowHandler that will implement the authenticate method.
77
+
78
+ The `authenticate` method will be stored as the callback in the ContextState.user_auth_callback
79
+ """
80
+
81
+ @abstractmethod
82
+ async def authenticate(self, config: AuthProviderBaseConfig, method: AuthFlowType) -> AuthenticatedContext:
83
+ """
84
+ Perform the authentication process for the client.
85
+
86
+ This method handles the necessary steps to authenticate the client with the
87
+ target API service, which may include obtaining tokens, refreshing credentials,
88
+ or completing multistep authentication flows.
89
+
90
+ Raises:
91
+ NotImplementedError: Must be implemented by subclasses.
92
+ """
93
+ pass
@@ -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,107 @@
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 datetime import datetime
17
+ from datetime import timezone
18
+
19
+ from authlib.integrations.httpx_client import OAuth2Client as AuthlibOAuth2Client
20
+ from pydantic import SecretStr
21
+
22
+ from nat.authentication.interfaces import AuthProviderBase
23
+ from nat.authentication.oauth2.oauth2_auth_code_flow_provider_config import OAuth2AuthCodeFlowProviderConfig
24
+ from nat.builder.context import Context
25
+ from nat.data_models.authentication import AuthFlowType
26
+ from nat.data_models.authentication import AuthResult
27
+ from nat.data_models.authentication import BearerTokenCred
28
+
29
+
30
+ class OAuth2AuthCodeFlowProvider(AuthProviderBase[OAuth2AuthCodeFlowProviderConfig]):
31
+
32
+ def __init__(self, config: OAuth2AuthCodeFlowProviderConfig):
33
+ super().__init__(config)
34
+ self._authenticated_tokens: dict[str, AuthResult] = {}
35
+ self._context = Context.get()
36
+
37
+ async def _attempt_token_refresh(self, user_id: str, auth_result: AuthResult) -> AuthResult | None:
38
+ refresh_token = auth_result.raw.get("refresh_token")
39
+ if not isinstance(refresh_token, str):
40
+ return None
41
+
42
+ with AuthlibOAuth2Client(
43
+ client_id=self.config.client_id,
44
+ client_secret=self.config.client_secret,
45
+ ) as client:
46
+ try:
47
+ new_token_data = client.refresh_token(self.config.token_url, refresh_token=refresh_token)
48
+ except Exception:
49
+ # On any failure, we'll fall back to the full auth flow.
50
+ return None
51
+
52
+ expires_at_ts = new_token_data.get("expires_at")
53
+ new_expires_at = datetime.fromtimestamp(expires_at_ts, tz=timezone.utc) if expires_at_ts else None
54
+
55
+ new_auth_result = AuthResult(
56
+ credentials=[BearerTokenCred(token=SecretStr(new_token_data["access_token"]))],
57
+ token_expires_at=new_expires_at,
58
+ raw=new_token_data,
59
+ )
60
+
61
+ self._authenticated_tokens[user_id] = new_auth_result
62
+
63
+ return new_auth_result
64
+
65
+ async def authenticate(self, user_id: str | None = None) -> AuthResult:
66
+ if user_id is None and hasattr(Context.get(), "metadata") and hasattr(
67
+ Context.get().metadata, "cookies") and Context.get().metadata.cookies is not None:
68
+ session_id = Context.get().metadata.cookies.get("nat-session", None)
69
+ if not session_id:
70
+ raise RuntimeError("Authentication failed. No session ID found. Cannot identify user.")
71
+
72
+ user_id = session_id
73
+
74
+ if user_id and user_id in self._authenticated_tokens:
75
+ auth_result = self._authenticated_tokens[user_id]
76
+ if not auth_result.is_expired():
77
+ return auth_result
78
+
79
+ refreshed_auth_result = await self._attempt_token_refresh(user_id, auth_result)
80
+ if refreshed_auth_result:
81
+ return refreshed_auth_result
82
+
83
+ auth_callback = self._context.user_auth_callback
84
+ if not auth_callback:
85
+ raise RuntimeError("Authentication callback not set on Context.")
86
+
87
+ try:
88
+ authenticated_context = await auth_callback(self.config, AuthFlowType.OAUTH2_AUTHORIZATION_CODE)
89
+ except Exception as e:
90
+ raise RuntimeError(f"Authentication callback failed: {e}") from e
91
+
92
+ auth_header = authenticated_context.headers.get("Authorization", "")
93
+ if not auth_header.startswith("Bearer "):
94
+ raise RuntimeError("Invalid Authorization header")
95
+
96
+ token = auth_header.split(" ")[1]
97
+
98
+ auth_result = AuthResult(
99
+ credentials=[BearerTokenCred(token=SecretStr(token))],
100
+ token_expires_at=authenticated_context.metadata.get("expires_at"),
101
+ raw=authenticated_context.metadata.get("raw_token"),
102
+ )
103
+
104
+ if user_id:
105
+ self._authenticated_tokens[user_id] = auth_result
106
+
107
+ return auth_result
@@ -0,0 +1,39 @@
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.
15
+
16
+ from pydantic import Field
17
+
18
+ from nat.data_models.authentication import AuthProviderBaseConfig
19
+
20
+
21
+ class OAuth2AuthCodeFlowProviderConfig(AuthProviderBaseConfig, name="oauth2_auth_code_flow"):
22
+
23
+ client_id: str = Field(description="The client ID for OAuth 2.0 authentication.")
24
+ client_secret: str = Field(description="The secret associated with the client_id.")
25
+ authorization_url: str = Field(description="The authorization URL for OAuth 2.0 authentication.")
26
+ token_url: str = Field(description="The token URL for OAuth 2.0 authentication.")
27
+ token_endpoint_auth_method: str | None = Field(
28
+ description=("The authentication method for the token endpoint. "
29
+ "Usually one of `client_secret_post` or `client_secret_basic`."),
30
+ default=None)
31
+ redirect_uri: str = Field(description="The redirect URI for OAuth 2.0 authentication. Must match the registered "
32
+ "redirect URI with the OAuth provider.")
33
+ scopes: list[str] = Field(description="The scopes for OAuth 2.0 authentication.", default_factory=list)
34
+ use_pkce: bool = Field(default=False,
35
+ description="Whether to use PKCE (Proof Key for Code Exchange) in the OAuth 2.0 flow.")
36
+
37
+ authorization_kwargs: dict[str, str] | None = Field(description=("Additional keyword arguments for the "
38
+ "authorization request."),
39
+ default=None)
@@ -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.authentication.oauth2.oauth2_auth_code_flow_provider_config import OAuth2AuthCodeFlowProviderConfig
17
+ from nat.builder.builder import Builder
18
+ from nat.cli.register_workflow import register_auth_provider
19
+
20
+
21
+ @register_auth_provider(config_type=OAuth2AuthCodeFlowProviderConfig)
22
+ async def oauth2_client(authentication_provider: OAuth2AuthCodeFlowProviderConfig, builder: Builder):
23
+ from nat.authentication.oauth2.oauth2_auth_code_flow_provider import OAuth2AuthCodeFlowProvider
24
+
25
+ yield OAuth2AuthCodeFlowProvider(authentication_provider)
@@ -0,0 +1,21 @@
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
+ # pylint: disable=unused-import
17
+ # flake8: noqa
18
+
19
+ from nat.authentication.api_key import register as register_api_key
20
+ from nat.authentication.http_basic_auth import register as register_http_basic_auth
21
+ from nat.authentication.oauth2 import register as register_oauth2
File without changes