nvidia-nat 1.4.0a20251112__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 +563 -1
  43. nat/builder/child_builder.py +385 -0
  44. nat/builder/component_utils.py +34 -4
  45. nat/builder/context.py +34 -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 +3 -1
  50. nat/builder/front_end.py +1 -1
  51. nat/builder/function.py +113 -5
  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 +5 -3
  61. nat/builder/workflow_builder.py +619 -378
  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 +233 -5
  101. nat/cli/type_registry.py +237 -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 +9 -1
  113. nat/data_models/component_ref.py +45 -1
  114. nat/data_models/config.py +78 -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 +15 -2
  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 +37 -0
  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 +26 -0
  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 +226 -81
  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 +23 -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 +318 -59
  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 +35 -0
  285. nat/middleware/cache/__init__.py +14 -0
  286. nat/middleware/cache/cache_middleware.py +253 -0
  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 +370 -0
  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 +298 -0
  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 +22 -0
  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 +32 -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 +12 -4
  415. nat/retriever/milvus/retriever.py +103 -41
  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 +13 -6
  425. nat/runtime/session.py +458 -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 +2 -2
  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.0a20251112.dist-info → nvidia_nat-1.4.0a20260113.dist-info}/METADATA +46 -15
  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 -151
  483. nat/front_ends/mcp/mcp_front_end_plugin_worker.py +0 -362
  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.0a20251112.dist-info/RECORD +0 -481
  488. nvidia_nat-1.4.0a20251112.dist-info/entry_points.txt +0 -22
  489. {nvidia_nat-1.4.0a20251112.dist-info → nvidia_nat-1.4.0a20260113.dist-info}/WHEEL +0 -0
  490. {nvidia_nat-1.4.0a20251112.dist-info → nvidia_nat-1.4.0a20260113.dist-info}/licenses/LICENSE-3rd-party.txt +0 -0
  491. {nvidia_nat-1.4.0a20251112.dist-info → nvidia_nat-1.4.0a20260113.dist-info}/licenses/LICENSE.md +0 -0
  492. {nvidia_nat-1.4.0a20251112.dist-info → nvidia_nat-1.4.0a20260113.dist-info}/top_level.txt +0 -0
nat/builder/builder.py CHANGED
@@ -1,4 +1,4 @@
1
- # SPDX-FileCopyrightText: Copyright (c) 2024-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
1
+ # SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
2
2
  # SPDX-License-Identifier: Apache-2.0
3
3
  #
4
4
  # Licensed under the Apache License, Version 2.0 (the "License");
@@ -18,6 +18,7 @@ import typing
18
18
  from abc import ABC
19
19
  from abc import abstractmethod
20
20
  from collections.abc import Sequence
21
+ from contextvars import ContextVar
21
22
  from pathlib import Path
22
23
 
23
24
  from nat.authentication.interfaces import AuthProviderBase
@@ -32,29 +33,44 @@ from nat.data_models.component_ref import FunctionGroupRef
32
33
  from nat.data_models.component_ref import FunctionRef
33
34
  from nat.data_models.component_ref import LLMRef
34
35
  from nat.data_models.component_ref import MemoryRef
36
+ from nat.data_models.component_ref import MiddlewareRef
35
37
  from nat.data_models.component_ref import ObjectStoreRef
36
38
  from nat.data_models.component_ref import RetrieverRef
39
+ from nat.data_models.component_ref import TrainerAdapterRef
40
+ from nat.data_models.component_ref import TrainerRef
41
+ from nat.data_models.component_ref import TrajectoryBuilderRef
37
42
  from nat.data_models.component_ref import TTCStrategyRef
38
43
  from nat.data_models.embedder import EmbedderBaseConfig
39
44
  from nat.data_models.evaluator import EvaluatorBaseConfig
45
+ from nat.data_models.finetuning import TrainerAdapterConfig
46
+ from nat.data_models.finetuning import TrainerConfig
47
+ from nat.data_models.finetuning import TrajectoryBuilderConfig
40
48
  from nat.data_models.function import FunctionBaseConfig
41
49
  from nat.data_models.function import FunctionGroupBaseConfig
