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,132 @@
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
+ """Configuration for dynamic middleware."""
16
+
17
+ from __future__ import annotations
18
+
19
+ from pydantic import BaseModel
20
+ from pydantic import Field
21
+ from pydantic import model_validator
22
+
23
+ from nat.data_models.component import ComponentGroup
24
+ from nat.data_models.component_ref import AuthenticationRef
25
+ from nat.data_models.component_ref import EmbedderRef
26
+ from nat.data_models.component_ref import LLMRef
27
+ from nat.data_models.component_ref import MemoryRef
28
+ from nat.data_models.component_ref import ObjectStoreRef
29
+ from nat.data_models.component_ref import RetrieverRef
30
+ from nat.data_models.middleware import FunctionMiddlewareBaseConfig
31
+
32
+
33
+ class AllowedComponentFunctions(BaseModel):
34
+ """Component functions allowed for auto-registration.
35
+
36
+ Default allowlists are provided for each component type. User-provided
37
+ values are automatically merged with defaults.
38
+ Set to None or omit to use only defaults.
39
+ """
40
+
41
+ llms: set[str] | None = Field(
42
+ default=None, description="Additional LLM functions that should be allowed to register with middleware.")
43
+ embedders: set[str] | None = Field(
44
+ default=None, description="Additional Embedder functions that should be allowed to register with middleware.")
45
+ retrievers: set[str] | None = Field(
46
+ default=None, description="Additional Retriever functions that should be allowed to register with middleware.")
47
+ memory: set[str] | None = Field(
48
+ default=None, description="Additional Memory functions that should be allowed to register with middleware.")
49
+ object_stores: set[str] | None = Field(
50
+ default=None,
51
+ description="Additional Object Store functions that should be allowed to register with middleware.")
52
+ authentication: set[str] | None = Field(
53
+ default=None,
54
+ description="Additional Authentication functions that should be allowed to register with middleware.")
55
+
56
+ @model_validator(mode='after')
57
+ def merge_with_defaults(self):
58
+ """Merge user-provided values with defaults from COMPONENT_FUNCTION_ALLOWLISTS."""
59
+ from nat.middleware.utils.workflow_inventory import COMPONENT_FUNCTION_ALLOWLISTS
60
+
61
+ def merge(component_group: ComponentGroup, user_set: set[str] | None) -> set[str]:
62
+ defaults = COMPONENT_FUNCTION_ALLOWLISTS[component_group]
63
+ if user_set is None:
64
+ return defaults.copy()
65
+ return defaults | user_set
66
+
67
+ self.llms = merge(ComponentGroup.LLMS, self.llms)
68
+ self.embedders = merge(ComponentGroup.EMBEDDERS, self.embedders)
69
+ self.retrievers = merge(ComponentGroup.RETRIEVERS, self.retrievers)
70
+ self.memory = merge(ComponentGroup.MEMORY, self.memory)
71
+ self.object_stores = merge(ComponentGroup.OBJECT_STORES, self.object_stores)
72
+ self.authentication = merge(ComponentGroup.AUTHENTICATION, self.authentication)
73
+
74
+ return self
75
+
76
+
77
+ class DynamicMiddlewareConfig(FunctionMiddlewareBaseConfig, name="dynamic_middleware"):
78
+ """Configuration for dynamic middleware.
79
+
80
+ Controls which components and functions to intercept, and which policies to apply.
81
+ Supports explicit component references and auto-discovery flags.
82
+ """
83
+
84
+ # === Component References ===
85
+
86
+ llms: list[LLMRef] | None = Field(default=None, description="LLMs to intercept")
87
+
88
+ embedders: list[EmbedderRef] | None = Field(default=None, description="Embedders component functions to intercept")
89
+
90
+ retrievers: list[RetrieverRef] | None = Field(default=None,
91
+ description="Retrievers component functions to intercept")
92
+
93
+ memory: list[MemoryRef] | None = Field(default=None, description="Memory component functions to intercept")
94
+
95
+ object_stores: list[ObjectStoreRef] | None = Field(default=None,
96
+ description="Object stores component functions to intercept")
97
+
98
+ auth_providers: list[AuthenticationRef] | None = Field(
99
+ default=None, description="Authentication providers component functions to intercept")
100
+
101
+ # === Component and Function Auto-Discovery Flags ===
102
+
103
+ register_llms: bool | None = Field(default=False,
104
+ description="Auto-discover and register all LLMs component functions")
105
+
106
+ register_embedders: bool | None = Field(default=False,
107
+ description="Auto-discover and register all embedders component functions")
108
+
109
+ register_retrievers: bool | None = Field(
110
+ default=False, description="Auto-discover and register all retrievers component functions")
111
+
112
+ register_memory: bool | None = Field(
113
+ default=False, description="Auto-discover and register all memory providers component functions")
114
+
115
+ register_object_stores: bool | None = Field(
116
+ default=False, description="Auto-discover and register all object stores component functions")
117
+
118
+ register_auth_providers: bool | None = Field(
119
+ default=False, description="Auto-discover and register all authentication providers component functions")
120
+
121
+ register_workflow_functions: bool | None = Field(default=False,
122
+ description="Auto-discover and register all workflow functions")
123
+
124
+ # === Enable/Disable ===
125
+
126
+ enabled: bool = Field(default=True, description="Whether this middleware is active")
127
+
128
+ # === Component Function Allowlists ===
129
+
130
+ allowed_component_functions: AllowedComponentFunctions | None = Field(
131
+ default=None,
132
+ description="Functions allowed for auto-registration. Omit to use defaults, provide to extend them")
@@ -0,0 +1,34 @@
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
+ """Registration for dynamic middleware."""
16
+
17
+ from nat.builder.builder import Builder
18
+ from nat.cli.register_workflow import register_middleware
19
+ from nat.middleware.dynamic.dynamic_function_middleware import DynamicFunctionMiddleware
20
+ from nat.middleware.dynamic.dynamic_middleware_config import DynamicMiddlewareConfig
21
+
22
+
23
+ @register_middleware(config_type=DynamicMiddlewareConfig)
24
+ async def dynamic_middleware(config: DynamicMiddlewareConfig, builder: Builder):
25
+ """Build a dynamic middleware from configuration.
26
+
27
+ Args:
28
+ config: The dynamic middleware configuration
29
+ builder: The workflow builder
30
+
31
+ Yields:
32
+ A configured dynamic middleware instance
33
+ """
34
+ yield DynamicFunctionMiddleware(config=config, builder=builder)
@@ -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");
@@ -37,81 +37,250 @@ from typing import Any
37
37
  from nat.middleware.middleware import CallNext
