nvidia-nat 1.2.0a20250813__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 (436) hide show
  1. nat/agent/__init__.py +0 -0
  2. nat/agent/base.py +239 -0
  3. nat/agent/dual_node.py +67 -0
  4. nat/agent/react_agent/__init__.py +0 -0
  5. nat/agent/react_agent/agent.py +355 -0
  6. nat/agent/react_agent/output_parser.py +104 -0
  7. nat/agent/react_agent/prompt.py +41 -0
  8. nat/agent/react_agent/register.py +149 -0
  9. nat/agent/reasoning_agent/__init__.py +0 -0
  10. nat/agent/reasoning_agent/reasoning_agent.py +225 -0
  11. nat/agent/register.py +23 -0
  12. nat/agent/rewoo_agent/__init__.py +0 -0
  13. nat/agent/rewoo_agent/agent.py +411 -0
  14. nat/agent/rewoo_agent/prompt.py +108 -0
  15. nat/agent/rewoo_agent/register.py +158 -0
  16. nat/agent/tool_calling_agent/__init__.py +0 -0
  17. nat/agent/tool_calling_agent/agent.py +119 -0
  18. nat/agent/tool_calling_agent/register.py +106 -0
  19. nat/authentication/__init__.py +14 -0
  20. nat/authentication/api_key/__init__.py +14 -0
  21. nat/authentication/api_key/api_key_auth_provider.py +96 -0
  22. nat/authentication/api_key/api_key_auth_provider_config.py +124 -0
  23. nat/authentication/api_key/register.py +26 -0
  24. nat/authentication/exceptions/__init__.py +14 -0
  25. nat/authentication/exceptions/api_key_exceptions.py +38 -0
  26. nat/authentication/http_basic_auth/__init__.py +0 -0
  27. nat/authentication/http_basic_auth/http_basic_auth_provider.py +81 -0
  28. nat/authentication/http_basic_auth/register.py +30 -0
  29. nat/authentication/interfaces.py +93 -0
  30. nat/authentication/oauth2/__init__.py +14 -0
  31. nat/authentication/oauth2/oauth2_auth_code_flow_provider.py +107 -0
  32. nat/authentication/oauth2/oauth2_auth_code_flow_provider_config.py +39 -0
  33. nat/authentication/oauth2/register.py +25 -0
  34. nat/authentication/register.py +21 -0
  35. nat/builder/__init__.py +0 -0
  36. nat/builder/builder.py +285 -0
  37. nat/builder/component_utils.py +316 -0
  38. nat/builder/context.py +270 -0
  39. nat/builder/embedder.py +24 -0
  40. nat/builder/eval_builder.py +161 -0
  41. nat/builder/evaluator.py +29 -0
  42. nat/builder/framework_enum.py +24 -0
  43. nat/builder/front_end.py +73 -0
  44. nat/builder/function.py +344 -0
  45. nat/builder/function_base.py +380 -0
  46. nat/builder/function_info.py +627 -0
  47. nat/builder/intermediate_step_manager.py +174 -0
  48. nat/builder/llm.py +25 -0
  49. nat/builder/retriever.py +25 -0
  50. nat/builder/user_interaction_manager.py +78 -0
  51. nat/builder/workflow.py +148 -0
  52. nat/builder/workflow_builder.py +1117 -0
  53. nat/cli/__init__.py +14 -0
  54. nat/cli/cli_utils/__init__.py +0 -0
  55. nat/cli/cli_utils/config_override.py +231 -0
  56. nat/cli/cli_utils/validation.py +37 -0
  57. nat/cli/commands/__init__.py +0 -0
  58. nat/cli/commands/configure/__init__.py +0 -0
  59. nat/cli/commands/configure/channel/__init__.py +0 -0
  60. nat/cli/commands/configure/channel/add.py +28 -0
  61. nat/cli/commands/configure/channel/channel.py +36 -0
  62. nat/cli/commands/configure/channel/remove.py +30 -0
  63. nat/cli/commands/configure/channel/update.py +30 -0
  64. nat/cli/commands/configure/configure.py +33 -0
  65. nat/cli/commands/evaluate.py +139 -0
  66. nat/cli/commands/info/__init__.py +14 -0
  67. nat/cli/commands/info/info.py +39 -0
  68. nat/cli/commands/info/list_channels.py +32 -0
  69. nat/cli/commands/info/list_components.py +129 -0
  70. nat/cli/commands/info/list_mcp.py +304 -0
  71. nat/cli/commands/registry/__init__.py +14 -0
  72. nat/cli/commands/registry/publish.py +88 -0
  73. nat/cli/commands/registry/pull.py +118 -0
  74. nat/cli/commands/registry/registry.py +38 -0
  75. nat/cli/commands/registry/remove.py +108 -0
  76. nat/cli/commands/registry/search.py +155 -0
  77. nat/cli/commands/sizing/__init__.py +14 -0
  78. nat/cli/commands/sizing/calc.py +297 -0
  79. nat/cli/commands/sizing/sizing.py +27 -0
  80. nat/cli/commands/start.py +246 -0
  81. nat/cli/commands/uninstall.py +81 -0
  82. nat/cli/commands/validate.py +47 -0
  83. nat/cli/commands/workflow/__init__.py +14 -0
  84. nat/cli/commands/workflow/templates/__init__.py.j2 +0 -0
  85. nat/cli/commands/workflow/templates/config.yml.j2 +16 -0
  86. nat/cli/commands/workflow/templates/pyproject.toml.j2 +22 -0
  87. nat/cli/commands/workflow/templates/register.py.j2 +5 -0
  88. nat/cli/commands/workflow/templates/workflow.py.j2 +36 -0
  89. nat/cli/commands/workflow/workflow.py +37 -0
  90. nat/cli/commands/workflow/workflow_commands.py +317 -0
  91. nat/cli/entrypoint.py +135 -0
  92. nat/cli/main.py +57 -0
  93. nat/cli/register_workflow.py +488 -0
  94. nat/cli/type_registry.py +1000 -0
  95. nat/data_models/__init__.py +14 -0
  96. nat/data_models/api_server.py +709 -0
  97. nat/data_models/authentication.py +231 -0
  98. nat/data_models/common.py +171 -0
  99. nat/data_models/component.py +58 -0
  100. nat/data_models/component_ref.py +168 -0
  101. nat/data_models/config.py +410 -0
  102. nat/data_models/dataset_handler.py +123 -0
  103. nat/data_models/discovery_metadata.py +334 -0
  104. nat/data_models/embedder.py +27 -0
  105. nat/data_models/evaluate.py +127 -0
  106. nat/data_models/evaluator.py +26 -0
  107. nat/data_models/front_end.py +26 -0
  108. nat/data_models/function.py +30 -0
  109. nat/data_models/function_dependencies.py +72 -0
  110. nat/data_models/interactive.py +246 -0
  111. nat/data_models/intermediate_step.py +302 -0
  112. nat/data_models/invocation_node.py +38 -0
  113. nat/data_models/llm.py +27 -0
  114. nat/data_models/logging.py +26 -0
  115. nat/data_models/memory.py +27 -0
  116. nat/data_models/object_store.py +44 -0
  117. nat/data_models/profiler.py +54 -0
  118. nat/data_models/registry_handler.py +26 -0
  119. nat/data_models/retriever.py +30 -0
  120. nat/data_models/retry_mixin.py +35 -0
  121. nat/data_models/span.py +190 -0
  122. nat/data_models/step_adaptor.py +64 -0
  123. nat/data_models/streaming.py +33 -0
  124. nat/data_models/swe_bench_model.py +54 -0
  125. nat/data_models/telemetry_exporter.py +26 -0
  126. nat/data_models/ttc_strategy.py +30 -0
  127. nat/embedder/__init__.py +0 -0
  128. nat/embedder/langchain_client.py +41 -0
  129. nat/embedder/nim_embedder.py +59 -0
  130. nat/embedder/openai_embedder.py +43 -0
  131. nat/embedder/register.py +24 -0
  132. nat/eval/__init__.py +14 -0
  133. nat/eval/config.py +60 -0
  134. nat/eval/dataset_handler/__init__.py +0 -0
  135. nat/eval/dataset_handler/dataset_downloader.py +106 -0
  136. nat/eval/dataset_handler/dataset_filter.py +52 -0
  137. nat/eval/dataset_handler/dataset_handler.py +254 -0
  138. nat/eval/evaluate.py +510 -0
  139. nat/eval/evaluator/__init__.py +14 -0
  140. nat/eval/evaluator/base_evaluator.py +77 -0
  141. nat/eval/evaluator/evaluator_model.py +45 -0
  142. nat/eval/intermediate_step_adapter.py +99 -0
  143. nat/eval/rag_evaluator/__init__.py +0 -0
  144. nat/eval/rag_evaluator/evaluate.py +178 -0
  145. nat/eval/rag_evaluator/register.py +143 -0
  146. nat/eval/register.py +23 -0
  147. nat/eval/remote_workflow.py +133 -0
  148. nat/eval/runners/__init__.py +14 -0
  149. nat/eval/runners/config.py +39 -0
  150. nat/eval/runners/multi_eval_runner.py +54 -0
  151. nat/eval/runtime_event_subscriber.py +52 -0
  152. nat/eval/swe_bench_evaluator/__init__.py +0 -0
  153. nat/eval/swe_bench_evaluator/evaluate.py +215 -0
  154. nat/eval/swe_bench_evaluator/register.py +36 -0
  155. nat/eval/trajectory_evaluator/__init__.py +0 -0
  156. nat/eval/trajectory_evaluator/evaluate.py +75 -0
  157. nat/eval/trajectory_evaluator/register.py +40 -0
  158. nat/eval/tunable_rag_evaluator/__init__.py +0 -0
  159. nat/eval/tunable_rag_evaluator/evaluate.py +245 -0
  160. nat/eval/tunable_rag_evaluator/register.py +52 -0
  161. nat/eval/usage_stats.py +41 -0
  162. nat/eval/utils/__init__.py +0 -0
  163. nat/eval/utils/output_uploader.py +140 -0
  164. nat/eval/utils/tqdm_position_registry.py +40 -0
  165. nat/eval/utils/weave_eval.py +184 -0
  166. nat/experimental/__init__.py +0 -0
  167. nat/experimental/decorators/__init__.py +0 -0
  168. nat/experimental/decorators/experimental_warning_decorator.py +134 -0
  169. nat/experimental/test_time_compute/__init__.py +0 -0
  170. nat/experimental/test_time_compute/editing/__init__.py +0 -0
  171. nat/experimental/test_time_compute/editing/iterative_plan_refinement_editor.py +147 -0
  172. nat/experimental/test_time_compute/editing/llm_as_a_judge_editor.py +204 -0
  173. nat/experimental/test_time_compute/editing/motivation_aware_summarization.py +107 -0
  174. nat/experimental/test_time_compute/functions/__init__.py +0 -0
  175. nat/experimental/test_time_compute/functions/execute_score_select_function.py +105 -0
  176. nat/experimental/test_time_compute/functions/plan_select_execute_function.py +224 -0
  177. nat/experimental/test_time_compute/functions/ttc_tool_orchestration_function.py +205 -0
  178. nat/experimental/test_time_compute/functions/ttc_tool_wrapper_function.py +146 -0
  179. nat/experimental/test_time_compute/models/__init__.py +0 -0
  180. nat/experimental/test_time_compute/models/editor_config.py +132 -0
  181. nat/experimental/test_time_compute/models/scoring_config.py +112 -0
  182. nat/experimental/test_time_compute/models/search_config.py +120 -0
  183. nat/experimental/test_time_compute/models/selection_config.py +154 -0
  184. nat/experimental/test_time_compute/models/stage_enums.py +43 -0
  185. nat/experimental/test_time_compute/models/strategy_base.py +66 -0
  186. nat/experimental/test_time_compute/models/tool_use_config.py +41 -0
  187. nat/experimental/test_time_compute/models/ttc_item.py +48 -0
  188. nat/experimental/test_time_compute/register.py +36 -0
  189. nat/experimental/test_time_compute/scoring/__init__.py +0 -0
  190. nat/experimental/test_time_compute/scoring/llm_based_agent_scorer.py +168 -0
  191. nat/experimental/test_time_compute/scoring/llm_based_plan_scorer.py +168 -0
  192. nat/experimental/test_time_compute/scoring/motivation_aware_scorer.py +111 -0
  193. nat/experimental/test_time_compute/search/__init__.py +0 -0
  194. nat/experimental/test_time_compute/search/multi_llm_planner.py +128 -0
  195. nat/experimental/test_time_compute/search/multi_query_retrieval_search.py +122 -0
  196. nat/experimental/test_time_compute/search/single_shot_multi_plan_planner.py +128 -0
  197. nat/experimental/test_time_compute/selection/__init__.py +0 -0
  198. nat/experimental/test_time_compute/selection/best_of_n_selector.py +63 -0
  199. nat/experimental/test_time_compute/selection/llm_based_agent_output_selector.py +131 -0
  200. nat/experimental/test_time_compute/selection/llm_based_output_merging_selector.py +159 -0
  201. nat/experimental/test_time_compute/selection/llm_based_plan_selector.py +128 -0
  202. nat/experimental/test_time_compute/selection/threshold_selector.py +58 -0
  203. nat/front_ends/__init__.py +14 -0
  204. nat/front_ends/console/__init__.py +14 -0
  205. nat/front_ends/console/authentication_flow_handler.py +233 -0
  206. nat/front_ends/console/console_front_end_config.py +32 -0
  207. nat/front_ends/console/console_front_end_plugin.py +96 -0
  208. nat/front_ends/console/register.py +25 -0
  209. nat/front_ends/cron/__init__.py +14 -0
  210. nat/front_ends/fastapi/__init__.py +14 -0
  211. nat/front_ends/fastapi/auth_flow_handlers/__init__.py +0 -0
  212. nat/front_ends/fastapi/auth_flow_handlers/http_flow_handler.py +27 -0
  213. nat/front_ends/fastapi/auth_flow_handlers/websocket_flow_handler.py +107 -0
  214. nat/front_ends/fastapi/fastapi_front_end_config.py +242 -0
  215. nat/front_ends/fastapi/fastapi_front_end_controller.py +68 -0
  216. nat/front_ends/fastapi/fastapi_front_end_plugin.py +116 -0
  217. nat/front_ends/fastapi/fastapi_front_end_plugin_worker.py +1087 -0
  218. nat/front_ends/fastapi/html_snippets/__init__.py +14 -0
  219. nat/front_ends/fastapi/html_snippets/auth_code_grant_success.py +35 -0
  220. nat/front_ends/fastapi/intermediate_steps_subscriber.py +80 -0
  221. nat/front_ends/fastapi/job_store.py +183 -0
  222. nat/front_ends/fastapi/main.py +72 -0
  223. nat/front_ends/fastapi/message_handler.py +309 -0
  224. nat/front_ends/fastapi/message_validator.py +354 -0
  225. nat/front_ends/fastapi/register.py +25 -0
  226. nat/front_ends/fastapi/response_helpers.py +195 -0
  227. nat/front_ends/fastapi/step_adaptor.py +319 -0
  228. nat/front_ends/mcp/__init__.py +14 -0
  229. nat/front_ends/mcp/mcp_front_end_config.py +35 -0
  230. nat/front_ends/mcp/mcp_front_end_plugin.py +81 -0
  231. nat/front_ends/mcp/mcp_front_end_plugin_worker.py +143 -0
  232. nat/front_ends/mcp/register.py +27 -0
  233. nat/front_ends/mcp/tool_converter.py +242 -0
  234. nat/front_ends/register.py +22 -0
  235. nat/front_ends/simple_base/__init__.py +14 -0
  236. nat/front_ends/simple_base/simple_front_end_plugin_base.py +54 -0
  237. nat/llm/__init__.py +0 -0
  238. nat/llm/aws_bedrock_llm.py +57 -0
  239. nat/llm/nim_llm.py +46 -0
  240. nat/llm/openai_llm.py +46 -0
  241. nat/llm/register.py +23 -0
  242. nat/llm/utils/__init__.py +14 -0
  243. nat/llm/utils/env_config_value.py +94 -0
  244. nat/llm/utils/error.py +17 -0
  245. nat/memory/__init__.py +20 -0
  246. nat/memory/interfaces.py +183 -0
  247. nat/memory/models.py +112 -0
  248. nat/meta/module_to_distro.json +4 -0
  249. nat/meta/pypi.md +58 -0
  250. nat/object_store/__init__.py +20 -0
  251. nat/object_store/in_memory_object_store.py +76 -0
  252. nat/object_store/interfaces.py +84 -0
  253. nat/object_store/models.py +36 -0
  254. nat/object_store/register.py +20 -0
  255. nat/observability/__init__.py +14 -0
  256. nat/observability/exporter/__init__.py +14 -0
  257. nat/observability/exporter/base_exporter.py +449 -0
  258. nat/observability/exporter/exporter.py +78 -0
  259. nat/observability/exporter/file_exporter.py +33 -0
  260. nat/observability/exporter/processing_exporter.py +322 -0
  261. nat/observability/exporter/raw_exporter.py +52 -0
  262. nat/observability/exporter/span_exporter.py +288 -0
  263. nat/observability/exporter_manager.py +335 -0
  264. nat/observability/mixin/__init__.py +14 -0
  265. nat/observability/mixin/batch_config_mixin.py +26 -0
  266. nat/observability/mixin/collector_config_mixin.py +23 -0
  267. nat/observability/mixin/file_mixin.py +288 -0
  268. nat/observability/mixin/file_mode.py +23 -0
  269. nat/observability/mixin/resource_conflict_mixin.py +134 -0
  270. nat/observability/mixin/serialize_mixin.py +61 -0
  271. nat/observability/mixin/type_introspection_mixin.py +183 -0
  272. nat/observability/processor/__init__.py +14 -0
  273. nat/observability/processor/batching_processor.py +310 -0
  274. nat/observability/processor/callback_processor.py +42 -0
  275. nat/observability/processor/intermediate_step_serializer.py +28 -0
  276. nat/observability/processor/processor.py +71 -0
  277. nat/observability/register.py +96 -0
  278. nat/observability/utils/__init__.py +14 -0
  279. nat/observability/utils/dict_utils.py +236 -0
  280. nat/observability/utils/time_utils.py +31 -0
  281. nat/plugins/.namespace +1 -0
  282. nat/profiler/__init__.py +0 -0
  283. nat/profiler/calc/__init__.py +14 -0
  284. nat/profiler/calc/calc_runner.py +627 -0
  285. nat/profiler/calc/calculations.py +288 -0
  286. nat/profiler/calc/data_models.py +188 -0
  287. nat/profiler/calc/plot.py +345 -0
  288. nat/profiler/callbacks/__init__.py +0 -0
  289. nat/profiler/callbacks/agno_callback_handler.py +295 -0
  290. nat/profiler/callbacks/base_callback_class.py +20 -0
  291. nat/profiler/callbacks/langchain_callback_handler.py +290 -0
  292. nat/profiler/callbacks/llama_index_callback_handler.py +205 -0
  293. nat/profiler/callbacks/semantic_kernel_callback_handler.py +238 -0
  294. nat/profiler/callbacks/token_usage_base_model.py +27 -0
  295. nat/profiler/data_frame_row.py +51 -0
  296. nat/profiler/data_models.py +24 -0
  297. nat/profiler/decorators/__init__.py +0 -0
  298. nat/profiler/decorators/framework_wrapper.py +131 -0
  299. nat/profiler/decorators/function_tracking.py +254 -0
  300. nat/profiler/forecasting/__init__.py +0 -0
  301. nat/profiler/forecasting/config.py +18 -0
  302. nat/profiler/forecasting/model_trainer.py +75 -0
  303. nat/profiler/forecasting/models/__init__.py +22 -0
  304. nat/profiler/forecasting/models/forecasting_base_model.py +40 -0
  305. nat/profiler/forecasting/models/linear_model.py +196 -0
  306. nat/profiler/forecasting/models/random_forest_regressor.py +268 -0
  307. nat/profiler/inference_metrics_model.py +28 -0
  308. nat/profiler/inference_optimization/__init__.py +0 -0
  309. nat/profiler/inference_optimization/bottleneck_analysis/__init__.py +0 -0
  310. nat/profiler/inference_optimization/bottleneck_analysis/nested_stack_analysis.py +460 -0
  311. nat/profiler/inference_optimization/bottleneck_analysis/simple_stack_analysis.py +258 -0
  312. nat/profiler/inference_optimization/data_models.py +386 -0
  313. nat/profiler/inference_optimization/experimental/__init__.py +0 -0
  314. nat/profiler/inference_optimization/experimental/concurrency_spike_analysis.py +468 -0
  315. nat/profiler/inference_optimization/experimental/prefix_span_analysis.py +405 -0
  316. nat/profiler/inference_optimization/llm_metrics.py +212 -0
  317. nat/profiler/inference_optimization/prompt_caching.py +163 -0
  318. nat/profiler/inference_optimization/token_uniqueness.py +107 -0
  319. nat/profiler/inference_optimization/workflow_runtimes.py +72 -0
  320. nat/profiler/intermediate_property_adapter.py +102 -0
  321. nat/profiler/profile_runner.py +473 -0
  322. nat/profiler/utils.py +184 -0
  323. nat/registry_handlers/__init__.py +0 -0
  324. nat/registry_handlers/local/__init__.py +0 -0
  325. nat/registry_handlers/local/local_handler.py +176 -0
  326. nat/registry_handlers/local/register_local.py +37 -0
  327. nat/registry_handlers/metadata_factory.py +60 -0
  328. nat/registry_handlers/package_utils.py +571 -0
  329. nat/registry_handlers/pypi/__init__.py +0 -0
  330. nat/registry_handlers/pypi/pypi_handler.py +251 -0
  331. nat/registry_handlers/pypi/register_pypi.py +40 -0
  332. nat/registry_handlers/register.py +21 -0
  333. nat/registry_handlers/registry_handler_base.py +157 -0
  334. nat/registry_handlers/rest/__init__.py +0 -0
  335. nat/registry_handlers/rest/register_rest.py +56 -0
  336. nat/registry_handlers/rest/rest_handler.py +237 -0
  337. nat/registry_handlers/schemas/__init__.py +0 -0
  338. nat/registry_handlers/schemas/headers.py +42 -0
  339. nat/registry_handlers/schemas/package.py +68 -0
  340. nat/registry_handlers/schemas/publish.py +68 -0
  341. nat/registry_handlers/schemas/pull.py +82 -0
  342. nat/registry_handlers/schemas/remove.py +36 -0
  343. nat/registry_handlers/schemas/search.py +91 -0
  344. nat/registry_handlers/schemas/status.py +47 -0
  345. nat/retriever/__init__.py +0 -0
  346. nat/retriever/interface.py +41 -0
  347. nat/retriever/milvus/__init__.py +14 -0
  348. nat/retriever/milvus/register.py +81 -0
  349. nat/retriever/milvus/retriever.py +228 -0
  350. nat/retriever/models.py +77 -0
  351. nat/retriever/nemo_retriever/__init__.py +14 -0
  352. nat/retriever/nemo_retriever/register.py +60 -0
  353. nat/retriever/nemo_retriever/retriever.py +190 -0
  354. nat/retriever/register.py +22 -0
  355. nat/runtime/__init__.py +14 -0
  356. nat/runtime/loader.py +219 -0
  357. nat/runtime/runner.py +195 -0
  358. nat/runtime/session.py +162 -0
  359. nat/runtime/user_metadata.py +130 -0
  360. nat/settings/__init__.py +0 -0
  361. nat/settings/global_settings.py +318 -0
  362. nat/test/.namespace +1 -0
  363. nat/tool/__init__.py +0 -0
  364. nat/tool/chat_completion.py +74 -0
  365. nat/tool/code_execution/README.md +151 -0
  366. nat/tool/code_execution/__init__.py +0 -0
  367. nat/tool/code_execution/code_sandbox.py +267 -0
  368. nat/tool/code_execution/local_sandbox/.gitignore +1 -0
  369. nat/tool/code_execution/local_sandbox/Dockerfile.sandbox +60 -0
  370. nat/tool/code_execution/local_sandbox/__init__.py +13 -0
  371. nat/tool/code_execution/local_sandbox/local_sandbox_server.py +198 -0
  372. nat/tool/code_execution/local_sandbox/sandbox.requirements.txt +6 -0
  373. nat/tool/code_execution/local_sandbox/start_local_sandbox.sh +50 -0
  374. nat/tool/code_execution/register.py +74 -0
  375. nat/tool/code_execution/test_code_execution_sandbox.py +414 -0
  376. nat/tool/code_execution/utils.py +100 -0
  377. nat/tool/datetime_tools.py +42 -0
  378. nat/tool/document_search.py +141 -0
  379. nat/tool/github_tools/__init__.py +0 -0
  380. nat/tool/github_tools/create_github_commit.py +133 -0
  381. nat/tool/github_tools/create_github_issue.py +87 -0
  382. nat/tool/github_tools/create_github_pr.py +106 -0
  383. nat/tool/github_tools/get_github_file.py +106 -0
  384. nat/tool/github_tools/get_github_issue.py +166 -0
  385. nat/tool/github_tools/get_github_pr.py +256 -0
  386. nat/tool/github_tools/update_github_issue.py +100 -0
  387. nat/tool/mcp/__init__.py +14 -0
  388. nat/tool/mcp/exceptions.py +142 -0
  389. nat/tool/mcp/mcp_client.py +255 -0
  390. nat/tool/mcp/mcp_tool.py +96 -0
  391. nat/tool/memory_tools/__init__.py +0 -0
  392. nat/tool/memory_tools/add_memory_tool.py +79 -0
  393. nat/tool/memory_tools/delete_memory_tool.py +67 -0
  394. nat/tool/memory_tools/get_memory_tool.py +72 -0
  395. nat/tool/nvidia_rag.py +95 -0
  396. nat/tool/register.py +38 -0
  397. nat/tool/retriever.py +94 -0
  398. nat/tool/server_tools.py +66 -0
  399. nat/utils/__init__.py +0 -0
  400. nat/utils/data_models/__init__.py +0 -0
  401. nat/utils/data_models/schema_validator.py +58 -0
  402. nat/utils/debugging_utils.py +43 -0
  403. nat/utils/dump_distro_mapping.py +32 -0
  404. nat/utils/exception_handlers/__init__.py +0 -0
  405. nat/utils/exception_handlers/automatic_retries.py +289 -0
  406. nat/utils/exception_handlers/mcp.py +211 -0
  407. nat/utils/exception_handlers/schemas.py +114 -0
  408. nat/utils/io/__init__.py +0 -0
  409. nat/utils/io/model_processing.py +28 -0
  410. nat/utils/io/yaml_tools.py +119 -0
  411. nat/utils/log_utils.py +37 -0
  412. nat/utils/metadata_utils.py +74 -0
  413. nat/utils/optional_imports.py +142 -0
  414. nat/utils/producer_consumer_queue.py +178 -0
  415. nat/utils/reactive/__init__.py +0 -0
  416. nat/utils/reactive/base/__init__.py +0 -0
  417. nat/utils/reactive/base/observable_base.py +65 -0
  418. nat/utils/reactive/base/observer_base.py +55 -0
  419. nat/utils/reactive/base/subject_base.py +79 -0
  420. nat/utils/reactive/observable.py +59 -0
  421. nat/utils/reactive/observer.py +76 -0
  422. nat/utils/reactive/subject.py +131 -0
  423. nat/utils/reactive/subscription.py +49 -0
  424. nat/utils/settings/__init__.py +0 -0
  425. nat/utils/settings/global_settings.py +197 -0
  426. nat/utils/string_utils.py +38 -0
  427. nat/utils/type_converter.py +290 -0
  428. nat/utils/type_utils.py +484 -0
  429. nat/utils/url_utils.py +27 -0
  430. nvidia_nat-1.2.0a20250813.dist-info/METADATA +363 -0
  431. nvidia_nat-1.2.0a20250813.dist-info/RECORD +436 -0
  432. nvidia_nat-1.2.0a20250813.dist-info/WHEEL +5 -0
  433. nvidia_nat-1.2.0a20250813.dist-info/entry_points.txt +21 -0
  434. nvidia_nat-1.2.0a20250813.dist-info/licenses/LICENSE-3rd-party.txt +3686 -0
  435. nvidia_nat-1.2.0a20250813.dist-info/licenses/LICENSE.md +201 -0
  436. nvidia_nat-1.2.0a20250813.dist-info/top_level.txt +1 -0
