nvidia-nat 1.4.0a20251120__py3-none-any.whl → 1.4.0a20260113__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 (492) hide show
  1. aiq/__init__.py +1 -1
  2. nat/{front_ends/mcp → agent/auto_memory_wrapper}/__init__.py +1 -1
  3. nat/agent/auto_memory_wrapper/agent.py +278 -0
  4. nat/agent/auto_memory_wrapper/register.py +227 -0
  5. nat/agent/auto_memory_wrapper/state.py +30 -0
  6. nat/agent/base.py +1 -1
  7. nat/agent/dual_node.py +1 -1
  8. nat/agent/prompt_optimizer/prompt.py +1 -1
  9. nat/agent/prompt_optimizer/register.py +1 -1
  10. nat/agent/react_agent/agent.py +16 -9
  11. nat/agent/react_agent/output_parser.py +2 -2
  12. nat/agent/react_agent/prompt.py +3 -2
  13. nat/agent/react_agent/register.py +2 -2
  14. nat/agent/react_agent/register_per_user_agent.py +104 -0
  15. nat/agent/reasoning_agent/reasoning_agent.py +1 -1
  16. nat/agent/register.py +3 -1
  17. nat/agent/responses_api_agent/__init__.py +1 -1
  18. nat/agent/responses_api_agent/register.py +1 -1
  19. nat/agent/rewoo_agent/agent.py +9 -4
  20. nat/agent/rewoo_agent/prompt.py +1 -1
  21. nat/agent/rewoo_agent/register.py +1 -1
  22. nat/agent/tool_calling_agent/agent.py +5 -4
  23. nat/agent/tool_calling_agent/register.py +1 -1
  24. nat/authentication/__init__.py +1 -1
  25. nat/authentication/api_key/__init__.py +1 -1
  26. nat/authentication/api_key/api_key_auth_provider.py +1 -1
  27. nat/authentication/api_key/api_key_auth_provider_config.py +22 -7
  28. nat/authentication/api_key/register.py +1 -1
  29. nat/authentication/credential_validator/__init__.py +1 -1
  30. nat/authentication/credential_validator/bearer_token_validator.py +1 -1
  31. nat/authentication/exceptions/__init__.py +1 -1
  32. nat/authentication/exceptions/api_key_exceptions.py +1 -1
  33. nat/authentication/http_basic_auth/http_basic_auth_provider.py +1 -1
  34. nat/authentication/http_basic_auth/register.py +1 -1
  35. nat/authentication/interfaces.py +1 -1
  36. nat/authentication/oauth2/__init__.py +1 -1
  37. nat/authentication/oauth2/oauth2_auth_code_flow_provider.py +1 -1
  38. nat/authentication/oauth2/oauth2_auth_code_flow_provider_config.py +1 -1
  39. nat/authentication/oauth2/oauth2_resource_server_config.py +1 -1
  40. nat/authentication/oauth2/register.py +1 -1
  41. nat/authentication/register.py +1 -1
  42. nat/builder/builder.py +511 -1
  43. nat/builder/child_builder.py +385 -0
  44. nat/builder/component_utils.py +28 -4
  45. nat/builder/context.py +17 -1
  46. nat/builder/embedder.py +1 -1
  47. nat/builder/eval_builder.py +19 -7
  48. nat/builder/evaluator.py +1 -1
  49. nat/builder/framework_enum.py +2 -1
  50. nat/builder/front_end.py +1 -1
  51. nat/builder/function.py +40 -3
  52. nat/builder/function_base.py +1 -1
  53. nat/builder/function_info.py +1 -1
  54. nat/builder/intermediate_step_manager.py +1 -1
  55. nat/builder/llm.py +1 -1
  56. nat/builder/per_user_workflow_builder.py +843 -0
  57. nat/builder/retriever.py +1 -1
  58. nat/builder/sync_builder.py +571 -0
  59. nat/builder/user_interaction_manager.py +1 -1
  60. nat/builder/workflow.py +1 -1
  61. nat/builder/workflow_builder.py +536 -424
  62. nat/cli/__init__.py +1 -1
  63. nat/cli/cli_utils/config_override.py +1 -1
  64. nat/cli/cli_utils/validation.py +32 -1
  65. nat/cli/commands/configure/channel/add.py +1 -1
  66. nat/cli/commands/configure/channel/channel.py +1 -1
  67. nat/cli/commands/configure/channel/remove.py +1 -1
  68. nat/cli/commands/configure/channel/update.py +1 -1
  69. nat/cli/commands/configure/configure.py +1 -1
  70. nat/cli/commands/evaluate.py +87 -13
  71. nat/cli/commands/finetune.py +132 -0
  72. nat/cli/commands/info/__init__.py +1 -1
  73. nat/cli/commands/info/info.py +1 -1
  74. nat/cli/commands/info/list_channels.py +1 -1
  75. nat/cli/commands/info/list_components.py +1 -1
  76. nat/cli/commands/object_store/__init__.py +1 -1
  77. nat/cli/commands/object_store/object_store.py +1 -1
  78. nat/cli/commands/optimize.py +1 -1
  79. nat/cli/commands/{mcp → red_teaming}/__init__.py +1 -1
  80. nat/cli/commands/red_teaming/red_teaming.py +138 -0
  81. nat/cli/commands/red_teaming/red_teaming_utils.py +73 -0
  82. nat/cli/commands/registry/__init__.py +1 -1
  83. nat/cli/commands/registry/publish.py +1 -1
  84. nat/cli/commands/registry/pull.py +1 -1
  85. nat/cli/commands/registry/registry.py +1 -1
  86. nat/cli/commands/registry/remove.py +1 -1
  87. nat/cli/commands/registry/search.py +1 -1
  88. nat/cli/commands/sizing/__init__.py +1 -1
  89. nat/cli/commands/sizing/calc.py +1 -1
  90. nat/cli/commands/sizing/sizing.py +1 -1
  91. nat/cli/commands/start.py +1 -1
  92. nat/cli/commands/uninstall.py +1 -1
  93. nat/cli/commands/validate.py +1 -1
  94. nat/cli/commands/workflow/__init__.py +1 -1
  95. nat/cli/commands/workflow/workflow.py +1 -1
  96. nat/cli/commands/workflow/workflow_commands.py +3 -2
  97. nat/cli/entrypoint.py +15 -37
  98. nat/cli/main.py +2 -2
  99. nat/cli/plugin_loader.py +69 -0
  100. nat/cli/register_workflow.py +183 -5
  101. nat/cli/type_registry.py +169 -3
  102. nat/control_flow/register.py +1 -1
  103. nat/control_flow/router_agent/agent.py +1 -1
  104. nat/control_flow/router_agent/prompt.py +1 -1
  105. nat/control_flow/router_agent/register.py +1 -1
  106. nat/control_flow/sequential_executor.py +28 -7
  107. nat/data_models/__init__.py +1 -1
  108. nat/data_models/agent.py +1 -1
  109. nat/data_models/api_server.py +38 -3
  110. nat/data_models/authentication.py +1 -1
  111. nat/data_models/common.py +1 -1
  112. nat/data_models/component.py +7 -1
  113. nat/data_models/component_ref.py +34 -1
  114. nat/data_models/config.py +62 -1
  115. nat/data_models/dataset_handler.py +15 -2
  116. nat/data_models/discovery_metadata.py +1 -1
  117. nat/data_models/embedder.py +1 -1
  118. nat/data_models/evaluate.py +6 -1
  119. nat/data_models/evaluator.py +1 -1
  120. nat/data_models/finetuning.py +260 -0
  121. nat/data_models/front_end.py +1 -1
  122. nat/data_models/function.py +1 -1
  123. nat/data_models/function_dependencies.py +1 -1
  124. nat/data_models/gated_field_mixin.py +1 -1
  125. nat/data_models/interactive.py +1 -1
  126. nat/data_models/intermediate_step.py +29 -2
  127. nat/data_models/invocation_node.py +1 -1
  128. nat/data_models/llm.py +1 -1
  129. nat/data_models/logging.py +1 -1
  130. nat/data_models/memory.py +1 -1
  131. nat/data_models/middleware.py +3 -1
  132. nat/data_models/object_store.py +1 -1
  133. nat/data_models/openai_mcp.py +1 -1
  134. nat/data_models/optimizable.py +1 -1
  135. nat/data_models/optimizer.py +1 -1
  136. nat/data_models/profiler.py +1 -1
  137. nat/data_models/registry_handler.py +1 -1
  138. nat/data_models/retriever.py +1 -1
  139. nat/data_models/retry_mixin.py +1 -1
  140. nat/data_models/runtime_enum.py +1 -1
  141. nat/data_models/span.py +1 -1
  142. nat/data_models/step_adaptor.py +1 -1
  143. nat/data_models/streaming.py +1 -1
  144. nat/data_models/swe_bench_model.py +1 -1
  145. nat/data_models/telemetry_exporter.py +1 -1
  146. nat/data_models/thinking_mixin.py +1 -1
  147. nat/data_models/ttc_strategy.py +1 -1
  148. nat/embedder/azure_openai_embedder.py +1 -1
  149. nat/embedder/nim_embedder.py +1 -1
  150. nat/embedder/openai_embedder.py +1 -1
  151. nat/embedder/register.py +1 -1
  152. nat/eval/__init__.py +1 -1
  153. nat/eval/config.py +8 -1
  154. nat/eval/dataset_handler/dataset_downloader.py +1 -1
  155. nat/eval/dataset_handler/dataset_filter.py +1 -1
  156. nat/eval/dataset_handler/dataset_handler.py +4 -2
  157. nat/eval/evaluate.py +217 -80
  158. nat/eval/evaluator/__init__.py +1 -1
  159. nat/eval/evaluator/base_evaluator.py +2 -2
  160. nat/eval/evaluator/evaluator_model.py +3 -2
  161. nat/eval/intermediate_step_adapter.py +1 -1
  162. nat/eval/llm_validator.py +336 -0
  163. nat/eval/rag_evaluator/evaluate.py +17 -10
  164. nat/eval/rag_evaluator/register.py +1 -1
  165. nat/eval/red_teaming_evaluator/__init__.py +14 -0
  166. nat/eval/red_teaming_evaluator/data_models.py +66 -0
  167. nat/eval/red_teaming_evaluator/evaluate.py +327 -0
  168. nat/eval/red_teaming_evaluator/filter_conditions.py +75 -0
  169. nat/eval/red_teaming_evaluator/register.py +55 -0
  170. nat/eval/register.py +2 -1
  171. nat/eval/remote_workflow.py +1 -1
  172. nat/eval/runners/__init__.py +1 -1
  173. nat/eval/runners/config.py +1 -1
  174. nat/eval/runners/multi_eval_runner.py +1 -1
  175. nat/eval/runners/red_teaming_runner/__init__.py +24 -0
  176. nat/eval/runners/red_teaming_runner/config.py +282 -0
  177. nat/eval/runners/red_teaming_runner/report_utils.py +707 -0
  178. nat/eval/runners/red_teaming_runner/runner.py +867 -0
  179. nat/eval/runtime_evaluator/__init__.py +1 -1
  180. nat/eval/runtime_evaluator/evaluate.py +1 -1
  181. nat/eval/runtime_evaluator/register.py +1 -1
  182. nat/eval/runtime_event_subscriber.py +1 -1
  183. nat/eval/swe_bench_evaluator/evaluate.py +1 -1
  184. nat/eval/swe_bench_evaluator/register.py +1 -1
  185. nat/eval/trajectory_evaluator/evaluate.py +2 -2
  186. nat/eval/trajectory_evaluator/register.py +1 -1
  187. nat/eval/tunable_rag_evaluator/evaluate.py +5 -5
  188. nat/eval/tunable_rag_evaluator/register.py +1 -1
  189. nat/eval/usage_stats.py +1 -1
  190. nat/eval/utils/eval_trace_ctx.py +1 -1
  191. nat/eval/utils/output_uploader.py +1 -1
  192. nat/eval/utils/tqdm_position_registry.py +1 -1
  193. nat/eval/utils/weave_eval.py +1 -1
  194. nat/experimental/decorators/experimental_warning_decorator.py +1 -1
  195. nat/experimental/test_time_compute/editing/iterative_plan_refinement_editor.py +1 -1
  196. nat/experimental/test_time_compute/editing/llm_as_a_judge_editor.py +1 -1
  197. nat/experimental/test_time_compute/editing/motivation_aware_summarization.py +1 -1
  198. nat/experimental/test_time_compute/functions/execute_score_select_function.py +1 -1
  199. nat/experimental/test_time_compute/functions/multi_llm_judge_function.py +88 -0
  200. nat/experimental/test_time_compute/functions/plan_select_execute_function.py +1 -1
  201. nat/experimental/test_time_compute/functions/ttc_tool_orchestration_function.py +1 -1
  202. nat/experimental/test_time_compute/functions/ttc_tool_wrapper_function.py +1 -1
  203. nat/experimental/test_time_compute/models/editor_config.py +1 -1
  204. nat/experimental/test_time_compute/models/scoring_config.py +1 -1
  205. nat/experimental/test_time_compute/models/search_config.py +20 -2
  206. nat/experimental/test_time_compute/models/selection_config.py +33 -2
  207. nat/experimental/test_time_compute/models/stage_enums.py +1 -1
  208. nat/experimental/test_time_compute/models/strategy_base.py +1 -1
  209. nat/experimental/test_time_compute/models/tool_use_config.py +1 -1
  210. nat/experimental/test_time_compute/models/ttc_item.py +1 -1
  211. nat/experimental/test_time_compute/register.py +4 -1
  212. nat/experimental/test_time_compute/scoring/llm_based_agent_scorer.py +1 -1
  213. nat/experimental/test_time_compute/scoring/llm_based_plan_scorer.py +1 -1
  214. nat/experimental/test_time_compute/scoring/motivation_aware_scorer.py +1 -1
  215. nat/experimental/test_time_compute/search/multi_llm_generation.py +115 -0
  216. nat/experimental/test_time_compute/search/multi_llm_planner.py +1 -1
  217. nat/experimental/test_time_compute/search/multi_query_retrieval_search.py +1 -1
  218. nat/experimental/test_time_compute/search/single_shot_multi_plan_planner.py +1 -1
  219. nat/experimental/test_time_compute/selection/best_of_n_selector.py +1 -1
  220. nat/experimental/test_time_compute/selection/llm_based_agent_output_selector.py +1 -1
  221. nat/experimental/test_time_compute/selection/llm_based_output_merging_selector.py +1 -1
  222. nat/experimental/test_time_compute/selection/llm_based_plan_selector.py +1 -1
  223. nat/experimental/test_time_compute/selection/llm_judge_selection.py +127 -0
  224. nat/experimental/test_time_compute/selection/threshold_selector.py +1 -1
  225. nat/finetuning/__init__.py +24 -0
  226. nat/finetuning/finetuning_runtime.py +143 -0
  227. nat/finetuning/interfaces/__init__.py +24 -0
  228. nat/finetuning/interfaces/finetuning_runner.py +261 -0
  229. nat/finetuning/interfaces/trainer_adapter.py +103 -0
  230. nat/finetuning/interfaces/trajectory_builder.py +115 -0
  231. nat/finetuning/utils/__init__.py +15 -0
  232. nat/finetuning/utils/parsers/__init__.py +15 -0
  233. nat/finetuning/utils/parsers/adk_parser.py +141 -0
  234. nat/finetuning/utils/parsers/base_parser.py +238 -0
  235. nat/finetuning/utils/parsers/common.py +91 -0
  236. nat/finetuning/utils/parsers/langchain_parser.py +267 -0
  237. nat/finetuning/utils/parsers/llama_index_parser.py +218 -0
  238. nat/front_ends/__init__.py +1 -1
  239. nat/front_ends/console/__init__.py +1 -1
  240. nat/front_ends/console/authentication_flow_handler.py +1 -1
  241. nat/front_ends/console/console_front_end_config.py +4 -1
  242. nat/front_ends/console/console_front_end_plugin.py +5 -4
  243. nat/front_ends/console/register.py +1 -1
  244. nat/front_ends/cron/__init__.py +1 -1
  245. nat/front_ends/fastapi/__init__.py +1 -1
  246. nat/front_ends/fastapi/async_job.py +128 -0
  247. nat/front_ends/fastapi/auth_flow_handlers/http_flow_handler.py +1 -1
  248. nat/front_ends/fastapi/auth_flow_handlers/websocket_flow_handler.py +13 -9
  249. nat/front_ends/fastapi/dask_client_mixin.py +1 -1
  250. nat/front_ends/fastapi/fastapi_front_end_config.py +1 -1
  251. nat/front_ends/fastapi/fastapi_front_end_controller.py +1 -1
  252. nat/front_ends/fastapi/fastapi_front_end_plugin.py +25 -30
  253. nat/front_ends/fastapi/fastapi_front_end_plugin_worker.py +195 -60
  254. nat/front_ends/fastapi/html_snippets/__init__.py +1 -1
  255. nat/front_ends/fastapi/html_snippets/auth_code_grant_success.py +1 -1
  256. nat/front_ends/fastapi/intermediate_steps_subscriber.py +12 -1
  257. nat/front_ends/fastapi/job_store.py +23 -11
  258. nat/front_ends/fastapi/main.py +1 -1
  259. nat/front_ends/fastapi/message_handler.py +27 -4
  260. nat/front_ends/fastapi/message_validator.py +54 -2
  261. nat/front_ends/fastapi/register.py +1 -1
  262. nat/front_ends/fastapi/response_helpers.py +16 -15
  263. nat/front_ends/fastapi/step_adaptor.py +1 -1
  264. nat/front_ends/fastapi/utils.py +1 -1
  265. nat/front_ends/register.py +1 -2
  266. nat/front_ends/simple_base/__init__.py +1 -1
  267. nat/front_ends/simple_base/simple_front_end_plugin_base.py +6 -4
  268. nat/llm/aws_bedrock_llm.py +1 -1
  269. nat/llm/azure_openai_llm.py +10 -1
  270. nat/llm/dynamo_llm.py +363 -0
  271. nat/llm/huggingface_llm.py +177 -0
  272. nat/llm/litellm_llm.py +1 -1
  273. nat/llm/nim_llm.py +1 -1
  274. nat/llm/openai_llm.py +1 -1
  275. nat/llm/register.py +3 -1
  276. nat/llm/utils/__init__.py +1 -1
  277. nat/llm/utils/env_config_value.py +1 -1
  278. nat/llm/utils/error.py +1 -1
  279. nat/llm/utils/thinking.py +1 -1
  280. nat/memory/__init__.py +1 -1
  281. nat/memory/interfaces.py +1 -1
  282. nat/memory/models.py +1 -1
  283. nat/meta/pypi.md +1 -1
  284. nat/middleware/__init__.py +5 -5
  285. nat/middleware/cache/__init__.py +14 -0
  286. nat/middleware/{cache_middleware.py → cache/cache_middleware.py} +39 -42
  287. nat/middleware/cache/cache_middleware_config.py +44 -0
  288. nat/middleware/cache/register.py +33 -0
  289. nat/middleware/defense/__init__.py +14 -0
  290. nat/middleware/defense/defense_middleware.py +362 -0
  291. nat/middleware/defense/defense_middleware_content_guard.py +455 -0
  292. nat/middleware/defense/defense_middleware_data_models.py +91 -0
  293. nat/middleware/defense/defense_middleware_output_verifier.py +440 -0
  294. nat/middleware/defense/defense_middleware_pii.py +356 -0
  295. nat/middleware/defense/register.py +82 -0
  296. nat/middleware/dynamic/__init__.py +14 -0
  297. nat/middleware/dynamic/dynamic_function_middleware.py +962 -0
  298. nat/middleware/dynamic/dynamic_middleware_config.py +132 -0
  299. nat/middleware/dynamic/register.py +34 -0
  300. nat/middleware/function_middleware.py +236 -52
  301. nat/middleware/logging/__init__.py +14 -0
  302. nat/middleware/logging/logging_middleware.py +67 -0
  303. nat/middleware/logging/logging_middleware_config.py +28 -0
  304. nat/middleware/logging/register.py +33 -0
  305. nat/middleware/middleware.py +142 -28
  306. nat/middleware/red_teaming/__init__.py +14 -0
  307. nat/middleware/red_teaming/red_teaming_middleware.py +344 -0
  308. nat/middleware/red_teaming/red_teaming_middleware_config.py +112 -0
  309. nat/middleware/red_teaming/register.py +47 -0
  310. nat/middleware/register.py +7 -20
  311. nat/middleware/utils/__init__.py +14 -0
  312. nat/middleware/utils/workflow_inventory.py +155 -0
  313. nat/object_store/__init__.py +1 -1
  314. nat/object_store/in_memory_object_store.py +1 -1
  315. nat/object_store/interfaces.py +1 -1
  316. nat/object_store/models.py +1 -1
  317. nat/object_store/register.py +1 -1
  318. nat/observability/__init__.py +1 -1
  319. nat/observability/exporter/__init__.py +1 -1
  320. nat/observability/exporter/base_exporter.py +1 -1
  321. nat/observability/exporter/exporter.py +1 -1
  322. nat/observability/exporter/file_exporter.py +1 -1
  323. nat/observability/exporter/processing_exporter.py +1 -1
  324. nat/observability/exporter/raw_exporter.py +1 -1
  325. nat/observability/exporter/span_exporter.py +7 -1
  326. nat/observability/exporter_manager.py +1 -1
  327. nat/observability/mixin/__init__.py +1 -1
  328. nat/observability/mixin/batch_config_mixin.py +1 -1
  329. nat/observability/mixin/collector_config_mixin.py +1 -1
  330. nat/observability/mixin/file_mixin.py +1 -1
  331. nat/observability/mixin/file_mode.py +1 -1
  332. nat/observability/mixin/redaction_config_mixin.py +1 -1
  333. nat/observability/mixin/resource_conflict_mixin.py +1 -1
  334. nat/observability/mixin/serialize_mixin.py +1 -1
  335. nat/observability/mixin/tagging_config_mixin.py +1 -1
  336. nat/observability/mixin/type_introspection_mixin.py +1 -1
  337. nat/observability/processor/__init__.py +1 -1
  338. nat/observability/processor/batching_processor.py +1 -1
  339. nat/observability/processor/callback_processor.py +1 -1
  340. nat/observability/processor/falsy_batch_filter_processor.py +1 -1
  341. nat/observability/processor/intermediate_step_serializer.py +1 -1
  342. nat/observability/processor/processor.py +1 -1
  343. nat/observability/processor/processor_factory.py +1 -1
  344. nat/observability/processor/redaction/__init__.py +1 -1
  345. nat/observability/processor/redaction/contextual_redaction_processor.py +1 -1
  346. nat/observability/processor/redaction/contextual_span_redaction_processor.py +1 -1
  347. nat/observability/processor/redaction/redaction_processor.py +1 -1
  348. nat/observability/processor/redaction/span_header_redaction_processor.py +1 -1
  349. nat/observability/processor/span_tagging_processor.py +1 -1
  350. nat/observability/register.py +1 -1
  351. nat/observability/utils/__init__.py +1 -1
  352. nat/observability/utils/dict_utils.py +1 -1
  353. nat/observability/utils/time_utils.py +1 -1
  354. nat/profiler/calc/__init__.py +1 -1
  355. nat/profiler/calc/calc_runner.py +3 -3
  356. nat/profiler/calc/calculations.py +1 -1
  357. nat/profiler/calc/data_models.py +1 -1
  358. nat/profiler/calc/plot.py +30 -3
  359. nat/profiler/callbacks/agno_callback_handler.py +1 -1
  360. nat/profiler/callbacks/base_callback_class.py +1 -1
  361. nat/profiler/callbacks/langchain_callback_handler.py +33 -3
  362. nat/profiler/callbacks/llama_index_callback_handler.py +13 -10
  363. nat/profiler/callbacks/semantic_kernel_callback_handler.py +1 -1
  364. nat/profiler/callbacks/token_usage_base_model.py +1 -1
  365. nat/profiler/data_frame_row.py +1 -1
  366. nat/profiler/data_models.py +1 -1
  367. nat/profiler/decorators/framework_wrapper.py +16 -1
  368. nat/profiler/decorators/function_tracking.py +1 -1
  369. nat/profiler/forecasting/config.py +1 -1
  370. nat/profiler/forecasting/model_trainer.py +1 -1
  371. nat/profiler/forecasting/models/__init__.py +1 -1
  372. nat/profiler/forecasting/models/forecasting_base_model.py +1 -1
  373. nat/profiler/forecasting/models/linear_model.py +1 -1
  374. nat/profiler/forecasting/models/random_forest_regressor.py +1 -1
  375. nat/profiler/inference_metrics_model.py +1 -1
  376. nat/profiler/inference_optimization/bottleneck_analysis/nested_stack_analysis.py +1 -1
  377. nat/profiler/inference_optimization/bottleneck_analysis/simple_stack_analysis.py +1 -1
  378. nat/profiler/inference_optimization/data_models.py +1 -1
  379. nat/profiler/inference_optimization/experimental/concurrency_spike_analysis.py +1 -1
  380. nat/profiler/inference_optimization/experimental/prefix_span_analysis.py +1 -1
  381. nat/profiler/inference_optimization/llm_metrics.py +1 -1
  382. nat/profiler/inference_optimization/prompt_caching.py +1 -1
  383. nat/profiler/inference_optimization/token_uniqueness.py +1 -1
  384. nat/profiler/inference_optimization/workflow_runtimes.py +1 -1
  385. nat/profiler/intermediate_property_adapter.py +1 -1
  386. nat/profiler/parameter_optimization/optimizable_utils.py +1 -1
  387. nat/profiler/parameter_optimization/optimizer_runtime.py +1 -1
  388. nat/profiler/parameter_optimization/parameter_optimizer.py +1 -1
  389. nat/profiler/parameter_optimization/parameter_selection.py +1 -1
  390. nat/profiler/parameter_optimization/pareto_visualizer.py +1 -1
  391. nat/profiler/parameter_optimization/prompt_optimizer.py +1 -1
  392. nat/profiler/parameter_optimization/update_helpers.py +1 -1
  393. nat/profiler/profile_runner.py +1 -1
  394. nat/profiler/utils.py +1 -1
  395. nat/registry_handlers/local/local_handler.py +1 -1
  396. nat/registry_handlers/local/register_local.py +1 -1
  397. nat/registry_handlers/metadata_factory.py +1 -1
  398. nat/registry_handlers/package_utils.py +1 -1
  399. nat/registry_handlers/pypi/pypi_handler.py +1 -1
  400. nat/registry_handlers/pypi/register_pypi.py +1 -1
  401. nat/registry_handlers/register.py +1 -1
  402. nat/registry_handlers/registry_handler_base.py +1 -1
  403. nat/registry_handlers/rest/register_rest.py +1 -1
  404. nat/registry_handlers/rest/rest_handler.py +1 -1
  405. nat/registry_handlers/schemas/headers.py +1 -1
  406. nat/registry_handlers/schemas/package.py +1 -1
  407. nat/registry_handlers/schemas/publish.py +1 -1
  408. nat/registry_handlers/schemas/pull.py +1 -1
  409. nat/registry_handlers/schemas/remove.py +1 -1
  410. nat/registry_handlers/schemas/search.py +1 -1
  411. nat/registry_handlers/schemas/status.py +1 -1
  412. nat/retriever/interface.py +1 -1
  413. nat/retriever/milvus/__init__.py +1 -1
  414. nat/retriever/milvus/register.py +1 -1
  415. nat/retriever/milvus/retriever.py +1 -1
  416. nat/retriever/models.py +1 -1
  417. nat/retriever/nemo_retriever/__init__.py +1 -1
  418. nat/retriever/nemo_retriever/register.py +1 -1
  419. nat/retriever/nemo_retriever/retriever.py +5 -5
  420. nat/retriever/register.py +1 -1
  421. nat/runtime/__init__.py +1 -1
  422. nat/runtime/loader.py +10 -3
  423. nat/runtime/metrics.py +180 -0
  424. nat/runtime/runner.py +1 -5
  425. nat/runtime/session.py +451 -32
  426. nat/runtime/user_metadata.py +1 -1
  427. nat/settings/global_settings.py +1 -1
  428. nat/tool/chat_completion.py +1 -1
  429. nat/tool/code_execution/README.md +1 -1
  430. nat/tool/code_execution/code_sandbox.py +1 -1
  431. nat/tool/code_execution/local_sandbox/Dockerfile.sandbox +1 -1
  432. nat/tool/code_execution/local_sandbox/__init__.py +1 -1
  433. nat/tool/code_execution/local_sandbox/local_sandbox_server.py +1 -1
  434. nat/tool/code_execution/local_sandbox/start_local_sandbox.sh +1 -1
  435. nat/tool/code_execution/register.py +1 -1
  436. nat/tool/code_execution/utils.py +1 -1
  437. nat/tool/datetime_tools.py +1 -1
  438. nat/tool/document_search.py +1 -1
  439. nat/tool/github_tools.py +1 -1
  440. nat/tool/memory_tools/add_memory_tool.py +1 -1
  441. nat/tool/memory_tools/delete_memory_tool.py +1 -1
  442. nat/tool/memory_tools/get_memory_tool.py +1 -1
  443. nat/tool/nvidia_rag.py +2 -2
  444. nat/tool/register.py +1 -1
  445. nat/tool/retriever.py +1 -1
  446. nat/tool/server_tools.py +1 -1
  447. nat/utils/__init__.py +8 -5
  448. nat/utils/callable_utils.py +1 -1
  449. nat/utils/data_models/schema_validator.py +1 -1
  450. nat/utils/debugging_utils.py +1 -1
  451. nat/utils/decorators.py +1 -1
  452. nat/utils/dump_distro_mapping.py +1 -1
  453. nat/utils/exception_handlers/automatic_retries.py +3 -3
  454. nat/utils/exception_handlers/schemas.py +1 -1
  455. nat/utils/io/model_processing.py +1 -1
  456. nat/utils/io/supress_logs.py +33 -0
  457. nat/utils/io/yaml_tools.py +1 -1
  458. nat/utils/log_levels.py +1 -1
  459. nat/utils/log_utils.py +13 -1
  460. nat/utils/metadata_utils.py +1 -1
  461. nat/utils/optional_imports.py +1 -1
  462. nat/utils/producer_consumer_queue.py +1 -1
  463. nat/utils/reactive/base/observable_base.py +1 -1
  464. nat/utils/reactive/base/observer_base.py +1 -1
  465. nat/utils/reactive/base/subject_base.py +1 -1
  466. nat/utils/reactive/observable.py +1 -1
  467. nat/utils/reactive/observer.py +1 -1
  468. nat/utils/reactive/subject.py +1 -1
  469. nat/utils/reactive/subscription.py +1 -1
  470. nat/utils/responses_api.py +1 -1
  471. nat/utils/settings/global_settings.py +1 -1
  472. nat/utils/string_utils.py +1 -1
  473. nat/utils/type_converter.py +18 -5
  474. nat/utils/type_utils.py +1 -1
  475. nat/utils/url_utils.py +1 -1
  476. {nvidia_nat-1.4.0a20251120.dist-info → nvidia_nat-1.4.0a20260113.dist-info}/METADATA +39 -14
  477. nvidia_nat-1.4.0a20260113.dist-info/RECORD +547 -0
  478. nvidia_nat-1.4.0a20260113.dist-info/entry_points.txt +38 -0
  479. nat/cli/commands/mcp/mcp.py +0 -986
  480. nat/front_ends/mcp/introspection_token_verifier.py +0 -73
  481. nat/front_ends/mcp/mcp_front_end_config.py +0 -109
  482. nat/front_ends/mcp/mcp_front_end_plugin.py +0 -155
  483. nat/front_ends/mcp/mcp_front_end_plugin_worker.py +0 -388
  484. nat/front_ends/mcp/memory_profiler.py +0 -320
  485. nat/front_ends/mcp/register.py +0 -27
  486. nat/front_ends/mcp/tool_converter.py +0 -321
  487. nvidia_nat-1.4.0a20251120.dist-info/RECORD +0 -488
  488. nvidia_nat-1.4.0a20251120.dist-info/entry_points.txt +0 -23
  489. {nvidia_nat-1.4.0a20251120.dist-info → nvidia_nat-1.4.0a20260113.dist-info}/WHEEL +0 -0
  490. {nvidia_nat-1.4.0a20251120.dist-info → nvidia_nat-1.4.0a20260113.dist-info}/licenses/LICENSE-3rd-party.txt +0 -0
  491. {nvidia_nat-1.4.0a20251120.dist-info → nvidia_nat-1.4.0a20260113.dist-info}/licenses/LICENSE.md +0 -0
  492. {nvidia_nat-1.4.0a20251120.dist-info → nvidia_nat-1.4.0a20260113.dist-info}/top_level.txt +0 -0
