aiqtoolkit 1.1.0rc6__py3-none-any.whl → 1.2.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 (319) hide show
  1. aiqtoolkit-1.2.0.dist-info/METADATA +29 -0
  2. aiqtoolkit-1.2.0.dist-info/RECORD +4 -0
  3. {aiqtoolkit-1.1.0rc6.dist-info → aiqtoolkit-1.2.0.dist-info}/WHEEL +1 -1
  4. aiqtoolkit-1.2.0.dist-info/top_level.txt +1 -0
  5. aiq/agent/__init__.py +0 -0
  6. aiq/agent/base.py +0 -76
  7. aiq/agent/dual_node.py +0 -67
  8. aiq/agent/react_agent/__init__.py +0 -0
  9. aiq/agent/react_agent/agent.py +0 -322
  10. aiq/agent/react_agent/output_parser.py +0 -104
  11. aiq/agent/react_agent/prompt.py +0 -46
  12. aiq/agent/react_agent/register.py +0 -148
  13. aiq/agent/reasoning_agent/__init__.py +0 -0
  14. aiq/agent/reasoning_agent/reasoning_agent.py +0 -224
  15. aiq/agent/register.py +0 -23
  16. aiq/agent/rewoo_agent/__init__.py +0 -0
  17. aiq/agent/rewoo_agent/agent.py +0 -410
  18. aiq/agent/rewoo_agent/prompt.py +0 -108
  19. aiq/agent/rewoo_agent/register.py +0 -158
  20. aiq/agent/tool_calling_agent/__init__.py +0 -0
  21. aiq/agent/tool_calling_agent/agent.py +0 -123
  22. aiq/agent/tool_calling_agent/register.py +0 -105
  23. aiq/builder/__init__.py +0 -0
  24. aiq/builder/builder.py +0 -223
  25. aiq/builder/component_utils.py +0 -303
  26. aiq/builder/context.py +0 -227
  27. aiq/builder/embedder.py +0 -24
  28. aiq/builder/eval_builder.py +0 -120
  29. aiq/builder/evaluator.py +0 -29
  30. aiq/builder/framework_enum.py +0 -24
  31. aiq/builder/front_end.py +0 -73
  32. aiq/builder/function.py +0 -297
  33. aiq/builder/function_base.py +0 -376
  34. aiq/builder/function_info.py +0 -627
  35. aiq/builder/intermediate_step_manager.py +0 -176
  36. aiq/builder/llm.py +0 -25
  37. aiq/builder/retriever.py +0 -25
  38. aiq/builder/user_interaction_manager.py +0 -71
  39. aiq/builder/workflow.py +0 -143
  40. aiq/builder/workflow_builder.py +0 -757
  41. aiq/cli/__init__.py +0 -14
  42. aiq/cli/cli_utils/__init__.py +0 -0
  43. aiq/cli/cli_utils/config_override.py +0 -231
  44. aiq/cli/cli_utils/validation.py +0 -37
  45. aiq/cli/commands/__init__.py +0 -0
  46. aiq/cli/commands/configure/__init__.py +0 -0
  47. aiq/cli/commands/configure/channel/__init__.py +0 -0
  48. aiq/cli/commands/configure/channel/add.py +0 -28
  49. aiq/cli/commands/configure/channel/channel.py +0 -36
  50. aiq/cli/commands/configure/channel/remove.py +0 -30
  51. aiq/cli/commands/configure/channel/update.py +0 -30
  52. aiq/cli/commands/configure/configure.py +0 -33
  53. aiq/cli/commands/evaluate.py +0 -139
  54. aiq/cli/commands/info/__init__.py +0 -14
  55. aiq/cli/commands/info/info.py +0 -39
  56. aiq/cli/commands/info/list_channels.py +0 -32
  57. aiq/cli/commands/info/list_components.py +0 -129
  58. aiq/cli/commands/info/list_mcp.py +0 -126
  59. aiq/cli/commands/registry/__init__.py +0 -14
  60. aiq/cli/commands/registry/publish.py +0 -88
  61. aiq/cli/commands/registry/pull.py +0 -118
  62. aiq/cli/commands/registry/registry.py +0 -38
  63. aiq/cli/commands/registry/remove.py +0 -108
  64. aiq/cli/commands/registry/search.py +0 -155
  65. aiq/cli/commands/start.py +0 -250
  66. aiq/cli/commands/uninstall.py +0 -83
  67. aiq/cli/commands/validate.py +0 -47
  68. aiq/cli/commands/workflow/__init__.py +0 -14
  69. aiq/cli/commands/workflow/templates/__init__.py.j2 +0 -0
  70. aiq/cli/commands/workflow/templates/config.yml.j2 +0 -16
  71. aiq/cli/commands/workflow/templates/pyproject.toml.j2 +0 -22
  72. aiq/cli/commands/workflow/templates/register.py.j2 +0 -5
  73. aiq/cli/commands/workflow/templates/workflow.py.j2 +0 -36
  74. aiq/cli/commands/workflow/workflow.py +0 -37
  75. aiq/cli/commands/workflow/workflow_commands.py +0 -313
  76. aiq/cli/entrypoint.py +0 -133
  77. aiq/cli/main.py +0 -44
  78. aiq/cli/register_workflow.py +0 -408
  79. aiq/cli/type_registry.py +0 -879
  80. aiq/data_models/__init__.py +0 -14
  81. aiq/data_models/api_server.py +0 -588
  82. aiq/data_models/common.py +0 -143
  83. aiq/data_models/component.py +0 -46
  84. aiq/data_models/component_ref.py +0 -135
  85. aiq/data_models/config.py +0 -349
  86. aiq/data_models/dataset_handler.py +0 -122
  87. aiq/data_models/discovery_metadata.py +0 -286
  88. aiq/data_models/embedder.py +0 -26
  89. aiq/data_models/evaluate.py +0 -104
  90. aiq/data_models/evaluator.py +0 -26
  91. aiq/data_models/front_end.py +0 -26
  92. aiq/data_models/function.py +0 -30
  93. aiq/data_models/function_dependencies.py +0 -64
  94. aiq/data_models/interactive.py +0 -237
  95. aiq/data_models/intermediate_step.py +0 -269
  96. aiq/data_models/invocation_node.py +0 -38
  97. aiq/data_models/llm.py +0 -26
  98. aiq/data_models/logging.py +0 -26
  99. aiq/data_models/memory.py +0 -26
  100. aiq/data_models/profiler.py +0 -53
  101. aiq/data_models/registry_handler.py +0 -26
  102. aiq/data_models/retriever.py +0 -30
  103. aiq/data_models/step_adaptor.py +0 -64
  104. aiq/data_models/streaming.py +0 -33
  105. aiq/data_models/swe_bench_model.py +0 -54
  106. aiq/data_models/telemetry_exporter.py +0 -26
  107. aiq/embedder/__init__.py +0 -0
  108. aiq/embedder/langchain_client.py +0 -41
  109. aiq/embedder/nim_embedder.py +0 -58
  110. aiq/embedder/openai_embedder.py +0 -42
  111. aiq/embedder/register.py +0 -24
  112. aiq/eval/__init__.py +0 -14
  113. aiq/eval/config.py +0 -42
  114. aiq/eval/dataset_handler/__init__.py +0 -0
  115. aiq/eval/dataset_handler/dataset_downloader.py +0 -106
  116. aiq/eval/dataset_handler/dataset_filter.py +0 -52
  117. aiq/eval/dataset_handler/dataset_handler.py +0 -169
  118. aiq/eval/evaluate.py +0 -325
  119. aiq/eval/evaluator/__init__.py +0 -14
  120. aiq/eval/evaluator/evaluator_model.py +0 -44
  121. aiq/eval/intermediate_step_adapter.py +0 -93
  122. aiq/eval/rag_evaluator/__init__.py +0 -0
  123. aiq/eval/rag_evaluator/evaluate.py +0 -138
  124. aiq/eval/rag_evaluator/register.py +0 -138
  125. aiq/eval/register.py +0 -23
  126. aiq/eval/remote_workflow.py +0 -128
  127. aiq/eval/runtime_event_subscriber.py +0 -52
  128. aiq/eval/swe_bench_evaluator/__init__.py +0 -0
  129. aiq/eval/swe_bench_evaluator/evaluate.py +0 -215
  130. aiq/eval/swe_bench_evaluator/register.py +0 -36
  131. aiq/eval/trajectory_evaluator/__init__.py +0 -0
  132. aiq/eval/trajectory_evaluator/evaluate.py +0 -118
  133. aiq/eval/trajectory_evaluator/register.py +0 -40
  134. aiq/eval/tunable_rag_evaluator/__init__.py +0 -0
  135. aiq/eval/tunable_rag_evaluator/evaluate.py +0 -263
  136. aiq/eval/tunable_rag_evaluator/register.py +0 -50
  137. aiq/eval/utils/__init__.py +0 -0
  138. aiq/eval/utils/output_uploader.py +0 -131
  139. aiq/eval/utils/tqdm_position_registry.py +0 -40
  140. aiq/front_ends/__init__.py +0 -14
  141. aiq/front_ends/console/__init__.py +0 -14
  142. aiq/front_ends/console/console_front_end_config.py +0 -32
  143. aiq/front_ends/console/console_front_end_plugin.py +0 -107
  144. aiq/front_ends/console/register.py +0 -25
  145. aiq/front_ends/cron/__init__.py +0 -14
  146. aiq/front_ends/fastapi/__init__.py +0 -14
  147. aiq/front_ends/fastapi/fastapi_front_end_config.py +0 -150
  148. aiq/front_ends/fastapi/fastapi_front_end_plugin.py +0 -103
  149. aiq/front_ends/fastapi/fastapi_front_end_plugin_worker.py +0 -607
  150. aiq/front_ends/fastapi/intermediate_steps_subscriber.py +0 -80
  151. aiq/front_ends/fastapi/job_store.py +0 -161
  152. aiq/front_ends/fastapi/main.py +0 -70
  153. aiq/front_ends/fastapi/message_handler.py +0 -279
  154. aiq/front_ends/fastapi/message_validator.py +0 -345
  155. aiq/front_ends/fastapi/register.py +0 -25
  156. aiq/front_ends/fastapi/response_helpers.py +0 -195
  157. aiq/front_ends/fastapi/step_adaptor.py +0 -320
  158. aiq/front_ends/fastapi/websocket.py +0 -148
  159. aiq/front_ends/mcp/__init__.py +0 -14
  160. aiq/front_ends/mcp/mcp_front_end_config.py +0 -32
  161. aiq/front_ends/mcp/mcp_front_end_plugin.py +0 -93
  162. aiq/front_ends/mcp/register.py +0 -27
  163. aiq/front_ends/mcp/tool_converter.py +0 -242
  164. aiq/front_ends/register.py +0 -22
  165. aiq/front_ends/simple_base/__init__.py +0 -14
  166. aiq/front_ends/simple_base/simple_front_end_plugin_base.py +0 -52
  167. aiq/llm/__init__.py +0 -0
  168. aiq/llm/nim_llm.py +0 -45
  169. aiq/llm/openai_llm.py +0 -45
  170. aiq/llm/register.py +0 -22
  171. aiq/llm/utils/__init__.py +0 -14
  172. aiq/llm/utils/env_config_value.py +0 -94
  173. aiq/llm/utils/error.py +0 -17
  174. aiq/memory/__init__.py +0 -20
  175. aiq/memory/interfaces.py +0 -183
  176. aiq/memory/models.py +0 -112
  177. aiq/meta/module_to_distro.json +0 -3
  178. aiq/meta/pypi.md +0 -58
  179. aiq/observability/__init__.py +0 -0
  180. aiq/observability/async_otel_listener.py +0 -429
  181. aiq/observability/register.py +0 -99
  182. aiq/plugins/.namespace +0 -1
  183. aiq/profiler/__init__.py +0 -0
  184. aiq/profiler/callbacks/__init__.py +0 -0
  185. aiq/profiler/callbacks/agno_callback_handler.py +0 -295
  186. aiq/profiler/callbacks/base_callback_class.py +0 -20
  187. aiq/profiler/callbacks/langchain_callback_handler.py +0 -278
  188. aiq/profiler/callbacks/llama_index_callback_handler.py +0 -205
  189. aiq/profiler/callbacks/semantic_kernel_callback_handler.py +0 -238
  190. aiq/profiler/callbacks/token_usage_base_model.py +0 -27
  191. aiq/profiler/data_frame_row.py +0 -51
  192. aiq/profiler/decorators/__init__.py +0 -0
  193. aiq/profiler/decorators/framework_wrapper.py +0 -131
  194. aiq/profiler/decorators/function_tracking.py +0 -254
  195. aiq/profiler/forecasting/__init__.py +0 -0
  196. aiq/profiler/forecasting/config.py +0 -18
  197. aiq/profiler/forecasting/model_trainer.py +0 -75
  198. aiq/profiler/forecasting/models/__init__.py +0 -22
  199. aiq/profiler/forecasting/models/forecasting_base_model.py +0 -40
  200. aiq/profiler/forecasting/models/linear_model.py +0 -196
  201. aiq/profiler/forecasting/models/random_forest_regressor.py +0 -268
  202. aiq/profiler/inference_metrics_model.py +0 -25
  203. aiq/profiler/inference_optimization/__init__.py +0 -0
  204. aiq/profiler/inference_optimization/bottleneck_analysis/__init__.py +0 -0
  205. aiq/profiler/inference_optimization/bottleneck_analysis/nested_stack_analysis.py +0 -452
  206. aiq/profiler/inference_optimization/bottleneck_analysis/simple_stack_analysis.py +0 -258
  207. aiq/profiler/inference_optimization/data_models.py +0 -386
  208. aiq/profiler/inference_optimization/experimental/__init__.py +0 -0
  209. aiq/profiler/inference_optimization/experimental/concurrency_spike_analysis.py +0 -468
  210. aiq/profiler/inference_optimization/experimental/prefix_span_analysis.py +0 -405
  211. aiq/profiler/inference_optimization/llm_metrics.py +0 -212
  212. aiq/profiler/inference_optimization/prompt_caching.py +0 -163
  213. aiq/profiler/inference_optimization/token_uniqueness.py +0 -107
  214. aiq/profiler/inference_optimization/workflow_runtimes.py +0 -72
  215. aiq/profiler/intermediate_property_adapter.py +0 -102
  216. aiq/profiler/profile_runner.py +0 -433
  217. aiq/profiler/utils.py +0 -184
  218. aiq/registry_handlers/__init__.py +0 -0
  219. aiq/registry_handlers/local/__init__.py +0 -0
  220. aiq/registry_handlers/local/local_handler.py +0 -176
  221. aiq/registry_handlers/local/register_local.py +0 -37
  222. aiq/registry_handlers/metadata_factory.py +0 -60
  223. aiq/registry_handlers/package_utils.py +0 -198
  224. aiq/registry_handlers/pypi/__init__.py +0 -0
  225. aiq/registry_handlers/pypi/pypi_handler.py +0 -251
  226. aiq/registry_handlers/pypi/register_pypi.py +0 -40
  227. aiq/registry_handlers/register.py +0 -21
  228. aiq/registry_handlers/registry_handler_base.py +0 -157
  229. aiq/registry_handlers/rest/__init__.py +0 -0
  230. aiq/registry_handlers/rest/register_rest.py +0 -56
  231. aiq/registry_handlers/rest/rest_handler.py +0 -237
  232. aiq/registry_handlers/schemas/__init__.py +0 -0
  233. aiq/registry_handlers/schemas/headers.py +0 -42
  234. aiq/registry_handlers/schemas/package.py +0 -68
  235. aiq/registry_handlers/schemas/publish.py +0 -63
  236. aiq/registry_handlers/schemas/pull.py +0 -82
  237. aiq/registry_handlers/schemas/remove.py +0 -36
  238. aiq/registry_handlers/schemas/search.py +0 -91
  239. aiq/registry_handlers/schemas/status.py +0 -47
  240. aiq/retriever/__init__.py +0 -0
  241. aiq/retriever/interface.py +0 -37
  242. aiq/retriever/milvus/__init__.py +0 -14
  243. aiq/retriever/milvus/register.py +0 -81
  244. aiq/retriever/milvus/retriever.py +0 -228
  245. aiq/retriever/models.py +0 -74
  246. aiq/retriever/nemo_retriever/__init__.py +0 -14
  247. aiq/retriever/nemo_retriever/register.py +0 -60
  248. aiq/retriever/nemo_retriever/retriever.py +0 -190
  249. aiq/retriever/register.py +0 -22
  250. aiq/runtime/__init__.py +0 -14
  251. aiq/runtime/loader.py +0 -188
  252. aiq/runtime/runner.py +0 -176
  253. aiq/runtime/session.py +0 -140
  254. aiq/runtime/user_metadata.py +0 -131
  255. aiq/settings/__init__.py +0 -0
  256. aiq/settings/global_settings.py +0 -318
  257. aiq/test/.namespace +0 -1
  258. aiq/tool/__init__.py +0 -0
  259. aiq/tool/code_execution/__init__.py +0 -0
  260. aiq/tool/code_execution/code_sandbox.py +0 -188
  261. aiq/tool/code_execution/local_sandbox/Dockerfile.sandbox +0 -60
  262. aiq/tool/code_execution/local_sandbox/__init__.py +0 -13
  263. aiq/tool/code_execution/local_sandbox/local_sandbox_server.py +0 -83
  264. aiq/tool/code_execution/local_sandbox/sandbox.requirements.txt +0 -4
  265. aiq/tool/code_execution/local_sandbox/start_local_sandbox.sh +0 -25
  266. aiq/tool/code_execution/register.py +0 -70
  267. aiq/tool/code_execution/utils.py +0 -100
  268. aiq/tool/datetime_tools.py +0 -42
  269. aiq/tool/document_search.py +0 -141
  270. aiq/tool/github_tools/__init__.py +0 -0
  271. aiq/tool/github_tools/create_github_commit.py +0 -133
  272. aiq/tool/github_tools/create_github_issue.py +0 -87
  273. aiq/tool/github_tools/create_github_pr.py +0 -106
  274. aiq/tool/github_tools/get_github_file.py +0 -106
  275. aiq/tool/github_tools/get_github_issue.py +0 -166
  276. aiq/tool/github_tools/get_github_pr.py +0 -256
  277. aiq/tool/github_tools/update_github_issue.py +0 -100
  278. aiq/tool/mcp/__init__.py +0 -14
  279. aiq/tool/mcp/mcp_client.py +0 -220
  280. aiq/tool/mcp/mcp_tool.py +0 -95
  281. aiq/tool/memory_tools/__init__.py +0 -0
  282. aiq/tool/memory_tools/add_memory_tool.py +0 -79
  283. aiq/tool/memory_tools/delete_memory_tool.py +0 -67
  284. aiq/tool/memory_tools/get_memory_tool.py +0 -72
  285. aiq/tool/nvidia_rag.py +0 -95
  286. aiq/tool/register.py +0 -37
  287. aiq/tool/retriever.py +0 -89
  288. aiq/tool/server_tools.py +0 -63
  289. aiq/utils/__init__.py +0 -0
  290. aiq/utils/data_models/__init__.py +0 -0
  291. aiq/utils/data_models/schema_validator.py +0 -58
  292. aiq/utils/debugging_utils.py +0 -43
  293. aiq/utils/exception_handlers/__init__.py +0 -0
  294. aiq/utils/exception_handlers/schemas.py +0 -114
  295. aiq/utils/io/__init__.py +0 -0
  296. aiq/utils/io/yaml_tools.py +0 -119
  297. aiq/utils/metadata_utils.py +0 -74
  298. aiq/utils/optional_imports.py +0 -142
  299. aiq/utils/producer_consumer_queue.py +0 -178
  300. aiq/utils/reactive/__init__.py +0 -0
  301. aiq/utils/reactive/base/__init__.py +0 -0
  302. aiq/utils/reactive/base/observable_base.py +0 -65
  303. aiq/utils/reactive/base/observer_base.py +0 -55
  304. aiq/utils/reactive/base/subject_base.py +0 -79
  305. aiq/utils/reactive/observable.py +0 -59
  306. aiq/utils/reactive/observer.py +0 -76
  307. aiq/utils/reactive/subject.py +0 -131
  308. aiq/utils/reactive/subscription.py +0 -49
  309. aiq/utils/settings/__init__.py +0 -0
  310. aiq/utils/settings/global_settings.py +0 -197
  311. aiq/utils/type_converter.py +0 -232
  312. aiq/utils/type_utils.py +0 -397
  313. aiq/utils/url_utils.py +0 -27
  314. aiqtoolkit-1.1.0rc6.dist-info/METADATA +0 -331
  315. aiqtoolkit-1.1.0rc6.dist-info/RECORD +0 -316
  316. aiqtoolkit-1.1.0rc6.dist-info/entry_points.txt +0 -17
  317. aiqtoolkit-1.1.0rc6.dist-info/licenses/LICENSE-3rd-party.txt +0 -3686
  318. aiqtoolkit-1.1.0rc6.dist-info/licenses/LICENSE.md +0 -201
  319. aiqtoolkit-1.1.0rc6.dist-info/top_level.txt +0 -1