42
50
  from nat.data_models.function_dependencies import FunctionDependencies
43
51
  from nat.data_models.llm import LLMBaseConfig
44
52
  from nat.data_models.memory import MemoryBaseConfig
53
+ from nat.data_models.middleware import MiddlewareBaseConfig
45
54
  from nat.data_models.object_store import ObjectStoreBaseConfig
46
55
  from nat.data_models.retriever import RetrieverBaseConfig
47
56
  from nat.data_models.ttc_strategy import TTCStrategyBaseConfig
48
57
  from nat.experimental.decorators.experimental_warning_decorator import experimental
49
58
  from nat.experimental.test_time_compute.models.stage_enums import PipelineTypeEnum
50
59
  from nat.experimental.test_time_compute.models.stage_enums import StageTypeEnum
60
+ from nat.finetuning.interfaces.finetuning_runner import Trainer
61
+ from nat.finetuning.interfaces.trainer_adapter import TrainerAdapter
62
+ from nat.finetuning.interfaces.trajectory_builder import TrajectoryBuilder
51
63
  from nat.memory.interfaces import MemoryEditor
64
+ from nat.middleware.middleware import Middleware
52
65
  from nat.object_store.interfaces import ObjectStore
53
66
  from nat.retriever.interface import Retriever
54
67
 
55
68
  if typing.TYPE_CHECKING:
69
+ from nat.builder.sync_builder import SyncBuilder
56
70
  from nat.experimental.test_time_compute.models.strategy_base import StrategyBase
57
71
 
72
+ _current_builder_context: ContextVar["Builder | None"] = ContextVar("current_builder", default=None)
73
+
58
74
 
59
75
  class UserManagerHolder:
60
76
 
@@ -67,71 +83,223 @@ class UserManagerHolder:
67
83
 
68
84
  class Builder(ABC):
69
85
 
86
+ @staticmethod
87
+ def current() -> "Builder":
88
+ """Get the Builder object from the current context.
89
+
90
+ Returns:
91
+ The Builder object stored in the ContextVar, or raises ValueError if not set.
92
+ """
93
+ builder = _current_builder_context.get()
94
+
95
+ if builder is None:
96
+ raise ValueError("Builder not set in context")
97
+
98
+ return builder
99
+
100
+ @property
101
+ @abstractmethod
102
+ def sync_builder(self) -> "SyncBuilder":
103
+ """Get the synchronous version of the builder.
104
+
105
+ Returns:
106
+ The SyncBuilder object (synchronous wrapper).
107
+ """
108
+ pass
109
+
70
110
  @abstractmethod
71
111
  async def add_function(self, name: str | FunctionRef, config: FunctionBaseConfig) -> Function:
112
+ """Add a function to the builder.
113
+
114
+ Args:
115
+ name: The name or reference for the function
116
+ config: The configuration for the function
117
+
118
+ Returns:
119
+ The built function instance
120
+ """
72
121
  pass
73
122
 
74
123
  @abstractmethod
75
124
  async def add_function_group(self, name: str | FunctionGroupRef, config: FunctionGroupBaseConfig) -> FunctionGroup:
125
+ """Add a function group to the builder.
126
+
127
+ Args:
128
+ name: The name or reference for the function group
129
+ config: The configuration for the function group
130
+
131
+ Returns:
132
+ The built function group instance
133
+ """
76
134
  pass
77
135
 
78
136
  @abstractmethod
79
137
  async def get_function(self, name: str | FunctionRef) -> Function:
138
+ """Get a function by name.
139
+
140
+ Args:
141
+ name: The name or reference of the function
142
+
143
+ Returns:
144
+ The built function instance
145
+ """
80
146
  pass
81
147
 
82
148
  @abstractmethod
83
149
  async def get_function_group(self, name: str | FunctionGroupRef) -> FunctionGroup:
150
+ """Get a function group by name.
151
+
152
+ Args:
153
+ name: The name or reference of the function group
154
+
155
+ Returns:
156
+ The built function group instance
157
+ """
84
158
  pass
85
159
 
86
160
  async def get_functions(self, function_names: Sequence[str | FunctionRef]) -> list[Function]:
