diff --git a/api/.importlinter b/api/.importlinter index 55a332cdb36f8b..81f386b66a4449 100644 --- a/api/.importlinter +++ b/api/.importlinter @@ -1,14 +1,16 @@ [importlinter] root_packages = - services.entities.external_knowledge_entities - enums core + enums errors core.rag.extractor.entity + core.tools.__base + services.entities.external_knowledge_entities constants context configs controllers + controllers.console.explore extensions factories fields @@ -48,9 +50,6 @@ 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 @@ -1264,6 +1263,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 @@ -1283,23 +1462,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 @@ -1328,47 +1490,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] @@ -1388,6 +1637,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 @@ -1403,25 +1669,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 @@ -1430,23 +1698,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] @@ -1467,10 +1791,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 @@ -1492,44 +1818,28 @@ 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.app.generation.adapters.chat_runner + services.app.generation.adapters.completion_runner + services.workflow.execution.adapters.knowledge_retrieval + services.workflow.execution.adapters.human_input_events 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 - [importlinter:contract:annotation-application-boundary] name = Annotation application services and contracts are framework and persistence neutral type = forbidden 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/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/rag/retrieval/dataset_retrieval.py b/api/core/rag/retrieval/dataset_retrieval.py deleted file mode 100644 index ee749a9806d34b..00000000000000 --- a/api/core/rag/retrieval/dataset_retrieval.py +++ /dev/null @@ -1,2146 +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, 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.retrieval_source_repository import get_retrieval_source_documents -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 = get_retrieval_source_documents( - tenant_id=request.tenant_id, - dataset_ids=available_datasets_ids, - document_ids=document_ids, - session=retrieval_session, - ) - - 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) - - inputs = {key: str(value) for key, value in inputs.items()} if inputs else {} - 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 - final_content = f"{record.summary}\n{segment_content}" if record.summary else 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.is_(True), - DatasetDocument.archived.is_(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: - # top_k=0 is a legal explicit value (the SINGLE-retrieval - # tool path reads the stored top_k as-is); only a missing - # entry falls back to the default. - dataset_top_k = retrieval_model.get("top_k") - # retrieval source - documents = RetrievalService.retrieve( - retrieval_method=retrieval_model["search_method"], - dataset_id=dataset.id, - query=query, - top_k=dataset_top_k if dataset_top_k is not None else 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=0 is a legal explicit value (the non-tool MULTIPLE path - # honors it); only None falls back to the default. - top_k=retrieve_config.top_k if retrieve_config.top_k is not None else 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, strict=False): - # 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.is_(True), - DatasetDocument.archived.is_(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], - 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( - or_( - json_field.is_(None), - 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(or_(json_field.is_(None), json_field != value)) - case int() | float(): - filters.append( - or_(json_field.is_(None), 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")) - elif condition == "in": - filters.append(json_field.in_(value_list)) - else: - filters.append(or_(json_field.is_(None), json_field.notin_(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: 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.is_(True), - DocumentModel.archived.is_(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 5c0f0b2e7760e3..00000000000000 --- a/api/core/repositories/human_input_repository.py +++ /dev/null @@ -1,702 +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, update -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 DeliveryMethodType, 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, - EmailExternalRecipientPayload, - EmailMemberRecipientPayload, - HumanInputDelivery, - HumanInputForm, - HumanInputFormRecipient, - InstantMessageRecipientPayload, - StandaloneWebAppRecipientPayload, -) -from models.human_input_contracts import ( - FormCreateParams, - FormNotFoundError, - HumanInputFormRecipientEntity, -) -from models.human_input_contracts import ( - HumanInputFormEntity as SharedHumanInputFormEntity, -) -from models.human_input_contracts import ( - HumanInputFormRecord as SharedHumanInputFormRecord, -) -from models.human_input_contracts import ( - HumanInputFormRepository as SharedHumanInputFormRepository, -) -from models.human_input_delivery import ( - BoundRecipient, - DeliveryChannelConfig, - EmailDeliveryMethod, - EmailRecipients, - ExternalRecipient, - InstantMessageChannelRecipient, - InstantMessageDeliveryConfig, - InstantMessageDeliveryMethod, - InstantMessageUserRecipient, - InteractiveSurfaceDeliveryMethod, -) -from models.human_input_entities import FormDefinition, HumanInputNodeData -from repositories.human_input_errors import FormAlreadyHandledError - - -@dataclasses.dataclass(frozen=True) -class _DeliveryAndRecipients: - delivery: HumanInputDelivery - recipients: Sequence[HumanInputFormRecipient] - - -@dataclasses.dataclass(frozen=True) -class _WorkspaceMemberInfo: - user_id: str - email: str - - -class HumanInputFormEntity(SharedHumanInputFormEntity, Protocol): - """Legacy form view with the submission token used by core callers.""" - - @property - def submission_token(self) -> str | None: ... - - -class HumanInputFormRepository(SharedHumanInputFormRepository, Protocol): - @override - def get_form(self, node_id: str, *, form_id: str | None = None) -> HumanInputFormEntity | None: ... - - @override - def create_form(self, params: FormCreateParams) -> HumanInputFormEntity: ... - - def mark_timeout(self, node_id: str, *, form_id: str) -> HumanInputFormEntity: - """Time out an overdue waiting form and return its current state.""" - ... - - -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(SharedHumanInputFormRecord): - """Legacy constructor kept on the shared human-input record contract.""" - - @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, - ) - ) - case InstantMessageDeliveryMethod(): - recipients.extend( - self._build_instant_message_recipients( - form_id=form_id, - delivery_id=delivery_id, - config=delivery_method.config, - ) - ) - - return _DeliveryAndRecipients(delivery=delivery_model, recipients=recipients) - - @staticmethod - def _build_instant_message_recipients( - *, - form_id: str, - delivery_id: str, - config: InstantMessageDeliveryConfig, - ) -> list[HumanInputFormRecipient]: - recipient_models: list[HumanInputFormRecipient] = [] - for recipient in config.recipients.items: - match recipient: - case InstantMessageChannelRecipient(): - payload = InstantMessageRecipientPayload( - provider=config.provider, - recipient_kind=recipient.type, - channel_id=recipient.channel_id, - ) - case InstantMessageUserRecipient(): - payload = InstantMessageRecipientPayload( - provider=config.provider, - recipient_kind=recipient.type, - user_id=recipient.user_id, - ) - recipient_models.append( - HumanInputFormRecipient.new( - form_id=form_id, - delivery_id=delivery_id, - payload=payload, - ) - ) - return recipient_models - - 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 - if params.form_kind == HumanInputFormKind.RUNTIME and workflow_execution_id is None: - raise ValueError("a runtime human input form requires a workflow_execution_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) - if self._app_id is not None: - form_query = form_query.where(HumanInputForm.app_id == self._app_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) - - def mark_timeout(self, node_id: str, *, form_id: str) -> HumanInputFormEntity: - if self._workflow_execution_id is None: - raise ValueError("workflow_execution_id is required to time out runtime human input forms") - - statement = ( - update(HumanInputForm) - .where( - HumanInputForm.id == form_id, - HumanInputForm.tenant_id == self._tenant_id, - HumanInputForm.workflow_run_id == self._workflow_execution_id, - HumanInputForm.node_id == node_id, - HumanInputForm.status == HumanInputFormStatus.WAITING, - HumanInputForm.submitted_at.is_(None), - HumanInputForm.expiration_time <= naive_utc_now(), - ) - .values(status=HumanInputFormStatus.TIMEOUT) - ) - if self._app_id is not None: - statement = statement.where(HumanInputForm.app_id == self._app_id) - with session_factory.create_session() as session, session.begin(): - session.execute(statement) - - # A submission or global expiration may have won the race to finish the form. - form = self.get_form(node_id, form_id=form_id) - if form is None: - raise FormNotFoundError(f"form not found, id={form_id}") - return form - - -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_ids( - self, form_ids: Sequence[str], *, tenant_id: str, app_id: str, workflow_run_id: str - ) -> dict[str, HumanInputFormRecord]: - """Materialize a batch of forms without retaining a database session.""" - if not form_ids: - return {} - with session_factory.create_session() as session: - forms = session.scalars( - select(HumanInputForm).where( - HumanInputForm.id.in_(form_ids), - HumanInputForm.tenant_id == tenant_id, - HumanInputForm.app_id == app_id, - HumanInputForm.workflow_run_id == workflow_run_id, - ) - ) - return {form.id: HumanInputFormRecord.from_models(form, None) for form in forms} - - 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, with_for_update=True) - if form_model is None: - raise FormNotFoundError(f"form not found, id={form_id}") - if form_model.status != HumanInputFormStatus.WAITING: - raise FormAlreadyHandledError(form_model.status) - - 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/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 c3437e5a1590a7..00000000000000 --- a/api/core/tools/utils/dataset_retriever/dataset_multi_retriever_tool.py +++ /dev/null @@ -1,215 +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.is_(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] = [] - for resource_number, segment in enumerate(sorted_segments, start=1): - dataset = session.get(Dataset, segment.dataset_id) - document_stmt = select(Document).where( - Document.id == segment.document_id, - Document.enabled.is_(True), - Document.archived.is_(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) - - 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 - dataset_top_k = retrieval_model.get("top_k") - - 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=0 is a legal explicit value; only None falls back to 4. - top_k=dataset_top_k if dataset_top_k is not None else 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=0 is a legal explicit value; only None falls back to 4. - top_k=dataset_top_k if dataset_top_k is not None else 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 ac8d12b19f2df5..00000000000000 --- a/api/core/tools/utils/dataset_retriever/dataset_retriever_tool.py +++ /dev/null @@ -1,242 +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, - ) - document_ids_filter = metadata_filter_document_ids.get(dataset.id, []) if metadata_filter_document_ids else 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 - final_content = f"{record.summary}\n{segment_content}" if record.summary else 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.is_(True), - DatasetDocument.archived.is_(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 424ccea86c10ef..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]: - """ - 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/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 55e47b1c37ef95..c2c242ee7b3cd5 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 @@ -713,6 +713,18 @@ def workflow_from_snapshot(snapshot: WorkflowSnapshot) -> Workflow: return 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 @@ -721,6 +733,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( @@ -750,12 +763,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: @@ -768,9 +779,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, @@ -851,6 +864,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) @@ -863,6 +879,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/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 ef3b67e5862a89..00000000000000 --- a/api/services/agent/legacy_retirement_service.py +++ /dev/null @@ -1,223 +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 repositories.agent.home_snapshot_repository import AgentHomeSnapshotRepository -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( - AgentHomeSnapshotRepository(session).retire_all_for_agent( - 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 8b705f55f1359a..00000000000000 --- a/api/services/agent/legacy_workflow_publish_service.py +++ /dev/null @@ -1,745 +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 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 repositories.agent.workflow_binding_repository import WorkflowAgentBindingRepository, workflow_binding_scope -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.agent.workflow_validator import WorkflowAgentNodeValidationError, WorkflowAgentNodeValidator -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, binding_version: str = Workflow.VERSION_DRAFT - ) -> dict[str, Any]: - """Return the workflow graph with persisted Agent binding fields projected into node data. - - Reads the binding rows of ``binding_version``. The default is the draft; - a DSL export passes the exported workflow's own version. - - 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 == binding_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: - validator = WorkflowAgentNodeValidator(repository=WorkflowAgentBindingRepository(session)) - validator.validate_published_workflow(workflow=workflow_binding_scope(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: - validator = WorkflowAgentNodeValidator(repository=WorkflowAgentBindingRepository(session)) - validator.validate_draft_workflow(workflow=workflow_binding_scope(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, - binding_type=resolved_binding_type, - 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]: - binding_repository = WorkflowAgentBindingRepository(session) - source_agent = binding_repository.get_agent(draft_workflow.tenant_id, source_agent_id) - if ( - source_agent is None - or source_agent.scope != AgentScope.WORKFLOW_ONLY - or source_agent.status != AgentStatus.ACTIVE - ): - raise ValueError(f"Workflow Agent node {node_id} references an unavailable inline agent.") - source_snapshot = binding_repository.get_snapshot( - draft_workflow.tenant_id, - source_agent.id, - source_snapshot_id, - ) - if source_snapshot is None: - raise ValueError(f"Workflow Agent node {node_id} references a missing inline agent config snapshot.") - - cloned_agent, snapshot_id = binding_repository.clone( - workflow=workflow_binding_scope(draft_workflow), - node_id=node_id, - source_agent=source_agent, - source_snapshot=source_snapshot, - account_id=account_id, - ) - agent = session.scalar( - select(Agent) - .where( - Agent.tenant_id == draft_workflow.tenant_id, - Agent.id == cloned_agent.id, - ) - .limit(1) - ) - if agent is None: - raise RuntimeError(f"Workflow Agent clone for node {node_id} was not persisted.") - 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, - ) -> None: - """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. - """ - 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() - if not bindings: - return - - for binding in bindings: - 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) - - @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/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/trigger/webhook_service.py b/api/services/trigger/webhook_service.py index d7fc3f345a4433..c215ff3206f232 100644 --- a/api/services/trigger/webhook_service.py +++ b/api/services/trigger/webhook_service.py @@ -10,6 +10,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 @@ -249,7 +250,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 @@ -364,6 +377,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}, {} @@ -379,6 +394,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/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 e428e109a226ad..00000000000000 --- a/api/services/workflow_draft_variable_service.py +++ /dev/null @@ -1,1210 +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, 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.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, ConversationStatus, DraftVariableType, ExecutionOffLoadType, InvokeFrom -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, - tenant_id=workflow.tenant_id, - ) - if node_exec is None or node_exec.app_id != workflow.app_id: - logger.warning( - "Node execution 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 - - output_offload = next( - (item for item in node_exec.offload_data if item.type_ == ExecutionOffLoadType.OUTPUTS), - None, - ) - if output_offload is None: - outputs_dict = node_exec.outputs_dict or {} - else: - if ( - output_offload.tenant_id != node_exec.tenant_id - or output_offload.app_id != node_exec.app_id - or output_offload.file is None - or output_offload.file.tenant_id != node_exec.tenant_id - ): - raise VariableResetError(f"Execution output file not found, execution_id={node_exec.id}") - outputs_dict = json.loads(storage.load(output_offload.file.key)) 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=ConversationStatus.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: - # Save nested output variables only for variable assigners. - return self._enclosing_node_id is None or self._node_type == BuiltinNodeTypes.VARIABLE_ASSIGNER - - 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 - value_seg = value if isinstance(value, Segment) else _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 - original_value = value_seg.value - if isinstance(original_value, (list, dict)): - original_length = len(original_value) - - # Prepare content for storage - original_content_serialized: str - if isinstance(value_seg, StringSegment): - # For string types, store as plain text - original_content_serialized = 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 - return not (node_id == SYSTEM_VARIABLE_NODE_ID and not is_system_variable_editable(name)) - - @staticmethod - def _should_variable_be_visible(node_id: str, node_type: NodeType, name: str) -> bool: - if node_type == BuiltinNodeTypes.IF_ELSE: - return False - return not (node_id == SYSTEM_VARIABLE_NODE_ID and not is_system_variable_editable(name)) - - 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 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/dataset_retriever/test_dataset_multi_retriever_tool.py b/api/tests/unit_tests/core/tools/utils/dataset_retriever/test_dataset_multi_retriever_tool.py deleted file mode 100644 index c7012917a57118..00000000000000 --- a/api/tests/unit_tests/core/tools/utils/dataset_retriever/test_dataset_multi_retriever_tool.py +++ /dev/null @@ -1,598 +0,0 @@ -from __future__ import annotations - -import uuid -from contextlib import nullcontext -from dataclasses import dataclass -from datetime import datetime -from pathlib import Path -from unittest.mock import patch - -import pytest -from sqlalchemy.orm import Session, scoped_session, sessionmaker -from yaml import YAMLError - -from core.app.app_config.entities import DatasetRetrieveConfigEntity -from core.callback_handler.index_tool_callback_handler import DatasetIndexToolCallbackHandler -from core.model_manager import ModelInstance, ModelManager -from core.provider_manager import ProviderManager -from core.rag.embedding.retrieval import RetrievalSegments -from core.rag.models.document import Document as RagDocument -from core.rag.rerank.rerank_model import RerankModelRunner -from core.tools.utils.dataset_retriever import dataset_multi_retriever_tool as multi_retriever_module -from core.tools.utils.dataset_retriever import dataset_retriever_tool as single_retriever_module -from core.tools.utils.dataset_retriever.dataset_multi_retriever_tool import DatasetMultiRetrieverTool -from core.tools.utils.dataset_retriever.dataset_retriever_tool import DatasetRetrieverTool as SingleDatasetRetrieverTool -from core.tools.utils.text_processing_utils import remove_leading_symbols -from core.tools.utils.uuid_utils import is_valid_uuid -from core.tools.utils.yaml_utils import _load_yaml_file, load_yaml_file_cached -from models.dataset import Dataset, Document, DocumentSegment -from models.enums import DataSourceType, DocumentCreatedFrom, SegmentStatus -from tests.helpers.unit.model_factories import upload_reader - - -def _retrieve_config() -> DatasetRetrieveConfigEntity: - return DatasetRetrieveConfigEntity(retrieve_strategy=DatasetRetrieveConfigEntity.RetrieveStrategy.SINGLE) - - -def _persist_dataset( - session: Session, - *, - dataset_id: str | None = None, - tenant_id: str | None = None, - name: str = "Knowledge Base", - provider: str = "vendor", - retrieval_model: dict | None = None, -) -> Dataset: - dataset = Dataset( - id=dataset_id or str(uuid.uuid4()), - tenant_id=tenant_id or str(uuid.uuid4()), - name=name, - provider=provider, - data_source_type=DataSourceType.UPLOAD_FILE, - indexing_technique="high_quality", - retrieval_model=retrieval_model, - created_by=str(uuid.uuid4()), - ) - session.add(dataset) - session.commit() - return dataset - - -def _persist_document( - session: Session, - *, - dataset: Dataset, - name: str, - data_source_type: DataSourceType = DataSourceType.UPLOAD_FILE, - doc_metadata: dict | None = None, -) -> Document: - document = Document( - id=str(uuid.uuid4()), - tenant_id=dataset.tenant_id, - dataset_id=dataset.id, - position=1, - data_source_type=data_source_type, - batch="batch-1", - name=name, - created_from=DocumentCreatedFrom.WEB, - created_by=dataset.created_by, - doc_metadata=doc_metadata, - ) - session.add(document) - session.commit() - return document - - -class _FakeFlaskApp: - def app_context(self): - return nullcontext() - - -class _FakeCurrentApp: - def _get_current_object(self) -> _FakeFlaskApp: - return _FakeFlaskApp() - - -@dataclass(eq=False, repr=False) -class _DatabaseWithSession: - session: scoped_session[Session] - - -class _UnusedProviderManager(ProviderManager): - def __init__(self) -> None: - pass - - -class _UnusedModelInstance(ModelInstance): - def __init__(self) -> None: - pass - - -class _ImmediateThread: - def __init__(self, target=None, kwargs=None, **_kwargs): - self._target = target - self._kwargs = kwargs or {} - - def start(self): - if self._target is not None: - self._target(**self._kwargs) - - def join(self): - return None - - -class _TestHitCallback(DatasetIndexToolCallbackHandler): - def __init__(self): - self.queries: list[tuple[str, str]] = [] - self.documents: list[RagDocument] | None = None - self.resources = None - - def on_query(self, query: str, dataset_id: str, session=None): - self.queries.append((query, dataset_id)) - - def on_tool_end(self, documents: list[RagDocument], session=None): - self.documents = documents - - def return_retriever_resource_info(self, resource): - self.resources = list(resource) - - -def test_remove_leading_symbols_preserves_markdown_link_and_strips_punctuation(): - markdown = "[Example](https://example.com) content" - assert remove_leading_symbols(markdown) == markdown - - assert remove_leading_symbols("...Hello world") == "Hello world" - - -def test_is_valid_uuid_handles_valid_invalid_and_empty_values(): - assert is_valid_uuid(str(uuid.uuid4())) is True - assert is_valid_uuid("not-a-uuid") is False - assert is_valid_uuid("") is False - assert is_valid_uuid(None) is False - - -def test_load_yaml_file_valid(tmp_path: Path): - valid_file = tmp_path / "valid.yaml" - valid_file.write_text("a: 1\nb: two\n", encoding="utf-8") - - loaded = _load_yaml_file(file_path=str(valid_file)) - - assert loaded == {"a": 1, "b": "two"} - - -def test_load_yaml_file_missing(tmp_path: Path): - with pytest.raises(FileNotFoundError): - _load_yaml_file(file_path=str(tmp_path / "missing.yaml")) - - -def test_load_yaml_file_invalid(tmp_path: Path): - invalid_file = tmp_path / "invalid.yaml" - invalid_file.write_text("a: [1, 2\n", encoding="utf-8") - - with pytest.raises(YAMLError): - _load_yaml_file(file_path=str(invalid_file)) - - -def test_load_yaml_file_cached_hits(tmp_path: Path): - valid_file = tmp_path / "valid.yaml" - valid_file.write_text("a: 1\nb: two\n", encoding="utf-8") - - load_yaml_file_cached.cache_clear() - assert load_yaml_file_cached(str(valid_file)) == {"a": 1, "b": "two"} - - assert load_yaml_file_cached(str(valid_file)) == {"a": 1, "b": "two"} - assert load_yaml_file_cached.cache_info().hits == 1 - - -def test_single_dataset_retriever_from_dataset_builds_name_and_description(sqlite_session: Session): - dataset = _persist_dataset( - sqlite_session, - dataset_id="dataset-1", - tenant_id="tenant-1", - name="Knowledge", - ) - - tool = SingleDatasetRetrieverTool.from_dataset( - dataset=dataset, - retrieve_config=_retrieve_config(), - return_resource=False, - retriever_from="prod", - inputs={}, - ) - - assert tool.name == "dataset_dataset_1" - assert tool.description == "useful for when you want to answer queries about the Knowledge" - - -def test_single_dataset_retriever_external_run_returns_content_and_resources(sqlite_session: Session): - dataset = _persist_dataset(sqlite_session, provider="external", retrieval_model={}) - callback = _TestHitCallback() - metadata_filter_result = ( - {dataset.id: ["doc-a"]}, - {"logical_operator": "and"}, - ) - external_documents = [ - {"content": "first", "metadata": {"document_id": "doc-a"}, "score": 0.9, "title": "Doc A"}, - {"content": "second", "metadata": {"document_id": "doc-b"}, "score": 0.8, "title": "Doc B"}, - ] - - tool = SingleDatasetRetrieverTool( - tenant_id=dataset.tenant_id, - dataset_id=dataset.id, - retrieve_config=_retrieve_config(), - return_resource=True, - retriever_from="dev", - hit_callbacks=[callback], - inputs={"x": 1}, - ) - - with patch.object( - single_retriever_module.DatasetRetrieval, - "get_metadata_filter_condition", - return_value=metadata_filter_result, - ): - with patch.object( - single_retriever_module.ExternalDatasetService, - "fetch_external_knowledge_retrieval", - return_value=external_documents, - ) as fetch_mock: - result = tool.run(session=sqlite_session, query="hello") - - assert result == "first\nsecond" - assert callback.queries == [("hello", dataset.id)] - assert callback.resources is not None - resource_info = callback.resources - assert [item.position for item in resource_info] == [1, 2] - assert resource_info[0].dataset_id == dataset.id - fetch_mock.assert_called_once() - assert fetch_mock.call_args.kwargs["session"] is sqlite_session - - -def test_single_dataset_retriever_returns_empty_when_metadata_filter_finds_no_documents( - sqlite_session: Session, -): - dataset = _persist_dataset(sqlite_session) - tool = SingleDatasetRetrieverTool( - tenant_id=dataset.tenant_id, - dataset_id=dataset.id, - retrieve_config=_retrieve_config(), - return_resource=False, - retriever_from="prod", - hit_callbacks=[_TestHitCallback()], - inputs={}, - ) - - with patch.object( - single_retriever_module.DatasetRetrieval, - "get_metadata_filter_condition", - return_value=({dataset.id: []}, {"logical_operator": "and"}), - ): - with patch.object(single_retriever_module.RetrievalService, "retrieve") as retrieve_mock: - result = tool.run(session=sqlite_session, query="hello") - - assert result == "" - retrieve_mock.assert_not_called() - - -def test_single_dataset_retriever_non_economy_run_sorts_context_and_resources(sqlite_session: Session): - dataset = _persist_dataset( - sqlite_session, - retrieval_model={ - "search_method": "semantic_search", - "score_threshold_enabled": True, - "score_threshold": 0.2, - "reranking_enable": True, - "reranking_model": {"reranking_provider_name": "provider", "reranking_model_name": "model"}, - "reranking_mode": "reranking_model", - "weights": {"vector_setting": {"vector_weight": 0.6}}, - }, - ) - document_low = _persist_document( - sqlite_session, - dataset=dataset, - name="Document Low", - doc_metadata={"lang": "en"}, - ) - document_high = _persist_document( - sqlite_session, - dataset=dataset, - name="Document High", - data_source_type=DataSourceType.NOTION_IMPORT, - doc_metadata={"lang": "fr"}, - ) - callback = _TestHitCallback() - low_segment = DocumentSegment( - tenant_id=dataset.tenant_id, - dataset_id=dataset.id, - document_id=document_low.id, - index_node_id="node-low", - content="raw low", - hit_count=1, - word_count=10, - position=3, - index_node_hash="hash-low", - tokens=10, - created_by=dataset.created_by, - answer="low answer", - status=SegmentStatus.COMPLETED, - completed_at=datetime.now(), - ) - high_segment = DocumentSegment( - tenant_id=dataset.tenant_id, - dataset_id=dataset.id, - document_id=document_high.id, - index_node_id="node-high", - content="raw high", - hit_count=9, - word_count=25, - position=1, - index_node_hash="hash-high", - tokens=25, - created_by=dataset.created_by, - status=SegmentStatus.COMPLETED, - completed_at=datetime.now(), - ) - sqlite_session.add_all([low_segment, high_segment]) - sqlite_session.commit() - records = [ - RetrievalSegments(segment=low_segment, score=0.2, summary="summary low"), - RetrievalSegments(segment=high_segment, score=0.9), - ] - documents = [ - RagDocument(page_content="first", metadata={"doc_id": "node-low", "score": 0.2}), - RagDocument(page_content="second", metadata={"doc_id": "node-high", "score": 0.9}), - ] - tool = SingleDatasetRetrieverTool( - tenant_id=dataset.tenant_id, - dataset_id=dataset.id, - retrieve_config=_retrieve_config(), - return_resource=True, - retriever_from="dev", - hit_callbacks=[callback], - inputs={}, - top_k=2, - ) - - with ( - patch.object( - single_retriever_module.DatasetRetrieval, - "get_metadata_filter_condition", - return_value=(None, None), - ), - patch.object(single_retriever_module.RetrievalService, "retrieve", return_value=documents), - patch.object( - single_retriever_module.RetrievalService, - "format_retrieval_documents", - return_value=records, - ), - patch.object( - single_retriever_module, - "sign_segment_content", - lambda segment, **_kwargs: segment.content.replace("raw", "signed"), - ), - ): - result = tool.run(session=sqlite_session, query="hello") - - assert result == "signed high\nsummary low\nquestion:signed low answer:low answer" - assert callback.documents == documents - assert callback.resources is not None - resource_info = callback.resources - assert [item.position for item in resource_info] == [1, 2] - assert resource_info[0].segment_id == high_segment.id - assert resource_info[0].hit_count == 9 - assert resource_info[1].summary == "summary low" - assert resource_info[1].content == "question:raw low \nanswer:low answer" - - -def test_multi_dataset_retriever_from_dataset_sets_tool_name(): - tool = DatasetMultiRetrieverTool.from_dataset( - dataset_ids=["dataset-1"], - tenant_id="tenant-1", - reranking_provider_name="provider", - reranking_model_name="model", - return_resource=False, - retriever_from="prod", - ) - - assert tool.name == "dataset_tenant_1" - - -def test_multi_dataset_retriever_retriever_returns_early_when_dataset_is_missing( - sqlite_session_factory: sessionmaker[Session], -): - callback = _TestHitCallback() - all_documents: list[RagDocument] = [] - db_session = scoped_session(sqlite_session_factory) - tool = DatasetMultiRetrieverTool( - tenant_id="tenant-1", - dataset_ids=["dataset-1"], - reranking_provider_name="provider", - reranking_model_name="model", - return_resource=False, - retriever_from="prod", - ) - - try: - with patch.object(multi_retriever_module, "db", _DatabaseWithSession(db_session)): - with patch.object(multi_retriever_module.RetrievalService, "retrieve") as retrieve_mock: - result = tool._retriever( - flask_app=_FakeFlaskApp(), - dataset_id=str(uuid.uuid4()), - query="hello", - all_documents=all_documents, - hit_callbacks=[callback], - ) - finally: - db_session.remove() - - assert result == [] - assert all_documents == [] - assert callback.queries == [] - retrieve_mock.assert_not_called() - - -def test_multi_dataset_retriever_retriever_non_economy_uses_retrieval_model( - sqlite_session: Session, - sqlite_session_factory: sessionmaker[Session], -): - dataset = _persist_dataset( - sqlite_session, - retrieval_model={ - "search_method": "semantic_search", - "top_k": 6, - "score_threshold_enabled": True, - "score_threshold": 0.4, - "reranking_enable": False, - "reranking_mode": None, - "weights": {"balanced": True}, - }, - ) - callback = _TestHitCallback() - documents = [RagDocument(page_content="retrieved", metadata={"doc_id": "node-1", "score": 0.4})] - all_documents: list[RagDocument] = [] - db_session = scoped_session(sqlite_session_factory) - tool = DatasetMultiRetrieverTool( - tenant_id=dataset.tenant_id, - dataset_ids=[dataset.id], - reranking_provider_name="provider", - reranking_model_name="model", - return_resource=False, - retriever_from="prod", - top_k=2, - ) - - try: - with patch.object(multi_retriever_module, "db", _DatabaseWithSession(db_session)): - with patch.object( - multi_retriever_module.RetrievalService, "retrieve", return_value=documents - ) as retrieve_mock: - tool._retriever( - flask_app=_FakeFlaskApp(), - dataset_id=dataset.id, - query="hello", - all_documents=all_documents, - hit_callbacks=[callback], - ) - finally: - db_session.remove() - - assert all_documents == documents - assert callback.queries == [("hello", dataset.id)] - retrieve_mock.assert_called_once_with( - retrieval_method="semantic_search", - dataset_id=dataset.id, - query="hello", - top_k=6, - score_threshold=0.4, - reranking_model=None, - reranking_mode="reranking_model", - weights={"balanced": True}, - ) - - -def test_multi_dataset_retriever_run_orders_segments_and_returns_resources(sqlite_session: Session): - dataset_one = _persist_dataset(sqlite_session, name="Dataset One") - dataset_two = _persist_dataset(sqlite_session, tenant_id=dataset_one.tenant_id, name="Dataset Two") - document_two = _persist_document( - sqlite_session, - dataset=dataset_one, - name="Doc Two", - data_source_type=DataSourceType.NOTION_IMPORT, - doc_metadata={"p": 2}, - ) - document_one = _persist_document( - sqlite_session, - dataset=dataset_two, - name="Doc One", - doc_metadata={"p": 1}, - ) - callback = _TestHitCallback() - tool = DatasetMultiRetrieverTool( - tenant_id=dataset_one.tenant_id, - dataset_ids=[dataset_one.id, dataset_two.id], - reranking_provider_name="provider", - reranking_model_name="model", - return_resource=True, - retriever_from="dev", - hit_callbacks=[callback], - top_k=2, - score_threshold=0.1, - ) - first_doc = RagDocument(page_content="first", metadata={"doc_id": "node-2", "score": 0.4}) - second_doc = RagDocument(page_content="second", metadata={"doc_id": "node-1", "score": 0.9}) - - def fake_retriever(**kwargs): - if kwargs["dataset_id"] == dataset_one.id: - kwargs["all_documents"].append(first_doc) - else: - kwargs["all_documents"].append(second_doc) - - segment_for_node_2 = DocumentSegment( - tenant_id=dataset_one.tenant_id, - dataset_id=dataset_one.id, - document_id=document_two.id, - index_node_id="node-2", - content="raw two", - hit_count=2, - word_count=20, - position=2, - index_node_hash="hash-2", - tokens=20, - created_by=dataset_one.created_by, - answer="answer two", - status=SegmentStatus.COMPLETED, - completed_at=datetime.now(), - ) - segment_for_node_1 = DocumentSegment( - tenant_id=dataset_two.tenant_id, - dataset_id=dataset_two.id, - document_id=document_one.id, - index_node_id="node-1", - content="raw one", - hit_count=7, - word_count=30, - position=1, - index_node_hash="hash-1", - tokens=30, - created_by=dataset_two.created_by, - status=SegmentStatus.COMPLETED, - completed_at=datetime.now(), - ) - sqlite_session.add_all([segment_for_node_2, segment_for_node_1]) - sqlite_session.commit() - model_manager = ModelManager(provider_manager=_UnusedProviderManager()) - model_instance = _UnusedModelInstance() - rerank_runner = RerankModelRunner( - model_instance, load_upload=upload_reader(sessionmaker(bind=sqlite_session.get_bind())) - ) - fake_current_app = _FakeCurrentApp() - - with ( - patch.object( - multi_retriever_module, - "sign_segment_content", - lambda segment, **_kwargs: segment.content.replace("raw", "signed"), - ), - patch.object(tool, "_retriever", side_effect=fake_retriever) as retriever_mock, - patch.object(multi_retriever_module, "current_app", fake_current_app), - patch.object(multi_retriever_module.threading, "Thread", _ImmediateThread), - patch.object(multi_retriever_module.ModelManager, "for_tenant", return_value=model_manager), - patch.object(model_manager, "get_model_instance", return_value=model_instance), - patch.object(multi_retriever_module, "RerankModelRunner", return_value=rerank_runner) as rerank_runner_class, - patch.object(rerank_runner, "run", return_value=[second_doc, first_doc]), - ): - result = tool.run(session=sqlite_session, query="hello") - - assert result == "signed one\nquestion:signed two answer:answer two" - rerank_runner_class.assert_called_once() - assert rerank_runner_class.call_args.args == (model_instance,) - assert rerank_runner_class.call_args.kwargs["load_upload"].keywords == {"workspace_id": dataset_one.tenant_id} - assert retriever_mock.call_count == 2 - assert callback.documents == [second_doc, first_doc] - assert callback.resources is not None - resource_info = callback.resources - assert [item.position for item in resource_info] == [1, 2] - assert resource_info[0].score == 0.9 - assert resource_info[0].content == "raw one" - assert resource_info[1].score == 0.4 - assert resource_info[1].content == "question:raw two \nanswer:answer two" diff --git a/api/tests/unit_tests/core/tools/utils/test_dataset_retriever_tool.py b/api/tests/unit_tests/core/tools/utils/test_dataset_retriever_tool.py index a9fceb1dd2bbbb..ad14773b35a3dc 100644 --- a/api/tests/unit_tests/core/tools/utils/test_dataset_retriever_tool.py +++ b/api/tests/unit_tests/core/tools/utils/test_dataset_retriever_tool.py @@ -1,89 +1,208 @@ """Dataset tools query injected records and preserve their dependencies when forked.""" +from collections.abc import Mapping, Sequence from dataclasses import dataclass, field +from typing import TypedDict, override from uuid import uuid4 import pytest -from sqlalchemy.orm import Session +from sqlalchemy.orm import Session, sessionmaker -from core.app.app_config.entities import DatasetRetrieveConfigEntity -from core.app.entities.app_invoke_entities import InvokeFrom +from core.app.app_config.entities import DatasetRetrieveConfigEntity, ModelConfig +from core.app.entities.app_invoke_entities import EasyUIBasedAppGenerateEntity, InvokeFrom +from core.rag.entities import MetadataFilteringCondition, RetrievalSourceMetadata from core.tools.__base.tool_runtime import ToolRuntime +from core.tools.entities.tool_entities import ToolInvokeMessage +from core.workflow.nodes.knowledge_retrieval.retrieval import KnowledgeRetrievalRequest, Source +from graphon.model_runtime.entities.llm_entities import LLMUsage from models.dataset import Dataset, Document from repositories.knowledge.retrieval_repository import KnowledgeRetrievalRepository +from services.knowledge.retrieval.adapters.resource_events import DatasetIndexToolCallbackHandler from services.tools.dataset.tool import DatasetRetrieverTool from tests.helpers.unit.model_factories import make_dataset +class RetrievalCall(TypedDict): + tenant_id: str + app_id: str + user_id: str + dataset_id: str + query: str + config: DatasetRetrieveConfigEntity + top_k: int + inputs: dict[str, object] + invoke_from: InvokeFrom + return_resource: bool + hit_callback: DatasetIndexToolCallbackHandler + + @dataclass class Retrieval: - calls: list = field(default_factory=list) + calls: list[RetrievalCall] = field(default_factory=list) + + @property + def llm_usage(self) -> LLMUsage: + return LLMUsage.empty_usage() + + def set_request_metadata(self, request_metadata: Mapping[str, object] | None) -> None: + del request_metadata + + def knowledge_retrieval(self, request: KnowledgeRetrievalRequest) -> list[Source]: + raise AssertionError(f"unexpected knowledge retrieval: {request}") + + def get_metadata_filter_condition( + self, + dataset_ids: list[str], + query: str, + tenant_id: str, + user_id: str, + metadata_filtering_mode: str, + metadata_model_config: ModelConfig | None, + metadata_filtering_conditions: MetadataFilteringCondition | None, + inputs: dict[str, object], + ) -> tuple[dict[str, list[str]] | None, MetadataFilteringCondition | None]: + raise AssertionError( + "unexpected metadata filtering: " + f"{dataset_ids}, {query}, {tenant_id}, {user_id}, {metadata_filtering_mode}, " + f"{metadata_model_config}, {metadata_filtering_conditions}, {inputs}" + ) + + def retrieve_dataset( + self, + *, + tenant_id: str, + app_id: str, + user_id: str, + dataset_id: str, + query: str, + config: DatasetRetrieveConfigEntity, + top_k: int, + inputs: dict[str, object], + invoke_from: InvokeFrom, + return_resource: bool, + hit_callback: DatasetIndexToolCallbackHandler, + ) -> str: + self.calls.append( + RetrievalCall( + tenant_id=tenant_id, + app_id=app_id, + user_id=user_id, + dataset_id=dataset_id, + query=query, + config=config, + top_k=top_k, + inputs=inputs, + invoke_from=invoke_from, + return_resource=return_resource, + hit_callback=hit_callback, + ) + ) + return f"result:{query}" + + +class NoResources(DatasetIndexToolCallbackHandler): + def __init__(self) -> None: + pass + + @override + def return_retriever_resource_info(self, resource: Sequence[RetrievalSourceMetadata]) -> None: + del resource + raise AssertionError("unexpected resources") - def retrieve_dataset(self, **kwargs): - self.calls.append(kwargs) - return f"result:{kwargs['query']}" +@dataclass(frozen=True) +class RetrievalFactory: + retrieval: Retrieval + + def __call__(self, application_generate_entity: EasyUIBasedAppGenerateEntity | None = None) -> Retrieval: + del application_generate_entity + return self.retrieval -class NoResources: - def return_retriever_resource_info(self, _resources): - raise AssertionError("unexpected resources") + +type DatasetTools = tuple[ + list[DatasetRetrieverTool], + Retrieval, + DatasetRetrieveConfigEntity, + str, + NoResources, +] @pytest.fixture -def dataset_tools(sqlite_session: Session, sqlite_session_factory): +def dataset_tools(sqlite_session: Session, sqlite_session_factory: sessionmaker[Session]) -> DatasetTools: dataset_id = str(uuid4()) sqlite_session.add( make_dataset(dataset_id=dataset_id, tenant_id="tenant", provider="external", description="one\ntwo") ) sqlite_session.commit() retrieval = Retrieval() + callback = NoResources() config = DatasetRetrieveConfigEntity(retrieve_strategy="multiple") tools = DatasetRetrieverTool.get_dataset_tools( records=KnowledgeRetrievalRepository(sqlite_session_factory), - retrieval=lambda: retrieval, + retrieval=RetrievalFactory(retrieval), tenant_id="tenant", app_id="app", dataset_ids=[dataset_id], retrieve_config=config, return_resource=True, invoke_from=InvokeFrom.DEBUGGER, - hit_callback=NoResources(), + hit_callback=callback, user_id="user", inputs={"x": 1}, ) - return tools, retrieval, config + return tools, retrieval, config, dataset_id, callback -def test_tool_metadata_and_fork_preserve_injected_retrieval(dataset_tools, unbound_session: Session): - tools, retrieval, config = dataset_tools +def test_tool_metadata_and_fork_preserve_injected_retrieval( + dataset_tools: DatasetTools, unbound_session: Session +) -> None: + tools, retrieval, config, dataset_id, callback = dataset_tools assert len(tools) == 1 tool = tools[0] + assert tool.entity.description is not None assert tool.entity.description.llm == "onetwo" assert [parameter.name for parameter in tool.get_runtime_parameters()] == ["query"] assert config.retrieve_strategy == "multiple" forked = tool.fork_tool_runtime(ToolRuntime(tenant_id="tenant")) result = list(forked.invoke(session=unbound_session, user_id="user", tool_parameters={"query": "hello"})) + assert isinstance(result[0].message, ToolInvokeMessage.TextMessage) assert result[0].message.text == "result:hello" - assert retrieval.calls[0]["inputs"] == {"x": 1} - assert retrieval.calls[0]["app_id"] == "app" - - -def test_empty_query_does_not_start_retrieval(dataset_tools, unbound_session: Session): - tools, retrieval, _ = dataset_tools - assert list(tools[0].invoke(session=unbound_session, user_id="user", tool_parameters={}))[0].message.text == ( - "please input query" - ) + assert retrieval.calls == [ + { + "tenant_id": "tenant", + "app_id": "app", + "user_id": "user", + "dataset_id": dataset_id, + "query": "hello", + "config": config, + "top_k": 2, + "inputs": {"x": 1}, + "invoke_from": InvokeFrom.DEBUGGER, + "return_resource": True, + "hit_callback": callback, + } + ] + + +def test_empty_query_does_not_start_retrieval(dataset_tools: DatasetTools, unbound_session: Session) -> None: + tools, retrieval, _, _, _ = dataset_tools + messages = list(tools[0].invoke(session=unbound_session, user_id="user", tool_parameters={})) + assert isinstance(messages[0].message, ToolInvokeMessage.TextMessage) + assert messages[0].message.text == "please input query" assert not retrieval.calls @pytest.mark.parametrize("sqlite_session", [(Dataset, Document)], indirect=True) @pytest.mark.parametrize("case", ["no-datasets", "no-config", "foreign-tenant"]) -def test_tool_preparation_rejects_unavailable_datasets(sqlite_session, sqlite_session_factory, case): +def test_tool_preparation_rejects_unavailable_datasets( + sqlite_session: Session, sqlite_session_factory: sessionmaker[Session], case: str +) -> None: sqlite_session.add(make_dataset(dataset_id="dataset", tenant_id="other", provider="external")) sqlite_session.commit() result = DatasetRetrieverTool.get_dataset_tools( records=KnowledgeRetrievalRepository(sqlite_session_factory), - retrieval=Retrieval, + retrieval=RetrievalFactory(Retrieval()), tenant_id="tenant", app_id="app", dataset_ids=[] if case == "no-datasets" else ["dataset"], 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 new file mode 100644 index 00000000000000..9d218af5290d98 --- /dev/null +++ b/api/tests/unit_tests/core/tools/utils/test_uuid_utils.py @@ -0,0 +1,22 @@ +import uuid + +import pytest + +from core.tools.utils.uuid_utils import is_valid_uuid + + +@pytest.mark.parametrize( + ("value", "expected"), + [ + ("00000000-0000-4000-8000-000000000001", True), + ("not-a-uuid", False), + ("", False), + (None, False), + ], +) +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/core/tools/utils/test_yaml_utils.py b/api/tests/unit_tests/core/tools/utils/test_yaml_utils.py index ac192cb6c50601..cebfa6c9f71550 100644 --- a/api/tests/unit_tests/core/tools/utils/test_yaml_utils.py +++ b/api/tests/unit_tests/core/tools/utils/test_yaml_utils.py @@ -4,7 +4,7 @@ import pytest from yaml import YAMLError -from core.tools.utils.yaml_utils import _load_yaml_file +from core.tools.utils.yaml_utils import _load_yaml_file, load_yaml_file_cached EXAMPLE_YAML_FILE = "example_yaml.yaml" INVALID_YAML_FILE = "invalid_yaml.yaml" @@ -75,6 +75,15 @@ def test_load_valid_yaml_file(prepare_example_yaml_file: str) -> None: assert yaml_data.get("non_existed_key") is None +def test_load_yaml_file_cached_reuses_parsed_value(prepare_example_yaml_file: str) -> None: + load_yaml_file_cached.cache_clear() + + assert load_yaml_file_cached(prepare_example_yaml_file)["age"] == 30 + assert load_yaml_file_cached(prepare_example_yaml_file)["age"] == 30 + + assert load_yaml_file_cached.cache_info().hits == 1 + + def test_load_invalid_yaml_file(prepare_invalid_yaml_file: str) -> None: # yaml syntax error with pytest.raises(YAMLError): diff --git a/api/tests/unit_tests/core/workflow/test_workflow_tool_container_handler.py b/api/tests/unit_tests/core/workflow/test_workflow_tool_container_handler.py index 3fef6d98ddbbf4..e90463c13b765e 100644 --- a/api/tests/unit_tests/core/workflow/test_workflow_tool_container_handler.py +++ b/api/tests/unit_tests/core/workflow/test_workflow_tool_container_handler.py @@ -29,10 +29,6 @@ from core.app.workflow.file_runtime import init_app from core.app.workflow.layers.observability import ObservabilityLayer from core.helper.code_executor.code_executor import CodeExecutionError -from core.repositories.human_input_repository import ( - HumanInputFormRepositoryImpl, - HumanInputFormSubmissionRepository, -) from core.repositories.sqlalchemy_workflow_execution_repository import SQLAlchemyWorkflowExecutionRepository from core.repositories.sqlalchemy_workflow_node_execution_repository import SQLAlchemyWorkflowNodeExecutionRepository from core.tools.__base.tool_runtime import ToolRuntime @@ -130,6 +126,10 @@ from models.tool_runtime_contracts import WorkflowToolQueries from models.workflow import WorkflowNodeExecutionModel, WorkflowNodeExecutionTriggeredFrom from repositories.agent.runtime_repository import WorkflowAgentBindingResolver +from repositories.human_input.form_repository import ( + HumanInputFormRepositoryImpl, + HumanInputFormSubmissionRepository, +) from services.tools.workflow.tool import WorkflowTool from services.workflow.execution.adapters import execution_context_layer as module from services.workflow.execution.adapters.execution_context_layer import ExecutionContextLayer 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 24fc69413e47b2..1ec2eabe829ff5 100644 --- a/api/tests/unit_tests/repositories/workflow/test_definition_repository.py +++ b/api/tests/unit_tests/repositories/workflow/test_definition_repository.py @@ -78,6 +78,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( @@ -136,7 +153,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), @@ -175,7 +192,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: @@ -298,7 +315,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)) ) @@ -377,7 +394,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)) ) @@ -527,7 +544,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) @@ -621,7 +638,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( @@ -680,7 +697,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), @@ -746,7 +763,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), @@ -761,7 +778,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), @@ -790,11 +807,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), @@ -938,7 +955,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), @@ -973,7 +990,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), @@ -994,7 +1011,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), @@ -1088,7 +1105,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), @@ -1120,7 +1137,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), @@ -1148,7 +1165,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), @@ -1184,7 +1201,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), @@ -1287,7 +1304,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), @@ -1310,7 +1327,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), @@ -1346,7 +1363,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), @@ -1375,7 +1392,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), @@ -1529,7 +1546,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), @@ -1546,9 +1563,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 @@ -1613,11 +1644,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 @@ -1637,9 +1669,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: @@ -1828,7 +1860,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)) ) @@ -1953,7 +1985,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) ) @@ -2022,7 +2054,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/agent/test_roster_service.py b/api/tests/unit_tests/services/agent/test_roster_service.py index f63f4a7df7c752..366f215ebe98bf 100644 --- a/api/tests/unit_tests/services/agent/test_roster_service.py +++ b/api/tests/unit_tests/services/agent/test_roster_service.py @@ -1,4 +1,5 @@ import json +from collections.abc import Iterable from contextlib import nullcontext from datetime import UTC, datetime from types import SimpleNamespace @@ -46,6 +47,7 @@ from models.tools import ToolFile from models.workflow import Workflow from repositories.agent.config_repository import AgentConfigRepository +from repositories.agent.retirement_repository import WorkflowAgentRetirementRepository from repositories.agent.workflow_binding_repository import workflow_binding_scope from repositories.agent_workspace_repository import AgentWorkspaceRepository from services.agent import composer_service, roster_service @@ -61,6 +63,7 @@ InvalidComposerConfigError, ) from services.agent.home_snapshot_service import AgentHomeSnapshotService +from services.agent.retirement_service import WorkflowAgentRetirementService from services.agent.roster_service import AgentRosterService from services.agent.workflow_publish_service import WorkflowAgentPublishService from services.agent.workflow_resources_gateway import WorkflowAgentSkillReader @@ -633,11 +636,20 @@ def test_save_workflow_composer_commits_before_retiring_replaced_inline_agent( monkeypatch.setattr(AgentComposerService, "collect_validation_findings", lambda **_kwargs: {}) event.listen(session, "after_commit", lambda _session: events.append("commit")) - def retire_unowned(**kwargs): - assert kwargs["agent_ids"] == {"old-inline-agent"} + def retire_unowned( + self: WorkflowAgentRetirementService, + *, + tenant_id: str, + agent_ids: Iterable[str], + account_id: str | None, + ) -> None: + assert isinstance(self._repository, WorkflowAgentRetirementRepository) + assert tenant_id == "tenant-1" + assert set(agent_ids) == {"old-inline-agent"} + assert account_id == "account-1" events.append("retire") - monkeypatch.setattr(composer_service.WorkflowAgentRetirementService, "retire_unowned", staticmethod(retire_unowned)) + monkeypatch.setattr(WorkflowAgentRetirementService, "retire_unowned", retire_unowned) payload = ComposerSavePayload.model_validate( { "variant": ComposerVariant.WORKFLOW, diff --git a/api/tests/unit_tests/services/knowledge/test_dataset_service.py b/api/tests/unit_tests/services/knowledge/test_dataset_service.py index 78185c8ce77f14..df80864ed772a4 100644 --- a/api/tests/unit_tests/services/knowledge/test_dataset_service.py +++ b/api/tests/unit_tests/services/knowledge/test_dataset_service.py @@ -26,7 +26,6 @@ IndexStructureType as _dataset_service_segment_IndexStructureType, ) from core.rag.index_processor.constant.index_type import IndexTechniqueType -from core.rag.retrieval.dataset_retrieval import DatasetRetrieval from enums import DeploymentEdition from graphon.model_runtime.entities.model_entities import ModelFeature, ModelType from models import Account as _dataset_service_dataset_Account @@ -55,6 +54,7 @@ from models.model import UploadFile from models.source import DataSourceOauthBinding from repositories.knowledge.collection_binding_repository import DatasetCollectionBindingRepository +from repositories.knowledge.retrieval_repository import KnowledgeRetrievalRepository from repositories.knowledge.segment_repository import SQLAlchemySegmentRepository from services.enterprise import rbac_service from services.entities.knowledge_entities.rag_pipeline_entities import ( @@ -353,7 +353,7 @@ def test_fractional_built_in_dates_keep_integer_second_boundary_semantics( with sqlite_session_factory.begin() as session: session.add(make_document(document_id="document-1", doc_metadata=doc_metadata)) - filters = DatasetRetrieval.process_metadata_filter_func(0, condition, field, CREATED_AT_EPOCH, []) + filters = KnowledgeRetrievalRepository.process_metadata_filter_func(0, condition, field, CREATED_AT_EPOCH, []) with sqlite_session_factory() as session: found = session.scalars(select(_metadata_built_in_fields_Document.id).where(*filters)).all() diff --git a/api/tests/unit_tests/services/test_workflow_draft_variable_service.py b/api/tests/unit_tests/services/test_workflow_draft_variable_service.py index 0f342c06657455..925410832c5677 100644 --- a/api/tests/unit_tests/services/test_workflow_draft_variable_service.py +++ b/api/tests/unit_tests/services/test_workflow_draft_variable_service.py @@ -57,7 +57,6 @@ ) from services.workflow.variable_file_gateway import WorkflowVariableFileGateway from services.workflow.variable_service import WorkflowVariableService -from services.workflow_draft_variable_service import DraftVarLoader """Simplified unit tests for DraftVarLoader focusing on core functionality.""" @@ -110,12 +109,15 @@ class TestDraftVarLoaderSimple: @pytest.fixture def draft_var_loader(self, sqlite_engine: Engine): """Create DraftVarLoader instance for testing.""" - return DraftVarLoader( - engine=sqlite_engine, + return _ApplicationDraftVarLoader( + repository=WorkflowDraftVariableRepository( + sessions=sessionmaker(bind=sqlite_engine, expire_on_commit=False) + ), app_id="test-app-id", tenant_id="test-tenant-id", user_id="test-user-id", - fallback_variables=[], + load_file=Mock(), + file_inputs=Mock(), ) def test_load_offloaded_variable_object_type_unit(self, draft_var_loader): @@ -162,9 +164,7 @@ def test_load_offloaded_variable_object_type_unit(self, draft_var_loader): test_object = {"key1": "value1", "key2": 42} test_json_content = json.dumps(test_object, ensure_ascii=False, separators=(",", ":")) - with patch("services.workflow_draft_variable_service.storage") as mock_storage: - mock_storage.load.return_value = test_json_content.encode() - + with patch.object(draft_var_loader, "_load_file", return_value=test_json_content.encode()) as load_file: # Execute the method selector_tuple, variable = draft_var_loader._load_offloaded_variable(draft_var) @@ -176,7 +176,7 @@ def test_load_offloaded_variable_object_type_unit(self, draft_var_loader): assert variable.value == test_object # Verify method calls - mock_storage.load.assert_called_once_with("storage/key/test.json") + load_file.assert_called_once_with("storage/key/test.json") def test_load_offloaded_variable_missing_variable_file_unit(self, draft_var_loader): """Test that assertion error is raised when variable_file is None.""" @@ -275,8 +275,7 @@ def test_load_offloaded_variable_array_type_unit(self, draft_var_loader): test_array = ["item1", 2, True] test_json_content = json.dumps(test_array) - with patch("services.workflow_draft_variable_service.storage") as mock_storage: - mock_storage.load.return_value = test_json_content.encode() + with patch.object(draft_var_loader, "_load_file", return_value=test_json_content.encode()) as load_file: # Execute the method selector_tuple, variable = draft_var_loader._load_offloaded_variable(draft_var) @@ -287,7 +286,7 @@ def test_load_offloaded_variable_array_type_unit(self, draft_var_loader): assert variable.description == "test array description" # Verify method calls - mock_storage.load.assert_called_once_with("storage/key/test_array.json") + load_file.assert_called_once_with("storage/key/test_array.json") def test_load_offloaded_variable_file_type_rebuilds_storage_backed_payload(self, draft_var_loader): upload_file = UploadFile( @@ -356,12 +355,9 @@ def test_load_offloaded_variable_file_type_rebuilds_storage_backed_payload(self, } with ( - patch("services.workflow_draft_variable_service.storage") as mock_storage, - patch("models.workflow._resolve_workflow_app_tenant_id", return_value="tenant-1"), - patch("models.workflow.build_file_from_stored_mapping", return_value=rebuilt_file) as rebuild_file, + patch.object(draft_var_loader, "_load_file", return_value=json.dumps(raw_file).encode()) as load_file, + patch.object(draft_var_loader._file_inputs, "restore", return_value=rebuilt_file) as rebuild_file, ): - mock_storage.load.return_value = json.dumps(raw_file).encode() - selector_tuple, variable = draft_var_loader._load_offloaded_variable(draft_var) assert selector_tuple == ("test-node-id", "test_file") @@ -369,7 +365,8 @@ def test_load_offloaded_variable_file_type_rebuilds_storage_backed_payload(self, assert variable.name == "test_file" assert variable.description == "test file description" assert variable.value == rebuilt_file - rebuild_file.assert_called_once_with(file_mapping=raw_file, tenant_id="tenant-1") + load_file.assert_called_once_with("storage/key/test_file.json") + rebuild_file.assert_called_once_with(mapping=raw_file, tenant_id="test-tenant-id") @pytest.mark.parametrize( "sqlite_session", @@ -378,7 +375,7 @@ def test_load_offloaded_variable_file_type_rebuilds_storage_backed_payload(self, ) def test_load_variables_with_offloaded_variables_unit( self, - draft_var_loader: DraftVarLoader, + draft_var_loader: _ApplicationDraftVarLoader, sqlite_session: Session, ): """Test load_variables method with mix of regular and offloaded variables.""" @@ -413,13 +410,12 @@ def test_load_variables_with_offloaded_variables_unit( offloaded_variable.selector = ["node2", "offloaded_var"] with ( - patch("services.workflow_draft_variable_service.StorageKeyLoader"), patch.object( draft_var_loader, "_load_offloaded_variable", return_value=(("node2", "offloaded_var"), offloaded_variable), ) as load_offloaded, - patch("services.workflow_draft_variable_service.ThreadPoolExecutor") as executor_cls, + patch("services.workflow.draft_variable_service.ThreadPoolExecutor") as executor_cls, ): executor = executor_cls.return_value.__enter__.return_value executor.map.side_effect = lambda function, values: [function(value) for value in values] @@ -442,7 +438,7 @@ def test_load_variables_with_offloaded_variables_unit( ) def test_load_variables_all_offloaded_variables_unit( self, - draft_var_loader: DraftVarLoader, + draft_var_loader: _ApplicationDraftVarLoader, sqlite_session: Session, ): """Test load_variables method with only offloaded variables.""" @@ -459,10 +455,7 @@ def test_load_variables_all_offloaded_variables_unit( ) sqlite_session.commit() - with ( - patch("services.workflow_draft_variable_service.StorageKeyLoader"), - patch("services.workflow_draft_variable_service.ThreadPoolExecutor") as executor_cls, - ): + with patch("services.workflow.draft_variable_service.ThreadPoolExecutor") as executor_cls: executor = executor_cls.return_value.__enter__.return_value executor.map.return_value = [ (("node1", "offloaded_var1"), Mock()), 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 09328eb39732a0..6f86da00a869f2 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.helpers.unit.model_factories import make_account from tests.unit_tests.services.agent.test_tool_invocation_service import ( TENANT_ID, 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 fec89732605f28..9f725f3e5dfbc7 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()