nvidia-nat 1.1.0a20251020__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 (480) hide show
  1. aiq/__init__.py +66 -0
  2. nat/agent/__init__.py +0 -0
  3. nat/agent/base.py +265 -0
  4. nat/agent/dual_node.py +72 -0
  5. nat/agent/prompt_optimizer/__init__.py +0 -0
  6. nat/agent/prompt_optimizer/prompt.py +68 -0
  7. nat/agent/prompt_optimizer/register.py +149 -0
  8. nat/agent/react_agent/__init__.py +0 -0
  9. nat/agent/react_agent/agent.py +394 -0
  10. nat/agent/react_agent/output_parser.py +104 -0
  11. nat/agent/react_agent/prompt.py +44 -0
  12. nat/agent/react_agent/register.py +168 -0
  13. nat/agent/reasoning_agent/__init__.py +0 -0
  14. nat/agent/reasoning_agent/reasoning_agent.py +227 -0
  15. nat/agent/register.py +23 -0
  16. nat/agent/rewoo_agent/__init__.py +0 -0
  17. nat/agent/rewoo_agent/agent.py +593 -0
  18. nat/agent/rewoo_agent/prompt.py +107 -0
  19. nat/agent/rewoo_agent/register.py +175 -0
  20. nat/agent/tool_calling_agent/__init__.py +0 -0
  21. nat/agent/tool_calling_agent/agent.py +246 -0
  22. nat/agent/tool_calling_agent/register.py +129 -0
  23. nat/authentication/__init__.py +14 -0
  24. nat/authentication/api_key/__init__.py +14 -0
  25. nat/authentication/api_key/api_key_auth_provider.py +96 -0
  26. nat/authentication/api_key/api_key_auth_provider_config.py +124 -0
  27. nat/authentication/api_key/register.py +26 -0
  28. nat/authentication/credential_validator/__init__.py +14 -0
  29. nat/authentication/credential_validator/bearer_token_validator.py +557 -0
  30. nat/authentication/exceptions/__init__.py +14 -0
  31. nat/authentication/exceptions/api_key_exceptions.py +38 -0
  32. nat/authentication/http_basic_auth/__init__.py +0 -0
  33. nat/authentication/http_basic_auth/http_basic_auth_provider.py +81 -0
  34. nat/authentication/http_basic_auth/register.py +30 -0
  35. nat/authentication/interfaces.py +96 -0
  36. nat/authentication/oauth2/__init__.py +14 -0
  37. nat/authentication/oauth2/oauth2_auth_code_flow_provider.py +140 -0
  38. nat/authentication/oauth2/oauth2_auth_code_flow_provider_config.py +39 -0
  39. nat/authentication/oauth2/oauth2_resource_server_config.py +124 -0
  40. nat/authentication/oauth2/register.py +25 -0
  41. nat/authentication/register.py +20 -0
  42. nat/builder/__init__.py +0 -0
  43. nat/builder/builder.py +317 -0
  44. nat/builder/component_utils.py +320 -0
  45. nat/builder/context.py +321 -0
  46. nat/builder/embedder.py +24 -0
  47. nat/builder/eval_builder.py +166 -0
  48. nat/builder/evaluator.py +29 -0
  49. nat/builder/framework_enum.py +25 -0
  50. nat/builder/front_end.py +73 -0
  51. nat/builder/function.py +714 -0
  52. nat/builder/function_base.py +380 -0
  53. nat/builder/function_info.py +625 -0
  54. nat/builder/intermediate_step_manager.py +206 -0
  55. nat/builder/llm.py +25 -0
  56. nat/builder/retriever.py +25 -0
  57. nat/builder/user_interaction_manager.py +78 -0
  58. nat/builder/workflow.py +160 -0
  59. nat/builder/workflow_builder.py +1365 -0
  60. nat/cli/__init__.py +14 -0
  61. nat/cli/cli_utils/__init__.py +0 -0
  62. nat/cli/cli_utils/config_override.py +231 -0
  63. nat/cli/cli_utils/validation.py +37 -0
  64. nat/cli/commands/__init__.py +0 -0
  65. nat/cli/commands/configure/__init__.py +0 -0
  66. nat/cli/commands/configure/channel/__init__.py +0 -0
  67. nat/cli/commands/configure/channel/add.py +28 -0
  68. nat/cli/commands/configure/channel/channel.py +34 -0
  69. nat/cli/commands/configure/channel/remove.py +30 -0
  70. nat/cli/commands/configure/channel/update.py +30 -0
  71. nat/cli/commands/configure/configure.py +33 -0
  72. nat/cli/commands/evaluate.py +139 -0
  73. nat/cli/commands/info/__init__.py +14 -0
  74. nat/cli/commands/info/info.py +47 -0
  75. nat/cli/commands/info/list_channels.py +32 -0
  76. nat/cli/commands/info/list_components.py +128 -0
  77. nat/cli/commands/mcp/__init__.py +14 -0
  78. nat/cli/commands/mcp/mcp.py +986 -0
  79. nat/cli/commands/object_store/__init__.py +14 -0
  80. nat/cli/commands/object_store/object_store.py +227 -0
  81. nat/cli/commands/optimize.py +90 -0
  82. nat/cli/commands/registry/__init__.py +14 -0
  83. nat/cli/commands/registry/publish.py +88 -0
  84. nat/cli/commands/registry/pull.py +118 -0
  85. nat/cli/commands/registry/registry.py +36 -0
  86. nat/cli/commands/registry/remove.py +108 -0
  87. nat/cli/commands/registry/search.py +153 -0
  88. nat/cli/commands/sizing/__init__.py +14 -0
  89. nat/cli/commands/sizing/calc.py +297 -0
  90. nat/cli/commands/sizing/sizing.py +27 -0
  91. nat/cli/commands/start.py +257 -0
  92. nat/cli/commands/uninstall.py +81 -0
  93. nat/cli/commands/validate.py +47 -0
  94. nat/cli/commands/workflow/__init__.py +14 -0
  95. nat/cli/commands/workflow/templates/__init__.py.j2 +0 -0
  96. nat/cli/commands/workflow/templates/config.yml.j2 +17 -0
  97. nat/cli/commands/workflow/templates/pyproject.toml.j2 +25 -0
  98. nat/cli/commands/workflow/templates/register.py.j2 +4 -0
  99. nat/cli/commands/workflow/templates/workflow.py.j2 +50 -0
  100. nat/cli/commands/workflow/workflow.py +37 -0
  101. nat/cli/commands/workflow/workflow_commands.py +403 -0
  102. nat/cli/entrypoint.py +141 -0
  103. nat/cli/main.py +60 -0
  104. nat/cli/register_workflow.py +522 -0
  105. nat/cli/type_registry.py +1069 -0
  106. nat/control_flow/__init__.py +0 -0
  107. nat/control_flow/register.py +20 -0
  108. nat/control_flow/router_agent/__init__.py +0 -0
  109. nat/control_flow/router_agent/agent.py +329 -0
  110. nat/control_flow/router_agent/prompt.py +48 -0
  111. nat/control_flow/router_agent/register.py +91 -0
  112. nat/control_flow/sequential_executor.py +166 -0
  113. nat/data_models/__init__.py +14 -0
  114. nat/data_models/agent.py +34 -0
  115. nat/data_models/api_server.py +843 -0
  116. nat/data_models/authentication.py +245 -0
  117. nat/data_models/common.py +171 -0
  118. nat/data_models/component.py +60 -0
  119. nat/data_models/component_ref.py +179 -0
  120. nat/data_models/config.py +434 -0
  121. nat/data_models/dataset_handler.py +169 -0
  122. nat/data_models/discovery_metadata.py +305 -0
  123. nat/data_models/embedder.py +27 -0
  124. nat/data_models/evaluate.py +130 -0
  125. nat/data_models/evaluator.py +26 -0
  126. nat/data_models/front_end.py +26 -0
  127. nat/data_models/function.py +64 -0
  128. nat/data_models/function_dependencies.py +80 -0
  129. nat/data_models/gated_field_mixin.py +242 -0
  130. nat/data_models/interactive.py +246 -0
  131. nat/data_models/intermediate_step.py +302 -0
  132. nat/data_models/invocation_node.py +38 -0
  133. nat/data_models/llm.py +27 -0
  134. nat/data_models/logging.py +26 -0
  135. nat/data_models/memory.py +27 -0
  136. nat/data_models/object_store.py +44 -0
  137. nat/data_models/optimizable.py +119 -0
  138. nat/data_models/optimizer.py +149 -0
  139. nat/data_models/profiler.py +54 -0
  140. nat/data_models/registry_handler.py +26 -0
  141. nat/data_models/retriever.py +30 -0
  142. nat/data_models/retry_mixin.py +35 -0
  143. nat/data_models/span.py +228 -0
  144. nat/data_models/step_adaptor.py +64 -0
  145. nat/data_models/streaming.py +33 -0
  146. nat/data_models/swe_bench_model.py +54 -0
  147. nat/data_models/telemetry_exporter.py +26 -0
  148. nat/data_models/temperature_mixin.py +44 -0
  149. nat/data_models/thinking_mixin.py +86 -0
  150. nat/data_models/top_p_mixin.py +44 -0
  151. nat/data_models/ttc_strategy.py +30 -0
  152. nat/embedder/__init__.py +0 -0
  153. nat/embedder/azure_openai_embedder.py +46 -0
  154. nat/embedder/nim_embedder.py +59 -0
  155. nat/embedder/openai_embedder.py +42 -0
  156. nat/embedder/register.py +22 -0
  157. nat/eval/__init__.py +14 -0
  158. nat/eval/config.py +62 -0
  159. nat/eval/dataset_handler/__init__.py +0 -0
  160. nat/eval/dataset_handler/dataset_downloader.py +106 -0
  161. nat/eval/dataset_handler/dataset_filter.py +52 -0
  162. nat/eval/dataset_handler/dataset_handler.py +431 -0
  163. nat/eval/evaluate.py +565 -0
  164. nat/eval/evaluator/__init__.py +14 -0
  165. nat/eval/evaluator/base_evaluator.py +77 -0
  166. nat/eval/evaluator/evaluator_model.py +58 -0
  167. nat/eval/intermediate_step_adapter.py +99 -0
  168. nat/eval/rag_evaluator/__init__.py +0 -0
  169. nat/eval/rag_evaluator/evaluate.py +178 -0
  170. nat/eval/rag_evaluator/register.py +143 -0
  171. nat/eval/register.py +26 -0
  172. nat/eval/remote_workflow.py +133 -0
  173. nat/eval/runners/__init__.py +14 -0
  174. nat/eval/runners/config.py +39 -0
  175. nat/eval/runners/multi_eval_runner.py +54 -0
  176. nat/eval/runtime_evaluator/__init__.py +14 -0
  177. nat/eval/runtime_evaluator/evaluate.py +123 -0
  178. nat/eval/runtime_evaluator/register.py +100 -0
  179. nat/eval/runtime_event_subscriber.py +52 -0
  180. nat/eval/swe_bench_evaluator/__init__.py +0 -0
  181. nat/eval/swe_bench_evaluator/evaluate.py +215 -0
  182. nat/eval/swe_bench_evaluator/register.py +36 -0
  183. nat/eval/trajectory_evaluator/__init__.py +0 -0
  184. nat/eval/trajectory_evaluator/evaluate.py +75 -0
  185. nat/eval/trajectory_evaluator/register.py +40 -0
  186. nat/eval/tunable_rag_evaluator/__init__.py +0 -0
  187. nat/eval/tunable_rag_evaluator/evaluate.py +242 -0
  188. nat/eval/tunable_rag_evaluator/register.py +52 -0
  189. nat/eval/usage_stats.py +41 -0
  190. nat/eval/utils/__init__.py +0 -0
  191. nat/eval/utils/eval_trace_ctx.py +89 -0
  192. nat/eval/utils/output_uploader.py +140 -0
  193. nat/eval/utils/tqdm_position_registry.py +40 -0
  194. nat/eval/utils/weave_eval.py +193 -0
  195. nat/experimental/__init__.py +0 -0
  196. nat/experimental/decorators/__init__.py +0 -0
  197. nat/experimental/decorators/experimental_warning_decorator.py +154 -0
  198. nat/experimental/test_time_compute/__init__.py +0 -0
  199. nat/experimental/test_time_compute/editing/__init__.py +0 -0
  200. nat/experimental/test_time_compute/editing/iterative_plan_refinement_editor.py +147 -0
  201. nat/experimental/test_time_compute/editing/llm_as_a_judge_editor.py +204 -0
  202. nat/experimental/test_time_compute/editing/motivation_aware_summarization.py +107 -0
  203. nat/experimental/test_time_compute/functions/__init__.py +0 -0
  204. nat/experimental/test_time_compute/functions/execute_score_select_function.py +105 -0
  205. nat/experimental/test_time_compute/functions/plan_select_execute_function.py +228 -0
  206. nat/experimental/test_time_compute/functions/ttc_tool_orchestration_function.py +205 -0
  207. nat/experimental/test_time_compute/functions/ttc_tool_wrapper_function.py +146 -0
  208. nat/experimental/test_time_compute/models/__init__.py +0 -0
  209. nat/experimental/test_time_compute/models/editor_config.py +132 -0
  210. nat/experimental/test_time_compute/models/scoring_config.py +112 -0
  211. nat/experimental/test_time_compute/models/search_config.py +120 -0
  212. nat/experimental/test_time_compute/models/selection_config.py +154 -0
  213. nat/experimental/test_time_compute/models/stage_enums.py +43 -0
  214. nat/experimental/test_time_compute/models/strategy_base.py +67 -0
  215. nat/experimental/test_time_compute/models/tool_use_config.py +41 -0
  216. nat/experimental/test_time_compute/models/ttc_item.py +48 -0
  217. nat/experimental/test_time_compute/register.py +35 -0
  218. nat/experimental/test_time_compute/scoring/__init__.py +0 -0
  219. nat/experimental/test_time_compute/scoring/llm_based_agent_scorer.py +168 -0
  220. nat/experimental/test_time_compute/scoring/llm_based_plan_scorer.py +168 -0
  221. nat/experimental/test_time_compute/scoring/motivation_aware_scorer.py +111 -0
  222. nat/experimental/test_time_compute/search/__init__.py +0 -0
  223. nat/experimental/test_time_compute/search/multi_llm_planner.py +128 -0
  224. nat/experimental/test_time_compute/search/multi_query_retrieval_search.py +122 -0
  225. nat/experimental/test_time_compute/search/single_shot_multi_plan_planner.py +128 -0
  226. nat/experimental/test_time_compute/selection/__init__.py +0 -0
  227. nat/experimental/test_time_compute/selection/best_of_n_selector.py +63 -0
  228. nat/experimental/test_time_compute/selection/llm_based_agent_output_selector.py +131 -0
  229. nat/experimental/test_time_compute/selection/llm_based_output_merging_selector.py +157 -0
  230. nat/experimental/test_time_compute/selection/llm_based_plan_selector.py +128 -0
  231. nat/experimental/test_time_compute/selection/threshold_selector.py +58 -0
  232. nat/front_ends/__init__.py +14 -0
  233. nat/front_ends/console/__init__.py +14 -0
  234. nat/front_ends/console/authentication_flow_handler.py +285 -0
  235. nat/front_ends/console/console_front_end_config.py +32 -0
  236. nat/front_ends/console/console_front_end_plugin.py +108 -0
  237. nat/front_ends/console/register.py +25 -0
  238. nat/front_ends/cron/__init__.py +14 -0
  239. nat/front_ends/fastapi/__init__.py +14 -0
  240. nat/front_ends/fastapi/auth_flow_handlers/__init__.py +0 -0
  241. nat/front_ends/fastapi/auth_flow_handlers/http_flow_handler.py +27 -0
  242. nat/front_ends/fastapi/auth_flow_handlers/websocket_flow_handler.py +142 -0
  243. nat/front_ends/fastapi/dask_client_mixin.py +65 -0
  244. nat/front_ends/fastapi/fastapi_front_end_config.py +272 -0
  245. nat/front_ends/fastapi/fastapi_front_end_controller.py +68 -0
  246. nat/front_ends/fastapi/fastapi_front_end_plugin.py +247 -0
  247. nat/front_ends/fastapi/fastapi_front_end_plugin_worker.py +1257 -0
  248. nat/front_ends/fastapi/html_snippets/__init__.py +14 -0
  249. nat/front_ends/fastapi/html_snippets/auth_code_grant_success.py +35 -0
  250. nat/front_ends/fastapi/intermediate_steps_subscriber.py +80 -0
  251. nat/front_ends/fastapi/job_store.py +602 -0
  252. nat/front_ends/fastapi/main.py +64 -0
  253. nat/front_ends/fastapi/message_handler.py +344 -0
  254. nat/front_ends/fastapi/message_validator.py +351 -0
  255. nat/front_ends/fastapi/register.py +25 -0
  256. nat/front_ends/fastapi/response_helpers.py +195 -0
  257. nat/front_ends/fastapi/step_adaptor.py +319 -0
  258. nat/front_ends/fastapi/utils.py +57 -0
  259. nat/front_ends/mcp/__init__.py +14 -0
  260. nat/front_ends/mcp/introspection_token_verifier.py +73 -0
  261. nat/front_ends/mcp/mcp_front_end_config.py +90 -0
  262. nat/front_ends/mcp/mcp_front_end_plugin.py +113 -0
  263. nat/front_ends/mcp/mcp_front_end_plugin_worker.py +268 -0
  264. nat/front_ends/mcp/memory_profiler.py +320 -0
  265. nat/front_ends/mcp/register.py +27 -0
  266. nat/front_ends/mcp/tool_converter.py +290 -0
  267. nat/front_ends/register.py +21 -0
  268. nat/front_ends/simple_base/__init__.py +14 -0
  269. nat/front_ends/simple_base/simple_front_end_plugin_base.py +56 -0
  270. nat/llm/__init__.py +0 -0
  271. nat/llm/aws_bedrock_llm.py +69 -0
  272. nat/llm/azure_openai_llm.py +57 -0
  273. nat/llm/litellm_llm.py +69 -0
  274. nat/llm/nim_llm.py +58 -0
  275. nat/llm/openai_llm.py +54 -0
  276. nat/llm/register.py +27 -0
  277. nat/llm/utils/__init__.py +14 -0
  278. nat/llm/utils/env_config_value.py +93 -0
  279. nat/llm/utils/error.py +17 -0
  280. nat/llm/utils/thinking.py +215 -0
  281. nat/memory/__init__.py +20 -0
  282. nat/memory/interfaces.py +183 -0
  283. nat/memory/models.py +112 -0
  284. nat/meta/pypi.md +58 -0
  285. nat/object_store/__init__.py +20 -0
  286. nat/object_store/in_memory_object_store.py +76 -0
  287. nat/object_store/interfaces.py +84 -0
  288. nat/object_store/models.py +38 -0
  289. nat/object_store/register.py +19 -0
  290. nat/observability/__init__.py +14 -0
  291. nat/observability/exporter/__init__.py +14 -0
  292. nat/observability/exporter/base_exporter.py +449 -0
  293. nat/observability/exporter/exporter.py +78 -0
  294. nat/observability/exporter/file_exporter.py +33 -0
  295. nat/observability/exporter/processing_exporter.py +550 -0
  296. nat/observability/exporter/raw_exporter.py +52 -0
  297. nat/observability/exporter/span_exporter.py +308 -0
  298. nat/observability/exporter_manager.py +335 -0
  299. nat/observability/mixin/__init__.py +14 -0
  300. nat/observability/mixin/batch_config_mixin.py +26 -0
  301. nat/observability/mixin/collector_config_mixin.py +23 -0
  302. nat/observability/mixin/file_mixin.py +288 -0
  303. nat/observability/mixin/file_mode.py +23 -0
  304. nat/observability/mixin/redaction_config_mixin.py +42 -0
  305. nat/observability/mixin/resource_conflict_mixin.py +134 -0
  306. nat/observability/mixin/serialize_mixin.py +61 -0
  307. nat/observability/mixin/tagging_config_mixin.py +62 -0
  308. nat/observability/mixin/type_introspection_mixin.py +496 -0
  309. nat/observability/processor/__init__.py +14 -0
  310. nat/observability/processor/batching_processor.py +308 -0
  311. nat/observability/processor/callback_processor.py +42 -0
  312. nat/observability/processor/falsy_batch_filter_processor.py +55 -0
  313. nat/observability/processor/intermediate_step_serializer.py +28 -0
  314. nat/observability/processor/processor.py +74 -0
  315. nat/observability/processor/processor_factory.py +70 -0
  316. nat/observability/processor/redaction/__init__.py +24 -0
  317. nat/observability/processor/redaction/contextual_redaction_processor.py +125 -0
  318. nat/observability/processor/redaction/contextual_span_redaction_processor.py +66 -0
  319. nat/observability/processor/redaction/redaction_processor.py +177 -0
  320. nat/observability/processor/redaction/span_header_redaction_processor.py +92 -0
  321. nat/observability/processor/span_tagging_processor.py +68 -0
  322. nat/observability/register.py +114 -0
  323. nat/observability/utils/__init__.py +14 -0
  324. nat/observability/utils/dict_utils.py +236 -0
  325. nat/observability/utils/time_utils.py +31 -0
  326. nat/plugins/.namespace +1 -0
  327. nat/profiler/__init__.py +0 -0
  328. nat/profiler/calc/__init__.py +14 -0
  329. nat/profiler/calc/calc_runner.py +626 -0
  330. nat/profiler/calc/calculations.py +288 -0
  331. nat/profiler/calc/data_models.py +188 -0
  332. nat/profiler/calc/plot.py +345 -0
  333. nat/profiler/callbacks/__init__.py +0 -0
  334. nat/profiler/callbacks/agno_callback_handler.py +295 -0
  335. nat/profiler/callbacks/base_callback_class.py +20 -0
  336. nat/profiler/callbacks/langchain_callback_handler.py +297 -0
  337. nat/profiler/callbacks/llama_index_callback_handler.py +205 -0
  338. nat/profiler/callbacks/semantic_kernel_callback_handler.py +238 -0
  339. nat/profiler/callbacks/token_usage_base_model.py +27 -0
  340. nat/profiler/data_frame_row.py +51 -0
  341. nat/profiler/data_models.py +24 -0
  342. nat/profiler/decorators/__init__.py +0 -0
  343. nat/profiler/decorators/framework_wrapper.py +180 -0
  344. nat/profiler/decorators/function_tracking.py +411 -0
  345. nat/profiler/forecasting/__init__.py +0 -0
  346. nat/profiler/forecasting/config.py +18 -0
  347. nat/profiler/forecasting/model_trainer.py +75 -0
  348. nat/profiler/forecasting/models/__init__.py +22 -0
  349. nat/profiler/forecasting/models/forecasting_base_model.py +42 -0
  350. nat/profiler/forecasting/models/linear_model.py +197 -0
  351. nat/profiler/forecasting/models/random_forest_regressor.py +269 -0
  352. nat/profiler/inference_metrics_model.py +28 -0
  353. nat/profiler/inference_optimization/__init__.py +0 -0
  354. nat/profiler/inference_optimization/bottleneck_analysis/__init__.py +0 -0
  355. nat/profiler/inference_optimization/bottleneck_analysis/nested_stack_analysis.py +460 -0
  356. nat/profiler/inference_optimization/bottleneck_analysis/simple_stack_analysis.py +258 -0
  357. nat/profiler/inference_optimization/data_models.py +386 -0
  358. nat/profiler/inference_optimization/experimental/__init__.py +0 -0
  359. nat/profiler/inference_optimization/experimental/concurrency_spike_analysis.py +468 -0
  360. nat/profiler/inference_optimization/experimental/prefix_span_analysis.py +404 -0
  361. nat/profiler/inference_optimization/llm_metrics.py +212 -0
  362. nat/profiler/inference_optimization/prompt_caching.py +163 -0
  363. nat/profiler/inference_optimization/token_uniqueness.py +107 -0
  364. nat/profiler/inference_optimization/workflow_runtimes.py +72 -0
  365. nat/profiler/intermediate_property_adapter.py +102 -0
  366. nat/profiler/parameter_optimization/__init__.py +0 -0
  367. nat/profiler/parameter_optimization/optimizable_utils.py +93 -0
  368. nat/profiler/parameter_optimization/optimizer_runtime.py +67 -0
  369. nat/profiler/parameter_optimization/parameter_optimizer.py +153 -0
  370. nat/profiler/parameter_optimization/parameter_selection.py +107 -0
  371. nat/profiler/parameter_optimization/pareto_visualizer.py +380 -0
  372. nat/profiler/parameter_optimization/prompt_optimizer.py +384 -0
  373. nat/profiler/parameter_optimization/update_helpers.py +66 -0
  374. nat/profiler/profile_runner.py +478 -0
  375. nat/profiler/utils.py +186 -0
  376. nat/registry_handlers/__init__.py +0 -0
  377. nat/registry_handlers/local/__init__.py +0 -0
  378. nat/registry_handlers/local/local_handler.py +176 -0
  379. nat/registry_handlers/local/register_local.py +37 -0
  380. nat/registry_handlers/metadata_factory.py +60 -0
  381. nat/registry_handlers/package_utils.py +570 -0
  382. nat/registry_handlers/pypi/__init__.py +0 -0
  383. nat/registry_handlers/pypi/pypi_handler.py +248 -0
  384. nat/registry_handlers/pypi/register_pypi.py +40 -0
  385. nat/registry_handlers/register.py +20 -0
  386. nat/registry_handlers/registry_handler_base.py +157 -0
  387. nat/registry_handlers/rest/__init__.py +0 -0
  388. nat/registry_handlers/rest/register_rest.py +56 -0
  389. nat/registry_handlers/rest/rest_handler.py +236 -0
  390. nat/registry_handlers/schemas/__init__.py +0 -0
  391. nat/registry_handlers/schemas/headers.py +42 -0
  392. nat/registry_handlers/schemas/package.py +68 -0
  393. nat/registry_handlers/schemas/publish.py +68 -0
  394. nat/registry_handlers/schemas/pull.py +82 -0
  395. nat/registry_handlers/schemas/remove.py +36 -0
  396. nat/registry_handlers/schemas/search.py +91 -0
  397. nat/registry_handlers/schemas/status.py +47 -0
  398. nat/retriever/__init__.py +0 -0
  399. nat/retriever/interface.py +41 -0
  400. nat/retriever/milvus/__init__.py +14 -0
  401. nat/retriever/milvus/register.py +81 -0
  402. nat/retriever/milvus/retriever.py +228 -0
  403. nat/retriever/models.py +77 -0
  404. nat/retriever/nemo_retriever/__init__.py +14 -0
  405. nat/retriever/nemo_retriever/register.py +60 -0
  406. nat/retriever/nemo_retriever/retriever.py +190 -0
  407. nat/retriever/register.py +21 -0
  408. nat/runtime/__init__.py +14 -0
  409. nat/runtime/loader.py +220 -0
  410. nat/runtime/runner.py +292 -0
  411. nat/runtime/session.py +223 -0
  412. nat/runtime/user_metadata.py +130 -0
  413. nat/settings/__init__.py +0 -0
  414. nat/settings/global_settings.py +329 -0
  415. nat/test/.namespace +1 -0
  416. nat/tool/__init__.py +0 -0
  417. nat/tool/chat_completion.py +77 -0
  418. nat/tool/code_execution/README.md +151 -0
  419. nat/tool/code_execution/__init__.py +0 -0
  420. nat/tool/code_execution/code_sandbox.py +267 -0
  421. nat/tool/code_execution/local_sandbox/.gitignore +1 -0
  422. nat/tool/code_execution/local_sandbox/Dockerfile.sandbox +60 -0
  423. nat/tool/code_execution/local_sandbox/__init__.py +13 -0
  424. nat/tool/code_execution/local_sandbox/local_sandbox_server.py +198 -0
  425. nat/tool/code_execution/local_sandbox/sandbox.requirements.txt +6 -0
  426. nat/tool/code_execution/local_sandbox/start_local_sandbox.sh +50 -0
  427. nat/tool/code_execution/register.py +74 -0
  428. nat/tool/code_execution/test_code_execution_sandbox.py +414 -0
  429. nat/tool/code_execution/utils.py +100 -0
  430. nat/tool/datetime_tools.py +82 -0
  431. nat/tool/document_search.py +141 -0
  432. nat/tool/github_tools.py +450 -0
  433. nat/tool/memory_tools/__init__.py +0 -0
  434. nat/tool/memory_tools/add_memory_tool.py +79 -0
  435. nat/tool/memory_tools/delete_memory_tool.py +66 -0
  436. nat/tool/memory_tools/get_memory_tool.py +72 -0
  437. nat/tool/nvidia_rag.py +95 -0
  438. nat/tool/register.py +31 -0
  439. nat/tool/retriever.py +95 -0
  440. nat/tool/server_tools.py +66 -0
  441. nat/utils/__init__.py +0 -0
  442. nat/utils/callable_utils.py +70 -0
  443. nat/utils/data_models/__init__.py +0 -0
  444. nat/utils/data_models/schema_validator.py +58 -0
  445. nat/utils/debugging_utils.py +43 -0
  446. nat/utils/decorators.py +210 -0
  447. nat/utils/dump_distro_mapping.py +32 -0
  448. nat/utils/exception_handlers/__init__.py +0 -0
  449. nat/utils/exception_handlers/automatic_retries.py +342 -0
  450. nat/utils/exception_handlers/schemas.py +114 -0
  451. nat/utils/io/__init__.py +0 -0
  452. nat/utils/io/model_processing.py +28 -0
  453. nat/utils/io/yaml_tools.py +119 -0
  454. nat/utils/log_levels.py +25 -0
  455. nat/utils/log_utils.py +37 -0
  456. nat/utils/metadata_utils.py +74 -0
  457. nat/utils/optional_imports.py +142 -0
  458. nat/utils/producer_consumer_queue.py +178 -0
  459. nat/utils/reactive/__init__.py +0 -0
  460. nat/utils/reactive/base/__init__.py +0 -0
  461. nat/utils/reactive/base/observable_base.py +65 -0
  462. nat/utils/reactive/base/observer_base.py +55 -0
  463. nat/utils/reactive/base/subject_base.py +79 -0
  464. nat/utils/reactive/observable.py +59 -0
  465. nat/utils/reactive/observer.py +76 -0
  466. nat/utils/reactive/subject.py +131 -0
  467. nat/utils/reactive/subscription.py +49 -0
  468. nat/utils/settings/__init__.py +0 -0
  469. nat/utils/settings/global_settings.py +195 -0
  470. nat/utils/string_utils.py +38 -0
  471. nat/utils/type_converter.py +299 -0
  472. nat/utils/type_utils.py +488 -0
  473. nat/utils/url_utils.py +27 -0
  474. nvidia_nat-1.1.0a20251020.dist-info/METADATA +195 -0
  475. nvidia_nat-1.1.0a20251020.dist-info/RECORD +480 -0
  476. nvidia_nat-1.1.0a20251020.dist-info/WHEEL +5 -0
  477. nvidia_nat-1.1.0a20251020.dist-info/entry_points.txt +22 -0
  478. nvidia_nat-1.1.0a20251020.dist-info/licenses/LICENSE-3rd-party.txt +5478 -0
  479. nvidia_nat-1.1.0a20251020.dist-info/licenses/LICENSE.md +201 -0
  480. nvidia_nat-1.1.0a20251020.dist-info/top_level.txt +2 -0