@@ -1,4 +1,4 @@
1
- # SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
1
+ # SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
2
2
  # SPDX-License-Identifier: Apache-2.0
3
3
  #
4
4
  # Licensed under the Apache License, Version 2.0 (the "License");
@@ -1,4 +1,4 @@
1
- # SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
1
+ # SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
2
2
  # SPDX-License-Identifier: Apache-2.0
3
3
  #
4
4
  # Licensed under the Apache License, Version 2.0 (the "License");
@@ -1,4 +1,4 @@
1
- # SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
1
+ # SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
2
2
  # SPDX-License-Identifier: Apache-2.0
3
3
  #
4
4
  # Licensed under the Apache License, Version 2.0 (the "License");
@@ -1,4 +1,4 @@
1
- # SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
1
+ # SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
2
2
  # SPDX-License-Identifier: Apache-2.0
3
3
  #
4
4
  # Licensed under the Apache License, Version 2.0 (the "License");
@@ -1,4 +1,4 @@
1
- # SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
1
+ # SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
2
2
  # SPDX-License-Identifier: Apache-2.0
3
3
  #
4
4
  # Licensed under the Apache License, Version 2.0 (the "License");
nat/cli/commands/start.py CHANGED
@@ -1,4 +1,4 @@
1
- # SPDX-FileCopyrightText: Copyright (c) 2024-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
1
+ # SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
2
2
  # SPDX-License-Identifier: Apache-2.0