aiq/cli/type_registry.py DELETED
@@ -1,879 +0,0 @@
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 logging
17
- import typing
18
- from collections.abc import AsyncIterator
19
- from collections.abc import Callable
20
- from contextlib import AbstractAsyncContextManager
21
- from contextlib import contextmanager
22
- from copy import deepcopy
23
- from functools import cached_property
24
- from logging import Handler
25
-
26
- from pydantic import BaseModel
27
- from pydantic import ConfigDict
28
- from pydantic import Field
29
- from pydantic import Tag
30
- from pydantic import computed_field
31
- from pydantic import field_validator
32
-
33
- from aiq.builder.builder import Builder
34
- from aiq.builder.builder import EvalBuilder
35
- from aiq.builder.embedder import EmbedderProviderInfo
36
- from aiq.builder.evaluator import EvaluatorInfo
37
- from aiq.builder.front_end import FrontEndBase
38
- from aiq.builder.function import Function
39
- from aiq.builder.function_base import FunctionBase
40
- from aiq.builder.function_info import FunctionInfo
41
- from aiq.builder.llm import LLMProviderInfo
42
- from aiq.builder.retriever import RetrieverProviderInfo
43
- from aiq.data_models.common import TypedBaseModelT
44
- from aiq.data_models.component import AIQComponentEnum
45
- from aiq.data_models.config import AIQConfig
46
- from aiq.data_models.discovery_metadata import DiscoveryMetadata
47
- from aiq.data_models.embedder import EmbedderBaseConfig
48
- from aiq.data_models.embedder import EmbedderBaseConfigT
49
- from aiq.data_models.evaluator import EvaluatorBaseConfig
50
- from aiq.data_models.evaluator import EvaluatorBaseConfigT
51
- from aiq.data_models.front_end import FrontEndBaseConfig
52
- from aiq.data_models.front_end import FrontEndConfigT
53
- from aiq.data_models.function import FunctionBaseConfig
54
- from aiq.data_models.function import FunctionConfigT
55
- from aiq.data_models.llm import LLMBaseConfig
56
- from aiq.data_models.llm import LLMBaseConfigT
57
- from aiq.data_models.logging import LoggingBaseConfig
58
- from aiq.data_models.logging import LoggingMethodConfigT
59
- from aiq.data_models.memory import MemoryBaseConfig
60
- from aiq.data_models.memory import MemoryBaseConfigT
61
- from aiq.data_models.registry_handler import RegistryHandlerBaseConfig
62
- from aiq.data_models.registry_handler import RegistryHandlerBaseConfigT
63
- from aiq.data_models.retriever import RetrieverBaseConfig
64
- from aiq.data_models.retriever import RetrieverBaseConfigT
65
- from aiq.data_models.telemetry_exporter import TelemetryExporterBaseConfig
66
- from aiq.data_models.telemetry_exporter import TelemetryExporterConfigT
67
- from aiq.memory.interfaces import MemoryEditor
68
- from aiq.registry_handlers.registry_handler_base import AbstractRegistryHandler
69
- from aiq.utils.optional_imports import TelemetryOptionalImportError
70
- from aiq.utils.optional_imports import try_import_opentelemetry
71
-
72
- # Try to import OpenTelemetry modules
73
- # If the dependencies are not installed, use a dummy span exporter here
74
- try:
75
- opentelemetry = try_import_opentelemetry()
76
- from opentelemetry.sdk.trace.export import SpanExporter
77
- except TelemetryOptionalImportError:
78
- from aiq.utils.optional_imports import DummySpanExporter # pylint: disable=ungrouped-imports
79
- SpanExporter = DummySpanExporter
80
-
81
- logger = logging.getLogger(__name__)
82
-
83
- FrontEndBuildCallableT = Callable[[FrontEndConfigT, AIQConfig], AsyncIterator[FrontEndBase]]
84
- TelemetryExporterBuildCallableT = Callable[[TelemetryExporterConfigT, Builder], AsyncIterator[SpanExporter]]
85
- LoggingMethodBuildCallableT = Callable[[LoggingMethodConfigT, Builder], AsyncIterator[Handler]]
86
- FunctionBuildCallableT = Callable[[FunctionConfigT, Builder], AsyncIterator[FunctionInfo | Callable | FunctionBase]]
87
- LLMProviderBuildCallableT = Callable[[LLMBaseConfigT, Builder], AsyncIterator[LLMProviderInfo]]
88
- LLMClientBuildCallableT = Callable[[LLMBaseConfigT, Builder], AsyncIterator[typing.Any]]
89
- EmbedderProviderBuildCallableT = Callable[[EmbedderBaseConfigT, Builder], AsyncIterator[EmbedderProviderInfo]]
90
- EmbedderClientBuildCallableT = Callable[[EmbedderBaseConfigT, Builder], AsyncIterator[typing.Any]]
91
- EvaluatorBuildCallableT = Callable[[EvaluatorBaseConfigT, EvalBuilder], AsyncIterator[EvaluatorInfo]]
92
- MemoryBuildCallableT = Callable[[MemoryBaseConfigT, Builder], AsyncIterator[MemoryEditor]]
93
- RetrieverProviderBuildCallableT = Callable[[RetrieverBaseConfigT, Builder], AsyncIterator[RetrieverProviderInfo]]
94
- RetrieverClientBuildCallableT = Callable[[RetrieverBaseConfigT, Builder], AsyncIterator[typing.Any]]
95
- RegistryHandlerBuildCallableT = Callable[[RegistryHandlerBaseConfigT], AsyncIterator[AbstractRegistryHandler]]
96
- ToolWrapperBuildCallableT = Callable[[str, Function, Builder], typing.Any]
97
-
98
- TeleExporterRegisteredCallableT = Callable[[TelemetryExporterConfigT, Builder], AbstractAsyncContextManager[typing.Any]]
99
- LoggingMethodRegisteredCallableT = Callable[[LoggingMethodConfigT, Builder], AbstractAsyncContextManager[typing.Any]]
100
- FrontEndRegisteredCallableT = Callable[[FrontEndConfigT, AIQConfig], AbstractAsyncContextManager[FrontEndBase]]
101
- FunctionRegisteredCallableT = Callable[[FunctionConfigT, Builder],
102
- AbstractAsyncContextManager[FunctionInfo | Callable | FunctionBase]]
103
- LLMProviderRegisteredCallableT = Callable[[LLMBaseConfigT, Builder], AbstractAsyncContextManager[LLMProviderInfo]]
104
- LLMClientRegisteredCallableT = Callable[[LLMBaseConfigT, Builder], AbstractAsyncContextManager[typing.Any]]
105
- EmbedderProviderRegisteredCallableT = Callable[[EmbedderBaseConfigT, Builder],
106
- AbstractAsyncContextManager[EmbedderProviderInfo]]
107
- EmbedderClientRegisteredCallableT = Callable[[EmbedderBaseConfigT, Builder], AbstractAsyncContextManager[typing.Any]]
108
- EvaluatorRegisteredCallableT = Callable[[EvaluatorBaseConfigT, EvalBuilder], AbstractAsyncContextManager[EvaluatorInfo]]
109
- MemoryRegisteredCallableT = Callable[[MemoryBaseConfigT, Builder], AbstractAsyncContextManager[MemoryEditor]]
110
- RetrieverProviderRegisteredCallableT = Callable[[RetrieverBaseConfigT, Builder],
111
- AbstractAsyncContextManager[RetrieverProviderInfo]]
112
- RetrieverClientRegisteredCallableT = Callable[[RetrieverBaseConfigT, Builder], AbstractAsyncContextManager[typing.Any]]
113
- RegistryHandlerRegisteredCallableT = Callable[[RegistryHandlerBaseConfigT],
114
- AbstractAsyncContextManager[AbstractRegistryHandler]]
115
-
116
-
117
- class RegisteredInfo(BaseModel, typing.Generic[TypedBaseModelT]):
118
-
119
- model_config = ConfigDict(frozen=True)
120
-
121
- full_type: str
122
- config_type: type[TypedBaseModelT]
123
- discovery_metadata: DiscoveryMetadata = DiscoveryMetadata()
124
-
125
- @computed_field
126
- @cached_property
127
- def module_name(self) -> str:
128
- return self.full_type.split("/")[0]
129
-
130
- @computed_field
131
- @cached_property
132
- def local_name(self) -> str:
133
- return self.full_type.split("/")[-1]
134
-
135
- @field_validator("full_type", mode="after")
136
- @classmethod
137
- def validate_full_type(cls, full_type: str) -> str:
138
- parts = full_type.split("/")
139
-
140
- if (len(parts) != 2):
141
- raise ValueError(f"Invalid full type: {full_type}. Expected format: `module_name/local_name`")
142
-
143
- return full_type
144
-
145
-
146
- class RegisteredTelemetryExporter(RegisteredInfo[TelemetryExporterBaseConfig]):
147
-
148
- build_fn: TeleExporterRegisteredCallableT = Field(repr=False)
149
-
150
-
151
- class RegisteredLoggingMethod(RegisteredInfo[LoggingBaseConfig]):
152
-
153
- build_fn: LoggingMethodRegisteredCallableT = Field(repr=False)
154
-
155
-
156
- class RegisteredFrontEndInfo(RegisteredInfo[FrontEndBaseConfig]):
157
- """
158
- Represents a registered front end. Front ends are the entry points to the workflow and are responsible for
159
- orchestrating the workflow.
160
- """
161
-
162
- build_fn: FrontEndRegisteredCallableT = Field(repr=False)
163
-
164
-
165
- class RegisteredFunctionInfo(RegisteredInfo[FunctionBaseConfig]):
166
- """
167
- Represents a registered function. Functions are the building blocks of the workflow with predefined inputs, outputs,
168
- and a description.
169
- """
170
-
171
- build_fn: FunctionRegisteredCallableT = Field(repr=False)
172
- framework_wrappers: list[str] = Field(default_factory=list)
173
-
174
-
175
- class RegisteredLLMProviderInfo(RegisteredInfo[LLMBaseConfig]):
176
- """
177
- Represents a registered LLM provider. LLM Providers are the operators of the LLMs. i.e. NIMs, OpenAI, Anthropic,
178
- etc.
179
- """
180
-
181
- build_fn: LLMProviderRegisteredCallableT = Field(repr=False)
182
-
183
-
184
- class RegisteredLLMClientInfo(RegisteredInfo[LLMBaseConfig]):
185
- """
186
- Represents a registered LLM client. LLM Clients are the clients that interact with the LLM providers and are
187
- specific to a particular LLM framework.
188
- """
189
-
190
- llm_framework: str
191
- build_fn: LLMClientRegisteredCallableT = Field(repr=False)
192
-
193
-
194
- class RegisteredEmbedderProviderInfo(RegisteredInfo[EmbedderBaseConfig]):
195
- """
196
- Represents a registered Embedder provider. Embedder Providers are the operators of the Embedder models. i.e. NIMs,
197
- OpenAI, Anthropic, etc.
198
- """
199
-
200
- build_fn: EmbedderProviderRegisteredCallableT = Field(repr=False)
201
-
202
-
203
- class RegisteredEmbedderClientInfo(RegisteredInfo[EmbedderBaseConfig]):
204
- """
205
- Represents a registered Embedder client. Embedder Clients are the clients that interact with the Embedder providers
206
- and are specific to a particular LLM framework.
207
- """
208
-
209
- llm_framework: str
210
- build_fn: EmbedderClientRegisteredCallableT = Field(repr=False)
211
-
212
-
213
- class RegisteredEvaluatorInfo(RegisteredInfo[EvaluatorBaseConfig]):
214
- """
215
- Represents a registered Evaluator e.g. RagEvaluator, TrajectoryEvaluator, etc.
216
- """
217
-
218
- build_fn: EvaluatorRegisteredCallableT = Field(repr=False)
219
-
220
-
221
- class RegisteredMemoryInfo(RegisteredInfo[MemoryBaseConfig]):
222
- """
223
- Represents a registered Memory object which adheres to the memory interface.
224
- """
225
-
226
- build_fn: MemoryRegisteredCallableT = Field(repr=False)
227
-
228
-
229
- class RegisteredToolWrapper(BaseModel):
230
- """
231
- Represents a registered tool wrapper. Tool wrappers are used to wrap the functions in a particular LLM framework.
232
- They do not have their own configuration, but they are used to wrap the functions in a particular LLM framework.
233
- """
234
-
235
- llm_framework: str
236
- build_fn: ToolWrapperBuildCallableT = Field(repr=False)
237
- discovery_metadata: DiscoveryMetadata
238
-
239
-
240
- class RegisteredRetrieverProviderInfo(RegisteredInfo[RetrieverBaseConfig]):
241
- """
242
- Represents a registered Retriever object which adheres to the retriever interface.
243
- """
244
-
245
- build_fn: RetrieverProviderRegisteredCallableT = Field(repr=False)
246
-
247
-
248
- class RegisteredRetrieverClientInfo(RegisteredInfo[RetrieverBaseConfig]):
249
- """
250
- Represents a registered Retriever Client. Retriever Clients are the LLM Framework-specific clients that expose an
251
- interface to the Retriever object.
252
- """
253
- llm_framework: str | None
254
- build_fn: RetrieverClientRegisteredCallableT = Field(repr=False)
255
-
256
-
257
- class RegisteredRegistryHandlerInfo(RegisteredInfo[RegistryHandlerBaseConfig]):
258
- """
259
- Represents a registered LLM client. LLM Clients are the clients that interact with the LLM providers and are
260
- specific to a particular LLM framework.
261
- """
262
-
263
- build_fn: RegistryHandlerRegisteredCallableT = Field(repr=False)
264
-
265
-
266
- class RegisteredPackage(BaseModel):
267
- package_name: str
268
- discovery_metadata: DiscoveryMetadata
269
-
270
-
271
- class TypeRegistry: # pylint: disable=too-many-public-methods
272
-
273
- def __init__(self) -> None:
274
- # Telemetry Exporters
275
- self._registered_telemetry_exporters: dict[type[TelemetryExporterBaseConfig], RegisteredTelemetryExporter] = {}
276
-
277
- # Logging Methods
278
- self._registered_logging_methods: dict[type[LoggingBaseConfig], RegisteredLoggingMethod] = {}
279
-
280
- # Front Ends
281
- self._registered_front_end_infos: dict[type[FrontEndBaseConfig], RegisteredFrontEndInfo] = {}
282
-
283
- # Functions
284
- self._registered_functions: dict[type[FunctionBaseConfig], RegisteredFunctionInfo] = {}
285
-
286
- # LLMs
287
- self._registered_llm_provider_infos: dict[type[LLMBaseConfig], RegisteredLLMProviderInfo] = {}
288
- self._llm_client_provider_to_framework: dict[type[LLMBaseConfig], dict[str, RegisteredLLMClientInfo]] = {}
289
- self._llm_client_framework_to_provider: dict[str, dict[type[LLMBaseConfig], RegisteredLLMClientInfo]] = {}
290
-
291
- # Embedders
292
- self._registered_embedder_provider_infos: dict[type[EmbedderBaseConfig], RegisteredEmbedderProviderInfo] = {}
293
- self._embedder_client_provider_to_framework: dict[type[EmbedderBaseConfig],
294
- dict[str, RegisteredEmbedderClientInfo]] = {}
295
- self._embedder_client_framework_to_provider: dict[str,
296
- dict[type[EmbedderBaseConfig],
297
- RegisteredEmbedderClientInfo]] = {}
298
-
299
- # Evaluators
300
- self._registered_evaluator_infos: dict[type[EvaluatorBaseConfig], RegisteredEvaluatorInfo] = {}
301
-
302
- # Memory
303
- self._registered_memory_infos: dict[type[MemoryBaseConfig], RegisteredMemoryInfo] = {}
304
-
305
- # Retrievers
306
- self._registered_retriever_provider_infos: dict[type[RetrieverBaseConfig], RegisteredRetrieverProviderInfo] = {}
307
- self._retriever_client_provider_to_framework: dict[type[RetrieverBaseConfig],
308
- dict[str | None, RegisteredRetrieverClientInfo]] = {}
309
- self._retriever_client_framework_to_provider: dict[str | None,
310
- dict[type[RetrieverBaseConfig],
311
- RegisteredRetrieverClientInfo]] = {}
312
-
313
- # Registry Handlers
314
- self._registered_registry_handler_infos: dict[type[RegistryHandlerBaseConfig],
315
- RegisteredRegistryHandlerInfo] = {}
316
-
317
- # Tool Wrappers
318
- self._registered_tool_wrappers: dict[str, RegisteredToolWrapper] = {}
319
-
320
- # Packages
321
- self._registered_packages: dict[str, RegisteredPackage] = {}
322
-
323
- self._registration_changed_hooks: list[Callable[[], None]] = []
324
- self._registration_changed_hooks_active: bool = True
325
-
326
- self._registered_channel_map = {}
327
-
328
- def _registration_changed(self):
329
-
330
- if (not self._registration_changed_hooks_active):
331
- return
332
-
333
- logger.debug("Registration changed. Notifying hooks.")
334
-
335
- for hook in self._registration_changed_hooks:
336
- hook()
337
-
338
- def add_registration_changed_hook(self, cb: Callable[[], typing.Any]) -> None:
339
-
340
- self._registration_changed_hooks.append(cb)
341
-
342
- @contextmanager
343
- def pause_registration_changed_hooks(self):
344
-
345
- self._registration_changed_hooks_active = False
346
-
347
- try:
348
- yield
349
- finally:
350
- self._registration_changed_hooks_active = True
351
-
352
- # Ensure that the registration changed hooks are called
353
- self._registration_changed()
354
-
355
- def register_telemetry_exporter(self, registration: RegisteredTelemetryExporter):
356
-
357
- if (registration.config_type in self._registered_telemetry_exporters):
358
- raise ValueError(f"A telemetry exporter with the same config type `{registration.config_type}` has already "
359
- "been registered.")
360
-
361
- self._registered_telemetry_exporters[registration.config_type] = registration
362
-
363
- self._registration_changed()
364
-
365
- def get_telemetry_exporter(self, config_type: type[TelemetryExporterBaseConfig]) -> RegisteredTelemetryExporter:
366
-
367
- try:
368
- return self._registered_telemetry_exporters[config_type]
369
- except KeyError as err:
370
- raise KeyError(f"Could not find a registered telemetry exporter for config `{config_type}`. "
371
- f"Registered configs: {set(self._registered_telemetry_exporters.keys())}") from err
372
-
373
- def get_registered_telemetry_exporters(self) -> list[RegisteredInfo[TelemetryExporterBaseConfig]]:
374
-
375
- return list(self._registered_telemetry_exporters.values())
376
-
377
- def register_logging_method(self, registration: RegisteredLoggingMethod):
378
-
379
- if (registration.config_type in self._registered_logging_methods):
380
- raise ValueError(f"A logging method with the same config type `{registration.config_type}` has already "
381
- "been registered.")
382
-
383
- self._registered_logging_methods[registration.config_type] = registration
384
-
385
- self._registration_changed()
386
-
387
- def get_logging_method(self, config_type: type[LoggingBaseConfig]) -> RegisteredLoggingMethod:
388
- try:
389
- return self._registered_logging_methods[config_type]
390
- except KeyError as err:
391
- raise KeyError(f"No logging method found for config `{config_type}`. "
392
- f"Known: {set(self._registered_logging_methods.keys())}") from err
393
-
394
- def get_registered_logging_method(self) -> list[RegisteredInfo[LoggingBaseConfig]]:
395
-
396
- return list(self._registered_logging_methods.values())
397
-
398
- def register_front_end(self, registration: RegisteredFrontEndInfo):
399
-
400
- if (registration.config_type in self._registered_front_end_infos):
401
- raise ValueError(f"A front end with the same config type `{registration.config_type}` has already been "
402
- "registered.")
403
-
404
- self._registered_front_end_infos[registration.config_type] = registration
405
-
406
- self._registration_changed()
407
-
408
- def get_front_end(self, config_type: type[FrontEndBaseConfig]) -> RegisteredFrontEndInfo:
409
-
410
- try:
411
- return self._registered_front_end_infos[config_type]
412
- except KeyError as err:
413
- raise KeyError(f"Could not find a registered front end for config `{config_type}`. "
414
- f"Registered configs: {set(self._registered_front_end_infos.keys())}") from err
415
-
416
- def get_registered_front_ends(self) -> list[RegisteredInfo[FrontEndBaseConfig]]:
417
-
418
- return list(self._registered_front_end_infos.values())
419
-
420
- def register_function(self, registration: RegisteredFunctionInfo):
421
-
422
- if (registration.config_type in self._registered_functions):
423
- raise ValueError(f"A function with the same config type `{registration.config_type}` has already been "
424
- "registered.")
425
-
426
- self._registered_functions[registration.config_type] = registration
427
-
428
- self._registration_changed()
429
-
430
- def get_function(self, config_type: type[FunctionBaseConfig]) -> RegisteredFunctionInfo:
431
-
432
- try:
433
- return self._registered_functions[config_type]
434
- except KeyError as err:
435
- raise KeyError(f"Could not find a registered function for config `{config_type}`. "
436
- f"Registered configs: {set(self._registered_functions.keys())}") from err
437
-
438
- def get_registered_functions(self) -> list[RegisteredInfo[FunctionBaseConfig]]:
439
-
440
- return list(self._registered_functions.values())
441
-
442
- def register_llm_provider(self, info: RegisteredLLMProviderInfo):
443
-
444
- if (info.config_type in self._registered_llm_provider_infos):
445
- raise ValueError(
446
- f"An LLM provider with the same config type `{info.config_type}` has already been registered.")
447
-
448
- self._registered_llm_provider_infos[info.config_type] = info
449
-
450
- self._registration_changed()
451
-
452
- def get_llm_provider(self, config_type: type[LLMBaseConfig]) -> RegisteredLLMProviderInfo:
453
-
454
- try:
455
- return self._registered_llm_provider_infos[config_type]
456
- except KeyError as err:
457
- raise KeyError(f"Could not find a registered LLM provider for config `{config_type}`. "
458
- f"Registered configs: {set(self._registered_llm_provider_infos.keys())}") from err
459
-
460
- def get_registered_llm_providers(self) -> list[RegisteredInfo[LLMBaseConfig]]:
461
-
462
- return list(self._registered_llm_provider_infos.values())
463
-
464
- def register_llm_client(self, info: RegisteredLLMClientInfo):
465
-
466
- if (info.config_type in self._llm_client_provider_to_framework
467
- and info.llm_framework in self._llm_client_provider_to_framework[info.config_type]):
468
- raise ValueError(f"An LLM client with the same config type `{info.config_type}` "
469
- f"and LLM framework `{info.llm_framework}` has already been registered.")
470
-
471
- self._llm_client_provider_to_framework.setdefault(info.config_type, {})[info.llm_framework] = info
472
- self._llm_client_framework_to_provider.setdefault(info.llm_framework, {})[info.config_type] = info
473
-
474
- self._registration_changed()
475
-
476
- def get_llm_client(self, config_type: type[LLMBaseConfig], wrapper_type: str) -> RegisteredLLMClientInfo:
477
-
478
- try:
479
- client_info = self._llm_client_provider_to_framework[config_type][wrapper_type]
480
- except KeyError as err:
481
- raise KeyError(f"An invalid LLM config and wrapper combination was supplied. Config: `{config_type}`, "
482
- f"Wrapper: `{wrapper_type}`. The workflow is requesting a {wrapper_type} LLM client but "
483
- f"there is no registered conversion from that LLM provider to LLM framework: "
484
- f"{wrapper_type}. "
485
- f"Please provide an LLM configuration from one of the following providers: "
486
- f"{set(self._llm_client_provider_to_framework.keys())}") from err
487
-
488
- return client_info
489
-
490
- def register_embedder_provider(self, info: RegisteredEmbedderProviderInfo):
491
-
492
- if (info.config_type in self._registered_embedder_provider_infos):
493
- raise ValueError(f"An Embedder provider with the same config type `{info.config_type}` has already been "
494
- "registered.")
495
-
496
- self._registered_embedder_provider_infos[info.config_type] = info
497
-
498
- self._registration_changed()
499
-
500
- def get_embedder_provider(self, config_type: type[EmbedderBaseConfig]) -> RegisteredEmbedderProviderInfo:
501
-
502
- try:
503
- return self._registered_embedder_provider_infos[config_type]
504
- except KeyError as err:
505
- raise KeyError(f"Could not find a registered Embedder provider for config `{config_type}`. "
506
- f"Registered configs: {set(self._registered_embedder_provider_infos.keys())}") from err
507
-
508
- def get_registered_embedder_providers(self) -> list[RegisteredInfo[EmbedderBaseConfig]]:
509
-
510
- return list(self._registered_embedder_provider_infos.values())
511
-
512
- def register_embedder_client(self, info: RegisteredEmbedderClientInfo):
513
-
514
- if (info.config_type in self._embedder_client_provider_to_framework
515
- and info.llm_framework in self._embedder_client_provider_to_framework[info.config_type]):
516
- raise ValueError(f"An Embedder client with the same config type `{info.config_type}` has already been "
517
- "registered.")
518
-
519
- self._embedder_client_provider_to_framework.setdefault(info.config_type, {})[info.llm_framework] = info
520
- self._embedder_client_framework_to_provider.setdefault(info.llm_framework, {})[info.config_type] = info
521
-
522
- self._registration_changed()
523
-
524
- def get_embedder_client(self, config_type: type[EmbedderBaseConfig],
525
- wrapper_type: str) -> RegisteredEmbedderClientInfo:
526
-
527
- try:
528
- client_info = self._embedder_client_provider_to_framework[config_type][wrapper_type]
529
- except KeyError as err:
530
- raise KeyError(
531
- f"An invalid Embedder config and wrapper combination was supplied. Config: `{config_type}`, "
532
- "Wrapper: `{wrapper_type}`. The workflow is requesting a {wrapper_type} Embedder client but "
533
- "there is no registered conversion from that Embedder provider to LLM framework: {wrapper_type}. "
534
- "Please provide an Embedder configuration from one of the following providers: "
535
- f"{set(self._embedder_client_provider_to_framework.keys())}") from err
536
-
537
- return client_info
538
-
539
- def register_evaluator(self, info: RegisteredEvaluatorInfo):
540
-
541
- if (info.config_type in self._registered_evaluator_infos):
542
- raise ValueError(f"An Evaluator with the same config type `{info.config_type}` has already been "
543
- "registered.")
544
-
545
- self._registered_evaluator_infos[info.config_type] = info
546
-
547
- self._registration_changed()
548
-
549
- def get_evaluator(self, config_type: type[EvaluatorBaseConfig]) -> RegisteredEvaluatorInfo:
550
-
551
- try:
552
- return self._registered_evaluator_infos[config_type]
553
- except KeyError as err:
554
- raise KeyError(f"Could not find a registered Evaluator for config `{config_type}`. "
555
- f"Registered configs: {set(self._registered_evaluator_infos.keys())}") from err
556
-
557
- def get_registered_evaluators(self) -> list[RegisteredInfo[EvaluatorBaseConfig]]:
558
-
559
- return list(self._registered_evaluator_infos.values())
560
-
561
- def register_memory(self, info: RegisteredMemoryInfo):
562
-
563
- if (info.config_type in self._registered_memory_infos):
564
- raise ValueError(
565
- f"A Memory client with the same config type `{info.config_type}` has already been registered.")
566
-
567
- self._registered_memory_infos[info.config_type] = info
568
-
569
- self._registration_changed()
570
-
571
- def get_memory(self, config_type: type[MemoryBaseConfig]) -> RegisteredMemoryInfo:
572
-
573
- try:
574
- return self._registered_memory_infos[config_type]
575
- except KeyError as err:
576
- raise KeyError(f"Could not find a registered Memory client for config `{config_type}`. "
577
- f"Registered configs: {set(self._registered_memory_infos.keys())}") from err
578
-
579
- def get_registered_memorys(self) -> list[RegisteredInfo[MemoryBaseConfig]]:
580
-
581
- return list(self._registered_memory_infos.values())
582
-
583
- def register_retriever_provider(self, info: RegisteredRetrieverProviderInfo):
584
-
585
- if (info.config_type in self._registered_retriever_provider_infos):
586
- raise ValueError(
587
- f"A Retriever provider with the same config type `{info.config_type}` has already been registered")
588
-
589
- self._registered_retriever_provider_infos[info.config_type] = info
590
-
591
- self._registration_changed()
592
-
593
- def get_retriever_provider(self, config_type: type[RetrieverBaseConfig]) -> RegisteredRetrieverProviderInfo:
594
-
595
- try:
596
- return self._registered_retriever_provider_infos[config_type]
597
- except KeyError as err:
598
- raise KeyError(f"Could not find a registered Retriever provider for config `{config_type}`. "
599
- f"Registered configs: {set(self._registered_retriever_provider_infos.keys())}") from err
600
-
601
- def get_registered_retriever_providers(self) -> list[RegisteredInfo[RetrieverBaseConfig]]:
602
-
603
- return list(self._registered_retriever_provider_infos.values())
604
-
605
- def register_retriever_client(self, info: RegisteredRetrieverClientInfo):
606
-
607
- if (info.config_type in self._retriever_client_provider_to_framework
608
- and info.llm_framework in self._retriever_client_provider_to_framework[info.config_type]):
609
- raise ValueError(f"A Retriever client with the same config type `{info.config_type}` "
610
- " and LLM framework `{info.llm_framework}` has already been registered.")
611
-
612
- self._retriever_client_provider_to_framework.setdefault(info.config_type, {})[info.llm_framework] = info
613
- self._retriever_client_framework_to_provider.setdefault(info.llm_framework, {})[info.config_type] = info
614
-
615
- self._registration_changed()
616
-
617
- def get_retriever_client(self, config_type: type[RetrieverBaseConfig],
618
- wrapper_type: str | None) -> RegisteredRetrieverClientInfo:
619
-
620
- try:
621
- client_info = self._retriever_client_provider_to_framework[config_type][wrapper_type]
622
- except KeyError as err:
623
- raise KeyError(
624
- f"An invalid Retriever config and wrapper combination was supplied. Config: `{config_type}`, "
625
- "Wrapper: `{wrapper_type}`. The workflow is requesting a {wrapper_type} Retriever client but "
626
- "there is no registered conversion from that Retriever provider to LLM framework: {wrapper_type}. "
627
- "Please provide a Retriever configuration from one of the following providers: "
628
- f"{set(self._retriever_client_provider_to_framework.keys())}") from err
629
-
630
- return client_info
631
-
632
- def register_tool_wrapper(self, registration: RegisteredToolWrapper):
633
-
634
- if (registration.llm_framework in self._registered_tool_wrappers):
635
- raise ValueError(f"A tool wrapper for the LLM framework `{registration.llm_framework}` has already been "
636
- "registered.")
637
-
638
- self._registered_tool_wrappers[registration.llm_framework] = registration
639
-
640
- self._registration_changed()
641
-
642
- def get_tool_wrapper(self, llm_framework: str) -> RegisteredToolWrapper:
643
-
644
- try:
645
- return self._registered_tool_wrappers[llm_framework]
646
- except KeyError as err:
647
- raise KeyError(f"Could not find a registered tool wrapper for LLM framework `{llm_framework}`. "
648
- f"Registered LLM frameworks: {set(self._registered_tool_wrappers.keys())}") from err
649
-
650
- def register_registry_handler(self, info: RegisteredRegistryHandlerInfo):
651
-
652
- if (info.config_type in self._registered_memory_infos):
653
- raise ValueError(
654
- f"A Registry Handler with the same config type `{info.config_type}` has already been registered.")
655
-
656
- self._registered_registry_handler_infos[info.config_type] = info
657
- self._registered_channel_map[info.config_type.static_type()] = info
658
-
659
- self._registration_changed()
660
-
661
- def get_registry_handler(self, config_type: type[RegistryHandlerBaseConfig]) -> RegisteredRegistryHandlerInfo:
662
-
663
- try:
664
- return self._registered_registry_handler_infos[config_type]
665
- except KeyError as err:
666
- raise KeyError(f"Could not find a registered Registry Handler for config `{config_type}`. "
667
- f"Registered configs: {set(self._registered_registry_handler_infos.keys())}") from err
668
-
669
- def get_registered_registry_handlers(self) -> list[RegisteredInfo[RegistryHandlerBaseConfig]]:
670
-
671
- return list(self._registered_registry_handler_infos.values())
672
-
673
- def register_package(self, package_name: str, package_version: str | None = None):
674
-
675
- discovery_metadata = DiscoveryMetadata.from_package_name(package_name=package_name,
676
- package_version=package_version)
677
- package = RegisteredPackage(discovery_metadata=discovery_metadata, package_name=package_name)
678
- self._registered_packages[package.package_name] = package
679
-
680
- self._registration_changed()
681
-
682
- def get_infos_by_type(self, component_type: AIQComponentEnum) -> dict: # pylint: disable=R0911
683
-
684
- if component_type == AIQComponentEnum.FRONT_END:
685
- return self._registered_front_end_infos
686
-
687
- if component_type == AIQComponentEnum.FUNCTION:
688
- return self._registered_functions
689
-
690
- if component_type == AIQComponentEnum.TOOL_WRAPPER:
691
- return self._registered_tool_wrappers
692
-
693
- if component_type == AIQComponentEnum.LLM_PROVIDER:
694
- return self._registered_llm_provider_infos
695
-
696
- if component_type == AIQComponentEnum.LLM_CLIENT:
697
- leaf_llm_client_infos = {}
698
- for framework in self._llm_client_provider_to_framework.values():
699
- for info in framework.values():
700
- leaf_llm_client_infos[info.discovery_metadata.component_name] = info
701
- return leaf_llm_client_infos
702
-
703
- if component_type == AIQComponentEnum.EMBEDDER_PROVIDER:
704
- return self._registered_embedder_provider_infos
705
-
706
- if component_type == AIQComponentEnum.EMBEDDER_CLIENT:
707
- leaf_embedder_client_infos = {}
708
- for framework in self._embedder_client_provider_to_framework.values():
709
- for info in framework.values():
710
- leaf_embedder_client_infos[info.discovery_metadata.component_name] = info
711
- return leaf_embedder_client_infos
712
-
713
- if component_type == AIQComponentEnum.RETRIEVER_PROVIDER:
714
- return self._registered_retriever_provider_infos
715
-
716
- if component_type == AIQComponentEnum.RETRIEVER_CLIENT:
717
- leaf_retriever_client_infos = {}
718
- for framework in self._retriever_client_provider_to_framework.values():
719
- for info in framework.values():
720
- leaf_retriever_client_infos[info.discovery_metadata.component_name] = info
721
- return leaf_retriever_client_infos
722
-
723
- if component_type == AIQComponentEnum.EVALUATOR:
724
- return self._registered_evaluator_infos
725
-
726
- if component_type == AIQComponentEnum.MEMORY:
727
- return self._registered_memory_infos
728
-
729
- if component_type == AIQComponentEnum.REGISTRY_HANDLER:
730
- return self._registered_registry_handler_infos
731
-
732
- if component_type == AIQComponentEnum.LOGGING:
733
- return self._registered_logging_methods
734
-
735
- if component_type == AIQComponentEnum.TRACING:
736
- return self._registered_telemetry_exporters
737
-
738
- if component_type == AIQComponentEnum.PACKAGE:
739
- return self._registered_packages
740
-
741
- raise ValueError(f"Supplied an unsupported component type {component_type}")
742
-
743
- def get_registered_types_by_component_type( # pylint: disable=R0911
744
- self, component_type: AIQComponentEnum) -> list[str]:
745
-
746
- if component_type == AIQComponentEnum.FUNCTION:
747
- return [i.static_type() for i in self._registered_functions]
748
-
749
- if component_type == AIQComponentEnum.TOOL_WRAPPER:
750
- return list(self._registered_tool_wrappers)
751
-
752
- if component_type == AIQComponentEnum.LLM_PROVIDER:
753
- return [i.static_type() for i in self._registered_llm_provider_infos]
754
-
755
- if component_type == AIQComponentEnum.LLM_CLIENT:
756
- leaf_client_provider_framework_types = []
757
- for framework in self._llm_client_provider_to_framework.values():
758
- for info in framework.values():
759
- leaf_client_provider_framework_types.append([info.discovery_metadata.component_name])
760
- return leaf_client_provider_framework_types
761
-
762
- if component_type == AIQComponentEnum.EMBEDDER_PROVIDER:
763
- return [i.static_type() for i in self._registered_embedder_provider_infos]
764
-
765
- if component_type == AIQComponentEnum.EMBEDDER_CLIENT:
766
- leaf_embedder_provider_framework_types = []
767
- for framework in self._embedder_client_provider_to_framework.values():
768
- for info in framework.values():
769
- leaf_embedder_provider_framework_types.append([info.discovery_metadata.component_name])
770
- return leaf_embedder_provider_framework_types
771
-
772
- if component_type == AIQComponentEnum.EVALUATOR:
773
- return [i.static_type() for i in self._registered_evaluator_infos]
774
-
775
- if component_type == AIQComponentEnum.MEMORY:
776
- return [i.static_type() for i in self._registered_memory_infos]
777
-
778
- if component_type == AIQComponentEnum.REGISTRY_HANDLER:
779
- return [i.static_type() for i in self._registered_registry_handler_infos]
780
-
781
- if component_type == AIQComponentEnum.LOGGING:
782
- return [i.static_type() for i in self._registered_logging_methods]
783
-
784
- if component_type == AIQComponentEnum.TRACING:
785
- return [i.static_type() for i in self._registered_telemetry_exporters]
786
-
787
- if component_type == AIQComponentEnum.PACKAGE:
788
- return list(self._registered_packages)
789
-
790
- raise ValueError(f"Supplied an unsupported component type {component_type}")
791
-
792
- def get_registered_channel_info_by_channel_type(self, channel_type: str) -> RegisteredRegistryHandlerInfo:
793
- return self._registered_channel_map[channel_type]
794
-
795
- def _do_compute_annotation(self, cls: type[TypedBaseModelT], registrations: list[RegisteredInfo[TypedBaseModelT]]):
796
-
797
- while (len(registrations) < 2):
798
- registrations.append(RegisteredInfo[TypedBaseModelT](full_type=f"_ignore/{len(registrations)}",
799
- config_type=cls))
800
-
801
- short_names: dict[str, int] = {}
802
- type_list: list[tuple[str, type[TypedBaseModelT]]] = []
803
-
804
- # For all keys in the list, split the key by / and increment the count of the last element
805
- for key in registrations:
806
- short_names[key.local_name] = short_names.get(key.local_name, 0) + 1
807
-
808
- type_list.append((key.full_type, key.config_type))
809
-
810
- # Now loop again and if the short name is unique, then create two entries, for the short and full name
811
- for key in registrations:
812
-
813
- if (short_names[key.local_name] == 1):
814
- type_list.append((key.local_name, key.config_type))
815
-
816
- # pylint: disable=consider-alternative-union-syntax
817
- return typing.Union[tuple(typing.Annotated[x_type, Tag(x_id)] for x_id, x_type in type_list)]
818
-
819
- def compute_annotation(self, cls: type[TypedBaseModelT]):
820
-
821
- if issubclass(cls, EmbedderBaseConfig):
822
- return self._do_compute_annotation(cls, self.get_registered_embedder_providers())
823
-
824
- if issubclass(cls, EvaluatorBaseConfig):
825
- return self._do_compute_annotation(cls, self.get_registered_evaluators())
826
-
827
- if issubclass(cls, FrontEndBaseConfig):
828
- return self._do_compute_annotation(cls, self.get_registered_front_ends())
829
-
830
- if issubclass(cls, FunctionBaseConfig):
831
- return self._do_compute_annotation(cls, self.get_registered_functions())
832
-
833
- if issubclass(cls, LLMBaseConfig):
834
- return self._do_compute_annotation(cls, self.get_registered_llm_providers())
835
-
836
- if issubclass(cls, MemoryBaseConfig):
837
- return self._do_compute_annotation(cls, self.get_registered_memorys())
838
-
839
- if issubclass(cls, RegistryHandlerBaseConfig):
840
- return self._do_compute_annotation(cls, self.get_registered_registry_handlers())
841
-
842
- if issubclass(cls, RetrieverBaseConfig):
843
- return self._do_compute_annotation(cls, self.get_registered_retriever_providers())
844
-
845
- if issubclass(cls, TelemetryExporterBaseConfig):
846
- return self._do_compute_annotation(cls, self.get_registered_telemetry_exporters())
847
-
848
- if issubclass(cls, LoggingBaseConfig):
849
- return self._do_compute_annotation(cls, self.get_registered_logging_method())
850
-
851
- raise ValueError(f"Supplied an unsupported component type {cls}")
852
-
853
-
854
- class GlobalTypeRegistry:
855
-
856
- _global_registry: TypeRegistry = TypeRegistry()
857
-
858
- @staticmethod
859
- def get() -> TypeRegistry:
860
- return GlobalTypeRegistry._global_registry
861
-
862
- @staticmethod
863
- @contextmanager
864
- def push():
865
-
866
- saved = GlobalTypeRegistry._global_registry
867
- registry = deepcopy(saved)
868
-
869
- try:
870
- GlobalTypeRegistry._global_registry = registry
871
-
872
- yield registry
873
- finally:
874
- GlobalTypeRegistry._global_registry = saved
875
- GlobalTypeRegistry._global_registry._registration_changed()
876
-
877
-
878
- # Finally, update the Config object each time the registry changes
879
- GlobalTypeRegistry.get().add_registration_changed_hook(lambda: AIQConfig.rebuild_annotations())