aiqtoolkit 1.2.0a20250706__py3-none-any.whl → 1.2.0a20250730__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 (197) hide show
  1. aiq/agent/base.py +171 -8
  2. aiq/agent/dual_node.py +1 -1
  3. aiq/agent/react_agent/agent.py +113 -113
  4. aiq/agent/react_agent/register.py +31 -14
  5. aiq/agent/rewoo_agent/agent.py +36 -35
  6. aiq/agent/rewoo_agent/register.py +2 -2
  7. aiq/agent/tool_calling_agent/agent.py +3 -7
  8. aiq/authentication/__init__.py +14 -0
  9. aiq/authentication/api_key/__init__.py +14 -0
  10. aiq/authentication/api_key/api_key_auth_provider.py +92 -0
  11. aiq/authentication/api_key/api_key_auth_provider_config.py +124 -0
  12. aiq/authentication/api_key/register.py +26 -0
  13. aiq/authentication/exceptions/__init__.py +14 -0
  14. aiq/authentication/exceptions/api_key_exceptions.py +38 -0
  15. aiq/authentication/exceptions/auth_code_grant_exceptions.py +86 -0
  16. aiq/authentication/exceptions/call_back_exceptions.py +38 -0
  17. aiq/authentication/exceptions/request_exceptions.py +54 -0
  18. aiq/authentication/http_basic_auth/__init__.py +0 -0
  19. aiq/authentication/http_basic_auth/http_basic_auth_provider.py +81 -0
  20. aiq/authentication/http_basic_auth/register.py +30 -0
  21. aiq/authentication/interfaces.py +93 -0
  22. aiq/authentication/oauth2/__init__.py +14 -0
  23. aiq/authentication/oauth2/oauth2_auth_code_flow_provider.py +107 -0
  24. aiq/authentication/oauth2/oauth2_auth_code_flow_provider_config.py +39 -0
  25. aiq/authentication/oauth2/register.py +25 -0
  26. aiq/authentication/register.py +21 -0
  27. aiq/builder/builder.py +64 -2
  28. aiq/builder/component_utils.py +16 -3
  29. aiq/builder/context.py +26 -0
  30. aiq/builder/eval_builder.py +43 -2
  31. aiq/builder/function.py +32 -4
  32. aiq/builder/function_base.py +1 -1
  33. aiq/builder/intermediate_step_manager.py +6 -8
  34. aiq/builder/user_interaction_manager.py +3 -0
  35. aiq/builder/workflow.py +23 -18
  36. aiq/builder/workflow_builder.py +420 -73
  37. aiq/cli/commands/info/list_mcp.py +103 -16
  38. aiq/cli/commands/sizing/__init__.py +14 -0
  39. aiq/cli/commands/sizing/calc.py +294 -0
  40. aiq/cli/commands/sizing/sizing.py +27 -0
  41. aiq/cli/commands/start.py +1 -0
  42. aiq/cli/entrypoint.py +2 -0
  43. aiq/cli/register_workflow.py +80 -0
  44. aiq/cli/type_registry.py +151 -30
  45. aiq/data_models/api_server.py +117 -11
  46. aiq/data_models/authentication.py +231 -0
  47. aiq/data_models/common.py +35 -7
  48. aiq/data_models/component.py +17 -9
  49. aiq/data_models/component_ref.py +33 -0
  50. aiq/data_models/config.py +60 -3
  51. aiq/data_models/embedder.py +1 -0
  52. aiq/data_models/function_dependencies.py +8 -0
  53. aiq/data_models/interactive.py +10 -1
  54. aiq/data_models/intermediate_step.py +15 -5
  55. aiq/data_models/its_strategy.py +30 -0
  56. aiq/data_models/llm.py +1 -0
  57. aiq/data_models/memory.py +1 -0
  58. aiq/data_models/object_store.py +44 -0
  59. aiq/data_models/retry_mixin.py +35 -0
  60. aiq/data_models/span.py +187 -0
  61. aiq/data_models/telemetry_exporter.py +2 -2
  62. aiq/embedder/nim_embedder.py +2 -1
  63. aiq/embedder/openai_embedder.py +2 -1
  64. aiq/eval/config.py +19 -1
  65. aiq/eval/dataset_handler/dataset_handler.py +75 -1
  66. aiq/eval/evaluate.py +53 -10
  67. aiq/eval/rag_evaluator/evaluate.py +23 -12
  68. aiq/eval/remote_workflow.py +7 -2
  69. aiq/eval/runners/__init__.py +14 -0
  70. aiq/eval/runners/config.py +39 -0
  71. aiq/eval/runners/multi_eval_runner.py +54 -0
  72. aiq/eval/usage_stats.py +6 -0
  73. aiq/eval/utils/weave_eval.py +5 -1
  74. aiq/experimental/__init__.py +0 -0
  75. aiq/experimental/decorators/__init__.py +0 -0
  76. aiq/experimental/decorators/experimental_warning_decorator.py +130 -0
  77. aiq/experimental/inference_time_scaling/__init__.py +0 -0
  78. aiq/experimental/inference_time_scaling/editing/__init__.py +0 -0
  79. aiq/experimental/inference_time_scaling/editing/iterative_plan_refinement_editor.py +147 -0
  80. aiq/experimental/inference_time_scaling/editing/llm_as_a_judge_editor.py +204 -0
  81. aiq/experimental/inference_time_scaling/editing/motivation_aware_summarization.py +107 -0
  82. aiq/experimental/inference_time_scaling/functions/__init__.py +0 -0
  83. aiq/experimental/inference_time_scaling/functions/execute_score_select_function.py +105 -0
  84. aiq/experimental/inference_time_scaling/functions/its_tool_orchestration_function.py +205 -0
  85. aiq/experimental/inference_time_scaling/functions/its_tool_wrapper_function.py +146 -0
  86. aiq/experimental/inference_time_scaling/functions/plan_select_execute_function.py +224 -0
  87. aiq/experimental/inference_time_scaling/models/__init__.py +0 -0
  88. aiq/experimental/inference_time_scaling/models/editor_config.py +132 -0
  89. aiq/experimental/inference_time_scaling/models/its_item.py +48 -0
  90. aiq/experimental/inference_time_scaling/models/scoring_config.py +112 -0
  91. aiq/experimental/inference_time_scaling/models/search_config.py +120 -0
  92. aiq/experimental/inference_time_scaling/models/selection_config.py +154 -0
  93. aiq/experimental/inference_time_scaling/models/stage_enums.py +43 -0
  94. aiq/experimental/inference_time_scaling/models/strategy_base.py +66 -0
  95. aiq/experimental/inference_time_scaling/models/tool_use_config.py +41 -0
  96. aiq/experimental/inference_time_scaling/register.py +36 -0
  97. aiq/experimental/inference_time_scaling/scoring/__init__.py +0 -0
  98. aiq/experimental/inference_time_scaling/scoring/llm_based_agent_scorer.py +168 -0
  99. aiq/experimental/inference_time_scaling/scoring/llm_based_plan_scorer.py +168 -0
  100. aiq/experimental/inference_time_scaling/scoring/motivation_aware_scorer.py +111 -0
  101. aiq/experimental/inference_time_scaling/search/__init__.py +0 -0
  102. aiq/experimental/inference_time_scaling/search/multi_llm_planner.py +128 -0
  103. aiq/experimental/inference_time_scaling/search/multi_query_retrieval_search.py +122 -0
  104. aiq/experimental/inference_time_scaling/search/single_shot_multi_plan_planner.py +128 -0
  105. aiq/experimental/inference_time_scaling/selection/__init__.py +0 -0
  106. aiq/experimental/inference_time_scaling/selection/best_of_n_selector.py +63 -0
  107. aiq/experimental/inference_time_scaling/selection/llm_based_agent_output_selector.py +131 -0
  108. aiq/experimental/inference_time_scaling/selection/llm_based_output_merging_selector.py +159 -0
  109. aiq/experimental/inference_time_scaling/selection/llm_based_plan_selector.py +128 -0
  110. aiq/experimental/inference_time_scaling/selection/threshold_selector.py +58 -0
  111. aiq/front_ends/console/authentication_flow_handler.py +233 -0
  112. aiq/front_ends/console/console_front_end_plugin.py +11 -2
  113. aiq/front_ends/fastapi/auth_flow_handlers/__init__.py +0 -0
  114. aiq/front_ends/fastapi/auth_flow_handlers/http_flow_handler.py +27 -0
  115. aiq/front_ends/fastapi/auth_flow_handlers/websocket_flow_handler.py +107 -0
  116. aiq/front_ends/fastapi/fastapi_front_end_config.py +20 -0
  117. aiq/front_ends/fastapi/fastapi_front_end_controller.py +68 -0
  118. aiq/front_ends/fastapi/fastapi_front_end_plugin.py +14 -1
  119. aiq/front_ends/fastapi/fastapi_front_end_plugin_worker.py +353 -31
  120. aiq/front_ends/fastapi/html_snippets/__init__.py +14 -0
  121. aiq/front_ends/fastapi/html_snippets/auth_code_grant_success.py +35 -0
  122. aiq/front_ends/fastapi/main.py +2 -0
  123. aiq/front_ends/fastapi/message_handler.py +102 -84
  124. aiq/front_ends/fastapi/step_adaptor.py +2 -1
  125. aiq/llm/aws_bedrock_llm.py +2 -1
  126. aiq/llm/nim_llm.py +2 -1
  127. aiq/llm/openai_llm.py +2 -1
  128. aiq/object_store/__init__.py +20 -0
  129. aiq/object_store/in_memory_object_store.py +74 -0
  130. aiq/object_store/interfaces.py +84 -0
  131. aiq/object_store/models.py +36 -0
  132. aiq/object_store/register.py +20 -0
  133. aiq/observability/__init__.py +14 -0
  134. aiq/observability/exporter/__init__.py +14 -0
  135. aiq/observability/exporter/base_exporter.py +449 -0
  136. aiq/observability/exporter/exporter.py +78 -0
  137. aiq/observability/exporter/file_exporter.py +33 -0
  138. aiq/observability/exporter/processing_exporter.py +269 -0
  139. aiq/observability/exporter/raw_exporter.py +52 -0
  140. aiq/observability/exporter/span_exporter.py +264 -0
  141. aiq/observability/exporter_manager.py +335 -0
  142. aiq/observability/mixin/__init__.py +14 -0
  143. aiq/observability/mixin/batch_config_mixin.py +26 -0
  144. aiq/observability/mixin/collector_config_mixin.py +23 -0
  145. aiq/observability/mixin/file_mixin.py +288 -0
  146. aiq/observability/mixin/file_mode.py +23 -0
  147. aiq/observability/mixin/resource_conflict_mixin.py +134 -0
  148. aiq/observability/mixin/serialize_mixin.py +61 -0
  149. aiq/observability/mixin/type_introspection_mixin.py +183 -0
  150. aiq/observability/processor/__init__.py +14 -0
  151. aiq/observability/processor/batching_processor.py +316 -0
  152. aiq/observability/processor/intermediate_step_serializer.py +28 -0
  153. aiq/observability/processor/processor.py +68 -0
  154. aiq/observability/register.py +32 -116
  155. aiq/observability/utils/__init__.py +14 -0
  156. aiq/observability/utils/dict_utils.py +236 -0
  157. aiq/observability/utils/time_utils.py +31 -0
  158. aiq/profiler/calc/__init__.py +14 -0
  159. aiq/profiler/calc/calc_runner.py +623 -0
  160. aiq/profiler/calc/calculations.py +288 -0
  161. aiq/profiler/calc/data_models.py +176 -0
  162. aiq/profiler/calc/plot.py +345 -0
  163. aiq/profiler/data_models.py +2 -0
  164. aiq/profiler/profile_runner.py +16 -13
  165. aiq/runtime/loader.py +8 -2
  166. aiq/runtime/runner.py +23 -9
  167. aiq/runtime/session.py +16 -5
  168. aiq/tool/chat_completion.py +74 -0
  169. aiq/tool/code_execution/README.md +152 -0
  170. aiq/tool/code_execution/code_sandbox.py +151 -72
  171. aiq/tool/code_execution/local_sandbox/.gitignore +1 -0
  172. aiq/tool/code_execution/local_sandbox/local_sandbox_server.py +139 -24
  173. aiq/tool/code_execution/local_sandbox/sandbox.requirements.txt +3 -1
  174. aiq/tool/code_execution/local_sandbox/start_local_sandbox.sh +27 -2
  175. aiq/tool/code_execution/register.py +7 -3
  176. aiq/tool/code_execution/test_code_execution_sandbox.py +414 -0
  177. aiq/tool/mcp/exceptions.py +142 -0
  178. aiq/tool/mcp/mcp_client.py +17 -3
  179. aiq/tool/mcp/mcp_tool.py +1 -1
  180. aiq/tool/register.py +1 -0
  181. aiq/tool/server_tools.py +2 -2
  182. aiq/utils/exception_handlers/automatic_retries.py +289 -0
  183. aiq/utils/exception_handlers/mcp.py +211 -0
  184. aiq/utils/io/model_processing.py +28 -0
  185. aiq/utils/log_utils.py +37 -0
  186. aiq/utils/string_utils.py +38 -0
  187. aiq/utils/type_converter.py +18 -2
  188. aiq/utils/type_utils.py +87 -0
  189. {aiqtoolkit-1.2.0a20250706.dist-info → aiqtoolkit-1.2.0a20250730.dist-info}/METADATA +37 -9
  190. {aiqtoolkit-1.2.0a20250706.dist-info → aiqtoolkit-1.2.0a20250730.dist-info}/RECORD +195 -80
  191. {aiqtoolkit-1.2.0a20250706.dist-info → aiqtoolkit-1.2.0a20250730.dist-info}/entry_points.txt +3 -0
  192. aiq/front_ends/fastapi/websocket.py +0 -153
  193. aiq/observability/async_otel_listener.py +0 -470
  194. {aiqtoolkit-1.2.0a20250706.dist-info → aiqtoolkit-1.2.0a20250730.dist-info}/WHEEL +0 -0
  195. {aiqtoolkit-1.2.0a20250706.dist-info → aiqtoolkit-1.2.0a20250730.dist-info}/licenses/LICENSE-3rd-party.txt +0 -0
  196. {aiqtoolkit-1.2.0a20250706.dist-info → aiqtoolkit-1.2.0a20250730.dist-info}/licenses/LICENSE.md +0 -0
  197. {aiqtoolkit-1.2.0a20250706.dist-info → aiqtoolkit-1.2.0a20250730.dist-info}/top_level.txt +0 -0
