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