161
+ """Get multiple functions by name.
162
+
163
+ Args:
164
+ function_names: The names or references of the functions
165
+
166
+ Returns:
167
+ List of built function instances
168
+ """
87
169
  tasks = [self.get_function(name) for name in function_names]
88
170
  return list(await asyncio.gather(*tasks, return_exceptions=False))
89
171
 
90
172
  async def get_function_groups(self, function_group_names: Sequence[str | FunctionGroupRef]) -> list[FunctionGroup]:
173
+ """Get multiple function groups by name.
174
+
175
+ Args:
176
+ function_group_names: The names or references of the function groups
177
+
178
+ Returns:
179
+ List of built function group instances
180
+ """
91
181
  tasks = [self.get_function_group(name) for name in function_group_names]
92
182
  return list(await asyncio.gather(*tasks, return_exceptions=False))
93
183
 
94
184
  @abstractmethod
95
185
  def get_function_config(self, name: str | FunctionRef) -> FunctionBaseConfig:
186
+ """Get the configuration for a function.
187
+
188
+ Args:
189
+ name: The name or reference of the function
190
+
191
+ Returns:
192
+ The configuration for the function
193
+ """
96
194
  pass
97
195
 
98
196
  @abstractmethod
99
197
  def get_function_group_config(self, name: str | FunctionGroupRef) -> FunctionGroupBaseConfig:
198
+ """Get the configuration for a function group.
199
+
200
+ Args:
201
+ name: The name or reference of the function group
202
+
203
+ Returns:
204
+ The configuration for the function group
205
+ """
100
206
  pass
101
207
 
102
208
  @abstractmethod
103
209
  async def set_workflow(self, config: FunctionBaseConfig) -> Function:
210
+ """Set the workflow function.
211
+
212
+ Args:
213
+ config: The configuration for the workflow function
214
+
215
+ Returns:
216
+ The built workflow function instance
217
+ """
104
218
  pass
105
219
 
106
220
  @abstractmethod
107
221
  def get_workflow(self) -> Function:
222
+ """Get the workflow function.
223
+
224
+ Returns:
225
+ The workflow function instance
226
+ """
108
227
  pass
109
228
 
110
229
  @abstractmethod
111
230
  def get_workflow_config(self) -> FunctionBaseConfig:
231
+ """Get the configuration for the workflow.
232
+
233
+ Returns:
234
+ The configuration for the workflow function
235
+ """
112
236
  pass
113
237
 
114
238
  @abstractmethod
115
239
  async def get_tools(self,
116
240
  tool_names: Sequence[str | FunctionRef | FunctionGroupRef],
117
241
  wrapper_type: LLMFrameworkEnum | str) -> list[typing.Any]:
242
+ """Get multiple tools by name wrapped in the specified framework type.
243
+
244
+ Args:
245
+ tool_names: The names or references of the tools (functions or function groups)
246
+ wrapper_type: The LLM framework type to wrap the tools in
247
+
248
+ Returns:
249
+ List of tools wrapped in the specified framework type
250
+ """
118
251
  pass
119
252
 
120
253
  @abstractmethod
121
254
  async def get_tool(self, fn_name: str | FunctionRef, wrapper_type: LLMFrameworkEnum | str) -> typing.Any:
255
+ """Get a tool by name wrapped in the specified framework type.
256
+
257
+ Args:
258
+ fn_name: The name or reference of the tool (function)
259
+ wrapper_type: The LLM framework type to wrap the tool in
260
+
261
+ Returns:
262
+ The tool wrapped in the specified framework type
263
+ """
122
264
  pass
123
265
 
124
266
  @abstractmethod
125
267
  async def add_llm(self, name: str | LLMRef, config: LLMBaseConfig) -> typing.Any:
268
+ """Add an LLM to the builder.
269
+
270
+ Args:
271
+ name: The name or reference for the LLM
272
+ config: The configuration for the LLM
273
+
274
+ Returns:
275
+ The built LLM instance
276
+ """
126
277
  pass
127
278
 
128
279
  @abstractmethod
129
280
  async def get_llm(self, llm_name: str | LLMRef, wrapper_type: LLMFrameworkEnum | str) -> typing.Any:
281
+ """Get an LLM by name wrapped in the specified framework type.
282
+
283
+ Args:
284
+ llm_name: The name or reference of the LLM
285
+ wrapper_type: The LLM framework type to wrap the LLM in
286
+
287
+ Returns:
288
+ The LLM wrapped in the specified framework type
289
+ """
130
290
  pass
131
291
 
132
292
  async def get_llms(self, llm_names: Sequence[str | LLMRef],
133
293
  wrapper_type: LLMFrameworkEnum | str) -> list[typing.Any]:
294
+ """Get multiple LLMs by name wrapped in the specified framework type.
295
+
296
+ Args:
297
+ llm_names: The names or references of the LLMs
298
+ wrapper_type: The LLM framework type to wrap the LLMs in
134
299
 
300
+ Returns:
301
+ List of LLMs wrapped in the specified framework type
302
+ """
135
303
  coros = [self.get_llm(llm_name=n, wrapper_type=wrapper_type) for n in llm_names]
136
304
 
137
305
  llms = await asyncio.gather(*coros, return_exceptions=False)
@@ -140,20 +308,52 @@ class Builder(ABC):
140
308
 
141
309
  @abstractmethod
142
310
  def get_llm_config(self, llm_name: str | LLMRef) -> LLMBaseConfig:
311
+ """Get the configuration for an LLM.
312
+
313
+ Args:
314
+ llm_name: The name or reference of the LLM
315
+
316
+ Returns:
317
+ The configuration for the LLM
318
+ """
143
319
  pass
144
320
 
145
321
  @abstractmethod
146
322
  @experimental(feature_name="Authentication")
147
323
  async def add_auth_provider(self, name: str | AuthenticationRef,
148
324
  config: AuthProviderBaseConfig) -> AuthProviderBase:
325
+ """Add an authentication provider to the builder.
326
+
327
+ Args:
328
+ name: The name or reference for the authentication provider
329
+ config: The configuration for the authentication provider
330
+
331
+ Returns:
332
+ The built authentication provider instance
333
+ """
149
334
  pass
150
335
 
151
336
  @abstractmethod
152
337
  async def get_auth_provider(self, auth_provider_name: str | AuthenticationRef) -> AuthProviderBase:
338
+ """Get an authentication provider by name.
339
+
340
+ Args:
341
+ auth_provider_name: The name or reference of the authentication provider
342
+
343
+ Returns:
344
+ The authentication provider instance
345
+ """
153
346
  pass
154
347
 
155
348
  async def get_auth_providers(self, auth_provider_names: list[str | AuthenticationRef]):
349
+ """Get multiple authentication providers by name.
350
+
351
+ Args:
352
+ auth_provider_names: The names or references of the authentication providers
156
353
 
354
+ Returns:
355
+ List of authentication provider instances
356
+ """
157
357
  coros = [self.get_auth_provider(auth_provider_name=n) for n in auth_provider_names]
158
358
 
159
359
  auth_providers = await asyncio.gather(*coros, return_exceptions=False)
@@ -162,6 +362,15 @@ class Builder(ABC):
162
362
 
163
363
  @abstractmethod
164
364
  async def add_object_store(self, name: str | ObjectStoreRef, config: ObjectStoreBaseConfig) -> ObjectStore:
365
+ """Add an object store to the builder.
366
+
367
+ Args:
368
+ name: The name or reference for the object store
369
+ config: The configuration for the object store
370
+
371
+ Returns:
372
+ The built object store instance
373
+ """
165
374
  pass
166
375
 
167
376
  async def get_object_store_clients(self, object_store_names: Sequence[str | ObjectStoreRef]) -> list[ObjectStore]:
@@ -172,19 +381,49 @@ class Builder(ABC):
172
381
 
173
382
  @abstractmethod
174
383
  async def get_object_store_client(self, object_store_name: str | ObjectStoreRef) -> ObjectStore:
384
+ """Get an object store client by name.
385
+
386
+ Args:
387
+ object_store_name: The name or reference of the object store
388
+
389
+ Returns:
390
+ The object store client instance
391
+ """
175
392
  pass