@@ -21,8 +21,10 @@ from contextlib import AbstractAsyncContextManager
21
21
  from contextlib import AsyncExitStack
22
22
  from contextlib import asynccontextmanager
23
23
 
24
+ from aiq.authentication.interfaces import AuthProviderBase
24
25
  from aiq.builder.builder import Builder
25
26
  from aiq.builder.builder import UserManagerHolder
27
+ from aiq.builder.component_utils import ComponentInstanceData
26
28
  from aiq.builder.component_utils import build_dependency_sequence
27
29
  from aiq.builder.context import AIQContext
28
30
  from aiq.builder.context import AIQContextState
@@ -36,43 +38,45 @@ from aiq.builder.retriever import RetrieverProviderInfo
36
38
  from aiq.builder.workflow import Workflow
37
39
  from aiq.cli.type_registry import GlobalTypeRegistry
38
40
  from aiq.cli.type_registry import TypeRegistry
41
+ from aiq.data_models.authentication import AuthProviderBaseConfig
39
42
  from aiq.data_models.component import ComponentGroup
43
+ from aiq.data_models.component_ref import AuthenticationRef
40
44
  from aiq.data_models.component_ref import EmbedderRef
41
45
  from aiq.data_models.component_ref import FunctionRef
46
+ from aiq.data_models.component_ref import ITSStrategyRef
42
47
  from aiq.data_models.component_ref import LLMRef
