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,757 @@
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 dataclasses
17
+ import inspect
18
+ import logging
19
+ import warnings
20
+ from contextlib import AbstractAsyncContextManager
21
+ from contextlib import AsyncExitStack
22
+ from contextlib import asynccontextmanager
23
+
24
+ from aiq.builder.builder import Builder
25
+ from aiq.builder.builder import UserManagerHolder
26
+ from aiq.builder.component_utils import build_dependency_sequence
27
+ from aiq.builder.context import AIQContext
28
+ from aiq.builder.context import AIQContextState
29
+ from aiq.builder.embedder import EmbedderProviderInfo
30
+ from aiq.builder.framework_enum import LLMFrameworkEnum
31
+ from aiq.builder.function import Function
32
+ from aiq.builder.function import LambdaFunction
33
+ from aiq.builder.function_info import FunctionInfo
34
+ from aiq.builder.llm import LLMProviderInfo
35
+ from aiq.builder.retriever import RetrieverProviderInfo
36
+ from aiq.builder.workflow import Workflow
37
+ from aiq.cli.type_registry import GlobalTypeRegistry
38
+ from aiq.cli.type_registry import TypeRegistry
39
+ from aiq.data_models.component import ComponentGroup
40
+ from aiq.data_models.component_ref import EmbedderRef
41
+ from aiq.data_models.component_ref import FunctionRef
42
+ from aiq.data_models.component_ref import LLMRef
43
+ from aiq.data_models.component_ref import MemoryRef
44
+ from aiq.data_models.component_ref import RetrieverRef
45
+ from aiq.data_models.config import AIQConfig
46
+ from aiq.data_models.config import GeneralConfig
47
+ from aiq.data_models.embedder import EmbedderBaseConfig
48
+ from aiq.data_models.function import FunctionBaseConfig
49
+ from aiq.data_models.function_dependencies import FunctionDependencies
50
+ from aiq.data_models.llm import LLMBaseConfig
51
+ from aiq.data_models.memory import MemoryBaseConfig
52
+ from aiq.data_models.retriever import RetrieverBaseConfig
53
+ from aiq.data_models.telemetry_exporter import TelemetryExporterBaseConfig
54
+ from aiq.memory.interfaces import MemoryEditor
55
+ from aiq.profiler.decorators.framework_wrapper import chain_wrapped_build_fn
56
+ from aiq.profiler.utils import detect_llm_frameworks_in_build_fn
57
+ from aiq.utils.optional_imports import TelemetryOptionalImportError
58
+ from aiq.utils.optional_imports import try_import_opentelemetry
59
+ from aiq.utils.type_utils import override
60
+
61
+ # SpanExporter is needed to define ConfiguredExporter. Handling when OpenTelemetry is not installed here.
62
+ try:
63
+ opentelemetry = try_import_opentelemetry()
64
+ from opentelemetry.sdk.trace.export import SpanExporter
65
+ except TelemetryOptionalImportError:
66
+ from aiq.utils.optional_imports import DummySpanExporter # pylint: disable=ungrouped-imports
67
+ SpanExporter = DummySpanExporter
68
+
69
+ logger = logging.getLogger(__name__)
70
+
71
+
72
+ @dataclasses.dataclass
73
+ class ConfiguredExporter:
74
+ config: TelemetryExporterBaseConfig
75
+ instance: SpanExporter
76
+
77
+
78
+ @dataclasses.dataclass
79
+ class ConfiguredFunction:
80
+ config: FunctionBaseConfig
81
+ instance: Function
82
+
83
+
84
+ @dataclasses.dataclass
85
+ class ConfiguredLLM:
86
+ config: LLMBaseConfig
87
+ instance: LLMProviderInfo
88
+
89
+
90
+ @dataclasses.dataclass
91
+ class ConfiguredEmbedder:
92
+ config: EmbedderBaseConfig
93
+ instance: EmbedderProviderInfo
94
+
95
+
96
+ @dataclasses.dataclass
97
+ class ConfiguredMemory:
98
+ config: MemoryBaseConfig
99
+ instance: MemoryEditor
100
+
101
+
102
+ @dataclasses.dataclass
103
+ class ConfiguredRetriever:
104
+ config: RetrieverBaseConfig
105
+ instance: RetrieverProviderInfo
106
+
107
+
108
+ # pylint: disable=too-many-public-methods
109
+ class WorkflowBuilder(Builder, AbstractAsyncContextManager):
110
+
111
+ def __init__(self, *, general_config: GeneralConfig | None = None, registry: TypeRegistry | None = None):
112
+
113
+ if general_config is None:
114
+ general_config = GeneralConfig()
115
+
116
+ if registry is None:
117
+ registry = GlobalTypeRegistry.get()
118
+
119
+ self.general_config = general_config
120
+
121
+ self._registry = registry
122
+
123
+ self._logging_handlers: dict[str, logging.Handler] = {}
124
+ self._exporters: dict[str, ConfiguredExporter] = {}
125
+
126
+ self._functions: dict[str, ConfiguredFunction] = {}
127
+ self._workflow: ConfiguredFunction | None = None
128
+
129
+ self._llms: dict[str, ConfiguredLLM] = {}
130
+ self._embedders: dict[str, ConfiguredEmbedder] = {}
131
+ self._memory_clients: dict[str, ConfiguredMemory] = {}
132
+ self._retrievers: dict[str, ConfiguredRetriever] = {}
133
+
134
+ self._context_state = AIQContextState.get()
135
+
136
+ self._exit_stack: AsyncExitStack | None = None
137
+
138
+ # Create a mapping to track function name -> other function names it depends on
139
+ self.function_dependencies: dict[str, FunctionDependencies] = {}
140
+ self.current_function_building: str | None = None
141
+
142
+ async def __aenter__(self):
143
+
144
+ self._exit_stack = AsyncExitStack()
145
+
146
+ # Get the exporter info from the config
147
+ telemetry_config = self.general_config.telemetry
148
+
149
+ for key, logging_config in telemetry_config.logging.items():
150
+ # Use the same pattern as tracing, but for logging
151
+ logging_info = self._registry.get_logging_method(type(logging_config))
152
+ handler = await self._exit_stack.enter_async_context(logging_info.build_fn(logging_config, self))
153
+
154
+ # Type check
155
+ if not isinstance(handler, logging.Handler):
156
+ raise TypeError(f"Expected a logging.Handler from {key}, got {type(handler)}")
157
+
158
+ # Store them in a dict so we can un-register them if needed
159
+ self._logging_handlers[key] = handler
160
+
161
+ # Now attach to AIQ Toolkit's root logger
162
+ logging.getLogger().addHandler(handler)
163
+
164
+ # If tracing is configured, try to import telemetry dependencies and set up tracing
165
+ if telemetry_config.tracing:
166
+ # If the dependencies are not installed, a TelemetryOptionalImportError will be raised
167
+
168
+ # pylint: disable=unused-variable,redefined-outer-name
169
+ opentelemetry = try_import_opentelemetry() # noqa: F841
170
+ from opentelemetry import trace
171
+ from opentelemetry.sdk.trace import TracerProvider
172
+ from opentelemetry.sdk.trace.export import BatchSpanProcessor
173
+
174
+ provider = TracerProvider()
175
+ trace.set_tracer_provider(provider)
176
+
177
+ for key, trace_exporter_config in telemetry_config.tracing.items():
178
+
179
+ exporter_info = self._registry.get_telemetry_exporter(type(trace_exporter_config))
180
+
181
+ instance = await self._exit_stack.enter_async_context(
182
+ exporter_info.build_fn(trace_exporter_config, self))
183
+
184
+ span_processor_instance = BatchSpanProcessor(instance)
185
+ provider.add_span_processor(span_processor_instance)
186
+
187
+ self._exporters[key] = ConfiguredExporter(config=trace_exporter_config, instance=instance)
188
+
189
+ return self
190
+
191
+ async def __aexit__(self, *exc_details):
192
+
193
+ assert self._exit_stack is not None, "Exit stack not initialized"
194
+
195
+ for _, handler in self._logging_handlers.items():
196
+ logging.getLogger().removeHandler(handler)
197
+
198
+ await self._exit_stack.__aexit__(*exc_details)
199
+
200
+ def build(self, entry_function: str | None = None) -> Workflow:
201
+ """
202
+ Creates an instance of a workflow object using the added components and the desired entry function.
203
+
204
+ Parameters
205
+ ----------
206
+ entry_function : str | None, optional
207
+ The function name to use as the entry point for the created workflow. If None, the entry point will be the
208
+ specified workflow function. By default None
209
+
210
+ Returns
211
+ -------
212
+ Workflow
213
+ A created workflow.
214
+
215
+ Raises
216
+ ------
217
+ ValueError
218
+ If the workflow has not been set before building.
219
+ """
220
+
221
+ if (self._workflow is None):
222
+ raise ValueError("Must set a workflow before building")
223
+
224
+ # Build the config from the added objects
225
+ config = AIQConfig(general=self.general_config,
226
+ functions={
227
+ k: v.config
228
+ for k, v in self._functions.items()
229
+ },
230
+ workflow=self._workflow.config,
231
+ llms={
232
+ k: v.config
233
+ for k, v in self._llms.items()
234
+ },
235
+ embedders={
236
+ k: v.config
237
+ for k, v in self._embedders.items()
238
+ },
239
+ memory={
240
+ k: v.config
241
+ for k, v in self._memory_clients.items()
242
+ },
243
+ retrievers={
244
+ k: v.config
245
+ for k, v in self._retrievers.items()
246
+ })
247
+
248
+ if (entry_function is None):
249
+ entry_fn_obj = self.get_workflow()
250
+ else:
251
+ entry_fn_obj = self.get_function(entry_function)
252
+
253
+ workflow = Workflow.from_entry_fn(config=config,
254
+ entry_fn=entry_fn_obj,
255
+ functions={
256
+ k: v.instance
257
+ for k, v in self._functions.items()
258
+ },
259
+ llms={
260
+ k: v.instance
261
+ for k, v in self._llms.items()
262
+ },
263
+ embeddings={
264
+ k: v.instance
265
+ for k, v in self._embedders.items()
266
+ },
267
+ memory={
268
+ k: v.instance
269
+ for k, v in self._memory_clients.items()
270
+ },
271
+ exporters={
272
+ k: v.instance
273
+ for k, v in self._exporters.items()
274
+ },
275
+ retrievers={
276
+ k: v.instance
277
+ for k, v in self._retrievers.items()
278
+ },
279
+ context_state=self._context_state)
280
+
281
+ return workflow
282
+
283
+ def _get_exit_stack(self) -> AsyncExitStack:
284
+
285
+ if self._exit_stack is None:
286
+ raise ValueError(
287
+ "Exit stack not initialized. Did you forget to call `async with WorkflowBuilder() as builder`?")
288
+
289
+ return self._exit_stack
290
+
291
+ async def _build_function(self, name: str, config: FunctionBaseConfig) -> ConfiguredFunction:
292
+ registration = self._registry.get_function(type(config))
293
+
294
+ inner_builder = ChildBuilder(self)
295
+
296
+ # We need to do this for every function because we don't know
297
+ # Where LLama Index Agents are Instantiated and Settings need to
298
+ # be set before the function is built
299
+ # It's only slower the first time because of the import
300
+ # So we can afford to do this for every function
301
+
302
+ llms = {k: v.instance for k, v in self._llms.items()}
303
+ function_frameworks = detect_llm_frameworks_in_build_fn(registration)
304
+
305
+ build_fn = chain_wrapped_build_fn(registration.build_fn, llms, function_frameworks)
306
+
307
+ # Set the currently building function so the ChildBuilder can track dependencies
308
+ self.current_function_building = config.type
309
+ # Empty set of dependencies for the current function
310
+ self.function_dependencies[config.type] = FunctionDependencies()
311
+
312
+ build_result = await self._get_exit_stack().enter_async_context(build_fn(config, inner_builder))
313
+
314
+ self.function_dependencies[name] = inner_builder.dependencies
315
+
316
+ # If the build result is a function, wrap it in a FunctionInfo
317
+ if inspect.isfunction(build_result):
318
+
319
+ build_result = FunctionInfo.from_fn(build_result)
320
+
321
+ if (isinstance(build_result, FunctionInfo)):
322
+ # Create the function object
323
+ build_result = LambdaFunction.from_info(config=config, info=build_result)
324
+
325
+ if (not isinstance(build_result, Function)):
326
+ raise ValueError("Expected a function, FunctionInfo object, or FunctionBase object to be "
327
+ f"returned from the function builder. Got {type(build_result)}")
328
+
329
+ return ConfiguredFunction(config=config, instance=build_result)
330
+
331
+ @override
332
+ async def add_function(self, name: str | FunctionRef, config: FunctionBaseConfig) -> Function:
333
+
334
+ if (name in self._functions):
335
+ raise ValueError(f"Function `{name}` already exists in the list of functions")
336
+
337
+ build_result = await self._build_function(name=name, config=config)
338
+
339
+ self._functions[name] = build_result
340
+
341
+ return build_result.instance
342
+
343
+ @override
344
+ def get_function(self, name: str | FunctionRef) -> Function:
345
+
346
+ if name not in self._functions:
347
+ raise ValueError(f"Function `{name}` not found")
348
+
349
+ return self._functions[name].instance
350
+
351
+ @override
352
+ def get_function_config(self, name: str | FunctionRef) -> FunctionBaseConfig:
353
+ if name not in self._functions:
354
+ raise ValueError(f"Function `{name}` not found")
355
+
356
+ return self._functions[name].config
357
+
358
+ @override
359
+ async def set_workflow(self, config: FunctionBaseConfig) -> Function:
360
+
361
+ if self._workflow is not None:
362
+ warnings.warn("Overwriting existing workflow")
363
+
364
+ build_result = await self._build_function(name="<workflow>", config=config)
365
+
366
+ self._workflow = build_result
367
+
368
+ return build_result.instance
369
+
370
+ @override
371
+ def get_workflow(self) -> Function:
372
+
373
+ if self._workflow is None:
374
+ raise ValueError("No workflow set")
375
+
376
+ return self._workflow.instance
377
+
378
+ @override
379
+ def get_workflow_config(self) -> FunctionBaseConfig:
380
+ if self._workflow is None:
381
+ raise ValueError("No workflow set")
382
+
383
+ return self._workflow.config
384
+
385
+ @override
386
+ def get_function_dependencies(self, fn_name: str | FunctionRef) -> FunctionDependencies:
387
+ return self.function_dependencies[fn_name]
388
+
389
+ @override
390
+ def get_tool(self, fn_name: str | FunctionRef, wrapper_type: LLMFrameworkEnum | str):
391
+
392
+ if fn_name not in self._functions:
393
+ raise ValueError(f"Function `{fn_name}` not found in list of functions")
394
+
395
+ fn = self._functions[fn_name]
396
+
397
+ try:
398
+ # Using the registry, get the tool wrapper for the requested framework
399
+ tool_wrapper_reg = self._registry.get_tool_wrapper(llm_framework=wrapper_type)
400
+
401
+ # Wrap in the correct wrapper
402
+ return tool_wrapper_reg.build_fn(fn_name, fn.instance, self)
403
+ except Exception as e:
404
+ logger.error("Error fetching tool `%s`", fn_name, exc_info=True)
405
+ raise e
406
+
407
+ @override
408
+ async def add_llm(self, name: str | LLMRef, config: LLMBaseConfig):
409
+
410
+ if (name in self._llms):
411
+ raise ValueError(f"LLM `{name}` already exists in the list of LLMs")
412
+
413
+ try:
414
+ llm_info = self._registry.get_llm_provider(type(config))
415
+
416
+ info_obj = await self._get_exit_stack().enter_async_context(llm_info.build_fn(config, self))
417
+
418
+ self._llms[name] = ConfiguredLLM(config=config, instance=info_obj)
419
+ except Exception as e:
420
+ logger.error("Error adding llm `%s` with config `%s`", name, config, exc_info=True)
421
+ raise e
422
+
423
+ @override
424
+ async def get_llm(self, llm_name: str | LLMRef, wrapper_type: LLMFrameworkEnum | str):
425
+
426
+ if (llm_name not in self._llms):
427
+ raise ValueError(f"LLM `{llm_name}` not found")
428
+
429
+ try:
430
+ # Get llm info
431
+ llm_info = self._llms[llm_name]
432
+
433
+ # Generate wrapped client from registered client info
434
+ client_info = self._registry.get_llm_client(config_type=type(llm_info.config), wrapper_type=wrapper_type)
435
+
436
+ client = await self._get_exit_stack().enter_async_context(client_info.build_fn(llm_info.config, self))
437
+
438
+ # Return a frameworks specific client
439
+ return client
440
+ except Exception as e:
441
+ logger.error("Error getting llm `%s` with wrapper `%s`", llm_name, wrapper_type, exc_info=True)
442
+ raise e
443
+
444
+ @override
445
+ def get_llm_config(self, llm_name: str | LLMRef) -> LLMBaseConfig:
446
+
447
+ if llm_name not in self._llms:
448
+ raise ValueError(f"LLM `{llm_name}` not found")
449
+
450
+ # Return the tool configuration object
451
+ return self._llms[llm_name].config
452
+
453
+ @override
454
+ async def add_embedder(self, name: str | EmbedderRef, config: EmbedderBaseConfig):
455
+
456
+ if (name in self._embedders):
457
+ raise ValueError(f"Embedder `{name}` already exists in the list of embedders")
458
+
459
+ try:
460
+ embedder_info = self._registry.get_embedder_provider(type(config))
461
+
462
+ info_obj = await self._get_exit_stack().enter_async_context(embedder_info.build_fn(config, self))
463
+
464
+ self._embedders[name] = ConfiguredEmbedder(config=config, instance=info_obj)
465
+ except Exception as e:
466
+ logger.error("Error adding embedder `%s` with config `%s`", name, config, exc_info=True)
467
+
468
+ raise e
469
+
470
+ @override
471
+ async def get_embedder(self, embedder_name: str | EmbedderRef, wrapper_type: LLMFrameworkEnum | str):
472
+
473
+ if (embedder_name not in self._embedders):
474
+ raise ValueError(f"Embedder `{embedder_name}` not found")
475
+
476
+ try:
477
+ # Get embedder info
478
+ embedder_info = self._embedders[embedder_name]
479
+
480
+ # Generate wrapped client from registered client info
481
+ client_info = self._registry.get_embedder_client(config_type=type(embedder_info.config),
482
+ wrapper_type=wrapper_type)
483
+ client = await self._get_exit_stack().enter_async_context(client_info.build_fn(embedder_info.config, self))
484
+
485
+ # Return a frameworks specific client
486
+ return client
487
+ except Exception as e:
488
+ logger.error("Error getting embedder `%s` with wrapper `%s`", embedder_name, wrapper_type, exc_info=True)
489
+ raise e
490
+
491
+ @override
492
+ def get_embedder_config(self, embedder_name: str | EmbedderRef) -> EmbedderBaseConfig:
493
+
494
+ if embedder_name not in self._embedders:
495
+ raise ValueError(f"Tool `{embedder_name}` not found")
496
+
497
+ # Return the tool configuration object
498
+ return self._embedders[embedder_name].config
499
+
500
+ @override
501
+ async def add_memory_client(self, name: str | MemoryRef, config: MemoryBaseConfig) -> MemoryEditor:
502
+
503
+ if (name in self._memory_clients):
504
+ raise ValueError(f"Memory `{name}` already exists in the list of memories")
505
+
506
+ memory_info = self._registry.get_memory(type(config))
507
+
508
+ info_obj = await self._get_exit_stack().enter_async_context(memory_info.build_fn(config, self))
509
+
510
+ self._memory_clients[name] = ConfiguredMemory(config=config, instance=info_obj)
511
+
512
+ return info_obj
513
+
514
+ @override
515
+ def get_memory_client(self, memory_name: str | MemoryRef) -> MemoryEditor:
516
+ """
517
+ Return the instantiated memory client for the given name.
518
+ """
519
+ if memory_name not in self._memory_clients:
520
+ raise ValueError(f"Memory `{memory_name}` not found")
521
+
522
+ return self._memory_clients[memory_name].instance
523
+
524
+ @override
525
+ def get_memory_client_config(self, memory_name: str | MemoryRef) -> MemoryBaseConfig:
526
+
527
+ if memory_name not in self._memory_clients:
528
+ raise ValueError(f"Memory `{memory_name}` not found")
529
+
530
+ # Return the tool configuration object
531
+ return self._memory_clients[memory_name].config
532
+
533
+ @override
534
+ async def add_retriever(self, name: str | RetrieverRef, config: RetrieverBaseConfig):
535
+
536
+ if (name in self._retrievers):
537
+ raise ValueError(f"Retriever '{name}' already exists in the list of retrievers")
538
+
539
+ try:
540
+ retriever_info = self._registry.get_retriever_provider(type(config))
541
+
542
+ info_obj = await self._get_exit_stack().enter_async_context(retriever_info.build_fn(config, self))
543
+
544
+ self._retrievers[name] = ConfiguredRetriever(config=config, instance=info_obj)
545
+
546
+ except Exception as e:
547
+ logger.error("Error adding retriever `%s` with config `%s`", name, config, exc_info=True)
548
+
549
+ raise e
550
+
551
+ # return info_obj
552
+
553
+ @override
554
+ async def get_retriever(self,
555
+ retriever_name: str | RetrieverRef,
556
+ wrapper_type: LLMFrameworkEnum | str | None = None):
557
+
558
+ if retriever_name not in self._retrievers:
559
+ raise ValueError(f"Retriever '{retriever_name}' not found")
560
+
561
+ try:
562
+ # Get retriever info
563
+ retriever_info = self._retrievers[retriever_name]
564
+
565
+ # Generate wrapped client from registered client info
566
+ client_info = self._registry.get_retriever_client(config_type=type(retriever_info.config),
567
+ wrapper_type=wrapper_type)
568
+
569
+ client = await self._get_exit_stack().enter_async_context(client_info.build_fn(retriever_info.config, self))
570
+
571
+ # Return a frameworks specific client
572
+ return client
573
+ except Exception as e:
574
+ logger.error("Error getting retriever `%s` with wrapper `%s`", retriever_name, wrapper_type, exc_info=True)
575
+ raise e
576
+
577
+ @override
578
+ async def get_retriever_config(self, retriever_name: str | RetrieverRef) -> RetrieverBaseConfig:
579
+
580
+ if retriever_name not in self._retrievers:
581
+ raise ValueError(f"Retriever `{retriever_name}` not found")
582
+
583
+ return self._retrievers[retriever_name].config
584
+
585
+ @override
586
+ def get_user_manager(self):
587
+ return UserManagerHolder(context=AIQContext(self._context_state))
588
+
589
+ async def populate_builder(self, config: AIQConfig, skip_workflow: bool = False):
590
+ """
591
+ Populate the builder with components and optionally set up the workflow.
592
+
593
+ Args:
594
+ config (AIQConfig): The configuration object containing component definitions.
595
+ skip_workflow (bool): If True, skips the workflow instantiation step. Defaults to False.
596
+
597
+ """
598
+ # Generate the build sequence
599
+ build_sequence = build_dependency_sequence(config)
600
+
601
+ # Loop over all objects and add to the workflow builder
602
+ for component_instance in build_sequence:
603
+ # Instantiate a the llm
604
+ if component_instance.component_group == ComponentGroup.LLMS:
605
+ await self.add_llm(component_instance.name, component_instance.config)
606
+ # Instantiate a the embedder
607
+ elif component_instance.component_group == ComponentGroup.EMBEDDERS:
608
+ await self.add_embedder(component_instance.name, component_instance.config)
609
+ # Instantiate a memory client
610
+ elif component_instance.component_group == ComponentGroup.MEMORY:
611
+ await self.add_memory_client(component_instance.name, component_instance.config)
612
+ # Instantiate a retriever client
613
+ elif component_instance.component_group == ComponentGroup.RETRIEVERS:
614
+ await self.add_retriever(component_instance.name, component_instance.config)
615
+ # Instantiate a function
616
+ elif component_instance.component_group == ComponentGroup.FUNCTIONS:
617
+ # If the function is the root, set it as the workflow later
618
+ if (not component_instance.is_root):
619
+ await self.add_function(component_instance.name, component_instance.config)
620
+ else:
621
+ raise ValueError(f"Unknown component group {component_instance.component_group}")
622
+
623
+ # Instantiate the workflow
624
+ if not skip_workflow:
625
+ await self.set_workflow(config.workflow)
626
+
627
+ @classmethod
628
+ @asynccontextmanager
629
+ async def from_config(cls, config: AIQConfig):
630
+
631
+ async with cls(general_config=config.general) as builder:
632
+ await builder.populate_builder(config)
633
+ yield builder
634
+
635
+
636
+ class ChildBuilder(Builder):
637
+
638
+ def __init__(self, workflow_builder: WorkflowBuilder) -> None:
639
+
640
+ self._workflow_builder = workflow_builder
641
+
642
+ self._dependencies = FunctionDependencies()
643
+
644
+ @property
645
+ def dependencies(self) -> FunctionDependencies:
646
+ return self._dependencies
647
+
648
+ @override
649
+ async def add_function(self, name: str, config: FunctionBaseConfig) -> Function:
650
+ return await self._workflow_builder.add_function(name, config)
651
+
652
+ @override
653
+ def get_function(self, name: str) -> Function:
654
+ # If a function tries to get another function, we assume it uses it
655
+ fn = self._workflow_builder.get_function(name)
656
+
657
+ self._dependencies.add_function(name)
658
+
659
+ return fn
660
+
661
+ @override
662
+ def get_function_config(self, name: str) -> FunctionBaseConfig:
663
+ return self._workflow_builder.get_function_config(name)
664
+
665
+ @override
666
+ async def set_workflow(self, config: FunctionBaseConfig) -> Function:
667
+ return await self._workflow_builder.set_workflow(config)
668
+
669
+ @override
670
+ def get_workflow(self) -> Function:
671
+ return self._workflow_builder.get_workflow()
672
+
673
+ @override
674
+ def get_workflow_config(self) -> FunctionBaseConfig:
675
+ return self._workflow_builder.get_workflow_config()
676
+
677
+ @override
678
+ def get_tool(self, fn_name: str, wrapper_type: LLMFrameworkEnum | str):
679
+ # If a function tries to get another function as a tool, we assume it uses it
680
+ fn = self._workflow_builder.get_tool(fn_name, wrapper_type)
681
+
682
+ self._dependencies.add_function(fn_name)
683
+
684
+ return fn
685
+
686
+ @override
687
+ async def add_llm(self, name: str, config: LLMBaseConfig):
688
+ return await self._workflow_builder.add_llm(name, config)
689
+
690
+ @override
691
+ async def get_llm(self, llm_name: str, wrapper_type: LLMFrameworkEnum | str):
692
+ llm = await self._workflow_builder.get_llm(llm_name, wrapper_type)
693
+
694
+ self._dependencies.add_llm(llm_name)
695
+
696
+ return llm
697
+
698
+ @override
699
+ def get_llm_config(self, llm_name: str) -> LLMBaseConfig:
700
+ return self._workflow_builder.get_llm_config(llm_name)
701
+
702
+ @override
703
+ async def add_embedder(self, name: str, config: EmbedderBaseConfig):
704
+ return await self._workflow_builder.add_embedder(name, config)
705
+
706
+ @override
707
+ async def get_embedder(self, embedder_name: str, wrapper_type: LLMFrameworkEnum | str):
708
+ embedder = await self._workflow_builder.get_embedder(embedder_name, wrapper_type)
709
+
710
+ self._dependencies.add_embedder(embedder_name)
711
+
712
+ return embedder
713
+
714
+ @override
715
+ def get_embedder_config(self, embedder_name: str) -> EmbedderBaseConfig:
716
+ return self._workflow_builder.get_embedder_config(embedder_name)
717
+
718
+ @override
719
+ async def add_memory_client(self, name: str, config: MemoryBaseConfig) -> MemoryEditor:
720
+ return await self._workflow_builder.add_memory_client(name, config)
721
+
722
+ @override
723
+ def get_memory_client(self, memory_name: str) -> MemoryEditor:
724
+ """
725
+ Return the instantiated memory client for the given name.
726
+ """
727
+ memory_client = self._workflow_builder.get_memory_client(memory_name)
728
+
729
+ self._dependencies.add_memory_client(memory_name)
730
+
731
+ return memory_client
732
+
733
+ @override
734
+ def get_memory_client_config(self, memory_name: str) -> MemoryBaseConfig:
735
+ return self._workflow_builder.get_memory_client_config(memory_name=memory_name)
736
+
737
+ @override
738
+ async def add_retriever(self, name: str, config: RetrieverBaseConfig):
739
+ return await self._workflow_builder.add_retriever(name, config)
740
+
741
+ @override
742
+ async def get_retriever(self, retriever_name: str, wrapper_type: LLMFrameworkEnum | str | None = None):
743
+ if not wrapper_type:
744
+ return await self._workflow_builder.get_retriever(retriever_name=retriever_name)
745
+ return await self._workflow_builder.get_retriever(retriever_name=retriever_name, wrapper_type=wrapper_type)
746
+
747
+ @override
748
+ async def get_retriever_config(self, retriever_name: str) -> RetrieverBaseConfig:
749
+ return await self._workflow_builder.get_retriever_config(retriever_name=retriever_name)
750
+
751
+ @override
752
+ def get_user_manager(self) -> UserManagerHolder:
753
+ return self._workflow_builder.get_user_manager()
754
+
755
+ @override
756
+ def get_function_dependencies(self, fn_name: str) -> FunctionDependencies:
757
+ return self._workflow_builder.get_function_dependencies(fn_name)