@@ -0,0 +1,709 @@
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 abc
17
+ import datetime
18
+ import typing
19
+ import uuid
20
+ from abc import abstractmethod
21
+ from enum import Enum
22
+
23
+ from pydantic import BaseModel
24
+ from pydantic import ConfigDict
25
+ from pydantic import Discriminator
26
+ from pydantic import Field
27
+ from pydantic import HttpUrl
28
+ from pydantic import conlist
29
+ from pydantic import field_serializer
30
+ from pydantic import field_validator
31
+ from pydantic_core.core_schema import ValidationInfo
32
+
33
+ from nat.data_models.interactive import HumanPrompt
34
+ from nat.utils.type_converter import GlobalTypeConverter
35
+
36
+
37
+ class Request(BaseModel):
38
+ """
39
+ Request is a data model that represents HTTP request attributes.
40
+ """
41
+ model_config = ConfigDict(extra="forbid")
42
+
43
+ method: str | None = Field(default=None,
44
+ description="HTTP method used for the request (e.g., GET, POST, PUT, DELETE).")
45
+ url_path: str | None = Field(default=None, description="URL request path.")
46
+ url_port: int | None = Field(default=None, description="URL request port number.")
47
+ url_scheme: str | None = Field(default=None, description="URL scheme indicating the protocol (e.g., http, https).")
48
+ headers: typing.Any | None = Field(default=None, description="HTTP headers associated with the request.")
49
+ query_params: typing.Any | None = Field(default=None, description="Query parameters included in the request URL.")
50
+ path_params: dict[str, str] | None = Field(default=None,
51
+ description="Path parameters extracted from the request URL.")
52
+ client_host: str | None = Field(default=None, description="Client host address from which the request originated.")
53
+ client_port: int | None = Field(default=None, description="Client port number from which the request originated.")
54
+ cookies: dict[str, str] | None = Field(
55
+ default=None, description="Cookies sent with the request, stored in a dictionary-like object.")
56
+
57
+
58
+ class ChatContentType(str, Enum):
59
+ """
60
+ ChatContentType is an Enum that represents the type of Chat content.
61
+ """
62
+ TEXT = "text"
63
+ IMAGE_URL = "image_url"
64
+ INPUT_AUDIO = "input_audio"
65
+
66
+
67
+ class InputAudio(BaseModel):
68
+ data: str = "default"
69
+ format: str = "default"
70
+
71
+
72
+ class AudioContent(BaseModel):
73
+ model_config = ConfigDict(extra="forbid")
74
+
75
+ type: typing.Literal[ChatContentType.INPUT_AUDIO] = ChatContentType.INPUT_AUDIO
76
+ input_audio: InputAudio = InputAudio()
77
+
78
+
79
+ class ImageUrl(BaseModel):
80
+ url: HttpUrl = HttpUrl(url="http://default.com")
81
+
82
+
83
+ class ImageContent(BaseModel):
84
+ model_config = ConfigDict(extra="forbid")
85
+
86
+ type: typing.Literal[ChatContentType.IMAGE_URL] = ChatContentType.IMAGE_URL
87
+ image_url: ImageUrl = ImageUrl()
88
+
89
+
90
+ class TextContent(BaseModel):
91
+ model_config = ConfigDict(extra="forbid")
92
+
93
+ type: typing.Literal[ChatContentType.TEXT] = ChatContentType.TEXT
94
+ text: str = "default"
95
+
96
+
97
+ class Security(BaseModel):
98
+ model_config = ConfigDict(extra="forbid")
99
+
100
+ api_key: str = "default"
101
+ token: str = "default"
102
+
103
+
104
+ UserContent = typing.Annotated[TextContent | ImageContent | AudioContent, Discriminator("type")]
105
+
106
+
107
+ class Message(BaseModel):
108
+ content: str | list[UserContent]
109
+ role: str
110
+
111
+
112
+ class ChatRequest(BaseModel):
113
+ """
114
+ ChatRequest is a data model that represents a request to the AIQ Toolkit chat API.
115
+ Fully compatible with OpenAI Chat Completions API specification.
116
+ """
117
+
118
+ # Required fields
119
+ messages: typing.Annotated[list[Message], conlist(Message, min_length=1)]
120
+
121
+ # Optional fields (OpenAI Chat Completions API compatible)
122
+ model: str | None = Field(default=None, description="name of the model to use")
123
+ frequency_penalty: float | None = Field(default=0.0,
124
+ description="Penalty for new tokens based on frequency in text")
125
+ logit_bias: dict[str, float] | None = Field(default=None,
126
+ description="Modify likelihood of specified tokens appearing")
127
+ logprobs: bool | None = Field(default=None, description="Whether to return log probabilities")
128
+ top_logprobs: int | None = Field(default=None, description="Number of most likely tokens to return")
129
+ max_tokens: int | None = Field(default=None, description="Maximum number of tokens to generate")
130
+ n: int | None = Field(default=1, description="Number of chat completion choices to generate")
131
+ presence_penalty: float | None = Field(default=0.0, description="Penalty for new tokens based on presence in text")
132
+ response_format: dict[str, typing.Any] | None = Field(default=None, description="Response format specification")
133
+ seed: int | None = Field(default=None, description="Random seed for deterministic sampling")
134
+ service_tier: typing.Literal["auto", "default"] | None = Field(default=None,
135
+ description="Service tier for the request")
136
+ stream: bool | None = Field(default=False, description="Whether to stream partial message deltas")
137
+ stream_options: dict[str, typing.Any] | None = Field(default=None, description="Options for streaming")
138
+ temperature: float | None = Field(default=1.0, description="Sampling temperature between 0 and 2")
139
+ top_p: float | None = Field(default=None, description="Nucleus sampling parameter")
140
+ tools: list[dict[str, typing.Any]] | None = Field(default=None, description="List of tools the model may call")
141
+ tool_choice: str | dict[str, typing.Any] | None = Field(default=None, description="Controls which tool is called")
142
+ parallel_tool_calls: bool | None = Field(default=True, description="Whether to enable parallel function calling")
143
+ user: str | None = Field(default=None, description="Unique identifier representing end-user")
144
+
145
+ model_config = ConfigDict(extra="allow",
146
+ json_schema_extra={
147
+ "example": {
148
+ "model": "nvidia/nemotron",
149
+ "messages": [{
150
+ "role": "user", "content": "who are you?"
151
+ }],
152
+ "temperature": 0.7,
153
+ "stream": False
154
+ }
155
+ })
156
+
157
+ @staticmethod
158
+ def from_string(data: str,
159
+ *,
160
+ model: str | None = None,
161
+ temperature: float | None = None,
162
+ max_tokens: int | None = None,
163
+ top_p: float | None = None) -> "ChatRequest":
164
+
165
+ return ChatRequest(messages=[Message(content=data, role="user")],
166
+ model=model,
167
+ temperature=temperature,
168
+ max_tokens=max_tokens,
169
+ top_p=top_p)
170
+
171
+ @staticmethod
172
+ def from_content(content: list[UserContent],
173
+ *,
174
+ model: str | None = None,
175
+ temperature: float | None = None,
176
+ max_tokens: int | None = None,
177
+ top_p: float | None = None) -> "ChatRequest":
178
+
179
+ return ChatRequest(messages=[Message(content=content, role="user")],
180
+ model=model,
181
+ temperature=temperature,
182
+ max_tokens=max_tokens,
183
+ top_p=top_p)
184
+
185
+
186
+ class ChoiceMessage(BaseModel):
187
+ content: str | None = None
188
+ role: str | None = None
189
+
190
+
191
+ class ChoiceDelta(BaseModel):
192
+ """Delta object for streaming responses (OpenAI-compatible)"""
193
+ content: str | None = None
194
+ role: str | None = None
195
+
196
+
197
+ class Choice(BaseModel):
198
+ model_config = ConfigDict(extra="allow")
199
+
200
+ message: ChoiceMessage | None = None
201
+ delta: ChoiceDelta | None = None
202
+ finish_reason: typing.Literal['stop', 'length', 'tool_calls', 'content_filter', 'function_call'] | None = None
203
+ index: int
204
+ # logprobs: AIQChoiceLogprobs | None = None
205
+
206
+
207
+ class Usage(BaseModel):
208
+ prompt_tokens: int
209
+ completion_tokens: int
210
+ total_tokens: int
211
+
212
+
213
+ class ResponseSerializable(abc.ABC):
214
+ """
215
+ AIQChatResponseSerializable is an abstract class that defines the interface for serializing output for the AIQ
216
+ Toolkit chat streaming API.
217
+ """
218
+
219
+ @abstractmethod
220
+ def get_stream_data(self) -> str:
221
+ pass
222
+
223
+
224
+ class ResponseBaseModelOutput(BaseModel, ResponseSerializable):
225
+
226
+ def get_stream_data(self) -> str:
227
+ return f"data: {self.model_dump_json()}\n\n"
228
+
229
+
230
+ class ResponseBaseModelIntermediate(BaseModel, ResponseSerializable):
231
+
232
+ def get_stream_data(self) -> str:
233
+ return f"intermediate_data: {self.model_dump_json()}\n\n"
234
+
235
+
236
+ class ChatResponse(ResponseBaseModelOutput):
237
+ """
238
+ ChatResponse is a data model that represents a response from the AIQ Toolkit chat API.
239
+ Fully compatible with OpenAI Chat Completions API specification.
240
+ """
241
+
242
+ # Allow extra fields in the model_config to support derived models
243
+ model_config = ConfigDict(extra="allow")
244
+ id: str
245
+ object: str = "chat.completion"
246
+ model: str = ""
247
+ created: datetime.datetime
248
+ choices: list[Choice]
249
+ usage: Usage | None = None
250
+ system_fingerprint: str | None = None
251
+ service_tier: typing.Literal["scale", "default"] | None = None
252
+
253
+ @field_serializer('created')
254
+ def serialize_created(self, created: datetime.datetime) -> int:
255
+ """Serialize datetime to Unix timestamp for OpenAI compatibility"""
256
+ return int(created.timestamp())
257
+
258
+ @staticmethod
259
+ def from_string(data: str,
260
+ *,
261
+ id_: str | None = None,
262
+ object_: str | None = None,
263
+ model: str | None = None,
264
+ created: datetime.datetime | None = None,
265
+ usage: Usage | None = None) -> "ChatResponse":
266
+
267
+ if id_ is None:
268
+ id_ = str(uuid.uuid4())
269
+ if object_ is None:
270
+ object_ = "chat.completion"
271
+ if model is None:
272
+ model = ""
273
+ if created is None:
274
+ created = datetime.datetime.now(datetime.timezone.utc)
275
+
276
+ return ChatResponse(id=id_,
277
+ object=object_,
278
+ model=model,
279
+ created=created,
280
+ choices=[Choice(index=0, message=ChoiceMessage(content=data), finish_reason="stop")],
281
+ usage=usage)
282
+
283
+
284
+ class ChatResponseChunk(ResponseBaseModelOutput):
285
+ """
286
+ ChatResponseChunk is a data model that represents a response chunk from the AIQ Toolkit chat streaming API.
287
+ Fully compatible with OpenAI Chat Completions API specification.
288
+ """
289
+
290
+ # Allow extra fields in the model_config to support derived models
291
+ model_config = ConfigDict(extra="allow")
292
+
293
+ id: str
294
+ choices: list[Choice]
295
+ created: datetime.datetime
296
+ model: str = ""
297
+ object: str = "chat.completion.chunk"
298
+ system_fingerprint: str | None = None
299
+ service_tier: typing.Literal["scale", "default"] | None = None
300
+ usage: Usage | None = None
301
+
302
+ @field_serializer('created')
303
+ def serialize_created(self, created: datetime.datetime) -> int:
304
+ """Serialize datetime to Unix timestamp for OpenAI compatibility"""
305
+ return int(created.timestamp())
306
+
307
+ @staticmethod
308
+ def from_string(data: str,
309
+ *,
310
+ id_: str | None = None,
311
+ created: datetime.datetime | None = None,
312
+ model: str | None = None,
313
+ object_: str | None = None) -> "ChatResponseChunk":
314
+
315
+ if id_ is None:
316
+ id_ = str(uuid.uuid4())
317
+ if created is None:
318
+ created = datetime.datetime.now(datetime.timezone.utc)
319
+ if model is None:
320
+ model = ""
321
+ if object_ is None:
322
+ object_ = "chat.completion.chunk"
323
+
324
+ return ChatResponseChunk(id=id_,
325
+ choices=[Choice(index=0, message=ChoiceMessage(content=data), finish_reason="stop")],
326
+ created=created,
327
+ model=model,
328
+ object=object_)
329
+
330
+ @staticmethod
331
+ def create_streaming_chunk(content: str,
332
+ *,
333
+ id_: str | None = None,
334
+ created: datetime.datetime | None = None,
335
+ model: str | None = None,
336
+ role: str | None = None,
337
+ finish_reason: str | None = None,
338
+ usage: Usage | None = None,
339
+ system_fingerprint: str | None = None) -> "ChatResponseChunk":
340
+ """Create an OpenAI-compatible streaming chunk"""
341
+ if id_ is None:
342
+ id_ = str(uuid.uuid4())
343
+ if created is None:
344
+ created = datetime.datetime.now(datetime.timezone.utc)
345
+ if model is None:
346
+ model = ""
347
+
348
+ delta = ChoiceDelta(content=content, role=role) if content is not None or role is not None else ChoiceDelta()
349
+
350
+ return ChatResponseChunk(id=id_,
351
+ choices=[Choice(index=0, message=None, delta=delta, finish_reason=finish_reason)],
352
+ created=created,
353
+ model=model,
354
+ object="chat.completion.chunk",
355
+ usage=usage,
356
+ system_fingerprint=system_fingerprint)
357
+
358
+
359
+ class ResponseIntermediateStep(ResponseBaseModelIntermediate):
360
+ """
361
+ ResponseSerializedStep is a data model that represents a serialized step in the AIQ Toolkit chat streaming API.
362
+ """
363
+
364
+ # Allow extra fields in the model_config to support derived models
365
+ model_config = ConfigDict(extra="allow")
366
+
367
+ id: str
368
+ parent_id: str | None = None
369
+ type: str = "markdown"
370
+ name: str
371
+ payload: str
372
+
373
+
374
+ class ResponsePayloadOutput(BaseModel, ResponseSerializable):
375
+
376
+ payload: typing.Any
377
+
378
+ def get_stream_data(self) -> str:
379
+
380
+ if (isinstance(self.payload, BaseModel)):
381
+ return f"data: {self.payload.model_dump_json()}\n\n"
382
+
383
+ return f"data: {self.payload}\n\n"
384
+
385
+
386
+ class GenerateResponse(BaseModel):
387
+ # Allow extra fields in the model_config to support derived models
388
+ model_config = ConfigDict(extra="allow")
389
+
390
+ # (fixme) define the intermediate step model
391
+ intermediate_steps: list[tuple] | None = None
392
+ output: str
393
+ value: str | None = "default"
394
+
395
+
396
+ class UserMessageContentRoleType(str, Enum):
397
+ USER = "user"
398
+ ASSISTANT = "assistant"
399
+
400
+
401
+ class WebSocketMessageType(str, Enum):
402
+ """
403
+ WebSocketMessageType is an Enum that represents WebSocket Message types.
404
+ """
405
+ USER_MESSAGE = "user_message"
406
+ RESPONSE_MESSAGE = "system_response_message"
407
+ INTERMEDIATE_STEP_MESSAGE = "system_intermediate_message"
408
+ SYSTEM_INTERACTION_MESSAGE = "system_interaction_message"
409
+ USER_INTERACTION_MESSAGE = "user_interaction_message"
410
+ ERROR_MESSAGE = "error_message"
411
+
412
+
413
+ class WorkflowSchemaType(str, Enum):
414
+ """
415
+ WorkflowSchemaType is an Enum that represents Workkflow response types.
416
+ """
417
+ GENERATE_STREAM = "generate_stream"
418
+ CHAT_STREAM = "chat_stream"
419
+ GENERATE = "generate"
420
+ CHAT = "chat"
421
+
422
+
423
+ class WebSocketMessageStatus(str, Enum):
424
+ """
425
+ WebSocketMessageStatus is an Enum that represents the status of a WebSocket message.
426
+ """
427
+ IN_PROGRESS = "in_progress"
428
+ COMPLETE = "complete"
429
+
430
+
431
+ class UserMessages(BaseModel):
432
+ model_config = ConfigDict(extra="forbid")
433
+
434
+ role: UserMessageContentRoleType
435
+ content: list[UserContent]
436
+
437
+
438
+ class UserMessageContent(BaseModel):
439
+ model_config = ConfigDict(extra="forbid")
440
+ messages: list[UserMessages]
441
+
442
+
443
+ class User(BaseModel):
444
+ model_config = ConfigDict(extra="forbid")
445
+
446
+ name: str = "default"
447
+ email: str = "default"
448
+
449
+
450
+ class ErrorTypes(str, Enum):
451
+ UNKNOWN_ERROR = "unknown_error"
452
+ INVALID_MESSAGE = "invalid_message"
453
+ INVALID_MESSAGE_TYPE = "invalid_message_type"
454
+ INVALID_USER_MESSAGE_CONTENT = "invalid_user_message_content"
455
+ INVALID_DATA_CONTENT = "invalid_data_content"
456
+
457
+
458
+ class Error(BaseModel):
459
+ model_config = ConfigDict(extra="forbid")
460
+
461
+ code: ErrorTypes = ErrorTypes.UNKNOWN_ERROR
462
+ message: str = "default"
463
+ details: str = "default"
464
+
465
+
466
+ class WebSocketUserMessage(BaseModel):
467
+ """
468
+ For more details, refer to the API documentation:
469
+ docs/source/developer_guide/websockets.md
470
+ """
471
+ # Allow extra fields in the model_config to support derived models
472
+ model_config = ConfigDict(extra="allow")
473
+
474
+ type: typing.Literal[WebSocketMessageType.USER_MESSAGE]
475
+ schema_type: WorkflowSchemaType
476
+ id: str = "default"
477
+ conversation_id: str | None = None
478
+ content: UserMessageContent
479
+ user: User = User()
480
+ security: Security = Security()
481
+ error: Error = Error()
482
+ schema_version: str = "1.0.0"
483
+ timestamp: str = str(datetime.datetime.now(datetime.timezone.utc))
484
+
485
+
486
+ class WebSocketUserInteractionResponseMessage(BaseModel):
487
+ """
488
+ For more details, refer to the API documentation:
489
+ docs/source/developer_guide/websockets.md
490
+ """
491
+ type: typing.Literal[WebSocketMessageType.USER_INTERACTION_MESSAGE]
492
+ id: str = "default"
493
+ thread_id: str = "default"
494
+ content: UserMessageContent
495
+ user: User = User()
496
+ security: Security = Security()
497
+ error: Error = Error()
498
+ schema_version: str = "1.0.0"
499
+ timestamp: str = str(datetime.datetime.now(datetime.timezone.utc))
500
+
501
+
502
+ class SystemIntermediateStepContent(BaseModel):
503
+ model_config = ConfigDict(extra="forbid")
504
+ name: str
505
+ payload: str
506
+
507
+
508
+ class WebSocketSystemIntermediateStepMessage(BaseModel):
509
+ """
510
+ For more details, refer to the API documentation:
511
+ docs/source/developer_guide/websockets.md
512
+ """
513
+ # Allow extra fields in the model_config to support derived models
514
+ model_config = ConfigDict(extra="allow")
515
+
516
+ type: typing.Literal[WebSocketMessageType.INTERMEDIATE_STEP_MESSAGE]
517
+ id: str = "default"
518
+ thread_id: str | None = "default"
519
+ parent_id: str = "default"
520
+ intermediate_parent_id: str | None = "default"
521
+ update_message_id: str | None = "default"
522
+ conversation_id: str | None = None
523
+ content: SystemIntermediateStepContent
524
+ status: WebSocketMessageStatus
525
+ timestamp: str = str(datetime.datetime.now(datetime.timezone.utc))
526
+
527
+
528
+ class SystemResponseContent(BaseModel):
529
+ model_config = ConfigDict(extra="forbid")
530
+
531
+ text: str | None = None
532
+
533
+
534
+ class WebSocketSystemResponseTokenMessage(BaseModel):
535
+ """
536
+ For more details, refer to the API documentation:
537
+ docs/source/developer_guide/websockets.md
538
+ """
539
+ # Allow extra fields in the model_config to support derived models
540
+ model_config = ConfigDict(extra="allow")
541
+
542
+ type: typing.Literal[WebSocketMessageType.RESPONSE_MESSAGE, WebSocketMessageType.ERROR_MESSAGE]
543
+ id: str | None = "default"
544
+ thread_id: str | None = "default"
545
+ parent_id: str = "default"
546
+ conversation_id: str | None = None
547
+ content: SystemResponseContent | Error | GenerateResponse
548
+ status: WebSocketMessageStatus
549
+ timestamp: str = str(datetime.datetime.now(datetime.timezone.utc))
550
+
551
+ @field_validator("content")
552
+ @classmethod
553
+ def validate_content_by_type(cls, value: SystemResponseContent | Error | GenerateResponse, info: ValidationInfo):
554
+ if info.data.get("type") == WebSocketMessageType.ERROR_MESSAGE and not isinstance(value, Error):
555
+ raise ValueError(f"Field: content must be 'Error' when type is {WebSocketMessageType.ERROR_MESSAGE}")
556
+
557
+ if info.data.get("type") == WebSocketMessageType.RESPONSE_MESSAGE and not isinstance(
558
+ value, (SystemResponseContent, GenerateResponse)):
559
+ raise ValueError(
560
+ f"Field: content must be 'SystemResponseContent' when type is {WebSocketMessageType.RESPONSE_MESSAGE}")
561
+ return value
562
+
563
+
564
+ class WebSocketSystemInteractionMessage(BaseModel):
565
+ """
566
+ For more details, refer to the API documentation:
567
+ docs/source/developer_guide/websockets.md
568
+ """
569
+ # Allow extra fields in the model_config to support derived models
570
+ model_config = ConfigDict(extra="allow")
571
+
572
+ type: typing.Literal[
573
+ WebSocketMessageType.SYSTEM_INTERACTION_MESSAGE] = WebSocketMessageType.SYSTEM_INTERACTION_MESSAGE
574
+ id: str | None = "default"
575
+ thread_id: str | None = "default"
576
+ parent_id: str = "default"
577
+ conversation_id: str | None = None
578
+ content: HumanPrompt
579
+ status: WebSocketMessageStatus
580
+ timestamp: str = str(datetime.datetime.now(datetime.timezone.utc))
581
+
582
+
583
+ # ======== AIQGenerateResponse Converters ========
584
+
585
+
586
+ def _generate_response_to_str(response: GenerateResponse) -> str:
587
+ return response.output
588
+
589
+
590
+ GlobalTypeConverter.register_converter(_generate_response_to_str)
591
+
592
+
593
+ def _generate_response_to_chat_response(response: GenerateResponse) -> ChatResponse:
594
+ data = response.output
595
+
596
+ # Simulate usage
597
+ prompt_tokens = 0
598
+ usage = Usage(prompt_tokens=prompt_tokens,
599
+ completion_tokens=len(data.split()),
600
+ total_tokens=prompt_tokens + len(data.split()))
601
+
602
+ # Build and return the response
603
+ return ChatResponse.from_string(data, usage=usage)
604
+
605
+
606
+ GlobalTypeConverter.register_converter(_generate_response_to_chat_response)
607
+
608
+
609
+ # ======== AIQChatRequest Converters ========
610
+ def _aiq_chat_request_to_string(data: ChatRequest) -> str:
611
+ if isinstance(data.messages[-1].content, str):
612
+ return data.messages[-1].content
613
+ return str(data.messages[-1].content)
614
+
615
+
616
+ GlobalTypeConverter.register_converter(_aiq_chat_request_to_string)
617
+
618
+
619
+ def _string_to_aiq_chat_request(data: str) -> ChatRequest:
620
+ return ChatRequest.from_string(data, model="")
621
+
622
+
623
+ GlobalTypeConverter.register_converter(_string_to_aiq_chat_request)
624
+
625
+
626
+ # ======== AIQChatResponse Converters ========
627
+ def _aiq_chat_response_to_string(data: ChatResponse) -> str:
628
+ return data.choices[0].message.content or ""
629
+
630
+
631
+ GlobalTypeConverter.register_converter(_aiq_chat_response_to_string)
632
+
633
+
634
+ def _string_to_aiq_chat_response(data: str) -> ChatResponse:
635
+ '''Converts a string to an AIQChatResponse object'''
636
+
637
+ # Simulate usage
638
+ prompt_tokens = 0
639
+ usage = Usage(prompt_tokens=prompt_tokens,
640
+ completion_tokens=len(data.split()),
641
+ total_tokens=prompt_tokens + len(data.split()))
642
+
643
+ # Build and return the response
644
+ return ChatResponse.from_string(data, usage=usage)
645
+
646
+
647
+ GlobalTypeConverter.register_converter(_string_to_aiq_chat_response)
648
+
649
+
650
+ def _chat_response_to_chat_response_chunk(data: ChatResponse) -> ChatResponseChunk:
651
+ # Preserve original message structure for backward compatibility
652
+ return ChatResponseChunk(id=data.id, choices=data.choices, created=data.created, model=data.model)
653
+
654
+
655
+ GlobalTypeConverter.register_converter(_chat_response_to_chat_response_chunk)
656
+
657
+
658
+ # ======== ChatResponseChunk Converters ========
659
+ def _chat_response_chunk_to_string(data: ChatResponseChunk) -> str:
660
+ if data.choices and len(data.choices) > 0:
661
+ choice = data.choices[0]
662
+ if choice.delta and choice.delta.content:
663
+ return choice.delta.content
664
+ if choice.message and choice.message.content:
665
+ return choice.message.content
666
+ return ""
667
+
668
+
669
+ GlobalTypeConverter.register_converter(_chat_response_chunk_to_string)
670
+
671
+
672
+ def _string_to_aiq_chat_response_chunk(data: str) -> ChatResponseChunk:
673
+ '''Converts a string to an AIQChatResponseChunk object'''
674
+
675
+ # Build and return the response
676
+ return ChatResponseChunk.from_string(data)
677
+
678
+
679
+ GlobalTypeConverter.register_converter(_string_to_aiq_chat_response_chunk)
680
+
681
+
682
+ # ======== AINodeMessageChunk Converters ========
683
+ def _ai_message_chunk_to_aiq_chat_response_chunk(data) -> ChatResponseChunk:
684
+ '''Converts LangChain AINodeMessageChunk to AIQChatResponseChunk'''
685
+ content = ""
686
+ if hasattr(data, 'content') and data.content is not None:
687
+ content = str(data.content)
688
+ elif hasattr(data, 'text') and data.text is not None:
689
+ content = str(data.text)
690
+ elif hasattr(data, 'message') and data.message is not None:
691
+ content = str(data.message)
692
+
693
+ return ChatResponseChunk.create_streaming_chunk(content=content, role="assistant", finish_reason=None)
694
+
695
+
696
+ # Compatibility aliases with previous releases
697
+ AIQChatRequest = ChatRequest
698
+ AIQChoiceMessage = ChoiceMessage
699
+ AIQChoiceDelta = ChoiceDelta
700
+ AIQChoice = Choice
701
+ AIQUsage = Usage
702
+ AIQResponseSerializable = ResponseSerializable
703
+ AIQResponseBaseModelOutput = ResponseBaseModelOutput
704
+ AIQResponseBaseModelIntermediate = ResponseBaseModelIntermediate
705
+ AIQChatResponse = ChatResponse
706
+ AIQChatResponseChunk = ChatResponseChunk
707
+ AIQResponseIntermediateStep = ResponseIntermediateStep
708
+ AIQResponsePayloadOutput = ResponsePayloadOutput
709
+ AIQGenerateResponse = GenerateResponse