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
nat/runtime/session.py CHANGED
@@ -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");
@@ -14,15 +14,20 @@
14
14
  # limitations under the License.
15
15
 
16
16
  import asyncio
17
- import contextvars
17
+ import logging
18
+ import time
18
19
  import typing
19
20
  import uuid
20
21
  from collections.abc import Awaitable
21
22
  from collections.abc import Callable
22
23
  from contextlib import asynccontextmanager
23
24
  from contextlib import nullcontext
25
+ from datetime import datetime
24
26
 
25
27
  from fastapi import WebSocket
28
+ from pydantic import BaseModel
29
+ from pydantic import ConfigDict
30
+ from pydantic import Field
26
31
  from starlette.requests import HTTPConnection
27
32
  from starlette.requests import Request
28
33
 
@@ -37,67 +42,392 @@ from nat.data_models.interactive import HumanResponse
37
42
  from nat.data_models.interactive import InteractionPrompt
38
43
  from nat.data_models.runtime_enum import RuntimeTypeEnum
39
44
 
40
- _T = typing.TypeVar("_T")
45
+ if typing.TYPE_CHECKING:
46
+ from nat.builder.per_user_workflow_builder import PerUserWorkflowBuilder
47
+ from nat.builder.workflow_builder import WorkflowBuilder
41
48
 
49
+ logger = logging.getLogger(__name__)
42
50
 
43
- class UserManagerBase:
44
- pass
45
51
 
52
+ class PerUserBuilderInfo(BaseModel):
53
+ """
54
+ Container for per-user builder data with activity tracking.
46
55
 
47
- class SessionManager:
56
+ Tracks lifecycle and usage of per-user builders for automatic cleanup.
57
+ """
58
+
59
+ model_config = ConfigDict(arbitrary_types_allowed=True, validate_assignment=False)
60
+
61
+ builder: typing.Any = Field(description="The per-user workflow builder instance")
62
+ workflow: typing.Any = Field(description="The cached per-user workflow instance")
63
+ semaphore: typing.Any = Field(description="Per-user semaphore for concurrency control")
64
+ last_activity: datetime = Field(default_factory=datetime.now,
65
+ description="The timestamp of the last access to this builder")
66
+ ref_count: int = Field(default=0, ge=0, description="The reference count of this builder")
67
+ lock: asyncio.Lock = Field(default_factory=asyncio.Lock, description="Lock for thread-safe ref_count updates")
68
+
69
+ # Monitoring metrics
70
+ created_at: datetime = Field(default_factory=datetime.now, description="When the per-user workflow was created")
71
+ total_requests: int = Field(default=0, ge=0, description="Total number of requests processed")
72
+ error_count: int = Field(default=0, ge=0, description="Total number of failed requests")
73
+ total_latency_ms: float = Field(default=0.0, ge=0, description="Total latency of all requests in milliseconds")
74
+
75
+ def record_request(self, latency_ms: float, success: bool) -> None:
76
+ """Record metrics for a completed request.
77
+
78
+ Args:
79
+ latency_ms: Request latency in milliseconds
80
+ success: Whether the request was successful
81
+ """
82
+ self.total_requests += 1
83
+ self.total_latency_ms += latency_ms
84
+ if not success:
85
+ self.error_count += 1
86
+
87
+
88
+ class Session:
89
+ """
90
+ Represents an active session with access to workflow and builders.
91
+
92
+ Each session is tied to a specific request, and provides access to the appropriate workflow
93
+ instance (shared or per-user).
94
+
95
+ Lifecycle:
96
+ - Created for each request via SessionManager.session()
97
+ - Automatically manages ref_count for per-user builder tracking
98
+ - Cleans up context variables on exit
99
+
100
+ Concurrency:
101
+ - Each session has its own semaphore for concurrency control
102
+ - For per-user workflows: each user has an independent concurrency limit
103
+ - For shared workflows: all sessions share the SessionManager's semaphore
104
+ """
48
105
 
