nvidia-nat 1.2.0rc5__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 (435) hide show
  1. aiq/agent/__init__.py +0 -0
  2. aiq/agent/base.py +239 -0
  3. aiq/agent/dual_node.py +67 -0
  4. aiq/agent/react_agent/__init__.py +0 -0
  5. aiq/agent/react_agent/agent.py +355 -0
  6. aiq/agent/react_agent/output_parser.py +104 -0
  7. aiq/agent/react_agent/prompt.py +41 -0
  8. aiq/agent/react_agent/register.py +149 -0
  9. aiq/agent/reasoning_agent/__init__.py +0 -0
  10. aiq/agent/reasoning_agent/reasoning_agent.py +225 -0
  11. aiq/agent/register.py +23 -0
  12. aiq/agent/rewoo_agent/__init__.py +0 -0
  13. aiq/agent/rewoo_agent/agent.py +411 -0
  14. aiq/agent/rewoo_agent/prompt.py +108 -0
  15. aiq/agent/rewoo_agent/register.py +158 -0
  16. aiq/agent/tool_calling_agent/__init__.py +0 -0
  17. aiq/agent/tool_calling_agent/agent.py +119 -0
  18. aiq/agent/tool_calling_agent/register.py +106 -0
  19. aiq/authentication/__init__.py +14 -0
  20. aiq/authentication/api_key/__init__.py +14 -0
  21. aiq/authentication/api_key/api_key_auth_provider.py +96 -0
  22. aiq/authentication/api_key/api_key_auth_provider_config.py +124 -0
  23. aiq/authentication/api_key/register.py +26 -0
  24. aiq/authentication/exceptions/__init__.py +14 -0
  25. aiq/authentication/exceptions/api_key_exceptions.py +38 -0
  26. aiq/authentication/http_basic_auth/__init__.py +0 -0
  27. aiq/authentication/http_basic_auth/http_basic_auth_provider.py +81 -0
  28. aiq/authentication/http_basic_auth/register.py +30 -0
  29. aiq/authentication/interfaces.py +93 -0
  30. aiq/authentication/oauth2/__init__.py +14 -0
  31. aiq/authentication/oauth2/oauth2_auth_code_flow_provider.py +107 -0
  32. aiq/authentication/oauth2/oauth2_auth_code_flow_provider_config.py +39 -0
  33. aiq/authentication/oauth2/register.py +25 -0
  34. aiq/authentication/register.py +21 -0
  35. aiq/builder/__init__.py +0 -0
  36. aiq/builder/builder.py +285 -0
  37. aiq/builder/component_utils.py +316 -0
  38. aiq/builder/context.py +264 -0
  39. aiq/builder/embedder.py +24 -0
  40. aiq/builder/eval_builder.py +161 -0
  41. aiq/builder/evaluator.py +29 -0
  42. aiq/builder/framework_enum.py +24 -0
  43. aiq/builder/front_end.py +73 -0
  44. aiq/builder/function.py +344 -0
  45. aiq/builder/function_base.py +380 -0
  46. aiq/builder/function_info.py +627 -0
  47. aiq/builder/intermediate_step_manager.py +174 -0
  48. aiq/builder/llm.py +25 -0
  49. aiq/builder/retriever.py +25 -0
  50. aiq/builder/user_interaction_manager.py +74 -0
  51. aiq/builder/workflow.py +148 -0
  52. aiq/builder/workflow_builder.py +1117 -0
  53. aiq/cli/__init__.py +14 -0
  54. aiq/cli/cli_utils/__init__.py +0 -0
  55. aiq/cli/cli_utils/config_override.py +231 -0
  56. aiq/cli/cli_utils/validation.py +37 -0
  57. aiq/cli/commands/__init__.py +0 -0
  58. aiq/cli/commands/configure/__init__.py +0 -0
  59. aiq/cli/commands/configure/channel/__init__.py +0 -0
  60. aiq/cli/commands/configure/channel/add.py +28 -0
  61. aiq/cli/commands/configure/channel/channel.py +36 -0
  62. aiq/cli/commands/configure/channel/remove.py +30 -0
  63. aiq/cli/commands/configure/channel/update.py +30 -0
  64. aiq/cli/commands/configure/configure.py +33 -0
  65. aiq/cli/commands/evaluate.py +139 -0
  66. aiq/cli/commands/info/__init__.py +14 -0
  67. aiq/cli/commands/info/info.py +39 -0
  68. aiq/cli/commands/info/list_channels.py +32 -0
  69. aiq/cli/commands/info/list_components.py +129 -0
  70. aiq/cli/commands/info/list_mcp.py +213 -0
  71. aiq/cli/commands/registry/__init__.py +14 -0
  72. aiq/cli/commands/registry/publish.py +88 -0
  73. aiq/cli/commands/registry/pull.py +118 -0
  74. aiq/cli/commands/registry/registry.py +38 -0
  75. aiq/cli/commands/registry/remove.py +108 -0
  76. aiq/cli/commands/registry/search.py +155 -0
  77. aiq/cli/commands/sizing/__init__.py +14 -0
  78. aiq/cli/commands/sizing/calc.py +297 -0
  79. aiq/cli/commands/sizing/sizing.py +27 -0
  80. aiq/cli/commands/start.py +246 -0
  81. aiq/cli/commands/uninstall.py +81 -0
  82. aiq/cli/commands/validate.py +47 -0
  83. aiq/cli/commands/workflow/__init__.py +14 -0
  84. aiq/cli/commands/workflow/templates/__init__.py.j2 +0 -0
  85. aiq/cli/commands/workflow/templates/config.yml.j2 +16 -0
  86. aiq/cli/commands/workflow/templates/pyproject.toml.j2 +22 -0
  87. aiq/cli/commands/workflow/templates/register.py.j2 +5 -0
  88. aiq/cli/commands/workflow/templates/workflow.py.j2 +36 -0
  89. aiq/cli/commands/workflow/workflow.py +37 -0
  90. aiq/cli/commands/workflow/workflow_commands.py +313 -0
  91. aiq/cli/entrypoint.py +135 -0
  92. aiq/cli/main.py +44 -0
  93. aiq/cli/register_workflow.py +488 -0
  94. aiq/cli/type_registry.py +1000 -0
  95. aiq/data_models/__init__.py +14 -0
  96. aiq/data_models/api_server.py +694 -0
  97. aiq/data_models/authentication.py +231 -0
  98. aiq/data_models/common.py +171 -0
  99. aiq/data_models/component.py +54 -0
  100. aiq/data_models/component_ref.py +168 -0
  101. aiq/data_models/config.py +406 -0
  102. aiq/data_models/dataset_handler.py +123 -0
  103. aiq/data_models/discovery_metadata.py +335 -0
  104. aiq/data_models/embedder.py +27 -0
  105. aiq/data_models/evaluate.py +127 -0
  106. aiq/data_models/evaluator.py +26 -0
  107. aiq/data_models/front_end.py +26 -0
  108. aiq/data_models/function.py +30 -0
  109. aiq/data_models/function_dependencies.py +72 -0
  110. aiq/data_models/interactive.py +246 -0
  111. aiq/data_models/intermediate_step.py +302 -0
  112. aiq/data_models/invocation_node.py +38 -0
  113. aiq/data_models/llm.py +27 -0
  114. aiq/data_models/logging.py +26 -0
  115. aiq/data_models/memory.py +27 -0
  116. aiq/data_models/object_store.py +44 -0
  117. aiq/data_models/profiler.py +54 -0
  118. aiq/data_models/registry_handler.py +26 -0
  119. aiq/data_models/retriever.py +30 -0
  120. aiq/data_models/retry_mixin.py +35 -0
  121. aiq/data_models/span.py +187 -0
  122. aiq/data_models/step_adaptor.py +64 -0
  123. aiq/data_models/streaming.py +33 -0
  124. aiq/data_models/swe_bench_model.py +54 -0
  125. aiq/data_models/telemetry_exporter.py +26 -0
  126. aiq/data_models/ttc_strategy.py +30 -0
  127. aiq/embedder/__init__.py +0 -0
  128. aiq/embedder/langchain_client.py +41 -0
  129. aiq/embedder/nim_embedder.py +59 -0
  130. aiq/embedder/openai_embedder.py +43 -0
  131. aiq/embedder/register.py +24 -0
  132. aiq/eval/__init__.py +14 -0
  133. aiq/eval/config.py +60 -0
  134. aiq/eval/dataset_handler/__init__.py +0 -0
  135. aiq/eval/dataset_handler/dataset_downloader.py +106 -0
  136. aiq/eval/dataset_handler/dataset_filter.py +52 -0
  137. aiq/eval/dataset_handler/dataset_handler.py +254 -0
  138. aiq/eval/evaluate.py +506 -0
  139. aiq/eval/evaluator/__init__.py +14 -0
  140. aiq/eval/evaluator/base_evaluator.py +73 -0
  141. aiq/eval/evaluator/evaluator_model.py +45 -0
  142. aiq/eval/intermediate_step_adapter.py +99 -0
  143. aiq/eval/rag_evaluator/__init__.py +0 -0
  144. aiq/eval/rag_evaluator/evaluate.py +178 -0
  145. aiq/eval/rag_evaluator/register.py +143 -0
  146. aiq/eval/register.py +23 -0
  147. aiq/eval/remote_workflow.py +133 -0
  148. aiq/eval/runners/__init__.py +14 -0
  149. aiq/eval/runners/config.py +39 -0
  150. aiq/eval/runners/multi_eval_runner.py +54 -0
  151. aiq/eval/runtime_event_subscriber.py +52 -0
  152. aiq/eval/swe_bench_evaluator/__init__.py +0 -0
  153. aiq/eval/swe_bench_evaluator/evaluate.py +215 -0
  154. aiq/eval/swe_bench_evaluator/register.py +36 -0
  155. aiq/eval/trajectory_evaluator/__init__.py +0 -0
  156. aiq/eval/trajectory_evaluator/evaluate.py +75 -0
  157. aiq/eval/trajectory_evaluator/register.py +40 -0
  158. aiq/eval/tunable_rag_evaluator/__init__.py +0 -0
  159. aiq/eval/tunable_rag_evaluator/evaluate.py +245 -0
  160. aiq/eval/tunable_rag_evaluator/register.py +52 -0
  161. aiq/eval/usage_stats.py +41 -0
  162. aiq/eval/utils/__init__.py +0 -0
  163. aiq/eval/utils/output_uploader.py +140 -0
  164. aiq/eval/utils/tqdm_position_registry.py +40 -0
  165. aiq/eval/utils/weave_eval.py +184 -0
  166. aiq/experimental/__init__.py +0 -0
  167. aiq/experimental/decorators/__init__.py +0 -0
  168. aiq/experimental/decorators/experimental_warning_decorator.py +130 -0
  169. aiq/experimental/test_time_compute/__init__.py +0 -0
  170. aiq/experimental/test_time_compute/editing/__init__.py +0 -0
  171. aiq/experimental/test_time_compute/editing/iterative_plan_refinement_editor.py +147 -0
  172. aiq/experimental/test_time_compute/editing/llm_as_a_judge_editor.py +204 -0
  173. aiq/experimental/test_time_compute/editing/motivation_aware_summarization.py +107 -0
  174. aiq/experimental/test_time_compute/functions/__init__.py +0 -0
  175. aiq/experimental/test_time_compute/functions/execute_score_select_function.py +105 -0
  176. aiq/experimental/test_time_compute/functions/its_tool_orchestration_function.py +205 -0
  177. aiq/experimental/test_time_compute/functions/its_tool_wrapper_function.py +146 -0
  178. aiq/experimental/test_time_compute/functions/plan_select_execute_function.py +224 -0
  179. aiq/experimental/test_time_compute/models/__init__.py +0 -0
  180. aiq/experimental/test_time_compute/models/editor_config.py +132 -0
  181. aiq/experimental/test_time_compute/models/scoring_config.py +112 -0
  182. aiq/experimental/test_time_compute/models/search_config.py +120 -0
  183. aiq/experimental/test_time_compute/models/selection_config.py +154 -0
  184. aiq/experimental/test_time_compute/models/stage_enums.py +43 -0
  185. aiq/experimental/test_time_compute/models/strategy_base.py +66 -0
  186. aiq/experimental/test_time_compute/models/tool_use_config.py +41 -0
  187. aiq/experimental/test_time_compute/models/ttc_item.py +48 -0
  188. aiq/experimental/test_time_compute/register.py +36 -0
  189. aiq/experimental/test_time_compute/scoring/__init__.py +0 -0
  190. aiq/experimental/test_time_compute/scoring/llm_based_agent_scorer.py +168 -0
  191. aiq/experimental/test_time_compute/scoring/llm_based_plan_scorer.py +168 -0
  192. aiq/experimental/test_time_compute/scoring/motivation_aware_scorer.py +111 -0
  193. aiq/experimental/test_time_compute/search/__init__.py +0 -0
  194. aiq/experimental/test_time_compute/search/multi_llm_planner.py +128 -0
  195. aiq/experimental/test_time_compute/search/multi_query_retrieval_search.py +122 -0
  196. aiq/experimental/test_time_compute/search/single_shot_multi_plan_planner.py +128 -0
  197. aiq/experimental/test_time_compute/selection/__init__.py +0 -0
  198. aiq/experimental/test_time_compute/selection/best_of_n_selector.py +63 -0
  199. aiq/experimental/test_time_compute/selection/llm_based_agent_output_selector.py +131 -0
  200. aiq/experimental/test_time_compute/selection/llm_based_output_merging_selector.py +159 -0
  201. aiq/experimental/test_time_compute/selection/llm_based_plan_selector.py +128 -0
  202. aiq/experimental/test_time_compute/selection/threshold_selector.py +58 -0
  203. aiq/front_ends/__init__.py +14 -0
  204. aiq/front_ends/console/__init__.py +14 -0
  205. aiq/front_ends/console/authentication_flow_handler.py +233 -0
  206. aiq/front_ends/console/console_front_end_config.py +32 -0
  207. aiq/front_ends/console/console_front_end_plugin.py +96 -0
  208. aiq/front_ends/console/register.py +25 -0
  209. aiq/front_ends/cron/__init__.py +14 -0
  210. aiq/front_ends/fastapi/__init__.py +14 -0
  211. aiq/front_ends/fastapi/auth_flow_handlers/__init__.py +0 -0
  212. aiq/front_ends/fastapi/auth_flow_handlers/http_flow_handler.py +27 -0
  213. aiq/front_ends/fastapi/auth_flow_handlers/websocket_flow_handler.py +107 -0
  214. aiq/front_ends/fastapi/fastapi_front_end_config.py +234 -0
  215. aiq/front_ends/fastapi/fastapi_front_end_controller.py +68 -0
  216. aiq/front_ends/fastapi/fastapi_front_end_plugin.py +116 -0
  217. aiq/front_ends/fastapi/fastapi_front_end_plugin_worker.py +1092 -0
  218. aiq/front_ends/fastapi/html_snippets/__init__.py +14 -0
  219. aiq/front_ends/fastapi/html_snippets/auth_code_grant_success.py +35 -0
  220. aiq/front_ends/fastapi/intermediate_steps_subscriber.py +80 -0
  221. aiq/front_ends/fastapi/job_store.py +183 -0
  222. aiq/front_ends/fastapi/main.py +72 -0
  223. aiq/front_ends/fastapi/message_handler.py +298 -0
  224. aiq/front_ends/fastapi/message_validator.py +345 -0
  225. aiq/front_ends/fastapi/register.py +25 -0
  226. aiq/front_ends/fastapi/response_helpers.py +195 -0
  227. aiq/front_ends/fastapi/step_adaptor.py +321 -0
  228. aiq/front_ends/mcp/__init__.py +14 -0
  229. aiq/front_ends/mcp/mcp_front_end_config.py +32 -0
  230. aiq/front_ends/mcp/mcp_front_end_plugin.py +93 -0
  231. aiq/front_ends/mcp/register.py +27 -0
  232. aiq/front_ends/mcp/tool_converter.py +242 -0
  233. aiq/front_ends/register.py +22 -0
  234. aiq/front_ends/simple_base/__init__.py +14 -0
  235. aiq/front_ends/simple_base/simple_front_end_plugin_base.py +54 -0
  236. aiq/llm/__init__.py +0 -0
  237. aiq/llm/aws_bedrock_llm.py +57 -0
  238. aiq/llm/nim_llm.py +46 -0
  239. aiq/llm/openai_llm.py +46 -0
  240. aiq/llm/register.py +23 -0
  241. aiq/llm/utils/__init__.py +14 -0
  242. aiq/llm/utils/env_config_value.py +94 -0
  243. aiq/llm/utils/error.py +17 -0
  244. aiq/memory/__init__.py +20 -0
  245. aiq/memory/interfaces.py +183 -0
  246. aiq/memory/models.py +112 -0
  247. aiq/meta/module_to_distro.json +3 -0
  248. aiq/meta/pypi.md +58 -0
  249. aiq/object_store/__init__.py +20 -0
  250. aiq/object_store/in_memory_object_store.py +76 -0
  251. aiq/object_store/interfaces.py +84 -0
  252. aiq/object_store/models.py +36 -0
  253. aiq/object_store/register.py +20 -0
  254. aiq/observability/__init__.py +14 -0
  255. aiq/observability/exporter/__init__.py +14 -0
  256. aiq/observability/exporter/base_exporter.py +449 -0
  257. aiq/observability/exporter/exporter.py +78 -0
  258. aiq/observability/exporter/file_exporter.py +33 -0
  259. aiq/observability/exporter/processing_exporter.py +322 -0
  260. aiq/observability/exporter/raw_exporter.py +52 -0
  261. aiq/observability/exporter/span_exporter.py +265 -0
  262. aiq/observability/exporter_manager.py +335 -0
  263. aiq/observability/mixin/__init__.py +14 -0
  264. aiq/observability/mixin/batch_config_mixin.py +26 -0
  265. aiq/observability/mixin/collector_config_mixin.py +23 -0
  266. aiq/observability/mixin/file_mixin.py +288 -0
  267. aiq/observability/mixin/file_mode.py +23 -0
  268. aiq/observability/mixin/resource_conflict_mixin.py +134 -0
  269. aiq/observability/mixin/serialize_mixin.py +61 -0
  270. aiq/observability/mixin/type_introspection_mixin.py +183 -0
  271. aiq/observability/processor/__init__.py +14 -0
  272. aiq/observability/processor/batching_processor.py +310 -0
  273. aiq/observability/processor/callback_processor.py +42 -0
  274. aiq/observability/processor/intermediate_step_serializer.py +28 -0
  275. aiq/observability/processor/processor.py +71 -0
  276. aiq/observability/register.py +96 -0
  277. aiq/observability/utils/__init__.py +14 -0
  278. aiq/observability/utils/dict_utils.py +236 -0
  279. aiq/observability/utils/time_utils.py +31 -0
  280. aiq/plugins/.namespace +1 -0
  281. aiq/profiler/__init__.py +0 -0
  282. aiq/profiler/calc/__init__.py +14 -0
  283. aiq/profiler/calc/calc_runner.py +627 -0
  284. aiq/profiler/calc/calculations.py +288 -0
  285. aiq/profiler/calc/data_models.py +188 -0
  286. aiq/profiler/calc/plot.py +345 -0
  287. aiq/profiler/callbacks/__init__.py +0 -0
  288. aiq/profiler/callbacks/agno_callback_handler.py +295 -0
  289. aiq/profiler/callbacks/base_callback_class.py +20 -0
  290. aiq/profiler/callbacks/langchain_callback_handler.py +290 -0
  291. aiq/profiler/callbacks/llama_index_callback_handler.py +205 -0
  292. aiq/profiler/callbacks/semantic_kernel_callback_handler.py +238 -0
  293. aiq/profiler/callbacks/token_usage_base_model.py +27 -0
  294. aiq/profiler/data_frame_row.py +51 -0
  295. aiq/profiler/data_models.py +24 -0
  296. aiq/profiler/decorators/__init__.py +0 -0
  297. aiq/profiler/decorators/framework_wrapper.py +131 -0
  298. aiq/profiler/decorators/function_tracking.py +254 -0
  299. aiq/profiler/forecasting/__init__.py +0 -0
  300. aiq/profiler/forecasting/config.py +18 -0
  301. aiq/profiler/forecasting/model_trainer.py +75 -0
  302. aiq/profiler/forecasting/models/__init__.py +22 -0
  303. aiq/profiler/forecasting/models/forecasting_base_model.py +40 -0
  304. aiq/profiler/forecasting/models/linear_model.py +196 -0
  305. aiq/profiler/forecasting/models/random_forest_regressor.py +268 -0
  306. aiq/profiler/inference_metrics_model.py +28 -0
  307. aiq/profiler/inference_optimization/__init__.py +0 -0
  308. aiq/profiler/inference_optimization/bottleneck_analysis/__init__.py +0 -0
  309. aiq/profiler/inference_optimization/bottleneck_analysis/nested_stack_analysis.py +460 -0
  310. aiq/profiler/inference_optimization/bottleneck_analysis/simple_stack_analysis.py +258 -0
  311. aiq/profiler/inference_optimization/data_models.py +386 -0
  312. aiq/profiler/inference_optimization/experimental/__init__.py +0 -0
  313. aiq/profiler/inference_optimization/experimental/concurrency_spike_analysis.py +468 -0
  314. aiq/profiler/inference_optimization/experimental/prefix_span_analysis.py +405 -0
  315. aiq/profiler/inference_optimization/llm_metrics.py +212 -0
  316. aiq/profiler/inference_optimization/prompt_caching.py +163 -0
  317. aiq/profiler/inference_optimization/token_uniqueness.py +107 -0
  318. aiq/profiler/inference_optimization/workflow_runtimes.py +72 -0
  319. aiq/profiler/intermediate_property_adapter.py +102 -0
  320. aiq/profiler/profile_runner.py +473 -0
  321. aiq/profiler/utils.py +184 -0
  322. aiq/registry_handlers/__init__.py +0 -0
  323. aiq/registry_handlers/local/__init__.py +0 -0
  324. aiq/registry_handlers/local/local_handler.py +176 -0
  325. aiq/registry_handlers/local/register_local.py +37 -0
  326. aiq/registry_handlers/metadata_factory.py +60 -0
  327. aiq/registry_handlers/package_utils.py +567 -0
  328. aiq/registry_handlers/pypi/__init__.py +0 -0
  329. aiq/registry_handlers/pypi/pypi_handler.py +251 -0
  330. aiq/registry_handlers/pypi/register_pypi.py +40 -0
  331. aiq/registry_handlers/register.py +21 -0
  332. aiq/registry_handlers/registry_handler_base.py +157 -0
  333. aiq/registry_handlers/rest/__init__.py +0 -0
  334. aiq/registry_handlers/rest/register_rest.py +56 -0
  335. aiq/registry_handlers/rest/rest_handler.py +237 -0
  336. aiq/registry_handlers/schemas/__init__.py +0 -0
  337. aiq/registry_handlers/schemas/headers.py +42 -0
  338. aiq/registry_handlers/schemas/package.py +68 -0
  339. aiq/registry_handlers/schemas/publish.py +63 -0
  340. aiq/registry_handlers/schemas/pull.py +82 -0
  341. aiq/registry_handlers/schemas/remove.py +36 -0
  342. aiq/registry_handlers/schemas/search.py +91 -0
  343. aiq/registry_handlers/schemas/status.py +47 -0
  344. aiq/retriever/__init__.py +0 -0
  345. aiq/retriever/interface.py +37 -0
  346. aiq/retriever/milvus/__init__.py +14 -0
  347. aiq/retriever/milvus/register.py +81 -0
  348. aiq/retriever/milvus/retriever.py +228 -0
  349. aiq/retriever/models.py +74 -0
  350. aiq/retriever/nemo_retriever/__init__.py +14 -0
  351. aiq/retriever/nemo_retriever/register.py +60 -0
  352. aiq/retriever/nemo_retriever/retriever.py +190 -0
  353. aiq/retriever/register.py +22 -0
  354. aiq/runtime/__init__.py +14 -0
  355. aiq/runtime/loader.py +215 -0
  356. aiq/runtime/runner.py +190 -0
  357. aiq/runtime/session.py +158 -0
  358. aiq/runtime/user_metadata.py +130 -0
  359. aiq/settings/__init__.py +0 -0
  360. aiq/settings/global_settings.py +318 -0
  361. aiq/test/.namespace +1 -0
  362. aiq/tool/__init__.py +0 -0
  363. aiq/tool/chat_completion.py +74 -0
  364. aiq/tool/code_execution/README.md +151 -0
  365. aiq/tool/code_execution/__init__.py +0 -0
  366. aiq/tool/code_execution/code_sandbox.py +267 -0
  367. aiq/tool/code_execution/local_sandbox/.gitignore +1 -0
  368. aiq/tool/code_execution/local_sandbox/Dockerfile.sandbox +60 -0
  369. aiq/tool/code_execution/local_sandbox/__init__.py +13 -0
  370. aiq/tool/code_execution/local_sandbox/local_sandbox_server.py +198 -0
  371. aiq/tool/code_execution/local_sandbox/sandbox.requirements.txt +6 -0
  372. aiq/tool/code_execution/local_sandbox/start_local_sandbox.sh +50 -0
  373. aiq/tool/code_execution/register.py +74 -0
  374. aiq/tool/code_execution/test_code_execution_sandbox.py +414 -0
  375. aiq/tool/code_execution/utils.py +100 -0
  376. aiq/tool/datetime_tools.py +42 -0
  377. aiq/tool/document_search.py +141 -0
  378. aiq/tool/github_tools/__init__.py +0 -0
  379. aiq/tool/github_tools/create_github_commit.py +133 -0
  380. aiq/tool/github_tools/create_github_issue.py +87 -0
  381. aiq/tool/github_tools/create_github_pr.py +106 -0
  382. aiq/tool/github_tools/get_github_file.py +106 -0
  383. aiq/tool/github_tools/get_github_issue.py +166 -0
  384. aiq/tool/github_tools/get_github_pr.py +256 -0
  385. aiq/tool/github_tools/update_github_issue.py +100 -0
  386. aiq/tool/mcp/__init__.py +14 -0
  387. aiq/tool/mcp/exceptions.py +142 -0
  388. aiq/tool/mcp/mcp_client.py +255 -0
  389. aiq/tool/mcp/mcp_tool.py +96 -0
  390. aiq/tool/memory_tools/__init__.py +0 -0
  391. aiq/tool/memory_tools/add_memory_tool.py +79 -0
  392. aiq/tool/memory_tools/delete_memory_tool.py +67 -0
  393. aiq/tool/memory_tools/get_memory_tool.py +72 -0
  394. aiq/tool/nvidia_rag.py +95 -0
  395. aiq/tool/register.py +38 -0
  396. aiq/tool/retriever.py +89 -0
  397. aiq/tool/server_tools.py +66 -0
  398. aiq/utils/__init__.py +0 -0
  399. aiq/utils/data_models/__init__.py +0 -0
  400. aiq/utils/data_models/schema_validator.py +58 -0
  401. aiq/utils/debugging_utils.py +43 -0
  402. aiq/utils/dump_distro_mapping.py +32 -0
  403. aiq/utils/exception_handlers/__init__.py +0 -0
  404. aiq/utils/exception_handlers/automatic_retries.py +289 -0
  405. aiq/utils/exception_handlers/mcp.py +211 -0
  406. aiq/utils/exception_handlers/schemas.py +114 -0
  407. aiq/utils/io/__init__.py +0 -0
  408. aiq/utils/io/model_processing.py +28 -0
  409. aiq/utils/io/yaml_tools.py +119 -0
  410. aiq/utils/log_utils.py +37 -0
  411. aiq/utils/metadata_utils.py +74 -0
  412. aiq/utils/optional_imports.py +142 -0
  413. aiq/utils/producer_consumer_queue.py +178 -0
  414. aiq/utils/reactive/__init__.py +0 -0
  415. aiq/utils/reactive/base/__init__.py +0 -0
  416. aiq/utils/reactive/base/observable_base.py +65 -0
  417. aiq/utils/reactive/base/observer_base.py +55 -0
  418. aiq/utils/reactive/base/subject_base.py +79 -0
  419. aiq/utils/reactive/observable.py +59 -0
  420. aiq/utils/reactive/observer.py +76 -0
  421. aiq/utils/reactive/subject.py +131 -0
  422. aiq/utils/reactive/subscription.py +49 -0
  423. aiq/utils/settings/__init__.py +0 -0
  424. aiq/utils/settings/global_settings.py +197 -0
  425. aiq/utils/string_utils.py +38 -0
  426. aiq/utils/type_converter.py +290 -0
  427. aiq/utils/type_utils.py +484 -0
  428. aiq/utils/url_utils.py +27 -0
  429. nvidia_nat-1.2.0rc5.dist-info/METADATA +363 -0
  430. nvidia_nat-1.2.0rc5.dist-info/RECORD +435 -0
  431. nvidia_nat-1.2.0rc5.dist-info/WHEEL +5 -0
  432. nvidia_nat-1.2.0rc5.dist-info/entry_points.txt +20 -0
  433. nvidia_nat-1.2.0rc5.dist-info/licenses/LICENSE-3rd-party.txt +3686 -0
  434. nvidia_nat-1.2.0rc5.dist-info/licenses/LICENSE.md +201 -0
  435. nvidia_nat-1.2.0rc5.dist-info/top_level.txt +1 -0
