nvidia-nat 1.2.0__py3-none-any.whl

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (435) hide show
  1. aiq/__init__.py +66 -0
  2. nat/agent/__init__.py +0 -0
  3. nat/agent/base.py +256 -0
  4. nat/agent/dual_node.py +67 -0
  5. nat/agent/react_agent/__init__.py +0 -0
  6. nat/agent/react_agent/agent.py +363 -0
  7. nat/agent/react_agent/output_parser.py +104 -0
  8. nat/agent/react_agent/prompt.py +44 -0
  9. nat/agent/react_agent/register.py +149 -0
  10. nat/agent/reasoning_agent/__init__.py +0 -0
  11. nat/agent/reasoning_agent/reasoning_agent.py +225 -0
  12. nat/agent/register.py +23 -0
  13. nat/agent/rewoo_agent/__init__.py +0 -0
  14. nat/agent/rewoo_agent/agent.py +415 -0
  15. nat/agent/rewoo_agent/prompt.py +110 -0
  16. nat/agent/rewoo_agent/register.py +157 -0
  17. nat/agent/tool_calling_agent/__init__.py +0 -0
  18. nat/agent/tool_calling_agent/agent.py +119 -0
  19. nat/agent/tool_calling_agent/register.py +106 -0
  20. nat/authentication/__init__.py +14 -0
  21. nat/authentication/api_key/__init__.py +14 -0
  22. nat/authentication/api_key/api_key_auth_provider.py +96 -0
  23. nat/authentication/api_key/api_key_auth_provider_config.py +124 -0
  24. nat/authentication/api_key/register.py +26 -0
  25. nat/authentication/exceptions/__init__.py +14 -0
  26. nat/authentication/exceptions/api_key_exceptions.py +38 -0
  27. nat/authentication/http_basic_auth/__init__.py +0 -0
  28. nat/authentication/http_basic_auth/http_basic_auth_provider.py +81 -0
  29. nat/authentication/http_basic_auth/register.py +30 -0
  30. nat/authentication/interfaces.py +93 -0
  31. nat/authentication/oauth2/__init__.py +14 -0
  32. nat/authentication/oauth2/oauth2_auth_code_flow_provider.py +107 -0
  33. nat/authentication/oauth2/oauth2_auth_code_flow_provider_config.py +39 -0
  34. nat/authentication/oauth2/register.py +25 -0
  35. nat/authentication/register.py +21 -0
  36. nat/builder/__init__.py +0 -0
  37. nat/builder/builder.py +285 -0
  38. nat/builder/component_utils.py +316 -0
  39. nat/builder/context.py +270 -0
  40. nat/builder/embedder.py +24 -0
  41. nat/builder/eval_builder.py +161 -0
  42. nat/builder/evaluator.py +29 -0
  43. nat/builder/framework_enum.py +24 -0
  44. nat/builder/front_end.py +73 -0
  45. nat/builder/function.py +344 -0
  46. nat/builder/function_base.py +380 -0
  47. nat/builder/function_info.py +627 -0
  48. nat/builder/intermediate_step_manager.py +174 -0
  49. nat/builder/llm.py +25 -0
  50. nat/builder/retriever.py +25 -0
  51. nat/builder/user_interaction_manager.py +78 -0
  52. nat/builder/workflow.py +148 -0
  53. nat/builder/workflow_builder.py +1117 -0
  54. nat/cli/__init__.py +14 -0
  55. nat/cli/cli_utils/__init__.py +0 -0
  56. nat/cli/cli_utils/config_override.py +231 -0
  57. nat/cli/cli_utils/validation.py +37 -0
  58. nat/cli/commands/__init__.py +0 -0
  59. nat/cli/commands/configure/__init__.py +0 -0
  60. nat/cli/commands/configure/channel/__init__.py +0 -0
  61. nat/cli/commands/configure/channel/add.py +28 -0
  62. nat/cli/commands/configure/channel/channel.py +34 -0
  63. nat/cli/commands/configure/channel/remove.py +30 -0
  64. nat/cli/commands/configure/channel/update.py +30 -0
  65. nat/cli/commands/configure/configure.py +33 -0
  66. nat/cli/commands/evaluate.py +139 -0
  67. nat/cli/commands/info/__init__.py +14 -0
  68. nat/cli/commands/info/info.py +37 -0
  69. nat/cli/commands/info/list_channels.py +32 -0
  70. nat/cli/commands/info/list_components.py +129 -0
  71. nat/cli/commands/info/list_mcp.py +304 -0
  72. nat/cli/commands/registry/__init__.py +14 -0
  73. nat/cli/commands/registry/publish.py +88 -0
  74. nat/cli/commands/registry/pull.py +118 -0
  75. nat/cli/commands/registry/registry.py +36 -0
  76. nat/cli/commands/registry/remove.py +108 -0
  77. nat/cli/commands/registry/search.py +155 -0
  78. nat/cli/commands/sizing/__init__.py +14 -0
  79. nat/cli/commands/sizing/calc.py +297 -0
  80. nat/cli/commands/sizing/sizing.py +27 -0
  81. nat/cli/commands/start.py +246 -0
  82. nat/cli/commands/uninstall.py +81 -0
  83. nat/cli/commands/validate.py +47 -0
  84. nat/cli/commands/workflow/__init__.py +14 -0
  85. nat/cli/commands/workflow/templates/__init__.py.j2 +0 -0
  86. nat/cli/commands/workflow/templates/config.yml.j2 +16 -0
  87. nat/cli/commands/workflow/templates/pyproject.toml.j2 +22 -0
  88. nat/cli/commands/workflow/templates/register.py.j2 +5 -0
  89. nat/cli/commands/workflow/templates/workflow.py.j2 +36 -0
  90. nat/cli/commands/workflow/workflow.py +37 -0
  91. nat/cli/commands/workflow/workflow_commands.py +317 -0
  92. nat/cli/entrypoint.py +135 -0
  93. nat/cli/main.py +57 -0
  94. nat/cli/register_workflow.py +488 -0
  95. nat/cli/type_registry.py +1000 -0
  96. nat/data_models/__init__.py +14 -0
  97. nat/data_models/api_server.py +716 -0
  98. nat/data_models/authentication.py +231 -0
  99. nat/data_models/common.py +171 -0
  100. nat/data_models/component.py +58 -0
  101. nat/data_models/component_ref.py +168 -0
  102. nat/data_models/config.py +410 -0
  103. nat/data_models/dataset_handler.py +169 -0
  104. nat/data_models/discovery_metadata.py +305 -0
  105. nat/data_models/embedder.py +27 -0
  106. nat/data_models/evaluate.py +127 -0
  107. nat/data_models/evaluator.py +26 -0
  108. nat/data_models/front_end.py +26 -0
  109. nat/data_models/function.py +30 -0
  110. nat/data_models/function_dependencies.py +72 -0
  111. nat/data_models/interactive.py +246 -0
  112. nat/data_models/intermediate_step.py +302 -0
  113. nat/data_models/invocation_node.py +38 -0
  114. nat/data_models/llm.py +27 -0
  115. nat/data_models/logging.py +26 -0
  116. nat/data_models/memory.py +27 -0
  117. nat/data_models/object_store.py +44 -0
  118. nat/data_models/profiler.py +54 -0
  119. nat/data_models/registry_handler.py +26 -0
  120. nat/data_models/retriever.py +30 -0
  121. nat/data_models/retry_mixin.py +35 -0
  122. nat/data_models/span.py +190 -0
  123. nat/data_models/step_adaptor.py +64 -0
  124. nat/data_models/streaming.py +33 -0
  125. nat/data_models/swe_bench_model.py +54 -0
  126. nat/data_models/telemetry_exporter.py +26 -0
  127. nat/data_models/ttc_strategy.py +30 -0
  128. nat/embedder/__init__.py +0 -0
  129. nat/embedder/nim_embedder.py +59 -0
  130. nat/embedder/openai_embedder.py +43 -0
  131. nat/embedder/register.py +22 -0
  132. nat/eval/__init__.py +14 -0
  133. nat/eval/config.py +60 -0
  134. nat/eval/dataset_handler/__init__.py +0 -0
  135. nat/eval/dataset_handler/dataset_downloader.py +106 -0
  136. nat/eval/dataset_handler/dataset_filter.py +52 -0
  137. nat/eval/dataset_handler/dataset_handler.py +367 -0
  138. nat/eval/evaluate.py +510 -0
  139. nat/eval/evaluator/__init__.py +14 -0
  140. nat/eval/evaluator/base_evaluator.py +77 -0
  141. nat/eval/evaluator/evaluator_model.py +45 -0
  142. nat/eval/intermediate_step_adapter.py +99 -0
  143. nat/eval/rag_evaluator/__init__.py +0 -0
  144. nat/eval/rag_evaluator/evaluate.py +178 -0
  145. nat/eval/rag_evaluator/register.py +143 -0
  146. nat/eval/register.py +23 -0
  147. nat/eval/remote_workflow.py +133 -0
  148. nat/eval/runners/__init__.py +14 -0
  149. nat/eval/runners/config.py +39 -0
  150. nat/eval/runners/multi_eval_runner.py +54 -0
  151. nat/eval/runtime_event_subscriber.py +52 -0
  152. nat/eval/swe_bench_evaluator/__init__.py +0 -0
  153. nat/eval/swe_bench_evaluator/evaluate.py +215 -0
  154. nat/eval/swe_bench_evaluator/register.py +36 -0
  155. nat/eval/trajectory_evaluator/__init__.py +0 -0
  156. nat/eval/trajectory_evaluator/evaluate.py +75 -0
  157. nat/eval/trajectory_evaluator/register.py +40 -0
  158. nat/eval/tunable_rag_evaluator/__init__.py +0 -0
  159. nat/eval/tunable_rag_evaluator/evaluate.py +245 -0
  160. nat/eval/tunable_rag_evaluator/register.py +52 -0
  161. nat/eval/usage_stats.py +41 -0
  162. nat/eval/utils/__init__.py +0 -0
  163. nat/eval/utils/output_uploader.py +140 -0
  164. nat/eval/utils/tqdm_position_registry.py +40 -0
  165. nat/eval/utils/weave_eval.py +184 -0
  166. nat/experimental/__init__.py +0 -0
  167. nat/experimental/decorators/__init__.py +0 -0
  168. nat/experimental/decorators/experimental_warning_decorator.py +134 -0
  169. nat/experimental/test_time_compute/__init__.py +0 -0
  170. nat/experimental/test_time_compute/editing/__init__.py +0 -0
  171. nat/experimental/test_time_compute/editing/iterative_plan_refinement_editor.py +147 -0
  172. nat/experimental/test_time_compute/editing/llm_as_a_judge_editor.py +204 -0
  173. nat/experimental/test_time_compute/editing/motivation_aware_summarization.py +107 -0
  174. nat/experimental/test_time_compute/functions/__init__.py +0 -0
  175. nat/experimental/test_time_compute/functions/execute_score_select_function.py +105 -0
  176. nat/experimental/test_time_compute/functions/plan_select_execute_function.py +224 -0
  177. nat/experimental/test_time_compute/functions/ttc_tool_orchestration_function.py +205 -0
  178. nat/experimental/test_time_compute/functions/ttc_tool_wrapper_function.py +146 -0
  179. nat/experimental/test_time_compute/models/__init__.py +0 -0
  180. nat/experimental/test_time_compute/models/editor_config.py +132 -0
  181. nat/experimental/test_time_compute/models/scoring_config.py +112 -0
  182. nat/experimental/test_time_compute/models/search_config.py +120 -0
  183. nat/experimental/test_time_compute/models/selection_config.py +154 -0
  184. nat/experimental/test_time_compute/models/stage_enums.py +43 -0
  185. nat/experimental/test_time_compute/models/strategy_base.py +66 -0
  186. nat/experimental/test_time_compute/models/tool_use_config.py +41 -0
  187. nat/experimental/test_time_compute/models/ttc_item.py +48 -0
  188. nat/experimental/test_time_compute/register.py +36 -0
  189. nat/experimental/test_time_compute/scoring/__init__.py +0 -0
  190. nat/experimental/test_time_compute/scoring/llm_based_agent_scorer.py +168 -0
  191. nat/experimental/test_time_compute/scoring/llm_based_plan_scorer.py +168 -0
  192. nat/experimental/test_time_compute/scoring/motivation_aware_scorer.py +111 -0
  193. nat/experimental/test_time_compute/search/__init__.py +0 -0
  194. nat/experimental/test_time_compute/search/multi_llm_planner.py +128 -0
  195. nat/experimental/test_time_compute/search/multi_query_retrieval_search.py +122 -0
  196. nat/experimental/test_time_compute/search/single_shot_multi_plan_planner.py +128 -0
  197. nat/experimental/test_time_compute/selection/__init__.py +0 -0
  198. nat/experimental/test_time_compute/selection/best_of_n_selector.py +63 -0
  199. nat/experimental/test_time_compute/selection/llm_based_agent_output_selector.py +131 -0
  200. nat/experimental/test_time_compute/selection/llm_based_output_merging_selector.py +159 -0
  201. nat/experimental/test_time_compute/selection/llm_based_plan_selector.py +128 -0
  202. nat/experimental/test_time_compute/selection/threshold_selector.py +58 -0
  203. nat/front_ends/__init__.py +14 -0
  204. nat/front_ends/console/__init__.py +14 -0
  205. nat/front_ends/console/authentication_flow_handler.py +233 -0
  206. nat/front_ends/console/console_front_end_config.py +32 -0
  207. nat/front_ends/console/console_front_end_plugin.py +96 -0
  208. nat/front_ends/console/register.py +25 -0
  209. nat/front_ends/cron/__init__.py +14 -0
  210. nat/front_ends/fastapi/__init__.py +14 -0
  211. nat/front_ends/fastapi/auth_flow_handlers/__init__.py +0 -0
  212. nat/front_ends/fastapi/auth_flow_handlers/http_flow_handler.py +27 -0
  213. nat/front_ends/fastapi/auth_flow_handlers/websocket_flow_handler.py +107 -0
  214. nat/front_ends/fastapi/fastapi_front_end_config.py +241 -0
  215. nat/front_ends/fastapi/fastapi_front_end_controller.py +68 -0
  216. nat/front_ends/fastapi/fastapi_front_end_plugin.py +116 -0
  217. nat/front_ends/fastapi/fastapi_front_end_plugin_worker.py +1087 -0
  218. nat/front_ends/fastapi/html_snippets/__init__.py +14 -0
  219. nat/front_ends/fastapi/html_snippets/auth_code_grant_success.py +35 -0
  220. nat/front_ends/fastapi/intermediate_steps_subscriber.py +80 -0
  221. nat/front_ends/fastapi/job_store.py +183 -0
  222. nat/front_ends/fastapi/main.py +72 -0
  223. nat/front_ends/fastapi/message_handler.py +320 -0
  224. nat/front_ends/fastapi/message_validator.py +352 -0
  225. nat/front_ends/fastapi/register.py +25 -0
  226. nat/front_ends/fastapi/response_helpers.py +195 -0
  227. nat/front_ends/fastapi/step_adaptor.py +319 -0
  228. nat/front_ends/mcp/__init__.py +14 -0
  229. nat/front_ends/mcp/mcp_front_end_config.py +36 -0
  230. nat/front_ends/mcp/mcp_front_end_plugin.py +81 -0
  231. nat/front_ends/mcp/mcp_front_end_plugin_worker.py +143 -0
  232. nat/front_ends/mcp/register.py +27 -0
  233. nat/front_ends/mcp/tool_converter.py +241 -0
  234. nat/front_ends/register.py +22 -0
  235. nat/front_ends/simple_base/__init__.py +14 -0
  236. nat/front_ends/simple_base/simple_front_end_plugin_base.py +54 -0
  237. nat/llm/__init__.py +0 -0
  238. nat/llm/aws_bedrock_llm.py +57 -0
  239. nat/llm/nim_llm.py +46 -0
  240. nat/llm/openai_llm.py +46 -0
  241. nat/llm/register.py +23 -0
  242. nat/llm/utils/__init__.py +14 -0
  243. nat/llm/utils/env_config_value.py +94 -0
  244. nat/llm/utils/error.py +17 -0
  245. nat/memory/__init__.py +20 -0
  246. nat/memory/interfaces.py +183 -0
  247. nat/memory/models.py +112 -0
  248. nat/meta/pypi.md +58 -0
  249. nat/object_store/__init__.py +20 -0
  250. nat/object_store/in_memory_object_store.py +76 -0
  251. nat/object_store/interfaces.py +84 -0
  252. nat/object_store/models.py +38 -0
  253. nat/object_store/register.py +20 -0
  254. nat/observability/__init__.py +14 -0
  255. nat/observability/exporter/__init__.py +14 -0
  256. nat/observability/exporter/base_exporter.py +449 -0
  257. nat/observability/exporter/exporter.py +78 -0
  258. nat/observability/exporter/file_exporter.py +33 -0
  259. nat/observability/exporter/processing_exporter.py +322 -0
  260. nat/observability/exporter/raw_exporter.py +52 -0
  261. nat/observability/exporter/span_exporter.py +288 -0
  262. nat/observability/exporter_manager.py +335 -0
  263. nat/observability/mixin/__init__.py +14 -0
  264. nat/observability/mixin/batch_config_mixin.py +26 -0
  265. nat/observability/mixin/collector_config_mixin.py +23 -0
  266. nat/observability/mixin/file_mixin.py +288 -0
  267. nat/observability/mixin/file_mode.py +23 -0
  268. nat/observability/mixin/resource_conflict_mixin.py +134 -0
  269. nat/observability/mixin/serialize_mixin.py +61 -0
  270. nat/observability/mixin/type_introspection_mixin.py +183 -0
  271. nat/observability/processor/__init__.py +14 -0
  272. nat/observability/processor/batching_processor.py +310 -0
  273. nat/observability/processor/callback_processor.py +42 -0
  274. nat/observability/processor/intermediate_step_serializer.py +28 -0
  275. nat/observability/processor/processor.py +71 -0
  276. nat/observability/register.py +96 -0
  277. nat/observability/utils/__init__.py +14 -0
  278. nat/observability/utils/dict_utils.py +236 -0
  279. nat/observability/utils/time_utils.py +31 -0
  280. nat/plugins/.namespace +1 -0
  281. nat/profiler/__init__.py +0 -0
  282. nat/profiler/calc/__init__.py +14 -0
  283. nat/profiler/calc/calc_runner.py +627 -0
  284. nat/profiler/calc/calculations.py +288 -0
  285. nat/profiler/calc/data_models.py +188 -0
  286. nat/profiler/calc/plot.py +345 -0
  287. nat/profiler/callbacks/__init__.py +0 -0
  288. nat/profiler/callbacks/agno_callback_handler.py +295 -0
  289. nat/profiler/callbacks/base_callback_class.py +20 -0
  290. nat/profiler/callbacks/langchain_callback_handler.py +290 -0
  291. nat/profiler/callbacks/llama_index_callback_handler.py +205 -0
  292. nat/profiler/callbacks/semantic_kernel_callback_handler.py +238 -0
  293. nat/profiler/callbacks/token_usage_base_model.py +27 -0
  294. nat/profiler/data_frame_row.py +51 -0
  295. nat/profiler/data_models.py +24 -0
  296. nat/profiler/decorators/__init__.py +0 -0
  297. nat/profiler/decorators/framework_wrapper.py +131 -0
  298. nat/profiler/decorators/function_tracking.py +254 -0
  299. nat/profiler/forecasting/__init__.py +0 -0
  300. nat/profiler/forecasting/config.py +18 -0
  301. nat/profiler/forecasting/model_trainer.py +75 -0
  302. nat/profiler/forecasting/models/__init__.py +22 -0
  303. nat/profiler/forecasting/models/forecasting_base_model.py +40 -0
  304. nat/profiler/forecasting/models/linear_model.py +197 -0
  305. nat/profiler/forecasting/models/random_forest_regressor.py +269 -0
  306. nat/profiler/inference_metrics_model.py +28 -0
  307. nat/profiler/inference_optimization/__init__.py +0 -0
  308. nat/profiler/inference_optimization/bottleneck_analysis/__init__.py +0 -0
  309. nat/profiler/inference_optimization/bottleneck_analysis/nested_stack_analysis.py +460 -0
  310. nat/profiler/inference_optimization/bottleneck_analysis/simple_stack_analysis.py +258 -0
  311. nat/profiler/inference_optimization/data_models.py +386 -0
  312. nat/profiler/inference_optimization/experimental/__init__.py +0 -0
  313. nat/profiler/inference_optimization/experimental/concurrency_spike_analysis.py +468 -0
  314. nat/profiler/inference_optimization/experimental/prefix_span_analysis.py +405 -0
  315. nat/profiler/inference_optimization/llm_metrics.py +212 -0
  316. nat/profiler/inference_optimization/prompt_caching.py +163 -0
  317. nat/profiler/inference_optimization/token_uniqueness.py +107 -0
  318. nat/profiler/inference_optimization/workflow_runtimes.py +72 -0
  319. nat/profiler/intermediate_property_adapter.py +102 -0
  320. nat/profiler/profile_runner.py +473 -0
  321. nat/profiler/utils.py +184 -0
  322. nat/registry_handlers/__init__.py +0 -0
  323. nat/registry_handlers/local/__init__.py +0 -0
  324. nat/registry_handlers/local/local_handler.py +176 -0
  325. nat/registry_handlers/local/register_local.py +37 -0
  326. nat/registry_handlers/metadata_factory.py +60 -0
  327. nat/registry_handlers/package_utils.py +571 -0
  328. nat/registry_handlers/pypi/__init__.py +0 -0
  329. nat/registry_handlers/pypi/pypi_handler.py +251 -0
  330. nat/registry_handlers/pypi/register_pypi.py +40 -0
  331. nat/registry_handlers/register.py +21 -0
  332. nat/registry_handlers/registry_handler_base.py +157 -0
  333. nat/registry_handlers/rest/__init__.py +0 -0
  334. nat/registry_handlers/rest/register_rest.py +56 -0
  335. nat/registry_handlers/rest/rest_handler.py +237 -0
  336. nat/registry_handlers/schemas/__init__.py +0 -0
  337. nat/registry_handlers/schemas/headers.py +42 -0
  338. nat/registry_handlers/schemas/package.py +68 -0
  339. nat/registry_handlers/schemas/publish.py +68 -0
  340. nat/registry_handlers/schemas/pull.py +82 -0
  341. nat/registry_handlers/schemas/remove.py +36 -0
  342. nat/registry_handlers/schemas/search.py +91 -0
  343. nat/registry_handlers/schemas/status.py +47 -0
  344. nat/retriever/__init__.py +0 -0
  345. nat/retriever/interface.py +41 -0
  346. nat/retriever/milvus/__init__.py +14 -0
  347. nat/retriever/milvus/register.py +81 -0
  348. nat/retriever/milvus/retriever.py +228 -0
  349. nat/retriever/models.py +77 -0
  350. nat/retriever/nemo_retriever/__init__.py +14 -0
  351. nat/retriever/nemo_retriever/register.py +60 -0
  352. nat/retriever/nemo_retriever/retriever.py +190 -0
  353. nat/retriever/register.py +22 -0
  354. nat/runtime/__init__.py +14 -0
  355. nat/runtime/loader.py +220 -0
  356. nat/runtime/runner.py +195 -0
  357. nat/runtime/session.py +162 -0
  358. nat/runtime/user_metadata.py +130 -0
  359. nat/settings/__init__.py +0 -0
  360. nat/settings/global_settings.py +318 -0
  361. nat/test/.namespace +1 -0
  362. nat/tool/__init__.py +0 -0
  363. nat/tool/chat_completion.py +74 -0
  364. nat/tool/code_execution/README.md +151 -0
  365. nat/tool/code_execution/__init__.py +0 -0
  366. nat/tool/code_execution/code_sandbox.py +267 -0
  367. nat/tool/code_execution/local_sandbox/.gitignore +1 -0
  368. nat/tool/code_execution/local_sandbox/Dockerfile.sandbox +60 -0
  369. nat/tool/code_execution/local_sandbox/__init__.py +13 -0
  370. nat/tool/code_execution/local_sandbox/local_sandbox_server.py +198 -0
  371. nat/tool/code_execution/local_sandbox/sandbox.requirements.txt +6 -0
  372. nat/tool/code_execution/local_sandbox/start_local_sandbox.sh +50 -0
  373. nat/tool/code_execution/register.py +74 -0
  374. nat/tool/code_execution/test_code_execution_sandbox.py +414 -0
  375. nat/tool/code_execution/utils.py +100 -0
  376. nat/tool/datetime_tools.py +42 -0
  377. nat/tool/document_search.py +141 -0
  378. nat/tool/github_tools/__init__.py +0 -0
  379. nat/tool/github_tools/create_github_commit.py +133 -0
  380. nat/tool/github_tools/create_github_issue.py +87 -0
  381. nat/tool/github_tools/create_github_pr.py +106 -0
  382. nat/tool/github_tools/get_github_file.py +106 -0
  383. nat/tool/github_tools/get_github_issue.py +166 -0
  384. nat/tool/github_tools/get_github_pr.py +256 -0
  385. nat/tool/github_tools/update_github_issue.py +100 -0
  386. nat/tool/mcp/__init__.py +14 -0
  387. nat/tool/mcp/exceptions.py +142 -0
  388. nat/tool/mcp/mcp_client.py +255 -0
  389. nat/tool/mcp/mcp_tool.py +96 -0
  390. nat/tool/memory_tools/__init__.py +0 -0
  391. nat/tool/memory_tools/add_memory_tool.py +79 -0
  392. nat/tool/memory_tools/delete_memory_tool.py +67 -0
  393. nat/tool/memory_tools/get_memory_tool.py +72 -0
  394. nat/tool/nvidia_rag.py +95 -0
  395. nat/tool/register.py +38 -0
  396. nat/tool/retriever.py +94 -0
  397. nat/tool/server_tools.py +66 -0
  398. nat/utils/__init__.py +0 -0
  399. nat/utils/data_models/__init__.py +0 -0
  400. nat/utils/data_models/schema_validator.py +58 -0
  401. nat/utils/debugging_utils.py +43 -0
  402. nat/utils/dump_distro_mapping.py +32 -0
  403. nat/utils/exception_handlers/__init__.py +0 -0
  404. nat/utils/exception_handlers/automatic_retries.py +289 -0
  405. nat/utils/exception_handlers/mcp.py +211 -0
  406. nat/utils/exception_handlers/schemas.py +114 -0
  407. nat/utils/io/__init__.py +0 -0
  408. nat/utils/io/model_processing.py +28 -0
  409. nat/utils/io/yaml_tools.py +119 -0
  410. nat/utils/log_utils.py +37 -0
  411. nat/utils/metadata_utils.py +74 -0
  412. nat/utils/optional_imports.py +142 -0
  413. nat/utils/producer_consumer_queue.py +178 -0
  414. nat/utils/reactive/__init__.py +0 -0
  415. nat/utils/reactive/base/__init__.py +0 -0
  416. nat/utils/reactive/base/observable_base.py +65 -0
  417. nat/utils/reactive/base/observer_base.py +55 -0
  418. nat/utils/reactive/base/subject_base.py +79 -0
  419. nat/utils/reactive/observable.py +59 -0
  420. nat/utils/reactive/observer.py +76 -0
  421. nat/utils/reactive/subject.py +131 -0
  422. nat/utils/reactive/subscription.py +49 -0
  423. nat/utils/settings/__init__.py +0 -0
  424. nat/utils/settings/global_settings.py +197 -0
  425. nat/utils/string_utils.py +38 -0
  426. nat/utils/type_converter.py +290 -0
  427. nat/utils/type_utils.py +484 -0
  428. nat/utils/url_utils.py +27 -0
  429. nvidia_nat-1.2.0.dist-info/METADATA +365 -0
  430. nvidia_nat-1.2.0.dist-info/RECORD +435 -0
  431. nvidia_nat-1.2.0.dist-info/WHEEL +5 -0
  432. nvidia_nat-1.2.0.dist-info/entry_points.txt +21 -0
  433. nvidia_nat-1.2.0.dist-info/licenses/LICENSE-3rd-party.txt +5478 -0
  434. nvidia_nat-1.2.0.dist-info/licenses/LICENSE.md +201 -0
  435. nvidia_nat-1.2.0.dist-info/top_level.txt +2 -0