49
106
  def __init__(self,
107
+ session_manager: "SessionManager",
50
108
  workflow: Workflow,
109
+ semaphore: asyncio.Semaphore | nullcontext,
110
+ user_id: str | None = None):
111
+ self._session_manager = session_manager
112
+ self._workflow = workflow
113
+ self._semaphore = semaphore
114
+ self._user_id = user_id
115
+
116
+ @property
117
+ def user_id(self) -> str | None:
118
+ return self._user_id
119
+
120
+ @property
121
+ def workflow(self) -> Workflow:
122
+ return self._workflow
123
+
124
+ @property
125
+ def session_manager(self) -> "SessionManager":
126
+ return self._session_manager
127
+
128
+ @asynccontextmanager
129
+ async def run(self, message, runtime_type: RuntimeTypeEnum = RuntimeTypeEnum.RUN_OR_SERVE):
130
+ """
131
+ Start a workflow run using this session's workflow.
132
+
133
+ Args:
134
+ message: Input message for the workflow
135
+ runtime_type: Runtime type (defaults to SessionManager's runtime_type)
136
+
137
+ Yields:
138
+ Runner instance for the workflow execution
139
+ """
140
+ async with self._semaphore:
141
+ async with self._workflow.run(message, runtime_type=runtime_type) as runner:
142
+ yield runner
143
+
144
+
145
+ class SessionManager:
146
+
147
+ def __init__(self,
148
+ config: Config,
149
+ shared_builder: "WorkflowBuilder",
150
+ entry_function: str | None = None,
151
+ shared_workflow: Workflow | None = None,
51
152
  max_concurrency: int = 8,
52
153
  runtime_type: RuntimeTypeEnum = RuntimeTypeEnum.RUN_OR_SERVE):
53
154
  """
54
- The SessionManager class is used to run and manage a user workflow session. It runs and manages the context,
55
- and configuration of a workflow with the specified concurrency.
155
+ The SessionManager class is used to manage workflow builders and sessions.
156
+ It manages workflow sessions and per-user builders with lifecycle management.
157
+
158
+ Architecture:
159
+ - One SessionManager per FastAPI server
160
+ - Creates/caches PerUserWorkflowBuilder instances per user
161
+ - Cleans up inactive builders based on timeout
56
162
 
57
163
  Parameters
58
164
  ----------
59
- workflow : Workflow
60
- The workflow to run
165
+ config : Config
166
+ The configuration for the workflow
167
+ shared_builder : WorkflowBuilder
168
+ The shared workflow builder
169
+ entry_function : str | None, optional
170
+ The entry function for this SessionManager's workflows, by default None
171
+ shared_workflow : Workflow, optional
172
+ The shared workflow, by default None
61
173
  max_concurrency : int, optional
62
174
  The maximum number of simultaneous workflow invocations, by default 8
63
175
  runtime_type : RuntimeTypeEnum, optional
64
176
  The type of runtime the session manager is operating in, by default RuntimeTypeEnum.RUN_OR_SERVE
65
177
  """
66
178
 
67
- if (workflow is None):
68
- raise ValueError("Workflow cannot be None")
69
-
70
- self._workflow: Workflow = workflow
179
+ from nat.cli.type_registry import GlobalTypeRegistry
71
180
 
181
+ self._config = config
72
182
  self._max_concurrency = max_concurrency
73
- self._context_state = ContextState.get()
74
- self._context = Context(self._context_state)
75
- self._runtime_type = runtime_type
76
-
77
- # We save the context because Uvicorn spawns a new process
78
- # for each request, and we need to restore the context vars
79
- self._saved_context = contextvars.copy_context()
183
+ self._entry_function = entry_function
80
184
 
81
- if (max_concurrency > 0):
185
+ # Semaphore for limiting concurrency
186
+ if max_concurrency > 0:
82
187
  self._semaphore = asyncio.Semaphore(max_concurrency)