43
48
  from aiq.data_models.component_ref import MemoryRef
49
+ from aiq.data_models.component_ref import ObjectStoreRef
44
50
  from aiq.data_models.component_ref import RetrieverRef
45
51
  from aiq.data_models.config import AIQConfig
46
52
  from aiq.data_models.config import GeneralConfig
47
53
  from aiq.data_models.embedder import EmbedderBaseConfig
48
54
  from aiq.data_models.function import FunctionBaseConfig
49
55
  from aiq.data_models.function_dependencies import FunctionDependencies
56
+ from aiq.data_models.its_strategy import ITSStrategyBaseConfig
50
57
  from aiq.data_models.llm import LLMBaseConfig
51
58
  from aiq.data_models.memory import MemoryBaseConfig
59
+ from aiq.data_models.object_store import ObjectStoreBaseConfig
52
60
  from aiq.data_models.retriever import RetrieverBaseConfig
53
61
  from aiq.data_models.telemetry_exporter import TelemetryExporterBaseConfig
62
+ from aiq.experimental.decorators.experimental_warning_decorator import aiq_experimental
63
+ from aiq.experimental.inference_time_scaling.models.stage_enums import PipelineTypeEnum
64
+ from aiq.experimental.inference_time_scaling.models.stage_enums import StageTypeEnum
65
+ from aiq.experimental.inference_time_scaling.models.strategy_base import StrategyBase
54
66
  from aiq.memory.interfaces import MemoryEditor