38
38
  from nat.middleware.middleware import CallNextStream
39
39
  from nat.middleware.middleware import FunctionMiddlewareContext
40
+ from nat.middleware.middleware import InvocationContext
40
41
  from nat.middleware.middleware import Middleware
41
42
 
42
43
 
43
44
  class FunctionMiddleware(Middleware):
44
- """Specialized middleware for function-specific wrapping.
45
-
46
- This class extends the base Middleware class and provides function-specific
47
- wrapping methods. Functions that use this middleware type will call
48
- ``function_middleware_invoke`` and ``function_middleware_stream`` instead of
49
- the base ``middleware_invoke`` and ``middleware_stream`` methods.
45
+ """Base class for function middleware with pre/post-invoke hooks.
46
+
47
+ Middleware intercepts function calls and can:
48
+ - Transform inputs before execution (pre_invoke)
49
+ - Transform outputs after execution (post_invoke)
50
+ - Override function_middleware_invoke for full control
51
+
52
+ Lifecycle:
53
+ - Framework checks ``enabled`` property before calling any methods
54
+ - If disabled, middleware is skipped entirely (no methods called)
55
+ - Users do NOT need to check ``enabled`` in their implementations
56
+
57
+ Inherited abstract members that must be implemented:
58
+ - enabled: Property that returns whether middleware should run
59
+ - pre_invoke: Transform inputs before function execution
60
+ - post_invoke: Transform outputs after function execution
61
+
62
+ Context Flow:
63
+ - FunctionMiddlewareContext (frozen): Static function metadata only
64
+ - InvocationContext: Unified context for both pre and post invoke phases
65
+ - Pre-invoke: output is None, modify modified_args/modified_kwargs
66
+ - Post-invoke: output has the result, modify output to transform
67
+
68
+ Example::
69
+
70
+ class LoggingMiddleware(FunctionMiddleware):
71
+ def __init__(self, config: LoggingConfig):
72
+ super().__init__()
73
+ self._config = config
74
+
75
+ @property
76
+ def enabled(self) -> bool:
77
+ return self._config.enabled
78
+
79
+ async def pre_invoke(self, context: InvocationContext) -> InvocationContext | None:
80
+ logger.info(f"Calling {context.function_context.name} with {context.modified_args}")
81
+ logger.info(f"Original args: {context.original_args}")
82
+ return None # Pass through unchanged
83
+
84
+ async def post_invoke(self, context: InvocationContext) -> InvocationContext | None:
85
+ logger.info(f"Result: {context.output}")
86
+ return None # Pass through unchanged
50
87
  """
51
88
 
52
- async def middleware_invoke(self, value: Any, call_next: CallNext, context: FunctionMiddlewareContext) -> Any:
89
+ @property
90
+ def enabled(self) -> bool:
91
+ """Check if this middleware is enabled.
92
+
93
+ Returns:
94
+ True if the middleware should be applied, False otherwise.
95
+ Default implementation always returns True.
96
+ """
97
+ return True
98
+
99
+ async def pre_invoke(self, context: InvocationContext) -> InvocationContext | None:
100
+ """Pre-invocation hook called before the function is invoked.
101
+
102
+ Args:
103
+ context: Invocation context containing function metadata and args
104
+
105
+ Returns:
106
+ InvocationContext if modified, or None to pass through unchanged.
107
+ Default implementation does nothing.
108
+ """
109
+ del context # Unused by default implementation
110
+ return None
111
+
112
+ async def post_invoke(self, context: InvocationContext) -> InvocationContext | None:
113
+ """Post-invocation hook called after the function returns.
114
+
115
+ Args:
116
+ context: Invocation context containing function metadata, args, and output
117
+
118
+ Returns:
119
+ InvocationContext if modified, or None to pass through unchanged.
120
+ Default implementation does nothing.
121
+ """
122
+ del context # Unused by default implementation
123
+ return None
124
+
125
+ # ==================== Middleware Delegation ====================
126
+ async def middleware_invoke(self,
127
+ *args: Any,
128
+ call_next: CallNext,
129
+ context: FunctionMiddlewareContext,
130
+ **kwargs: Any) -> Any:
53
131
  """Delegate to function_middleware_invoke for function-specific handling."""
54
- return await self.function_middleware_invoke(value, call_next, context)
132
+ return await self.function_middleware_invoke(*args, call_next=call_next, context=context, **kwargs)
55
133
 
56
- async def middleware_stream(self, value: Any, call_next: CallNextStream,
57
- context: FunctionMiddlewareContext) -> AsyncIterator[Any]:
134
+ async def middleware_stream(self,
135
+ *args: Any,
136
+ call_next: CallNextStream,
137
+ context: FunctionMiddlewareContext,
138
+ **kwargs: Any) -> AsyncIterator[Any]:
58
139
  """Delegate to function_middleware_stream for function-specific handling."""
59
- async for chunk in self.function_middleware_stream(value, call_next, context):
140
+ async for chunk in self.function_middleware_stream(*args, call_next=call_next, context=context, **kwargs):
60
141
  yield chunk
61
142
 
62
- async def function_middleware_invoke(self, value: Any, call_next: CallNext,
63
- context: FunctionMiddlewareContext) -> Any:
64
- """Function-specific middleware for single-output invocations.
143
+ # ==================== Orchestration ====================
144
+
145
+ async def function_middleware_invoke(
146
+ self,
147
+ *args: Any,
148
+ call_next: CallNext,
149
+ context: FunctionMiddlewareContext,
150
+ **kwargs: Any,
151
+ ) -> Any:
152
+ """Execute middleware hooks around function call.
153
+
154
+ Default implementation orchestrates: pre_invoke → call_next → post_invoke
155
+
156
+ Override for full control over execution flow (e.g., caching,
157
+ retry logic, conditional execution).
158
+
159
+ Note: Framework checks ``enabled`` before calling this method.
160
+ You do NOT need to check ``enabled`` yourself.
65
161
 
66
162
  Args:
67
- value: The input value to process
68
- call_next: Callable to invoke the next middleware or function
69
- context: Metadata about the function being wrapped
163
+ args: Positional arguments for the function (first arg is typically the input value).
164
+ call_next: Callable to invoke next middleware or target function.
165
+ context: Static function metadata.
166
+ kwargs: Keyword arguments for the function.
70
167
 
71
168
  Returns:
72
- The (potentially modified) output from the function
73
-
74
- The default implementation simply delegates to ``call_next``. Override this
75
- in subclasses to add function-specific preprocessing and postprocessing.
169
+ The (potentially transformed) function output.
76
170
  """