@@ -0,0 +1,305 @@
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 importlib.metadata
17
+ import inspect
18
+ import logging
19
+ import typing
20
+ from enum import Enum
21
+ from functools import lru_cache
22
+ from types import ModuleType
23
+ from typing import TYPE_CHECKING
24
+
25
+ from pydantic import BaseModel
26
+ from pydantic import field_validator
27
+
28
+ from nat.builder.framework_enum import LLMFrameworkEnum
29
+ from nat.data_models.component import ComponentEnum
30
+ from nat.utils.metadata_utils import generate_config_type_docs
31
+
32
+ if TYPE_CHECKING:
33
+ from nat.cli.type_registry import ToolWrapperBuildCallableT
34
+ from nat.data_models.common import TypedBaseModelT
35
+
36
+ logger = logging.getLogger(__name__)
37
+
38
+
39
+ class DiscoveryStatusEnum(str, Enum):
40
+ SUCCESS = "success"
41
+ FAILURE = "failure"
42
+
43
+
44
+ class DiscoveryContractFieldsEnum(str, Enum):
45
+ PACKAGE = "package"
46
+ VERSION = "version"
47
+ COMPONENT_TYPE = "component_type"
48
+ COMPONENT_NAME = "component_name"
49
+ DESCRIPTION = "description"
50
+ DEVELOPER_NOTES = "developer_notes"
51
+
52
+
53
+ class DiscoveryMetadata(BaseModel):
54
+ """A data model representing metadata about each registered component to faciliate its discovery.
55
+
56
+ Args:
57
+ package (str): The name of the package containing the NAT component.
58
+ version (str): The version number of the package containing the NAT component.
59
+ component_type (ComponentEnum): The type of NAT component this metadata represents.
60
+ component_name (str): The registered name of the NAT component.
61
+ description (str): Description of the NAT component pulled from its config objects docstrings.
62
+ developer_notes (str): Other notes to a developers to aid in the use of the component.
63
+ status (DiscoveryStatusEnum): Provides the status of the metadata discovery process.
64
+ """
65
+
66
+ package: str = ""
67
+ version: str = ""
68
+ component_type: ComponentEnum = ComponentEnum.UNDEFINED
69
+ component_name: str = ""
70
+ description: str = ""
71
+ developer_notes: str = ""
72
+ status: DiscoveryStatusEnum = DiscoveryStatusEnum.SUCCESS
73
+
74
+ @field_validator("description", mode="before")
75
+ @classmethod
76
+ def ensure_description_string(cls, v: typing.Any):
77
+ if not isinstance(v, str):
78
+ return ""
79
+ return v
80
+
81
+ @staticmethod
82
+ def get_preferred_item(items: list, preferred: str) -> str:
83
+ return preferred if preferred in items else items[0]
84
+
85
+ @staticmethod
86
+ @lru_cache
87
+ def get_distribution_name_from_metadata(root_package_name: str) -> str | None:
88
+ """
89
+ This is not performant and is only present to be used (not used
90
+ currently) as a fallback when the distro name doesn't match the
91
+ module name and private_data is not available to map it.
92
+ """
93
+ mapping = importlib.metadata.packages_distributions()
94
+ try:
95
+ distro_names = mapping.get(root_package_name, [None])
96
+ distro_name = DiscoveryMetadata.get_preferred_item(distro_names, "nvidia-nat")
97
+ except KeyError:
98
+ return root_package_name
99
+
100
+ return distro_name if distro_name else root_package_name
101
+
102
+ @staticmethod
103
+ @lru_cache
104
+ def get_distribution_name_from_module(module: ModuleType | None) -> str:
105
+ """Get the distribution name from the config type using the mapping of module names to distro names.
106
+
107
+ Args:
108
+ module (ModuleType): A registered component's module.
109
+
110
+ Returns:
111
+ str: The distribution name of the NAT component.
112
+ """
113
+ from nat.runtime.loader import get_all_entrypoints_distro_mapping
114
+
115
+ if module is None:
116
+ return "nvidia-nat"
117
+
118
+ # Get the mapping of module names to distro names
119
+ mapping = get_all_entrypoints_distro_mapping()
120
+ module_package = module.__package__
121
+
122
+ if module_package is None:
123
+ return "nvidia-nat"
124
+
125
+ # Traverse the module package parts in reverse order to find the distro name
126
+ # This is because the module package is the root package for the NAT component
127
+ # and the distro name is the name of the package that contains the component
128
+ module_package_parts = module_package.split(".")
129
+ for part_idx in range(len(module_package_parts), 0, -1):
130
+ candidate_module_name = ".".join(module_package_parts[0:part_idx])
131
+ candidate_distro_name = mapping.get(candidate_module_name, None)
132
+ if candidate_distro_name is not None:
133
+ return candidate_distro_name
134
+
135
+ return "nvidia-nat"
136
+
137
+ @staticmethod
138
+ @lru_cache
139
+ def get_distribution_name_from_config_type(config_type: type["TypedBaseModelT"]) -> str:
140
+ """Get the distribution name from the config type using the mapping of module names to distro names.
141
+
142
+ Args:
143
+ config_type (type[TypedBaseModelT]): A registered component's configuration object.
144
+
145
+ Returns:
146
+ str: The distribution name of the NAT component.
147
+ """
148
+ module = inspect.getmodule(config_type)
149
+ return DiscoveryMetadata.get_distribution_name_from_module(module)
150
+
151
+ @staticmethod
152
+ def from_config_type(config_type: type["TypedBaseModelT"],
153
+ component_type: ComponentEnum = ComponentEnum.UNDEFINED) -> "DiscoveryMetadata":
154
+ """Generates discovery metadata from a NAT config object.
155
+
156
+ Args:
157
+ config_type (type[TypedBaseModelT]): A registered component's configuration object.
158
+ component_type (ComponentEnum, optional): The type of the registered component. Defaults to
159
+ ComponentEnum.UNDEFINED.
160
+
161
+ Returns:
162
+ DiscoveryMetadata: A an object containing component metadata to facilitate discovery and reuse.
163
+ """
164
+
165
+ try:
166
+ module = inspect.getmodule(config_type)
167
+ distro_name = DiscoveryMetadata.get_distribution_name_from_config_type(config_type)
168
+
169
+ if not distro_name:
170
+ # raise an exception
171
+ logger.error("Encountered issue getting distro_name for module %s", module.__name__)
172
+ return DiscoveryMetadata(status=DiscoveryStatusEnum.FAILURE)
173
+
174
+ try:
175
+ version = importlib.metadata.version(distro_name) if distro_name != "" else ""
176
+ except importlib.metadata.PackageNotFoundError:
177
+ logger.warning("Package metadata not found for %s", distro_name)
178
+ version = ""
179
+ except Exception as e:
180
+ logger.exception("Encountered issue extracting module metadata for %s: %s", config_type, e)
181
+ return DiscoveryMetadata(status=DiscoveryStatusEnum.FAILURE)
182
+
183
+ description = generate_config_type_docs(config_type=config_type)
184
+
185
+ return DiscoveryMetadata(package=distro_name,
186
+ version=version,
187
+ component_type=component_type,
188
+ component_name=config_type.static_type(),
189
+ description=description)
190
+
191
+ @staticmethod
192
+ def from_fn_wrapper(fn: "ToolWrapperBuildCallableT",
193
+ wrapper_type: LLMFrameworkEnum | str,
194
+ component_type: ComponentEnum = ComponentEnum.TOOL_WRAPPER) -> "DiscoveryMetadata":
195
+ """Generates discovery metadata from function with specified wrapper type.
196
+
197
+ Args:
198
+ fn (ToolWrapperBuildCallableT): A tool wrapper callable to source component metadata.
199
+ wrapper_type (LLMFrameworkEnum): The wrapper to apply to the callable to faciliate inter-framwork
200
+ interoperability.
201
+
202
+ component_type (ComponentEnum, optional): The type of the registered component. Defaults to
203
+ ComponentEnum.TOOL_WRAPPER.
204
+
205
+ Returns:
206
+ DiscoveryMetadata: A an object containing component metadata to facilitate discovery and reuse.
207
+ """
208
+
209
+ try:
210
+ module = inspect.getmodule(fn)
211
+ distro_name = DiscoveryMetadata.get_distribution_name_from_module(module)
212
+
213
+ try:
214
+ # version = importlib.metadata.version(root_package) if root_package != "" else ""
215
+ version = importlib.metadata.version(distro_name) if distro_name != "" else ""
216
+ except importlib.metadata.PackageNotFoundError:
217
+ logger.warning("Package metadata not found for %s", distro_name)
218
+ version = ""
219
+ except Exception as e:
220
+ logger.exception("Encountered issue extracting module metadata for %s: %s", fn, e)
221
+ return DiscoveryMetadata(status=DiscoveryStatusEnum.FAILURE)
222
+
223
+ if isinstance(wrapper_type, LLMFrameworkEnum):
224
+ wrapper_type = wrapper_type.value
225
+
226
+ return DiscoveryMetadata(package=distro_name,
227
+ version=version,
228
+ component_type=component_type,
229
+ component_name=wrapper_type,
230
+ description=fn.__doc__ or "")
231
+
232
+ @staticmethod
233
+ def from_package_name(package_name: str, package_version: str | None) -> "DiscoveryMetadata":
234
+ """Generates discovery metadata from an installed package name.
235
+
236
+ Args:
237
+ package_name (str): The name of the NAT plugin package containing registered components.
238
+ package_version (str, optional): The version of the package, Defaults to None.
239
+
240
+ Returns:
241
+ DiscoveryMetadata: A an object containing component metadata to facilitate discovery and reuse.
242
+ """
243
+
244
+ try:
245
+ try:
246
+ metadata = importlib.metadata.metadata(package_name)
247
+ description = metadata.get("Summary", "")
248
+ if (package_version is None):
249
+ package_version = importlib.metadata.version(package_name)
250
+ except importlib.metadata.PackageNotFoundError:
251
+ logger.warning("Package metadata not found for %s", package_name)
252
+ description = ""
253
+ package_version = package_version or ""
254
+ except Exception as e:
255
+ logger.exception("Encountered issue extracting module metadata for %s: %s", package_name, e)
256
+ return DiscoveryMetadata(status=DiscoveryStatusEnum.FAILURE)
257
+
258
+ return DiscoveryMetadata(package=package_name,
259
+ version=package_version,
260
+ component_type=ComponentEnum.PACKAGE,
261
+ component_name=package_name,
262
+ description=description)
263
+
264
+ @staticmethod
265
+ def from_provider_framework_map(config_type: type["TypedBaseModelT"],
266
+ wrapper_type: LLMFrameworkEnum | str | None,
267
+ provider_type: ComponentEnum,
268
+ component_type: ComponentEnum = ComponentEnum.UNDEFINED) -> "DiscoveryMetadata":
269
+ """Generates discovery metadata from provider and framework mapping information.
270
+
271
+ Args:
272
+ config_type (type[TypedBaseModelT]): A registered component's configuration object.
273
+ wrapper_type (LLMFrameworkEnum | str): The wrapper to apply to the callable to faciliate inter-framwork
274
+ interoperability.
275
+
276
+ provider_type (ComponentEnum): The type of provider the registered component supports.
277
+ component_type (ComponentEnum, optional): The type of the registered component. Defaults to
278
+ ComponentEnum.UNDEFINED.
279
+
280
+ Returns:
281
+ DiscoveryMetadata: A an object containing component metadata to facilitate discovery and reuse.
282
+ """
283
+
284
+ try:
285
+ module = inspect.getmodule(config_type)
286
+ distro_name = DiscoveryMetadata.get_distribution_name_from_module(module)
287
+ try:
288
+ version = importlib.metadata.version(distro_name) if distro_name != "" else ""
289
+ except importlib.metadata.PackageNotFoundError:
290
+ logger.warning("Package metadata not found for %s", distro_name)
291
+ version = ""
292
+ except Exception as e:
293
+ logger.exception("Encountered issue extracting module metadata for %s: %s", config_type, e)
294
+ return DiscoveryMetadata(status=DiscoveryStatusEnum.FAILURE)
295
+
296
+ wrapper_type = wrapper_type.value if isinstance(wrapper_type, LLMFrameworkEnum) else wrapper_type
297
+ component_name = f"{config_type.static_type()} ({provider_type.value}) - {wrapper_type}"
298
+
299
+ description = generate_config_type_docs(config_type=config_type)
300
+
301
+ return DiscoveryMetadata(package=distro_name,
302
+ version=version,
303
+ component_type=component_type,
304
+ component_name=component_name,
305
+ description=description)
@@ -0,0 +1,27 @@
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
+
18
+ from .common import BaseModelRegistryTag
19
+ from .common import TypedBaseModel
20
+
21
+
22
+ class EmbedderBaseConfig(TypedBaseModel, BaseModelRegistryTag):
23
+ """ Base configuration for embedding model providers. """
24
+ pass
25
+
26
+
27
+ EmbedderBaseConfigT = typing.TypeVar("EmbedderBaseConfigT", bound=EmbedderBaseConfig)
@@ -0,0 +1,130 @@
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 enum import Enum
18
+ from pathlib import Path
19
+
20
+ from pydantic import BaseModel
21
+ from pydantic import Discriminator
22
+ from pydantic import model_validator
23
+
24
+ from nat.data_models.common import TypedBaseModel
25
+ from nat.data_models.dataset_handler import EvalDatasetConfig
26
+ from nat.data_models.dataset_handler import EvalS3Config
27
+ from nat.data_models.evaluator import EvaluatorBaseConfig
28
+ from nat.data_models.intermediate_step import IntermediateStepType
29
+ from nat.data_models.profiler import ProfilerConfig
30
+
31
+
32
+ class JobEvictionPolicy(str, Enum):
33
+ """Policy for evicting old jobs when max_jobs is exceeded."""
34
+ TIME_CREATED = "time_created"
35
+ TIME_MODIFIED = "time_modified"
36
+
37
+
38
+ class EvalCustomScriptConfig(BaseModel):
39
+ # Path to the script to run
40
+ script: Path
41
+ # Keyword arguments to pass to the script
42
+ kwargs: dict[str, str] = {}
43
+
44
+
45
+ class JobManagementConfig(BaseModel):
46
+ # Whether to append a unique job ID to the output directory for each run
47
+ append_job_id_to_output_dir: bool = False
48
+ # Maximum number of jobs to keep in the output directory. Oldest jobs will be evicted.
49
+ # A value of 0 means no limit.
50
+ max_jobs: int = 0
51
+ # Policy for evicting old jobs. Defaults to using time_created.
52
+ eviction_policy: JobEvictionPolicy = JobEvictionPolicy.TIME_CREATED
53
+
54
+
55
+ class EvalOutputConfig(BaseModel):
56
+ # Output directory for the workflow and evaluation results
57
+ dir: Path = Path("./.tmp/nat/examples/default/")
58
+ # S3 prefix for the workflow and evaluation results
59
+ remote_dir: str | None = None
60
+ # Custom function to pre-evaluation process the eval input
61
+ # Format: "module.path.function_name"
62
+ custom_pre_eval_process_function: str | None = None
63
+ # Custom scripts to run after the workflow and evaluation results are saved
64
+ custom_scripts: dict[str, EvalCustomScriptConfig] = {}
65
+ # S3 config for uploading the contents of the output directory
66
+ s3: EvalS3Config | None = None
67
+ # Whether to cleanup the output directory before running the workflow
68
+ cleanup: bool = True
69
+ # Job management configuration (job id, eviction, etc.)
70
+ job_management: JobManagementConfig = JobManagementConfig()
71
+ # Filter for the workflow output steps
72
+ workflow_output_step_filter: list[IntermediateStepType] | None = None
73
+
74
+
75
+ class EvalGeneralConfig(BaseModel):
76
+ max_concurrency: int = 8
77
+
78
+ # Workflow alias for displaying in evaluation UI, if not provided,
79
+ # the workflow type will be used
80
+ workflow_alias: str | None = None
81
+
82
+ # Output directory for the workflow and evaluation results
83
+ output_dir: Path = Path("./.tmp/nat/examples/default/")
84
+
85
+ # If present overrides output_dir
86
+ output: EvalOutputConfig | None = None
87
+
88
+ # Dataset for running the workflow and evaluating
89
+ dataset: EvalDatasetConfig | None = None
90
+
91
+ # Inference profiler
92
+ profiler: ProfilerConfig | None = None
93
+
94
+ # overwrite the output_dir with the output config if present
95
+ @model_validator(mode="before")
96
+ @classmethod
97
+ def override_output_dir(cls, values):
98
+ if values.get("output") and values["output"].get("dir"):
99
+ values["output_dir"] = values["output"]["dir"]
100
+ return values
101
+
102
+
103
+ class EvalConfig(BaseModel):
104
+
105
+ # General Evaluation Options
106
+ general: EvalGeneralConfig = EvalGeneralConfig()
107
+
108
+ # Evaluators
109
+ evaluators: dict[str, EvaluatorBaseConfig] = {}
110
+
111
+ @classmethod
112
+ def rebuild_annotations(cls):
113
+
114
+ from nat.cli.type_registry import GlobalTypeRegistry
115
+
116
+ type_registry = GlobalTypeRegistry.get()
117
+
118
+ EvaluatorsAnnotation = dict[str,
119
+ typing.Annotated[type_registry.compute_annotation(EvaluatorBaseConfig),
120
+ Discriminator(TypedBaseModel.discriminator)]]
121
+
122
+ should_rebuild = False
123
+
124
+ evaluators_field = cls.model_fields.get("evaluators")
125
+ if evaluators_field is not None and evaluators_field.annotation != EvaluatorsAnnotation:
126
+ evaluators_field.annotation = EvaluatorsAnnotation
127
+ should_rebuild = True
128
+
129
+ if (should_rebuild):
130
+ cls.model_rebuild(force=True)
@@ -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
+ import typing
17
+
18
+ from .common import BaseModelRegistryTag
19
+ from .common import TypedBaseModel
20
+
21
+
22
+ class EvaluatorBaseConfig(TypedBaseModel, BaseModelRegistryTag):
23
+ pass
24
+
25
+
26
+ EvaluatorBaseConfigT = typing.TypeVar("EvaluatorBaseConfigT", bound=EvaluatorBaseConfig)
@@ -0,0 +1,26 @@
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
+
18
+ from .common import BaseModelRegistryTag
19
+ from .common import TypedBaseModel
20
+
21
+
22
+ class FrontEndBaseConfig(TypedBaseModel, BaseModelRegistryTag):
23
+ pass
24
+
25
+
26
+ FrontEndConfigT = typing.TypeVar("FrontEndConfigT", bound=FrontEndBaseConfig)
@@ -0,0 +1,64 @@
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 typing
17
+
18
+ from pydantic import Field
19
+ from pydantic import field_validator
20
+ from pydantic import model_validator
21
+
22
+ from .common import BaseModelRegistryTag
23
+ from .common import TypedBaseModel
24
+
25
+
26
+ class FunctionBaseConfig(TypedBaseModel, BaseModelRegistryTag):
27
+ pass
28
+
29
+
30
+ class FunctionGroupBaseConfig(TypedBaseModel, BaseModelRegistryTag):
31
+ """Base configuration for function groups.
32
+
33
+ Function groups enable sharing of configurations and resources across multiple functions.
34
+ """
35
+ include: list[str] = Field(
36
+ default_factory=list,
37
+ description="The list of function names which should be added to the global Function registry",
38
+ )
39
+ exclude: list[str] = Field(
40
+ default_factory=list,
41
+ description="The list of function names which should be excluded from default access to the group",
42
+ )
43
+
44
+ @field_validator("include", "exclude")
45
+ @classmethod
46
+ def _validate_fields_include_exclude(cls, value: list[str]) -> list[str]:
47
+ if len(set(value)) != len(value):
48
+ raise ValueError("Function names must be unique")
49
+ return sorted(value)
50
+
51
+ @model_validator(mode="after")
52
+ def _validate_include_exclude(self):
53
+ if self.include and self.exclude:
54
+ raise ValueError("include and exclude cannot be used together")
55
+ return self
56
+
57
+
58
+ class EmptyFunctionConfig(FunctionBaseConfig, name="EmptyFunctionConfig"):
59
+ pass
60
+
61
+
62
+ FunctionConfigT = typing.TypeVar("FunctionConfigT", bound=FunctionBaseConfig)
63
+
64
+ FunctionGroupConfigT = typing.TypeVar("FunctionGroupConfigT", bound=FunctionGroupBaseConfig)
@@ -0,0 +1,80 @@
1
+ # SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
2
+ # SPDX-License-Identifier: Apache-2.0
3
+ #
4
+ # Licensed under the Apache License, Version 2.0 (the "License");
5
+ # you may not use this file except in compliance with the License.
6
+ # You may obtain a copy of the License at
7
+ #
8
+ # http://www.apache.org/licenses/LICENSE-2.0
9
+ #
10
+ # Unless required by applicable law or agreed to in writing, software
11
+ # distributed under the License is distributed on an "AS IS" BASIS,
12
+ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13
+ # See the License for the specific language governing permissions and
14
+ # limitations under the License.
15
+
16
+ from pydantic import BaseModel
17
+ from pydantic import Field
18
+ from pydantic import field_serializer
19
+
20
+
21
+ class FunctionDependencies(BaseModel):
22
+ """
23
+ A class to represent the dependencies of a function.
24
+ """
25
+ functions: set[str] = Field(default_factory=set)
26
+ function_groups: set[str] = Field(default_factory=set)
27
+ llms: set[str] = Field(default_factory=set)
28
+ embedders: set[str] = Field(default_factory=set)
29
+ memory_clients: set[str] = Field(default_factory=set)
30
+ object_stores: set[str] = Field(default_factory=set)
31
+ retrievers: set[str] = Field(default_factory=set)
32
+
33
+ @field_serializer("functions", when_used="json")
34
+ def serialize_functions(self, v: set[str]) -> list[str]:
35
+ return list(v)
36
+
37
+ @field_serializer("function_groups", when_used="json")
38
+ def serialize_function_groups(self, v: set[str]) -> list[str]:
39
+ return list(v)
40
+
41
+ @field_serializer("llms", when_used="json")
42
+ def serialize_llms(self, v: set[str]) -> list[str]:
43
+ return list(v)
44
+
45
+ @field_serializer("embedders", when_used="json")
46
+ def serialize_embedders(self, v: set[str]) -> list[str]:
47
+ return list(v)
48
+
49
+ @field_serializer("memory_clients", when_used="json")
50
+ def serialize_memory_clients(self, v: set[str]) -> list[str]:
51
+ return list(v)
52
+
53
+ @field_serializer("object_stores", when_used="json")
54
+ def serialize_object_stores(self, v: set[str]) -> list[str]:
55
+ return list(v)
56
+
57
+ @field_serializer("retrievers", when_used="json")
58
+ def serialize_retrievers(self, v: set[str]) -> list[str]:
59
+ return list(v)
60
+
61
+ def add_function(self, function: str):
62
+ self.functions.add(function)
63
+
64
+ def add_function_group(self, function_group: str):
65
+ self.function_groups.add(function_group) # pylint: disable=no-member
66
+
67
+ def add_llm(self, llm: str):
68
+ self.llms.add(llm)
69
+
70
+ def add_embedder(self, embedder: str):
71
+ self.embedders.add(embedder)
72
+
73
+ def add_memory_client(self, memory_client: str):
74
+ self.memory_clients.add(memory_client)
75
+
76
+ def add_object_store(self, object_store: str):
77
+ self.object_stores.add(object_store)
78
+
79
+ def add_retriever(self, retriever: str):
80
+ self.retrievers.add(retriever)