From 841302de63e362b358a2b0f8fe9b8adf83c5dffd Mon Sep 17 00:00:00 2001 From: Byron Wang Date: Thu, 8 Oct 2026 18:36:54 +0800 Subject: [PATCH 1/5] refactor(api): retire legacy tools and service compatibility modules --- api/.importlinter | 501 +++- .../index_tool_callback_handler.py | 102 - api/core/plugin/backwards_invocation/app.py | 329 --- api/core/plugin/backwards_invocation/tool.py | 55 - api/core/rag/retrieval/dataset_retrieval.py | 2134 ----------------- .../repositories/human_input_repository.py | 676 ------ api/core/tools/custom_tool/provider.py | 218 -- api/core/tools/tool_engine.py | 396 --- api/core/tools/tool_label_manager.py | 135 -- api/core/tools/tool_manager.py | 1208 ---------- .../dataset_multi_retriever_tool.py | 214 -- .../dataset_retriever_base_tool.py | 32 - .../dataset_retriever_tool.py | 248 -- .../tools/utils/dataset_retriever_tool.py | 144 -- api/core/tools/workflow_as_tool/provider.py | 259 -- api/core/tools/workflow_as_tool/tool.py | 436 ---- .../sync_plugin_trigger_when_app_created.py | 34 - ...nc_workflow_schedule_when_app_published.py | 83 - ...oin_when_app_published_workflow_updated.py | 69 - ...ers_when_app_published_workflow_updated.py | 112 - .../workflow_app_log_query_repository.py | 261 -- .../workflow_collaboration_repository.py | 272 --- api/services/agent/agent_soul_state.py | 6 - .../agent/legacy_retirement_service.py | 224 -- .../agent/legacy_workflow_publish_service.py | 744 ------ api/services/agent_tool_inner_service.py | 151 -- api/services/app_creation_records.py | 45 - api/services/conversation_variable_updater.py | 29 - .../knowledge/retrieval/attachments.py | 33 - .../tools/legacy_api_tools_manage_service.py | 570 ----- .../legacy_builtin_tools_manage_service.py | 794 ------ .../tools/legacy_mcp_tools_manage_service.py | 884 ------- .../tools/legacy_tools_manage_service.py | 26 - .../tools/legacy_tools_transform_service.py | 538 ----- .../legacy_workflow_tools_manage_service.py | 419 ---- .../workflow_app_log_query_service.py | 174 -- .../workflow_draft_variable_service.py | 1202 ---------- .../workflow_version_number_service.py | 59 - 38 files changed, 404 insertions(+), 13412 deletions(-) delete mode 100644 api/core/callback_handler/index_tool_callback_handler.py delete mode 100644 api/core/plugin/backwards_invocation/app.py delete mode 100644 api/core/plugin/backwards_invocation/tool.py delete mode 100644 api/core/rag/retrieval/dataset_retrieval.py delete mode 100644 api/core/repositories/human_input_repository.py delete mode 100644 api/core/tools/custom_tool/provider.py delete mode 100644 api/core/tools/tool_engine.py delete mode 100644 api/core/tools/tool_label_manager.py delete mode 100644 api/core/tools/tool_manager.py delete mode 100644 api/core/tools/utils/dataset_retriever/dataset_multi_retriever_tool.py delete mode 100644 api/core/tools/utils/dataset_retriever/dataset_retriever_base_tool.py delete mode 100644 api/core/tools/utils/dataset_retriever/dataset_retriever_tool.py delete mode 100644 api/core/tools/utils/dataset_retriever_tool.py delete mode 100644 api/core/tools/workflow_as_tool/provider.py delete mode 100644 api/core/tools/workflow_as_tool/tool.py delete mode 100644 api/events/event_handlers/sync_plugin_trigger_when_app_created.py delete mode 100644 api/events/event_handlers/sync_workflow_schedule_when_app_published.py delete mode 100644 api/events/event_handlers/update_app_dataset_join_when_app_published_workflow_updated.py delete mode 100644 api/events/event_handlers/update_app_triggers_when_app_published_workflow_updated.py delete mode 100644 api/repositories/workflow_app_log_query_repository.py delete mode 100644 api/repositories/workflow_collaboration_repository.py delete mode 100644 api/services/agent/agent_soul_state.py delete mode 100644 api/services/agent/legacy_retirement_service.py delete mode 100644 api/services/agent/legacy_workflow_publish_service.py delete mode 100644 api/services/agent_tool_inner_service.py delete mode 100644 api/services/app_creation_records.py delete mode 100644 api/services/conversation_variable_updater.py delete mode 100644 api/services/knowledge/retrieval/attachments.py delete mode 100644 api/services/tools/legacy_api_tools_manage_service.py delete mode 100644 api/services/tools/legacy_builtin_tools_manage_service.py delete mode 100644 api/services/tools/legacy_mcp_tools_manage_service.py delete mode 100644 api/services/tools/legacy_tools_manage_service.py delete mode 100644 api/services/tools/legacy_tools_transform_service.py delete mode 100644 api/services/tools/legacy_workflow_tools_manage_service.py delete mode 100644 api/services/workflow_app_log_query_service.py delete mode 100644 api/services/workflow_draft_variable_service.py delete mode 100644 api/services/workflow_version_number_service.py diff --git a/api/.importlinter b/api/.importlinter index 80e5240ed65288..27af83471d72f6 100644 --- a/api/.importlinter +++ b/api/.importlinter @@ -1,13 +1,15 @@ [importlinter] root_packages = - services.entities.external_knowledge_entities - enums core + enums core.rag.extractor.entity + core.tools.__base + services.entities.external_knowledge_entities constants context configs controllers + controllers.console.explore extensions factories fields @@ -45,17 +47,11 @@ ignore_imports = core.rag.index_processor.processor.paragraph_index_processor -> services.knowledge.summaries.adapters core.rag.index_processor.processor.parent_child_index_processor -> services.knowledge.summaries.adapters core.rag.index_processor.processor.qa_index_processor -> services.knowledge.summaries.adapters - core.rag.retrieval.dataset_retrieval -> services.knowledge.external.service - core.rag.retrieval.dataset_retrieval -> services.knowledge.retrieval.attachments - core.rag.retrieval.dataset_retrieval -> services.feature_service core.rag.summary_index.summary_index -> services.knowledge.summaries.adapters core.rag.summary_index.summary_index -> tasks.generate_summary_index_task core.repositories.sqlalchemy_workflow_node_execution_repository -> services.file_service core.repositories.sqlalchemy_workflow_node_execution_repository -> services.variable_truncator - core.tools.tool_manager -> services.tools.legacy_builtin_tools_manage_service - core.tools.tool_manager -> services.tools.legacy_mcp_tools_manage_service - core.tools.tool_manager -> services.tools.legacy_tools_transform_service - # Node construction delegates datasource dependency assembly to the composition root. + # Node construction receives configured datasource and Agent dependencies from composition. core.workflow.nodes.agent_v2.runtime_feature_manifest -> services.agent.knowledge_datasets core.workflow.nodes.agent_v2.session_store -> services.agent.workspace_service core.workflow.nodes.agent_v2.workspace_retirement_layer -> tasks.collect_agent_resources_task @@ -1046,6 +1042,186 @@ forbidden_modules = sqlalchemy werkzeug +[importlinter:contract:console-workflow-application-boundary] +name = Console workflow use cases and contracts are framework and persistence neutral +type = forbidden +source_modules = + services.workflow.console_service + services.workflow.publication_service + services.rag_pipeline.publication_service + services.rag_pipeline.publication_contracts + services.workflow.draft_service + services.workflow.contracts + services.trigger.workflow_policy + services.trigger.schedule_policy +forbidden_modules = + controllers + extensions + fields + flask + models + repositories + sqlalchemy + werkzeug + services.workflow.runtime_gateway + services.human_input.file_gateway + services.workflow.access_gateway + services.workflow.definition_gateway + services.workflow_service + +[importlinter:contract:console-workflow-controller-boundary] +name = Console workflow controllers reach persistence and runtime implementations through composition +type = forbidden +source_modules = + controllers.console.app.workflow + controllers.console.app.workflow_admission +forbidden_modules = + repositories + sqlalchemy + extensions.ext_database + extensions.ext_redis + services.workflow_service + services.app_generate_service + services.workflow.runtime_gateway + services.human_input.file_gateway + services.workflow.access_gateway + services.workflow.definition_gateway + services.workflow.draft_service +allow_indirect_imports = True + +[importlinter:contract:workflow-definition-persistence-boundary] +name = Workflow definition persistence does not delegate back to workflow orchestration +type = forbidden +source_modules = + repositories.workflow.definition_repository + repositories.workflow.draft_repository + repositories.workflow.pipeline_publication_repository + repositories.workflow.debug_reservation_repository +forbidden_modules = + services.workflow_service + services.workflow.console_service + services.workflow.definition_gateway + services.rag_pipeline.publication_service + services.rag_pipeline.publication_gateway + services.rag_pipeline.rag_pipeline + services.knowledge.dataset_service +allow_indirect_imports = True + +[importlinter:contract:workflow-agent-persistence-boundary] +name = Definition persistence does not coordinate Agent services +type = forbidden +source_modules = + repositories.workflow.definition_repository + repositories.workflow.snippet_publication_repository + repositories.workflow.draft_repository + repositories.workflow.pipeline_publication_repository + repositories.workflow.debug_reservation_repository +forbidden_modules = + services.agent +allow_indirect_imports = True + +[importlinter:contract:workflow-execution-dependency-boundary] +name = Workflow and DSL services receive dependencies without calling composition roots +type = forbidden +source_modules = + services.app_dsl_service + services.snippet_dsl_service + services.data_migration.import_service + services.workflow_service + services.app_generate_service + services.snippet_generate_service + services.rag_pipeline.rag_pipeline + services.rag_pipeline.execution_service +forbidden_modules = + extensions.application_services + extensions.ext_application_services +allow_indirect_imports = True + +[importlinter:contract:human-input-debug-application-boundary] +name = Human input debug orchestration is independent of transport and persistence implementations +type = forbidden +source_modules = + services.human_input.debug_service + services.human_input.contracts +forbidden_modules = + controllers + extensions + flask + sqlalchemy + models.human_input + models.workflow + models.model + models.base + models.engine + repositories + services.human_input.file_gateway + services.human_input_service + services.human_input_delivery_test_service +allow_indirect_imports = True + + +[importlinter:contract:workflow-variable-persistence-boundary] +name = Variable persistence does not own variable policy or file uploads +type = forbidden +source_modules = + repositories.workflow.draft_variable_repository +forbidden_modules = + services.workflow.draft_variable_service + services.workflow.variable_service + services.file_service + services.variable_truncator + extensions.ext_storage +allow_indirect_imports = True + +[importlinter:contract:workflow-variable-core-boundary] +name = App execution receives variable loader and saver ports from its caller +type = forbidden +source_modules = + services.app.generation.input_adapter + services.workflow.execution.adapters.chatflow.app_generator + services.workflow.execution.adapters.chatflow.generate_task_pipeline + services.workflow.execution.adapters.response_converter + services.app.generation.adapters.message_cycle + services.workflow.execution.adapters.workflow.app_generator + services.workflow.execution.adapters.workflow.generate_task_pipeline + services.workflow.execution.adapters.pipeline.pipeline_generator + services.tools.workflow.tool + services.tools.workflow.provider +forbidden_modules = + repositories.workflow.draft_variable_repository + services.workflow.draft_variable_service + services.workflow.variable_service + extensions.application_services.workflow_variables +allow_indirect_imports = True + +[importlinter:contract:workflow-variable-policy-boundary] +name = Workflow variable merge policy is independent of persistence and transport +type = forbidden +source_modules = + services.workflow.variable_policy +forbidden_modules = + repositories + models + sqlalchemy + flask + controllers + extensions +allow_indirect_imports = True + +[importlinter:contract:workflow-variable-application-boundary] +name = Workflow variable services consume injected persistence and storage ports +type = forbidden +source_modules = + services.workflow.variable_service + services.workflow.draft_variable_service +forbidden_modules = + repositories + sqlalchemy + flask + controllers + extensions +allow_indirect_imports = True + [importlinter:contract:human-input-contract-boundary] name = Shared Human Input contracts are independent of application and persistence layers type = forbidden @@ -1065,23 +1241,6 @@ forbidden_modules = sqlalchemy werkzeug -[importlinter:contract:retrieval-persistence-boundary] -name = Retrieval use cases consume explicit persistence ports -type = forbidden -source_modules = - services.knowledge.retrieval.dataset_retrieval - core.rag.rerank.rerank_model - services.knowledge.retrieval.reranking -forbidden_modules = - sqlalchemy - repositories - extensions.ext_database - core.db - flask - extensions.application_services - extensions.ext_application_services -allow_indirect_imports = True - [importlinter:contract:workflow-agent-application-boundary] name = Workflow Agent binding rules consume ports without ORM or infrastructure adapters type = forbidden @@ -1110,47 +1269,134 @@ forbidden_modules = services.agent allow_indirect_imports = True -[importlinter:contract:workflow-definition-persistence-boundary] -name = Workflow definition persistence does not delegate back to workflow orchestration +[importlinter:contract:workflow-worker-lifetime-boundary] +name = Workflow worker lifetime is independent of transport and persistence type = forbidden source_modules = - repositories.workflow.definition_repository - repositories.workflow.draft_repository - repositories.workflow.pipeline_publication_repository + services.workflow.execution.generation_service forbidden_modules = - services.workflow_service - services.workflow.console_service - services.rag_pipeline.publication_service - services.rag_pipeline.publication_gateway - services.rag_pipeline.rag_pipeline - services.knowledge.dataset_service + flask + werkzeug + sqlalchemy + extensions + repositories + models + services.workflow.execution.adapters allow_indirect_imports = True -[importlinter:contract:workflow-variable-persistence-boundary] -name = Variable persistence does not own variable policy or file uploads +[importlinter:contract:workflow-invocation-read-boundary] +name = Workflow, Chatflow, Pipeline and Agent invocation adapters receive persistence through ports type = forbidden source_modules = - repositories.workflow.draft_variable_repository + services.workflow.execution.adapters.workflow.app_generator + services.workflow.execution.adapters.workflow.generate_task_pipeline + services.workflow.execution.adapters.chatflow.app_generator + services.workflow.execution.adapters.chatflow.generate_task_pipeline + services.workflow.execution.adapters.response_converter + services.app.generation.adapters.message_cycle + services.workflow.execution.adapters.graph + services.workflow.execution.adapters.pipeline.pipeline_generator + services.workflow.execution.adapters.pipeline.pipeline_runner + services.workflow.execution.adapters.chatflow.app_runner + services.app.generation.message_records + services.app.generation.adapters.message_pipeline + services.app.generation.adapters.base_runner + services.app.generation.adapters.agent_runner + services.app.generation.adapters.agent_app + services.app.generation.agent_config + services.app.generation.ports + services.app.generation.runtime + services.workflow.execution.ports + services.workflow.execution.chatflow_ports + services.workflow.execution.adapters.chatflow.conversation_variables + services.tools.workflow.contracts + services.annotation.reply_service forbidden_modules = - services.workflow.draft_variable_service - services.workflow.variable_service - services.file_service - services.variable_truncator - extensions.ext_storage + sqlalchemy + extensions.ext_database + core.db.session_factory + repositories allow_indirect_imports = True -[importlinter:contract:workflow-variable-policy-boundary] -name = Workflow variable merge policy is independent of persistence and transport +[importlinter:contract:message-persistence-contract-boundary] +name = Message persistence ports do not depend on execution services or runtime wiring type = forbidden source_modules = - services.workflow.variable_policy + services.app.generation.ports +forbidden_modules = + services.workflow + services.agent + services.app.generation.runtime + extensions + flask + werkzeug +allow_indirect_imports = True + +[importlinter:contract:refactored-route-error-boundary] +name = Refactored routes map domain failures through application HTTP errors +type = forbidden +source_modules = + controllers.trigger.webhook + controllers.trigger.trigger + controllers.console.explore.message + controllers.console.explore.trial + controllers.console.datasets.hit_testing_base + controllers.openapi.app_run +forbidden_modules = + werkzeug +allow_indirect_imports = True + +[importlinter:contract:agent-tool-controller-boundary] +name = Agent tool controller does not own persistence +type = forbidden +source_modules = + controllers.inner_api.agent.tools forbidden_modules = + sqlalchemy repositories - models + extensions.ext_database + core.db + controllers.console.app.wraps +allow_indirect_imports = True + +[importlinter:contract:agent-tool-application-boundary] +name = Agent tool invocation coordinates explicit ports +type = forbidden +source_modules = + services.agent.tool_invocation_service +forbidden_modules = sqlalchemy flask - controllers + werkzeug extensions + repositories + core + services.tools.agent_invocation_gateway +allow_indirect_imports = True + +[importlinter:contract:agent-tool-gateway-boundary] +name = Agent tool execution does not wrap external calls in a database session +type = forbidden +source_modules = + services.tools.agent_invocation_gateway +forbidden_modules = + sqlalchemy + extensions.ext_database + core.db +allow_indirect_imports = True + +[importlinter:contract:workflow-conversion-boundary] +name = Workflow conversion prepares values and persists through an explicit port +type = forbidden +source_modules = + services.workflow.conversion_service + services.workflow.workflow_converter +forbidden_modules = + sqlalchemy + extensions.ext_database + core.db + repositories + flask allow_indirect_imports = True [importlinter:contract:api-tool-provider-boundary] @@ -1170,6 +1416,23 @@ forbidden_modules = extensions.ext_application_services allow_indirect_imports = True +[importlinter:contract:workflow-node-repository-boundary] +name = Workflow node factories and runtimes receive repositories through ports +type = forbidden +source_modules = + services.workflow.execution.adapters.node_factory + services.workflow.execution.adapters.node_runtime + services.workflow.execution.adapters.agent_runtime + services.workflow.execution.adapters.agent_v2.dify_tools_builder +forbidden_modules = + repositories + sqlalchemy + extensions.ext_database + core.db + extensions.application_services + extensions.ext_application_services +allow_indirect_imports = True + [importlinter:contract:builtin-tool-credential-boundary] name = Builtin tool credential resolution owns no database or composition dependencies type = forbidden @@ -1185,25 +1448,27 @@ forbidden_modules = extensions.ext_application_services allow_indirect_imports = True -[importlinter:contract:workflow-variable-application-boundary] -name = Workflow variable services consume injected persistence and storage ports +[importlinter:contract:agent-log-controller-boundary] +name = Agent log controller delegates through admission and an application service type = forbidden source_modules = - services.workflow.variable_service - services.workflow.draft_variable_service + controllers.console.app.agent forbidden_modules = - repositories sqlalchemy - flask - controllers - extensions + extensions.ext_database + repositories + models + controllers.common.session + controllers.console.app.wraps + services.agent_service allow_indirect_imports = True -[importlinter:contract:console-draft-variable-use-case-boundary] -name = Console variable use cases receive persistence and file resolution through ports +[importlinter:contract:agent-log-service-boundary] +name = Agent log use cases and contracts own no transport or database resources type = forbidden source_modules = - services.workflow.console_variable_service + services.agent.log_service + services.agent.log_contracts forbidden_modules = sqlalchemy extensions.ext_database @@ -1212,23 +1477,79 @@ forbidden_modules = werkzeug controllers repositories + libs.login + services.tools.tool_manager extensions.application_services extensions.ext_application_services - factories.file_factory allow_indirect_imports = True -[importlinter:contract:message-persistence-contract-boundary] -name = Message persistence ports do not depend on execution services or runtime wiring +[importlinter:contract:domain-errors-without-http] +name = Annotation, document and trigger domain errors have no HTTP dependencies type = forbidden source_modules = - services.app.generation.ports + services.annotation.errors + services.knowledge.dataset_service + services.trigger.errors + services.errors.dataset + services.errors.document + tasks.annotation.batch_import_annotations_task forbidden_modules = - services.workflow - services.agent - services.app.generation.runtime - extensions + werkzeug + libs.exception + controllers +allow_indirect_imports = True + +[importlinter:contract:workflow-controller-error-boundary] +name = Workflow controllers translate failures with project error contracts +type = forbidden +source_modules = + controllers.console.app.workflow + controllers.console.app.workflow_admission + controllers.console.snippets.snippet_workflow + controllers.console.datasets.rag_pipeline.rag_pipeline_workflow + controllers.service_api.app.workflow + controllers.web.workflow +forbidden_modules = + werkzeug +allow_indirect_imports = True + +[importlinter:contract:console-draft-variable-controller-boundary] +name = App, Pipeline and Snippet variable controllers delegate without persistence or runtime construction +type = forbidden +source_modules = + controllers.console.app.workflow_draft_variable + controllers.console.datasets.rag_pipeline.rag_pipeline_draft_variable + controllers.console.snippets.snippet_workflow_draft_variable +forbidden_modules = + sqlalchemy + extensions.ext_database + core + repositories + models + factories + controllers.common.session + services.workflow_service + services.rag_pipeline.rag_pipeline + services.snippet_service + werkzeug +allow_indirect_imports = True + +[importlinter:contract:console-draft-variable-use-case-boundary] +name = Console variable use cases receive persistence and file resolution through ports +type = forbidden +source_modules = + services.workflow.console_variable_service +forbidden_modules = + sqlalchemy + extensions.ext_database + core.db flask werkzeug + controllers + repositories + extensions.application_services + extensions.ext_application_services + factories.file_factory allow_indirect_imports = True [importlinter:contract:agent-chat-runner-persistence-boundary] @@ -1262,40 +1583,26 @@ forbidden_modules = extensions.ext_application_services allow_indirect_imports = True -[importlinter:contract:workflow-node-repository-boundary] -name = Workflow node factories and runtimes receive repositories through ports +[importlinter:contract:retrieval-persistence-boundary] +name = Retrieval use cases and tools consume explicit persistence ports type = forbidden source_modules = - services.workflow.execution.adapters.node_factory - services.workflow.execution.adapters.node_runtime - services.workflow.execution.adapters.agent_runtime - services.workflow.execution.adapters.agent_v2.dify_tools_builder + services.knowledge.retrieval.dataset_retrieval + core.rag.rerank.rerank_model + services.knowledge.retrieval.reranking + services.knowledge_retrieval_inner_service + services.tools.dataset.tool + services.app.generation.adapters.chat_runner + services.app.generation.adapters.completion_runner + services.workflow.execution.adapters.knowledge_retrieval + services.workflow.execution.adapters.human_input_events + core.tools.__base.tool forbidden_modules = - repositories sqlalchemy + repositories extensions.ext_database core.db + flask extensions.application_services extensions.ext_application_services allow_indirect_imports = True - -[importlinter:contract:workflow-variable-core-boundary] -name = App execution receives variable loader and saver ports from its caller -type = forbidden -source_modules = - services.app.generation.input_adapter - services.workflow.execution.adapters.chatflow.app_generator - services.workflow.execution.adapters.chatflow.generate_task_pipeline - services.workflow.execution.adapters.response_converter - services.app.generation.adapters.message_cycle - services.workflow.execution.adapters.workflow.app_generator - services.workflow.execution.adapters.workflow.generate_task_pipeline - services.workflow.execution.adapters.pipeline.pipeline_generator - services.tools.workflow.tool - services.tools.workflow.provider -forbidden_modules = - repositories.workflow.draft_variable_repository - services.workflow.draft_variable_service - services.workflow.variable_service - extensions.application_services.workflow_variables -allow_indirect_imports = True diff --git a/api/core/callback_handler/index_tool_callback_handler.py b/api/core/callback_handler/index_tool_callback_handler.py deleted file mode 100644 index d2024454a680ae..00000000000000 --- a/api/core/callback_handler/index_tool_callback_handler.py +++ /dev/null @@ -1,102 +0,0 @@ -import logging -from collections.abc import Sequence - -from sqlalchemy import select, update -from sqlalchemy.orm import Session, sessionmaker - -from core.app.apps.base_app_queue_manager import AppQueueManager, PublishFrom -from core.app.entities.app_invoke_entities import InvokeFrom -from core.app.entities.queue_entities import QueueRetrieverResourcesEvent -from core.rag.entities import RetrievalSourceMetadata -from core.rag.index_processor.constant.index_type import IndexStructureType -from core.rag.models.document import Document -from extensions.ext_database import db -from models.dataset import ChildChunk, DatasetQuery, DocumentSegment -from models.dataset import Document as DatasetDocument -from models.enums import CreatorUserRole, DatasetQuerySource - -_logger = logging.getLogger(__name__) - - -class DatasetIndexToolCallbackHandler: - """Callback handler for dataset tool.""" - - def __init__( - self, queue_manager: AppQueueManager, app_id: str, message_id: str, user_id: str, invoke_from: InvokeFrom - ): - self._queue_manager = queue_manager - self._app_id = app_id - self._message_id = message_id - self._user_id = user_id - self._invoke_from = invoke_from - - def on_query(self, query: str, dataset_id: str, session: Session): - """ - Handle query. - """ - dataset_query = DatasetQuery( - dataset_id=dataset_id, - content=query, - source=DatasetQuerySource.APP, - source_app_id=self._app_id, - created_by_role=( - CreatorUserRole.ACCOUNT - if self._invoke_from in {InvokeFrom.EXPLORE, InvokeFrom.DEBUGGER} - else CreatorUserRole.END_USER - ), - created_by=self._user_id, - ) - - # Use an independent session so this audit-log side effect does - # not commit or close the caller's request-scoped session. - with sessionmaker(bind=db.engine, expire_on_commit=False).begin() as independent_session: - independent_session.add(dataset_query) - - def on_tool_end(self, documents: list[Document], session: Session): - """Handle tool end.""" - # Use an independent session so hit-count updates do not - # interfere with the caller's request-scoped session. - with Session(db.engine, expire_on_commit=False) as independent_session: - for document in documents: - if document.metadata is not None: - document_id = document.metadata["document_id"] - dataset_document_stmt = select(DatasetDocument).where(DatasetDocument.id == document_id) - dataset_document = independent_session.scalar(dataset_document_stmt) - if not dataset_document: - _logger.warning( - "Expected DatasetDocument record to exist, but none was found, document_id=%s", - document_id, - ) - continue - if dataset_document.doc_form == IndexStructureType.PARENT_CHILD_INDEX: - child_chunk_stmt = select(ChildChunk).where( - ChildChunk.index_node_id == document.metadata["doc_id"], - ChildChunk.dataset_id == dataset_document.dataset_id, - ChildChunk.document_id == dataset_document.id, - ) - child_chunk = independent_session.scalar(child_chunk_stmt) - if child_chunk: - independent_session.execute( - update(DocumentSegment) - .where(DocumentSegment.id == child_chunk.segment_id) - .values(hit_count=DocumentSegment.hit_count + 1) - ) - else: - conditions = [DocumentSegment.index_node_id == document.metadata["doc_id"]] - - if "dataset_id" in document.metadata: - conditions.append(DocumentSegment.dataset_id == document.metadata["dataset_id"]) - - # add hit count to document segment - independent_session.execute( - update(DocumentSegment).where(*conditions).values(hit_count=DocumentSegment.hit_count + 1) - ) - - independent_session.commit() - - # TODO(-LAN-): Improve type check - def return_retriever_resource_info(self, resource: Sequence[RetrievalSourceMetadata]): - """Handle return_retriever_resource_info.""" - self._queue_manager.publish( - QueueRetrieverResourcesEvent(retriever_resources=resource), PublishFrom.APPLICATION_MANAGER - ) diff --git a/api/core/plugin/backwards_invocation/app.py b/api/core/plugin/backwards_invocation/app.py deleted file mode 100644 index 66a1c547eeae71..00000000000000 --- a/api/core/plugin/backwards_invocation/app.py +++ /dev/null @@ -1,329 +0,0 @@ -import uuid -from collections.abc import Generator, Mapping -from typing import Any, Protocol, cast - -from sqlalchemy import select -from sqlalchemy.orm import Session - -from core.app.app_config.common.parameters_mapping import get_parameters_from_feature_dict -from core.app.apps.advanced_chat.app_generator import AdvancedChatAppGenerator -from core.app.apps.agent_chat.app_generator import AgentChatAppGenerator -from core.app.apps.chat.app_generator import ChatAppGenerator -from core.app.apps.completion.app_generator import CompletionAppGenerator -from core.app.apps.workflow.app_generator import WorkflowAppGenerator -from core.app.entities.app_invoke_entities import InvokeFrom -from core.app.layers.pause_state_persist_layer import PauseStateLayerConfig -from core.db.session_factory import create_session -from core.plugin.backwards_invocation.base import BaseBackwardsInvocation -from extensions.ext_database import db -from models import Account, TenantAccountJoin -from models.model import ( - App, - AppMode, - AppModelConfig, - AppModelConfigDict, - EndUser, - load_annotation_reply_config, -) -from models.workflow import Workflow - - -class AppScopedEndUserProvisioner(Protocol): - def get_or_create_end_user( - self, - tenant_id: str, - app_id: str, - user_id: str | None = None, - ) -> EndUser: ... - - -class PluginAppBackwardsInvocation(BaseBackwardsInvocation): - @classmethod - def fetch_app_info(cls, app_id: str, tenant_id: str) -> Mapping: - """ - Fetch app info - """ - app = cls._get_app(app_id, tenant_id) - - """Retrieve app parameters.""" - if app.mode in {AppMode.ADVANCED_CHAT, AppMode.WORKFLOW}: - workflow = cls._get_workflow(app) - if workflow is None: - raise ValueError("unexpected app type") - - features_dict: dict[str, Any] = workflow.features_dict - user_input_form = workflow.user_input_form(to_old_structure=True) - else: - app_model_config_dict = cls._get_app_model_config_dict(app) - if app_model_config_dict is None: - raise ValueError("unexpected app type") - - features_dict = cast(dict[str, Any], app_model_config_dict) - - user_input_form = features_dict.get("user_input_form", []) - - return { - "data": get_parameters_from_feature_dict(features_dict=features_dict, user_input_form=user_input_form), - } - - @classmethod - def invoke_app( - cls, - app_id: str, - user_id: str, - tenant_id: str, - conversation_id: str | None, - query: str | None, - stream: bool, - inputs: Mapping, - files: list[dict], - session: Session, - end_users: AppScopedEndUserProvisioner, - ) -> Generator[Mapping | str, None, None] | Mapping: - """ - invoke app - """ - app = cls._get_app(app_id, tenant_id) - if not user_id: - user = end_users.get_or_create_end_user(app.tenant_id, app.id, None) - else: - try: - user = cls._get_user(user_id, app) - except ValueError: - # Plugins such as WeCom Bot pass external sender IDs rather than EndUser UUIDs. - user = end_users.get_or_create_end_user(app.tenant_id, app.id, user_id) - - conversation_id = conversation_id or "" - - match app.mode: - case AppMode.ADVANCED_CHAT | AppMode.AGENT_CHAT | AppMode.CHAT: - if not query: - raise ValueError("missing query") - - return cls.invoke_chat_app(app, user, conversation_id, query, stream, inputs, files, session) - case AppMode.WORKFLOW: - workflow = cls._get_workflow(app) - if not workflow: - raise ValueError("unexpected app type") - return cls.invoke_workflow_app(app, workflow, user, stream, inputs, files) - case AppMode.COMPLETION: - return cls.invoke_completion_app(app, user, stream, inputs, files, session) - case _: - raise ValueError("unexpected app type") - - @classmethod - def invoke_chat_app( - cls, - app: App, - user: Account | EndUser, - conversation_id: str, - query: str, - stream: bool, - inputs: Mapping, - files: list[dict], - session: Session, - ) -> Generator[Mapping | str, None, None] | Mapping: - """ - invoke chat app - """ - match app.mode: - case AppMode.ADVANCED_CHAT: - workflow = cls._get_workflow(app) - if not workflow: - raise ValueError("unexpected app type") - - pause_config = PauseStateLayerConfig( - session_factory=db.engine, - state_owner_user_id=workflow.created_by, - ) - - return AdvancedChatAppGenerator().generate( - app_model=app, - workflow=workflow, - user=user, - args={ - "inputs": inputs, - "query": query, - "files": files, - "conversation_id": conversation_id, - }, - invoke_from=InvokeFrom.SERVICE_API, - workflow_run_id=str(uuid.uuid4()), - streaming=stream, - pause_state_config=pause_config, - session=session, - ) - case AppMode.AGENT_CHAT: - return AgentChatAppGenerator().generate( - app_model=app, - user=user, - args={ - "inputs": inputs, - "query": query, - "files": files, - "conversation_id": conversation_id, - }, - invoke_from=InvokeFrom.SERVICE_API, - streaming=stream, - session=session, - ) - case AppMode.CHAT: - return ChatAppGenerator().generate( - app_model=app, - user=user, - args={ - "inputs": inputs, - "query": query, - "files": files, - "conversation_id": conversation_id, - }, - invoke_from=InvokeFrom.SERVICE_API, - streaming=stream, - session=session, - ) - case _: - raise ValueError("unexpected app type") - - @classmethod - def invoke_workflow_app( - cls, - app: App, - workflow: Workflow, - user: EndUser | Account, - stream: bool, - inputs: Mapping, - files: list[dict], - ) -> Generator[Mapping | str, None, None] | Mapping: - """ - invoke workflow app - """ - pause_config = PauseStateLayerConfig( - session_factory=db.engine, - state_owner_user_id=workflow.created_by, - ) - - return WorkflowAppGenerator().generate( - app_model=app, - workflow=workflow, - user=user, - args={"inputs": inputs, "files": files}, - invoke_from=InvokeFrom.SERVICE_API, - streaming=stream, - call_depth=1, - pause_state_config=pause_config, - ) - - @classmethod - def invoke_completion_app( - cls, - app: App, - user: EndUser | Account, - stream: bool, - inputs: Mapping, - files: list[dict], - session: Session, - ) -> Generator[Mapping | str, None, None] | Mapping: - """ - invoke completion app - """ - return CompletionAppGenerator().generate( - app_model=app, - user=user, - args={"inputs": inputs, "files": files}, - invoke_from=InvokeFrom.SERVICE_API, - streaming=stream, - session=session, - ) - - @classmethod - def _get_user(cls, user_id: str, app: App) -> EndUser | Account: - """ - get the user by user id - """ - with create_session() as session: - stmt = select(EndUser).where( - EndUser.id == user_id, - EndUser.tenant_id == app.tenant_id, - EndUser.app_id == app.id, - ) - user = session.scalar(stmt) - if not user: - stmt = select(EndUser).where( - EndUser.session_id == user_id, - EndUser.tenant_id == app.tenant_id, - EndUser.app_id == app.id, - ) - user = session.scalar(stmt) - if not user: - stmt = select(Account).where( - Account.id == user_id, - Account.id == TenantAccountJoin.account_id, - TenantAccountJoin.tenant_id == app.tenant_id, - ) - user = session.scalar(stmt) - if user: - session.expunge(user) - - if not user: - raise ValueError("user not found") - - return user - - @classmethod - def _get_app(cls, app_id: str, tenant_id: str) -> App: - """ - get app - """ - try: - with create_session() as session: - app = session.scalar(select(App).where(App.id == app_id, App.tenant_id == tenant_id).limit(1)) - if app: - session.expunge(app) - except Exception as e: - raise ValueError("app not found") from e - - if not app: - raise ValueError("app not found") - - return app - - @classmethod - def _get_workflow(cls, app: App) -> Workflow | None: - """ - get workflow with an owned session, additionally scoped by tenant and app - - `App.workflow_with_session` resolves by `workflow_id` alone; a plugin - invocation must also confirm the workflow belongs to this tenant and app. - """ - if not app.workflow_id: - return None - - with create_session() as session: - workflow = session.scalar( - select(Workflow) - .where(Workflow.id == app.workflow_id, Workflow.tenant_id == app.tenant_id, Workflow.app_id == app.id) - .limit(1) - ) - if workflow: - session.expunge(workflow) - return workflow - - @classmethod - def _get_app_model_config_dict(cls, app: App) -> AppModelConfigDict | None: - """ - get app model config features without relying on request-scoped session-backed model properties - """ - if not app.app_model_config_id: - return None - - with create_session() as session: - app_model_config = session.scalar( - select(AppModelConfig) - .where(AppModelConfig.id == app.app_model_config_id, AppModelConfig.app_id == app.id) - .limit(1) - ) - if app_model_config is None: - return None - - annotation_reply = load_annotation_reply_config(session, app_model_config.app_id) - return app_model_config.to_dict(annotation_reply=annotation_reply) diff --git a/api/core/plugin/backwards_invocation/tool.py b/api/core/plugin/backwards_invocation/tool.py deleted file mode 100644 index 0a296dfea10bf0..00000000000000 --- a/api/core/plugin/backwards_invocation/tool.py +++ /dev/null @@ -1,55 +0,0 @@ -from collections.abc import Generator -from typing import Any - -from sqlalchemy.orm import Session - -from core.callback_handler.workflow_tool_callback_handler import DifyWorkflowCallbackHandler -from core.plugin.backwards_invocation.base import BaseBackwardsInvocation -from core.tools.entities.tool_entities import ToolInvokeMessage, ToolProviderType -from core.tools.tool_engine import ToolEngine -from core.tools.tool_manager import ToolManager -from core.tools.utils.message_transformer import ToolFileMessageTransformer - - -class PluginToolBackwardsInvocation(BaseBackwardsInvocation): - """ - Backwards invocation for plugin tools. - """ - - @classmethod - def invoke_tool( - cls, - session: Session, - tenant_id: str, - user_id: str, - tool_type: ToolProviderType, - provider: str, - tool_name: str, - tool_parameters: dict[str, Any], - credential_id: str | None = None, - ) -> Generator[ToolInvokeMessage, None, None]: - """ - invoke tool - """ - # get tool runtime - try: - tool_runtime = ToolManager.get_tool_runtime_from_plugin( - tool_type, - tenant_id, - provider, - tool_name, - tool_parameters, - user_id=user_id, - credential_id=credential_id, - ) - response = ToolEngine.generic_invoke( - session, tool_runtime, tool_parameters, user_id, DifyWorkflowCallbackHandler(), workflow_call_depth=1 - ) - - response = ToolFileMessageTransformer.transform_tool_invoke_messages( - response, user_id=user_id, tenant_id=tenant_id - ) - - return response - except Exception as e: - raise e diff --git a/api/core/rag/retrieval/dataset_retrieval.py b/api/core/rag/retrieval/dataset_retrieval.py deleted file mode 100644 index 97595c35403aae..00000000000000 --- a/api/core/rag/retrieval/dataset_retrieval.py +++ /dev/null @@ -1,2134 +0,0 @@ -import json -import logging -import math -import re -import threading -import time -from collections import Counter, defaultdict -from collections.abc import Generator, Mapping -from functools import partial -from typing import Any, Union, cast - -from flask import Flask, current_app -from opentelemetry.trace import get_current_span -from sqlalchemy import and_, func, literal, or_, select, update -from sqlalchemy.exc import DBAPIError -from sqlalchemy.orm import Session, sessionmaker -from tenacity import before_sleep_log, retry, retry_if_exception, stop_after_attempt, wait_exponential - -from core.app.app_config.entities import ( - DatasetEntity, - DatasetRetrieveConfigEntity, - MetadataFilteringCondition, - ModelConfig, -) -from core.app.entities.app_invoke_entities import ( - CreditUsageCreatedBy, - EasyUIBasedAppGenerateEntity, - InvokeFrom, - ModelConfigWithCredentialsEntity, - get_credit_usage_app_type, -) -from core.callback_handler.index_tool_callback_handler import DatasetIndexToolCallbackHandler -from core.db.session_factory import session_factory -from core.entities.agent_entities import PlanningStrategy -from core.entities.model_entities import ModelStatus -from core.memory.token_buffer_memory import TokenBufferMemory -from core.model_context import with_credit_usage_created_by, with_credit_usage_metadata -from core.model_manager import ModelInstance, ModelManager -from core.ops.entities.trace_entity import TraceTaskName -from core.ops.ops_trace_manager import TraceQueueManager, TraceTask -from core.ops.utils import measure_time -from core.prompt.advanced_prompt_transform import AdvancedPromptTransform -from core.prompt.entities.advanced_prompt_entities import ChatModelMessage, CompletionModelPromptTemplate -from core.prompt.simple_prompt_transform import ModelMode -from core.rag.data_post_processor.data_post_processor import DataPostProcessor, RerankingModelDict, WeightsDict -from core.rag.datasource.keyword.jieba.jieba_keyword_table_handler import JiebaKeywordTableHandler -from core.rag.datasource.retrieval_service import DefaultRetrievalModelDict, RetrievalService -from core.rag.entities import Condition, DocumentContext, RetrievalSourceMetadata -from core.rag.index_processor.constant.doc_type import DocType -from core.rag.index_processor.constant.index_type import IndexStructureType, IndexTechniqueType -from core.rag.index_processor.constant.query_type import QueryType -from core.rag.models.document import Document -from core.rag.rerank.rerank_type import RerankMode -from core.rag.retrieval.retrieval_methods import RetrievalMethod -from core.rag.retrieval.router.multi_dataset_function_call_router import FunctionCallMultiDatasetRouter -from core.rag.retrieval.router.multi_dataset_react_route import ReactMultiDatasetRouter -from core.rag.retrieval.template_prompts import ( - METADATA_FILTER_ASSISTANT_PROMPT_1, - METADATA_FILTER_ASSISTANT_PROMPT_2, - METADATA_FILTER_COMPLETION_PROMPT, - METADATA_FILTER_SYSTEM_PROMPT, - METADATA_FILTER_USER_PROMPT_1, - METADATA_FILTER_USER_PROMPT_2, - METADATA_FILTER_USER_PROMPT_3, -) -from core.tools.signature import sign_upload_file_preview_url -from core.tools.utils.dataset_retriever.dataset_retriever_base_tool import DatasetRetrieverBaseTool -from core.workflow.file_reference import build_file_reference -from core.workflow.nodes.knowledge_retrieval import exc -from core.workflow.nodes.knowledge_retrieval.retrieval import ( - KnowledgeRetrievalRequest, - Source, - SourceChildChunk, - SourceMetadata, -) -from extensions.ext_database import db -from extensions.ext_redis import redis_client -from extensions.otel import propagate_context, trace_span -from graphon.file import File, FileTransferMethod, FileType -from graphon.model_runtime.entities.llm_entities import LLMMode, LLMResult, LLMUsage -from graphon.model_runtime.entities.message_entities import PromptMessage, PromptMessageRole, PromptMessageTool -from graphon.model_runtime.entities.model_entities import ModelFeature, ModelType -from graphon.model_runtime.model_providers.base.large_language_model import LargeLanguageModel -from libs.helper import parse_uuid_str_or_none -from libs.json_in_md_parser import parse_and_check_json_markdown -from models.dataset import ( - ChildChunk, - Dataset, - DatasetMetadata, - DatasetQuery, - DocumentSegment, - RateLimitLog, - SegmentAttachmentBinding, -) -from models.dataset import Document as DatasetDocument -from models.dataset import Document as DocumentModel -from models.enums import CreatorUserRole, DatasetQuerySource -from repositories.knowledge.dataset_read_repository import get_dataset_available_document_count -from repositories.knowledge.segment_read_adapter import sign_segment_content -from repositories.knowledge.upload_file_repository import SQLAlchemyKnowledgeUploadRepository -from services.feature_service import FeatureService -from services.knowledge.external.service import ExternalDatasetService -from services.knowledge.retrieval.attachments import authorize_retrieved_segment - -default_retrieval_model: DefaultRetrievalModelDict = { - "search_method": RetrievalMethod.SEMANTIC_SEARCH, - "reranking_enable": False, - "reranking_model": {"reranking_provider_name": "", "reranking_model_name": ""}, - "top_k": 4, - "score_threshold_enabled": False, -} - -logger = logging.getLogger(__name__) - -_POSTGRES_DEADLOCK_SQLSTATE = "40P01" -_HIT_COUNT_UPDATE_MAX_ATTEMPTS = 3 - - -def _is_postgres_deadlock_error(exc: BaseException) -> bool: - if not isinstance(exc, DBAPIError) or exc.orig is None: - return False - orig = cast(Any, exc.orig) - return (hasattr(orig, "sqlstate") and orig.sqlstate == _POSTGRES_DEADLOCK_SQLSTATE) or ( - hasattr(orig, "pgcode") and orig.pgcode == _POSTGRES_DEADLOCK_SQLSTATE - ) - - -class DatasetRetrieval: - def __init__(self, application_generate_entity: EasyUIBasedAppGenerateEntity | None = None): - self.application_generate_entity = application_generate_entity - self._llm_usage = LLMUsage.empty_usage() - self._request_metadata: dict[str, object] | None = None - if application_generate_entity is not None: - app_config = application_generate_entity.app_config - self._request_metadata = { - "app_type": get_credit_usage_app_type(app_config.app_mode), - "app_id": app_config.app_id, - } - - def set_request_metadata(self, request_metadata: Mapping[str, object] | None) -> None: - self._request_metadata = dict(request_metadata) if request_metadata else None - - @property - def llm_usage(self) -> LLMUsage: - return self._llm_usage.model_copy() - - def _record_usage(self, usage: LLMUsage | None) -> None: - if usage is None or usage.total_tokens <= 0: - return - if self._llm_usage.total_tokens == 0: - self._llm_usage = usage - else: - self._llm_usage = self._llm_usage.plus(usage) - - @trace_span() - @with_credit_usage_metadata - @with_credit_usage_created_by(CreditUsageCreatedBy.KNOWLEDGE_RETRIEVAL) - def knowledge_retrieval(self, session: Session, request: KnowledgeRetrievalRequest) -> list[Source]: - self._check_knowledge_rate_limit(request.tenant_id) - available_datasets = self._get_available_datasets(request.tenant_id, request.dataset_ids) - available_datasets_ids = [i.id for i in available_datasets] - if not available_datasets_ids: - return [] - - if not request.query and not request.attachment_ids: - return [] - - metadata_filter_document_ids, metadata_condition = None, None - - if request.metadata_filtering_mode != "disabled": - app_metadata_model_config = ModelConfig(provider="", name="", mode=LLMMode.CHAT, completion_params={}) - if request.metadata_filtering_mode == "automatic": - if not request.metadata_model_config: - raise ValueError("metadata_model_config is required for this method") - - app_metadata_model_config = ModelConfig.model_validate(request.metadata_model_config.model_dump()) - - app_metadata_filtering_conditions = None - if request.metadata_filtering_conditions is not None: - app_metadata_filtering_conditions = MetadataFilteringCondition.model_validate( - request.metadata_filtering_conditions.model_dump() - ) - - query = request.query if request.query is not None else "" - - metadata_filter_document_ids, metadata_condition = self.get_metadata_filter_condition( - session=session, - dataset_ids=available_datasets_ids, - query=query, - tenant_id=request.tenant_id, - user_id=request.user_id, - metadata_filtering_mode=request.metadata_filtering_mode, - metadata_model_config=app_metadata_model_config, - metadata_filtering_conditions=app_metadata_filtering_conditions, - inputs={}, - ) - - if request.retrieval_mode == DatasetRetrieveConfigEntity.RetrieveStrategy.SINGLE: - planning_strategy = PlanningStrategy.REACT_ROUTER - # Ensure required fields are not None for single retrieval mode - if request.model_provider is None or request.model_name is None or request.query is None: - raise ValueError("model_provider, model_name, and query are required for single retrieval mode") - - model_manager = ModelManager.for_tenant(tenant_id=request.tenant_id, user_id=request.user_id) - model_instance = model_manager.get_model_instance( - tenant_id=request.tenant_id, - model_type=ModelType.LLM, - provider=request.model_provider, - model=request.model_name, - ) - - provider_model_bundle = model_instance.provider_model_bundle - model_type_instance = model_instance.model_type_instance - model_type_instance = cast(LargeLanguageModel, model_type_instance) - - model_credentials = model_instance.credentials - - # check model - provider_model = provider_model_bundle.configuration.get_provider_model( - model=request.model_name, model_type=ModelType.LLM - ) - - if provider_model is None: - raise exc.ModelNotExistError(f"Model {request.model_name} not exist.") - - if provider_model.status == ModelStatus.NO_CONFIGURE: - raise exc.ModelCredentialsNotInitializedError( - f"Model {request.model_name} credentials is not initialized." - ) - elif provider_model.status == ModelStatus.NO_PERMISSION: - raise exc.ModelNotSupportedError(f"Dify Hosted OpenAI {request.model_name} currently not support.") - elif provider_model.status == ModelStatus.QUOTA_EXCEEDED: - raise exc.ModelQuotaExceededError(f"Model provider {request.model_provider} quota exceeded.") - - stop = [] - completion_params = (request.completion_params or {}).copy() - if "stop" in completion_params: - stop = completion_params["stop"] - del completion_params["stop"] - - model_schema = model_type_instance.get_model_schema(request.model_name, model_credentials) - - if not model_schema: - raise exc.ModelNotExistError(f"Model {request.model_name} not exist.") - - model_config = ModelConfigWithCredentialsEntity( - provider=request.model_provider, - model=request.model_name, - model_schema=model_schema, - mode=request.model_mode or "chat", - provider_model_bundle=provider_model_bundle, - credentials=model_credentials, - parameters=completion_params, - stop=stop, - ) - all_documents = self.single_retrieve( - session, - request.app_id, - request.tenant_id, - request.user_id, - request.user_from, - request.query, - available_datasets, - model_instance, - model_config, - planning_strategy, - None, # message_id - metadata_filter_document_ids, - metadata_condition, - ) - else: - all_documents = self.multiple_retrieve( - app_id=request.app_id, - tenant_id=request.tenant_id, - user_id=request.user_id, - user_from=request.user_from, - available_datasets=available_datasets, - query=request.query, - top_k=request.top_k, - score_threshold=request.score_threshold, - reranking_mode=request.reranking_mode, - reranking_model=request.reranking_model, - weights=request.weights, - reranking_enable=request.reranking_enable, - metadata_filter_document_ids=metadata_filter_document_ids, - metadata_condition=metadata_condition, - attachment_ids=request.attachment_ids, - ) - - dify_documents = [item for item in all_documents if item.provider == "dify"] - external_documents = [item for item in all_documents if item.provider == "external"] - retrieval_resource_list = [] - # deal with external documents - for item in external_documents: - ext_meta = item.metadata or {} - title = ext_meta.get("title") or "" - doc_id = ext_meta.get("document_id") or title - source = Source( - metadata=SourceMetadata( - source="knowledge", - dataset_id=ext_meta.get("dataset_id") or "", - dataset_name=ext_meta.get("dataset_name") or "", - document_id=str(doc_id), - document_name=ext_meta.get("title") or "", - data_source_type="external", - retriever_from="workflow", - score=float(ext_meta.get("score") or 0.0), - doc_metadata=ext_meta, - ), - title=title, - content=item.page_content, - ) - retrieval_resource_list.append(source) - # deal with dify documents - if dify_documents: - # Materialize authorized, signed response values before closing the owned session. - with Session(bind=session.get_bind()) as retrieval_session: - records = RetrievalService.format_retrieval_documents(retrieval_session, dify_documents) - dataset_ids = [i.segment.dataset_id for i in records] - document_ids = [i.segment.document_id for i in records] - datasets = retrieval_session.scalars(select(Dataset).where(Dataset.id.in_(dataset_ids))).all() - documents = retrieval_session.scalars( - select(DatasetDocument).where(DatasetDocument.id.in_(document_ids)) - ).all() - - dataset_map = {i.id: i for i in datasets} - document_map = {i.id: i for i in documents} - - for record in records: - segment = record.segment - if ( - authorize_retrieved_segment( - segment, - tenant_id=request.tenant_id, - dataset_ids=available_datasets_ids, - session=retrieval_session, - ) - is None - ): - continue - dataset = dataset_map.get(segment.dataset_id) - document = document_map.get(segment.document_id) - - if dataset and document: - content = sign_segment_content(segment, session=retrieval_session) - source = Source( - metadata=SourceMetadata( - source="knowledge", - dataset_id=dataset.id, - dataset_name=dataset.name, - document_id=document.id, - document_name=document.name, - data_source_type=document.data_source_type, - segment_id=segment.id, - retriever_from="workflow", - score=record.score or 0.0, - segment_hit_count=segment.hit_count, - segment_word_count=segment.word_count, - segment_position=segment.position, - segment_index_node_hash=segment.index_node_hash, - doc_metadata=document.doc_metadata, - child_chunks=[ - SourceChildChunk( - id=str(getattr(chunk, "id", "")), - content=str(getattr(chunk, "content", "")), - position=int(getattr(chunk, "position", 0)), - score=float(getattr(chunk, "score", 0.0)), - ) - for chunk in (record.child_chunks or []) - ], - position=None, - ), - title=document.name, - files=list(record.files) if record.files else None, - content=content, - ) - if segment.answer: - source.content = f"question:{content} \nanswer:{segment.answer}" - - if record.summary: - source.summary = record.summary - - retrieval_resource_list.append(source) - - if retrieval_resource_list: - - def _score(item: Source) -> float: - meta = item.metadata - score = meta.score - if isinstance(score, (int, float)): - return float(score) - return 0.0 - - retrieval_resource_list = sorted( - retrieval_resource_list, - key=_score, # type: ignore[arg-type, return-value] - reverse=True, - ) - for position, item in enumerate(retrieval_resource_list, start=1): - item.metadata.position = position # type: ignore[index] - return retrieval_resource_list - - @with_credit_usage_metadata - @with_credit_usage_created_by(CreditUsageCreatedBy.KNOWLEDGE_RETRIEVAL) - def retrieve( - self, - session: Session, - app_id: str, - user_id: str, - tenant_id: str, - model_config: ModelConfigWithCredentialsEntity, - config: DatasetEntity, - query: str, - invoke_from: InvokeFrom, - show_retrieve_source: bool, - hit_callback: DatasetIndexToolCallbackHandler, - message_id: str, - memory: TokenBufferMemory | None = None, - inputs: Mapping[str, Any] | None = None, - vision_enabled: bool = False, - ) -> tuple[str | None, list[File] | None]: - """ - Retrieve dataset. - :param app_id: app_id - :param user_id: user_id - :param tenant_id: tenant id - :param model_config: model config - :param config: dataset config - :param query: query - :param invoke_from: invoke from - :param show_retrieve_source: show retrieve source - :param hit_callback: hit callback - :param message_id: message id - :param memory: memory - :param inputs: inputs - :return: - """ - dataset_ids = config.dataset_ids - if len(dataset_ids) == 0: - return None, [] - retrieve_config = config.retrieve_config - - model_manager = ModelManager.for_tenant(tenant_id=tenant_id, user_id=user_id) - model_instance = model_manager.get_model_instance( - tenant_id=tenant_id, model_type=ModelType.LLM, provider=model_config.provider, model=model_config.model - ) - model_type_instance = cast(LargeLanguageModel, model_instance.model_type_instance) - - # Reuse the caller-bound model instance for both schema resolution and - # downstream planner/invoke calls so a single request never mixes - # tenant-scope and request-bound runtimes. - model_schema = model_type_instance.get_model_schema( - model=model_instance.model_name, - credentials=model_instance.credentials, - ) - - if not model_schema: - return None, [] - - model_config.provider_model_bundle = model_instance.provider_model_bundle - model_config.credentials = model_instance.credentials - model_config.model_schema = model_schema - - planning_strategy = PlanningStrategy.REACT_ROUTER - features = model_schema.features - if features: - if ModelFeature.TOOL_CALL in features or ModelFeature.MULTI_TOOL_CALL in features: - planning_strategy = PlanningStrategy.ROUTER - available_datasets = self._get_available_datasets(tenant_id, dataset_ids) - - if inputs: - inputs = {key: str(value) for key, value in inputs.items()} - else: - inputs = {} - available_datasets_ids = [dataset.id for dataset in available_datasets] - metadata_filter_document_ids, metadata_condition = self.get_metadata_filter_condition( - session, - available_datasets_ids, - query, - tenant_id, - user_id, - retrieve_config.metadata_filtering_mode, # type: ignore - retrieve_config.metadata_model_config, # type: ignore - retrieve_config.metadata_filtering_conditions, - inputs, - ) - - all_documents = [] - user_from = "account" if invoke_from in {InvokeFrom.EXPLORE, InvokeFrom.DEBUGGER} else "end_user" - if retrieve_config.retrieve_strategy == DatasetRetrieveConfigEntity.RetrieveStrategy.SINGLE: - all_documents = self.single_retrieve( - session, - app_id, - tenant_id, - user_id, - user_from, - query, - available_datasets, - model_instance, - model_config, - planning_strategy, - message_id, - metadata_filter_document_ids, - metadata_condition, - ) - elif retrieve_config.retrieve_strategy == DatasetRetrieveConfigEntity.RetrieveStrategy.MULTIPLE: - all_documents = self.multiple_retrieve( - app_id, - tenant_id, - user_id, - user_from, - available_datasets, - query, - retrieve_config.top_k or 0, - retrieve_config.score_threshold or 0, - retrieve_config.rerank_mode or "reranking_model", - retrieve_config.reranking_model, - retrieve_config.weights, - True if retrieve_config.reranking_enabled is None else retrieve_config.reranking_enabled, - message_id, - metadata_filter_document_ids, - metadata_condition, - ) - - dify_documents = [item for item in all_documents if item.provider == "dify"] - external_documents = [item for item in all_documents if item.provider == "external"] - document_context_list: list[DocumentContext] = [] - context_files: list[File] = [] - retrieval_resource_list: list[RetrievalSourceMetadata] = [] - # deal with external documents - for item in external_documents: - document_context_list.append(DocumentContext(content=item.page_content, score=item.metadata.get("score"))) - source = RetrievalSourceMetadata( - dataset_id=item.metadata.get("dataset_id"), - dataset_name=item.metadata.get("dataset_name"), - document_id=item.metadata.get("document_id") or item.metadata.get("title"), - document_name=item.metadata.get("title"), - data_source_type="external", - retriever_from=invoke_from.to_source(), - score=item.metadata.get("score"), - content=item.page_content, - ) - retrieval_resource_list.append(source) - # deal with dify documents - if dify_documents: - with Session(bind=session.get_bind()) as format_session: - records = RetrievalService.format_retrieval_documents(format_session, dify_documents) - if records: - authorized_records = [] - for record in records: - segment = record.segment - attachments = authorize_retrieved_segment( - segment, tenant_id=tenant_id, dataset_ids=available_datasets_ids, session=session - ) - if attachments is None: - continue - authorized_records.append(record) - # Build content: if summary exists, add it before the segment content - if segment.answer: - segment_content = ( - f"question:{sign_segment_content(segment, session=session)} answer:{segment.answer}" - ) - else: - segment_content = sign_segment_content(segment, session=session) - - # If summary exists, prepend it to the content - if record.summary: - final_content = f"{record.summary}\n{segment_content}" - else: - final_content = segment_content - - document_context_list.append( - DocumentContext( - content=final_content, - score=record.score, - ) - ) - if vision_enabled: - if attachments: - for upload_file in attachments: - attachment_info = File( - file_id=upload_file.id, - filename=upload_file.name, - extension="." + upload_file.extension, - mime_type=upload_file.mime_type, - file_type=FileType.IMAGE, - transfer_method=FileTransferMethod.LOCAL_FILE, - remote_url=upload_file.source_url, - reference=build_file_reference( - record_id=str(upload_file.id), - ), - size=upload_file.size, - storage_key=upload_file.key, - url=sign_upload_file_preview_url(upload_file.id, upload_file.extension), - ) - context_files.append(attachment_info) - if show_retrieve_source: - dataset_ids = [record.segment.dataset_id for record in authorized_records] - document_ids = [record.segment.document_id for record in authorized_records] - dataset_document_stmt = select(DatasetDocument).where( - DatasetDocument.id.in_(document_ids), - DatasetDocument.enabled == True, - DatasetDocument.archived == False, - ) - documents = session.execute(dataset_document_stmt).scalars().all() # type: ignore - dataset_stmt = select(Dataset).where( - Dataset.id.in_(dataset_ids), - ) - datasets = session.execute(dataset_stmt).scalars().all() # type: ignore - dataset_map = {i.id: i for i in datasets} - document_map = {i.id: i for i in documents} - for record in authorized_records: - segment = record.segment - dataset_item = dataset_map.get(segment.dataset_id) - document_item = document_map.get(segment.document_id) - if dataset_item and document_item: - source = RetrievalSourceMetadata( - dataset_id=dataset_item.id, - dataset_name=dataset_item.name, - document_id=document_item.id, - document_name=document_item.name, - data_source_type=document_item.data_source_type, - segment_id=segment.id, - retriever_from=invoke_from.to_source(), - score=record.score or 0.0, - doc_metadata=document_item.doc_metadata, - ) - - if invoke_from.to_source() == "dev": - source.hit_count = segment.hit_count - source.word_count = segment.word_count - source.segment_position = segment.position - source.index_node_hash = segment.index_node_hash - if segment.answer: - source.content = f"question:{segment.content} \nanswer:{segment.answer}" - else: - source.content = segment.content - # Add summary if this segment was retrieved via summary - if hasattr(record, "summary") and record.summary: - source.summary = record.summary - retrieval_resource_list.append(source) - if hit_callback and retrieval_resource_list: - retrieval_resource_list = sorted(retrieval_resource_list, key=lambda x: x.score or 0.0, reverse=True) - for position, item in enumerate(retrieval_resource_list, start=1): - item.position = position - hit_callback.return_retriever_resource_info(retrieval_resource_list) - if document_context_list: - document_context_list = sorted(document_context_list, key=lambda x: x.score or 0.0, reverse=True) - return "\n".join([document_context.content for document_context in document_context_list]), context_files - return "", context_files - - @trace_span() - def single_retrieve( - self, - session: Session, - app_id: str, - tenant_id: str, - user_id: str, - user_from: str, - query: str, - available_datasets: list[Dataset], - model_instance: ModelInstance, - model_config: ModelConfigWithCredentialsEntity, - planning_strategy: PlanningStrategy, - message_id: str | None = None, - metadata_filter_document_ids: dict[str, list[str]] | None = None, - metadata_condition: MetadataFilteringCondition | None = None, - ): - tools: list[PromptMessageTool] = [] - for dataset in available_datasets: - description = dataset.description - if not description: - description = "useful for when you want to answer queries about the " + dataset.name - - description = description.replace("\n", "").replace("\r", "") - message_tool = PromptMessageTool( - name=dataset.id, - description=description, - parameters={ - "type": "object", - "properties": {}, - "required": [], - }, - ) - tools.append(message_tool) - dataset_id = None - router_usage = LLMUsage.empty_usage() - if planning_strategy == PlanningStrategy.REACT_ROUTER: - react_multi_dataset_router = ReactMultiDatasetRouter() - dataset_id, router_usage = react_multi_dataset_router.invoke( - query, tools, model_config, model_instance, user_id, tenant_id - ) - - elif planning_strategy == PlanningStrategy.ROUTER: - function_call_router = FunctionCallMultiDatasetRouter() - dataset_id, router_usage = function_call_router.invoke(query, tools, model_config, model_instance) - - self._record_usage(router_usage) - timer = None - if dataset_id: - allowed_dataset = next((dataset for dataset in available_datasets if dataset.id == dataset_id), None) - selected_dataset = ( - session.scalar(select(Dataset).where(Dataset.id == allowed_dataset.id, Dataset.tenant_id == tenant_id)) - if allowed_dataset - else None - ) - if selected_dataset: - results = [] - if selected_dataset.provider == "external": - external_documents = ExternalDatasetService.fetch_external_knowledge_retrieval( - session=session, - tenant_id=selected_dataset.tenant_id, - dataset_id=selected_dataset.id, - query=query, - external_retrieval_parameters=selected_dataset.retrieval_model, - metadata_condition=metadata_condition, - ) - for external_document in external_documents: - document = Document( - page_content=external_document.get("content"), - metadata=external_document.get("metadata"), - provider="external", - ) - if document.metadata is not None: - document.metadata["score"] = external_document.get("score") - document.metadata["title"] = external_document.get("title") - document.metadata["dataset_id"] = selected_dataset.id - document.metadata["dataset_name"] = selected_dataset.name - results.append(document) - else: - if metadata_condition and not metadata_filter_document_ids: - return [] - document_ids_filter = None - if metadata_filter_document_ids: - document_ids = metadata_filter_document_ids.get(selected_dataset.id, []) - if document_ids: - document_ids_filter = document_ids - else: - return [] - retrieval_model_config: DefaultRetrievalModelDict = ( - cast(DefaultRetrievalModelDict, selected_dataset.retrieval_model) - if selected_dataset.retrieval_model - else default_retrieval_model - ) - - # get top k - top_k = retrieval_model_config["top_k"] - # get retrieval method - if selected_dataset.indexing_technique == IndexTechniqueType.ECONOMY: - retrieval_method = RetrievalMethod.KEYWORD_SEARCH - else: - retrieval_method = retrieval_model_config["search_method"] - # get reranking model - reranking_model = ( - retrieval_model_config["reranking_model"] - if retrieval_model_config["reranking_enable"] - else None - ) - # get score threshold - score_threshold = 0.0 - score_threshold_enabled = retrieval_model_config.get("score_threshold_enabled") - if score_threshold_enabled: - score_threshold = retrieval_model_config.get("score_threshold", 0.0) - - with measure_time() as timer: - results = RetrievalService.retrieve( - retrieval_method=retrieval_method, - dataset_id=selected_dataset.id, - query=query, - top_k=top_k, - score_threshold=score_threshold, - reranking_model=reranking_model, - reranking_mode=retrieval_model_config.get("reranking_mode", "reranking_model"), - weights=retrieval_model_config.get("weights", None), - document_ids_filter=document_ids_filter, - ) - self._on_query(query, None, [selected_dataset.id], app_id, user_from, user_id) - - if results: - thread = threading.Thread( - target=propagate_context(self._on_retrieval_end), - kwargs={ - "flask_app": current_app._get_current_object(), # type: ignore - "documents": results, - "message_id": message_id, - "timer": timer, - }, - ) - thread.start() - - return results - return [] - - @trace_span() - def multiple_retrieve( - self, - app_id: str, - tenant_id: str, - user_id: str, - user_from: str, - available_datasets: list[Dataset], - query: str | None, - top_k: int, - score_threshold: float, - reranking_mode: str, - reranking_model: RerankingModelDict | None = None, - weights: WeightsDict | None = None, - reranking_enable: bool = True, - message_id: str | None = None, - metadata_filter_document_ids: dict[str, list[str]] | None = None, - metadata_condition: MetadataFilteringCondition | None = None, - attachment_ids: list[str] | None = None, - ): - if not available_datasets: - return [] - all_threads = [] - all_documents: list[Document] = [] - dataset_ids = [dataset.id for dataset in available_datasets] - index_type_check = all( - item.indexing_technique == available_datasets[0].indexing_technique for item in available_datasets - ) - if not index_type_check and (not reranking_enable or reranking_mode != RerankMode.RERANKING_MODEL): - raise ValueError( - "The configured knowledge base list have different indexing technique, please set reranking model." - ) - index_type = available_datasets[0].indexing_technique - if index_type == IndexTechniqueType.HIGH_QUALITY: - embedding_model_check = all( - item.embedding_model == available_datasets[0].embedding_model for item in available_datasets - ) - embedding_model_provider_check = all( - item.embedding_model_provider == available_datasets[0].embedding_model_provider - for item in available_datasets - ) - if ( - reranking_enable - and reranking_mode == "weighted_score" - and (not embedding_model_check or not embedding_model_provider_check) - ): - raise ValueError( - "The configured knowledge base list have different embedding model, please set reranking model." - ) - if reranking_enable and reranking_mode == RerankMode.WEIGHTED_SCORE: - if weights is not None: - weights["vector_setting"]["embedding_provider_name"] = available_datasets[ - 0 - ].embedding_model_provider - weights["vector_setting"]["embedding_model_name"] = available_datasets[0].embedding_model - dataset_count = len(available_datasets) - with measure_time() as timer: - cancel_event = threading.Event() - thread_exceptions: list[Exception] = [] - - if query: - query_thread = threading.Thread( - target=propagate_context(self._multiple_retrieve_thread_safely), - kwargs={ - "flask_app": current_app._get_current_object(), # type: ignore - "available_datasets": available_datasets, - "metadata_condition": metadata_condition, - "metadata_filter_document_ids": metadata_filter_document_ids, - "all_documents": all_documents, - "tenant_id": tenant_id, - "reranking_enable": reranking_enable, - "reranking_mode": reranking_mode, - "reranking_model": reranking_model, - "weights": weights, - "top_k": top_k, - "score_threshold": score_threshold, - "query": query, - "attachment_id": None, - "dataset_count": dataset_count, - "cancel_event": cancel_event, - "thread_exceptions": thread_exceptions, - }, - ) - all_threads.append(query_thread) - query_thread.start() - if attachment_ids: - for attachment_id in attachment_ids: - attachment_thread = threading.Thread( - target=propagate_context(self._multiple_retrieve_thread_safely), - kwargs={ - "flask_app": current_app._get_current_object(), # type: ignore - "available_datasets": available_datasets, - "metadata_condition": metadata_condition, - "metadata_filter_document_ids": metadata_filter_document_ids, - "all_documents": all_documents, - "tenant_id": tenant_id, - "reranking_enable": reranking_enable, - "reranking_mode": reranking_mode, - "reranking_model": reranking_model, - "weights": weights, - "top_k": top_k, - "score_threshold": score_threshold, - "query": None, - "attachment_id": attachment_id, - "dataset_count": dataset_count, - "cancel_event": cancel_event, - "thread_exceptions": thread_exceptions, - }, - ) - all_threads.append(attachment_thread) - attachment_thread.start() - - # Poll threads with short timeout to detect errors quickly (fail-fast) - while any(t.is_alive() for t in all_threads): - for thread in all_threads: - thread.join(timeout=0.1) - if thread_exceptions: - cancel_event.set() - break - if thread_exceptions: - break - - if thread_exceptions: - raise thread_exceptions[0] - self._on_query(query, attachment_ids, dataset_ids, app_id, user_from, user_id) - - if all_documents: - # add thread to call _on_retrieval_end - retrieval_end_thread = threading.Thread( - target=propagate_context(self._on_retrieval_end), - kwargs={ - "flask_app": current_app._get_current_object(), # type: ignore - "documents": all_documents, - "message_id": message_id, - "timer": timer, - }, - ) - retrieval_end_thread.start() - retrieval_resource_list = [] - doc_ids_filter = [] - for document in all_documents: - if document.provider == "dify": - doc_id = document.metadata.get("doc_id") - if doc_id and doc_id not in doc_ids_filter: - doc_ids_filter.append(doc_id) - retrieval_resource_list.append(document) - elif document.provider == "external": - retrieval_resource_list.append(document) - return retrieval_resource_list - - @retry( - retry=retry_if_exception(_is_postgres_deadlock_error), - stop=stop_after_attempt(_HIT_COUNT_UPDATE_MAX_ATTEMPTS), - wait=wait_exponential(multiplier=0.05, min=0.05, max=0.1), - before_sleep=before_sleep_log(logger, logging.WARNING), - reraise=True, - ) - def _on_retrieval_end( - self, - flask_app: Flask, - documents: list[Document], - message_id: str | None = None, - timer: dict[str, Any] | None = None, - ): - """Handle retrieval end.""" - with flask_app.app_context(): - dify_documents = [document for document in documents if document.provider == "dify"] - if not dify_documents: - self._send_trace_task(message_id, documents, timer) - return - - with sessionmaker(bind=db.engine).begin() as session: - # Collect all document_ids and batch fetch DatasetDocuments - document_ids = { - doc.metadata["document_id"] - for doc in dify_documents - if doc.metadata and "document_id" in doc.metadata - } - if not document_ids: - self._send_trace_task(message_id, documents, timer) - return - - dataset_docs_stmt = select(DatasetDocument).where(DatasetDocument.id.in_(document_ids)) - dataset_docs = session.scalars(dataset_docs_stmt).all() - dataset_doc_map = {str(doc.id): doc for doc in dataset_docs} - - # Categorize documents by type and collect necessary IDs - parent_child_text_docs: list[tuple[Document, DatasetDocument]] = [] - parent_child_image_docs: list[tuple[Document, DatasetDocument]] = [] - normal_text_docs: list[tuple[Document, DatasetDocument]] = [] - normal_image_docs: list[tuple[Document, DatasetDocument]] = [] - - for doc in dify_documents: - if not doc.metadata or "document_id" not in doc.metadata: - continue - dataset_doc = dataset_doc_map.get(doc.metadata["document_id"]) - if not dataset_doc: - continue - - is_image = doc.metadata.get("doc_type") == DocType.IMAGE - is_parent_child = dataset_doc.doc_form == IndexStructureType.PARENT_CHILD_INDEX - - if is_parent_child: - if is_image: - parent_child_image_docs.append((doc, dataset_doc)) - else: - parent_child_text_docs.append((doc, dataset_doc)) - else: - if is_image: - normal_image_docs.append((doc, dataset_doc)) - else: - normal_text_docs.append((doc, dataset_doc)) - - segment_ids_to_update: set[str] = set() - - # Process PARENT_CHILD_INDEX text documents - batch fetch ChildChunks - if parent_child_text_docs: - index_node_ids = [doc.metadata["doc_id"] for doc, _ in parent_child_text_docs if doc.metadata] - if index_node_ids: - child_chunks_stmt = select(ChildChunk).where(ChildChunk.index_node_id.in_(index_node_ids)) - child_chunks = session.scalars(child_chunks_stmt).all() - child_chunk_map = {chunk.index_node_id: chunk.segment_id for chunk in child_chunks} - for doc, _ in parent_child_text_docs: - if doc.metadata: - segment_id = child_chunk_map.get(doc.metadata["doc_id"]) - if segment_id: - segment_ids_to_update.add(str(segment_id)) - - # Process non-PARENT_CHILD_INDEX text documents - batch fetch DocumentSegments - if normal_text_docs: - index_node_ids = [doc.metadata["doc_id"] for doc, _ in normal_text_docs if doc.metadata] - if index_node_ids: - segments_stmt = select(DocumentSegment).where(DocumentSegment.index_node_id.in_(index_node_ids)) - segments = session.scalars(segments_stmt).all() - segment_map = {seg.index_node_id: seg.id for seg in segments} - for doc, _ in normal_text_docs: - if doc.metadata: - segment_id = segment_map.get(doc.metadata["doc_id"]) - if segment_id: - segment_ids_to_update.add(str(segment_id)) - - # Process IMAGE documents - batch fetch SegmentAttachmentBindings - all_image_docs = parent_child_image_docs + normal_image_docs - if all_image_docs: - attachment_ids = [ - doc.metadata["doc_id"] - for doc, _ in all_image_docs - if doc.metadata and doc.metadata.get("doc_id") - ] - if attachment_ids: - bindings_stmt = select(SegmentAttachmentBinding).where( - SegmentAttachmentBinding.attachment_id.in_(attachment_ids) - ) - bindings = session.scalars(bindings_stmt).all() - segment_ids_to_update.update(str(binding.segment_id) for binding in bindings) - - # Batch update hit_count for all segments - if segment_ids_to_update: - # PostgreSQL does not guarantee that an IN predicate is visited in parameter order. - # Lock every target row explicitly and consistently before the multi-row update. - session.scalars( - select(DocumentSegment.id) - .where(DocumentSegment.id.in_(segment_ids_to_update)) - .order_by(DocumentSegment.id) - .with_for_update() - ).all() - session.execute( - update(DocumentSegment) - .where(DocumentSegment.id.in_(segment_ids_to_update)) - .values(hit_count=DocumentSegment.hit_count + 1) - .execution_options(synchronize_session=False) - ) - - self._send_trace_task(message_id, documents, timer) - - def _send_trace_task(self, message_id: str | None, documents: list[Document], timer: dict[str, Any] | None): - """Send trace task if trace manager is available.""" - trace_manager: TraceQueueManager | None = ( - self.application_generate_entity.trace_manager if self.application_generate_entity else None - ) - if trace_manager: - trace_manager.add_trace_task( - TraceTask( - TraceTaskName.DATASET_RETRIEVAL_TRACE, message_id=message_id, documents=documents, timer=timer - ) - ) - - @staticmethod - def _resolve_creator_user_role(user_from: str) -> CreatorUserRole | None: - """Map runtime user source values to dataset query audit roles. - - Workflow run context uses the hyphenated ``end-user`` value, while - ``DatasetQuery.created_by_role`` persists the underscore-based - ``CreatorUserRole.END_USER`` enum. Query logging is a side effect, so an - unsupported value should be skipped instead of aborting retrieval. - """ - normalized_user_from = str(user_from).strip().lower().replace("-", "_") - if normalized_user_from == CreatorUserRole.ACCOUNT.value: - return CreatorUserRole.ACCOUNT - if normalized_user_from == CreatorUserRole.END_USER.value: - return CreatorUserRole.END_USER - - logger.warning("Skipping dataset query audit log for unsupported user_from=%r", user_from) - return None - - def _on_query( - self, - query: str | None, - attachment_ids: list[str] | None, - dataset_ids: list[str], - app_id: str, - user_from: str, - user_id: str, - ): - """ - Persist dataset query audit rows for retrieval requests. - - Query audit logging is a side effect of retrieval. Keep it in an - independent transaction so failures or commits here do not affect the - request/workflow transaction that called the retriever. - """ - if not query and not attachment_ids: - return - created_by = parse_uuid_str_or_none(user_id) - if created_by is None: - logger.debug( - "Skipping dataset query log: empty created_by user_id (user_from=%s, app_id=%s)", - user_from, - app_id, - ) - return - created_by_role = self._resolve_creator_user_role(user_from) - if created_by_role is None: - return - dataset_queries = [] - for dataset_id in dataset_ids: - contents = [] - if query: - contents.append({"content_type": QueryType.TEXT_QUERY, "content": query}) - if attachment_ids: - for attachment_id in attachment_ids: - contents.append({"content_type": QueryType.IMAGE_QUERY, "content": attachment_id}) - if contents: - dataset_query = DatasetQuery( - dataset_id=dataset_id, - content=json.dumps(contents), - source=DatasetQuerySource.APP, - source_app_id=app_id, - created_by_role=created_by_role, - created_by=created_by, - ) - dataset_queries.append(dataset_query) - - if not dataset_queries: - return - - with sessionmaker(bind=db.engine, expire_on_commit=False).begin() as session: - session.add_all(dataset_queries) - - def _retriever( - self, - flask_app: Flask, - session: Session, - dataset_id: str, - query: str, - top_k: int, - all_documents: list[Document], - document_ids_filter: list[str] | None = None, - metadata_condition: MetadataFilteringCondition | None = None, - attachment_ids: list[str] | None = None, - ): - with flask_app.app_context(): - dataset_stmt = select(Dataset).where(Dataset.id == dataset_id) - dataset = session.scalar(dataset_stmt) - - if not dataset: - return [] - - if dataset.provider == "external" and query: - external_documents = ExternalDatasetService.fetch_external_knowledge_retrieval( - session=session, - tenant_id=dataset.tenant_id, - dataset_id=dataset_id, - query=query, - external_retrieval_parameters=dataset.retrieval_model, - metadata_condition=metadata_condition, - ) - for external_document in external_documents: - document = Document( - page_content=external_document.get("content"), - metadata=external_document.get("metadata"), - provider="external", - ) - if document.metadata is not None: - document.metadata["score"] = external_document.get("score") - document.metadata["title"] = external_document.get("title") - document.metadata["dataset_id"] = dataset_id - document.metadata["dataset_name"] = dataset.name - all_documents.append(document) - else: - # get retrieval model , if the model is not setting , using default - retrieval_model: DefaultRetrievalModelDict = ( - cast(DefaultRetrievalModelDict, dataset.retrieval_model) - if dataset.retrieval_model - else default_retrieval_model - ) - - if dataset.indexing_technique == IndexTechniqueType.ECONOMY: - # use keyword table query - documents = RetrievalService.retrieve( - retrieval_method=RetrievalMethod.KEYWORD_SEARCH, - dataset_id=dataset.id, - query=query, - top_k=top_k, - document_ids_filter=document_ids_filter, - ) - if documents: - all_documents.extend(documents) - else: - if top_k > 0: - # retrieval source - documents = RetrievalService.retrieve( - retrieval_method=retrieval_model["search_method"], - dataset_id=dataset.id, - query=query, - top_k=retrieval_model.get("top_k") or 4, - score_threshold=retrieval_model.get("score_threshold", 0.0) - if retrieval_model["score_threshold_enabled"] - else 0.0, - reranking_model=retrieval_model.get("reranking_model", None) - if retrieval_model["reranking_enable"] - else None, - reranking_mode=retrieval_model.get("reranking_mode") or "reranking_model", - weights=retrieval_model.get("weights", None), - document_ids_filter=document_ids_filter, - attachment_ids=attachment_ids, - ) - - all_documents.extend(documents) - - @trace_span() - def _run_retriever_thread( - self, - *, - flask_app: Flask, - dataset_id: str, - query: str | None, - top_k: int, - all_documents: list[Document], - document_ids_filter: list[str] | None, - metadata_condition: MetadataFilteringCondition | None, - attachment_ids: list[str] | None, - ) -> None: - with session_factory.create_session() as session: - self._retriever( - flask_app=flask_app, - session=session, - dataset_id=dataset_id, - query=query or "", - top_k=top_k, - all_documents=all_documents, - document_ids_filter=document_ids_filter, - metadata_condition=metadata_condition, - attachment_ids=attachment_ids, - ) - - def _run_retriever_thread_safely( - self, - *, - flask_app: Flask, - dataset_id: str, - query: str | None, - top_k: int, - all_documents: list[Document], - document_ids_filter: list[str] | None, - metadata_condition: MetadataFilteringCondition | None, - attachment_ids: list[str] | None, - cancel_event: threading.Event | None, - thread_exceptions: list[Exception] | None, - skip_on_error: bool = False, - ) -> None: - """Collect errors after tracing, or skip dataset-level failures when requested.""" - try: - self._run_retriever_thread( - flask_app=flask_app, - dataset_id=dataset_id, - query=query, - top_k=top_k, - all_documents=all_documents, - document_ids_filter=document_ids_filter, - metadata_condition=metadata_condition, - attachment_ids=attachment_ids, - ) - except Exception as exc: - if skip_on_error: - logger.warning( - "Skipping dataset retrieval because retriever failed, dataset_id=%s, error_type=%s, error=%s", - dataset_id, - type(exc).__name__, - str(exc), - ) - span = get_current_span() - if span and span.is_recording(): - span.add_event( - "dataset_retrieval.skipped", - attributes={ - "dataset_id": dataset_id, - "error.type": type(exc).__name__, - "error.message": str(exc), - }, - ) - return - - if cancel_event: - cancel_event.set() - if thread_exceptions is not None: - thread_exceptions.append(exc) - - def to_dataset_retriever_tool( - self, - session: Session, - tenant_id: str, - dataset_ids: list[str], - retrieve_config: DatasetRetrieveConfigEntity, - return_resource: bool, - invoke_from: InvokeFrom, - hit_callback: DatasetIndexToolCallbackHandler, - user_id: str, - inputs: dict[str, Any], - ) -> list[DatasetRetrieverBaseTool] | None: - """ - A dataset tool is a tool that can be used to retrieve information from a dataset - :param tenant_id: tenant id - :param dataset_ids: dataset ids - :param retrieve_config: retrieve config - :param return_resource: return resource - :param invoke_from: invoke from - :param hit_callback: hit callback - """ - tools: list[DatasetRetrieverBaseTool] = [] - available_datasets = [] - for dataset_id in dataset_ids: - # get dataset from dataset id - dataset_stmt = select(Dataset).where(Dataset.tenant_id == tenant_id, Dataset.id == dataset_id) - dataset = session.scalar(dataset_stmt) - - # pass if dataset is not available - if not dataset: - continue - - # pass if dataset is not available - if ( - dataset - and dataset.provider != "external" - and get_dataset_available_document_count(dataset, session=session) == 0 - ): - continue - - available_datasets.append(dataset) - - if retrieve_config.retrieve_strategy == DatasetRetrieveConfigEntity.RetrieveStrategy.SINGLE: - # get retrieval model config - default_retrieval_model: DefaultRetrievalModelDict = { - "search_method": RetrievalMethod.SEMANTIC_SEARCH, - "reranking_enable": False, - "reranking_model": {"reranking_provider_name": "", "reranking_model_name": ""}, - "top_k": 2, - "score_threshold_enabled": False, - } - - for dataset in available_datasets: - retrieval_model_config: DefaultRetrievalModelDict = ( - cast(DefaultRetrievalModelDict, dataset.retrieval_model) - if dataset.retrieval_model - else default_retrieval_model - ) - - # get top k - top_k = retrieval_model_config["top_k"] - - # get score threshold - score_threshold = None - score_threshold_enabled = retrieval_model_config.get("score_threshold_enabled") - if score_threshold_enabled: - score_threshold = retrieval_model_config.get("score_threshold") - - from core.tools.utils.dataset_retriever.dataset_retriever_tool import DatasetRetrieverTool - - tool = DatasetRetrieverTool.from_dataset( - dataset=dataset, - top_k=top_k, - score_threshold=score_threshold, - hit_callbacks=[hit_callback], - return_resource=return_resource, - retriever_from=invoke_from.to_source(), - retrieve_config=retrieve_config, - user_id=user_id, - inputs=inputs, - ) - - tools.append(tool) - elif retrieve_config.retrieve_strategy == DatasetRetrieveConfigEntity.RetrieveStrategy.MULTIPLE: - from core.tools.utils.dataset_retriever.dataset_multi_retriever_tool import DatasetMultiRetrieverTool - - if retrieve_config.reranking_model is None: - raise ValueError("Reranking model is required for multiple retrieval") - - tool = DatasetMultiRetrieverTool.from_dataset( - dataset_ids=[dataset.id for dataset in available_datasets], - tenant_id=tenant_id, - top_k=retrieve_config.top_k or 4, - score_threshold=retrieve_config.score_threshold, - hit_callbacks=[hit_callback], - return_resource=return_resource, - retriever_from=invoke_from.to_source(), - reranking_provider_name=retrieve_config.reranking_model["reranking_provider_name"], - reranking_model_name=retrieve_config.reranking_model["reranking_model_name"], - ) - - tools.append(tool) - - return tools - - def calculate_keyword_score(self, query: str, documents: list[Document], top_k: int) -> list[Document]: - """ - Calculate keywords scores - :param query: search query - :param documents: documents for reranking - :param top_k: top k - - :return: - """ - keyword_table_handler = JiebaKeywordTableHandler() - query_keywords = keyword_table_handler.extract_keywords(query, None) - documents_keywords = [] - for document in documents: - if document.metadata is not None: - # get the document keywords - document_keywords = keyword_table_handler.extract_keywords(document.page_content, None) - document.metadata["keywords"] = document_keywords - documents_keywords.append(document_keywords) - - # Counter query keywords(TF) - query_keyword_counts = Counter(query_keywords) - - # total documents - total_documents = len(documents) - - # calculate all documents' keywords IDF - all_keywords = set() - for document_keywords in documents_keywords: - all_keywords.update(document_keywords) - - keyword_idf = {} - for keyword in all_keywords: - # calculate include query keywords' documents - doc_count_containing_keyword = sum(1 for doc_keywords in documents_keywords if keyword in doc_keywords) - # IDF - keyword_idf[keyword] = math.log((1 + total_documents) / (1 + doc_count_containing_keyword)) + 1 - - query_tfidf = {} - - for keyword, count in query_keyword_counts.items(): - tf = count - idf = keyword_idf.get(keyword, 0) - query_tfidf[keyword] = tf * idf - - # calculate all documents' TF-IDF - documents_tfidf = [] - for document_keywords in documents_keywords: - document_keyword_counts = Counter(document_keywords) - document_tfidf = {} - for keyword, count in document_keyword_counts.items(): - tf = count - idf = keyword_idf.get(keyword, 0) - document_tfidf[keyword] = tf * idf - documents_tfidf.append(document_tfidf) - - def cosine_similarity(vec1, vec2): - intersection = set(vec1.keys()) & set(vec2.keys()) - numerator = sum(vec1[x] * vec2[x] for x in intersection) - - sum1 = sum(vec1[x] ** 2 for x in vec1) - sum2 = sum(vec2[x] ** 2 for x in vec2) - denominator = math.sqrt(sum1) * math.sqrt(sum2) - - if not denominator: - return 0.0 - else: - return float(numerator) / denominator - - similarities = [] - for document_tfidf in documents_tfidf: - similarity = cosine_similarity(query_tfidf, document_tfidf) - similarities.append(similarity) - - for document, score in zip(documents, similarities): - # format document - if document.metadata is not None: - document.metadata["score"] = score - documents = sorted(documents, key=lambda x: x.metadata.get("score", 0) if x.metadata else 0, reverse=True) - return documents[:top_k] if top_k else documents - - def calculate_vector_score( - self, all_documents: list[Document], top_k: int, score_threshold: float - ) -> list[Document]: - filter_documents = [] - for document in all_documents: - if score_threshold is None or (document.metadata and document.metadata.get("score", 0) >= score_threshold): - filter_documents.append(document) - - if not filter_documents: - return [] - filter_documents = sorted( - filter_documents, key=lambda x: x.metadata.get("score", 0) if x.metadata else 0, reverse=True - ) - return filter_documents[:top_k] if top_k else filter_documents - - @with_credit_usage_created_by(CreditUsageCreatedBy.KNOWLEDGE_RETRIEVAL) - def get_metadata_filter_condition( - self, - session: Session, - dataset_ids: list[str], - query: str, - tenant_id: str, - user_id: str, - metadata_filtering_mode: str, - metadata_model_config: ModelConfig, - metadata_filtering_conditions: MetadataFilteringCondition | None, - inputs: dict[str, Any], - ) -> tuple[dict[str, list[str]] | None, MetadataFilteringCondition | None]: - document_query = select(DatasetDocument).where( - DatasetDocument.dataset_id.in_(dataset_ids), - DatasetDocument.indexing_status == "completed", - DatasetDocument.enabled == True, - DatasetDocument.archived == False, - ) - filters = [] # type: ignore - metadata_condition = None - if metadata_filtering_mode == "disabled": - return None, None - elif metadata_filtering_mode == "automatic": - automatic_metadata_filters = self._automatic_metadata_filter_func( - session, dataset_ids, query, tenant_id, user_id, metadata_model_config - ) - if automatic_metadata_filters: - conditions = [] - for sequence, filter in enumerate(automatic_metadata_filters): - self.process_metadata_filter_func( - sequence, - filter.get("condition"), # type: ignore - filter.get("metadata_name"), # type: ignore - filter.get("value"), - filters, # type: ignore - ) - conditions.append( - Condition( - name=filter.get("metadata_name"), # type: ignore - comparison_operator=filter.get("condition"), # type: ignore - value=filter.get("value"), - ) - ) - metadata_condition = MetadataFilteringCondition( - logical_operator=metadata_filtering_conditions.logical_operator - if metadata_filtering_conditions - else "or", # type: ignore - conditions=conditions, - ) - elif metadata_filtering_mode == "manual": - if metadata_filtering_conditions: - conditions = [] - for sequence, condition in enumerate(metadata_filtering_conditions.conditions): # type: ignore - metadata_name = condition.name - expected_value = condition.value - if expected_value is not None and condition.comparison_operator not in ("empty", "not empty"): - if isinstance(expected_value, str): - expected_value = self._replace_metadata_filter_value(expected_value, inputs) - conditions.append( - Condition( - name=metadata_name, - comparison_operator=condition.comparison_operator, - value=expected_value, - ) - ) - filters = self.process_metadata_filter_func( - sequence, - condition.comparison_operator, - metadata_name, - expected_value, - filters, - ) - metadata_condition = MetadataFilteringCondition( - logical_operator=metadata_filtering_conditions.logical_operator, - conditions=conditions, - ) - else: - raise ValueError("Invalid metadata filtering mode") - if filters: - if metadata_filtering_conditions and metadata_filtering_conditions.logical_operator == "and": # type: ignore - document_query = document_query.where(and_(*filters)) - else: - document_query = document_query.where(or_(*filters)) - documents = session.scalars(document_query).all() - # group by dataset_id - metadata_filter_document_ids = defaultdict(list) if documents else None # type: ignore - for document in documents: - metadata_filter_document_ids[document.dataset_id].append(document.id) # type: ignore - return metadata_filter_document_ids, metadata_condition - - def _replace_metadata_filter_value(self, text: str, inputs: dict[str, Any]) -> str: - if not inputs: - return text - - def replacer(match): - key = match.group(1) - return str(inputs.get(key, f"{{{{{key}}}}}")) - - pattern = re.compile(r"\{\{(\w+)\}\}") - output = pattern.sub(replacer, text) - if isinstance(output, str): - output = re.sub(r"[\r\n\t]+", " ", output).strip() - return output - - def _automatic_metadata_filter_func( - self, - session: Session, - dataset_ids: list[str], - query: str, - tenant_id: str, - user_id: str, - metadata_model_config: ModelConfig, - ) -> list[dict[str, Any]] | None: - # get all metadata field - metadata_stmt = select(DatasetMetadata).where(DatasetMetadata.dataset_id.in_(dataset_ids)) - metadata_fields = session.scalars(metadata_stmt).all() - all_metadata_fields = [metadata_field.name for metadata_field in metadata_fields] - # get metadata model config - if metadata_model_config is None: - raise ValueError("metadata_model_config is required") - # get metadata model instance - # fetch model config - model_instance, model_config = self._fetch_model_config(tenant_id, metadata_model_config, user_id=user_id) - - # fetch prompt messages - prompt_messages, stop = self._get_prompt_template( - model_config=model_config, - mode=metadata_model_config.mode, - metadata_fields=all_metadata_fields, - query=query or "", - ) - - try: - # handle invoke result - invoke_result = cast( - Generator[LLMResult, None, None], - model_instance.invoke_llm( - prompt_messages=prompt_messages, - model_parameters=model_config.parameters, - stop=stop, - stream=True, - ), - ) - - # handle invoke result - result_text, usage = self._handle_invoke_result(invoke_result=invoke_result) - self._record_usage(usage) - - result_text_json = parse_and_check_json_markdown(result_text, []) - automatic_metadata_filters = [] - if "metadata_map" in result_text_json: - metadata_map = result_text_json["metadata_map"] - for item in metadata_map: - if item.get("metadata_field_name") in all_metadata_fields: - automatic_metadata_filters.append( - { - "metadata_name": item.get("metadata_field_name"), - "value": item.get("metadata_field_value"), - "condition": item.get("comparison_operator"), - } - ) - except Exception as e: - logger.warning(e, exc_info=True) - return None - return automatic_metadata_filters - - @classmethod - def process_metadata_filter_func( - cls, sequence: int, condition: str, metadata_name: str, value: Any | None, filters: list - ): - if value is None and condition not in ("empty", "not empty"): - return filters - - json_field = DatasetDocument.doc_metadata[metadata_name].as_string() - - from libs.helper import escape_like_pattern - - match condition: - case "contains": - escaped_value = escape_like_pattern(str(value)) - filters.append(json_field.like(f"%{escaped_value}%", escape="\\")) - - case "not contains": - escaped_value = escape_like_pattern(str(value)) - filters.append(json_field.notlike(f"%{escaped_value}%", escape="\\")) - - case "start with": - escaped_value = escape_like_pattern(str(value)) - filters.append(json_field.like(f"{escaped_value}%", escape="\\")) - - case "end with": - escaped_value = escape_like_pattern(str(value)) - filters.append(json_field.like(f"%{escaped_value}", escape="\\")) - - case "is" | "=": - match value: - case str(): - filters.append(json_field == value) - case int() | float(): - filters.append(DatasetDocument.doc_metadata[metadata_name].as_float() == value) - - case "is not" | "≠": - match value: - case str(): - filters.append(json_field != value) - case int() | float(): - filters.append(DatasetDocument.doc_metadata[metadata_name].as_float() != value) - - case "empty": - filters.append(DatasetDocument.doc_metadata[metadata_name].is_(None)) - - case "not empty": - filters.append(DatasetDocument.doc_metadata[metadata_name].isnot(None)) - - case "before" | "<": - filters.append(DatasetDocument.doc_metadata[metadata_name].as_float() < value) - - case "after" | ">": - filters.append(DatasetDocument.doc_metadata[metadata_name].as_float() > value) - - case "≤" | "<=": - filters.append(DatasetDocument.doc_metadata[metadata_name].as_float() <= value) - - case "≥" | ">=": - filters.append(DatasetDocument.doc_metadata[metadata_name].as_float() >= value) - case "in" | "not in": - match value: - case str(): - value_list = [v.strip() for v in value.split(",") if v.strip()] - case list() | tuple(): - value_list = [str(v) for v in value if v is not None] - case _: - value_list = [str(value)] if value is not None else [] - - if not value_list: - # `field in []` is False, `field not in []` is True - filters.append(literal(condition == "not in")) - else: - op = json_field.in_ if condition == "in" else json_field.notin_ - filters.append(op(value_list)) - case _: - pass - - return filters - - def _fetch_model_config( - self, tenant_id: str, model: ModelConfig, user_id: str | None = None - ) -> tuple[ModelInstance, ModelConfigWithCredentialsEntity]: - """ - Fetch model config - """ - if model is None: - raise ValueError("single_retrieval_config is required") - model_name = model.name - provider_name = model.provider - - model_manager = ModelManager.for_tenant(tenant_id=tenant_id, user_id=user_id) - model_instance = model_manager.get_model_instance( - tenant_id=tenant_id, model_type=ModelType.LLM, provider=provider_name, model=model_name - ) - - provider_model_bundle = model_instance.provider_model_bundle - model_type_instance = model_instance.model_type_instance - model_type_instance = cast(LargeLanguageModel, model_type_instance) - - model_credentials = model_instance.credentials - - # check model - provider_model = provider_model_bundle.configuration.get_provider_model( - model=model_name, model_type=ModelType.LLM - ) - - if provider_model is None: - raise ValueError(f"Model {model_name} not exist.") - - if provider_model.status == ModelStatus.NO_CONFIGURE: - raise ValueError(f"Model {model_name} credentials is not initialized.") - elif provider_model.status == ModelStatus.NO_PERMISSION: - raise ValueError(f"Dify Hosted OpenAI {model_name} currently not support.") - elif provider_model.status == ModelStatus.QUOTA_EXCEEDED: - raise ValueError(f"Model provider {provider_name} quota exceeded.") - - # model config - completion_params = model.completion_params - stop = [] - if "stop" in completion_params: - stop = completion_params["stop"] - del completion_params["stop"] - - # get model mode - model_mode = model.mode - if not model_mode: - raise ValueError("LLM mode is required.") - - model_schema = model_type_instance.get_model_schema(model_name, model_credentials) - - if not model_schema: - raise ValueError(f"Model {model_name} not exist.") - - return model_instance, ModelConfigWithCredentialsEntity( - provider=provider_name, - model=model_name, - model_schema=model_schema, - mode=model_mode, - provider_model_bundle=provider_model_bundle, - credentials=model_credentials, - parameters=completion_params, - stop=stop, - ) - - def _get_prompt_template( - self, model_config: ModelConfigWithCredentialsEntity, mode: str, metadata_fields: list[str], query: str - ): - model_mode = ModelMode(mode) - input_text = query - - prompt_template: Union[CompletionModelPromptTemplate, list[ChatModelMessage]] - if model_mode == ModelMode.CHAT: - prompt_template = [] - system_prompt_messages = ChatModelMessage(role=PromptMessageRole.SYSTEM, text=METADATA_FILTER_SYSTEM_PROMPT) - prompt_template.append(system_prompt_messages) - user_prompt_message_1 = ChatModelMessage(role=PromptMessageRole.USER, text=METADATA_FILTER_USER_PROMPT_1) - prompt_template.append(user_prompt_message_1) - assistant_prompt_message_1 = ChatModelMessage( - role=PromptMessageRole.ASSISTANT, text=METADATA_FILTER_ASSISTANT_PROMPT_1 - ) - prompt_template.append(assistant_prompt_message_1) - user_prompt_message_2 = ChatModelMessage(role=PromptMessageRole.USER, text=METADATA_FILTER_USER_PROMPT_2) - prompt_template.append(user_prompt_message_2) - assistant_prompt_message_2 = ChatModelMessage( - role=PromptMessageRole.ASSISTANT, text=METADATA_FILTER_ASSISTANT_PROMPT_2 - ) - prompt_template.append(assistant_prompt_message_2) - user_prompt_message_3 = ChatModelMessage( - role=PromptMessageRole.USER, - text=METADATA_FILTER_USER_PROMPT_3.format( - input_text=input_text, - metadata_fields=json.dumps(metadata_fields, ensure_ascii=False), - ), - ) - prompt_template.append(user_prompt_message_3) - elif model_mode == ModelMode.COMPLETION: - prompt_template = CompletionModelPromptTemplate( - text=METADATA_FILTER_COMPLETION_PROMPT.format( - input_text=input_text, - metadata_fields=json.dumps(metadata_fields, ensure_ascii=False), - ) - ) - - else: - raise ValueError(f"Model mode {model_mode} not support.") - - prompt_transform = AdvancedPromptTransform() - prompt_messages = prompt_transform.get_prompt( - prompt_template=prompt_template, - inputs={}, - query=query or "", - files=[], - context=None, - memory_config=None, - memory=None, - model_config=model_config, - ) - stop = model_config.stop - - return prompt_messages, stop - - def _handle_invoke_result(self, invoke_result: Generator) -> tuple[str, LLMUsage]: - """ - Handle invoke result - :param invoke_result: invoke result - :return: - """ - model = None - prompt_messages: list[PromptMessage] = [] - full_text = "" - usage = None - for result in invoke_result: - text = result.delta.message.content - match text: - case str(): - full_text += text - case list(): - for i in text: - if i.data: - full_text += i.data - - if not model: - model = result.model - - if not prompt_messages: - prompt_messages = result.prompt_messages - - if not usage and result.delta.usage: - usage = result.delta.usage - - if not usage: - usage = LLMUsage.empty_usage() - - return full_text, usage - - @trace_span() - def _multiple_retrieve_thread( - self, - flask_app: Flask, - available_datasets: list[Dataset], - metadata_condition: MetadataFilteringCondition | None, - metadata_filter_document_ids: dict[str, list[str]] | None, - all_documents: list[Document], - tenant_id: str, - reranking_enable: bool, - reranking_mode: str, - reranking_model: RerankingModelDict | None, - weights: WeightsDict | None, - top_k: int, - score_threshold: float, - query: str | None, - attachment_id: str | None, - dataset_count: int, - cancel_event: threading.Event | None = None, - ) -> None: - try: - with flask_app.app_context(): - threads = [] - retrieval_thread_exceptions: list[Exception] = [] - all_documents_item: list[Document] = [] - index_type = None - for dataset in available_datasets: - # Check for cancellation signal - if cancel_event and cancel_event.is_set(): - break - index_type = dataset.indexing_technique - document_ids_filter = None - if dataset.provider != "external": - if metadata_condition and not metadata_filter_document_ids: - continue - if metadata_filter_document_ids: - document_ids = metadata_filter_document_ids.get(dataset.id, []) - if document_ids: - document_ids_filter = document_ids - else: - continue - retrieval_thread = threading.Thread( - target=propagate_context(self._run_retriever_thread_safely), - kwargs={ - "flask_app": flask_app, - "dataset_id": dataset.id, - "query": query, - "top_k": top_k, - "all_documents": all_documents_item, - "document_ids_filter": document_ids_filter, - "metadata_condition": metadata_condition, - "attachment_ids": [attachment_id] if attachment_id else None, - "cancel_event": cancel_event, - "thread_exceptions": retrieval_thread_exceptions, - "skip_on_error": True, - }, - ) - threads.append(retrieval_thread) - retrieval_thread.start() - - # Poll threads with short timeout to respond quickly to cancellation - while any(t.is_alive() for t in threads): - for thread in threads: - thread.join(timeout=0.1) - if cancel_event and cancel_event.is_set(): - break - if cancel_event and cancel_event.is_set(): - break - - if retrieval_thread_exceptions: - raise retrieval_thread_exceptions[0] - - # Skip second reranking when there is only one dataset - if reranking_enable and dataset_count > 1: - # do rerank for searched documents - with session_factory.create_session() as session: - data_post_processor = DataPostProcessor( - tenant_id, - reranking_mode, - reranking_model, - weights, - False, - load_upload=partial( - SQLAlchemyKnowledgeUploadRepository( - session_factory=sessionmaker(bind=session.get_bind()) - ).get_by_id, - workspace_id=tenant_id, - ), - ) - if query: - all_documents_item = data_post_processor.invoke( - query=query, - documents=all_documents_item, - score_threshold=score_threshold, - top_n=top_k, - query_type=QueryType.TEXT_QUERY, - ) - if attachment_id: - all_documents_item = data_post_processor.invoke( - documents=all_documents_item, - score_threshold=score_threshold, - top_n=top_k, - query_type=QueryType.IMAGE_QUERY, - query=attachment_id, - ) - else: - if index_type == IndexTechniqueType.ECONOMY: - if not query: - all_documents_item = [] - else: - all_documents_item = self.calculate_keyword_score(query, all_documents_item, top_k) - elif index_type == IndexTechniqueType.HIGH_QUALITY: - all_documents_item = self.calculate_vector_score(all_documents_item, top_k, score_threshold) - else: - all_documents_item = all_documents_item[:top_k] if top_k else all_documents_item - if all_documents_item: - all_documents.extend(all_documents_item) - except Exception: - raise - - def _multiple_retrieve_thread_safely( - self, - *, - flask_app: Flask, - available_datasets: list[Dataset], - metadata_condition: MetadataFilteringCondition | None, - metadata_filter_document_ids: dict[str, list[str]] | None, - all_documents: list[Document], - tenant_id: str, - reranking_enable: bool, - reranking_mode: str, - reranking_model: RerankingModelDict | None, - weights: WeightsDict | None, - top_k: int, - score_threshold: float, - query: str | None, - attachment_id: str | None, - dataset_count: int, - cancel_event: threading.Event | None = None, - thread_exceptions: list[Exception] | None = None, - ) -> None: - """Collect errors only after they pass through the traced multi-retrieval method.""" - try: - self._multiple_retrieve_thread( - flask_app=flask_app, - available_datasets=available_datasets, - metadata_condition=metadata_condition, - metadata_filter_document_ids=metadata_filter_document_ids, - all_documents=all_documents, - tenant_id=tenant_id, - reranking_enable=reranking_enable, - reranking_mode=reranking_mode, - reranking_model=reranking_model, - weights=weights, - top_k=top_k, - score_threshold=score_threshold, - query=query, - attachment_id=attachment_id, - dataset_count=dataset_count, - cancel_event=cancel_event, - ) - except Exception as exc: - if cancel_event: - cancel_event.set() - if thread_exceptions is not None: - thread_exceptions.append(exc) - - def _get_available_datasets(self, tenant_id: str, dataset_ids: list[str]) -> list[Dataset]: - with session_factory.create_session() as session: - subquery = ( - select(DocumentModel.dataset_id, func.count(DocumentModel.id).label("available_document_count")) - .where( - DocumentModel.indexing_status == "completed", - DocumentModel.enabled == True, - DocumentModel.archived == False, - DocumentModel.dataset_id.in_(dataset_ids), - ) - .group_by(DocumentModel.dataset_id) - .having(func.count(DocumentModel.id) > 0) - .subquery() - ) - - results = session.scalars( - select(Dataset) - .outerjoin(subquery, Dataset.id == subquery.c.dataset_id) - .where( - Dataset.tenant_id == tenant_id, - Dataset.id.in_(dataset_ids), - (subquery.c.available_document_count > 0) | (Dataset.provider == "external"), - ) - ).all() - - available_datasets = [] - for dataset in results: - if not dataset: - continue - available_datasets.append(dataset) - return available_datasets - - def _check_knowledge_rate_limit(self, tenant_id: str): - knowledge_rate_limit = FeatureService.get_knowledge_rate_limit(tenant_id) - if knowledge_rate_limit.enabled: - current_time = int(time.time() * 1000) - key = f"rate_limit_{tenant_id}" - redis_client.zadd(key, {current_time: current_time}) - redis_client.zremrangebyscore(key, 0, current_time - 60000) - request_count = redis_client.zcard(key) - if request_count > knowledge_rate_limit.limit: - # The rate-limit exception is raised after this block, so commit the audit row - # explicitly instead of relying on the Session context, which only closes it. - with session_factory.create_session() as session: - rate_limit_log = RateLimitLog( - tenant_id=tenant_id, - subscription_plan=knowledge_rate_limit.subscription_plan, - operation="knowledge", - ) - session.add(rate_limit_log) - session.commit() - raise exc.RateLimitExceededError( - "you have reached the knowledge base request rate limit of your subscription." - ) diff --git a/api/core/repositories/human_input_repository.py b/api/core/repositories/human_input_repository.py deleted file mode 100644 index f58dccc75ef553..00000000000000 --- a/api/core/repositories/human_input_repository.py +++ /dev/null @@ -1,676 +0,0 @@ -import dataclasses -import json -from collections.abc import Mapping, Sequence -from datetime import datetime -from typing import Any, Protocol, override - -from sqlalchemy import select -from sqlalchemy.orm import Session, selectinload - -from core.db.session_factory import session_factory -from core.workflow.human_input_adapter import is_human_input_webapp_enabled -from enums.human_input import HumanInputFormKind, HumanInputFormStatus, RecipientType -from libs.datetime_utils import naive_utc_now -from libs.uuid_utils import uuidv7 -from models.account import Account, TenantAccountJoin -from models.human_input import ( - BackstageRecipientPayload, - ConsoleDeliveryPayload, - ConsoleRecipientPayload, - DeliveryMethodType, - EmailExternalRecipientPayload, - EmailMemberRecipientPayload, - HumanInputDelivery, - HumanInputForm, - HumanInputFormRecipient, - StandaloneWebAppRecipientPayload, -) -from models.human_input_delivery import ( - BoundRecipient, - DeliveryChannelConfig, - EmailDeliveryMethod, - EmailRecipients, - ExternalRecipient, - InteractiveSurfaceDeliveryMethod, -) -from models.human_input_entities import FormDefinition, HumanInputNodeData - - -@dataclasses.dataclass(frozen=True) -class _DeliveryAndRecipients: - delivery: HumanInputDelivery - recipients: Sequence[HumanInputFormRecipient] - - -@dataclasses.dataclass(frozen=True) -class _WorkspaceMemberInfo: - user_id: str - email: str - - -class FormNotFoundError(Exception): - pass - - -@dataclasses.dataclass -class FormCreateParams: - workflow_execution_id: str | None - node_id: str - form_config: HumanInputNodeData - rendered_content: str - delivery_methods: Sequence[DeliveryChannelConfig] - display_in_ui: bool - resolved_default_values: Mapping[str, Any] - form_kind: HumanInputFormKind = HumanInputFormKind.RUNTIME - # ENG-635: the conversation this form belongs to. Set together with - # workflow_execution_id for chatflow runs; set alone (workflow_execution_id None) - # for Agent v2 chat ask_human forms, which have no workflow run. - conversation_id: str | None = None - form_id: str | None = None - - -class HumanInputFormRecipientEntity(Protocol): - @property - def id(self) -> str: ... - - @property - def token(self) -> str: ... - - -class HumanInputFormEntity(Protocol): - @property - def id(self) -> str: ... - - @property - def submission_token(self) -> str | None: ... - - @property - def recipients(self) -> list[HumanInputFormRecipientEntity]: ... - - @property - def rendered_content(self) -> str: ... - - @property - def selected_action_id(self) -> str | None: ... - - @property - def created_at(self) -> datetime: ... - - @property - def submitted_data(self) -> Mapping[str, Any] | None: ... - - @property - def submitted(self) -> bool: ... - - @property - def status(self) -> HumanInputFormStatus: ... - - @property - def expiration_time(self) -> datetime: ... - - -class HumanInputFormRepository(Protocol): - def get_form(self, node_id: str, *, form_id: str | None = None) -> HumanInputFormEntity | None: ... - - def create_form(self, params: FormCreateParams) -> HumanInputFormEntity: ... - - -class _HumanInputFormRecipientEntityImpl(HumanInputFormRecipientEntity): - def __init__(self, recipient_model: HumanInputFormRecipient): - self._recipient_model = recipient_model - - @property - @override - def id(self) -> str: - return self._recipient_model.id - - @property - @override - def token(self) -> str: - if self._recipient_model.access_token is None: - raise AssertionError(f"access_token should not be None for recipient {self._recipient_model.id}") - return self._recipient_model.access_token - - -class _HumanInputFormEntityImpl(HumanInputFormEntity): - def __init__(self, form_model: HumanInputForm, recipient_models: Sequence[HumanInputFormRecipient]): - self._form_model = form_model - self._recipients = [_HumanInputFormRecipientEntityImpl(recipient) for recipient in recipient_models] - self._interactive_surface_recipient = next( - ( - recipient - for recipient in recipient_models - if recipient.recipient_type == RecipientType.STANDALONE_WEB_APP - ), - None, - ) - self._console_recipient = next( - (recipient for recipient in recipient_models if recipient.recipient_type == RecipientType.CONSOLE), - None, - ) - self._submitted_data: Mapping[str, Any] | None = ( - json.loads(form_model.submitted_data) if form_model.submitted_data is not None else None - ) - - @property - @override - def id(self) -> str: - return self._form_model.id - - @property - @override - def submission_token(self) -> str | None: - if self._console_recipient is not None: - return self._console_recipient.access_token - if self._interactive_surface_recipient is None: - return None - return self._interactive_surface_recipient.access_token - - @property - @override - def recipients(self) -> list[HumanInputFormRecipientEntity]: - return list(self._recipients) - - @property - @override - def rendered_content(self) -> str: - return self._form_model.rendered_content - - @property - @override - def selected_action_id(self) -> str | None: - return self._form_model.selected_action_id - - @property - @override - def created_at(self) -> datetime: - return self._form_model.created_at - - @property - @override - def submitted_data(self) -> Mapping[str, Any] | None: - return self._submitted_data - - @property - @override - def submitted(self) -> bool: - return self._form_model.submitted_at is not None - - @property - @override - def status(self) -> HumanInputFormStatus: - return self._form_model.status - - @property - @override - def expiration_time(self) -> datetime: - return self._form_model.expiration_time - - -@dataclasses.dataclass(frozen=True) -class HumanInputFormRecord: - form_id: str - workflow_run_id: str | None - node_id: str - tenant_id: str - app_id: str - form_kind: HumanInputFormKind - definition: FormDefinition - rendered_content: str - created_at: datetime - expiration_time: datetime - status: HumanInputFormStatus - selected_action_id: str | None - submitted_data: Mapping[str, Any] | None - submitted_at: datetime | None - submission_user_id: str | None - submission_end_user_id: str | None - completed_by_recipient_id: str | None - recipient_id: str | None - recipient_type: RecipientType | None - access_token: str | None - # ENG-635: Agent v2 chat owner (NULL for workflow-owned forms). Trailing + - # defaulted so existing record constructions stay source-compatible. - conversation_id: str | None = None - - @property - def submitted(self) -> bool: - return self.submitted_at is not None - - @classmethod - def from_models( - cls, form_model: HumanInputForm, recipient_model: HumanInputFormRecipient | None - ) -> "HumanInputFormRecord": - definition_payload = json.loads(form_model.form_definition) - if "expiration_time" not in definition_payload: - definition_payload["expiration_time"] = form_model.expiration_time - return cls( - form_id=form_model.id, - workflow_run_id=form_model.workflow_run_id, - conversation_id=form_model.conversation_id, - node_id=form_model.node_id, - tenant_id=form_model.tenant_id, - app_id=form_model.app_id, - form_kind=form_model.form_kind, - definition=FormDefinition.model_validate(definition_payload), - rendered_content=form_model.rendered_content, - created_at=form_model.created_at, - expiration_time=form_model.expiration_time, - status=form_model.status, - selected_action_id=form_model.selected_action_id, - submitted_data=json.loads(form_model.submitted_data) if form_model.submitted_data else None, - submitted_at=form_model.submitted_at, - submission_user_id=form_model.submission_user_id, - submission_end_user_id=form_model.submission_end_user_id, - completed_by_recipient_id=form_model.completed_by_recipient_id, - recipient_id=recipient_model.id if recipient_model else None, - recipient_type=recipient_model.recipient_type if recipient_model else None, - access_token=recipient_model.access_token if recipient_model else None, - ) - - -class _InvalidTimeoutStatusError(ValueError): - pass - - -class HumanInputFormRepositoryImpl: - def __init__( - self, - *, - tenant_id: str, - app_id: str | None = None, - workflow_execution_id: str | None = None, - invoke_source: str | None = None, - submission_actor_id: str | None = None, - ) -> None: - self._tenant_id = tenant_id - self._app_id = app_id - self._workflow_execution_id = workflow_execution_id - self._invoke_source = invoke_source - self._submission_actor_id = submission_actor_id - - def _delivery_method_to_model( - self, - session: Session, - form_id: str, - delivery_method: DeliveryChannelConfig, - ) -> _DeliveryAndRecipients: - delivery_id = str(uuidv7()) - delivery_model = HumanInputDelivery( - id=delivery_id, - form_id=form_id, - delivery_method_type=delivery_method.type, - delivery_config_id=str(delivery_method.id), - channel_payload=delivery_method.model_dump_json(), - ) - recipients: list[HumanInputFormRecipient] = [] - match delivery_method: - case InteractiveSurfaceDeliveryMethod(): - recipient_model = HumanInputFormRecipient( - form_id=form_id, - delivery_id=delivery_id, - recipient_type=RecipientType.STANDALONE_WEB_APP, - recipient_payload=StandaloneWebAppRecipientPayload().model_dump_json(), - ) - recipients.append(recipient_model) - case EmailDeliveryMethod(): - email_recipients_config = delivery_method.config.recipients - recipients.extend( - self._build_email_recipients( - session=session, - form_id=form_id, - delivery_id=delivery_id, - recipients_config=email_recipients_config, - ) - ) - - return _DeliveryAndRecipients(delivery=delivery_model, recipients=recipients) - - def _build_email_recipients( - self, - session: Session, - form_id: str, - delivery_id: str, - recipients_config: EmailRecipients, - ) -> list[HumanInputFormRecipient]: - bound_reference_ids = [ - recipient.reference_id for recipient in recipients_config.items if isinstance(recipient, BoundRecipient) - ] - external_emails = [ - recipient.email for recipient in recipients_config.items if isinstance(recipient, ExternalRecipient) - ] - if recipients_config.include_bound_group: - members = self._query_all_workspace_members(session=session) - else: - members = self._query_workspace_members_by_ids(session=session, restrict_to_user_ids=bound_reference_ids) - - return self._create_email_recipients_from_resolved( - form_id=form_id, - delivery_id=delivery_id, - members=members, - external_emails=external_emails, - ) - - @staticmethod - def _create_email_recipients_from_resolved( - *, - form_id: str, - delivery_id: str, - members: Sequence[_WorkspaceMemberInfo], - external_emails: Sequence[str], - ) -> list[HumanInputFormRecipient]: - recipient_models: list[HumanInputFormRecipient] = [] - seen_emails: set[str] = set() - - for member in members: - if not member.email: - continue - if member.email in seen_emails: - continue - seen_emails.add(member.email) - payload = EmailMemberRecipientPayload(user_id=member.user_id, email=member.email) - recipient_models.append( - HumanInputFormRecipient.new( - form_id=form_id, - delivery_id=delivery_id, - payload=payload, - ) - ) - - for email in external_emails: - if not email: - continue - if email in seen_emails: - continue - seen_emails.add(email) - recipient_models.append( - HumanInputFormRecipient.new( - form_id=form_id, - delivery_id=delivery_id, - payload=EmailExternalRecipientPayload(email=email), - ) - ) - - return recipient_models - - def _query_all_workspace_members( - self, - session: Session, - ) -> list[_WorkspaceMemberInfo]: - stmt = ( - select(Account.id, Account.email) - .join(TenantAccountJoin, TenantAccountJoin.account_id == Account.id) - .where(TenantAccountJoin.tenant_id == self._tenant_id) - ) - rows = session.execute(stmt).all() - return [_WorkspaceMemberInfo(user_id=account_id, email=email) for account_id, email in rows] - - def _query_workspace_members_by_ids( - self, - session: Session, - restrict_to_user_ids: Sequence[str], - ) -> list[_WorkspaceMemberInfo]: - unique_ids = {user_id for user_id in restrict_to_user_ids if user_id} - if not unique_ids: - return [] - - stmt = ( - select(Account.id, Account.email) - .join(TenantAccountJoin, TenantAccountJoin.account_id == Account.id) - .where(TenantAccountJoin.tenant_id == self._tenant_id) - ) - stmt = stmt.where(Account.id.in_(unique_ids)) - - rows = session.execute(stmt).all() - return [_WorkspaceMemberInfo(user_id=account_id, email=email) for account_id, email in rows] - - def _should_create_console_recipient( - self, - *, - form_config: HumanInputNodeData, - form_kind: HumanInputFormKind, - ) -> bool: - if form_kind != HumanInputFormKind.RUNTIME: - return False - if self._invoke_source == "debugger": - return True - if self._invoke_source == "explore": - return is_human_input_webapp_enabled(form_config) - return False - - def _should_create_backstage_recipient(self, *, form_kind: HumanInputFormKind) -> bool: - return form_kind == HumanInputFormKind.RUNTIME and ( - self._invoke_source is not None or self._submission_actor_id is not None - ) - - def create_form(self, params: FormCreateParams) -> HumanInputFormEntity: - form_config: HumanInputNodeData = params.form_config - app_id = self._app_id - if not app_id: - raise ValueError("app_id is required to create a human input form") - workflow_execution_id = params.workflow_execution_id or self._workflow_execution_id - # A RUNTIME form must be owned by at least one of: a workflow run (workflow / - # Human-Input / agent node) or a conversation turn (ENG-635: Agent v2 chat - # ask_human; chatflow runs set both — workflow_run_id and conversation_id). - if ( - params.form_kind == HumanInputFormKind.RUNTIME - and workflow_execution_id is None - and params.conversation_id is None - ): - raise ValueError("a runtime human input form requires a workflow_execution_id or conversation_id") - - with session_factory.create_session() as session, session.begin(): - form_id = params.form_id or str(uuidv7()) - start_time = naive_utc_now() - node_expiration = form_config.expiration_time(start_time) - form_definition = FormDefinition( - form_content=form_config.form_content, - inputs=form_config.inputs, - user_actions=form_config.user_actions, - rendered_content=params.rendered_content, - expiration_time=node_expiration, - default_values=dict(params.resolved_default_values), - display_in_ui=params.display_in_ui, - node_title=form_config.title, - ) - form_model = HumanInputForm( - id=form_id, - tenant_id=self._tenant_id, - app_id=app_id, - workflow_run_id=workflow_execution_id, - conversation_id=params.conversation_id, - form_kind=params.form_kind, - node_id=params.node_id, - form_definition=form_definition.model_dump_json(), - rendered_content=params.rendered_content, - expiration_time=node_expiration, - created_at=start_time, - ) - session.add(form_model) - recipient_models: list[HumanInputFormRecipient] = [] - for delivery in params.delivery_methods: - delivery_and_recipients = self._delivery_method_to_model( - session=session, - form_id=form_id, - delivery_method=delivery, - ) - session.add(delivery_and_recipients.delivery) - session.add_all(delivery_and_recipients.recipients) - recipient_models.extend(delivery_and_recipients.recipients) - if self._should_create_console_recipient(form_config=form_config, form_kind=params.form_kind) and not any( - recipient.recipient_type == RecipientType.CONSOLE for recipient in recipient_models - ): - console_delivery_id = str(uuidv7()) - console_delivery = HumanInputDelivery( - id=console_delivery_id, - form_id=form_id, - delivery_method_type=DeliveryMethodType.WEBAPP, - delivery_config_id=None, - channel_payload=ConsoleDeliveryPayload().model_dump_json(), - ) - console_recipient = HumanInputFormRecipient( - form_id=form_id, - delivery_id=console_delivery_id, - recipient_type=RecipientType.CONSOLE, - recipient_payload=ConsoleRecipientPayload( - account_id=self._submission_actor_id, - ).model_dump_json(), - ) - session.add(console_delivery) - session.add(console_recipient) - recipient_models.append(console_recipient) - if self._should_create_backstage_recipient(form_kind=params.form_kind) and not any( - recipient.recipient_type == RecipientType.BACKSTAGE for recipient in recipient_models - ): - backstage_delivery_id = str(uuidv7()) - backstage_delivery = HumanInputDelivery( - id=backstage_delivery_id, - form_id=form_id, - delivery_method_type=DeliveryMethodType.WEBAPP, - delivery_config_id=None, - channel_payload=ConsoleDeliveryPayload().model_dump_json(), - ) - backstage_recipient = HumanInputFormRecipient( - form_id=form_id, - delivery_id=backstage_delivery_id, - recipient_type=RecipientType.BACKSTAGE, - recipient_payload=BackstageRecipientPayload( - account_id=self._submission_actor_id, - ).model_dump_json(), - ) - session.add(backstage_delivery) - session.add(backstage_recipient) - recipient_models.append(backstage_recipient) - session.flush() - - return _HumanInputFormEntityImpl(form_model=form_model, recipient_models=recipient_models) - - def get_form(self, node_id: str, *, form_id: str | None = None) -> HumanInputFormEntity | None: - if self._workflow_execution_id is None: - raise ValueError("workflow_execution_id is required to load runtime human input forms") - - form_query = select(HumanInputForm).where( - HumanInputForm.workflow_run_id == self._workflow_execution_id, - HumanInputForm.node_id == node_id, - HumanInputForm.tenant_id == self._tenant_id, - ) - if form_id is not None: - form_query = form_query.where(HumanInputForm.id == form_id) - with session_factory.create_session() as session: - form_model: HumanInputForm | None = session.scalars(form_query).first() - if form_model is None: - return None - - recipient_query = select(HumanInputFormRecipient).where(HumanInputFormRecipient.form_id == form_model.id) - recipient_models = session.scalars(recipient_query).all() - return _HumanInputFormEntityImpl(form_model=form_model, recipient_models=recipient_models) - - -class HumanInputFormSubmissionRepository: - """Repository for fetching and submitting human input forms.""" - - def get_by_token(self, form_token: str) -> HumanInputFormRecord | None: - query = ( - select(HumanInputFormRecipient) - .options(selectinload(HumanInputFormRecipient.form)) - .where(HumanInputFormRecipient.access_token == form_token) - ) - with session_factory.create_session() as session: - recipient_model = session.scalars(query).first() - if recipient_model is None or recipient_model.form is None: - return None - return HumanInputFormRecord.from_models(recipient_model.form, recipient_model) - - def get_by_form_id(self, form_id: str) -> HumanInputFormRecord | None: - with session_factory.create_session() as session: - form_model = session.get(HumanInputForm, form_id) - if form_model is None: - return None - return HumanInputFormRecord.from_models(form_model, None) - - def get_by_form_id_and_recipient_type( - self, - form_id: str, - recipient_type: RecipientType, - ) -> HumanInputFormRecord | None: - query = ( - select(HumanInputFormRecipient) - .options(selectinload(HumanInputFormRecipient.form)) - .where( - HumanInputFormRecipient.form_id == form_id, - HumanInputFormRecipient.recipient_type == recipient_type, - ) - ) - with session_factory.create_session() as session: - recipient_model = session.scalars(query).first() - if recipient_model is None or recipient_model.form is None: - return None - return HumanInputFormRecord.from_models(recipient_model.form, recipient_model) - - def mark_submitted( - self, - *, - form_id: str, - recipient_id: str | None, - selected_action_id: str, - form_data: Mapping[str, Any], - submission_user_id: str | None, - submission_end_user_id: str | None, - ) -> HumanInputFormRecord: - with session_factory.create_session() as session, session.begin(): - form_model = session.get(HumanInputForm, form_id) - if form_model is None: - raise FormNotFoundError(f"form not found, id={form_id}") - - recipient_model = session.get(HumanInputFormRecipient, recipient_id) if recipient_id else None - - form_model.selected_action_id = selected_action_id - form_model.submitted_data = json.dumps(form_data) - form_model.submitted_at = naive_utc_now() - form_model.status = HumanInputFormStatus.SUBMITTED - form_model.submission_user_id = submission_user_id - form_model.submission_end_user_id = submission_end_user_id - form_model.completed_by_recipient_id = recipient_id - - session.add(form_model) - session.flush() - session.refresh(form_model) - if recipient_model is not None: - session.refresh(recipient_model) - - return HumanInputFormRecord.from_models(form_model, recipient_model) - - def mark_timeout( - self, - *, - form_id: str, - timeout_status: HumanInputFormStatus, - reason: str | None = None, - ) -> HumanInputFormRecord: - with session_factory.create_session() as session, session.begin(): - form_model = session.get(HumanInputForm, form_id) - if form_model is None: - raise FormNotFoundError(f"form not found, id={form_id}") - - if timeout_status not in {HumanInputFormStatus.TIMEOUT, HumanInputFormStatus.EXPIRED}: - raise _InvalidTimeoutStatusError(f"invalid timeout status: {timeout_status}") - - # already handled or submitted - if form_model.status in {HumanInputFormStatus.TIMEOUT, HumanInputFormStatus.EXPIRED}: - return HumanInputFormRecord.from_models(form_model, None) - - if form_model.submitted_at is not None or form_model.status == HumanInputFormStatus.SUBMITTED: - raise FormNotFoundError(f"form already submitted, id={form_id}") - - form_model.status = timeout_status - form_model.selected_action_id = None - form_model.submitted_data = None - form_model.submission_user_id = None - form_model.submission_end_user_id = None - form_model.completed_by_recipient_id = None - # Reason is recorded in status/error downstream; not stored on form. - session.add(form_model) - session.flush() - session.refresh(form_model) - - return HumanInputFormRecord.from_models(form_model, None) diff --git a/api/core/tools/custom_tool/provider.py b/api/core/tools/custom_tool/provider.py deleted file mode 100644 index 41cd038cc5fb33..00000000000000 --- a/api/core/tools/custom_tool/provider.py +++ /dev/null @@ -1,218 +0,0 @@ -from __future__ import annotations - -from typing import override - -from pydantic import Field -from sqlalchemy import select -from sqlalchemy.orm import Session - -from core.entities.provider_entities import ProviderConfig, ProviderConfigType -from core.tools.__base.tool_provider import ToolProviderController -from core.tools.__base.tool_runtime import ToolRuntime -from core.tools.custom_tool.tool import ApiTool -from core.tools.entities.common_entities import I18nObject -from core.tools.entities.tool_bundle import ApiToolBundle -from core.tools.entities.tool_entities import ( - ApiProviderAuthType, - ToolDescription, - ToolEntity, - ToolIdentity, - ToolProviderEntity, - ToolProviderIdentity, - ToolProviderType, -) -from extensions.ext_database import db -from models.tools import ApiToolProvider - - -class ApiToolProviderController(ToolProviderController[ToolProviderEntity, ApiTool]): - provider_id: str - tenant_id: str - tools: list[ApiTool] = Field(default_factory=list) - - def __init__(self, entity: ToolProviderEntity, provider_id: str, tenant_id: str): - super().__init__(entity) - self.provider_id = provider_id - self.tenant_id = tenant_id - self.tools = [] - - @classmethod - def from_db( - cls, db_provider: ApiToolProvider, auth_type: ApiProviderAuthType, *, session: Session - ) -> ApiToolProviderController: - credentials_schema = [ - ProviderConfig( - name="auth_type", - required=True, - type=ProviderConfigType.SELECT, - options=[ - ProviderConfig.Option(value="none", label=I18nObject(en_US="None", zh_Hans="无")), - ProviderConfig.Option(value="api_key_header", label=I18nObject(en_US="Header", zh_Hans="请求头")), - ProviderConfig.Option( - value="api_key_query", label=I18nObject(en_US="Query Param", zh_Hans="查询参数") - ), - ], - default="none", - help=I18nObject(en_US="The auth type of the api provider", zh_Hans="api provider 的认证类型"), - ) - ] - if auth_type == ApiProviderAuthType.API_KEY_HEADER: - credentials_schema = [ - *credentials_schema, - ProviderConfig( - name="api_key_header", - required=False, - default="Authorization", - type=ProviderConfigType.TEXT_INPUT, - help=I18nObject(en_US="The header name of the api key", zh_Hans="携带 api key 的 header 名称"), - ), - ProviderConfig( - name="api_key_value", - required=True, - type=ProviderConfigType.SECRET_INPUT, - help=I18nObject(en_US="The api key", zh_Hans="api key 的值"), - ), - ProviderConfig( - name="api_key_header_prefix", - required=False, - default="basic", - type=ProviderConfigType.SELECT, - help=I18nObject(en_US="The prefix of the api key header", zh_Hans="api key header 的前缀"), - options=[ - ProviderConfig.Option(value="basic", label=I18nObject(en_US="Basic", zh_Hans="Basic")), - ProviderConfig.Option(value="bearer", label=I18nObject(en_US="Bearer", zh_Hans="Bearer")), - ProviderConfig.Option(value="custom", label=I18nObject(en_US="Custom", zh_Hans="Custom")), - ], - ), - ] - elif auth_type == ApiProviderAuthType.API_KEY_QUERY: - credentials_schema = [ - *credentials_schema, - ProviderConfig( - name="api_key_query_param", - required=False, - default="key", - type=ProviderConfigType.TEXT_INPUT, - help=I18nObject( - en_US="The query parameter name of the api key", zh_Hans="携带 api key 的查询参数名称" - ), - ), - ProviderConfig( - name="api_key_value", - required=True, - type=ProviderConfigType.SECRET_INPUT, - help=I18nObject(en_US="The api key", zh_Hans="api key 的值"), - ), - ] - elif auth_type == ApiProviderAuthType.NONE: - pass - - user = db_provider.user(session=session) - user_name = user.name if user else "" - - return ApiToolProviderController( - entity=ToolProviderEntity( - identity=ToolProviderIdentity( - author=user_name, - name=db_provider.name, - label=I18nObject(en_US=db_provider.name, zh_Hans=db_provider.name), - description=I18nObject(en_US=db_provider.description, zh_Hans=db_provider.description), - icon=db_provider.icon, - ), - credentials_schema=credentials_schema, - plugin_id=None, - ), - provider_id=db_provider.id or "", - tenant_id=db_provider.tenant_id or "", - ) - - @property - @override - def provider_type(self) -> ToolProviderType: - return ToolProviderType.API - - def _parse_tool_bundle(self, tool_bundle: ApiToolBundle) -> ApiTool: - """ - parse tool bundle to tool - - :param tool_bundle: the tool bundle - :return: the tool - """ - return ApiTool( - api_bundle=tool_bundle, - provider_id=self.provider_id, - entity=ToolEntity( - identity=ToolIdentity( - author=tool_bundle.author, - name=tool_bundle.operation_id or "default_tool", - label=I18nObject( - en_US=tool_bundle.operation_id or "default_tool", - zh_Hans=tool_bundle.operation_id or "default_tool", - ), - icon=self.entity.identity.icon, - provider=self.provider_id, - ), - description=ToolDescription( - human=I18nObject(en_US=tool_bundle.summary or "", zh_Hans=tool_bundle.summary or ""), - llm=tool_bundle.summary or "", - ), - parameters=tool_bundle.parameters or [], - ), - runtime=ToolRuntime(tenant_id=self.tenant_id), - ) - - def load_bundled_tools(self, tools: list[ApiToolBundle]): - """ - load bundled tools - - :param tools: the bundled tools - :return: the tools - """ - self.tools = [self._parse_tool_bundle(tool) for tool in tools] - - return self.tools - - def get_tools(self, tenant_id: str) -> list[ApiTool]: - """ - fetch tools from database - - :param tenant_id: the tenant id - :return: the tools - """ - if len(self.tools) > 0: - return self.tools - - tools: list[ApiTool] = [] - - # get tenant api providers - db_providers = db.session.scalars( - select(ApiToolProvider).where( - ApiToolProvider.tenant_id == tenant_id, ApiToolProvider.name == self.entity.identity.name - ) - ).all() - - if db_providers and len(db_providers) != 0: - for db_provider in db_providers: - for tool in db_provider.tools: - assistant_tool = self._parse_tool_bundle(tool) - tools.append(assistant_tool) - - self.tools = tools - return tools - - @override - def get_tool(self, tool_name: str) -> ApiTool: - """ - get tool by name - - :param tool_name: the name of the tool - :return: the tool - """ - if self.tools is None: - self.get_tools(self.tenant_id) - - for tool in self.tools: - if tool.entity.identity.name == tool_name: - return tool - - raise ValueError(f"tool {tool_name} not found") diff --git a/api/core/tools/tool_engine.py b/api/core/tools/tool_engine.py deleted file mode 100644 index 9c747acc2f883c..00000000000000 --- a/api/core/tools/tool_engine.py +++ /dev/null @@ -1,396 +0,0 @@ -import contextlib -import json -import logging -from collections.abc import Generator, Iterable -from copy import deepcopy -from datetime import UTC, datetime -from mimetypes import guess_type -from typing import Any, Union, cast - -from sqlalchemy.orm import Session, sessionmaker -from yarl import URL - -from core.app.entities.app_invoke_entities import InvokeFrom -from core.callback_handler.agent_tool_callback_handler import DifyAgentCallbackHandler -from core.callback_handler.workflow_tool_callback_handler import DifyWorkflowCallbackHandler -from core.ops.ops_trace_manager import TraceQueueManager -from core.tools.__base.tool import Tool -from core.tools.entities.tool_entities import ( - ToolInvokeMessage, - ToolInvokeMessageBinary, - ToolInvokeMeta, - ToolParameter, -) -from core.tools.errors import ( - ToolEngineInvokeError, - ToolInvokeError, - ToolNotFoundError, - ToolNotSupportedError, - ToolParameterValidationError, - ToolProviderCredentialValidationError, - ToolProviderNotFoundError, -) -from core.tools.utils.message_transformer import ToolFileMessageTransformer, safe_json_value -from core.tools.workflow_as_tool.tool import WorkflowTool -from extensions.ext_database import db -from graphon.file import FileTransferMethod, FileType -from models.enums import CreatorUserRole, MessageFileBelongsTo -from models.model import Message, MessageFile - -logger = logging.getLogger(__name__) - - -class ToolEngine: - """ - Tool runtime engine take care of the tool executions. - """ - - @staticmethod - def agent_invoke( - session: Session, - tool: Tool, - tool_parameters: Union[str, dict[str, Any]], - user_id: str, - tenant_id: str, - message: Message, - invoke_from: InvokeFrom, - agent_tool_callback: DifyAgentCallbackHandler, - trace_manager: TraceQueueManager | None = None, - conversation_id: str | None = None, - app_id: str | None = None, - message_id: str | None = None, - ) -> tuple[str, list[str], ToolInvokeMeta]: - """ - Agent invokes the tool with the given arguments. - """ - # check if arguments is a string - if isinstance(tool_parameters, str): - # check if this tool has only one parameter - parameters = [ - parameter - for parameter in tool.get_runtime_parameters() - if parameter.form == ToolParameter.ToolParameterForm.LLM - ] - if parameters and len(parameters) == 1: - tool_parameters = {parameters[0].name: tool_parameters} - else: - with contextlib.suppress(Exception): - tool_parameters = json.loads(tool_parameters) - if not isinstance(tool_parameters, dict): - raise ValueError(f"tool_parameters should be a dict, but got a string: {tool_parameters}") - - try: - # hit the callback handler - agent_tool_callback.on_tool_start(tool_name=tool.entity.identity.name, tool_inputs=tool_parameters) - - messages = ToolEngine._invoke(session, tool, tool_parameters, user_id, conversation_id, app_id, message_id) - invocation_meta_dict: dict[str, ToolInvokeMeta] = {} - - def message_callback( - invocation_meta_dict: dict[str, ToolInvokeMeta], - messages: Generator[ToolInvokeMessage | ToolInvokeMeta, None, None], - ): - for message in messages: - if isinstance(message, ToolInvokeMeta): - invocation_meta_dict["meta"] = message - else: - yield message - - messages = ToolFileMessageTransformer.transform_tool_invoke_messages( - messages=message_callback(invocation_meta_dict, messages), - user_id=user_id, - tenant_id=tenant_id, - conversation_id=message.conversation_id, - ) - - message_list = list(messages) - - # extract binary data from tool invoke message - binary_files = ToolEngine._extract_tool_response_binary_and_text(message_list) - # create message file - message_files = ToolEngine._create_message_files( - tool_messages=binary_files, agent_message=message, invoke_from=invoke_from, user_id=user_id - ) - - plain_text = ToolEngine.tool_response_to_str(message_list) - - meta = invocation_meta_dict["meta"] - - # hit the callback handler - agent_tool_callback.on_tool_end( - tool_name=tool.entity.identity.name, - tool_inputs=tool_parameters, - tool_outputs=plain_text, - message_id=message.id, - trace_manager=trace_manager, - ) - - # transform tool invoke message to get LLM friendly message - return plain_text, message_files, meta - except ToolProviderCredentialValidationError as e: - logger.error(e, exc_info=True) - error_response = "Please check your tool provider credentials" - agent_tool_callback.on_tool_error(e) - except (ToolNotFoundError, ToolNotSupportedError, ToolProviderNotFoundError) as e: - error_response = f"there is not a tool named {tool.entity.identity.name}" - logger.error(e, exc_info=True) - agent_tool_callback.on_tool_error(e) - except ToolParameterValidationError as e: - error_response = f"tool parameters validation error: {e}, please check your tool parameters" - agent_tool_callback.on_tool_error(e) - logger.error(e, exc_info=True) - except ToolInvokeError as e: - error_response = f"tool invoke error: {e}" - agent_tool_callback.on_tool_error(e) - logger.error(e, exc_info=True) - except ToolEngineInvokeError as e: - meta = e.meta - error_response = f"tool invoke error: {meta.error}" - agent_tool_callback.on_tool_error(e) - logger.error(e, exc_info=True) - return error_response, [], meta - except Exception as e: - error_response = f"unknown error: {e}" - agent_tool_callback.on_tool_error(e) - logger.error(e, exc_info=True) - - return error_response, [], ToolInvokeMeta.error_instance(error_response) - - @staticmethod - def generic_invoke( - session: Session, - tool: Tool, - tool_parameters: dict[str, Any], - user_id: str, - workflow_tool_callback: DifyWorkflowCallbackHandler, - workflow_call_depth: int, - conversation_id: str | None = None, - app_id: str | None = None, - message_id: str | None = None, - ) -> Generator[ToolInvokeMessage, None, None]: - """ - Workflow invokes the tool with the given arguments. - """ - try: - # hit the callback handler - workflow_tool_callback.on_tool_start(tool_name=tool.entity.identity.name, tool_inputs=tool_parameters) - - if isinstance(tool, WorkflowTool): - tool.workflow_call_depth = workflow_call_depth + 1 - - if tool.runtime and tool.runtime.runtime_parameters: - tool_parameters = {**tool.runtime.runtime_parameters, **tool_parameters} - - response = tool.invoke( - session=session, - user_id=user_id, - tool_parameters=tool_parameters, - conversation_id=conversation_id, - app_id=app_id, - message_id=message_id, - ) - - # hit the callback handler - response = workflow_tool_callback.on_tool_execution( - tool_name=tool.entity.identity.name, - tool_inputs=tool_parameters, - tool_outputs=response, - ) - - return response - except Exception as e: - workflow_tool_callback.on_tool_error(e) - raise e - - @staticmethod - def _invoke( - session: Session, - tool: Tool, - tool_parameters: dict[str, Any], - user_id: str, - conversation_id: str | None = None, - app_id: str | None = None, - message_id: str | None = None, - ) -> Generator[ToolInvokeMessage | ToolInvokeMeta, None, None]: - """ - Invoke the tool with the given arguments. - """ - started_at = datetime.now(UTC) - meta = ToolInvokeMeta( - time_cost=0.0, - error=None, - tool_config={ - "tool_name": tool.entity.identity.name, - "tool_provider": tool.entity.identity.provider, - "tool_provider_type": tool.tool_provider_type().value, - "tool_parameters": deepcopy(tool.runtime.runtime_parameters), - "tool_icon": tool.entity.identity.icon, - }, - ) - try: - yield from tool.invoke(session, user_id, tool_parameters, conversation_id, app_id, message_id) - except Exception as e: - meta.error = str(e) - raise ToolEngineInvokeError(meta) - finally: - ended_at = datetime.now(UTC) - meta.time_cost = (ended_at - started_at).total_seconds() - yield meta - - @staticmethod - def tool_response_to_str(tool_response: list[ToolInvokeMessage]) -> str: - """Convert tool invoke messages into the plain-text observation shown to the model/user.""" - parts: list[str] = [] - json_parts: list[str] = [] - - for response in tool_response: - if response.type == ToolInvokeMessage.MessageType.TEXT: - parts.append(cast(ToolInvokeMessage.TextMessage, response.message).text) - elif response.type == ToolInvokeMessage.MessageType.LINK: - parts.append( - f"result link: {cast(ToolInvokeMessage.TextMessage, response.message).text}." - + " please tell user to check it." - ) - elif response.type in {ToolInvokeMessage.MessageType.IMAGE_LINK, ToolInvokeMessage.MessageType.IMAGE}: - parts.append( - "image has been created and sent to user already, " - + "you do not need to create it, just tell the user to check it now." - ) - elif response.type == ToolInvokeMessage.MessageType.JSON: - json_message = cast(ToolInvokeMessage.JsonMessage, response.message) - if json_message.suppress_output: - continue - json_parts.append( - json.dumps( - safe_json_value(cast(ToolInvokeMessage.JsonMessage, response.message).json_object), - ensure_ascii=False, - ) - ) - elif response.type == ToolInvokeMessage.MessageType.VARIABLE: - continue - else: - parts.append(str(response.message)) - - # Add JSON parts, avoiding duplicates from text parts. - if json_parts: - - def normalize_json(text: str) -> str: - try: - return json.dumps(json.loads(text), ensure_ascii=False, sort_keys=True, separators=(",", ":")) - except (ValueError, RecursionError): - return text - - existing_parts = {normalize_json(p) for p in parts} - parts.extend(p for p in json_parts if normalize_json(p) not in existing_parts) - - return "".join(parts) - - @staticmethod - def _extract_tool_response_binary_and_text( - tool_response: list[ToolInvokeMessage], - ) -> Generator[ToolInvokeMessageBinary, None, None]: - """ - Extract tool response binary - """ - for response in tool_response: - if response.type in { - ToolInvokeMessage.MessageType.IMAGE_LINK, - ToolInvokeMessage.MessageType.IMAGE, - ToolInvokeMessage.MessageType.BINARY_LINK, - }: - mimetype = None - if not response.meta: - raise ValueError("missing meta data") - if response.meta.get("mime_type"): - mimetype = response.meta.get("mime_type") - else: - with contextlib.suppress(Exception): - url = URL(cast(ToolInvokeMessage.TextMessage, response.message).text) - extension = url.suffix - guess_type_result, _ = guess_type(f"a{extension}") - if guess_type_result: - mimetype = guess_type_result - - if not mimetype: - mimetype = ( - "image/jpeg" - if response.type != ToolInvokeMessage.MessageType.BINARY_LINK - else "application/octet-stream" - ) - - yield ToolInvokeMessageBinary( - mimetype=response.meta.get("mime_type", mimetype), - url=cast(ToolInvokeMessage.TextMessage, response.message).text, - ) - elif response.type == ToolInvokeMessage.MessageType.BLOB: - if not response.meta: - raise ValueError("missing meta data") - - yield ToolInvokeMessageBinary( - mimetype=response.meta.get("mime_type", "application/octet-stream"), - url=cast(ToolInvokeMessage.TextMessage, response.message).text, - ) - elif response.type == ToolInvokeMessage.MessageType.LINK: - # check if there is a mime type in meta - if response.meta and "mime_type" in response.meta: - yield ToolInvokeMessageBinary( - mimetype=response.meta.get("mime_type", "application/octet-stream") - if response.meta - else "application/octet-stream", - url=cast(ToolInvokeMessage.TextMessage, response.message).text, - ) - - @staticmethod - def _create_message_files( - tool_messages: Iterable[ToolInvokeMessageBinary], - agent_message: Message, - invoke_from: InvokeFrom, - user_id: str, - ) -> list[str]: - """ - Create message files produced by a tool call. - - Tool file persistence is a side effect of agent execution. Use an - independent transaction so this helper never commits or closes the - caller's request-scoped session. - - :return: message file ids - """ - result = [] - - with sessionmaker(bind=db.engine, expire_on_commit=False).begin() as session: - for message in tool_messages: - # extract tool file id from url - tool_file_id = message.url.split("/")[-1].split(".")[0] - message_file = MessageFile( - message_id=agent_message.id, - type=ToolEngine._resolve_tool_file_type(message), - transfer_method=FileTransferMethod.TOOL_FILE, - belongs_to=MessageFileBelongsTo.ASSISTANT, - url=message.url, - upload_file_id=tool_file_id, - created_by_role=( - CreatorUserRole.ACCOUNT - if invoke_from in {InvokeFrom.EXPLORE, InvokeFrom.DEBUGGER} - else CreatorUserRole.END_USER - ), - created_by=user_id, - ) - - session.add(message_file) - result.append(message_file.id) - - return result - - @staticmethod - def _resolve_tool_file_type(message: ToolInvokeMessageBinary) -> FileType: - if "image" in message.mimetype: - return FileType.IMAGE - elif "video" in message.mimetype: - return FileType.VIDEO - elif "audio" in message.mimetype: - return FileType.AUDIO - elif "text" in message.mimetype or "pdf" in message.mimetype: - return FileType.DOCUMENT - else: - return FileType.CUSTOM diff --git a/api/core/tools/tool_label_manager.py b/api/core/tools/tool_label_manager.py deleted file mode 100644 index 301ed3e1084d1c..00000000000000 --- a/api/core/tools/tool_label_manager.py +++ /dev/null @@ -1,135 +0,0 @@ -from sqlalchemy import delete, select -from sqlalchemy.orm import Session, sessionmaker - -from core.tools.__base.tool_provider import ToolProviderController -from core.tools.builtin_tool.provider import BuiltinToolProviderController -from core.tools.custom_tool.provider import ApiToolProviderController -from core.tools.entities.values import default_tool_label_name_list -from core.tools.workflow_as_tool.provider import WorkflowToolProviderController -from extensions.ext_database import db -from models.tools import ToolLabelBinding - - -class ToolLabelManager: - @classmethod - def filter_tool_labels(cls, tool_labels: list[str]) -> list[str]: - """ - Filter tool labels - """ - tool_labels = [label for label in tool_labels if label in default_tool_label_name_list] - return list(set(tool_labels)) - - @classmethod - def update_tool_labels( - cls, controller: ToolProviderController, labels: list[str], session: Session | None = None - ) -> None: - """ - Update tool labels - - :param controller: tool provider controller - :param labels: list of tool labels - :param session: database session, if None, a new session will be created - :return: None - """ - - labels = cls.filter_tool_labels(labels) - - if isinstance(controller, ApiToolProviderController | WorkflowToolProviderController): - provider_id = controller.provider_id - else: - raise ValueError("Unsupported tool type") - - if session is not None: - cls._update_tool_labels_logics(session, provider_id, controller, labels) - else: - with sessionmaker(db.engine).begin() as _session: - cls._update_tool_labels_logics(_session, provider_id, controller, labels) - - @classmethod - def _update_tool_labels_logics( - cls, session: Session, provider_id: str, controller: ToolProviderController, labels: list[str] - ) -> None: - """ - Update tool labels logics - - :param session: database session - :param provider_id: tool provider ID - :param controller: tool provider controller - :param labels: list of tool labels - :return: None - """ - - # delete old labels - _ = session.execute( - delete(ToolLabelBinding).where( - ToolLabelBinding.tool_id == provider_id, ToolLabelBinding.tool_type == controller.provider_type - ) - ) - - # insert new labels - for label in labels: - session.add(ToolLabelBinding(tool_id=provider_id, tool_type=controller.provider_type, label_name=label)) - - @classmethod - def get_tool_labels(cls, controller: ToolProviderController) -> list[str]: - """ - Get tool labels - - :param controller: tool provider controller - :return: list of tool labels (str) - """ - - match controller: - case ApiToolProviderController() | WorkflowToolProviderController(): - provider_id = controller.provider_id - case BuiltinToolProviderController(): - return controller.tool_labels - case _: - raise ValueError("Unsupported tool type") - stmt = select(ToolLabelBinding.label_name).where( - ToolLabelBinding.tool_id == provider_id, - ToolLabelBinding.tool_type == controller.provider_type, - ) - - with sessionmaker(db.engine, expire_on_commit=False).begin() as _session: - labels: list[str] = list(_session.scalars(stmt).all()) - - return labels - - @classmethod - def get_tools_labels(cls, tool_providers: list[ToolProviderController]) -> dict[str, list[str]]: - """ - Get tools labels - - :param tool_providers: list of tool providers - - :return: dict of tool labels - :key: tool id - :value: list of tool labels - """ - if not tool_providers: - return {} - - provider_ids: list[str] = [] - provider_types: set[str] = set() - - for controller in tool_providers: - if not isinstance(controller, ApiToolProviderController | WorkflowToolProviderController): - raise ValueError("Unsupported tool type") - provider_ids.append(controller.provider_id) - provider_types.add(controller.provider_type) - - labels: list[ToolLabelBinding] = [] - - with sessionmaker(db.engine, expire_on_commit=False).begin() as _session: - stmt = select(ToolLabelBinding).where( - ToolLabelBinding.tool_id.in_(provider_ids), ToolLabelBinding.tool_type.in_(list(provider_types)) - ) - labels = list(_session.scalars(stmt).all()) - - tool_labels: dict[str, list[str]] = {label.tool_id: [] for label in labels} - - for label in labels: - tool_labels[label.tool_id].append(label.label_name) - - return tool_labels diff --git a/api/core/tools/tool_manager.py b/api/core/tools/tool_manager.py deleted file mode 100644 index b492948f85dab9..00000000000000 --- a/api/core/tools/tool_manager.py +++ /dev/null @@ -1,1208 +0,0 @@ -import json -import logging -import mimetypes -import time -from collections.abc import Generator, Mapping -from os import listdir, path -from threading import Lock -from typing import TYPE_CHECKING, Any, Literal, Protocol, cast - -import sqlalchemy as sa -from pydantic import TypeAdapter -from sqlalchemy import select -from sqlalchemy.orm import Session -from typing_extensions import TypedDict -from yarl import URL - -import contexts -from configs import dify_config -from core.agent.entities import AgentToolEntity -from core.app.entities.app_invoke_entities import InvokeFrom -from core.entities import PluginCredentialType -from core.helper.module_import_helper import load_single_subclass_from_source -from core.helper.position_helper import is_filtered -from core.helper.provider_cache import ToolProviderCredentialsCache -from core.plugin.impl.exc import PluginDaemonNotFoundError, PluginNotFoundError -from core.plugin.impl.tool import PluginToolManager -from core.tools.__base.tool import Tool -from core.tools.__base.tool_provider import ToolProviderController -from core.tools.__base.tool_runtime import ToolRuntime -from core.tools.builtin_tool.provider import BuiltinToolProviderController -from core.tools.builtin_tool.providers._positions import BuiltinToolProviderSort -from core.tools.builtin_tool.tool import BuiltinTool -from core.tools.custom_tool.provider import ApiToolProviderController -from core.tools.custom_tool.tool import ApiTool -from core.tools.entities.api_entities import ToolProviderApiEntity, ToolProviderTypeApiLiteral -from core.tools.entities.common_entities import I18nObject -from core.tools.entities.tool_entities import ( - ApiProviderAuthType, - ApiProviderSchemaType, - EmojiIconDict, - ToolInvokeFrom, - ToolParameter, - ToolProviderType, - emoji_icon_adapter, -) -from core.tools.errors import ToolProviderCredentialValidationError, ToolProviderNotFoundError -from core.tools.mcp_tool.provider import MCPToolProviderController -from core.tools.mcp_tool.tool import MCPTool -from core.tools.plugin_tool.provider import PluginToolProviderController -from core.tools.plugin_tool.tool import PluginTool -from core.tools.tool_label_manager import ToolLabelManager -from core.tools.utils.configuration import ToolParameterConfigurationManager -from core.tools.utils.encryption import create_provider_encrypter, create_tool_provider_encrypter -from core.tools.utils.uuid_utils import is_valid_uuid -from core.tools.workflow_as_tool.provider import WorkflowToolProviderController -from core.tools.workflow_as_tool.tool import WorkflowTool -from extensions.ext_database import db -from graphon.runtime import VariablePool -from graphon.variables.template_resolution import convert_template -from models.provider_ids import ToolProviderID -from models.tools import ApiToolProvider, BuiltinToolProvider, WorkflowToolProvider -from services.tools.legacy_mcp_tools_manage_service import MCPToolManageService -from services.tools.legacy_tools_transform_service import ToolTransformService - -if TYPE_CHECKING: - pass - -logger = logging.getLogger(__name__) - - -class ApiProviderControllerItem(TypedDict): - provider: ApiToolProvider - controller: ApiToolProviderController - - -_credentials_adapter: TypeAdapter[dict[str, Any]] = TypeAdapter(dict[str, Any]) - - -class WorkflowToolRuntimeSpec(Protocol): - @property - def provider_type(self) -> ToolProviderType: ... - - @property - def provider_id(self) -> str: ... - - @property - def tool_name(self) -> str: ... - - @property - def tool_configurations(self) -> Mapping[str, Any]: ... - - @property - def credential_id(self) -> str | None: ... - - -class ToolManager: - _builtin_provider_lock = Lock() - _hardcoded_providers: dict[str, BuiltinToolProviderController] = {} - _builtin_providers_loaded = False - _builtin_tools_labels: dict[str, I18nObject | None] = {} - - @classmethod - def get_hardcoded_provider(cls, provider: str) -> BuiltinToolProviderController: - """ - - get the hardcoded provider - - """ - - if len(cls._hardcoded_providers) == 0: - # init the builtin providers - cls.load_hardcoded_providers_cache() - - return cls._hardcoded_providers[provider] - - @classmethod - def get_builtin_provider( - cls, provider: str, tenant_id: str - ) -> BuiltinToolProviderController | PluginToolProviderController: - """ - get the builtin provider - - :param provider: the name of the provider - :param tenant_id: the id of the tenant - :return: the provider - """ - # split provider to - - if len(cls._hardcoded_providers) == 0: - # init the builtin providers - cls.load_hardcoded_providers_cache() - - if provider not in cls._hardcoded_providers: - # get plugin provider - plugin_provider = cls.get_plugin_provider(provider, tenant_id) - if plugin_provider: - return plugin_provider - - return cls._hardcoded_providers[provider] - - @classmethod - def get_plugin_provider(cls, provider: str, tenant_id: str) -> PluginToolProviderController: - """ - get the plugin provider - """ - # check if context is set - - try: - contexts.plugin_tool_providers.get() - except LookupError: - contexts.plugin_tool_providers.set({}) - contexts.plugin_tool_providers_lock.set(Lock()) - - plugin_tool_providers = contexts.plugin_tool_providers.get() - if provider in plugin_tool_providers: - return plugin_tool_providers[provider] - - with contexts.plugin_tool_providers_lock.get(): - # double check - plugin_tool_providers = contexts.plugin_tool_providers.get() - if provider in plugin_tool_providers: - return plugin_tool_providers[provider] - - manager = PluginToolManager() - try: - provider_entity = manager.fetch_tool_provider(tenant_id, provider) - except (PluginNotFoundError, PluginDaemonNotFoundError) as exc: - # The plugin daemon rejected the provider name (probably an - # unknown or non-builtin identifier). Translate to the - # console's domain error so the API returns a 4xx instead - # of leaking the daemon's ``PluginNotFoundError`` as a 500. - raise ToolProviderNotFoundError(f"plugin provider {provider} not found") from exc - if not provider_entity: - raise ToolProviderNotFoundError(f"plugin provider {provider} not found") - - controller = PluginToolProviderController( - entity=provider_entity.declaration, - plugin_id=provider_entity.plugin_id, - plugin_unique_identifier=provider_entity.plugin_unique_identifier, - tenant_id=tenant_id, - ) - - plugin_tool_providers[provider] = controller - return controller - - @classmethod - def get_tool_runtime( - cls, - provider_type: ToolProviderType, - provider_id: str, - tool_name: str, - tenant_id: str, - user_id: str | None = None, - invoke_from: InvokeFrom = InvokeFrom.DEBUGGER, - tool_invoke_from: ToolInvokeFrom = ToolInvokeFrom.AGENT, - credential_id: str | None = None, - ) -> BuiltinTool | PluginTool | ApiTool | WorkflowTool | MCPTool: - """ - get the tool runtime - - :param provider_type: the type of the provider - :param provider_id: the id of the provider - :param tool_name: the name of the tool - :param tenant_id: the tenant id - :param user_id: the caller id bound to runtime-scoped model/tool lookups - :param invoke_from: invoke from - :param tool_invoke_from: the tool invoke from - :param credential_id: the credential id - - :return: the tool - """ - match provider_type: - case ToolProviderType.BUILT_IN: - provider_controller = cls.get_builtin_provider(provider_id, tenant_id) - - builtin_tool = provider_controller.get_tool(tool_name) - if not builtin_tool: - raise ToolProviderNotFoundError(f"builtin tool {tool_name} not found") - - if not provider_controller.need_credentials: - return builtin_tool.fork_tool_runtime( - runtime=ToolRuntime( - tenant_id=tenant_id, - user_id=user_id, - credentials={}, - invoke_from=invoke_from, - tool_invoke_from=tool_invoke_from, - ) - ) - builtin_provider = None - if isinstance(provider_controller, PluginToolProviderController): - provider_id_entity = ToolProviderID(provider_id) - if is_valid_uuid(credential_id): - try: - builtin_provider_stmt = select(BuiltinToolProvider).where( - BuiltinToolProvider.tenant_id == tenant_id, - BuiltinToolProvider.id == credential_id, - ) - builtin_provider = db.session.scalar(builtin_provider_stmt) - except Exception as e: - builtin_provider = None - logger.info("Error getting builtin provider %s:%s", credential_id, e, exc_info=True) - if builtin_provider is None: - raise ToolProviderCredentialValidationError( - f"Tool credential {credential_id} has been deleted. " - "Select or authorize another credential." - ) - - if builtin_provider is None: - with Session(db.engine) as session: - builtin_provider = session.scalar( - sa.select(BuiltinToolProvider) - .where( - BuiltinToolProvider.tenant_id == tenant_id, - (BuiltinToolProvider.provider == str(provider_id_entity)) - | (BuiltinToolProvider.provider == provider_id_entity.provider_name), - ) - .order_by(BuiltinToolProvider.is_default.desc(), BuiltinToolProvider.created_at.asc()) - ) - if builtin_provider is None: - raise ToolProviderCredentialValidationError( - f"No workspace credential is configured for tool provider {provider_id}. " - "Authorize the provider or select a credential." - ) - else: - builtin_provider = db.session.scalar( - select(BuiltinToolProvider) - .where( - BuiltinToolProvider.tenant_id == tenant_id, (BuiltinToolProvider.provider == provider_id) - ) - .order_by(BuiltinToolProvider.is_default.desc(), BuiltinToolProvider.created_at.asc()) - .limit(1) - ) - - if builtin_provider is None: - raise ToolProviderCredentialValidationError( - f"No credential is configured for built-in tool provider {provider_id}. " - "Authorize the provider or select a credential." - ) - - from core.helper.credential_utils import runtime_check_credential_policy_compliance - - runtime_check_credential_policy_compliance( - credential_id=builtin_provider.id, - provider=provider_id, - credential_type=PluginCredentialType.TOOL, - check_existence=False, - ) - - encrypter, cache = create_provider_encrypter( - tenant_id=tenant_id, - config=[ - x.to_basic_provider_config() - for x in provider_controller.get_credentials_schema_by_type(builtin_provider.credential_type) - ], - cache=ToolProviderCredentialsCache( - tenant_id=tenant_id, provider=provider_id, credential_id=builtin_provider.id - ), - ) - - decrypted_credentials: Mapping[str, Any] = encrypter.decrypt(builtin_provider.credentials) - - if builtin_provider.expires_at != -1 and (builtin_provider.expires_at - 60) < int(time.time()): - # TODO: circular import - from core.plugin.impl.oauth import OAuthHandler - from services.tools.legacy_builtin_tools_manage_service import BuiltinToolManageService - - tool_provider = ToolProviderID(provider_id) - provider_name = tool_provider.provider_name - redirect_uri = f"{dify_config.CONSOLE_API_URL}/console/api/oauth/plugin/{provider_id}/tool/callback" - system_credentials = BuiltinToolManageService.get_oauth_client(tenant_id, provider_id) - - oauth_handler = OAuthHandler() - try: - refreshed_credentials = oauth_handler.refresh_credentials( - tenant_id=tenant_id, - user_id=builtin_provider.user_id, - plugin_id=tool_provider.plugin_id, - provider=provider_name, - redirect_uri=redirect_uri, - system_credentials=system_credentials or {}, - credentials=decrypted_credentials, - ) - except Exception as exc: - logger.warning( - "Failed to refresh OAuth credentials for tool provider %s", provider_id, exc_info=True - ) - raise ToolProviderCredentialValidationError( - f"OAuth credential for tool provider {provider_id} could not be refreshed. " - "Reauthorize or select another credential." - ) from exc - # update the credentials - builtin_provider.encrypted_credentials = json.dumps( - encrypter.encrypt(refreshed_credentials.credentials) - ) - builtin_provider.expires_at = refreshed_credentials.expires_at - db.session.commit() - decrypted_credentials = refreshed_credentials.credentials - cache.delete() - - return builtin_tool.fork_tool_runtime( - runtime=ToolRuntime( - tenant_id=tenant_id, - user_id=user_id, - credentials=dict(decrypted_credentials), - credential_type=builtin_provider.credential_type, - runtime_parameters={}, - invoke_from=invoke_from, - tool_invoke_from=tool_invoke_from, - ) - ) - - case ToolProviderType.API: - api_provider, credentials = cls.get_api_provider_controller(tenant_id, provider_id) - encrypter, _ = create_tool_provider_encrypter( - tenant_id=tenant_id, - controller=api_provider, - ) - return api_provider.get_tool(tool_name).fork_tool_runtime( - runtime=ToolRuntime( - tenant_id=tenant_id, - user_id=user_id, - credentials=dict(encrypter.decrypt(credentials)), - invoke_from=invoke_from, - tool_invoke_from=tool_invoke_from, - ) - ) - case ToolProviderType.WORKFLOW: - workflow_provider_stmt = select(WorkflowToolProvider).where( - WorkflowToolProvider.tenant_id == tenant_id, WorkflowToolProvider.id == provider_id - ) - with Session(db.engine, expire_on_commit=False) as session, session.begin(): - workflow_provider = session.scalar(workflow_provider_stmt) - - if workflow_provider is None: - raise ToolProviderNotFoundError(f"workflow provider {provider_id} not found") - - controller = ToolTransformService.workflow_provider_to_controller(db_provider=workflow_provider) - controller_tools: list[WorkflowTool] = controller.get_tools(tenant_id=workflow_provider.tenant_id) - if controller_tools is None or len(controller_tools) == 0: - raise ToolProviderNotFoundError(f"workflow provider {provider_id} not found") - - return controller.get_tools(tenant_id=workflow_provider.tenant_id)[0].fork_tool_runtime( - runtime=ToolRuntime( - tenant_id=tenant_id, - user_id=user_id, - credentials={}, - invoke_from=invoke_from, - tool_invoke_from=tool_invoke_from, - ) - ) - case ToolProviderType.APP: - raise NotImplementedError("app provider not implemented") - case ToolProviderType.PLUGIN: - plugin_tool = cls.get_plugin_provider(provider_id, tenant_id).get_tool(tool_name) - runtime = getattr(plugin_tool, "runtime", None) - if runtime is not None: - runtime.user_id = user_id - runtime.invoke_from = invoke_from - runtime.tool_invoke_from = tool_invoke_from - return plugin_tool - case ToolProviderType.MCP: - mcp_tool = cls.get_mcp_provider_controller(tenant_id, provider_id).get_tool(tool_name) - runtime = getattr(mcp_tool, "runtime", None) - if runtime is not None: - runtime.user_id = user_id - runtime.invoke_from = invoke_from - runtime.tool_invoke_from = tool_invoke_from - return mcp_tool - case ToolProviderType.DATASET_RETRIEVAL: - raise ToolProviderNotFoundError(f"provider type {provider_type.value} not found") - case _: - raise ToolProviderNotFoundError(f"provider type {provider_type} not found") - - @classmethod - def get_agent_tool_runtime( - cls, - tenant_id: str, - app_id: str, - agent_tool: AgentToolEntity, - user_id: str | None = None, - invoke_from: InvokeFrom = InvokeFrom.DEBUGGER, - variable_pool: "VariablePool | None" = None, - allow_file_parameters: bool = False, - use_default_for_missing_form_parameters: bool = False, - ) -> Tool: - """ - get the agent tool runtime - """ - tool_entity = cls.get_tool_runtime( - provider_type=agent_tool.provider_type, - provider_id=agent_tool.provider_id, - tool_name=agent_tool.tool_name, - tenant_id=tenant_id, - user_id=user_id, - invoke_from=invoke_from, - tool_invoke_from=ToolInvokeFrom.AGENT, - credential_id=agent_tool.credential_id, - ) - runtime_parameters: dict[str, Any] = {} - parameters = tool_entity.get_merged_runtime_parameters() - runtime_parameters = cls._convert_tool_parameters_type( - parameters, - variable_pool, - agent_tool.tool_parameters, - typ="agent", - allow_file_parameters=allow_file_parameters, - use_default_for_missing_form_parameters=use_default_for_missing_form_parameters, - ) - # decrypt runtime parameters - encryption_manager = ToolParameterConfigurationManager( - tenant_id=tenant_id, - tool_runtime=tool_entity, - provider_name=agent_tool.provider_id, - provider_type=agent_tool.provider_type, - identity_id=f"AGENT.{app_id}", - ) - runtime_parameters = encryption_manager.decrypt_tool_parameters(runtime_parameters) - if tool_entity.runtime is None or tool_entity.runtime.runtime_parameters is None: - raise ValueError("runtime not found or runtime parameters not found") - - tool_entity.runtime.runtime_parameters.update(runtime_parameters) - return tool_entity - - @classmethod - def get_workflow_tool_runtime( - cls, - tenant_id: str, - app_id: str, - node_id: str, - workflow_tool: WorkflowToolRuntimeSpec, - user_id: str | None = None, - invoke_from: InvokeFrom = InvokeFrom.DEBUGGER, - variable_pool: "VariablePool | None" = None, - ) -> Tool: - """ - get the workflow tool runtime - """ - - tool_runtime = cls.get_tool_runtime( - provider_type=workflow_tool.provider_type, - provider_id=workflow_tool.provider_id, - tool_name=workflow_tool.tool_name, - tenant_id=tenant_id, - user_id=user_id, - invoke_from=invoke_from, - tool_invoke_from=ToolInvokeFrom.WORKFLOW, - credential_id=workflow_tool.credential_id, - ) - - parameters = tool_runtime.get_merged_runtime_parameters() - runtime_parameters = cls._convert_tool_parameters_type( - parameters, variable_pool, workflow_tool.tool_configurations, typ="workflow" - ) - # decrypt runtime parameters - encryption_manager = ToolParameterConfigurationManager( - tenant_id=tenant_id, - tool_runtime=tool_runtime, - provider_name=workflow_tool.provider_id, - provider_type=workflow_tool.provider_type, - identity_id=f"WORKFLOW.{app_id}.{node_id}", - ) - - if runtime_parameters: - runtime_parameters = encryption_manager.decrypt_tool_parameters(runtime_parameters) - - tool_runtime.runtime.runtime_parameters.update(runtime_parameters) - return tool_runtime - - @classmethod - def get_tool_runtime_from_plugin( - cls, - tool_type: ToolProviderType, - tenant_id: str, - provider: str, - tool_name: str, - tool_parameters: dict[str, Any], - user_id: str | None = None, - credential_id: str | None = None, - ) -> Tool: - """ - get tool runtime from plugin - """ - tool_entity = cls.get_tool_runtime( - provider_type=tool_type, - provider_id=provider, - tool_name=tool_name, - tenant_id=tenant_id, - user_id=user_id, - invoke_from=InvokeFrom.SERVICE_API, - tool_invoke_from=ToolInvokeFrom.PLUGIN, - credential_id=credential_id, - ) - runtime_parameters: dict[str, Any] = {} - parameters = tool_entity.get_merged_runtime_parameters() - for parameter in parameters: - if parameter.form == ToolParameter.ToolParameterForm.FORM: - # save tool parameter to tool entity memory - value = parameter.init_frontend_parameter(tool_parameters.get(parameter.name)) - runtime_parameters[parameter.name] = value - - tool_entity.runtime.runtime_parameters.update(runtime_parameters) - return tool_entity - - @classmethod - def get_hardcoded_provider_icon(cls, provider: str) -> tuple[str, str]: - """ - get the absolute path of the icon of the hardcoded provider - - :param provider: the name of the provider - :return: the absolute path of the icon, the mime type of the icon - """ - # get provider - provider_controller = cls.get_hardcoded_provider(provider) - - absolute_path = path.join( - path.dirname(path.realpath(__file__)), - "builtin_tool", - "providers", - provider, - "_assets", - provider_controller.entity.identity.icon, - ) - # check if the icon exists - if not path.exists(absolute_path): - raise ToolProviderNotFoundError(f"builtin provider {provider} icon not found") - - # get the mime type - mime_type, _ = mimetypes.guess_type(absolute_path) - mime_type = mime_type or "application/octet-stream" - - return absolute_path, mime_type - - @classmethod - def list_hardcoded_providers(cls): - # use cache first - if cls._builtin_providers_loaded: - yield from list(cls._hardcoded_providers.values()) - return - - with cls._builtin_provider_lock: - if cls._builtin_providers_loaded: - yield from list(cls._hardcoded_providers.values()) - return - - yield from cls._list_hardcoded_providers() - - @classmethod - def list_plugin_providers(cls, tenant_id: str) -> list[PluginToolProviderController]: - """ - list all the plugin providers - """ - - manager = PluginToolManager() - provider_entities = manager.fetch_tool_providers(tenant_id) - return [ - PluginToolProviderController( - entity=provider.declaration, - plugin_id=provider.plugin_id, - plugin_unique_identifier=provider.plugin_unique_identifier, - tenant_id=tenant_id, - ) - for provider in provider_entities - ] - - @classmethod - def list_builtin_providers( - cls, tenant_id: str - ) -> Generator[BuiltinToolProviderController | PluginToolProviderController, None, None]: - """ - list all the builtin providers - """ - yield from cls.list_hardcoded_providers() - # get plugin providers - yield from cls.list_plugin_providers(tenant_id) - - @classmethod - def _list_hardcoded_providers(cls) -> Generator[BuiltinToolProviderController, None, None]: - """ - list all the builtin providers - """ - for provider_path in listdir(path.join(path.dirname(path.realpath(__file__)), "builtin_tool", "providers")): - if provider_path.startswith("__"): - continue - - if path.isdir(path.join(path.dirname(path.realpath(__file__)), "builtin_tool", "providers", provider_path)): - if provider_path.startswith("__"): - continue - - # init provider - try: - provider_class = load_single_subclass_from_source( - module_name=f"core.tools.builtin_tool.providers.{provider_path}.{provider_path}", - script_path=path.join( - path.dirname(path.realpath(__file__)), - "builtin_tool", - "providers", - provider_path, - f"{provider_path}.py", - ), - parent_type=BuiltinToolProviderController, - ) - provider: BuiltinToolProviderController = provider_class() - cls._hardcoded_providers[provider.entity.identity.name] = provider - for tool in provider.get_tools(): - cls._builtin_tools_labels[tool.entity.identity.name] = tool.entity.identity.label - yield provider - - except Exception: - logger.exception("load builtin provider %s", provider_path) - continue - # set builtin providers loaded - cls._builtin_providers_loaded = True - - @classmethod - def load_hardcoded_providers_cache(cls): - for _ in cls.list_hardcoded_providers(): - pass - - @classmethod - def clear_hardcoded_providers_cache(cls): - cls._hardcoded_providers = {} - cls._builtin_providers_loaded = False - - @classmethod - def get_tool_label(cls, tool_name: str) -> I18nObject | None: - """ - get the tool label - - :param tool_name: the name of the tool - - :return: the label of the tool - """ - if len(cls._builtin_tools_labels) == 0: - # init the builtin providers - cls.load_hardcoded_providers_cache() - - if tool_name not in cls._builtin_tools_labels: - return None - - return cls._builtin_tools_labels[tool_name] - - @classmethod - def list_default_builtin_providers(cls, tenant_id: str) -> list[BuiltinToolProvider]: - """ - list all the builtin providers - """ - # according to multi credentials, select the one with is_default=True first, then created_at oldest - # for compatibility with old version - if dify_config.SQLALCHEMY_DATABASE_URI_SCHEME == "postgresql": - # PostgreSQL: Use DISTINCT ON - sql = """ - SELECT DISTINCT ON (tenant_id, provider) id - FROM tool_builtin_providers - WHERE tenant_id = :tenant_id - ORDER BY tenant_id, provider, is_default DESC, created_at DESC - """ - else: - # MySQL: Use window function to achieve same result - sql = """ - SELECT id FROM ( - SELECT id, - ROW_NUMBER() OVER ( - PARTITION BY tenant_id, provider - ORDER BY is_default DESC, created_at DESC - ) as rn - FROM tool_builtin_providers - WHERE tenant_id = :tenant_id - ) ranked WHERE rn = 1 - """ - - with Session(db.engine, autoflush=False) as session: - ids = [row.id for row in session.execute(sa.text(sql), {"tenant_id": tenant_id}).all()] - return list(session.scalars(select(BuiltinToolProvider).where(BuiltinToolProvider.id.in_(ids)))) - - @classmethod - def list_providers_from_api( - cls, user_id: str, tenant_id: str, typ: ToolProviderTypeApiLiteral - ) -> list[ToolProviderApiEntity]: - result_providers: dict[str, ToolProviderApiEntity] = {} - - filters = [] - if not typ: - filters.extend(["builtin", "api", "workflow", "mcp"]) - else: - filters.append(typ) - - # Use a single session for all database operations to reduce connection overhead - with Session(db.engine) as session: - if "builtin" in filters: - builtin_providers = list(cls.list_builtin_providers(tenant_id)) - - # key: provider name, value: provider - db_builtin_providers = { - str(ToolProviderID(provider.provider)): provider - for provider in cls.list_default_builtin_providers(tenant_id) - } - - # append builtin providers - for provider in builtin_providers: - # handle include, exclude - if is_filtered( - include_set=dify_config.POSITION_TOOL_INCLUDES_SET, - exclude_set=dify_config.POSITION_TOOL_EXCLUDES_SET, - data=provider, - name_func=lambda x: x.entity.identity.name, - ): - continue - user_provider = ToolTransformService.builtin_provider_to_user_provider( - provider_controller=provider, - db_provider=db_builtin_providers.get(provider.entity.identity.name), - decrypt_credentials=False, - ) - - if isinstance(provider, PluginToolProviderController): - result_providers[f"plugin_provider.{user_provider.name}"] = user_provider - else: - result_providers[f"builtin_provider.{user_provider.name}"] = user_provider - - # get db api providers - if "api" in filters: - db_api_providers = session.scalars( - select(ApiToolProvider).where(ApiToolProvider.tenant_id == tenant_id) - ).all() - - # Batch create controllers - api_provider_controllers: list[ApiProviderControllerItem] = [] - for api_provider in db_api_providers: - try: - controller = ToolTransformService.api_provider_to_controller(api_provider) - api_provider_controllers.append({"provider": api_provider, "controller": controller}) - except Exception: - # Skip invalid providers but continue processing others - logger.warning("Failed to create controller for API provider %s", api_provider.id) - - # Batch get labels for all API providers - if api_provider_controllers: - controllers = cast( - list[ToolProviderController], [item["controller"] for item in api_provider_controllers] - ) - labels = ToolLabelManager.get_tools_labels(controllers) - - for item in api_provider_controllers: - provider_controller = item["controller"] - db_provider = item["provider"] - provider_labels = labels.get(provider_controller.provider_id, []) - user_provider = ToolTransformService.api_provider_to_user_provider( - provider_controller=provider_controller, - db_provider=db_provider, - decrypt_credentials=False, - labels=provider_labels, - session=db.session(), - ) - result_providers[f"api_provider.{user_provider.name}"] = user_provider - - if "workflow" in filters: - # get workflow providers - workflow_providers = session.scalars( - select(WorkflowToolProvider).where(WorkflowToolProvider.tenant_id == tenant_id) - ).all() - - workflow_provider_controllers: list[WorkflowToolProviderController] = [] - for workflow_provider in workflow_providers: - try: - workflow_controller: WorkflowToolProviderController = ( - ToolTransformService.workflow_provider_to_controller(db_provider=workflow_provider) - ) - workflow_provider_controllers.append(workflow_controller) - except Exception: - # app has been deleted - logger.exception("Failed to transform workflow provider %s to controller", workflow_provider.id) - continue - # Batch get labels for workflow providers - if workflow_provider_controllers: - workflow_controllers: list[ToolProviderController] = [ - cast(ToolProviderController, controller) for controller in workflow_provider_controllers - ] - labels = ToolLabelManager.get_tools_labels(workflow_controllers) - - for workflow_provider_controller in workflow_provider_controllers: - provider_labels = labels.get(workflow_provider_controller.provider_id, []) - user_provider = ToolTransformService.workflow_provider_to_user_provider( - provider_controller=workflow_provider_controller, - labels=provider_labels, - ) - result_providers[f"workflow_provider.{user_provider.name}"] = user_provider - - if "mcp" in filters: - mcp_service = MCPToolManageService(session=session) - mcp_providers = mcp_service.list_providers(tenant_id=tenant_id) - for mcp_provider in mcp_providers: - result_providers[f"mcp_provider.{mcp_provider.name}"] = mcp_provider - - return BuiltinToolProviderSort.sort(list(result_providers.values())) - - @classmethod - def get_api_provider_controller( - cls, tenant_id: str, provider_id: str - ) -> tuple[ApiToolProviderController, dict[str, Any]]: - """ - get the api provider - - :param tenant_id: the id of the tenant - :param provider_id: the id of the provider - - :return: the provider controller, the credentials - """ - provider: ApiToolProvider | None = db.session.scalar( - select(ApiToolProvider) - .where( - ApiToolProvider.id == provider_id, - ApiToolProvider.tenant_id == tenant_id, - ) - .limit(1) - ) - - if provider is None: - raise ToolProviderNotFoundError(f"api provider {provider_id} not found") - - auth_type = ApiProviderAuthType.NONE - provider_auth_type = provider.credentials.get("auth_type") - if provider_auth_type in ("api_key_header", "api_key"): # backward compatibility - auth_type = ApiProviderAuthType.API_KEY_HEADER - elif provider_auth_type == "api_key_query": - auth_type = ApiProviderAuthType.API_KEY_QUERY - - controller = ApiToolProviderController.from_db( - provider, - auth_type, - session=db.session(), - ) - controller.load_bundled_tools(provider.tools) - - return controller, provider.credentials - - @classmethod - def get_mcp_provider_controller(cls, tenant_id: str, provider_id: str) -> MCPToolProviderController: - """ - get the api provider - - :param tenant_id: the id of the tenant - :param provider_id: the persisted reference of the provider, normally its - server identifier, or the primary key for graphs written before that - convention - - :return: the provider controller, the credentials - """ - with Session(db.engine) as session: - mcp_service = MCPToolManageService(session=session) - try: - provider = mcp_service.get_provider_by_persisted_reference( - id_or_server_identifier=provider_id, tenant_id=tenant_id - ) - except ValueError: - raise ToolProviderNotFoundError(f"mcp provider {provider_id} not found") - - controller = MCPToolProviderController.from_db(provider) - - return controller - - @classmethod - def user_get_api_provider(cls, provider: str, tenant_id: str, mask: bool = True): - """ - get api provider - """ - provider_name = provider - provider_obj: ApiToolProvider | None = db.session.scalar( - select(ApiToolProvider) - .where( - ApiToolProvider.tenant_id == tenant_id, - ApiToolProvider.name == provider, - ) - .limit(1) - ) - - if provider_obj is None: - raise ValueError(f"you have not added provider {provider_name}") - - try: - credentials = _credentials_adapter.validate_json(provider_obj.credentials_str) or {} - except Exception: - credentials = {} - - # package tool provider controller - auth_type = ApiProviderAuthType.NONE - credentials_auth_type = credentials.get("auth_type") - if credentials_auth_type in ("api_key_header", "api_key"): # backward compatibility - auth_type = ApiProviderAuthType.API_KEY_HEADER - elif credentials_auth_type == "api_key_query": - auth_type = ApiProviderAuthType.API_KEY_QUERY - - controller = ApiToolProviderController.from_db( - provider_obj, - auth_type, - session=db.session(), - ) - # init tool configuration - encrypter, _ = create_tool_provider_encrypter( - tenant_id=tenant_id, - controller=controller, - ) - if mask: - masked_credentials = encrypter.mask_plugin_credentials(encrypter.decrypt(credentials)) - else: - masked_credentials = encrypter.decrypt(credentials) - - try: - icon = emoji_icon_adapter.validate_json(provider_obj.icon) - except Exception: - icon = {"background": "#252525", "content": "\ud83d\ude01"} - - # add tool labels - labels = ToolLabelManager.get_tool_labels(controller) - schema_type = provider_obj.schema_type - schema_type_value = schema_type.value if isinstance(schema_type, ApiProviderSchemaType) else schema_type - - return { - "schema_type": schema_type_value, - "schema": provider_obj.schema, - "tools": [tool.model_dump(mode="json") for tool in provider_obj.tools], - "icon": icon, - "description": provider_obj.description, - "credentials": masked_credentials, - "privacy_policy": provider_obj.privacy_policy, - "custom_disclaimer": provider_obj.custom_disclaimer, - "labels": labels, - } - - @classmethod - def generate_builtin_tool_icon_url(cls, provider_id: str) -> str: - return str( - URL(dify_config.CONSOLE_API_URL or "/") - / "console" - / "api" - / "workspaces" - / "current" - / "tool-provider" - / "builtin" - / provider_id - / "icon" - ) - - @classmethod - def generate_plugin_tool_icon_url(cls, tenant_id: str, filename: str) -> str: - return str( - URL(dify_config.CONSOLE_API_URL or "/") - / "console" - / "api" - / "workspaces" - / "current" - / "plugin" - / "icon" - % {"tenant_id": tenant_id, "filename": filename} - ) - - @classmethod - def generate_workflow_tool_icon_url(cls, tenant_id: str, provider_id: str) -> EmojiIconDict: - try: - # Use a short-lived session to avoid holding a database transaction - # during long-running nested workflow execution. - # Fixes: idle in transaction when Workflow Tool runs (#36902) - with Session(db.engine, expire_on_commit=False) as session: - workflow_provider: WorkflowToolProvider | None = session.scalar( - select(WorkflowToolProvider) - .where(WorkflowToolProvider.tenant_id == tenant_id, WorkflowToolProvider.id == provider_id) - .limit(1) - ) - - if workflow_provider is None: - raise ToolProviderNotFoundError(f"workflow provider {provider_id} not found") - - icon = emoji_icon_adapter.validate_json(workflow_provider.icon) - return icon - except Exception: - return {"background": "#252525", "content": "\ud83d\ude01"} - - @classmethod - def generate_api_tool_icon_url(cls, tenant_id: str, provider_id: str) -> EmojiIconDict: - try: - # Use a short-lived session to avoid holding a database transaction - # during long-running tool execution. - with Session(db.engine, expire_on_commit=False) as session: - api_provider: ApiToolProvider | None = session.scalar( - select(ApiToolProvider) - .where(ApiToolProvider.tenant_id == tenant_id, ApiToolProvider.id == provider_id) - .limit(1) - ) - - if api_provider is None: - raise ToolProviderNotFoundError(f"api provider {provider_id} not found") - - icon = emoji_icon_adapter.validate_json(api_provider.icon) - return icon - except Exception: - return {"background": "#252525", "content": "\ud83d\ude01"} - - @classmethod - def generate_mcp_tool_icon_url(cls, tenant_id: str, provider_id: str) -> EmojiIconDict | str: - try: - with Session(db.engine) as session: - mcp_service = MCPToolManageService(session=session) - try: - mcp_provider = mcp_service.get_provider_entity_by_persisted_reference( - id_or_server_identifier=provider_id, tenant_id=tenant_id - ) - return cast(EmojiIconDict | str, mcp_provider.provider_icon) - except ValueError: - raise ToolProviderNotFoundError(f"mcp provider {provider_id} not found") - except Exception: - return {"background": "#252525", "content": "\ud83d\ude01"} - - @classmethod - def get_tool_icon( - cls, - tenant_id: str, - provider_type: ToolProviderType, - provider_id: str, - ) -> str | EmojiIconDict: - """ - get the tool icon - - :param tenant_id: the id of the tenant - :param provider_type: the type of the provider - :param provider_id: the id of the provider - :return: - """ - match provider_type: - case ToolProviderType.BUILT_IN: - provider = ToolManager.get_builtin_provider(provider_id, tenant_id) - if isinstance(provider, PluginToolProviderController): - try: - return cls.generate_plugin_tool_icon_url(tenant_id, provider.entity.identity.icon) - except Exception: - return {"background": "#252525", "content": "\ud83d\ude01"} - return cls.generate_builtin_tool_icon_url(provider_id) - case ToolProviderType.API: - return cls.generate_api_tool_icon_url(tenant_id, provider_id) - case ToolProviderType.WORKFLOW: - return cls.generate_workflow_tool_icon_url(tenant_id, provider_id) - case ToolProviderType.PLUGIN: - provider = ToolManager.get_plugin_provider(provider_id, tenant_id) - try: - return cls.generate_plugin_tool_icon_url(tenant_id, provider.entity.identity.icon) - except Exception: - return {"background": "#252525", "content": "\ud83d\ude01"} - case ToolProviderType.MCP: - return cls.generate_mcp_tool_icon_url(tenant_id, provider_id) - case ToolProviderType.APP | ToolProviderType.DATASET_RETRIEVAL: - raise ValueError(f"provider type {provider_type} not found") - case _: - raise ValueError(f"provider type {provider_type} not found") - - @classmethod - def _convert_tool_parameters_type( - cls, - parameters: list[ToolParameter], - variable_pool: "VariablePool | None", - tool_configurations: Mapping[str, Any], - typ: Literal["agent", "workflow", "tool"] = "workflow", - allow_file_parameters: bool = False, - use_default_for_missing_form_parameters: bool = False, - ) -> dict[str, Any]: - """ - Convert tool parameters type - """ - from graphon.nodes.tool.entities import ToolNodeData - from graphon.nodes.tool.exc import ToolParameterError - - runtime_parameters: dict[str, Any] = {} - for parameter in parameters: - if ( - parameter.type - in { - ToolParameter.ToolParameterType.SYSTEM_FILES, - ToolParameter.ToolParameterType.FILE, - ToolParameter.ToolParameterType.FILES, - } - and parameter.required - and typ == "agent" - and not allow_file_parameters - ): - raise ValueError(f"file type parameter {parameter.name} not supported in agent") - # save tool parameter to tool entity memory - if parameter.form == ToolParameter.ToolParameterForm.FORM: - if variable_pool: - config = tool_configurations.get(parameter.name, {}) - - selector_value = cls._extract_runtime_selector_value(parameter, config) - if selector_value is not None: - # Selector parameters carry structured dictionaries, not scalar ToolInput values. - runtime_parameters[parameter.name] = selector_value - continue - - if not (config and isinstance(config, dict) and config.get("value") is not None): - continue - tool_input = ToolNodeData.ToolInput.model_validate(tool_configurations.get(parameter.name, {})) - if tool_input.type == "variable": - variable_selector = tool_input.value - if not isinstance(variable_selector, list) or not all( - isinstance(selector_part, str) for selector_part in variable_selector - ): - raise ToolParameterError("Variable tool input must be a variable selector") - variable = variable_pool.get(variable_selector) - if variable is None: - raise ToolParameterError(f"Variable {tool_input.value} does not exist") - parameter_value = variable.value - elif tool_input.type == "constant": - parameter_value = tool_input.value - elif tool_input.type == "mixed": - segment_group = convert_template(variable_pool, str(tool_input.value)) - parameter_value = segment_group.text - else: - raise ToolParameterError(f"Unknown tool input type '{tool_input.type}'") - runtime_parameters[parameter.name] = parameter_value - - else: - parameter_value = tool_configurations.get(parameter.name) - if use_default_for_missing_form_parameters and parameter_value is None: - if parameter.default is not None: - parameter_value = parameter.default - elif ( - parameter.required - and parameter.type == ToolParameter.ToolParameterType.SELECT - and parameter.options - ): - parameter_value = parameter.options[0].value - else: - continue - value = parameter.init_frontend_parameter(parameter_value) - runtime_parameters[parameter.name] = value - return runtime_parameters - - @classmethod - def _extract_runtime_selector_value(cls, parameter: ToolParameter, config: Any) -> dict[str, Any] | None: - if parameter.type not in { - ToolParameter.ToolParameterType.MODEL_SELECTOR, - ToolParameter.ToolParameterType.APP_SELECTOR, - }: - return None - if not isinstance(config, dict): - return None - - input_value = config.get("value") - if isinstance(input_value, dict) and cls._is_selector_value(parameter, input_value): - return cast("dict[str, Any]", parameter.init_frontend_parameter(input_value)) - - if cls._is_selector_value(parameter, config): - selector_value = dict(config) - selector_value.pop("type", None) - selector_value.pop("value", None) - return cast("dict[str, Any]", parameter.init_frontend_parameter(selector_value)) - - return None - - @classmethod - def _is_selector_value(cls, parameter: ToolParameter, value: Mapping[str, Any]) -> bool: - if parameter.type == ToolParameter.ToolParameterType.MODEL_SELECTOR: - return ( - isinstance(value.get("provider"), str) - and isinstance(value.get("model"), str) - and isinstance(value.get("model_type"), str) - ) - if parameter.type == ToolParameter.ToolParameterType.APP_SELECTOR: - return isinstance(value.get("app_id"), str) - return False - - -ToolManager.load_hardcoded_providers_cache() diff --git a/api/core/tools/utils/dataset_retriever/dataset_multi_retriever_tool.py b/api/core/tools/utils/dataset_retriever/dataset_multi_retriever_tool.py deleted file mode 100644 index fbe4a3f8ffa772..00000000000000 --- a/api/core/tools/utils/dataset_retriever/dataset_multi_retriever_tool.py +++ /dev/null @@ -1,214 +0,0 @@ -import threading -from functools import partial -from typing import override - -from flask import Flask, current_app -from pydantic import BaseModel, Field -from sqlalchemy import select -from sqlalchemy.orm import Session, sessionmaker - -from core.callback_handler.index_tool_callback_handler import DatasetIndexToolCallbackHandler -from core.model_manager import ModelManager -from core.rag.datasource.retrieval_service import DefaultRetrievalModelDict, RetrievalService -from core.rag.entities import RetrievalSourceMetadata -from core.rag.index_processor.constant.index_type import IndexTechniqueType -from core.rag.models.document import Document as RagDocument -from core.rag.rerank.rerank_model import RerankModelRunner -from core.rag.retrieval.retrieval_methods import RetrievalMethod -from core.tools.utils.dataset_retriever.dataset_retriever_base_tool import DatasetRetrieverBaseTool -from extensions.ext_database import db -from graphon.model_runtime.entities.model_entities import ModelType -from models.dataset import Dataset, Document, DocumentSegment -from repositories.knowledge.segment_read_adapter import sign_segment_content -from repositories.knowledge.upload_file_repository import SQLAlchemyKnowledgeUploadRepository - -default_retrieval_model: DefaultRetrievalModelDict = { - "search_method": RetrievalMethod.SEMANTIC_SEARCH, - "reranking_enable": False, - "reranking_model": {"reranking_provider_name": "", "reranking_model_name": ""}, - "top_k": 2, - "score_threshold_enabled": False, -} - - -class DatasetMultiRetrieverToolInput(BaseModel): - query: str = Field(..., description="dataset multi retriever and rerank") - - -class DatasetMultiRetrieverTool(DatasetRetrieverBaseTool): - """Tool for querying multi dataset.""" - - name: str = "dataset_" - args_schema: type[BaseModel] = DatasetMultiRetrieverToolInput - description: str = "dataset multi retriever and rerank. " - dataset_ids: list[str] - reranking_provider_name: str - reranking_model_name: str - - @classmethod - def from_dataset(cls, dataset_ids: list[str], tenant_id: str, **kwargs): - return cls( - name=f"dataset_{tenant_id.replace('-', '_')}", tenant_id=tenant_id, dataset_ids=dataset_ids, **kwargs - ) - - @override - def _run(self, session: Session, query: str) -> str: - threads = [] - all_documents: list[RagDocument] = [] - for dataset_id in self.dataset_ids: - retrieval_thread = threading.Thread( - target=self._retriever, - kwargs={ - "flask_app": current_app._get_current_object(), # type: ignore - "dataset_id": dataset_id, - "query": query, - "all_documents": all_documents, - "hit_callbacks": self.hit_callbacks, - }, - ) - threads.append(retrieval_thread) - retrieval_thread.start() - for thread in threads: - thread.join() - # do rerank for searched documents - model_manager = ModelManager.for_tenant(tenant_id=self.tenant_id) - rerank_model_instance = model_manager.get_model_instance( - tenant_id=self.tenant_id, - provider=self.reranking_provider_name, - model_type=ModelType.RERANK, - model=self.reranking_model_name, - ) - - rerank_runner = RerankModelRunner( - rerank_model_instance, - load_upload=partial( - SQLAlchemyKnowledgeUploadRepository(session_factory=sessionmaker(bind=session.get_bind())).get_by_id, - workspace_id=self.tenant_id, - ), - ) - all_documents = rerank_runner.run(query, all_documents, self.score_threshold, self.top_k) - - for hit_callback in self.hit_callbacks: - hit_callback.on_tool_end(all_documents, session) - - document_score_list = {} - for item in all_documents: - if item.metadata and item.metadata.get("score"): - document_score_list[item.metadata["doc_id"]] = item.metadata["score"] - - document_context_list = [] - index_node_ids = [document.metadata["doc_id"] for document in all_documents if document.metadata] - document_segment_stmt = select(DocumentSegment).where( - DocumentSegment.dataset_id.in_(self.dataset_ids), - DocumentSegment.completed_at.isnot(None), - DocumentSegment.status == "completed", - DocumentSegment.enabled == True, - DocumentSegment.index_node_id.in_(index_node_ids), - ) - segments = session.scalars(document_segment_stmt).all() - - if segments: - index_node_id_to_position = {id: position for position, id in enumerate(index_node_ids)} - sorted_segments = sorted( - segments, key=lambda segment: index_node_id_to_position.get(segment.index_node_id, float("inf")) - ) - for segment in sorted_segments: - if segment.answer: - document_context_list.append( - f"question:{sign_segment_content(segment, session=session)} answer:{segment.answer}" - ) - else: - document_context_list.append(sign_segment_content(segment, session=session)) - if self.return_resource: - context_list: list[RetrievalSourceMetadata] = [] - resource_number = 1 - for segment in sorted_segments: - dataset = session.get(Dataset, segment.dataset_id) - document_stmt = select(Document).where( - Document.id == segment.document_id, - Document.enabled == True, - Document.archived == False, - ) - document = session.scalar(document_stmt) - if dataset and document: - source = RetrievalSourceMetadata( - position=resource_number, - dataset_id=dataset.id, - dataset_name=dataset.name, - document_id=document.id, - document_name=document.name, - data_source_type=document.data_source_type, - segment_id=segment.id, - retriever_from=self.retriever_from, - score=document_score_list.get(segment.index_node_id), - doc_metadata=document.doc_metadata, - ) - - if self.retriever_from == "dev": - source.hit_count = segment.hit_count - source.word_count = segment.word_count - source.segment_position = segment.position - source.index_node_hash = segment.index_node_hash - if segment.answer: - source.content = f"question:{segment.content} \nanswer:{segment.answer}" - else: - source.content = segment.content - context_list.append(source) - resource_number += 1 - - for hit_callback in self.hit_callbacks: - hit_callback.return_retriever_resource_info(context_list) - - return "\n".join(document_context_list) - return "" - - def _retriever( - self, - flask_app: Flask, - dataset_id: str, - query: str, - all_documents: list, - hit_callbacks: list[DatasetIndexToolCallbackHandler], - ): - with flask_app.app_context(): - stmt = select(Dataset).where(Dataset.tenant_id == self.tenant_id, Dataset.id == dataset_id) - dataset = db.session.scalar(stmt) - - if not dataset: - return [] - - for hit_callback in hit_callbacks: - hit_callback.on_query(query, dataset.id, db.session()) - - # get retrieval model , if the model is not setting , using default - retrieval_model = dataset.retrieval_model or default_retrieval_model - - if dataset.indexing_technique == IndexTechniqueType.ECONOMY: - # use keyword table query - documents = RetrievalService.retrieve( - retrieval_method=RetrievalMethod.KEYWORD_SEARCH, - dataset_id=dataset.id, - query=query, - top_k=retrieval_model.get("top_k") or 4, - ) - if documents: - all_documents.extend(documents) - else: - if self.top_k > 0: - # retrieval source - documents = RetrievalService.retrieve( - retrieval_method=retrieval_model["search_method"], - dataset_id=dataset.id, - query=query, - top_k=retrieval_model.get("top_k") or 4, - score_threshold=retrieval_model.get("score_threshold", 0.0) - if retrieval_model["score_threshold_enabled"] - else 0.0, - reranking_model=retrieval_model.get("reranking_model", None) - if retrieval_model["reranking_enable"] - else None, - reranking_mode=retrieval_model.get("reranking_mode") or "reranking_model", - weights=retrieval_model.get("weights", None), - ) - - all_documents.extend(documents) diff --git a/api/core/tools/utils/dataset_retriever/dataset_retriever_base_tool.py b/api/core/tools/utils/dataset_retriever/dataset_retriever_base_tool.py deleted file mode 100644 index 2d92761bc90506..00000000000000 --- a/api/core/tools/utils/dataset_retriever/dataset_retriever_base_tool.py +++ /dev/null @@ -1,32 +0,0 @@ -from abc import ABC, abstractmethod - -from pydantic import BaseModel, ConfigDict -from sqlalchemy.orm import Session - -from core.callback_handler.index_tool_callback_handler import DatasetIndexToolCallbackHandler - - -class DatasetRetrieverBaseTool(BaseModel, ABC): - """Tool for querying a Dataset.""" - - name: str = "dataset" - description: str = "use this to retrieve a dataset. " - tenant_id: str - top_k: int = 4 - score_threshold: float | None = None - hit_callbacks: list[DatasetIndexToolCallbackHandler] = [] - return_resource: bool - retriever_from: str - model_config = ConfigDict(arbitrary_types_allowed=True) - - def run(self, session: Session, query: str) -> str: - """Use the tool.""" - return self._run(session, query) - - @abstractmethod - def _run(self, session: Session, query: str) -> str: - """Use the tool. - - Add run_manager: Optional[CallbackManagerForToolRun] = None - to child implementations to enable tracing, - """ diff --git a/api/core/tools/utils/dataset_retriever/dataset_retriever_tool.py b/api/core/tools/utils/dataset_retriever/dataset_retriever_tool.py deleted file mode 100644 index ff0a1c16930b31..00000000000000 --- a/api/core/tools/utils/dataset_retriever/dataset_retriever_tool.py +++ /dev/null @@ -1,248 +0,0 @@ -from typing import Any, cast, override - -from pydantic import BaseModel, Field -from sqlalchemy import select -from sqlalchemy.orm import Session - -from core.app.app_config.entities import DatasetRetrieveConfigEntity, ModelConfig -from core.rag.datasource.retrieval_service import DefaultRetrievalModelDict, RetrievalService -from core.rag.entities import DocumentContext, RetrievalSourceMetadata -from core.rag.index_processor.constant.index_type import IndexTechniqueType -from core.rag.models.document import Document as RetrievalDocument -from core.rag.retrieval.dataset_retrieval import DatasetRetrieval -from core.rag.retrieval.retrieval_methods import RetrievalMethod -from core.tools.utils.dataset_retriever.dataset_retriever_base_tool import DatasetRetrieverBaseTool -from models.dataset import Dataset -from models.dataset import Document as DatasetDocument -from repositories.knowledge.segment_read_adapter import sign_segment_content -from services.knowledge.external.service import ExternalDatasetService - -default_retrieval_model: DefaultRetrievalModelDict = { - "search_method": RetrievalMethod.SEMANTIC_SEARCH, - "reranking_enable": False, - "reranking_model": {"reranking_provider_name": "", "reranking_model_name": ""}, - "reranking_mode": "reranking_model", - "top_k": 2, - "score_threshold_enabled": False, -} - - -class DatasetRetrieverToolInput(BaseModel): - query: str = Field(..., description="Query for the dataset to be used to retrieve the dataset.") - - -class DatasetRetrieverTool(DatasetRetrieverBaseTool): - """Tool for querying a Dataset.""" - - name: str = "dataset" - args_schema: type[BaseModel] = DatasetRetrieverToolInput - description: str = "use this to retrieve a dataset. " - dataset_id: str - user_id: str | None = None - retrieve_config: DatasetRetrieveConfigEntity - inputs: dict[str, Any] - - @classmethod - def from_dataset(cls, dataset: Dataset, **kwargs): - description = dataset.description - if not description: - description = "useful for when you want to answer queries about the " + dataset.name - - description = description.replace("\n", "").replace("\r", "") - return cls( - name=f"dataset_{dataset.id.replace('-', '_')}", - tenant_id=dataset.tenant_id, - dataset_id=dataset.id, - description=description, - **kwargs, - ) - - @override - def _run(self, session: Session, query: str) -> str: - dataset_stmt = select(Dataset).where(Dataset.tenant_id == self.tenant_id, Dataset.id == self.dataset_id) - dataset = session.scalar(dataset_stmt) - - if not dataset: - return "" - for hit_callback in self.hit_callbacks: - hit_callback.on_query(query, dataset.id, session) - dataset_retrieval = DatasetRetrieval() - metadata_filter_document_ids, metadata_condition = dataset_retrieval.get_metadata_filter_condition( - session, - [dataset.id], - query, - self.tenant_id, - self.user_id or "unknown", - cast(str, self.retrieve_config.metadata_filtering_mode), - cast(ModelConfig, self.retrieve_config.metadata_model_config), - self.retrieve_config.metadata_filtering_conditions, - self.inputs, - ) - if metadata_filter_document_ids: - document_ids_filter = metadata_filter_document_ids.get(dataset.id, []) - else: - document_ids_filter = None - if dataset.provider == "external": - results: list[RetrievalDocument] = [] - external_documents = ExternalDatasetService.fetch_external_knowledge_retrieval( - session=session, - tenant_id=dataset.tenant_id, - dataset_id=dataset.id, - query=query, - external_retrieval_parameters=dataset.retrieval_model, - metadata_condition=metadata_condition, - ) - for external_document in external_documents: - document = RetrievalDocument( - page_content=external_document.get("content"), - metadata=external_document.get("metadata"), - provider="external", - ) - if document.metadata is not None: - document.metadata["score"] = external_document.get("score") - document.metadata["title"] = external_document.get("title") - document.metadata["dataset_id"] = dataset.id - document.metadata["dataset_name"] = dataset.name - results.append(document) - # deal with external documents - context_list: list[RetrievalSourceMetadata] = [] - for position, item in enumerate(results, start=1): - if item.metadata is not None: - source = RetrievalSourceMetadata( - position=position, - dataset_id=item.metadata.get("dataset_id"), - dataset_name=item.metadata.get("dataset_name"), - document_id=item.metadata.get("document_id") or item.metadata.get("title"), - document_name=item.metadata.get("title"), - data_source_type="external", - retriever_from=self.retriever_from, - score=item.metadata.get("score"), - title=item.metadata.get("title"), - content=item.page_content, - ) - context_list.append(source) - for hit_callback in self.hit_callbacks: - hit_callback.return_retriever_resource_info(context_list) - - return "\n".join([item.page_content for item in results]) - else: - if metadata_condition and not document_ids_filter: - return "" - # get retrieval model , if the model is not setting , using default - retrieval_model = dataset.retrieval_model or default_retrieval_model - retrieval_resource_list: list[RetrievalSourceMetadata] = [] - if dataset.indexing_technique == IndexTechniqueType.ECONOMY: - # use keyword table query - documents = RetrievalService.retrieve( - retrieval_method=RetrievalMethod.KEYWORD_SEARCH, - dataset_id=dataset.id, - query=query, - top_k=self.top_k, - document_ids_filter=document_ids_filter, - ) - return "\n".join([document.page_content for document in documents]) - else: - if self.top_k > 0: - # retrieval source - documents = RetrievalService.retrieve( - retrieval_method=retrieval_model.get("search_method", "semantic_search"), - dataset_id=dataset.id, - query=query, - top_k=self.top_k, - score_threshold=retrieval_model.get("score_threshold", 0.0) - if retrieval_model["score_threshold_enabled"] - else 0.0, - reranking_model=retrieval_model.get("reranking_model") - if retrieval_model["reranking_enable"] - else None, - reranking_mode=retrieval_model.get("reranking_mode") or "reranking_model", - weights=retrieval_model.get("weights"), - document_ids_filter=document_ids_filter, - ) - else: - documents = [] - for hit_callback in self.hit_callbacks: - hit_callback.on_tool_end(documents, session) - document_score_list = {} - if dataset.indexing_technique != IndexTechniqueType.ECONOMY: - for item in documents: - if item.metadata is not None and item.metadata.get("score"): - document_score_list[item.metadata["doc_id"]] = item.metadata["score"] - document_context_list: list[DocumentContext] = [] - with Session(bind=session.get_bind()) as format_session: - records = RetrievalService.format_retrieval_documents(format_session, documents) - if records: - for record in records: - segment = record.segment - # Build content: if summary exists, add it before the segment content - if segment.answer: - segment_content = ( - f"question:{sign_segment_content(segment, session=session)} answer:{segment.answer}" - ) - else: - segment_content = sign_segment_content(segment, session=session) - - # If summary exists, prepend it to the content - if record.summary: - final_content = f"{record.summary}\n{segment_content}" - else: - final_content = segment_content - - document_context_list.append( - DocumentContext( - content=final_content, - score=record.score, - ) - ) - - if self.return_resource: - for record in records: - segment = record.segment - dataset = session.get(Dataset, segment.dataset_id) - dataset_document_stmt = select(DatasetDocument).where( - DatasetDocument.id == segment.document_id, - DatasetDocument.enabled == True, - DatasetDocument.archived == False, - ) - document = session.scalar(dataset_document_stmt) - if dataset and document: - source = RetrievalSourceMetadata( - dataset_id=dataset.id, - dataset_name=dataset.name, - document_id=document.id, - document_name=document.name, - data_source_type=document.data_source_type, - segment_id=segment.id, - retriever_from=self.retriever_from, - score=record.score or 0.0, - doc_metadata=document.doc_metadata, - ) - - if self.retriever_from == "dev": - source.hit_count = segment.hit_count - source.word_count = segment.word_count - source.segment_position = segment.position - source.index_node_hash = segment.index_node_hash - if segment.answer: - source.content = f"question:{segment.content} \nanswer:{segment.answer}" - else: - source.content = segment.content - # Add summary if this segment was retrieved via summary - if hasattr(record, "summary") and record.summary: - source.summary = record.summary - retrieval_resource_list.append(source) - - if self.return_resource and retrieval_resource_list: - retrieval_resource_list = sorted( - retrieval_resource_list, - key=lambda x: x.score or 0.0, - reverse=True, - ) - for position, item in enumerate(retrieval_resource_list, start=1): # type: ignore - item.position = position # type: ignore - for hit_callback in self.hit_callbacks: - hit_callback.return_retriever_resource_info(retrieval_resource_list) - if document_context_list: - document_context_list = sorted(document_context_list, key=lambda x: x.score or 0.0, reverse=True) - return "\n".join([document_context.content for document_context in document_context_list]) - return "" diff --git a/api/core/tools/utils/dataset_retriever_tool.py b/api/core/tools/utils/dataset_retriever_tool.py deleted file mode 100644 index cd4f2849df9039..00000000000000 --- a/api/core/tools/utils/dataset_retriever_tool.py +++ /dev/null @@ -1,144 +0,0 @@ -from collections.abc import Generator -from typing import Any, override - -from sqlalchemy.orm import Session - -from core.app.app_config.entities import DatasetRetrieveConfigEntity -from core.app.entities.app_invoke_entities import InvokeFrom -from core.callback_handler.index_tool_callback_handler import DatasetIndexToolCallbackHandler -from core.rag.retrieval.dataset_retrieval import DatasetRetrieval -from core.tools.__base.tool import Tool -from core.tools.__base.tool_runtime import ToolRuntime -from core.tools.entities.common_entities import I18nObject -from core.tools.entities.tool_entities import ( - ToolDescription, - ToolEntity, - ToolIdentity, - ToolInvokeMessage, - ToolParameter, - ToolProviderType, -) -from core.tools.utils.dataset_retriever.dataset_retriever_base_tool import DatasetRetrieverBaseTool - - -class DatasetRetrieverTool(Tool): - def __init__(self, entity: ToolEntity, runtime: ToolRuntime, retrieval_tool: DatasetRetrieverBaseTool): - super().__init__(entity, runtime) - self.retrieval_tool = retrieval_tool - - @staticmethod - def get_dataset_tools( - session: Session, - tenant_id: str, - dataset_ids: list[str], - retrieve_config: DatasetRetrieveConfigEntity | None, - return_resource: bool, - invoke_from: InvokeFrom, - hit_callback: DatasetIndexToolCallbackHandler, - user_id: str, - inputs: dict[str, Any], - ) -> list["DatasetRetrieverTool"]: - """ - get dataset tool - """ - # check if retrieve_config is valid - if dataset_ids is None or len(dataset_ids) == 0: - return [] - if retrieve_config is None: - return [] - - feature = DatasetRetrieval() - - # save original retrieve strategy, and set retrieve strategy to SINGLE - # Agent only support SINGLE mode - original_retriever_mode = retrieve_config.retrieve_strategy - retrieve_config.retrieve_strategy = DatasetRetrieveConfigEntity.RetrieveStrategy.SINGLE - retrieval_tools = feature.to_dataset_retriever_tool( - session=session, - tenant_id=tenant_id, - dataset_ids=dataset_ids, - retrieve_config=retrieve_config, - return_resource=return_resource, - invoke_from=invoke_from, - hit_callback=hit_callback, - user_id=user_id, - inputs=inputs, - ) - if retrieval_tools is None or len(retrieval_tools) == 0: - return [] - - # restore retrieve strategy - retrieve_config.retrieve_strategy = original_retriever_mode - - # convert retrieval tools to Tools - tools = [] - for retrieval_tool in retrieval_tools: - tool = DatasetRetrieverTool( - retrieval_tool=retrieval_tool, - entity=ToolEntity( - identity=ToolIdentity( - provider="", author="", name=retrieval_tool.name, label=I18nObject(en_US="", zh_Hans="") - ), - parameters=[], - description=ToolDescription(human=I18nObject(en_US="", zh_Hans=""), llm=retrieval_tool.description), - ), - runtime=ToolRuntime(tenant_id=tenant_id), - ) - - tools.append(tool) - - return tools - - @override - def get_runtime_parameters( - self, - conversation_id: str | None = None, - app_id: str | None = None, - message_id: str | None = None, - ) -> list[ToolParameter]: - return [ - ToolParameter( - name="query", - label=I18nObject(en_US="", zh_Hans=""), - human_description=I18nObject(en_US="", zh_Hans=""), - type=ToolParameter.ToolParameterType.STRING, - form=ToolParameter.ToolParameterForm.LLM, - llm_description="Query for the dataset to be used to retrieve the dataset.", - required=True, - default="", - placeholder=I18nObject(en_US="", zh_Hans=""), - ), - ] - - @override - def tool_provider_type(self) -> ToolProviderType: - return ToolProviderType.DATASET_RETRIEVAL - - @override - def _invoke( - self, - session: Session, - user_id: str, - tool_parameters: dict[str, Any], - conversation_id: str | None = None, - app_id: str | None = None, - message_id: str | None = None, - ) -> Generator[ToolInvokeMessage, None, None]: - """ - invoke dataset retriever tool - """ - query = tool_parameters.get("query") - if not query: - yield self.create_text_message(text="please input query") - else: - # invoke dataset retriever tool - result = self.retrieval_tool.run(session=session, query=query) - yield self.create_text_message(text=result) - - def validate_credentials( - self, credentials: dict[str, Any], parameters: dict[str, Any], format_only: bool = False - ) -> str | None: - """ - validate the credentials for dataset retriever tool - """ - pass diff --git a/api/core/tools/workflow_as_tool/provider.py b/api/core/tools/workflow_as_tool/provider.py deleted file mode 100644 index 1611cd1d63a46d..00000000000000 --- a/api/core/tools/workflow_as_tool/provider.py +++ /dev/null @@ -1,259 +0,0 @@ -from __future__ import annotations - -from collections.abc import Mapping -from typing import override - -from sqlalchemy import select -from sqlalchemy.orm import Session - -from core.app.apps.workflow.app_config_manager import WorkflowAppConfigManager -from core.db.session_factory import session_factory -from core.plugin.entities.parameters import PluginParameterOption -from core.tools.__base.tool_provider import ToolProviderController -from core.tools.__base.tool_runtime import ToolRuntime -from core.tools.entities.common_entities import I18nObject -from core.tools.entities.tool_entities import ( - ToolDescription, - ToolEntity, - ToolIdentity, - ToolParameter, - ToolProviderEntity, - ToolProviderIdentity, - ToolProviderType, -) -from core.tools.utils.workflow_configuration_sync import WorkflowToolConfigurationUtils -from core.tools.workflow_as_tool.tool import WorkflowTool -from extensions.ext_database import db -from graphon.variables.input_entities import VariableEntity, VariableEntityType -from models.account import Account -from models.model import App, AppMode -from models.tools import WorkflowToolProvider -from models.workflow import Workflow - -VARIABLE_TO_PARAMETER_TYPE_MAPPING = { - VariableEntityType.TEXT_INPUT: ToolParameter.ToolParameterType.STRING, - VariableEntityType.PARAGRAPH: ToolParameter.ToolParameterType.STRING, - VariableEntityType.SELECT: ToolParameter.ToolParameterType.SELECT, - VariableEntityType.NUMBER: ToolParameter.ToolParameterType.NUMBER, - VariableEntityType.CHECKBOX: ToolParameter.ToolParameterType.BOOLEAN, - VariableEntityType.FILE: ToolParameter.ToolParameterType.FILE, - VariableEntityType.FILE_LIST: ToolParameter.ToolParameterType.FILES, - VariableEntityType.JSON_OBJECT: ToolParameter.ToolParameterType.OBJECT, -} - - -class WorkflowToolProviderController(ToolProviderController[ToolProviderEntity, WorkflowTool | None]): - provider_id: str - tools: list[WorkflowTool] | None - - def __init__(self, entity: ToolProviderEntity, provider_id: str): - super().__init__(entity=entity) - self.provider_id = provider_id - self.tools = None - - @classmethod - def from_db(cls, db_provider: WorkflowToolProvider) -> WorkflowToolProviderController: - with session_factory.create_session() as session, session.begin(): - app = session.get(App, db_provider.app_id) - if not app: - raise ValueError("app not found") - - user = session.get(Account, db_provider.user_id) if db_provider.user_id else None - controller = WorkflowToolProviderController( - entity=ToolProviderEntity( - identity=ToolProviderIdentity( - author=user.name if user else "", - name=db_provider.label, - label=I18nObject(en_US=db_provider.label, zh_Hans=db_provider.label), - description=I18nObject(en_US=db_provider.description, zh_Hans=db_provider.description), - icon=db_provider.icon, - ), - credentials_schema=[], - plugin_id=None, - ), - provider_id=db_provider.id, - ) - - controller.tools = [ - controller._get_db_provider_tool(db_provider, app, session=session, user=user), - ] - - return controller - - @property - @override - def provider_type(self) -> ToolProviderType: - return ToolProviderType.WORKFLOW - - def _get_db_provider_tool( - self, - db_provider: WorkflowToolProvider, - app: App, - *, - session: Session, - user: Account | None = None, - ) -> WorkflowTool: - """ - get db provider tool - :param db_provider: the db provider - :param app: the app - :return: the tool - """ - workflow: Workflow | None = session.scalar( - select(Workflow) - .where(Workflow.app_id == db_provider.app_id, Workflow.version == db_provider.version) - .limit(1) - ) - - if not workflow: - raise ValueError("workflow not found") - - # fetch start node - graph: Mapping = workflow.graph_dict - features_dict: Mapping = workflow.features_dict - features = WorkflowAppConfigManager.convert_features(config_dict=features_dict, app_mode=AppMode.WORKFLOW) - - parameters = db_provider.parameter_configurations - variables = WorkflowToolConfigurationUtils.get_workflow_graph_variables(graph) - - def fetch_workflow_variable(variable_name: str) -> VariableEntity | None: - return next(filter(lambda x: x.variable == variable_name, variables), None) - - workflow_tool_parameters = [] - for parameter in parameters: - variable = fetch_workflow_variable(parameter.name) - if variable: - parameter_type = None - options = [] - if variable.type not in VARIABLE_TO_PARAMETER_TYPE_MAPPING: - raise ValueError(f"unsupported variable type {variable.type}") - parameter_type = VARIABLE_TO_PARAMETER_TYPE_MAPPING[variable.type] - - if variable.type == VariableEntityType.SELECT and variable.options: - options = [ - PluginParameterOption(value=option, label=I18nObject(en_US=option, zh_Hans=option)) - for option in variable.options - ] - - workflow_tool_parameters.append( - ToolParameter( - name=parameter.name, - label=I18nObject(en_US=variable.label, zh_Hans=variable.label), - human_description=I18nObject(en_US=parameter.description, zh_Hans=parameter.description), - type=parameter_type, - form=parameter.form, - llm_description=parameter.description, - required=variable.required, - default=variable.default, - options=options, - placeholder=I18nObject(en_US="", zh_Hans=""), - ) - ) - elif features.file_upload: - workflow_tool_parameters.append( - ToolParameter( - name=parameter.name, - label=I18nObject(en_US=parameter.name, zh_Hans=parameter.name), - human_description=I18nObject(en_US=parameter.description, zh_Hans=parameter.description), - type=ToolParameter.ToolParameterType.SYSTEM_FILES, - llm_description=parameter.description, - required=False, - form=parameter.form, - placeholder=I18nObject(en_US="", zh_Hans=""), - ) - ) - else: - raise ValueError("variable not found") - - # get output schema from workflow - outputs = WorkflowToolConfigurationUtils.get_workflow_graph_output(graph) - - reserved_keys = {"json", "text", "files"} - - properties = {} - for output in outputs: - if output.variable not in reserved_keys: - properties[output.variable] = { - "type": output.value_type, - "description": "", - } - output_schema = {"type": "object", "properties": properties} - - return WorkflowTool( - workflow_as_tool_id=db_provider.id, - entity=ToolEntity( - identity=ToolIdentity( - author=user.name if user else "", - name=db_provider.name, - label=I18nObject(en_US=db_provider.label, zh_Hans=db_provider.label), - provider=self.provider_id, - icon=db_provider.icon, - ), - description=ToolDescription( - human=I18nObject(en_US=db_provider.description, zh_Hans=db_provider.description), - llm=db_provider.description, - ), - parameters=workflow_tool_parameters, - output_schema=output_schema, - ), - runtime=ToolRuntime( - tenant_id=db_provider.tenant_id, - ), - workflow_app_id=app.id, - workflow_entities={ - "app": app, - "workflow": workflow, - }, - version=db_provider.version, - workflow_call_depth=0, - label=db_provider.label, - ) - - def get_tools(self, tenant_id: str) -> list[WorkflowTool]: - """ - fetch tools from database - - :param tenant_id: the tenant id - :return: the tools - """ - if self.tools is not None: - return self.tools - - with Session(db.engine, expire_on_commit=False) as session, session.begin(): - db_provider: WorkflowToolProvider | None = session.scalar( - select(WorkflowToolProvider) - .where( - WorkflowToolProvider.tenant_id == tenant_id, - WorkflowToolProvider.id == self.provider_id, - ) - .limit(1) - ) - - if not db_provider: - return [] - - app = session.get(App, db_provider.app_id) - if not app: - raise ValueError("app not found") - - user = session.get(Account, db_provider.user_id) if db_provider.user_id else None - self.tools = [self._get_db_provider_tool(db_provider, app, session=session, user=user)] - - return self.tools - - @override - def get_tool(self, tool_name: str) -> WorkflowTool | None: - """ - get tool by name - - :param tool_name: the name of the tool - :return: the tool - """ - if self.tools is None: - return None - - for tool in self.tools: - if tool.entity.identity.name == tool_name: - return tool - - return None diff --git a/api/core/tools/workflow_as_tool/tool.py b/api/core/tools/workflow_as_tool/tool.py deleted file mode 100644 index b5a70148b8b77c..00000000000000 --- a/api/core/tools/workflow_as_tool/tool.py +++ /dev/null @@ -1,436 +0,0 @@ -from __future__ import annotations - -import json -import logging -from collections.abc import Generator, Mapping, Sequence -from typing import Any, cast, override - -from sqlalchemy import select -from sqlalchemy.orm import Session - -from core.app.file_access import DatabaseFileAccessController -from core.db.session_factory import session_factory -from core.helper.trace_id_helper import ( - ParentTraceContext, - extract_parent_trace_context_from_args, - extract_trace_session_id_from_args, -) -from core.tools.__base.tool import Tool -from core.tools.__base.tool_runtime import ToolRuntime -from core.tools.entities.tool_entities import ( - ToolEntity, - ToolInvokeMessage, - ToolParameter, - ToolProviderType, -) -from core.tools.errors import ToolInvokeError -from core.workflow.file_reference import resolve_file_record_id -from factories.file_factory import build_from_mapping -from graphon.file import FILE_MODEL_IDENTITY, File, FileTransferMethod -from graphon.model_runtime.entities.llm_entities import LLMUsage, LLMUsageMetadata -from models import Account, Tenant -from models.model import App, EndUser -from models.utils.file_input_compat import build_file_from_stored_mapping -from models.workflow import Workflow - -logger = logging.getLogger(__name__) -_file_access_controller = DatabaseFileAccessController() - - -class WorkflowTool(Tool): - """ - Workflow tool. - """ - - _parent_trace_context: ParentTraceContext | None - _trace_session_id: str | None - - def __init__( - self, - workflow_app_id: str, - workflow_as_tool_id: str, - version: str, - workflow_entities: dict[str, Any], - workflow_call_depth: int, - entity: ToolEntity, - runtime: ToolRuntime, - label: str = "Workflow", - ): - self.workflow_app_id = workflow_app_id - self.workflow_as_tool_id = workflow_as_tool_id - self.version = version - self.workflow_entities = workflow_entities - self.workflow_call_depth = workflow_call_depth - self.label = label - self._latest_usage = LLMUsage.empty_usage() - self._parent_trace_context = None - self._trace_session_id = None - - super().__init__(entity=entity, runtime=runtime) - - @override - def tool_provider_type(self) -> ToolProviderType: - """ - get the tool provider type - - :return: the tool provider type - """ - return ToolProviderType.WORKFLOW - - @override - def _invoke( - self, - session: Session, - user_id: str, - tool_parameters: dict[str, Any], - conversation_id: str | None = None, - app_id: str | None = None, - message_id: str | None = None, - ) -> Generator[ToolInvokeMessage, None, None]: - """ - invoke the tool - """ - app = self._get_app(app_id=self.workflow_app_id) - workflow = self._get_workflow(app_id=self.workflow_app_id, version=self.version) - - # transform the tool parameters - tool_parameters, files = self._transform_args(tool_parameters=tool_parameters) - - from core.app.apps.workflow.app_generator import WorkflowAppGenerator - - generator = WorkflowAppGenerator() - assert self.runtime is not None - assert self.runtime.invoke_from is not None - - user = self._resolve_user(user_id=user_id) - if user is None: - raise ToolInvokeError("User not found") - - self._latest_usage = LLMUsage.empty_usage() - - generator_args: dict[str, Any] = {"inputs": tool_parameters, "files": files} - if self._parent_trace_context: - generator_args.update( - extract_parent_trace_context_from_args({"parent_trace_context": self._parent_trace_context}) - ) - if self._trace_session_id: - generator_args.update(extract_trace_session_id_from_args({"trace_session_id": self._trace_session_id})) - - result = generator.generate( - app_model=app, - workflow=workflow, - user=user, - args=generator_args, - invoke_from=self.runtime.invoke_from, - streaming=False, - call_depth=self.workflow_call_depth + 1, - # NOTE(QuantumGhost): We explicitly set `pause_state_config` to `None` - # because workflow pausing mechanisms (such as HumanInput) are not - # supported within WorkflowTool execution context. - pause_state_config=None, - ) - assert isinstance(result, dict) - data = result.get("data", {}) - - if err := data.get("error"): - raise ToolInvokeError(err) - - outputs = data.get("outputs") - if outputs is None: - outputs = {} - else: - outputs, files = self._extract_files(outputs) # type: ignore - for file in files: - yield self.create_file_message(file) # type: ignore - - # traverse `outputs` field and create variable messages - for key, value in outputs.items(): - if key not in {"text", "json", "files"}: - yield self.create_variable_message(variable_name=key, variable_value=value) - - self._latest_usage = self._derive_usage_from_result(data) - - yield self.create_text_message(json.dumps(outputs, ensure_ascii=False)) - yield self.create_json_message(outputs, suppress_output=True) - - @property - def latest_usage(self) -> LLMUsage: - return self._latest_usage - - @classmethod - def _derive_usage_from_result(cls, data: Mapping[str, Any]) -> LLMUsage: - usage_dict = cls._extract_usage_dict(data) - if usage_dict is not None: - return LLMUsage.from_metadata(cast(LLMUsageMetadata, dict(usage_dict))) - - total_tokens = data.get("total_tokens") - total_price = data.get("total_price") - if total_tokens is None and total_price is None: - return LLMUsage.empty_usage() - - usage_metadata: dict[str, Any] = {} - if total_tokens is not None: - try: - usage_metadata["total_tokens"] = int(str(total_tokens)) - except (TypeError, ValueError): - pass - if total_price is not None: - usage_metadata["total_price"] = str(total_price) - currency = data.get("currency") - if currency is not None: - usage_metadata["currency"] = currency - - if not usage_metadata: - return LLMUsage.empty_usage() - - return LLMUsage.from_metadata(cast(LLMUsageMetadata, usage_metadata)) - - @classmethod - def _extract_usage_dict(cls, payload: Mapping[str, Any]) -> Mapping[str, Any] | None: - usage_candidate = payload.get("usage") - if isinstance(usage_candidate, Mapping): - return usage_candidate - - metadata_candidate = payload.get("metadata") - if isinstance(metadata_candidate, Mapping): - usage_candidate = metadata_candidate.get("usage") - if isinstance(usage_candidate, Mapping): - return usage_candidate - - for value in payload.values(): - match value: - case _ if isinstance(value, Mapping): - found = cls._extract_usage_dict(value) - if found is not None: - return found - case _ if isinstance(value, Sequence) and not isinstance(value, (str, bytes, bytearray)): - for item in value: - if isinstance(item, Mapping): - found = cls._extract_usage_dict(item) - if found is not None: - return found - return None - - @override - def fork_tool_runtime(self, runtime: ToolRuntime) -> WorkflowTool: - """ - fork a new tool with metadata - - :return: the new tool - """ - forked = self.__class__( - entity=self.entity.model_copy(), - runtime=runtime, - workflow_app_id=self.workflow_app_id, - workflow_as_tool_id=self.workflow_as_tool_id, - workflow_entities=self.workflow_entities, - workflow_call_depth=self.workflow_call_depth, - version=self.version, - label=self.label, - ) - forked._parent_trace_context = self._parent_trace_context.model_copy() if self._parent_trace_context else None - forked._trace_session_id = self._trace_session_id - return forked - - def set_parent_trace_context( - self, - *, - parent_workflow_run_id: str, - parent_node_execution_id: str, - ) -> None: - """Attach outer workflow trace context without exposing it as tool input.""" - self._parent_trace_context = ParentTraceContext( - parent_workflow_run_id=parent_workflow_run_id, - parent_node_execution_id=parent_node_execution_id, - ) - - def clear_parent_trace_context(self) -> None: - """Remove parent trace context before invoking this tool outside a nested workflow.""" - self._parent_trace_context = None - - def set_trace_session_id(self, trace_session_id: str) -> None: - """Attach parent trace session ID without exposing it as tool input.""" - self._trace_session_id = trace_session_id - - def clear_trace_session_id(self) -> None: - """Remove trace session ID before invoking this tool outside a traced session.""" - self._trace_session_id = None - - def _resolve_user(self, user_id: str) -> Account | EndUser | None: - """ - Resolve user object in both HTTP and worker contexts. - - In HTTP context: dereference the current_user LocalProxy (can return Account or EndUser). - In worker context: load Account(knowledge pipeline) or EndUser(trigger) from database by user_id. - - Returns: - Account | EndUser | None: The resolved user object, or None if resolution fails. - """ - return self._resolve_user_from_database(user_id=user_id) - - def _resolve_user_from_database(self, user_id: str) -> Account | EndUser | None: - """ - Resolve user from database (worker/Celery context). - """ - with session_factory.create_session() as session: - tenant_stmt = select(Tenant).where(Tenant.id == self.runtime.tenant_id) - tenant = session.scalar(tenant_stmt) - if not tenant: - return None - - user_stmt = select(Account).where(Account.id == user_id) - user = session.scalar(user_stmt) - if user: - user.set_current_tenant_with_session(tenant, session=session) - session.expunge(user) - return user - - end_user_stmt = select(EndUser).where(EndUser.id == user_id, EndUser.tenant_id == tenant.id) - end_user = session.scalar(end_user_stmt) - if end_user: - session.expunge(end_user) - return end_user - - return None - - def _get_workflow(self, app_id: str, version: str) -> Workflow: - """ - get the workflow by app id and version - """ - with session_factory.create_session() as session, session.begin(): - if not version: - stmt = ( - select(Workflow) - .where(Workflow.app_id == app_id, Workflow.version != Workflow.VERSION_DRAFT) - .order_by(Workflow.created_at.desc()) - ) - workflow = session.scalars(stmt).first() - else: - stmt = select(Workflow).where(Workflow.app_id == app_id, Workflow.version == version) - workflow = session.scalar(stmt) - - if not workflow: - raise ValueError("workflow not found or not published") - - session.expunge(workflow) - return workflow - - def _get_app(self, app_id: str) -> App: - """ - get the app by app id - """ - stmt = select(App).where(App.id == app_id) - with session_factory.create_session() as session, session.begin(): - app = session.scalar(stmt) - if not app: - raise ValueError("app not found") - - session.expunge(app) - return app - - def _transform_args(self, tool_parameters: dict[str, Any]) -> tuple[dict[str, Any], list[dict[str, str | None]]]: - """ - transform the tool parameters - - :param tool_parameters: the tool parameters - :return: tool_parameters, files - """ - parameter_rules = self.get_merged_runtime_parameters() - parameters_result = {} - files = [] - for parameter in parameter_rules: - if parameter.type == ToolParameter.ToolParameterType.SYSTEM_FILES: - file = tool_parameters.get(parameter.name) - if file: - try: - file_var_list = [ - build_file_from_stored_mapping( - file_mapping=cast(Mapping[str, Any], f), - tenant_id=str(self.runtime.tenant_id), - ) - for f in file - if isinstance(f, Mapping) - ] - for file in file_var_list: - file_dict: dict[str, str | None] = { - "transfer_method": file.transfer_method.value, - "type": file.type.value, - } - match file.transfer_method: - case FileTransferMethod.TOOL_FILE: - file_dict["tool_file_id"] = resolve_file_record_id(file.reference) - case FileTransferMethod.LOCAL_FILE: - file_dict["upload_file_id"] = resolve_file_record_id(file.reference) - case FileTransferMethod.DATASOURCE_FILE: - file_dict["datasource_file_id"] = resolve_file_record_id(file.reference) - case FileTransferMethod.REMOTE_URL: - file_dict["url"] = file.generate_url() - - files.append(file_dict) - except Exception: - logger.exception("Failed to transform file %s", file) - elif parameter.type == ToolParameter.ToolParameterType.FILES: - value = tool_parameters.get(parameter.name) - if not parameter.required and self._is_empty_files_parameter_value(value): - value = [] - parameters_result[parameter.name] = value - else: - parameters_result[parameter.name] = tool_parameters.get(parameter.name) - - return parameters_result, files - - @staticmethod - def _is_empty_files_parameter_value(value: Any) -> bool: - """Identify empty optional file-list placeholders before workflow input validation.""" - - if value is None or value == "": - return True - return isinstance(value, list) and all(item is None or item == "" for item in value) - - def _extract_files(self, outputs: dict[str, Any]) -> tuple[dict[str, Any], list[File]]: - """ - extract files from the result - - :return: the result, files - """ - files: list[File] = [] - result = {} - for key, value in outputs.items(): - match value: - case list(): - for item in value: - if isinstance(item, dict) and item.get("dify_model_identity") == FILE_MODEL_IDENTITY: - item = self._update_file_mapping(item) - file = build_from_mapping( - mapping=item, - tenant_id=str(self.runtime.tenant_id), - access_controller=_file_access_controller, - ) - files.append(file) - case dict() if value.get("dify_model_identity") == FILE_MODEL_IDENTITY: - value = self._update_file_mapping(value) - file = build_from_mapping( - mapping=value, - tenant_id=str(self.runtime.tenant_id), - access_controller=_file_access_controller, - ) - files.append(file) - - result[key] = value - - return result, files - - def _update_file_mapping(self, file_dict: dict[str, Any]) -> dict[str, Any]: - file_id = resolve_file_record_id(file_dict.get("reference") or file_dict.get("related_id")) - transfer_method_value = file_dict.get("transfer_method") - if not isinstance(transfer_method_value, str): - raise ValueError("Workflow file mapping is missing a valid transfer_method") - transfer_method = FileTransferMethod.value_of(transfer_method_value) - match transfer_method: - case FileTransferMethod.TOOL_FILE: - file_dict["tool_file_id"] = file_id - case FileTransferMethod.LOCAL_FILE: - file_dict["upload_file_id"] = file_id - case FileTransferMethod.REMOTE_URL | FileTransferMethod.DATASOURCE_FILE: - pass - return file_dict diff --git a/api/events/event_handlers/sync_plugin_trigger_when_app_created.py b/api/events/event_handlers/sync_plugin_trigger_when_app_created.py deleted file mode 100644 index cc6e851ad870c2..00000000000000 --- a/api/events/event_handlers/sync_plugin_trigger_when_app_created.py +++ /dev/null @@ -1,34 +0,0 @@ -import logging - -from events.app_event import app_draft_workflow_was_synced, app_published_workflow_was_updated -from models.model import App, AppMode -from models.workflow import Workflow -from services.trigger.trigger_service import TriggerService - -logger = logging.getLogger(__name__) - - -@app_draft_workflow_was_synced.connect -@app_published_workflow_was_updated.connect -def handle( - sender, - synced_draft_workflow: Workflow | None = None, - published_workflow: Workflow | None = None, - **kwargs, -): - """ - Sync plugin trigger relationships when a draft changes or is published. - - The published workflow must be reconciled as well because production trigger - dispatch relies on these relationships, while debug dispatch does not. - """ - app: App = sender - if app.mode != AppMode.WORKFLOW.value: - # only handle workflow app, chatflow is not supported yet - return - - workflow = published_workflow if published_workflow is not None else synced_draft_workflow - if workflow is None: - return - - TriggerService.sync_plugin_trigger_relationships(app, workflow) diff --git a/api/events/event_handlers/sync_workflow_schedule_when_app_published.py b/api/events/event_handlers/sync_workflow_schedule_when_app_published.py deleted file mode 100644 index 5f8fcd86170b7a..00000000000000 --- a/api/events/event_handlers/sync_workflow_schedule_when_app_published.py +++ /dev/null @@ -1,83 +0,0 @@ -import logging -from typing import cast - -from sqlalchemy import select -from sqlalchemy.orm import sessionmaker - -from core.workflow.nodes.trigger_schedule.entities import SchedulePlanUpdate -from events.app_event import app_published_workflow_was_updated -from extensions.ext_database import db -from models import AppMode, Workflow, WorkflowSchedulePlan -from services.trigger.schedule_service import ScheduleService - -logger = logging.getLogger(__name__) - - -@app_published_workflow_was_updated.connect -def handle(sender, **kwargs): - """ - Handle app published workflow update event to sync workflow_schedule_plans table. - - When a workflow is published, this handler will: - 1. Extract schedule trigger nodes from the workflow graph - 2. Compare with existing workflow_schedule_plans records - 3. Create/update/delete schedule plans as needed - """ - app = sender - if app.mode != AppMode.WORKFLOW.value: - return - - published_workflow = kwargs.get("published_workflow") - published_workflow = cast(Workflow, published_workflow) - - sync_schedule_from_workflow(tenant_id=app.tenant_id, app_id=app.id, workflow=published_workflow) - - -def sync_schedule_from_workflow(tenant_id: str, app_id: str, workflow: Workflow) -> WorkflowSchedulePlan | None: - """ - Sync schedule plan from workflow graph configuration. - - Args: - tenant_id: Tenant ID - app_id: App ID - workflow: Published workflow instance - - Returns: - Updated or created WorkflowSchedulePlan, or None if no schedule node - """ - with sessionmaker(db.engine).begin() as session: - schedule_config = ScheduleService.extract_schedule_config(workflow) - - existing_plan = session.scalar( - select(WorkflowSchedulePlan).where( - WorkflowSchedulePlan.tenant_id == tenant_id, - WorkflowSchedulePlan.app_id == app_id, - ) - ) - - if not schedule_config: - if existing_plan: - logger.info("No schedule node in workflow for app %s, removing schedule plan", app_id) - ScheduleService.delete_schedule(session=session, schedule_id=existing_plan.id) - return None - - if existing_plan: - updates = SchedulePlanUpdate( - node_id=schedule_config.node_id, - cron_expression=schedule_config.cron_expression, - timezone=schedule_config.timezone, - ) - updated_plan = ScheduleService.update_schedule( - session=session, - schedule_id=existing_plan.id, - updates=updates, - ) - return updated_plan - else: - new_plan = ScheduleService.create_schedule( - session=session, - tenant_id=tenant_id, - app_id=app_id, - config=schedule_config, - ) - return new_plan diff --git a/api/events/event_handlers/update_app_dataset_join_when_app_published_workflow_updated.py b/api/events/event_handlers/update_app_dataset_join_when_app_published_workflow_updated.py deleted file mode 100644 index 6769b94cde0093..00000000000000 --- a/api/events/event_handlers/update_app_dataset_join_when_app_published_workflow_updated.py +++ /dev/null @@ -1,69 +0,0 @@ -from typing import cast - -from sqlalchemy import delete, select - -from core.workflow.nodes.knowledge_retrieval.entities import KnowledgeRetrievalNodeData -from events.app_event import app_published_workflow_was_updated -from extensions.ext_database import db -from graphon.nodes import BuiltinNodeTypes -from models.dataset import AppDatasetJoin -from models.workflow import Workflow - - -@app_published_workflow_was_updated.connect -def handle(sender, **kwargs): - app = sender - published_workflow = kwargs.get("published_workflow") - published_workflow = cast(Workflow, published_workflow) - - dataset_ids = get_dataset_ids_from_workflow(published_workflow) - app_dataset_joins = db.session.scalars(select(AppDatasetJoin).where(AppDatasetJoin.app_id == app.id)).all() - - removed_dataset_ids: set[str] = set() - if not app_dataset_joins: - added_dataset_ids = dataset_ids - else: - old_dataset_ids: set[str] = set() - old_dataset_ids.update(app_dataset_join.dataset_id for app_dataset_join in app_dataset_joins) - - added_dataset_ids = dataset_ids - old_dataset_ids - removed_dataset_ids = old_dataset_ids - dataset_ids - - if removed_dataset_ids: - for dataset_id in removed_dataset_ids: - db.session.execute( - delete(AppDatasetJoin).where(AppDatasetJoin.app_id == app.id, AppDatasetJoin.dataset_id == dataset_id) - ) - - if added_dataset_ids: - for dataset_id in added_dataset_ids: - app_dataset_join = AppDatasetJoin(app_id=app.id, dataset_id=dataset_id) - db.session.add(app_dataset_join) - - db.session.commit() - - -def get_dataset_ids_from_workflow(published_workflow: Workflow) -> set[str]: - dataset_ids: set[str] = set() - graph = published_workflow.graph_dict - if not graph: - return dataset_ids - - nodes = graph.get("nodes", []) - - # fetch all knowledge retrieval nodes - knowledge_retrieval_nodes = [ - node for node in nodes if node.get("data", {}).get("type") == BuiltinNodeTypes.KNOWLEDGE_RETRIEVAL - ] - - if not knowledge_retrieval_nodes: - return dataset_ids - - for node in knowledge_retrieval_nodes: - try: - node_data = KnowledgeRetrievalNodeData.model_validate(node.get("data", {})) - dataset_ids.update(dataset_id for dataset_id in node_data.dataset_ids) - except Exception: - continue - - return dataset_ids diff --git a/api/events/event_handlers/update_app_triggers_when_app_published_workflow_updated.py b/api/events/event_handlers/update_app_triggers_when_app_published_workflow_updated.py deleted file mode 100644 index d55fe262fb25a3..00000000000000 --- a/api/events/event_handlers/update_app_triggers_when_app_published_workflow_updated.py +++ /dev/null @@ -1,112 +0,0 @@ -from typing import cast - -from sqlalchemy import select -from sqlalchemy.orm import sessionmaker - -from core.trigger.constants import TRIGGER_NODE_TYPES -from events.app_event import app_published_workflow_was_updated -from extensions.ext_database import db -from models import AppMode -from models.enums import AppTriggerStatus -from models.trigger import AppTrigger -from models.workflow import Workflow - - -@app_published_workflow_was_updated.connect -def handle(sender, **kwargs): - """ - Handle app published workflow update event to sync app_triggers table. - - When a workflow is published, this handler will: - 1. Extract trigger nodes from the workflow graph - 2. Compare with existing app_triggers records - 3. Add new triggers and remove obsolete ones - """ - app = sender - if app.mode != AppMode.WORKFLOW.value: - return - - published_workflow = kwargs.get("published_workflow") - published_workflow = cast(Workflow, published_workflow) - # Extract trigger info from workflow - trigger_infos = get_trigger_infos_from_workflow(published_workflow) - - with sessionmaker(db.engine).begin() as session: - # Get existing app triggers - existing_triggers = ( - session.execute( - select(AppTrigger).where(AppTrigger.tenant_id == app.tenant_id, AppTrigger.app_id == app.id) - ) - .scalars() - .all() - ) - - # Convert existing triggers to dict for easy lookup - existing_triggers_map = {trigger.node_id: trigger for trigger in existing_triggers} - - # Get current and new node IDs - existing_node_ids = set(existing_triggers_map.keys()) - new_node_ids = {info["node_id"] for info in trigger_infos} - - # Calculate changes - added_node_ids = new_node_ids - existing_node_ids - removed_node_ids = existing_node_ids - new_node_ids - - # Remove obsolete triggers - for node_id in removed_node_ids: - session.delete(existing_triggers_map[node_id]) - - for trigger_info in trigger_infos: - node_id = trigger_info["node_id"] - - if node_id in added_node_ids: - # Create new trigger - app_trigger = AppTrigger( - tenant_id=app.tenant_id, - app_id=app.id, - trigger_type=trigger_info["node_type"], - title=trigger_info["node_title"], - node_id=node_id, - provider_name=trigger_info.get("node_provider_name", ""), - status=AppTriggerStatus.ENABLED, - ) - session.add(app_trigger) - elif node_id in existing_node_ids: - # Update existing trigger if needed - existing_trigger = existing_triggers_map[node_id] - new_title = trigger_info["node_title"] - if new_title and existing_trigger.title != new_title: - existing_trigger.title = new_title - session.add(existing_trigger) - - -def get_trigger_infos_from_workflow(published_workflow: Workflow) -> list[dict]: - """ - Extract trigger node information from the workflow graph. - - Returns: - List of trigger info dictionaries containing: - - node_type: The type of the trigger node ('trigger-webhook', 'trigger-schedule', 'trigger-plugin') - - node_id: The node ID in the workflow - - node_title: The title of the node - - node_provider_name: The name of the node's provider, only for plugin - """ - graph = published_workflow.graph_dict - if not graph: - return [] - - nodes = graph.get("nodes", []) - trigger_types = TRIGGER_NODE_TYPES - - trigger_infos = [ - { - "node_type": node.get("data", {}).get("type"), - "node_id": node.get("id"), - "node_title": node.get("data", {}).get("title"), - "node_provider_name": node.get("data", {}).get("provider_name"), - } - for node in nodes - if node.get("data", {}).get("type") in trigger_types - ] - - return trigger_infos diff --git a/api/repositories/workflow_app_log_query_repository.py b/api/repositories/workflow_app_log_query_repository.py deleted file mode 100644 index 440337ede47fce..00000000000000 --- a/api/repositories/workflow_app_log_query_repository.py +++ /dev/null @@ -1,261 +0,0 @@ -import uuid -from datetime import datetime -from enum import Enum -from typing import override - -from sqlalchemy import and_, func, literal, or_, select -from sqlalchemy.engine import RowMapping -from sqlalchemy.orm import Session, aliased, sessionmaker - -from graphon.enums import WorkflowExecutionStatus -from libs.helper import escape_like_pattern -from models import Account, EndUser, TenantAccountJoin, WorkflowAppLog, WorkflowRun -from models.enums import CreatorUserRole -from models.trigger import WorkflowTriggerLog -from services.workflow_app_log_query_service import ( - WorkflowAppLogAccount, - WorkflowAppLogEndUser, - WorkflowAppLogItem, - WorkflowAppLogPage, - WorkflowAppLogQuery, - WorkflowAppLogRunSummary, -) - - -class WorkflowAppLogQueryRepository(WorkflowAppLogQuery): - def __init__(self, *, session_factory: sessionmaker[Session]) -> None: - self._session_factory = session_factory - - @override - def get_paginated( - self, - *, - tenant_id: str, - app_id: str, - keyword: str | None = None, - status: WorkflowExecutionStatus | None = None, - created_at_before: datetime | None = None, - created_at_after: datetime | None = None, - page: int = 1, - limit: int = 20, - detail: bool = False, - created_by_end_user_session_id: str | None = None, - created_by_account: str | None = None, - ) -> WorkflowAppLogPage: - with self._session_factory() as session: - workflow_run = aliased(WorkflowRun) - stmt = ( - select( - WorkflowAppLog.id.label("log_id"), - WorkflowAppLog.created_from.label("log_created_from"), - WorkflowAppLog.created_by_role.label("log_created_by_role"), - WorkflowAppLog.created_at.label("log_created_at"), - workflow_run.id.label("run_id"), - workflow_run.version.label("run_version"), - workflow_run.status.label("run_status"), - workflow_run.triggered_from.label("run_triggered_from"), - workflow_run.error.label("run_error"), - workflow_run.elapsed_time.label("run_elapsed_time"), - workflow_run.total_tokens.label("run_total_tokens"), - workflow_run.total_steps.label("run_total_steps"), - workflow_run.created_at.label("run_created_at"), - workflow_run.finished_at.label("run_finished_at"), - workflow_run.exceptions_count.label("run_exceptions_count"), - ) - .select_from(WorkflowAppLog) - .outerjoin( - workflow_run, - and_( - workflow_run.id == WorkflowAppLog.workflow_run_id, - workflow_run.tenant_id == WorkflowAppLog.tenant_id, - workflow_run.app_id == WorkflowAppLog.app_id, - ), - ) - .where( - WorkflowAppLog.tenant_id == tenant_id, - WorkflowAppLog.app_id == app_id, - ) - ) - - if detail: - workflow_trigger_log = aliased(WorkflowTriggerLog) - stmt = stmt.outerjoin( - workflow_trigger_log, - and_( - workflow_trigger_log.tenant_id == tenant_id, - workflow_trigger_log.app_id == app_id, - workflow_trigger_log.workflow_run_id == WorkflowAppLog.workflow_run_id, - ), - ).add_columns(workflow_trigger_log.trigger_metadata.label("trigger_metadata")) - else: - stmt = stmt.add_columns(literal(None).label("trigger_metadata")) - - if keyword: - escaped_keyword = escape_like_pattern(keyword[:30]) - keyword_like_value = f"%{escaped_keyword}%" - run_creator = aliased(EndUser) - keyword_conditions = ( - workflow_run.inputs.ilike(keyword_like_value, escape="\\"), - workflow_run.outputs.ilike(keyword_like_value, escape="\\"), - select(1) - .select_from(run_creator) - .where( - workflow_run.created_by_role == CreatorUserRole.END_USER, - run_creator.id == workflow_run.created_by, - run_creator.session_id.ilike(keyword_like_value, escape="\\"), - ) - .exists(), - ) - keyword_uuid = self._safe_parse_uuid(keyword) - if keyword_uuid is not None: - stmt = stmt.where(or_(*keyword_conditions, workflow_run.id == keyword_uuid)) - else: - stmt = stmt.where(or_(*keyword_conditions)) - - if status is not None: - stmt = stmt.where(workflow_run.status == status) - if created_at_before is not None: - stmt = stmt.where(WorkflowAppLog.created_at <= created_at_before) - if created_at_after is not None: - stmt = stmt.where(WorkflowAppLog.created_at >= created_at_after) - - if created_by_end_user_session_id: - log_creator = aliased(EndUser) - stmt = stmt.where( - WorkflowAppLog.created_by_role == CreatorUserRole.END_USER, - select(1) - .select_from(log_creator) - .where( - log_creator.id == WorkflowAppLog.created_by, - log_creator.session_id == created_by_end_user_session_id, - ) - .exists(), - ) - if created_by_account: - account_id = session.scalar( - select(Account.id) - .join(TenantAccountJoin, TenantAccountJoin.account_id == Account.id) - .where( - Account.email == created_by_account, - TenantAccountJoin.tenant_id == tenant_id, - ) - ) - if account_id is None: - raise ValueError(f"Account not found: {created_by_account}") - - stmt = stmt.where( - WorkflowAppLog.created_by_role == CreatorUserRole.ACCOUNT, - WorkflowAppLog.created_by == account_id, - ) - - total = session.scalar(select(func.count()).select_from(stmt.subquery())) or 0 - - actor_account = aliased(Account) - actor_end_user = aliased(EndUser) - paginated_stmt = ( - stmt.outerjoin( - actor_account, - and_( - WorkflowAppLog.created_by_role == CreatorUserRole.ACCOUNT, - actor_account.id == WorkflowAppLog.created_by, - ), - ) - .outerjoin( - actor_end_user, - and_( - WorkflowAppLog.created_by_role == CreatorUserRole.END_USER, - actor_end_user.id == WorkflowAppLog.created_by, - ), - ) - .add_columns( - actor_account.id.label("account_id"), - actor_account.name.label("account_name"), - actor_account.email.label("account_email"), - actor_end_user.id.label("end_user_id"), - actor_end_user.type.label("end_user_type"), - actor_end_user.session_id.label("end_user_session_id"), - ) - .order_by(WorkflowAppLog.created_at.desc()) - .offset((page - 1) * limit) - .limit(limit) - ) - rows = session.execute(paginated_stmt).mappings().all() - - return WorkflowAppLogPage( - page=page, - limit=limit, - total=total, - has_more=total > page * limit, - data=tuple(self._to_item(row=row, detail=detail) for row in rows), - ) - - @staticmethod - def _to_item(*, row: RowMapping, detail: bool) -> WorkflowAppLogItem: - run_id = row["run_id"] - workflow_run = ( - WorkflowAppLogRunSummary( - id=run_id, - version=row["run_version"], - status=WorkflowAppLogQueryRepository._enum_value(row["run_status"]), - triggered_from=WorkflowAppLogQueryRepository._enum_value(row["run_triggered_from"]), - error=row["run_error"], - elapsed_time=row["run_elapsed_time"], - total_tokens=row["run_total_tokens"], - total_steps=row["run_total_steps"], - created_at=row["run_created_at"], - finished_at=row["run_finished_at"], - exceptions_count=row["run_exceptions_count"], - ) - if run_id is not None - else None - ) - - account_id = row["account_id"] - account = ( - WorkflowAppLogAccount( - id=account_id, - name=row["account_name"], - email=row["account_email"], - ) - if account_id is not None - else None - ) - end_user_id = row["end_user_id"] - end_user = ( - WorkflowAppLogEndUser( - id=end_user_id, - type=WorkflowAppLogQueryRepository._enum_value(row["end_user_type"]), - is_anonymous=False, - session_id=row["end_user_session_id"], - ) - if end_user_id is not None - else None - ) - - return WorkflowAppLogItem( - id=row["log_id"], - workflow_run=workflow_run, - details={"trigger_metadata": row["trigger_metadata"]} if detail else None, - created_from=WorkflowAppLogQueryRepository._enum_value(row["log_created_from"]), - created_by_role=WorkflowAppLogQueryRepository._enum_value(row["log_created_by_role"]), - created_by_account=account, - created_by_end_user=end_user, - created_at=row["log_created_at"], - ) - - @staticmethod - def _enum_value(value: Enum | str | None) -> str: - if value is None: - raise ValueError("Required enum value is missing") - if isinstance(value, str): - return value - return str(value.value) - - @staticmethod - def _safe_parse_uuid(value: str) -> uuid.UUID | None: - if len(value) < 32: - return None - try: - return uuid.UUID(value) - except ValueError: - return None diff --git a/api/repositories/workflow_collaboration_repository.py b/api/repositories/workflow_collaboration_repository.py deleted file mode 100644 index b3184528bfc20f..00000000000000 --- a/api/repositories/workflow_collaboration_repository.py +++ /dev/null @@ -1,272 +0,0 @@ -from __future__ import annotations - -import json -from typing import NotRequired, TypedDict, override - -from redis.lock import Lock - -from extensions.ext_redis import redis_client -from extensions.redis_names import serialize_redis_name - -SESSION_STATE_TTL_SECONDS = 3600 -SERVER_HEARTBEAT_TTL_SECONDS = 90 -WORKFLOW_ONLINE_USERS_PREFIX = "workflow_online_users:" -WORKFLOW_LEADER_PREFIX = "workflow_leader:" -WS_SID_MAP_PREFIX = "ws_sid_map:" -WS_SERVER_HEARTBEAT_PREFIX = "ws_server_heartbeat:" -WS_SERVER_SESSIONS_PREFIX = "ws_server_sessions:" -GRAPH_VIEW_STATE_LOCK_PREFIX = "workflow_graph_view_state_lock:" -GRAPH_VIEW_STATE_LOCK_TIMEOUT_SECONDS = 15 - -_UPDATE_SESSION_GRAPH_ACTIVE_LUA = """ -local raw = redis.call('HGET', KEYS[1], ARGV[1]) -if not raw then - return 0 -end - -local decoded_ok, session_info = pcall(cjson.decode, raw) -if not decoded_ok or type(session_info) ~= 'table' then - return 0 -end - -local incoming_sequence = tonumber(ARGV[3]) -local current_sequence = tonumber(session_info.graph_active_sequence) -if current_sequence and incoming_sequence <= current_sequence then - return 0 -end - -session_info.graph_active = ARGV[2] == '1' -session_info.graph_active_sequence = incoming_sequence -redis.call('HSET', KEYS[1], ARGV[1], cjson.encode(session_info)) -return 1 -""" - - -class WorkflowSessionInfo(TypedDict): - user_id: str - username: str - avatar: str | None - sid: str - connected_at: int - server_id: NotRequired[str] - graph_active: NotRequired[bool] - - -class SidMapping(TypedDict): - workflow_id: str - user_id: str - server_id: NotRequired[str] - - -class WorkflowCollaborationRepository: - def __init__(self) -> None: - self._redis = redis_client - - @override - def __repr__(self) -> str: - return f"{self.__class__.__name__}(redis_client={self._redis})" - - @staticmethod - def workflow_key(workflow_id: str) -> str: - return f"{WORKFLOW_ONLINE_USERS_PREFIX}{workflow_id}" - - @staticmethod - def leader_key(workflow_id: str) -> str: - return f"{WORKFLOW_LEADER_PREFIX}{workflow_id}" - - @staticmethod - def sid_key(sid: str) -> str: - return f"{WS_SID_MAP_PREFIX}{sid}" - - @staticmethod - def server_key(server_id: str) -> str: - return f"{WS_SERVER_HEARTBEAT_PREFIX}{server_id}" - - @staticmethod - def server_sessions_key(server_id: str) -> str: - return f"{WS_SERVER_SESSIONS_PREFIX}{server_id}" - - @staticmethod - def _decode(value: str | bytes | None) -> str | None: - if value is None: - return None - if isinstance(value, bytes): - return value.decode("utf-8") - return value - - def refresh_session_state(self, workflow_id: str, sid: str) -> None: - workflow_key = self.workflow_key(workflow_id) - sid_key = self.sid_key(sid) - if self._redis.exists(workflow_key): - self._redis.expire(workflow_key, SESSION_STATE_TTL_SECONDS) - if self._redis.exists(sid_key): - self._redis.expire(sid_key, SESSION_STATE_TTL_SECONDS) - - def set_session_info(self, workflow_id: str, session_info: WorkflowSessionInfo) -> None: - workflow_key = self.workflow_key(workflow_id) - sid_mapping: SidMapping = {"workflow_id": workflow_id, "user_id": session_info["user_id"]} - if server_id := session_info.get("server_id"): - sid_mapping["server_id"] = server_id - - self._redis.hset(workflow_key, session_info["sid"], json.dumps(session_info)) - if server_id: - self._redis.hset(self.server_sessions_key(server_id), session_info["sid"], workflow_id) - self._redis.expire(self.server_sessions_key(server_id), SESSION_STATE_TTL_SECONDS) - self._redis.set( - self.sid_key(session_info["sid"]), - json.dumps(sid_mapping), - ex=SESSION_STATE_TTL_SECONDS, - ) - self.refresh_session_state(workflow_id, session_info["sid"]) - - def get_sid_mapping(self, sid: str) -> SidMapping | None: - raw = self._redis.get(self.sid_key(sid)) - if not raw: - return None - value = self._decode(raw) - if not value: - return None - try: - return json.loads(value) - except (TypeError, json.JSONDecodeError): - return None - - def delete_session(self, workflow_id: str, sid: str) -> None: - mapping = self.get_sid_mapping(sid) - if mapping and (server_id := mapping.get("server_id")): - self._redis.hdel(self.server_sessions_key(server_id), sid) - self._redis.hdel(self.workflow_key(workflow_id), sid) - self._redis.delete(self.sid_key(sid)) - - def session_exists(self, workflow_id: str, sid: str) -> bool: - return self._redis.hexists(self.workflow_key(workflow_id), sid) - - def sid_mapping_exists(self, sid: str) -> bool: - return self._redis.exists(self.sid_key(sid)) - - def get_session_sids(self, workflow_id: str) -> list[str]: - raw_sids = self._redis.hkeys(self.workflow_key(workflow_id)) - decoded_sids: list[str] = [] - for sid in raw_sids: - decoded = self._decode(sid) - if decoded: - decoded_sids.append(decoded) - return decoded_sids - - def list_sessions(self, workflow_id: str) -> list[WorkflowSessionInfo]: - sessions_json = self._redis.hgetall(self.workflow_key(workflow_id)) - users: list[WorkflowSessionInfo] = [] - - for session_info_json in sessions_json.values(): - value = self._decode(session_info_json) - if not value: - continue - try: - session_info = json.loads(value) - except (TypeError, json.JSONDecodeError): - continue - - if not isinstance(session_info, dict): - continue - if "user_id" not in session_info or "username" not in session_info or "sid" not in session_info: - continue - - user: WorkflowSessionInfo = { - "user_id": str(session_info["user_id"]), - "username": str(session_info["username"]), - "avatar": session_info.get("avatar"), - "sid": str(session_info["sid"]), - "connected_at": int(session_info.get("connected_at") or 0), - } - if isinstance(session_info.get("graph_active"), bool): - user["graph_active"] = session_info["graph_active"] - users.append(user) - - return users - - def get_session_info(self, workflow_id: str, sid: str) -> WorkflowSessionInfo | None: - raw = self._redis.hget(self.workflow_key(workflow_id), sid) - value = self._decode(raw) - if not value: - return None - try: - session_info = json.loads(value) - except (TypeError, json.JSONDecodeError): - return None - - if not isinstance(session_info, dict): - return None - if "user_id" not in session_info or "username" not in session_info or "sid" not in session_info: - return None - - user: WorkflowSessionInfo = { - "user_id": str(session_info["user_id"]), - "username": str(session_info["username"]), - "avatar": session_info.get("avatar"), - "sid": str(session_info["sid"]), - "connected_at": int(session_info.get("connected_at") or 0), - } - if isinstance(session_info.get("server_id"), str): - user["server_id"] = session_info["server_id"] - if isinstance(session_info.get("graph_active"), bool): - user["graph_active"] = session_info["graph_active"] - return user - - def update_session_graph_active(self, workflow_id: str, sid: str, active: bool, sequence: int) -> bool: - """Atomically apply a graph visibility update when its client sequence is newer.""" - # RedisClientWrapper prefixes regular hash calls, but eval is delegated to the raw client. - workflow_key = serialize_redis_name(self.workflow_key(workflow_id)) - result = self._redis.eval( - _UPDATE_SESSION_GRAPH_ACTIVE_LUA, - 1, - workflow_key, - sid, - "1" if active else "0", - sequence, - ) - return bool(result) - - def graph_view_state_lock(self, workflow_id: str) -> Lock: - """Serialize visibility state changes and their leader-election side effects.""" - return self._redis.lock( - f"{GRAPH_VIEW_STATE_LOCK_PREFIX}{workflow_id}", - timeout=GRAPH_VIEW_STATE_LOCK_TIMEOUT_SECONDS, - ) - - def refresh_server_heartbeat(self, server_id: str) -> None: - self._redis.set(self.server_key(server_id), "1", ex=SERVER_HEARTBEAT_TTL_SECONDS) - - def server_heartbeat_exists(self, server_id: str) -> bool: - return bool(self._redis.exists(self.server_key(server_id))) # tests assert `is True` - - def refresh_server_sessions(self, server_id: str) -> None: - """Refresh Redis TTLs for sessions owned by a live websocket worker.""" - server_sessions_key = self.server_sessions_key(server_id) - sessions = self._redis.hgetall(server_sessions_key) - for raw_sid, raw_workflow_id in sessions.items(): - sid = self._decode(raw_sid) - workflow_id = self._decode(raw_workflow_id) - if not sid or not workflow_id: - continue - if not self.sid_mapping_exists(sid): - self._redis.hdel(server_sessions_key, sid) - continue - self.refresh_session_state(workflow_id, sid) - if sessions: - self._redis.expire(server_sessions_key, SESSION_STATE_TTL_SECONDS) - - def get_current_leader(self, workflow_id: str) -> str | None: - raw = self._redis.get(self.leader_key(workflow_id)) - return self._decode(raw) - - def set_leader_if_absent(self, workflow_id: str, sid: str) -> bool: - return bool(self._redis.set(self.leader_key(workflow_id), sid, nx=True, ex=SESSION_STATE_TTL_SECONDS)) - - def set_leader(self, workflow_id: str, sid: str) -> None: - self._redis.set(self.leader_key(workflow_id), sid, ex=SESSION_STATE_TTL_SECONDS) - - def delete_leader(self, workflow_id: str) -> None: - self._redis.delete(self.leader_key(workflow_id)) - - def expire_leader(self, workflow_id: str) -> None: - self._redis.expire(self.leader_key(workflow_id), SESSION_STATE_TTL_SECONDS) diff --git a/api/services/agent/agent_soul_state.py b/api/services/agent/agent_soul_state.py deleted file mode 100644 index dfc0a0335e47db..00000000000000 --- a/api/services/agent/agent_soul_state.py +++ /dev/null @@ -1,6 +0,0 @@ -from models.agent_config_entities import AgentSoulConfig - - -def agent_soul_has_model(agent_soul: AgentSoulConfig) -> bool: - """Return whether the Agent Soul has the minimum model config required for runtime.""" - return agent_soul.model is not None diff --git a/api/services/agent/legacy_retirement_service.py b/api/services/agent/legacy_retirement_service.py deleted file mode 100644 index c800a8ea51bd58..00000000000000 --- a/api/services/agent/legacy_retirement_service.py +++ /dev/null @@ -1,224 +0,0 @@ -"""Workflow-only Agent ownership retirement after product transactions commit.""" - -from __future__ import annotations - -import logging -from collections.abc import Iterable - -from sqlalchemy import delete, select -from sqlalchemy.orm import Session - -from core.db.session_factory import session_factory -from libs.datetime_utils import naive_utc_now -from models.agent import ( - Agent, - AgentScope, - AgentStatus, - AgentWorkingResourceStatus, - AgentWorkspace, - AgentWorkspaceBinding, - WorkflowAgentNodeBinding, -) -from models.model import App, AppMode -from models.workflow import Workflow -from services.agent.home_snapshot_service import AgentHomeSnapshotService -from services.agent.workspace_service import AgentWorkspaceService -from tasks.collect_agent_resources_task import enqueue_agent_resource_collection -from tasks.remove_app_and_related_data_task import remove_app_and_related_data_task - -logger = logging.getLogger(__name__) - - -class WorkflowAgentRetirementService: - """Delete workflow-only Agent aggregates after their last Workflow owner is gone.""" - - @classmethod - def retire_unowned( - cls, - *, - tenant_id: str, - agent_ids: Iterable[str], - account_id: str | None, - ) -> None: - """Retire unowned workflow-only Agents in an independent transaction. - - This method returns ``None``. It archives orphan Agents, retires their - working resources, and deletes their hidden Apps before committing. It - then publishes every hidden-App cleanup before publishing the Agent - resource collector; database and task-publication errors propagate. - - Archived Agents, missing hidden App rows, and already-retired resources - remain cleanup candidates. A retry can therefore publish duplicate - cleanup tasks, which are expected to be idempotent. - """ - - candidates = tuple(sorted({agent_id for agent_id in agent_ids if agent_id})) - if not candidates: - return - backing_app_ids: list[str] = [] - retired_bindings: list[str] = [] - retired_workspaces: list[str] = [] - retired_snapshots: list[str] = [] - try: - with session_factory.create_session() as session: - retired_agent_ids = cls.archive_unowned( - session=session, - tenant_id=tenant_id, - agent_ids=candidates, - account_id=account_id, - ) - retired_agents = session.scalars( - select(Agent).where( - Agent.tenant_id == tenant_id, - Agent.id.in_(retired_agent_ids), - ) - ).all() - backing_app_ids = sorted({agent.backing_app_id for agent in retired_agents if agent.backing_app_id}) - for app_id in backing_app_ids: - AgentWorkspaceService.retire_all_for_app( - session=session, - tenant_id=tenant_id, - app_id=app_id, - ) - retired_workspaces.extend( - session.scalars( - select(AgentWorkspace.id).where( - AgentWorkspace.tenant_id == tenant_id, - AgentWorkspace.app_id == app_id, - AgentWorkspace.status == AgentWorkingResourceStatus.RETIRED, - ) - ).all() - ) - for agent_id in retired_agent_ids: - bindings = session.scalars( - select(AgentWorkspaceBinding).where( - AgentWorkspaceBinding.tenant_id == tenant_id, - AgentWorkspaceBinding.agent_id == agent_id, - ) - ).all() - for binding in bindings: - if binding.status == AgentWorkingResourceStatus.ACTIVE: - AgentWorkspaceService.retire_binding( - session=session, - tenant_id=tenant_id, - binding_id=binding.id, - ) - retired_bindings.append(binding.id) - retired_snapshots.extend( - AgentHomeSnapshotService.retire_all_for_agent( - session=session, - tenant_id=tenant_id, - agent_id=agent_id, - ) - ) - if backing_app_ids: - session.execute( - delete(App).where( - App.tenant_id == tenant_id, - App.id.in_(backing_app_ids), - App.mode == AppMode.AGENT, - ) - ) - session.commit() - except Exception: - logger.exception( - "Failed to retire unowned Workflow Agents", - extra={ - "tenant_id": tenant_id, - "agent_ids": candidates, - }, - ) - raise - - for app_id in backing_app_ids: - try: - remove_app_and_related_data_task.delay(tenant_id=tenant_id, app_id=app_id) - except Exception: - logger.exception( - "Failed to enqueue hidden Agent App cleanup", - extra={"tenant_id": tenant_id, "app_id": app_id}, - ) - raise - enqueue_agent_resource_collection( - tenant_id=tenant_id, - workspace_ids=retired_workspaces, - binding_ids=retired_bindings, - home_snapshot_ids=retired_snapshots, - purge_agent_ids=retired_agent_ids, - ) - - @classmethod - def archive_unowned( - cls, - *, - session: Session, - tenant_id: str, - agent_ids: Iterable[str], - account_id: str | None, - ) -> list[str]: - """Archive active orphans and return complete aggregate purge candidates.""" - candidates = tuple(sorted({agent_id for agent_id in agent_ids if agent_id})) - if not candidates: - return [] - agents = session.scalars( - select(Agent).where( - Agent.tenant_id == tenant_id, - Agent.id.in_(candidates), - Agent.scope == AgentScope.WORKFLOW_ONLY, - Agent.status.in_((AgentStatus.ACTIVE, AgentStatus.ARCHIVED)), - ) - ).all() - retained_agent_ids = cls.retained_agent_ids( - session=session, - tenant_id=tenant_id, - agent_ids=[agent.id for agent in agents], - ) - now = naive_utc_now() - cleanup_candidates: list[str] = [] - for agent in agents: - if agent.id in retained_agent_ids: - continue - if agent.status == AgentStatus.ACTIVE: - agent.status = AgentStatus.ARCHIVED - agent.archived_by = account_id - agent.archived_at = now - agent.updated_by = account_id or agent.updated_by - agent.updated_at = now - cleanup_candidates.append(agent.id) - session.flush() - return cleanup_candidates - - @staticmethod - def retained_agent_ids( - *, - session: Session, - tenant_id: str, - agent_ids: list[str], - ) -> set[str]: - """Return Agents that still have an exact persisted Workflow owner. - - The owner key is tenant, App, Workflow, and Workflow version. Draft and - every published version, whether current or historical, count equally; - the App's current-Workflow pointer is not part of ownership. - """ - if not agent_ids: - return set() - values = session.scalars( - select(WorkflowAgentNodeBinding.agent_id) - .join( - Workflow, - (Workflow.tenant_id == WorkflowAgentNodeBinding.tenant_id) - & (Workflow.app_id == WorkflowAgentNodeBinding.app_id) - & (Workflow.id == WorkflowAgentNodeBinding.workflow_id) - & (Workflow.version == WorkflowAgentNodeBinding.workflow_version), - ) - .where( - WorkflowAgentNodeBinding.tenant_id == tenant_id, - WorkflowAgentNodeBinding.agent_id.in_(agent_ids), - ) - .distinct() - ).all() - return {agent_id for agent_id in values if agent_id} - - -__all__ = ["WorkflowAgentRetirementService"] diff --git a/api/services/agent/legacy_workflow_publish_service.py b/api/services/agent/legacy_workflow_publish_service.py deleted file mode 100644 index 358b49e80a18dc..00000000000000 --- a/api/services/agent/legacy_workflow_publish_service.py +++ /dev/null @@ -1,744 +0,0 @@ -from __future__ import annotations - -import copy -from collections.abc import Mapping -from typing import Any, cast - -from pydantic import ValidationError -from sqlalchemy import select -from sqlalchemy.orm import Session - -from core.agent.publish_visibility import workflow_callable_active_snapshot_filter -from core.workflow.nodes.agent_v2.validators import WorkflowAgentNodeValidationError, WorkflowAgentNodeValidator -from models.agent import ( - Agent, - AgentConfigSnapshot, - AgentScope, - AgentStatus, - WorkflowAgentBindingType, - WorkflowAgentNodeBinding, -) -from models.agent_config_entities import ( - AgentSoulConfig, - WorkflowNodeJobConfig, - WorkflowOutputRoutes, - WorkflowPreviousNodeOutputRef, -) -from models.workflow import Workflow -from services.agent.composer_validator import ComposerConfigValidator -from services.agent.prompt_mentions import ( - extract_workflow_node_output_selectors, - workflow_previous_node_output_refs_from_selectors, -) -from services.entities.agent_entities import ( - ComposerSavePayload, - ComposerSaveStrategy, - ComposerSoulLockPayload, - ComposerVariant, -) - - -class _InlineAgentUnavailableError(ValueError): - pass - - -class _InlineAgentOwnershipError(ValueError): - pass - - -class _InlineAgentSnapshotError(ValueError): - pass - - -class WorkflowAgentPublishService: - """Validate and freeze Workflow Agent v2 bindings during workflow publish.""" - - _DRAFT_WORKFLOW_VERSION = Workflow.VERSION_DRAFT - _AGENT_BINDING_KEY = "agent_binding" - _AGENT_TASK_KEY = "agent_task" - _AGENT_DECLARED_OUTPUTS_KEY = "agent_declared_outputs" - _AGENT_OUTPUT_ROUTES_KEY = "agent_output_routes" - - @classmethod - def project_draft_bindings_to_graph(cls, *, session: Session, draft_workflow: Workflow) -> dict[str, Any]: - """Return draft graph with persisted Agent binding fields projected into node data. - - Workflow draft graph is the front-end's editing source of truth, while - runtime/publish reads WorkflowAgentNodeBinding. This - response-only projection keeps reads aligned without writing binding - details back into the stored graph JSON. - """ - graph = cast(dict[str, Any], copy.deepcopy(draft_workflow.graph_dict)) - agent_nodes = dict(WorkflowAgentNodeValidator.iter_agent_v2_nodes(graph)) - if not agent_nodes: - return graph - - bindings = session.scalars( - select(WorkflowAgentNodeBinding).where( - WorkflowAgentNodeBinding.tenant_id == draft_workflow.tenant_id, - WorkflowAgentNodeBinding.app_id == draft_workflow.app_id, - WorkflowAgentNodeBinding.workflow_id == draft_workflow.id, - WorkflowAgentNodeBinding.workflow_version == cls._DRAFT_WORKFLOW_VERSION, - WorkflowAgentNodeBinding.node_id.in_(list(agent_nodes.keys())), - ) - ).all() - for binding in bindings: - node_data = agent_nodes.get(binding.node_id) - if not isinstance(node_data, dict): - continue - graph_binding = node_data.get(cls._AGENT_BINDING_KEY) - is_pending_inline_graph_binding = ( - isinstance(graph_binding, Mapping) - and graph_binding.get("binding_type") == WorkflowAgentBindingType.INLINE_AGENT.value - and (not graph_binding.get("agent_id") or not graph_binding.get("current_snapshot_id")) - ) - if not is_pending_inline_graph_binding or binding.binding_type == WorkflowAgentBindingType.INLINE_AGENT: - node_data[cls._AGENT_BINDING_KEY] = { - "binding_type": binding.binding_type.value, - "agent_id": binding.agent_id, - "current_snapshot_id": binding.current_snapshot_id, - } - node_job = WorkflowNodeJobConfig.model_validate(binding.node_job_config_dict) - if node_job.workflow_prompt is not None: - node_data[cls._AGENT_TASK_KEY] = node_job.workflow_prompt - node_data[cls._AGENT_OUTPUT_ROUTES_KEY] = node_job.output_routes.model_dump(mode="json") - node_data[cls._AGENT_DECLARED_OUTPUTS_KEY] = [ - output.model_dump(mode="json") for output in node_job.declared_outputs - ] - return graph - - @classmethod - def validate_agent_nodes_for_publish(cls, *, session: Session, draft_workflow: Workflow) -> None: - WorkflowAgentNodeValidator.validate_published_workflow(session=session, workflow=draft_workflow) - cls._validate_composer_configs_for_publish(session=session, draft_workflow=draft_workflow) - - @classmethod - def validate_agent_nodes_for_draft_sync(cls, *, session: Session, draft_workflow: Workflow) -> None: - WorkflowAgentNodeValidator.validate_draft_workflow(session=session, workflow=draft_workflow) - - @classmethod - def _validate_composer_configs_for_publish(cls, *, session: Session, draft_workflow: Workflow) -> None: - node_ids = { - node_id for node_id, _node_data in WorkflowAgentNodeValidator.iter_agent_v2_nodes(draft_workflow.graph_dict) - } - if not node_ids: - return - - bindings = session.scalars( - select(WorkflowAgentNodeBinding).where( - WorkflowAgentNodeBinding.tenant_id == draft_workflow.tenant_id, - WorkflowAgentNodeBinding.app_id == draft_workflow.app_id, - WorkflowAgentNodeBinding.workflow_id == draft_workflow.id, - WorkflowAgentNodeBinding.workflow_version == draft_workflow.version, - WorkflowAgentNodeBinding.node_id.in_(node_ids), - ) - ).all() - for binding in bindings: - cls._validate_binding_composer_config_for_publish(session=session, binding=binding) - - @classmethod - def _validate_binding_composer_config_for_publish( - cls, - *, - session: Session, - binding: WorkflowAgentNodeBinding, - ) -> None: - if not binding.agent_id: - return - - agent = session.scalar( - select(Agent) - .where( - Agent.tenant_id == binding.tenant_id, - Agent.id == binding.agent_id, - ) - .limit(1) - ) - if agent is None: - return - - snapshot_id = ( - agent.active_config_snapshot_id - if binding.binding_type == WorkflowAgentBindingType.ROSTER_AGENT - else binding.current_snapshot_id - ) - if snapshot_id is None: - return - - snapshot = session.scalar( - select(AgentConfigSnapshot) - .where( - AgentConfigSnapshot.tenant_id == binding.tenant_id, - AgentConfigSnapshot.agent_id == agent.id, - AgentConfigSnapshot.id == snapshot_id, - ) - .limit(1) - ) - if snapshot is None: - return - - agent_soul = AgentSoulConfig.model_validate(snapshot.config_snapshot_dict) - node_job = WorkflowNodeJobConfig.model_validate(binding.node_job_config_dict) - payload = ComposerSavePayload.model_construct( - variant=ComposerVariant.WORKFLOW, - save_strategy=ComposerSaveStrategy.NODE_JOB_ONLY, - soul_lock=ComposerSoulLockPayload(locked=False), - agent_soul=agent_soul, - node_job=node_job, - ) - ComposerConfigValidator.validate_publish_payload(payload) - cls._require_config_asset_refs_resolved_for_publish( - session=session, - binding=binding, - snapshot_id=snapshot_id, - agent_soul=agent_soul, - ) - - @classmethod - def _require_config_asset_refs_resolved_for_publish( - cls, - *, - session: Session, - binding: WorkflowAgentNodeBinding, - snapshot_id: str, - agent_soul: AgentSoulConfig, - ) -> None: - from services.agent.prompt_mentions import MentionKind, parse_prompt_mentions - from services.skill_management_service import SkillManagementService - - mentions = parse_prompt_mentions(agent_soul.prompt.system_prompt) - configured_skill_names = {item.name for item in agent_soul.config_skills if not item.is_missing} - has_unresolved_skill_ref = any( - mention.kind == MentionKind.SKILL and mention.ref_id not in configured_skill_names for mention in mentions - ) - if has_unresolved_skill_ref and binding.agent_id is not None: - configured_skill_names.update( - str(item["name"]) - for item in SkillManagementService(session=session).list_runtime_agent_skills( - tenant_id=binding.tenant_id, - agent_id=binding.agent_id, - config_snapshot_id=snapshot_id, - ) - ) - configured_file_names = {item.name for item in agent_soul.config_files if not item.is_missing} - missing_refs: list[str] = [] - for mention in mentions: - if mention.kind not in {MentionKind.SKILL, MentionKind.FILE}: - continue - ref_name = mention.ref_id - if mention.kind == MentionKind.SKILL and ref_name not in configured_skill_names: - missing_refs.append(f"skill_ref_dangling: skill '{mention.label or ref_name}' is not configured.") - if mention.kind == MentionKind.FILE and ref_name not in configured_file_names: - missing_refs.append(f"file_ref_dangling: file '{mention.label or ref_name}' is not configured.") - if missing_refs: - raise WorkflowAgentNodeValidationError( - f"Workflow Agent node {binding.node_id} has invalid Agent Soul config refs: {'; '.join(missing_refs)}" - ) - - @classmethod - def sync_agent_bindings_for_draft( - cls, - *, - session: Session, - draft_workflow: Workflow, - account_id: str, - ) -> set[str]: - agent_nodes = dict(WorkflowAgentNodeValidator.iter_agent_v2_nodes(draft_workflow.graph_dict)) - existing_bindings = list( - session.scalars( - select(WorkflowAgentNodeBinding).where( - WorkflowAgentNodeBinding.tenant_id == draft_workflow.tenant_id, - WorkflowAgentNodeBinding.app_id == draft_workflow.app_id, - WorkflowAgentNodeBinding.workflow_id == draft_workflow.id, - WorkflowAgentNodeBinding.workflow_version == cls._DRAFT_WORKFLOW_VERSION, - ) - ).all() - ) - existing_by_node_id = {binding.node_id: binding for binding in existing_bindings} - retirement_candidates: set[str] = set() - - for binding in existing_bindings: - if binding.node_id not in agent_nodes: - if binding.binding_type == WorkflowAgentBindingType.INLINE_AGENT and binding.agent_id: - retirement_candidates.add(binding.agent_id) - session.delete(binding) - - for node_id, node_data in agent_nodes.items(): - binding_payload = node_data.get(cls._AGENT_BINDING_KEY) - if binding_payload is None: - continue - if not isinstance(binding_payload, Mapping): - raise ValueError(f"Workflow Agent node {node_id} has invalid agent_binding.") - if binding_payload.get("binding_type") == WorkflowAgentBindingType.INLINE_AGENT.value and ( - not binding_payload.get("agent_id") or not binding_payload.get("current_snapshot_id") - ): - continue - existing_binding = existing_by_node_id.get(node_id) - replaced_inline_agent_id = ( - existing_binding.agent_id - if existing_binding is not None - and existing_binding.binding_type == WorkflowAgentBindingType.INLINE_AGENT - and existing_binding.agent_id - else None - ) - cls._sync_agent_binding_for_node( - session=session, - draft_workflow=draft_workflow, - node_id=node_id, - node_data=node_data, - node_binding=binding_payload, - existing_binding=existing_binding, - account_id=account_id, - ) - if ( - replaced_inline_agent_id - and existing_binding is not None - and ( - existing_binding.binding_type != WorkflowAgentBindingType.INLINE_AGENT - or existing_binding.agent_id != replaced_inline_agent_id - ) - ): - retirement_candidates.add(replaced_inline_agent_id) - session.flush() - return retirement_candidates - - @classmethod - def sync_roster_agent_bindings_for_draft( - cls, - *, - session: Session, - draft_workflow: Workflow, - account_id: str, - ) -> set[str]: - return cls.sync_agent_bindings_for_draft( - session=session, - draft_workflow=draft_workflow, - account_id=account_id, - ) - - @classmethod - def _sync_agent_binding_for_node( - cls, - *, - session: Session, - draft_workflow: Workflow, - node_id: str, - node_data: Mapping[str, Any], - node_binding: Mapping[str, Any], - existing_binding: WorkflowAgentNodeBinding | None, - account_id: str, - ) -> None: - binding_type = node_binding.get("binding_type") - agent_id = node_binding.get("agent_id") - if not isinstance(agent_id, str) or not agent_id: - raise ValueError(f"Workflow Agent node {node_id} agent binding requires agent_id.") - - node_job_config = cls._node_job_config_from_node_data( - existing_binding=existing_binding, - node_data=node_data, - ) - - if binding_type == WorkflowAgentBindingType.ROSTER_AGENT.value: - agent, current_snapshot_id = cls._resolve_roster_agent_graph_binding( - session=session, - draft_workflow=draft_workflow, - node_id=node_id, - agent_id=agent_id, - ) - resolved_binding_type = WorkflowAgentBindingType.ROSTER_AGENT - elif binding_type == WorkflowAgentBindingType.INLINE_AGENT.value: - raw_current_snapshot_id = node_binding.get("current_snapshot_id") - if not isinstance(raw_current_snapshot_id, str) or not raw_current_snapshot_id: - raise ValueError(f"Workflow Agent node {node_id} inline_agent binding requires current_snapshot_id.") - current_snapshot_id = raw_current_snapshot_id - try: - agent = cls._resolve_inline_agent_graph_binding( - session=session, - draft_workflow=draft_workflow, - node_id=node_id, - agent_id=agent_id, - current_snapshot_id=current_snapshot_id, - ) - except (_InlineAgentUnavailableError, _InlineAgentOwnershipError): - existing_agent = cls._resolve_existing_inline_binding_agent( - session=session, - draft_workflow=draft_workflow, - node_id=node_id, - existing_binding=existing_binding, - ) - if existing_agent is not None and existing_binding is not None: - agent = existing_agent - current_snapshot_id = existing_binding.current_snapshot_id or current_snapshot_id - else: - agent, current_snapshot_id = cls._clone_inline_graph_binding_for_node( - session=session, - draft_workflow=draft_workflow, - node_id=node_id, - source_agent_id=agent_id, - source_snapshot_id=current_snapshot_id, - account_id=account_id, - ) - resolved_binding_type = WorkflowAgentBindingType.INLINE_AGENT - else: - raise ValueError(f"Workflow Agent node {node_id} has unsupported agent_binding type.") - - binding = existing_binding - if binding is None: - binding = WorkflowAgentNodeBinding( - tenant_id=draft_workflow.tenant_id, - app_id=draft_workflow.app_id, - workflow_id=draft_workflow.id, - workflow_version=cls._DRAFT_WORKFLOW_VERSION, - node_id=node_id, - node_job_config=node_job_config, - created_by=account_id, - ) - session.add(binding) - else: - binding.node_job_config = node_job_config - - binding.binding_type = resolved_binding_type - binding.agent_id = agent.id - binding.current_snapshot_id = current_snapshot_id - binding.updated_by = account_id - - @classmethod - def _resolve_existing_inline_binding_agent( - cls, - *, - session: Session, - draft_workflow: Workflow, - node_id: str, - existing_binding: WorkflowAgentNodeBinding | None, - ) -> Agent | None: - if ( - existing_binding is None - or existing_binding.binding_type != WorkflowAgentBindingType.INLINE_AGENT - or not existing_binding.agent_id - or not existing_binding.current_snapshot_id - ): - return None - try: - return cls._resolve_inline_agent_graph_binding( - session=session, - draft_workflow=draft_workflow, - node_id=node_id, - agent_id=existing_binding.agent_id, - current_snapshot_id=existing_binding.current_snapshot_id, - ) - except ValueError: - return None - - @classmethod - def _clone_inline_graph_binding_for_node( - cls, - *, - session: Session, - draft_workflow: Workflow, - node_id: str, - source_agent_id: str, - source_snapshot_id: str, - account_id: str, - ) -> tuple[Agent, str]: - source_agent = session.scalar( - select(Agent) - .where( - Agent.tenant_id == draft_workflow.tenant_id, - Agent.id == source_agent_id, - Agent.scope == AgentScope.WORKFLOW_ONLY, - Agent.status == AgentStatus.ACTIVE, - ) - .limit(1) - ) - if source_agent is None: - raise ValueError(f"Workflow Agent node {node_id} references an unavailable inline agent.") - source_snapshot = session.scalar( - select(AgentConfigSnapshot) - .where( - AgentConfigSnapshot.tenant_id == draft_workflow.tenant_id, - AgentConfigSnapshot.agent_id == source_agent.id, - AgentConfigSnapshot.id == source_snapshot_id, - ) - .limit(1) - ) - if source_snapshot is None: - raise ValueError(f"Workflow Agent node {node_id} references a missing inline agent config snapshot.") - - from services.agent.dsl_service import AgentDslService - - agent, snapshot = AgentDslService(session).clone_inline_binding_for_node( - workflow=draft_workflow, - node_id=node_id, - source_agent=source_agent, - source_snapshot=source_snapshot, - account_id=account_id, - ) - return agent, snapshot.id - - @classmethod - def _resolve_roster_agent_graph_binding( - cls, - *, - session: Session, - draft_workflow: Workflow, - node_id: str, - agent_id: str, - ) -> tuple[Agent, str]: - """Resolve an active roster Agent whose published snapshot is callable.""" - - agent = session.scalar( - select(Agent) - .where( - Agent.tenant_id == draft_workflow.tenant_id, - Agent.id == agent_id, - Agent.scope == AgentScope.ROSTER, - Agent.status == AgentStatus.ACTIVE, - workflow_callable_active_snapshot_filter(), - ) - .limit(1) - ) - if agent is None: - raise ValueError(f"Workflow Agent node {node_id} references an unavailable or unpublished roster agent.") - if agent.scope != AgentScope.ROSTER: - raise ValueError(f"Workflow Agent node {node_id} roster_agent binding must reference a roster agent.") - if not agent.active_config_snapshot_id: - raise ValueError(f"Workflow Agent node {node_id} roster agent has no active config snapshot.") - return agent, agent.active_config_snapshot_id - - @classmethod - def _resolve_inline_agent_graph_binding( - cls, - *, - session: Session, - draft_workflow: Workflow, - node_id: str, - agent_id: str, - current_snapshot_id: str, - ) -> Agent: - agent = session.scalar( - select(Agent) - .where( - Agent.tenant_id == draft_workflow.tenant_id, - Agent.id == agent_id, - Agent.scope == AgentScope.WORKFLOW_ONLY, - Agent.app_id == draft_workflow.app_id, - Agent.workflow_id == draft_workflow.id, - Agent.workflow_node_id == node_id, - Agent.status == AgentStatus.ACTIVE, - ) - .limit(1) - ) - if agent is None: - raise _InlineAgentUnavailableError(f"Workflow Agent node {node_id} references an unavailable inline agent.") - if ( - agent.scope != AgentScope.WORKFLOW_ONLY - or agent.app_id != draft_workflow.app_id - or agent.workflow_id != draft_workflow.id - or agent.workflow_node_id != node_id - ): - raise _InlineAgentOwnershipError( - f"Workflow Agent node {node_id} inline_agent binding does not belong to this node." - ) - - snapshot = session.scalar( - select(AgentConfigSnapshot) - .where( - AgentConfigSnapshot.tenant_id == draft_workflow.tenant_id, - AgentConfigSnapshot.agent_id == agent.id, - AgentConfigSnapshot.id == current_snapshot_id, - ) - .limit(1) - ) - if snapshot is None or snapshot.agent_id != agent.id: - raise _InlineAgentSnapshotError( - f"Workflow Agent node {node_id} references a missing inline agent config snapshot." - ) - return agent - - @classmethod - def _node_job_config_from_node_data( - cls, - *, - existing_binding: WorkflowAgentNodeBinding | None, - node_data: Mapping[str, Any], - ) -> WorkflowNodeJobConfig: - if existing_binding and existing_binding.node_job_config: - node_job = WorkflowNodeJobConfig.model_validate(existing_binding.node_job_config_dict) - else: - node_job = WorkflowNodeJobConfig() - - agent_task = node_data.get(cls._AGENT_TASK_KEY) - if isinstance(agent_task, str): - node_job.workflow_prompt = agent_task - node_job.previous_node_output_refs = cls._previous_node_output_refs_from_prompt(agent_task) - - declared_outputs_payload = node_data.get(cls._AGENT_DECLARED_OUTPUTS_KEY) - if declared_outputs_payload is not None: - if not isinstance(declared_outputs_payload, list): - raise ValueError("Workflow Agent node agent_declared_outputs must be a list.") - try: - node_job = WorkflowNodeJobConfig.model_validate( - { - **node_job.model_dump(mode="python"), - "declared_outputs": declared_outputs_payload, - } - ) - except ValidationError as exc: - raise ValueError("Workflow Agent node has invalid agent_declared_outputs.") from exc - - node_job.output_routes = WorkflowOutputRoutes.model_validate(node_data.get(cls._AGENT_OUTPUT_ROUTES_KEY, {})) - return node_job - - @classmethod - def _previous_node_output_refs_from_prompt(cls, prompt: str) -> list[WorkflowPreviousNodeOutputRef]: - """Derive persisted refs from the current frontend workflow markers only.""" - return workflow_previous_node_output_refs_from_selectors(extract_workflow_node_output_selectors(prompt)) - - @classmethod - def copy_agent_node_bindings_to_published( - cls, - *, - session: Session, - draft_workflow: Workflow, - published_workflow: Workflow, - ) -> bool: - """Copy all draft Roster and inline bindings to a published version. - - Only copied inline bindings add owners for workflow-only Agents. - Publishing does not release existing draft or historical inline owners, - produces no retirement candidates. The return value reports whether - the published Workflow contains an inline Agent binding. - """ - node_ids = { - node_id for node_id, _node_data in WorkflowAgentNodeValidator.iter_agent_v2_nodes(draft_workflow.graph_dict) - } - if not node_ids: - return False - - bindings = session.scalars( - select(WorkflowAgentNodeBinding).where( - WorkflowAgentNodeBinding.tenant_id == draft_workflow.tenant_id, - WorkflowAgentNodeBinding.app_id == draft_workflow.app_id, - WorkflowAgentNodeBinding.workflow_id == draft_workflow.id, - WorkflowAgentNodeBinding.workflow_version == draft_workflow.version, - WorkflowAgentNodeBinding.node_id.in_(node_ids), - ) - ).all() - if not bindings: - return False - - has_inline_agent = False - for binding in bindings: - has_inline_agent = has_inline_agent or ( - binding.binding_type == WorkflowAgentBindingType.INLINE_AGENT - and binding.agent_id is not None - and binding.current_snapshot_id is not None - ) - current_snapshot_id = binding.current_snapshot_id - if binding.binding_type == WorkflowAgentBindingType.ROSTER_AGENT and binding.agent_id: - _, current_snapshot_id = cls._resolve_roster_agent_graph_binding( - session=session, - draft_workflow=draft_workflow, - node_id=binding.node_id, - agent_id=binding.agent_id, - ) - copied = WorkflowAgentNodeBinding( - tenant_id=binding.tenant_id, - app_id=binding.app_id, - workflow_id=published_workflow.id, - workflow_version=published_workflow.version, - node_id=binding.node_id, - binding_type=binding.binding_type, - agent_id=binding.agent_id, - current_snapshot_id=current_snapshot_id, - node_job_config=WorkflowNodeJobConfig.model_validate(binding.node_job_config_dict), - created_by=binding.created_by, - updated_by=binding.updated_by, - ) - session.add(copied) - return has_inline_agent - - @classmethod - def restore_agent_node_bindings_to_draft( - cls, - *, - session: Session, - source_workflow: Workflow, - draft_workflow: Workflow, - account_id: str, - ) -> set[str]: - """Replace draft bindings with the frozen bindings of a published workflow.""" - - existing = session.scalars( - select(WorkflowAgentNodeBinding).where( - WorkflowAgentNodeBinding.tenant_id == draft_workflow.tenant_id, - WorkflowAgentNodeBinding.app_id == draft_workflow.app_id, - WorkflowAgentNodeBinding.workflow_id == draft_workflow.id, - WorkflowAgentNodeBinding.workflow_version == cls._DRAFT_WORKFLOW_VERSION, - ) - ).all() - retirement_candidates = { - binding.agent_id - for binding in existing - if binding.binding_type == WorkflowAgentBindingType.INLINE_AGENT and binding.agent_id - } - source_bindings = session.scalars( - select(WorkflowAgentNodeBinding).where( - WorkflowAgentNodeBinding.tenant_id == source_workflow.tenant_id, - WorkflowAgentNodeBinding.app_id == source_workflow.app_id, - WorkflowAgentNodeBinding.workflow_id == source_workflow.id, - WorkflowAgentNodeBinding.workflow_version == source_workflow.version, - ) - ).all() - for source in source_bindings: - if source.binding_type == WorkflowAgentBindingType.ROSTER_AGENT and source.agent_id: - cls._resolve_roster_agent_graph_binding( - session=session, - draft_workflow=draft_workflow, - node_id=source.node_id, - agent_id=source.agent_id, - ) - - for binding in existing: - session.delete(binding) - session.flush() - - for source in source_bindings: - agent_id = source.agent_id - snapshot_id = source.current_snapshot_id - if source.binding_type == WorkflowAgentBindingType.INLINE_AGENT and agent_id and snapshot_id: - try: - cls._resolve_inline_agent_graph_binding( - session=session, - draft_workflow=draft_workflow, - node_id=source.node_id, - agent_id=agent_id, - current_snapshot_id=snapshot_id, - ) - except ValueError: - agent, snapshot_id = cls._clone_inline_graph_binding_for_node( - session=session, - draft_workflow=draft_workflow, - node_id=source.node_id, - source_agent_id=agent_id, - source_snapshot_id=snapshot_id, - account_id=account_id, - ) - agent_id = agent.id - session.add( - WorkflowAgentNodeBinding( - tenant_id=draft_workflow.tenant_id, - app_id=draft_workflow.app_id, - workflow_id=draft_workflow.id, - workflow_version=cls._DRAFT_WORKFLOW_VERSION, - node_id=source.node_id, - binding_type=source.binding_type, - agent_id=agent_id, - current_snapshot_id=snapshot_id, - node_job_config=WorkflowNodeJobConfig.model_validate(source.node_job_config_dict), - created_by=account_id, - updated_by=account_id, - ) - ) - session.flush() - return retirement_candidates diff --git a/api/services/agent_tool_inner_service.py b/api/services/agent_tool_inner_service.py deleted file mode 100644 index 633ca8930079c1..00000000000000 --- a/api/services/agent_tool_inner_service.py +++ /dev/null @@ -1,151 +0,0 @@ -"""Service layer for Agent-backend core tool invocations. - -The `dify.core.tools` layer executes inside the API service boundary so Dify -Agent can expose API-owned tool providers (`plugin`, `builtin`, `api`, -`workflow`, and `mcp`) without learning their storage, credential, or runtime -internals. Production workflow and agent-app request builders still keep -existing plugin tool configs on the direct `dify.plugin.tools` route by -default; this service is the lower-level API execution path used when a caller -explicitly submits a `dify.core.tools` declaration. The service validates -tenant/app ownership, reuses `ToolManager.get_agent_tool_runtime(...)` to build -the runtime with `ToolInvokeFrom.AGENT`, invokes through -`ToolEngine.generic_invoke(...)`, and formats observations through the shared -public `ToolEngine.tool_response_to_str` helper so Agent and workflow paths -stay consistent. -""" - -from __future__ import annotations - -from sqlalchemy.orm import Session - -from core.agent.entities import AgentToolEntity -from core.app.entities.app_invoke_entities import InvokeFrom -from core.callback_handler.workflow_tool_callback_handler import DifyWorkflowCallbackHandler -from core.tools.entities.tool_entities import ToolProviderType -from core.tools.errors import ( - ToolInvokeError, - ToolNotFoundError, - ToolNotSupportedError, - ToolParameterValidationError, - ToolProviderCredentialValidationError, - ToolProviderNotFoundError, -) -from core.tools.tool_engine import ToolEngine -from core.tools.tool_manager import ToolManager -from core.tools.utils.message_transformer import ToolFileMessageTransformer -from models.model import App -from services.entities.agent_tool_inner import AgentToolInvokeRequest, AgentToolInvokeResponse -from services.errors.agent_tool_inner import AgentToolInnerServiceError - - -class AgentToolInnerService: - """Invoke one API-owned Agent tool declaration, including explicit plugin-via-core calls.""" - - def invoke(self, request: AgentToolInvokeRequest, *, session: Session) -> AgentToolInvokeResponse: - app = session.get(App, request.caller.app_id) - if app is None: - raise AgentToolInnerServiceError( - error_code="app_not_found", - description="App not found.", - status_code=404, - ) - if app.tenant_id != request.caller.tenant_id: - raise AgentToolInnerServiceError( - error_code="app_tenant_mismatch", - description="App does not belong to the caller tenant.", - status_code=403, - ) - - agent_tool = AgentToolEntity( - provider_type=ToolProviderType.value_of(request.tool.provider_type), - provider_id=request.tool.provider_id, - tool_name=request.tool.tool_name, - tool_parameters=dict(request.tool.runtime_parameters), - credential_id=request.tool.credential_id, - ) - try: - tool_runtime = ToolManager.get_agent_tool_runtime( - tenant_id=request.caller.tenant_id, - app_id=request.caller.app_id, - agent_tool=agent_tool, - user_id=request.caller.user_id, - invoke_from=InvokeFrom.value_of(request.caller.invoke_from), - variable_pool=None, - allow_file_parameters=True, - use_default_for_missing_form_parameters=True, - ) - messages = ToolEngine.generic_invoke( - session=session, - tool=tool_runtime, - tool_parameters=dict(request.tool.tool_parameters), - user_id=request.caller.user_id, - workflow_tool_callback=DifyWorkflowCallbackHandler(), - workflow_call_depth=0, - conversation_id=request.caller.conversation_id, - app_id=request.caller.app_id, - ) - transformed_messages = list( - ToolFileMessageTransformer.transform_tool_invoke_messages( - messages=messages, - user_id=request.caller.user_id, - tenant_id=request.caller.tenant_id, - conversation_id=request.caller.conversation_id, - ) - ) - except ToolProviderNotFoundError as exc: - raise AgentToolInnerServiceError( - error_code="agent_tool_declaration_not_found", - description=str(exc), - status_code=404, - ) from exc - except ToolProviderCredentialValidationError as exc: - raise AgentToolInnerServiceError( - error_code="agent_tool_credential_invalid", - description=str(exc), - status_code=422, - ) from exc - except ToolParameterValidationError as exc: - raise AgentToolInnerServiceError( - error_code="tool_parameters_invalid", - description=str(exc), - status_code=422, - ) from exc - except (ToolInvokeError, ToolNotFoundError, ToolNotSupportedError) as exc: - raise AgentToolInnerServiceError( - error_code="agent_tool_invoke_failed", - description=str(exc), - status_code=422, - ) from exc - except ValueError as exc: - raise _map_value_error(exc) from exc - except Exception as exc: - raise AgentToolInnerServiceError( - error_code="agent_tool_invoke_unexpected_error", - description=str(exc), - status_code=500, - ) from exc - - return AgentToolInvokeResponse( - messages=[message.model_dump(mode="json") for message in transformed_messages], - observation=ToolEngine.tool_response_to_str(transformed_messages), - metadata={ - "provider_type": request.tool.provider_type, - "provider_id": request.tool.provider_id, - "tool_name": request.tool.tool_name, - }, - ) - - -def _map_value_error(error: ValueError) -> AgentToolInnerServiceError: - description = str(error) - if description == "app not found": - return AgentToolInnerServiceError( - error_code="app_not_found", - description="App not found.", - status_code=404, - ) - return AgentToolInnerServiceError( - error_code="agent_tool_invoke_failed", - description=description, - status_code=422, - ) diff --git a/api/services/app_creation_records.py b/api/services/app_creation_records.py deleted file mode 100644 index f79d42d3d6fba0..00000000000000 --- a/api/services/app_creation_records.py +++ /dev/null @@ -1,45 +0,0 @@ -"""Database records required for a newly created App.""" - -from sqlalchemy import select -from sqlalchemy.orm import Session - -from models import Account -from models.enums import CustomizeTokenStrategy -from models.model import App, InstalledApp, Site - - -def create_site_record(*, app: App, account: Account, session: Session) -> None: - site = Site( - app_id=app.id, - title=app.name, - icon_type=app.icon_type, - icon=app.icon, - icon_background=app.icon_background, - default_language=account.interface_language or "en-US", - customize_token_strategy=CustomizeTokenStrategy.NOT_ALLOW, - code=Site.generate_code(16, session=session), - created_by=app.created_by, - updated_by=app.updated_by, - ) - session.add(site) - session.flush() - - -def create_installed_app_record(*, app: App, session: Session) -> None: - installed_app_id = session.scalar( - select(InstalledApp.id).where(InstalledApp.tenant_id == app.tenant_id, InstalledApp.app_id == app.id).limit(1) - ) - if installed_app_id: - return - - session.add( - InstalledApp( - tenant_id=app.tenant_id, - app_id=app.id, - app_owner_tenant_id=app.tenant_id, - ) - ) - session.flush() - - -__all__ = ["create_installed_app_record", "create_site_record"] diff --git a/api/services/conversation_variable_updater.py b/api/services/conversation_variable_updater.py deleted file mode 100644 index a821b612b41cf1..00000000000000 --- a/api/services/conversation_variable_updater.py +++ /dev/null @@ -1,29 +0,0 @@ -from sqlalchemy import select -from sqlalchemy.orm import Session, sessionmaker - -from graphon.variables.variables import VariableBase -from models import ConversationVariable - - -class ConversationVariableNotFoundError(Exception): - pass - - -class ConversationVariableUpdater: - def __init__(self, session_maker: sessionmaker[Session]) -> None: - self._session_maker: sessionmaker[Session] = session_maker - - def update(self, conversation_id: str, variable: VariableBase) -> None: - stmt = select(ConversationVariable).where( - ConversationVariable.id == ConversationVariable.storage_id(variable), - ConversationVariable.conversation_id == conversation_id, - ) - with self._session_maker() as session: - row = session.scalar(stmt) - if not row: - raise ConversationVariableNotFoundError("conversation variable not found in the database") - row.data = variable.model_dump_json() - session.commit() - - def flush(self) -> None: - pass diff --git a/api/services/knowledge/retrieval/attachments.py b/api/services/knowledge/retrieval/attachments.py deleted file mode 100644 index 9c5ae792bec5f3..00000000000000 --- a/api/services/knowledge/retrieval/attachments.py +++ /dev/null @@ -1,33 +0,0 @@ -"""Grant execution-local file access only for validated knowledge retrieval results.""" - -from collections.abc import Sequence - -from sqlalchemy.orm import Session - -from core.app.file_access import ( - get_current_file_access_scope, - grant_retriever_segment_access, - grant_upload_file_access, -) -from models.dataset import DocumentSegment -from models.model import UploadFile -from repositories.knowledge.dataset_read_repository import get_segment_attachment_files, is_retrieved_segment_owned - - -def authorize_retrieved_segment( - segment: DocumentSegment, *, tenant_id: str, dataset_ids: Sequence[str], session: Session -) -> Sequence[UploadFile] | None: - """Authorize the owner chain and attachments before content links are signed. - - A valid segment with no attachments returns an empty sequence; an invalid - retrieval result returns None and must not be included in the response. - """ - scope = get_current_file_access_scope() - if scope is not None and scope.tenant_id != tenant_id: - return None - if not is_retrieved_segment_owned(segment, tenant_id=tenant_id, dataset_ids=dataset_ids, session=session): - return None - attachments = get_segment_attachment_files(segment, session=session) - grant_retriever_segment_access([segment.id]) - grant_upload_file_access(file.id for file in attachments) - return attachments diff --git a/api/services/tools/legacy_api_tools_manage_service.py b/api/services/tools/legacy_api_tools_manage_service.py deleted file mode 100644 index a904ce14fd391e..00000000000000 --- a/api/services/tools/legacy_api_tools_manage_service.py +++ /dev/null @@ -1,570 +0,0 @@ -import json -import logging -from typing import Any, Literal, TypedDict, cast - -from pydantic import TypeAdapter -from sqlalchemy import select -from sqlalchemy.orm import sessionmaker - -from core.entities.provider_entities import ProviderConfig, ProviderConfigType -from core.helper import ssrf_proxy -from core.tools.__base.tool_runtime import ToolRuntime -from core.tools.custom_tool.provider import ApiToolProviderController -from core.tools.entities.api_entities import ToolApiEntity, ToolProviderApiEntity -from core.tools.entities.common_entities import I18nObject -from core.tools.entities.tool_bundle import ApiToolBundle -from core.tools.entities.tool_entities import ( - ApiProviderAuthType, - ApiProviderSchemaType, -) -from core.tools.errors import ApiToolProviderNotFoundError -from core.tools.tool_label_manager import ToolLabelManager -from core.tools.tool_manager import ToolManager -from core.tools.utils.encryption import create_tool_provider_encrypter -from core.tools.utils.parser import ApiBasedToolSchemaParser -from extensions.ext_database import db -from models.tools import ApiToolProvider -from services.tools.legacy_tools_transform_service import ToolTransformService - -logger = logging.getLogger(__name__) - - -class ApiSchemaParseResult(TypedDict): - schema_type: str - parameters_schema: list[dict[str, Any]] - credentials_schema: list[dict[str, Any]] - warning: dict[str, str] - - -class ApiToolPreviewResult(TypedDict, total=False): - result: str - error: str - - -class RemoteSchemaResult(TypedDict): - schema: str - - -class SimpleSuccessResult(TypedDict): - result: Literal["success"] - - -def _dump_api_tool_bundles(tool_bundles: list[ApiToolBundle]) -> list[dict[str, Any]]: - return cast(list[dict[str, Any]], TypeAdapter(list[ApiToolBundle]).dump_python(tool_bundles, mode="json")) - - -def _dump_provider_configs(configs: list[ProviderConfig]) -> list[dict[str, Any]]: - return cast(list[dict[str, Any]], TypeAdapter(list[ProviderConfig]).dump_python(configs, mode="json")) - - -class ApiToolManageService: - @staticmethod - def parser_api_schema(schema: str) -> ApiSchemaParseResult: - """ - parse api schema to tool bundle - """ - try: - warnings: dict[str, str] = {} - try: - tool_bundles, schema_type = ApiBasedToolSchemaParser.auto_parse_to_tool_bundle(schema, warning=warnings) - except Exception as e: - raise ValueError(f"invalid schema: {str(e)}") - - credentials_schema = [ - ProviderConfig( - name="auth_type", - type=ProviderConfigType.SELECT, - required=True, - default="none", - options=[ - ProviderConfig.Option(value="none", label=I18nObject(en_US="None", zh_Hans="无")), - ProviderConfig.Option(value="api_key", label=I18nObject(en_US="Api Key", zh_Hans="Api Key")), - ], - placeholder=I18nObject(en_US="Select auth type", zh_Hans="选择认证方式"), - ), - ProviderConfig( - name="api_key_header", - type=ProviderConfigType.TEXT_INPUT, - required=False, - placeholder=I18nObject(en_US="Enter api key header", zh_Hans="输入 api key header，如：X-API-KEY"), - default="api_key", - help=I18nObject(en_US="HTTP header name for api key", zh_Hans="HTTP 头部字段名，用于传递 api key"), - ), - ProviderConfig( - name="api_key_value", - type=ProviderConfigType.TEXT_INPUT, - required=False, - placeholder=I18nObject(en_US="Enter api key", zh_Hans="输入 api key"), - default="", - ), - ] - - return cast( - ApiSchemaParseResult, - { - "schema_type": schema_type.value, - "parameters_schema": _dump_api_tool_bundles(tool_bundles), - "credentials_schema": _dump_provider_configs(credentials_schema), - "warning": warnings, - }, - ) - except Exception as e: - raise ValueError(f"invalid schema: {str(e)}") - - @staticmethod - def convert_schema_to_tool_bundles( - schema: str, extra_info: dict[str, Any] | None = None - ) -> tuple[list[ApiToolBundle], ApiProviderSchemaType]: - """ - convert schema to tool bundles - - :return: the list of tool bundles, description - """ - try: - return ApiBasedToolSchemaParser.auto_parse_to_tool_bundle(schema, extra_info=extra_info) - except Exception as e: - raise ValueError(f"invalid schema: {str(e)}") - - @staticmethod - def create_api_tool_provider( - user_id: str, - tenant_id: str, - provider_name: str, - icon: dict[str, Any], - credentials: dict[str, Any], - schema_type: ApiProviderSchemaType, - schema: str, - privacy_policy: str, - custom_disclaimer: str, - labels: list[str], - ) -> SimpleSuccessResult: - """ - Create a new API tool provider. - - :param user_id: The ID of the user creating the provider. - :param tenant_id: The ID of the workspace/tenant. - :param provider_name: The name of the API tool provider. - :param icon: The icon configuration for the provider. - :param credentials: The credentials for the provider. - :param schema_type: The type of schema (e.g., OpenAPI). - :param schema: The raw schema string. - :param privacy_policy: The privacy policy URL or text. - :param custom_disclaimer: Custom disclaimer text. - :param labels: A list of labels for the provider. - :return: A dictionary indicating the result status. - """ - - provider_name = provider_name.strip() - - # check if the provider exists - # Create new session with automatic transaction management - with sessionmaker(db.engine, expire_on_commit=False).begin() as _session: - provider: ApiToolProvider | None = _session.scalar( - select(ApiToolProvider) - .where( - ApiToolProvider.tenant_id == tenant_id, - ApiToolProvider.name == provider_name, - ) - .limit(1) - ) - - if provider is not None: - raise ValueError(f"provider {provider_name} already exists") - - # parse openapi to tool bundle - extra_info: dict[str, str] = {} - # extra info like description will be set here - tool_bundles, schema_type = ApiToolManageService.convert_schema_to_tool_bundles(schema, extra_info) - - if len(tool_bundles) > 100: - raise ValueError("the number of apis should be less than 100") - - # create API tool provider - api_tool_provider = ApiToolProvider( - tenant_id=tenant_id, - user_id=user_id, - name=provider_name, - icon=json.dumps(icon), - schema=schema, - description=extra_info.get("description", ""), - schema_type_str=schema_type, - tools_str=json.dumps(_dump_api_tool_bundles(tool_bundles)), - credentials_str="{}", - privacy_policy=privacy_policy, - custom_disclaimer=custom_disclaimer, - ) - - if "auth_type" not in credentials: - raise ValueError("auth_type is required") - - # get auth type, none or api key - auth_type = ApiProviderAuthType.value_of(credentials["auth_type"]) - - # create provider entity - provider_controller = ApiToolProviderController.from_db(api_tool_provider, auth_type, session=_session) - # load tools into provider entity - provider_controller.load_bundled_tools(tool_bundles) - - # encrypt credentials - encrypter, _ = create_tool_provider_encrypter( - tenant_id=tenant_id, - controller=provider_controller, - ) - api_tool_provider.credentials_str = json.dumps(encrypter.encrypt(credentials)) - - _session.add(api_tool_provider) - - # update labels - ToolLabelManager.update_tool_labels(provider_controller, labels, _session) - - return {"result": "success"} - - @staticmethod - def get_api_tool_provider_remote_schema(user_id: str, tenant_id: str, url: str) -> RemoteSchemaResult: - """ - get api tool provider remote schema - """ - headers = { - "User-Agent": "Mozilla/5.0 (Macintosh; Intel Mac OS X 10_15_7) AppleWebKit/537.36 (KHTML, like Gecko)" - " Chrome/120.0.0.0 Safari/537.36 Edg/120.0.0.0", - "Accept": "*/*", - } - - try: - response = ssrf_proxy.get(url, headers=headers, timeout=10) - if response.status_code != 200: - raise ValueError(f"Got status code {response.status_code}") - schema = response.text - - # try to parse schema, avoid SSRF attack - ApiToolManageService.parser_api_schema(schema) - except Exception as e: - logger.exception("parse api schema error") - raise ValueError("invalid schema, please check the url you provided") from e - - return {"schema": schema} - - @staticmethod - def list_api_tool_provider_tools(user_id: str, tenant_id: str, provider_name: str) -> list[ToolApiEntity]: - """ - List tools provided by a specific API tool provider. - - :param user_id: The ID of the user requesting the list. - :param tenant_id: The ID of the workspace/tenant. - :param provider_name: The name of the API tool provider. - :return: A list of ToolApiEntity objects. - """ - - # create new session with automatic transaction management - provider: ApiToolProvider | None = None - with sessionmaker(db.engine, expire_on_commit=False).begin() as _session: - provider = _session.scalar( - select(ApiToolProvider) - .where( - ApiToolProvider.tenant_id == tenant_id, - ApiToolProvider.name == provider_name, - ) - .limit(1) - ) - - if provider is None: - raise ValueError(f"you have not added provider {provider_name}") - - controller = ToolTransformService.api_provider_to_controller(db_provider=provider) - labels = ToolLabelManager.get_tool_labels(controller) - - return [ - ToolTransformService.convert_tool_entity_to_api_entity( - tool_bundle, - tenant_id=tenant_id, - labels=labels, - ) - for tool_bundle in provider.tools - ] - - @staticmethod - def update_api_tool_provider( - user_id: str, - tenant_id: str, - provider_name: str, - original_provider: str, - icon: dict[str, Any], - credentials: dict[str, Any], - _schema_type: ApiProviderSchemaType, - schema: str, - privacy_policy: str | None, - custom_disclaimer: str, - labels: list[str], - ) -> SimpleSuccessResult: - """ - Update an existing API tool provider. - - :param user_id: The ID of the user updating the provider. - :param tenant_id: The ID of the workspace/tenant. - :param provider_name: The new name of the API tool provider. - :param original_provider: The original name of the API tool provider. - :param icon: The icon configuration for the provider. - :param credentials: The credentials for the provider. - :param _schema_type: The type of schema (e.g., OpenAPI). - :param schema: The raw schema string. - :param privacy_policy: The privacy policy URL or text. - :param custom_disclaimer: Custom disclaimer text. - :param labels: A list of labels for the provider. - :return: A dictionary indicating the result status. - """ - - provider_name = provider_name.strip() - - # check if the provider exists - # create new session with automatic transaction management - with sessionmaker(db.engine, expire_on_commit=False).begin() as _session: - provider: ApiToolProvider | None = _session.scalar( - select(ApiToolProvider) - .where( - ApiToolProvider.tenant_id == tenant_id, - ApiToolProvider.name == original_provider, - ) - .limit(1) - ) - - if provider is None: - raise ApiToolProviderNotFoundError(provider_name=original_provider, tenant_id=tenant_id) - - # parse openapi to tool bundle - extra_info: dict[str, str] = {} - # extra info like description will be set here - tool_bundles, schema_type = ApiToolManageService.convert_schema_to_tool_bundles(schema, extra_info) - - # update db provider - provider.name = provider_name - provider.icon = json.dumps(icon) - provider.schema = schema - provider.description = extra_info.get("description", "") - provider.schema_type_str = schema_type - provider.tools_str = json.dumps(_dump_api_tool_bundles(tool_bundles)) - provider.privacy_policy = privacy_policy - provider.custom_disclaimer = custom_disclaimer - - if "auth_type" not in credentials: - raise ValueError("auth_type is required") - - # get auth type, none or api key - auth_type = ApiProviderAuthType.value_of(credentials["auth_type"]) - - # create provider entity - provider_controller = ApiToolProviderController.from_db(provider, auth_type, session=_session) - # load tools into provider entity - provider_controller.load_bundled_tools(tool_bundles) - - # get original credentials if exists - encrypter, cache = create_tool_provider_encrypter( - tenant_id=tenant_id, - controller=provider_controller, - ) - - original_credentials = encrypter.decrypt(provider.credentials) - masked_credentials = encrypter.mask_plugin_credentials(original_credentials) - - # check if the credential has changed, save the original credential - for name, value in credentials.items(): - if name in masked_credentials and value == masked_credentials[name]: - credentials[name] = original_credentials[name] - - credentials = dict(encrypter.encrypt(credentials)) - provider.credentials_str = json.dumps(credentials) - - _session.add(provider) - - # update labels - ToolLabelManager.update_tool_labels(provider_controller, labels, _session) - - # delete cache - cache.delete() - - return {"result": "success"} - - @staticmethod - def delete_api_tool_provider(user_id: str, tenant_id: str, provider_name: str) -> SimpleSuccessResult: - """ - Delete an API tool provider. - - :param user_id: The ID of the user performing the deletion operation. - :param tenant_id: The ID of the workspace/tenant where the provider belongs. - :param provider_name: The unique name of the API tool provider to be deleted. - :raises ValueError: If the specified provider does not exist in the tenant. - :return: A dictionary indicating the result status. - """ - - # create new session with automatic transaction management - with sessionmaker(db.engine, expire_on_commit=False).begin() as _session: - provider: ApiToolProvider | None = _session.scalar( - select(ApiToolProvider) - .where( - ApiToolProvider.tenant_id == tenant_id, - ApiToolProvider.name == provider_name, - ) - .limit(1) - ) - - if provider is None: - raise ValueError(f"you have not added provider {provider_name}") - - _session.delete(provider) - - return {"result": "success"} - - @staticmethod - def get_api_tool_provider(user_id: str, tenant_id: str, provider: str) -> dict[str, Any]: - """ - Get API tool provider details. - - :param user_id: The ID of the user requesting the provider. - :param tenant_id: The ID of the workspace/tenant. - :param provider: The name of the API tool provider. - :return: A dictionary containing the provider details. - """ - return ToolManager.user_get_api_provider(provider=provider, tenant_id=tenant_id) - - @staticmethod - def test_api_tool_preview( - tenant_id: str, - provider_name: str, - tool_name: str, - credentials: dict[str, Any], - parameters: dict[str, Any], - schema_type: ApiProviderSchemaType | str, - schema: str, - ) -> ApiToolPreviewResult: - """ - Test an API tool before adding the API tool provider. - - :param tenant_id: The ID of the workspace/tenant. - :param provider_name: The name of the API tool provider. - :param tool_name: The name of the specific tool to test. - :param credentials: The credentials for the provider. - :param parameters: The parameters to pass to the tool. - :param schema_type: The type of schema (e.g., OpenAPI). - :param schema: The raw schema string. - :return: A dictionary containing the result or error message. - """ - - if schema_type not in [member.value for member in ApiProviderSchemaType]: - raise ValueError(f"invalid schema type {schema_type}") - - try: - tool_bundles, _ = ApiBasedToolSchemaParser.auto_parse_to_tool_bundle(schema) - except Exception as e: - raise ValueError("invalid schema") from e - - # get tool bundle - tool_bundle = next(filter(lambda tb: tb.operation_id == tool_name, tool_bundles), None) - if tool_bundle is None: - raise ValueError(f"invalid tool name {tool_name}") - - # create new session with automatic transaction management to get the provider - provider: ApiToolProvider | None = None - with sessionmaker(db.engine, expire_on_commit=False).begin() as _session: - provider = _session.scalar( - select(ApiToolProvider) - .where( - ApiToolProvider.tenant_id == tenant_id, - ApiToolProvider.name == provider_name, - ) - .limit(1) - ) - - if provider is None: - # create a fake db provider - provider = ApiToolProvider( - tenant_id="", - user_id="", - name="", - icon="", - schema=schema, - description="", - schema_type_str=ApiProviderSchemaType.OPENAPI, - tools_str=json.dumps(_dump_api_tool_bundles(tool_bundles)), - credentials_str=json.dumps(credentials), - ) - - if "auth_type" not in credentials: - raise ValueError("auth_type is required") - - # get auth type, none or api key - auth_type = ApiProviderAuthType.value_of(credentials["auth_type"]) - - # create provider entity - provider_controller = ApiToolProviderController.from_db(provider, auth_type, session=db.session()) - # load tools into provider entity - provider_controller.load_bundled_tools(tool_bundles) - - # decrypt credentials - if provider.id: - encrypter, _ = create_tool_provider_encrypter( - tenant_id=tenant_id, - controller=provider_controller, - ) - decrypted_credentials = encrypter.decrypt(credentials) - # check if the credential has changed, save the original credential - masked_credentials = encrypter.mask_plugin_credentials(decrypted_credentials) - for name, value in credentials.items(): - if name in masked_credentials and value == masked_credentials[name]: - credentials[name] = decrypted_credentials[name] - - try: - provider_controller.validate_credentials_format(credentials) - # get tool - tool = provider_controller.get_tool(tool_name) - tool = tool.fork_tool_runtime( - runtime=ToolRuntime( - credentials=credentials, - tenant_id=tenant_id, - ) - ) - result = tool.validate_credentials(credentials, parameters) - except Exception as e: - return {"error": str(e)} - - return {"result": result or "empty response"} - - @staticmethod - def list_api_tools(tenant_id: str) -> list[ToolProviderApiEntity]: - """ - List all API tools for a specific tenant. - - :param tenant_id: The ID of the workspace/tenant. - :return: A list of ToolProviderApiEntity objects. - """ - # get all api providers - # create new session with automatic transaction management - providers: list[ApiToolProvider] = [] - with sessionmaker(db.engine, expire_on_commit=False).begin() as _session: - providers = list( - _session.scalars(select(ApiToolProvider).where(ApiToolProvider.tenant_id == tenant_id)).all() - ) - - result: list[ToolProviderApiEntity] = [] - for provider in providers: - # convert provider controller to user provider - provider_controller = ToolTransformService.api_provider_to_controller(db_provider=provider) - labels = ToolLabelManager.get_tool_labels(provider_controller) - user_provider = ToolTransformService.api_provider_to_user_provider( - provider_controller, db_provider=provider, decrypt_credentials=True, session=db.session() - ) - user_provider.labels = labels - - # add icon - ToolTransformService.repack_provider(tenant_id=tenant_id, provider=user_provider) - - tools = provider_controller.get_tools(tenant_id=tenant_id) - - for tool in tools or []: - user_provider.tools.append( - ToolTransformService.convert_tool_entity_to_api_entity( - tenant_id=tenant_id, tool=tool, labels=labels - ) - ) - - result.append(user_provider) - - return result diff --git a/api/services/tools/legacy_builtin_tools_manage_service.py b/api/services/tools/legacy_builtin_tools_manage_service.py deleted file mode 100644 index d557387f91a758..00000000000000 --- a/api/services/tools/legacy_builtin_tools_manage_service.py +++ /dev/null @@ -1,794 +0,0 @@ -import json -import logging -from collections.abc import Mapping -from pathlib import Path -from typing import Any - -from sqlalchemy import delete, exists, func, select, update -from sqlalchemy.orm import Session, sessionmaker - -from configs import dify_config -from constants import HIDDEN_VALUE, UNKNOWN_VALUE -from core.helper.name_generator import generate_incremental_name -from core.helper.position_helper import is_filtered -from core.helper.provider_cache import NoOpProviderCredentialCache, ToolProviderCredentialsCache -from core.plugin.entities.plugin_daemon import CredentialType -from core.plugin.plugin_service import PluginService -from core.tools.builtin_tool.provider import BuiltinToolProviderController -from core.tools.builtin_tool.providers._positions import BuiltinToolProviderSort -from core.tools.entities.api_entities import ( - ToolApiEntity, - ToolProviderApiEntity, - ToolProviderCredentialApiEntity, - ToolProviderCredentialInfoApiEntity, -) -from core.tools.errors import ToolProviderNotFoundError -from core.tools.plugin_tool.provider import PluginToolProviderController -from core.tools.tool_label_manager import ToolLabelManager -from core.tools.tool_manager import ToolManager -from core.tools.utils.encryption import create_provider_encrypter -from core.tools.utils.system_encryption import decrypt_system_params -from extensions.ext_database import db -from extensions.ext_redis import redis_client -from models.account import Account -from models.enums import PermissionEnum -from models.provider_ids import ToolProviderID -from models.tools import BuiltinToolProvider, ToolOAuthSystemClient, ToolOAuthTenantClient -from services.credentials.query import CredentialQuery, ToolCredentialRecord -from services.tools.legacy_tools_transform_service import ToolTransformService - -logger = logging.getLogger(__name__) - - -class BuiltinToolManageService: - __MAX_BUILTIN_TOOL_PROVIDER_COUNT__ = 100 - - @staticmethod - def delete_custom_oauth_client_params(tenant_id: str, provider: str): - """ - delete custom oauth client params - """ - tool_provider = ToolProviderID(provider) - with sessionmaker(bind=db.engine).begin() as session: - session.execute( - delete(ToolOAuthTenantClient) - .where( - ToolOAuthTenantClient.tenant_id == tenant_id, - ToolOAuthTenantClient.provider == tool_provider.provider_name, - ToolOAuthTenantClient.plugin_id == tool_provider.plugin_id, - ) - .execution_options(synchronize_session=False) - ) - return {"result": "success"} - - @staticmethod - def get_builtin_tool_provider_oauth_client_schema(tenant_id: str, provider_name: str): - """ - get builtin tool provider oauth client schema - """ - provider = ToolManager.get_builtin_provider(provider_name, tenant_id) - verified = not isinstance(provider, PluginToolProviderController) or PluginService.is_plugin_verified( - tenant_id, provider.plugin_unique_identifier - ) - - is_oauth_custom_client_enabled = BuiltinToolManageService.is_oauth_custom_client_enabled( - tenant_id, provider_name - ) - is_system_oauth_params_exists = verified and BuiltinToolManageService.is_oauth_system_client_exists( - provider_name - ) - result = { - "schema": provider.get_oauth_client_schema(), - "is_oauth_custom_client_enabled": is_oauth_custom_client_enabled, - "is_system_oauth_params_exists": is_system_oauth_params_exists, - "client_params": BuiltinToolManageService.get_custom_oauth_client_params(tenant_id, provider_name), - "redirect_uri": f"{dify_config.CONSOLE_API_URL}/console/api/oauth/plugin/{provider_name}/tool/callback", - } - return result - - @staticmethod - def list_builtin_tool_provider_tools(tenant_id: str, provider: str) -> list[ToolApiEntity]: - """ - list builtin tool provider tools - - :param tenant_id: the id of the tenant - :param provider: the name of the provider - - :return: the list of tools - """ - provider_controller = ToolManager.get_builtin_provider(provider, tenant_id) - tools = provider_controller.get_tools() - - result: list[ToolApiEntity] = [] - for tool in tools or []: - result.append( - ToolTransformService.convert_tool_entity_to_api_entity( - tool=tool, - tenant_id=tenant_id, - labels=ToolLabelManager.get_tool_labels(provider_controller), - ) - ) - - return result - - @staticmethod - def get_builtin_tool_provider_info(tenant_id: str, provider: str): - """ - get builtin tool provider info - """ - provider_controller = ToolManager.get_builtin_provider(provider, tenant_id) - # check if user has added the provider - builtin_provider = BuiltinToolManageService.get_builtin_provider(provider, tenant_id) - if builtin_provider is None: - raise ValueError(f"you have not added provider {provider}") - - entity = ToolTransformService.builtin_provider_to_user_provider( - provider_controller=provider_controller, - db_provider=builtin_provider, - decrypt_credentials=True, - ) - - entity.original_credentials = {} - return entity - - @staticmethod - def list_builtin_provider_credentials_schema(provider_name: str, credential_type: CredentialType, tenant_id: str): - """ - list builtin provider credentials schema - - :param credential_type: credential type - :param provider_name: the name of the provider - :param tenant_id: the id of the tenant - :return: the list of tool providers - """ - provider = ToolManager.get_builtin_provider(provider_name, tenant_id) - return provider.get_credentials_schema_by_type(credential_type) - - @staticmethod - def update_builtin_tool_provider( - user_id: str, - tenant_id: str, - provider: str, - credential_id: str, - credentials: dict[str, Any] | None = None, - name: str | None = None, - ): - """ - update builtin tool provider - """ - with sessionmaker(bind=db.engine).begin() as session: - # get if the provider exists - db_provider = session.scalar( - select(BuiltinToolProvider) - .where( - BuiltinToolProvider.tenant_id == tenant_id, - BuiltinToolProvider.id == credential_id, - ) - .limit(1) - ) - if db_provider is None: - raise ValueError(f"you have not added provider {provider}") - - try: - if CredentialType.of(db_provider.credential_type).is_editable() and credentials: - provider_controller = ToolManager.get_builtin_provider(provider, tenant_id) - if not provider_controller.need_credentials: - raise ValueError(f"provider {provider} does not need credentials") - - encrypter, cache = BuiltinToolManageService.create_tool_encrypter( - tenant_id, db_provider, provider, provider_controller - ) - - original_credentials = encrypter.decrypt(db_provider.credentials) - new_credentials: dict[str, Any] = { - key: value if value != HIDDEN_VALUE else original_credentials.get(key, UNKNOWN_VALUE) - for key, value in credentials.items() - } - - if CredentialType.of(db_provider.credential_type).is_validate_allowed(): - provider_controller.validate_credentials(user_id, new_credentials) - - # encrypt credentials - db_provider.encrypted_credentials = json.dumps(encrypter.encrypt(new_credentials)) - - cache.delete() - - # update name if provided - if name and name != db_provider.name: - # check if the name is already used - if session.scalar( - select( - exists().where( - BuiltinToolProvider.tenant_id == tenant_id, - BuiltinToolProvider.provider == provider, - BuiltinToolProvider.name == name, - ) - ) - ): - raise ValueError(f"the credential name '{name}' is already used") - - db_provider.name = name - - # Visibility is immutable after creation — no update path. To change scope, - # create a new credential. partial-member access is handled by RBAC. - - except Exception as e: - raise ValueError(str(e)) - return {"result": "success"} - - @staticmethod - def add_builtin_tool_provider( - user_id: str, - api_type: CredentialType, - tenant_id: str, - provider: str, - credentials: dict[str, Any], - expires_at: int = -1, - name: str | None = None, - visibility: str | None = None, - ): - """ - add builtin tool provider - """ - with sessionmaker(bind=db.engine).begin() as session: - try: - lock = f"builtin_tool_provider_create_lock:{tenant_id}_{provider}" - with redis_client.lock(lock, timeout=20): - provider_controller = ToolManager.get_builtin_provider(provider, tenant_id) - if not provider_controller.need_credentials: - raise ValueError(f"provider {provider} does not need credentials") - - provider_count = ( - session.scalar( - select(func.count(BuiltinToolProvider.id)).where( - BuiltinToolProvider.tenant_id == tenant_id, - BuiltinToolProvider.provider == provider, - ) - ) - or 0 - ) - - # check if the provider count is reached the limit - if provider_count >= BuiltinToolManageService.__MAX_BUILTIN_TOOL_PROVIDER_COUNT__: - raise ValueError(f"you have reached the maximum number of providers for {provider}") - - # validate credentials if allowed - if CredentialType.of(api_type).is_validate_allowed(): - provider_controller.validate_credentials(user_id, credentials) - - # generate name if not provided - if name is None or name == "": - name = BuiltinToolManageService.generate_builtin_tool_provider_name( - session=session, tenant_id=tenant_id, provider=provider, credential_type=api_type - ) - else: - # check if the name is already used - if session.scalar( - select( - exists().where( - BuiltinToolProvider.tenant_id == tenant_id, - BuiltinToolProvider.provider == provider, - BuiltinToolProvider.name == name, - ) - ) - ): - raise ValueError(f"the credential name '{name}' is already used") - - # create encrypter - encrypter, _ = create_provider_encrypter( - tenant_id=tenant_id, - config=[ - x.to_basic_provider_config() - for x in provider_controller.get_credentials_schema_by_type(api_type) - ], - cache=NoOpProviderCredentialCache(), - ) - - visibility_enum = PermissionEnum(visibility) if visibility else PermissionEnum.ALL_TEAM - # Plugin credentials only expose only_me / all_team_members at creation; - # partial-member access is handled by workspace RBAC, not per-credential. - if visibility_enum == PermissionEnum.PARTIAL_TEAM: - raise ValueError("partial_members visibility is no longer supported for plugin credentials") - - db_provider = BuiltinToolProvider( - tenant_id=tenant_id, - user_id=user_id, - provider=provider, - encrypted_credentials=json.dumps(encrypter.encrypt(credentials)), - credential_type=api_type, - name=name, - expires_at=expires_at if expires_at is not None else -1, - visibility=visibility_enum, - ) - - session.add(db_provider) - session.flush() - except Exception as e: - raise ValueError(str(e)) - - return {"result": "success"} - - @staticmethod - def create_tool_encrypter( - tenant_id: str, - db_provider: BuiltinToolProvider | ToolCredentialRecord, - provider: str, - provider_controller: BuiltinToolProviderController, - ): - encrypter, cache = create_provider_encrypter( - tenant_id=tenant_id, - config=[ - x.to_basic_provider_config() - for x in provider_controller.get_credentials_schema_by_type(db_provider.credential_type) - ], - cache=ToolProviderCredentialsCache(tenant_id=tenant_id, provider=provider, credential_id=db_provider.id), - ) - return encrypter, cache - - @staticmethod - def generate_builtin_tool_provider_name( - tenant_id: str, provider: str, credential_type: CredentialType, *, session: Session - ) -> str: - db_providers = session.scalars( - select(BuiltinToolProvider) - .where( - BuiltinToolProvider.tenant_id == tenant_id, - BuiltinToolProvider.provider == provider, - BuiltinToolProvider.credential_type == credential_type, - ) - .order_by(BuiltinToolProvider.created_at.desc()) - ).all() - return generate_incremental_name( - [provider.name for provider in db_providers], - f"{credential_type.get_name()}", - ) - - @staticmethod - def get_builtin_tool_provider_credentials( - tenant_id: str, - provider_name: str, - user: Account | None = None, - include_credential_ids: list[str] | None = None, - *, - credential_query: CredentialQuery, - ) -> list[ToolProviderCredentialApiEntity]: - """ - get builtin tool provider credentials, filtered by visibility. - - ``user`` is used to filter the result list by per-credential visibility - (only_me / all_team_members / legacy partial_members). When ``None`` the - query returns every credential for the tenant — meant for internal / - background callers that don't act on behalf of a specific user. - - ``include_credential_ids`` lets callers request specific credential IDs that should be - returned even if the visibility filter would normally hide them (e.g. an only_me credential - owned by another member which the current workflow/agent node still references). Those - rows are marked with ``from_other_member=True`` so the UI can render them as - borrowed-from-teammate (selectable but not editable). - """ - providers = credential_query.list_tools( - workspace_id=tenant_id, - provider=provider_name, - actor_id=user.id if user else None, - include_credential_ids=include_credential_ids or (), - ) - if not providers: - return [] - provider_controller = ToolManager.get_builtin_provider(providers[0].provider, tenant_id) - credentials: list[ToolProviderCredentialApiEntity] = [] - for provider in providers: - encrypter, _ = BuiltinToolManageService.create_tool_encrypter( - tenant_id, provider, provider.provider, provider_controller - ) - decrypted = encrypter.mask_plugin_credentials(encrypter.decrypt(provider.credentials)) - credentials.append( - ToolProviderCredentialApiEntity( - id=provider.id, - name=provider.name, - provider=provider.provider, - credential_type=provider.credential_type, - is_default=provider.is_default, - credentials=dict(decrypted), - visibility=provider.visibility, - created_by=provider.created_by, - partial_member_list=list(provider.partial_member_ids), - from_other_member=provider.from_other_member, - ) - ) - return credentials - - @staticmethod - def get_builtin_tool_provider_credential_info( - tenant_id: str, - provider: str, - user: Account | None = None, - include_credential_ids: list[str] | None = None, - *, - credential_query: CredentialQuery, - ) -> ToolProviderCredentialInfoApiEntity: - """ - get builtin tool provider credential info - """ - provider_controller = ToolManager.get_builtin_provider(provider, tenant_id) - supported_credential_types = provider_controller.get_supported_credential_types() - credentials = BuiltinToolManageService.get_builtin_tool_provider_credentials( - tenant_id, - provider, - user=user, - include_credential_ids=include_credential_ids, - credential_query=credential_query, - ) - credential_info = ToolProviderCredentialInfoApiEntity( - supported_credential_types=supported_credential_types, - is_oauth_custom_client_enabled=BuiltinToolManageService.is_oauth_custom_client_enabled(tenant_id, provider), - credentials=credentials, - ) - - return credential_info - - @staticmethod - def delete_builtin_tool_provider(tenant_id: str, provider: str, credential_id: str): - """ - delete tool provider - """ - with sessionmaker(bind=db.engine).begin() as session: - db_provider = session.scalar( - select(BuiltinToolProvider) - .where( - BuiltinToolProvider.tenant_id == tenant_id, - BuiltinToolProvider.id == credential_id, - ) - .limit(1) - ) - - if db_provider is None: - raise ValueError(f"you have not added provider {provider}") - - session.delete(db_provider) - - # delete cache - provider_controller = ToolManager.get_builtin_provider(provider, tenant_id) - _, cache = BuiltinToolManageService.create_tool_encrypter( - tenant_id, db_provider, provider, provider_controller - ) - cache.delete() - - return {"result": "success"} - - @staticmethod - def set_default_provider(tenant_id: str, provider: str, id: str): - """ - set default provider - """ - with sessionmaker(bind=db.engine).begin() as session: - # get provider - target_provider = session.scalar( - select(BuiltinToolProvider) - .where(BuiltinToolProvider.id == id, BuiltinToolProvider.tenant_id == tenant_id) - .limit(1) - ) - if target_provider is None: - raise ValueError("provider not found") - - # clear default provider (tenant-scoped: only one default per provider per workspace) - session.execute( - update(BuiltinToolProvider) - .where( - BuiltinToolProvider.tenant_id == tenant_id, - BuiltinToolProvider.provider == provider, - BuiltinToolProvider.is_default.is_(True), - ) - .values(is_default=False) - .execution_options(synchronize_session=False) - ) - - # set new default provider - target_provider.is_default = True - - return {"result": "success"} - - @staticmethod - def is_oauth_system_client_exists(provider_name: str) -> bool: - """ - check if oauth system client exists - """ - tool_provider = ToolProviderID(provider_name) - with Session(db.engine, autoflush=False) as session: - system_client = session.scalar( - select(ToolOAuthSystemClient) - .where( - ToolOAuthSystemClient.plugin_id == tool_provider.plugin_id, - ToolOAuthSystemClient.provider == tool_provider.provider_name, - ) - .limit(1) - ) - return system_client is not None - - @staticmethod - def is_oauth_custom_client_enabled(tenant_id: str, provider: str) -> bool: - """ - check if oauth custom client is enabled - """ - tool_provider = ToolProviderID(provider) - with Session(db.engine, autoflush=False) as session: - user_client = session.scalar( - select(ToolOAuthTenantClient) - .where( - ToolOAuthTenantClient.tenant_id == tenant_id, - ToolOAuthTenantClient.provider == tool_provider.provider_name, - ToolOAuthTenantClient.plugin_id == tool_provider.plugin_id, - ToolOAuthTenantClient.enabled.is_(True), - ) - .limit(1) - ) - return user_client is not None and user_client.enabled - - @staticmethod - def get_oauth_client(tenant_id: str, provider: str) -> Mapping[str, Any] | None: - """ - get builtin tool provider - """ - tool_provider = ToolProviderID(provider) - provider_controller = ToolManager.get_builtin_provider(provider, tenant_id) - encrypter, _ = create_provider_encrypter( - tenant_id=tenant_id, - config=[x.to_basic_provider_config() for x in provider_controller.get_oauth_client_schema()], - cache=NoOpProviderCredentialCache(), - ) - with Session(db.engine, autoflush=False) as session: - user_client = session.scalar( - select(ToolOAuthTenantClient) - .where( - ToolOAuthTenantClient.tenant_id == tenant_id, - ToolOAuthTenantClient.provider == tool_provider.provider_name, - ToolOAuthTenantClient.plugin_id == tool_provider.plugin_id, - ToolOAuthTenantClient.enabled.is_(True), - ) - .limit(1) - ) - oauth_params: Mapping[str, Any] | None = None - if user_client: - oauth_params = encrypter.decrypt(user_client.oauth_params) - return oauth_params - - # only verified provider can use official oauth client - is_verified = not isinstance( - provider_controller, PluginToolProviderController - ) or PluginService.is_plugin_verified(tenant_id, provider_controller.plugin_unique_identifier) - if not is_verified: - return oauth_params - - system_client = session.scalar( - select(ToolOAuthSystemClient) - .where( - ToolOAuthSystemClient.plugin_id == tool_provider.plugin_id, - ToolOAuthSystemClient.provider == tool_provider.provider_name, - ) - .limit(1) - ) - if system_client: - try: - oauth_params = decrypt_system_params(system_client.encrypted_oauth_params) - except Exception as e: - raise ValueError(f"Error decrypting system oauth params: {e}") - - return oauth_params - - @staticmethod - def get_builtin_tool_provider_icon(provider: str): - """ - get tool provider icon and it's mimetype - """ - icon_path, mime_type = ToolManager.get_hardcoded_provider_icon(provider) - icon_bytes = Path(icon_path).read_bytes() - - return icon_bytes, mime_type - - @staticmethod - def list_builtin_tools(user_id: str, tenant_id: str) -> list[ToolProviderApiEntity]: - """ - list builtin tools - """ - # get all builtin providers - provider_controllers = ToolManager.list_builtin_providers(tenant_id) - - # get all user added providers - db_providers: list[BuiltinToolProvider] = ToolManager.list_default_builtin_providers(tenant_id) - - # rewrite db_providers - for db_provider in db_providers: - db_provider.provider = str(ToolProviderID(db_provider.provider)) - - # find provider - def find_provider(provider): - return next(filter(lambda db_provider: db_provider.provider == provider, db_providers), None) - - result: list[ToolProviderApiEntity] = [] - - for provider_controller in provider_controllers: - try: - # handle include, exclude - if is_filtered( - include_set=dify_config.POSITION_TOOL_INCLUDES_SET, - exclude_set=dify_config.POSITION_TOOL_EXCLUDES_SET, - data=provider_controller, - name_func=lambda x: x.entity.identity.name, - ): - continue - - # convert provider controller to user provider - user_builtin_provider = ToolTransformService.builtin_provider_to_user_provider( - provider_controller=provider_controller, - db_provider=find_provider(provider_controller.entity.identity.name), - decrypt_credentials=True, - ) - - # add icon - ToolTransformService.repack_provider(tenant_id=tenant_id, provider=user_builtin_provider) - - tools = provider_controller.get_tools() - for tool in tools or []: - user_builtin_provider.tools.append( - ToolTransformService.convert_tool_entity_to_api_entity( - tenant_id=tenant_id, - tool=tool, - labels=ToolLabelManager.get_tool_labels(provider_controller), - ) - ) - - result.append(user_builtin_provider) - except Exception as e: - raise e - - return BuiltinToolProviderSort.sort(result) - - @staticmethod - def get_builtin_provider(provider_name: str, tenant_id: str) -> BuiltinToolProvider | None: - """ - This method is used to fetch the builtin provider from the database - 1.if the default provider exists, return the default provider - 2.if the default provider does not exist, return the oldest provider - """ - with Session(db.engine, autoflush=False) as session: - try: - full_provider_name = provider_name - provider_id_entity = ToolProviderID(provider_name) - provider_name = provider_id_entity.provider_name - - if provider_id_entity.organization != "langgenius": - provider = session.scalar( - select(BuiltinToolProvider) - .where( - BuiltinToolProvider.tenant_id == tenant_id, - BuiltinToolProvider.provider == full_provider_name, - ) - .order_by( - BuiltinToolProvider.is_default.desc(), # default=True first - BuiltinToolProvider.created_at.asc(), # oldest first - ) - .limit(1) - ) - else: - provider = session.scalar( - select(BuiltinToolProvider) - .where( - BuiltinToolProvider.tenant_id == tenant_id, - (BuiltinToolProvider.provider == provider_name) - | (BuiltinToolProvider.provider == full_provider_name), - ) - .order_by( - BuiltinToolProvider.is_default.desc(), # default=True first - BuiltinToolProvider.created_at.asc(), # oldest first - ) - .limit(1) - ) - - if provider is None: - return None - - provider.provider = ToolProviderID(provider.provider).to_string() - return provider - except Exception: - # it's an old provider without organization - return session.scalar( - select(BuiltinToolProvider) - .where(BuiltinToolProvider.tenant_id == tenant_id, BuiltinToolProvider.provider == provider_name) - .order_by( - BuiltinToolProvider.is_default.desc(), # default=True first - BuiltinToolProvider.created_at.asc(), # oldest first - ) - .limit(1) - ) - - @staticmethod - def save_custom_oauth_client_params( - tenant_id: str, - provider: str, - client_params: dict[str, Any] | None = None, - enable_oauth_custom_client: bool | None = None, - ): - """ - setup oauth custom client - """ - if client_params is None and enable_oauth_custom_client is None: - return {"result": "success"} - - tool_provider = ToolProviderID(provider) - provider_controller = ToolManager.get_builtin_provider(provider, tenant_id) - if not provider_controller: - raise ToolProviderNotFoundError(f"Provider {provider} not found") - - if not isinstance(provider_controller, (BuiltinToolProviderController, PluginToolProviderController)): - raise ValueError(f"Provider {provider} is not a builtin or plugin provider") - - with sessionmaker(bind=db.engine).begin() as session: - custom_client_params = session.scalar( - select(ToolOAuthTenantClient) - .where( - ToolOAuthTenantClient.tenant_id == tenant_id, - ToolOAuthTenantClient.plugin_id == tool_provider.plugin_id, - ToolOAuthTenantClient.provider == tool_provider.provider_name, - ) - .limit(1) - ) - - # if the record does not exist, create a basic record - if custom_client_params is None: - custom_client_params = ToolOAuthTenantClient( - tenant_id=tenant_id, - plugin_id=tool_provider.plugin_id, - provider=tool_provider.provider_name, - ) - session.add(custom_client_params) - - if client_params is not None: - encrypter, _ = create_provider_encrypter( - tenant_id=tenant_id, - config=[x.to_basic_provider_config() for x in provider_controller.get_oauth_client_schema()], - cache=NoOpProviderCredentialCache(), - ) - original_params = encrypter.decrypt(custom_client_params.oauth_params) - new_params = { - key: value if value != HIDDEN_VALUE else original_params.get(key, UNKNOWN_VALUE) - for key, value in client_params.items() - } - custom_client_params.encrypted_oauth_params = json.dumps(encrypter.encrypt(new_params)) - - if enable_oauth_custom_client is not None: - custom_client_params.enabled = enable_oauth_custom_client - - return {"result": "success"} - - @staticmethod - def get_custom_oauth_client_params(tenant_id: str, provider: str): - """ - get custom oauth client params - """ - with Session(db.engine) as session: - tool_provider = ToolProviderID(provider) - custom_oauth_client_params = session.scalar( - select(ToolOAuthTenantClient) - .where( - ToolOAuthTenantClient.tenant_id == tenant_id, - ToolOAuthTenantClient.plugin_id == tool_provider.plugin_id, - ToolOAuthTenantClient.provider == tool_provider.provider_name, - ) - .limit(1) - ) - if custom_oauth_client_params is None: - return {} - - provider_controller = ToolManager.get_builtin_provider(provider, tenant_id) - if not provider_controller: - raise ToolProviderNotFoundError(f"Provider {provider} not found") - - if not isinstance(provider_controller, BuiltinToolProviderController): - raise ValueError(f"Provider {provider} is not a builtin or plugin provider") - - encrypter, _ = create_provider_encrypter( - tenant_id=tenant_id, - config=[x.to_basic_provider_config() for x in provider_controller.get_oauth_client_schema()], - cache=NoOpProviderCredentialCache(), - ) - - return encrypter.mask_plugin_credentials(encrypter.decrypt(custom_oauth_client_params.oauth_params)) diff --git a/api/services/tools/legacy_mcp_tools_manage_service.py b/api/services/tools/legacy_mcp_tools_manage_service.py deleted file mode 100644 index 357dec80ae45d1..00000000000000 --- a/api/services/tools/legacy_mcp_tools_manage_service.py +++ /dev/null @@ -1,884 +0,0 @@ -import hashlib -import json -import logging -import uuid -from collections.abc import Mapping -from datetime import datetime -from enum import StrEnum -from typing import Any -from urllib.parse import urlparse - -from pydantic import BaseModel, Field -from sqlalchemy import Select, or_, select -from sqlalchemy.exc import IntegrityError -from sqlalchemy.orm import Session - -from core.entities.mcp_provider import IdentityMode, MCPAuthentication, MCPConfiguration, MCPProviderEntity -from core.helper import encrypter -from core.helper.provider_cache import NoOpProviderCredentialCache -from core.mcp.auth.auth_flow import auth -from core.mcp.auth_client import MCPClientWithAuthRetry -from core.mcp.entities import AuthActionType, AuthResult -from core.mcp.error import MCPAuthError, MCPError -from core.mcp.types import Tool as MCPTool -from core.tools.entities.api_entities import ToolProviderApiEntity -from core.tools.utils.encryption import ProviderConfigEncrypter -from models.tools import MCPToolProvider -from services.tools.legacy_tools_transform_service import ToolTransformService - -logger = logging.getLogger(__name__) - -# Constants -UNCHANGED_SERVER_URL_PLACEHOLDER = "[__HIDDEN__]" -CLIENT_NAME = "Dify" -EMPTY_TOOLS_JSON = "[]" -EMPTY_CREDENTIALS_JSON = "{}" - - -class OAuthDataType(StrEnum): - """Types of OAuth data that can be saved.""" - - TOKENS = "tokens" - CLIENT_INFO = "client_info" - CODE_VERIFIER = "code_verifier" - MIXED = "mixed" - - -class ReconnectResult(BaseModel): - """Result of reconnecting to an MCP provider""" - - authed: bool = Field(description="Whether the provider is authenticated") - tools: str = Field(description="JSON string of tool list") - encrypted_credentials: str = Field(description="JSON string of encrypted credentials") - - -class ServerUrlValidationResult(BaseModel): - """Result of server URL validation check""" - - needs_validation: bool - validation_passed: bool = False - reconnect_result: ReconnectResult | None = None - encrypted_server_url: str | None = None - server_url_hash: str | None = None - - @property - def should_update_server_url(self) -> bool: - """Check if server URL should be updated based on validation result""" - return self.needs_validation and self.validation_passed and self.reconnect_result is not None - - -class ProviderUrlValidationData(BaseModel): - """Data required for URL validation, extracted from database to perform network operations outside of session""" - - current_server_url_hash: str - headers: dict[str, str] - timeout: float | None - sse_read_timeout: float | None - - -class MCPToolManageService: - """Service class for managing MCP tools and providers.""" - - def __init__(self, session: Session): - self._session = session - - # ========== Provider CRUD Operations ========== - - def get_provider_by_id(self, *, provider_id: str, tenant_id: str) -> MCPToolProvider: - """ - Get MCP provider by its primary key. - - Args: - provider_id: Provider primary key (UUID). Server identifiers are not - accepted here — use get_provider_by_server_identifier instead. - tenant_id: Tenant ID - - Returns: - MCPToolProvider instance - - Raises: - ValueError: If provider_id is not a UUID, or the provider does not exist - """ - try: - uuid.UUID(provider_id) - except (AttributeError, TypeError, ValueError): - raise ValueError("expected a valid UUID, but got: " + provider_id) - - stmt = select(MCPToolProvider).where(MCPToolProvider.tenant_id == tenant_id, MCPToolProvider.id == provider_id) - return self._scalar_one_provider(stmt) - - def get_provider_by_server_identifier(self, *, server_identifier: str, tenant_id: str) -> MCPToolProvider: - """ - Get MCP provider by its tenant-scoped server identifier. - - Args: - server_identifier: Server identifier (the runtime-facing reference) - tenant_id: Tenant ID - - Returns: - MCPToolProvider instance - - Raises: - ValueError: If provider not found - """ - stmt = select(MCPToolProvider).where( - MCPToolProvider.tenant_id == tenant_id, MCPToolProvider.server_identifier == server_identifier - ) - return self._scalar_one_provider(stmt) - - def _scalar_one_provider(self, stmt: Select[tuple[MCPToolProvider]]) -> MCPToolProvider: - provider = self._session.scalar(stmt) - if not provider: - raise ValueError("MCP tool not found") - return provider - - def get_provider_entity_by_id(self, *, provider_id: str, tenant_id: str) -> MCPProviderEntity: - """Get provider entity by primary key.""" - return self.get_provider_by_id(provider_id=provider_id, tenant_id=tenant_id).to_entity() - - def get_provider_entity_by_server_identifier(self, *, server_identifier: str, tenant_id: str) -> MCPProviderEntity: - """Get provider entity by server identifier.""" - return self.get_provider_by_server_identifier( - server_identifier=server_identifier, tenant_id=tenant_id - ).to_entity() - - def get_provider_by_persisted_reference(self, *, id_or_server_identifier: str, tenant_id: str) -> MCPToolProvider: - """ - Resolve an MCP provider from a reference persisted in a graph or agent config. - - For backward compatibility. Previously, the id returned by the - server is non-deterministic, depending on the query parameter, - so we accept both primary key (id) and server identifier for any existing DSLs - - Raises: - ValueError: If no provider matches the reference - """ - try: - uuid.UUID(id_or_server_identifier) - except (AttributeError, TypeError, ValueError): - # Not uuid-shaped, so it can only ever be a server identifier. Skipping - # the primary-key lookup also keeps its "expected a valid UUID" error, - # which would misdescribe a legitimate reference, out of this path. - return self.get_provider_by_server_identifier( - server_identifier=id_or_server_identifier, tenant_id=tenant_id - ) - - try: - return self.get_provider_by_id(provider_id=id_or_server_identifier, tenant_id=tenant_id) - except ValueError: - # A server identifier is free text, so it may itself look like a uuid. - return self.get_provider_by_server_identifier( - server_identifier=id_or_server_identifier, tenant_id=tenant_id - ) - - def get_provider_entity_by_persisted_reference( - self, *, id_or_server_identifier: str, tenant_id: str - ) -> MCPProviderEntity: - """Resolve a provider entity from a persisted graph or agent config reference. - - Accepts both meanings — see get_provider_by_persisted_reference. - """ - return self.get_provider_by_persisted_reference( - id_or_server_identifier=id_or_server_identifier, tenant_id=tenant_id - ).to_entity() - - def create_provider( - self, - *, - tenant_id: str, - name: str, - server_url: str, - user_id: str, - icon: str, - icon_type: str, - icon_background: str, - server_identifier: str, - configuration: MCPConfiguration, - authentication: MCPAuthentication | None = None, - headers: dict[str, str] | None = None, - identity_mode: IdentityMode = IdentityMode.OFF, - ) -> ToolProviderApiEntity: - """Create a new MCP provider.""" - # Validate URL format - if not self._is_valid_url(server_url): - raise ValueError("Server URL is not valid.") - - server_url_hash = hashlib.sha256(server_url.encode()).hexdigest() - - # Check for existing provider - self._check_provider_exists(tenant_id, name, server_url_hash, server_identifier) - - # Encrypt sensitive data - encrypted_server_url = encrypter.encrypt_token(tenant_id, server_url) - encrypted_headers = self._prepare_encrypted_dict(headers, tenant_id) if headers else None - encrypted_credentials = None - if authentication is not None and authentication.client_id: - encrypted_credentials = self._build_and_encrypt_credentials( - authentication.client_id, authentication.client_secret, tenant_id - ) - - # Create provider - mcp_tool = MCPToolProvider( - tenant_id=tenant_id, - name=name, - server_url=encrypted_server_url, - server_url_hash=server_url_hash, - user_id=user_id, - authed=False, - tools=EMPTY_TOOLS_JSON, - icon=self._prepare_icon(icon, icon_type, icon_background), - server_identifier=server_identifier, - timeout=configuration.timeout, - sse_read_timeout=configuration.sse_read_timeout, - encrypted_headers=encrypted_headers, - encrypted_credentials=encrypted_credentials, - identity_mode=identity_mode, - ) - - self._session.add(mcp_tool) - self._session.flush() - - mcp_providers = ToolTransformService.mcp_provider_to_user_provider(mcp_tool) - return mcp_providers - - def update_provider( - self, - *, - tenant_id: str, - provider_id: str, - name: str, - server_url: str, - icon: str, - icon_type: str, - icon_background: str, - server_identifier: str, - headers: dict[str, str] | None = None, - configuration: MCPConfiguration, - authentication: MCPAuthentication | None = None, - validation_result: ServerUrlValidationResult | None = None, - identity_mode: IdentityMode = IdentityMode.OFF, - ) -> None: - """ - Update an MCP provider. - - Args: - validation_result: Pre-validation result from validate_server_url_standalone. - If provided and contains reconnect_result, it will be used - instead of performing network operations. - """ - mcp_provider = self.get_provider_by_id(provider_id=provider_id, tenant_id=tenant_id) - - # Check for duplicate name (excluding current provider) - if name != mcp_provider.name: - stmt = select(MCPToolProvider).where( - MCPToolProvider.tenant_id == tenant_id, - MCPToolProvider.name == name, - MCPToolProvider.id != provider_id, - ) - existing_provider = self._session.scalar(stmt) - if existing_provider: - raise ValueError(f"MCP tool {name} already exists") - - # Get URL update data from validation result - encrypted_server_url = None - server_url_hash = None - reconnect_result = None - - if validation_result and validation_result.encrypted_server_url: - # Use all data from validation result - encrypted_server_url = validation_result.encrypted_server_url - server_url_hash = validation_result.server_url_hash - reconnect_result = validation_result.reconnect_result - - try: - # Update basic fields - mcp_provider.updated_at = datetime.now() - mcp_provider.name = name - mcp_provider.icon = self._prepare_icon(icon, icon_type, icon_background) - mcp_provider.server_identifier = server_identifier - - # Update server URL if changed - if encrypted_server_url and server_url_hash: - mcp_provider.server_url = encrypted_server_url - mcp_provider.server_url_hash = server_url_hash - - if reconnect_result: - mcp_provider.authed = reconnect_result.authed - mcp_provider.tools = reconnect_result.tools - mcp_provider.encrypted_credentials = reconnect_result.encrypted_credentials - - # Update optional configuration fields - self._update_optional_fields(mcp_provider, configuration) - - # Update headers if provided - if headers is not None: - mcp_provider.encrypted_headers = self._process_headers(headers, mcp_provider, tenant_id) - - # Update credentials if provided - if authentication and authentication.client_id: - mcp_provider.encrypted_credentials = self._process_credentials(authentication, mcp_provider, tenant_id) - - # Update user-identity forwarding mode. The controller has already - # resolved "leave unchanged" and applied the Enterprise-edition gate, - # so this is always a concrete, vetted value. - mcp_provider.identity_mode = identity_mode - - # Flush changes to database - self._session.flush() - - except IntegrityError as e: - self._handle_integrity_error(e, name, server_url, server_identifier) - - def delete_provider(self, *, tenant_id: str, provider_id: str) -> None: - """Delete an MCP provider.""" - mcp_tool = self.get_provider_by_id(provider_id=provider_id, tenant_id=tenant_id) - self._session.delete(mcp_tool) - - def list_providers(self, *, tenant_id: str, include_sensitive: bool = True) -> list[ToolProviderApiEntity]: - """List all MCP providers for a tenant. - - Args: - tenant_id: Tenant ID - include_sensitive: If False, skip expensive decryption operations (default: True for backward compatibility) - """ - from models.account import Account - - stmt = select(MCPToolProvider).where(MCPToolProvider.tenant_id == tenant_id).order_by(MCPToolProvider.name) - mcp_providers = self._session.scalars(stmt).all() - - if not mcp_providers: - return [] - - # Batch query all users to avoid N+1 problem - user_ids = {provider.user_id for provider in mcp_providers} - users = self._session.scalars(select(Account).where(Account.id.in_(user_ids))).all() - user_name_map = {user.id: user.name for user in users} - - return [ - ToolTransformService.mcp_provider_to_user_provider( - provider, - user_name=user_name_map.get(provider.user_id), - include_sensitive=include_sensitive, - ) - for provider in mcp_providers - ] - - # ========== Tool Operations ========== - - def list_provider_tools(self, *, tenant_id: str, provider_id: str) -> ToolProviderApiEntity: - """List tools from remote MCP server.""" - # Load provider and convert to entity - db_provider = self.get_provider_by_id(provider_id=provider_id, tenant_id=tenant_id) - provider_entity = db_provider.to_entity() - - # Verify authentication - if not provider_entity.authed: - raise ValueError("Please auth the tool first") - - # Prepare headers with auth token - headers = self._prepare_auth_headers(provider_entity) - - # Retrieve tools from remote server - server_url = provider_entity.decrypt_server_url() - try: - tools = self._retrieve_remote_mcp_tools(server_url, headers, provider_entity) - except MCPError as e: - raise ValueError(f"Failed to connect to MCP server: {e}") - - # Update database with retrieved tools (ensure description is a non-null string) - tools_payload = [] - for tool in tools: - data = tool.model_dump() - if data.get("description") is None: - data["description"] = "" - tools_payload.append(data) - db_provider.tools = json.dumps(tools_payload) - db_provider.authed = True - db_provider.updated_at = datetime.now() - self._session.flush() - - # Build API response - return self._build_tool_provider_response(db_provider, provider_entity, tools) - - # ========== OAuth and Credentials Operations ========== - - def update_provider_credentials( - self, *, provider_id: str, tenant_id: str, credentials: dict[str, Any], authed: bool | None = None - ) -> None: - """ - Update provider credentials with encryption. - - Args: - provider_id: Provider ID - tenant_id: Tenant ID - credentials: Credentials to save - authed: Whether provider is authenticated (None means keep current state) - """ - from core.tools.mcp_tool.provider import MCPToolProviderController - - # Get provider from current session - provider = self.get_provider_by_id(provider_id=provider_id, tenant_id=tenant_id) - - # Encrypt new credentials - provider_controller = MCPToolProviderController.from_db(provider) - tool_configuration = ProviderConfigEncrypter( - tenant_id=provider.tenant_id, - config=list(provider_controller.get_credentials_schema()), - provider_config_cache=NoOpProviderCredentialCache(), - ) - encrypted_credentials = tool_configuration.encrypt(credentials) - - # Update provider - provider.updated_at = datetime.now() - provider.encrypted_credentials = json.dumps({**provider.credentials, **encrypted_credentials}) - - if authed is not None: - provider.authed = authed - if not authed: - provider.tools = EMPTY_TOOLS_JSON - - # Flush changes to database - self._session.flush() - - def save_oauth_data( - self, provider_id: str, tenant_id: str, data: dict[str, Any], data_type: OAuthDataType = OAuthDataType.MIXED - ) -> None: - """ - Save OAuth-related data (tokens, client info, code verifier). - - Args: - provider_id: Provider ID - tenant_id: Tenant ID - data: Data to save (tokens, client info, or code verifier) - data_type: Type of OAuth data to save - """ - # Determine if this makes the provider authenticated - authed = ( - data_type == OAuthDataType.TOKENS or (data_type == OAuthDataType.MIXED and "access_token" in data) or None - ) - - # update_provider_credentials will validate provider existence - self.update_provider_credentials(provider_id=provider_id, tenant_id=tenant_id, credentials=data, authed=authed) - - def clear_provider_credentials(self, *, provider_id: str, tenant_id: str) -> None: - """ - Clear all credentials for a provider. - - Args: - provider_id: Provider ID - tenant_id: Tenant ID - """ - # Get provider from current session - provider = self.get_provider_by_id(provider_id=provider_id, tenant_id=tenant_id) - - provider.tools = EMPTY_TOOLS_JSON - provider.encrypted_credentials = EMPTY_CREDENTIALS_JSON - provider.updated_at = datetime.now() - provider.authed = False - - # ========== Private Helper Methods ========== - - def _check_provider_exists(self, tenant_id: str, name: str, server_url_hash: str, server_identifier: str) -> None: - """Check if provider with same attributes already exists.""" - stmt = select(MCPToolProvider).where( - MCPToolProvider.tenant_id == tenant_id, - or_( - MCPToolProvider.name == name, - MCPToolProvider.server_url_hash == server_url_hash, - MCPToolProvider.server_identifier == server_identifier, - ), - ) - existing_provider = self._session.scalar(stmt) - - if existing_provider: - if existing_provider.name == name: - raise ValueError(f"MCP tool {name} already exists") - if existing_provider.server_url_hash == server_url_hash: - raise ValueError("MCP tool with this server URL already exists") - if existing_provider.server_identifier == server_identifier: - raise ValueError(f"MCP tool {server_identifier} already exists") - - def _prepare_icon(self, icon: str, icon_type: str, icon_background: str) -> str: - """Prepare icon data for storage.""" - if icon_type == "emoji": - return json.dumps({"content": icon, "background": icon_background}) - return icon - - def _encrypt_dict_fields(self, data: dict[str, Any], secret_fields: list[str], tenant_id: str) -> Mapping[str, str]: - """Encrypt specified fields in a dictionary. - - Args: - data: Dictionary containing data to encrypt - secret_fields: List of field names to encrypt - tenant_id: Tenant ID for encryption - - Returns: - JSON string of encrypted data - """ - from core.entities.provider_entities import BasicProviderConfig, ProviderConfigType - from core.tools.utils.encryption import create_provider_encrypter - - # Create config for secret fields - config = [BasicProviderConfig(type=ProviderConfigType.SECRET_INPUT, name=field) for field in secret_fields] - - encrypter_instance, _ = create_provider_encrypter( - tenant_id=tenant_id, - config=config, - cache=NoOpProviderCredentialCache(), - ) - - encrypted_data = encrypter_instance.encrypt(data) - return encrypted_data - - def _prepare_encrypted_dict(self, headers: dict[str, str], tenant_id: str) -> str: - """Encrypt headers and prepare for storage.""" - # All headers are treated as secret - return json.dumps(self._encrypt_dict_fields(headers, list(headers.keys()), tenant_id)) - - def _prepare_auth_headers(self, provider_entity: MCPProviderEntity) -> dict[str, str]: - """Prepare headers with OAuth token if available.""" - headers = provider_entity.decrypt_headers() - tokens = provider_entity.retrieve_tokens() - if tokens: - headers["Authorization"] = f"{tokens.token_type.capitalize()} {tokens.access_token}" - return headers - - def _retrieve_remote_mcp_tools( - self, - server_url: str, - headers: dict[str, str], - provider_entity: MCPProviderEntity, - ): - """Retrieve tools from remote MCP server.""" - with MCPClientWithAuthRetry( - server_url=server_url, - headers=headers, - timeout=provider_entity.timeout, - sse_read_timeout=provider_entity.sse_read_timeout, - provider_entity=provider_entity, - ) as mcp_client: - return mcp_client.list_tools() - - _ACTION_TO_OAUTH: dict[AuthActionType, OAuthDataType] = { - AuthActionType.SAVE_CLIENT_INFO: OAuthDataType.CLIENT_INFO, - AuthActionType.SAVE_TOKENS: OAuthDataType.TOKENS, - AuthActionType.SAVE_CODE_VERIFIER: OAuthDataType.CODE_VERIFIER, - } - - def execute_auth_actions(self, auth_result: AuthResult) -> dict[str, str]: - """ - Execute the actions returned by the auth function. - - This method processes the AuthResult and performs the necessary database operations. - - Args: - auth_result: The result from the auth function - - Returns: - The response from the auth result - """ - for action in auth_result.actions: - if action.provider_id is None or action.tenant_id is None: - continue - - oauth_type = self._ACTION_TO_OAUTH.get(action.action_type) - if oauth_type is not None: - self.save_oauth_data(action.provider_id, action.tenant_id, action.data, oauth_type) - - return auth_result.response - - def auth_with_actions( - self, - provider_entity: MCPProviderEntity, - authorization_code: str | None = None, - resource_metadata_url: str | None = None, - scope_hint: str | None = None, - ) -> dict[str, str]: - """ - Perform authentication and execute all resulting actions. - - This method is used by MCPClientWithAuthRetry for automatic re-authentication. - - Args: - provider_entity: The MCP provider entity - authorization_code: Optional authorization code - resource_metadata_url: Optional Protected Resource Metadata URL from WWW-Authenticate - scope_hint: Optional scope hint from WWW-Authenticate header - - Returns: - Response dictionary from auth result - """ - auth_result = auth( - provider_entity, - authorization_code, - resource_metadata_url=resource_metadata_url, - scope_hint=scope_hint, - ) - return self.execute_auth_actions(auth_result) - - def get_provider_for_url_validation(self, *, tenant_id: str, provider_id: str) -> ProviderUrlValidationData: - """ - Get provider data required for URL validation. - This method performs database read and should be called within a session. - - Returns: - ProviderUrlValidationData: Data needed for standalone URL validation - """ - provider = self.get_provider_by_id(provider_id=provider_id, tenant_id=tenant_id) - provider_entity = provider.to_entity() - return ProviderUrlValidationData( - current_server_url_hash=provider.server_url_hash, - headers=provider_entity.headers, - timeout=provider_entity.timeout, - sse_read_timeout=provider_entity.sse_read_timeout, - ) - - @staticmethod - def validate_server_url_standalone( - *, - tenant_id: str, - new_server_url: str, - validation_data: ProviderUrlValidationData, - ) -> ServerUrlValidationResult: - """ - Validate server URL change by attempting to connect to the new server. - This method performs network operations and MUST be called OUTSIDE of any database session - to avoid holding locks during network I/O. - - Args: - tenant_id: Tenant ID for encryption - new_server_url: The new server URL to validate - validation_data: Provider data obtained from get_provider_for_url_validation - - Returns: - ServerUrlValidationResult: Validation result with connection status and tools if successful - """ - # Handle hidden/unchanged URL - if UNCHANGED_SERVER_URL_PLACEHOLDER in new_server_url: - return ServerUrlValidationResult(needs_validation=False) - - # Validate URL format - parsed = urlparse(new_server_url) - if not all([parsed.scheme, parsed.netloc]) or parsed.scheme not in ["http", "https"]: - raise ValueError("Server URL is not valid.") - - # Always encrypt and hash the URL - encrypted_server_url = encrypter.encrypt_token(tenant_id, new_server_url) - new_server_url_hash = hashlib.sha256(new_server_url.encode()).hexdigest() - - # Check if URL is actually different - if new_server_url_hash == validation_data.current_server_url_hash: - # URL hasn't changed, but still return the encrypted data - return ServerUrlValidationResult( - needs_validation=False, - encrypted_server_url=encrypted_server_url, - server_url_hash=new_server_url_hash, - ) - - # Perform network validation - this is the expensive operation that should be outside session - reconnect_result = MCPToolManageService._reconnect_with_url( - server_url=new_server_url, - headers=validation_data.headers, - timeout=validation_data.timeout, - sse_read_timeout=validation_data.sse_read_timeout, - ) - return ServerUrlValidationResult( - needs_validation=True, - validation_passed=True, - reconnect_result=reconnect_result, - encrypted_server_url=encrypted_server_url, - server_url_hash=new_server_url_hash, - ) - - @staticmethod - def reconnect_with_url( - *, - server_url: str, - headers: dict[str, str], - timeout: float | None, - sse_read_timeout: float | None, - ) -> ReconnectResult: - return MCPToolManageService._reconnect_with_url( - server_url=server_url, - headers=headers, - timeout=timeout, - sse_read_timeout=sse_read_timeout, - ) - - @staticmethod - def _reconnect_with_url( - *, - server_url: str, - headers: dict[str, str], - timeout: float | None, - sse_read_timeout: float | None, - ) -> ReconnectResult: - """ - Attempt to connect to MCP server with given URL. - This is a static method that performs network I/O without database access. - """ - from core.mcp.mcp_client import MCPClient - - try: - with MCPClient( - server_url=server_url, - headers=headers, - timeout=timeout, - sse_read_timeout=sse_read_timeout, - ) as mcp_client: - tools = mcp_client.list_tools() - # Ensure tool descriptions are non-null in payload - tools_payload = [] - for t in tools: - d = t.model_dump() - if d.get("description") is None: - d["description"] = "" - tools_payload.append(d) - return ReconnectResult( - authed=True, - tools=json.dumps(tools_payload), - encrypted_credentials=EMPTY_CREDENTIALS_JSON, - ) - except MCPAuthError: - return ReconnectResult(authed=False, tools=EMPTY_TOOLS_JSON, encrypted_credentials=EMPTY_CREDENTIALS_JSON) - except MCPError as e: - raise ValueError(f"Failed to re-connect MCP server: {e}") from e - - def _build_tool_provider_response( - self, db_provider: MCPToolProvider, provider_entity: MCPProviderEntity, tools: list[MCPTool] - ) -> ToolProviderApiEntity: - """Build API response for tool provider.""" - user = db_provider.load_user(self._session) - user_name = user.name if user else None - response = provider_entity.to_api_response( - user_name=user_name, - ) - response["tools"] = ToolTransformService.mcp_tool_to_user_tool(db_provider, tools, user_name=user_name) - response["plugin_unique_identifier"] = provider_entity.server_identifier - return ToolProviderApiEntity(**response) - - def _handle_integrity_error( - self, error: IntegrityError, name: str, server_url: str, server_identifier: str - ) -> None: - """Handle database integrity errors with user-friendly messages.""" - error_msg = str(error.orig) - if "unique_mcp_provider_name" in error_msg: - raise ValueError(f"MCP tool {name} already exists") - if "unique_mcp_provider_server_url" in error_msg: - raise ValueError(f"MCP tool {server_url} already exists") - if "unique_mcp_provider_server_identifier" in error_msg: - raise ValueError(f"MCP tool {server_identifier} already exists") - raise error - - def _is_valid_url(self, url: str) -> bool: - """Validate URL format.""" - if not url: - return False - try: - parsed = urlparse(url) - return all([parsed.scheme, parsed.netloc]) and parsed.scheme in ["http", "https"] - except (ValueError, TypeError): - return False - - def _update_optional_fields(self, mcp_provider: MCPToolProvider, configuration: MCPConfiguration) -> None: - """Update optional configuration fields using setattr for cleaner code.""" - field_mapping = {"timeout": configuration.timeout, "sse_read_timeout": configuration.sse_read_timeout} - - for field, value in field_mapping.items(): - if value is not None: - setattr(mcp_provider, field, value) - - def _process_headers(self, headers: dict[str, str], mcp_provider: MCPToolProvider, tenant_id: str) -> str | None: - """Process headers update, handling empty dict to clear headers.""" - if not headers: - return None - - # Merge with existing headers to preserve masked values - final_headers = self._merge_headers_with_masked(incoming_headers=headers, mcp_provider=mcp_provider) - return self._prepare_encrypted_dict(final_headers, tenant_id) - - def _process_credentials( - self, authentication: MCPAuthentication, mcp_provider: MCPToolProvider, tenant_id: str - ) -> str: - """Process credentials update, handling masked values.""" - # Merge with existing credentials - final_client_id, final_client_secret = self._merge_credentials_with_masked( - authentication.client_id, authentication.client_secret, mcp_provider - ) - - # Build and encrypt - return self._build_and_encrypt_credentials(final_client_id, final_client_secret, tenant_id) - - def _merge_headers_with_masked( - self, incoming_headers: dict[str, str], mcp_provider: MCPToolProvider - ) -> dict[str, str]: - """Merge incoming headers with existing ones, preserving unchanged masked values. - - Args: - incoming_headers: Headers from frontend (may contain masked values) - mcp_provider: The MCP provider instance - - Returns: - Final headers dict with proper values (original for unchanged masked, new for changed) - """ - mcp_provider_entity = mcp_provider.to_entity() - existing_decrypted = mcp_provider_entity.decrypt_headers() - existing_masked = mcp_provider_entity.masked_headers() - - return { - key: (str(existing_decrypted[key]) if key in existing_masked and value == existing_masked[key] else value) - for key, value in incoming_headers.items() - if key in existing_decrypted or value != existing_masked.get(key) - } - - def _merge_credentials_with_masked( - self, - client_id: str, - client_secret: str | None, - mcp_provider: MCPToolProvider, - ) -> tuple[ - str, - str | None, - ]: - """Merge incoming credentials with existing ones, preserving unchanged masked values. - - Args: - client_id: Client ID from frontend (may be masked) - client_secret: Client secret from frontend (may be masked) - mcp_provider: The MCP provider instance - - Returns: - Tuple of (final_client_id, final_client_secret) - """ - mcp_provider_entity = mcp_provider.to_entity() - existing_decrypted = mcp_provider_entity.decrypt_credentials() - existing_masked = mcp_provider_entity.masked_credentials() - - # Check if client_id is masked and unchanged - final_client_id = client_id - if existing_masked.get("client_id") and client_id == existing_masked["client_id"]: - # Use existing decrypted value - final_client_id = existing_decrypted.get("client_id", client_id) - - # Check if client_secret is masked and unchanged - final_client_secret = client_secret - if existing_masked.get("client_secret") and client_secret == existing_masked["client_secret"]: - # Use existing decrypted value - final_client_secret = existing_decrypted.get("client_secret", client_secret) - - return final_client_id, final_client_secret - - def _build_and_encrypt_credentials(self, client_id: str, client_secret: str | None, tenant_id: str) -> str: - """Build credentials and encrypt sensitive fields.""" - # Create a flat structure with all credential data - credentials_data = { - "client_id": client_id, - "client_name": CLIENT_NAME, - "is_dynamic_registration": False, - } - secret_fields = [] - if client_secret is not None: - credentials_data["encrypted_client_secret"] = client_secret - secret_fields = ["encrypted_client_secret"] - client_info = self._encrypt_dict_fields(credentials_data, secret_fields, tenant_id) - return json.dumps({"client_information": client_info}) diff --git a/api/services/tools/legacy_tools_manage_service.py b/api/services/tools/legacy_tools_manage_service.py deleted file mode 100644 index 44db6a97e36a59..00000000000000 --- a/api/services/tools/legacy_tools_manage_service.py +++ /dev/null @@ -1,26 +0,0 @@ -import logging - -from core.tools.entities.api_entities import ToolProviderTypeApiLiteral -from core.tools.tool_manager import ToolManager -from services.tools.legacy_tools_transform_service import ToolTransformService - -logger = logging.getLogger(__name__) - - -class ToolCommonService: - @staticmethod - def list_tool_providers(user_id: str, tenant_id: str, typ: ToolProviderTypeApiLiteral | None = None): - """ - list tool providers - - :return: the list of tool providers - """ - providers = ToolManager.list_providers_from_api(user_id, tenant_id, typ) - - # add icon - for provider in providers: - ToolTransformService.repack_provider(tenant_id=tenant_id, provider=provider) - - result = [provider.to_dict() for provider in providers] - - return result diff --git a/api/services/tools/legacy_tools_transform_service.py b/api/services/tools/legacy_tools_transform_service.py deleted file mode 100644 index ee5729a0c4cbd4..00000000000000 --- a/api/services/tools/legacy_tools_transform_service.py +++ /dev/null @@ -1,538 +0,0 @@ -import logging -from collections.abc import Mapping -from typing import Any - -from pydantic import TypeAdapter, ValidationError -from sqlalchemy.orm import Session -from yarl import URL - -from configs import dify_config -from core.helper.provider_cache import ToolProviderCredentialsCache -from core.mcp.types import Tool as MCPTool -from core.plugin.entities.plugin_daemon import CredentialType, PluginDatasourceProviderEntity -from core.plugin.plugin_service import PluginService -from core.tools.__base.tool import Tool -from core.tools.__base.tool_runtime import ToolRuntime -from core.tools.builtin_tool.provider import BuiltinToolProviderController -from core.tools.custom_tool.provider import ApiToolProviderController -from core.tools.entities.api_entities import ToolApiEntity, ToolProviderApiEntity, ToolProviderCredentialApiEntity -from core.tools.entities.common_entities import I18nObject -from core.tools.entities.tool_bundle import ApiToolBundle -from core.tools.entities.tool_entities import ( - ApiProviderAuthType, - ToolParameter, - ToolProviderType, - emoji_icon_adapter, -) -from core.tools.plugin_tool.provider import PluginToolProviderController -from core.tools.utils.encryption import create_provider_encrypter, create_tool_provider_encrypter -from core.tools.workflow_as_tool.provider import WorkflowToolProviderController -from core.tools.workflow_as_tool.tool import WorkflowTool -from extensions.ext_database import db -from models.tools import ApiToolProvider, BuiltinToolProvider, MCPToolProvider, WorkflowToolProvider - -logger = logging.getLogger(__name__) - -_mcp_tools_adapter: TypeAdapter[list[MCPTool]] = TypeAdapter(list[MCPTool]) - - -class ToolTransformService: - _MCP_SCHEMA_TYPE_RESOLUTION_MAX_DEPTH = 10 - - @classmethod - def get_tool_provider_icon_url( - cls, provider_type: str, provider_name: str, icon: str | Mapping[str, str] - ) -> str | Mapping[str, str]: - """ - get tool provider icon url - """ - url_prefix = ( - URL(dify_config.CONSOLE_API_URL or "/") / "console" / "api" / "workspaces" / "current" / "tool-provider" - ) - - match provider_type: - case ToolProviderType.BUILT_IN: - return str(url_prefix / "builtin" / provider_name / "icon") - case ToolProviderType.API | ToolProviderType.WORKFLOW: - try: - if isinstance(icon, str): - parsed = emoji_icon_adapter.validate_json(icon) - return {"background": parsed["background"], "content": parsed["content"]} - return {"background": icon["background"], "content": icon["content"]} - except (ValueError, ValidationError, KeyError): - return {"background": "#252525", "content": "\ud83d\ude01"} - case ToolProviderType.MCP: - if isinstance(icon, Mapping): - return {"background": icon.get("background", ""), "content": icon.get("content", "")} - return icon - case ToolProviderType.PLUGIN | ToolProviderType.APP | ToolProviderType.DATASET_RETRIEVAL: - return "" - case _: - return "" - - @staticmethod - def repack_provider( - tenant_id: str, provider: dict[str, Any] | ToolProviderApiEntity | PluginDatasourceProviderEntity - ): - """ - repack provider - - :param tenant_id: the tenant id - :param provider: the provider dict - """ - match provider: - case dict() if "icon" in provider: - provider["icon"] = ToolTransformService.get_tool_provider_icon_url( - provider_type=provider["type"], provider_name=provider["name"], icon=provider["icon"] - ) - case ToolProviderApiEntity(): - if provider.plugin_id: - if isinstance(provider.icon, str): - provider.icon = PluginService.get_plugin_icon_url(tenant_id=tenant_id, filename=provider.icon) - if isinstance(provider.icon_dark, str) and provider.icon_dark: - provider.icon_dark = PluginService.get_plugin_icon_url( - tenant_id=tenant_id, filename=provider.icon_dark - ) - else: - provider.icon = ToolTransformService.get_tool_provider_icon_url( - provider_type=provider.type.value, provider_name=provider.name, icon=provider.icon - ) - if provider.icon_dark: - provider.icon_dark = ToolTransformService.get_tool_provider_icon_url( - provider_type=provider.type.value, provider_name=provider.name, icon=provider.icon_dark - ) - case PluginDatasourceProviderEntity(): - if provider.plugin_id: - if isinstance(provider.declaration.identity.icon, str): - provider.declaration.identity.icon = PluginService.get_plugin_icon_url( - tenant_id=tenant_id, filename=provider.declaration.identity.icon - ) - - @classmethod - def builtin_provider_to_user_provider( - cls, - provider_controller: BuiltinToolProviderController | PluginToolProviderController, - db_provider: BuiltinToolProvider | None, - decrypt_credentials: bool = True, - ) -> ToolProviderApiEntity: - """ - convert provider controller to user provider - """ - result = ToolProviderApiEntity( - id=provider_controller.entity.identity.name, - author=provider_controller.entity.identity.author, - name=provider_controller.entity.identity.name, - description=provider_controller.entity.identity.description, - icon=provider_controller.entity.identity.icon, - icon_dark=provider_controller.entity.identity.icon_dark or "", - label=provider_controller.entity.identity.label, - type=ToolProviderType.BUILT_IN, - masked_credentials={}, - is_team_authorization=False, - plugin_id=None, - tools=[], - labels=provider_controller.tool_labels, - ) - - if isinstance(provider_controller, PluginToolProviderController): - result.plugin_id = provider_controller.plugin_id - result.plugin_unique_identifier = provider_controller.plugin_unique_identifier - - # get credentials schema - schema = { - x.to_basic_provider_config().name: x - for x in provider_controller.get_credentials_schema_by_type( - CredentialType.of(db_provider.credential_type) if db_provider else CredentialType.API_KEY - ) - } - - masked_creds = {} - for name in schema: - masked_creds[name] = "" - result.masked_credentials = masked_creds - - # check if the provider need credentials - if not provider_controller.need_credentials: - result.is_team_authorization = True - result.allow_delete = False - elif db_provider: - result.is_team_authorization = True - - if decrypt_credentials: - credentials = db_provider.credentials - if not db_provider.tenant_id: - raise ValueError(f"Required tenant_id is missing for BuiltinToolProvider with id {db_provider.id}") - # init tool configuration - encrypter, _ = create_provider_encrypter( - tenant_id=db_provider.tenant_id, - config=[ - x.to_basic_provider_config() - for x in provider_controller.get_credentials_schema_by_type( - CredentialType.of(db_provider.credential_type) - ) - ], - cache=ToolProviderCredentialsCache( - tenant_id=db_provider.tenant_id, - provider=db_provider.provider, - credential_id=db_provider.id, - ), - ) - # decrypt the credentials and mask the credentials - decrypted_credentials = encrypter.decrypt(data=credentials) - masked_credentials = encrypter.mask_plugin_credentials(data=decrypted_credentials) - - result.masked_credentials = masked_credentials - result.original_credentials = decrypted_credentials - - return result - - @staticmethod - def api_provider_to_controller( - db_provider: ApiToolProvider, - ) -> ApiToolProviderController: - """ - convert provider controller to user provider - """ - # package tool provider controller - auth_type = ApiProviderAuthType.NONE - credentials_auth_type = db_provider.credentials.get("auth_type") - if credentials_auth_type in ("api_key_header", "api_key"): # backward compatibility - auth_type = ApiProviderAuthType.API_KEY_HEADER - elif credentials_auth_type == "api_key_query": - auth_type = ApiProviderAuthType.API_KEY_QUERY - - controller = ApiToolProviderController.from_db( - db_provider=db_provider, - auth_type=auth_type, - session=db.session(), - ) - - return controller - - @staticmethod - def workflow_provider_to_controller(db_provider: WorkflowToolProvider) -> WorkflowToolProviderController: - """ - convert provider controller to provider - """ - return WorkflowToolProviderController.from_db(db_provider) - - @staticmethod - def workflow_provider_to_user_provider( - provider_controller: WorkflowToolProviderController, - labels: list[str] | None = None, - workflow_app_id: str | None = None, - ): - """ - convert provider controller to user provider - """ - return ToolProviderApiEntity( - id=provider_controller.provider_id, - author=provider_controller.entity.identity.author, - name=provider_controller.entity.identity.name, - description=provider_controller.entity.identity.description, - icon=provider_controller.entity.identity.icon, - icon_dark=provider_controller.entity.identity.icon_dark or "", - label=provider_controller.entity.identity.label, - type=ToolProviderType.WORKFLOW, - masked_credentials={}, - is_team_authorization=True, - plugin_id=None, - plugin_unique_identifier=None, - tools=[], - labels=labels or [], - workflow_app_id=workflow_app_id, - ) - - @staticmethod - def mcp_provider_to_user_provider( - db_provider: MCPToolProvider, - user_name: str | None = None, - include_sensitive: bool = True, - ) -> ToolProviderApiEntity: - from core.entities.mcp_provider import MCPConfiguration - - # Use provided user_name to avoid N+1 query, fallback to load_user() if not provided - if user_name is None: - user = db_provider.load_user(db.session()) - user_name = user.name if user else None - - # Convert to entity and use its API response method - provider_entity = db_provider.to_entity() - - response = provider_entity.to_api_response(user_name=user_name, include_sensitive=include_sensitive) - try: - mcp_tools = _mcp_tools_adapter.validate_json(db_provider.tools) - except (ValidationError, ValueError): - mcp_tools = [] - # Add additional fields specific to the transform - response["id"] = db_provider.id - response["tools"] = ToolTransformService.mcp_tool_to_user_tool(db_provider, mcp_tools, user_name=user_name) - response["server_identifier"] = db_provider.server_identifier - - # Convert configuration dict to MCPConfiguration object - if "configuration" in response and isinstance(response["configuration"], dict): - response["configuration"] = MCPConfiguration( - timeout=float(response["configuration"]["timeout"]), - sse_read_timeout=float(response["configuration"]["sse_read_timeout"]), - ) - - return ToolProviderApiEntity(**response) - - @staticmethod - def mcp_tool_to_user_tool( - mcp_provider: MCPToolProvider, tools: list[MCPTool], user_name: str | None = None - ) -> list[ToolApiEntity]: - # Use provided user_name to avoid N+1 query, fallback to load_user() if not provided - if user_name is None: - user = mcp_provider.load_user(db.session()) - user_name = user.name if user else "Anonymous" - - return [ - ToolApiEntity( - author=user_name or "Anonymous", - name=tool.name, - # Prefer the server-provided human-readable title and fall back - # to the programmatic name when the server omitted the title - # or returned ``title: null`` (e.g. Exa's MCP server). - label=I18nObject(en_US=tool.title or tool.name, zh_Hans=tool.title or tool.name), - description=I18nObject(en_US=tool.description or "", zh_Hans=tool.description or ""), - parameters=ToolTransformService.convert_mcp_schema_to_parameter(tool.inputSchema), - labels=[], - output_schema=tool.outputSchema or {}, - ) - for tool in tools - ] - - @classmethod - def api_provider_to_user_provider( - cls, - provider_controller: ApiToolProviderController, - db_provider: ApiToolProvider, - decrypt_credentials: bool = True, - labels: list[str] | None = None, - *, - session: Session, - ) -> ToolProviderApiEntity: - """ - convert provider controller to user provider - """ - username = "Anonymous" - user = db_provider.user(session=session) - if user is None: - raise ValueError(f"user is None for api provider {db_provider.id}") - try: - if not user: - raise ValueError("user not found") - - username = user.name - except Exception: - logger.exception("failed to get user name for api provider %s", db_provider.id) - # add provider into providers - credentials = db_provider.credentials - result = ToolProviderApiEntity( - id=db_provider.id, - author=username, - name=db_provider.name, - description=I18nObject( - en_US=db_provider.description, - zh_Hans=db_provider.description, - ), - icon=db_provider.icon, - label=I18nObject( - en_US=db_provider.name, - zh_Hans=db_provider.name, - ), - type=ToolProviderType.API, - plugin_id=None, - plugin_unique_identifier=None, - masked_credentials={}, - is_team_authorization=True, - tools=[], - labels=labels or [], - ) - - if decrypt_credentials: - # init tool configuration - encrypter, _ = create_tool_provider_encrypter( - tenant_id=db_provider.tenant_id, - controller=provider_controller, - ) - - # decrypt the credentials and mask the credentials - decrypted_credentials = encrypter.decrypt(data=credentials) - masked_credentials = encrypter.mask_plugin_credentials(data=decrypted_credentials) - - result.masked_credentials = masked_credentials - - return result - - @staticmethod - def convert_tool_entity_to_api_entity( - tool: ApiToolBundle | WorkflowTool | Tool, - tenant_id: str, - labels: list[str] | None = None, - ) -> ToolApiEntity: - """ - convert tool to user tool - """ - if isinstance(tool, Tool): - # fork tool runtime - tool = tool.fork_tool_runtime( - runtime=ToolRuntime( - credentials={}, - tenant_id=tenant_id, - ) - ) - - # get tool parameters - base_parameters = tool.entity.parameters or [] - # get tool runtime parameters - runtime_parameters = tool.get_runtime_parameters() - - # merge parameters using a functional approach to avoid type issues - merged_parameters: list[ToolParameter] = [] - - # create a mapping of runtime parameters for quick lookup - runtime_param_map = {(rp.name, rp.form): rp for rp in runtime_parameters} - - # process base parameters, replacing with runtime versions if they exist - for base_param in base_parameters: - key = (base_param.name, base_param.form) - if key in runtime_param_map: - merged_parameters.append(runtime_param_map[key]) - else: - merged_parameters.append(base_param) - - # add any runtime parameters that weren't in base parameters - for runtime_parameter in runtime_parameters: - if runtime_parameter.form == ToolParameter.ToolParameterForm.FORM: - # check if this parameter is already in merged_parameters - already_exists = any( - p.name == runtime_parameter.name and p.form == runtime_parameter.form for p in merged_parameters - ) - if not already_exists: - merged_parameters.append(runtime_parameter) - - return ToolApiEntity( - author=tool.entity.identity.author, - name=tool.entity.identity.name, - label=tool.entity.identity.label, - description=tool.entity.description.human if tool.entity.description else I18nObject(en_US=""), - output_schema=tool.entity.output_schema, - parameters=merged_parameters, - labels=labels or [], - ) - else: - assert tool.operation_id - return ToolApiEntity( - author=tool.author, - name=tool.operation_id or "", - label=I18nObject(en_US=tool.operation_id, zh_Hans=tool.operation_id), - description=I18nObject(en_US=tool.summary or "", zh_Hans=tool.summary or ""), - output_schema=tool.output_schema, - parameters=tool.parameters, - labels=labels or [], - ) - - @staticmethod - def convert_builtin_provider_to_credential_entity( - provider: BuiltinToolProvider, credentials: dict[str, Any] - ) -> ToolProviderCredentialApiEntity: - return ToolProviderCredentialApiEntity( - id=provider.id, - name=provider.name, - provider=provider.provider, - credential_type=provider.credential_type, - is_default=provider.is_default, - credentials=credentials, - ) - - @staticmethod - def convert_mcp_schema_to_parameter(schema: dict[str, Any]) -> list["ToolParameter"]: - """ - Convert MCP JSON schema to tool parameters - - :param schema: JSON schema dictionary - :return: list of ToolParameter instances - """ - - def resolve_property_type(prop: dict[str, Any], depth: int = 0) -> str: - """ - Resolve a JSON schema property type while guarding against cyclic or deeply nested unions. - """ - if depth >= ToolTransformService._MCP_SCHEMA_TYPE_RESOLUTION_MAX_DEPTH: - return "string" - prop_type = prop.get("type") - match prop_type: - case list(): - non_null_types = [type_name for type_name in prop_type if type_name != "null"] - if non_null_types: - return non_null_types[0] - if prop_type: - return "string" - case str(): - if prop_type == "null": - return "string" - return prop_type - - for union_key in ("anyOf", "oneOf"): - union_schemas = prop.get(union_key) - if not isinstance(union_schemas, list): - continue - - for union_schema in union_schemas: - if not isinstance(union_schema, dict): - continue - union_type = resolve_property_type(union_schema, depth + 1) - if union_type != "null": - return union_type - - all_of_schemas = prop.get("allOf") - if isinstance(all_of_schemas, list): - for all_of_schema in all_of_schemas: - if not isinstance(all_of_schema, dict): - continue - all_of_type = resolve_property_type(all_of_schema, depth + 1) - if all_of_type != "null": - return all_of_type - return "string" - - def create_parameter( - name: str, description: str, param_type: str, required: bool, input_schema: dict[str, Any] | None = None - ) -> ToolParameter: - """Create a ToolParameter instance with given attributes""" - input_schema_dict: dict[str, Any] = {"input_schema": input_schema} if input_schema else {} - return ToolParameter( - name=name, - llm_description=description, - label=I18nObject(en_US=name), - form=ToolParameter.ToolParameterForm.LLM, - required=required, - type=ToolParameter.ToolParameterType(param_type), - human_description=I18nObject(en_US=description), - **input_schema_dict, - ) - - def process_properties( - props: dict[str, dict[str, Any]], required: list[str], prefix: str = "" - ) -> list[ToolParameter]: - """Process properties recursively""" - TYPE_MAPPING = {"integer": "number", "float": "number"} - COMPLEX_TYPES = ["array", "object"] - - parameters = [] - for name, prop in props.items(): - current_description = prop.get("description", "") - prop_type = resolve_property_type(prop) - if prop_type in TYPE_MAPPING: - prop_type = TYPE_MAPPING[prop_type] - input_schema = prop if prop_type in COMPLEX_TYPES else None - parameters.append( - create_parameter(name, current_description, prop_type, name in required, input_schema) - ) - - return parameters - - if schema.get("type") == "object" and "properties" in schema: - return process_properties(schema["properties"], schema.get("required", [])) - return [] diff --git a/api/services/tools/legacy_workflow_tools_manage_service.py b/api/services/tools/legacy_workflow_tools_manage_service.py deleted file mode 100644 index 46cdaffbc9018b..00000000000000 --- a/api/services/tools/legacy_workflow_tools_manage_service.py +++ /dev/null @@ -1,419 +0,0 @@ -import json -import logging -from datetime import datetime -from typing import Any - -from sqlalchemy import delete, or_, select -from sqlalchemy.orm import sessionmaker - -from core.tools.__base.tool_provider import ToolProviderController -from core.tools.entities.api_entities import ToolApiEntity, ToolProviderApiEntity -from core.tools.entities.tool_entities import WorkflowToolParameterConfiguration, emoji_icon_adapter -from core.tools.tool_label_manager import ToolLabelManager -from core.tools.utils.workflow_configuration_sync import WorkflowToolConfigurationUtils -from core.tools.workflow_as_tool.provider import WorkflowToolProviderController -from core.tools.workflow_as_tool.tool import WorkflowTool -from extensions.ext_database import db -from graphon.model_runtime.utils.encoders import jsonable_encoder -from models.model import App -from models.tools import WorkflowToolProvider -from models.workflow import Workflow -from services.tools.legacy_tools_transform_service import ToolTransformService - -logger = logging.getLogger(__name__) - - -class WorkflowToolManageService: - """ - Service class for managing workflow tools. - """ - - @staticmethod - def create_workflow_tool( - *, - user_id: str, - tenant_id: str, - workflow_app_id: str, - name: str, - label: str, - icon: dict[str, Any], - description: str, - parameters: list[WorkflowToolParameterConfiguration], - privacy_policy: str = "", - labels: list[str] | None = None, - import_id: str = "", - ): - # check if the name is unique - existing_workflow_tool_provider: WorkflowToolProvider | None = None - with sessionmaker(db.engine, expire_on_commit=False).begin() as _session: - # query if the name or app_id exists - existing_workflow_tool_provider = _session.scalar( - select(WorkflowToolProvider) - .where( - WorkflowToolProvider.tenant_id == tenant_id, - # name or app_id - or_(WorkflowToolProvider.name == name, WorkflowToolProvider.app_id == workflow_app_id), - ) - .limit(1) - ) - - # if the name or app_id exists raise error - if existing_workflow_tool_provider is not None: - raise ValueError(f"Tool with name {name} or app_id {workflow_app_id} already exists") - - # query the app and its published workflow in the same session - app: App | None = None - workflow: Workflow | None = None - with sessionmaker(db.engine, expire_on_commit=False).begin() as _session: - app = _session.scalar(select(App).where(App.id == workflow_app_id, App.tenant_id == tenant_id).limit(1)) - if app is not None: - workflow = app.workflow_with_session(session=_session) - - # if not found raise error - if app is None: - raise ValueError(f"App {workflow_app_id} not found") - - # if not found raise error - if workflow is None: - raise ValueError(f"Workflow not found for app {workflow_app_id}") - - # check if workflow configuration is synced - WorkflowToolConfigurationUtils.ensure_no_human_input_nodes(workflow.graph_dict) - - # create workflow tool provider - workflow_tool_provider = WorkflowToolProvider( - tenant_id=tenant_id, - user_id=user_id, - app_id=workflow_app_id, - name=name, - label=label, - icon=json.dumps(icon), - description=description, - parameter_configuration=json.dumps([p.model_dump() for p in parameters]), - privacy_policy=privacy_policy, - version=workflow.version, - ) - if import_id: - workflow_tool_provider.id = import_id - try: - WorkflowToolProviderController.from_db(workflow_tool_provider) - except Exception as e: - logger.warning(e, exc_info=True) - raise ValueError(str(e)) - - with sessionmaker(db.engine, expire_on_commit=False).begin() as _session: - _session.add(workflow_tool_provider) - - # keep the session open to make orm instances in the same session - if labels is not None: - ToolLabelManager.update_tool_labels( - ToolTransformService.workflow_provider_to_controller(workflow_tool_provider), labels - ) - - return {"result": "success"} - - @classmethod - def update_workflow_tool( - cls, - user_id: str, - tenant_id: str, - workflow_tool_id: str, - name: str, - label: str, - icon: dict[str, Any], - description: str, - parameters: list[WorkflowToolParameterConfiguration], - privacy_policy: str = "", - labels: list[str] | None = None, - ): - """ - Update a workflow tool. - - :param user_id: the user id - :param tenant_id: the tenant id - :param workflow_tool_id: workflow tool id - :param name: name - :param label: label - :param icon: icon - :param description: description - :param parameters: parameters - :param privacy_policy: privacy policy - :param labels: labels - :return: the updated tool - """ - - existing_workflow_tool_provider: WorkflowToolProvider | None = None - with sessionmaker(db.engine, expire_on_commit=False).begin() as _session: - # query if the name exists for other tools - existing_workflow_tool_provider = _session.scalar( - select(WorkflowToolProvider) - .where( - WorkflowToolProvider.tenant_id == tenant_id, - WorkflowToolProvider.name == name, - WorkflowToolProvider.id != workflow_tool_id, - ) - .limit(1) - ) - - # if the name exists raise error - if existing_workflow_tool_provider is not None: - raise ValueError(f"Tool with name {name} already exists") - - # query the workflow tool provider - workflow_tool_provider: WorkflowToolProvider | None = None - with sessionmaker(db.engine, expire_on_commit=False).begin() as _session: - workflow_tool_provider = _session.scalar( - select(WorkflowToolProvider) - .where(WorkflowToolProvider.tenant_id == tenant_id, WorkflowToolProvider.id == workflow_tool_id) - .limit(1) - ) - - # if not found raise error - if workflow_tool_provider is None: - raise ValueError(f"Tool {workflow_tool_id} not found") - - # query the app and its published workflow in the same session - app: App | None = None - workflow: Workflow | None = None - with sessionmaker(db.engine, expire_on_commit=False).begin() as _session: - app = _session.scalar( - select(App).where(App.id == workflow_tool_provider.app_id, App.tenant_id == tenant_id).limit(1) - ) - if app is not None: - workflow = app.workflow_with_session(session=_session) - - # if not found raise error - if app is None: - raise ValueError(f"App {workflow_tool_provider.app_id} not found") - - # if not found raise error - if workflow is None: - raise ValueError(f"Workflow not found for app {workflow_tool_provider.app_id}") - - # check if workflow configuration is synced - WorkflowToolConfigurationUtils.ensure_no_human_input_nodes(workflow.graph_dict) - - with sessionmaker(db.engine).begin() as _session: - _session.add(workflow_tool_provider) - - # update workflow tool provider - workflow_tool_provider.name = name - workflow_tool_provider.label = label - workflow_tool_provider.icon = json.dumps(icon) - workflow_tool_provider.description = description - workflow_tool_provider.parameter_configuration = json.dumps([p.model_dump() for p in parameters]) - workflow_tool_provider.privacy_policy = privacy_policy - workflow_tool_provider.version = workflow.version - workflow_tool_provider.updated_at = datetime.now() - - try: - WorkflowToolProviderController.from_db(workflow_tool_provider) - except Exception as e: - raise ValueError(str(e)) - - if labels is not None: - ToolLabelManager.update_tool_labels( - ToolTransformService.workflow_provider_to_controller(workflow_tool_provider), - labels, - session=_session, - ) - - return {"result": "success"} - - @classmethod - def list_tenant_workflow_tools(cls, user_id: str, tenant_id: str) -> list[ToolProviderApiEntity]: - """ - List workflow tools. - - :param user_id: the user id - :param tenant_id: the tenant id - :return: the list of tools - """ - - providers: list[WorkflowToolProvider] = [] - with sessionmaker(db.engine, expire_on_commit=False).begin() as _session: - providers = list( - _session.scalars(select(WorkflowToolProvider).where(WorkflowToolProvider.tenant_id == tenant_id)).all() - ) - - # Create a mapping from provider_id to app_id - provider_id_to_app_id = {provider.id: provider.app_id for provider in providers} - - tools: list[WorkflowToolProviderController] = [] - for provider in providers: - try: - tools.append(ToolTransformService.workflow_provider_to_controller(provider)) - except Exception: - # skip deleted tools - logger.exception("Failed to load workflow tool provider %s", provider.id) - - labels = ToolLabelManager.get_tools_labels([tool for tool in tools if isinstance(tool, ToolProviderController)]) - - result: list[ToolProviderApiEntity] = [] - - for tool in tools: - workflow_app_id = provider_id_to_app_id.get(tool.provider_id) - user_tool_provider = ToolTransformService.workflow_provider_to_user_provider( - provider_controller=tool, - labels=labels.get(tool.provider_id, []), - workflow_app_id=workflow_app_id, - ) - ToolTransformService.repack_provider(tenant_id=tenant_id, provider=user_tool_provider) - user_tool_provider.tools = [ - ToolTransformService.convert_tool_entity_to_api_entity( - tool=tool.get_tools(tenant_id)[0], - labels=labels.get(tool.provider_id, []), - tenant_id=tenant_id, - ) - ] - result.append(user_tool_provider) - - return result - - @classmethod - def delete_workflow_tool(cls, user_id: str, tenant_id: str, workflow_tool_id: str): - """ - Delete a workflow tool. - - :param user_id: the user id - :param tenant_id: the tenant id - :param workflow_tool_id: the workflow tool id - """ - - with sessionmaker(db.engine).begin() as _session: - _ = _session.execute( - delete(WorkflowToolProvider).where( - WorkflowToolProvider.tenant_id == tenant_id, WorkflowToolProvider.id == workflow_tool_id - ) - ) - - return {"result": "success"} - - @classmethod - def get_workflow_tool_by_tool_id(cls, user_id: str, tenant_id: str, workflow_tool_id: str): - """ - Get a workflow tool. - - :param user_id: the user id - :param tenant_id: the tenant id - :param workflow_tool_id: the workflow tool id - :return: the tool - """ - - tool_provider: WorkflowToolProvider | None = None - with sessionmaker(db.engine, expire_on_commit=False).begin() as _session: - tool_provider = _session.scalar( - select(WorkflowToolProvider) - .where(WorkflowToolProvider.tenant_id == tenant_id, WorkflowToolProvider.id == workflow_tool_id) - .limit(1) - ) - - return cls._get_workflow_tool(tenant_id, tool_provider) - - @classmethod - def get_workflow_tool_by_app_id(cls, user_id: str, tenant_id: str, workflow_app_id: str): - """ - Get a workflow tool. - - :param user_id: the user id - :param tenant_id: the tenant id - :param workflow_app_id: the workflow app id - :return: the tool - """ - - with sessionmaker(db.engine, expire_on_commit=False).begin() as _session: - tool_provider: WorkflowToolProvider | None = _session.scalar( - select(WorkflowToolProvider) - .where(WorkflowToolProvider.tenant_id == tenant_id, WorkflowToolProvider.app_id == workflow_app_id) - .limit(1) - ) - - return cls._get_workflow_tool(tenant_id, tool_provider) - - @classmethod - def _get_workflow_tool(cls, tenant_id: str, db_tool: WorkflowToolProvider | None): - """ - Get a workflow tool. - - :db_tool: the database tool - :return: the tool - """ - if db_tool is None: - raise ValueError("Tool not found") - - workflow_app: App | None = None - workflow: Workflow | None = None - with sessionmaker(db.engine, expire_on_commit=False).begin() as _session: - workflow_app = _session.scalar( - select(App).where(App.id == db_tool.app_id, App.tenant_id == db_tool.tenant_id).limit(1) - ) - if workflow_app is not None: - workflow = workflow_app.workflow_with_session(session=_session) - - if workflow_app is None: - raise ValueError(f"App {db_tool.app_id} not found") - - if not workflow: - raise ValueError("Workflow not found") - - tool = ToolTransformService.workflow_provider_to_controller(db_tool) - workflow_tools: list[WorkflowTool] = tool.get_tools(tenant_id) - if len(workflow_tools) == 0: - raise ValueError(f"Tool {db_tool.id} not found") - - tool_entity = workflow_tools[0].entity - # get output schema from workflow tool entity - output_schema = tool_entity.output_schema - - return { - "name": db_tool.name, - "label": db_tool.label, - "workflow_tool_id": db_tool.id, - "workflow_app_id": db_tool.app_id, - "icon": emoji_icon_adapter.validate_json(db_tool.icon), - "description": db_tool.description, - "parameters": jsonable_encoder(db_tool.parameter_configurations), - "output_schema": output_schema, - "tool": ToolTransformService.convert_tool_entity_to_api_entity( - tool=tool.get_tools(db_tool.tenant_id)[0], - labels=ToolLabelManager.get_tool_labels(tool), - tenant_id=tenant_id, - ), - "synced": workflow.version == db_tool.version, - "privacy_policy": db_tool.privacy_policy, - } - - @classmethod - def list_single_workflow_tools(cls, user_id: str, tenant_id: str, workflow_tool_id: str) -> list[ToolApiEntity]: - """ - List workflow tool provider tools. - - :param user_id: the user id - :param tenant_id: the tenant id - :param workflow_tool_id: the workflow tool id - :return: the list of tools - """ - - provider: WorkflowToolProvider | None = None - with sessionmaker(db.engine, expire_on_commit=False).begin() as _session: - provider = _session.scalar( - select(WorkflowToolProvider) - .where(WorkflowToolProvider.tenant_id == tenant_id, WorkflowToolProvider.id == workflow_tool_id) - .limit(1) - ) - - if provider is None: - raise ValueError(f"Tool {workflow_tool_id} not found") - - tool = ToolTransformService.workflow_provider_to_controller(provider) - workflow_tools: list[WorkflowTool] = tool.get_tools(tenant_id) - if len(workflow_tools) == 0: - raise ValueError(f"Tool {workflow_tool_id} not found") - - return [ - ToolTransformService.convert_tool_entity_to_api_entity( - tool=tool.get_tools(provider.tenant_id)[0], - labels=ToolLabelManager.get_tool_labels(tool), - tenant_id=tenant_id, - ) - ] diff --git a/api/services/workflow_app_log_query_service.py b/api/services/workflow_app_log_query_service.py deleted file mode 100644 index 4eb1df365d8f65..00000000000000 --- a/api/services/workflow_app_log_query_service.py +++ /dev/null @@ -1,174 +0,0 @@ -import json -from collections.abc import Mapping -from dataclasses import dataclass, replace -from datetime import datetime -from typing import Any, Protocol - -from core.plugin.plugin_service import PluginService -from graphon.enums import WorkflowExecutionStatus -from models.enums import AppTriggerType -from services.workflow.entities import TriggerMetadata - - -@dataclass(frozen=True, slots=True) -class WorkflowAppLogAccount: - id: str - name: str - email: str - - -@dataclass(frozen=True, slots=True) -class WorkflowAppLogEndUser: - id: str - type: str - is_anonymous: bool - session_id: str | None - - -@dataclass(frozen=True, slots=True) -class WorkflowAppLogRunSummary: - id: str - version: str | None - status: str - triggered_from: str - error: str | None - elapsed_time: float | None - total_tokens: int | None - total_steps: int | None - created_at: datetime | None - finished_at: datetime | None - exceptions_count: int | None - - -@dataclass(frozen=True, slots=True) -class WorkflowAppLogItem: - id: str - workflow_run: WorkflowAppLogRunSummary | None - details: dict[str, Any] | None - created_from: str - created_by_role: str - created_by_account: WorkflowAppLogAccount | None - created_by_end_user: WorkflowAppLogEndUser | None - created_at: datetime - - -@dataclass(frozen=True, slots=True) -class WorkflowAppLogPage: - page: int - limit: int - total: int - has_more: bool - data: tuple[WorkflowAppLogItem, ...] - - -class WorkflowAppLogQuery(Protocol): - def get_paginated( - self, - *, - tenant_id: str, - app_id: str, - keyword: str | None = None, - status: WorkflowExecutionStatus | None = None, - created_at_before: datetime | None = None, - created_at_after: datetime | None = None, - page: int = 1, - limit: int = 20, - detail: bool = False, - created_by_end_user_session_id: str | None = None, - created_by_account: str | None = None, - ) -> WorkflowAppLogPage: ... - - -class WorkflowAppLogQueryService: - def __init__( - self, - *, - logs: WorkflowAppLogQuery, - ) -> None: - self._logs = logs - - def list_logs( - self, - *, - tenant_id: str, - app_id: str, - keyword: str | None = None, - status: WorkflowExecutionStatus | None = None, - created_at_before: datetime | None = None, - created_at_after: datetime | None = None, - page: int = 1, - limit: int = 20, - detail: bool = False, - created_by_end_user_session_id: str | None = None, - created_by_account: str | None = None, - ) -> WorkflowAppLogPage: - result = self._logs.get_paginated( - tenant_id=tenant_id, - app_id=app_id, - keyword=keyword, - status=status, - created_at_before=created_at_before, - created_at_after=created_at_after, - page=page, - limit=limit, - detail=detail, - created_by_end_user_session_id=created_by_end_user_session_id, - created_by_account=created_by_account, - ) - - return WorkflowAppLogPage( - page=result.page, - limit=result.limit, - total=result.total, - has_more=result.has_more, - data=tuple(self._resolve_details(tenant_id=tenant_id, item=item) for item in result.data), - ) - - def _resolve_details( - self, - *, - tenant_id: str, - item: WorkflowAppLogItem, - ) -> WorkflowAppLogItem: - if item.details is None: - return item - - return replace( - item, - details={ - "trigger_metadata": self._resolve_trigger_metadata( - tenant_id, - item.details.get("trigger_metadata"), - ) - }, - ) - - @staticmethod - def _resolve_trigger_metadata( - tenant_id: str, - value: str | Mapping[str, Any] | None, - ) -> dict[str, Any]: - metadata = WorkflowAppLogQueryService._safe_json_loads(value) - if not metadata: - return {} - - trigger_metadata = TriggerMetadata.model_validate(metadata) - if trigger_metadata.type == AppTriggerType.TRIGGER_PLUGIN: - icon = metadata.get("icon_filename") - icon_dark = metadata.get("icon_dark_filename") - metadata["icon"] = PluginService.get_plugin_icon_url(tenant_id=tenant_id, filename=icon) if icon else None - metadata["icon_dark"] = ( - PluginService.get_plugin_icon_url(tenant_id=tenant_id, filename=icon_dark) if icon_dark else None - ) - return metadata - - @staticmethod - def _safe_json_loads(value: Any) -> Any: - if not value: - return None - if isinstance(value, str): - try: - return json.loads(value) - except Exception: - return None - return value diff --git a/api/services/workflow_draft_variable_service.py b/api/services/workflow_draft_variable_service.py deleted file mode 100644 index e0a0c18d5c6fbe..00000000000000 --- a/api/services/workflow_draft_variable_service.py +++ /dev/null @@ -1,1202 +0,0 @@ -import dataclasses -import json -import logging -from collections.abc import Mapping, Sequence, Set -from concurrent.futures import ThreadPoolExecutor -from datetime import datetime -from enum import StrEnum -from typing import Any, ClassVar, NotRequired, TypedDict, cast, override - -from sqlalchemy import Engine, delete, orm, select -from sqlalchemy.dialects.mysql import insert as mysql_insert -from sqlalchemy.dialects.postgresql import insert as pg_insert -from sqlalchemy.orm import Session, sessionmaker -from sqlalchemy.sql.expression import and_, or_ - -from configs import dify_config -from core.app.entities.app_invoke_entities import InvokeFrom -from core.app.file_access import DatabaseFileAccessController -from core.trigger.constants import is_trigger_node_type -from core.workflow.system_variables import SystemVariableKey -from core.workflow.variable_prefixes import ( - CONVERSATION_VARIABLE_NODE_ID, - ENVIRONMENT_VARIABLE_NODE_ID, - RAG_PIPELINE_VARIABLE_NODE_ID, - SYSTEM_VARIABLE_NODE_ID, -) -from extensions.ext_storage import storage -from factories.file_factory import StorageKeyLoader -from factories.variable_factory import build_segment, segment_to_variable -from graphon.enums import NodeType -from graphon.file import File -from graphon.nodes import BuiltinNodeTypes -from graphon.nodes.variable_assigner.common.helpers import get_updated_variables -from graphon.variable_loader import VariableLoader -from graphon.variables import Segment, StringSegment, VariableBase -from graphon.variables.consts import SELECTORS_LENGTH -from graphon.variables.segments import ( - ArrayFileSegment, - FileSegment, -) -from graphon.variables.types import SegmentType -from graphon.variables.utils import dumps_with_segments -from libs.datetime_utils import naive_utc_now -from libs.uuid_utils import uuidv7 -from models import Account, App, Conversation -from models.enums import ConversationFromSource, DraftVariableType -from models.utils.file_input_compat import build_file_from_stored_mapping -from models.workflow import Workflow, WorkflowDraftVariable, WorkflowDraftVariableFile, is_system_variable_editable -from repositories.factory import DifyAPIRepositoryFactory -from services.file_service import FileService -from services.variable_truncator import VariableTruncator - -logger = logging.getLogger(__name__) - - -@dataclasses.dataclass(frozen=True) -class WorkflowDraftVariableList: - variables: list[WorkflowDraftVariable] - total: int | None = None - - -@dataclasses.dataclass(frozen=True) -class DraftVarFileDeletion: - draft_var_id: str - draft_var_file_id: str - - -class WorkflowDraftVariableError(Exception): - pass - - -class VariableResetError(WorkflowDraftVariableError): - pass - - -class UpdateNotSupportedError(WorkflowDraftVariableError): - pass - - -class DraftVarLoader(VariableLoader): - # This implements the VariableLoader interface for loading draft variables. - # - # ref: graphon.variable_loader.VariableLoader - - # Database engine used for loading variables. - _engine: Engine - # Application ID for which variables are being loaded. - _app_id: str - _user_id: str - _tenant_id: str - _fallback_variables: Sequence[VariableBase] - - def __init__( - self, - engine: Engine, - app_id: str, - tenant_id: str, - user_id: str, - fallback_variables: Sequence[VariableBase] | None = None, - ): - self._engine = engine - self._app_id = app_id - self._user_id = user_id - self._tenant_id = tenant_id - self._fallback_variables = fallback_variables or [] - - def _selector_to_tuple(self, selector: Sequence[str]) -> tuple[str, str]: - return (selector[0], selector[1]) - - @override - def load_variables(self, selectors: list[list[str]]) -> list[VariableBase]: - if not selectors: - return [] - - # Map each selector (as a tuple via `_selector_to_tuple`) to its corresponding variable instance. - variable_by_selector: dict[tuple[str, str], VariableBase] = {} - - with Session(bind=self._engine, expire_on_commit=False) as session: - srv = WorkflowDraftVariableService(session) - draft_vars = srv.get_draft_variables_by_selectors(self._app_id, selectors, user_id=self._user_id) - - # Important: - files: list[File] = [] - # FileSegment and ArrayFileSegment are not subject to offloading, so their values - # can be safely accessed before any offloading logic is applied. - for draft_var in draft_vars: - value = draft_var.get_value() - match value: - case FileSegment(): - files.append(value.value) - case ArrayFileSegment(): - files.extend(value.value) - with Session(bind=self._engine) as session: - storage_key_loader = StorageKeyLoader( - session, - tenant_id=self._tenant_id, - access_controller=DatabaseFileAccessController(), - ) - storage_key_loader.load_storage_keys(files) - - offloaded_draft_vars = [] - for draft_var in draft_vars: - if draft_var.is_truncated(): - offloaded_draft_vars.append(draft_var) - continue - - segment = draft_var.get_value() - variable = segment_to_variable( - segment=segment, - selector=draft_var.get_selector(), - variable_id=draft_var.id, - name=draft_var.name, - description=draft_var.description, - ) - selector_tuple = self._selector_to_tuple(variable.selector) - variable_by_selector[selector_tuple] = variable - - # Load offloaded variables using multithreading. - # This approach reduces loading time by querying external systems concurrently. - with ThreadPoolExecutor(max_workers=10) as executor: - offloaded_variables = executor.map(self._load_offloaded_variable, offloaded_draft_vars) - for selector, offloaded_variable in offloaded_variables: - variable_by_selector[selector] = offloaded_variable - - return list(variable_by_selector.values()) - - def _load_offloaded_variable(self, draft_var: WorkflowDraftVariable) -> tuple[tuple[str, str], VariableBase]: - # This logic is closely tied to `WorkflowDraftVaribleService._try_offload_large_variable` - # and must remain synchronized with it. - # Ideally, these should be co-located for better maintainability. - # However, due to the current code structure, this is not straightforward. - - variable_file = draft_var.variable_file - assert variable_file is not None - upload_file = variable_file.upload_file - assert upload_file is not None - content = storage.load(upload_file.key) - if variable_file.value_type == SegmentType.STRING: - # The inferenced type is StringSegment, which is not correct inside this function. - segment: Segment = StringSegment(value=content.decode()) - - variable = segment_to_variable( - segment=segment, - selector=draft_var.get_selector(), - variable_id=draft_var.id, - name=draft_var.name, - description=draft_var.description, - ) - return (draft_var.node_id, draft_var.name), variable - - deserialized = json.loads(content) - segment = draft_var.build_segment_from_serialized_value(variable_file.value_type, deserialized) - variable = segment_to_variable( - segment=segment, - selector=draft_var.get_selector(), - variable_id=draft_var.id, - name=draft_var.name, - description=draft_var.description, - ) - # No special handling needed for ArrayFileSegment, as we do not offload ArrayFileSegment - return (draft_var.node_id, draft_var.name), variable - - -class WorkflowDraftVariableService: - _session: Session - - def __init__(self, session: Session): - """ - Initialize the WorkflowDraftVariableService with a SQLAlchemy session. - - Args: - session (Session): The SQLAlchemy session used to execute database queries. - The provided session must be bound to an `Engine` object, not a specific `Connection`. - - Raises: - AssertionError: If the provided session is not bound to an `Engine` object. - """ - self._session = session - engine = session.get_bind() - # Ensure the session is bound to a engine. - assert isinstance(engine, Engine) - session_maker = sessionmaker(bind=engine, expire_on_commit=False) - self._api_node_execution_repo = DifyAPIRepositoryFactory.create_api_workflow_node_execution_repository( - session_maker - ) - - def get_variable(self, variable_id: str) -> WorkflowDraftVariable | None: - return self._session.scalar( - select(WorkflowDraftVariable) - .options(orm.selectinload(WorkflowDraftVariable.variable_file)) - .where(WorkflowDraftVariable.id == variable_id) - ) - - def get_draft_variables_by_selectors( - self, - app_id: str, - selectors: Sequence[list[str]], - user_id: str, - ) -> list[WorkflowDraftVariable]: - """ - Retrieve WorkflowDraftVariable instances based on app_id and selectors. - - The returned WorkflowDraftVariable objects are guaranteed to have their - associated variable_file and variable_file.upload_file relationships preloaded. - """ - ors = [] - for selector in selectors: - assert len(selector) >= SELECTORS_LENGTH, f"Invalid selector to get: {selector}" - node_id, name = selector[:2] - ors.append(and_(WorkflowDraftVariable.node_id == node_id, WorkflowDraftVariable.name == name)) - - # NOTE(QuantumGhost): Although the number of `or` expressions may be large, as long as - # each expression includes conditions on both `node_id` and `name` (which are covered by the unique index), - # PostgreSQL can efficiently retrieve the results using a bitmap index scan. - # - # Alternatively, a `SELECT` statement could be constructed for each selector and - # combined using `UNION` to fetch all rows. - # Benchmarking indicates that both approaches yield comparable performance. - return list( - self._session.scalars( - select(WorkflowDraftVariable) - .options( - orm.selectinload(WorkflowDraftVariable.variable_file).selectinload( - WorkflowDraftVariableFile.upload_file - ) - ) - .where( - WorkflowDraftVariable.app_id == app_id, - WorkflowDraftVariable.user_id == user_id, - or_(*ors), - ) - ) - ) - - def list_variables_without_values( - self, - app_id: str, - page: int, - limit: int, - user_id: str, - *, - exclude_node_ids: Set[str] | None = None, - ) -> WorkflowDraftVariableList: - criteria = [ - WorkflowDraftVariable.app_id == app_id, - WorkflowDraftVariable.user_id == user_id, - ] - if exclude_node_ids: - criteria.append(WorkflowDraftVariable.node_id.notin_(list(exclude_node_ids))) - total = None - base_stmt = select(WorkflowDraftVariable).where(*criteria) - if page == 1: - from sqlalchemy import func as sa_func - - total = self._session.scalar(select(sa_func.count()).select_from(base_stmt.subquery())) - variables = list( - self._session.scalars( - # Do not load the `value` field - base_stmt.options( - orm.defer(WorkflowDraftVariable.value, raiseload=True), - ) - .order_by(WorkflowDraftVariable.created_at.desc()) - .limit(limit) - .offset((page - 1) * limit) - ) - ) - - return WorkflowDraftVariableList(variables=variables, total=total) - - def _list_node_variables(self, app_id: str, node_id: str, user_id: str) -> WorkflowDraftVariableList: - criteria = [ - WorkflowDraftVariable.app_id == app_id, - WorkflowDraftVariable.node_id == node_id, - WorkflowDraftVariable.user_id == user_id, - ] - variables = list( - self._session.scalars( - select(WorkflowDraftVariable) - .options(orm.selectinload(WorkflowDraftVariable.variable_file)) - .where(*criteria) - .order_by(WorkflowDraftVariable.created_at.desc()) - ) - ) - return WorkflowDraftVariableList(variables=variables) - - def list_node_variables(self, app_id: str, node_id: str, user_id: str) -> WorkflowDraftVariableList: - return self._list_node_variables(app_id, node_id, user_id=user_id) - - def list_conversation_variables(self, app_id: str, user_id: str) -> WorkflowDraftVariableList: - return self._list_node_variables(app_id, CONVERSATION_VARIABLE_NODE_ID, user_id=user_id) - - def list_system_variables(self, app_id: str, user_id: str) -> WorkflowDraftVariableList: - return self._list_node_variables(app_id, SYSTEM_VARIABLE_NODE_ID, user_id=user_id) - - def get_conversation_variable(self, app_id: str, name: str, user_id: str) -> WorkflowDraftVariable | None: - return self._get_variable(app_id=app_id, node_id=CONVERSATION_VARIABLE_NODE_ID, name=name, user_id=user_id) - - def get_system_variable(self, app_id: str, name: str, user_id: str) -> WorkflowDraftVariable | None: - return self._get_variable(app_id=app_id, node_id=SYSTEM_VARIABLE_NODE_ID, name=name, user_id=user_id) - - def get_node_variable(self, app_id: str, node_id: str, name: str, user_id: str) -> WorkflowDraftVariable | None: - return self._get_variable(app_id, node_id, name, user_id=user_id) - - def _get_variable(self, app_id: str, node_id: str, name: str, user_id: str) -> WorkflowDraftVariable | None: - return self._session.scalar( - select(WorkflowDraftVariable) - .options(orm.selectinload(WorkflowDraftVariable.variable_file)) - .where( - WorkflowDraftVariable.app_id == app_id, - WorkflowDraftVariable.node_id == node_id, - WorkflowDraftVariable.name == name, - WorkflowDraftVariable.user_id == user_id, - ) - ) - - def update_variable( - self, - variable: WorkflowDraftVariable, - name: str | None = None, - value: Segment | None = None, - ) -> WorkflowDraftVariable: - if not variable.editable: - raise UpdateNotSupportedError(f"variable not support updating, id={variable.id}") - if name is not None: - variable.set_name(name) - if value is not None: - variable.set_value(value) - variable.last_edited_at = naive_utc_now() - self._session.flush() - return variable - - def _reset_conv_var(self, workflow: Workflow, variable: WorkflowDraftVariable) -> WorkflowDraftVariable | None: - conv_var_by_name = {i.name: i for i in workflow.conversation_variables} - conv_var = conv_var_by_name.get(variable.name) - - if conv_var is None: - self._session.delete(instance=variable) - self._session.flush() - logger.warning( - "Conversation variable not found for draft variable, id=%s, name=%s", variable.id, variable.name - ) - return None - - variable.set_value(conv_var) - variable.last_edited_at = None - self._session.add(variable) - self._session.flush() - return variable - - def _reset_node_var_or_sys_var( - self, workflow: Workflow, variable: WorkflowDraftVariable - ) -> WorkflowDraftVariable | None: - # If a variable does not allow updating, it makes no sense to reset it. - if not variable.editable: - return variable - # No execution record for this variable, delete the variable instead. - if variable.node_execution_id is None: - self._session.delete(instance=variable) - self._session.flush() - logger.warning("draft variable has no node_execution_id, id=%s, name=%s", variable.id, variable.name) - return None - - node_exec = self._api_node_execution_repo.get_execution_by_id(variable.node_execution_id) - if node_exec is None: - logger.warning( - "Node exectution not found for draft variable, id=%s, name=%s, node_execution_id=%s", - variable.id, - variable.name, - variable.node_execution_id, - ) - self._session.delete(instance=variable) - self._session.flush() - return None - - outputs_dict = node_exec.load_full_outputs(self._session, storage) or {} - # a sentinel value used to check the absent of the output variable key. - absent = object() - - if variable.get_variable_type() == DraftVariableType.NODE: - # Get node type for proper value extraction - node_config = workflow.get_node_config_by_id(variable.node_id) - node_type = workflow.get_node_type_from_node_config(node_config) - - # Note: Based on the implementation in `_build_from_variable_assigner_mapping`, - # VariableAssignerNode (both v1 and v2) can only create conversation draft variables. - # For consistency, we should simply return when processing VARIABLE_ASSIGNER nodes. - # - # This implementation must remain synchronized with the `_build_from_variable_assigner_mapping` - # and `save` methods. - if node_type == BuiltinNodeTypes.VARIABLE_ASSIGNER: - return variable - output_value = outputs_dict.get(variable.name, absent) - else: - output_value = outputs_dict.get(f"sys.{variable.name}", absent) - - # We cannot use `is None` to check the existence of an output variable here as - # the value of the output may be `None`. - if output_value is absent: - # If variable not found in execution data, delete the variable - self._session.delete(instance=variable) - self._session.flush() - return None - value_seg = WorkflowDraftVariable.build_segment_with_type(variable.value_type, output_value) - # Extract variable value using unified logic - variable.set_value(value_seg) - variable.last_edited_at = None # Reset to indicate this is a reset operation - self._session.flush() - return variable - - def reset_variable(self, workflow: Workflow, variable: WorkflowDraftVariable) -> WorkflowDraftVariable | None: - variable_type = variable.get_variable_type() - if variable_type == DraftVariableType.SYS and not is_system_variable_editable(variable.name): - raise VariableResetError(f"cannot reset system variable, variable_id={variable.id}") - if variable_type == DraftVariableType.CONVERSATION: - return self._reset_conv_var(workflow, variable) - else: - return self._reset_node_var_or_sys_var(workflow, variable) - - def delete_variable(self, variable: WorkflowDraftVariable): - if not variable.is_truncated(): - self._session.delete(variable) - return - - variable_query = ( - select(WorkflowDraftVariable) - .options( - orm.selectinload(WorkflowDraftVariable.variable_file).selectinload( - WorkflowDraftVariableFile.upload_file - ), - ) - .where(WorkflowDraftVariable.id == variable.id) - ) - variable_reloaded = self._session.execute(variable_query).scalars().first() - if variable_reloaded is None: - logger.warning("Associated WorkflowDraftVariable not found, draft_var_id=%s", variable.id) - self._session.delete(variable) - return - variable_file = variable_reloaded.variable_file - if variable_file is None: - logger.warning( - "Associated WorkflowDraftVariableFile not found, draft_var_id=%s, file_id=%s", - variable_reloaded.id, - variable_reloaded.file_id, - ) - self._session.delete(variable) - return - - upload_file = variable_file.upload_file - if upload_file is None: - logger.warning( - "Associated UploadFile not found, draft_var_id=%s, file_id=%s, upload_file_id=%s", - variable_reloaded.id, - variable_reloaded.file_id, - variable_file.upload_file_id, - ) - self._session.delete(variable) - self._session.delete(variable_file) - return - - storage.delete(upload_file.key) - self._session.delete(upload_file) - self._session.delete(upload_file) - self._session.delete(variable) - - def delete_user_workflow_variables(self, app_id: str, user_id: str): - self._session.execute( - delete(WorkflowDraftVariable) - .where( - WorkflowDraftVariable.app_id == app_id, - WorkflowDraftVariable.user_id == user_id, - ) - .execution_options(synchronize_session=False) - ) - - def delete_app_workflow_variables(self, app_id: str): - self._session.execute( - delete(WorkflowDraftVariable) - .where(WorkflowDraftVariable.app_id == app_id) - .execution_options(synchronize_session=False) - ) - - def delete_workflow_draft_variable_file(self, deletions: list[DraftVarFileDeletion]): - variable_files_query = ( - select(WorkflowDraftVariableFile) - .options(orm.selectinload(WorkflowDraftVariableFile.upload_file)) - .where(WorkflowDraftVariableFile.id.in_([i.draft_var_file_id for i in deletions])) - ) - variable_files = self._session.execute(variable_files_query).scalars().all() - variable_files_by_id = {i.id: i for i in variable_files} - for i in deletions: - variable_file = variable_files_by_id.get(i.draft_var_file_id) - if variable_file is None: - logger.warning( - "Associated WorkflowDraftVariableFile not found, draft_var_id=%s, file_id=%s", - i.draft_var_id, - i.draft_var_file_id, - ) - continue - - upload_file = variable_file.upload_file - if upload_file is None: - logger.warning( - "Associated UploadFile not found, draft_var_id=%s, file_id=%s, upload_file_id=%s", - i.draft_var_id, - i.draft_var_file_id, - variable_file.upload_file_id, - ) - self._session.delete(variable_file) - else: - storage.delete(upload_file.key) - self._session.delete(upload_file) - self._session.delete(variable_file) - - def delete_node_variables(self, app_id: str, node_id: str, user_id: str): - return self._delete_node_variables(app_id, node_id, user_id=user_id) - - def _delete_node_variables(self, app_id: str, node_id: str, user_id: str): - self._session.execute( - delete(WorkflowDraftVariable) - .where( - WorkflowDraftVariable.app_id == app_id, - WorkflowDraftVariable.node_id == node_id, - WorkflowDraftVariable.user_id == user_id, - ) - .execution_options(synchronize_session=False) - ) - - def _get_conversation_id_from_draft_variable(self, app_id: str, user_id: str) -> str | None: - draft_var = self._get_variable( - app_id=app_id, - node_id=SYSTEM_VARIABLE_NODE_ID, - name=str(SystemVariableKey.CONVERSATION_ID), - user_id=user_id, - ) - if draft_var is None: - return None - segment = draft_var.get_value() - if not isinstance(segment, StringSegment): - logger.warning( - "sys.conversation_id variable is not a string: app_id=%s, user_id=%s, id=%s", - app_id, - user_id, - draft_var.id, - ) - return None - return segment.value - - def get_or_create_conversation( - self, - account_id: str, - app: App, - workflow: Workflow, - ) -> str: - """ - get_or_create_conversation creates and returns the ID of a conversation for debugging. - - If a conversation already exists, as determined by the following criteria, its ID is returned: - - The system variable `sys.conversation_id` exists in the draft variable table, and - - A corresponding conversation record is found in the database. - - If no such conversation exists, a new conversation is created and its ID is returned. - """ - conv_id = self._get_conversation_id_from_draft_variable(workflow.app_id, account_id) - - if conv_id is not None: - conversation = self._session.scalar( - select(Conversation).where( - Conversation.id == conv_id, - Conversation.app_id == workflow.app_id, - ) - ) - # Only return the conversation ID if it exists and is valid (has a correspond conversation record in DB). - if conversation is not None: - return conv_id - conversation = Conversation( - app_id=workflow.app_id, - app_model_config_id=app.app_model_config_id, - model_provider=None, - model_id="", - override_model_configs=None, - mode=app.mode, - name="Draft Debugging Conversation", - inputs={}, - introduction="", - system_instruction="", - system_instruction_tokens=0, - status="normal", - invoke_from=InvokeFrom.DEBUGGER, - from_source=ConversationFromSource.CONSOLE, - from_end_user_id=None, - from_account_id=account_id, - ) - - self._session.add(conversation) - self._session.flush() - return conversation.id - - def prefill_conversation_variable_default_values(self, workflow: Workflow, user_id: str): - """""" - draft_conv_vars: list[WorkflowDraftVariable] = [] - for conv_var in workflow.conversation_variables: - draft_var = WorkflowDraftVariable.new_conversation_variable( - app_id=workflow.app_id, - user_id=user_id, - name=conv_var.name, - value=conv_var, - description=conv_var.description, - ) - draft_conv_vars.append(draft_var) - _batch_upsert_draft_variable( - self._session, - draft_conv_vars, - policy=_UpsertPolicy.IGNORE, - ) - - -class _UpsertPolicy(StrEnum): - IGNORE = "ignore" - OVERWRITE = "overwrite" - - -def _batch_upsert_draft_variable( - session: Session, - draft_vars: Sequence[WorkflowDraftVariable], - policy: _UpsertPolicy = _UpsertPolicy.OVERWRITE, -): - if not draft_vars: - return None - # Although we could use SQLAlchemy ORM operations here, we choose not to for several reasons: - # - # 1. The variable saving process involves writing multiple rows to the - # `workflow_draft_variables` table. Batch insertion significantly improves performance. - # 2. Using the ORM would require either: - # - # a. Checking for the existence of each variable before insertion, - # resulting in 2n SQL statements for n variables and potential concurrency issues. - # b. Attempting insertion first, then updating if a unique index violation occurs, - # which still results in n to 2n SQL statements. - # - # Both approaches are inefficient and suboptimal. - # 3. We do not need to retrieve the results of the SQL execution or populate ORM - # model instances with the returned values. - # 4. Batch insertion with `ON CONFLICT DO UPDATE` allows us to insert or update all - # variables in a single SQL statement, avoiding the issues above. - # - # For these reasons, we use the SQLAlchemy query builder and rely on dialect-specific - # insert operations instead of the ORM layer. - - # Use different insert statements based on database type - if dify_config.SQLALCHEMY_DATABASE_URI_SCHEME == "postgresql": - stmt = pg_insert(WorkflowDraftVariable).values([_model_to_insertion_dict(v) for v in draft_vars]) - if policy == _UpsertPolicy.OVERWRITE: - stmt = stmt.on_conflict_do_update( - index_elements=WorkflowDraftVariable.unique_app_id_user_id_node_id_name(), - set_={ - # Refresh creation timestamp to ensure updated variables - # appear first in chronologically sorted result sets. - "created_at": stmt.excluded.created_at, - "updated_at": stmt.excluded.updated_at, - "last_edited_at": stmt.excluded.last_edited_at, - "description": stmt.excluded.description, - "value_type": stmt.excluded.value_type, - "value": stmt.excluded.value, - "visible": stmt.excluded.visible, - "editable": stmt.excluded.editable, - "node_execution_id": stmt.excluded.node_execution_id, - "file_id": stmt.excluded.file_id, - }, - ) - elif policy == _UpsertPolicy.IGNORE: - stmt = stmt.on_conflict_do_nothing( - index_elements=WorkflowDraftVariable.unique_app_id_user_id_node_id_name() - ) - else: - stmt = mysql_insert(WorkflowDraftVariable).values([_model_to_insertion_dict(v) for v in draft_vars]) # type: ignore[assignment] - if policy == _UpsertPolicy.OVERWRITE: - stmt = stmt.on_duplicate_key_update( # type: ignore[attr-defined] - # Refresh creation timestamp to ensure updated variables - # appear first in chronologically sorted result sets. - created_at=stmt.inserted.created_at, # type: ignore[attr-defined] - updated_at=stmt.inserted.updated_at, # type: ignore[attr-defined] - last_edited_at=stmt.inserted.last_edited_at, # type: ignore[attr-defined] - description=stmt.inserted.description, # type: ignore[attr-defined] - value_type=stmt.inserted.value_type, # type: ignore[attr-defined] - value=stmt.inserted.value, # type: ignore[attr-defined] - visible=stmt.inserted.visible, # type: ignore[attr-defined] - editable=stmt.inserted.editable, # type: ignore[attr-defined] - node_execution_id=stmt.inserted.node_execution_id, # type: ignore[attr-defined] - file_id=stmt.inserted.file_id, # type: ignore[attr-defined] - ) - elif policy == _UpsertPolicy.IGNORE: - stmt = stmt.prefix_with("IGNORE") - - if policy not in [_UpsertPolicy.OVERWRITE, _UpsertPolicy.IGNORE]: - raise Exception("Invalid value for update policy.") - session.execute(stmt) - - -class _InsertionDict(TypedDict): - id: str - app_id: str - user_id: str | None - last_edited_at: datetime | None - node_id: str - name: str - selector: str - value_type: SegmentType - value: str - node_execution_id: str | None - file_id: str | None - visible: NotRequired[bool] - editable: NotRequired[bool] - is_default_value: NotRequired[bool] - created_at: NotRequired[datetime] - updated_at: NotRequired[datetime] - description: NotRequired[str] - - -def _model_to_insertion_dict(model: WorkflowDraftVariable) -> _InsertionDict: - d: _InsertionDict = { - "id": model.id, - "app_id": model.app_id, - "user_id": model.user_id, - "last_edited_at": None, - "node_id": model.node_id, - "name": model.name, - "selector": model.selector, - "value_type": model.value_type, - "value": model.value, - "node_execution_id": model.node_execution_id, - "file_id": model.file_id, - } - if model.visible is not None: - d["visible"] = model.visible - if model.editable is not None: - d["editable"] = model.editable - if model.created_at is not None: - d["created_at"] = model.created_at - if model.updated_at is not None: - d["updated_at"] = model.updated_at - if model.description is not None: - d["description"] = model.description - if model.is_default_value is not None: - d["is_default_value"] = model.is_default_value - return d - - -def _build_segment_for_serialized_values(v: Any) -> Segment: - """ - Reconstructs Segment objects from serialized values, with special handling - for FileSegment and ArrayFileSegment types. - - This function should only be used when: - 1. No explicit type information is available - 2. The input value is in serialized form (dict or list) - - It detects potential file objects in the serialized data and properly rebuilds the - appropriate segment type. - """ - return build_segment(WorkflowDraftVariable.rebuild_file_types(v)) - - -def _make_filename_trans_table() -> dict[int, str]: - linux_chars = ["/", "\x00"] - windows_chars = [ - "<", - ">", - ":", - '"', - "/", - "\\", - "|", - "?", - "*", - ] - windows_chars.extend(chr(i) for i in range(32)) - - trans_table = dict.fromkeys(linux_chars + windows_chars, "_") - return str.maketrans(trans_table) - - -_FILENAME_TRANS_TABLE = _make_filename_trans_table() - - -class DraftVariableSaver: - """Persist draft outputs under the tenant that owns the app or pipeline.""" - - # _DUMMY_OUTPUT_IDENTITY is a placeholder output for workflow nodes. - # Its sole possible value is `None`. - # - # This is used to signal the execution of a workflow node when it has no other outputs. - _DUMMY_OUTPUT_IDENTITY: ClassVar[str] = "__dummy__" - _DUMMY_OUTPUT_VALUE: ClassVar[None] = None - - # _EXCLUDE_VARIABLE_NAMES_MAPPING maps node types and versions to variable names that - # should be excluded when saving draft variables. This prevents certain internal or - # technical variables from being exposed in the draft environment, particularly those - # that aren't meant to be directly edited or viewed by users. - _EXCLUDE_VARIABLE_NAMES_MAPPING: dict[NodeType, frozenset[str]] = { - BuiltinNodeTypes.LLM: frozenset(("finish_reason",)), - BuiltinNodeTypes.LOOP: frozenset(("loop_round",)), - } - - # Database session used for persisting draft variables. - _session: Session - - # Resource owner tenant. An account's current tenant may be unset or point elsewhere - # when draft variables are persisted by an asynchronous workflow execution. - _tenant_id: str - - # The application ID associated with the draft variables. - # This should match the `Workflow.app_id` of the workflow to which the current node belongs. - _app_id: str - - # The ID of the node for which DraftVariableSaver is saving output variables. - _node_id: str - - # The type of the current node (see NodeType). - _node_type: NodeType - - # - _node_execution_id: str - - # _enclosing_node_id identifies the container node that the current node belongs to. - # For example, if the current node is an LLM node inside an Iteration node - # or Loop node, then `_enclosing_node_id` refers to the ID of - # the containing Iteration or Loop node. - # - # If the current node is not nested within another node, `_enclosing_node_id` is - # `None`. - _enclosing_node_id: str | None - - def __init__( - self, - session: Session, - tenant_id: str, - app_id: str, - node_id: str, - node_type: NodeType, - node_execution_id: str, - user: Account, - enclosing_node_id: str | None = None, - ): - # Important: `node_execution_id` parameter refers to the primary key (`id`) of the - # WorkflowNodeExecutionModel/WorkflowNodeExecution, not their `node_execution_id` - # field. These are distinct database fields with different purposes. - self._session = session - self._tenant_id = tenant_id - self._app_id = app_id - self._node_id = node_id - self._node_type = node_type - self._node_execution_id = node_execution_id - self._user = user - self._enclosing_node_id = enclosing_node_id - - def _create_dummy_output_variable(self): - return WorkflowDraftVariable.new_node_variable( - app_id=self._app_id, - user_id=self._user.id, - node_id=self._node_id, - name=self._DUMMY_OUTPUT_IDENTITY, - node_execution_id=self._node_execution_id, - value=build_segment(self._DUMMY_OUTPUT_VALUE), - visible=False, - editable=False, - ) - - def _should_save_output_variables_for_draft(self) -> bool: - if self._enclosing_node_id is not None and self._node_type != BuiltinNodeTypes.VARIABLE_ASSIGNER: - # Currently we do not save output variables for nodes inside loop or iteration. - return False - return True - - def _build_from_variable_assigner_mapping(self, process_data: Mapping[str, Any]) -> list[WorkflowDraftVariable]: - draft_vars: list[WorkflowDraftVariable] = [] - updated_variables = get_updated_variables(process_data) or [] - - for item in updated_variables: - selector = item.selector - if len(selector) < SELECTORS_LENGTH: - raise Exception("selector too short") - # NOTE(QuantumGhost): only the following two kinds of variable could be updated by - # VariableAssigner: ConversationVariable and iteration variable. - # We only save conversation variable here. - if selector[0] != CONVERSATION_VARIABLE_NODE_ID: - continue - # Conversation variables are exposed as NUMBER in the UI even if their - # persisted type is INTEGER. Allow float updates by loosening the type - # to NUMBER here so downstream storage infers the precise subtype. - segment_type = SegmentType.NUMBER if item.value_type == SegmentType.INTEGER else item.value_type - segment = WorkflowDraftVariable.build_segment_with_type(segment_type=segment_type, value=item.new_value) - draft_vars.append( - WorkflowDraftVariable.new_conversation_variable( - app_id=self._app_id, - user_id=self._user.id, - name=item.name, - value=segment, - ) - ) - # Add a dummy output variable to indicate that this node is executed. - draft_vars.append(self._create_dummy_output_variable()) - return draft_vars - - def _build_variables_from_start_mapping(self, output: Mapping[str, Any]) -> list[WorkflowDraftVariable]: - draft_vars = [] - has_non_sys_variables = False - for name, value in output.items(): - value_seg = _build_segment_for_serialized_values(value) - node_id, name = self._normalize_variable_for_start_node(name) - if node_id == SYSTEM_VARIABLE_NODE_ID: - if name == SystemVariableKey.FILES: - # Here we know the type of variable must be `array[file]`, we - # just rebuild files from the serialized payload. - files = [ - build_file_from_stored_mapping( - file_mapping=v, - tenant_id=self._tenant_id, - ) - for v in value - ] - if files: - value_seg = WorkflowDraftVariable.build_segment_with_type(SegmentType.ARRAY_FILE, files) - else: - value_seg = ArrayFileSegment(value=[]) - - draft_vars.append( - WorkflowDraftVariable.new_sys_variable( - app_id=self._app_id, - user_id=self._user.id, - name=name, - node_execution_id=self._node_execution_id, - value=value_seg, - editable=self._should_variable_be_editable(node_id, name), - ) - ) - elif node_id == CONVERSATION_VARIABLE_NODE_ID: - draft_vars.append( - WorkflowDraftVariable.new_conversation_variable( - app_id=self._app_id, - user_id=self._user.id, - name=name, - value=value_seg, - ) - ) - has_non_sys_variables = True - else: - draft_vars.append( - WorkflowDraftVariable.new_node_variable( - app_id=self._app_id, - user_id=self._user.id, - node_id=node_id, - name=name, - node_execution_id=self._node_execution_id, - value=value_seg, - visible=self._should_variable_be_visible(node_id, self._node_type, name), - editable=self._should_variable_be_editable(node_id, name), - ) - ) - has_non_sys_variables = True - if not has_non_sys_variables: - draft_vars.append(self._create_dummy_output_variable()) - return draft_vars - - def _normalize_variable_for_start_node(self, name: str) -> tuple[str, str]: - for reserved_node_id in ( - SYSTEM_VARIABLE_NODE_ID, - ENVIRONMENT_VARIABLE_NODE_ID, - CONVERSATION_VARIABLE_NODE_ID, - RAG_PIPELINE_VARIABLE_NODE_ID, - ): - prefix = f"{reserved_node_id}." - if name.startswith(prefix): - _, name_ = name.split(".", maxsplit=1) - return reserved_node_id, name_ - - return self._node_id, name - - def _build_variables_from_mapping(self, output: Mapping[str, Any]) -> list[WorkflowDraftVariable]: - draft_vars = [] - for name, value in output.items(): - if not self._should_variable_be_saved(name): - logger.debug( - "Skip saving variable as it has been excluded by its node_type, name=%s, node_type=%s", - name, - self._node_type, - ) - continue - if isinstance(value, Segment): - value_seg = value - else: - value_seg = _build_segment_for_serialized_values(value) - draft_vars.append( - self._create_draft_variable( - name=name, - value=value_seg, - visible=True, - editable=True, - ), - # WorkflowDraftVariable.new_node_variable( - # app_id=self._app_id, - # node_id=self._node_id, - # name=name, - # node_execution_id=self._node_execution_id, - # value=value_seg, - # visible=self._should_variable_be_visible(self._node_id, self._node_type, name), - # ) - ) - return draft_vars - - def _generate_filename(self, name: str): - node_id_escaped = self._node_id.translate(_FILENAME_TRANS_TABLE) - return f"{node_id_escaped}-{name}" - - def _try_offload_large_variable( - self, - name: str, - value_seg: Segment, - ) -> tuple[Segment, WorkflowDraftVariableFile] | None: - # This logic is closely tied to `DraftVarLoader._load_offloaded_variable` and must remain - # synchronized with it. - # Ideally, these should be co-located for better maintainability. - # However, due to the current code structure, this is not straightforward. - truncator = VariableTruncator( - max_size_bytes=dify_config.WORKFLOW_VARIABLE_TRUNCATION_MAX_SIZE, - array_element_limit=dify_config.WORKFLOW_VARIABLE_TRUNCATION_ARRAY_LENGTH, - string_length_limit=dify_config.WORKFLOW_VARIABLE_TRUNCATION_STRING_LENGTH, - ) - truncation_result = truncator.truncate(value_seg) - if not truncation_result.truncated: - return None - - original_length = None - if isinstance(value_seg.value, (list, dict)): - original_length = len(value_seg.value) - - # Prepare content for storage - original_content_serialized: str - if isinstance(value_seg, StringSegment): - # For string types, store as plain text - original_content_serialized = cast(str, value_seg.value) - content_type = "text/plain" - filename = f"{self._generate_filename(name)}.txt" - else: - # For other types, store as JSON - original_content_serialized = dumps_with_segments(value_seg.value) - content_type = "application/json" - filename = f"{self._generate_filename(name)}.json" - - original_size = len(original_content_serialized.encode("utf-8")) - - bind = self._session.get_bind() - assert isinstance(bind, Engine) - file_srv = FileService(bind) - - upload_file = file_srv.upload_file( - filename=filename, - content=original_content_serialized.encode(), - mimetype=content_type, - user=self._user, - tenant_id=self._tenant_id, - ) - # Create WorkflowDraftVariableFile record - variable_file = WorkflowDraftVariableFile( - upload_file_id=upload_file.id, - size=original_size, - length=original_length, - value_type=value_seg.value_type, - app_id=self._app_id, - tenant_id=self._tenant_id, - user_id=self._user.id, - ) - variable_file.id = str(uuidv7()) - engine = bind = self._session.get_bind() - assert isinstance(engine, Engine) - with sessionmaker(bind=engine, expire_on_commit=False).begin() as session: - session.add(variable_file) - - return truncation_result.result, variable_file - - def _create_draft_variable( - self, - *, - name: str, - value: Segment, - visible: bool = True, - editable: bool = True, - ) -> WorkflowDraftVariable: - """Create a draft variable with large variable handling and truncation.""" - # Handle Segment values - - offload_result = self._try_offload_large_variable(name, value) - - if offload_result is None: - # Create the draft variable - draft_var = WorkflowDraftVariable.new_node_variable( - app_id=self._app_id, - user_id=self._user.id, - node_id=self._node_id, - name=name, - node_execution_id=self._node_execution_id, - value=value, - visible=visible, - editable=editable, - ) - return draft_var - else: - truncated, var_file = offload_result - # Create the draft variable - draft_var = WorkflowDraftVariable.new_node_variable( - app_id=self._app_id, - user_id=self._user.id, - node_id=self._node_id, - name=name, - node_execution_id=self._node_execution_id, - value=truncated, - visible=visible, - editable=False, - file_id=var_file.id, - ) - return draft_var - - def save( - self, - process_data: Mapping[str, Any] | None = None, - outputs: Mapping[str, Any] | None = None, - ): - draft_vars: list[WorkflowDraftVariable] = [] - if outputs is None: - outputs = {} - if process_data is None: - process_data = {} - if not self._should_save_output_variables_for_draft(): - return - if self._node_type == BuiltinNodeTypes.VARIABLE_ASSIGNER: - draft_vars = self._build_from_variable_assigner_mapping(process_data=process_data) - elif self._node_type == BuiltinNodeTypes.START or is_trigger_node_type(self._node_type): - draft_vars = self._build_variables_from_start_mapping(outputs) - else: - draft_vars = self._build_variables_from_mapping(outputs) - _batch_upsert_draft_variable(self._session, draft_vars) - - @staticmethod - def _should_variable_be_editable(node_id: str, name: str) -> bool: - if node_id in (CONVERSATION_VARIABLE_NODE_ID, ENVIRONMENT_VARIABLE_NODE_ID): - return False - if node_id == SYSTEM_VARIABLE_NODE_ID and not is_system_variable_editable(name): - return False - return True - - @staticmethod - def _should_variable_be_visible(node_id: str, node_type: NodeType, name: str) -> bool: - if node_type == BuiltinNodeTypes.IF_ELSE: - return False - if node_id == SYSTEM_VARIABLE_NODE_ID and not is_system_variable_editable(name): - return False - return True - - def _should_variable_be_saved(self, name: str) -> bool: - exclude_var_names = self._EXCLUDE_VARIABLE_NAMES_MAPPING.get(self._node_type) - if exclude_var_names is None: - return True - return name not in exclude_var_names diff --git a/api/services/workflow_version_number_service.py b/api/services/workflow_version_number_service.py deleted file mode 100644 index 506702ab58d153..00000000000000 --- a/api/services/workflow_version_number_service.py +++ /dev/null @@ -1,59 +0,0 @@ -"""Allocation of user-facing workflow version numbers (`#N`). - -Numbers are unique and monotonically increasing within an app, and are never -reused: `workflow_version_counters` keeps the highest number handed out so far, -so deleting a published version does not free its number. - -The counter is keyed by `Workflow.app_id`, which is polymorphic — it holds an app -id, a pipeline id or a snippet id depending on the workflow kind. UUID uniqueness -across those tables is why the same counter table serves all of them. -""" - -from sqlalchemy import select -from sqlalchemy.dialects.mysql import insert as mysql_insert -from sqlalchemy.dialects.postgresql import insert as pg_insert -from sqlalchemy.orm import Session - -from configs import dify_config -from models.workflow import WorkflowVersionCounter - - -def allocate_version_number(*, session: Session, app_id: str) -> int: - """Reserve and return the next version number for `app_id`. - - The upsert acquires a row lock that is held until the caller's transaction - commits, so concurrent publishes of the same app serialize and never receive - the same number. Callers must run inside a transaction; if it rolls back, - both the counter update and workflow creation roll back, leaving the number - available for the next successful publish. - """ - # Dialect-specific upsert, mirroring `workflow_draft_variable_service`: the - # ORM cannot express "insert or increment" and a read-then-write would race. - # PostgreSQL returns the new value inline; MySQL has no RETURNING, so the - # value is read back within the same transaction while the row is still - # locked by the upsert. - if dify_config.SQLALCHEMY_DATABASE_URI_SCHEME == "postgresql": - stmt = ( - pg_insert(WorkflowVersionCounter) - .values(app_id=app_id, last_version_number=1) - .on_conflict_do_update( - index_elements=[WorkflowVersionCounter.app_id], - set_={"last_version_number": WorkflowVersionCounter.last_version_number + 1}, - ) - .returning(WorkflowVersionCounter.last_version_number) - ) - version_number = session.scalar(stmt) - else: - insert_stmt = mysql_insert(WorkflowVersionCounter).values(app_id=app_id, last_version_number=1) - session.execute( - insert_stmt.on_duplicate_key_update( # type: ignore[attr-defined] - last_version_number=WorkflowVersionCounter.last_version_number + 1, - ) - ) - version_number = session.scalar( - select(WorkflowVersionCounter.last_version_number).where(WorkflowVersionCounter.app_id == app_id) - ) - - if version_number is None: - raise ValueError(f"Failed to allocate a workflow version number for app {app_id}.") - return version_number From 6795e319d9ac301474dfeebc615078d4edfe78ea Mon Sep 17 00:00:00 2001 From: "autofix-ci[bot]" <114827586+autofix-ci[bot]@users.noreply.github.com> Date: Thu, 8 Oct 2026 10:50:40 +0000 Subject: [PATCH 2/5] [autofix.ci] apply automated fixes --- api/openapi/markdown/console-openapi.md | 94 +++++++------------ api/openapi/markdown/service-openapi.md | 2 +- api/openapi/markdown/web-openapi.md | 2 +- .../generated/api/console/apps/orpc.gen.ts | 94 ++----------------- .../generated/api/console/apps/types.gen.ts | 12 +-- .../generated/api/console/apps/zod.gen.ts | 8 +- .../generated/api/console/rag/orpc.gen.ts | 30 +++--- .../generated/api/console/rag/types.gen.ts | 44 ++++++--- .../generated/api/console/rag/zod.gen.ts | 45 +++++++-- .../api/console/snippets/orpc.gen.ts | 4 +- 10 files changed, 136 insertions(+), 199 deletions(-) diff --git a/api/openapi/markdown/console-openapi.md b/api/openapi/markdown/console-openapi.md index 2f4d95bc57d21f..759c421256b5af 100644 --- a/api/openapi/markdown/console-openapi.md +++ b/api/openapi/markdown/console-openapi.md @@ -1730,8 +1730,6 @@ Get advanced chat workflow run list | 200 | Workflow runs count retrieved successfully | **application/json**: [WorkflowRunCountResponse](#workflowruncountresponse)
| ### [POST] /apps/{app_id}/advanced-chat/workflows/draft/human-input/nodes/{node_id}/form/preview -**Preview human input form content and placeholders** - Get human input form preview for advanced chat workflow #### Parameters @@ -1754,8 +1752,6 @@ Get human input form preview for advanced chat workflow | 200 | Human input form preview | **application/json**: [HumanInputFormPreviewResponse](#humaninputformpreviewresponse)
| ### [POST] /apps/{app_id}/advanced-chat/workflows/draft/human-input/nodes/{node_id}/form/run -**Submit human input form preview** - Submit human input form preview for advanced chat workflow #### Parameters @@ -1778,8 +1774,6 @@ Submit human input form preview for advanced chat workflow | 200 | Human input form submission result | **application/json**: [HumanInputFormSubmitResponse](#humaninputformsubmitresponse)
| ### [POST] /apps/{app_id}/advanced-chat/workflows/draft/iteration/nodes/{node_id}/run -**Run draft workflow iteration node** - Run draft workflow iteration node for advanced chat #### Parameters @@ -1804,8 +1798,6 @@ Run draft workflow iteration node for advanced chat | 404 | Node not found | | ### [POST] /apps/{app_id}/advanced-chat/workflows/draft/loop/nodes/{node_id}/run -**Run draft workflow loop node** - Run draft workflow loop node for advanced chat #### Parameters @@ -1830,8 +1822,6 @@ Run draft workflow loop node for advanced chat | 404 | Node not found | | ### [POST] /apps/{app_id}/advanced-chat/workflows/draft/run -**Run draft workflow** - Run draft workflow for advanced chat application #### Parameters @@ -2650,11 +2640,7 @@ Get conversation variables for an application | 200 | Conversation variables retrieved successfully | **application/json**: [PaginatedConversationVariableResponse](#paginatedconversationvariableresponse)
| ### [POST] /apps/{app_id}/convert-to-workflow -**Convert basic mode of chatbot app to workflow mode** - Convert application to workflow mode -Convert expert mode of chatbot app to workflow mode -Convert Completion App to Workflow App #### Parameters @@ -3443,8 +3429,6 @@ Get workflow application execution logs | 200 | Workflow runs count retrieved successfully | **application/json**: [WorkflowRunCountResponse](#workflowruncountresponse)
| ### [POST] /apps/{app_id}/workflow-runs/tasks/{task_id}/stop -**Stop workflow task** - Stop running workflow task #### Parameters @@ -3809,8 +3793,6 @@ Get workflow daily token cost statistics | 200 | Daily token cost statistics retrieved successfully | **application/json**: [WorkflowDailyTokenCostStatisticResponse](#workflowdailytokencoststatisticresponse)
| ### [GET] /apps/{app_id}/workflows -**Get published workflows** - Get all published workflows for an application #### Parameters @@ -3830,8 +3812,6 @@ Get all published workflows for an application | 200 | Published workflows retrieved successfully | **application/json**: [WorkflowPaginationResponse](#workflowpaginationresponse)
| ### [GET] /apps/{app_id}/workflows/default-workflow-block-configs -**Get default block config** - Get default block configurations for workflow #### Parameters @@ -3847,8 +3827,6 @@ Get default block configurations for workflow | 200 | Default block configurations retrieved successfully | **application/json**: [DefaultBlockConfigsResponse](#defaultblockconfigsresponse)
| ### [GET] /apps/{app_id}/workflows/default-workflow-block-configs/{block_type} -**Get default block config** - Get default block configuration by type #### Parameters @@ -3867,8 +3845,6 @@ Get default block configuration by type | 404 | Block type not found | | ### [GET] /apps/{app_id}/workflows/draft -**Get draft workflow** - Get draft workflow for an application #### Parameters @@ -3885,8 +3861,6 @@ Get draft workflow for an application | 404 | Draft workflow not found | | ### [POST] /apps/{app_id}/workflows/draft -**Sync draft workflow** - Sync draft workflow configuration #### Parameters @@ -3947,7 +3921,7 @@ Update conversation variables for workflow draft | 200 | Conversation variables updated successfully | **application/json**: [SimpleResultResponse](#simpleresultresponse)
| ### [GET] /apps/{app_id}/workflows/draft/environment-variables -**Get draft workflow** +**Get environment variables** Get environment variables for workflow @@ -4007,8 +3981,6 @@ Update draft workflow features | 200 | Workflow features updated successfully | **application/json**: [SimpleResultResponse](#simpleresultresponse)
| ### [POST] /apps/{app_id}/workflows/draft/human-input/nodes/{node_id}/delivery-test -**Test human input delivery** - Test human input delivery for workflow #### Parameters @@ -4031,8 +4003,6 @@ Test human input delivery for workflow | 200 | Human input delivery test result | **application/json**: [EmptyObjectResponse](#emptyobjectresponse)
| ### [POST] /apps/{app_id}/workflows/draft/human-input/nodes/{node_id}/form/preview -**Preview human input form content and placeholders** - Get human input form preview for workflow #### Parameters @@ -4055,8 +4025,6 @@ Get human input form preview for workflow | 200 | Human input form preview | **application/json**: [HumanInputFormPreviewResponse](#humaninputformpreviewresponse)
| ### [POST] /apps/{app_id}/workflows/draft/human-input/nodes/{node_id}/form/run -**Submit human input form preview** - Submit human input form preview for workflow #### Parameters @@ -4079,7 +4047,7 @@ Submit human input form preview for workflow | 200 | Human input form submission result | **application/json**: [HumanInputFormSubmitResponse](#humaninputformsubmitresponse)
| ### [POST] /apps/{app_id}/workflows/draft/iteration/nodes/{node_id}/run -**Run draft workflow iteration node** +Run draft workflow iteration node #### Parameters @@ -4103,7 +4071,7 @@ Submit human input form preview for workflow | 404 | Node not found | | ### [POST] /apps/{app_id}/workflows/draft/loop/nodes/{node_id}/run -**Run draft workflow loop node** +Run draft workflow loop node #### Parameters @@ -4274,7 +4242,7 @@ Get last run result for draft workflow node | 404 | Node last run not found | | ### [POST] /apps/{app_id}/workflows/draft/nodes/{node_id}/run -**Run draft workflow node** +Run draft workflow node #### Parameters @@ -4298,7 +4266,7 @@ Get last run result for draft workflow node | 404 | Node not found | | ### [POST] /apps/{app_id}/workflows/draft/nodes/{node_id}/trigger/run -**Poll for trigger events and execute single node when event arrives** +Poll for trigger events and execute single node when event arrives #### Parameters @@ -4348,7 +4316,7 @@ Get variables for a specific node | 200 | Node variables retrieved successfully | **application/json**: [WorkflowDraftVariableListResponse](#workflowdraftvariablelistresponse)
| ### [POST] /apps/{app_id}/workflows/draft/run -**Run draft workflow** +Run draft workflow #### Parameters @@ -4456,7 +4424,7 @@ Get system variables for workflow | 200 | System variables retrieved successfully | **application/json**: [WorkflowDraftVariableListResponse](#workflowdraftvariablelistresponse)
| ### [POST] /apps/{app_id}/workflows/draft/trigger/run -**Poll for trigger events and execute full workflow when event arrives** +Poll for trigger events and execute full workflow when event arrives #### Parameters @@ -4479,7 +4447,7 @@ Get system variables for workflow | 500 | Internal server error | | ### [POST] /apps/{app_id}/workflows/draft/trigger/run-all -**Full workflow debug when the start node is a trigger** +Full workflow debug when the start node is a trigger #### Parameters @@ -4517,7 +4485,7 @@ Delete all draft workflow variables | 204 | Workflow variables deleted successfully | ### [GET] /apps/{app_id}/workflows/draft/variables -**Get draft workflow** +**List draft workflow variables without loading their values** Get draft workflow variables @@ -4611,8 +4579,6 @@ Reset a workflow variable to its default value | 404 | Variable not found | | ### [GET] /apps/{app_id}/workflows/publish -**Get published workflow** - Get published workflow for an application #### Parameters @@ -4628,8 +4594,6 @@ Get published workflow for an application | 200 | Published workflow retrieved successfully, or null if not found | **application/json**: [WorkflowResponse](#workflowresponse)
| ### [POST] /apps/{app_id}/workflows/publish -**Publish workflow** - #### Parameters | Name | Located in | Description | Required | Schema | @@ -4736,8 +4700,6 @@ Full value for one declared output of a published run. | 200 | Success | **application/json**: [WebhookTriggerResponse](#webhooktriggerresponse)
| ### [DELETE] /apps/{app_id}/workflows/{workflow_id} -**Delete workflow** - #### Parameters | Name | Located in | Description | Required | Schema | @@ -4749,11 +4711,9 @@ Full value for one declared output of a published run. | Code | Description | | ---- | ----------- | -| 200 | Success | +| 204 | Workflow deleted successfully | ### [PATCH] /apps/{app_id}/workflows/{workflow_id} -**Update workflow attributes** - Update workflow by ID #### Parameters @@ -7827,11 +7787,17 @@ Update account-level Step-by-step Tour state | ---- | ---------- | ----------- | -------- | ------ | | pipeline_id | path | | Yes | string (uuid) | +#### Request Body + +| Required | Schema | +| -------- | ------ | +| Yes | **application/json**: [DraftWorkflowSyncPayload](#draftworkflowsyncpayload)
| + #### Responses -| Code | Description | -| ---- | ----------- | -| 200 | Success | +| Code | Description | Schema | +| ---- | ----------- | ------ | +| 200 | Success | **application/json**: [RagPipelineWorkflowSyncResponse](#ragpipelineworkflowsyncresponse)
| ### [POST] /rag/pipelines/{pipeline_id}/workflows/draft/datasource/nodes/{node_id}/run **Run rag pipeline datasource** @@ -7877,7 +7843,7 @@ Update account-level Step-by-step Tour state | 200 | Datasource variables set successfully | **application/json**: [WorkflowRunNodeExecutionResponse](#workflowrunnodeexecutionresponse)
| ### [GET] /rag/pipelines/{pipeline_id}/workflows/draft/environment-variables -**Get draft workflow** +**Get environment variables** #### Parameters @@ -8079,7 +8045,7 @@ Update account-level Step-by-step Tour state | 204 | Workflow variables deleted successfully | ### [GET] /rag/pipelines/{pipeline_id}/workflows/draft/variables -**Get draft workflow** +**List draft pipeline variables without loading their values** #### Parameters @@ -8131,6 +8097,12 @@ Update account-level Step-by-step Tour state | pipeline_id | path | | Yes | string (uuid) | | variable_id | path | | Yes | string (uuid) | +#### Request Body + +| Required | Schema | +| -------- | ------ | +| Yes | **application/json**: [WorkflowDraftVariablePatchPayload](#workflowdraftvariablepatchpayload)
| + #### Responses | Code | Description | Schema | @@ -8168,7 +8140,7 @@ Update account-level Step-by-step Tour state | 200 | Published workflow retrieved successfully, or null if not exist | **application/json**: [WorkflowResponse](#workflowresponse)
| ### [POST] /rag/pipelines/{pipeline_id}/workflows/publish -**Publish workflow** +**Publish a prepared Pipeline and its dataset settings atomically** #### Parameters @@ -8280,8 +8252,6 @@ Update account-level Step-by-step Tour state | 200 | Success | **application/json**: [RagPipelineOpaqueResponse](#ragpipelineopaqueresponse)
| ### [DELETE] /rag/pipelines/{pipeline_id}/workflows/{workflow_id} -**Delete a published workflow version that is not currently active on the pipeline** - #### Parameters | Name | Located in | Description | Required | Schema | @@ -9074,7 +9044,9 @@ Reset a draft workflow variable to its default value (snippet scope) | 400 | No draft workflow found | | ### [DELETE] /snippets/{snippet_id}/workflows/{workflow_id} -**Delete a published snippet workflow version** +**Delete a version and retire its unowned Agents through the shared use case** + +Delete a published snippet workflow version #### Parameters @@ -15533,10 +15505,10 @@ Legacy Chat App model config used only for follow-up question generation. | status | string | | Yes | | time_cost | number
integer | | Yes | | tool_icon | | | No | -| tool_input | object | | Yes | +| tool_input | | | Yes | | tool_label | string | | Yes | | tool_name | string | | Yes | -| tool_output | object | | Yes | +| tool_output | | | Yes | | tool_parameters | object | | Yes | #### AgentUserSatisfactionRateStatisticResponse diff --git a/api/openapi/markdown/service-openapi.md b/api/openapi/markdown/service-openapi.md index a30ee05de6fa0a..86f95c03781c20 100644 --- a/api/openapi/markdown/service-openapi.md +++ b/api/openapi/markdown/service-openapi.md @@ -1018,7 +1018,7 @@ Execute the full knowledge pipeline for a knowledge base. Published runs are que | 200 | Pipeline execution result. Published runs return a JSON object containing `batch`, `dataset`, and `documents`. Draft runs return `text/event-stream` for streaming mode or a workflow result JSON object for blocking mode. | **application/json**: [PipelineRunJsonResponse](#pipelinerunjsonresponse)
**text/event-stream**: string
| | 400 | Bad request - invalid payload or pipeline is not configured | | | 401 | Unauthorized - invalid API token | | -| 403 | `forbidden` : Forbidden. | | +| 403 | `forbidden` : AccessDeniedError. | | | 404 | `not_found` : Dataset not found. | | | 500 | `pipeline_run_error` : Pipeline execution failed. | | diff --git a/api/openapi/markdown/web-openapi.md b/api/openapi/markdown/web-openapi.md index 1f6d7a8e723090..cdd295149db3c4 100644 --- a/api/openapi/markdown/web-openapi.md +++ b/api/openapi/markdown/web-openapi.md @@ -361,7 +361,7 @@ GET /api/form/human_input/ | Code | Description | Schema | | ---- | ----------- | ------ | | 200 | Form retrieved successfully | **application/json**: [HumanInputFormDefinitionResponse](#humaninputformdefinitionresponse)
| -| 403 | Forbidden | | +| 403 | AccessDeniedError | | | 404 | Form not found | | | 412 | Form already submitted or expired | | | 429 | Too many requests | | diff --git a/packages/contracts/generated/api/console/apps/orpc.gen.ts b/packages/contracts/generated/api/console/apps/orpc.gen.ts index 8df754ae8f6ee2..8197443bd438b7 100644 --- a/packages/contracts/generated/api/console/apps/orpc.gen.ts +++ b/packages/contracts/generated/api/console/apps/orpc.gen.ts @@ -632,8 +632,6 @@ export const workflowRuns = { } /** - * Preview human input form content and placeholders - * * Get human input form preview for advanced chat workflow */ export const post4 = oc @@ -643,7 +641,6 @@ export const post4 = oc method: 'POST', operationId: 'postAppsByAppIdAdvancedChatWorkflowsDraftHumanInputNodesByNodeIdFormPreview', path: '/apps/{app_id}/advanced-chat/workflows/draft/human-input/nodes/{node_id}/form/preview', - summary: 'Preview human input form content and placeholders', tags: ['console'], }) .input( @@ -659,8 +656,6 @@ export const preview = { } /** - * Submit human input form preview - * * Submit human input form preview for advanced chat workflow */ export const post5 = oc @@ -670,7 +665,6 @@ export const post5 = oc method: 'POST', operationId: 'postAppsByAppIdAdvancedChatWorkflowsDraftHumanInputNodesByNodeIdFormRun', path: '/apps/{app_id}/advanced-chat/workflows/draft/human-input/nodes/{node_id}/form/run', - summary: 'Submit human input form preview', tags: ['console'], }) .input( @@ -703,8 +697,6 @@ export const humanInput = { } /** - * Run draft workflow iteration node - * * Run draft workflow iteration node for advanced chat */ export const post6 = oc @@ -714,7 +706,6 @@ export const post6 = oc method: 'POST', operationId: 'postAppsByAppIdAdvancedChatWorkflowsDraftIterationNodesByNodeIdRun', path: '/apps/{app_id}/advanced-chat/workflows/draft/iteration/nodes/{node_id}/run', - summary: 'Run draft workflow iteration node', tags: ['console'], }) .input( @@ -742,8 +733,6 @@ export const iteration = { } /** - * Run draft workflow loop node - * * Run draft workflow loop node for advanced chat */ export const post7 = oc @@ -753,7 +742,6 @@ export const post7 = oc method: 'POST', operationId: 'postAppsByAppIdAdvancedChatWorkflowsDraftLoopNodesByNodeIdRun', path: '/apps/{app_id}/advanced-chat/workflows/draft/loop/nodes/{node_id}/run', - summary: 'Run draft workflow loop node', tags: ['console'], }) .input( @@ -781,8 +769,6 @@ export const loop = { } /** - * Run draft workflow - * * Run draft workflow for advanced chat application */ export const post8 = oc @@ -792,7 +778,6 @@ export const post8 = oc method: 'POST', operationId: 'postAppsByAppIdAdvancedChatWorkflowsDraftRun', path: '/apps/{app_id}/advanced-chat/workflows/draft/run', - summary: 'Run draft workflow', tags: ['console'], }) .input( @@ -1750,21 +1735,15 @@ export const conversationVariables = { } /** - * Convert basic mode of chatbot app to workflow mode - * * Convert application to workflow mode - * Convert expert mode of chatbot app to workflow mode - * Convert Completion App to Workflow App */ export const post21 = oc .route({ - description: - 'Convert application to workflow mode\nConvert expert mode of chatbot app to workflow mode\nConvert Completion App to Workflow App', + description: 'Convert application to workflow mode', inputStructure: 'detailed', method: 'POST', operationId: 'postAppsByAppIdConvertToWorkflow', path: '/apps/{app_id}/convert-to-workflow', - summary: 'Convert basic mode of chatbot app to workflow mode', tags: ['console'], }) .input( @@ -2600,8 +2579,6 @@ export const count3 = { } /** - * Stop workflow task - * * Stop running workflow task */ export const post38 = oc @@ -2611,7 +2588,6 @@ export const post38 = oc method: 'POST', operationId: 'postAppsByAppIdWorkflowRunsTasksByTaskIdStop', path: '/apps/{app_id}/workflow-runs/tasks/{task_id}/stop', - summary: 'Stop workflow task', tags: ['console'], }) .input(z.object({ params: zPostAppsByAppIdWorkflowRunsTasksByTaskIdStopPath })) @@ -3142,8 +3118,6 @@ export const workflow = { } /** - * Get default block config - * * Get default block configuration by type */ export const get61 = oc @@ -3153,7 +3127,6 @@ export const get61 = oc method: 'GET', operationId: 'getAppsByAppIdWorkflowsDefaultWorkflowBlockConfigsByBlockType', path: '/apps/{app_id}/workflows/default-workflow-block-configs/{block_type}', - summary: 'Get default block config', tags: ['console'], }) .input( @@ -3169,8 +3142,6 @@ export const byBlockType = { } /** - * Get default block config - * * Get default block configurations for workflow */ export const get62 = oc @@ -3180,7 +3151,6 @@ export const get62 = oc method: 'GET', operationId: 'getAppsByAppIdWorkflowsDefaultWorkflowBlockConfigs', path: '/apps/{app_id}/workflows/default-workflow-block-configs', - summary: 'Get default block config', tags: ['console'], }) .input(z.object({ params: zGetAppsByAppIdWorkflowsDefaultWorkflowBlockConfigsPath })) @@ -3232,7 +3202,7 @@ export const conversationVariables2 = { } /** - * Get draft workflow + * Get environment variables * * Get environment variables for workflow */ @@ -3243,7 +3213,7 @@ export const get64 = oc method: 'GET', operationId: 'getAppsByAppIdWorkflowsDraftEnvironmentVariables', path: '/apps/{app_id}/workflows/draft/environment-variables', - summary: 'Get draft workflow', + summary: 'Get environment variables', tags: ['console'], }) .input(z.object({ params: zGetAppsByAppIdWorkflowsDraftEnvironmentVariablesPath })) @@ -3299,8 +3269,6 @@ export const features = { } /** - * Test human input delivery - * * Test human input delivery for workflow */ export const post46 = oc @@ -3310,7 +3278,6 @@ export const post46 = oc method: 'POST', operationId: 'postAppsByAppIdWorkflowsDraftHumanInputNodesByNodeIdDeliveryTest', path: '/apps/{app_id}/workflows/draft/human-input/nodes/{node_id}/delivery-test', - summary: 'Test human input delivery', tags: ['console'], }) .input( @@ -3326,8 +3293,6 @@ export const deliveryTest = { } /** - * Preview human input form content and placeholders - * * Get human input form preview for workflow */ export const post47 = oc @@ -3337,7 +3302,6 @@ export const post47 = oc method: 'POST', operationId: 'postAppsByAppIdWorkflowsDraftHumanInputNodesByNodeIdFormPreview', path: '/apps/{app_id}/workflows/draft/human-input/nodes/{node_id}/form/preview', - summary: 'Preview human input form content and placeholders', tags: ['console'], }) .input( @@ -3353,8 +3317,6 @@ export const preview4 = { } /** - * Submit human input form preview - * * Submit human input form preview for workflow */ export const post48 = oc @@ -3364,7 +3326,6 @@ export const post48 = oc method: 'POST', operationId: 'postAppsByAppIdWorkflowsDraftHumanInputNodesByNodeIdFormRun', path: '/apps/{app_id}/workflows/draft/human-input/nodes/{node_id}/form/run', - summary: 'Submit human input form preview', tags: ['console'], }) .input( @@ -3398,8 +3359,6 @@ export const humanInput2 = { } /** - * Run draft workflow iteration node - * * Run draft workflow iteration node */ export const post49 = oc @@ -3409,7 +3368,6 @@ export const post49 = oc method: 'POST', operationId: 'postAppsByAppIdWorkflowsDraftIterationNodesByNodeIdRun', path: '/apps/{app_id}/workflows/draft/iteration/nodes/{node_id}/run', - summary: 'Run draft workflow iteration node', tags: ['console'], }) .input( @@ -3437,8 +3395,6 @@ export const iteration2 = { } /** - * Run draft workflow loop node - * * Run draft workflow loop node */ export const post50 = oc @@ -3448,7 +3404,6 @@ export const post50 = oc method: 'POST', operationId: 'postAppsByAppIdWorkflowsDraftLoopNodesByNodeIdRun', path: '/apps/{app_id}/workflows/draft/loop/nodes/{node_id}/run', - summary: 'Run draft workflow loop node', tags: ['console'], }) .input( @@ -3634,8 +3589,6 @@ export const lastRun = { } /** - * Run draft workflow node - * * Run draft workflow node */ export const post55 = oc @@ -3645,7 +3598,6 @@ export const post55 = oc method: 'POST', operationId: 'postAppsByAppIdWorkflowsDraftNodesByNodeIdRun', path: '/apps/{app_id}/workflows/draft/nodes/{node_id}/run', - summary: 'Run draft workflow node', tags: ['console'], }) .input( @@ -3661,8 +3613,6 @@ export const run8 = { } /** - * Poll for trigger events and execute single node when event arrives - * * Poll for trigger events and execute single node when event arrives */ export const post56 = oc @@ -3672,7 +3622,6 @@ export const post56 = oc method: 'POST', operationId: 'postAppsByAppIdWorkflowsDraftNodesByNodeIdTriggerRun', path: '/apps/{app_id}/workflows/draft/nodes/{node_id}/trigger/run', - summary: 'Poll for trigger events and execute single node when event arrives', tags: ['console'], }) .input(z.object({ params: zPostAppsByAppIdWorkflowsDraftNodesByNodeIdTriggerRunPath })) @@ -3735,8 +3684,6 @@ export const nodes7 = { } /** - * Run draft workflow - * * Run draft workflow */ export const post57 = oc @@ -3746,7 +3693,6 @@ export const post57 = oc method: 'POST', operationId: 'postAppsByAppIdWorkflowsDraftRun', path: '/apps/{app_id}/workflows/draft/run', - summary: 'Run draft workflow', tags: ['console'], }) .input( @@ -3876,8 +3822,6 @@ export const systemVariables = { } /** - * Poll for trigger events and execute full workflow when event arrives - * * Poll for trigger events and execute full workflow when event arrives */ export const post58 = oc @@ -3887,7 +3831,6 @@ export const post58 = oc method: 'POST', operationId: 'postAppsByAppIdWorkflowsDraftTriggerRun', path: '/apps/{app_id}/workflows/draft/trigger/run', - summary: 'Poll for trigger events and execute full workflow when event arrives', tags: ['console'], }) .input( @@ -3903,8 +3846,6 @@ export const run11 = { } /** - * Full workflow debug when the start node is a trigger - * * Full workflow debug when the start node is a trigger */ export const post59 = oc @@ -3914,7 +3855,6 @@ export const post59 = oc method: 'POST', operationId: 'postAppsByAppIdWorkflowsDraftTriggerRunAll', path: '/apps/{app_id}/workflows/draft/trigger/run-all', - summary: 'Full workflow debug when the start node is a trigger', tags: ['console'], }) .input( @@ -4028,7 +3968,7 @@ export const delete13 = oc .output(zDeleteAppsByAppIdWorkflowsDraftVariablesResponse) /** - * Get draft workflow + * List draft workflow variables without loading their values * * Get draft workflow variables */ @@ -4039,7 +3979,7 @@ export const get75 = oc method: 'GET', operationId: 'getAppsByAppIdWorkflowsDraftVariables', path: '/apps/{app_id}/workflows/draft/variables', - summary: 'Get draft workflow', + summary: 'List draft workflow variables without loading their values', tags: ['console'], }) .input( @@ -4057,8 +3997,6 @@ export const variables2 = { } /** - * Get draft workflow - * * Get draft workflow for an application */ export const get76 = oc @@ -4068,15 +4006,12 @@ export const get76 = oc method: 'GET', operationId: 'getAppsByAppIdWorkflowsDraft', path: '/apps/{app_id}/workflows/draft', - summary: 'Get draft workflow', tags: ['console'], }) .input(z.object({ params: zGetAppsByAppIdWorkflowsDraftPath })) .output(zGetAppsByAppIdWorkflowsDraftResponse) /** - * Sync draft workflow - * * Sync draft workflow configuration */ export const post60 = oc @@ -4086,7 +4021,6 @@ export const post60 = oc method: 'POST', operationId: 'postAppsByAppIdWorkflowsDraft', path: '/apps/{app_id}/workflows/draft', - summary: 'Sync draft workflow', tags: ['console'], }) .input( @@ -4115,8 +4049,6 @@ export const draft2 = { } /** - * Get published workflow - * * Get published workflow for an application */ export const get77 = oc @@ -4126,22 +4058,17 @@ export const get77 = oc method: 'GET', operationId: 'getAppsByAppIdWorkflowsPublish', path: '/apps/{app_id}/workflows/publish', - summary: 'Get published workflow', tags: ['console'], }) .input(z.object({ params: zGetAppsByAppIdWorkflowsPublishPath })) .output(zGetAppsByAppIdWorkflowsPublishResponse) -/** - * Publish workflow - */ export const post61 = oc .route({ inputStructure: 'detailed', method: 'POST', operationId: 'postAppsByAppIdWorkflowsPublish', path: '/apps/{app_id}/workflows/publish', - summary: 'Publish workflow', tags: ['console'], }) .input( @@ -4307,24 +4234,19 @@ export const restore = { post: post62, } -/** - * Delete workflow - */ export const delete14 = oc .route({ inputStructure: 'detailed', method: 'DELETE', operationId: 'deleteAppsByAppIdWorkflowsByWorkflowId', path: '/apps/{app_id}/workflows/{workflow_id}', - summary: 'Delete workflow', + successStatus: 204, tags: ['console'], }) .input(z.object({ params: zDeleteAppsByAppIdWorkflowsByWorkflowIdPath })) .output(zDeleteAppsByAppIdWorkflowsByWorkflowIdResponse) /** - * Update workflow attributes - * * Update workflow by ID */ export const patch3 = oc @@ -4334,7 +4256,6 @@ export const patch3 = oc method: 'PATCH', operationId: 'patchAppsByAppIdWorkflowsByWorkflowId', path: '/apps/{app_id}/workflows/{workflow_id}', - summary: 'Update workflow attributes', tags: ['console'], }) .input( @@ -4352,8 +4273,6 @@ export const byWorkflowId = { } /** - * Get published workflows - * * Get all published workflows for an application */ export const get83 = oc @@ -4363,7 +4282,6 @@ export const get83 = oc method: 'GET', operationId: 'getAppsByAppIdWorkflows', path: '/apps/{app_id}/workflows', - summary: 'Get published workflows', tags: ['console'], }) .input( diff --git a/packages/contracts/generated/api/console/apps/types.gen.ts b/packages/contracts/generated/api/console/apps/types.gen.ts index 2d4fbb1b6c22ad..bec2b43430dfe0 100644 --- a/packages/contracts/generated/api/console/apps/types.gen.ts +++ b/packages/contracts/generated/api/console/apps/types.gen.ts @@ -2496,14 +2496,10 @@ export type AgentToolCallResponse = { status: string time_cost: number | number tool_icon?: unknown - tool_input: { - [key: string]: unknown - } + tool_input: unknown tool_label: string tool_name: string - tool_output: { - [key: string]: unknown - } + tool_output: unknown tool_parameters: { [key: string]: unknown } @@ -7596,9 +7592,7 @@ export type DeleteAppsByAppIdWorkflowsByWorkflowIdData = { } export type DeleteAppsByAppIdWorkflowsByWorkflowIdResponses = { - 200: { - [key: string]: unknown - } + 204: void } export type DeleteAppsByAppIdWorkflowsByWorkflowIdResponse = diff --git a/packages/contracts/generated/api/console/apps/zod.gen.ts b/packages/contracts/generated/api/console/apps/zod.gen.ts index eb7445c1fc7415..2570468168d888 100644 --- a/packages/contracts/generated/api/console/apps/zod.gen.ts +++ b/packages/contracts/generated/api/console/apps/zod.gen.ts @@ -2195,10 +2195,10 @@ export const zAgentToolCallResponse = z.object({ status: z.string(), time_cost: z.union([z.number(), z.int()]), tool_icon: z.unknown().optional(), - tool_input: z.record(z.string(), z.unknown()), + tool_input: z.unknown(), tool_label: z.string(), tool_name: z.string(), - tool_output: z.record(z.string(), z.unknown()), + tool_output: z.unknown(), tool_parameters: z.record(z.string(), z.unknown()), }) @@ -7508,9 +7508,9 @@ export const zDeleteAppsByAppIdWorkflowsByWorkflowIdPath = z.object({ }) /** - * Success + * Workflow deleted successfully */ -export const zDeleteAppsByAppIdWorkflowsByWorkflowIdResponse = z.record(z.string(), z.unknown()) +export const zDeleteAppsByAppIdWorkflowsByWorkflowIdResponse = z.void() export const zPatchAppsByAppIdWorkflowsByWorkflowIdBody = zWorkflowUpdatePayload diff --git a/packages/contracts/generated/api/console/rag/orpc.gen.ts b/packages/contracts/generated/api/console/rag/orpc.gen.ts index d2a27ca561d997..223139073f73ca 100644 --- a/packages/contracts/generated/api/console/rag/orpc.gen.ts +++ b/packages/contracts/generated/api/console/rag/orpc.gen.ts @@ -75,6 +75,7 @@ import { zPatchRagPipelineCustomizedTemplatesByTemplateIdResponse, zPatchRagPipelinesByPipelineIdWorkflowsByWorkflowIdPath, zPatchRagPipelinesByPipelineIdWorkflowsByWorkflowIdResponse, + zPatchRagPipelinesByPipelineIdWorkflowsDraftVariablesByVariableIdBody, zPatchRagPipelinesByPipelineIdWorkflowsDraftVariablesByVariableIdPath, zPatchRagPipelinesByPipelineIdWorkflowsDraftVariablesByVariableIdResponse, zPostRagPipelineCustomizedTemplatesByTemplateIdPath, @@ -89,6 +90,7 @@ import { zPostRagPipelinesByPipelineIdWorkflowRunsTasksByTaskIdStopResponse, zPostRagPipelinesByPipelineIdWorkflowsByWorkflowIdRestorePath, zPostRagPipelinesByPipelineIdWorkflowsByWorkflowIdRestoreResponse, + zPostRagPipelinesByPipelineIdWorkflowsDraftBody, zPostRagPipelinesByPipelineIdWorkflowsDraftDatasourceNodesByNodeIdRunBody, zPostRagPipelinesByPipelineIdWorkflowsDraftDatasourceNodesByNodeIdRunPath, zPostRagPipelinesByPipelineIdWorkflowsDraftDatasourceNodesByNodeIdRunResponse, @@ -609,7 +611,7 @@ export const datasource = { } /** - * Get draft workflow + * Get environment variables */ export const get12 = oc .route({ @@ -617,7 +619,7 @@ export const get12 = oc method: 'GET', operationId: 'getRagPipelinesByPipelineIdWorkflowsDraftEnvironmentVariables', path: '/rag/pipelines/{pipeline_id}/workflows/draft/environment-variables', - summary: 'Get draft workflow', + summary: 'Get environment variables', tags: ['console'], }) .input(z.object({ params: zGetRagPipelinesByPipelineIdWorkflowsDraftEnvironmentVariablesPath })) @@ -924,7 +926,10 @@ export const patch2 = oc tags: ['console'], }) .input( - z.object({ params: zPatchRagPipelinesByPipelineIdWorkflowsDraftVariablesByVariableIdPath }), + z.object({ + body: zPatchRagPipelinesByPipelineIdWorkflowsDraftVariablesByVariableIdBody, + params: zPatchRagPipelinesByPipelineIdWorkflowsDraftVariablesByVariableIdPath, + }), ) .output(zPatchRagPipelinesByPipelineIdWorkflowsDraftVariablesByVariableIdResponse) @@ -948,7 +953,7 @@ export const delete4 = oc .output(zDeleteRagPipelinesByPipelineIdWorkflowsDraftVariablesResponse) /** - * Get draft workflow + * List draft pipeline variables without loading their values */ export const get19 = oc .route({ @@ -956,7 +961,7 @@ export const get19 = oc method: 'GET', operationId: 'getRagPipelinesByPipelineIdWorkflowsDraftVariables', path: '/rag/pipelines/{pipeline_id}/workflows/draft/variables', - summary: 'Get draft workflow', + summary: 'List draft pipeline variables without loading their values', tags: ['console'], }) .input( @@ -1000,7 +1005,12 @@ export const post15 = oc summary: 'Sync draft workflow', tags: ['console'], }) - .input(z.object({ params: zPostRagPipelinesByPipelineIdWorkflowsDraftPath })) + .input( + z.object({ + body: zPostRagPipelinesByPipelineIdWorkflowsDraftBody, + params: zPostRagPipelinesByPipelineIdWorkflowsDraftPath, + }), + ) .output(zPostRagPipelinesByPipelineIdWorkflowsDraftResponse) export const draft = { @@ -1034,7 +1044,7 @@ export const get21 = oc .output(zGetRagPipelinesByPipelineIdWorkflowsPublishResponse) /** - * Publish workflow + * Publish a prepared Pipeline and its dataset settings atomically */ export const post16 = oc .route({ @@ -1042,7 +1052,7 @@ export const post16 = oc method: 'POST', operationId: 'postRagPipelinesByPipelineIdWorkflowsPublish', path: '/rag/pipelines/{pipeline_id}/workflows/publish', - summary: 'Publish workflow', + summary: 'Publish a prepared Pipeline and its dataset settings atomically', tags: ['console'], }) .input(z.object({ params: zPostRagPipelinesByPipelineIdWorkflowsPublishPath })) @@ -1216,9 +1226,6 @@ export const restore = { post: post20, } -/** - * Delete a published workflow version that is not currently active on the pipeline - */ export const delete5 = oc .route({ inputStructure: 'detailed', @@ -1226,7 +1233,6 @@ export const delete5 = oc operationId: 'deleteRagPipelinesByPipelineIdWorkflowsByWorkflowId', path: '/rag/pipelines/{pipeline_id}/workflows/{workflow_id}', successStatus: 204, - summary: 'Delete a published workflow version that is not currently active on the pipeline', tags: ['console'], }) .input(z.object({ params: zDeleteRagPipelinesByPipelineIdWorkflowsByWorkflowIdPath })) diff --git a/packages/contracts/generated/api/console/rag/types.gen.ts b/packages/contracts/generated/api/console/rag/types.gen.ts index 70953beb7726cf..e6186ad9b34a82 100644 --- a/packages/contracts/generated/api/console/rag/types.gen.ts +++ b/packages/contracts/generated/api/console/rag/types.gen.ts @@ -180,6 +180,31 @@ export type WorkflowResponse = { version_number?: number | null } +export type DraftWorkflowSyncPayload = { + conversation_variables?: Array<{ + [key: string]: unknown + }> | null + environment_variables?: Array<{ + [key: string]: unknown + }> | null + features?: { + [key: string]: unknown + } | null + graph: { + [key: string]: unknown + } + hash?: string | null + rag_pipeline_variables?: Array<{ + [key: string]: unknown + }> | null +} + +export type RagPipelineWorkflowSyncResponse = { + hash: string + result: string + updated_at: number +} + export type DatasourceNodeRunPayload = { credential_id?: string | null datasource_type: string @@ -277,6 +302,11 @@ export type WorkflowDraftVariableResponse = { visible: boolean } +export type WorkflowDraftVariablePatchPayload = { + name?: string | null + value?: unknown | null +} + export type RagPipelineWorkflowPublishResponse = { created_at: number result: string @@ -304,12 +334,6 @@ export type PublishedWorkflowRunPayload = { start_node_id: string } -export type RagPipelineWorkflowSyncResponse = { - hash: string - result: string - updated_at: number -} - export type ImportStatus = 'completed' | 'completed-with-warnings' | 'failed' | 'pending' export type DatasetDocMetadataResponse = { @@ -1122,7 +1146,7 @@ export type GetRagPipelinesByPipelineIdWorkflowsDraftResponse = GetRagPipelinesByPipelineIdWorkflowsDraftResponses[keyof GetRagPipelinesByPipelineIdWorkflowsDraftResponses] export type PostRagPipelinesByPipelineIdWorkflowsDraftData = { - body?: never + body: DraftWorkflowSyncPayload path: { pipeline_id: string } @@ -1131,9 +1155,7 @@ export type PostRagPipelinesByPipelineIdWorkflowsDraftData = { } export type PostRagPipelinesByPipelineIdWorkflowsDraftResponses = { - 200: { - [key: string]: unknown - } + 200: RagPipelineWorkflowSyncResponse } export type PostRagPipelinesByPipelineIdWorkflowsDraftResponse = @@ -1428,7 +1450,7 @@ export type GetRagPipelinesByPipelineIdWorkflowsDraftVariablesByVariableIdRespon GetRagPipelinesByPipelineIdWorkflowsDraftVariablesByVariableIdResponses[keyof GetRagPipelinesByPipelineIdWorkflowsDraftVariablesByVariableIdResponses] export type PatchRagPipelinesByPipelineIdWorkflowsDraftVariablesByVariableIdData = { - body?: never + body: WorkflowDraftVariablePatchPayload path: { pipeline_id: string variable_id: string diff --git a/packages/contracts/generated/api/console/rag/zod.gen.ts b/packages/contracts/generated/api/console/rag/zod.gen.ts index a5afe6f494881d..3d92e1996df244 100644 --- a/packages/contracts/generated/api/console/rag/zod.gen.ts +++ b/packages/contracts/generated/api/console/rag/zod.gen.ts @@ -76,6 +76,27 @@ export const zDefaultBlockConfigsResponse = z.array(z.record(z.string(), z.unkno */ export const zDefaultBlockConfigResponse = z.record(z.string(), z.unknown()) +/** + * DraftWorkflowSyncPayload + */ +export const zDraftWorkflowSyncPayload = z.object({ + conversation_variables: z.array(z.record(z.string(), z.unknown())).nullish(), + environment_variables: z.array(z.record(z.string(), z.unknown())).nullish(), + features: z.record(z.string(), z.unknown()).nullish(), + graph: z.record(z.string(), z.unknown()), + hash: z.string().nullish(), + rag_pipeline_variables: z.array(z.record(z.string(), z.unknown())).nullish(), +}) + +/** + * RagPipelineWorkflowSyncResponse + */ +export const zRagPipelineWorkflowSyncResponse = z.object({ + hash: z.string(), + result: z.string(), + updated_at: z.int(), +}) + /** * DatasourceNodeRunPayload */ @@ -126,6 +147,14 @@ export const zDraftWorkflowRunPayload = z.object({ start_node_id: z.string(), }) +/** + * WorkflowDraftVariablePatchPayload + */ +export const zWorkflowDraftVariablePatchPayload = z.object({ + name: z.string().nullish(), + value: z.unknown().nullish(), +}) + /** * RagPipelineWorkflowPublishResponse */ @@ -156,15 +185,6 @@ export const zPublishedWorkflowRunPayload = z.object({ start_node_id: z.string(), }) -/** - * RagPipelineWorkflowSyncResponse - */ -export const zRagPipelineWorkflowSyncResponse = z.object({ - hash: z.string(), - result: z.string(), - updated_at: z.int(), -}) - /** * ImportStatus */ @@ -1169,6 +1189,8 @@ export const zGetRagPipelinesByPipelineIdWorkflowsDraftPath = z.object({ */ export const zGetRagPipelinesByPipelineIdWorkflowsDraftResponse = zWorkflowResponse +export const zPostRagPipelinesByPipelineIdWorkflowsDraftBody = zDraftWorkflowSyncPayload + export const zPostRagPipelinesByPipelineIdWorkflowsDraftPath = z.object({ pipeline_id: z.uuid(), }) @@ -1176,7 +1198,7 @@ export const zPostRagPipelinesByPipelineIdWorkflowsDraftPath = z.object({ /** * Success */ -export const zPostRagPipelinesByPipelineIdWorkflowsDraftResponse = z.record(z.string(), z.unknown()) +export const zPostRagPipelinesByPipelineIdWorkflowsDraftResponse = zRagPipelineWorkflowSyncResponse export const zPostRagPipelinesByPipelineIdWorkflowsDraftDatasourceNodesByNodeIdRunBody = zDatasourceNodeRunPayload @@ -1382,6 +1404,9 @@ export const zGetRagPipelinesByPipelineIdWorkflowsDraftVariablesByVariableIdPath export const zGetRagPipelinesByPipelineIdWorkflowsDraftVariablesByVariableIdResponse = zWorkflowDraftVariableResponse +export const zPatchRagPipelinesByPipelineIdWorkflowsDraftVariablesByVariableIdBody = + zWorkflowDraftVariablePatchPayload + export const zPatchRagPipelinesByPipelineIdWorkflowsDraftVariablesByVariableIdPath = z.object({ pipeline_id: z.uuid(), variable_id: z.uuid(), diff --git a/packages/contracts/generated/api/console/snippets/orpc.gen.ts b/packages/contracts/generated/api/console/snippets/orpc.gen.ts index 4e664987a77a71..219e988ec2b34a 100644 --- a/packages/contracts/generated/api/console/snippets/orpc.gen.ts +++ b/packages/contracts/generated/api/console/snippets/orpc.gen.ts @@ -877,7 +877,7 @@ export const restore = { } /** - * Delete a published snippet workflow version + * Delete a version and retire its unowned Agents through the shared use case * * Delete a published snippet workflow version */ @@ -889,7 +889,7 @@ export const delete4 = oc operationId: 'deleteSnippetsBySnippetIdWorkflowsByWorkflowId', path: '/snippets/{snippet_id}/workflows/{workflow_id}', successStatus: 204, - summary: 'Delete a published snippet workflow version', + summary: 'Delete a version and retire its unowned Agents through the shared use case', tags: ['console'], }) .input(z.object({ params: zDeleteSnippetsBySnippetIdWorkflowsByWorkflowIdPath })) From 7584d952745179b952635c7db1f65ac58fbedc15 Mon Sep 17 00:00:00 2001 From: Byron Wang Date: Fri, 9 Oct 2026 14:23:25 +0800 Subject: [PATCH 3/5] test: use the active API provider after legacy retirement --- .../unit_tests/services/tools/test_api_provider_runtime.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/api/tests/unit_tests/services/tools/test_api_provider_runtime.py b/api/tests/unit_tests/services/tools/test_api_provider_runtime.py index d64f362dd2c0b9..1649a2372106ab 100644 --- a/api/tests/unit_tests/services/tools/test_api_provider_runtime.py +++ b/api/tests/unit_tests/services/tools/test_api_provider_runtime.py @@ -11,7 +11,6 @@ from sqlalchemy.pool import QueuePool from core.db import session_factory -from core.tools.custom_tool.provider import ApiToolProviderController from core.tools.entities.tool_bundle import ApiToolBundle from core.tools.entities.tool_entities import ApiProviderSchemaType from core.tools.errors import ToolInvokeError @@ -23,6 +22,7 @@ from models.tools import ApiToolProvider from services.entities.agent_tool_inner import AgentToolInvokeRequest from services.errors.agent_tool_inner import AgentToolInnerServiceError +from services.tools.api.provider import ApiToolProviderController from tests.unit_tests.model_factories import make_account from tests.unit_tests.services.agent.test_tool_invocation_service import ( TENANT_ID, From e0d363e4407e2078d840c93aa1374f01e89d3f37 Mon Sep 17 00:00:00 2001 From: Byron Wang Date: Fri, 9 Oct 2026 14:25:03 +0800 Subject: [PATCH 4/5] fix: align tool boundary rules with explicit session injection --- api/.importlinter | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/api/.importlinter b/api/.importlinter index 12e8d70ef9bcdd..4267578bee6b57 100644 --- a/api/.importlinter +++ b/api/.importlinter @@ -1599,10 +1599,12 @@ forbidden_modules = allow_indirect_imports = True [importlinter:contract:tool-engine-persistence-boundary] -name = Tool Engine receives Sessions but cannot source persistence dependencies +name = Tool invocation receives Sessions but cannot source persistence dependencies type = forbidden source_modules = services.tools.tool_engine + core.tools.__base.tool + services.tools.dataset.tool forbidden_modules = extensions.ext_database core.db @@ -1632,12 +1634,10 @@ source_modules = core.rag.rerank.rerank_model services.knowledge.retrieval.reranking services.knowledge_retrieval_inner_service - services.tools.dataset.tool services.app.generation.adapters.chat_runner services.app.generation.adapters.completion_runner services.workflow.execution.adapters.knowledge_retrieval services.workflow.execution.adapters.human_input_events - core.tools.__base.tool forbidden_modules = sqlalchemy repositories From 192450403e5ea3f4f9d4aeae6c7d35f9b9491cc9 Mon Sep 17 00:00:00 2001 From: Byron Wang Date: Mon, 26 Oct 2026 23:57:47 +0800 Subject: [PATCH 5/5] Preserve publication and webhook boundaries after legacy retirement --- api/controllers/trigger/webhook.py | 3 - .../application_services/workflow.py | 29 +++++- .../workflow/definition_repository.py | 27 +++++- api/services/trigger/webhook_service.py | 19 +++- .../controllers/trigger/test_webhook.py | 40 ++++++-- .../core/tools/utils/test_uuid_utils.py | 6 +- .../workflow/test_definition_repository.py | 96 ++++++++++++------- .../services/trigger/test_webhook_service.py | 24 ++++- 8 files changed, 192 insertions(+), 52 deletions(-) diff --git a/api/controllers/trigger/webhook.py b/api/controllers/trigger/webhook.py index c753e3553e60d5..b2a4018cd2a4a7 100644 --- a/api/controllers/trigger/webhook.py +++ b/api/controllers/trigger/webhook.py @@ -2,7 +2,6 @@ import time from flask import jsonify, request -from werkzeug.exceptions import RequestEntityTooLarge from controllers.common.access_response import webhook_not_found_response from controllers.common.errors import RequestBodyTooLargeError @@ -164,8 +163,6 @@ def handle_webhook_debug(webhook_id: str): return webhook_not_found_response() except WebhookBodyTooLargeError as error: raise RequestBodyTooLargeError(str(error)) from error - except RequestEntityTooLarge: - raise except Exception: logger.exception("Webhook debug processing failed for %s", webhook_id) return jsonify({"error": "Internal server error", "message": "An internal error has occurred."}), 500 diff --git a/api/extensions/application_services/workflow.py b/api/extensions/application_services/workflow.py index da933383c68830..c81c4e316f69ea 100644 --- a/api/extensions/application_services/workflow.py +++ b/api/extensions/application_services/workflow.py @@ -29,6 +29,7 @@ from services.rbac.members import MemberService from services.tools.provider_queries import ToolProviders from services.tools.tool_engine import ToolEngine +from services.workflow.contracts import WorkflowSnapshot from services.workflow.execution.chatflow_ports import ChatflowRecords, ConversationVariables from services.workflow.execution.contracts import WorkflowRunDriver from services.workflow.execution.node_queries import ConversationHistory, DatasourceCredentials, RetrieverAttachments @@ -222,6 +223,23 @@ def build_workflow_execution_dependencies(database_client: sessionmaker[Session] from services.workflow_service import WorkflowService +def _sync_publication_knowledge_spaces( + *, + session: Session, + session_maker: sessionmaker[Session], + workflow: WorkflowSnapshot, +) -> None: + """Bind the KnowledgeFS publication effect without reversing repository dependencies.""" + from repositories.workflow.definition_repository import workflow_from_snapshot + from services.agent.knowledge_spaces import sync_workflow_app_knowledge_bindings + + sync_workflow_app_knowledge_bindings( + session=session, + session_maker=session_maker, + workflow=workflow_from_snapshot(workflow), + ) + + def build_console_workflow_service( *, database_client: sessionmaker[Session], @@ -253,7 +271,10 @@ def build_console_workflow_service( session_maker=database_client, runtime=build_workflow_execution_dependencies(database_client) ) - definitions = WorkflowDefinitionRepository(session_factory=database_client) + definitions = WorkflowDefinitionRepository( + session_factory=database_client, + publication_knowledge_spaces=_sync_publication_knowledge_spaces, + ) drafts = WorkflowDraftRepository(database_client) lifecycle = build_workflow_definition_gateway(database_client, definitions, workflows, drafts=drafts) apps = ConsoleAppRepository(session_factory=database_client) @@ -345,7 +366,11 @@ def build_workflow_drafts( from repositories.workflow.draft_repository import WorkflowDraftRepository from services.workflow_service import WorkflowService - definitions = WorkflowDefinitionRepository(session_factory=database_client, import_session=import_session) + definitions = WorkflowDefinitionRepository( + session_factory=database_client, + import_session=import_session, + publication_knowledge_spaces=_sync_publication_knowledge_spaces, + ) drafts = WorkflowDraftRepository(database_client, import_session=import_session) lifecycle = build_workflow_definition_gateway( database_client, diff --git a/api/repositories/workflow/definition_repository.py b/api/repositories/workflow/definition_repository.py index d926f5af0343ee..9633705637e103 100644 --- a/api/repositories/workflow/definition_repository.py +++ b/api/repositories/workflow/definition_repository.py @@ -10,7 +10,7 @@ from collections.abc import Generator, Mapping, Sequence from contextlib import contextmanager, nullcontext from dataclasses import dataclass -from typing import Any +from typing import Any, Protocol from sqlalchemy import delete, select from sqlalchemy.orm import Session, sessionmaker @@ -709,6 +709,18 @@ def workflow_from_snapshot(snapshot: WorkflowSnapshot) -> Workflow: ) +class WorkflowKnowledgeSpaceSync(Protocol): + """Publication effect supplied by the application composition root.""" + + def __call__( + self, + *, + session: Session, + session_maker: sessionmaker[Session], + workflow: WorkflowSnapshot, + ) -> None: ... + + @dataclass class _PublicationTransaction: session: Session @@ -717,6 +729,7 @@ class _PublicationTransaction: actor: Account draft: WorkflowSnapshot bindings: WorkflowAgentBindingRepository + knowledge_spaces: WorkflowKnowledgeSpaceSync def create_version(self, *, marked_name: str, marked_comment: str) -> WorkflowSnapshot: workflow = WorkflowDefinitionStore.publish_workflow( @@ -746,12 +759,10 @@ def sync_datasets(self, dataset_ids: set[str]) -> None: WorkflowDefinitionStore.sync_dataset_relationships(self.session, self.app, dataset_ids) def sync_knowledge_spaces(self, workflow: WorkflowSnapshot) -> None: - from services.agent.knowledge_spaces import sync_workflow_app_knowledge_bindings - - sync_workflow_app_knowledge_bindings( + self.knowledge_spaces( session=self.session, session_maker=self.session_maker, - workflow=workflow_from_snapshot(workflow), + workflow=workflow, ) def activate(self, workflow: WorkflowSnapshot) -> None: @@ -764,9 +775,11 @@ def __init__( *, session_factory: sessionmaker[Session], import_session: Session | None = None, + publication_knowledge_spaces: WorkflowKnowledgeSpaceSync | None = None, ) -> None: self._sessions = session_factory self._import_session = import_session + self._publication_knowledge_spaces = publication_knowledge_spaces def debug_context( self, @@ -847,6 +860,9 @@ def publication( ) -> Generator[WorkflowPublicationTransaction[WorkflowAgentBindingRepository]]: from repositories.app.console_repository import console_app_actor + knowledge_spaces = self._publication_knowledge_spaces + if knowledge_spaces is None: + raise RuntimeError("workflow publication requires a KnowledgeFS synchronization callback") with self._sessions.begin() as session: app = WorkflowDefinitionStore.lock_app(session, context, app_id) current = WorkflowDefinitionStore.get_draft_workflow(app, session=session, for_update=True) @@ -859,6 +875,7 @@ def publication( console_app_actor(session, context), workflow_snapshot(current), WorkflowAgentBindingRepository(session), + knowledge_spaces, ) def debug_node(self, context: RequestContext, app_id: str, node_id: str) -> WorkflowDebugNode: diff --git a/api/services/trigger/webhook_service.py b/api/services/trigger/webhook_service.py index 241955835652f8..c4e9cf5b3236be 100644 --- a/api/services/trigger/webhook_service.py +++ b/api/services/trigger/webhook_service.py @@ -9,6 +9,7 @@ from sqlalchemy import select from sqlalchemy.orm import Session, sessionmaker from werkzeug.datastructures import FileStorage +from werkzeug.exceptions import RequestEntityTooLarge from configs import dify_config from core.app.file_access import DatabaseFileAccessController @@ -248,7 +249,19 @@ def extract_webhook_data(cls, webhook_trigger: WorkflowWebhookTrigger) -> RawWeb extractor = cls._extract_text_body # Extract body and files - body_data, files_data = extractor() + try: + body_data, files_data = extractor() + except RequestEntityTooLarge as error: + content_length = request.content_length + maximum_size = request.max_content_length + if content_length is not None and maximum_size is not None: + message = ( + f"Webhook request too large: {content_length} bytes exceeds maximum allowed size " + f"of {maximum_size} bytes" + ) + else: + message = "Webhook request too large" + raise WebhookBodyTooLargeError(message) from error data["body"] = body_data data["files"] = files_data @@ -363,6 +376,8 @@ def _extract_octet_stream_body( return {"raw": file_obj.to_dict()}, {} else: return {"raw": None}, {} + except RequestEntityTooLarge: + raise except Exception: logger.exception("Failed to process octet-stream data") return {"raw": None}, {} @@ -378,6 +393,8 @@ def _extract_text_body(cls) -> tuple[dict[str, Any], dict[str, Any]]: """ try: body = {"raw": request.get_data(as_text=True)} + except RequestEntityTooLarge: + raise except Exception: logger.warning("Failed to extract text body") body = {"raw": ""} diff --git a/api/tests/unit_tests/controllers/trigger/test_webhook.py b/api/tests/unit_tests/controllers/trigger/test_webhook.py index d2fc8d90144bcc..bf77948203a607 100644 --- a/api/tests/unit_tests/controllers/trigger/test_webhook.py +++ b/api/tests/unit_tests/controllers/trigger/test_webhook.py @@ -4,12 +4,14 @@ import pytest from flask import Flask -from werkzeug.exceptions import NotFound, RequestEntityTooLarge +from werkzeug.exceptions import NotFound import controllers.trigger.webhook as module +from controllers.common.errors import RequestBodyTooLargeError from models.trigger import WorkflowWebhookTrigger from models.workflow import Workflow from services.errors.app import QuotaExceededError +from services.trigger.errors import WebhookBodyTooLargeError from tests.helpers.unit.model_factories import make_workflow @@ -214,11 +216,23 @@ def test_value_error_not_found(self, mock_get): assert response.status_code == 404 assert response.data == NotFound().get_body().encode() - @patch.object(module.WebhookService, "get_webhook_trigger_and_workflow", side_effect=RequestEntityTooLarge()) - def test_request_entity_too_large(self, mock_get): - with pytest.raises(RequestEntityTooLarge): + @patch.object( + module.WebhookService, + "get_webhook_trigger_and_workflow", + side_effect=WebhookBodyTooLargeError( + "Webhook request too large: 9 bytes exceeds maximum allowed size of 8 bytes" + ), + ) + def test_webhook_body_too_large(self, mock_get): + with pytest.raises(RequestBodyTooLargeError) as error_info: module.handle_webhook("wh-1") + assert error_info.value.data == { + "code": "request_entity_too_large", + "message": "Webhook request too large: 9 bytes exceeds maximum allowed size of 8 bytes", + "status": 413, + } + @patch.object(module.WebhookService, "get_webhook_trigger_and_workflow", side_effect=Exception("boom")) def test_internal_error(self, mock_get): app = Flask(__name__) @@ -295,11 +309,23 @@ def test_debug_not_found(self, mock_get): assert response.status_code == 404 assert response.data == NotFound().get_body().encode() - @patch.object(module.WebhookService, "get_webhook_trigger_and_workflow", side_effect=RequestEntityTooLarge()) - def test_debug_request_entity_too_large(self, mock_get): - with pytest.raises(RequestEntityTooLarge): + @patch.object( + module.WebhookService, + "get_webhook_trigger_and_workflow", + side_effect=WebhookBodyTooLargeError( + "Webhook request too large: 9 bytes exceeds maximum allowed size of 8 bytes" + ), + ) + def test_debug_webhook_body_too_large(self, mock_get): + with pytest.raises(RequestBodyTooLargeError) as error_info: module.handle_webhook_debug("wh-1") + assert error_info.value.data == { + "code": "request_entity_too_large", + "message": "Webhook request too large: 9 bytes exceeds maximum allowed size of 8 bytes", + "status": 413, + } + @patch.object(module.WebhookService, "get_webhook_trigger_and_workflow", side_effect=Exception("boom")) def test_debug_internal_error(self, mock_get): app = Flask(__name__) diff --git a/api/tests/unit_tests/core/tools/utils/test_uuid_utils.py b/api/tests/unit_tests/core/tools/utils/test_uuid_utils.py index 3c02c40838480a..9d218af5290d98 100644 --- a/api/tests/unit_tests/core/tools/utils/test_uuid_utils.py +++ b/api/tests/unit_tests/core/tools/utils/test_uuid_utils.py @@ -8,7 +8,7 @@ @pytest.mark.parametrize( ("value", "expected"), [ - (str(uuid.uuid4()), True), + ("00000000-0000-4000-8000-000000000001", True), ("not-a-uuid", False), ("", False), (None, False), @@ -16,3 +16,7 @@ ) def test_is_valid_uuid(value: str | None, expected: bool) -> None: assert is_valid_uuid(value) is expected + + +def test_is_valid_uuid_accepts_runtime_generated_uuid4() -> None: + assert is_valid_uuid(str(uuid.uuid4())) is True diff --git a/api/tests/unit_tests/repositories/workflow/test_definition_repository.py b/api/tests/unit_tests/repositories/workflow/test_definition_repository.py index 547e8474dcbd8d..6f7bf900553bd1 100644 --- a/api/tests/unit_tests/repositories/workflow/test_definition_repository.py +++ b/api/tests/unit_tests/repositories/workflow/test_definition_repository.py @@ -77,6 +77,23 @@ CONTEXT = RequestContext("request", None, "account-1", "tenant-1") +def _ignore_knowledge_space_sync( + *, + session: Session, + session_maker: sessionmaker[Session], + workflow: WorkflowSnapshot, +) -> None: + del session, session_maker, workflow + + +def publication_repository(sessions: sessionmaker[Session]) -> WorkflowDefinitionRepository: + """Compose a repository for publication tests that do not exercise KnowledgeFS.""" + return WorkflowDefinitionRepository( + session_factory=sessions, + publication_knowledge_spaces=_ignore_knowledge_space_sync, + ) + + @pytest.mark.parametrize("kind", ["tool", "agent", "llm"]) @pytest.mark.parametrize("denied", [False, True]) def test_publication_credential_policy_runs_after_read_transactions_close( @@ -135,7 +152,7 @@ def check(credential_id: str, *_args: object, **_kwargs: object) -> None: lifecycle = WorkflowDefinitionGateway( session_factory=seeded, drafts=WorkflowDraftRepository(seeded), - definitions=WorkflowDefinitionRepository(session_factory=seeded), + definitions=publication_repository(seeded), workflows=workflows, credentials=CredentialQueryRepository(session_factory=seeded), skills=create_autospec(AgentSkillReader, instance=True), @@ -174,7 +191,7 @@ def seeded(sqlite_session_factory: sessionmaker[Session]) -> sessionmaker[Sessio @pytest.fixture def repository(seeded: sessionmaker[Session]) -> WorkflowDefinitionRepository: - return WorkflowDefinitionRepository(session_factory=seeded) + return publication_repository(seeded) def test_published_result_is_complete_after_session_closes(repository: WorkflowDefinitionRepository) -> None: @@ -297,7 +314,7 @@ def wrap(**_kwargs: object) -> str: app.mode = AppMode.ADVANCED_CHAT workflows = WorkflowService(seeded) monkeypatch.setattr(workflows, "validate_graph_structure", lambda **_kwargs: None) - definitions = WorkflowDefinitionRepository(session_factory=seeded) + definitions = publication_repository(seeded) service = application( build_workflow_definition_gateway(seeded, definitions, workflows, drafts=WorkflowDraftRepository(seeded)) ) @@ -376,7 +393,7 @@ def encrypt(*, tenant_id: str, token: str) -> str: workflows = WorkflowService(seeded) monkeypatch.setattr(workflows, "validate_graph_structure", lambda *_args, **_kwargs: None) monkeypatch.setattr(workflows, "validate_features_structure", lambda *_args: None) - definitions = WorkflowDefinitionRepository(session_factory=seeded) + definitions = publication_repository(seeded) service = application( build_workflow_definition_gateway(seeded, definitions, workflows, drafts=WorkflowDraftRepository(seeded)) ) @@ -526,7 +543,7 @@ def test_workflow_and_bindings_delete_atomically_before_agent_retirement( ("roster", WorkflowAgentBindingType.ROSTER_AGENT, "roster-agent"), ] ) - definitions = WorkflowDefinitionRepository(session_factory=seeded) + definitions = publication_repository(seeded) service = application( build_workflow_definition_gateway( seeded, definitions, WorkflowService(seeded), drafts=WorkflowDraftRepository(seeded) @@ -620,7 +637,7 @@ def test_draft_projection_does_not_mutate_persisted_graph( "project_draft_bindings_to_graph", lambda _self, **_kwargs: {"nodes": [{"id": "projected"}]}, ) - repository = WorkflowDefinitionRepository(session_factory=seeded) + repository = publication_repository(seeded) with repository.draft(CONTEXT, "app-1") as draft: assert draft is not None projected = WorkflowAgentPublishService(repository=draft.bindings).project_draft_bindings_to_graph( @@ -679,7 +696,7 @@ def test_publication_and_app_pointer_commit_atomically(seeded: sessionmaker[Sess lifecycle = WorkflowDefinitionGateway( session_factory=seeded, drafts=WorkflowDraftRepository(seeded), - definitions=WorkflowDefinitionRepository(session_factory=seeded), + definitions=publication_repository(seeded), workflows=WorkflowService(session_maker=seeded), credentials=CredentialQueryRepository(session_factory=seeded), skills=create_autospec(AgentSkillReader, instance=True), @@ -745,7 +762,7 @@ def test_sync_prepares_patch_without_replacing_server_environment(seeded: sessio WorkflowDefinitionGateway( session_factory=seeded, drafts=WorkflowDraftRepository(seeded), - definitions=WorkflowDefinitionRepository(session_factory=seeded), + definitions=publication_repository(seeded), workflows=WorkflowService(session_maker=seeded), credentials=CredentialQueryRepository(session_factory=seeded), skills=create_autospec(AgentSkillReader, instance=True), @@ -760,7 +777,7 @@ def test_sync_prepares_patch_without_replacing_server_environment(seeded: sessio WorkflowDefinitionGateway( session_factory=seeded, drafts=WorkflowDraftRepository(seeded), - definitions=WorkflowDefinitionRepository(session_factory=seeded), + definitions=publication_repository(seeded), workflows=WorkflowService(session_maker=seeded), credentials=CredentialQueryRepository(session_factory=seeded), skills=create_autospec(AgentSkillReader, instance=True), @@ -789,11 +806,11 @@ def test_app_maintainers_reuses_tenant_scoped_app_repository(seeded: sessionmake def test_use_case_notifies_only_after_real_transaction_commits( seeded: sessionmaker[Session], monkeypatch: pytest.MonkeyPatch, operation: str, fail_commit: bool ) -> None: - repository = WorkflowDefinitionRepository(session_factory=seeded) + repository = publication_repository(seeded) lifecycle = WorkflowDefinitionGateway( session_factory=seeded, drafts=WorkflowDraftRepository(seeded), - definitions=WorkflowDefinitionRepository(session_factory=seeded), + definitions=publication_repository(seeded), workflows=WorkflowService(session_maker=seeded), credentials=CredentialQueryRepository(session_factory=seeded), skills=create_autospec(AgentSkillReader, instance=True), @@ -937,7 +954,7 @@ def test_publish_rejects_revision_modified_after_validation(seeded: sessionmaker lifecycle = WorkflowDefinitionGateway( session_factory=seeded, drafts=WorkflowDraftRepository(seeded), - definitions=WorkflowDefinitionRepository(session_factory=seeded), + definitions=publication_repository(seeded), workflows=WorkflowService(session_maker=seeded), credentials=CredentialQueryRepository(session_factory=seeded), skills=create_autospec(AgentSkillReader, instance=True), @@ -972,7 +989,7 @@ def test_trigger_limit_rejects_publish_before_version_or_pointer_changes( lifecycle = WorkflowDefinitionGateway( session_factory=seeded, drafts=WorkflowDraftRepository(seeded), - definitions=WorkflowDefinitionRepository(session_factory=seeded), + definitions=publication_repository(seeded), workflows=WorkflowService(session_maker=seeded), credentials=CredentialQueryRepository(session_factory=seeded), skills=create_autospec(AgentSkillReader, instance=True), @@ -993,7 +1010,7 @@ def test_publication_versions_increase_without_reusing_deleted_version(seeded: s lifecycle = WorkflowDefinitionGateway( session_factory=seeded, drafts=WorkflowDraftRepository(seeded), - definitions=WorkflowDefinitionRepository(session_factory=seeded), + definitions=publication_repository(seeded), workflows=WorkflowService(session_maker=seeded), credentials=CredentialQueryRepository(session_factory=seeded), skills=create_autospec(AgentSkillReader, instance=True), @@ -1087,7 +1104,7 @@ def test_publication_retains_blocking_domain_validations( lifecycle = WorkflowDefinitionGateway( session_factory=seeded, drafts=WorkflowDraftRepository(seeded), - definitions=WorkflowDefinitionRepository(session_factory=seeded), + definitions=publication_repository(seeded), workflows=WorkflowService(session_maker=seeded), credentials=CredentialQueryRepository(session_factory=seeded), skills=create_autospec(AgentSkillReader, instance=True), @@ -1119,7 +1136,7 @@ def test_restore_preserves_historical_features_without_normalizing_persisted_sou lifecycle = WorkflowDefinitionGateway( session_factory=seeded, drafts=WorkflowDraftRepository(seeded), - definitions=WorkflowDefinitionRepository(session_factory=seeded), + definitions=publication_repository(seeded), workflows=WorkflowService(session_maker=seeded), credentials=CredentialQueryRepository(session_factory=seeded), skills=create_autospec(AgentSkillReader, instance=True), @@ -1147,7 +1164,7 @@ def capture(execution: ORMExecuteState) -> None: lifecycle = WorkflowDefinitionGateway( session_factory=seeded, drafts=WorkflowDraftRepository(seeded), - definitions=WorkflowDefinitionRepository(session_factory=seeded), + definitions=publication_repository(seeded), workflows=WorkflowService(session_maker=seeded), credentials=CredentialQueryRepository(session_factory=seeded), skills=create_autospec(AgentSkillReader, instance=True), @@ -1183,7 +1200,7 @@ def test_migration_uses_the_same_atomic_publication_boundary( lifecycle = WorkflowDefinitionGateway( session_factory=seeded, drafts=WorkflowDraftRepository(seeded), - definitions=WorkflowDefinitionRepository(session_factory=seeded), + definitions=publication_repository(seeded), workflows=WorkflowService(session_maker=seeded), credentials=CredentialQueryRepository(session_factory=seeded), skills=create_autospec(AgentSkillReader, instance=True), @@ -1286,7 +1303,7 @@ def test_redis_outage_cannot_leave_a_successful_publication_without_runtime_conf lifecycle = WorkflowDefinitionGateway( session_factory=seeded, drafts=WorkflowDraftRepository(seeded), - definitions=WorkflowDefinitionRepository(session_factory=seeded), + definitions=publication_repository(seeded), workflows=WorkflowService(session_maker=seeded), credentials=CredentialQueryRepository(session_factory=seeded), skills=create_autospec(AgentSkillReader, instance=True), @@ -1309,7 +1326,7 @@ def test_trigger_storage_failure_rolls_back_version_pointer_and_every_projection lifecycle = WorkflowDefinitionGateway( session_factory=seeded, drafts=WorkflowDraftRepository(seeded), - definitions=WorkflowDefinitionRepository(session_factory=seeded), + definitions=publication_repository(seeded), workflows=WorkflowService(session_maker=seeded), credentials=CredentialQueryRepository(session_factory=seeded), skills=create_autospec(AgentSkillReader, instance=True), @@ -1345,7 +1362,7 @@ def test_republication_replaces_all_trigger_configuration( lifecycle = WorkflowDefinitionGateway( session_factory=seeded, drafts=WorkflowDraftRepository(seeded), - definitions=WorkflowDefinitionRepository(session_factory=seeded), + definitions=publication_repository(seeded), workflows=WorkflowService(session_maker=seeded), credentials=CredentialQueryRepository(session_factory=seeded), skills=create_autospec(AgentSkillReader, instance=True), @@ -1374,7 +1391,7 @@ def test_draft_webhook_receiver_uses_injected_database( lifecycle = WorkflowDefinitionGateway( session_factory=seeded, drafts=WorkflowDraftRepository(seeded), - definitions=WorkflowDefinitionRepository(session_factory=seeded), + definitions=publication_repository(seeded), workflows=WorkflowService(session_maker=seeded), credentials=CredentialQueryRepository(session_factory=seeded), skills=create_autospec(AgentSkillReader, instance=True), @@ -1528,7 +1545,7 @@ def test_republish_without_schedule_removes_the_runtime_plan(seeded: sessionmake lifecycle = WorkflowDefinitionGateway( session_factory=seeded, drafts=WorkflowDraftRepository(seeded), - definitions=WorkflowDefinitionRepository(session_factory=seeded), + definitions=publication_repository(seeded), workflows=WorkflowService(session_maker=seeded), credentials=CredentialQueryRepository(session_factory=seeded), skills=create_autospec(AgentSkillReader, instance=True), @@ -1545,9 +1562,23 @@ def test_republish_without_schedule_removes_the_runtime_plan(seeded: sessionmake assert session.scalar(select(WorkflowSchedulePlan).where(WorkflowSchedulePlan.app_id == "app-1")) is None +def test_publication_requires_knowledge_space_sync(seeded: sessionmaker[Session]) -> None: + with seeded() as session: + draft = session.get(Workflow, "draft") + assert draft is not None + snapshot = workflow_snapshot(draft) + + repository = WorkflowDefinitionRepository( + session_factory=seeded, + ) + with pytest.raises(RuntimeError, match="requires a KnowledgeFS synchronization callback"): + with repository.publication(CONTEXT, "app-1", snapshot): + pytest.fail("publication transaction must not start without its required effect") + + @pytest.mark.parametrize("rollback_at", [None, "knowledge", "after-activation"]) def test_agent_bindings_knowledge_spaces_and_active_version_share_the_publication_transaction( - seeded: sessionmaker[Session], monkeypatch: pytest.MonkeyPatch, rollback_at: str | None + seeded: sessionmaker[Session], rollback_at: str | None ) -> None: from models.agent import Agent, AgentConfigSnapshot, AgentScope, AgentSource from models.agent_config_entities import AgentSoulConfig, AgentSoulModelConfig @@ -1612,11 +1643,12 @@ def test_agent_bindings_knowledge_spaces_and_active_version_share_the_publicatio draft = session.get(Workflow, "draft") assert draft is not None snapshot = workflow_snapshot(draft) - repository = WorkflowDefinitionRepository(session_factory=seeded) published_id = "" knowledge_sync_sessions: list[Session] = [] - def sync_knowledge_spaces(*, session: Session, session_maker: sessionmaker[Session], workflow: Workflow) -> None: + def sync_knowledge_spaces( + *, session: Session, session_maker: sessionmaker[Session], workflow: WorkflowSnapshot + ) -> None: nonlocal published_id assert session_maker is seeded published_id = workflow.id @@ -1636,9 +1668,9 @@ def sync_knowledge_spaces(*, session: Session, session_maker: sessionmaker[Sessi if rollback_at == "knowledge": raise RuntimeError("rollback during KnowledgeFS synchronization") - monkeypatch.setattr( - "services.agent.knowledge_spaces.sync_workflow_app_knowledge_bindings", - sync_knowledge_spaces, + repository = WorkflowDefinitionRepository( + session_factory=seeded, + publication_knowledge_spaces=sync_knowledge_spaces, ) with repository.binding_reader() as bindings: @@ -1827,7 +1859,7 @@ def load(key: str) -> bytes: app.mode = AppMode.ADVANCED_CHAT workflows = WorkflowService(seeded) monkeypatch.setattr(workflows, "validate_graph_structure", lambda **_kwargs: None) - definitions = WorkflowDefinitionRepository(session_factory=seeded) + definitions = publication_repository(seeded) service = application( build_workflow_definition_gateway(seeded, definitions, workflows, drafts=WorkflowDraftRepository(seeded)) ) @@ -1952,7 +1984,7 @@ def test_shared_draft_writer_preserves_hash_and_variable_rules(seeded: sessionma previous_hash = draft.unique_hash if state == "new": session.delete(draft) - definitions = WorkflowDefinitionRepository(session_factory=seeded) + definitions = publication_repository(seeded) lifecycle = build_workflow_definition_gateway( seeded, definitions, WorkflowService(seeded), drafts=WorkflowDraftRepository(seeded) ) @@ -2021,7 +2053,7 @@ def test_publication_repository_is_independent_of_agent_service_and_commits_requ assert row.is_published is True assert session.get(Workflow, workflow.id) is not None else: - definitions = WorkflowDefinitionRepository(session_factory=seeded) + definitions = publication_repository(seeded) with definitions.draft(CONTEXT, "app-1") as read: assert read is not None assert isinstance(read.bindings, WorkflowAgentBindingRepository) diff --git a/api/tests/unit_tests/services/trigger/test_webhook_service.py b/api/tests/unit_tests/services/trigger/test_webhook_service.py index e226c9d9098aa3..5ebb4d64203b9a 100644 --- a/api/tests/unit_tests/services/trigger/test_webhook_service.py +++ b/api/tests/unit_tests/services/trigger/test_webhook_service.py @@ -10,6 +10,7 @@ from sqlalchemy.engine import Engine as _webhook_service_Engine from sqlalchemy.orm import Session, sessionmaker from werkzeug.datastructures import FileStorage +from werkzeug.exceptions import RequestEntityTooLarge from core.workflow.nodes.trigger_webhook.entities import ( ContentType, @@ -1026,9 +1027,30 @@ def test_extract_webhook_data_should_raise_for_request_too_large( apply_config_overrides(monkeypatch, WEBHOOK_REQUEST_BODY_MAX_SIZE=1) with flask_app.test_request_context("/webhook", method="POST", data="ab"): - with pytest.raises(WebhookBodyTooLargeError): + with pytest.raises(WebhookBodyTooLargeError) as error_info: WebhookService.extract_webhook_data(_workflow_trigger()) + assert str(error_info.value) == ( + "Webhook request too large: 2 bytes exceeds maximum allowed size of 1 bytes" + ) + + def test_extract_webhook_data_translates_framework_body_limit(self, flask_app: Flask) -> None: + flask_app.config["MAX_CONTENT_LENGTH"] = 1 + + with flask_app.test_request_context( + "/webhook", + method="POST", + content_type="text/plain", + data="ab", + ): + with pytest.raises(WebhookBodyTooLargeError) as error_info: + WebhookService.extract_webhook_data(_workflow_trigger()) + + assert isinstance(error_info.value.__cause__, RequestEntityTooLarge) + assert str(error_info.value) == ( + "Webhook request too large: 2 bytes exceeds maximum allowed size of 1 bytes" + ) + def test_extract_octet_stream_body_should_return_none_when_empty_payload(self, flask_app: Flask) -> None: webhook_trigger = _workflow_trigger()