3
3
  #
4
4
  # Licensed under the Apache License, Version 2.0 (the "License");
@@ -1,4 +1,4 @@
1
- # SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
1
+ # SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
2
2
  # SPDX-License-Identifier: Apache-2.0
3
3
  #
4
4
  # Licensed under the Apache License, Version 2.0 (the "License");
@@ -1,4 +1,4 @@
1
- # SPDX-FileCopyrightText: Copyright (c) 2024-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
1
+ # SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
2
2
  # SPDX-License-Identifier: Apache-2.0
3
3
  #
4
4
  # Licensed under the Apache License, Version 2.0 (the "License");
@@ -1,4 +1,4 @@
1
- # SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
1
+ # SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
2
2
  # SPDX-License-Identifier: Apache-2.0
3
3
  #
4
4
  # Licensed under the Apache License, Version 2.0 (the "License");
@@ -1,4 +1,4 @@
1
- # SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
1
+ # SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
2
2
  # SPDX-License-Identifier: Apache-2.0
3
3
  #
4
4
  # Licensed under the Apache License, Version 2.0 (the "License");
@@ -1,4 +1,4 @@
1
- # SPDX-FileCopyrightText: Copyright (c) 2024-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
1
+ # SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
2
2
  # SPDX-License-Identifier: Apache-2.0