77
- return await call_next(value)
78
-
79
- async def function_middleware_stream(self,
80
- value: Any,
81
- call_next: CallNextStream,
82
- context: FunctionMiddlewareContext) -> AsyncIterator[Any]:
83
- """Function-specific middleware for streaming invocations.
171
+ # Build invocation context with frozen originals + mutable current
172
+ # output starts as None (pre-invoke phase)
173
+ ctx = InvocationContext(
174
+ function_context=context,
175
+ original_args=args,
176
+ original_kwargs=dict(kwargs),
177
+ modified_args=args,
178
+ modified_kwargs=dict(kwargs),
179
+ output=None,
180
+ )
181
+
182
+ # Pre-invoke transformation (output is None at this phase)
183
+ result = await self.pre_invoke(ctx)
184
+ if result is not None:
185
+ ctx = result
186
+
187
+ # Execute function with (potentially modified) args/kwargs
188
+ ctx.output = await call_next(*ctx.modified_args, **ctx.modified_kwargs)
189
+
190
+ # Post-invoke transformation (output now has the result)
191
+ result = await self.post_invoke(ctx)
192
+ if result is not None:
193
+ ctx = result
194
+
195
+ return ctx.output
196
+
197
+ async def function_middleware_stream(
198
+ self,
199
+ *args: Any,
200
+ call_next: CallNextStream,
201
+ context: FunctionMiddlewareContext,
202
+ **kwargs: Any,
203
+ ) -> AsyncIterator[Any]:
204
+ """Execute middleware hooks around streaming function call.
205
+
206
+ Pre-invoke runs once before streaming starts.
207
+ Post-invoke runs per-chunk as they stream through.
208
+
209
+ Override for custom streaming behavior (e.g., buffering,
210
+ aggregation, chunk filtering).
211
+
212
+ Note: Framework checks ``enabled`` before calling this method.
213
+ You do NOT need to check ``enabled`` yourself.
84
214
 
