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
@@ -0,0 +1,843 @@
1
+ # SPDX-FileCopyrightText: Copyright (c) 2025-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
+
16
+ import asyncio
17
+ import logging
18
+ import typing
19
+ from collections.abc import Sequence
20
+ from contextlib import AbstractAsyncContextManager
21
+ from contextlib import AsyncExitStack
22
+ from contextlib import asynccontextmanager
23
+ from typing import cast
24
+
25
+ from nat.authentication.interfaces import AuthProviderBase
26
+ from nat.builder.builder import Builder
27
+ from nat.builder.builder import UserManagerHolder
28
+ from nat.builder.child_builder import ChildBuilder
29
+ from nat.builder.component_utils import WORKFLOW_COMPONENT_NAME
30
+ from nat.builder.component_utils import build_dependency_sequence
31
+ from nat.builder.framework_enum import LLMFrameworkEnum
32
+ from nat.builder.function import Function
33
+ from nat.builder.function import FunctionGroup
34
+ from nat.builder.sync_builder import SyncBuilder
35
+ from nat.builder.workflow import Workflow
36
+ from nat.builder.workflow_builder import ConfiguredFunction
37
+ from nat.builder.workflow_builder import ConfiguredFunctionGroup
38
+ from nat.builder.workflow_builder import WorkflowBuilder
39
+ from nat.builder.workflow_builder import _build_function_group_impl
40
+ from nat.builder.workflow_builder import _build_function_impl
41
+ from nat.builder.workflow_builder import _log_build_failure
42
+ from nat.cli.type_registry import GlobalTypeRegistry
43
+ from nat.cli.type_registry import TypeRegistry
44
+ from nat.data_models.authentication import AuthProviderBaseConfig
45
+ from nat.data_models.component import ComponentGroup
46
+ from nat.data_models.component_ref import FunctionGroupRef
47
+ from nat.data_models.component_ref import FunctionRef
48
+ from nat.data_models.component_ref import MiddlewareRef
49
+ from nat.data_models.component_ref import RetrieverRef
50
+ from nat.data_models.component_ref import TrainerAdapterRef
51
+ from nat.data_models.component_ref import TrainerRef
52
+ from nat.data_models.component_ref import TrajectoryBuilderRef
53
+ from nat.data_models.component_ref import TTCStrategyRef
54
+ from nat.data_models.config import Config
55
+ from nat.data_models.embedder import EmbedderBaseConfig
56
+ from nat.data_models.finetuning import TrainerAdapterConfig
57
+ from nat.data_models.finetuning import TrainerConfig
58
+ from nat.data_models.finetuning import TrajectoryBuilderConfig
59
+ from nat.data_models.function import FunctionBaseConfig
60
+ from nat.data_models.function import FunctionGroupBaseConfig
61
+ from nat.data_models.function_dependencies import FunctionDependencies
62
+ from nat.data_models.llm import LLMBaseConfig
63
+ from nat.data_models.memory import MemoryBaseConfig
64
+ from nat.data_models.middleware import MiddlewareBaseConfig
65
+ from nat.data_models.object_store import ObjectStoreBaseConfig
66
+ from nat.data_models.retriever import RetrieverBaseConfig
67
+ from nat.data_models.ttc_strategy import TTCStrategyBaseConfig
68
+ from nat.experimental.decorators.experimental_warning_decorator import experimental
69
+ from nat.experimental.test_time_compute.models.stage_enums import PipelineTypeEnum
70
+ from nat.experimental.test_time_compute.models.stage_enums import StageTypeEnum
71
+ from nat.experimental.test_time_compute.models.strategy_base import StrategyBase
72
+ from nat.finetuning.interfaces.finetuning_runner import Trainer
73
+ from nat.finetuning.interfaces.trainer_adapter import TrainerAdapter
74
+ from nat.finetuning.interfaces.trajectory_builder import TrajectoryBuilder
75
+ from nat.memory.interfaces import MemoryEditor
76
+ from nat.middleware.function_middleware import FunctionMiddleware
77
+ from nat.middleware.middleware import Middleware
78
+ from nat.object_store.interfaces import ObjectStore
79
+ from nat.retriever.interface import Retriever
80
+ from nat.utils.type_utils import override
81
+
82
+ logger = logging.getLogger(__name__)
83
+
84
+
85
+ class PerUserWorkflowBuilder(Builder, AbstractAsyncContextManager):
86
+ """
87
+ Builder for per-user components that are lazily instantiated.
88
+
89
+ This builder is created per-user and only builds functions/function_groups
90
+ that are marked as per-user. It delegates to a shared WorkflowBuilder for
91
+ all shared components (LLMs, embedders, memory, etc.).
92
+
93
+ Lifecycle:
94
+ - Created when a user first makes a request
95
+ - Kept alive while the user is active
96
+ - Cleaned up after user inactivity timeout
97
+ """
98
+
99
+ def __init__(self, user_id: str, shared_builder: WorkflowBuilder, registry: TypeRegistry | None = None):
100
+
101
+ self._user_id = user_id
102
+ self._shared_builder = shared_builder
103
+ self._workflow: ConfiguredFunction | None = None
104
+
105
+ if registry is None:
106
+ registry = GlobalTypeRegistry.get()
107
+ self._registry = registry
108
+
109
+ self._per_user_functions: dict[str, ConfiguredFunction] = {}
110
+ self._per_user_function_groups: dict[str, ConfiguredFunctionGroup] = {}
111
+
112
+ self._exit_stack: AsyncExitStack | None = None
113
+
114
+ self.per_user_function_dependencies: dict[str, FunctionDependencies] = {}
115
+ self.per_user_function_group_dependencies: dict[str, FunctionDependencies] = {}
116
+
117
+ # Copy the completed and remaining components from the shared builder
118
+ self.completed_components: list[tuple[str, str]] = shared_builder.completed_components.copy()
119
+ self.remaining_components: list[tuple[str, str]] = shared_builder.remaining_components.copy()
120
+
121
+ async def __aenter__(self):
122
+
123
+ self._exit_stack = AsyncExitStack()
124
+ return self
125
+
126
+ async def __aexit__(self, *exc_details):
127
+
128
+ assert self._exit_stack is not None, "Exit stack not initialized"
129
+ await self._exit_stack.__aexit__(*exc_details)
130
+
131
+ def _get_exit_stack(self) -> AsyncExitStack:
132
+ if self._exit_stack is None:
133
+ raise ValueError(
134
+ "Exit stack not initialized. Did you forget to call `async with PerUserWorkflowBuilder() as builder`?")
135
+ return self._exit_stack
136
+
137
+ @override
138
+ @property
139
+ def sync_builder(self) -> SyncBuilder:
140
+ return SyncBuilder(self)
141
+
142
+ @property
143
+ def user_id(self) -> str:
144
+ return self._user_id
145
+
146
+ async def _resolve_middleware_instances_from_shared_builder(self,
147
+ middleware_names: Sequence[str],
148
+ component_type: str = "function"
149
+ ) -> list[FunctionMiddleware]:
150
+ """
151
+ Resolve middleware names to FunctionMiddleware instances from the shared builder.
152
+ """
153
+ middleware_instances: list[FunctionMiddleware] = []
154
+ for middleware_name in middleware_names:
155
+ middleware_obj = await self._shared_builder.get_middleware(middleware_name)
156
+ if not isinstance(middleware_obj, FunctionMiddleware):
157
+ raise TypeError(f"Middleware `{middleware_name}` is not a FunctionMiddleware and cannot be used "
158
+ f"with {component_type}s. "
159
+ f"Only FunctionMiddleware types support function-specific wrapping.")
160
+ middleware_instances.append(middleware_obj)
161
+ return middleware_instances
162
+
163
+ async def _build_per_user_function(self, name: str, config: FunctionBaseConfig) -> ConfiguredFunction:
164
+ registration = self._registry.get_function(type(config))
165
+
166
+ if not registration.is_per_user:
167
+ raise ValueError(f"Function `{name}` is not a per-user function")
168
+
169
+ inner_builder = ChildBuilder(self)
170
+
171
+ llms = {k: v.instance for k, v in self._shared_builder._llms.items()}
172
+ middleware_instances = await self._resolve_middleware_instances_from_shared_builder(
173
+ config.middleware, "function")
174
+ return await _build_function_impl(name=name,
175
+ config=config,
176
+ registry=self._registry,
177
+ exit_stack=self._get_exit_stack(),
178
+ inner_builder=inner_builder,
179
+ llms=llms,
180
+ dependencies=self.per_user_function_dependencies,
181
+ middleware_instances=middleware_instances)
182
+
183
+ async def _build_per_user_function_group(self, name: str,
184
+ config: FunctionGroupBaseConfig) -> ConfiguredFunctionGroup:
185
+ registration = self._registry.get_function_group(type(config))
186
+
187
+ if not registration.is_per_user:
188
+ raise ValueError(f"Function group `{name}` is not a per-user function group")
189
+
190
+ inner_builder = ChildBuilder(self)
191
+
192
+ llms = {k: v.instance for k, v in self._shared_builder._llms.items()}
193
+ middleware_instances = await self._resolve_middleware_instances_from_shared_builder(
194
+ config.middleware, "function group")
195
+
196
+ return await _build_function_group_impl(name=name,
197
+ config=config,
198
+ registry=self._registry,
199
+ exit_stack=self._get_exit_stack(),
200
+ inner_builder=inner_builder,
201
+ llms=llms,
202
+ dependencies=self.per_user_function_group_dependencies,
203
+ middleware_instances=middleware_instances)
204
+
205
+ @override
206
+ async def add_function(self, name: str | FunctionRef, config: FunctionBaseConfig) -> Function:
207
+ if isinstance(name, FunctionRef):
208
+ name = str(name)
209
+
210
+ if (name in self._per_user_functions) or \
211
+ (name in self._per_user_function_groups) or \
212
+ (name in self._shared_builder._functions) or \
213
+ (name in self._shared_builder._function_groups):
214
+ raise ValueError(f"Function `{name}` already exists in the list of functions or function groups")
215
+ if any(name.startswith(k + FunctionGroup.SEPARATOR) for k in self._per_user_function_groups.keys()) or \
216
+ any(name.startswith(k + FunctionGroup.SEPARATOR) for k in self._shared_builder._function_groups.keys()):
217
+ raise ValueError(f"A Function name starts with a Function Group name: `{name}`")
218
+
219
+ registration = self._registry.get_function(type(config))
220
+ if registration.is_per_user:
221
+ build_result = await self._build_per_user_function(name, config)
222
+ self._per_user_functions[name] = build_result
223
+
224
+ return build_result.instance
225
+
226
+ return await self._shared_builder.add_function(name, config)
227
+
228
+ def _check_backwards_compatibility_function_name(self, name: str) -> str:
229
+ if name in self._per_user_functions:
230
+ return name
231
+ new_name = name.replace(FunctionGroup.LEGACY_SEPARATOR, FunctionGroup.SEPARATOR)
232
+ if new_name in self._per_user_functions:
233
+ logger.warning(
234
+ f"Function `{name}` is deprecated and will be removed in a future release. Use `{new_name}` instead.")
235
+ return new_name
236
+ return name
237
+
238
+ @override
239
+ async def get_function(self, name: str | FunctionRef) -> Function:
240
+ if isinstance(name, FunctionRef):
241
+ name = str(name)
242
+
243
+ old_name = name
244
+ # Check for backwards compatibility with the old function name format
245
+ name = self._check_backwards_compatibility_function_name(name)
246
+
247
+ # Check per-user cache first
248
+ if name in self._per_user_functions:
249
+ return self._per_user_functions[name].instance
250
+
251
+ # Delegate to shared builder
252
+ return await self._shared_builder.get_function(old_name)
253
+
254
+ @override
255
+ def get_function_config(self, name: str | FunctionRef) -> FunctionBaseConfig:
256
+ if isinstance(name, FunctionRef):
257
+ name = str(name)
258
+
259
+ old_name = name
260
+ # Check for backwards compatibility with the old function name format
261
+ name = self._check_backwards_compatibility_function_name(name)
262
+
263
+ if name in self._per_user_functions:
264
+ return self._per_user_functions[name].config
265
+
266
+ return self._shared_builder.get_function_config(old_name)
267
+
268
+ @override
269
+ async def add_function_group(self, name: str | FunctionGroupRef, config: FunctionGroupBaseConfig) -> FunctionGroup:
270
+ if isinstance(name, FunctionGroupRef):
271
+ name = str(name)
272
+
273
+ if (name in self._per_user_function_groups) or \
274
+ (name in self._per_user_functions) or \
275
+ (name in self._shared_builder._function_groups) or \
276
+ (name in self._shared_builder._functions):
277
+ raise ValueError(f"Function group `{name}` already exists in the list of function groups or functions")
278
+ if any(k.startswith(name + FunctionGroup.SEPARATOR) for k in self._per_user_functions.keys()) or \
279
+ any(k.startswith(name + FunctionGroup.SEPARATOR) for k in self._shared_builder._functions.keys()):
280
+ raise ValueError(f"A Function name starts with a Function Group name: `{name}`")
281
+
282
+ registration = self._registry.get_function_group(type(config))
283
+ if registration.is_per_user:
284
+ # Build the per-user function group
285
+ build_result = await self._build_per_user_function_group(name=name, config=config)
286
+
287
+ self._per_user_function_groups[name] = build_result
288
+
289
+ # If the function group exposes functions, add them to the per-user function registry
290
+ included_functions = await build_result.instance.get_included_functions()
291
+ for k in included_functions:
292
+ if k in self._per_user_functions or k in self._shared_builder._functions:
293
+ raise ValueError(f"Exposed function `{k}` from group `{name}` conflicts with an existing function")
294
+ self._per_user_functions.update({
295
+ k: ConfiguredFunction(config=v.config, instance=v)
296
+ for k, v in included_functions.items()
297
+ })
298
+
299
+ return build_result.instance
300
+ else:
301
+ # Shared function group - delegate to shared builder
302
+ return await self._shared_builder.add_function_group(name, config)
303
+
304
+ @override
305
+ async def get_function_group(self, name: str | FunctionGroupRef) -> FunctionGroup:
306
+ if isinstance(name, FunctionGroupRef):
307
+ name = str(name)
308
+
309
+ # Check per-user function groups first
310
+ if name in self._per_user_function_groups:
311
+ return self._per_user_function_groups[name].instance
312
+
313
+ # Fall back to shared builder for shared function groups
314
+ return await self._shared_builder.get_function_group(name)
315
+
316
+ @override
317
+ def get_function_group_config(self, name: str | FunctionGroupRef) -> FunctionGroupBaseConfig:
318
+ if isinstance(name, FunctionGroupRef):
319
+ name = str(name)
320
+
321
+ # Check per-user function groups first
322
+ if name in self._per_user_function_groups:
323
+ return self._per_user_function_groups[name].config
324
+
325
+ # Fall back to shared builder
326
+ return self._shared_builder.get_function_group_config(name)
327
+
328
+ @override
329
+ async def set_workflow(self, config: FunctionBaseConfig) -> Function:
330
+ if self._workflow is not None:
331
+ logger.warning("Overwriting existing workflow")
332
+
333
+ build_result = await self._build_per_user_function(name=WORKFLOW_COMPONENT_NAME, config=config)
334
+
335
+ self._workflow = build_result
336
+
337
+ return build_result.instance
338
+
339
+ @override
340
+ def get_workflow(self) -> Function:
341
+ # If we have a per-user workflow, return it
342
+ if self._workflow is not None:
343
+ return self._workflow.instance
344
+
345
+ # Otherwise, delegate to shared builder
346
+ return self._shared_builder.get_workflow()
347
+
348
+ @override
349
+ def get_workflow_config(self) -> FunctionBaseConfig:
350
+ # If we have a per-user workflow config, return it
351
+ if self._workflow is not None:
352
+ return self._workflow.config
353
+
354
+ # Otherwise, delegate to shared builder
355
+ return self._shared_builder.get_workflow_config()
356
+
357
+ @override
358
+ def get_function_dependencies(self, fn_name: str | FunctionRef) -> FunctionDependencies:
359
+ if isinstance(fn_name, FunctionRef):
360
+ fn_name = str(fn_name)
361
+
362
+ old_fn_name = fn_name
363
+ # Check for backwards compatibility with the old function name format
364
+ fn_name = self._check_backwards_compatibility_function_name(fn_name)
365
+
366
+ if fn_name in self.per_user_function_dependencies:
367
+ return self.per_user_function_dependencies[fn_name]
368
+ return self._shared_builder.get_function_dependencies(old_fn_name)
369
+
370
+ @override
371
+ def get_function_group_dependencies(self, fn_name: str | FunctionGroupRef) -> FunctionDependencies:
372
+ if isinstance(fn_name, FunctionGroupRef):
373
+ fn_name = str(fn_name)
374
+
375
+ # Check per-user dependencies first
376
+ if fn_name in self.per_user_function_group_dependencies:
377
+ return self.per_user_function_group_dependencies[fn_name]
378
+
379
+ # Fall back to shared builder
380
+ return self._shared_builder.get_function_group_dependencies(fn_name)
381
+
382
+ @override
383
+ async def get_tools(self,
384
+ tool_names: Sequence[str | FunctionRef | FunctionGroupRef],
385
+ wrapper_type: LLMFrameworkEnum | str) -> list[typing.Any]:
386
+ unique = set(tool_names)
387
+ if len(unique) != len(tool_names):
388
+ raise ValueError("Tool names must be unique")
389
+
390
+ async def _get_tools(n: str | FunctionRef | FunctionGroupRef):
391
+ tools = []
392
+ is_function_group_ref = isinstance(n, FunctionGroupRef)
393
+ if isinstance(n, FunctionRef) or is_function_group_ref:
394
+ n = str(n)
395
+
396
+ # Check per-user function groups first
397
+ if n not in self._per_user_function_groups:
398
+ # Check shared function groups
399
+ if n not in self._shared_builder._function_groups:
400
+ # The passed tool name is probably a function, but first check if it's a function group
401
+ if is_function_group_ref:
402
+ raise ValueError(f"Function group `{n}` not found in the list of function groups")
403
+ tools.append(await self.get_tool(n, wrapper_type))
404
+ else:
405
+ # It's a shared function group
406
+ tool_wrapper_reg = self._registry.get_tool_wrapper(llm_framework=wrapper_type)
407
+ current_function_group = self._shared_builder._function_groups[n]
408
+ for fn_name, fn_instance in \
409
+ (await current_function_group.instance.get_accessible_functions()).items():
410
+ try:
411
+ tools.append(tool_wrapper_reg.build_fn(fn_name, fn_instance, self))
412
+ except Exception:
413
+ logger.error("Error fetching tool `%s`", fn_name, exc_info=True)
414
+ raise
415
+ else:
416
+ # It's a per-user function group
417
+ tool_wrapper_reg = self._registry.get_tool_wrapper(llm_framework=wrapper_type)
418
+ current_function_group = self._per_user_function_groups[n]
419
+ for fn_name, fn_instance in (await current_function_group.instance.get_accessible_functions()).items():
420
+ try:
421
+ tools.append(tool_wrapper_reg.build_fn(fn_name, fn_instance, self))
422
+ except Exception:
423
+ logger.error("Error fetching tool `%s`", fn_name, exc_info=True)
424
+ raise
425
+ return tools
426
+
427
+ tool_lists = await asyncio.gather(*[_get_tools(n) for n in tool_names])
428
+ # Flatten the list of lists into a single list
429
+ return [tool for sublist in tool_lists for tool in sublist]
430
+
431
+ @override
432
+ async def get_tool(self, fn_name: str | FunctionRef, wrapper_type: LLMFrameworkEnum | str) -> typing.Any:
433
+ if isinstance(fn_name, FunctionRef):
434
+ fn_name = str(fn_name)
435
+
436
+ old_fn_name = fn_name
437
+ # Check for backwards compatibility with the old function name format
438
+ fn_name = self._check_backwards_compatibility_function_name(fn_name)
439
+
440
+ if fn_name in self._per_user_functions:
441
+ fn = self._per_user_functions[fn_name]
442
+ try:
443
+ tool_wrapper_reg = self._registry.get_tool_wrapper(llm_framework=wrapper_type)
444
+ return tool_wrapper_reg.build_fn(fn_name, fn.instance, self)
445
+ except Exception as e:
446
+ logger.error("Error fetching tool `%s`: %s", fn_name, e)
447
+ raise
448
+ return await self._shared_builder.get_tool(old_fn_name, wrapper_type)
449
+
450
+ @override
451
+ async def add_llm(self, name: str, config: LLMBaseConfig) -> None:
452
+ return await self._shared_builder.add_llm(name, config)
453
+
454
+ @override
455
+ async def get_llm(self, llm_name: str, wrapper_type: LLMFrameworkEnum | str) -> typing.Any:
456
+ return await self._shared_builder.get_llm(llm_name, wrapper_type)
457
+
458
+ @override
459
+ def get_llm_config(self, llm_name: str) -> LLMBaseConfig:
460
+ return self._shared_builder.get_llm_config(llm_name)
461
+
462
+ @experimental(feature_name="Authentication")
463
+ @override
464
+ async def add_auth_provider(self, name: str, config: AuthProviderBaseConfig) -> AuthProviderBase:
465
+ return await self._shared_builder.add_auth_provider(name, config)
466
+
467
+ @override
468
+ async def get_auth_provider(self, auth_provider_name: str) -> AuthProviderBase:
469
+ return await self._shared_builder.get_auth_provider(auth_provider_name)
470
+
471
+ @override
472
+ async def add_embedder(self, name: str, config: EmbedderBaseConfig) -> None:
473
+ return await self._shared_builder.add_embedder(name, config)
474
+
475
+ @override
476
+ async def get_embedder(self, embedder_name: str, wrapper_type: LLMFrameworkEnum | str) -> typing.Any:
477
+ return await self._shared_builder.get_embedder(embedder_name, wrapper_type)
478
+
479
+ @override
480
+ def get_embedder_config(self, embedder_name: str) -> EmbedderBaseConfig:
481
+ return self._shared_builder.get_embedder_config(embedder_name)
482
+
483
+ @override
484
+ async def add_memory_client(self, name: str, config: MemoryBaseConfig) -> MemoryEditor:
485
+ return await self._shared_builder.add_memory_client(name, config)
486
+
487
+ @override
488
+ async def get_memory_client(self, memory_name: str) -> MemoryEditor:
489
+ return await self._shared_builder.get_memory_client(memory_name)
490
+
491
+ @override
492
+ def get_memory_client_config(self, memory_name: str) -> MemoryBaseConfig:
493
+ return self._shared_builder.get_memory_client_config(memory_name)
494
+
495
+ @override
496
+ async def add_object_store(self, name: str, config: ObjectStoreBaseConfig) -> ObjectStore:
497
+ return await self._shared_builder.add_object_store(name, config)
498
+
499
+ @override
500
+ async def get_object_store_client(self, object_store_name: str) -> ObjectStore:
501
+ return await self._shared_builder.get_object_store_client(object_store_name)
502
+
503
+ @override
504
+ def get_object_store_config(self, object_store_name: str) -> ObjectStoreBaseConfig:
505
+ return self._shared_builder.get_object_store_config(object_store_name)
506
+
507
+ @override
508
+ async def add_retriever(self, name: str | RetrieverRef, config: RetrieverBaseConfig) -> None:
509
+ return await self._shared_builder.add_retriever(name, config)
510
+
511
+ @override
512
+ async def get_retriever(self,
513
+ retriever_name: str | RetrieverRef,
514
+ wrapper_type: LLMFrameworkEnum | str | None = None) -> Retriever:
515
+ return await self._shared_builder.get_retriever(retriever_name, wrapper_type)
516
+
517
+ @override
518
+ async def get_retriever_config(self, retriever_name: str | RetrieverRef) -> RetrieverBaseConfig:
519
+ return await self._shared_builder.get_retriever_config(retriever_name)
520
+
521
+ @experimental(feature_name="TTC")
522
+ @override
523
+ async def add_ttc_strategy(self, name: str | TTCStrategyRef, config: TTCStrategyBaseConfig) -> None:
524
+ return await self._shared_builder.add_ttc_strategy(name, config)
525
+
526
+ @override
527
+ async def get_ttc_strategy(self,
528
+ strategy_name: str | TTCStrategyRef,
529
+ pipeline_type: PipelineTypeEnum,
530
+ stage_type: StageTypeEnum) -> StrategyBase:
531
+ return await self._shared_builder.get_ttc_strategy(strategy_name, pipeline_type, stage_type)
532
+
533
+ @override
534
+ async def get_ttc_strategy_config(self,
535
+ strategy_name: str | TTCStrategyRef,
536
+ pipeline_type: PipelineTypeEnum,
537
+ stage_type: StageTypeEnum) -> TTCStrategyBaseConfig:
538
+ return await self._shared_builder.get_ttc_strategy_config(strategy_name, pipeline_type, stage_type)
539
+
540
+ @override
541
+ def get_user_manager(self) -> UserManagerHolder:
542
+ return self._shared_builder.get_user_manager()
543
+
544
+ @override
545
+ async def add_middleware(self, name: str | MiddlewareRef, config: MiddlewareBaseConfig) -> Middleware:
546
+ return await self._shared_builder.add_middleware(name, config)
547
+
548
+ @override
549
+ async def get_middleware(self, middleware_name: str | MiddlewareRef) -> Middleware:
550
+ return await self._shared_builder.get_middleware(middleware_name)
551
+
552
+ @override
553
+ def get_middleware_config(self, middleware_name: str | MiddlewareRef) -> MiddlewareBaseConfig:
554
+ return self._shared_builder.get_middleware_config(middleware_name)
555
+
556
+ @experimental(feature_name="Finetuning")
557
+ @override
558
+ async def add_trainer(self, name: str | TrainerRef, config: TrainerConfig) -> Trainer:
559
+ return await self._shared_builder.add_trainer(name, config)
560
+
561
+ @experimental(feature_name="Finetuning")
562
+ @override
563
+ async def add_trainer_adapter(self, name: str | TrainerAdapterRef, config: TrainerAdapterConfig) -> TrainerAdapter:
564
+ return await self._shared_builder.add_trainer_adapter(name, config)
565
+
566
+ @experimental(feature_name="Finetuning")
567
+ @override
568
+ async def add_trajectory_builder(self, name: str | TrajectoryBuilderRef,
569
+ config: TrajectoryBuilderConfig) -> TrajectoryBuilder:
570
+ return await self._shared_builder.add_trajectory_builder(name, config)
571
+
572
+ @override
573
+ async def get_trainer(self,
574
+ trainer_name: str | TrainerRef,
575
+ trajectory_builder: TrajectoryBuilder,
576
+ trainer_adapter: TrainerAdapter) -> Trainer:
577
+ return await self._shared_builder.get_trainer(trainer_name, trajectory_builder, trainer_adapter)
578
+
579
+ @override
580
+ async def get_trainer_adapter(self, trainer_adapter_name: str | TrainerAdapterRef) -> TrainerAdapter:
581
+ return await self._shared_builder.get_trainer_adapter(trainer_adapter_name)
582
+
583
+ @override
584
+ async def get_trajectory_builder(self, trajectory_builder_name: str | TrajectoryBuilderRef) -> TrajectoryBuilder:
585
+ return await self._shared_builder.get_trajectory_builder(trajectory_builder_name)
586
+
587
+ @override
588
+ async def get_trainer_config(self, trainer_name: str | TrainerRef) -> TrainerConfig:
589
+ return await self._shared_builder.get_trainer_config(trainer_name)
590
+
591
+ @override
592
+ async def get_trainer_adapter_config(self, trainer_adapter_name: str | TrainerAdapterRef) -> TrainerAdapterConfig:
593
+ return await self._shared_builder.get_trainer_adapter_config(trainer_adapter_name)
594
+
595
+ @override
596
+ async def get_trajectory_builder_config(
597
+ self, trajectory_builder_name: str | TrajectoryBuilderRef) -> TrajectoryBuilderConfig:
598
+ return await self._shared_builder.get_trajectory_builder_config(trajectory_builder_name)
599
+
600
+ async def populate_builder(self, config: Config, skip_workflow: bool = False):
601
+ """
602
+ Populate the per-user builder with per-user components from config.
603
+
604
+ Only builds components that are marked as per-user.
605
+ Builds in dependency order to handle per-user functions depending on other per-user functions.
606
+
607
+ Args:
608
+ config: The full configuration object
609
+ skip_workflow: If True, skips the workflow instantiation step. Defaults to False.
610
+ Raises:
611
+ ValueError: If a per-user component has invalid dependencies
612
+ """
613
+ # Generate build sequence using the same dependency resolution as shared builder
614
+ build_sequence = build_dependency_sequence(config)
615
+
616
+ if not skip_workflow:
617
+ if (WORKFLOW_COMPONENT_NAME, "workflow") not in self.remaining_components:
618
+ self.remaining_components.append((WORKFLOW_COMPONENT_NAME, "workflow"))
619
+
620
+ # Filter to only per-user functions and function groups and build them in dependency order
621
+ for component_instance in build_sequence:
622
+ try:
623
+ if component_instance.component_group == ComponentGroup.FUNCTION_GROUPS:
624
+ config_obj = cast(FunctionGroupBaseConfig, component_instance.config)
625
+ registration = self._registry.get_function_group(type(config_obj))
626
+ if registration.is_per_user:
627
+ # Build the per-user function group
628
+ logger.debug(
629
+ f"Building per-user function group '{component_instance.name}' for user {self._user_id}")
630
+ await self.add_function_group(component_instance.name, config_obj)
631
+ self.remaining_components.remove(
632
+ (str(component_instance.name), component_instance.component_group.value))
633
+ self.completed_components.append(
634
+ (str(component_instance.name), component_instance.component_group.value))
635
+ else:
636
+ continue
637
+
638
+ elif component_instance.component_group == ComponentGroup.FUNCTIONS:
639
+ config_obj = cast(FunctionBaseConfig, component_instance.config)
640
+ registration = self._registry.get_function(type(config_obj))
641
+ if registration.is_per_user:
642
+ if not component_instance.is_root:
643
+ logger.debug(
644
+ f"Building per-user function '{component_instance.name}' for user {self._user_id}")
645
+ await self.add_function(component_instance.name, config_obj)
646
+ self.remaining_components.remove(
647
+ (str(component_instance.name), component_instance.component_group.value))
648
+ self.completed_components.append(
649
+ (str(component_instance.name), component_instance.component_group.value))
650
+ else:
651
+ continue
652
+
653
+ except Exception as e:
654
+ _log_build_failure(str(component_instance.name),
655
+ component_instance.component_group.value,
656
+ self.completed_components,
657
+ self.remaining_components,
658
+ e)
659
+ raise
660
+
661
+ if not skip_workflow:
662
+ try:
663
+ registration = self._registry.get_function(type(config.workflow))
664
+ if registration.is_per_user:
665
+ self.remaining_components.remove((WORKFLOW_COMPONENT_NAME, "workflow"))
666
+ await self.set_workflow(config.workflow)
667
+ self.completed_components.append((WORKFLOW_COMPONENT_NAME, "workflow"))
668
+ except Exception as e:
669
+ _log_build_failure(WORKFLOW_COMPONENT_NAME,
670
+ "workflow",
671
+ self.completed_components,
672
+ self.remaining_components,
673
+ e)
674
+ raise
675
+
676
+ async def build(self, entry_function: str | None = None) -> Workflow:
677
+ """
678
+ Creates a workflow instance for this specific user.
679
+
680
+ Combines per-user functions with shared components from the shared builder.
681
+
682
+ Parameters
683
+ ----------
684
+ entry_function : str | None, optional
685
+ The function name to use as the entry point. If None, uses the workflow.
686
+ By default None
687
+
688
+ Returns
689
+ -------
690
+ Workflow
691
+ A per-user workflow instance
692
+
693
+ Raises
694
+ ------
695
+ ValueError
696
+ If no workflow is set (neither per-user nor shared)
697
+ """
698
+ # Determine entry function
699
+ if entry_function is None:
700
+ # Use workflow (could be per-user or shared)
701
+ entry_fn_obj = self.get_workflow()
702
+ else:
703
+ # Use specified function (could be per-user or shared)
704
+ entry_fn_obj = await self.get_function(entry_function)
705
+
706
+ # Collect function names that are included by function groups (shared + per-user)
707
+ # These will be skipped when populating function_configs and all_functions
708
+ included_functions: set[str] = set()
709
+ for configured_fg in self._shared_builder._function_groups.values():
710
+ included_functions.update((await configured_fg.instance.get_included_functions()).keys())
711
+ for configured_fg in self._per_user_function_groups.values():
712
+ included_functions.update((await configured_fg.instance.get_included_functions()).keys())
713
+
714
+ # Collect all functions (per-user + shared), excluding those already in function groups
715
+ all_functions = {}
716
+
717
+ # Add shared functions (skip those included by function groups)
718
+ for name, configured_fn in self._shared_builder._functions.items():
719
+ if name not in included_functions:
720
+ all_functions[name] = configured_fn.instance
721
+
722
+ # Override with per-user functions (skip those included by function groups)
723
+ for name, configured_fn in self._per_user_functions.items():
724
+ if name not in included_functions:
725
+ all_functions[name] = configured_fn.instance
726
+
727
+ # Collect all function groups (shared + per-user)
728
+ all_function_groups = {}
729
+ # Add shared function groups
730
+ for name, configured_fg in self._shared_builder._function_groups.items():
731
+ all_function_groups[name] = configured_fg.instance
732
+ # Override with per-user function groups
733
+ for name, configured_fg in self._per_user_function_groups.items():
734
+ all_function_groups[name] = configured_fg.instance
735
+
736
+ # Build function configs (per-user + shared), excluding those already in function groups
737
+ function_configs = {}
738
+ for name, configured_fn in self._shared_builder._functions.items():
739
+ if name not in included_functions:
740
+ function_configs[name] = configured_fn.config
741
+ for name, configured_fn in self._per_user_functions.items():
742
+ if name not in included_functions:
743
+ function_configs[name] = configured_fn.config
744
+
745
+ # Build function group configs (shared + per-user)
746
+ function_group_configs = {}
747
+ for name, configured_fg in self._shared_builder._function_groups.items():
748
+ function_group_configs[name] = configured_fg.config
749
+ for name, configured_fg in self._per_user_function_groups.items():
750
+ function_group_configs[name] = configured_fg.config
751
+
752
+ # Determine workflow config
753
+ if self._workflow is not None:
754
+ workflow_config = self._workflow.config
755
+ else:
756
+ workflow_config = self._shared_builder.get_workflow_config()
757
+
758
+ # Build the Config object
759
+ per_user_config = Config(general=self._shared_builder.general_config,
760
+ functions=function_configs,
761
+ function_groups=function_group_configs,
762
+ workflow=workflow_config,
763
+ llms={
764
+ k: v.config
765
+ for k, v in self._shared_builder._llms.items()
766
+ },
767
+ embedders={
768
+ k: v.config
769
+ for k, v in self._shared_builder._embedders.items()
770
+ },
771
+ memory={
772
+ k: v.config
773
+ for k, v in self._shared_builder._memory_clients.items()
774
+ },
775
+ object_stores={
776
+ k: v.config
777
+ for k, v in self._shared_builder._object_stores.items()
778
+ },
779
+ retrievers={
780
+ k: v.config
781
+ for k, v in self._shared_builder._retrievers.items()
782
+ },
783
+ ttc_strategies={
784
+ k: v.config
785
+ for k, v in self._shared_builder._ttc_strategies.items()
786
+ })
787
+
788
+ # Create the Workflow instance
789
+ workflow = Workflow.from_entry_fn(config=per_user_config,
790
+ entry_fn=entry_fn_obj,
791
+ functions=all_functions,
792
+ function_groups=all_function_groups,
793
+ llms={
794
+ k: v.instance
795
+ for k, v in self._shared_builder._llms.items()
796
+ },
797
+ embeddings={
798
+ k: v.instance
799
+ for k, v in self._shared_builder._embedders.items()
800
+ },
801
+ memory={
802
+ k: v.instance
803
+ for k, v in self._shared_builder._memory_clients.items()
804
+ },
805
+ object_stores={
806
+ k: v.instance
807
+ for k, v in self._shared_builder._object_stores.items()
808
+ },
809
+ telemetry_exporters={
810
+ k: v.instance
811
+ for k, v in self._shared_builder._telemetry_exporters.items()
812
+ },
813
+ retrievers={
814
+ k: v.instance
815
+ for k, v in self._shared_builder._retrievers.items()
816
+ },
817
+ ttc_strategies={
818
+ k: v.instance
819
+ for k, v in self._shared_builder._ttc_strategies.items()
820
+ },
821
+ context_state=self._shared_builder._context_state)
822
+
823
+ return workflow
824
+
825
+ @classmethod
826
+ @asynccontextmanager
827
+ async def from_config(cls, user_id: str, config: Config, shared_builder: WorkflowBuilder):
828
+ """
829
+ Create and populate a PerUserWorkflowBuilder from config.
830
+
831
+ This is the primary entry point for creating per-user builders.
832
+
833
+ Args:
834
+ user_id: Unique identifier for the user
835
+ config: Full configuration object
836
+ shared_builder: The shared WorkflowBuilder instance
837
+
838
+ Yields:
839
+ PerUserWorkflowBuilder: Populated per-user builder instance
840
+ """
841
+ async with cls(user_id=user_id, shared_builder=shared_builder) as builder:
842
+ await builder.populate_builder(config)
843
+ yield builder