3
3
  #
4
4
  # Licensed under the Apache License, Version 2.0 (the "License");
@@ -19,6 +19,7 @@ import shutil
19
19
  import subprocess
20
20
  from pathlib import Path
21
21
  from urllib.parse import urlparse
22
+ from urllib.request import url2pathname
22
23
 
23
24
  import click
24
25
  from jinja2 import Environment
@@ -151,7 +152,7 @@ def find_package_root(package_name: str) -> Path | None:
151
152
  logger.error("Invalid URL scheme in direct_url.json: %s", url)
152
153
  return None
153
154
 
154
- package_root = Path(parsed_url.path).resolve()
155
+ package_root = Path(url2pathname(parsed_url.path)).resolve()
155
156
 
156
157
  # Ensure the path exists and is within an allowed base directory
157
158
  if not package_root.exists() or not package_root.is_dir():
nat/cli/entrypoint.py CHANGED
@@ -1,4 +1,4 @@
1
- # SPDX-FileCopyrightText: Copyright (c) 2024-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
1
+ # SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
2
2
  # SPDX-License-Identifier: Apache-2.0
3
3
  #
4
4
  # Licensed under the Apache License, Version 2.0 (the "License");
@@ -13,7 +13,7 @@
13
13
  # See the License for the specific language governing permissions and