85
215
  Args:
86
- value: The input value to process
87
- call_next: Callable to invoke the next middleware or function stream
88
- context: Metadata about the function being wrapped
216
+ args: Positional arguments for the function (first arg is typically the input value).
217
+ call_next: Callable to invoke next middleware or target stream.
218
+ context: Static function metadata.
219
+ kwargs: Keyword arguments for the function.
89
220
 
90
221
  Yields:
91
- Chunks from the stream (potentially modified)
92
-
93
- The default implementation forwards to ``call_next`` untouched. Override this
94
- in subclasses to add function-specific preprocessing and chunk transformations.
222
+ Stream chunks (potentially transformed by post_invoke).
95
223
  """
96
- async for chunk in call_next(value):
97
- yield chunk
224
+ # Build invocation context with frozen originals + mutable current
225
+ # output starts as None (pre-invoke phase)
226
+ ctx = InvocationContext(
227
+ function_context=context,
228
+ original_args=args,
229
+ original_kwargs=dict(kwargs),
230
+ modified_args=args,
231
+ modified_kwargs=dict(kwargs),
232
+ output=None,
233
+ )
234
+
235
+ # Pre-invoke transformation (once before streaming)
236
+ result = await self.pre_invoke(ctx)
237
+ if result is not None:
238
+ ctx = result
239
+
240
+ # Stream with per-chunk post-invoke
241
+ async for chunk in call_next(*ctx.modified_args, **ctx.modified_kwargs):
242
+ # Set output for this chunk
243
+ ctx.output = chunk
244
+
245
+ # Post-invoke transformation per chunk
246
+ result = await self.post_invoke(ctx)
247
+ if result is not None:
248
+ ctx = result
249
+
250
+ yield ctx.output
98
251
 
99
252
 
100
253
  class FunctionMiddlewareChain:
101
- """Utility that composes middleware-style callables.
254
+ """Composes middleware into an execution chain.
255
+
256
+ The chain builder checks each middleware's ``enabled`` property.
257
+ Disabled middleware is skipped entirely—no methods are called.
102
258
 
103
- This class builds a chain of middleware that executes in order,
104
- with each middleware able to preprocess inputs, call the next middleware,
105
- and postprocess outputs.
259
+ Execution order:
260
+ - Pre-invoke: first middleware last middleware function
261
+ - Post-invoke: function → last middleware → first middleware
262
+
263
+ Context:
264
+ - FunctionMiddlewareContext contains only static function metadata
265
+ - Original args/kwargs are captured by the orchestration layer
266
+ - Middleware receives InvocationContext with frozen originals and mutable args/output
106
267
  """