176
393
 
177
394
  @abstractmethod
178
395
  def get_object_store_config(self, object_store_name: str | ObjectStoreRef) -> ObjectStoreBaseConfig:
396
+ """Get the configuration for an object store.
397
+
398
+ Args:
399
+ object_store_name: The name or reference of the object store
400
+
401
+ Returns:
402
+ The configuration for the object store
403
+ """
179
404
  pass
180
405
 
181
406
  @abstractmethod
182
407
  async def add_embedder(self, name: str | EmbedderRef, config: EmbedderBaseConfig) -> None:
408
+ """Add an embedder to the builder.
409
+
410
+ Args:
411
+ name: The name or reference for the embedder
412
+ config: The configuration for the embedder
413
+ """
183
414
  pass
184
415
 
185
416
  async def get_embedders(self, embedder_names: Sequence[str | EmbedderRef],
186
417
  wrapper_type: LLMFrameworkEnum | str) -> list[typing.Any]:
418
+ """Get multiple embedders by name wrapped in the specified framework type.
419
+
420
+ Args:
421
+ embedder_names: The names or references of the embedders
422
+ wrapper_type: The LLM framework type to wrap the embedders in
187
423
 
424
+ Returns:
425
+ List of embedders wrapped in the specified framework type
426
+ """
188
427
  coros = [self.get_embedder(embedder_name=n, wrapper_type=wrapper_type) for n in embedder_names]
189
428
 
190
429
  embedders = await asyncio.gather(*coros, return_exceptions=False)
@@ -193,14 +432,40 @@ class Builder(ABC):
193
432
 
194
433
  @abstractmethod
195
434
  async def get_embedder(self, embedder_name: str | EmbedderRef, wrapper_type: LLMFrameworkEnum | str) -> typing.Any:
435
+ """Get an embedder by name wrapped in the specified framework type.
436
+
437
+ Args:
438
+ embedder_name: The name or reference of the embedder
439
+ wrapper_type: The LLM framework type to wrap the embedder in
440
+
441
+ Returns:
442
+ The embedder wrapped in the specified framework type
443
+ """
196
444
  pass
197
445
 
198
446
  @abstractmethod
199
447
  def get_embedder_config(self, embedder_name: str | EmbedderRef) -> EmbedderBaseConfig:
448
+ """Get the configuration for an embedder.
449
+
450
+ Args:
451
+ embedder_name: The name or reference of the embedder
452
+
453
+ Returns:
454
+ The configuration for the embedder
455
+ """
200
456
  pass
201
457
 
202
458
  @abstractmethod
203
459
  async def add_memory_client(self, name: str | MemoryRef, config: MemoryBaseConfig) -> MemoryEditor:
460
+ """Add a memory client to the builder.
461
+
462
+ Args:
463
+ name: The name or reference for the memory client
464
+ config: The configuration for the memory client
465
+
466
+ Returns:
467
+ The built memory client instance
468
+ """
204
469
  pass
205
470
 
206
471
  async def get_memory_clients(self, memory_names: Sequence[str | MemoryRef]) -> list[MemoryEditor]:
@@ -219,16 +484,38 @@ class Builder(ABC):
219
484
 
220
485
  @abstractmethod
221
486
  def get_memory_client_config(self, memory_name: str | MemoryRef) -> MemoryBaseConfig:
487
+ """Get the configuration for a memory client.
488
+
489
+ Args:
490
+ memory_name: The name or reference of the memory client
491
+
492
+ Returns:
493
+ The configuration for the memory client
494
+ """
222
495
  pass
223
496
 
224
497
  @abstractmethod
225
498
  async def add_retriever(self, name: str | RetrieverRef, config: RetrieverBaseConfig) -> None:
499
+ """Add a retriever to the builder.
500
+
501
+ Args:
502
+ name: The name or reference for the retriever
503
+ config: The configuration for the retriever
504
+ """
226
505
  pass
227
506
 
228
507
  async def get_retrievers(self,
229
508
  retriever_names: Sequence[str | RetrieverRef],
230
509
  wrapper_type: LLMFrameworkEnum | str | None = None) -> list[Retriever]:
510
+ """Get multiple retrievers by name.
231
511
 
512
+ Args:
513
+ retriever_names: The names or references of the retrievers
514
+ wrapper_type: Optional LLM framework type to wrap the retrievers in
515
+
516
+ Returns:
517
+ List of retriever instances
518
+ """
232
519
  tasks = [self.get_retriever(n, wrapper_type=wrapper_type) for n in retriever_names]
233
520
 
234
521
  retrievers = await asyncio.gather(*tasks, return_exceptions=False)
@@ -252,15 +539,159 @@ class Builder(ABC):
252
539
  async def get_retriever(self,
253
540
  retriever_name: str | RetrieverRef,
254
541
  wrapper_type: LLMFrameworkEnum | str | None = None) -> typing.Any:
542
+ """Get a retriever by name.
543
+
544
+ Args:
545
+ retriever_name: The name or reference of the retriever
546
+ wrapper_type: Optional LLM framework type to wrap the retriever in
547
+
548
+ Returns:
549
+ The retriever instance, optionally wrapped in the specified framework type
550
+ """
255
551
  pass
256
552
 
257
553
  @abstractmethod
258
554
  async def get_retriever_config(self, retriever_name: str | RetrieverRef) -> RetrieverBaseConfig:
555
+ """Get the configuration for a retriever.
556
+
557
+ Args:
558
+ retriever_name: The name or reference of the retriever
559
+
560
+ Returns:
561
+ The configuration for the retriever
562
+ """
563
+ pass
564
+
565
+ @abstractmethod
566
+ @experimental(feature_name="Finetuning")
567
+ async def add_trainer(self, name: str | TrainerRef, config: TrainerConfig) -> Trainer:
568
+ """Add a trainer to the builder.
569
+
570
+ Args:
571
+ name: The name or reference for the trainer
572
+ config: The configuration for the trainer
573
+
574
+ Returns:
575
+ The built trainer instance
576
+ """
577
+ pass
578
+
579
+ @abstractmethod
580
+ @experimental(feature_name="Finetuning")
581
+ async def add_trainer_adapter(self, name: str | TrainerAdapterRef, config: TrainerAdapterConfig) -> TrainerAdapter:
582
+ """Add a trainer adapter to the builder.
583
+
584
+ Args:
585
+ name: The name or reference for the trainer adapter
586
+ config: The configuration for the trainer adapter
587
+
588
+ Returns:
589
+ The built trainer adapter instance
590
+ """
591
+ pass
592
+
593
+ @abstractmethod
594
+ @experimental(feature_name="Finetuning")
595
+ async def add_trajectory_builder(self, name: str | TrajectoryBuilderRef,
596
+ config: TrajectoryBuilderConfig) -> TrajectoryBuilder:
597
+ """Add a trajectory builder to the builder.
598
+
599
+ Args:
600
+ name: The name or reference for the trajectory builder
601
+ config: The configuration for the trajectory builder
602
+
603
+ Returns:
604
+ The built trajectory builder instance
605
+ """
606
+ pass
607
+
608
+ @abstractmethod
609
+ async def get_trainer(self,
610
+ trainer_name: str | TrainerRef,
611
+ trajectory_builder: TrajectoryBuilder,
612
+ trainer_adapter: TrainerAdapter) -> Trainer:
613
+ """Get a trainer by name with the specified trajectory builder and trainer adapter.
614
+
615
+ Args:
616
+ trainer_name: The name or reference of the trainer
617
+ trajectory_builder: The trajectory builder instance
618
+ trainer_adapter: The trainer adapter instance
619
+
620
+ Returns:
621
+ The trainer instance
622
+ """
623
+ pass
624
+
625
+ @abstractmethod
626
+ async def get_trainer_adapter(self, trainer_adapter_name: str | TrainerAdapterRef) -> TrainerAdapter:
627
+ """Get a trainer adapter by name.
628
+
629
+ Args:
630
+ trainer_adapter_name: The name or reference of the trainer adapter
631
+
632
+ Returns:
633
+ The trainer adapter instance
634
+ """
635
+ pass
636
+
637
+ @abstractmethod
638
+ async def get_trajectory_builder(self, trajectory_builder_name: str | TrajectoryBuilderRef) -> TrajectoryBuilder:
639
+ """Get a trajectory builder by name.
640
+
641
+ Args:
642
+ trajectory_builder_name: The name or reference of the trajectory builder
643
+
644
+ Returns:
645
+ The trajectory builder instance
646
+ """
647
+ pass
648
+
649
+ @abstractmethod
650
+ async def get_trainer_config(self, trainer_name: str | TrainerRef) -> TrainerConfig:
651
+ """Get the configuration for a trainer.
652
+
653
+ Args:
654
+ trainer_name: The name or reference of the trainer
655
+
656
+ Returns:
657
+ The configuration for the trainer
658
+ """
659
+ pass
660
+
661
+ @abstractmethod
662
+ async def get_trainer_adapter_config(self, trainer_adapter_name: str | TrainerAdapterRef) -> TrainerAdapterConfig:
663
+ """Get the configuration for a trainer adapter.
664
+
665
+ Args:
666
+ trainer_adapter_name: The name or reference of the trainer adapter
667
+
668
+ Returns:
669
+ The configuration for the trainer adapter
670
+ """
671
+ pass
672
+
673
+ @abstractmethod
674
+ async def get_trajectory_builder_config(
675
+ self, trajectory_builder_name: str | TrajectoryBuilderRef) -> (TrajectoryBuilderConfig):
676
+ """Get the configuration for a trajectory builder.
677
+
678
+ Args:
679
+ trajectory_builder_name: The name or reference of the trajectory builder
680
+
681
+ Returns:
682
+ The configuration for the trajectory builder
683
+ """
259
684
  pass