14
14
  # limitations under the License.
15
15
 
16
- # SPDX-FileCopyrightText: Copyright (c) 2024-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
16
+ # SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
17
17
  # SPDX-License-Identifier: LicenseRef-NvidiaProprietary
18
18
  #
19
19
  # NVIDIA CORPORATION, its affiliates and licensors retain all intellectual
@@ -28,39 +28,25 @@ import sys
28
28
  import time
29
29
 
30
30
  import click
31
- import nest_asyncio
31
+ import nest_asyncio2
32
32
  from dotenv import load_dotenv
33
33
 
34
34
  from nat.utils.log_levels import LOG_LEVELS
35
+ from nat.utils.log_utils import setup_logging as log_utils_setup_logging
35
36
 
36
- from .commands.configure.configure import configure_command
37
- from .commands.evaluate import eval_command
38
- from .commands.info.info import info_command
39
- from .commands.mcp.mcp import mcp_command
40
- from .commands.object_store.object_store import object_store_command
41
- from .commands.optimize import optimizer_command
42
- from .commands.registry.registry import registry_command
43
- from .commands.sizing.sizing import sizing
44
- from .commands.start import start_command
45
- from .commands.uninstall import uninstall_command
46
- from .commands.validate import validate_command
47
- from .commands.workflow.workflow import workflow_command
37
+ from .plugin_loader import discover_and_load_cli_plugins
48
38
 