67
+ from aiq.object_store.interfaces import ObjectStore
68
+ from aiq.observability.exporter.base_exporter import BaseExporter
55
69
  from aiq.profiler.decorators.framework_wrapper import chain_wrapped_build_fn
56
70
  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
71
  from aiq.utils.type_utils import override
60
72
 
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
73
  logger = logging.getLogger(__name__)
70
74
 
71
75
 
72
76
  @dataclasses.dataclass
73
- class ConfiguredExporter:
77
+ class ConfiguredTelemetryExporter:
74
78
  config: TelemetryExporterBaseConfig
75
- instance: SpanExporter
79
+ instance: BaseExporter
76
80
 
77
81
 
78
82
  @dataclasses.dataclass
@@ -99,12 +103,30 @@ class ConfiguredMemory:
99
103
  instance: MemoryEditor
100
104
 
101
105
 
106
+ @dataclasses.dataclass
107
+ class ConfiguredObjectStore:
108
+ config: ObjectStoreBaseConfig
109
+ instance: ObjectStore
110
+
111
+
102
112
  @dataclasses.dataclass
103
113
  class ConfiguredRetriever:
104
114
  config: RetrieverBaseConfig
105
115
  instance: RetrieverProviderInfo
106
116
 
107
117
 
118
+ @dataclasses.dataclass
119
+ class ConfiguredAuthProvider:
120
+ config: AuthProviderBaseConfig
121
+ instance: AuthProviderBase
122
+
123
+
124
+ @dataclasses.dataclass
125
+ class ConfiguredITSStrategy:
126
+ config: ITSStrategyBaseConfig
127
+ instance: StrategyBase
128
+
129
+
108
130
  # pylint: disable=too-many-public-methods
109
131
  class WorkflowBuilder(Builder, AbstractAsyncContextManager):
110
132
 
@@ -121,15 +143,18 @@ class WorkflowBuilder(Builder, AbstractAsyncContextManager):
121
143
  self._registry = registry
122
144
 
123
145
  self._logging_handlers: dict[str, logging.Handler] = {}
124
- self._exporters: dict[str, ConfiguredExporter] = {}
146
+ self._telemetry_exporters: dict[str, ConfiguredTelemetryExporter] = {}
125
147
 
126
148
  self._functions: dict[str, ConfiguredFunction] = {}
127
149
  self._workflow: ConfiguredFunction | None = None
128
150
 
129
151
  self._llms: dict[str, ConfiguredLLM] = {}
152
+ self._auth_providers: dict[str, ConfiguredAuthProvider] = {}
130
153
  self._embedders: dict[str, ConfiguredEmbedder] = {}
131
154
  self._memory_clients: dict[str, ConfiguredMemory] = {}
155
+ self._object_stores: dict[str, ConfiguredObjectStore] = {}
132
156
  self._retrievers: dict[str, ConfiguredRetriever] = {}
157
+ self._its_strategies: dict[str, ConfiguredITSStrategy] = {}
133
158
 
134
159
  self._context_state = AIQContextState.get()
135
160
 
@@ -143,7 +168,7 @@ class WorkflowBuilder(Builder, AbstractAsyncContextManager):
143
168
 
144
169
  self._exit_stack = AsyncExitStack()
145
170
 
146
- # Get the exporter info from the config
171
+ # Get the telemetry info from the config
147
172
  telemetry_config = self.general_config.telemetry
148
173
 
149
174
  for key, logging_config in telemetry_config.logging.items():
@@ -161,43 +186,9 @@ class WorkflowBuilder(Builder, AbstractAsyncContextManager):
161
186
  # Now attach to AIQ Toolkit's root logger
162
187
  logging.getLogger().addHandler(handler)