@@ -0,0 +1,473 @@
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 json
17
+ import logging
18
+ import math
19
+ import os
20
+ import statistics
21
+ from pathlib import Path
22
+ from typing import Any
23
+
24
+ from pydantic import BaseModel
25
+
26
+ from nat.data_models.evaluate import ProfilerConfig
27
+ from nat.data_models.intermediate_step import IntermediateStep
28
+ from nat.profiler.data_models import ProfilerResults
29
+ from nat.profiler.forecasting.model_trainer import ModelTrainer
30
+ from nat.profiler.inference_metrics_model import InferenceMetricsModel
31
+ from nat.profiler.utils import create_standardized_dataframe
32
+ from nat.utils.type_converter import TypeConverter
33
+
34
+ logger = logging.getLogger(__name__)
35
+
36
+
37
+ class SimpleMetricsHolder(BaseModel):
38
+ workflow_run_time_confidence_intervals: Any
39
+ llm_latency_confidence_intervals: Any
40
+ throughput_estimate_confidence_interval: Any
41
+
42
+
43
+ class InferenceOptimizationHolder(BaseModel):
44
+ confidence_intervals: SimpleMetricsHolder
45
+ common_prefixes: Any
46
+ token_uniqueness: Any
47
+ workflow_runtimes: Any
48
+
49
+
50
+ class ProfilerRunner:
51
+ """
52
+ A utility to run a series of prompts through a NAT workflow for profiling:
53
+
54
+ - can load prompts from a file
55
+ - or generate them via an LLM
56
+ - collect usage stats for each run
57
+ - store them in a configured directory
58
+
59
+ Updated version with additional metrics:
60
+
61
+ - For each request, we collect a list of UsageStatistic objects, store them individually,
62
+ and also keep a final large JSON of all requests.
63
+ - We then compute:
64
+ 1. 90, 95, 99% confidence intervals for the mean total workflow run time.
65
+ 2. 90, 95, 99% confidence intervals for the mean LLM latency.
66
+ 3. 90, 95, 99% estimates of throughput.
67
+
68
+ All computed metrics are saved to a metrics JSON file at the end.
69
+ """
70
+
71
+ def __init__(self, profiler_config: ProfilerConfig, output_dir: Path, write_output: bool = True):
72
+ self.profile_config = profiler_config
73
+ self.output_dir = output_dir
74
+ self.write_output = write_output
75
+ self._converter = TypeConverter([])
76
+
77
+ # Holds per-request data (prompt, output, usage_stats, etc.)
78
+ # This will be saved at the end to a big JSON file
79
+ self.all_requests_data: list[dict] = []
80
+ self.all_steps = []
81
+
82
+ # Ensure output directory
83
+ os.makedirs(output_dir, exist_ok=True)
84
+
85
+ async def run(self, all_steps: list[list[IntermediateStep]]) -> ProfilerResults:
86
+ """
87
+ Main entrypoint: Works on Input DataFrame generated from eval to fit forecasting model,
88
+ writes out combined requests JSON, then computes and saves additional metrics,
89
+ and optionally fits a forecasting model.
90
+ """
91
+ from nat.profiler.inference_optimization.bottleneck_analysis.nested_stack_analysis import \
92
+ multi_example_call_profiling
93
+ from nat.profiler.inference_optimization.bottleneck_analysis.simple_stack_analysis import \
94
+ profile_workflow_bottlenecks
95
+ from nat.profiler.inference_optimization.experimental.concurrency_spike_analysis import \
96
+ concurrency_spike_analysis
97
+ from nat.profiler.inference_optimization.experimental.prefix_span_analysis import \
98
+ prefixspan_subworkflow_with_text
99
+ from nat.profiler.inference_optimization.llm_metrics import LLMMetrics
100
+ from nat.profiler.inference_optimization.prompt_caching import get_common_prefixes
101
+ from nat.profiler.inference_optimization.token_uniqueness import compute_inter_query_token_uniqueness_by_llm
102
+ from nat.profiler.inference_optimization.workflow_runtimes import compute_workflow_runtime_metrics
103
+ from nat.profiler.intermediate_property_adapter import IntermediatePropertyAdaptor
104
+
105
+ # Convert the incoming DataFrame to a list of dicts and store
106
+ all_steps = [[IntermediatePropertyAdaptor.from_intermediate_step(step) for step in steps]
107
+ for steps in all_steps] # Add adapter properties to each step
108
+
109
+ self.all_steps = all_steps
110
+ self.all_requests_data = []
111
+ for i, steps in enumerate(all_steps):
112
+ request_data = []
113
+ for step in steps:
114
+ request_data.append(step.model_dump())
115
+ self.all_requests_data.append({"request_number": i, "intermediate_steps": request_data})
116
+
117
+ # Write the final big JSON (all requests)
118
+ if self.write_output:
119
+ final_path = os.path.join(self.output_dir, "all_requests_profiler_traces.json")
120
+ with open(final_path, 'w', encoding='utf-8') as f:
121
+ json.dump(self.all_requests_data, f, indent=2, default=str)
122
+ logger.info("Wrote combined data to: %s", final_path)
123
+
124
+ # ------------------------------------------------------------
125
+ # Generate one standardized dataframe for all usage stats
126
+ # ------------------------------------------------------------
127
+ merged_df = create_standardized_dataframe(all_steps)
128
+
129
+ if self.profile_config.compute_llm_metrics and not merged_df.empty:
130
+ merged_df = LLMMetrics.compute_profiling_metrics(all_steps)
131
+
132
+ output_df = merged_df.copy()
133
+
134
+ if self.profile_config.csv_exclude_io_text and not output_df.empty:
135
+ # Exclude text fields from CSV
136
+ output_df = output_df.drop(columns=['llm_text_input', 'llm_text_output', 'llm_new_token'])
137
+
138
+ # Write this single CSV
139
+ csv_path = os.path.join(self.output_dir, "standardized_data_all.csv")
140
+ output_df.to_csv(csv_path, index=False, encoding='utf-8')
141
+ logger.info("Wrote merged standardized DataFrame to %s", csv_path)
142
+
143
+ # ------------------------------------------------------------
144
+ # Compute and save additional performance metrics
145
+ # ------------------------------------------------------------
146
+ workflow_run_time_ci: InferenceMetricsModel = self._compute_workflow_run_time_confidence_intervals()
147
+
148
+ # 2. 90, 95, 99% confidence intervals of mean LLM latency
149
+ llm_latency_ci: InferenceMetricsModel = self._compute_llm_latency_confidence_intervals()
150
+
151
+ # 3. 90, 95, 99% estimates of throughput
152
+ throughput_ci: InferenceMetricsModel = self._compute_throughput_estimates()
153
+
154
+ # Collect all computed metrics
155
+ simple_metrics = SimpleMetricsHolder(workflow_run_time_confidence_intervals=workflow_run_time_ci.model_dump(),
156
+ llm_latency_confidence_intervals=llm_latency_ci.model_dump(),
157
+ throughput_estimate_confidence_interval=throughput_ci.model_dump())
158
+
159
+ common_prefix_results = token_uniqueness_results = workflow_runtimes_results = None
160
+
161
+ if self.profile_config.prompt_caching_prefixes.enable:
162
+ # ------------------------------------------------------------
163
+ # Compute and save common prefixes
164
+ # ------------------------------------------------------------
165
+
166
+ prefixes = get_common_prefixes(all_steps, self.profile_config.prompt_caching_prefixes.min_frequency)
167
+ common_prefix_results = prefixes
168
+
169
+ if self.profile_config.token_uniqueness_forecast:
170
+ # ------------------------------------------------------------
171
+ # Compute and save inter-query token uniqueness
172
+ # ------------------------------------------------------------
173
+
174
+ uniqueness = compute_inter_query_token_uniqueness_by_llm(all_steps)
175
+ token_uniqueness_results = uniqueness
176
+
177
+ if self.profile_config.workflow_runtime_forecast or self.profile_config.base_metrics:
178
+ # ------------------------------------------------------------
179
+ # Compute and save workflow runtime metrics
180
+ # ------------------------------------------------------------
181
+
182
+ workflow_runtimes = compute_workflow_runtime_metrics(all_steps)
183
+ workflow_runtimes_results = workflow_runtimes
184
+
185
+ inference_optimization_results = InferenceOptimizationHolder(confidence_intervals=simple_metrics,
186
+ common_prefixes=common_prefix_results,
187
+ token_uniqueness=token_uniqueness_results,
188
+ workflow_runtimes=workflow_runtimes_results)
189
+
190
+ if self.write_output and inference_optimization_results:
191
+ # Save to JSON
192
+ optimization_results_path = os.path.join(self.output_dir, "inference_optimization.json")
193
+ with open(optimization_results_path, 'w', encoding='utf-8') as f:
194
+ json.dump(inference_optimization_results.model_dump(), f, indent=2)
195
+ logger.info("Wrote inference optimization results to: %s", optimization_results_path)
196
+
197
+ workflow_profiling_reports = ""
198
+ workflow_profiling_metrics = {}
199
+
200
+ if self.profile_config.bottleneck_analysis.enable_simple_stack:
201
+ # ------------------------------------------------------------
202
+ # Profile workflow bottlenecks
203
+ # ------------------------------------------------------------
204
+
205
+ workflow_bottlenecks = profile_workflow_bottlenecks(all_steps)
206
+ workflow_bottlenecks = workflow_bottlenecks.model_dump()
207
+ workflow_profiling_reports += "\n\n\n" + workflow_bottlenecks["summary"]
208
+ workflow_profiling_metrics["simple_stack_analysis"] = workflow_bottlenecks["stats"]
209
+ logger.info("Simple stack analysis complete")
210
+
211
+ if self.profile_config.bottleneck_analysis.enable_nested_stack:
212
+ # ------------------------------------------------------------
213
+ # Profile workflow bottlenecks with nested stack analysis
214
+ # ------------------------------------------------------------
215
+ nested_bottlenecks = multi_example_call_profiling(all_steps, output_dir=str(self.output_dir))
216
+ workflow_profiling_reports += "\n\n\n" + nested_bottlenecks.textual_report
217
+ workflow_profiling_metrics["nested_stack_analysis"] = nested_bottlenecks.model_dump(
218
+ exclude=["textual_report"])
219
+ logger.info("Nested stack analysis complete")
220
+
221
+ if self.profile_config.concurrency_spike_analysis.enable:
222
+ # ------------------------------------------------------------
223
+ # Profile concurrency spikes
224
+ # ------------------------------------------------------------
225
+ concurrency_metrics = concurrency_spike_analysis(
226
+ all_steps, self.profile_config.concurrency_spike_analysis.spike_threshold)
227
+ workflow_profiling_reports += "\n\n\n" + concurrency_metrics.textual_report
228
+ workflow_profiling_metrics["concurrency_spike_analysis"] = concurrency_metrics.model_dump(
229
+ exclude=["textual_report"])
230
+ logger.info("Concurrency spike analysis complete")
231
+
232
+ if self.profile_config.prefix_span_analysis.enable:
233
+ # ------------------------------------------------------------
234
+ # Profile prefix span analysis
235
+ # ------------------------------------------------------------
236
+ prefix_list = []
237
+ if (self.profile_config.prefix_span_analysis.chain_with_common_prefixes
238
+ and "common_prefixes" in inference_optimization_results):
239
+ logger.info("Using common prefixes for prefix span analysis")
240
+ for _, llm_data in inference_optimization_results["common_prefixes"].items():
241
+ for prefix_data in llm_data["prefix_info"]:
242
+ prefix_list.append(prefix_data["prefix"])
243
+
244
+ prefix_span_analysis = prefixspan_subworkflow_with_text(
245
+ all_steps,
246
+ **self.profile_config.prefix_span_analysis.model_dump(exclude=["enable", "chain_with_common_prefixes"]),
247
+ prefix_list=prefix_list)
248
+
249
+ workflow_profiling_reports += "\n\n\n" + prefix_span_analysis.textual_report
250
+ workflow_profiling_metrics["prefix_span_analysis"] = prefix_span_analysis.model_dump(
251
+ exclude=["textual_report"])
252
+ logger.info("Prefix span analysis complete")
253
+
254
+ if self.write_output and workflow_profiling_reports:
255
+ # Save to text file
256
+ profiling_report_path = os.path.join(self.output_dir, "workflow_profiling_report.txt")
257
+ with open(profiling_report_path, 'w', encoding='utf-8') as f:
258
+ f.write(workflow_profiling_reports)
259
+ logger.info("Wrote workflow profiling report to: %s", profiling_report_path)
260
+
261
+ if self.write_output and workflow_profiling_metrics:
262
+ # Save to JSON
263
+ profiling_metrics_path = os.path.join(self.output_dir, "workflow_profiling_metrics.json")
264
+ with open(profiling_metrics_path, 'w', encoding='utf-8') as f:
265
+ json.dump(workflow_profiling_metrics, f, indent=2)
266
+ logger.info("Wrote workflow profiling metrics to: %s", profiling_metrics_path)
267
+
268
+ if self.profile_config.token_usage_forecast:
269
+ # ------------------------------------------------------------
270
+ # Fit forecasting model and save
271
+ # ------------------------------------------------------------
272
+
273
+ logger.info("Fitting model for forecasting.")
274
+ model_trainer = ModelTrainer()
275
+
276
+ try:
277
+ fitted_model = model_trainer.train(all_steps)
278
+ logger.info("Fitted model for forecasting.")
279
+ except Exception as e:
280
+ logger.exception("Fitting model failed. %s", e, exc_info=True)
281
+ return ProfilerResults()
282
+
283
+ if self.write_output:
284
+ os.makedirs(self.output_dir, exist_ok=True)
285
+
286
+ import pickle
287
+ with open(os.path.join(self.output_dir, "fitted_model.pkl"), 'wb') as f:
288
+ pickle.dump(fitted_model, f)
289
+
290
+ logger.info("Saved fitted model to disk.")
291
+
292
+ return ProfilerResults(workflow_runtime_metrics=workflow_runtimes_results, llm_latency_ci=llm_latency_ci)
293
+
294
+ # -------------------------------------------------------------------
295
+ # Confidence Intervals / Metrics
296
+ # -------------------------------------------------------------------
297
+ def _compute_workflow_run_time_confidence_intervals(self) -> InferenceMetricsModel:
298
+ """
299
+ Computes 90, 95, 99% confidence intervals for the mean total workflow run time (in seconds).
300
+ The total workflow run time for each request is the difference between the last and first
301
+ event timestamps in usage_stats.
302
+ """
303
+ run_times = []
304
+ for req_data in self.all_steps:
305
+ # Find the min and max event_timestamp
306
+ timestamps = [u.event_timestamp for u in req_data]
307
+ if not timestamps:
308
+ continue
309
+
310
+ start_time = min(timestamps)
311
+ end_time = max(timestamps)
312
+ run_times.append(end_time - start_time)
313
+
314
+ return self._compute_confidence_intervals(run_times, "Workflow Run Time")
315
+
316
+ def _compute_llm_latency_confidence_intervals(self) -> InferenceMetricsModel:
317
+ """
318
+ Computes 90, 95, 99% confidence intervals for the mean LLM latency.
319
+ LLM latency is defined as the difference between an LLM_END event_timestamp and
320
+ the immediately preceding LLM_START event_timestamp, across all usage_stats.
321
+ """
322
+ latencies = []
323
+ for req_data in self.all_steps:
324
+
325
+ usage_stats_sorted = sorted(req_data, key=lambda x: x.event_timestamp)
326
+
327
+ previous_llm_start_time = None
328
+ for u in usage_stats_sorted:
329
+ event_type = u.event_type.value
330
+ ts = u.event_timestamp
331
+ if event_type == "LLM_START":
332
+ previous_llm_start_time = ts
333
+ elif event_type == "LLM_END" and previous_llm_start_time is not None:
334
+ latencies.append(ts - previous_llm_start_time)
335
+ previous_llm_start_time = None
336
+
337
+ return self._compute_confidence_intervals(latencies, "LLM Latency")
338
+
339
+ def _compute_throughput_estimates(self) -> InferenceMetricsModel:
340
+ """
341
+ Computes 90, 95, 99% confidence intervals for throughput, defined as:
342
+
343
+ | throughput = (total number of requests) / (total time window),
344
+
345
+ where total time window is from the earliest usage_stats event across all requests
346
+ to the latest usage_stats event.
347
+ Note: This is a simple approximate measure of overall throughput for the entire run.
348
+ """
349
+ # Gather min timestamp and max timestamp across ALL requests
350
+ all_timestamps = []
351
+ for req_data in self.all_steps:
352
+ for u in req_data:
353
+ all_timestamps.append(u.event_timestamp)
354
+
355
+ if not all_timestamps:
356
+ return InferenceMetricsModel()
357
+
358
+ min_ts = min(all_timestamps)
359
+ max_ts = max(all_timestamps)
360
+ total_time = max_ts - min_ts
361
+ if total_time <= 0:
362
+ # Can't compute a meaningful throughput if time <= 0
363
+ return InferenceMetricsModel()
364
+
365
+ total_requests = len(self.all_requests_data)
366
+ # Single estimate of throughput
367
+ throughput_value = total_requests / total_time
368
+
369
+ # For confidence intervals of throughput, we do a simplistic assumption:
370
+ # We treat each request's contribution as 1 occurrence, and approximate
371
+ # the distribution as if these arrivals were uniform. This is quite simplified.
372
+ # We can compute a standard error: SE = sqrt(throughput_value / total_time)
373
+ # However, a more accurate approach might require a different method (e.g., Poisson).
374
+ # We'll do a naive normal approximation here.
375
+
376
+ # We'll guess that the standard deviation of #requests is sqrt(N), so stdev_n ~ sqrt(N).
377
+ # stdev_time is quite small though. We'll do a naive approach:
378
+ # We'll treat the throughput as a sample mean with n=total_requests.
379
+ # Then standard error is (throughput_value / sqrt(n)).
380
+ # This is purely heuristic.
381
+ n = total_requests
382
+ if n <= 1:
383
+ return InferenceMetricsModel()
384
+
385
+ # A rough standard error for throughput:
386
+ standard_error = throughput_value / math.sqrt(n)
387
+
388
+ # Build confidence intervals using z-scores for 90%, 95%, 99%
389
+ intervals = {'n': total_requests, 'mean': throughput_value}
390
+ for confidence, zvalue in \
391
+ [("ninetieth_interval", 1.645), ("ninety_fifth_interval", 1.96), ("ninety_ninth_interval", 2.576)]:
392
+ ci_lower = throughput_value - zvalue * standard_error
393
+ ci_upper = throughput_value + zvalue * standard_error
394
+ intervals[confidence] = (max(ci_lower, 0.0), ci_upper)
395
+
396
+ return InferenceMetricsModel(**intervals)
397
+
398
+ def _compute_confidence_intervals(self, data: list[float], metric_name: str) -> InferenceMetricsModel:
399
+ """
400
+ Helper to compute 90, 95, 99 % confidence intervals **and** the empirical
401
+ 90th/95th/99th percentiles (p90/p95/p99) for the mean of a dataset.
402
+ Uses a z-score from the normal approximation for large samples.
403
+
404
+ Returns a dict like::
405
+
406
+ {
407
+ 'ninetieth_interval': (lower, upper),
408
+ 'ninety_fifth_interval': (lower, upper),
409
+ 'ninety_ninth_interval': (lower, upper),
410
+ }
411
+ """
412
+ if not data:
413
+ logger.warning("No data points for %s, cannot compute intervals.", metric_name)
414
+ return InferenceMetricsModel()
415
+
416
+ n = len(data)
417
+ mean_val = statistics.mean(data)
418
+ if n <= 1:
419
+ return InferenceMetricsModel(
420
+ n=n,
421
+ mean=mean_val,
422
+ ninetieth_interval=(mean_val, mean_val),
423
+ ninety_fifth_interval=(mean_val, mean_val),
424
+ ninety_ninth_interval=(mean_val, mean_val),
425
+ p90=mean_val,
426
+ p95=mean_val,
427
+ p99=mean_val,
428
+ )
429
+
430
+ stdev_val = statistics.pstdev(data) # population stdev or use stdev for sample
431
+ # standard error
432
+ se = stdev_val / math.sqrt(n)
433
+
434
+ intervals = {}
435
+ for confidence, zvalue in \
436
+ [("ninetieth_interval", 1.645), ("ninety_fifth_interval", 1.96), ("ninety_ninth_interval", 2.576)]:
437
+ margin = zvalue * se
438
+ lower = mean_val - margin
439
+ upper = mean_val + margin
440
+ intervals[confidence] = (lower, upper)
441
+
442
+ # Optionally, store more info
443
+ intervals["n"] = n
444
+ intervals["mean"] = mean_val
445
+
446
+ # ------------------------------------------------------------------
447
+ # Percentiles
448
+ # ------------------------------------------------------------------
449
+ sorted_data = sorted(data)
450
+
451
+ def _percentile(arr: list[float], pct: float) -> float:
452
+ """
453
+ Linear interpolation between closest ranks.
454
+ pct is given from 0‑100 (e.g. 90 for p90).
455
+ """
456
+ if not arr:
457
+ return 0.0
458
+ k = (len(arr) - 1) * (pct / 100.0)
459
+ f = math.floor(k)
460
+ c = math.ceil(k)
461
+ if f == c:
462
+ return arr[int(k)]
463
+ return arr[f] + (arr[c] - arr[f]) * (k - f)
464
+
465
+ p90_val = _percentile(sorted_data, 90)
466
+ p95_val = _percentile(sorted_data, 95)
467
+ p99_val = _percentile(sorted_data, 99)
468
+
469
+ intervals["p90"] = p90_val
470
+ intervals["p95"] = p95_val
471
+ intervals["p99"] = p99_val
472
+
473
+ return InferenceMetricsModel(**intervals)
nat/profiler/utils.py ADDED
@@ -0,0 +1,184 @@
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 inspect
17
+ import logging
18
+ import re
19
+ from collections.abc import Callable
20
+ from typing import Any
21
+
22
+ import pandas as pd
23
+
24
+ from nat.builder.framework_enum import LLMFrameworkEnum
25
+ from nat.cli.type_registry import RegisteredFunctionInfo
26
+ from nat.data_models.intermediate_step import IntermediateStep
27
+ from nat.profiler.data_frame_row import DataFrameRow
28
+
29
+ # A simple set of regex patterns to scan for direct references to LLMFrameworkEnum
30
+ _FRAMEWORK_REGEX_MAP = {t: fr'\b{t._name_}\b' for t in LLMFrameworkEnum}
31
+
32
+ logger = logging.getLogger(__name__)
33
+
34
+
35
+ def detect_llm_frameworks_in_build_fn(registration: RegisteredFunctionInfo) -> list[LLMFrameworkEnum]:
36
+ """
37
+ Analyze a function's source (the build_fn) to see which LLM frameworks it uses. Also recurses
38
+ into any additional Python functions that the build_fn calls while passing `builder`, so that
39
+ references to LLMFrameworkEnum in those helper calls are also detected.
40
+
41
+ 1. If `registration.framework_wrappers` is non-empty, we return that first.
42
+ (We do convert them to LLMFrameworkEnum if possible.)
43
+ 2. Otherwise, we attempt to:
44
+
45
+ - Get the build_fn's source via `inspect.getsource(...)`
46
+ - Parse it for references to LLMFrameworkEnum
47
+ - Find any function calls that include the word "builder" in the arguments
48
+
49
+ - Recursively parse those functions' source code for frameworks
50
+
51
+ 3. If we cannot parse the source at all (e.g. OSError), we return a list of all frameworks.
52
+ """
53
+ # ----------------------------------------------------------------
54
+ # 1) If frameworks were explicitly declared in registration.framework_wrappers, use them:
55
+ if registration.framework_wrappers:
56
+ results: list[LLMFrameworkEnum] = []
57
+ for fw_str in registration.framework_wrappers:
58
+ try:
59
+ results.append(LLMFrameworkEnum(fw_str))
60
+ except ValueError:
61
+ # If it's not recognized, ignore or log
62
+ logger.warning("Unrecognized framework %s in registration.framework_wrappers", fw_str)
63
+
64
+ return list(set(results)) # unique
65
+ # ----------------------------------------------------------------
66
+
67
+ # Because we want to recursively parse code, we'll keep track of visited function objects
68
+ visited_fns: set[Callable[..., Any]] = set()
69
+ # We also need a place to store discovered frameworks
70
+ discovered: set[LLMFrameworkEnum] = set()
71
+
72
+ def _parse_source_for_frameworks(src: str) -> None:
73
+ """Check lines for any direct references to LLMFrameworkEnum.* or placeholders in the map."""
74
+ for fw_enum_member, pattern in _FRAMEWORK_REGEX_MAP.items():
75
+ if re.search(pattern, src):
76
+ discovered.add(fw_enum_member)
77
+
78
+ def _find_builder_func_calls(src: str) -> list[str]:
79
+ """
80
+ Look for calls of the form: some_func(..., builder, ...)
81
+ or some_func(..., builder=..., ...)
82
+
83
+ This returns the name of each function we found being called, e.g. 'some_func'.
84
+ It's a naive best-effort approach
85
+ and group(1) is the function name.
86
+ """
87
+ # E.g. foo(builder) or foo( param=..., builder=builder )
88
+ pattern = r'(\w+)\s*\([^)]*\bbuilder\b[^)]*\)'
89
+ return re.findall(pattern, src)
90
+
91
+ def _recurse_parse(fn: Callable[..., Any], visited: set[Callable[..., Any]]) -> None:
92
+ """Recursively parse the source code of `fn`, add discovered frameworks,
93
+ and parse any new functions that get called with 'builder'."""
94
+ if fn in visited:
95
+ return
96
+ visited.add(fn)
97
+
98
+ try:
99
+ src = inspect.getsource(fn)
100
+ except OSError:
101
+ # If we can't parse source, we add all frameworks and bail
102
+ discovered.update([k for k, v in _FRAMEWORK_REGEX_MAP.items()])
103
+ return
104
+
105
+ # parse direct references
106
+ _parse_source_for_frameworks(src)
107
+
108
+ # parse any function calls that pass in "builder"
109
+ child_func_names = _find_builder_func_calls(src)
110
+ if not child_func_names:
111
+ return
112
+
113
+ # We'll try to find these child functions in the same module as `fn`
114
+ mod = inspect.getmodule(fn)
115
+ if not mod:
116
+ return
117
+ # We'll see if the child function is a top-level in that module
118
+ for child_name in child_func_names:
119
+ # get the function object if it exists in the module
120
+ child_obj = getattr(mod, child_name, None)
121
+ if callable(child_obj):
122
+ _recurse_parse(child_obj, visited)
123
+
124
+ # ----------------------------------------------------------------
125
+ # 2) Actually do the BFS/DFS parse on `registration.build_fn`
126
+ main_fn = registration.build_fn
127
+
128
+ try:
129
+ _recurse_parse(main_fn, visited_fns)
130
+ except Exception:
131
+ # If an unexpected error occurs, fallback to "all frameworks"
132
+ discovered.update([k for k, v in _FRAMEWORK_REGEX_MAP.items()])
133
+ # ----------------------------------------------------------------
134
+ if len(discovered) > 0:
135
+ logger.warning(
136
+ "Discovered frameworks: %s in function %s by inspecting "
137
+ "source. It is recommended and more reliable to instead add the used LLMFrameworkEnum "
138
+ "types in the framework_wrappers argument when calling @register_function.",
139
+ discovered,
140
+ main_fn.__name__)
141
+
142
+ return list(discovered)
143
+
144
+
145
+ # -------------------------------------------------------------------
146
+ # Create a single standardized DataFrame for all usage stats
147
+ # -------------------------------------------------------------------
148
+ def create_standardized_dataframe(requests_data: list[list[IntermediateStep]]) -> pd.DataFrame:
149
+ """
150
+ Merge usage stats for *all* requests into one DataFrame, each row representing a usage_stats entry.
151
+ - Include a column 'example_number' to mark which request it originated from.
152
+ """
153
+ all_rows = []
154
+ try:
155
+ for i, steps in enumerate(requests_data):
156
+ for step in steps:
157
+ # Create a DataFrameRow
158
+ all_rows.append(
159
+ DataFrameRow(event_timestamp=step.event_timestamp,
160
+ example_number=i,
161
+ prompt_tokens=step.token_usage.prompt_tokens,
162
+ completion_tokens=step.token_usage.completion_tokens,
163
+ total_tokens=step.token_usage.total_tokens,
164
+ llm_text_input=step.llm_text_input,
165
+ llm_text_output=step.llm_text_output,
166
+ llm_new_token=step.llm_text_chunk,
167
+ llm_name=step.llm_name,
168
+ tool_name=step.tool_name,
169
+ function_name=step.function_name,
170
+ function_id=step.function_id,
171
+ parent_function_name=step.parent_function_name,
172
+ parent_function_id=step.parent_function_id,
173
+ UUID=step.payload.UUID,
174
+ framework=step.framework,
175
+ event_type=step.event_type).model_dump(), )
176
+
177
+ except Exception as e:
178
+ logger.exception("Error creating standardized DataFrame: %s", e, exc_info=True)
179
+ return pd.DataFrame()
180
+
181
+ if not all_rows:
182
+ return pd.DataFrame()
183
+
184
+ return pd.DataFrame.from_records(all_rows)
File without changes
File without changes