83
188
  else:
84
189
  # If max_concurrency is 0, then we don't need to limit the concurrency but we still need a context
85
190
  self._semaphore = nullcontext()
86
191
 
192
+ self._runtime_type = runtime_type
193
+
194
+ # Context state for per-request context variables
195
+ self._context_state = ContextState.get()
196
+ self._context = Context(self._context_state)
197
+
198
+ # Track if workflow is shared or per-user
199
+ workflow_registration = GlobalTypeRegistry.get().get_function(type(config.workflow))
200
+ self._is_workflow_per_user = workflow_registration.is_per_user
201
+
202
+ # Shared components
203
+ self._shared_builder = shared_builder
204
+ self._shared_workflow = shared_workflow
205
+
206
+ # Per-user management
207
+ self._per_user_builders: dict[str, PerUserBuilderInfo] = {}
208
+ self._per_user_builders_lock = asyncio.Lock()
209
+ self._per_user_builders_cleanup_task: asyncio.Task | None = None
210
+ self._per_user_session_timeout = config.general.per_user_workflow_timeout
211
+ self._per_user_session_cleanup_interval = config.general.per_user_workflow_cleanup_interval
212
+ self._shutdown_event = asyncio.Event()
213
+
214
+ # Cache schemas for per-user workflows
215
+ if self._is_workflow_per_user:
216
+ self._per_user_workflow_input_schema = workflow_registration.per_user_function_input_schema
217
+ self._per_user_workflow_single_output_schema = workflow_registration.per_user_function_single_output_schema
218
+ self._per_user_workflow_streaming_output_schema = \
219
+ workflow_registration.per_user_function_streaming_output_schema
220
+ else:
221
+ self._per_user_workflow_input_schema = None
222
+ self._per_user_workflow_single_output_schema = None
223
+ self._per_user_workflow_streaming_output_schema = None
224
+
87
225
  @property
88
226
  def config(self) -> Config:
89
- return self._workflow.config
227
+ return self._config
90
228
 
91
229
  @property
92
230
  def workflow(self) -> Workflow:
93
- return self._workflow
231
+ """
232
+ Get workflow for backward compatibility.
233
+
234
+ Only works for shared workflows. For per-user workflows, use session.workflow.
235
+
236
+ Raises:
237
+ ValueError: If workflow is per-user
238
+ """
239
+ if self._is_workflow_per_user:
240
+ raise ValueError("Workflow is per-user. Access workflow through session.workflow instead.")
241
+ if self._shared_workflow is None:
242
+ raise ValueError("No shared workflow available")
243
+ return self._shared_workflow
244
+
245
+ @property
246
+ def shared_builder(self) -> "WorkflowBuilder":
247
+ return self._shared_builder
94
248
 
95
249
  @property