163
188
 
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 openinference.semconv.resource import ResourceAttributes
171
- from opentelemetry import trace
172
- from opentelemetry.sdk.resources import Resource
173
- from opentelemetry.sdk.trace import TracerProvider
174
- from opentelemetry.sdk.trace.export import BatchSpanProcessor
175
-
176
- from aiq.observability.register import PhoenixTelemetryExporter
177
-
178
- # Create a default provider first
179
- provider = TracerProvider()
180
-
181
- # Check if we have a phoenix telemetry exporter and use its project name
182
- for key, trace_exporter_config in telemetry_config.tracing.items():
183
- if isinstance(trace_exporter_config, PhoenixTelemetryExporter):
184
- provider = TracerProvider(resource=Resource(
185
- attributes={ResourceAttributes.PROJECT_NAME: trace_exporter_config.project}))
186
- break
187
-
188
- trace.set_tracer_provider(provider)
189
-
190
- for key, trace_exporter_config in telemetry_config.tracing.items():
191
-
192
- exporter_info = self._registry.get_telemetry_exporter(type(trace_exporter_config))
193
-
194
- instance = await self._exit_stack.enter_async_context(
195
- exporter_info.build_fn(trace_exporter_config, self))
196
-
197
- span_processor_instance = BatchSpanProcessor(instance)
198
- provider.add_span_processor(span_processor_instance)
199
-
200
- self._exporters[key] = ConfiguredExporter(config=trace_exporter_config, instance=instance)
189
+ # Add the telemetry exporters
190
+ for key, telemetry_exporter_config in telemetry_config.tracing.items():
191
+ await self.add_telemetry_exporter(key, telemetry_exporter_config)
201
192
 
202
193
  return self
203
194
 
@@ -253,9 +244,17 @@ class WorkflowBuilder(Builder, AbstractAsyncContextManager):
253
244
  k: v.config
254
245
  for k, v in self._memory_clients.items()
255
246
  },
247
+ object_stores={
248
+ k: v.config
249
+ for k, v in self._object_stores.items()
250
+ },
256
251
  retrievers={
257
252
  k: v.config
258
253
  for k, v in self._retrievers.items()
254
+ },
255
+ its_strategies={
256
+ k: v.config
257
+ for k, v in self._its_strategies.items()
259
258
  })
260
259
 
261
260
  if (entry_function is None):
@@ -281,14 +280,22 @@ class WorkflowBuilder(Builder, AbstractAsyncContextManager):
281
280
  k: v.instance
282
281
  for k, v in self._memory_clients.items()
283
282
  },
284
- exporters={
283
+ object_stores={
284
+ k: v.instance
285
+ for k, v in self._object_stores.items()
286
+ },
287
+ telemetry_exporters={
285
288
  k: v.instance
286
- for k, v in self._exporters.items()
289
+ for k, v in self._telemetry_exporters.items()
287
290
  },
288
291
  retrievers={
289
292
  k: v.instance
290
293
  for k, v in self._retrievers.items()
291
294
  },
295
+ its_strategies={
296
+ k: v.instance
297
+ for k, v in self._its_strategies.items()
298
+ },
292
299
  context_state=self._context_state)
293
300
 
294
301
  return workflow
@@ -463,6 +470,76 @@ class WorkflowBuilder(Builder, AbstractAsyncContextManager):
463
470
  # Return the tool configuration object
464
471
  return self._llms[llm_name].config
465
472
 
473
+ @aiq_experimental(feature_name="Authentication")
474
+ @override
475
+ async def add_auth_provider(self, name: str | AuthenticationRef,
476
+ config: AuthProviderBaseConfig) -> AuthProviderBase:
477
+ """
478
+ Add an authentication provider to the workflow by constructing it from a configuration object.
479
+
480
+ Note: The Authentication Provider API is experimental and the API may change in future releases.
481
+
482
+ Parameters
483
+ ----------
484
+ name : str | AuthenticationRef
485
+ The name of the authentication provider to add.
486
+ config : AuthProviderBaseConfig
487
+ The configuration for the authentication provider.
488
+
489
+ Returns
490
+ -------
491
+ AuthProviderBase
492
+ The authentication provider instance.
493
+
494
+ Raises
495
+ ------
496
+ ValueError
497
+ If the authentication provider is already in the list of authentication providers.
498
+ """
499
+
500
+ if (name in self._auth_providers):
501
+ raise ValueError(f"Authentication `{name}` already exists in the list of Authentication Providers")
502
+
503
+ try:
504
+ authentication_info = self._registry.get_auth_provider(type(config))
505
+
506
+ info_obj = await self._get_exit_stack().enter_async_context(authentication_info.build_fn(config, self))
507
+
508
+ self._auth_providers[name] = ConfiguredAuthProvider(config=config, instance=info_obj)
509
+
510
+ return info_obj
511
+ except Exception as e:
512
+ logger.error("Error adding authentication `%s` with config `%s`", name, config, exc_info=True)
513
+ raise e
514
+
515
+ @override
516
+ async def get_auth_provider(self, auth_provider_name: str) -> AuthProviderBase:
517
+ """
518
+ Get the authentication provider instance for the given name.
519
+
520
+ Note: The Authentication Provider API is experimental and the API may change in future releases.
521
+
522
+ Parameters
523
+ ----------
524
+ auth_provider_name : str
525
+ The name of the authentication provider to get.
526
+
527
+ Returns
528
+ -------
529
+ AuthProviderBase
530
+ The authentication provider instance.
531
+
532
+ Raises
533
+ ------
534
+ ValueError
535
+ If the authentication provider is not found.
536
+ """
537
+
538
+ if auth_provider_name not in self._auth_providers:
539
+ raise ValueError(f"Authentication `{auth_provider_name}` not found")
540
+
541
+ return self._auth_providers[auth_provider_name].instance
542
+
466
543
  @override
467
544
  async def add_embedder(self, name: str | EmbedderRef, config: EmbedderBaseConfig):
468
545
 
@@ -543,6 +620,33 @@ class WorkflowBuilder(Builder, AbstractAsyncContextManager):
543
620
  # Return the tool configuration object
544
621
  return self._memory_clients[memory_name].config
545
622
 