260
685
 
261
686
  @abstractmethod
262
687
  @experimental(feature_name="TTC")
263
688
  async def add_ttc_strategy(self, name: str | TTCStrategyRef, config: TTCStrategyBaseConfig):
689
+ """Add a test-time compute strategy to the builder.
690
+
691
+ Args:
692
+ name: The name or reference for the TTC strategy
693
+ config: The configuration for the TTC strategy
694
+ """
264
695
  pass
265
696
 
266
697
  @abstractmethod
@@ -268,6 +699,16 @@ class Builder(ABC):
268
699
  strategy_name: str | TTCStrategyRef,
269
700
  pipeline_type: PipelineTypeEnum,
270
701
  stage_type: StageTypeEnum) -> "StrategyBase":
702
+ """Get a test-time compute strategy by name.
703
+
704
+ Args:
705
+ strategy_name: The name or reference of the TTC strategy
706
+ pipeline_type: The pipeline type for the strategy
707
+ stage_type: The stage type for the strategy
708
+
709
+ Returns:
710
+ The TTC strategy instance
711
+ """
271
712
  pass
272
713
 
273
714
  @abstractmethod
@@ -275,43 +716,164 @@ class Builder(ABC):
275
716
  strategy_name: str | TTCStrategyRef,
276
717
  pipeline_type: PipelineTypeEnum,
277
718
  stage_type: StageTypeEnum) -> TTCStrategyBaseConfig:
719
+ """Get the configuration for a test-time compute strategy.
720
+
721
+ Args:
722
+ strategy_name: The name or reference of the TTC strategy
723
+ pipeline_type: The pipeline type for the strategy
724
+ stage_type: The stage type for the strategy
725
+
726
+ Returns:
727
+ The configuration for the TTC strategy
728
+ """
278
729
  pass
279
730
 
280
731
  @abstractmethod
281
732
  def get_user_manager(self) -> UserManagerHolder:
733
+ """Get the user manager holder.
734
+
735
+ Returns:
736
+ The user manager holder instance
737
+ """
282
738
  pass
283
739
 
284
740
  @abstractmethod
285
741
  def get_function_dependencies(self, fn_name: str) -> FunctionDependencies:
742
+ """Get the dependencies for a function.
743
+
744
+ Args:
745
+ fn_name: The name of the function
746
+
747
+ Returns:
748
+ The function dependencies
749
+ """
286
750
  pass
287
751
 
288
752
  @abstractmethod