96
- def context(self) -> Context:
97
- return self._context
250
+ def is_workflow_per_user(self) -> bool:
251
+ return self._is_workflow_per_user
252
+
253
+ def get_workflow_input_schema(self) -> type[BaseModel]:
254
+ """Get workflow input schema for OpenAPI documentation."""
255
+
256
+ if self._is_workflow_per_user:
257
+ return self._per_user_workflow_input_schema
258
+
259
+ return self._shared_workflow.input_schema
260
+
261
+ def get_workflow_single_output_schema(self) -> type[BaseModel]:
262
+ """Get workflow single output schema for OpenAPI documentation."""
263
+
264
+ if self._is_workflow_per_user:
265
+ return self._per_user_workflow_single_output_schema
266
+
267
+ return self._shared_workflow.single_output_schema
268
+
269
+ def get_workflow_streaming_output_schema(self) -> type[BaseModel]:
270
+ """Get workflow streaming output schema for OpenAPI documentation."""
271
+
272
+ if self._is_workflow_per_user:
273
+ return self._per_user_workflow_streaming_output_schema
274
+
275
+ return self._shared_workflow.streaming_output_schema
276
+
277
+ @classmethod
278
+ async def create(cls,
279
+ config: Config,
280
+ shared_builder: "WorkflowBuilder",
281
+ entry_function: str | None = None,
282
+ max_concurrency: int = 8,
283
+ runtime_type: RuntimeTypeEnum = RuntimeTypeEnum.RUN_OR_SERVE) -> "SessionManager":
284
+ """
285
+ Create a SessionManager. This is the preferred way to instantiate.
286
+
287
+ Handles async workflow building and starts cleanup task if per-user.
288
+ """
289
+ from nat.cli.type_registry import GlobalTypeRegistry
290
+
291
+ workflow_registration = GlobalTypeRegistry.get().get_function(type(config.workflow))
292
+
293
+ if workflow_registration.is_per_user:
294
+ shared_workflow = None
295
+ logger.info(f"Workflow is per-user (entry_function={entry_function})")
296
+ else:
297
+ shared_workflow = await shared_builder.build(entry_function=entry_function)
298
+ logger.info(f"Shared workflow built (entry_function={entry_function})")
299
+
300
+ session_manager = cls(config=config,
301
+ shared_builder=shared_builder,
302
+ entry_function=entry_function,
303
+ shared_workflow=shared_workflow,
304
+ max_concurrency=max_concurrency,
305
+ runtime_type=runtime_type)
306
+
307
+ # Start cleanup task for per-user workflows
308
+ if session_manager._is_workflow_per_user:
309
+ session_manager._per_user_builders_cleanup_task = asyncio.create_task(
310
+ session_manager._run_periodic_cleanup())
311
+
312
+ return session_manager
313
+
314
+ async def _run_periodic_cleanup(self):
315
+
316
+ logger.debug("Running periodic cleanup of per-user builders")
317
+ while not self._shutdown_event.is_set():
318
+ try:
319
+ # Wait for either cleanup interval or shutdown
320
+ await asyncio.wait_for(self._shutdown_event.wait(),
321
+ timeout=self._per_user_session_cleanup_interval.total_seconds())
322
+ # If we get here, shutdown was signaled
323
+ break
324
+ except TimeoutError:
325
+ # Timeout means it's time to run cleanup
326
+ try:
327
+ await self._cleanup_inactive_per_user_builders()
328
+ except Exception:
329
+ logger.exception("Error during periodic cleanup")
330
+
331
+ logger.debug("Periodic cleanup task shutting down")
332
+
333
+ async def _cleanup_inactive_per_user_builders(self) -> int:
334
+
335
+ now = datetime.now()
336
+ threshold = now - self._per_user_session_timeout
337
+ builders_to_cleanup: list[tuple[str, PerUserBuilderInfo]] = []
338
+
339
+ # Identify builders to cleanup (under lock)
340
+ async with self._per_user_builders_lock:
341
+ for user_id, builder_info in list(self._per_user_builders.items()):
342
+ if builder_info.ref_count == 0 and builder_info.last_activity < threshold:
343
+ # Remove from dict and add to cleanup list
344
+ builders_to_cleanup.append((user_id, builder_info))
345
+ del self._per_user_builders[user_id]
346
+ logger.debug(f"Marked per-user builder for user {user_id} for cleanup "
347
+ f"(inactive since {builder_info.last_activity.isoformat()})")
348
+ # Cleanup builders (outside lock to avoid blocking)
349
+ for user_id, builder_info in builders_to_cleanup:
350
+ try:
351
+ await builder_info.builder.__aexit__(None, None, None)
352
+ logger.info(f"Cleaned up inactive per-user builder for user={user_id} "
353
+ f"(remaining users: {len(self._per_user_builders)})")
354
+ except Exception:
355
+ logger.exception(f"Error cleaning up per-user builder for user {user_id}")
356
+
357
+ return len(builders_to_cleanup)
358
+
359
+ def _get_user_id_from_context(self) -> str | None:
360
+ """
361
+ Get user ID from current context.
362
+
363
+ Extraction order:
364
+ 1. From context user_id (set from nat-session cookie)
365
+ 2. From context user_manager if set
366
+ 3. None (for shared workflow or unauthenticated access)
367
+
368
+ """
369
+ try:
370
+ # Primary: Get from context user_id (already extracted from nat-session cookie)
371
+ user_id = self._context.user_id
372
+ if user_id:
373
+ return user_id
374
+
375
+ # Fallback: Get from user_manager if set
376
+ user_manager = self._context.user_manager
377
+ if user_manager:
378
+ return user_manager.get_id()
379
+ return None
380
+ except Exception as e:
381
+ logger.debug(f"Could not extract user_id from context: {e}")
382
+ return None
383
+
384
+ async def _get_or_create_per_user_builder(self, user_id: str) -> tuple["PerUserWorkflowBuilder", Workflow]:
385
+ from nat.builder.per_user_workflow_builder import PerUserWorkflowBuilder
386
+
387
+ async with self._per_user_builders_lock:
388
+ if user_id in self._per_user_builders:
389
+ builder_info = self._per_user_builders[user_id]
390
+ builder_info.last_activity = datetime.now()
391
+
392
+ return builder_info.builder, builder_info.workflow
393
+
394
+ logger.info(f"Creating per-user builder for user={user_id}, entry_function={self._entry_function}")
395
+ builder = PerUserWorkflowBuilder(user_id=user_id, shared_builder=self._shared_builder)
396
+ # Enter the builder's context manually to avoid exiting the context manager
397
+ # Exit the context when cleaning up the builder
398
+ await builder.__aenter__()
399
+
400
+ try:
401
+ await builder.populate_builder(self._config)
402
+ workflow = await builder.build(entry_function=self._entry_function)
403
+
404
+ # Create per-user semaphore for concurrency control
405
+ if self._max_concurrency > 0:
406
+ per_user_semaphore = asyncio.Semaphore(self._max_concurrency)
407
+ else:
408
+ per_user_semaphore = nullcontext()
409
+
410
+ builder_info = PerUserBuilderInfo(builder=builder,
411
+ workflow=workflow,
412
+ semaphore=per_user_semaphore,
413
+ last_activity=datetime.now(),
414
+ ref_count=0,
415
+ lock=asyncio.Lock())
416
+ self._per_user_builders[user_id] = builder_info
417
+ logger.info(
418
+ f"Created per-user builder for user={user_id} (total users: {len(self._per_user_builders)})")
419
+ return builder_info.builder, builder_info.workflow
420
+ except Exception:
421
+ logger.exception(f"Error creating per-user builder for user {user_id}")
422
+ try:
423
+ await builder.__aexit__(None, None, None)
424
+ except Exception:
425
+ logger.exception("Error during builder cleanup after failed creation")
426
+ raise
98
427
 