49
39
  # Load environment variables from .env file, if it exists
50
40
  load_dotenv()
51
41
 
52
42
  # Apply at the beginning of the file to avoid issues with asyncio
53
- nest_asyncio.apply()
43
+ nest_asyncio2.apply()
54
44
 
55
45
 
56
46
  def setup_logging(log_level: str):
57
47
  """Configure logging with the specified level"""
58
48
  numeric_level = LOG_LEVELS.get(log_level.upper(), logging.INFO)
59
- logging.basicConfig(
60
- level=numeric_level,
61
- format="%(asctime)s - %(levelname)-8s - %(name)s:%(lineno)d - %(message)s",
62
- datefmt="%Y-%m-%d %H:%M:%S",
63
- )
49
+ log_utils_setup_logging(numeric_level)
64
50
  return numeric_level
65
51
 
66
52
 
@@ -102,22 +88,14 @@ def cli(ctx: click.Context, log_level: str):
102
88
  ctx_dict["log_level"] = log_level
103
89
 
104
90
 
105
- cli.add_command(configure_command, name="configure")
106
- cli.add_command(eval_command, name="eval")
107
- cli.add_command(info_command, name="info")
108
- cli.add_command(registry_command, name="registry")
109
- cli.add_command(start_command, name="start")
110
- cli.add_command(uninstall_command, name="uninstall")
111
- cli.add_command(validate_command, name="validate")
112
- cli.add_command(workflow_command, name="workflow")
113
- cli.add_command(sizing, name="sizing")
114
- cli.add_command(optimizer_command, name="optimize")
115
- cli.add_command(object_store_command, name="object-store")
116
- cli.add_command(mcp_command, name="mcp")
117
-
118
- # Aliases
119
- cli.add_command(start_command.get_command(None, "console"), name="run") # type: ignore
120
- cli.add_command(start_command.get_command(None, "fastapi"), name="serve") # type: ignore
91
+ # Discover and load ALL CLI commands (core + plugins) via entry points
92
+ discover_and_load_cli_plugins(cli)
93
+
94
+ # Aliases - need to get start_command from the loaded commands
95
+ start_cmd = cli.commands.get("start")
96
+ if start_cmd and hasattr(start_cmd, "get_command"):
97
+ cli.add_command(start_cmd.get_command(None, "console"), name="run") # type: ignore
98
+ cli.add_command(start_cmd.get_command(None, "fastapi"), name="serve") # type: ignore
121
99
 
122
100
 
123
101
  @cli.result_callback()
nat/cli/main.py CHANGED
@@ -1,4 +1,4 @@
1
- # SPDX-FileCopyrightText: Copyright (c) 2024-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
1
+ # SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
2
2
  # SPDX-License-Identifier: Apache-2.0
3
3
  #
4
4
  # Licensed under the Apache License, Version 2.0 (the "License");
@@ -13,7 +13,7 @@
13
13
  # See the License for the specific language governing permissions and
14
14
  # limitations under the License.
15
15
 
16
- # SPDX-FileCopyrightText: Copyright (c) 2024-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
16
+ # SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
17
17
  # SPDX-License-Identifier: LicenseRef-NvidiaProprietary
18
18
  #
19
19
  # NVIDIA CORPORATION, its affiliates and licensors retain all intellectual
