aiqtoolkit 1.1.0__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.

Potentially problematic release.


This version of aiqtoolkit might be problematic. Click here for more details.

Files changed (316) hide show
  1. aiq/agent/__init__.py +0 -0
  2. aiq/agent/base.py +76 -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 +322 -0
  6. aiq/agent/react_agent/output_parser.py +104 -0
  7. aiq/agent/react_agent/prompt.py +46 -0
  8. aiq/agent/react_agent/register.py +148 -0
  9. aiq/agent/reasoning_agent/__init__.py +0 -0
  10. aiq/agent/reasoning_agent/reasoning_agent.py +224 -0
  11. aiq/agent/register.py +23 -0
  12. aiq/agent/rewoo_agent/__init__.py +0 -0
  13. aiq/agent/rewoo_agent/agent.py +410 -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 +123 -0
  18. aiq/agent/tool_calling_agent/register.py +105 -0
  19. aiq/builder/__init__.py +0 -0
  20. aiq/builder/builder.py +223 -0
  21. aiq/builder/component_utils.py +303 -0
  22. aiq/builder/context.py +227 -0
  23. aiq/builder/embedder.py +24 -0
  24. aiq/builder/eval_builder.py +120 -0
  25. aiq/builder/evaluator.py +29 -0
  26. aiq/builder/framework_enum.py +24 -0
  27. aiq/builder/front_end.py +73 -0
  28. aiq/builder/function.py +297 -0
  29. aiq/builder/function_base.py +376 -0
  30. aiq/builder/function_info.py +627 -0
  31. aiq/builder/intermediate_step_manager.py +176 -0
  32. aiq/builder/llm.py +25 -0
  33. aiq/builder/retriever.py +25 -0
  34. aiq/builder/user_interaction_manager.py +71 -0
  35. aiq/builder/workflow.py +143 -0
  36. aiq/builder/workflow_builder.py +757 -0
  37. aiq/cli/__init__.py +14 -0
  38. aiq/cli/cli_utils/__init__.py +0 -0
  39. aiq/cli/cli_utils/config_override.py +231 -0
  40. aiq/cli/cli_utils/validation.py +37 -0
  41. aiq/cli/commands/__init__.py +0 -0
  42. aiq/cli/commands/configure/__init__.py +0 -0
  43. aiq/cli/commands/configure/channel/__init__.py +0 -0
  44. aiq/cli/commands/configure/channel/add.py +28 -0
  45. aiq/cli/commands/configure/channel/channel.py +36 -0
  46. aiq/cli/commands/configure/channel/remove.py +30 -0
  47. aiq/cli/commands/configure/channel/update.py +30 -0
  48. aiq/cli/commands/configure/configure.py +33 -0
  49. aiq/cli/commands/evaluate.py +139 -0
  50. aiq/cli/commands/info/__init__.py +14 -0
  51. aiq/cli/commands/info/info.py +39 -0
  52. aiq/cli/commands/info/list_channels.py +32 -0
  53. aiq/cli/commands/info/list_components.py +129 -0
  54. aiq/cli/commands/info/list_mcp.py +126 -0
  55. aiq/cli/commands/registry/__init__.py +14 -0
  56. aiq/cli/commands/registry/publish.py +88 -0
  57. aiq/cli/commands/registry/pull.py +118 -0
  58. aiq/cli/commands/registry/registry.py +38 -0
  59. aiq/cli/commands/registry/remove.py +108 -0
  60. aiq/cli/commands/registry/search.py +155 -0
  61. aiq/cli/commands/start.py +250 -0
  62. aiq/cli/commands/uninstall.py +83 -0
  63. aiq/cli/commands/validate.py +47 -0
  64. aiq/cli/commands/workflow/__init__.py +14 -0
  65. aiq/cli/commands/workflow/templates/__init__.py.j2 +0 -0
  66. aiq/cli/commands/workflow/templates/config.yml.j2 +16 -0
  67. aiq/cli/commands/workflow/templates/pyproject.toml.j2 +22 -0
  68. aiq/cli/commands/workflow/templates/register.py.j2 +5 -0
  69. aiq/cli/commands/workflow/templates/workflow.py.j2 +36 -0
  70. aiq/cli/commands/workflow/workflow.py +37 -0
  71. aiq/cli/commands/workflow/workflow_commands.py +313 -0
  72. aiq/cli/entrypoint.py +133 -0
  73. aiq/cli/main.py +44 -0
  74. aiq/cli/register_workflow.py +408 -0
  75. aiq/cli/type_registry.py +879 -0
  76. aiq/data_models/__init__.py +14 -0
  77. aiq/data_models/api_server.py +588 -0
  78. aiq/data_models/common.py +143 -0
  79. aiq/data_models/component.py +46 -0
  80. aiq/data_models/component_ref.py +135 -0
  81. aiq/data_models/config.py +349 -0
  82. aiq/data_models/dataset_handler.py +122 -0
  83. aiq/data_models/discovery_metadata.py +286 -0
  84. aiq/data_models/embedder.py +26 -0
  85. aiq/data_models/evaluate.py +104 -0
  86. aiq/data_models/evaluator.py +26 -0
  87. aiq/data_models/front_end.py +26 -0
  88. aiq/data_models/function.py +30 -0
  89. aiq/data_models/function_dependencies.py +64 -0
  90. aiq/data_models/interactive.py +237 -0
  91. aiq/data_models/intermediate_step.py +269 -0
  92. aiq/data_models/invocation_node.py +38 -0
  93. aiq/data_models/llm.py +26 -0
  94. aiq/data_models/logging.py +26 -0
  95. aiq/data_models/memory.py +26 -0
  96. aiq/data_models/profiler.py +53 -0
  97. aiq/data_models/registry_handler.py +26 -0
  98. aiq/data_models/retriever.py +30 -0
  99. aiq/data_models/step_adaptor.py +64 -0
  100. aiq/data_models/streaming.py +33 -0
  101. aiq/data_models/swe_bench_model.py +54 -0
  102. aiq/data_models/telemetry_exporter.py +26 -0
  103. aiq/embedder/__init__.py +0 -0
  104. aiq/embedder/langchain_client.py +41 -0
  105. aiq/embedder/nim_embedder.py +58 -0
  106. aiq/embedder/openai_embedder.py +42 -0
  107. aiq/embedder/register.py +24 -0
  108. aiq/eval/__init__.py +14 -0
  109. aiq/eval/config.py +42 -0
  110. aiq/eval/dataset_handler/__init__.py +0 -0
  111. aiq/eval/dataset_handler/dataset_downloader.py +106 -0
  112. aiq/eval/dataset_handler/dataset_filter.py +52 -0
  113. aiq/eval/dataset_handler/dataset_handler.py +169 -0
  114. aiq/eval/evaluate.py +325 -0
  115. aiq/eval/evaluator/__init__.py +14 -0
  116. aiq/eval/evaluator/evaluator_model.py +44 -0
  117. aiq/eval/intermediate_step_adapter.py +93 -0
  118. aiq/eval/rag_evaluator/__init__.py +0 -0
  119. aiq/eval/rag_evaluator/evaluate.py +138 -0
  120. aiq/eval/rag_evaluator/register.py +138 -0
  121. aiq/eval/register.py +23 -0
  122. aiq/eval/remote_workflow.py +128 -0
  123. aiq/eval/runtime_event_subscriber.py +52 -0
  124. aiq/eval/swe_bench_evaluator/__init__.py +0 -0
  125. aiq/eval/swe_bench_evaluator/evaluate.py +215 -0
  126. aiq/eval/swe_bench_evaluator/register.py +36 -0
  127. aiq/eval/trajectory_evaluator/__init__.py +0 -0
  128. aiq/eval/trajectory_evaluator/evaluate.py +118 -0
  129. aiq/eval/trajectory_evaluator/register.py +40 -0
  130. aiq/eval/tunable_rag_evaluator/__init__.py +0 -0
  131. aiq/eval/tunable_rag_evaluator/evaluate.py +263 -0
  132. aiq/eval/tunable_rag_evaluator/register.py +50 -0
  133. aiq/eval/utils/__init__.py +0 -0
  134. aiq/eval/utils/output_uploader.py +131 -0
  135. aiq/eval/utils/tqdm_position_registry.py +40 -0
  136. aiq/front_ends/__init__.py +14 -0
  137. aiq/front_ends/console/__init__.py +14 -0
  138. aiq/front_ends/console/console_front_end_config.py +32 -0
  139. aiq/front_ends/console/console_front_end_plugin.py +107 -0
  140. aiq/front_ends/console/register.py +25 -0
  141. aiq/front_ends/cron/__init__.py +14 -0
  142. aiq/front_ends/fastapi/__init__.py +14 -0
  143. aiq/front_ends/fastapi/fastapi_front_end_config.py +150 -0
  144. aiq/front_ends/fastapi/fastapi_front_end_plugin.py +103 -0
  145. aiq/front_ends/fastapi/fastapi_front_end_plugin_worker.py +607 -0
  146. aiq/front_ends/fastapi/intermediate_steps_subscriber.py +80 -0
  147. aiq/front_ends/fastapi/job_store.py +161 -0
  148. aiq/front_ends/fastapi/main.py +70 -0
  149. aiq/front_ends/fastapi/message_handler.py +279 -0
  150. aiq/front_ends/fastapi/message_validator.py +345 -0
  151. aiq/front_ends/fastapi/register.py +25 -0
  152. aiq/front_ends/fastapi/response_helpers.py +195 -0
  153. aiq/front_ends/fastapi/step_adaptor.py +320 -0
  154. aiq/front_ends/fastapi/websocket.py +148 -0
  155. aiq/front_ends/mcp/__init__.py +14 -0
  156. aiq/front_ends/mcp/mcp_front_end_config.py +32 -0
  157. aiq/front_ends/mcp/mcp_front_end_plugin.py +93 -0
  158. aiq/front_ends/mcp/register.py +27 -0
  159. aiq/front_ends/mcp/tool_converter.py +242 -0
  160. aiq/front_ends/register.py +22 -0
  161. aiq/front_ends/simple_base/__init__.py +14 -0
  162. aiq/front_ends/simple_base/simple_front_end_plugin_base.py +52 -0
  163. aiq/llm/__init__.py +0 -0
  164. aiq/llm/nim_llm.py +45 -0
  165. aiq/llm/openai_llm.py +45 -0
  166. aiq/llm/register.py +22 -0
  167. aiq/llm/utils/__init__.py +14 -0
  168. aiq/llm/utils/env_config_value.py +94 -0
  169. aiq/llm/utils/error.py +17 -0
  170. aiq/memory/__init__.py +20 -0
  171. aiq/memory/interfaces.py +183 -0
  172. aiq/memory/models.py +112 -0
  173. aiq/meta/module_to_distro.json +3 -0
  174. aiq/meta/pypi.md +58 -0
  175. aiq/observability/__init__.py +0 -0
  176. aiq/observability/async_otel_listener.py +429 -0
  177. aiq/observability/register.py +99 -0
  178. aiq/plugins/.namespace +1 -0
  179. aiq/profiler/__init__.py +0 -0
  180. aiq/profiler/callbacks/__init__.py +0 -0
  181. aiq/profiler/callbacks/agno_callback_handler.py +295 -0
  182. aiq/profiler/callbacks/base_callback_class.py +20 -0
  183. aiq/profiler/callbacks/langchain_callback_handler.py +278 -0
  184. aiq/profiler/callbacks/llama_index_callback_handler.py +205 -0
  185. aiq/profiler/callbacks/semantic_kernel_callback_handler.py +238 -0
  186. aiq/profiler/callbacks/token_usage_base_model.py +27 -0
  187. aiq/profiler/data_frame_row.py +51 -0
  188. aiq/profiler/decorators/__init__.py +0 -0
  189. aiq/profiler/decorators/framework_wrapper.py +131 -0
  190. aiq/profiler/decorators/function_tracking.py +254 -0
  191. aiq/profiler/forecasting/__init__.py +0 -0
  192. aiq/profiler/forecasting/config.py +18 -0
  193. aiq/profiler/forecasting/model_trainer.py +75 -0
  194. aiq/profiler/forecasting/models/__init__.py +22 -0
  195. aiq/profiler/forecasting/models/forecasting_base_model.py +40 -0
  196. aiq/profiler/forecasting/models/linear_model.py +196 -0
  197. aiq/profiler/forecasting/models/random_forest_regressor.py +268 -0
  198. aiq/profiler/inference_metrics_model.py +25 -0
  199. aiq/profiler/inference_optimization/__init__.py +0 -0
  200. aiq/profiler/inference_optimization/bottleneck_analysis/__init__.py +0 -0
  201. aiq/profiler/inference_optimization/bottleneck_analysis/nested_stack_analysis.py +452 -0
  202. aiq/profiler/inference_optimization/bottleneck_analysis/simple_stack_analysis.py +258 -0
  203. aiq/profiler/inference_optimization/data_models.py +386 -0
  204. aiq/profiler/inference_optimization/experimental/__init__.py +0 -0
  205. aiq/profiler/inference_optimization/experimental/concurrency_spike_analysis.py +468 -0
  206. aiq/profiler/inference_optimization/experimental/prefix_span_analysis.py +405 -0
  207. aiq/profiler/inference_optimization/llm_metrics.py +212 -0
  208. aiq/profiler/inference_optimization/prompt_caching.py +163 -0
  209. aiq/profiler/inference_optimization/token_uniqueness.py +107 -0
  210. aiq/profiler/inference_optimization/workflow_runtimes.py +72 -0
  211. aiq/profiler/intermediate_property_adapter.py +102 -0
  212. aiq/profiler/profile_runner.py +433 -0
  213. aiq/profiler/utils.py +184 -0
  214. aiq/registry_handlers/__init__.py +0 -0
  215. aiq/registry_handlers/local/__init__.py +0 -0
  216. aiq/registry_handlers/local/local_handler.py +176 -0
  217. aiq/registry_handlers/local/register_local.py +37 -0
  218. aiq/registry_handlers/metadata_factory.py +60 -0
  219. aiq/registry_handlers/package_utils.py +198 -0
  220. aiq/registry_handlers/pypi/__init__.py +0 -0
  221. aiq/registry_handlers/pypi/pypi_handler.py +251 -0
  222. aiq/registry_handlers/pypi/register_pypi.py +40 -0
  223. aiq/registry_handlers/register.py +21 -0
  224. aiq/registry_handlers/registry_handler_base.py +157 -0
  225. aiq/registry_handlers/rest/__init__.py +0 -0
  226. aiq/registry_handlers/rest/register_rest.py +56 -0
  227. aiq/registry_handlers/rest/rest_handler.py +237 -0
  228. aiq/registry_handlers/schemas/__init__.py +0 -0
  229. aiq/registry_handlers/schemas/headers.py +42 -0
  230. aiq/registry_handlers/schemas/package.py +68 -0
  231. aiq/registry_handlers/schemas/publish.py +63 -0
  232. aiq/registry_handlers/schemas/pull.py +82 -0
  233. aiq/registry_handlers/schemas/remove.py +36 -0
  234. aiq/registry_handlers/schemas/search.py +91 -0
  235. aiq/registry_handlers/schemas/status.py +47 -0
  236. aiq/retriever/__init__.py +0 -0
  237. aiq/retriever/interface.py +37 -0
  238. aiq/retriever/milvus/__init__.py +14 -0
  239. aiq/retriever/milvus/register.py +81 -0
  240. aiq/retriever/milvus/retriever.py +228 -0
  241. aiq/retriever/models.py +74 -0
  242. aiq/retriever/nemo_retriever/__init__.py +14 -0
  243. aiq/retriever/nemo_retriever/register.py +60 -0
  244. aiq/retriever/nemo_retriever/retriever.py +190 -0
  245. aiq/retriever/register.py +22 -0
  246. aiq/runtime/__init__.py +14 -0
  247. aiq/runtime/loader.py +188 -0
  248. aiq/runtime/runner.py +176 -0
  249. aiq/runtime/session.py +140 -0
  250. aiq/runtime/user_metadata.py +131 -0
  251. aiq/settings/__init__.py +0 -0
  252. aiq/settings/global_settings.py +318 -0
  253. aiq/test/.namespace +1 -0
  254. aiq/tool/__init__.py +0 -0
  255. aiq/tool/code_execution/__init__.py +0 -0
  256. aiq/tool/code_execution/code_sandbox.py +188 -0
  257. aiq/tool/code_execution/local_sandbox/Dockerfile.sandbox +60 -0
  258. aiq/tool/code_execution/local_sandbox/__init__.py +13 -0
  259. aiq/tool/code_execution/local_sandbox/local_sandbox_server.py +83 -0
  260. aiq/tool/code_execution/local_sandbox/sandbox.requirements.txt +4 -0
  261. aiq/tool/code_execution/local_sandbox/start_local_sandbox.sh +25 -0
  262. aiq/tool/code_execution/register.py +70 -0
  263. aiq/tool/code_execution/utils.py +100 -0
  264. aiq/tool/datetime_tools.py +42 -0
  265. aiq/tool/document_search.py +141 -0
  266. aiq/tool/github_tools/__init__.py +0 -0
  267. aiq/tool/github_tools/create_github_commit.py +133 -0
  268. aiq/tool/github_tools/create_github_issue.py +87 -0
  269. aiq/tool/github_tools/create_github_pr.py +106 -0
  270. aiq/tool/github_tools/get_github_file.py +106 -0
  271. aiq/tool/github_tools/get_github_issue.py +166 -0
  272. aiq/tool/github_tools/get_github_pr.py +256 -0
  273. aiq/tool/github_tools/update_github_issue.py +100 -0
  274. aiq/tool/mcp/__init__.py +14 -0
  275. aiq/tool/mcp/mcp_client.py +220 -0
  276. aiq/tool/mcp/mcp_tool.py +95 -0
  277. aiq/tool/memory_tools/__init__.py +0 -0
  278. aiq/tool/memory_tools/add_memory_tool.py +79 -0
  279. aiq/tool/memory_tools/delete_memory_tool.py +67 -0
  280. aiq/tool/memory_tools/get_memory_tool.py +72 -0
  281. aiq/tool/nvidia_rag.py +95 -0
  282. aiq/tool/register.py +37 -0
  283. aiq/tool/retriever.py +89 -0
  284. aiq/tool/server_tools.py +63 -0
  285. aiq/utils/__init__.py +0 -0
  286. aiq/utils/data_models/__init__.py +0 -0
  287. aiq/utils/data_models/schema_validator.py +58 -0
  288. aiq/utils/debugging_utils.py +43 -0
  289. aiq/utils/exception_handlers/__init__.py +0 -0
  290. aiq/utils/exception_handlers/schemas.py +114 -0
  291. aiq/utils/io/__init__.py +0 -0
  292. aiq/utils/io/yaml_tools.py +119 -0
  293. aiq/utils/metadata_utils.py +74 -0
  294. aiq/utils/optional_imports.py +142 -0
  295. aiq/utils/producer_consumer_queue.py +178 -0
  296. aiq/utils/reactive/__init__.py +0 -0
  297. aiq/utils/reactive/base/__init__.py +0 -0
  298. aiq/utils/reactive/base/observable_base.py +65 -0
  299. aiq/utils/reactive/base/observer_base.py +55 -0
  300. aiq/utils/reactive/base/subject_base.py +79 -0
  301. aiq/utils/reactive/observable.py +59 -0
  302. aiq/utils/reactive/observer.py +76 -0
  303. aiq/utils/reactive/subject.py +131 -0
  304. aiq/utils/reactive/subscription.py +49 -0
  305. aiq/utils/settings/__init__.py +0 -0
  306. aiq/utils/settings/global_settings.py +197 -0
  307. aiq/utils/type_converter.py +232 -0
  308. aiq/utils/type_utils.py +397 -0
  309. aiq/utils/url_utils.py +27 -0
  310. aiqtoolkit-1.1.0.dist-info/METADATA +331 -0
  311. aiqtoolkit-1.1.0.dist-info/RECORD +316 -0
  312. aiqtoolkit-1.1.0.dist-info/WHEEL +5 -0
  313. aiqtoolkit-1.1.0.dist-info/entry_points.txt +17 -0
  314. aiqtoolkit-1.1.0.dist-info/licenses/LICENSE-3rd-party.txt +3686 -0
  315. aiqtoolkit-1.1.0.dist-info/licenses/LICENSE.md +201 -0
  316. aiqtoolkit-1.1.0.dist-info/top_level.txt +1 -0