107
268
 
108
- def __init__(self, *, middleware: Sequence[Middleware], context: FunctionMiddlewareContext) -> None:
269
+ def __init__(self, *, middleware: Sequence[FunctionMiddleware], context: FunctionMiddlewareContext) -> None:
270
+ """Initialize the middleware chain.
271
+
272
+ Args:
273
+ middleware: Sequence of middleware to chain (order matters)
274
+ context: Static function metadata
275
+ """
109
276
  self._middleware = tuple(middleware)
110
277
  self._context = context
111
278
 
112
279
  def build_single(self, final_call: CallNext) -> CallNext:
113
280
  """Build the middleware chain for single-output invocations.
114
281
 
282
+ Disabled middleware (enabled=False) is skipped entirely.
283
+
115
284
  Args:
116
285
  final_call: The final function to call (the actual function implementation)
117
286
 
@@ -121,18 +290,28 @@ class FunctionMiddlewareChain:
121
290
  call = final_call
122
291
 
123
292
  for mw in reversed(self._middleware):
293
+ # Framework-enforced: skip disabled middleware
294
+ if not mw.enabled:
295
+ continue
296
+
124
297
  call_next = call
125
298
 
126
- async def wrapped(value: Any, *, _middleware: Middleware = mw, _call_next: CallNext = call_next) -> Any:
127
- return await _middleware.middleware_invoke(value, _call_next, self._context)
299
+ async def wrapped(*args: Any,
300
+ _middleware: FunctionMiddleware = mw,
301
+ _call_next: CallNext = call_next,
302
+ _context: FunctionMiddlewareContext = self._context,
303
+ **kwargs: Any) -> Any:
304
+ return await _middleware.middleware_invoke(*args, call_next=_call_next, context=_context, **kwargs)
128
305
 
129
- call = wrapped
306
+ call = wrapped # type: ignore[assignment]
130
307
 
131
308
  return call
132
309
 
133
310
  def build_stream(self, final_call: CallNextStream) -> CallNextStream:
134
311
  """Build the middleware chain for streaming invocations.
135
312
 
313
+ Disabled middleware (enabled=False) is skipped entirely.
314
+
136
315
  Args:
137
316
  final_call: The final function to call (the actual function implementation)
138
317
 
@@ -142,16 +321,21 @@ class FunctionMiddlewareChain:
142
321
  call = final_call
143
322
 
144
323
  for mw in reversed(self._middleware):
324
+ if not mw.enabled:
325
+ continue
326
+
145
327
  call_next = call
146
328
 
147
- async def wrapped(value: Any,
148
- *,
149
- _middleware: Middleware = mw,
150
- _call_next: CallNextStream = call_next) -> AsyncIterator[Any]:
151
- async for chunk in _middleware.middleware_stream(value, _call_next, self._context):
329
+ async def wrapped(*args: Any,
330
+ _middleware: FunctionMiddleware = mw,
331
+ _call_next: CallNextStream = call_next,
332
+ _context: FunctionMiddlewareContext = self._context,
333
+ **kwargs: Any) -> AsyncIterator[Any]:
334
+ stream = _middleware.middleware_stream(*args, call_next=_call_next, context=_context, **kwargs)
335
+ async for chunk in stream:
152
336
  yield chunk
