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
nat/eval/evaluate.py ADDED
@@ -0,0 +1,565 @@
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 asyncio
17
+ import logging
18
+ import shutil
19
+ from pathlib import Path
20
+ from typing import Any
21
+ from uuid import uuid4
22
+
23
+ from pydantic import BaseModel
24
+ from tqdm import tqdm
25
+
26
+ from nat.data_models.evaluate import EvalConfig
27
+ from nat.data_models.evaluate import JobEvictionPolicy
28
+ from nat.eval.config import EvaluationRunConfig
29
+ from nat.eval.config import EvaluationRunOutput
30
+ from nat.eval.dataset_handler.dataset_handler import DatasetHandler
31
+ from nat.eval.evaluator.evaluator_model import EvalInput
32
+ from nat.eval.evaluator.evaluator_model import EvalInputItem
33
+ from nat.eval.evaluator.evaluator_model import EvalOutput
34
+ from nat.eval.usage_stats import UsageStats
35
+ from nat.eval.usage_stats import UsageStatsItem
36
+ from nat.eval.usage_stats import UsageStatsLLM
37
+ from nat.eval.utils.output_uploader import OutputUploader
38
+ from nat.eval.utils.weave_eval import WeaveEvaluationIntegration
39
+ from nat.profiler.data_models import ProfilerResults
40
+ from nat.runtime.session import SessionManager
41
+
42
+ logger = logging.getLogger(__name__)
43
+
44
+
45
+ class EvaluationRun:
46
+ """
47
+ Instantiated for each evaluation run and used to store data for that single run.
48
+
49
+ .. warning::
50
+ **Experimental Feature**: The Evaluation API is experimental and may change in future releases.
51
+ Future versions may introduce breaking changes without notice.
52
+ """
53
+
54
+ def __init__(self, config: EvaluationRunConfig):
55
+ """
56
+ Initialize an EvaluationRun with configuration.
57
+ """
58
+ from nat.eval.intermediate_step_adapter import IntermediateStepAdapter
59
+
60
+ # Run-specific configuration
61
+ self.config: EvaluationRunConfig = config
62
+ self.eval_config: EvalConfig | None = None
63
+
64
+ # Helpers
65
+ self.intermediate_step_adapter: IntermediateStepAdapter = IntermediateStepAdapter()
66
+
67
+ # Create evaluation trace context
68
+ try:
69
+ from nat.eval.utils.eval_trace_ctx import WeaveEvalTraceContext
70
+ self.eval_trace_context = WeaveEvalTraceContext()
71
+ except Exception:
72
+ from nat.eval.utils.eval_trace_ctx import EvalTraceContext
73
+ self.eval_trace_context = EvalTraceContext()
74
+
75
+ self.weave_eval: WeaveEvaluationIntegration = WeaveEvaluationIntegration(self.eval_trace_context)
76
+ # Metadata
77
+ self.eval_input: EvalInput | None = None
78
+ self.workflow_interrupted: bool = False
79
+
80
+ # evaluation_results is list of tuples (evaluator_name, EvalOutput)
81
+ self.evaluation_results: list[tuple[str, EvalOutput]] = []
82
+
83
+ # usage stats
84
+ self.usage_stats: UsageStats = UsageStats()
85
+
86
+ # workflow output file
87
+ self.workflow_output_file: Path | None = None
88
+
89
+ # evaluation output files
90
+ self.evaluator_output_files: list[Path] = []
91
+
92
+ def _compute_usage_stats(self, item: EvalInputItem):
93
+ """Compute usage stats for a single item using the intermediate steps"""
94
+ # get the prompt and completion tokens from the intermediate steps
95
+ from nat.profiler.intermediate_property_adapter import IntermediatePropertyAdaptor
96
+ steps = [IntermediatePropertyAdaptor.from_intermediate_step(step) for step in item.trajectory]
97
+ usage_stats_per_llm = {}
98
+ total_tokens = 0
99
+ for step in steps:
100
+ if step.event_type == "LLM_END":
101
+ llm_name = step.llm_name
102
+ if llm_name not in usage_stats_per_llm:
103
+ usage_stats_per_llm[llm_name] = UsageStatsLLM()
104
+ usage_stats_per_llm[llm_name].prompt_tokens += step.token_usage.prompt_tokens
105
+ usage_stats_per_llm[llm_name].completion_tokens += step.token_usage.completion_tokens
106
+ usage_stats_per_llm[llm_name].total_tokens += step.token_usage.total_tokens
107
+ total_tokens += step.token_usage.total_tokens
108
+
109
+ # find min and max event timestamps
110
+ if item.trajectory:
111
+ min_timestamp = min(step.event_timestamp for step in item.trajectory)
112
+ max_timestamp = max(step.event_timestamp for step in item.trajectory)
113
+ runtime = max_timestamp - min_timestamp
114
+ else:
115
+ min_timestamp = 0.0
116
+ max_timestamp = 0.0
117
+ runtime = 0.0
118
+
119
+ # find llm latency by calculating p95 of all llm calls
120
+ llm_latencies = []
121
+ previous_llm_start_time = None
122
+ for step in steps:
123
+ if step.event_type == "LLM_START":
124
+ previous_llm_start_time = step.event_timestamp
125
+ elif step.event_type == "LLM_END" and previous_llm_start_time is not None:
126
+ llm_latencies.append(step.event_timestamp - previous_llm_start_time)
127
+ previous_llm_start_time = None
128
+
129
+ # Calculate p95 LLM latency (or 0 if no LLM calls)
130
+ if llm_latencies:
131
+ import numpy as np
132
+ llm_latency = float(np.percentile(llm_latencies, 95))
133
+ else:
134
+ llm_latency = 0.0
135
+
136
+ # add the usage stats to the usage stats dict
137
+ self.usage_stats.usage_stats_items[item.id] = UsageStatsItem(usage_stats_per_llm=usage_stats_per_llm,
138
+ runtime=runtime,
139
+ total_tokens=total_tokens,
140
+ min_timestamp=min_timestamp,
141
+ max_timestamp=max_timestamp,
142
+ llm_latency=llm_latency)
143
+ return self.usage_stats.usage_stats_items[item.id]
144
+
145
+ async def run_workflow_local(self, session_manager: SessionManager):
146
+ '''
147
+ Launch the workflow with the specified questions and extract the output using the jsonpath
148
+ '''
149
+ # import function level dependencies
150
+ from jsonpath_ng import parse
151
+
152
+ from nat.eval.runtime_event_subscriber import pull_intermediate
153
+
154
+ # Run the workflow
155
+ jsonpath_expr = parse(self.config.result_json_path)
156
+ stop_event = asyncio.Event()
157
+
158
+ async def run_one(item: EvalInputItem):
159
+ if stop_event.is_set():
160
+ return "", []
161
+
162
+ async with session_manager.run(item.input_obj) as runner:
163
+ if not session_manager.workflow.has_single_output:
164
+ # raise an error if the workflow has multiple outputs
165
+ raise NotImplementedError("Multiple outputs are not supported")
166
+
167
+ runner_result = None
168
+ intermediate_future = None
169
+
170
+ try:
171
+ # Start usage stats and intermediate steps collection in parallel
172
+ intermediate_future = pull_intermediate()
173
+ runner_result = runner.result()
174
+ base_output = await runner_result
175
+ intermediate_steps = await intermediate_future
176
+ except NotImplementedError as e:
177
+ logger.error("Failed to run the workflow: %s", e)
178
+ # raise original error
179
+ raise
180
+ except Exception as e:
181
+ logger.exception("Failed to run the workflow: %s", e)
182
+ # stop processing if a workflow error occurs
183
+ self.workflow_interrupted = True
184
+
185
+ # Cancel any coroutines that are still running, avoiding a warning about unawaited coroutines
186
+ # (typically one of these two is what raised the exception and the other is still running)
187
+ for coro in (runner_result, intermediate_future):
188
+ if coro is not None:
189
+ asyncio.ensure_future(coro).cancel()
190
+
191
+ stop_event.set()
192
+ return
193
+
194
+ try:
195
+ base_output = runner.convert(base_output, to_type=str)
196
+ except ValueError:
197
+ pass
198
+
199
+ # if base_output is a pydantic model dump it to json
200
+ if isinstance(base_output, BaseModel):
201
+ output = base_output.model_dump_json(indent=2)
202
+ else:
203
+ m = jsonpath_expr.find(base_output)
204
+ if (not m):
205
+ raise RuntimeError(f"Failed to extract output using jsonpath: {self.config.result_json_path}")
206
+ if (len(m) > 1):
207
+ logger.warning("Multiple matches found for jsonpath at row '%s'. Matches: %s. Using the first",
208
+ base_output,
209
+ m)
210
+ output = m[0].value
211
+
212
+ item.output_obj = output
213
+ item.trajectory = self.intermediate_step_adapter.validate_intermediate_steps(intermediate_steps)
214
+ usage_stats_item = self._compute_usage_stats(item)
215
+
216
+ self.weave_eval.log_prediction(item, output)
217
+ await self.weave_eval.log_usage_stats(item, usage_stats_item)
218
+
219
+ async def wrapped_run(item: EvalInputItem) -> None:
220
+ await run_one(item)
221
+ pbar.update(1)
222
+
223
+ # if self.config.skip_complete is set skip eval_input_items with a non-empty output_obj
224
+ if self.config.skip_completed_entries:
225
+ eval_input_items = [item for item in self.eval_input.eval_input_items if not item.output_obj]
226
+ if not eval_input_items:
227
+ logger.warning("All items have a non-empty output. Skipping workflow pass altogether.")
228
+ return
229
+ else:
230
+ eval_input_items = self.eval_input.eval_input_items
231
+ pbar = tqdm(total=len(eval_input_items), desc="Running workflow")
232
+ await asyncio.gather(*[wrapped_run(item) for item in eval_input_items])
233
+ pbar.close()
234
+
235
+ async def run_workflow_remote(self):
236
+ from nat.eval.remote_workflow import EvaluationRemoteWorkflowHandler
237
+ handler = EvaluationRemoteWorkflowHandler(self.config, self.eval_config.general.max_concurrency)
238
+ await handler.run_workflow_remote(self.eval_input)
239
+ for item in self.eval_input.eval_input_items:
240
+ usage_stats_item = self._compute_usage_stats(item)
241
+ self.weave_eval.log_prediction(item, item.output_obj)
242
+ await self.weave_eval.log_usage_stats(item, usage_stats_item)
243
+
244
+ async def profile_workflow(self) -> ProfilerResults:
245
+ """
246
+ Profile a dataset
247
+ """
248
+
249
+ if not self.eval_config.general.profiler:
250
+ logger.info("Profiler is not enabled. Skipping profiling.")
251
+ return ProfilerResults()
252
+
253
+ from nat.profiler.profile_runner import ProfilerRunner
254
+
255
+ all_stats = []
256
+ for input_item in self.eval_input.eval_input_items:
257
+ all_stats.append(input_item.trajectory)
258
+
259
+ profiler_runner = ProfilerRunner(self.eval_config.general.profiler,
260
+ self.eval_config.general.output_dir,
261
+ write_output=self.config.write_output)
262
+
263
+ return await profiler_runner.run(all_stats)
264
+
265
+ def cleanup_output_directory(self):
266
+ '''Remove contents of the output directory if it exists'''
267
+ output_config = self.eval_config.general.output
268
+ output_dir = output_config.dir
269
+
270
+ if not (output_config and output_dir.exists()):
271
+ return
272
+
273
+ # If cleanup is true, remove the entire directory and we are done
274
+ if output_config.cleanup:
275
+ logger.info("Cleaning up entire output directory: %s", output_config.dir)
276
+ shutil.rmtree(output_config.dir)
277
+ return
278
+
279
+ if output_config.job_management.max_jobs == 0:
280
+ # No eviction policy
281
+ return
282
+
283
+ base_dir = output_dir / "jobs"
284
+ if not base_dir.exists():
285
+ return
286
+
287
+ # Get all subdirectories, which represent individual job runs
288
+ job_dirs = [d for d in base_dir.iterdir() if d.is_dir()]
289
+ if len(job_dirs) <= output_config.job_management.max_jobs:
290
+ return
291
+
292
+ # Determine sort key based on eviction_policy, defaulting to creation time
293
+ if output_config.job_management.eviction_policy == JobEvictionPolicy.TIME_MODIFIED:
294
+
295
+ def sort_key(x):
296
+ return x.stat().st_mtime
297
+
298
+ logger.info("Using last modified time for job eviction policy.")
299
+ else:
300
+
301
+ def sort_key(x):
302
+ return x.stat().st_ctime
303
+
304
+ logger.info("Using creation time for job eviction policy.")
305
+
306
+ # Sort directories (oldest first)
307
+ job_dirs.sort(key=sort_key)
308
+ num_to_delete = len(job_dirs) - output_config.job_management.max_jobs
309
+
310
+ logger.info("Found %d jobs, exceeding limit of %d. Removing %d oldest jobs.",
311
+ len(job_dirs),
312
+ output_config.job_management.max_jobs,
313
+ num_to_delete)
314
+
315
+ for dir_to_delete in job_dirs[:num_to_delete]:
316
+ try:
317
+ logger.info("Deleting old job directory: %s", dir_to_delete)
318
+ shutil.rmtree(dir_to_delete)
319
+ except Exception as e:
320
+ logger.exception("Failed to delete old job directory: %s: %s", dir_to_delete, e)
321
+
322
+ def write_output(self, dataset_handler: DatasetHandler, profiler_results: ProfilerResults):
323
+ workflow_output_file = self.eval_config.general.output_dir / "workflow_output.json"
324
+ workflow_output_file.parent.mkdir(parents=True, exist_ok=True)
325
+
326
+ # Write the workflow output to a file (this can be used for re-running the evaluation)
327
+
328
+ step_filter = self.eval_config.general.output.workflow_output_step_filter \
329
+ if self.eval_config.general.output else None
330
+ workflow_output = dataset_handler.publish_eval_input(self.eval_input, step_filter)
331
+ with open(workflow_output_file, "w", encoding="utf-8") as f:
332
+ # set indent to 2 for pretty printing
333
+ f.write(workflow_output)
334
+ self.workflow_output_file = workflow_output_file
335
+ logger.info("Workflow output written to %s", workflow_output_file)
336
+
337
+ # Write the output of each evaluator to a separate json file
338
+ for evaluator_name, eval_output in self.evaluation_results:
339
+ output_file = self.eval_config.general.output_dir / f"{evaluator_name}_output.json"
340
+ output_file.parent.mkdir(parents=True, exist_ok=True)
341
+ # create json content using the evaluation results
342
+ output = eval_output.model_dump_json(indent=2)
343
+ with open(output_file, "w", encoding="utf-8") as f:
344
+ f.write(output)
345
+ self.evaluator_output_files.append(output_file)
346
+ logger.info("Evaluation results written to %s", output_file)
347
+
348
+ def publish_output(self, dataset_handler: DatasetHandler, profiler_results: ProfilerResults):
349
+ """Publish the output"""
350
+ if self.config.write_output:
351
+ self.write_output(dataset_handler, profiler_results)
352
+
353
+ if self.workflow_interrupted:
354
+ # Issue a warning if the workflow was not completed on all datasets
355
+ msg = ("Workflow execution was interrupted due to an error. The results may be incomplete. "
356
+ "You can re-execute evaluation for incomplete results by running "
357
+ "`eval` with the --skip_completed_entries flag.")
358
+ logger.warning(msg)
359
+
360
+ self.weave_eval.log_summary(self.usage_stats, self.evaluation_results, profiler_results)
361
+
362
+ async def run_single_evaluator(self, evaluator_name: str, evaluator: Any):
363
+ """Run a single evaluator and store its results."""
364
+ try:
365
+ eval_output = await evaluator.evaluate_fn(self.eval_input)
366
+ self.evaluation_results.append((evaluator_name, eval_output))
367
+
368
+ await self.weave_eval.alog_score(eval_output, evaluator_name)
369
+ except Exception as e:
370
+ logger.exception("An error occurred while running evaluator %s: %s", evaluator_name, e)
371
+
372
+ async def run_evaluators(self, evaluators: dict[str, Any]):
373
+ """Run all configured evaluators asynchronously."""
374
+ tasks = [self.run_single_evaluator(name, evaluator) for name, evaluator in evaluators.items() if evaluator]
375
+
376
+ if not tasks:
377
+ logger.warning("All evaluators were empty or invalid.")
378
+ return
379
+
380
+ try:
381
+ await asyncio.gather(*tasks)
382
+ except Exception as e:
383
+ logger.error("An error occurred while running evaluators: %s", e)
384
+ raise
385
+ finally:
386
+ # Finish prediction loggers in Weave
387
+ await self.weave_eval.afinish_loggers()
388
+
389
+ def apply_overrides(self):
390
+ from nat.cli.cli_utils.config_override import load_and_override_config
391
+ from nat.data_models.config import Config
392
+ from nat.runtime.loader import PluginTypes
393
+ from nat.runtime.loader import discover_and_register_plugins
394
+ from nat.utils.data_models.schema_validator import validate_schema
395
+
396
+ # Register plugins before validation
397
+ discover_and_register_plugins(PluginTypes.CONFIG_OBJECT)
398
+
399
+ config_dict = load_and_override_config(self.config.config_file, self.config.override)
400
+ config = validate_schema(config_dict, Config)
401
+ return config
402
+
403
+ def _get_workflow_alias(self, workflow_type: str | None = None):
404
+ """Get the workflow alias for displaying in evaluation UI."""
405
+ if self.eval_config.general.workflow_alias:
406
+ return self.eval_config.general.workflow_alias
407
+
408
+ if not workflow_type or workflow_type == "EmptyFunctionConfig":
409
+ return "nat-eval"
410
+
411
+ return workflow_type
412
+
413
+ async def wait_for_all_export_tasks_local(self, session_manager: SessionManager, timeout: float) -> None:
414
+ """Wait for all trace export tasks to complete for local workflows.
415
+
416
+ This only works for local workflows where we have direct access to the
417
+ SessionManager and its underlying workflow with exporter manager.
418
+ """
419
+ try:
420
+ workflow = session_manager.workflow
421
+ all_exporters = await workflow.get_all_exporters()
422
+ if not all_exporters:
423
+ logger.debug("No exporters to wait for")
424
+ return
425
+
426
+ logger.info("Waiting for export tasks from %d local exporters (timeout: %ds)", len(all_exporters), timeout)
427
+
428
+ for name, exporter in all_exporters.items():
429
+ try:
430
+ await exporter.wait_for_tasks(timeout=timeout)
431
+ logger.info("Export tasks completed for exporter: %s", name)
432
+ except Exception as e:
433
+ logger.warning("Error waiting for export tasks from %s: %s", name, e)
434
+
435
+ logger.info("All local export task waiting completed")
436
+
437
+ except Exception as e:
438
+ logger.warning("Failed to wait for local export tasks: %s", e)
439
+
440
+ async def run_and_evaluate(self,
441
+ session_manager: SessionManager | None = None,
442
+ job_id: str | None = None) -> EvaluationRunOutput:
443
+ """
444
+ Run the workflow with the specified config file and evaluate the dataset
445
+ """
446
+ logger.info("Starting evaluation run with config file: %s", self.config.config_file)
447
+
448
+ from nat.builder.eval_builder import WorkflowEvalBuilder
449
+ from nat.runtime.loader import load_config
450
+
451
+ # Load and override the config
452
+ config = None
453
+ if isinstance(self.config.config_file, BaseModel):
454
+ config = self.config.config_file
455
+ elif self.config.override:
456
+ config = self.apply_overrides()
457
+ else:
458
+ config = load_config(self.config.config_file)
459
+
460
+ self.eval_config = config.eval
461
+ workflow_alias = self._get_workflow_alias(config.workflow.type)
462
+ logger.debug("Loaded %s evaluation configuration: %s", workflow_alias, self.eval_config)
463
+
464
+ # Cleanup the output directory
465
+ if self.eval_config.general.output:
466
+ self.cleanup_output_directory()
467
+
468
+ # Generate a job_id if append_job_id_to_output_dir is enabled and no job_id provided
469
+ if (self.eval_config.general.output
470
+ and self.eval_config.general.output.job_management.append_job_id_to_output_dir and not job_id):
471
+ job_id = "job_" + str(uuid4())
472
+ logger.info("Generated job ID for output directory: %s", job_id)
473
+
474
+ # If a job id is provided keep the data per-job
475
+ if job_id:
476
+ self.eval_config.general.output_dir = self.eval_config.general.output_dir / f"jobs/{job_id}"
477
+ if self.eval_config.general.output:
478
+ self.eval_config.general.output.dir = self.eval_config.general.output_dir
479
+
480
+ # Load the input dataset
481
+ # For multiple datasets, one handler per dataset can be created
482
+ dataset_config = self.eval_config.general.dataset # Currently only one dataset is supported
483
+ if not dataset_config:
484
+ logger.info("No dataset found, nothing to evaluate")
485
+ return EvaluationRunOutput(workflow_output_file=self.workflow_output_file,
486
+ evaluator_output_files=self.evaluator_output_files,
487
+ workflow_interrupted=self.workflow_interrupted,
488
+ eval_input=EvalInput(eval_input_items=[]),
489
+ evaluation_results=[],
490
+ usage_stats=UsageStats(),
491
+ profiler_results=ProfilerResults())
492
+
493
+ custom_pre_eval_process_function = self.eval_config.general.output.custom_pre_eval_process_function \
494
+ if self.eval_config.general.output else None
495
+ dataset_handler = DatasetHandler(dataset_config=dataset_config,
496
+ reps=self.config.reps,
497
+ concurrency=self.eval_config.general.max_concurrency,
498
+ num_passes=self.config.num_passes,
499
+ adjust_dataset_size=self.config.adjust_dataset_size,
500
+ custom_pre_eval_process_function=custom_pre_eval_process_function)
501
+ self.eval_input = dataset_handler.get_eval_input_from_dataset(self.config.dataset)
502
+ if not self.eval_input.eval_input_items:
503
+ logger.info("Dataset is empty. Nothing to evaluate.")
504
+ return EvaluationRunOutput(workflow_output_file=self.workflow_output_file,
505
+ evaluator_output_files=self.evaluator_output_files,
506
+ workflow_interrupted=self.workflow_interrupted,
507
+ eval_input=self.eval_input,
508
+ evaluation_results=self.evaluation_results,
509
+ usage_stats=self.usage_stats,
510
+ profiler_results=ProfilerResults())
511
+
512
+ # Run workflow and evaluate
513
+ async with WorkflowEvalBuilder.from_config(config=config) as eval_workflow:
514
+ # Initialize Weave integration
515
+ self.weave_eval.initialize_logger(workflow_alias, self.eval_input, config)
516
+
517
+ with self.eval_trace_context.evaluation_context():
518
+ # Run workflow
519
+ if self.config.endpoint:
520
+ await self.run_workflow_remote()
521
+ elif not self.config.skip_workflow:
522
+ if session_manager is None:
523
+ workflow = await eval_workflow.build()
524
+ session_manager = SessionManager(workflow,
525
+ max_concurrency=self.eval_config.general.max_concurrency)
526
+ await self.run_workflow_local(session_manager)
527
+
528
+ # Pre-evaluation process the workflow output
529
+ self.eval_input = dataset_handler.pre_eval_process_eval_input(self.eval_input)
530
+
531
+ # Evaluate
532
+ evaluators = {name: eval_workflow.get_evaluator(name) for name in self.eval_config.evaluators}
533
+ await self.run_evaluators(evaluators)
534
+
535
+ # Wait for all trace export tasks to complete (local workflows only)
536
+ if session_manager and not self.config.endpoint:
537
+ await self.wait_for_all_export_tasks_local(session_manager, timeout=self.config.export_timeout)
538
+
539
+ # Profile the workflow
540
+ profiler_results = await self.profile_workflow()
541
+
542
+ # compute total runtime
543
+ if self.usage_stats.usage_stats_items:
544
+ self.usage_stats.total_runtime = max(self.usage_stats.usage_stats_items.values(),
545
+ key=lambda x: x.max_timestamp).max_timestamp - \
546
+ min(self.usage_stats.usage_stats_items.values(), key=lambda x: x.min_timestamp).min_timestamp
547
+ else:
548
+ self.usage_stats.total_runtime = 0.0
549
+
550
+ # Publish the results
551
+ self.publish_output(dataset_handler, profiler_results)
552
+
553
+ # Run custom scripts and upload evaluation outputs to S3
554
+ if self.eval_config.general.output:
555
+ output_uploader = OutputUploader(self.eval_config.general.output, job_id=job_id)
556
+ output_uploader.run_custom_scripts()
557
+ await output_uploader.upload_directory()
558
+
559
+ return EvaluationRunOutput(workflow_output_file=self.workflow_output_file,
560
+ evaluator_output_files=self.evaluator_output_files,
561
+ workflow_interrupted=self.workflow_interrupted,
562
+ eval_input=self.eval_input,
563
+ evaluation_results=self.evaluation_results,
564
+ usage_stats=self.usage_stats,
565
+ profiler_results=profiler_results)
@@ -0,0 +1,14 @@
1
+ # SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
2
+ # SPDX-License-Identifier: Apache-2.0
3
+ #
4
+ # Licensed under the Apache License, Version 2.0 (the "License");
5
+ # you may not use this file except in compliance with the License.
6
+ # You may obtain a copy of the License at
7
+ #
8
+ # http://www.apache.org/licenses/LICENSE-2.0
9
+ #
10
+ # Unless required by applicable law or agreed to in writing, software
11
+ # distributed under the License is distributed on an "AS IS" BASIS,
12
+ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13
+ # See the License for the specific language governing permissions and
14
+ # limitations under the License.
@@ -0,0 +1,77 @@
1
+ # SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
2
+ # SPDX-License-Identifier: Apache-2.0
3
+ #
4
+ # Licensed under the Apache License, Version 2.0 (the "License");
5
+ # you may not use this file except in compliance with the License.
6
+ # You may obtain a copy of the License at
7
+ #
8
+ # http://www.apache.org/licenses/LICENSE-2.0
9
+ #
10
+ # Unless required by applicable law or agreed to in writing, software
11
+ # distributed under the License is distributed on an "AS IS" BASIS,
12
+ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13
+ # See the License for the specific language governing permissions and
14
+ # limitations under the License.
15
+
16
+ import asyncio
17
+ from abc import ABC
18
+ from abc import abstractmethod
19
+
20
+ from tqdm import tqdm
21
+
22
+ from nat.eval.evaluator.evaluator_model import EvalInput
23
+ from nat.eval.evaluator.evaluator_model import EvalInputItem
24
+ from nat.eval.evaluator.evaluator_model import EvalOutput
25
+ from nat.eval.evaluator.evaluator_model import EvalOutputItem
26
+ from nat.eval.utils.tqdm_position_registry import TqdmPositionRegistry
27
+
28
+
29
+ class BaseEvaluator(ABC):
30
+ """
31
+ Base class for custom evaluators.
32
+
33
+ .. warning::
34
+ **Experimental Feature**: The Evaluation API is experimental and may change in future releases.
35
+ Future versions may introduce breaking changes without notice.
36
+
37
+ Each custom evaluator must implement the `evaluate_item` method which is used to evaluate a
38
+ single EvalInputItem.
39
+ """
40
+
41
+ def __init__(self, max_concurrency: int = 4, tqdm_desc: str = "Evaluating"):
42
+ self.max_concurrency = max_concurrency
43
+ self.semaphore = asyncio.Semaphore(max_concurrency)
44
+ self.tqdm_desc = tqdm_desc
45
+
46
+ @abstractmethod
47
+ async def evaluate_item(self, item: EvalInputItem) -> EvalOutputItem:
48
+ """Each evaluator must implement this for item-level evaluation"""
49
+ pass
50
+
51
+ async def evaluate(self, eval_input: EvalInput) -> EvalOutput:
52
+ pbar = None
53
+ try:
54
+ tqdm_position = TqdmPositionRegistry.claim()
55
+ pbar = tqdm(total=len(eval_input.eval_input_items), desc=self.tqdm_desc, position=tqdm_position)
56
+
57
+ async def wrapped(item):
58
+ async with self.semaphore:
59
+ try:
60
+ output_item = await self.evaluate_item(item)
61
+ pbar.update(1)
62
+ return output_item
63
+ except Exception as e:
64
+ # If the evaluator fails, return an error item with a score of 0.0
65
+ pbar.update(1)
66
+ return EvalOutputItem(id=item.id, score=0.0, reasoning={"error": f"Evaluator error: {str(e)}"})
67
+
68
+ output_items = await asyncio.gather(*[wrapped(item) for item in eval_input.eval_input_items])
69
+ finally:
70
+ pbar.close()
71
+ TqdmPositionRegistry.release(tqdm_position)
72
+
73
+ # Compute average if possible
74
+ numeric_scores = [item.score for item in output_items if isinstance(item.score, int | float)]
75
+ avg_score = round(sum(numeric_scores) / len(numeric_scores), 2) if numeric_scores else None
76
+
77
+ return EvalOutput(average_score=avg_score, eval_output_items=output_items)