623
+ @override
624
+ async def add_object_store(self, name: str | ObjectStoreRef, config: ObjectStoreBaseConfig) -> ObjectStore:
625
+ if name in self._object_stores:
626
+ raise ValueError(f"Object store `{name}` already exists in the list of object stores")
627
+
628
+ object_store_info = self._registry.get_object_store(type(config))
629
+
630
+ info_obj = await self._get_exit_stack().enter_async_context(object_store_info.build_fn(config, self))
631
+
632
+ self._object_stores[name] = ConfiguredObjectStore(config=config, instance=info_obj)
633
+
634
+ return info_obj
635
+
636
+ @override
637
+ async def get_object_store_client(self, object_store_name: str | ObjectStoreRef) -> ObjectStore:
638
+ if object_store_name not in self._object_stores:
639
+ raise ValueError(f"Object store `{object_store_name}` not found")
640
+
641
+ return self._object_stores[object_store_name].instance
642
+
643
+ @override
644
+ def get_object_store_config(self, object_store_name: str | ObjectStoreRef) -> ObjectStoreBaseConfig:
645
+ if object_store_name not in self._object_stores:
646
+ raise ValueError(f"Object store `{object_store_name}` not found")
647
+
648
+ return self._object_stores[object_store_name].config
649
+
546
650
  @override
547
651
  async def add_retriever(self, name: str | RetrieverRef, config: RetrieverBaseConfig):
548
652
 
@@ -595,10 +699,167 @@ class WorkflowBuilder(Builder, AbstractAsyncContextManager):
595
699
 
596
700
  return self._retrievers[retriever_name].config
597
701
 
702
+ @aiq_experimental(feature_name="ITS")
703
+ @override
704
+ async def add_its_strategy(self, name: str | str, config: ITSStrategyBaseConfig):
705
+ if (name in self._its_strategies):
706
+ raise ValueError(f"ITS strategy '{name}' already exists in the list of ITS strategies")
707
+
708
+ try:
709
+ its_strategy_info = self._registry.get_its_strategy(type(config))
710
+
711
+ info_obj = await self._get_exit_stack().enter_async_context(its_strategy_info.build_fn(config, self))
712
+
713
+ self._its_strategies[name] = ConfiguredITSStrategy(config=config, instance=info_obj)
714
+
715
+ except Exception as e:
716
+ logger.error("Error adding ITS strategy `%s` with config `%s`", name, config, exc_info=True)
717
+
718
+ raise e
719
+
720
+ @override
721
+ async def get_its_strategy(self,
722
+ strategy_name: str | ITSStrategyRef,
723
+ pipeline_type: PipelineTypeEnum,
724
+ stage_type: StageTypeEnum) -> StrategyBase:
725
+
726
+ if strategy_name not in self._its_strategies:
727
+ raise ValueError(f"ITS strategy '{strategy_name}' not found")
728
+
729
+ try:
730
+ # Get strategy info
731
+ its_strategy_info = self._its_strategies[strategy_name]
732
+
733
+ instance = its_strategy_info.instance
734
+
735
+ if not stage_type == instance.stage_type():
736
+ raise ValueError(f"ITS strategy '{strategy_name}' is not compatible with stage type '{stage_type}'")
737
+
738
+ if pipeline_type not in instance.supported_pipeline_types():
739
+ raise ValueError(
740
+ f"ITS strategy '{strategy_name}' is not compatible with pipeline type '{pipeline_type}'")
741
+
742
+ instance.set_pipeline_type(pipeline_type)
743
+
744
+ return instance
745
+ except Exception as e:
746
+ logger.error("Error getting ITS strategy `%s`", strategy_name, exc_info=True)
747
+ raise e
748
+
749
+ @override
750
+ async def get_its_strategy_config(self,
751
+ strategy_name: str | ITSStrategyRef,
752
+ pipeline_type: PipelineTypeEnum,
753
+ stage_type: StageTypeEnum) -> ITSStrategyBaseConfig:
754
+ if strategy_name not in self._its_strategies:
755
+ raise ValueError(f"ITS strategy '{strategy_name}' not found")
756
+
757
+ strategy_info = self._its_strategies[strategy_name]
758
+ instance = strategy_info.instance
759
+ config = strategy_info.config
760
+
761
+ if not stage_type == instance.stage_type():
762
+ raise ValueError(f"ITS strategy '{strategy_name}' is not compatible with stage type '{stage_type}'")
763
+
764
+ if pipeline_type not in instance.supported_pipeline_types():
765
+ raise ValueError(f"ITS strategy '{strategy_name}' is not compatible with pipeline type '{pipeline_type}'")
766
+
767
+ return config
768
+
598
769
  @override
599
770
  def get_user_manager(self):
600
771
  return UserManagerHolder(context=AIQContext(self._context_state))
601
772
 