153
337
 
154
- call = wrapped
338
+ call = wrapped # type: ignore[assignment]
155
339
 
156
340
  return call
157
341
 
@@ -0,0 +1,14 @@
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.
@@ -0,0 +1,67 @@
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 logging
17
+
18
+ from nat.builder.builder import Builder
19
+ from nat.middleware.dynamic.dynamic_function_middleware import DynamicFunctionMiddleware
20
+ from nat.middleware.logging.logging_middleware_config import LoggingMiddlewareConfig
21
+ from nat.middleware.middleware import InvocationContext
22
+
23
+ logger = logging.getLogger(__name__)
24
+
25
+
26
+ class LoggingMiddleware(DynamicFunctionMiddleware):
27
+ """Middleware that logs intercepted function inputs/outputs.
28
+
29
+ This middleware extends DynamicFunctionMiddleware to get automatic chain
30
+ orchestration and dynamic discovery features. Custom logic is implemented
31
+ through the pre_invoke and post_invoke hooks.
32
+ """
33
+
34
+ def __init__(self, config: LoggingMiddlewareConfig, builder: Builder):
35
+ """Initialize logging middleware.
36
+
37
+ Args:
38
+ config: Logging middleware configuration
39
+ builder: Workflow builder
40
+ """
41
+ super().__init__(config=config, builder=builder)
42
+ self._config: LoggingMiddlewareConfig = config
43
+
44
+ async def pre_invoke(self, context: InvocationContext) -> InvocationContext | None:
45
+ """Log inputs before function execution.
46
+
47
+ Returns:
48
+ None to pass through unchanged (logging only, no modification)
49
+ """
50
+ log_level = getattr(logging, self._config.log_level.upper(), logging.INFO)
51
+ logger.log(log_level, f"Calling {context.function_context.name} with args: {context.modified_args}")
52
+
53
+ # Log if args were modified by prior middleware
54
+ if context.modified_args != context.original_args:
55
+ logger.log(log_level, f" (original args were: {context.original_args})")
56
+
57
+ return None # Pass through unchanged
58
+
59
+ async def post_invoke(self, context: InvocationContext) -> InvocationContext | None:
60
+ """Log outputs after function execution.
61
+
62
+ Returns:
63
+ None to pass through unchanged (logging only, no modification)
64
+ """
65
+ log_level = getattr(logging, self._config.log_level.upper(), logging.INFO)
66
+ logger.log(log_level, f"Function {context.function_context.name} returned: {context.output}")
67
+ return None # Pass through unchanged
@@ -0,0 +1,28 @@
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
+ from pydantic import Field
17
+
18
+ from nat.middleware.dynamic.dynamic_middleware_config import DynamicMiddlewareConfig
19
+
20
+
21
+ class LoggingMiddlewareConfig(DynamicMiddlewareConfig, name="logging_middleware"):
22
+ """Configuration for logging middleware.
23
+
24
+ Extends DynamicMiddlewareConfig to inherit all dynamic discovery features
25
+ (register_llms, register_workflow_functions, and so on) and the enabled toggle.
26
+ """
27
+
28
+ log_level: str = Field(default="INFO", description="Logging level (DEBUG, INFO, WARNING, ERROR)")
@@ -0,0 +1,33 @@
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
+ from nat.builder.builder import Builder
17
+ from nat.cli.register_workflow import register_middleware
18
+ from nat.middleware.logging.logging_middleware import LoggingMiddleware
19
+ from nat.middleware.logging.logging_middleware_config import LoggingMiddlewareConfig
20
+
21
+
22
+ @register_middleware(config_type=LoggingMiddlewareConfig)
23
+ async def logging_middleware(config: LoggingMiddlewareConfig, builder: Builder):
24
+ """Build a logging middleware from configuration.
25
+
26
+ Args:
27
+ config: The logging middleware configuration
28
+ builder: The workflow builder
29
+
30
+ Yields:
31
+ A configured logging middleware instance
32
+ """
33
+ yield LoggingMiddleware(config=config, builder=builder)