@@ -0,0 +1,69 @@
1
+ # SPDX-FileCopyrightText: Copyright (c) 2026, 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
+ """CLI plugin discovery system for loading plugin-specific commands."""
16
+
17
+ import logging
18
+ from importlib.metadata import entry_points
19
+
20
+ import click
21
+
22
+ logger = logging.getLogger(__name__)
23
+
24
+
25
+ def discover_and_load_cli_plugins(cli_group: click.Group) -> None:
26
+ """Discover and load CLI command plugins from installed packages.
27
+
28
+ This function uses Python entry points to discover CLI commands provided by
29
+ plugin packages. Plugins register their commands under the 'nat.cli' entry
30
+ point group in their pyproject.toml.
31
+
32
+ The function handles import errors gracefully - if a plugin cannot be loaded
33
+ (e.g., due to missing dependencies), it logs a debug message but continues
34
+ loading other plugins.
35
+
36
+ Args:
37
+ cli_group: The Click group to add discovered commands to
38
+
39
+ Example plugin registration in pyproject.toml:
40
+ [project.entry-points.'nat.cli']
41
+ mcp = "nat.plugins.mcp.cli.commands:mcp_command"
42
+ """
43
+ discovered_eps = entry_points(group='nat.cli')
44
+
45
+ for ep in discovered_eps:
46
+ try:
47
+ # Load the command from the entry point
48
+ command = ep.load()
49
+
50
+ # Verify it's a Click command or group
51
+ if not isinstance(command, click.Command | click.Group):
52
+ logger.warning("CLI plugin '%s' from '%s' is not a Click command/group, skipping", ep.name, ep.value)
53
+ continue
54
+
55
+ # Add the command to the CLI group
56
+ cli_group.add_command(command, name=ep.name)
57
+ logger.debug("Loaded CLI plugin: %s from %s", ep.name, ep.value)
58
+
59
+ except ImportError as e:
60
+ # Plugin package not installed or missing dependencies - this is expected
61
+ logger.debug(
62
+ "Could not load CLI plugin '%s' from '%s': %s. "
63
+ "This is expected if the plugin package is not installed.",
64
+ ep.name,
65
+ ep.value,
66
+ e)
67
+ except Exception as e: # noqa: BLE001
68
+ # Unexpected error - log as warning but continue
69
+ logger.warning("Error loading CLI plugin '%s' from '%s': %s", ep.name, ep.value, e, exc_info=True)
@@ -1,4 +1,4 @@
1
- # SPDX-FileCopyrightText: Copyright (c) 2024-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
1
+ # SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
2
2
  # SPDX-License-Identifier: Apache-2.0
3
3
  #
4
4
  # Licensed under the Apache License, Version 2.0 (the "License");
@@ -15,6 +15,8 @@
15
15
 
16
16
  from contextlib import asynccontextmanager
17
17
 
18
+ from pydantic import BaseModel
19
+
18
20
  from nat.builder.framework_enum import LLMFrameworkEnum
19
21
  from nat.cli.type_registry import AuthProviderBuildCallableT
20
22
  from nat.cli.type_registry import AuthProviderRegisteredCallableT
@@ -55,13 +57,22 @@ from nat.cli.type_registry import TeleExporterRegisteredCallableT
55
57
  from nat.cli.type_registry import TelemetryExporterBuildCallableT
56
58
  from nat.cli.type_registry import TelemetryExporterConfigT
57
59
  from nat.cli.type_registry import ToolWrapperBuildCallableT
60
+ from nat.cli.type_registry import TrainerAdapterBuildCallableT
61
+ from nat.cli.type_registry import TrainerAdapterRegisteredCallableT
62
+ from nat.cli.type_registry import TrainerBuildCallableT
63
+ from nat.cli.type_registry import TrainerRegisteredCallableT
64
+ from nat.cli.type_registry import TrajectoryBuilderBuildCallableT
65
+ from nat.cli.type_registry import TrajectoryBuilderRegisteredCallableT
58
66
  from nat.cli.type_registry import TTCStrategyBuildCallableT
59
- from nat.cli.type_registry import TTCStrategyRegisterCallableT
67
+ from nat.cli.type_registry import TTCStrategyRegisteredCallableT
60
68
  from nat.data_models.authentication import AuthProviderBaseConfigT
61
69
  from nat.data_models.component import ComponentEnum
62
70
  from nat.data_models.discovery_metadata import DiscoveryMetadata
63
71
  from nat.data_models.embedder import EmbedderBaseConfigT
64
72
  from nat.data_models.evaluator import EvaluatorBaseConfigT
73
+ from nat.data_models.finetuning import TrainerAdapterConfigT
74
+ from nat.data_models.finetuning import TrainerConfigT
75
+ from nat.data_models.finetuning import TrajectoryBuilderConfigT
65
76
  from nat.data_models.front_end import FrontEndConfigT
66
77
  from nat.data_models.function import FunctionConfigT
67
78
  from nat.data_models.function import FunctionGroupConfigT
@@ -71,6 +82,8 @@ from nat.data_models.middleware import MiddlewareBaseConfigT
71
82
  from nat.data_models.object_store import ObjectStoreBaseConfigT
72
83
  from nat.data_models.registry_handler import RegistryHandlerBaseConfigT
73
84
  from nat.data_models.retriever import RetrieverBaseConfigT
85
+ from nat.data_models.ttc_strategy import TTCStrategyBaseConfigT
86
+ from nat.utils.type_utils import DecomposedType
74
87
 
75
88
 
76
89
  def register_telemetry_exporter(config_type: type[TelemetryExporterConfigT]):