@@ -0,0 +1,30 @@
1
+ # SPDX-FileCopyrightText: Copyright (c) 2025, 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 aiq.builder.builder import Builder
17
+ from aiq.cli.register_workflow import register_auth_provider
18
+ from aiq.data_models.authentication import AuthProviderBaseConfig
19
+
20
+
21
+ class HTTPBasicAuthProviderConfig(AuthProviderBaseConfig, name="http_basic"):
22
+ pass
23
+
24
+
25
+ @register_auth_provider(config_type=HTTPBasicAuthProviderConfig)
26
+ async def http_basic_auth_provider(config: HTTPBasicAuthProviderConfig, builder: Builder):
27
+
28
+ from aiq.authentication.http_basic_auth.http_basic_auth_provider import HTTPBasicAuthProvider
29
+
30
+ yield HTTPBasicAuthProvider(config)
@@ -0,0 +1,93 @@
1
+ # SPDX-FileCopyrightText: Copyright (c) 2024-2025, 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 typing
17
+ from abc import ABC
18
+ from abc import abstractmethod
19
+
20
+ from aiq.data_models.authentication import AuthenticatedContext
21
+ from aiq.data_models.authentication import AuthFlowType
22
+ from aiq.data_models.authentication import AuthProviderBaseConfig
23
+ from aiq.data_models.authentication import AuthProviderBaseConfigT
24
+ from aiq.data_models.authentication import AuthResult
25
+
26
+ AUTHORIZATION_HEADER = "Authorization"
27
+
28
+
29
+ class AuthProviderBase(typing.Generic[AuthProviderBaseConfigT], ABC):
30
+ """
31
+ Base class for authenticating to API services.
32
+ This class provides an interface for authenticating to API services.
33
+ """
34
+
35
+ def __init__(self, config: AuthProviderBaseConfigT):
36
+ """
37
+ Initialize the AuthProviderBase with the given configuration.
38
+
39
+ Args:
40
+ config (AuthProviderBaseConfig): Configuration items for authentication.
41
+ """
42
+ self._config = config
43
+
44
+ @property
45
+ def config(self) -> AuthProviderBaseConfigT:
46
+ """
47
+ Returns the auth provider configuration object.
48
+
49
+ Returns
50
+ -------
51
+ AuthProviderBaseConfigT
52
+ The auth provider configuration object.
53
+ """
54
+ return self._config
55
+
56
+ @abstractmethod
57
+ async def authenticate(self, user_id: str | None = None) -> AuthResult:
58
+ """
59
+ Perform the authentication process for the client.
60
+
61
+ This method handles the necessary steps to authenticate the client with the
62
+ target API service, which may include obtaining tokens, refreshing credentials,
63
+ or completing multi-step authentication flows.
64
+
65
+ Raises:
66
+ NotImplementedError: Must be implemented by subclasses.
67
+ """
68
+ # This method will call the frontend FlowHandlerBase `authenticate` method
69
+ pass
70
+
71
+
72
+ class FlowHandlerBase(ABC):
73
+ """
74
+ Handles front-end specifc flows for authentication clients.
75
+
76
+ Each front end will define a FlowHandler that will implement the authenticate method.
77
+
78
+ The `authenticate` method will be stored as the callback in the AIQContextState.user_auth_callback
79
+ """
80
+
81
+ @abstractmethod
82
+ async def authenticate(self, config: AuthProviderBaseConfig, method: AuthFlowType) -> AuthenticatedContext:
83
+ """
84
+ Perform the authentication process for the client.
85
+
86
+ This method handles the necessary steps to authenticate the client with the
87
+ target API service, which may include obtaining tokens, refreshing credentials,
88
+ or completing multistep authentication flows.
89
+
90
+ Raises:
91
+ NotImplementedError: Must be implemented by subclasses.
92
+ """
93
+ pass
@@ -0,0 +1,14 @@
1
+ # SPDX-FileCopyrightText: Copyright (c) 2024-2025, 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,107 @@
1
+ # SPDX-FileCopyrightText: Copyright (c) 2025, 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 datetime import datetime
17
+ from datetime import timezone
18
+
19
+ from authlib.integrations.httpx_client import OAuth2Client as AuthlibOAuth2Client
20
+ from pydantic import SecretStr
21
+
22
+ from aiq.authentication.interfaces import AuthProviderBase
23
+ from aiq.authentication.oauth2.oauth2_auth_code_flow_provider_config import OAuth2AuthCodeFlowProviderConfig
24
+ from aiq.builder.context import AIQContext
25
+ from aiq.data_models.authentication import AuthFlowType
26
+ from aiq.data_models.authentication import AuthResult
27
+ from aiq.data_models.authentication import BearerTokenCred
28
+
29
+
30
+ class OAuth2AuthCodeFlowProvider(AuthProviderBase[OAuth2AuthCodeFlowProviderConfig]):
31
+
32
+ def __init__(self, config: OAuth2AuthCodeFlowProviderConfig):
33
+ super().__init__(config)
34
+ self._authenticated_tokens: dict[str, AuthResult] = {}
35
+ self._context = AIQContext.get()
36
+
37
+ async def _attempt_token_refresh(self, user_id: str, auth_result: AuthResult) -> AuthResult | None:
38
+ refresh_token = auth_result.raw.get("refresh_token")
39
+ if not isinstance(refresh_token, str):
40
+ return None
41
+
42
+ with AuthlibOAuth2Client(
43
+ client_id=self.config.client_id,
44
+ client_secret=self.config.client_secret,
45
+ ) as client:
46
+ try:
47
+ new_token_data = client.refresh_token(self.config.token_url, refresh_token=refresh_token)
48
+ except Exception:
49
+ # On any failure, we'll fall back to the full auth flow.
50
+ return None
51
+
52
+ expires_at_ts = new_token_data.get("expires_at")
53
+ new_expires_at = datetime.fromtimestamp(expires_at_ts, tz=timezone.utc) if expires_at_ts else None
54
+
55
+ new_auth_result = AuthResult(
56
+ credentials=[BearerTokenCred(token=SecretStr(new_token_data["access_token"]))],
57
+ token_expires_at=new_expires_at,
58
+ raw=new_token_data,
59
+ )
60
+
61
+ self._authenticated_tokens[user_id] = new_auth_result
62
+
63
+ return new_auth_result
64
+
65
+ async def authenticate(self, user_id: str | None = None) -> AuthResult:
66
+ if user_id is None and hasattr(AIQContext.get(), "metadata") and hasattr(
67
+ AIQContext.get().metadata, "cookies") and AIQContext.get().metadata.cookies is not None:
68
+ session_id = AIQContext.get().metadata.cookies.get("aiqtoolkit-session", None)
69
+ if not session_id:
70
+ raise RuntimeError("Authentication failed. No session ID found. Cannot identify user.")
71
+
72
+ user_id = session_id
73
+
74
+ if user_id and user_id in self._authenticated_tokens:
75
+ auth_result = self._authenticated_tokens[user_id]
76
+ if not auth_result.is_expired():
77
+ return auth_result
78
+
79
+ refreshed_auth_result = await self._attempt_token_refresh(user_id, auth_result)
80
+ if refreshed_auth_result:
81
+ return refreshed_auth_result
82
+
83
+ auth_callback = self._context.user_auth_callback
84
+ if not auth_callback:
85
+ raise RuntimeError("Authentication callback not set on AIQContext.")
86
+
87
+ try:
88
+ authenticated_context = await auth_callback(self.config, AuthFlowType.OAUTH2_AUTHORIZATION_CODE)
89
+ except Exception as e:
90
+ raise RuntimeError(f"Authentication callback failed: {e}") from e
91
+
92
+ auth_header = authenticated_context.headers.get("Authorization", "")
93
+ if not auth_header.startswith("Bearer "):
94
+ raise RuntimeError("Invalid Authorization header")
95
+
96
+ token = auth_header.split(" ")[1]
97
+
98
+ auth_result = AuthResult(
99
+ credentials=[BearerTokenCred(token=SecretStr(token))],
100
+ token_expires_at=authenticated_context.metadata.get("expires_at"),
101
+ raw=authenticated_context.metadata.get("raw_token"),
102
+ )
103
+
104
+ if user_id:
105
+ self._authenticated_tokens[user_id] = auth_result
106
+
107
+ return auth_result
@@ -0,0 +1,39 @@
1
+ # SPDX-FileCopyrightText: Copyright (c) 2024-2025, 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 aiq.data_models.authentication import AuthProviderBaseConfig
19
+
20
+
21
+ class OAuth2AuthCodeFlowProviderConfig(AuthProviderBaseConfig, name="oauth2_auth_code_flow"):
22
+
23
+ client_id: str = Field(description="The client ID for OAuth 2.0 authentication.")
24
+ client_secret: str = Field(description="The secret associated with the client_id.")
25
+ authorization_url: str = Field(description="The authorization URL for OAuth 2.0 authentication.")
26
+ token_url: str = Field(description="The token URL for OAuth 2.0 authentication.")
27
+ token_endpoint_auth_method: str | None = Field(
28
+ description=("The authentication method for the token endpoint. "
29
+ "Usually one of `client_secret_post` or `client_secret_basic`."),
30
+ default=None)
31
+ redirect_uri: str = Field(description="The redirect URI for OAuth 2.0 authentication. Must match the registered "
32
+ "redirect URI with the OAuth provider.")
33
+ scopes: list[str] = Field(description="The scopes for OAuth 2.0 authentication.", default_factory=list)
34
+ use_pkce: bool = Field(default=False,
35
+ description="Whether to use PKCE (Proof Key for Code Exchange) in the OAuth 2.0 flow.")
36
+
37
+ authorization_kwargs: dict[str, str] | None = Field(description=("Additional keyword arguments for the "
38
+ "authorization request."),
39
+ default=None)
@@ -0,0 +1,25 @@
1
+ # SPDX-FileCopyrightText: Copyright (c) 2025, 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 aiq.authentication.oauth2.oauth2_auth_code_flow_provider_config import OAuth2AuthCodeFlowProviderConfig
17
+ from aiq.builder.builder import Builder
18
+ from aiq.cli.register_workflow import register_auth_provider
19
+
20
+
21
+ @register_auth_provider(config_type=OAuth2AuthCodeFlowProviderConfig)
22
+ async def oauth2_client(authentication_provider: OAuth2AuthCodeFlowProviderConfig, builder: Builder):
23
+ from aiq.authentication.oauth2.oauth2_auth_code_flow_provider import OAuth2AuthCodeFlowProvider
24
+
25
+ yield OAuth2AuthCodeFlowProvider(authentication_provider)
@@ -0,0 +1,21 @@
1
+ # SPDX-FileCopyrightText: Copyright (c) 2025, 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
+ # pylint: disable=unused-import
17
+ # flake8: noqa
18
+
19
+ from aiq.authentication.api_key import register as register_api_key
20
+ from aiq.authentication.http_basic_auth import register as register_http_basic_auth
21
+ from aiq.authentication.oauth2 import register as register_oauth2
File without changes
aiq/builder/builder.py ADDED
@@ -0,0 +1,285 @@
1
+ # SPDX-FileCopyrightText: Copyright (c) 2024-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
2
+ # SPDX-License-Identifier: Apache-2.0
3
+ #
4
+ # Licensed under the Apache License, Version 2.0 (the "License");
5
+ # you may not use this file except in compliance with the License.
6
+ # You may obtain a copy of the License at
7
+ #
8
+ # http://www.apache.org/licenses/LICENSE-2.0
9
+ #
10
+ # Unless required by applicable law or agreed to in writing, software
11
+ # distributed under the License is distributed on an "AS IS" BASIS,
12
+ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13
+ # See the License for the specific language governing permissions and
14
+ # limitations under the License.
15
+
16
+ import asyncio
17
+ import typing
18
+ from abc import ABC
19
+ from abc import abstractmethod
20
+ from collections.abc import Sequence
21
+ from pathlib import Path
22
+
23
+ from aiq.authentication.interfaces import AuthProviderBase
24
+ from aiq.builder.context import AIQContext
25
+ from aiq.builder.framework_enum import LLMFrameworkEnum
26
+ from aiq.builder.function import Function
27
+ from aiq.data_models.authentication import AuthProviderBaseConfig
28
+ from aiq.data_models.component_ref import AuthenticationRef
29
+ from aiq.data_models.component_ref import EmbedderRef
30
+ from aiq.data_models.component_ref import FunctionRef
31
+ from aiq.data_models.component_ref import LLMRef
32
+ from aiq.data_models.component_ref import MemoryRef
33
+ from aiq.data_models.component_ref import ObjectStoreRef
34
+ from aiq.data_models.component_ref import RetrieverRef
35
+ from aiq.data_models.component_ref import TTCStrategyRef
36
+ from aiq.data_models.embedder import EmbedderBaseConfig
37
+ from aiq.data_models.evaluator import EvaluatorBaseConfig
38
+ from aiq.data_models.function import FunctionBaseConfig
39
+ from aiq.data_models.function_dependencies import FunctionDependencies
40
+ from aiq.data_models.llm import LLMBaseConfig
41
+ from aiq.data_models.memory import MemoryBaseConfig
42
+ from aiq.data_models.object_store import ObjectStoreBaseConfig
43
+ from aiq.data_models.retriever import RetrieverBaseConfig
44
+ from aiq.data_models.ttc_strategy import TTCStrategyBaseConfig
45
+ from aiq.experimental.test_time_compute.models.stage_enums import PipelineTypeEnum
46
+ from aiq.experimental.test_time_compute.models.stage_enums import StageTypeEnum
47
+ from aiq.memory.interfaces import MemoryEditor
48
+ from aiq.object_store.interfaces import ObjectStore
49
+ from aiq.retriever.interface import AIQRetriever
50
+
51
+
52
+ class UserManagerHolder():
53
+
54
+ def __init__(self, context: AIQContext) -> None:
55
+ self._context = context
56
+
57
+ def get_id(self):
58
+ return self._context.user_manager.get_id()
59
+
60
+
61
+ class Builder(ABC): # pylint: disable=too-many-public-methods
62
+
63
+ @abstractmethod
64
+ async def add_function(self, name: str | FunctionRef, config: FunctionBaseConfig) -> Function:
65
+ pass
66
+
67
+ @abstractmethod
68
+ def get_function(self, name: str | FunctionRef) -> Function:
69
+ pass
70
+
71
+ def get_functions(self, function_names: Sequence[str | FunctionRef]) -> list[Function]:
72
+
73
+ return [self.get_function(name) for name in function_names]
74
+
75
+ @abstractmethod
76
+ def get_function_config(self, name: str | FunctionRef) -> FunctionBaseConfig:
77
+ pass
78
+
79
+ @abstractmethod
80
+ async def set_workflow(self, config: FunctionBaseConfig) -> Function:
81
+ pass
82
+
83
+ @abstractmethod
84
+ def get_workflow(self) -> Function:
85
+ pass
86
+
87
+ @abstractmethod
88
+ def get_workflow_config(self) -> FunctionBaseConfig:
89
+ pass
90
+
91
+ def get_tools(self, tool_names: Sequence[str | FunctionRef],
92
+ wrapper_type: LLMFrameworkEnum | str) -> list[typing.Any]:
93
+
94
+ return [self.get_tool(fn_name=n, wrapper_type=wrapper_type) for n in tool_names]
95
+
96
+ @abstractmethod
97
+ def get_tool(self, fn_name: str | FunctionRef, wrapper_type: LLMFrameworkEnum | str) -> typing.Any:
98
+ pass
99
+
100
+ @abstractmethod
101
+ async def add_llm(self, name: str | LLMRef, config: LLMBaseConfig):
102
+ pass
103
+
104
+ @abstractmethod
105
+ async def get_llm(self, llm_name: str | LLMRef, wrapper_type: LLMFrameworkEnum | str) -> typing.Any:
106
+ pass
107
+
108
+ async def get_llms(self, llm_names: Sequence[str | LLMRef],
109
+ wrapper_type: LLMFrameworkEnum | str) -> list[typing.Any]:
110
+
111
+ coros = [self.get_llm(llm_name=n, wrapper_type=wrapper_type) for n in llm_names]
112
+
113
+ llms = await asyncio.gather(*coros, return_exceptions=False)
114
+
115
+ return list(llms)
116
+
117
+ @abstractmethod
118
+ def get_llm_config(self, llm_name: str | LLMRef) -> LLMBaseConfig:
119
+ pass
120
+
121
+ @abstractmethod
122
+ async def add_auth_provider(self, name: str | AuthenticationRef, config: AuthProviderBaseConfig):
123
+ pass
124
+
125
+ @abstractmethod
126
+ async def get_auth_provider(self, auth_provider_name: str | AuthenticationRef) -> AuthProviderBase:
127
+ pass
128
+
129
+ async def get_auth_providers(self, auth_provider_names: list[str | AuthenticationRef]):
130
+
131
+ coros = [self.get_auth_provider(auth_provider_name=n) for n in auth_provider_names]
132
+
133
+ auth_providers = await asyncio.gather(*coros, return_exceptions=False)
134
+
135
+ return list(auth_providers)
136
+
137
+ @abstractmethod
138
+ async def add_object_store(self, name: str | ObjectStoreRef, config: ObjectStoreBaseConfig):
139
+ pass
140
+
141
+ async def get_object_store_clients(self, object_store_names: Sequence[str | ObjectStoreRef]) -> list[ObjectStore]:
142
+ """
143
+ Return a list of all object store clients.
144
+ """
145
+ return list(await asyncio.gather(*[self.get_object_store_client(name) for name in object_store_names]))
146
+
147
+ @abstractmethod
148
+ async def get_object_store_client(self, object_store_name: str | ObjectStoreRef) -> ObjectStore:
149
+ pass
150
+
151
+ @abstractmethod
152
+ def get_object_store_config(self, object_store_name: str | ObjectStoreRef) -> ObjectStoreBaseConfig:
153
+ pass
154
+
155
+ @abstractmethod
156
+ async def add_embedder(self, name: str | EmbedderRef, config: EmbedderBaseConfig):
157
+ pass
158
+
159
+ async def get_embedders(self, embedder_names: Sequence[str | EmbedderRef],
160
+ wrapper_type: LLMFrameworkEnum | str) -> list[typing.Any]:
161
+
162
+ coros = [self.get_embedder(embedder_name=n, wrapper_type=wrapper_type) for n in embedder_names]
163
+
164
+ embedders = await asyncio.gather(*coros, return_exceptions=False)
165
+
166
+ return list(embedders)
167
+
168
+ @abstractmethod
169
+ async def get_embedder(self, embedder_name: str | EmbedderRef, wrapper_type: LLMFrameworkEnum | str) -> typing.Any:
170
+ pass
171
+
172
+ @abstractmethod
173
+ def get_embedder_config(self, embedder_name: str | EmbedderRef) -> EmbedderBaseConfig:
174
+ pass
175
+
176
+ @abstractmethod
177
+ async def add_memory_client(self, name: str | MemoryRef, config: MemoryBaseConfig):
178
+ pass
179
+
180
+ def get_memory_clients(self, memory_names: Sequence[str | MemoryRef]) -> list[MemoryEditor]:
181
+ """
182
+ Return a list of memory clients for the specified names.
183
+ """
184
+ return [self.get_memory_client(n) for n in memory_names]
185
+
186
+ @abstractmethod
187
+ def get_memory_client(self, memory_name: str | MemoryRef) -> MemoryEditor:
188
+ """
189
+ Return the instantiated memory client for the given name.
190
+ """
191
+ pass
192
+
193
+ @abstractmethod
194
+ def get_memory_client_config(self, memory_name: str | MemoryRef) -> MemoryBaseConfig:
195
+ pass
196
+
197
+ @abstractmethod
198
+ async def add_retriever(self, name: str | RetrieverRef, config: RetrieverBaseConfig):
199
+ pass
200
+
201
+ async def get_retrievers(self,
202
+ retriever_names: Sequence[str | RetrieverRef],
203
+ wrapper_type: LLMFrameworkEnum | str | None = None):
204
+
205
+ tasks = [self.get_retriever(n, wrapper_type=wrapper_type) for n in retriever_names]
206
+
207
+ retrievers = await asyncio.gather(*tasks, return_exceptions=False)
208
+
209
+ return list(retrievers)
210
+
211
+ @typing.overload
212
+ async def get_retriever(self, retriever_name: str | RetrieverRef,
213
+ wrapper_type: LLMFrameworkEnum | str) -> typing.Any:
214
+ ...
215
+
216
+ @typing.overload
217
+ async def get_retriever(self, retriever_name: str | RetrieverRef, wrapper_type: None) -> AIQRetriever:
218
+ ...
219
+
220
+ @typing.overload
221
+ async def get_retriever(self, retriever_name: str | RetrieverRef) -> AIQRetriever:
222
+ ...
223
+
224
+ @abstractmethod
225
+ async def get_retriever(self,
226
+ retriever_name: str | RetrieverRef,
227
+ wrapper_type: LLMFrameworkEnum | str | None = None) -> typing.Any:
228
+ pass
229
+
230
+ @abstractmethod
231
+ async def get_retriever_config(self, retriever_name: str | RetrieverRef) -> RetrieverBaseConfig:
232
+ pass
233
+
234
+ @abstractmethod
235
+ async def add_ttc_strategy(self, name: str | str, config: TTCStrategyBaseConfig):
236
+ pass
237
+
238
+ @abstractmethod
239
+ async def get_ttc_strategy(self,
240
+ strategy_name: str | TTCStrategyRef,
241
+ pipeline_type: PipelineTypeEnum,
242
+ stage_type: StageTypeEnum):
243
+ pass
244
+
245
+ @abstractmethod
246
+ async def get_ttc_strategy_config(self,
247
+ strategy_name: str | TTCStrategyRef,
248
+ pipeline_type: PipelineTypeEnum,
249
+ stage_type: StageTypeEnum) -> TTCStrategyBaseConfig:
250
+ pass
251
+
252
+ @abstractmethod
253
+ def get_user_manager(self) -> UserManagerHolder:
254
+ pass
255
+
256
+ @abstractmethod
257
+ def get_function_dependencies(self, fn_name: str) -> FunctionDependencies:
258
+ pass
259
+
260
+
261
+ class EvalBuilder(Builder):
262
+
263
+ @abstractmethod
264
+ async def add_evaluator(self, name: str, config: EvaluatorBaseConfig):
265
+ pass
266
+
267
+ @abstractmethod
268
+ def get_evaluator(self, evaluator_name: str) -> typing.Any:
269
+ pass
270
+
271
+ @abstractmethod
272
+ def get_evaluator_config(self, evaluator_name: str) -> EvaluatorBaseConfig:
273
+ pass
274
+
275
+ @abstractmethod
276
+ def get_max_concurrency(self) -> int:
277
+ pass
278
+
279
+ @abstractmethod
280
+ def get_output_dir(self) -> Path:
281
+ pass
282
+
283
+ @abstractmethod
284
+ def get_all_tools(self, wrapper_type: LLMFrameworkEnum | str) -> list[typing.Any]:
285
+ pass