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,14 @@
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.
@@ -0,0 +1,35 @@
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
+ AUTH_REDIRECT_SUCCESS_HTML = """
17
+ <!DOCTYPE html>
18
+ <html>
19
+ <head>
20
+ <title>Authentication Complete</title>
21
+ <script>
22
+ (function () {
23
+ window.history.replaceState(null, "", window.location.pathname);
24
+
25
+ window.opener?.postMessage({ type: 'AUTH_SUCCESS' }, '*');
26
+
27
+ window.close();
28
+ })();
29
+ </script>
30
+ </head>
31
+ <body>
32
+ <p>Authentication complete. You may now close this window.</p>
33
+ </body>
34
+ </html>
35
+ """
@@ -0,0 +1,80 @@
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 logging
18
+
19
+ from aiq.builder.context import AIQContext
20
+ from aiq.data_models.api_server import AIQResponseIntermediateStep
21
+ from aiq.data_models.intermediate_step import IntermediateStep
22
+
23
+ logger = logging.getLogger(__name__)
24
+
25
+
26
+ async def pull_intermediate(_q, adapter):
27
+ """
28
+ Subscribes to the runner's event stream (which is now a simplified Observable)
29
+ using direct callbacks. Processes each event with the adapter and enqueues
30
+ results to `_q`.
31
+ """
32
+ intermediate_done = asyncio.Event()
33
+ context = AIQContext.get()
34
+ loop = asyncio.get_running_loop()
35
+
36
+ async def set_intermediate_done():
37
+ intermediate_done.set()
38
+
39
+ def on_next_cb(item: IntermediateStep):
40
+ """
41
+ Synchronously called whenever the runner publishes an event.
42
+ We process it, then place it into the async queue (via a small async task).
43
+ If adapter is None, convert the raw IntermediateStep into the complete
44
+ AIQResponseIntermediateStep and place it into the queue.
45
+ """
46
+ if adapter is None:
47
+ adapted = AIQResponseIntermediateStep(id=item.UUID,
48
+ type=item.event_type,
49
+ name=item.name or "",
50
+ parent_id=item.parent_id,
51
+ payload=item.payload.model_dump_json())
52
+ else:
53
+ adapted = adapter.process(item)
54
+
55
+ if adapted is not None:
56
+ loop.create_task(_q.put(adapted))
57
+
58
+ def on_error_cb(exc: Exception):
59
+ """
60
+ Called if the runner signals an error. We log it and unblock our wait.
61
+ """
62
+ logger.error("Hit on_error: %s", exc)
63
+
64
+ loop.create_task(set_intermediate_done())
65
+
66
+ def on_complete_cb():
67
+ """
68
+ Called once the runner signals no more items. We unblock our wait.
69
+ """
70
+ logger.debug("Completed reading intermediate steps")
71
+
72
+ loop.create_task(set_intermediate_done())
73
+
74
+ # Subscribe to the runner's "reactive_event_stream" (now a simple Observable)
75
+ _ = context.intermediate_step_manager.subscribe(on_next=on_next_cb,
76
+ on_error=on_error_cb,
77
+ on_complete=on_complete_cb)
78
+
79
+ # Wait until on_complete or on_error sets intermediate_done
80
+ return intermediate_done
@@ -0,0 +1,183 @@
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
+ import os
18
+ import shutil
19
+ import threading
20
+ from datetime import UTC
21
+ from datetime import datetime
22
+ from datetime import timedelta
23
+ from enum import Enum
24
+ from uuid import uuid4
25
+
26
+ from pydantic import BaseModel
27
+
28
+ logger = logging.getLogger(__name__)
29
+
30
+
31
+ class JobStatus(str, Enum):
32
+ SUBMITTED = "submitted"
33
+ RUNNING = "running"
34
+ SUCCESS = "success"
35
+ FAILURE = "failure"
36
+ INTERRUPTED = "interrupted"
37
+ NOT_FOUND = "not_found"
38
+
39
+
40
+ # pydantic model for the job status
41
+ class JobInfo(BaseModel):
42
+ job_id: str
43
+ status: JobStatus
44
+ config_file: str | None
45
+ error: str | None
46
+ output_path: str | None
47
+ created_at: datetime
48
+ updated_at: datetime
49
+ expiry_seconds: int
50
+ output: BaseModel | None = None
51
+
52
+
53
+ class JobStore:
54
+
55
+ MIN_EXPIRY = 600 # 10 minutes
56
+ MAX_EXPIRY = 86400 # 24 hours
57
+ DEFAULT_EXPIRY = 3600 # 1 hour
58
+
59
+ # active jobs are exempt from expiry
60
+ ACTIVE_STATUS = {"running", "submitted"}
61
+
62
+ def __init__(self):
63
+ self._jobs = {}
64
+ self._lock = threading.Lock() # Ensure thread safety for job operations
65
+
66
+ def create_job(self,
67
+ config_file: str | None = None,
68
+ job_id: str | None = None,
69
+ expiry_seconds: int = DEFAULT_EXPIRY) -> str:
70
+ if job_id is None:
71
+ job_id = str(uuid4())
72
+
73
+ clamped_expiry = max(self.MIN_EXPIRY, min(expiry_seconds, self.MAX_EXPIRY))
74
+ if expiry_seconds != clamped_expiry:
75
+ logger.info("Clamped expiry_seconds from %d to %d for job %s", expiry_seconds, clamped_expiry, job_id)
76
+
77
+ job = JobInfo(job_id=job_id,
78
+ status=JobStatus.SUBMITTED,
79
+ config_file=config_file,
80
+ created_at=datetime.now(UTC),
81
+ updated_at=datetime.now(UTC),
82
+ error=None,
83
+ output_path=None,
84
+ expiry_seconds=clamped_expiry)
85
+
86
+ with self._lock:
87
+ self._jobs[job_id] = job
88
+
89
+ logger.info("Created new job %s with config %s", job_id, config_file)
90
+ return job_id
91
+
92
+ def update_status(self,
93
+ job_id: str,
94
+ status: str,
95
+ error: str | None = None,
96
+ output_path: str | None = None,
97
+ output: BaseModel | None = None):
98
+ if job_id not in self._jobs:
99
+ raise ValueError(f"Job {job_id} not found")
100
+
101
+ with self._lock:
102
+ job = self._jobs[job_id]
103
+ job.status = status
104
+ job.error = error
105
+ job.output_path = output_path
106
+ job.updated_at = datetime.now(UTC)
107
+ job.output = output
108
+
109
+ def get_status(self, job_id: str) -> JobInfo | None:
110
+ with self._lock:
111
+ return self._jobs.get(job_id)
112
+
113
+ def list_jobs(self):
114
+ with self._lock:
115
+ return self._jobs
116
+
117
+ def get_job(self, job_id: str) -> JobInfo | None:
118
+ """Get a job by its ID."""
119
+ with self._lock:
120
+ return self._jobs.get(job_id)
121
+
122
+ def get_last_job(self) -> JobInfo | None:
123
+ """Get the last created job."""
124
+ with self._lock:
125
+ if not self._jobs:
126
+ logger.info("No jobs found in job store")
127
+ return None
128
+ last_job = max(self._jobs.values(), key=lambda job: job.created_at)
129
+ logger.info("Retrieved last job %s created at %s", last_job.job_id, last_job.created_at)
130
+ return last_job
131
+
132
+ def get_jobs_by_status(self, status: str) -> list[JobInfo]:
133
+ """Get all jobs with the specified status."""
134
+ with self._lock:
135
+ return [job for job in self._jobs.values() if job.status == status]
136
+
137
+ def get_all_jobs(self) -> list[JobInfo]:
138
+ """Get all jobs in the store."""
139
+ with self._lock:
140
+ return list(self._jobs.values())
141
+
142
+ def get_expires_at(self, job: JobInfo) -> datetime | None:
143
+ """Get the time for a job to expire."""
144
+ if job.status in self.ACTIVE_STATUS:
145
+ return None
146
+ return job.updated_at + timedelta(seconds=job.expiry_seconds)
147
+
148
+ def cleanup_expired_jobs(self):
149
+ """
150
+ Cleanup expired jobs, keeping the most recent one.
151
+ Updated_at is used instead of created_at to determine the most recent job.
152
+ This is because jobs may not be processed in the order they are created.
153
+ """
154
+ now = datetime.now(UTC)
155
+
156
+ # Filter out active jobs
157
+ with self._lock:
158
+ finished_jobs = {job_id: job for job_id, job in self._jobs.items() if job.status not in self.ACTIVE_STATUS}
159
+
160
+ # Sort finished jobs by updated_at descending
161
+ sorted_finished = sorted(finished_jobs.items(), key=lambda item: item[1].updated_at, reverse=True)
162
+
163
+ # Always keep the most recent finished job
164
+ jobs_to_check = sorted_finished[1:]
165
+
166
+ expired_ids = []
167
+ for job_id, job in jobs_to_check:
168
+ expires_at = self.get_expires_at(job)
169
+ if expires_at and now > expires_at:
170
+ expired_ids.append(job_id)
171
+ # cleanup output dir if present
172
+ if job.output_path:
173
+ logger.info("Cleaning up output directory for job %s at %s", job_id, job.output_path)
174
+ # If it is a file remove it
175
+ if os.path.isfile(job.output_path):
176
+ os.remove(job.output_path)
177
+ # If it is a directory remove it
178
+ elif os.path.isdir(job.output_path):
179
+ shutil.rmtree(job.output_path)
180
+
181
+ with self._lock:
182
+ for job_id in expired_ids:
183
+ del self._jobs[job_id]
@@ -0,0 +1,72 @@
1
+ # SPDX-FileCopyrightText: Copyright (c) 2024-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
2
+ # SPDX-License-Identifier: Apache-2.0
3
+ #
4
+ # Licensed under the Apache License, Version 2.0 (the "License");
5
+ # you may not use this file except in compliance with the License.
6
+ # You may obtain a copy of the License at
7
+ #
8
+ # http://www.apache.org/licenses/LICENSE-2.0
9
+ #
10
+ # Unless required by applicable law or agreed to in writing, software
11
+ # distributed under the License is distributed on an "AS IS" BASIS,
12
+ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13
+ # See the License for the specific language governing permissions and
14
+ # limitations under the License.
15
+
16
+ import importlib
17
+ import logging
18
+ import os
19
+
20
+ from aiq.front_ends.fastapi.fastapi_front_end_plugin_worker import FastApiFrontEndPluginWorkerBase
21
+ from aiq.runtime.loader import load_config
22
+
23
+ logger = logging.getLogger(__name__)
24
+
25
+
26
+ def get_app():
27
+
28
+ config_file_path = os.getenv("AIQ_CONFIG_FILE")
29
+ front_end_worker_full_name = os.getenv("AIQ_FRONT_END_WORKER")
30
+
31
+ if (not config_file_path):
32
+ raise ValueError("Config file not found in environment variable AIQ_CONFIG_FILE.")
33
+
34
+ if (not front_end_worker_full_name):
35
+ raise ValueError("Front end worker not found in environment variable AIQ_FRONT_END_WORKER.")
36
+
37
+ # Try to import the front end worker class
38
+ try:
39
+ # Split the package from the class
40
+ front_end_worker_parts = front_end_worker_full_name.split(".")
41
+
42
+ front_end_worker_module_name = ".".join(front_end_worker_parts[:-1])
43
+ front_end_worker_class_name = front_end_worker_parts[-1]
44
+
45
+ front_end_worker_module = importlib.import_module(front_end_worker_module_name)
46
+
47
+ if not hasattr(front_end_worker_module, front_end_worker_class_name):
48
+ raise ValueError(f"Front end worker {front_end_worker_full_name} not found.")
49
+
50
+ front_end_worker_class: type[FastApiFrontEndPluginWorkerBase] = getattr(front_end_worker_module,
51
+ front_end_worker_class_name)
52
+
53
+ if (not issubclass(front_end_worker_class, FastApiFrontEndPluginWorkerBase)):
54
+ raise ValueError(
55
+ f"Front end worker {front_end_worker_full_name} is not a subclass of FastApiFrontEndPluginWorker.")
56
+
57
+ # Load the config
58
+ abs_config_file_path = os.path.abspath(config_file_path)
59
+
60
+ config = load_config(abs_config_file_path)
61
+
62
+ # Create an instance of the front end worker class
63
+ front_end_worker = front_end_worker_class(config)
64
+
65
+ aiq_app = front_end_worker.build_app()
66
+
67
+ return aiq_app
68
+
69
+ except ImportError as e:
70
+ raise ValueError(f"Front end worker {front_end_worker_full_name} not found.") from e
71
+ except Exception as e:
72
+ raise ValueError(f"Error loading front end worker {front_end_worker_full_name}: {e}") from e
@@ -0,0 +1,298 @@
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 logging
18
+ import typing
19
+ import uuid
20
+ from typing import Any
21
+
22
+ from fastapi import WebSocket
23
+ from pydantic import BaseModel
24
+ from pydantic import ValidationError
25
+ from starlette.websockets import WebSocketDisconnect
26
+
27
+ from aiq.authentication.interfaces import FlowHandlerBase
28
+ from aiq.data_models.api_server import AIQChatResponse
29
+ from aiq.data_models.api_server import AIQResponsePayloadOutput
30
+ from aiq.data_models.api_server import AIQResponseSerializable
31
+ from aiq.data_models.api_server import Error
32
+ from aiq.data_models.api_server import ErrorTypes
33
+ from aiq.data_models.api_server import SystemResponseContent
34
+ from aiq.data_models.api_server import TextContent
35
+ from aiq.data_models.api_server import WebSocketMessageStatus
36
+ from aiq.data_models.api_server import WebSocketMessageType
37
+ from aiq.data_models.api_server import WebSocketSystemInteractionMessage
38
+ from aiq.data_models.api_server import WebSocketSystemIntermediateStepMessage
39
+ from aiq.data_models.api_server import WebSocketSystemResponseTokenMessage
40
+ from aiq.data_models.api_server import WebSocketUserInteractionResponseMessage
41
+ from aiq.data_models.api_server import WebSocketUserMessage
42
+ from aiq.data_models.interactive import HumanPromptNotification
43
+ from aiq.data_models.interactive import HumanResponse
44
+ from aiq.data_models.interactive import HumanResponseNotification
45
+ from aiq.data_models.interactive import InteractionPrompt
46
+ from aiq.front_ends.fastapi.message_validator import MessageValidator
47
+ from aiq.front_ends.fastapi.response_helpers import generate_streaming_response
48
+ from aiq.front_ends.fastapi.step_adaptor import StepAdaptor
49
+ from aiq.runtime.session import AIQSessionManager
50
+
51
+ logger = logging.getLogger(__name__)
52
+
53
+
54
+ class WebSocketMessageHandler:
55
+
56
+ def __init__(self, socket: WebSocket, session_manager: AIQSessionManager, step_adaptor: StepAdaptor):
57
+ self._socket: WebSocket = socket
58
+ self._session_manager: AIQSessionManager = session_manager
59
+ self._step_adaptor: StepAdaptor = step_adaptor
60
+
61
+ self._message_validator: MessageValidator = MessageValidator()
62
+ self._running_workflow_task: asyncio.Task | None = None
63
+ self._message_parent_id: str = "default_id"
64
+ self._workflow_schema_type: str = None
65
+ self._user_interaction_response: asyncio.Future[HumanResponse] | None = None
66
+
67
+ self._flow_handler: FlowHandlerBase | None = None
68
+
69
+ def set_flow_handler(self, flow_handler: FlowHandlerBase) -> None:
70
+ self._flow_handler = flow_handler
71
+
72
+ async def __aenter__(self) -> "WebSocketMessageHandler":
73
+ await self._socket.accept()
74
+
75
+ return self
76
+
77
+ async def __aexit__(self, exc_type, exc_value, traceback) -> None:
78
+
79
+ # TODO: Handle the exit
80
+ pass
81
+
82
+ async def run(self) -> None:
83
+ """
84
+ Processes received messages from websocket and routes them appropriately.
85
+ """
86
+ while True:
87
+
88
+ try:
89
+
90
+ message: dict[str, Any] = await self._socket.receive_json()
91
+
92
+ validated_message: BaseModel = await self._message_validator.validate_message(message)
93
+
94
+ # Received a request to start a workflow
95
+ if (isinstance(validated_message, WebSocketUserMessage)):
96
+ await self.process_workflow_request(validated_message)
97
+
98
+ elif isinstance(
99
+ validated_message,
100
+ ( # noqa: E131
101
+ WebSocketSystemResponseTokenMessage,
102
+ WebSocketSystemIntermediateStepMessage,
103
+ WebSocketSystemInteractionMessage)):
104
+ # These messages are already handled by self.create_websocket_message(data_model=value, …)
105
+ # No further processing is needed here.
106
+ pass
107
+
108
+ elif (isinstance(validated_message, WebSocketUserInteractionResponseMessage)):
109
+ user_content = await self.process_user_message_content(validated_message)
110
+ self._user_interaction_response.set_result(user_content)
111
+ except (asyncio.CancelledError, WebSocketDisconnect):
112
+ # TODO: Handle the disconnect
113
+ break
114
+
115
+ return None
116
+
117
+ async def process_user_message_content(
118
+ self, user_content: WebSocketUserMessage | WebSocketUserInteractionResponseMessage) -> BaseModel | None:
119
+ """
120
+ Processes the contents of a user message.
121
+
122
+ :param user_content: Incoming content data model.
123
+ :return: A validated Pydantic user content model or None if not found.
124
+ """
125
+
126
+ for user_message in user_content.content.messages[::-1]:
127
+ if (user_message.role == "user"):
128
+
129
+ for attachment in user_message.content:
130
+
131
+ if isinstance(attachment, TextContent):
132
+ return attachment
133
+
134
+ return None
135
+
136
+ async def process_workflow_request(self, user_message_as_validated_type: WebSocketUserMessage) -> None:
137
+ """
138
+ Process user messages and routes them appropriately.
139
+
140
+ :param user_message_as_validated_type: A WebSocketUserMessage Data Model instance.
141
+ """
142
+
143
+ try:
144
+ self._message_parent_id = user_message_as_validated_type.id
145
+ self._workflow_schema_type = user_message_as_validated_type.schema_type
146
+ conversation_id: str = user_message_as_validated_type.conversation_id
147
+
148
+ content: BaseModel | None = await self.process_user_message_content(user_message_as_validated_type)
149
+
150
+ if content is None:
151
+ raise ValueError(f"User message content could not be found: {user_message_as_validated_type}")
152
+
153
+ if isinstance(content, TextContent) and (self._running_workflow_task is None):
154
+
155
+ def _done_callback(task: asyncio.Task):
156
+ self._running_workflow_task = None
157
+
158
+ # await self._process_response()
159
+ self._running_workflow_task = asyncio.create_task(
160
+ self._run_workflow(content.text, conversation_id,
161
+ result_type=AIQChatResponse)).add_done_callback(_done_callback)
162
+
163
+ except ValueError as e:
164
+ logger.error("User message content not found: %s", str(e), exc_info=True)
165
+ await self.create_websocket_message(data_model=Error(code=ErrorTypes.INVALID_USER_MESSAGE_CONTENT,
166
+ message="User message content could not be found",
167
+ details=str(e)),
168
+ message_type=WebSocketMessageType.ERROR_MESSAGE,
169
+ status=WebSocketMessageStatus.IN_PROGRESS)
170
+
171
+ async def create_websocket_message(self,
172
+ data_model: BaseModel,
173
+ message_type: str | None = None,
174
+ status: str = WebSocketMessageStatus.IN_PROGRESS) -> None:
175
+ """
176
+ Creates a websocket message that will be ready for routing based on message type or data model.
177
+
178
+ :param data_model: Message content model.
179
+ :param message_type: Message content model.
180
+ :param status: Message content model.
181
+ """
182
+ try:
183
+ message: BaseModel | None = None
184
+
185
+ if message_type is None:
186
+ message_type = await self._message_validator.resolve_message_type_by_data(data_model)
187
+
188
+ message_schema: type[BaseModel] = await self._message_validator.get_message_schema_by_type(message_type)
189
+
190
+ if 'id' in data_model.model_fields:
191
+ message_id: str = data_model.id
192
+ else:
193
+ message_id = str(uuid.uuid4())
194
+
195
+ content: BaseModel = await self._message_validator.convert_data_to_message_content(data_model)
196
+
197
+ if issubclass(message_schema, WebSocketSystemResponseTokenMessage):
198
+ message = await self._message_validator.create_system_response_token_message(
199
+ message_id=message_id, parent_id=self._message_parent_id, content=content, status=status)
200
+
201
+ elif issubclass(message_schema, WebSocketSystemIntermediateStepMessage):
202
+ message = await self._message_validator.create_system_intermediate_step_message(
203
+ message_id=message_id,
204
+ parent_id=await self._message_validator.get_intermediate_step_parent_id(data_model),
205
+ content=content,
206
+ status=status)
207
+
208
+ elif issubclass(message_schema, WebSocketSystemInteractionMessage):
209
+ message = await self._message_validator.create_system_interaction_message(
210
+ message_id=message_id, parent_id=self._message_parent_id, content=content, status=status)
211
+
212
+ elif isinstance(content, Error):
213
+ raise ValidationError(f"Invalid input data creating websocket message. {data_model.model_dump_json()}")
214
+
215
+ elif issubclass(message_schema, Error):
216
+ raise TypeError(f"Invalid message type: {message_type}")
217
+
218
+ elif (message is None):
219
+ raise ValueError(
220
+ f"Message type could not be resolved by input data model: {data_model.model_dump_json()}")
221
+
222
+ except (ValidationError, TypeError, ValueError) as e:
223
+ logger.error("A data vaidation error ocurred creating websocket message: %s", str(e), exc_info=True)
224
+ message = await self._message_validator.create_system_response_token_message(
225
+ message_type=WebSocketMessageType.ERROR_MESSAGE,
226
+ content=Error(code=ErrorTypes.UNKNOWN_ERROR, message="default", details=str(e)))
227
+
228
+ finally:
229
+ if (message is not None):
230
+ await self._socket.send_json(message.model_dump())
231
+
232
+ async def human_interaction_callback(self, prompt: InteractionPrompt) -> HumanResponse:
233
+ """
234
+ Registered human interaction callback that processes human interactions and returns
235
+ responses from websocket connection.
236
+
237
+ :param prompt: Incoming interaction content data model.
238
+ :return: A Text Content Base Pydantic model.
239
+ """
240
+
241
+ # First create a future from the loop for the human response
242
+ human_response_future: asyncio.Future[HumanResponse] = asyncio.get_running_loop().create_future()
243
+
244
+ # Then add the future to the outstanding human prompts dictionary
245
+ self._user_interaction_response = human_response_future
246
+
247
+ try:
248
+
249
+ await self.create_websocket_message(data_model=prompt.content,
250
+ message_type=WebSocketMessageType.SYSTEM_INTERACTION_MESSAGE,
251
+ status=WebSocketMessageStatus.IN_PROGRESS)
252
+
253
+ if (isinstance(prompt.content, HumanPromptNotification)):
254
+
255
+ return HumanResponseNotification()
256
+
257
+ # Wait for the human response future to complete
258
+ interaction_response: HumanResponse = await human_response_future
259
+
260
+ interaction_response: HumanResponse = await self._message_validator.convert_text_content_to_human_response(
261
+ interaction_response, prompt.content)
262
+
263
+ return interaction_response
264
+
265
+ finally:
266
+ # Delete the future from the outstanding human prompts dictionary
267
+ self._user_interaction_response = None
268
+
269
+ async def _run_workflow(self,
270
+ payload: typing.Any,
271
+ conversation_id: str | None = None,
272
+ result_type: type | None = None,
273
+ output_type: type | None = None) -> None:
274
+
275
+ try:
276
+ async with self._session_manager.session(
277
+ conversation_id=conversation_id,
278
+ request=self._socket,
279
+ user_input_callback=self.human_interaction_callback,
280
+ user_authentication_callback=(self._flow_handler.authenticate
281
+ if self._flow_handler else None)) as session:
282
+
283
+ async for value in generate_streaming_response(payload,
284
+ session_manager=session,
285
+ streaming=True,
286
+ step_adaptor=self._step_adaptor,
287
+ result_type=result_type,
288
+ output_type=output_type):
289
+
290
+ if not isinstance(value, AIQResponseSerializable):
291
+ value = AIQResponsePayloadOutput(payload=value)
292
+
293
+ await self.create_websocket_message(data_model=value, status=WebSocketMessageStatus.IN_PROGRESS)
294
+
295
+ finally:
296
+ await self.create_websocket_message(data_model=SystemResponseContent(),
297
+ message_type=WebSocketMessageType.RESPONSE_MESSAGE,
298
+ status=WebSocketMessageStatus.COMPLETE)