@@ -184,6 +197,64 @@ def register_function(config_type: type[FunctionConfigT],
184
197
  return register_function_inner
185
198
 
186
199
 
200
+ def register_per_user_function(config_type: type[FunctionConfigT],
201
+ input_type: type | type[BaseModel],
202
+ single_output_type: type | type[BaseModel] | None = None,
203
+ streaming_output_type: type | type[BaseModel] | None = None,
204
+ framework_wrappers: list[LLMFrameworkEnum | str] | None = None):
205
+ """
206
+ Register a per-user function with optional framework_wrappers for automatic profiler hooking.
207
+
208
+ The per-user function is instantiated lazily on user's first invocation. Each user will have a separate instance of
209
+ the function. Schemas must be provided to enable OpenAPI documentation generation without a concrete instance.
210
+
211
+ Args:
212
+ config_type: The function configuration type
213
+ input_type: The input type for the function (can be a type or a Pydantic model)
214
+ single_output_type: The single output type for the function (can be a type or a Pydantic model)
215
+ streaming_output_type: The streaming output type for the function (can be a type or a Pydantic model)
216
+ framework_wrappers: Optional list of framework wrappers for automatic profiler hooking
217
+ """
218
+
219
+ def register_per_user_function_inner(
220
+ fn: FunctionBuildCallableT[FunctionConfigT]) -> FunctionRegisteredCallableT[FunctionConfigT]:
221
+ from .type_registry import GlobalTypeRegistry
222
+ from .type_registry import RegisteredFunctionInfo
223
+
224
+ if not input_type:
225
+ raise ValueError("input_type must be provided to register a per-user function")
226
+
227
+ context_manager_fn = asynccontextmanager(fn)
228
+
229
+ framework_wrappers_list = list(framework_wrappers or [])
230
+
231
+ discovery_metadata = DiscoveryMetadata.from_config_type(config_type=config_type,
232
+ component_type=ComponentEnum.FUNCTION)
233
+
234
+ # Convert types to Pydantic models if they are not already
235
+ input_schema = DecomposedType(input_type).get_pydantic_schema()
236
+ single_output_schema = DecomposedType(single_output_type).get_pydantic_schema() if single_output_type else None
237
+ streaming_output_schema = DecomposedType(
238
+ streaming_output_type).get_pydantic_schema() if streaming_output_type else None
239
+
240
+ GlobalTypeRegistry.get().register_function(
241
+ RegisteredFunctionInfo(
242
+ full_type=config_type.full_type,
243
+ config_type=config_type,
244
+ build_fn=context_manager_fn,
245
+ framework_wrappers=framework_wrappers_list,
246
+ discovery_metadata=discovery_metadata,
247
+ is_per_user=True,
248
+ per_user_function_input_schema=input_schema,
249
+ per_user_function_single_output_schema=single_output_schema,
250
+ per_user_function_streaming_output_schema=streaming_output_schema,
251
+ ))
252
+
253
+ return context_manager_fn
254
+
255
+ return register_per_user_function_inner
256
+
257
+
187
258
  def register_function_group(config_type: type[FunctionGroupConfigT],
188
259
  framework_wrappers: list[LLMFrameworkEnum | str] | None = None):
189
260
  """
@@ -218,6 +289,43 @@ def register_function_group(config_type: type[FunctionGroupConfigT],
218
289
  return register_function_group_inner
219
290
 
220
291
 
292
+ def register_per_user_function_group(config_type: type[FunctionGroupConfigT],
293
+ framework_wrappers: list[LLMFrameworkEnum | str] | None = None):
294
+ """
295
+ Register a per-user function group with optional framework_wrappers for automatic profiler hooking.
296
+
297
+ Per-user function groups are instantiated separately for each user, allowing for user-specific
298
+ shared state across multiple functions within the group.
299
+ """
300
+
301
+ def register_per_user_function_group_inner(
302
+ fn: FunctionGroupBuildCallableT[FunctionGroupConfigT]
303
+ ) -> FunctionGroupRegisteredCallableT[FunctionGroupConfigT]:
304
+ from .type_registry import GlobalTypeRegistry
305
+ from .type_registry import RegisteredFunctionGroupInfo
306
+
307
+ context_manager_fn = asynccontextmanager(fn)
308
+
309
+ framework_wrappers_list = list(framework_wrappers or [])
310
+
311
+ discovery_metadata = DiscoveryMetadata.from_config_type(config_type=config_type,
312
+ component_type=ComponentEnum.FUNCTION_GROUP)
313
+
314
+ GlobalTypeRegistry.get().register_function_group(
315
+ RegisteredFunctionGroupInfo(
316
+ full_type=config_type.full_type,
317
+ config_type=config_type,
318
+ build_fn=context_manager_fn,
319
+ framework_wrappers=framework_wrappers_list,
320
+ discovery_metadata=discovery_metadata,
321
+ is_per_user=True,
322
+ ))
323
+
324
+ return context_manager_fn
325
+
326
+ return register_per_user_function_group_inner
327
+
328
+
221
329
  def register_middleware(config_type: type[MiddlewareBaseConfigT]):
222
330
  """
223
331
  Register a middleware component.
@@ -455,11 +563,11 @@ def register_object_store(config_type: type[ObjectStoreBaseConfigT]):
455
563
  return register_kv_store_inner
456
564
 
457
565
 
458
- def register_ttc_strategy(config_type: type[TTCStrategyRegisterCallableT]):
566
+ def register_ttc_strategy(config_type: type[TTCStrategyBaseConfigT]):
459
567
 
460
568
  def register_ttc_strategy_inner(
461
- fn: TTCStrategyBuildCallableT[TTCStrategyRegisterCallableT]
462
- ) -> TTCStrategyRegisterCallableT[TTCStrategyRegisterCallableT]:
569
+ fn: TTCStrategyBuildCallableT[TTCStrategyBaseConfigT]
570
+ ) -> TTCStrategyRegisteredCallableT[TTCStrategyBaseConfigT]:
463
571
  from .type_registry import GlobalTypeRegistry
464
572
  from .type_registry import RegisteredTTCStrategyInfo
465
573
 
@@ -479,6 +587,76 @@ def register_ttc_strategy(config_type: type[TTCStrategyRegisterCallableT]):
479
587
  return register_ttc_strategy_inner
480
588
 
481
589
 
590
+ def register_trainer(config_type: type[TrainerConfigT]):
591
+
592
+ def register_trainer_inner(fn: TrainerBuildCallableT[TrainerConfigT]) -> TrainerRegisteredCallableT[TrainerConfigT]:
593
+ from .type_registry import GlobalTypeRegistry
594
+ from .type_registry import RegisteredTrainerInfo
595
+
596
+ context_manager_fn = asynccontextmanager(fn)
597
+
598
+ discovery_metadata = DiscoveryMetadata.from_config_type(config_type=config_type,
599
+ component_type=ComponentEnum.TRAINER)
600
+
601
+ GlobalTypeRegistry.get().register_trainer(
602
+ RegisteredTrainerInfo(full_type=config_type.full_type,
603
+ config_type=config_type,
604
+ build_fn=context_manager_fn,
605
+ discovery_metadata=discovery_metadata))
606
+
607
+ return context_manager_fn
608
+
609
+ return register_trainer_inner
610
+
611
+
612
+ def register_trainer_adapter(config_type: type[TrainerAdapterConfigT]):
613
+
614
+ def register_trainer_adapter_inner(
615
+ fn: TrainerAdapterBuildCallableT[TrainerAdapterConfigT]
616
+ ) -> TrainerAdapterRegisteredCallableT[TrainerAdapterConfigT]:
617
+ from .type_registry import GlobalTypeRegistry
618
+ from .type_registry import RegisteredTrainerAdapterInfo
619
+
620
+ context_manager_fn = asynccontextmanager(fn)
621
+
622
+ discovery_metadata = DiscoveryMetadata.from_config_type(config_type=config_type,
623
+ component_type=ComponentEnum.TRAINER_ADAPTER)
624
+
625
+ GlobalTypeRegistry.get().register_trainer_adapter(
626
+ RegisteredTrainerAdapterInfo(full_type=config_type.full_type,
627
+ config_type=config_type,
628
+ build_fn=context_manager_fn,
629
+ discovery_metadata=discovery_metadata))
630
+
631
+ return context_manager_fn
632
+
633
+ return register_trainer_adapter_inner
634
+
635
+
636
+ def register_trajectory_builder(config_type: type[TrajectoryBuilderConfigT]):
637
+
638
+ def register_trajectory_builder_inner(
639
+ fn: TrajectoryBuilderBuildCallableT[TrajectoryBuilderConfigT]
640
+ ) -> TrajectoryBuilderRegisteredCallableT[TrajectoryBuilderConfigT]:
641
+ from .type_registry import GlobalTypeRegistry
642
+ from .type_registry import RegisteredTrajectoryBuilderInfo
643
+
644
+ context_manager_fn = asynccontextmanager(fn)
645
+
646
+ discovery_metadata = DiscoveryMetadata.from_config_type(config_type=config_type,
647
+ component_type=ComponentEnum.TRAJECTORY_BUILDER)
648
+
649
+ GlobalTypeRegistry.get().register_trajectory_builder(
650
+ RegisteredTrajectoryBuilderInfo(full_type=config_type.full_type,
651
+ config_type=config_type,
652
+ build_fn=context_manager_fn,
653
+ discovery_metadata=discovery_metadata))
654
+
655
+ return context_manager_fn
656
+
657
+ return register_trajectory_builder_inner
658
+
659
+
482
660
  def register_retriever_provider(config_type: type[RetrieverBaseConfigT]):
483
661
 
484
662
  def register_retriever_provider_inner(