@@ -0,0 +1,303 @@
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 typing
18
+ from collections.abc import Generator
19
+ from collections.abc import Iterable
20
+
21
+ import networkx as nx
22
+ from pydantic import BaseModel
23
+
24
+ from aiq.data_models.common import TypedBaseModel
25
+ from aiq.data_models.component import ComponentGroup
26
+ from aiq.data_models.component_ref import ComponentRef
27
+ from aiq.data_models.component_ref import ComponentRefNode
28
+ from aiq.data_models.component_ref import generate_instance_id
29
+ from aiq.data_models.config import AIQConfig
30
+ from aiq.data_models.embedder import EmbedderBaseConfig
31
+ from aiq.data_models.function import FunctionBaseConfig
32
+ from aiq.data_models.llm import LLMBaseConfig
33
+ from aiq.data_models.memory import MemoryBaseConfig
34
+ from aiq.data_models.retriever import RetrieverBaseConfig
35
+ from aiq.utils.type_utils import DecomposedType
36
+
37
+ logger = logging.getLogger(__name__)
38
+
39
+ # Order in which we want to process the component groups
40
+ _component_group_order = [
41
+ ComponentGroup.EMBEDDERS,
42
+ ComponentGroup.LLMS,
43
+ ComponentGroup.MEMORY,
44
+ ComponentGroup.RETRIEVERS,
45
+ ComponentGroup.FUNCTIONS
46
+ ]
47
+
48
+
49
+ class ComponentInstanceData(BaseModel):
50
+ """A data model to hold component runtime instance metadata to support generating build sequences.
51
+
52
+ Args:
53
+ component_group (ComponentGroup): The component group in an AIQ Toolkit configuration object.
54
+ name (ComponentRef): The name of the component runtime instance.
55
+ config (TypedBaseModel): The runtime instance's configuration object.
56
+ instance_id (str): Unique identifier for each runtime instance.
57
+ is_root (bool): A flag to indicate if the runtime instance is the root of the workflow.
58
+ """
59
+
60
+ component_group: ComponentGroup
61
+ name: ComponentRef
62
+ config: TypedBaseModel
63
+ instance_id: str
64
+ is_root: bool = False
65
+
66
+
67
+ def iterate_leaf_to_root(graph: nx.DiGraph) -> Generator[ComponentRefNode]:
68
+ """A recursive generator that yields leaf nodes from the bottom to the root of a directed graph.
69
+
70
+ Args:
71
+ graph (nx.DiGraph): A networkx directed graph object.
72
+
73
+ Yields:
74
+ ComponentRefNode: An object contain a ComponentRef and its component group.
75
+ """
76
+
77
+ leaf_nodes = [node for node, degree in graph.out_degree() if degree == 0]
78
+
79
+ if len(leaf_nodes) > 0:
80
+ for leaf_node in leaf_nodes:
81
+ yield leaf_node
82
+ graph.remove_node(leaf_node)
83
+
84
+ yield from iterate_leaf_to_root(graph)
85
+
86
+
87
+ def group_from_component(component: TypedBaseModel) -> ComponentGroup | None:
88
+ """Determines the component group from a runtime instance configuration object.
89
+
90
+ Args:
91
+ component (TypedBaseModel): A runtime instance configuration object.
92
+
93
+ Returns:
94
+ ComponentGroup | None: The component group of the runtime instance configuration object. If the
95
+ component is not a valid runtime instance, None is returned.
96
+ """
97
+
98
+ if (isinstance(component, EmbedderBaseConfig)):
99
+ return ComponentGroup.EMBEDDERS
100
+ if (isinstance(component, FunctionBaseConfig)):
101
+ return ComponentGroup.FUNCTIONS
102
+ if (isinstance(component, LLMBaseConfig)):
103
+ return ComponentGroup.LLMS
104
+ if (isinstance(component, MemoryBaseConfig)):
105
+ return ComponentGroup.MEMORY
106
+ if (isinstance(component, RetrieverBaseConfig)):
107
+ return ComponentGroup.RETRIEVERS
108
+
109
+ return None
110
+
111
+
112
+ def recursive_componentref_discovery(cls: TypedBaseModel, value: typing.Any,
113
+ type_hint: type[typing.Any]) -> Generator[tuple[str, ComponentRefNode]]:
114
+ """Discovers instances of ComponentRefs in a configuration object and updates the dependency graph.
115
+
116
+ Args:
117
+ cls (TypedBaseModel): A configuration object for a runtime instance.
118
+ value (typing.Any): The current traversed value from the configuration object.
119
+ type_hint (type[typing.Any]): The type of the current traversed value from the configuration object.
120
+ """
121
+
122
+ decomposed_type = DecomposedType(type_hint)
123
+
124
+ if (value is None):
125
+ return
126
+
127
+ if ((decomposed_type.origin is None) and (not issubclass(type(value), BaseModel))):
128
+ if issubclass(type(value), ComponentRef):
129
+ instance_id = generate_instance_id(cls)
130
+ value_node = ComponentRefNode(ref_name=value, component_group=value.component_group)
131
+ yield instance_id, value_node
132
+
133
+ elif ((decomposed_type.origin in (tuple, list, set)) and (isinstance(value, Iterable))):
134
+ for v in value:
135
+ yield from recursive_componentref_discovery(cls, v, decomposed_type.args[0])
136
+ elif ((decomposed_type.origin in (dict, type(typing.TypedDict))) and (isinstance(value, dict))):
137
+ for v in value.values():
138
+ yield from recursive_componentref_discovery(cls, v, decomposed_type.args[1])
139
+ elif (issubclass(type(value), BaseModel)):
140
+ for field, field_info in value.model_fields.items():
141
+ field_data = getattr(value, field)
142
+ yield from recursive_componentref_discovery(cls, field_data, field_info.annotation)
143
+ if (decomposed_type.is_union):
144
+ for arg in decomposed_type.args:
145
+ if (isinstance(value, DecomposedType(arg).root)):
146
+ yield from recursive_componentref_discovery(cls, value, arg)
147
+ else:
148
+ for arg in decomposed_type.args:
149
+ yield from recursive_componentref_discovery(cls, value, arg)
150
+
151
+
152
+ def update_dependency_graph(config: "AIQConfig", instance_config: TypedBaseModel,
153
+ dependency_graph: nx.DiGraph) -> nx.DiGraph:
154
+ """Updates the hierarchical component instance dependency graph from a configuration runtime instance.
155
+
156
+ Args:
157
+ config (AIQConfig): An AIQ Toolkit configuration object with runtime instance details.
158
+ instance_config (TypedBaseModel): A component's runtime instance configuration object.
159
+ dependency_graph (nx.DiGraph): A graph tracking runtime instance component dependencies.
160
+
161
+ Returns:
162
+ nx.DiGraph: An dependency graph that has been updated with the provided runtime instance.
163
+ """
164
+
165
+ for field_name, field_info in instance_config.model_fields.items():
166
+
167
+ for instance_id, value_node in recursive_componentref_discovery(
168
+ instance_config,
169
+ getattr(instance_config, field_name),
170
+ field_info.annotation): # type: ignore
171
+
172
+ # add immediate edge
173
+ dependency_graph.add_edge(instance_id, value_node)
174
+ # add dependency edge to ensure connections to leaf nodes exist
175
+ dependency_component_dict = getattr(config, value_node.component_group)
176
+ dependency_component_instance_config = dependency_component_dict.get(value_node.ref_name)
177
+ dependency_component_instance_id = generate_instance_id(dependency_component_instance_config)
178
+ dependency_graph.add_edge(value_node, dependency_component_instance_id)
179
+
180
+ return dependency_graph
181
+
182
+
183
+ def config_to_dependency_objects(config: "AIQConfig") -> tuple[dict[str, ComponentInstanceData], nx.DiGraph]:
184
+ """Generates a map of component runtime instance IDs to use when generating a build sequence.
185
+
186
+ Args:
187
+ config (AIQConfig): The AIQ Toolkit workflow configuration object.
188
+
189
+ Returns:
190
+ tuple[dict[str, ComponentInstanceData], nx.DiGraph]: A tuple containing a map of component runtime instance
191
+ IDs to a component object containing its metadata and a dependency graph of nested components.
192
+ """
193
+
194
+ # Build map of every runtime instances
195
+ dependency_map: dict[str, ComponentInstanceData] = {}
196
+ dependency_graph: nx.DiGraph = nx.DiGraph()
197
+
198
+ # Create the dependency map preserving as much order as we can
199
+ for group in _component_group_order:
200
+
201
+ component_dict = getattr(config, group.value)
202
+
203
+ assert isinstance(component_dict, dict), "Config components must be a dictionary"
204
+
205
+ for component_instance_name, component_instance_config in component_dict.items():
206
+
207
+ instance_id = generate_instance_id(component_instance_config)
208
+ dependency_map[instance_id] = ComponentInstanceData(component_group=group,
209
+ instance_id=instance_id,
210
+ name=component_instance_name,
211
+ config=component_instance_config)
212
+
213
+ dependency_graph = update_dependency_graph(config=config,
214
+ instance_config=component_instance_config,
215
+ dependency_graph=dependency_graph)
216
+
217
+ # Set the workflow flag on the workflow instance (must be last)
218
+ workflow_instance_id = generate_instance_id(config.workflow)
219
+
220
+ dependency_map[workflow_instance_id] = ComponentInstanceData(
221
+ component_group=ComponentGroup.FUNCTIONS,
222
+ instance_id=workflow_instance_id,
223
+ name="<workflow>", # type: ignore
224
+ config=config.workflow,
225
+ is_root=True)
226
+
227
+ dependency_graph = update_dependency_graph(config=config,
228
+ instance_config=config.workflow,
229
+ dependency_graph=dependency_graph)
230
+
231
+ return dependency_map, dependency_graph
232
+
233
+
234
+ def build_dependency_sequence(config: "AIQConfig") -> list[ComponentInstanceData]:
235
+ """Generates the depencency sequence from an AIQ Toolkit configuration object
236
+
237
+ Args:
238
+ config (AIQConfig): An AIQ Toolkit configuration object.
239
+
240
+ Returns:
241
+ list[ComponentInstanceData]: A list representing the instatiation sequence to ensure all valid
242
+ runtime instance references.
243
+ """
244
+
245
+ total_node_count = len(config.embedders) + len(config.functions) + len(config.llms) + len(config.memory) + len(
246
+ config.retrievers) + 1 # +1 for the workflow
247
+
248
+ dependency_map: dict
249
+ dependency_graph: nx.DiGraph
250
+ dependency_map, dependency_graph = config_to_dependency_objects(config=config)
251
+
252
+ dependency_sequence: list[ComponentInstanceData] = []
253
+ instance_ids = set()
254
+ for node in iterate_leaf_to_root(dependency_graph.copy()): # type: ignore
255
+
256
+ if (node not in dependency_sequence):
257
+
258
+ # Convert node to id
259
+ if (isinstance(node, ComponentRefNode) and issubclass(type(node.ref_name), ComponentRef)):
260
+
261
+ component_group_configs = getattr(config, node.component_group.value)
262
+ node_config = component_group_configs.get(node.ref_name, None)
263
+
264
+ # Only add nodes that are valid in the current instance configuration
265
+ if (node_config is None):
266
+ continue
267
+
268
+ component_instance = ComponentInstanceData(
269
+ name=node.ref_name,
270
+ component_group=node.component_group.value, # type: ignore
271
+ config=node_config,
272
+ instance_id=generate_instance_id(node_config))
273
+
274
+ else:
275
+
276
+ component_instance = dependency_map.get(node, None)
277
+
278
+ # Only add nodes that are valid in the current instance configuration
279
+ if (component_instance is None):
280
+ continue
281
+
282
+ if (component_instance.instance_id not in instance_ids):
283
+
284
+ dependency_sequence.append(component_instance)
285
+ instance_ids.add(component_instance.instance_id)
286
+
287
+ remaining_dependency_sequence: list[ComponentInstanceData] = []
288
+
289
+ # Find the remaining nodes that are not in the sequence preserving order
290
+ for instance_id, instance in dependency_map.items():
291
+ if (instance_id not in instance_ids):
292
+ remaining_dependency_sequence.append(instance)
293
+
294
+ # Add the remaining at the front of the sequence
295
+ dependency_sequence = remaining_dependency_sequence + dependency_sequence
296
+
297
+ # Find the root node and make sure it is the last node in the sequence
298
+ dependency_sequence = [x for x in dependency_sequence if not x.is_root
299
+ ] + [x for x in dependency_sequence if x.is_root]
300
+
301
+ assert len(dependency_sequence) == total_node_count, "Dependency sequence generation failed. Report as bug."
302
+
303
+ return dependency_sequence
aiq/builder/context.py ADDED
@@ -0,0 +1,227 @@
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 typing
17
+ import uuid
18
+ from collections.abc import Awaitable
19
+ from collections.abc import Callable
20
+ from contextlib import contextmanager
21
+ from contextvars import ContextVar
22
+
23
+ from aiq.builder.intermediate_step_manager import IntermediateStepManager
24
+ from aiq.builder.user_interaction_manager import AIQUserInteractionManager
25
+ from aiq.data_models.interactive import HumanResponse
26
+ from aiq.data_models.interactive import InteractionPrompt
27
+ from aiq.data_models.intermediate_step import IntermediateStep
28
+ from aiq.data_models.intermediate_step import IntermediateStepPayload
29
+ from aiq.data_models.intermediate_step import IntermediateStepType
30
+ from aiq.data_models.intermediate_step import StreamEventData
31
+ from aiq.data_models.invocation_node import InvocationNode
32
+ from aiq.runtime.user_metadata import RequestAttributes
33
+ from aiq.utils.reactive.subject import Subject
34
+
35
+
36
+ class Singleton(type):
37
+
38
+ def __init__(cls, name, bases, dict): # pylint: disable=W0622
39
+ super(Singleton, cls).__init__(name, bases, dict)
40
+ cls.instance = None
41
+
42
+ def __call__(cls, *args, **kw):
43
+ if cls.instance is None:
44
+ cls.instance = super(Singleton, cls).__call__(*args, **kw)
45
+ return cls.instance
46
+
47
+
48
+ class ActiveFunctionContextManager:
49
+
50
+ def __init__(self):
51
+ self._output: typing.Any | None = None
52
+
53
+ @property
54
+ def output(self) -> typing.Any | None:
55
+ return self._output
56
+
57
+ def set_output(self, output: typing.Any):
58
+ self._output = output
59
+
60
+
61
+ class AIQContextState(metaclass=Singleton):
62
+
63
+ def __init__(self):
64
+ self.input_message: ContextVar[typing.Any] = ContextVar("input_message", default=None)
65
+ self.user_manager: ContextVar[typing.Any] = ContextVar("user_manager", default=None)
66
+ self.metadata: ContextVar[RequestAttributes] = ContextVar("request_attributes", default=RequestAttributes())
67
+ self.event_stream: ContextVar[Subject[IntermediateStep] | None] = ContextVar("event_stream", default=Subject())
68
+ self.active_function: ContextVar[InvocationNode] = ContextVar("active_function",
69
+ default=InvocationNode(function_id="root",
70
+ function_name="root"))
71
+ self.active_span_id_stack: ContextVar[list[str]] = ContextVar("active_span_id_stack", default=["root"])
72
+
73
+ # Default is a lambda no-op which returns NoneType
74
+ self.user_input_callback: ContextVar[Callable[[InteractionPrompt], Awaitable[HumanResponse | None]]
75
+ | None] = ContextVar(
76
+ "user_input_callback",
77
+ default=AIQUserInteractionManager.default_callback_handler)
78
+
79
+ @staticmethod
80
+ def get() -> "AIQContextState":
81
+ return AIQContextState()
82
+
83
+
84
+ class AIQContext:
85
+
86
+ def __init__(self, context: AIQContextState):
87
+ self._context_state = context
88
+
89
+ @property
90
+ def input_message(self):
91
+ """
92
+ Retrieves the input message from the context state.
93
+
94
+ The input_message property is used to access the message stored in the
95
+ context state. This property returns the message as it is currently
96
+ maintained in the context.
97
+
98
+ Returns:
99
+ str: The input message retrieved from the context state.
100
+ """
101
+ return self._context_state.input_message.get()
102
+
103
+ @property
104
+ def user_manager(self):
105
+ """
106
+ Retrieves the user manager instance from the current context state.
107
+
108
+ This property provides access to the user manager through the context
109
+ state, allowing interaction with user management functionalities.
110
+
111
+ Returns:
112
+ UserManager: The instance of the user manager retrieved from the
113
+ context state.
114
+ """
115
+ return self._context_state.user_manager.get()
116
+
117
+ @property
118
+ def metadata(self):
119
+ """
120
+ Retrieves the request attributes instance from the current context state
121
+ providing access to user-defined metadata.
122
+
123
+ Returns:
124
+ RequestAttributes: The instance of the request attributes
125
+ retrieved from the context state.
126
+ """
127
+ return self._context_state.metadata.get()
128
+
129
+ @property
130
+ def user_interaction_manager(self) -> AIQUserInteractionManager:
131
+ """
132
+ Return an instance of AIQUserInteractionManager that uses
133
+ the current context's user_input_callback.
134
+ """
135
+ return AIQUserInteractionManager(self._context_state)
136
+
137
+ @property
138
+ def intermediate_step_manager(self) -> IntermediateStepManager:
139
+ """
140
+ Retrieves the intermediate step manager instance from the current context state.
141
+
142
+ This property provides access to the intermediate step manager through the context
143
+ state, allowing interaction with intermediate step management functionalities.
144
+
145
+ Returns:
146
+ IntermediateStepManager: The instance of the intermediate step manager retrieved
147
+ from the context state.
148
+ """
149
+ return IntermediateStepManager(self._context_state)
150
+
151
+ @contextmanager
152
+ def push_active_function(self, function_name: str, input_data: typing.Any | None):
153
+ """
154
+ Set the 'active_function' in context, push an invocation node,
155
+ AND create an OTel child span for that function call.
156
+ """
157
+ parent_function_node = self._context_state.active_function.get()
158
+ current_function_id = str(uuid.uuid4())
159
+ current_function_node = InvocationNode(function_id=current_function_id,
160
+ function_name=function_name,
161
+ parent_id=parent_function_node.function_id,
162
+ parent_name=parent_function_node.function_name)
163
+
164
+ # 1) Set the active function in the contextvar
165
+ fn_token = self._context_state.active_function.set(current_function_node)
166
+
167
+ # 2) Optionally record function start as an intermediate step
168
+ step_manager = self.intermediate_step_manager
169
+ step_manager.push_intermediate_step(
170
+ IntermediateStepPayload(UUID=current_function_id,
171
+ event_type=IntermediateStepType.FUNCTION_START,
172
+ name=function_name,
173
+ data=StreamEventData(input=input_data)))
174
+
175
+ manager = ActiveFunctionContextManager()
176
+
177
+ try:
178
+ yield manager # run the function body
179
+ finally:
180
+ # 3) Record function end
181
+
182
+ data = StreamEventData(input=input_data, output=manager.output)
183
+
184
+ step_manager.push_intermediate_step(
185
+ IntermediateStepPayload(UUID=current_function_id,
186
+ event_type=IntermediateStepType.FUNCTION_END,
187
+ name=function_name,
188
+ data=data))
189
+
190
+ # 4) Unset the function contextvar
191
+ self._context_state.active_function.reset(fn_token)
192
+
193
+ @property
194
+ def active_function(self) -> InvocationNode:
195
+ """
196
+ Retrieves the active function from the context state.
197
+
198
+ This property is used to access the active function stored in the context
199
+ state. The active function is the function that is currently being executed.
200
+ """
201
+ return self._context_state.active_function.get()
202
+
203
+ @property
204
+ def active_span_id(self) -> str:
205
+ """
206
+ Retrieves the active span ID from the context state.
207
+
208
+ This property provides access to the active span ID stored in the context state. The active span ID represents
209
+ the currently running function/tool/llm/agent/etc and can be used to group telemetry data together.
210
+
211
+ Returns:
212
+ str: The active span ID.
213
+ """
214
+ return self._context_state.active_span_id_stack.get()[-1]
215
+
216
+ @staticmethod
217
+ def get() -> "AIQContext":
218
+ """
219
+ Static method to retrieve the current AIQContext instance.
220
+
221
+ This method creates and returns an instance of the AIQContext class
222
+ by obtaining the current state from the AIQContextState.
223
+
224
+ Returns:
225
+ AIQContext: The created AIQContext instance.
226
+ """
227
+ return AIQContext(AIQContextState.get())
@@ -0,0 +1,24 @@
1
+ # SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
2
+ # SPDX-License-Identifier: Apache-2.0
3
+ #
4
+ # Licensed under the Apache License, Version 2.0 (the "License");
5
+ # you may not use this file except in compliance with the License.
6
+ # You may obtain a copy of the License at
7
+ #
8
+ # http://www.apache.org/licenses/LICENSE-2.0
9
+ #
10
+ # Unless required by applicable law or agreed to in writing, software
11
+ # distributed under the License is distributed on an "AS IS" BASIS,
12
+ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13
+ # See the License for the specific language governing permissions and
14
+ # limitations under the License.
15
+
16
+ from aiq.data_models.embedder import EmbedderBaseConfig
17
+
18
+
19
+ class EmbedderProviderInfo:
20
+
21
+ def __init__(self, *, config: EmbedderBaseConfig, description: str):
22
+ self.config = config
23
+ self.provider_type = type(config).static_type()
24
+ self.description = description
@@ -0,0 +1,120 @@
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 dataclasses
17
+ import logging
18
+ from contextlib import asynccontextmanager
19
+ from pathlib import Path
20
+
21
+ from aiq.builder.builder import EvalBuilder
22
+ from aiq.builder.evaluator import EvaluatorInfo
23
+ from aiq.builder.framework_enum import LLMFrameworkEnum
24
+ from aiq.builder.workflow_builder import WorkflowBuilder
25
+ from aiq.cli.type_registry import TypeRegistry
26
+ from aiq.data_models.config import AIQConfig
27
+ from aiq.data_models.config import GeneralConfig
28
+ from aiq.data_models.evaluate import EvalGeneralConfig
29
+ from aiq.data_models.evaluator import EvaluatorBaseConfig
30
+ from aiq.data_models.function import EmptyFunctionConfig
31
+ from aiq.utils.type_utils import override
32
+
33
+ logger = logging.getLogger(__name__)
34
+
35
+
36
+ @dataclasses.dataclass
37
+ class ConfiguredEvaluator:
38
+ config: EvaluatorBaseConfig
39
+ instance: EvaluatorInfo
40
+
41
+
42
+ class WorkflowEvalBuilder(WorkflowBuilder, EvalBuilder):
43
+
44
+ def __init__(self,
45
+ general_config: GeneralConfig | None = None,
46
+ eval_general_config: EvalGeneralConfig | None = None,
47
+ registry: TypeRegistry | None = None):
48
+ super().__init__(general_config=general_config, registry=registry)
49
+ self.eval_general_config = eval_general_config
50
+ self._evaluators: dict[str, ConfiguredEvaluator] = {}
51
+
52
+ @override
53
+ async def add_evaluator(self, name: str, config: EvaluatorBaseConfig):
54
+ if name in self._evaluators:
55
+ raise ValueError(f"Evaluator `{name}` already exists in the list of evaluators")
56
+
57
+ try:
58
+ evaluator_info = self._registry.get_evaluator(type(config))
59
+ info_obj = await self._get_exit_stack().enter_async_context(evaluator_info.build_fn(config, self))
60
+
61
+ # Store the evaluator
62
+ self._evaluators[name] = ConfiguredEvaluator(config=config, instance=info_obj)
63
+ except Exception as e:
64
+ logger.error("Error %s adding evaluator `%s` with config `%s`", e, name, config, exc_info=True)
65
+ raise
66
+
67
+ @override
68
+ def get_evaluator(self, evaluator_name: str) -> EvaluatorInfo:
69
+
70
+ if (evaluator_name not in self._evaluators):
71
+ raise ValueError(f"Evaluator `{evaluator_name}` not found")
72
+
73
+ return self._evaluators[evaluator_name].instance
74
+
75
+ @override
76
+ def get_evaluator_config(self, evaluator_name: str) -> EvaluatorBaseConfig:
77
+
78
+ if evaluator_name not in self._evaluators:
79
+ raise ValueError(f"Evaluator `{evaluator_name}` not found")
80
+
81
+ # Return the tool configuration object
82
+ return self._evaluators[evaluator_name].config
83
+
84
+ @override
85
+ def get_max_concurrency(self) -> int:
86
+ return self.eval_general_config.max_concurrency
87
+
88
+ @override
89
+ def get_output_dir(self) -> Path:
90
+ return self.eval_general_config.output_dir
91
+
92
+ @override
93
+ def get_all_tools(self, wrapper_type: LLMFrameworkEnum | str):
94
+ tools = []
95
+ tool_wrapper_reg = self._registry.get_tool_wrapper(llm_framework=wrapper_type)
96
+ for fn_name in self._functions:
97
+ fn = self.get_function(fn_name)
98
+ try:
99
+ tools.append(tool_wrapper_reg.build_fn(fn_name, fn, self))
100
+ except Exception:
101
+ logger.exception("Error fetching tool `%s`", fn_name, exc_info=True)
102
+
103
+ return tools
104
+
105
+ async def populate_builder(self, config: AIQConfig):
106
+ # Skip setting workflow if workflow config is EmptyFunctionConfig
107
+ skip_workflow = isinstance(config.workflow, EmptyFunctionConfig)
108
+
109
+ await super().populate_builder(config, skip_workflow)
110
+ # Instantiate the evaluators
111
+ for name, evaluator_config in config.eval.evaluators.items():
112
+ await self.add_evaluator(name, evaluator_config)
113
+
114
+ @classmethod
115
+ @asynccontextmanager
116
+ async def from_config(cls, config: AIQConfig):
117
+
118
+ async with cls(config.general, config.eval.general, registry=None) as builder:
119
+ await builder.populate_builder(config)
120
+ yield builder