99
428
  @asynccontextmanager
100
429
  async def session(self,
430
+ user_id: str | None = None,
101
431
  user_manager=None,
102
432
  http_connection: HTTPConnection | None = None,
103
433
  user_message_id: str | None = None,
@@ -124,9 +454,65 @@ class SessionManager:
124
454
  if isinstance(http_connection, Request):
125
455
  self.set_metadata_from_http_request(http_connection)
126
456
 
457
+ builder_info: PerUserBuilderInfo | None = None
458
+ request_start_time: float | None = None
459
+ request_success = True
460
+
461
+ if self._is_workflow_per_user:
462
+ # Resolve user_id: explicit param > context
463
+ if user_id is None:
464
+ user_id = self._get_user_id_from_context()
465
+ if user_id is None:
466
+ raise ValueError("user_id is required for per-user workflow but could not be determined. "
467
+ "Ensure 'nat-session' cookie is set or pass user_id explicitly.")
468
+
469
+ # To ensure the user_id is set in the context before the per-user builder is created
470
+ self._context_state.user_id.set(user_id)
471
+
472
+ # Get or create per-user builder
473
+ logger.debug(f"Getting or creating per-user builder for user {user_id}")
474
+ _, workflow = await self._get_or_create_per_user_builder(user_id)
475
+ builder_info = self._per_user_builders[user_id]
476
+ async with builder_info.lock:
477
+ builder_info.ref_count += 1
478
+ logger.debug(f"Incremented ref_count for user {user_id} to {builder_info.ref_count}")
479
+ # Use per-user semaphore for concurrency control
480
+ semaphore = builder_info.semaphore
481
+ # Start request timing for metrics
482
+ request_start_time = time.perf_counter()
483
+ else:
484
+ workflow = self._shared_workflow
485
+ # Use shared semaphore for concurrency control
486
+ semaphore = self._semaphore
487
+
488
+ # TODO: this logic needs to be cleaned up since it is a duplicated setting of the user_id
489
+ # But we need to keep it for now to maintain the token_user_id
490
+ token_user_id = None
491
+ if user_id is not None:
492
+ token_user_id = self._context_state.user_id.set(user_id)
493
+
127
494
  try:
128
- yield self
495
+ session = Session(session_manager=self, user_id=user_id, workflow=workflow, semaphore=semaphore)
496
+
497
+ yield session
498
+
499
+ except Exception:
500
+ request_success = False
501
+ raise
502
+
129
503
  finally:
504
+ if builder_info is not None:
505
+ async with builder_info.lock:
506
+ builder_info.ref_count -= 1
507
+ builder_info.last_activity = datetime.now()
508
+
509
+ # Record request metrics
510
+ if request_start_time is not None:
511
+ latency_ms = (time.perf_counter() - request_start_time) * 1000
512
+ builder_info.record_request(latency_ms, request_success)
513
+
514
+ if token_user_id is not None:
515
+ self._context_state.user_id.reset(token_user_id)
130
516
  if token_user_manager is not None:
131
517
  self._context_state.user_manager.reset(token_user_manager)
132
518
  if token_user_input is not None:
@@ -139,14 +525,39 @@ class SessionManager:
139
525
  """
140
526
  Start a workflow run
141
527
  """
528
+ if self._is_workflow_per_user:
529
+ raise ValueError("Cannot use SessionManager.run() with per-user workflows. "
530
+ "Use 'async with session_manager.session() as session' then 'session.run()' instead.")
142
531
  async with self._semaphore:
143
- # Apply the saved context
144
- for k, v in self._saved_context.items():
145
- k.set(v)
146
-
147
- async with self._workflow.run(message, runtime_type=runtime_type) as runner:
532
+ async with self._shared_workflow.run(message, runtime_type=runtime_type) as runner:
148
533
  yield runner
149
534
 
535
+ async def shutdown(self) -> None:
536
+ """
537
+ Shutdown the SessionManager and cleanup resources.
538
+
539
+ Call this when the SessionManager is no longer needed.
540
+ """
541
+ if self._is_workflow_per_user:
542
+ # Shutdown cleanup task
543
+ self._shutdown_event.set()
544
+ if self._per_user_builders_cleanup_task:
545
+ try:
546
+ await asyncio.wait_for(self._per_user_builders_cleanup_task, timeout=5.0)
547
+ except TimeoutError:
548
+ logger.warning("Cleanup task did not finish in time, cancelling")
549
+ self._per_user_builders_cleanup_task.cancel()
550
+
551
+ # Cleanup all per-user builders
552
+ async with self._per_user_builders_lock:
553
+ for user_id, builder_info in list(self._per_user_builders.items()):
554
+ logger.debug(f"Cleaning up per-user builder for user {user_id}")
555
+ try:
556
+ await builder_info.builder.__aexit__(None, None, None)
557
+ except Exception:
558
+ logger.exception(f"Error cleaning up builder for user {user_id}")
559
+ self._per_user_builders.clear()
560
+
150
561
  def set_metadata_from_http_request(self, request: Request) -> None:
151
562
  """
152
563
  Extracts and sets user metadata request attributes from a HTTP request.
@@ -169,6 +580,10 @@ class SessionManager:
169
580
  if request.headers.get("user-message-id"):
170
581
  self._context_state.user_message_id.set(request.headers["user-message-id"])
171
582
 
583
+ # Set user_id from nat-session cookie
584
+ if request.cookies.get("nat-session"):
585
+ self._context_state.user_id.set(request.cookies["nat-session"])
586
+
172
587
  # W3C Trace Context header: traceparent: 00-<trace-id>-<span-id>-<flags>
173
588
  traceparent = request.headers.get("traceparent")
174
589
  if traceparent:
@@ -219,6 +634,10 @@ class SessionManager:
219
634
  self._context.metadata._request.cookies = cookies
220
635
  self._context_state.metadata.set(self._context.metadata)
221
636
 
637
+ # Set user_id from nat-session cookie
638
+ if cookies.get("nat-session"):
639
+ self._context_state.user_id.set(cookies["nat-session"])
640
+
222
641
  if conversation_id is not None:
223
642
  self._context_state.conversation_id.set(conversation_id)
224
643
 
@@ -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) 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) 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,5 +1,5 @@
1
1
  <!--
2
- SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
2
+ SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
3
3
  SPDX-License-Identifier: Apache-2.0
4
4
 
5
5
  Licensed under the Apache License, Version 2.0 (the "License");
@@ -1,4 +1,4 @@
1
- # Copyright (c) 2024-2025, NVIDIA CORPORATION. All rights reserved.
1
+ # Copyright (c) 2024-2026, NVIDIA CORPORATION. All rights reserved.
2
2
  #
3
3
  # Licensed under the Apache License, Version 2.0 (the "License");
4
4
  # you may not use this file except in compliance with the License.
@@ -1,4 +1,4 @@
1
- # Copyright (c) 2024-2025, NVIDIA CORPORATION. All rights reserved.
1
+ # Copyright (c) 2024-2026, NVIDIA CORPORATION. All rights reserved.
2
2
  #
3
3
  # Licensed under the Apache License, Version 2.0 (the "License");
4
4
  # you may not use this file except in compliance with the License.
@@ -1,4 +1,4 @@
1
- # Copyright (c) 2024-2025, NVIDIA CORPORATION. All rights reserved.
1
+ # Copyright (c) 2024-2026, NVIDIA CORPORATION. All rights reserved.
2
2
  #
3
3
  # Licensed under the Apache License, Version 2.0 (the "License");
4
4
  # you may not use this file except in compliance with the License.
@@ -1,4 +1,4 @@
1
- # Copyright (c) 2024-2025, NVIDIA CORPORATION. All rights reserved.
1
+ # Copyright (c) 2024-2026, NVIDIA CORPORATION. All rights reserved.
2
2
  #
3
3
  # Licensed under the Apache License, Version 2.0 (the "License");
4
4
  # you may not use this file except in compliance with the License.
@@ -1,6 +1,6 @@
1
1
  #!/bin/bash
2
2
 
3
- # Copyright (c) 2024-2025, NVIDIA CORPORATION. All rights reserved.
3
+ # Copyright (c) 2024-2026, NVIDIA CORPORATION. All rights reserved.
4
4
  #
5
5
  # Licensed under the Apache License, Version 2.0 (the "License");
6
6
  # you may not use this file except in compliance with 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
- # Copyright (c) 2024-2025, NVIDIA CORPORATION. All rights reserved.
1
+ # Copyright (c) 2024-2026, NVIDIA CORPORATION. All rights reserved.
2
2
  #
3
3
  # Licensed under the Apache License, Version 2.0 (the "License");
4
4
  # you may not use this file except in compliance with 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) 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");