773
+ async def add_telemetry_exporter(self, name: str, config: TelemetryExporterBaseConfig) -> None:
774
+ """Add an configured telemetry exporter to the builder.
775
+
776
+ Args:
777
+ name (str): The name of the telemetry exporter
778
+ config (TelemetryExporterBaseConfig): The configuration for the exporter
779
+ """
780
+ if (name in self._telemetry_exporters):
781
+ raise ValueError(f"Telemetry exporter '{name}' already exists in the list of telemetry exporters")
782
+
783
+ exporter_info = self._registry.get_telemetry_exporter(type(config))
784
+
785
+ # Build the exporter outside the lock (parallel)
786
+ exporter_context_manager = exporter_info.build_fn(config, self)
787
+
788
+ # Only protect the shared state modifications (serialized)
789
+ exporter = await self._get_exit_stack().enter_async_context(exporter_context_manager)
790
+ self._telemetry_exporters[name] = ConfiguredTelemetryExporter(config=config, instance=exporter)
791
+
792
+ def _log_build_failure(self,
793
+ component_name: str,
794
+ component_type: str,
795
+ completed_components: list[tuple[str, str]],
796
+ remaining_components: list[tuple[str, str]],
797
+ original_error: Exception) -> None:
798
+ """
799
+ Common method to log comprehensive build failure information.
800
+
801
+ Args:
802
+ component_name (str): The name of the component that failed to build
803
+ component_type (str): The type of the component that failed to build
804
+ completed_components (list[tuple[str, str]]): List of (name, type) tuples for successfully built components
805
+ remaining_components (list[tuple[str, str]]): List of (name, type) tuples for components still to be built
806
+ original_error (Exception): The original exception that caused the failure
807
+ """
808
+ logger.error("Failed to initialize component %s (%s)", component_name, component_type)
809
+
810
+ if completed_components:
811
+ logger.error("Successfully built components:")
812
+ for name, comp_type in completed_components:
813
+ logger.error("- %s (%s)", name, comp_type)
814
+ else:
815
+ logger.error("No components were successfully built before this failure")
816
+
817
+ if remaining_components:
818
+ logger.error("Remaining components to build:")
819
+ for name, comp_type in remaining_components:
820
+ logger.error("- %s (%s)", name, comp_type)
821
+ else:
822
+ logger.error("No remaining components to build")
823
+
824
+ logger.error("Original error:", exc_info=original_error)
825
+
826
+ def _log_build_failure_component(self,
827
+ failing_component: ComponentInstanceData,
828
+ completed_components: list[tuple[str, str]],
829
+ remaining_components: list[tuple[str, str]],
830
+ original_error: Exception) -> None:
831
+ """
832
+ Log comprehensive component build failure information.
833
+
834
+ Args:
835
+ failing_component (ComponentInstanceData): The ComponentInstanceData that failed to build
836
+ completed_components (list[tuple[str, str]]): List of (name, type) tuples for successfully built components
837
+ remaining_components (list[tuple[str, str]]): List of (name, type) tuples for components still to be built
838
+ original_error (Exception): The original exception that caused the failure
839
+ """
840
+ component_name = failing_component.name
841
+ component_type = failing_component.component_group.value
842
+
843
+ self._log_build_failure(component_name,
844
+ component_type,
845
+ completed_components,
846
+ remaining_components,
847
+ original_error)
848
+
849
+ def _log_build_failure_workflow(self,
850
+ completed_components: list[tuple[str, str]],
851
+ remaining_components: list[tuple[str, str]],
852
+ original_error: Exception) -> None:
853
+ """
854
+ Log comprehensive workflow build failure information.
855
+
856
+ Args:
857
+ completed_components (list[tuple[str, str]]): List of (name, type) tuples for successfully built components
858
+ remaining_components (list[tuple[str, str]]): List of (name, type) tuples for components still to be built
859
+ original_error (Exception): The original exception that caused the failure
860
+ """
861
+ self._log_build_failure("<workflow>", "workflow", completed_components, remaining_components, original_error)
862
+
602
863
  async def populate_builder(self, config: AIQConfig, skip_workflow: bool = False):
