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
aiq/agent/__init__.py ADDED
File without changes
aiq/agent/base.py ADDED
@@ -0,0 +1,239 @@
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 asyncio
17
+ import json
18
+ import logging
19
+ from abc import ABC
20
+ from abc import abstractmethod
21
+ from enum import Enum
22
+ from typing import Any
23
+
24
+ from colorama import Fore
25
+ from langchain_core.callbacks import AsyncCallbackHandler
26
+ from langchain_core.language_models import BaseChatModel
27
+ from langchain_core.messages import AIMessage
28
+ from langchain_core.messages import BaseMessage
29
+ from langchain_core.messages import ToolMessage
30
+ from langchain_core.runnables import RunnableConfig
31
+ from langchain_core.tools import BaseTool
32
+ from langgraph.graph.graph import CompiledGraph
33
+
34
+ logger = logging.getLogger(__name__)
35
+
36
+ TOOL_NOT_FOUND_ERROR_MESSAGE = "There is no tool named {tool_name}. Tool must be one of {tools}."
37
+ INPUT_SCHEMA_MESSAGE = ". Arguments must be provided as a valid JSON object following this format: {schema}"
38
+ NO_INPUT_ERROR_MESSAGE = "No human input received to the agent, Please ask a valid question."
39
+
40
+ AGENT_LOG_PREFIX = "[AGENT]"
41
+ AGENT_CALL_LOG_MESSAGE = f"\n{'-' * 30}\n" + \
42
+ AGENT_LOG_PREFIX + "\n" + \
43
+ Fore.YELLOW + \
44
+ "Agent input: %s\n" + \
45
+ Fore.CYAN + \
46
+ "Agent's thoughts: \n%s" + \
47
+ Fore.RESET + \
48
+ f"\n{'-' * 30}"
49
+
50
+ TOOL_CALL_LOG_MESSAGE = f"\n{'-' * 30}\n" + \
51
+ AGENT_LOG_PREFIX + "\n" + \
52
+ Fore.WHITE + \
53
+ "Calling tools: %s\n" + \
54
+ Fore.YELLOW + \
55
+ "Tool's input: %s\n" + \
56
+ Fore.CYAN + \
57
+ "Tool's response: \n%s" + \
58
+ Fore.RESET + \
59
+ f"\n{'-' * 30}"
60
+
61
+
62
+ class AgentDecision(Enum):
63
+ TOOL = "tool"
64
+ END = "finished"
65
+
66
+
67
+ class BaseAgent(ABC):
68
+
69
+ def __init__(self,
70
+ llm: BaseChatModel,
71
+ tools: list[BaseTool],
72
+ callbacks: list[AsyncCallbackHandler] | None = None,
73
+ detailed_logs: bool = False) -> None:
74
+ logger.debug("Initializing Agent Graph")
75
+ self.llm = llm
76
+ self.tools = tools
77
+ self.callbacks = callbacks or []
78
+ self.detailed_logs = detailed_logs
79
+ self.graph = None
80
+
81
+ async def _stream_llm(self,
82
+ runnable: Any,
83
+ inputs: dict[str, Any],
84
+ config: RunnableConfig | None = None) -> AIMessage:
85
+ """
86
+ Stream from LLM runnable. Retry logic is handled automatically by the underlying LLM client.
87
+
88
+ Parameters
89
+ ----------
90
+ runnable : Any
91
+ The LLM runnable (prompt | llm or similar)
92
+ inputs : Dict[str, Any]
93
+ The inputs to pass to the runnable
94
+ config : RunnableConfig | None
95
+ The config to pass to the runnable (should include callbacks)
96
+
97
+ Returns
98
+ -------
99
+ AIMessage
100
+ The LLM response
101
+ """
102
+ output_message = ""
103
+ async for event in runnable.astream(inputs, config=config):
104
+ output_message += event.content
105
+
106
+ return AIMessage(content=output_message)
107
+
108
+ async def _call_llm(self, messages: list[BaseMessage]) -> AIMessage:
109
+ """
110
+ Call the LLM directly. Retry logic is handled automatically by the underlying LLM client.
111
+
112
+ Parameters
113
+ ----------
114
+ messages : list[BaseMessage]
115
+ The messages to send to the LLM
116
+
117
+ Returns
118
+ -------
119
+ AIMessage
120
+ The LLM response
121
+ """
122
+ response = await self.llm.ainvoke(messages)
123
+ return AIMessage(content=str(response.content))
124
+
125
+ async def _call_tool(self,
126
+ tool: BaseTool,
127
+ tool_input: dict[str, Any] | str,
128
+ config: RunnableConfig | None = None,
129
+ max_retries: int = 3) -> ToolMessage:
130
+ """
131
+ Call a tool with retry logic and error handling.
132
+
133
+ Parameters
134
+ ----------
135
+ tool : BaseTool
136
+ The tool to call
137
+ tool_input : Union[Dict[str, Any], str]
138
+ The input to pass to the tool
139
+ config : RunnableConfig | None
140
+ The config to pass to the tool
141
+ max_retries : int
142
+ Maximum number of retry attempts (default: 3)
143
+
144
+ Returns
145
+ -------
146
+ ToolMessage
147
+ The tool response
148
+ """
149
+ last_exception = None
150
+
151
+ for attempt in range(1, max_retries + 1):
152
+ try:
153
+ response = await tool.ainvoke(tool_input, config=config)
154
+
155
+ # Handle empty responses
156
+ if response is None or (isinstance(response, str) and response == ""):
157
+ return ToolMessage(name=tool.name,
158
+ tool_call_id=tool.name,
159
+ content=f"The tool {tool.name} provided an empty response.")
160
+
161
+ return ToolMessage(name=tool.name, tool_call_id=tool.name, content=response)
162
+
163
+ except Exception as e:
164
+ last_exception = e
165
+
166
+ # If this was the last attempt, don't sleep
167
+ if attempt == max_retries:
168
+ break
169
+
170
+ logger.warning("%s Tool call attempt %d/%d failed for tool %s: %s",
171
+ AGENT_LOG_PREFIX,
172
+ attempt,
173
+ max_retries,
174
+ tool.name,
175
+ str(e))
176
+
177
+ # Exponential backoff: 2^attempt seconds
178
+ sleep_time = 2**attempt
179
+ logger.debug("%s Retrying tool call for %s in %d seconds...", AGENT_LOG_PREFIX, tool.name, sleep_time)
180
+ await asyncio.sleep(sleep_time)
181
+
182
+ # All retries exhausted, return error message
183
+ error_content = "Tool call failed after all retry attempts. Last error: %s" % str(last_exception)
184
+ logger.error("%s %s", AGENT_LOG_PREFIX, error_content)
185
+ return ToolMessage(name=tool.name, tool_call_id=tool.name, content=error_content, status="error")
186
+
187
+ def _log_tool_response(self, tool_name: str, tool_input: Any, tool_response: str, max_chars: int = 1000) -> None:
188
+ """
189
+ Log tool response with consistent formatting and length limits.
190
+
191
+ Parameters
192
+ ----------
193
+ tool_name : str
194
+ The name of the tool that was called
195
+ tool_input : Any
196
+ The input that was passed to the tool
197
+ tool_response : str
198
+ The response from the tool
199
+ max_chars : int
200
+ Maximum number of characters to log (default: 1000)
201
+ """
202
+ if self.detailed_logs:
203
+ # Truncate tool response if too long
204
+ display_response = tool_response[:max_chars] + "...(rest of response truncated)" if len(
205
+ tool_response) > max_chars else tool_response
206
+
207
+ # Format the tool input for display
208
+ tool_input_str = str(tool_input)
209
+
210
+ tool_response_log_message = TOOL_CALL_LOG_MESSAGE % (tool_name, tool_input_str, display_response)
211
+ logger.info(tool_response_log_message)
212
+
213
+ def _parse_json(self, json_string: str) -> dict[str, Any]:
214
+ """
215
+ Safely parse JSON with graceful error handling.
216
+ If JSON parsing fails, returns an empty dict or error info.
217
+
218
+ Parameters
219
+ ----------
220
+ json_string : str
221
+ The JSON string to parse
222
+
223
+ Returns
224
+ -------
225
+ Dict[str, Any]
226
+ The parsed JSON or error information
227
+ """
228
+ try:
229
+ return json.loads(json_string)
230
+ except json.JSONDecodeError as e:
231
+ logger.warning("%s JSON parsing failed, returning the original string: %s", AGENT_LOG_PREFIX, str(e))
232
+ return {"error": f"JSON parsing failed: {str(e)}", "original_string": json_string}
233
+ except Exception as e:
234
+ logger.warning("%s Unexpected error during JSON parsing: %s", AGENT_LOG_PREFIX, str(e))
235
+ return {"error": f"Unexpected parsing error: {str(e)}", "original_string": json_string}
236
+
237
+ @abstractmethod
238
+ async def _build_graph(self, state_schema: type) -> CompiledGraph:
239
+ pass
aiq/agent/dual_node.py ADDED
@@ -0,0 +1,67 @@
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
+ from abc import abstractmethod
18
+
19
+ from langchain_core.callbacks import AsyncCallbackHandler
20
+ from langchain_core.language_models import BaseChatModel
21
+ from langchain_core.tools import BaseTool
22
+ from langgraph.graph import StateGraph
23
+ from langgraph.graph.graph import CompiledGraph
24
+ from pydantic import BaseModel
25
+
26
+ from .base import AgentDecision
27
+ from .base import BaseAgent
28
+
29
+ log = logging.getLogger(__name__)
30
+
31
+
32
+ class DualNodeAgent(BaseAgent):
33
+
34
+ def __init__(self,
35
+ llm: BaseChatModel,
36
+ tools: list[BaseTool],
37
+ callbacks: list[AsyncCallbackHandler] | None = None,
38
+ detailed_logs: bool = False):
39
+ super().__init__(llm=llm, tools=tools, callbacks=callbacks, detailed_logs=detailed_logs)
40
+
41
+ @abstractmethod
42
+ async def agent_node(self, state: BaseModel) -> BaseModel:
43
+ pass
44
+
45
+ @abstractmethod
46
+ async def tool_node(self, state: BaseModel) -> BaseModel:
47
+ pass
48
+
49
+ @abstractmethod
50
+ async def conditional_edge(self, state: BaseModel) -> str:
51
+ pass
52
+
53
+ async def _build_graph(self, state_schema) -> CompiledGraph:
54
+ log.debug("Building and compiling the Agent Graph")
55
+
56
+ graph = StateGraph(state_schema)
57
+ graph.add_node("agent", self.agent_node)
58
+ graph.add_node("tool", self.tool_node)
59
+ graph.add_edge("tool", "agent")
60
+
61
+ conditional_edge_possible_outputs = {AgentDecision.TOOL: "tool", AgentDecision.END: "__end__"}
62
+ graph.add_conditional_edges("agent", self.conditional_edge, conditional_edge_possible_outputs)
63
+
64
+ graph.set_entry_point("agent")
65
+ self.graph = graph.compile()
66
+
67
+ return self.graph
File without changes
@@ -0,0 +1,355 @@
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.agents import AgentAction
22
+ from langchain_core.agents import AgentFinish
23
+ from langchain_core.callbacks.base import AsyncCallbackHandler
24
+ from langchain_core.language_models import BaseChatModel
25
+ from langchain_core.messages.ai import AIMessage
26
+ from langchain_core.messages.base import BaseMessage
27
+ from langchain_core.messages.human import HumanMessage
28
+ from langchain_core.messages.tool import ToolMessage
29
+ from langchain_core.prompts import ChatPromptTemplate
30
+ from langchain_core.prompts import MessagesPlaceholder
31
+ from langchain_core.runnables.config import RunnableConfig
32
+ from langchain_core.tools import BaseTool
33
+ from pydantic import BaseModel
34
+ from pydantic import Field
35
+
36
+ from aiq.agent.base import AGENT_CALL_LOG_MESSAGE
37
+ from aiq.agent.base import AGENT_LOG_PREFIX
38
+ from aiq.agent.base import INPUT_SCHEMA_MESSAGE
39
+ from aiq.agent.base import NO_INPUT_ERROR_MESSAGE
40
+ from aiq.agent.base import TOOL_NOT_FOUND_ERROR_MESSAGE
41
+ from aiq.agent.base import AgentDecision
42
+ from aiq.agent.dual_node import DualNodeAgent
43
+ from aiq.agent.react_agent.output_parser import ReActOutputParser
44
+ from aiq.agent.react_agent.output_parser import ReActOutputParserException
45
+ from aiq.agent.react_agent.prompt import SYSTEM_PROMPT
46
+ from aiq.agent.react_agent.prompt import USER_PROMPT
47
+ from aiq.agent.react_agent.register import ReActAgentWorkflowConfig
48
+
49
+ logger = logging.getLogger(__name__)
50
+
51
+
52
+ class ReActGraphState(BaseModel):
53
+ """State schema for the ReAct Agent Graph"""
54
+ messages: list[BaseMessage] = Field(default_factory=list) # input and output of the ReAct Agent
55
+ agent_scratchpad: list[AgentAction] = Field(default_factory=list) # agent thoughts / intermediate steps
56
+ tool_responses: list[BaseMessage] = Field(default_factory=list) # the responses from any tool calls
57
+
58
+
59
+ class ReActAgentGraph(DualNodeAgent):
60
+ """Configurable LangGraph ReAct Agent. A ReAct Agent performs reasoning inbetween tool calls, and utilizes the tool
61
+ names and descriptions to select the optimal tool. Supports retrying on output parsing errors. Argument
62
+ "detailed_logs" toggles logging of inputs, outputs, and intermediate steps."""
63
+
64
+ def __init__(self,
65
+ llm: BaseChatModel,
66
+ prompt: ChatPromptTemplate,
67
+ tools: list[BaseTool],
68
+ use_tool_schema: bool = True,
69
+ callbacks: list[AsyncCallbackHandler] | None = None,
70
+ detailed_logs: bool = False,
71
+ retry_agent_response_parsing_errors: bool = True,
72
+ parse_agent_response_max_retries: int = 1,
73
+ tool_call_max_retries: int = 1,
74
+ pass_tool_call_errors_to_agent: bool = True):
75
+ super().__init__(llm=llm, tools=tools, callbacks=callbacks, detailed_logs=detailed_logs)
76
+ self.parse_agent_response_max_retries = (parse_agent_response_max_retries
77
+ if retry_agent_response_parsing_errors else 1)
78
+ self.tool_call_max_retries = tool_call_max_retries
79
+ self.pass_tool_call_errors_to_agent = pass_tool_call_errors_to_agent
80
+ logger.debug(
81
+ "%s Filling the prompt variables 'tools' and 'tool_names', using the tools provided in the config.",
82
+ AGENT_LOG_PREFIX)
83
+ tool_names = ",".join([tool.name for tool in tools[:-1]]) + ',' + tools[-1].name # prevent trailing ","
84
+ if not use_tool_schema:
85
+ tool_names_and_descriptions = "\n".join(
86
+ [f"{tool.name}: {tool.description}"
87
+ for tool in tools[:-1]]) + "\n" + f"{tools[-1].name}: {tools[-1].description}" # prevent trailing "\n"
88
+ else:
89
+ logger.debug("%s Adding the tools' input schema to the tools' description", AGENT_LOG_PREFIX)
90
+ tool_names_and_descriptions = "\n".join([
91
+ f"{tool.name}: {tool.description}. {INPUT_SCHEMA_MESSAGE.format(schema=tool.input_schema.model_fields)}"
92
+ for tool in tools[:-1]
93
+ ]) + "\n" + (f"{tools[-1].name}: {tools[-1].description}. "
94
+ f"{INPUT_SCHEMA_MESSAGE.format(schema=tools[-1].input_schema.model_fields)}")
95
+ prompt = prompt.partial(tools=tool_names_and_descriptions, tool_names=tool_names)
96
+ # construct the ReAct Agent
97
+ bound_llm = llm.bind(stop=["Observation:"]) # type: ignore
98
+ self.agent = prompt | bound_llm
99
+ self.tools_dict = {tool.name: tool for tool in tools}
100
+ logger.debug("%s Initialized ReAct Agent Graph", AGENT_LOG_PREFIX)
101
+
102
+ def _get_tool(self, tool_name: str):
103
+ try:
104
+ return self.tools_dict.get(tool_name)
105
+ except Exception as ex:
106
+ logger.exception("%s Unable to find tool with the name %s\n%s",
107
+ AGENT_LOG_PREFIX,
108
+ tool_name,
109
+ ex,
110
+ exc_info=True)
111
+ raise ex
112
+
113
+ async def agent_node(self, state: ReActGraphState):
114
+ try:
115
+ logger.debug("%s Starting the ReAct Agent Node", AGENT_LOG_PREFIX)
116
+ # keeping a working state allows us to resolve parsing errors without polluting the agent scratchpad
117
+ # the agent "forgets" about the parsing error after solving it - prevents hallucinations in next cycles
118
+ working_state = []
119
+ # Starting from attempt 1 instead of 0 for logging
120
+ for attempt in range(1, self.parse_agent_response_max_retries + 1):
121
+ # the first time we are invoking the ReAct Agent, it won't have any intermediate steps / agent thoughts
122
+ if len(state.agent_scratchpad) == 0 and len(working_state) == 0:
123
+ # the user input comes from the "messages" state channel
124
+ if len(state.messages) == 0:
125
+ raise RuntimeError('No input received in state: "messages"')
126
+ # to check is any human input passed or not, if no input passed Agent will return the state
127
+ content = str(state.messages[0].content)
128
+ if content.strip() == "":
129
+ logger.error("%s No human input passed to the agent.", AGENT_LOG_PREFIX)
130
+ state.messages += [AIMessage(content=NO_INPUT_ERROR_MESSAGE)]
131
+ return state
132
+ question = content
133
+ logger.debug("%s Querying agent, attempt: %s", AGENT_LOG_PREFIX, attempt)
134
+
135
+ output_message = await self._stream_llm(
136
+ self.agent,
137
+ {"question": question},
138
+ RunnableConfig(callbacks=self.callbacks) # type: ignore
139
+ )
140
+
141
+ if self.detailed_logs:
142
+ logger.info(AGENT_CALL_LOG_MESSAGE, question, output_message.content)
143
+ else:
144
+ # ReAct Agents require agentic cycles
145
+ # in an agentic cycle, preserve the agent's thoughts from the previous cycles,
146
+ # and give the agent the response from the tool it called
147
+ agent_scratchpad = []
148
+ for index, intermediate_step in enumerate(state.agent_scratchpad):
149
+ agent_thoughts = AIMessage(content=intermediate_step.log)
150
+ agent_scratchpad.append(agent_thoughts)
151
+ tool_response_content = str(state.tool_responses[index].content)
152
+ tool_response = HumanMessage(content=tool_response_content)
153
+ agent_scratchpad.append(tool_response)
154
+ agent_scratchpad += working_state
155
+ question = str(state.messages[0].content)
156
+ logger.debug("%s Querying agent, attempt: %s", AGENT_LOG_PREFIX, attempt)
157
+
158
+ output_message = await self._stream_llm(self.agent, {
159
+ "question": question, "agent_scratchpad": agent_scratchpad
160
+ },
161
+ RunnableConfig(callbacks=self.callbacks))
162
+
163
+ if self.detailed_logs:
164
+ logger.info(AGENT_CALL_LOG_MESSAGE, question, output_message.content)
165
+ logger.debug("%s The agent's scratchpad (with tool result) was:\n%s",
166
+ AGENT_LOG_PREFIX,
167
+ agent_scratchpad)
168
+ try:
169
+ # check if the agent has the final answer yet
170
+ logger.debug("%s Successfully obtained agent response. Parsing agent's response", AGENT_LOG_PREFIX)
171
+ agent_output = await ReActOutputParser().aparse(output_message.content)
172
+ logger.debug("%s Successfully parsed agent response after %s attempts", AGENT_LOG_PREFIX, attempt)
173
+ if isinstance(agent_output, AgentFinish):
174
+ final_answer = agent_output.return_values.get('output', output_message.content)
175
+ logger.debug("%s The agent has finished, and has the final answer", AGENT_LOG_PREFIX)
176
+ # this is where we handle the final output of the Agent, we can clean-up/format/postprocess here
177
+ # the final answer goes in the "messages" state channel
178
+ state.messages += [AIMessage(content=final_answer)]
179
+ else:
180
+ # the agent wants to call a tool, ensure the thoughts are preserved for the next agentic cycle
181
+ agent_output.log = output_message.content
182
+ logger.debug("%s The agent wants to call a tool: %s", AGENT_LOG_PREFIX, agent_output.tool)
183
+ state.agent_scratchpad += [agent_output]
184
+
185
+ return state
186
+ except ReActOutputParserException as ex:
187
+ # the agent output did not meet the expected ReAct output format. This can happen for a few reasons:
188
+ # the agent mentioned a tool, but already has the final answer, this can happen with Llama models
189
+ # - the ReAct Agent already has the answer, and is reflecting on how it obtained the answer
190
+ # the agent might have also missed Action or Action Input in its output
191
+ logger.debug("%s Error parsing agent output\nObservation:%s\nAgent Output:\n%s",
192
+ AGENT_LOG_PREFIX,
193
+ ex.observation,
194
+ output_message.content)
195
+ if attempt == self.parse_agent_response_max_retries:
196
+ logger.warning(
197
+ "%s Failed to parse agent output after %d attempts, consider enabling or "
198
+ "increasing parse_agent_response_max_retries",
199
+ AGENT_LOG_PREFIX,
200
+ attempt)
201
+ # the final answer goes in the "messages" state channel
202
+ combined_content = str(ex.observation) + '\n' + str(output_message.content)
203
+ output_message.content = combined_content
204
+ state.messages += [output_message]
205
+ return state
206
+ # retry parsing errors, if configured
207
+ logger.info("%s Retrying ReAct Agent, including output parsing Observation", AGENT_LOG_PREFIX)
208
+ working_state.append(output_message)
209
+ working_state.append(HumanMessage(content=str(ex.observation)))
210
+ except Exception as ex:
211
+ logger.exception("%s Failed to call agent_node: %s", AGENT_LOG_PREFIX, ex, exc_info=True)
212
+ raise ex
213
+
214
+ async def conditional_edge(self, state: ReActGraphState):
215
+ try:
216
+ logger.debug("%s Starting the ReAct Conditional Edge", AGENT_LOG_PREFIX)
217
+ if len(state.messages) > 1:
218
+ # the ReAct Agent has finished executing, the last agent output was AgentFinish
219
+ last_message_content = str(state.messages[-1].content)
220
+ logger.debug("%s Final answer:\n%s", AGENT_LOG_PREFIX, last_message_content)
221
+ return AgentDecision.END
222
+ # else the agent wants to call a tool
223
+ agent_output = state.agent_scratchpad[-1]
224
+ logger.debug("%s The agent wants to call: %s with input: %s",
225
+ AGENT_LOG_PREFIX,
226
+ agent_output.tool,
227
+ agent_output.tool_input)
228
+ return AgentDecision.TOOL
229
+ except Exception as ex:
230
+ logger.exception("Failed to determine whether agent is calling a tool: %s", ex, exc_info=True)
231
+ logger.warning("%s Ending graph traversal", AGENT_LOG_PREFIX)
232
+ return AgentDecision.END
233
+
234
+ async def tool_node(self, state: ReActGraphState):
235
+
236
+ logger.debug("%s Starting the Tool Call Node", AGENT_LOG_PREFIX)
237
+ if len(state.agent_scratchpad) == 0:
238
+ raise RuntimeError('No tool input received in state: "agent_scratchpad"')
239
+ agent_thoughts = state.agent_scratchpad[-1]
240
+ # the agent can run any installed tool, simply install the tool and add it to the config file
241
+ requested_tool = self._get_tool(agent_thoughts.tool)
242
+ if not requested_tool:
243
+ configured_tool_names = list(self.tools_dict.keys())
244
+ logger.warning(
245
+ "%s ReAct Agent wants to call tool %s. In the ReAct Agent's configuration within the config file,"
246
+ "there is no tool with that name: %s",
247
+ AGENT_LOG_PREFIX,
248
+ agent_thoughts.tool,
249
+ configured_tool_names)
250
+ tool_response = ToolMessage(name='agent_error',
251
+ tool_call_id='agent_error',
252
+ content=TOOL_NOT_FOUND_ERROR_MESSAGE.format(tool_name=agent_thoughts.tool,
253
+ tools=configured_tool_names))
254
+ state.tool_responses += [tool_response]
255
+ return state
256
+
257
+ logger.debug("%s Calling tool %s with input: %s",
258
+ AGENT_LOG_PREFIX,
259
+ requested_tool.name,
260
+ agent_thoughts.tool_input)
261
+
262
+ # Run the tool. Try to use structured input, if possible.
263
+ try:
264
+ tool_input_str = str(agent_thoughts.tool_input).strip().replace("'", '"')
265
+ tool_input_dict = json.loads(tool_input_str) if tool_input_str != 'None' else tool_input_str
266
+ logger.debug("%s Successfully parsed structured tool input from Action Input", AGENT_LOG_PREFIX)
267
+
268
+ tool_response = await self._call_tool(requested_tool,
269
+ tool_input_dict,
270
+ RunnableConfig(callbacks=self.callbacks),
271
+ max_retries=self.tool_call_max_retries)
272
+
273
+ if self.detailed_logs:
274
+ self._log_tool_response(requested_tool.name, tool_input_dict, str(tool_response.content))
275
+
276
+ except JSONDecodeError as ex:
277
+ logger.debug(
278
+ "%s Unable to parse structured tool input from Action Input. Using Action Input as is."
279
+ "\nParsing error: %s",
280
+ AGENT_LOG_PREFIX,
281
+ ex,
282
+ exc_info=True)
283
+ tool_input_str = str(agent_thoughts.tool_input)
284
+
285
+ tool_response = await self._call_tool(requested_tool,
286
+ tool_input_str,
287
+ RunnableConfig(callbacks=self.callbacks),
288
+ max_retries=self.tool_call_max_retries)
289
+
290
+ if self.detailed_logs:
291
+ self._log_tool_response(requested_tool.name, tool_input_str, str(tool_response.content))
292
+
293
+ if not self.pass_tool_call_errors_to_agent:
294
+ if tool_response.status == "error":
295
+ logger.error("%s Tool %s failed: %s", AGENT_LOG_PREFIX, requested_tool.name, tool_response.content)
296
+ raise RuntimeError("Tool call failed: " + str(tool_response.content))
297
+
298
+ state.tool_responses += [tool_response]
299
+ return state
300
+
301
+ async def build_graph(self):
302
+ try:
303
+ await super()._build_graph(state_schema=ReActGraphState)
304
+ logger.debug("%s ReAct Graph built and compiled successfully", AGENT_LOG_PREFIX)
305
+ return self.graph
306
+ except Exception as ex:
307
+ logger.exception("%s Failed to build ReAct Graph: %s", AGENT_LOG_PREFIX, ex, exc_info=ex)
308
+ raise ex
309
+
310
+ @staticmethod
311
+ def validate_system_prompt(system_prompt: str) -> bool:
312
+ errors = []
313
+ if not system_prompt:
314
+ errors.append("The system prompt cannot be empty.")
315
+ required_prompt_variables = {
316
+ "{tools}": "The system prompt must contain {tools} so the agent knows about configured tools.",
317
+ "{tool_names}": "The system prompt must contain {tool_names} so the agent knows tool names."
318
+ }
319
+ for variable_name, error_message in required_prompt_variables.items():
320
+ if variable_name not in system_prompt:
321
+ errors.append(error_message)
322
+ if errors:
323
+ error_text = "\n".join(errors)
324
+ logger.exception("%s %s", AGENT_LOG_PREFIX, error_text)
325
+ raise ValueError(error_text)
326
+ return True
327
+
328
+
329
+ def create_react_agent_prompt(config: ReActAgentWorkflowConfig) -> ChatPromptTemplate:
330
+ """
331
+ Create a ReAct Agent prompt from the config.
332
+
333
+ Args:
334
+ config (ReActAgentWorkflowConfig): The config to use for the prompt.
335
+
336
+ Returns:
337
+ ChatPromptTemplate: The ReAct Agent prompt.
338
+ """
339
+ # the ReAct Agent prompt can be customized via config option system_prompt and additional_instructions.
340
+
341
+ if config.system_prompt:
342
+ prompt_str = config.system_prompt
343
+ else:
344
+ prompt_str = SYSTEM_PROMPT
345
+
346
+ if config.additional_instructions:
347
+ prompt_str += f" {config.additional_instructions}"
348
+
349
+ valid_prompt = ReActAgentGraph.validate_system_prompt(prompt_str)
350
+ if not valid_prompt:
351
+ logger.exception("%s Invalid system_prompt", AGENT_LOG_PREFIX)
352
+ raise ValueError("Invalid system_prompt")
353
+ prompt = ChatPromptTemplate([("system", prompt_str), ("user", USER_PROMPT),
354
+ MessagesPlaceholder(variable_name='agent_scratchpad', optional=True)])
355
+ return prompt