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,411 @@
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
+ import json
17
+ # pylint: disable=R0917
18
+ import logging
19
+ from json import JSONDecodeError
20
+
21
+ from langchain_core.callbacks.base import AsyncCallbackHandler
22
+ from langchain_core.language_models import BaseChatModel
23
+ from langchain_core.messages.ai import AIMessage
24
+ from langchain_core.messages.human import HumanMessage
25
+ from langchain_core.messages.tool import ToolMessage
26
+ from langchain_core.prompts.chat import ChatPromptTemplate
27
+ from langchain_core.runnables.config import RunnableConfig
28
+ from langchain_core.tools import BaseTool
29
+ from langgraph.graph import StateGraph
30
+ from pydantic import BaseModel
31
+ from pydantic import Field
32
+
33
+ from aiq.agent.base import AGENT_CALL_LOG_MESSAGE
34
+ from aiq.agent.base import AGENT_LOG_PREFIX
35
+ from aiq.agent.base import INPUT_SCHEMA_MESSAGE
36
+ from aiq.agent.base import NO_INPUT_ERROR_MESSAGE
37
+ from aiq.agent.base import TOOL_NOT_FOUND_ERROR_MESSAGE
38
+ from aiq.agent.base import AgentDecision
39
+ from aiq.agent.base import BaseAgent
40
+
41
+ logger = logging.getLogger(__name__)
42
+
43
+
44
+ class ReWOOGraphState(BaseModel):
45
+ """State schema for the ReWOO Agent Graph"""
46
+ task: HumanMessage = Field(default_factory=lambda: HumanMessage(content="")) # the task provided by user
47
+ plan: AIMessage = Field(
48
+ default_factory=lambda: AIMessage(content="")) # the plan generated by the planner to solve the task
49
+ steps: AIMessage = Field(
50
+ default_factory=lambda: AIMessage(content="")) # the steps to solve the task, parsed from the plan
51
+ intermediate_results: dict[str, ToolMessage] = Field(default_factory=dict) # the intermediate results of each step
52
+ result: AIMessage = Field(
53
+ default_factory=lambda: AIMessage(content="")) # the final result of the task, generated by the solver
54
+
55
+
56
+ class ReWOOAgentGraph(BaseAgent):
57
+ """Configurable LangGraph ReWOO Agent. A ReWOO Agent performs reasoning by interacting with other objects or tools
58
+ and utilizes their outputs to make decisions. Supports retrying on output parsing errors. Argument
59
+ "detailed_logs" toggles logging of inputs, outputs, and intermediate steps."""
60
+
61
+ def __init__(self,
62
+ llm: BaseChatModel,
63
+ planner_prompt: ChatPromptTemplate,
64
+ solver_prompt: ChatPromptTemplate,
65
+ tools: list[BaseTool],
66
+ use_tool_schema: bool = True,
67
+ callbacks: list[AsyncCallbackHandler] | None = None,
68
+ detailed_logs: bool = False):
69
+ super().__init__(llm=llm, tools=tools, callbacks=callbacks, detailed_logs=detailed_logs)
70
+
71
+ logger.debug(
72
+ "%s Filling the prompt variables 'tools' and 'tool_names', using the tools provided in the config.",
73
+ AGENT_LOG_PREFIX)
74
+ tool_names = ",".join([tool.name for tool in tools[:-1]]) + ',' + tools[-1].name # prevent trailing ","
75
+ if not use_tool_schema:
76
+ tool_names_and_descriptions = "\n".join(
77
+ [f"{tool.name}: {tool.description}"
78
+ for tool in tools[:-1]]) + "\n" + f"{tools[-1].name}: {tools[-1].description}" # prevent trailing "\n"
79
+ else:
80
+ logger.debug("%s Adding the tools' input schema to the tools' description", AGENT_LOG_PREFIX)
81
+ tool_names_and_descriptions = "\n".join([
82
+ f"{tool.name}: {tool.description}. {INPUT_SCHEMA_MESSAGE.format(schema=tool.input_schema.model_fields)}"
83
+ for tool in tools[:-1]
84
+ ]) + "\n" + (f"{tools[-1].name}: {tools[-1].description}. "
85
+ f"{INPUT_SCHEMA_MESSAGE.format(schema=tools[-1].input_schema.model_fields)}")
86
+
87
+ self.planner_prompt = planner_prompt.partial(tools=tool_names_and_descriptions, tool_names=tool_names)
88
+ self.solver_prompt = solver_prompt
89
+ self.tools_dict = {tool.name: tool for tool in tools}
90
+
91
+ logger.debug("%s Initialized ReWOO Agent Graph", AGENT_LOG_PREFIX)
92
+
93
+ def _get_tool(self, tool_name: str):
94
+ try:
95
+ return self.tools_dict.get(tool_name)
96
+ except Exception as ex:
97
+ logger.exception("%s Unable to find tool with the name %s\n%s",
98
+ AGENT_LOG_PREFIX,
99
+ tool_name,
100
+ ex,
101
+ exc_info=True)
102
+ raise ex
103
+
104
+ @staticmethod
105
+ def _get_current_step(state: ReWOOGraphState) -> int:
106
+ steps = state.steps.content
107
+ if len(steps) == 0:
108
+ raise RuntimeError('No steps received in ReWOOGraphState')
109
+
110
+ if len(state.intermediate_results) == len(steps):
111
+ # all steps are done
112
+ return -1
113
+
114
+ return len(state.intermediate_results)
115
+
116
+ @staticmethod
117
+ def _parse_planner_output(planner_output: str) -> AIMessage:
118
+
119
+ try:
120
+ steps = json.loads(planner_output)
121
+ except json.JSONDecodeError as ex:
122
+ raise ValueError(f"The output of planner is invalid JSON format: {planner_output}") from ex
123
+
124
+ return AIMessage(content=steps)
125
+
126
+ @staticmethod
127
+ def _replace_placeholder(placeholder: str, tool_input: str | dict, tool_output: str | dict) -> str | dict:
128
+
129
+ # Replace the placeholders in the tool input with the previous tool output
130
+ if isinstance(tool_input, dict):
131
+ for key, value in tool_input.items():
132
+ if value is not None:
133
+ if value == placeholder:
134
+ tool_input[key] = tool_output
135
+ elif placeholder in value:
136
+ # If the placeholder is part of the value, replace it with the stringified output
137
+ tool_input[key] = value.replace(placeholder, str(tool_output))
138
+
139
+ elif isinstance(tool_input, str):
140
+ tool_input = tool_input.replace(placeholder, str(tool_output))
141
+
142
+ else:
143
+ assert False, f"Unexpected type for tool_input: {type(tool_input)}"
144
+ return tool_input
145
+
146
+ @staticmethod
147
+ def _parse_tool_input(tool_input: str | dict):
148
+
149
+ # If the input is already a dictionary, return it as is
150
+ if isinstance(tool_input, dict):
151
+ logger.debug("%s Tool input is already a dictionary. Use the tool input as is.", AGENT_LOG_PREFIX)
152
+ return tool_input
153
+
154
+ # If the input is a string, attempt to parse it as JSON
155
+ try:
156
+ tool_input = tool_input.strip()
157
+ # If the input is already a valid JSON string, load it
158
+ tool_input_parsed = json.loads(tool_input)
159
+ logger.debug("%s Successfully parsed structured tool input", AGENT_LOG_PREFIX)
160
+
161
+ except JSONDecodeError:
162
+ try:
163
+ # Replace single quotes with double quotes and attempt parsing again
164
+ tool_input_fixed = tool_input.replace("'", '"')
165
+ tool_input_parsed = json.loads(tool_input_fixed)
166
+ logger.debug(
167
+ "%s Successfully parsed structured tool input after replacing single quotes with double quotes",
168
+ AGENT_LOG_PREFIX)
169
+
170
+ except JSONDecodeError:
171
+ # If it still fails, fall back to using the input as a raw string
172
+ tool_input_parsed = tool_input
173
+ logger.debug("%s Unable to parse structured tool input. Using raw tool input as is.", AGENT_LOG_PREFIX)
174
+
175
+ return tool_input_parsed
176
+
177
+ async def planner_node(self, state: ReWOOGraphState):
178
+ try:
179
+ logger.debug("%s Starting the ReWOO Planner Node", AGENT_LOG_PREFIX)
180
+
181
+ planner = self.planner_prompt | self.llm
182
+ task = str(state.task.content)
183
+ if not task:
184
+ logger.error("%s No task provided to the ReWOO Agent. Please provide a valid task.", AGENT_LOG_PREFIX)
185
+ return {"result": NO_INPUT_ERROR_MESSAGE}
186
+
187
+ plan = await self._stream_llm(
188
+ planner,
189
+ {"task": task},
190
+ RunnableConfig(callbacks=self.callbacks) # type: ignore
191
+ )
192
+
193
+ steps = self._parse_planner_output(str(plan.content))
194
+
195
+ if self.detailed_logs:
196
+ agent_response_log_message = AGENT_CALL_LOG_MESSAGE % (task, str(plan.content))
197
+ logger.info("ReWOO agent planner output: %s", agent_response_log_message)
198
+
199
+ return {"plan": plan, "steps": steps}
200
+
201
+ except Exception as ex:
202
+ logger.exception("%s Failed to call planner_node: %s", AGENT_LOG_PREFIX, ex, exc_info=True)
203
+ raise ex
204
+
205
+ async def executor_node(self, state: ReWOOGraphState):
206
+ try:
207
+ logger.debug("%s Starting the ReWOO Executor Node", AGENT_LOG_PREFIX)
208
+
209
+ current_step = self._get_current_step(state)
210
+ # The executor node should not be invoked after all steps are finished
211
+ if current_step < 0:
212
+ logger.error("%s ReWOO Executor is invoked with an invalid step number: %s",
213
+ AGENT_LOG_PREFIX,
214
+ current_step)
215
+ raise RuntimeError(f"ReWOO Executor is invoked with an invalid step number: {current_step}")
216
+
217
+ steps_content = state.steps.content
218
+ if isinstance(steps_content, list) and current_step < len(steps_content):
219
+ step = steps_content[current_step]
220
+ if isinstance(step, dict) and "evidence" in step:
221
+ step_info = step["evidence"]
222
+ placeholder = step_info.get("placeholder", "")
223
+ tool = step_info.get("tool", "")
224
+ tool_input = step_info.get("tool_input", "")
225
+ else:
226
+ logger.error("%s Invalid step format at index %s", AGENT_LOG_PREFIX, current_step)
227
+ return {"intermediate_results": state.intermediate_results}
228
+ else:
229
+ logger.error("%s Invalid steps content or index %s", AGENT_LOG_PREFIX, current_step)
230
+ return {"intermediate_results": state.intermediate_results}
231
+
232
+ intermediate_results = state.intermediate_results
233
+
234
+ # Replace the placeholder in the tool input with the previous tool output
235
+ for _placeholder, _tool_output in intermediate_results.items():
236
+ _tool_output = _tool_output.content
237
+ # If the content is a list, get the first element which should be a dict
238
+ if isinstance(_tool_output, list):
239
+ _tool_output = _tool_output[0]
240
+ assert isinstance(_tool_output, dict)
241
+
242
+ tool_input = self._replace_placeholder(_placeholder, tool_input, _tool_output)
243
+
244
+ requested_tool = self._get_tool(tool)
245
+ if not requested_tool:
246
+ configured_tool_names = list(self.tools_dict.keys())
247
+ logger.warning(
248
+ "%s ReWOO Agent wants to call tool %s. In the ReWOO Agent's configuration within the config file,"
249
+ "there is no tool with that name: %s",
250
+ AGENT_LOG_PREFIX,
251
+ tool,
252
+ configured_tool_names)
253
+
254
+ intermediate_results[placeholder] = ToolMessage(content=TOOL_NOT_FOUND_ERROR_MESSAGE.format(
255
+ tool_name=tool, tools=configured_tool_names),
256
+ tool_call_id=tool)
257
+ return {"intermediate_results": intermediate_results}
258
+
259
+ if self.detailed_logs:
260
+ logger.debug("%s Calling tool %s with input: %s", AGENT_LOG_PREFIX, requested_tool.name, tool_input)
261
+
262
+ # Run the tool. Try to use structured input, if possible
263
+ tool_input_parsed = self._parse_tool_input(tool_input)
264
+ tool_response = await self._call_tool(requested_tool,
265
+ tool_input_parsed,
266
+ RunnableConfig(callbacks=self.callbacks),
267
+ max_retries=3)
268
+
269
+ # ToolMessage only accepts str or list[str | dict] as content.
270
+ # Convert into list if the response is a dict.
271
+ if isinstance(tool_response, dict):
272
+ tool_response = [tool_response]
273
+
274
+ tool_response_message = ToolMessage(name=tool, tool_call_id=tool, content=tool_response)
275
+
276
+ if self.detailed_logs:
277
+ self._log_tool_response(requested_tool.name, tool_input_parsed, str(tool_response))
278
+
279
+ intermediate_results[placeholder] = tool_response_message
280
+ return {"intermediate_results": intermediate_results}
281
+
282
+ except Exception as ex:
283
+ logger.exception("%s Failed to call executor_node: %s", AGENT_LOG_PREFIX, ex, exc_info=True)
284
+ raise ex
285
+
286
+ async def solver_node(self, state: ReWOOGraphState):
287
+ try:
288
+ logger.debug("%s Starting the ReWOO Solver Node", AGENT_LOG_PREFIX)
289
+
290
+ plan = ""
291
+ # Add the tool outputs of each step to the plan
292
+ for step in state.steps.content:
293
+ step_info = step["evidence"]
294
+ placeholder = step_info.get("placeholder", "")
295
+ tool_input = step_info.get("tool_input", "")
296
+
297
+ intermediate_results = state.intermediate_results
298
+ for _placeholder, _tool_output in intermediate_results.items():
299
+ _tool_output = _tool_output.content
300
+ # If the content is a list, get the first element which should be a dict
301
+ if isinstance(_tool_output, list):
302
+ _tool_output = _tool_output[0]
303
+ assert isinstance(_tool_output, dict)
304
+
305
+ tool_input = self._replace_placeholder(_placeholder, tool_input, _tool_output)
306
+
307
+ placeholder = placeholder.replace(_placeholder, str(_tool_output))
308
+
309
+ _plan = step.get("plan")
310
+ tool = step_info.get("tool")
311
+ plan += f"Plan: {_plan}\n{placeholder} = {tool}[{tool_input}]"
312
+
313
+ task = str(state.task.content)
314
+ solver_prompt = self.solver_prompt.partial(plan=plan)
315
+ solver = solver_prompt | self.llm
316
+
317
+ output_message = await self._stream_llm(solver, {"task": task},
318
+ RunnableConfig(callbacks=self.callbacks)) # type: ignore
319
+
320
+ if self.detailed_logs:
321
+ solver_output_log_message = AGENT_CALL_LOG_MESSAGE % (task, str(output_message.content))
322
+ logger.info("ReWOO agent solver output: %s", solver_output_log_message)
323
+
324
+ return {"result": output_message}
325
+
326
+ except Exception as ex:
327
+ logger.exception("%s Failed to call solver_node: %s", AGENT_LOG_PREFIX, ex, exc_info=True)
328
+ raise ex
329
+
330
+ async def conditional_edge(self, state: ReWOOGraphState):
331
+ try:
332
+ logger.debug("%s Starting the ReWOO Conditional Edge", AGENT_LOG_PREFIX)
333
+
334
+ current_step = self._get_current_step(state)
335
+ if current_step == -1:
336
+ logger.debug("%s The ReWOO Executor has finished its task", AGENT_LOG_PREFIX)
337
+ return AgentDecision.END
338
+
339
+ logger.debug("%s The ReWOO Executor is still working on the task", AGENT_LOG_PREFIX)
340
+ return AgentDecision.TOOL
341
+
342
+ except Exception as ex:
343
+ logger.exception("%s Failed to determine whether agent is calling a tool: %s",
344
+ AGENT_LOG_PREFIX,
345
+ ex,
346
+ exc_info=True)
347
+ logger.warning("%s Ending graph traversal", AGENT_LOG_PREFIX)
348
+ return AgentDecision.END
349
+
350
+ async def _build_graph(self, state_schema):
351
+ try:
352
+ logger.debug("%s Building and compiling the ReWOO Graph", AGENT_LOG_PREFIX)
353
+
354
+ graph = StateGraph(state_schema)
355
+ graph.add_node("planner", self.planner_node)
356
+ graph.add_node("executor", self.executor_node)
357
+ graph.add_node("solver", self.solver_node)
358
+
359
+ graph.add_edge("planner", "executor")
360
+ conditional_edge_possible_outputs = {AgentDecision.TOOL: "executor", AgentDecision.END: "solver"}
361
+ graph.add_conditional_edges("executor", self.conditional_edge, conditional_edge_possible_outputs)
362
+
363
+ graph.set_entry_point("planner")
364
+ graph.set_finish_point("solver")
365
+
366
+ self.graph = graph.compile()
367
+ logger.debug("%s ReWOO Graph built and compiled successfully", AGENT_LOG_PREFIX)
368
+
369
+ return self.graph
370
+
371
+ except Exception as ex:
372
+ logger.exception("%s Failed to build ReWOO Graph: %s", AGENT_LOG_PREFIX, ex, exc_info=ex)
373
+ raise ex
374
+
375
+ async def build_graph(self):
376
+ try:
377
+ await self._build_graph(state_schema=ReWOOGraphState)
378
+ logger.debug("%s ReWOO Graph built and compiled successfully", AGENT_LOG_PREFIX)
379
+ return self.graph
380
+ except Exception as ex:
381
+ logger.exception("%s Failed to build ReWOO Graph: %s", AGENT_LOG_PREFIX, ex, exc_info=ex)
382
+ raise ex
383
+
384
+ @staticmethod
385
+ def validate_planner_prompt(planner_prompt: str) -> bool:
386
+ errors = []
387
+ if not planner_prompt:
388
+ errors.append("The planner prompt cannot be empty.")
389
+ required_prompt_variables = {
390
+ "{tools}": "The planner prompt must contain {tools} so the planner agent knows about configured tools.",
391
+ "{tool_names}": "The planner prompt must contain {tool_names} so the planner agent knows tool names."
392
+ }
393
+ for variable_name, error_message in required_prompt_variables.items():
394
+ if variable_name not in planner_prompt:
395
+ errors.append(error_message)
396
+ if errors:
397
+ error_text = "\n".join(errors)
398
+ logger.exception("%s %s", AGENT_LOG_PREFIX, error_text)
399
+ raise ValueError(error_text)
400
+ return True
401
+
402
+ @staticmethod
403
+ def validate_solver_prompt(solver_prompt: str) -> bool:
404
+ errors = []
405
+ if not solver_prompt:
406
+ errors.append("The solver prompt cannot be empty.")
407
+ if errors:
408
+ error_text = "\n".join(errors)
409
+ logger.exception("%s %s", AGENT_LOG_PREFIX, error_text)
410
+ raise ValueError(error_text)
411
+ return True
@@ -0,0 +1,108 @@
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
+ # flake8: noqa
17
+ from langchain_core.prompts.chat import ChatPromptTemplate
18
+
19
+ PLANNER_SYSTEM_PROMPT = """
20
+ For the following task, make plans that can solve the problem step by step. For each plan, indicate \
21
+ which external tool together with tool input to retrieve evidence. You can store the evidence into a \
22
+ placeholder #E that can be called by later tools. (Plan, #E1, Plan, #E2, Plan, ...)
23
+
24
+ You may ask the human to the following tools:
25
+
26
+ {tools}
27
+
28
+ The tools should be one of the following: [{tool_names}]
29
+
30
+ You are not required to use all the tools listed. Choose only the ones that best fit the needs of each plan step.
31
+
32
+ Your output must be a JSON array where each element represents one planning step. Each step must be an object with exactly two keys:
33
+
34
+ 1. "plan": A string that describes in detail the action or reasoning for that step.
35
+
36
+ 2. "evidence": An object representing the external tool call associated with that plan step. This object must have the following keys:
37
+
38
+ -"placeholder": A string that identifies the evidence placeholder (e.g., "#E1", "#E2", etc.). The numbering should be sequential based on the order of steps.
39
+
40
+ -"tool": A string specifying the name of the external tool used.
41
+
42
+ -"tool_input": The input to the tool. This can be a string, array, or object, depending on the requirements of the tool.
43
+
44
+ Do not include any additional keys or characters in your output, and do not wrap your response with markdown formatting. Your output must be strictly valid JSON.
45
+
46
+ Important instructions:
47
+
48
+ Do not output any additional text, comments, or markdown formatting.
49
+
50
+ Do not include any explanation or reasoning text outside of the JSON array.
51
+
52
+ The output must be a valid JSON array that can be parsed directly.
53
+
54
+ Here is an example of how a valid JSON output should look:
55
+
56
+ [
57
+ \'{{
58
+ "plan": "Calculate the result of 2023 minus 25.",
59
+ "evidence": \'{{
60
+ "placeholder": "#E1",
61
+ "tool": "calculator_subtract",
62
+ "tool_input": [2023, 25]
63
+ }}\'
64
+ }}\',
65
+ \'{{
66
+ "plan": "Retrieve the year represented by the result stored in #E1.",
67
+ "evidence": \'{{
68
+ "placeholder": "#E2",
69
+ "tool": "haystack_chitchat_agent",
70
+ "tool_input": "Response with the result number contained in #E1"
71
+ }}\'
72
+ }}\',
73
+ \'{{
74
+ "plan": "Search for the CEO of Golden State Warriors in the year stored in #E2.",
75
+ "evidence": \'{{
76
+ "placeholder": "#E3",
77
+ "tool": "internet_search",
78
+ "tool_input": "Who was the CEO of Golden State Warriors in the year #E2?"
79
+ }}\'
80
+ }}\'
81
+ ]
82
+
83
+ Begin!
84
+ """
85
+
86
+ PLANNER_USER_PROMPT = """
87
+ task: {task}
88
+ """
89
+
90
+ rewoo_planner_prompt = ChatPromptTemplate([("system", PLANNER_SYSTEM_PROMPT), ("user", PLANNER_USER_PROMPT)])
91
+
92
+ SOLVER_SYSTEM_PROMPT = """
93
+ Solve the following task or problem. To solve the problem, we have made step-by-step Plan and \
94
+ retrieved corresponding Evidence to each Plan. Use them with caution since long evidence might \
95
+ contain irrelevant information.
96
+
97
+ Now solve the question or task according to provided Evidence above. Respond with the answer
98
+ directly with no extra words.
99
+
100
+ """
101
+ SOLVER_USER_PROMPT = """
102
+ plan: {plan}
103
+ task: {task}
104
+
105
+ Response:
106
+ """
107
+
108
+ rewoo_solver_prompt = ChatPromptTemplate([("system", SOLVER_SYSTEM_PROMPT), ("user", SOLVER_USER_PROMPT)])
@@ -0,0 +1,158 @@
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
+ import logging
17
+
18
+ from pydantic import Field
19
+
20
+ from aiq.builder.builder import Builder
21
+ from aiq.builder.framework_enum import LLMFrameworkEnum
22
+ from aiq.builder.function_info import FunctionInfo
23
+ from aiq.cli.register_workflow import register_function
24
+ from aiq.data_models.api_server import AIQChatRequest
25
+ from aiq.data_models.api_server import AIQChatResponse
26
+ from aiq.data_models.component_ref import FunctionRef
27
+ from aiq.data_models.component_ref import LLMRef
28
+ from aiq.data_models.function import FunctionBaseConfig
29
+ from aiq.utils.type_converter import GlobalTypeConverter
30
+
31
+ logger = logging.getLogger(__name__)
32
+
33
+
34
+ class ReWOOAgentWorkflowConfig(FunctionBaseConfig, name="rewoo_agent"):
35
+ """
36
+ Defines an AIQ Toolkit function that uses a ReWOO Agent performs reasoning inbetween tool calls, and utilizes the
37
+ tool names and descriptions to select the optimal tool.
38
+ """
39
+
40
+ tool_names: list[FunctionRef] = Field(default_factory=list,
41
+ description="The list of tools to provide to the rewoo agent.")
42
+ llm_name: LLMRef = Field(description="The LLM model to use with the rewoo agent.")
43
+ verbose: bool = Field(default=False, description="Set the verbosity of the rewoo agent's logging.")
44
+ include_tool_input_schema_in_tool_description: bool = Field(
45
+ default=True, description="Specify inclusion of tool input schemas in the prompt.")
46
+ description: str = Field(default="ReWOO Agent Workflow", description="The description of this functions use.")
47
+ planner_prompt: str | None = Field(
48
+ default=None,
49
+ description="Provides the PLANNER_PROMPT to use with the agent") # defaults to PLANNER_PROMPT in prompt.py
50
+ solver_prompt: str | None = Field(
51
+ default=None,
52
+ description="Provides the SOLVER_PROMPT to use with the agent") # defaults to SOLVER_PROMPT in prompt.py
53
+ max_history: int = Field(default=15, description="Maximum number of messages to keep in the conversation history.")
54
+ use_openai_api: bool = Field(default=False,
55
+ description=("Use OpenAI API for the input/output types to the function. "
56
+ "If False, strings will be used."))
57
+ additional_instructions: str | None = Field(
58
+ default=None, description="Additional instructions to provide to the agent in addition to the base prompt.")
59
+
60
+
61
+ @register_function(config_type=ReWOOAgentWorkflowConfig, framework_wrappers=[LLMFrameworkEnum.LANGCHAIN])
62
+ async def ReWOO_agent_workflow(config: ReWOOAgentWorkflowConfig, builder: Builder):
63
+ from langchain.schema import BaseMessage
64
+ from langchain_core.messages import trim_messages
65
+ from langchain_core.messages.human import HumanMessage
66
+ from langchain_core.prompts import ChatPromptTemplate
67
+ from langgraph.graph.graph import CompiledGraph
68
+
69
+ from aiq.agent.rewoo_agent.prompt import PLANNER_USER_PROMPT
70
+ from aiq.agent.rewoo_agent.prompt import SOLVER_USER_PROMPT
71
+
72
+ from .agent import ReWOOAgentGraph
73
+ from .agent import ReWOOGraphState
74
+ from .prompt import rewoo_planner_prompt
75
+ from .prompt import rewoo_solver_prompt
76
+
77
+ # the ReWOO Agent prompt comes from prompt.py, and can be customized there or via config option planner_prompt and
78
+ # solver_prompt.
79
+ if config.planner_prompt:
80
+ planner_prompt = config.planner_prompt
81
+ if config.additional_instructions:
82
+ planner_prompt += f"{config.additional_instructions}"
83
+ valid = ReWOOAgentGraph.validate_planner_prompt(config.planner_prompt)
84
+ if not valid:
85
+ logger.exception("Invalid planner_prompt")
86
+ raise ValueError("Invalid planner_prompt")
87
+ planner_prompt = ChatPromptTemplate([("system", config.planner_prompt), ("user", PLANNER_USER_PROMPT)])
88
+ else:
89
+ planner_prompt = rewoo_planner_prompt
90
+
91
+ if config.solver_prompt:
92
+ solver_prompt = config.solver_prompt
93
+ if config.additional_instructions:
94
+ solver_prompt += f"{config.additional_instructions}"
95
+ valid = ReWOOAgentGraph.validate_solver_prompt(config.solver_prompt)
96
+ if not valid:
97
+ logger.exception("Invalid solver_prompt")
98
+ raise ValueError("Invalid solver_prompt")
99
+ solver_prompt = ChatPromptTemplate([("system", config.solver_prompt), ("user", SOLVER_USER_PROMPT)])
100
+ else:
101
+ solver_prompt = rewoo_solver_prompt
102
+
103
+ # we can choose an LLM for the ReWOO agent in the config file
104
+ llm = await builder.get_llm(config.llm_name, wrapper_type=LLMFrameworkEnum.LANGCHAIN)
105
+
106
+ # the agent can run any installed tool, simply install the tool and add it to the config file
107
+ # the sample tool provided can easily be copied or changed
108
+ tools = builder.get_tools(tool_names=config.tool_names, wrapper_type=LLMFrameworkEnum.LANGCHAIN)
109
+ if not tools:
110
+ raise ValueError(f"No tools specified for ReWOO Agent '{config.llm_name}'")
111
+
112
+ # construct the ReWOO Agent Graph from the configured llm, prompt, and tools
113
+ graph: CompiledGraph = await ReWOOAgentGraph(llm=llm,
114
+ planner_prompt=planner_prompt,
115
+ solver_prompt=solver_prompt,
116
+ tools=tools,
117
+ use_tool_schema=config.include_tool_input_schema_in_tool_description,
118
+ detailed_logs=config.verbose).build_graph()
119
+
120
+ async def _response_fn(input_message: AIQChatRequest) -> AIQChatResponse:
121
+ try:
122
+ # initialize the starting state with the user query
123
+ messages: list[BaseMessage] = trim_messages(messages=[m.model_dump() for m in input_message.messages],
124
+ max_tokens=config.max_history,
125
+ strategy="last",
126
+ token_counter=len,
127
+ start_on="human",
128
+ include_system=True)
129
+ task = HumanMessage(content=messages[0].content)
130
+ state = ReWOOGraphState(task=task)
131
+
132
+ # run the ReWOO Agent Graph
133
+ state = await graph.ainvoke(state)
134
+
135
+ # get and return the output from the state
136
+ state = ReWOOGraphState(**state)
137
+ output_message = state.result.content # pylint: disable=E1101
138
+ return AIQChatResponse.from_string(output_message)
139
+
140
+ except Exception as ex:
141
+ logger.exception("ReWOO Agent failed with exception: %s", ex, exc_info=ex)
142
+ # here, we can implement custom error messages
143
+ if config.verbose:
144
+ return AIQChatResponse.from_string(str(ex))
145
+ return AIQChatResponse.from_string("I seem to be having a problem.")
146
+
147
+ if (config.use_openai_api):
148
+ yield FunctionInfo.from_fn(_response_fn, description=config.description)
149
+
150
+ else:
151
+
152
+ async def _str_api_fn(input_message: str) -> str:
153
+ oai_input = GlobalTypeConverter.get().try_convert(input_message, to_type=AIQChatRequest)
154
+ oai_output = await _response_fn(oai_input)
155
+
156
+ return GlobalTypeConverter.get().try_convert(oai_output, to_type=str)
157
+
158
+ yield FunctionInfo.from_fn(_str_api_fn, description=config.description)
File without changes