603
864
  """
604
865
  Populate the builder with components and optionally set up the workflow.
@@ -611,31 +872,68 @@ class WorkflowBuilder(Builder, AbstractAsyncContextManager):
611
872
  # Generate the build sequence
612
873
  build_sequence = build_dependency_sequence(config)
613
874
 
875
+ # Initialize progress tracking
876
+ completed_components = []
877
+ remaining_components = [(str(comp.name), comp.component_group.value) for comp in build_sequence
878
+ if not comp.is_root]
879
+ if not skip_workflow:
880
+ remaining_components.append(("<workflow>", "workflow"))
881
+
614
882
  # Loop over all objects and add to the workflow builder
615
883
  for component_instance in build_sequence:
616
- # Instantiate a the llm
617
- if component_instance.component_group == ComponentGroup.LLMS:
618
- await self.add_llm(component_instance.name, component_instance.config)
619
- # Instantiate a the embedder
620
- elif component_instance.component_group == ComponentGroup.EMBEDDERS:
621
- await self.add_embedder(component_instance.name, component_instance.config)
622
- # Instantiate a memory client
623
- elif component_instance.component_group == ComponentGroup.MEMORY:
624
- await self.add_memory_client(component_instance.name, component_instance.config)
625
- # Instantiate a retriever client
626
- elif component_instance.component_group == ComponentGroup.RETRIEVERS:
627
- await self.add_retriever(component_instance.name, component_instance.config)
628
- # Instantiate a function
629
- elif component_instance.component_group == ComponentGroup.FUNCTIONS:
630
- # If the function is the root, set it as the workflow later
631
- if (not component_instance.is_root):
632
- await self.add_function(component_instance.name, component_instance.config)
633
- else:
634
- raise ValueError(f"Unknown component group {component_instance.component_group}")
884
+ try:
885
+ # Remove from remaining as we start building (if not root)
886
+ if not component_instance.is_root:
887
+ remaining_components.remove(
888
+ (str(component_instance.name), component_instance.component_group.value))
889
+
890
+ # Instantiate a the llm
891
+ if component_instance.component_group == ComponentGroup.LLMS:
892
+ await self.add_llm(component_instance.name, component_instance.config)
893
+ # Instantiate a the embedder
894
+ elif component_instance.component_group == ComponentGroup.EMBEDDERS:
895
+ await self.add_embedder(component_instance.name, component_instance.config)
896
+ # Instantiate a memory client
897
+ elif component_instance.component_group == ComponentGroup.MEMORY:
898
+ await self.add_memory_client(component_instance.name, component_instance.config)
899
+ # Instantiate a object store client
900
+ elif component_instance.component_group == ComponentGroup.OBJECT_STORES:
901
+ await self.add_object_store(component_instance.name, component_instance.config)
902
+ # Instantiate a retriever client
903
+ elif component_instance.component_group == ComponentGroup.RETRIEVERS:
904
+ await self.add_retriever(component_instance.name, component_instance.config)
905
+ # Instantiate a function
906
+ elif component_instance.component_group == ComponentGroup.FUNCTIONS:
907
+ # If the function is the root, set it as the workflow later
908
+ if (not component_instance.is_root):
909
+ await self.add_function(component_instance.name, component_instance.config)
910
+ elif component_instance.component_group == ComponentGroup.ITS_STRATEGIES:
911
+ await self.add_its_strategy(component_instance.name, component_instance.config)
912
+
913
+ elif component_instance.component_group == ComponentGroup.AUTHENTICATION:
914
+ await self.add_auth_provider(component_instance.name, component_instance.config)
915
+ else:
916
+ raise ValueError(f"Unknown component group {component_instance.component_group}")
917
+
918
+ # Add to completed after successful build (if not root)
919
+ if not component_instance.is_root:
920
+ completed_components.append(
921
+ (str(component_instance.name), component_instance.component_group.value))
922
+
923
+ except Exception as e:
924
+ self._log_build_failure_component(component_instance, completed_components, remaining_components, e)
925
+ raise
635
926
 
636
927
  # Instantiate the workflow
637
928
  if not skip_workflow:
638
- await self.set_workflow(config.workflow)
929
+ try:
930
+ # Remove workflow from remaining as we start building
931
+ remaining_components.remove(("<workflow>", "workflow"))
932
+ await self.set_workflow(config.workflow)
933
+ completed_components.append(("<workflow>", "workflow"))
934
+ except Exception as e:
935
+ self._log_build_failure_workflow(completed_components, remaining_components, e)
936
+ raise
639
937
 
640
938
  @classmethod
641
939
  @asynccontextmanager
@@ -700,6 +998,14 @@ class ChildBuilder(Builder):
700
998
  async def add_llm(self, name: str, config: LLMBaseConfig):
701
999
  return await self._workflow_builder.add_llm(name, config)
702
1000
 
1001
+ @override
1002
+ async def add_auth_provider(self, name: str, config: AuthProviderBaseConfig):
1003
+ return await self._workflow_builder.add_auth_provider(name, config)
1004
+
1005
+ @override
1006
+ async def get_auth_provider(self, auth_provider_name: str):
1007
+ return await self._workflow_builder.get_auth_provider(auth_provider_name)
1008
+
703
1009
  @override
704
1010
  async def get_llm(self, llm_name: str, wrapper_type: LLMFrameworkEnum | str):
705
1011
  llm = await self._workflow_builder.get_llm(llm_name, wrapper_type)
@@ -747,6 +1053,47 @@ class ChildBuilder(Builder):
747
1053
  def get_memory_client_config(self, memory_name: str) -> MemoryBaseConfig:
748
1054
  return self._workflow_builder.get_memory_client_config(memory_name=memory_name)
749
1055
 
1056
+ @override
1057
+ async def add_object_store(self, name: str, config: ObjectStoreBaseConfig):
1058
+ return await self._workflow_builder.add_object_store(name, config)
1059
+
1060
+ @override
1061
+ async def get_object_store_client(self, object_store_name: str) -> ObjectStore:
1062
+ """
1063
+ Return the instantiated object store client for the given name.
1064
+ """
1065
+ object_store_client = await self._workflow_builder.get_object_store_client(object_store_name)
1066
+
1067
+ self._dependencies.add_object_store(object_store_name)
1068
+
1069
+ return object_store_client
1070
+
1071
+ @override
1072
+ def get_object_store_config(self, object_store_name: str) -> ObjectStoreBaseConfig:
1073
+ return self._workflow_builder.get_object_store_config(object_store_name)
1074
+
1075
+ @override
1076
+ async def add_its_strategy(self, name: str, config: ITSStrategyBaseConfig):
1077
+ return await self._workflow_builder.add_its_strategy(name, config)
1078
+
1079
+ @override
1080
+ async def get_its_strategy(self,
1081
+ strategy_name: str | ITSStrategyRef,
1082
+ pipeline_type: PipelineTypeEnum,
1083
+ stage_type: StageTypeEnum) -> StrategyBase:
1084
+ return await self._workflow_builder.get_its_strategy(strategy_name=strategy_name,
1085
+ pipeline_type=pipeline_type,
1086
+ stage_type=stage_type)
1087
+
1088
+ @override
1089
+ async def get_its_strategy_config(self,
1090
+ strategy_name: str | ITSStrategyRef,
1091
+ pipeline_type: PipelineTypeEnum,
1092
+ stage_type: StageTypeEnum) -> ITSStrategyBaseConfig:
1093
+ return await self._workflow_builder.get_its_strategy_config(strategy_name=strategy_name,
1094
+ pipeline_type=pipeline_type,
1095
+ stage_type=stage_type)
1096
+
750
1097
  @override
751
1098
  async def add_retriever(self, name: str, config: RetrieverBaseConfig):
752
1099
  return await self._workflow_builder.add_retriever(name, config)