289
753
  def get_function_group_dependencies(self, fn_name: str) -> FunctionDependencies:
754
+ """Get the dependencies for a function group.
755
+
756
+ Args:
757
+ fn_name: The name of the function group
758
+
759
+ Returns:
760
+ The function group dependencies
761
+ """
762
+ pass
763
+
764
+ @abstractmethod
765
+ async def add_middleware(self, name: str | MiddlewareRef, config: MiddlewareBaseConfig) -> Middleware:
766
+ """Add middleware to the builder.
767
+
768
+ Args:
769
+ name: The name or reference for the middleware
770
+ config: The configuration for the middleware
771
+
772
+ Returns:
773
+ The built middleware instance
774
+ """
775
+ pass
776
+
777
+ @abstractmethod
778
+ async def get_middleware(self, middleware_name: str | MiddlewareRef) -> Middleware:
779
+ """Get built middleware by name.
780
+
781
+ Args:
782
+ middleware_name: The name or reference of the middleware
783
+
784
+ Returns:
785
+ The built middleware instance
786
+ """
290
787
  pass
291
788
 
789
+ @abstractmethod
790
+ def get_middleware_config(self, middleware_name: str | MiddlewareRef) -> MiddlewareBaseConfig:
791
+ """Get the configuration for middleware.
792
+
793
+ Args:
794
+ middleware_name: The name or reference of the middleware
795
+
796
+ Returns:
797
+ The configuration for the middleware
798
+ """
799
+ pass
800
+
801
+ async def get_middleware_list(self, middleware_names: Sequence[str | MiddlewareRef]) -> list[Middleware]:
802
+ """Get multiple middleware by name.
803
+
804
+ Args:
805
+ middleware_names: The names or references of the middleware
806
+
807
+ Returns:
808
+ List of built middleware instances
809
+ """
810
+ tasks = [self.get_middleware(name) for name in middleware_names]
811
+ return list(await asyncio.gather(*tasks, return_exceptions=False))
812
+
292
813
 
293
814
  class EvalBuilder(ABC):
815
+ """Abstract base class for evaluation builder functionality."""
294
816
 
295
817
  @abstractmethod
296
818
  async def add_evaluator(self, name: str, config: EvaluatorBaseConfig):
819
+ """Add an evaluator to the builder.
820
+
821
+ Args:
822
+ name: The name for the evaluator
823
+ config: The configuration for the evaluator
824
+ """
297
825
  pass
298
826
 
299
827
  @abstractmethod
300
828
  def get_evaluator(self, evaluator_name: str) -> typing.Any:
829
+ """Get an evaluator by name.
830
+
831
+ Args:
832
+ evaluator_name: The name of the evaluator
833
+
834
+ Returns:
835
+ The evaluator instance
836
+ """
301
837
  pass
302
838
 
303
839
  @abstractmethod
304
840
  def get_evaluator_config(self, evaluator_name: str) -> EvaluatorBaseConfig:
841
+ """Get the configuration for an evaluator.
842
+
843
+ Args:
844
+ evaluator_name: The name of the evaluator
845
+
846
+ Returns:
847
+ The configuration for the evaluator
848
+ """
305
849
  pass
306
850
 
307
851
  @abstractmethod
308
852
  def get_max_concurrency(self) -> int:
853
+ """Get the maximum concurrency for evaluation.
854
+
855
+ Returns:
856
+ The maximum concurrency value
857
+ """
309
858
  pass
310
859
 
311
860
  @abstractmethod
312
861
  def get_output_dir(self) -> Path:
862
+ """Get the output directory for evaluation results.
863
+
864
+ Returns:
865
+ The output directory path
866
+ """
313
867
  pass
314
868
 
315
869
  @abstractmethod
316
870
  async def get_all_tools(self, wrapper_type: LLMFrameworkEnum | str) -> list[typing.Any]:
871
+ """Get all tools wrapped in the specified framework type.
872
+
873
+ Args:
874
+ wrapper_type: The LLM framework type to wrap the tools in
875
+
876
+ Returns:
877
+ List of all tools wrapped in the specified framework type
878
+ """
317
879
  pass