test(api): cover remaining workflow typing branches

Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com>

.gemini/config.yaml

@@ -0,0 +1,13 @@
+have_fun: false
+memory_config:
+  disabled: false
+code_review:
+  disable: true
+  comment_severity_threshold: MEDIUM
+  max_review_comments: -1
+  pull_request_opened:
+    help: false
+    summary: false
+    code_review: false
+    include_drafts: false
+ignore_patterns: []

.github/actions/setup-web/action.yml

@@ -4,10 +4,9 @@ runs:
   using: composite
   steps:
     - name: Setup Vite+
-      uses: voidzero-dev/setup-vp@4a524139920f87f9f7080d3b8545acac019e1852 # v1.0.0
+      uses: voidzero-dev/setup-vp@20553a7a7429c429a74894104a2835d7fed28a72 # v1.3.0
       with:
-        node-version-file: web/.nvmrc
+        working-directory: web
+        node-version-file: .nvmrc
         cache: true
-        cache-dependency-path: web/pnpm-lock.yaml
-        run-install: |
-          cwd: ./web
+        run-install: true

.github/workflows/autofix.yml

@@ -94,11 +94,6 @@ jobs:
           find . -name "*.py" -type f -exec sed -i.bak -E 's/"([^"]+)" \| None/Optional["\1"]/g; s/'"'"'([^'"'"']+)'"'"' \| None/Optional['"'"'\1'"'"']/g' {} \;
           find . -name "*.py.bak" -type f -delete
 
-      # mdformat breaks YAML front matter in markdown files. Add --exclude for directories containing YAML front matter.
-      - name: mdformat
-        run: |
-          uvx --python 3.13 mdformat . --exclude ".agents/skills/**"
-
       - name: Setup web environment
         if: steps.web-changes.outputs.any_changed == 'true'
         uses: ./.github/actions/setup-web

.github/workflows/style.yml

@@ -84,20 +84,20 @@ jobs:
         if: steps.changed-files.outputs.any_changed == 'true'
         uses: ./.github/actions/setup-web
 
+      - name: Restore ESLint cache
+        if: steps.changed-files.outputs.any_changed == 'true'
+        id: eslint-cache-restore
+        uses: actions/cache/restore@668228422ae6a00e4ad889ee87cd7109ec5666a7 # v5.0.4
+        with:
+          path: web/.eslintcache
+          key: ${{ runner.os }}-web-eslint-${{ hashFiles('web/package.json', 'web/pnpm-lock.yaml', 'web/eslint.config.mjs', 'web/eslint.constants.mjs', 'web/plugins/eslint/**') }}-${{ github.sha }}
+          restore-keys: |
+            ${{ runner.os }}-web-eslint-${{ hashFiles('web/package.json', 'web/pnpm-lock.yaml', 'web/eslint.config.mjs', 'web/eslint.constants.mjs', 'web/plugins/eslint/**') }}-
+
       - name: Web style check
         if: steps.changed-files.outputs.any_changed == 'true'
         working-directory: ./web
-        run: |
-          vp run lint:ci
-        # pnpm run lint:report
-        # continue-on-error: true
-
-      # - name: Annotate Code
-      #   if: steps.changed-files.outputs.any_changed == 'true' && github.event_name == 'pull_request'
-      #   uses: DerLev/eslint-annotations@51347b3a0abfb503fc8734d5ae31c4b151297fae
-      #   with:
-      #     eslint-report: web/eslint_report.json
-      #     github-token: ${{ secrets.GITHUB_TOKEN }}
+        run: vp run lint:ci
 
       - name: Web tsslint
         if: steps.changed-files.outputs.any_changed == 'true'
@@ -114,6 +114,13 @@ jobs:
         working-directory: ./web
         run: vp run knip
 
+      - name: Save ESLint cache
+        if: steps.changed-files.outputs.any_changed == 'true' && success() && steps.eslint-cache-restore.outputs.cache-hit != 'true'
+        uses: actions/cache/save@668228422ae6a00e4ad889ee87cd7109ec5666a7 # v5.0.4
+        with:
+          path: web/.eslintcache
+          key: ${{ steps.eslint-cache-restore.outputs.cache-primary-key }}
+
   superlinter:
     name: SuperLinter
     runs-on: ubuntu-latest

.github/workflows/translate-i18n-claude.yml

@@ -120,7 +120,7 @@ jobs:
 
       - name: Run Claude Code for Translation Sync
         if: steps.detect_changes.outputs.CHANGED_FILES != ''
-        uses: anthropics/claude-code-action@df37d2f0760a4b5683a6e617c9325bc1a36443f6 # v1.0.75
+        uses: anthropics/claude-code-action@ff9acae5886d41a99ed4ec14b7dc147d55834722 # v1.0.77
         with:
           anthropic_api_key: ${{ secrets.ANTHROPIC_API_KEY }}
           github_token: ${{ secrets.GITHUB_TOKEN }}

CONTRIBUTING.md

@@ -97,3 +97,8 @@ Feel free to reach out if you encounter any issues during the setup process.
 ## Getting Help
 
 If you ever get stuck or get a burning question while contributing, simply shoot your queries our way via the related GitHub issue, or hop onto our [Discord](https://discord.gg/8Tpq4AcN9c) for a quick chat.
+
+## Automated Agent Contributions
+
+> [!NOTE]
+> If you are an automated agent， just add 🤖🤖🤖 to the end of the PR title to opt-in.

api/.env.example

@@ -353,6 +353,9 @@ BAIDU_VECTOR_DB_SHARD=1
 BAIDU_VECTOR_DB_REPLICAS=3
 BAIDU_VECTOR_DB_INVERTED_INDEX_ANALYZER=DEFAULT_ANALYZER
 BAIDU_VECTOR_DB_INVERTED_INDEX_PARSER_MODE=COARSE_MODE
+BAIDU_VECTOR_DB_AUTO_BUILD_ROW_COUNT_INCREMENT=500
+BAIDU_VECTOR_DB_AUTO_BUILD_ROW_COUNT_INCREMENT_RATIO=0.05
+BAIDU_VECTOR_DB_REBUILD_INDEX_TIMEOUT_IN_SECONDS=300
 
 # Upstash configuration
 UPSTASH_VECTOR_URL=your-server-url

api/commands/vector.py

@@ -10,6 +10,7 @@ from configs import dify_config
 from core.rag.datasource.vdb.vector_factory import Vector
 from core.rag.datasource.vdb.vector_type import VectorType
 from core.rag.index_processor.constant.built_in_field import BuiltInField
+from core.rag.index_processor.constant.index_type import IndexStructureType, IndexTechniqueType
 from core.rag.models.document import ChildDocument, Document
 from extensions.ext_database import db
 from models.dataset import Dataset, DatasetCollectionBinding, DatasetMetadata, DatasetMetadataBinding, DocumentSegment
@@ -85,7 +86,7 @@ def migrate_annotation_vector_database():
                 dataset = Dataset(
                     id=app.id,
                     tenant_id=app.tenant_id,
-                    indexing_technique="high_quality",
+                    indexing_technique=IndexTechniqueType.HIGH_QUALITY,
                     embedding_model_provider=dataset_collection_binding.provider_name,
                     embedding_model=dataset_collection_binding.model_name,
                     collection_binding_id=dataset_collection_binding.id,
@@ -177,7 +178,9 @@ def migrate_knowledge_vector_database():
     while True:
         try:
             stmt = (
-                select(Dataset).where(Dataset.indexing_technique == "high_quality").order_by(Dataset.created_at.desc())
+                select(Dataset)
+                .where(Dataset.indexing_technique == IndexTechniqueType.HIGH_QUALITY)
+                .order_by(Dataset.created_at.desc())
             )
 
             datasets = db.paginate(select=stmt, page=page, per_page=50, max_per_page=50, error_out=False)
@@ -269,7 +272,7 @@ def migrate_knowledge_vector_database():
                                 "dataset_id": segment.dataset_id,
                             },
                         )
-                        if dataset_document.doc_form == "hierarchical_model":
+                        if dataset_document.doc_form == IndexStructureType.PARENT_CHILD_INDEX:
                             child_chunks = segment.get_child_chunks()
                             if child_chunks:
                                 child_documents = []

api/configs/middleware/vdb/baidu_vector_config.py

@@ -51,3 +51,18 @@ class BaiduVectorDBConfig(BaseSettings):
         description="Parser mode for inverted index in Baidu Vector Database (default is COARSE_MODE)",
         default="COARSE_MODE",
     )
+
+    BAIDU_VECTOR_DB_AUTO_BUILD_ROW_COUNT_INCREMENT: int = Field(
+        description="Auto build row count increment threshold (default is 500)",
+        default=500,
+    )
+
+    BAIDU_VECTOR_DB_AUTO_BUILD_ROW_COUNT_INCREMENT_RATIO: float = Field(
+        description="Auto build row count increment ratio threshold (default is 0.05)",
+        default=0.05,
+    )
+
+    BAIDU_VECTOR_DB_REBUILD_INDEX_TIMEOUT_IN_SECONDS: int = Field(
+        description="Timeout in seconds for rebuilding the index in Baidu Vector Database (default is 3600 seconds)",
+        default=300,
+    )

api/controllers/console/apikey.py

@@ -1,7 +1,7 @@
 import flask_restx
 from flask_restx import Resource, fields, marshal_with
 from flask_restx._http import HTTPStatus
-from sqlalchemy import select
+from sqlalchemy import delete, func, select
 from sqlalchemy.orm import Session
 from werkzeug.exceptions import Forbidden
 
@@ -9,6 +9,7 @@ from extensions.ext_database import db
 from libs.helper import TimestampField
 from libs.login import current_account_with_tenant, login_required
 from models.dataset import Dataset
+from models.enums import ApiTokenType
 from models.model import ApiToken, App
 from services.api_token_service import ApiTokenCache
 
@@ -33,16 +34,10 @@ api_key_list_model = console_ns.model(
 
 
 def _get_resource(resource_id, tenant_id, resource_model):
-    if resource_model == App:
-        with Session(db.engine) as session:
-            resource = session.execute(
-                select(resource_model).filter_by(id=resource_id, tenant_id=tenant_id)
-            ).scalar_one_or_none()
-    else:
-        with Session(db.engine) as session:
-            resource = session.execute(
-                select(resource_model).filter_by(id=resource_id, tenant_id=tenant_id)
-            ).scalar_one_or_none()
+    with Session(db.engine) as session:
+        resource = session.execute(
+            select(resource_model).filter_by(id=resource_id, tenant_id=tenant_id)
+        ).scalar_one_or_none()
 
     if resource is None:
         flask_restx.abort(HTTPStatus.NOT_FOUND, message=f"{resource_model.__name__} not found.")
@@ -53,7 +48,7 @@ def _get_resource(resource_id, tenant_id, resource_model):
 class BaseApiKeyListResource(Resource):
     method_decorators = [account_initialization_required, login_required, setup_required]
 
-    resource_type: str | None = None
+    resource_type: ApiTokenType | None = None
     resource_model: type | None = None
     resource_id_field: str | None = None
     token_prefix: str | None = None
@@ -80,10 +75,13 @@ class BaseApiKeyListResource(Resource):
         resource_id = str(resource_id)
         _, current_tenant_id = current_account_with_tenant()
         _get_resource(resource_id, current_tenant_id, self.resource_model)
-        current_key_count = (
-            db.session.query(ApiToken)
-            .where(ApiToken.type == self.resource_type, getattr(ApiToken, self.resource_id_field) == resource_id)
-            .count()
+        current_key_count: int = (
+            db.session.scalar(
+                select(func.count(ApiToken.id)).where(
+                    ApiToken.type == self.resource_type, getattr(ApiToken, self.resource_id_field) == resource_id
+                )
+            )
+            or 0
         )
 
         if current_key_count >= self.max_keys:
@@ -94,6 +92,7 @@ class BaseApiKeyListResource(Resource):
             )
 
         key = ApiToken.generate_api_key(self.token_prefix or "", 24)
+        assert self.resource_type is not None, "resource_type must be set"
         api_token = ApiToken()
         setattr(api_token, self.resource_id_field, resource_id)
         api_token.tenant_id = current_tenant_id
@@ -107,7 +106,7 @@ class BaseApiKeyListResource(Resource):
 class BaseApiKeyResource(Resource):
     method_decorators = [account_initialization_required, login_required, setup_required]
 
-    resource_type: str | None = None
+    resource_type: ApiTokenType | None = None
     resource_model: type | None = None
     resource_id_field: str | None = None
 
@@ -119,14 +118,14 @@ class BaseApiKeyResource(Resource):
         if not current_user.is_admin_or_owner:
             raise Forbidden()
 
-        key = (
-            db.session.query(ApiToken)
+        key = db.session.scalar(
+            select(ApiToken)
             .where(
                 getattr(ApiToken, self.resource_id_field) == resource_id,
                 ApiToken.type == self.resource_type,
                 ApiToken.id == api_key_id,
             )
-            .first()
+            .limit(1)
         )
 
         if key is None:
@@ -137,7 +136,7 @@ class BaseApiKeyResource(Resource):
         assert key is not None  # nosec - for type checker only
         ApiTokenCache.delete(key.token, key.type)
 
-        db.session.query(ApiToken).where(ApiToken.id == api_key_id).delete()
+        db.session.execute(delete(ApiToken).where(ApiToken.id == api_key_id))
         db.session.commit()
 
         return {"result": "success"}, 204
@@ -162,7 +161,7 @@ class AppApiKeyListResource(BaseApiKeyListResource):
         """Create a new API key for an app"""
         return super().post(resource_id)
 
-    resource_type = "app"
+    resource_type = ApiTokenType.APP
     resource_model = App
     resource_id_field = "app_id"
     token_prefix = "app-"
@@ -178,7 +177,7 @@ class AppApiKeyResource(BaseApiKeyResource):
         """Delete an API key for an app"""
         return super().delete(resource_id, api_key_id)
 
-    resource_type = "app"
+    resource_type = ApiTokenType.APP
     resource_model = App
     resource_id_field = "app_id"
 
@@ -202,7 +201,7 @@ class DatasetApiKeyListResource(BaseApiKeyListResource):
         """Create a new API key for a dataset"""
         return super().post(resource_id)
 
-    resource_type = "dataset"
+    resource_type = ApiTokenType.DATASET
     resource_model = Dataset
     resource_id_field = "dataset_id"
     token_prefix = "ds-"
@@ -218,6 +217,6 @@ class DatasetApiKeyResource(BaseApiKeyResource):
         """Delete an API key for a dataset"""
         return super().delete(resource_id, api_key_id)
 
-    resource_type = "dataset"
+    resource_type = ApiTokenType.DATASET
     resource_model = Dataset
     resource_id_field = "dataset_id"

api/controllers/console/app/app.py

@@ -95,7 +95,7 @@ class CreateAppPayload(BaseModel):
     name: str = Field(..., min_length=1, description="App name")
     description: str | None = Field(default=None, description="App description (max 400 chars)", max_length=400)
     mode: Literal["chat", "agent-chat", "advanced-chat", "workflow", "completion"] = Field(..., description="App mode")
-    icon_type: str | None = Field(default=None, description="Icon type")
+    icon_type: IconType | None = Field(default=None, description="Icon type")
     icon: str | None = Field(default=None, description="Icon")
     icon_background: str | None = Field(default=None, description="Icon background color")
 
@@ -103,7 +103,7 @@ class CreateAppPayload(BaseModel):
 class UpdateAppPayload(BaseModel):
     name: str = Field(..., min_length=1, description="App name")
     description: str | None = Field(default=None, description="App description (max 400 chars)", max_length=400)
-    icon_type: str | None = Field(default=None, description="Icon type")
+    icon_type: IconType | None = Field(default=None, description="Icon type")
     icon: str | None = Field(default=None, description="Icon")
     icon_background: str | None = Field(default=None, description="Icon background color")
     use_icon_as_answer_icon: bool | None = Field(default=None, description="Use icon as answer icon")
@@ -113,7 +113,7 @@ class UpdateAppPayload(BaseModel):
 class CopyAppPayload(BaseModel):
     name: str | None = Field(default=None, description="Name for the copied app")
     description: str | None = Field(default=None, description="Description for the copied app", max_length=400)
-    icon_type: str | None = Field(default=None, description="Icon type")
+    icon_type: IconType | None = Field(default=None, description="Icon type")
     icon: str | None = Field(default=None, description="Icon")
     icon_background: str | None = Field(default=None, description="Icon background color")
 
@@ -594,7 +594,7 @@ class AppApi(Resource):
         args_dict: AppService.ArgsDict = {
             "name": args.name,
             "description": args.description or "",
-            "icon_type": args.icon_type or "",
+            "icon_type": args.icon_type,
             "icon": args.icon or "",
             "icon_background": args.icon_background or "",
             "use_icon_as_answer_icon": args.use_icon_as_answer_icon or False,

api/controllers/console/app/conversation.py

@@ -5,7 +5,7 @@ from flask import abort, request
 from flask_restx import Resource, fields, marshal_with
 from pydantic import BaseModel, Field, field_validator
 from sqlalchemy import func, or_
-from sqlalchemy.orm import joinedload
+from sqlalchemy.orm import selectinload
 from werkzeug.exceptions import NotFound
 
 from controllers.console import console_ns
@@ -376,8 +376,12 @@ class CompletionConversationApi(Resource):
 
         # FIXME, the type ignore in this file
         if args.annotation_status == "annotated":
-            query = query.options(joinedload(Conversation.message_annotations)).join(  # type: ignore
-                MessageAnnotation, MessageAnnotation.conversation_id == Conversation.id
+            query = (
+                query.options(selectinload(Conversation.message_annotations))  # type: ignore[arg-type]
+                .join(  # type: ignore
+                    MessageAnnotation, MessageAnnotation.conversation_id == Conversation.id
+                )
+                .distinct()
             )
         elif args.annotation_status == "not_annotated":
             query = (
@@ -454,9 +458,7 @@ class ChatConversationApi(Resource):
         args = ChatConversationQuery.model_validate(request.args.to_dict(flat=True))  # type: ignore
 
         subquery = (
-            db.session.query(
-                Conversation.id.label("conversation_id"), EndUser.session_id.label("from_end_user_session_id")
-            )
+            sa.select(Conversation.id.label("conversation_id"), EndUser.session_id.label("from_end_user_session_id"))
             .outerjoin(EndUser, Conversation.from_end_user_id == EndUser.id)
             .subquery()
         )
@@ -511,8 +513,12 @@ class ChatConversationApi(Resource):
 
         match args.annotation_status:
             case "annotated":
-                query = query.options(joinedload(Conversation.message_annotations)).join(  # type: ignore
-                    MessageAnnotation, MessageAnnotation.conversation_id == Conversation.id
+                query = (
+                    query.options(selectinload(Conversation.message_annotations))  # type: ignore[arg-type]
+                    .join(  # type: ignore
+                        MessageAnnotation, MessageAnnotation.conversation_id == Conversation.id
+                    )
+                    .distinct()
                 )
             case "not_annotated":
                 query = (
@@ -587,10 +593,8 @@ class ChatConversationDetailApi(Resource):
 
 def _get_conversation(app_model, conversation_id):
     current_user, _ = current_account_with_tenant()
-    conversation = (
-        db.session.query(Conversation)
-        .where(Conversation.id == conversation_id, Conversation.app_id == app_model.id)
-        .first()
+    conversation = db.session.scalar(
+        sa.select(Conversation).where(Conversation.id == conversation_id, Conversation.app_id == app_model.id).limit(1)
     )
 
     if not conversation:

api/controllers/console/app/generator.py

@@ -168,7 +168,7 @@ class InstructionGenerateApi(Resource):
         try:
             # Generate from nothing for a workflow node
             if (args.current in (code_template, "")) and args.node_id != "":
-                app = db.session.query(App).where(App.id == args.flow_id).first()
+                app = db.session.get(App, args.flow_id)
                 if not app:
                     return {"error": f"app {args.flow_id} not found"}, 400
                 workflow = WorkflowService().get_draft_workflow(app_model=app)

api/controllers/console/app/mcp_server.py

@@ -2,6 +2,7 @@ import json
 
 from flask_restx import Resource, marshal_with
 from pydantic import BaseModel, Field
+from sqlalchemy import select
 from werkzeug.exceptions import NotFound
 
 from controllers.console import console_ns
@@ -47,7 +48,7 @@ class AppMCPServerController(Resource):
     @get_app_model
     @marshal_with(app_server_model)
     def get(self, app_model):
-        server = db.session.query(AppMCPServer).where(AppMCPServer.app_id == app_model.id).first()
+        server = db.session.scalar(select(AppMCPServer).where(AppMCPServer.app_id == app_model.id).limit(1))
         return server
 
     @console_ns.doc("create_app_mcp_server")
@@ -98,7 +99,7 @@ class AppMCPServerController(Resource):
     @edit_permission_required
     def put(self, app_model):
         payload = MCPServerUpdatePayload.model_validate(console_ns.payload or {})
-        server = db.session.query(AppMCPServer).where(AppMCPServer.id == payload.id).first()
+        server = db.session.get(AppMCPServer, payload.id)
         if not server:
             raise NotFound()
 
@@ -135,11 +136,10 @@ class AppMCPServerRefreshController(Resource):
     @edit_permission_required
     def get(self, server_id):
         _, current_tenant_id = current_account_with_tenant()
-        server = (
-            db.session.query(AppMCPServer)
-            .where(AppMCPServer.id == server_id)
-            .where(AppMCPServer.tenant_id == current_tenant_id)
-            .first()
+        server = db.session.scalar(
+            select(AppMCPServer)
+            .where(AppMCPServer.id == server_id, AppMCPServer.tenant_id == current_tenant_id)
+            .limit(1)
         )
         if not server:
             raise NotFound()

api/controllers/console/app/message.py

@@ -4,7 +4,7 @@ from typing import Literal
 from flask import request
 from flask_restx import Resource, fields, marshal_with
 from pydantic import BaseModel, Field, field_validator
-from sqlalchemy import exists, select
+from sqlalchemy import exists, func, select
 from werkzeug.exceptions import InternalServerError, NotFound
 
 from controllers.common.schema import register_schema_models
@@ -244,27 +244,25 @@ class ChatMessageListApi(Resource):
     def get(self, app_model):
         args = ChatMessagesQuery.model_validate(request.args.to_dict())
 
-        conversation = (
-            db.session.query(Conversation)
+        conversation = db.session.scalar(
+            select(Conversation)
             .where(Conversation.id == args.conversation_id, Conversation.app_id == app_model.id)
-            .first()
+            .limit(1)
         )
 
         if not conversation:
             raise NotFound("Conversation Not Exists.")
 
         if args.first_id:
-            first_message = (
-                db.session.query(Message)
-                .where(Message.conversation_id == conversation.id, Message.id == args.first_id)
-                .first()
+            first_message = db.session.scalar(
+                select(Message).where(Message.conversation_id == conversation.id, Message.id == args.first_id).limit(1)
             )
 
             if not first_message:
                 raise NotFound("First message not found")
 
-            history_messages = (
-                db.session.query(Message)
+            history_messages = db.session.scalars(
+                select(Message)
                 .where(
                     Message.conversation_id == conversation.id,
                     Message.created_at < first_message.created_at,
@@ -272,16 +270,14 @@ class ChatMessageListApi(Resource):
                 )
                 .order_by(Message.created_at.desc())
                 .limit(args.limit)
-                .all()
-            )
+            ).all()
         else:
-            history_messages = (
-                db.session.query(Message)
+            history_messages = db.session.scalars(
+                select(Message)
                 .where(Message.conversation_id == conversation.id)
                 .order_by(Message.created_at.desc())
                 .limit(args.limit)
-                .all()
-            )
+            ).all()
 
         # Initialize has_more based on whether we have a full page
         if len(history_messages) == args.limit:
@@ -326,7 +322,9 @@ class MessageFeedbackApi(Resource):
 
         message_id = str(args.message_id)
 
-        message = db.session.query(Message).where(Message.id == message_id, Message.app_id == app_model.id).first()
+        message = db.session.scalar(
+            select(Message).where(Message.id == message_id, Message.app_id == app_model.id).limit(1)
+        )
 
         if not message:
             raise NotFound("Message Not Exists.")
@@ -375,7 +373,9 @@ class MessageAnnotationCountApi(Resource):
     @login_required
     @account_initialization_required
     def get(self, app_model):
-        count = db.session.query(MessageAnnotation).where(MessageAnnotation.app_id == app_model.id).count()
+        count = db.session.scalar(
+            select(func.count(MessageAnnotation.id)).where(MessageAnnotation.app_id == app_model.id)
+        )
 
         return {"count": count}
 
@@ -479,7 +479,9 @@ class MessageApi(Resource):
     def get(self, app_model, message_id: str):
         message_id = str(message_id)
 
-        message = db.session.query(Message).where(Message.id == message_id, Message.app_id == app_model.id).first()
+        message = db.session.scalar(
+            select(Message).where(Message.id == message_id, Message.app_id == app_model.id).limit(1)
+        )
 
         if not message:
             raise NotFound("Message Not Exists.")

api/controllers/console/app/model_config.py

@@ -69,9 +69,7 @@ class ModelConfigResource(Resource):
 
         if app_model.mode == AppMode.AGENT_CHAT or app_model.is_agent:
             # get original app model config
-            original_app_model_config = (
-                db.session.query(AppModelConfig).where(AppModelConfig.id == app_model.app_model_config_id).first()
-            )
+            original_app_model_config = db.session.get(AppModelConfig, app_model.app_model_config_id)
             if original_app_model_config is None:
                 raise ValueError("Original app model config not found")
             agent_mode = original_app_model_config.agent_mode_dict

api/controllers/console/app/site.py

@@ -2,6 +2,7 @@ from typing import Literal
 
 from flask_restx import Resource, marshal_with
 from pydantic import BaseModel, Field, field_validator
+from sqlalchemy import select
 from werkzeug.exceptions import NotFound
 
 from constants.languages import supported_language
@@ -75,7 +76,7 @@ class AppSite(Resource):
     def post(self, app_model):
         args = AppSiteUpdatePayload.model_validate(console_ns.payload or {})
         current_user, _ = current_account_with_tenant()
-        site = db.session.query(Site).where(Site.app_id == app_model.id).first()
+        site = db.session.scalar(select(Site).where(Site.app_id == app_model.id).limit(1))
         if not site:
             raise NotFound
 
@@ -124,7 +125,7 @@ class AppSiteAccessTokenReset(Resource):
     @marshal_with(app_site_model)
     def post(self, app_model):
         current_user, _ = current_account_with_tenant()
-        site = db.session.query(Site).where(Site.app_id == app_model.id).first()
+        site = db.session.scalar(select(Site).where(Site.app_id == app_model.id).limit(1))
 
         if not site:
             raise NotFound

api/controllers/console/app/workflow.py

@@ -7,7 +7,7 @@ from flask import abort, request
 from flask_restx import Resource, fields, marshal_with
 from pydantic import BaseModel, Field, field_validator
 from sqlalchemy.orm import Session
-from werkzeug.exceptions import Forbidden, InternalServerError, NotFound
+from werkzeug.exceptions import BadRequest, Forbidden, InternalServerError, NotFound
 
 import services
 from controllers.console import console_ns
@@ -46,13 +46,14 @@ from models import App
 from models.model import AppMode
 from models.workflow import Workflow
 from services.app_generate_service import AppGenerateService
-from services.errors.app import WorkflowHashNotEqualError
+from services.errors.app import IsDraftWorkflowError, WorkflowHashNotEqualError, WorkflowNotFoundError
 from services.errors.llm import InvokeRateLimitError
 from services.workflow_service import DraftWorkflowDeletionError, WorkflowInUseError, WorkflowService
 
 logger = logging.getLogger(__name__)
 LISTENING_RETRY_IN = 2000
 DEFAULT_REF_TEMPLATE_SWAGGER_2_0 = "#/definitions/{model}"
+RESTORE_SOURCE_WORKFLOW_MUST_BE_PUBLISHED_MESSAGE = "source workflow must be published"
 
 # Register models for flask_restx to avoid dict type issues in Swagger
 # Register in dependency order: base models first, then dependent models
@@ -284,7 +285,9 @@ class DraftWorkflowApi(Resource):
         workflow_service = WorkflowService()
 
         try:
-            environment_variables_list = args.get("environment_variables") or []
+            environment_variables_list = Workflow.normalize_environment_variable_mappings(
+                args.get("environment_variables") or [],
+            )
             environment_variables = [
                 variable_factory.build_environment_variable_from_mapping(obj) for obj in environment_variables_list
             ]
@@ -994,6 +997,43 @@ class PublishedAllWorkflowApi(Resource):
             }
 
 
+@console_ns.route("/apps/<uuid:app_id>/workflows/<string:workflow_id>/restore")
+class DraftWorkflowRestoreApi(Resource):
+    @console_ns.doc("restore_workflow_to_draft")
+    @console_ns.doc(description="Restore a published workflow version into the draft workflow")
+    @console_ns.doc(params={"app_id": "Application ID", "workflow_id": "Published workflow ID"})
+    @console_ns.response(200, "Workflow restored successfully")
+    @console_ns.response(400, "Source workflow must be published")
+    @console_ns.response(404, "Workflow not found")
+    @setup_required
+    @login_required
+    @account_initialization_required
+    @get_app_model(mode=[AppMode.ADVANCED_CHAT, AppMode.WORKFLOW])
+    @edit_permission_required
+    def post(self, app_model: App, workflow_id: str):
+        current_user, _ = current_account_with_tenant()
+        workflow_service = WorkflowService()
+
+        try:
+            workflow = workflow_service.restore_published_workflow_to_draft(
+                app_model=app_model,
+                workflow_id=workflow_id,
+                account=current_user,
+            )
+        except IsDraftWorkflowError as exc:
+            raise BadRequest(RESTORE_SOURCE_WORKFLOW_MUST_BE_PUBLISHED_MESSAGE) from exc
+        except WorkflowNotFoundError as exc:
+            raise NotFound(str(exc)) from exc
+        except ValueError as exc:
+            raise BadRequest(str(exc)) from exc
+
+        return {
+            "result": "success",
+            "hash": workflow.unique_hash,
+            "updated_at": TimestampField().format(workflow.updated_at or workflow.created_at),
+        }
+
+
 @console_ns.route("/apps/<uuid:app_id>/workflows/<string:workflow_id>")
 class WorkflowByIdApi(Resource):
     @console_ns.doc("update_workflow_by_id")

api/controllers/console/app/wraps.py

@@ -2,6 +2,8 @@ from collections.abc import Callable
 from functools import wraps
 from typing import ParamSpec, TypeVar, Union
 
+from sqlalchemy import select
+
 from controllers.console.app.error import AppNotFoundError
 from extensions.ext_database import db
 from libs.login import current_account_with_tenant
@@ -15,16 +17,14 @@ R1 = TypeVar("R1")
 
 def _load_app_model(app_id: str) -> App | None:
     _, current_tenant_id = current_account_with_tenant()
-    app_model = (
-        db.session.query(App)
-        .where(App.id == app_id, App.tenant_id == current_tenant_id, App.status == "normal")
-        .first()
+    app_model = db.session.scalar(
+        select(App).where(App.id == app_id, App.tenant_id == current_tenant_id, App.status == "normal").limit(1)
     )
     return app_model
 
 
 def _load_app_model_with_trial(app_id: str) -> App | None:
-    app_model = db.session.query(App).where(App.id == app_id, App.status == "normal").first()
+    app_model = db.session.scalar(select(App).where(App.id == app_id, App.status == "normal").limit(1))
     return app_model
 
 

api/controllers/console/auth/email_register.py

@@ -1,7 +1,7 @@
 from flask import request
 from flask_restx import Resource
 from pydantic import BaseModel, Field, field_validator
-from sqlalchemy.orm import Session
+from sqlalchemy.orm import sessionmaker
 
 from configs import dify_config
 from constants.languages import languages
@@ -73,7 +73,7 @@ class EmailRegisterSendEmailApi(Resource):
         if dify_config.BILLING_ENABLED and BillingService.is_email_in_freeze(normalized_email):
             raise AccountInFreezeError()
 
-        with Session(db.engine) as session:
+        with sessionmaker(db.engine).begin() as session:
             account = AccountService.get_account_by_email_with_case_fallback(args.email, session=session)
         token = AccountService.send_email_register_email(email=normalized_email, account=account, language=language)
         return {"result": "success", "data": token}
@@ -145,7 +145,7 @@ class EmailRegisterResetApi(Resource):
         email = register_data.get("email", "")
         normalized_email = email.lower()
 
-        with Session(db.engine) as session:
+        with sessionmaker(db.engine).begin() as session:
             account = AccountService.get_account_by_email_with_case_fallback(email, session=session)
 
             if account:

api/controllers/console/auth/forgot_password.py

@@ -4,7 +4,7 @@ import secrets
 from flask import request
 from flask_restx import Resource
 from pydantic import BaseModel, Field, field_validator
-from sqlalchemy.orm import Session
+from sqlalchemy.orm import sessionmaker
 
 from controllers.common.schema import register_schema_models
 from controllers.console import console_ns
@@ -102,7 +102,7 @@ class ForgotPasswordSendEmailApi(Resource):
         else:
             language = "en-US"
 
-        with Session(db.engine) as session:
+        with sessionmaker(db.engine).begin() as session:
             account = AccountService.get_account_by_email_with_case_fallback(args.email, session=session)
 
         token = AccountService.send_reset_password_email(
@@ -201,7 +201,7 @@ class ForgotPasswordResetApi(Resource):
         password_hashed = hash_password(args.new_password, salt)
 
         email = reset_data.get("email", "")
-        with Session(db.engine) as session:
+        with sessionmaker(db.engine).begin() as session:
             account = AccountService.get_account_by_email_with_case_fallback(email, session=session)
 
             if account:
@@ -215,7 +215,6 @@ class ForgotPasswordResetApi(Resource):
         # Update existing account credentials
         account.password = base64.b64encode(password_hashed).decode()
         account.password_salt = base64.b64encode(salt).decode()
-        session.commit()
 
         # Create workspace if needed
         if (

api/controllers/console/auth/oauth.py

@@ -1,9 +1,10 @@
 import logging
+import urllib.parse
 
 import httpx
 from flask import current_app, redirect, request
 from flask_restx import Resource
-from sqlalchemy.orm import Session
+from sqlalchemy.orm import sessionmaker
 from werkzeug.exceptions import Unauthorized
 
 from configs import dify_config
@@ -112,6 +113,9 @@ class OAuthCallback(Resource):
                 error_text = e.response.text
             logger.exception("An error occurred during the OAuth process with %s: %s", provider, error_text)
             return {"error": "OAuth process failed"}, 400
+        except ValueError as e:
+            logger.warning("OAuth error with %s", provider, exc_info=True)
+            return redirect(f"{dify_config.CONSOLE_WEB_URL}/signin?message={urllib.parse.quote(str(e))}")
 
         if invite_token and RegisterService.is_valid_invite_token(invite_token):
             invitation = RegisterService.get_invitation_by_token(token=invite_token)
@@ -176,7 +180,7 @@ def _get_account_by_openid_or_email(provider: str, user_info: OAuthUserInfo) ->
     account: Account | None = Account.get_by_openid(provider, user_info.id)
 
     if not account:
-        with Session(db.engine) as session:
+        with sessionmaker(db.engine).begin() as session:
             account = AccountService.get_account_by_email_with_case_fallback(user_info.email, session=session)
 
     return account

api/controllers/console/datasets/datasets.py

@@ -3,7 +3,7 @@ from typing import Any, cast
 from flask import request
 from flask_restx import Resource, fields, marshal, marshal_with
 from pydantic import BaseModel, Field, field_validator
-from sqlalchemy import select
+from sqlalchemy import func, select
 from werkzeug.exceptions import Forbidden, NotFound
 
 import services
@@ -29,6 +29,7 @@ from core.provider_manager import ProviderManager
 from core.rag.datasource.vdb.vector_type import VectorType
 from core.rag.extractor.entity.datasource_type import DatasourceType
 from core.rag.extractor.entity.extract_setting import ExtractSetting, NotionInfo, WebsiteInfo
+from core.rag.index_processor.constant.index_type import IndexTechniqueType
 from core.rag.retrieval.retrieval_methods import RetrievalMethod
 from dify_graph.model_runtime.entities.model_entities import ModelType
 from extensions.ext_database import db
@@ -54,7 +55,7 @@ from fields.document_fields import document_status_fields
 from libs.login import current_account_with_tenant, login_required
 from models import ApiToken, Dataset, Document, DocumentSegment, UploadFile
 from models.dataset import DatasetPermission, DatasetPermissionEnum
-from models.enums import SegmentStatus
+from models.enums import ApiTokenType, SegmentStatus
 from models.provider_ids import ModelProviderID
 from services.api_token_service import ApiTokenCache
 from services.dataset_service import DatasetPermissionService, DatasetService, DocumentService
@@ -355,7 +356,7 @@ class DatasetListApi(Resource):
 
         for item in data:
             # convert embedding_model_provider to plugin standard format
-            if item["indexing_technique"] == "high_quality" and item["embedding_model_provider"]:
+            if item["indexing_technique"] == IndexTechniqueType.HIGH_QUALITY and item["embedding_model_provider"]:
                 item["embedding_model_provider"] = str(ModelProviderID(item["embedding_model_provider"]))
                 item_model = f"{item['embedding_model']}:{item['embedding_model_provider']}"
                 if item_model in model_names:
@@ -436,7 +437,7 @@ class DatasetApi(Resource):
         except services.errors.account.NoPermissionError as e:
             raise Forbidden(str(e))
         data = cast(dict[str, Any], marshal(dataset, dataset_detail_fields))
-        if dataset.indexing_technique == "high_quality":
+        if dataset.indexing_technique == IndexTechniqueType.HIGH_QUALITY:
             if dataset.embedding_model_provider:
                 provider_id = ModelProviderID(dataset.embedding_model_provider)
                 data["embedding_model_provider"] = str(provider_id)
@@ -454,7 +455,7 @@ class DatasetApi(Resource):
         for embedding_model in embedding_models:
             model_names.append(f"{embedding_model.model}:{embedding_model.provider.provider}")
 
-        if data["indexing_technique"] == "high_quality":
+        if data["indexing_technique"] == IndexTechniqueType.HIGH_QUALITY:
             item_model = f"{data['embedding_model']}:{data['embedding_model_provider']}"
             if item_model in model_names:
                 data["embedding_available"] = True
@@ -485,7 +486,7 @@ class DatasetApi(Resource):
         current_user, current_tenant_id = current_account_with_tenant()
         # check embedding model setting
         if (
-            payload.indexing_technique == "high_quality"
+            payload.indexing_technique == IndexTechniqueType.HIGH_QUALITY
             and payload.embedding_model_provider is not None
             and payload.embedding_model is not None
         ):
@@ -738,20 +739,23 @@ class DatasetIndexingStatusApi(Resource):
         documents_status = []
         for document in documents:
             completed_segments = (
-                db.session.query(DocumentSegment)
-                .where(
-                    DocumentSegment.completed_at.isnot(None),
-                    DocumentSegment.document_id == str(document.id),
-                    DocumentSegment.status != SegmentStatus.RE_SEGMENT,
+                db.session.scalar(
+                    select(func.count(DocumentSegment.id)).where(
+                        DocumentSegment.completed_at.isnot(None),
+                        DocumentSegment.document_id == str(document.id),
+                        DocumentSegment.status != SegmentStatus.RE_SEGMENT,
+                    )
                 )
-                .count()
+                or 0
             )
             total_segments = (
-                db.session.query(DocumentSegment)
-                .where(
-                    DocumentSegment.document_id == str(document.id), DocumentSegment.status != SegmentStatus.RE_SEGMENT
+                db.session.scalar(
+                    select(func.count(DocumentSegment.id)).where(
+                        DocumentSegment.document_id == str(document.id),
+                        DocumentSegment.status != SegmentStatus.RE_SEGMENT,
+                    )
                 )
-                .count()
+                or 0
             )
             # Create a dictionary with document attributes and additional fields
             document_dict = {
@@ -777,7 +781,7 @@ class DatasetIndexingStatusApi(Resource):
 class DatasetApiKeyApi(Resource):
     max_keys = 10
     token_prefix = "dataset-"
-    resource_type = "dataset"
+    resource_type = ApiTokenType.DATASET
 
     @console_ns.doc("get_dataset_api_keys")
     @console_ns.doc(description="Get dataset API keys")
@@ -802,9 +806,12 @@ class DatasetApiKeyApi(Resource):
         _, current_tenant_id = current_account_with_tenant()
 
         current_key_count = (
-            db.session.query(ApiToken)
-            .where(ApiToken.type == self.resource_type, ApiToken.tenant_id == current_tenant_id)
-            .count()
+            db.session.scalar(
+                select(func.count(ApiToken.id)).where(
+                    ApiToken.type == self.resource_type, ApiToken.tenant_id == current_tenant_id
+                )
+            )
+            or 0
         )
 
         if current_key_count >= self.max_keys:
@@ -826,7 +833,7 @@ class DatasetApiKeyApi(Resource):
 
 @console_ns.route("/datasets/api-keys/<uuid:api_key_id>")
 class DatasetApiDeleteApi(Resource):
-    resource_type = "dataset"
+    resource_type = ApiTokenType.DATASET
 
     @console_ns.doc("delete_dataset_api_key")
     @console_ns.doc(description="Delete dataset API key")
@@ -839,14 +846,14 @@ class DatasetApiDeleteApi(Resource):
     def delete(self, api_key_id):
         _, current_tenant_id = current_account_with_tenant()
         api_key_id = str(api_key_id)
-        key = (
-            db.session.query(ApiToken)
+        key = db.session.scalar(
+            select(ApiToken)
             .where(
                 ApiToken.tenant_id == current_tenant_id,
                 ApiToken.type == self.resource_type,
                 ApiToken.id == api_key_id,
             )
-            .first()
+            .limit(1)
         )
 
         if key is None:
@@ -857,7 +864,7 @@ class DatasetApiDeleteApi(Resource):
         assert key is not None  # nosec - for type checker only
         ApiTokenCache.delete(key.token, key.type)
 
-        db.session.query(ApiToken).where(ApiToken.id == api_key_id).delete()
+        db.session.delete(key)
         db.session.commit()
 
         return {"result": "success"}, 204

api/controllers/console/datasets/datasets_document.py

@@ -10,7 +10,7 @@ import sqlalchemy as sa
 from flask import request, send_file
 from flask_restx import Resource, fields, marshal, marshal_with
 from pydantic import BaseModel, Field
-from sqlalchemy import asc, desc, select
+from sqlalchemy import asc, desc, func, select
 from werkzeug.exceptions import Forbidden, NotFound
 
 import services
@@ -27,6 +27,7 @@ from core.model_manager import ModelManager
 from core.plugin.impl.exc import PluginDaemonClientSideError
 from core.rag.extractor.entity.datasource_type import DatasourceType
 from core.rag.extractor.entity.extract_setting import ExtractSetting, NotionInfo, WebsiteInfo
+from core.rag.index_processor.constant.index_type import IndexTechniqueType
 from dify_graph.model_runtime.entities.model_entities import ModelType
 from dify_graph.model_runtime.errors.invoke import InvokeAuthorizationError
 from extensions.ext_database import db
@@ -211,12 +212,11 @@ class GetProcessRuleApi(Resource):
                 raise Forbidden(str(e))
 
             # get the latest process rule
-            dataset_process_rule = (
-                db.session.query(DatasetProcessRule)
+            dataset_process_rule = db.session.scalar(
+                select(DatasetProcessRule)
                 .where(DatasetProcessRule.dataset_id == document.dataset_id)
                 .order_by(DatasetProcessRule.created_at.desc())
                 .limit(1)
-                .one_or_none()
             )
             if dataset_process_rule:
                 mode = dataset_process_rule.mode
@@ -330,21 +330,23 @@ class DatasetDocumentListApi(Resource):
         if fetch:
             for document in documents:
                 completed_segments = (
-                    db.session.query(DocumentSegment)
-                    .where(
-                        DocumentSegment.completed_at.isnot(None),
-                        DocumentSegment.document_id == str(document.id),
-                        DocumentSegment.status != SegmentStatus.RE_SEGMENT,
+                    db.session.scalar(
+                        select(func.count(DocumentSegment.id)).where(
+                            DocumentSegment.completed_at.isnot(None),
+                            DocumentSegment.document_id == str(document.id),
+                            DocumentSegment.status != SegmentStatus.RE_SEGMENT,
+                        )
                     )
-                    .count()
+                    or 0
                 )
                 total_segments = (
-                    db.session.query(DocumentSegment)
-                    .where(
-                        DocumentSegment.document_id == str(document.id),
-                        DocumentSegment.status != SegmentStatus.RE_SEGMENT,
+                    db.session.scalar(
+                        select(func.count(DocumentSegment.id)).where(
+                            DocumentSegment.document_id == str(document.id),
+                            DocumentSegment.status != SegmentStatus.RE_SEGMENT,
+                        )
                     )
-                    .count()
+                    or 0
                 )
                 document.completed_segments = completed_segments
                 document.total_segments = total_segments
@@ -448,7 +450,7 @@ class DatasetInitApi(Resource):
             raise Forbidden()
 
         knowledge_config = KnowledgeConfig.model_validate(console_ns.payload or {})
-        if knowledge_config.indexing_technique == "high_quality":
+        if knowledge_config.indexing_technique == IndexTechniqueType.HIGH_QUALITY:
             if knowledge_config.embedding_model is None or knowledge_config.embedding_model_provider is None:
                 raise ValueError("embedding model and embedding model provider are required for high quality indexing.")
             try:
@@ -462,7 +464,7 @@ class DatasetInitApi(Resource):
                 is_multimodal = DatasetService.check_is_multimodal_model(
                     current_tenant_id, knowledge_config.embedding_model_provider, knowledge_config.embedding_model
                 )
-                knowledge_config.is_multimodal = is_multimodal
+                knowledge_config.is_multimodal = is_multimodal  # pyrefly: ignore[bad-assignment]
             except InvokeAuthorizationError:
                 raise ProviderNotInitializeError(
                     "No Embedding Model available. Please configure a valid provider in the Settings -> Model Provider."
@@ -521,10 +523,10 @@ class DocumentIndexingEstimateApi(DocumentResource):
             if data_source_info and "upload_file_id" in data_source_info:
                 file_id = data_source_info["upload_file_id"]
 
-                file = (
-                    db.session.query(UploadFile)
+                file = db.session.scalar(
+                    select(UploadFile)
                     .where(UploadFile.tenant_id == document.tenant_id, UploadFile.id == file_id)
-                    .first()
+                    .limit(1)
                 )
 
                 # raise error if file not found
@@ -586,10 +588,10 @@ class DocumentBatchIndexingEstimateApi(DocumentResource):
                     if not data_source_info:
                         continue
                     file_id = data_source_info["upload_file_id"]
-                    file_detail = (
-                        db.session.query(UploadFile)
+                    file_detail = db.session.scalar(
+                        select(UploadFile)
                         .where(UploadFile.tenant_id == current_tenant_id, UploadFile.id == file_id)
-                        .first()
+                        .limit(1)
                     )
 
                     if file_detail is None:
@@ -672,20 +674,23 @@ class DocumentBatchIndexingStatusApi(DocumentResource):
         documents_status = []
         for document in documents:
             completed_segments = (
-                db.session.query(DocumentSegment)
-                .where(
-                    DocumentSegment.completed_at.isnot(None),
-                    DocumentSegment.document_id == str(document.id),
-                    DocumentSegment.status != SegmentStatus.RE_SEGMENT,
+                db.session.scalar(
+                    select(func.count(DocumentSegment.id)).where(
+                        DocumentSegment.completed_at.isnot(None),
+                        DocumentSegment.document_id == str(document.id),
+                        DocumentSegment.status != SegmentStatus.RE_SEGMENT,
+                    )
                 )
-                .count()
+                or 0
             )
             total_segments = (
-                db.session.query(DocumentSegment)
-                .where(
-                    DocumentSegment.document_id == str(document.id), DocumentSegment.status != SegmentStatus.RE_SEGMENT
+                db.session.scalar(
+                    select(func.count(DocumentSegment.id)).where(
+                        DocumentSegment.document_id == str(document.id),
+                        DocumentSegment.status != SegmentStatus.RE_SEGMENT,
+                    )
                 )
-                .count()
+                or 0
             )
             # Create a dictionary with document attributes and additional fields
             document_dict = {
@@ -723,18 +728,23 @@ class DocumentIndexingStatusApi(DocumentResource):
         document = self.get_document(dataset_id, document_id)
 
         completed_segments = (
-            db.session.query(DocumentSegment)
-            .where(
-                DocumentSegment.completed_at.isnot(None),
-                DocumentSegment.document_id == str(document_id),
-                DocumentSegment.status != SegmentStatus.RE_SEGMENT,
+            db.session.scalar(
+                select(func.count(DocumentSegment.id)).where(
+                    DocumentSegment.completed_at.isnot(None),
+                    DocumentSegment.document_id == str(document_id),
+                    DocumentSegment.status != SegmentStatus.RE_SEGMENT,
+                )
             )
-            .count()
+            or 0
         )
         total_segments = (
-            db.session.query(DocumentSegment)
-            .where(DocumentSegment.document_id == str(document_id), DocumentSegment.status != SegmentStatus.RE_SEGMENT)
-            .count()
+            db.session.scalar(
+                select(func.count(DocumentSegment.id)).where(
+                    DocumentSegment.document_id == str(document_id),
+                    DocumentSegment.status != SegmentStatus.RE_SEGMENT,
+                )
+            )
+            or 0
         )
 
         # Create a dictionary with document attributes and additional fields
@@ -1258,11 +1268,11 @@ class DocumentPipelineExecutionLogApi(DocumentResource):
         document = DocumentService.get_document(dataset.id, document_id)
         if not document:
             raise NotFound("Document not found.")
-        log = (
-            db.session.query(DocumentPipelineExecutionLog)
-            .filter_by(document_id=document_id)
+        log = db.session.scalar(
+            select(DocumentPipelineExecutionLog)
+            .where(DocumentPipelineExecutionLog.document_id == document_id)
             .order_by(DocumentPipelineExecutionLog.created_at.desc())
-            .first()
+            .limit(1)
         )
         if not log:
             return {
@@ -1328,7 +1338,7 @@ class DocumentGenerateSummaryApi(Resource):
             raise BadRequest("document_list cannot be empty.")
 
         # Check if dataset configuration supports summary generation
-        if dataset.indexing_technique != "high_quality":
+        if dataset.indexing_technique != IndexTechniqueType.HIGH_QUALITY:
             raise ValueError(
                 f"Summary generation is only available for 'high_quality' indexing technique. "
                 f"Current indexing technique: {dataset.indexing_technique}"

api/controllers/console/datasets/datasets_segments.py

@@ -26,6 +26,7 @@ from controllers.console.wraps import (
 )
 from core.errors.error import LLMBadRequestError, ProviderTokenNotInitError
 from core.model_manager import ModelManager
+from core.rag.index_processor.constant.index_type import IndexTechniqueType
 from dify_graph.model_runtime.entities.model_entities import ModelType
 from extensions.ext_database import db
 from extensions.ext_redis import redis_client
@@ -45,7 +46,7 @@ def _get_segment_with_summary(segment, dataset_id):
     """Helper function to marshal segment and add summary information."""
     from services.summary_index_service import SummaryIndexService
 
-    segment_dict = dict(marshal(segment, segment_fields))
+    segment_dict = dict(marshal(segment, segment_fields))  # type: ignore
     # Query summary for this segment (only enabled summaries)
     summary = SummaryIndexService.get_segment_summary(segment_id=segment.id, dataset_id=dataset_id)
     segment_dict["summary"] = summary.summary_content if summary else None
@@ -206,7 +207,7 @@ class DatasetDocumentSegmentListApi(Resource):
         # Add summary to each segment
         segments_with_summary = []
         for segment in segments.items:
-            segment_dict = dict(marshal(segment, segment_fields))
+            segment_dict = dict(marshal(segment, segment_fields))  # type: ignore
             segment_dict["summary"] = summaries.get(segment.id)
             segments_with_summary.append(segment_dict)
 
@@ -279,7 +280,7 @@ class DatasetDocumentSegmentApi(Resource):
             DatasetService.check_dataset_permission(dataset, current_user)
         except services.errors.account.NoPermissionError as e:
             raise Forbidden(str(e))
-        if dataset.indexing_technique == "high_quality":
+        if dataset.indexing_technique == IndexTechniqueType.HIGH_QUALITY:
             # check embedding model setting
             try:
                 model_manager = ModelManager()
@@ -333,7 +334,7 @@ class DatasetDocumentSegmentAddApi(Resource):
         if not current_user.is_dataset_editor:
             raise Forbidden()
         # check embedding model setting
-        if dataset.indexing_technique == "high_quality":
+        if dataset.indexing_technique == IndexTechniqueType.HIGH_QUALITY:
             try:
                 model_manager = ModelManager()
                 model_manager.get_model_instance(
@@ -383,7 +384,7 @@ class DatasetDocumentSegmentUpdateApi(Resource):
         document = DocumentService.get_document(dataset_id, document_id)
         if not document:
             raise NotFound("Document not found.")
-        if dataset.indexing_technique == "high_quality":
+        if dataset.indexing_technique == IndexTechniqueType.HIGH_QUALITY:
             # check embedding model setting
             try:
                 model_manager = ModelManager()
@@ -401,10 +402,10 @@ class DatasetDocumentSegmentUpdateApi(Resource):
                 raise ProviderNotInitializeError(ex.description)
             # check segment
         segment_id = str(segment_id)
-        segment = (
-            db.session.query(DocumentSegment)
+        segment = db.session.scalar(
+            select(DocumentSegment)
             .where(DocumentSegment.id == str(segment_id), DocumentSegment.tenant_id == current_tenant_id)
-            .first()
+            .limit(1)
         )
         if not segment:
             raise NotFound("Segment not found.")
@@ -447,10 +448,10 @@ class DatasetDocumentSegmentUpdateApi(Resource):
             raise NotFound("Document not found.")
         # check segment
         segment_id = str(segment_id)
-        segment = (
-            db.session.query(DocumentSegment)
+        segment = db.session.scalar(
+            select(DocumentSegment)
             .where(DocumentSegment.id == str(segment_id), DocumentSegment.tenant_id == current_tenant_id)
-            .first()
+            .limit(1)
         )
         if not segment:
             raise NotFound("Segment not found.")
@@ -494,7 +495,7 @@ class DatasetDocumentSegmentBatchImportApi(Resource):
         payload = BatchImportPayload.model_validate(console_ns.payload or {})
         upload_file_id = payload.upload_file_id
 
-        upload_file = db.session.query(UploadFile).where(UploadFile.id == upload_file_id).first()
+        upload_file = db.session.scalar(select(UploadFile).where(UploadFile.id == upload_file_id).limit(1))
         if not upload_file:
             raise NotFound("UploadFile not found.")
 
@@ -559,17 +560,17 @@ class ChildChunkAddApi(Resource):
             raise NotFound("Document not found.")
         # check segment
         segment_id = str(segment_id)
-        segment = (
-            db.session.query(DocumentSegment)
+        segment = db.session.scalar(
+            select(DocumentSegment)
             .where(DocumentSegment.id == str(segment_id), DocumentSegment.tenant_id == current_tenant_id)
-            .first()
+            .limit(1)
         )
         if not segment:
             raise NotFound("Segment not found.")
         if not current_user.is_dataset_editor:
             raise Forbidden()
         # check embedding model setting
-        if dataset.indexing_technique == "high_quality":
+        if dataset.indexing_technique == IndexTechniqueType.HIGH_QUALITY:
             try:
                 model_manager = ModelManager()
                 model_manager.get_model_instance(
@@ -616,10 +617,10 @@ class ChildChunkAddApi(Resource):
             raise NotFound("Document not found.")
         # check segment
         segment_id = str(segment_id)
-        segment = (
-            db.session.query(DocumentSegment)
+        segment = db.session.scalar(
+            select(DocumentSegment)
             .where(DocumentSegment.id == str(segment_id), DocumentSegment.tenant_id == current_tenant_id)
-            .first()
+            .limit(1)
         )
         if not segment:
             raise NotFound("Segment not found.")
@@ -666,10 +667,10 @@ class ChildChunkAddApi(Resource):
             raise NotFound("Document not found.")
             # check segment
         segment_id = str(segment_id)
-        segment = (
-            db.session.query(DocumentSegment)
+        segment = db.session.scalar(
+            select(DocumentSegment)
             .where(DocumentSegment.id == str(segment_id), DocumentSegment.tenant_id == current_tenant_id)
-            .first()
+            .limit(1)
         )
         if not segment:
             raise NotFound("Segment not found.")
@@ -714,24 +715,24 @@ class ChildChunkUpdateApi(Resource):
             raise NotFound("Document not found.")
         # check segment
         segment_id = str(segment_id)
-        segment = (
-            db.session.query(DocumentSegment)
+        segment = db.session.scalar(
+            select(DocumentSegment)
             .where(DocumentSegment.id == str(segment_id), DocumentSegment.tenant_id == current_tenant_id)
-            .first()
+            .limit(1)
         )
         if not segment:
             raise NotFound("Segment not found.")
         # check child chunk
         child_chunk_id = str(child_chunk_id)
-        child_chunk = (
-            db.session.query(ChildChunk)
+        child_chunk = db.session.scalar(
+            select(ChildChunk)
             .where(
                 ChildChunk.id == str(child_chunk_id),
                 ChildChunk.tenant_id == current_tenant_id,
                 ChildChunk.segment_id == segment.id,
                 ChildChunk.document_id == document_id,
             )
-            .first()
+            .limit(1)
         )
         if not child_chunk:
             raise NotFound("Child chunk not found.")
@@ -771,24 +772,24 @@ class ChildChunkUpdateApi(Resource):
             raise NotFound("Document not found.")
             # check segment
         segment_id = str(segment_id)
-        segment = (
-            db.session.query(DocumentSegment)
+        segment = db.session.scalar(
+            select(DocumentSegment)
             .where(DocumentSegment.id == str(segment_id), DocumentSegment.tenant_id == current_tenant_id)
-            .first()
+            .limit(1)
         )
         if not segment:
             raise NotFound("Segment not found.")
         # check child chunk
         child_chunk_id = str(child_chunk_id)
-        child_chunk = (
-            db.session.query(ChildChunk)
+        child_chunk = db.session.scalar(
+            select(ChildChunk)
             .where(
                 ChildChunk.id == str(child_chunk_id),
                 ChildChunk.tenant_id == current_tenant_id,
                 ChildChunk.segment_id == segment.id,
                 ChildChunk.document_id == document_id,
             )
-            .first()
+            .limit(1)
         )
         if not child_chunk:
             raise NotFound("Child chunk not found.")

api/controllers/console/datasets/rag_pipeline/rag_pipeline_workflow.py

@@ -6,7 +6,7 @@ from flask import abort, request
 from flask_restx import Resource, marshal_with  # type: ignore
 from pydantic import BaseModel, Field
 from sqlalchemy.orm import Session
-from werkzeug.exceptions import Forbidden, InternalServerError, NotFound
+from werkzeug.exceptions import BadRequest, Forbidden, InternalServerError, NotFound
 
 import services
 from controllers.common.schema import register_schema_models
@@ -16,7 +16,11 @@ from controllers.console.app.error import (
     DraftWorkflowNotExist,
     DraftWorkflowNotSync,
 )
-from controllers.console.app.workflow import workflow_model, workflow_pagination_model
+from controllers.console.app.workflow import (
+    RESTORE_SOURCE_WORKFLOW_MUST_BE_PUBLISHED_MESSAGE,
+    workflow_model,
+    workflow_pagination_model,
+)
 from controllers.console.app.workflow_run import (
     workflow_run_detail_model,
     workflow_run_node_execution_list_model,
@@ -42,7 +46,8 @@ from libs.login import current_account_with_tenant, current_user, login_required
 from models import Account
 from models.dataset import Pipeline
 from models.model import EndUser
-from services.errors.app import WorkflowHashNotEqualError
+from models.workflow import Workflow
+from services.errors.app import IsDraftWorkflowError, WorkflowHashNotEqualError, WorkflowNotFoundError
 from services.errors.llm import InvokeRateLimitError
 from services.rag_pipeline.pipeline_generate_service import PipelineGenerateService
 from services.rag_pipeline.rag_pipeline import RagPipelineService
@@ -203,9 +208,12 @@ class DraftRagPipelineApi(Resource):
             abort(415)
 
         payload = DraftWorkflowSyncPayload.model_validate(payload_dict)
+        rag_pipeline_service = RagPipelineService()
 
         try:
-            environment_variables_list = payload.environment_variables or []
+            environment_variables_list = Workflow.normalize_environment_variable_mappings(
+                payload.environment_variables or [],
+            )
             environment_variables = [
                 variable_factory.build_environment_variable_from_mapping(obj) for obj in environment_variables_list
             ]
@@ -213,7 +221,6 @@ class DraftRagPipelineApi(Resource):
             conversation_variables = [
                 variable_factory.build_conversation_variable_from_mapping(obj) for obj in conversation_variables_list
             ]
-            rag_pipeline_service = RagPipelineService()
             workflow = rag_pipeline_service.sync_draft_workflow(
                 pipeline=pipeline,
                 graph=payload.graph,
@@ -705,6 +712,36 @@ class PublishedAllRagPipelineApi(Resource):
             }
 
 
+@console_ns.route("/rag/pipelines/<uuid:pipeline_id>/workflows/<string:workflow_id>/restore")
+class RagPipelineDraftWorkflowRestoreApi(Resource):
+    @setup_required
+    @login_required
+    @account_initialization_required
+    @edit_permission_required
+    @get_rag_pipeline
+    def post(self, pipeline: Pipeline, workflow_id: str):
+        current_user, _ = current_account_with_tenant()
+        rag_pipeline_service = RagPipelineService()
+
+        try:
+            workflow = rag_pipeline_service.restore_published_workflow_to_draft(
+                pipeline=pipeline,
+                workflow_id=workflow_id,
+                account=current_user,
+            )
+        except IsDraftWorkflowError as exc:
+            # Use a stable, predefined message to keep the 400 response consistent
+            raise BadRequest(RESTORE_SOURCE_WORKFLOW_MUST_BE_PUBLISHED_MESSAGE) from exc
+        except WorkflowNotFoundError as exc:
+            raise NotFound(str(exc)) from exc
+
+        return {
+            "result": "success",
+            "hash": workflow.unique_hash,
+            "updated_at": TimestampField().format(workflow.updated_at or workflow.created_at),
+        }
+
+
 @console_ns.route("/rag/pipelines/<uuid:pipeline_id>/workflows/<string:workflow_id>")
 class RagPipelineByIdApi(Resource):
     @setup_required

api/controllers/console/datasets/wraps.py

@@ -2,6 +2,8 @@ from collections.abc import Callable
 from functools import wraps
 from typing import ParamSpec, TypeVar
 
+from sqlalchemy import select
+
 from controllers.console.datasets.error import PipelineNotFoundError
 from extensions.ext_database import db
 from libs.login import current_account_with_tenant
@@ -24,10 +26,8 @@ def get_rag_pipeline(view_func: Callable[P, R]):
 
         del kwargs["pipeline_id"]
 
-        pipeline = (
-            db.session.query(Pipeline)
-            .where(Pipeline.id == pipeline_id, Pipeline.tenant_id == current_tenant_id)
-            .first()
+        pipeline = db.session.scalar(
+            select(Pipeline).where(Pipeline.id == pipeline_id, Pipeline.tenant_id == current_tenant_id).limit(1)
         )
 
         if not pipeline:

api/controllers/console/explore/banner.py

@@ -1,5 +1,6 @@
 from flask import request
 from flask_restx import Resource
+from sqlalchemy import select
 
 from controllers.console import api
 from controllers.console.explore.wraps import explore_banner_enabled
@@ -17,14 +18,18 @@ class BannerApi(Resource):
         language = request.args.get("language", "en-US")
 
         # Build base query for enabled banners
-        base_query = db.session.query(ExporleBanner).where(ExporleBanner.status == BannerStatus.ENABLED)
+        base_query = select(ExporleBanner).where(ExporleBanner.status == BannerStatus.ENABLED)
 
         # Try to get banners in the requested language
-        banners = base_query.where(ExporleBanner.language == language).order_by(ExporleBanner.sort).all()
+        banners = db.session.scalars(
+            base_query.where(ExporleBanner.language == language).order_by(ExporleBanner.sort)
+        ).all()
 
         # Fallback to en-US if no banners found and language is not en-US
         if not banners and language != "en-US":
-            banners = base_query.where(ExporleBanner.language == "en-US").order_by(ExporleBanner.sort).all()
+            banners = db.session.scalars(
+                base_query.where(ExporleBanner.language == "en-US").order_by(ExporleBanner.sort)
+            ).all()
         # Convert banners to serializable format
         result = []
         for banner in banners:

api/controllers/console/explore/installed_app.py

@@ -133,13 +133,15 @@ class InstalledAppsListApi(Resource):
     def post(self):
         payload = InstalledAppCreatePayload.model_validate(console_ns.payload or {})
 
-        recommended_app = db.session.query(RecommendedApp).where(RecommendedApp.app_id == payload.app_id).first()
+        recommended_app = db.session.scalar(
+            select(RecommendedApp).where(RecommendedApp.app_id == payload.app_id).limit(1)
+        )
         if recommended_app is None:
             raise NotFound("Recommended app not found")
 
         _, current_tenant_id = current_account_with_tenant()
 
-        app = db.session.query(App).where(App.id == payload.app_id).first()
+        app = db.session.get(App, payload.app_id)
 
         if app is None:
             raise NotFound("App entity not found")
@@ -147,10 +149,10 @@ class InstalledAppsListApi(Resource):
         if not app.is_public:
             raise Forbidden("You can't install a non-public app")
 
-        installed_app = (
-            db.session.query(InstalledApp)
+        installed_app = db.session.scalar(
+            select(InstalledApp)
             .where(and_(InstalledApp.app_id == payload.app_id, InstalledApp.tenant_id == current_tenant_id))
-            .first()
+            .limit(1)
         )
 
         if installed_app is None:

api/controllers/console/explore/trial.py

@@ -4,6 +4,7 @@ from typing import Any, Literal, cast
 from flask import request
 from flask_restx import Resource, fields, marshal, marshal_with
 from pydantic import BaseModel
+from sqlalchemy import select
 from werkzeug.exceptions import Forbidden, InternalServerError, NotFound
 
 import services
@@ -476,7 +477,7 @@ class TrialSitApi(Resource):
 
         Returns the site configuration for the application including theme, icons, and text.
         """
-        site = db.session.query(Site).where(Site.app_id == app_model.id).first()
+        site = db.session.scalar(select(Site).where(Site.app_id == app_model.id).limit(1))
 
         if not site:
             raise Forbidden()
@@ -541,13 +542,7 @@ class AppWorkflowApi(Resource):
         if not app_model.workflow_id:
             raise AppUnavailableError()
 
-        workflow = (
-            db.session.query(Workflow)
-            .where(
-                Workflow.id == app_model.workflow_id,
-            )
-            .first()
-        )
+        workflow = db.session.get(Workflow, app_model.workflow_id)
         return workflow
 
 

api/controllers/console/explore/wraps.py

@@ -4,6 +4,7 @@ from typing import Concatenate, ParamSpec, TypeVar
 
 from flask import abort
 from flask_restx import Resource
+from sqlalchemy import select
 from werkzeug.exceptions import NotFound
 
 from controllers.console.explore.error import AppAccessDeniedError, TrialAppLimitExceeded, TrialAppNotAllowed
@@ -24,10 +25,10 @@ def installed_app_required(view: Callable[Concatenate[InstalledApp, P], R] | Non
         @wraps(view)
         def decorated(installed_app_id: str, *args: P.args, **kwargs: P.kwargs):
             _, current_tenant_id = current_account_with_tenant()
-            installed_app = (
-                db.session.query(InstalledApp)
+            installed_app = db.session.scalar(
+                select(InstalledApp)
                 .where(InstalledApp.id == str(installed_app_id), InstalledApp.tenant_id == current_tenant_id)
-                .first()
+                .limit(1)
             )
 
             if installed_app is None:
@@ -78,7 +79,7 @@ def trial_app_required(view: Callable[Concatenate[App, P], R] | None = None):
         def decorated(app_id: str, *args: P.args, **kwargs: P.kwargs):
             current_user, _ = current_account_with_tenant()
 
-            trial_app = db.session.query(TrialApp).where(TrialApp.app_id == str(app_id)).first()
+            trial_app = db.session.scalar(select(TrialApp).where(TrialApp.app_id == str(app_id)).limit(1))
 
             if trial_app is None:
                 raise TrialAppNotAllowed()
@@ -87,10 +88,10 @@ def trial_app_required(view: Callable[Concatenate[App, P], R] | None = None):
             if app is None:
                 raise TrialAppNotAllowed()
 
-            account_trial_app_record = (
-                db.session.query(AccountTrialAppRecord)
+            account_trial_app_record = db.session.scalar(
+                select(AccountTrialAppRecord)
                 .where(AccountTrialAppRecord.account_id == current_user.id, AccountTrialAppRecord.app_id == app_id)
-                .first()
+                .limit(1)
             )
             if account_trial_app_record:
                 if account_trial_app_record.count >= trial_app.trial_limit:

api/controllers/console/setup.py

@@ -2,6 +2,7 @@ from typing import Literal
 
 from flask import request
 from pydantic import BaseModel, Field, field_validator
+from sqlalchemy import select
 
 from configs import dify_config
 from controllers.fastopenapi import console_router
@@ -100,6 +101,6 @@ def setup_system(payload: SetupRequestPayload) -> SetupResponse:
 
 def get_setup_status() -> DifySetup | bool | None:
     if dify_config.EDITION == "SELF_HOSTED":
-        return db.session.query(DifySetup).first()
+        return db.session.scalar(select(DifySetup).limit(1))
 
     return True

api/controllers/console/workspace/account.py

@@ -212,13 +212,13 @@ class AccountInitApi(Resource):
                 raise ValueError("invitation_code is required")
 
             # check invitation code
-            invitation_code = (
-                db.session.query(InvitationCode)
+            invitation_code = db.session.scalar(
+                select(InvitationCode)
                 .where(
                     InvitationCode.code == args.invitation_code,
                     InvitationCode.status == InvitationCodeStatus.UNUSED,
                 )
-                .first()
+                .limit(1)
             )
 
             if not invitation_code:

api/controllers/console/workspace/members.py

@@ -171,7 +171,7 @@ class MemberCancelInviteApi(Resource):
         current_user, _ = current_account_with_tenant()
         if not current_user.current_tenant:
             raise ValueError("No current tenant")
-        member = db.session.query(Account).where(Account.id == str(member_id)).first()
+        member = db.session.get(Account, str(member_id))
         if member is None:
             abort(404)
         else:

api/controllers/console/workspace/workspace.py

@@ -7,6 +7,7 @@ from sqlalchemy import select
 from werkzeug.exceptions import Unauthorized
 
 import services
+from configs import dify_config
 from controllers.common.errors import (
     FilenameNotExistsError,
     FileTooLargeError,
@@ -29,6 +30,7 @@ from libs.helper import TimestampField
 from libs.login import current_account_with_tenant, login_required
 from models.account import Tenant, TenantStatus
 from services.account_service import TenantService
+from services.billing_service import BillingService, SubscriptionPlan
 from services.enterprise.enterprise_service import EnterpriseService
 from services.feature_service import FeatureService
 from services.file_service import FileService
@@ -108,9 +110,29 @@ class TenantListApi(Resource):
         current_user, current_tenant_id = current_account_with_tenant()
         tenants = TenantService.get_join_tenants(current_user)
         tenant_dicts = []
+        is_enterprise_only = dify_config.ENTERPRISE_ENABLED and not dify_config.BILLING_ENABLED
+        is_saas = dify_config.EDITION == "CLOUD" and dify_config.BILLING_ENABLED
+        tenant_plans: dict[str, SubscriptionPlan] = {}
+
+        if is_saas:
+            tenant_ids = [tenant.id for tenant in tenants]
+            if tenant_ids:
+                tenant_plans = BillingService.get_plan_bulk(tenant_ids)
+                if not tenant_plans:
+                    logger.warning("get_plan_bulk returned empty result, falling back to legacy feature path")
 
         for tenant in tenants:
-            features = FeatureService.get_features(tenant.id)
+            plan: str = CloudPlan.SANDBOX
+            if is_saas:
+                tenant_plan = tenant_plans.get(tenant.id)
+                if tenant_plan:
+                    plan = tenant_plan["plan"] or CloudPlan.SANDBOX
+                else:
+                    features = FeatureService.get_features(tenant.id)
+                    plan = features.billing.subscription.plan or CloudPlan.SANDBOX
+            elif not is_enterprise_only:
+                features = FeatureService.get_features(tenant.id)
+                plan = features.billing.subscription.plan or CloudPlan.SANDBOX
 
             # Create a dictionary with tenant attributes
             tenant_dict = {
@@ -118,7 +140,7 @@ class TenantListApi(Resource):
                 "name": tenant.name,
                 "status": tenant.status,
                 "created_at": tenant.created_at,
-                "plan": features.billing.subscription.plan if features.billing.enabled else CloudPlan.SANDBOX,
+                "plan": plan,
                 "current": tenant.id == current_tenant_id if current_tenant_id else False,
             }
 
@@ -198,7 +220,7 @@ class SwitchWorkspaceApi(Resource):
         except Exception:
             raise AccountNotLinkTenantError("Account not link tenant")
 
-        new_tenant = db.session.query(Tenant).get(args.tenant_id)  # Get new tenant
+        new_tenant = db.session.get(Tenant, args.tenant_id)  # Get new tenant
         if new_tenant is None:
             raise ValueError("Tenant not found")
 

api/controllers/console/wraps.py

@@ -7,6 +7,7 @@ from functools import wraps
 from typing import ParamSpec, TypeVar
 
 from flask import abort, request
+from sqlalchemy import select
 
 from configs import dify_config
 from controllers.console.auth.error import AuthenticationFailedError, EmailCodeError
@@ -218,13 +219,9 @@ def setup_required(view: Callable[P, R]) -> Callable[P, R]:
     @wraps(view)
     def decorated(*args: P.args, **kwargs: P.kwargs) -> R:
         # check setup
-        if (
-            dify_config.EDITION == "SELF_HOSTED"
-            and os.environ.get("INIT_PASSWORD")
-            and not db.session.query(DifySetup).first()
-        ):
-            raise NotInitValidateError()
-        elif dify_config.EDITION == "SELF_HOSTED" and not db.session.query(DifySetup).first():
+        if dify_config.EDITION == "SELF_HOSTED" and not db.session.scalar(select(DifySetup).limit(1)):
+            if os.environ.get("INIT_PASSWORD"):
+                raise NotInitValidateError()
             raise NotSetupError()
 
         return view(*args, **kwargs)

api/controllers/inner_api/plugin/wraps.py

@@ -5,6 +5,7 @@ from typing import ParamSpec, TypeVar
 from flask import current_app, request
 from flask_login import user_logged_in
 from pydantic import BaseModel
+from sqlalchemy import select
 from sqlalchemy.orm import Session
 
 from extensions.ext_database import db
@@ -36,23 +37,16 @@ def get_user(tenant_id: str, user_id: str | None) -> EndUser:
             user_model = None
 
             if is_anonymous:
-                user_model = (
-                    session.query(EndUser)
+                user_model = session.scalar(
+                    select(EndUser)
                     .where(
                         EndUser.session_id == user_id,
                         EndUser.tenant_id == tenant_id,
                     )
-                    .first()
+                    .limit(1)
                 )
             else:
-                user_model = (
-                    session.query(EndUser)
-                    .where(
-                        EndUser.id == user_id,
-                        EndUser.tenant_id == tenant_id,
-                    )
-                    .first()
-                )
+                user_model = session.get(EndUser, user_id)
 
             if not user_model:
                 user_model = EndUser(
@@ -85,16 +79,7 @@ def get_user_tenant(view_func: Callable[P, R]):
         if not user_id:
             user_id = DefaultEndUserSessionID.DEFAULT_SESSION_ID
 
-        try:
-            tenant_model = (
-                db.session.query(Tenant)
-                .where(
-                    Tenant.id == tenant_id,
-                )
-                .first()
-            )
-        except Exception:
-            raise ValueError("tenant not found")
+        tenant_model = db.session.get(Tenant, tenant_id)
 
         if not tenant_model:
             raise ValueError("tenant not found")

api/controllers/inner_api/workspace/workspace.py

@@ -2,6 +2,7 @@ import json
 
 from flask_restx import Resource
 from pydantic import BaseModel
+from sqlalchemy import select
 
 from controllers.common.schema import register_schema_models
 from controllers.console.wraps import setup_required
@@ -42,7 +43,7 @@ class EnterpriseWorkspace(Resource):
     def post(self):
         args = WorkspaceCreatePayload.model_validate(inner_api_ns.payload or {})
 
-        account = db.session.query(Account).filter_by(email=args.owner_email).first()
+        account = db.session.scalar(select(Account).where(Account.email == args.owner_email).limit(1))
         if account is None:
             return {"message": "owner account not found."}, 404
 

api/controllers/inner_api/wraps.py

@@ -75,7 +75,7 @@ def enterprise_inner_api_user_auth(view: Callable[P, R]):
         if signature_base64 != token:
             return view(*args, **kwargs)
 
-        kwargs["user"] = db.session.query(EndUser).where(EndUser.id == user_id).first()
+        kwargs["user"] = db.session.get(EndUser, user_id)
 
         return view(*args, **kwargs)
 

api/controllers/service_api/dataset/dataset.py

@@ -15,6 +15,7 @@ from controllers.service_api.wraps import (
     cloud_edition_billing_rate_limit_check,
 )
 from core.provider_manager import ProviderManager
+from core.rag.index_processor.constant.index_type import IndexTechniqueType
 from dify_graph.model_runtime.entities.model_entities import ModelType
 from fields.dataset_fields import dataset_detail_fields
 from fields.tag_fields import DataSetTag
@@ -153,15 +154,20 @@ class DatasetListApi(DatasetApiResource):
 
         data = marshal(datasets, dataset_detail_fields)
         for item in data:
-            if item["indexing_technique"] == "high_quality" and item["embedding_model_provider"]:
-                item["embedding_model_provider"] = str(ModelProviderID(item["embedding_model_provider"]))
-                item_model = f"{item['embedding_model']}:{item['embedding_model_provider']}"
+            if (
+                item["indexing_technique"] == IndexTechniqueType.HIGH_QUALITY  # pyrefly: ignore[bad-index]
+                and item["embedding_model_provider"]  # pyrefly: ignore[bad-index]
+            ):
+                item["embedding_model_provider"] = str(  # pyrefly: ignore[unsupported-operation]
+                    ModelProviderID(item["embedding_model_provider"])  # pyrefly: ignore[bad-index]
+                )
+                item_model = f"{item['embedding_model']}:{item['embedding_model_provider']}"  # pyrefly: ignore[bad-index]
                 if item_model in model_names:
-                    item["embedding_available"] = True
+                    item["embedding_available"] = True  # type: ignore
                 else:
-                    item["embedding_available"] = False
+                    item["embedding_available"] = False  # type: ignore
             else:
-                item["embedding_available"] = True
+                item["embedding_available"] = True  # type: ignore
         response = {
             "data": data,
             "has_more": len(datasets) == query.limit,
@@ -265,7 +271,7 @@ class DatasetApi(DatasetApiResource):
         for embedding_model in embedding_models:
             model_names.append(f"{embedding_model.model}:{embedding_model.provider.provider}")
 
-        if data.get("indexing_technique") == "high_quality":
+        if data.get("indexing_technique") == IndexTechniqueType.HIGH_QUALITY:
             item_model = f"{data.get('embedding_model')}:{data.get('embedding_model_provider')}"
             if item_model in model_names:
                 data["embedding_available"] = True
@@ -315,7 +321,7 @@ class DatasetApi(DatasetApiResource):
         # check embedding model setting
         embedding_model_provider = payload.embedding_model_provider
         embedding_model = payload.embedding_model
-        if payload.indexing_technique == "high_quality" or embedding_model_provider:
+        if payload.indexing_technique == IndexTechniqueType.HIGH_QUALITY or embedding_model_provider:
             if embedding_model_provider and embedding_model:
                 DatasetService.check_embedding_model_setting(
                     dataset.tenant_id, embedding_model_provider, embedding_model

api/controllers/service_api/dataset/segment.py

@@ -17,6 +17,7 @@ from controllers.service_api.wraps import (
 )
 from core.errors.error import LLMBadRequestError, ProviderTokenNotInitError
 from core.model_manager import ModelManager
+from core.rag.index_processor.constant.index_type import IndexTechniqueType
 from dify_graph.model_runtime.entities.model_entities import ModelType
 from extensions.ext_database import db
 from fields.segment_fields import child_chunk_fields, segment_fields
@@ -103,7 +104,7 @@ class SegmentApi(DatasetApiResource):
         if not document.enabled:
             raise NotFound("Document is disabled.")
         # check embedding model setting
-        if dataset.indexing_technique == "high_quality":
+        if dataset.indexing_technique == IndexTechniqueType.HIGH_QUALITY:
             try:
                 model_manager = ModelManager()
                 model_manager.get_model_instance(
@@ -157,7 +158,7 @@ class SegmentApi(DatasetApiResource):
         if not document:
             raise NotFound("Document not found.")
         # check embedding model setting
-        if dataset.indexing_technique == "high_quality":
+        if dataset.indexing_technique == IndexTechniqueType.HIGH_QUALITY:
             try:
                 model_manager = ModelManager()
                 model_manager.get_model_instance(
@@ -262,7 +263,7 @@ class DatasetSegmentApi(DatasetApiResource):
         document = DocumentService.get_document(dataset_id, document_id)
         if not document:
             raise NotFound("Document not found.")
-        if dataset.indexing_technique == "high_quality":
+        if dataset.indexing_technique == IndexTechniqueType.HIGH_QUALITY:
             # check embedding model setting
             try:
                 model_manager = ModelManager()
@@ -358,7 +359,7 @@ class ChildChunkApi(DatasetApiResource):
             raise NotFound("Segment not found.")
 
         # check embedding model setting
-        if dataset.indexing_technique == "high_quality":
+        if dataset.indexing_technique == IndexTechniqueType.HIGH_QUALITY:
             try:
                 model_manager = ModelManager()
                 model_manager.get_model_instance(

api/controllers/web/human_input_form.py

@@ -8,6 +8,7 @@ from datetime import datetime
 
 from flask import Response, request
 from flask_restx import Resource, reqparse
+from sqlalchemy import select
 from werkzeug.exceptions import Forbidden
 
 from configs import dify_config
@@ -147,11 +148,11 @@ class HumanInputFormApi(Resource):
 
 def _get_app_site_from_form(form: Form) -> tuple[App, Site]:
     """Resolve App/Site for the form's app and validate tenant status."""
-    app_model = db.session.query(App).where(App.id == form.app_id).first()
+    app_model = db.session.get(App, form.app_id)
     if app_model is None or app_model.tenant_id != form.tenant_id:
         raise NotFoundError("Form not found")
 
-    site = db.session.query(Site).where(Site.app_id == app_model.id).first()
+    site = db.session.scalar(select(Site).where(Site.app_id == app_model.id).limit(1))
     if site is None:
         raise Forbidden()
 

api/controllers/web/site.py

@@ -1,6 +1,7 @@
 from typing import cast
 
 from flask_restx import fields, marshal, marshal_with
+from sqlalchemy import select
 from werkzeug.exceptions import Forbidden
 
 from configs import dify_config
@@ -72,7 +73,7 @@ class AppSiteApi(WebApiResource):
     def get(self, app_model, end_user):
         """Retrieve app site info."""
         # get site
-        site = db.session.query(Site).where(Site.app_id == app_model.id).first()
+        site = db.session.scalar(select(Site).where(Site.app_id == app_model.id).limit(1))
 
         if not site:
             raise Forbidden()

api/core/app/apps/advanced_chat/app_generator.py

@@ -5,7 +5,7 @@ import logging
 import threading
 import uuid
 from collections.abc import Generator, Mapping, Sequence
-from typing import TYPE_CHECKING, Any, Literal, TypeVar, Union, overload
+from typing import TYPE_CHECKING, Any, Literal, Union, overload
 
 from flask import Flask, current_app
 from pydantic import ValidationError
@@ -47,7 +47,6 @@ from extensions.ext_database import db
 from factories import file_factory
 from libs.flask_utils import preserve_flask_contexts
 from models import Account, App, Conversation, EndUser, Message, Workflow, WorkflowNodeExecutionTriggeredFrom
-from models.base import Base
 from models.enums import WorkflowRunTriggeredFrom
 from services.conversation_service import ConversationService
 from services.workflow_draft_variable_service import (
@@ -522,8 +521,10 @@ class AdvancedChatAppGenerator(MessageBasedAppGenerator):
 
         # release database connection, because the following new thread operations may take a long time
         with Session(bind=db.engine, expire_on_commit=False) as session:
-            workflow = _refresh_model(session, workflow)
-            message = _refresh_model(session, message)
+            workflow = _refresh_model(session=session, model=workflow)
+            message = _refresh_model(session=session, model=message)
+        if message is None:
+            raise RuntimeError("Failed to refresh Message; _refresh_model returned None.")
         #     workflow_ = session.get(Workflow, workflow.id)
         #     assert workflow_ is not None
         #     workflow = workflow_
@@ -690,11 +691,21 @@ class AdvancedChatAppGenerator(MessageBasedAppGenerator):
                 raise e
 
 
-_T = TypeVar("_T", bound=Base)
+@overload
+def _refresh_model(*, session: Session | None = None, model: Workflow) -> Workflow: ...
 
 
-def _refresh_model(session, model: _T) -> _T:
-    with Session(bind=db.engine, expire_on_commit=False) as session:
-        detach_model = session.get(type(model), model.id)
-        assert detach_model is not None
-        return detach_model
+@overload
+def _refresh_model(*, session: Session | None = None, model: Message) -> Message: ...
+
+
+def _refresh_model(*, session: Session | None = None, model: Any) -> Any:
+    if session is not None:
+        detached_model = session.get(type(model), model.id)
+        assert detached_model is not None
+        return detached_model
+
+    with Session(bind=db.engine, expire_on_commit=False) as refresh_session:
+        detached_model = refresh_session.get(type(model), model.id)
+        assert detached_model is not None
+        return detached_model

api/core/app/apps/advanced_chat/generate_response_converter.py

@@ -1,4 +1,4 @@
-from collections.abc import Generator
+from collections.abc import Generator, Iterator
 from typing import Any, cast
 
 from core.app.apps.base_app_generate_response_converter import AppGenerateResponseConverter
@@ -56,8 +56,8 @@ class AdvancedChatAppGenerateResponseConverter(AppGenerateResponseConverter):
 
     @classmethod
     def convert_stream_full_response(
-        cls, stream_response: Generator[AppStreamResponse, None, None]
-    ) -> Generator[dict | str, Any, None]:
+        cls, stream_response: Iterator[AppStreamResponse]
+    ) -> Generator[dict | str, None, None]:
         """
         Convert stream full response.
         :param stream_response: stream response
@@ -87,8 +87,8 @@ class AdvancedChatAppGenerateResponseConverter(AppGenerateResponseConverter):
 
     @classmethod
     def convert_stream_simple_response(
-        cls, stream_response: Generator[AppStreamResponse, None, None]
-    ) -> Generator[dict | str, Any, None]:
+        cls, stream_response: Iterator[AppStreamResponse]
+    ) -> Generator[dict | str, None, None]:
         """
         Convert stream simple response.
         :param stream_response: stream response

api/core/app/apps/agent_chat/generate_response_converter.py

@@ -1,4 +1,4 @@
-from collections.abc import Generator
+from collections.abc import Generator, Iterator
 from typing import cast
 
 from core.app.apps.base_app_generate_response_converter import AppGenerateResponseConverter
@@ -55,7 +55,7 @@ class AgentChatAppGenerateResponseConverter(AppGenerateResponseConverter):
 
     @classmethod
     def convert_stream_full_response(
-        cls, stream_response: Generator[AppStreamResponse, None, None]
+        cls, stream_response: Iterator[AppStreamResponse]
     ) -> Generator[dict | str, None, None]:
         """
         Convert stream full response.
@@ -86,7 +86,7 @@ class AgentChatAppGenerateResponseConverter(AppGenerateResponseConverter):
 
     @classmethod
     def convert_stream_simple_response(
-        cls, stream_response: Generator[AppStreamResponse, None, None]
+        cls, stream_response: Iterator[AppStreamResponse]
     ) -> Generator[dict | str, None, None]:
         """
         Convert stream simple response.

api/core/app/apps/base_app_generate_response_converter.py

@@ -1,7 +1,7 @@
 import logging
 from abc import ABC, abstractmethod
-from collections.abc import Generator, Mapping
-from typing import Any, Union
+from collections.abc import Generator, Iterator, Mapping
+from typing import Any
 
 from core.app.entities.app_invoke_entities import InvokeFrom
 from core.app.entities.task_entities import AppBlockingResponse, AppStreamResponse
@@ -16,24 +16,26 @@ class AppGenerateResponseConverter(ABC):
 
     @classmethod
     def convert(
-        cls, response: Union[AppBlockingResponse, Generator[AppStreamResponse, Any, None]], invoke_from: InvokeFrom
-    ) -> Mapping[str, Any] | Generator[str | Mapping[str, Any], Any, None]:
+        cls, response: AppBlockingResponse | Iterator[AppStreamResponse], invoke_from: InvokeFrom
+    ) -> Mapping[str, Any] | Generator[str | Mapping[str, Any], None, None]:
         if invoke_from in {InvokeFrom.DEBUGGER, InvokeFrom.SERVICE_API}:
             if isinstance(response, AppBlockingResponse):
                 return cls.convert_blocking_full_response(response)
             else:
+                stream_response = response
 
-                def _generate_full_response() -> Generator[dict | str, Any, None]:
-                    yield from cls.convert_stream_full_response(response)
+                def _generate_full_response() -> Generator[dict[str, Any] | str, None, None]:
+                    yield from cls.convert_stream_full_response(stream_response)
 
                 return _generate_full_response()
         else:
             if isinstance(response, AppBlockingResponse):
                 return cls.convert_blocking_simple_response(response)
             else:
+                stream_response = response
 
-                def _generate_simple_response() -> Generator[dict | str, Any, None]:
-                    yield from cls.convert_stream_simple_response(response)
+                def _generate_simple_response() -> Generator[dict[str, Any] | str, None, None]:
+                    yield from cls.convert_stream_simple_response(stream_response)
 
                 return _generate_simple_response()
 
@@ -50,14 +52,14 @@ class AppGenerateResponseConverter(ABC):
     @classmethod
     @abstractmethod
     def convert_stream_full_response(
-        cls, stream_response: Generator[AppStreamResponse, None, None]
+        cls, stream_response: Iterator[AppStreamResponse]
     ) -> Generator[dict | str, None, None]:
         raise NotImplementedError
 
     @classmethod
     @abstractmethod
     def convert_stream_simple_response(
-        cls, stream_response: Generator[AppStreamResponse, None, None]
+        cls, stream_response: Iterator[AppStreamResponse]
     ) -> Generator[dict | str, None, None]:
         raise NotImplementedError
 
@@ -74,11 +76,22 @@ class AppGenerateResponseConverter(ABC):
             for resource in metadata["retriever_resources"]:
                 updated_resources.append(
                     {
+                        "dataset_id": resource.get("dataset_id"),
+                        "dataset_name": resource.get("dataset_name"),
+                        "document_id": resource.get("document_id"),
                         "segment_id": resource.get("segment_id", ""),
                         "position": resource["position"],
+                        "data_source_type": resource.get("data_source_type"),
                         "document_name": resource["document_name"],
                         "score": resource["score"],
+                        "hit_count": resource.get("hit_count"),
+                        "word_count": resource.get("word_count"),
+                        "segment_position": resource.get("segment_position"),
+                        "index_node_hash": resource.get("index_node_hash"),
                         "content": resource["content"],
+                        "page": resource.get("page"),
+                        "title": resource.get("title"),
+                        "files": resource.get("files"),
                         "summary": resource.get("summary"),
                     }
                 )

api/core/app/apps/base_app_generator.py

@@ -224,6 +224,7 @@ class BaseAppGenerator:
     def _get_draft_var_saver_factory(invoke_from: InvokeFrom, account: Account | EndUser) -> DraftVariableSaverFactory:
         if invoke_from == InvokeFrom.DEBUGGER:
             assert isinstance(account, Account)
+            debug_account = account
 
             def draft_var_saver_factory(
                 session: Session,
@@ -240,7 +241,7 @@ class BaseAppGenerator:
                     node_type=node_type,
                     node_execution_id=node_execution_id,
                     enclosing_node_id=enclosing_node_id,
-                    user=account,
+                    user=debug_account,
                 )
         else:
 

api/core/app/apps/chat/app_generator.py

@@ -166,15 +166,19 @@ class ChatAppGenerator(MessageBasedAppGenerator):
 
         # init generate records
         (conversation, message) = self._init_generate_records(application_generate_entity, conversation)
+        if conversation is None or message is None:
+            raise RuntimeError("_init_generate_records() returned None for conversation or message")
+        generated_conversation_id = str(conversation.id)
+        generated_message_id = str(message.id)
 
         # init queue manager
         queue_manager = MessageBasedAppQueueManager(
             task_id=application_generate_entity.task_id,
             user_id=application_generate_entity.user_id,
             invoke_from=application_generate_entity.invoke_from,
-            conversation_id=conversation.id,
+            conversation_id=generated_conversation_id,
             app_mode=conversation.mode,
-            message_id=message.id,
+            message_id=generated_message_id,
         )
 
         # new thread with request context
@@ -184,8 +188,8 @@ class ChatAppGenerator(MessageBasedAppGenerator):
                 flask_app=current_app._get_current_object(),  # type: ignore
                 application_generate_entity=application_generate_entity,
                 queue_manager=queue_manager,
-                conversation_id=conversation.id,
-                message_id=message.id,
+                conversation_id=generated_conversation_id,
+                message_id=generated_message_id,
             )
 
         worker_thread = threading.Thread(target=worker_with_context)

api/core/app/apps/chat/generate_response_converter.py

@@ -1,4 +1,4 @@
-from collections.abc import Generator
+from collections.abc import Generator, Iterator
 from typing import cast
 
 from core.app.apps.base_app_generate_response_converter import AppGenerateResponseConverter
@@ -55,7 +55,7 @@ class ChatAppGenerateResponseConverter(AppGenerateResponseConverter):
 
     @classmethod
     def convert_stream_full_response(
-        cls, stream_response: Generator[AppStreamResponse, None, None]
+        cls, stream_response: Iterator[AppStreamResponse]
     ) -> Generator[dict | str, None, None]:
         """
         Convert stream full response.
@@ -86,7 +86,7 @@ class ChatAppGenerateResponseConverter(AppGenerateResponseConverter):
 
     @classmethod
     def convert_stream_simple_response(
-        cls, stream_response: Generator[AppStreamResponse, None, None]
+        cls, stream_response: Iterator[AppStreamResponse]
     ) -> Generator[dict | str, None, None]:
         """
         Convert stream simple response.

api/core/app/apps/completion/app_generator.py

@@ -149,6 +149,11 @@ class CompletionAppGenerator(MessageBasedAppGenerator):
 
         # init generate records
         (conversation, message) = self._init_generate_records(application_generate_entity)
+        if conversation is None or message is None:
+            raise RuntimeError(
+                "_init_generate_records() returned None for conversation or message, "
+                "which is required to proceed with generation."
+            )
 
         # init queue manager
         queue_manager = MessageBasedAppQueueManager(
@@ -312,15 +317,19 @@ class CompletionAppGenerator(MessageBasedAppGenerator):
 
         # init generate records
         (conversation, message) = self._init_generate_records(application_generate_entity)
+        assert conversation is not None
+        assert message is not None
+        conversation_id = str(conversation.id)
+        message_id = str(message.id)
 
         # init queue manager
         queue_manager = MessageBasedAppQueueManager(
             task_id=application_generate_entity.task_id,
             user_id=application_generate_entity.user_id,
             invoke_from=application_generate_entity.invoke_from,
-            conversation_id=conversation.id,
+            conversation_id=conversation_id,
             app_mode=conversation.mode,
-            message_id=message.id,
+            message_id=message_id,
         )
 
         # new thread with request context
@@ -330,7 +339,7 @@ class CompletionAppGenerator(MessageBasedAppGenerator):
                 flask_app=current_app._get_current_object(),  # type: ignore
                 application_generate_entity=application_generate_entity,
                 queue_manager=queue_manager,
-                message_id=message.id,
+                message_id=message_id,
             )
 
         worker_thread = threading.Thread(target=worker_with_context)

api/core/app/apps/completion/generate_response_converter.py

@@ -1,4 +1,4 @@
-from collections.abc import Generator
+from collections.abc import Generator, Iterator
 from typing import cast
 
 from core.app.apps.base_app_generate_response_converter import AppGenerateResponseConverter
@@ -54,7 +54,7 @@ class CompletionAppGenerateResponseConverter(AppGenerateResponseConverter):
 
     @classmethod
     def convert_stream_full_response(
-        cls, stream_response: Generator[AppStreamResponse, None, None]
+        cls, stream_response: Iterator[AppStreamResponse]
     ) -> Generator[dict | str, None, None]:
         """
         Convert stream full response.
@@ -84,7 +84,7 @@ class CompletionAppGenerateResponseConverter(AppGenerateResponseConverter):
 
     @classmethod
     def convert_stream_simple_response(
-        cls, stream_response: Generator[AppStreamResponse, None, None]
+        cls, stream_response: Iterator[AppStreamResponse]
     ) -> Generator[dict | str, None, None]:
         """
         Convert stream simple response.

api/core/app/apps/message_based_app_generator.py

@@ -33,7 +33,7 @@ from extensions.ext_redis import get_pubsub_broadcast_channel
 from libs.broadcast_channel.channel import Topic
 from libs.datetime_utils import naive_utc_now
 from models import Account
-from models.enums import CreatorUserRole, MessageFileBelongsTo
+from models.enums import ConversationFromSource, CreatorUserRole, MessageFileBelongsTo
 from models.model import App, AppMode, AppModelConfig, Conversation, EndUser, Message, MessageFile
 from services.errors.app_model_config import AppModelConfigBrokenError
 from services.errors.conversation import ConversationNotExistsError
@@ -130,10 +130,10 @@ class MessageBasedAppGenerator(BaseAppGenerator):
         end_user_id = None
         account_id = None
         if application_generate_entity.invoke_from in {InvokeFrom.WEB_APP, InvokeFrom.SERVICE_API}:
-            from_source = "api"
+            from_source = ConversationFromSource.API
             end_user_id = application_generate_entity.user_id
         else:
-            from_source = "console"
+            from_source = ConversationFromSource.CONSOLE
             account_id = application_generate_entity.user_id
 
         if isinstance(application_generate_entity, AdvancedChatAppGenerateEntity):

api/core/app/apps/pipeline/generate_response_converter.py

@@ -1,4 +1,4 @@
-from collections.abc import Generator
+from collections.abc import Generator, Iterator
 from typing import cast
 
 from core.app.apps.base_app_generate_response_converter import AppGenerateResponseConverter
@@ -36,7 +36,7 @@ class WorkflowAppGenerateResponseConverter(AppGenerateResponseConverter):
 
     @classmethod
     def convert_stream_full_response(
-        cls, stream_response: Generator[AppStreamResponse, None, None]
+        cls, stream_response: Iterator[AppStreamResponse]
     ) -> Generator[dict | str, None, None]:
         """
         Convert stream full response.
@@ -65,7 +65,7 @@ class WorkflowAppGenerateResponseConverter(AppGenerateResponseConverter):
 
     @classmethod
     def convert_stream_simple_response(
-        cls, stream_response: Generator[AppStreamResponse, None, None]
+        cls, stream_response: Iterator[AppStreamResponse]
     ) -> Generator[dict | str, None, None]:
         """
         Convert stream simple response.

api/core/app/apps/workflow/generate_response_converter.py

@@ -1,4 +1,4 @@
-from collections.abc import Generator
+from collections.abc import Generator, Iterator
 from typing import cast
 
 from core.app.apps.base_app_generate_response_converter import AppGenerateResponseConverter
@@ -36,7 +36,7 @@ class WorkflowAppGenerateResponseConverter(AppGenerateResponseConverter):
 
     @classmethod
     def convert_stream_full_response(
-        cls, stream_response: Generator[AppStreamResponse, None, None]
+        cls, stream_response: Iterator[AppStreamResponse]
     ) -> Generator[dict | str, None, None]:
         """
         Convert stream full response.
@@ -65,7 +65,7 @@ class WorkflowAppGenerateResponseConverter(AppGenerateResponseConverter):
 
     @classmethod
     def convert_stream_simple_response(
-        cls, stream_response: Generator[AppStreamResponse, None, None]
+        cls, stream_response: Iterator[AppStreamResponse]
     ) -> Generator[dict | str, None, None]:
         """
         Convert stream simple response.

api/core/app/apps/workflow/generate_task_pipeline.py

@@ -705,7 +705,7 @@ class WorkflowAppGenerateTaskPipeline(GraphRuntimeStateSupport):
             app_id=self._application_generate_entity.app_config.app_id,
             workflow_id=self._workflow.id,
             workflow_run_id=workflow_run_id,
-            created_from=created_from.value,
+            created_from=created_from,
             created_by_role=self._created_by_role,
             created_by=self._user_id,
         )

api/core/app/apps/workflow_app_runner.py

@@ -1,13 +1,17 @@
 import logging
 import time
 from collections.abc import Mapping, Sequence
-from typing import Any, cast
+from typing import Protocol, TypeAlias
 
 from pydantic import ValidationError
 
 from core.app.apps.base_app_queue_manager import AppQueueManager, PublishFrom
 from core.app.entities.agent_strategy import AgentStrategyInfo
-from core.app.entities.app_invoke_entities import InvokeFrom, UserFrom, build_dify_run_context
+from core.app.entities.app_invoke_entities import (
+    InvokeFrom,
+    UserFrom,
+    build_dify_run_context,
+)
 from core.app.entities.queue_entities import (
     AppQueueEvent,
     QueueAgentLogEvent,
@@ -36,7 +40,7 @@ from core.rag.entities.citation_metadata import RetrievalSourceMetadata
 from core.workflow.node_factory import DifyNodeFactory, get_default_root_node_id, resolve_workflow_node_class
 from core.workflow.workflow_entry import WorkflowEntry
 from dify_graph.entities import GraphInitParams
-from dify_graph.entities.graph_config import NodeConfigDictAdapter
+from dify_graph.entities.graph_config import NodeConfigDict, NodeConfigDictAdapter
 from dify_graph.entities.pause_reason import HumanInputRequired
 from dify_graph.graph import Graph
 from dify_graph.graph_engine.layers.base import GraphEngineLayer
@@ -75,6 +79,14 @@ from tasks.mail_human_input_delivery_task import dispatch_human_input_email_task
 
 logger = logging.getLogger(__name__)
 
+GraphConfigObject: TypeAlias = dict[str, object]
+GraphConfigMapping: TypeAlias = Mapping[str, object]
+
+
+class SingleNodeRunEntity(Protocol):
+    node_id: str
+    inputs: Mapping[str, object]
+
 
 class WorkflowBasedAppRunner:
     def __init__(
@@ -98,7 +110,7 @@ class WorkflowBasedAppRunner:
 
     def _init_graph(
         self,
-        graph_config: Mapping[str, Any],
+        graph_config: GraphConfigMapping,
         graph_runtime_state: GraphRuntimeState,
         user_from: UserFrom,
         invoke_from: InvokeFrom,
@@ -154,8 +166,8 @@ class WorkflowBasedAppRunner:
     def _prepare_single_node_execution(
         self,
         workflow: Workflow,
-        single_iteration_run: Any | None = None,
-        single_loop_run: Any | None = None,
+        single_iteration_run: SingleNodeRunEntity | None = None,
+        single_loop_run: SingleNodeRunEntity | None = None,
     ) -> tuple[Graph, VariablePool, GraphRuntimeState]:
         """
         Prepare graph, variable pool, and runtime state for single node execution
@@ -208,11 +220,88 @@ class WorkflowBasedAppRunner:
         # This ensures all nodes in the graph reference the same GraphRuntimeState instance
         return graph, variable_pool, graph_runtime_state
 
+    @staticmethod
+    def _get_graph_items(graph_config: GraphConfigMapping) -> tuple[list[GraphConfigMapping], list[GraphConfigMapping]]:
+        nodes = graph_config.get("nodes")
+        edges = graph_config.get("edges")
+        if not isinstance(nodes, list):
+            raise ValueError("nodes in workflow graph must be a list")
+        if not isinstance(edges, list):
+            raise ValueError("edges in workflow graph must be a list")
+
+        validated_nodes: list[GraphConfigMapping] = []
+        for node in nodes:
+            if not isinstance(node, Mapping):
+                raise ValueError("nodes in workflow graph must be mappings")
+            validated_nodes.append(node)
+
+        validated_edges: list[GraphConfigMapping] = []
+        for edge in edges:
+            if not isinstance(edge, Mapping):
+                raise ValueError("edges in workflow graph must be mappings")
+            validated_edges.append(edge)
+
+        return validated_nodes, validated_edges
+
+    @staticmethod
+    def _extract_start_node_id(node_config: GraphConfigMapping | None) -> str | None:
+        if node_config is None:
+            return None
+        node_data = node_config.get("data")
+        if not isinstance(node_data, Mapping):
+            return None
+        start_node_id = node_data.get("start_node_id")
+        return start_node_id if isinstance(start_node_id, str) else None
+
+    @classmethod
+    def _build_single_node_graph_config(
+        cls,
+        *,
+        graph_config: GraphConfigMapping,
+        node_id: str,
+        node_type_filter_key: str,
+    ) -> tuple[GraphConfigObject, NodeConfigDict]:
+        node_configs, edge_configs = cls._get_graph_items(graph_config)
+        main_node_config = next((node for node in node_configs if node.get("id") == node_id), None)
+        start_node_id = cls._extract_start_node_id(main_node_config)
+
+        filtered_node_configs = [
+            dict(node)
+            for node in node_configs
+            if node.get("id") == node_id
+            or (isinstance(node_data := node.get("data"), Mapping) and node_data.get(node_type_filter_key) == node_id)
+            or (start_node_id and node.get("id") == start_node_id)
+        ]
+        if not filtered_node_configs:
+            raise ValueError(f"node id {node_id} not found in workflow graph")
+
+        filtered_node_ids = {
+            str(node_id_value) for node in filtered_node_configs if isinstance((node_id_value := node.get("id")), str)
+        }
+        filtered_edge_configs = [
+            dict(edge)
+            for edge in edge_configs
+            if (edge.get("source") is None or edge.get("source") in filtered_node_ids)
+            and (edge.get("target") is None or edge.get("target") in filtered_node_ids)
+        ]
+
+        target_node_config = next((node for node in filtered_node_configs if node.get("id") == node_id), None)
+        if target_node_config is None:
+            raise ValueError(f"node id {node_id} not found in workflow graph")
+
+        return (
+            {
+                "nodes": filtered_node_configs,
+                "edges": filtered_edge_configs,
+            },
+            NodeConfigDictAdapter.validate_python(target_node_config),
+        )
+
     def _get_graph_and_variable_pool_for_single_node_run(
         self,
         workflow: Workflow,
         node_id: str,
-        user_inputs: dict[str, Any],
+        user_inputs: Mapping[str, object],
         graph_runtime_state: GraphRuntimeState,
         node_type_filter_key: str,  # 'iteration_id' or 'loop_id'
         node_type_label: str = "node",  # 'iteration' or 'loop' for error messages
@@ -236,41 +325,14 @@ class WorkflowBasedAppRunner:
         if not graph_config:
             raise ValueError("workflow graph not found")
 
-        graph_config = cast(dict[str, Any], graph_config)
-
         if "nodes" not in graph_config or "edges" not in graph_config:
             raise ValueError("nodes or edges not found in workflow graph")
 
-        if not isinstance(graph_config.get("nodes"), list):
-            raise ValueError("nodes in workflow graph must be a list")
-
-        if not isinstance(graph_config.get("edges"), list):
-            raise ValueError("edges in workflow graph must be a list")
-
-        # filter nodes only in the specified node type (iteration or loop)
-        main_node_config = next((n for n in graph_config.get("nodes", []) if n.get("id") == node_id), None)
-        start_node_id = main_node_config.get("data", {}).get("start_node_id") if main_node_config else None
-        node_configs = [
-            node
-            for node in graph_config.get("nodes", [])
-            if node.get("id") == node_id
-            or node.get("data", {}).get(node_type_filter_key, "") == node_id
-            or (start_node_id and node.get("id") == start_node_id)
-        ]
-
-        graph_config["nodes"] = node_configs
-
-        node_ids = [node.get("id") for node in node_configs]
-
-        # filter edges only in the specified node type
-        edge_configs = [
-            edge
-            for edge in graph_config.get("edges", [])
-            if (edge.get("source") is None or edge.get("source") in node_ids)
-            and (edge.get("target") is None or edge.get("target") in node_ids)
-        ]
-
-        graph_config["edges"] = edge_configs
+        graph_config, target_node_config = self._build_single_node_graph_config(
+            graph_config=graph_config,
+            node_id=node_id,
+            node_type_filter_key=node_type_filter_key,
+        )
 
         # Create required parameters for Graph.init
         graph_init_params = GraphInitParams(
@@ -299,18 +361,6 @@ class WorkflowBasedAppRunner:
         if not graph:
             raise ValueError("graph not found in workflow")
 
-        # fetch node config from node id
-        target_node_config = None
-        for node in node_configs:
-            if node.get("id") == node_id:
-                target_node_config = node
-                break
-
-        if not target_node_config:
-            raise ValueError(f"{node_type_label} node id not found in workflow graph")
-
-        target_node_config = NodeConfigDictAdapter.validate_python(target_node_config)
-
         # Get node class
         node_type = target_node_config["data"].type
         node_version = str(target_node_config["data"].version)

api/core/app/entities/app_invoke_entities.py

@@ -213,7 +213,7 @@ class AdvancedChatAppGenerateEntity(ConversationAppGenerateEntity):
         """
 
         node_id: str
-        inputs: Mapping
+        inputs: Mapping[str, object]
 
     single_iteration_run: SingleIterationRunEntity | None = None
 
@@ -223,7 +223,7 @@ class AdvancedChatAppGenerateEntity(ConversationAppGenerateEntity):
         """
 
         node_id: str
-        inputs: Mapping
+        inputs: Mapping[str, object]
 
     single_loop_run: SingleLoopRunEntity | None = None
 
@@ -243,7 +243,7 @@ class WorkflowAppGenerateEntity(AppGenerateEntity):
         """
 
         node_id: str
-        inputs: dict
+        inputs: Mapping[str, object]
 
     single_iteration_run: SingleIterationRunEntity | None = None
 
@@ -253,7 +253,7 @@ class WorkflowAppGenerateEntity(AppGenerateEntity):
         """
 
         node_id: str
-        inputs: dict
+        inputs: Mapping[str, object]
 
     single_loop_run: SingleLoopRunEntity | None = None
 

api/core/app/features/annotation_reply/annotation_reply.py

@@ -4,9 +4,10 @@ from sqlalchemy import select
 
 from core.app.entities.app_invoke_entities import InvokeFrom
 from core.rag.datasource.vdb.vector_factory import Vector
+from core.rag.index_processor.constant.index_type import IndexTechniqueType
 from extensions.ext_database import db
 from models.dataset import Dataset
-from models.enums import CollectionBindingType
+from models.enums import CollectionBindingType, ConversationFromSource
 from models.model import App, AppAnnotationSetting, Message, MessageAnnotation
 from services.annotation_service import AppAnnotationService
 from services.dataset_service import DatasetCollectionBindingService
@@ -50,7 +51,7 @@ class AnnotationReplyFeature:
             dataset = Dataset(
                 id=app_record.id,
                 tenant_id=app_record.tenant_id,
-                indexing_technique="high_quality",
+                indexing_technique=IndexTechniqueType.HIGH_QUALITY,
                 embedding_model_provider=embedding_provider_name,
                 embedding_model=embedding_model_name,
                 collection_binding_id=dataset_collection_binding.id,
@@ -68,9 +69,9 @@ class AnnotationReplyFeature:
                 annotation = AppAnnotationService.get_annotation_by_id(annotation_id)
                 if annotation:
                     if invoke_from in {InvokeFrom.SERVICE_API, InvokeFrom.WEB_APP}:
-                        from_source = "api"
+                        from_source = ConversationFromSource.API
                     else:
-                        from_source = "console"
+                        from_source = ConversationFromSource.CONSOLE
 
                     # insert annotation history
                     AppAnnotationService.add_annotation_history(

api/core/app/features/rate_limiting/rate_limit.py

@@ -19,6 +19,7 @@ class RateLimit:
     _REQUEST_MAX_ALIVE_TIME = 10 * 60  # 10 minutes
     _ACTIVE_REQUESTS_COUNT_FLUSH_INTERVAL = 5 * 60  # recalculate request_count from request_detail every 5 minutes
     _instance_dict: dict[str, "RateLimit"] = {}
+    max_active_requests: int
 
     def __new__(cls, client_id: str, max_active_requests: int):
         if client_id not in cls._instance_dict:
@@ -27,7 +28,13 @@ class RateLimit:
         return cls._instance_dict[client_id]
 
     def __init__(self, client_id: str, max_active_requests: int):
+        flush_cache = hasattr(self, "max_active_requests") and self.max_active_requests != max_active_requests
         self.max_active_requests = max_active_requests
+        # Only flush here if this instance has already been fully initialized,
+        # i.e. the Redis key attributes exist. Otherwise, rely on the flush at
+        # the end of initialization below.
+        if flush_cache and hasattr(self, "active_requests_key") and hasattr(self, "max_active_requests_key"):
+            self.flush_cache(use_local_value=True)
         # must be called after max_active_requests is set
         if self.disabled():
             return
@@ -41,8 +48,6 @@ class RateLimit:
         self.flush_cache(use_local_value=True)
 
     def flush_cache(self, use_local_value=False):
-        if self.disabled():
-            return
         self.last_recalculate_time = time.time()
         # flush max active requests
         if use_local_value or not redis_client.exists(self.max_active_requests_key):
@@ -50,7 +55,8 @@ class RateLimit:
         else:
             self.max_active_requests = int(redis_client.get(self.max_active_requests_key).decode("utf-8"))
             redis_client.expire(self.max_active_requests_key, timedelta(days=1))
-
+        if self.disabled():
+            return
         # flush max active requests (in-transit request list)
         if not redis_client.exists(self.active_requests_key):
             return

api/core/app/layers/suspend_layer.py

@@ -6,16 +6,23 @@ from dify_graph.graph_events.graph import GraphRunPausedEvent
 class SuspendLayer(GraphEngineLayer):
     """ """
 
+    def __init__(self) -> None:
+        super().__init__()
+        self._paused = False
+
     def on_graph_start(self):
-        pass
+        self._paused = False
 
     def on_event(self, event: GraphEngineEvent):
         """
         Handle the paused event, stash runtime state into storage and wait for resume.
         """
         if isinstance(event, GraphRunPausedEvent):
-            pass
+            self._paused = True
 
     def on_graph_end(self, error: Exception | None):
         """ """
-        pass
+        self._paused = False
+
+    def is_paused(self) -> bool:
+        return self._paused

api/core/app/workflow/layers/persistence.py

@@ -128,14 +128,14 @@ class WorkflowPersistenceLayer(GraphEngineLayer):
             self._handle_graph_run_paused(event)
             return
 
-        if isinstance(event, NodeRunStartedEvent):
-            self._handle_node_started(event)
-            return
-
         if isinstance(event, NodeRunRetryEvent):
             self._handle_node_retry(event)
             return
 
+        if isinstance(event, NodeRunStartedEvent):
+            self._handle_node_started(event)
+            return
+
         if isinstance(event, NodeRunSucceededEvent):
             self._handle_node_succeeded(event)
             return

api/core/indexing_runner.py

@@ -21,7 +21,7 @@ from core.rag.datasource.keyword.keyword_factory import Keyword
 from core.rag.docstore.dataset_docstore import DatasetDocumentStore
 from core.rag.extractor.entity.datasource_type import DatasourceType
 from core.rag.extractor.entity.extract_setting import ExtractSetting, NotionInfo, WebsiteInfo
-from core.rag.index_processor.constant.index_type import IndexStructureType
+from core.rag.index_processor.constant.index_type import IndexStructureType, IndexTechniqueType
 from core.rag.index_processor.index_processor_base import BaseIndexProcessor
 from core.rag.index_processor.index_processor_factory import IndexProcessorFactory
 from core.rag.models.document import ChildDocument, Document
@@ -271,7 +271,7 @@ class IndexingRunner:
         doc_form: str | None = None,
         doc_language: str = "English",
         dataset_id: str | None = None,
-        indexing_technique: str = "economy",
+        indexing_technique: str = IndexTechniqueType.ECONOMY,
     ) -> IndexingEstimate:
         """
         Estimate the indexing for the document.
@@ -289,7 +289,7 @@ class IndexingRunner:
             dataset = db.session.query(Dataset).filter_by(id=dataset_id).first()
             if not dataset:
                 raise ValueError("Dataset not found.")
-            if dataset.indexing_technique == "high_quality" or indexing_technique == "high_quality":
+            if IndexTechniqueType.HIGH_QUALITY in {dataset.indexing_technique, indexing_technique}:
                 if dataset.embedding_model_provider:
                     embedding_model_instance = self.model_manager.get_model_instance(
                         tenant_id=tenant_id,
@@ -303,7 +303,7 @@ class IndexingRunner:
                         model_type=ModelType.TEXT_EMBEDDING,
                     )
         else:
-            if indexing_technique == "high_quality":
+            if indexing_technique == IndexTechniqueType.HIGH_QUALITY:
                 embedding_model_instance = self.model_manager.get_default_model_instance(
                     tenant_id=tenant_id,
                     model_type=ModelType.TEXT_EMBEDDING,
@@ -573,7 +573,7 @@ class IndexingRunner:
         """
 
         embedding_model_instance = None
-        if dataset.indexing_technique == "high_quality":
+        if dataset.indexing_technique == IndexTechniqueType.HIGH_QUALITY:
             embedding_model_instance = self.model_manager.get_model_instance(
                 tenant_id=dataset.tenant_id,
                 provider=dataset.embedding_model_provider,
@@ -587,7 +587,7 @@ class IndexingRunner:
         create_keyword_thread = None
         if (
             dataset_document.doc_form != IndexStructureType.PARENT_CHILD_INDEX
-            and dataset.indexing_technique == "economy"
+            and dataset.indexing_technique == IndexTechniqueType.ECONOMY
         ):
             # create keyword index
             create_keyword_thread = threading.Thread(
@@ -597,7 +597,7 @@ class IndexingRunner:
             create_keyword_thread.start()
 
         max_workers = 10
-        if dataset.indexing_technique == "high_quality":
+        if dataset.indexing_technique == IndexTechniqueType.HIGH_QUALITY:
             with concurrent.futures.ThreadPoolExecutor(max_workers=max_workers) as executor:
                 futures = []
 
@@ -628,7 +628,7 @@ class IndexingRunner:
                     tokens += future.result()
         if (
             dataset_document.doc_form != IndexStructureType.PARENT_CHILD_INDEX
-            and dataset.indexing_technique == "economy"
+            and dataset.indexing_technique == IndexTechniqueType.ECONOMY
             and create_keyword_thread is not None
         ):
             create_keyword_thread.join()
@@ -654,7 +654,7 @@ class IndexingRunner:
                 raise ValueError("no dataset found")
             keyword = Keyword(dataset)
             keyword.create(documents)
-            if dataset.indexing_technique != "high_quality":
+            if dataset.indexing_technique != IndexTechniqueType.HIGH_QUALITY:
                 document_ids = [document.metadata["doc_id"] for document in documents]
                 db.session.query(DocumentSegment).where(
                     DocumentSegment.document_id == document_id,
@@ -764,7 +764,7 @@ class IndexingRunner:
     ) -> list[Document]:
         # get embedding model instance
         embedding_model_instance = None
-        if dataset.indexing_technique == "high_quality":
+        if dataset.indexing_technique == IndexTechniqueType.HIGH_QUALITY:
             if dataset.embedding_model_provider:
                 embedding_model_instance = self.model_manager.get_model_instance(
                     tenant_id=dataset.tenant_id,

api/core/ops/arize_phoenix_trace/arize_phoenix_trace.py

@@ -181,10 +181,6 @@ class ArizePhoenixDataTrace(BaseTraceInstance):
         arize_phoenix_config: ArizeConfig | PhoenixConfig,
     ):
         super().__init__(arize_phoenix_config)
-        import logging
-
-        logging.basicConfig()
-        logging.getLogger().setLevel(logging.DEBUG)
         self.arize_phoenix_config = arize_phoenix_config
         self.tracer, self.processor = setup_tracer(arize_phoenix_config)
         self.project = arize_phoenix_config.project

api/core/ops/weave_trace/entities/weave_trace_entity.py

@@ -67,7 +67,8 @@ class WeaveTraceModel(WeaveTokenUsage, WeaveMultiModel):
                 if field_name == "inputs":
                     data = {
                         "messages": [
-                            dict(msg, **{"usage_metadata": usage_metadata, "file_list": file_list}) for msg in v
+                            dict(msg, **{"usage_metadata": usage_metadata, "file_list": file_list})  # type: ignore
+                            for msg in v
                         ]
                         if isinstance(v, list)
                         else v,

api/core/plugin/impl/plugin.py

@@ -209,8 +209,7 @@ class PluginInstaller(BasePluginClient):
             "GET",
             f"plugin/{tenant_id}/management/decode/from_identifier",
             PluginDecodeResponse,
-            data={"plugin_unique_identifier": plugin_unique_identifier},
-            headers={"Content-Type": "application/json"},
+            params={"plugin_unique_identifier": plugin_unique_identifier},
         )
 
     def fetch_plugin_installation_by_ids(

api/core/provider_manager.py

@@ -918,11 +918,11 @@ class ProviderManager:
 
             trail_pool = CreditPoolService.get_pool(
                 tenant_id=tenant_id,
-                pool_type=ProviderQuotaType.TRIAL.value,
+                pool_type=ProviderQuotaType.TRIAL,
             )
             paid_pool = CreditPoolService.get_pool(
                 tenant_id=tenant_id,
-                pool_type=ProviderQuotaType.PAID.value,
+                pool_type=ProviderQuotaType.PAID,
             )
         else:
             trail_pool = None

api/core/rag/cleaner/clean_processor.py

@@ -1,9 +1,10 @@
 import re
+from typing import Any
 
 
 class CleanProcessor:
     @classmethod
-    def clean(cls, text: str, process_rule: dict) -> str:
+    def clean(cls, text: str, process_rule: dict[str, Any] | None) -> str:
         # default clean
         # remove invalid symbol
         text = re.sub(r"<\|", "<", text)

api/core/rag/datasource/keyword/jieba/jieba.py

@@ -4,6 +4,7 @@ from typing import Any
 import orjson
 from pydantic import BaseModel
 from sqlalchemy import select
+from typing_extensions import TypedDict
 
 from configs import dify_config
 from core.rag.datasource.keyword.jieba.jieba_keyword_table_handler import JiebaKeywordTableHandler
@@ -15,6 +16,11 @@ from extensions.ext_storage import storage
 from models.dataset import Dataset, DatasetKeywordTable, DocumentSegment
 
 
+class PreSegmentData(TypedDict):
+    segment: DocumentSegment
+    keywords: list[str]
+
+
 class KeywordTableConfig(BaseModel):
     max_keywords_per_chunk: int = 10
 
@@ -128,7 +134,7 @@ class Jieba(BaseKeyword):
                     file_key = "keyword_files/" + self.dataset.tenant_id + "/" + self.dataset.id + ".txt"
                     storage.delete(file_key)
 
-    def _save_dataset_keyword_table(self, keyword_table):
+    def _save_dataset_keyword_table(self, keyword_table: dict[str, set[str]] | None):
         keyword_table_dict = {
             "__type__": "keyword_table",
             "__data__": {"index_id": self.dataset.id, "summary": None, "table": keyword_table},
@@ -144,7 +150,7 @@ class Jieba(BaseKeyword):
                 storage.delete(file_key)
             storage.save(file_key, dumps_with_sets(keyword_table_dict).encode("utf-8"))
 
-    def _get_dataset_keyword_table(self) -> dict | None:
+    def _get_dataset_keyword_table(self) -> dict[str, set[str]] | None:
         dataset_keyword_table = self.dataset.dataset_keyword_table
         if dataset_keyword_table:
             keyword_table_dict = dataset_keyword_table.keyword_table_dict
@@ -169,14 +175,16 @@ class Jieba(BaseKeyword):
 
         return {}
 
-    def _add_text_to_keyword_table(self, keyword_table: dict, id: str, keywords: list[str]):
+    def _add_text_to_keyword_table(
+        self, keyword_table: dict[str, set[str]], id: str, keywords: list[str]
+    ) -> dict[str, set[str]]:
         for keyword in keywords:
             if keyword not in keyword_table:
                 keyword_table[keyword] = set()
             keyword_table[keyword].add(id)
         return keyword_table
 
-    def _delete_ids_from_keyword_table(self, keyword_table: dict, ids: list[str]):
+    def _delete_ids_from_keyword_table(self, keyword_table: dict[str, set[str]], ids: list[str]) -> dict[str, set[str]]:
         # get set of ids that correspond to node
         node_idxs_to_delete = set(ids)
 
@@ -193,7 +201,7 @@ class Jieba(BaseKeyword):
 
         return keyword_table
 
-    def _retrieve_ids_by_query(self, keyword_table: dict, query: str, k: int = 4):
+    def _retrieve_ids_by_query(self, keyword_table: dict[str, set[str]], query: str, k: int = 4) -> list[str]:
         keyword_table_handler = JiebaKeywordTableHandler()
         keywords = keyword_table_handler.extract_keywords(query)
 
@@ -228,7 +236,7 @@ class Jieba(BaseKeyword):
         keyword_table = self._add_text_to_keyword_table(keyword_table or {}, node_id, keywords)
         self._save_dataset_keyword_table(keyword_table)
 
-    def multi_create_segment_keywords(self, pre_segment_data_list: list):
+    def multi_create_segment_keywords(self, pre_segment_data_list: list[PreSegmentData]):
         keyword_table_handler = JiebaKeywordTableHandler()
         keyword_table = self._get_dataset_keyword_table()
         for pre_segment_data in pre_segment_data_list:

api/core/rag/datasource/retrieval_service.py

@@ -103,7 +103,7 @@ class RetrievalService:
         reranking_mode: str = "reranking_model",
         weights: WeightsDict | None = None,
         document_ids_filter: list[str] | None = None,
-        attachment_ids: list | None = None,
+        attachment_ids: list[str] | None = None,
     ):
         if not query and not attachment_ids:
             return []
@@ -250,8 +250,8 @@ class RetrievalService:
         dataset_id: str,
         query: str,
         top_k: int,
-        all_documents: list,
-        exceptions: list,
+        all_documents: list[Document],
+        exceptions: list[str],
         document_ids_filter: list[str] | None = None,
     ):
         with flask_app.app_context():
@@ -279,9 +279,9 @@ class RetrievalService:
         top_k: int,
         score_threshold: float | None,
         reranking_model: RerankingModelDict | None,
-        all_documents: list,
+        all_documents: list[Document],
         retrieval_method: RetrievalMethod,
-        exceptions: list,
+        exceptions: list[str],
         document_ids_filter: list[str] | None = None,
         query_type: QueryType = QueryType.TEXT_QUERY,
     ):
@@ -373,9 +373,9 @@ class RetrievalService:
         top_k: int,
         score_threshold: float | None,
         reranking_model: RerankingModelDict | None,
-        all_documents: list,
+        all_documents: list[Document],
         retrieval_method: str,
-        exceptions: list,
+        exceptions: list[str],
         document_ids_filter: list[str] | None = None,
     ):
         with flask_app.app_context():

api/core/rag/datasource/vdb/baidu/baidu_vector.py

@@ -13,6 +13,7 @@ from pymochow.exception import ServerError  # type: ignore
 from pymochow.model.database import Database
 from pymochow.model.enum import FieldType, IndexState, IndexType, MetricType, ServerErrCode, TableState  # type: ignore
 from pymochow.model.schema import (
+    AutoBuildRowCountIncrement,
     Field,
     FilteringIndex,
     HNSWParams,
@@ -51,6 +52,9 @@ class BaiduConfig(BaseModel):
     replicas: int = 3
     inverted_index_analyzer: str = "DEFAULT_ANALYZER"
     inverted_index_parser_mode: str = "COARSE_MODE"
+    auto_build_row_count_increment: int = 500
+    auto_build_row_count_increment_ratio: float = 0.05
+    rebuild_index_timeout_in_seconds: int = 300
 
     @model_validator(mode="before")
     @classmethod
@@ -107,18 +111,6 @@ class BaiduVector(BaseVector):
                 rows.append(row)
             table.upsert(rows=rows)
 
-        # rebuild vector index after upsert finished
-        table.rebuild_index(self.vector_index)
-        timeout = 3600  # 1 hour timeout
-        start_time = time.time()
-        while True:
-            time.sleep(1)
-            index = table.describe_index(self.vector_index)
-            if index.state == IndexState.NORMAL:
-                break
-            if time.time() - start_time > timeout:
-                raise TimeoutError(f"Index rebuild timeout after {timeout} seconds")
-
     def text_exists(self, id: str) -> bool:
         res = self._db.table(self._collection_name).query(primary_key={VDBField.PRIMARY_KEY: id})
         if res and res.code == 0:
@@ -232,8 +224,14 @@ class BaiduVector(BaseVector):
             return self._client.database(self._client_config.database)
 
     def _table_existed(self) -> bool:
-        tables = self._db.list_table()
-        return any(table.table_name == self._collection_name for table in tables)
+        try:
+            table = self._db.table(self._collection_name)
+        except ServerError as e:
+            if e.code == ServerErrCode.TABLE_NOT_EXIST:
+                return False
+            else:
+                raise
+        return True
 
     def _create_table(self, dimension: int):
         # Try to grab distributed lock and create table
@@ -287,6 +285,11 @@ class BaiduVector(BaseVector):
                     field=VDBField.VECTOR,
                     metric_type=metric_type,
                     params=HNSWParams(m=16, efconstruction=200),
+                    auto_build=True,
+                    auto_build_index_policy=AutoBuildRowCountIncrement(
+                        row_count_increment=self._client_config.auto_build_row_count_increment,
+                        row_count_increment_ratio=self._client_config.auto_build_row_count_increment_ratio,
+                    ),
                 )
             )
 
@@ -335,7 +338,7 @@ class BaiduVector(BaseVector):
             )
 
             # Wait for table created
-            timeout = 300  # 5 minutes timeout
+            timeout = self._client_config.rebuild_index_timeout_in_seconds  # default 5 minutes timeout
             start_time = time.time()
             while True:
                 time.sleep(1)
@@ -345,6 +348,20 @@ class BaiduVector(BaseVector):
                 if time.time() - start_time > timeout:
                     raise TimeoutError(f"Table creation timeout after {timeout} seconds")
             redis_client.set(table_exist_cache_key, 1, ex=3600)
+            # rebuild vector index immediately after table created, make sure index is ready
+            table.rebuild_index(self.vector_index)
+            timeout = 3600  # 1 hour timeout
+            self._wait_for_index_ready(table, timeout)
+
+    def _wait_for_index_ready(self, table, timeout: int = 3600):
+        start_time = time.time()
+        while True:
+            time.sleep(1)
+            index = table.describe_index(self.vector_index)
+            if index.state == IndexState.NORMAL:
+                break
+            if time.time() - start_time > timeout:
+                raise TimeoutError(f"Index rebuild timeout after {timeout} seconds")
 
 
 class BaiduVectorFactory(AbstractVectorFactory):
@@ -369,5 +386,8 @@ class BaiduVectorFactory(AbstractVectorFactory):
                 replicas=dify_config.BAIDU_VECTOR_DB_REPLICAS,
                 inverted_index_analyzer=dify_config.BAIDU_VECTOR_DB_INVERTED_INDEX_ANALYZER,
                 inverted_index_parser_mode=dify_config.BAIDU_VECTOR_DB_INVERTED_INDEX_PARSER_MODE,
+                auto_build_row_count_increment=dify_config.BAIDU_VECTOR_DB_AUTO_BUILD_ROW_COUNT_INCREMENT,
+                auto_build_row_count_increment_ratio=dify_config.BAIDU_VECTOR_DB_AUTO_BUILD_ROW_COUNT_INCREMENT_RATIO,
+                rebuild_index_timeout_in_seconds=dify_config.BAIDU_VECTOR_DB_REBUILD_INDEX_TIMEOUT_IN_SECONDS,
             ),
         )

api/core/rag/datasource/vdb/huawei/huawei_cloud_vector.py

@@ -124,13 +124,13 @@ class HuaweiCloudVector(BaseVector):
                 )
             )
 
+        score_threshold = float(kwargs.get("score_threshold") or 0.0)
         docs = []
         for doc, score in docs_and_scores:
-            score_threshold = float(kwargs.get("score_threshold") or 0.0)
             if score >= score_threshold:
                 if doc.metadata is not None:
                     doc.metadata["score"] = score
-            docs.append(doc)
+                    docs.append(doc)
 
         return docs
 

api/core/rag/datasource/vdb/tidb_on_qdrant/tidb_on_qdrant_vector.py

@@ -33,6 +33,7 @@ from core.rag.models.document import Document
 from extensions.ext_database import db
 from extensions.ext_redis import redis_client
 from models.dataset import Dataset, TidbAuthBinding
+from models.enums import TidbAuthBindingStatus
 
 if TYPE_CHECKING:
     from qdrant_client import grpc  # noqa
@@ -284,27 +285,29 @@ class TidbOnQdrantVector(BaseVector):
         from qdrant_client.http import models
         from qdrant_client.http.exceptions import UnexpectedResponse
 
-        for node_id in ids:
-            try:
-                filter = models.Filter(
-                    must=[
-                        models.FieldCondition(
-                            key="metadata.doc_id",
-                            match=models.MatchValue(value=node_id),
-                        ),
-                    ],
-                )
-                self._client.delete(
-                    collection_name=self._collection_name,
-                    points_selector=FilterSelector(filter=filter),
-                )
-            except UnexpectedResponse as e:
-                # Collection does not exist, so return
-                if e.status_code == 404:
-                    return
-                # Some other error occurred, so re-raise the exception
-                else:
-                    raise e
+        if not ids:
+            return
+
+        try:
+            filter = models.Filter(
+                must=[
+                    models.FieldCondition(
+                        key="metadata.doc_id",
+                        match=models.MatchAny(any=ids),
+                    ),
+                ],
+            )
+            self._client.delete(
+                collection_name=self._collection_name,
+                points_selector=FilterSelector(filter=filter),
+            )
+        except UnexpectedResponse as e:
+            # Collection does not exist, so return
+            if e.status_code == 404:
+                return
+            # Some other error occurred, so re-raise the exception
+            else:
+                raise e
 
     def text_exists(self, id: str) -> bool:
         all_collection_name = []
@@ -450,7 +453,7 @@ class TidbOnQdrantVectorFactory(AbstractVectorFactory):
                             password=new_cluster["password"],
                             tenant_id=dataset.tenant_id,
                             active=True,
-                            status="ACTIVE",
+                            status=TidbAuthBindingStatus.ACTIVE,
                         )
                         db.session.add(new_tidb_auth_binding)
                         db.session.commit()

api/core/rag/datasource/vdb/tidb_on_qdrant/tidb_service.py

@@ -9,6 +9,7 @@ from configs import dify_config
 from extensions.ext_database import db
 from extensions.ext_redis import redis_client
 from models.dataset import TidbAuthBinding
+from models.enums import TidbAuthBindingStatus
 
 
 class TidbService:
@@ -170,7 +171,7 @@ class TidbService:
                 userPrefix = item["userPrefix"]
                 if state == "ACTIVE" and len(userPrefix) > 0:
                     cluster_info = tidb_serverless_list_map[item["clusterId"]]
-                    cluster_info.status = "ACTIVE"
+                    cluster_info.status = TidbAuthBindingStatus.ACTIVE
                     cluster_info.account = f"{userPrefix}.root"
                     db.session.add(cluster_info)
             db.session.commit()

api/core/rag/docstore/dataset_docstore.py

@@ -6,6 +6,7 @@ from typing import Any
 from sqlalchemy import func, select
 
 from core.model_manager import ModelManager
+from core.rag.index_processor.constant.index_type import IndexTechniqueType
 from core.rag.models.document import AttachmentDocument, Document
 from dify_graph.model_runtime.entities.model_entities import ModelType
 from extensions.ext_database import db
@@ -71,7 +72,7 @@ class DatasetDocumentStore:
         if max_position is None:
             max_position = 0
         embedding_model = None
-        if self._dataset.indexing_technique == "high_quality":
+        if self._dataset.indexing_technique == IndexTechniqueType.HIGH_QUALITY:
             model_manager = ModelManager()
             embedding_model = model_manager.get_model_instance(
                 tenant_id=self._dataset.tenant_id,

api/core/rag/extractor/firecrawl/firecrawl_app.py

@@ -95,15 +95,11 @@ class FirecrawlApp:
         if response.status_code == 200:
             crawl_status_response = response.json()
             if crawl_status_response.get("status") == "completed":
-                total = crawl_status_response.get("total", 0)
-                if total == 0:
+                # Normalize to avoid None bypassing the zero-guard when the API returns null.
+                total = crawl_status_response.get("total") or 0
+                if total <= 0:
                     raise Exception("Failed to check crawl status. Error: No page found")
-                data = crawl_status_response.get("data", [])
-                url_data_list: list[FirecrawlDocumentData] = []
-                for item in data:
-                    if isinstance(item, dict) and "metadata" in item and "markdown" in item:
-                        url_data = self._extract_common_fields(item)
-                        url_data_list.append(url_data)
+                url_data_list = self._collect_all_crawl_pages(crawl_status_response, headers)
                 if url_data_list:
                     file_key = "website_files/" + job_id + ".txt"
                     try:
@@ -120,6 +116,36 @@ class FirecrawlApp:
         self._handle_error(response, "check crawl status")
         raise RuntimeError("unreachable: _handle_error always raises")
 
+    def _collect_all_crawl_pages(
+        self, first_page: dict[str, Any], headers: dict[str, str]
+    ) -> list[FirecrawlDocumentData]:
+        """Collect all crawl result pages by following pagination links.
+
+        Raises an exception if any paginated request fails, to avoid returning
+        partial data that is inconsistent with the reported total.
+
+        The number of pages processed is capped at ``total`` (the
+        server-reported page count) to guard against infinite loops caused by
+        a misbehaving server that keeps returning a ``next`` URL.
+        """
+        total: int = first_page.get("total") or 0
+        url_data_list: list[FirecrawlDocumentData] = []
+        current_page = first_page
+        pages_processed = 0
+        while True:
+            for item in current_page.get("data", []):
+                if isinstance(item, dict) and "metadata" in item and "markdown" in item:
+                    url_data_list.append(self._extract_common_fields(item))
+            next_url: str | None = current_page.get("next")
+            pages_processed += 1
+            if not next_url or pages_processed >= total:
+                break
+            response = self._get_request(next_url, headers)
+            if response.status_code != 200:
+                self._handle_error(response, "fetch next crawl page")
+            current_page = response.json()
+        return url_data_list
+
     def _format_crawl_status_response(
         self,
         status: str,

api/core/rag/extractor/word_extractor.py

@@ -366,7 +366,7 @@ class WordExtractor(BaseExtractor):
             paragraph_content = []
             # State for legacy HYPERLINK fields
             hyperlink_field_url = None
-            hyperlink_field_text_parts: list = []
+            hyperlink_field_text_parts: list[str] = []
             is_collecting_field_text = False
             # Iterate through paragraph elements in document order
             for child in paragraph._element:

api/core/rag/index_processor/index_processor.py

@@ -9,6 +9,7 @@ from flask import current_app
 from sqlalchemy import delete, func, select
 
 from core.db.session_factory import session_factory
+from core.rag.index_processor.constant.index_type import IndexTechniqueType
 from core.rag.index_processor.index_processor_base import SummaryIndexSettingDict
 from core.workflow.nodes.knowledge_index.exc import KnowledgeIndexNodeError
 from core.workflow.nodes.knowledge_index.protocols import Preview, PreviewItem, QaPreview
@@ -159,7 +160,7 @@ class IndexProcessor:
             tenant_id = dataset.tenant_id
 
         preview_output = self.format_preview(chunk_structure, chunks)
-        if indexing_technique != "high_quality":
+        if indexing_technique != IndexTechniqueType.HIGH_QUALITY:
             return preview_output
 
         if not summary_index_setting or not summary_index_setting.get("enable"):

api/core/rag/index_processor/processor/paragraph_index_processor.py

@@ -22,7 +22,7 @@ from core.rag.docstore.dataset_docstore import DatasetDocumentStore
 from core.rag.extractor.entity.extract_setting import ExtractSetting
 from core.rag.extractor.extract_processor import ExtractProcessor
 from core.rag.index_processor.constant.doc_type import DocType
-from core.rag.index_processor.constant.index_type import IndexStructureType
+from core.rag.index_processor.constant.index_type import IndexStructureType, IndexTechniqueType
 from core.rag.index_processor.index_processor_base import BaseIndexProcessor, SummaryIndexSettingDict
 from core.rag.models.document import AttachmentDocument, Document, MultimodalGeneralStructureChunk
 from core.rag.retrieval.retrieval_methods import RetrievalMethod
@@ -117,7 +117,7 @@ class ParagraphIndexProcessor(BaseIndexProcessor):
         with_keywords: bool = True,
         **kwargs,
     ) -> None:
-        if dataset.indexing_technique == "high_quality":
+        if dataset.indexing_technique == IndexTechniqueType.HIGH_QUALITY:
             vector = Vector(dataset)
             vector.create(documents)
             if multimodal_documents and dataset.is_multimodal:
@@ -155,7 +155,7 @@ class ParagraphIndexProcessor(BaseIndexProcessor):
                 # Delete all summaries for the dataset
                 SummaryIndexService.delete_summaries_for_segments(dataset, None)
 
-        if dataset.indexing_technique == "high_quality":
+        if dataset.indexing_technique == IndexTechniqueType.HIGH_QUALITY:
             vector = Vector(dataset)
             if node_ids:
                 vector.delete_by_ids(node_ids)
@@ -253,12 +253,12 @@ class ParagraphIndexProcessor(BaseIndexProcessor):
             doc_store = DatasetDocumentStore(dataset=dataset, user_id=document.created_by, document_id=document.id)
             # add document segments
             doc_store.add_documents(docs=documents, save_child=False)
-            if dataset.indexing_technique == "high_quality":
+            if dataset.indexing_technique == IndexTechniqueType.HIGH_QUALITY:
                 vector = Vector(dataset)
                 vector.create(documents)
                 if all_multimodal_documents and dataset.is_multimodal:
                     vector.create_multimodal(all_multimodal_documents)
-            elif dataset.indexing_technique == "economy":
+            elif dataset.indexing_technique == IndexTechniqueType.ECONOMY:
                 keyword = Keyword(dataset)
                 keyword.add_texts(documents)
 

api/core/rag/index_processor/processor/parent_child_index_processor.py

@@ -18,7 +18,7 @@ from core.rag.docstore.dataset_docstore import DatasetDocumentStore
 from core.rag.extractor.entity.extract_setting import ExtractSetting
 from core.rag.extractor.extract_processor import ExtractProcessor
 from core.rag.index_processor.constant.doc_type import DocType
-from core.rag.index_processor.constant.index_type import IndexStructureType
+from core.rag.index_processor.constant.index_type import IndexStructureType, IndexTechniqueType
 from core.rag.index_processor.index_processor_base import BaseIndexProcessor, SummaryIndexSettingDict
 from core.rag.models.document import AttachmentDocument, ChildDocument, Document, ParentChildStructureChunk
 from core.rag.retrieval.retrieval_methods import RetrievalMethod
@@ -128,7 +128,7 @@ class ParentChildIndexProcessor(BaseIndexProcessor):
         with_keywords: bool = True,
         **kwargs,
     ) -> None:
-        if dataset.indexing_technique == "high_quality":
+        if dataset.indexing_technique == IndexTechniqueType.HIGH_QUALITY:
             vector = Vector(dataset)
             for document in documents:
                 child_documents = document.children
@@ -166,7 +166,7 @@ class ParentChildIndexProcessor(BaseIndexProcessor):
                 # Delete all summaries for the dataset
                 SummaryIndexService.delete_summaries_for_segments(dataset, None)
 
-        if dataset.indexing_technique == "high_quality":
+        if dataset.indexing_technique == IndexTechniqueType.HIGH_QUALITY:
             delete_child_chunks = kwargs.get("delete_child_chunks") or False
             precomputed_child_node_ids = kwargs.get("precomputed_child_node_ids")
             vector = Vector(dataset)
@@ -332,7 +332,7 @@ class ParentChildIndexProcessor(BaseIndexProcessor):
             doc_store = DatasetDocumentStore(dataset=dataset, user_id=document.created_by, document_id=document.id)
             # add document segments
             doc_store.add_documents(docs=documents, save_child=True)
-            if dataset.indexing_technique == "high_quality":
+            if dataset.indexing_technique == IndexTechniqueType.HIGH_QUALITY:
                 all_child_documents = []
                 all_multimodal_documents = []
                 for doc in documents:

api/core/rag/index_processor/processor/qa_index_processor.py

@@ -21,7 +21,7 @@ from core.rag.datasource.vdb.vector_factory import Vector
 from core.rag.docstore.dataset_docstore import DatasetDocumentStore
 from core.rag.extractor.entity.extract_setting import ExtractSetting
 from core.rag.extractor.extract_processor import ExtractProcessor
-from core.rag.index_processor.constant.index_type import IndexStructureType
+from core.rag.index_processor.constant.index_type import IndexStructureType, IndexTechniqueType
 from core.rag.index_processor.index_processor_base import BaseIndexProcessor, SummaryIndexSettingDict
 from core.rag.models.document import AttachmentDocument, Document, QAStructureChunk
 from core.rag.retrieval.retrieval_methods import RetrievalMethod
@@ -141,7 +141,7 @@ class QAIndexProcessor(BaseIndexProcessor):
         with_keywords: bool = True,
         **kwargs,
     ) -> None:
-        if dataset.indexing_technique == "high_quality":
+        if dataset.indexing_technique == IndexTechniqueType.HIGH_QUALITY:
             vector = Vector(dataset)
             vector.create(documents)
             if multimodal_documents and dataset.is_multimodal:
@@ -224,7 +224,7 @@ class QAIndexProcessor(BaseIndexProcessor):
             # save node to document segment
             doc_store = DatasetDocumentStore(dataset=dataset, user_id=document.created_by, document_id=document.id)
             doc_store.add_documents(docs=documents, save_child=False)
-            if dataset.indexing_technique == "high_quality":
+            if dataset.indexing_technique == IndexTechniqueType.HIGH_QUALITY:
                 vector = Vector(dataset)
                 vector.create(documents)
             else:

api/core/rag/retrieval/dataset_retrieval.py

@@ -591,7 +591,7 @@ class DatasetRetrieval:
         user_id: str,
         user_from: str,
         query: str,
-        available_datasets: list,
+        available_datasets: list[Dataset],
         model_instance: ModelInstance,
         model_config: ModelConfigWithCredentialsEntity,
         planning_strategy: PlanningStrategy,
@@ -633,15 +633,15 @@ class DatasetRetrieval:
         if dataset_id:
             # get retrieval model config
             dataset_stmt = select(Dataset).where(Dataset.id == dataset_id)
-            dataset = db.session.scalar(dataset_stmt)
-            if dataset:
+            selected_dataset = db.session.scalar(dataset_stmt)
+            if selected_dataset:
                 results = []
-                if dataset.provider == "external":
+                if selected_dataset.provider == "external":
                     external_documents = ExternalDatasetService.fetch_external_knowledge_retrieval(
-                        tenant_id=dataset.tenant_id,
+                        tenant_id=selected_dataset.tenant_id,
                         dataset_id=dataset_id,
                         query=query,
-                        external_retrieval_parameters=dataset.retrieval_model,
+                        external_retrieval_parameters=selected_dataset.retrieval_model,
                         metadata_condition=metadata_condition,
                     )
                     for external_document in external_documents:
@@ -654,28 +654,28 @@ class DatasetRetrieval:
                             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
+                            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(dataset.id, [])
+                        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, dataset.retrieval_model)
-                        if dataset.retrieval_model
+                        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 dataset.indexing_technique == "economy":
+                    if selected_dataset.indexing_technique == IndexTechniqueType.ECONOMY:
                         retrieval_method = RetrievalMethod.KEYWORD_SEARCH
                     else:
                         retrieval_method = retrieval_model_config["search_method"]
@@ -694,7 +694,7 @@ class DatasetRetrieval:
                     with measure_time() as timer:
                         results = RetrievalService.retrieve(
                             retrieval_method=retrieval_method,
-                            dataset_id=dataset.id,
+                            dataset_id=selected_dataset.id,
                             query=query,
                             top_k=top_k,
                             score_threshold=score_threshold,
@@ -726,7 +726,7 @@ class DatasetRetrieval:
         tenant_id: str,
         user_id: str,
         user_from: str,
-        available_datasets: list,
+        available_datasets: list[Dataset],
         query: str | None,
         top_k: int,
         score_threshold: float,
@@ -752,7 +752,7 @@ class DatasetRetrieval:
                 "The configured knowledge base list have different indexing technique, please set reranking model."
             )
         index_type = available_datasets[0].indexing_technique
-        if index_type == "high_quality":
+        if index_type == IndexTechniqueType.HIGH_QUALITY:
             embedding_model_check = all(
                 item.embedding_model == available_datasets[0].embedding_model for item in available_datasets
             )
@@ -1028,7 +1028,7 @@ class DatasetRetrieval:
         dataset_id: str,
         query: str,
         top_k: int,
-        all_documents: list,
+        all_documents: list[Document],
         document_ids_filter: list[str] | None = None,
         metadata_condition: MetadataCondition | None = None,
         attachment_ids: list[str] | None = None,
@@ -1068,7 +1068,7 @@ class DatasetRetrieval:
                     else default_retrieval_model
                 )
 
-                if dataset.indexing_technique == "economy":
+                if dataset.indexing_technique == IndexTechniqueType.ECONOMY:
                     # use keyword table query
                     documents = RetrievalService.retrieve(
                         retrieval_method=RetrievalMethod.KEYWORD_SEARCH,
@@ -1298,7 +1298,7 @@ class DatasetRetrieval:
 
     def get_metadata_filter_condition(
         self,
-        dataset_ids: list,
+        dataset_ids: list[str],
         query: str,
         tenant_id: str,
         user_id: str,
@@ -1400,7 +1400,7 @@ class DatasetRetrieval:
         return output
 
     def _automatic_metadata_filter_func(
-        self, dataset_ids: list, query: str, tenant_id: str, user_id: str, metadata_model_config: ModelConfig
+        self, 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))
@@ -1598,7 +1598,7 @@ class DatasetRetrieval:
         )
 
     def _get_prompt_template(
-        self, model_config: ModelConfigWithCredentialsEntity, mode: str, metadata_fields: list, query: str
+        self, model_config: ModelConfigWithCredentialsEntity, mode: str, metadata_fields: list[str], query: str
     ):
         model_mode = ModelMode(mode)
         input_text = query
@@ -1690,7 +1690,7 @@ class DatasetRetrieval:
     def _multiple_retrieve_thread(
         self,
         flask_app: Flask,
-        available_datasets: list,
+        available_datasets: list[Dataset],
         metadata_condition: MetadataCondition | None,
         metadata_filter_document_ids: dict[str, list[str]] | None,
         all_documents: list[Document],

api/core/rag/summary_index/summary_index.py

@@ -2,6 +2,7 @@ import concurrent.futures
 import logging
 
 from core.db.session_factory import session_factory
+from core.rag.index_processor.constant.index_type import IndexTechniqueType
 from core.rag.index_processor.index_processor_base import SummaryIndexSettingDict
 from models.dataset import Dataset, Document, DocumentSegment, DocumentSegmentSummary
 from services.summary_index_service import SummaryIndexService
@@ -21,7 +22,7 @@ class SummaryIndex:
         if is_preview:
             with session_factory.create_session() as session:
                 dataset = session.query(Dataset).filter_by(id=dataset_id).first()
-                if not dataset or dataset.indexing_technique != "high_quality":
+                if not dataset or dataset.indexing_technique != IndexTechniqueType.HIGH_QUALITY:
                     return
 
                 if summary_index_setting is None:

api/core/tools/builtin_tool/tool.py

@@ -50,7 +50,7 @@ class BuiltinTool(Tool):
         return ModelInvocationUtils.invoke(
             user_id=user_id,
             tenant_id=self.runtime.tenant_id or "",
-            tool_type="builtin",
+            tool_type=ToolProviderType.BUILT_IN,
             tool_name=self.entity.identity.name,
             prompt_messages=prompt_messages,
         )

api/core/tools/tool_label_manager.py

@@ -38,7 +38,7 @@ class ToolLabelManager:
             db.session.add(
                 ToolLabelBinding(
                     tool_id=provider_id,
-                    tool_type=controller.provider_type.value,
+                    tool_type=controller.provider_type,
                     label_name=label,
                 )
             )
@@ -58,7 +58,7 @@ class ToolLabelManager:
             raise ValueError("Unsupported tool type")
         stmt = select(ToolLabelBinding.label_name).where(
             ToolLabelBinding.tool_id == provider_id,
-            ToolLabelBinding.tool_type == controller.provider_type.value,
+            ToolLabelBinding.tool_type == controller.provider_type,
         )
         labels = db.session.scalars(stmt).all()
 

api/core/tools/tool_manager.py

@@ -1045,9 +1045,10 @@ class ToolManager:
                         continue
                     tool_input = ToolNodeData.ToolInput.model_validate(tool_configurations.get(parameter.name, {}))
                     if tool_input.type == "variable":
-                        variable = variable_pool.get(tool_input.value)
+                        variable_selector = tool_input.require_variable_selector()
+                        variable = variable_pool.get(variable_selector)
                         if variable is None:
-                            raise ToolParameterError(f"Variable {tool_input.value} does not exist")
+                            raise ToolParameterError(f"Variable {variable_selector} does not exist")
                         parameter_value = variable.value
                     elif tool_input.type == "constant":
                         parameter_value = tool_input.value

api/core/tools/utils/dataset_retriever/dataset_multi_retriever_tool.py

@@ -8,6 +8,7 @@ from core.callback_handler.index_tool_callback_handler import DatasetIndexToolCa
 from core.model_manager import ModelManager
 from core.rag.datasource.retrieval_service import RetrievalService
 from core.rag.entities.citation_metadata 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
@@ -169,7 +170,7 @@ class DatasetMultiRetrieverTool(DatasetRetrieverBaseTool):
             # get retrieval model , if the model is not setting , using default
             retrieval_model = dataset.retrieval_model or default_retrieval_model
 
-            if dataset.indexing_technique == "economy":
+            if dataset.indexing_technique == IndexTechniqueType.ECONOMY:
                 # use keyword table query
                 documents = RetrievalService.retrieve(
                     retrieval_method=RetrievalMethod.KEYWORD_SEARCH,

api/core/tools/utils/dataset_retriever/dataset_retriever_tool.py

@@ -8,6 +8,7 @@ from core.rag.data_post_processor.data_post_processor import RerankingModelDict,
 from core.rag.datasource.retrieval_service import RetrievalService
 from core.rag.entities.citation_metadata import RetrievalSourceMetadata
 from core.rag.entities.context_entities import DocumentContext
+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
@@ -140,7 +141,7 @@ class DatasetRetrieverTool(DatasetRetrieverBaseTool):
             # 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 == "economy":
+            if dataset.indexing_technique == IndexTechniqueType.ECONOMY:
                 # use keyword table query
                 documents = RetrievalService.retrieve(
                     retrieval_method=RetrievalMethod.KEYWORD_SEARCH,
@@ -173,7 +174,7 @@ class DatasetRetrieverTool(DatasetRetrieverBaseTool):
                 for hit_callback in self.hit_callbacks:
                     hit_callback.on_tool_end(documents)
                 document_score_list = {}
-                if dataset.indexing_technique != "economy":
+                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"]

api/core/tools/utils/model_invocation_utils.py

@@ -9,6 +9,7 @@ from decimal import Decimal
 from typing import cast
 
 from core.model_manager import ModelManager
+from core.tools.entities.tool_entities import ToolProviderType
 from dify_graph.model_runtime.entities.llm_entities import LLMResult
 from dify_graph.model_runtime.entities.message_entities import PromptMessage
 from dify_graph.model_runtime.entities.model_entities import ModelPropertyKey, ModelType
@@ -78,7 +79,7 @@ class ModelInvocationUtils:
 
     @staticmethod
     def invoke(
-        user_id: str, tenant_id: str, tool_type: str, tool_name: str, prompt_messages: list[PromptMessage]
+        user_id: str, tenant_id: str, tool_type: ToolProviderType, tool_name: str, prompt_messages: list[PromptMessage]
     ) -> LLMResult:
         """
         invoke model with parameters in user's own context

api/core/workflow/nodes/agent/entities.py

@@ -1,13 +1,24 @@
-from enum import IntEnum, StrEnum, auto
-from typing import Any, Literal, Union
+from __future__ import annotations
 
-from pydantic import BaseModel
+from enum import IntEnum, StrEnum, auto
+from typing import Literal, TypeAlias
+
+from pydantic import BaseModel, TypeAdapter, field_validator
+from pydantic_core.core_schema import ValidationInfo
 
 from core.prompt.entities.advanced_prompt_entities import MemoryConfig
 from core.tools.entities.tool_entities import ToolSelector
 from dify_graph.entities.base_node_data import BaseNodeData
 from dify_graph.enums import BuiltinNodeTypes, NodeType
 
+AgentInputConstantValue: TypeAlias = (
+    list[ToolSelector] | str | int | float | bool | dict[str, object] | list[object] | None
+)
+VariableSelector: TypeAlias = list[str]
+
+_AGENT_INPUT_VALUE_ADAPTER: TypeAdapter[AgentInputConstantValue] = TypeAdapter(AgentInputConstantValue)
+_AGENT_VARIABLE_SELECTOR_ADAPTER: TypeAdapter[VariableSelector] = TypeAdapter(VariableSelector)
+
 
 class AgentNodeData(BaseNodeData):
     type: NodeType = BuiltinNodeTypes.AGENT
@@ -21,8 +32,20 @@ class AgentNodeData(BaseNodeData):
     tool_node_version: str | None = None
 
     class AgentInput(BaseModel):
-        value: Union[list[str], list[ToolSelector], Any]
         type: Literal["mixed", "variable", "constant"]
+        value: AgentInputConstantValue | VariableSelector
+
+        @field_validator("value", mode="before")
+        @classmethod
+        def validate_value(
+            cls, value: object, validation_info: ValidationInfo
+        ) -> AgentInputConstantValue | VariableSelector:
+            input_type = validation_info.data.get("type")
+            if input_type == "variable":
+                return _AGENT_VARIABLE_SELECTOR_ADAPTER.validate_python(value)
+            if input_type in {"mixed", "constant"}:
+                return _AGENT_INPUT_VALUE_ADAPTER.validate_python(value)
+            raise ValueError(f"Unknown agent input type: {input_type}")
 
     agent_parameters: dict[str, AgentInput]
 

api/core/workflow/nodes/agent/runtime_support.py

@@ -1,16 +1,17 @@
 from __future__ import annotations
 
 import json
-from collections.abc import Sequence
-from typing import Any, cast
+from collections.abc import Mapping, Sequence
+from typing import TypeAlias
 
 from packaging.version import Version
-from pydantic import ValidationError
+from pydantic import TypeAdapter, ValidationError
 from sqlalchemy import select
 from sqlalchemy.orm import Session
 
 from core.agent.entities import AgentToolEntity
 from core.agent.plugin_entities import AgentStrategyParameter
+from core.app.entities.app_invoke_entities import InvokeFrom
 from core.memory.token_buffer_memory import TokenBufferMemory
 from core.model_manager import ModelInstance, ModelManager
 from core.plugin.entities.request import InvokeCredentials
@@ -28,6 +29,14 @@ from .entities import AgentNodeData, AgentOldVersionModelFeatures, ParamsAutoGen
 from .exceptions import AgentInputTypeError, AgentVariableNotFoundError
 from .strategy_protocols import ResolvedAgentStrategy
 
+JsonObject: TypeAlias = dict[str, object]
+JsonObjectList: TypeAlias = list[JsonObject]
+VariableSelector: TypeAlias = list[str]
+
+_JSON_OBJECT_ADAPTER = TypeAdapter(JsonObject)
+_JSON_OBJECT_LIST_ADAPTER = TypeAdapter(JsonObjectList)
+_VARIABLE_SELECTOR_ADAPTER = TypeAdapter(VariableSelector)
+
 
 class AgentRuntimeSupport:
     def build_parameters(
@@ -39,12 +48,12 @@ class AgentRuntimeSupport:
         strategy: ResolvedAgentStrategy,
         tenant_id: str,
         app_id: str,
-        invoke_from: Any,
+        invoke_from: InvokeFrom,
         for_log: bool = False,
-    ) -> dict[str, Any]:
+    ) -> dict[str, object]:
         agent_parameters_dictionary = {parameter.name: parameter for parameter in agent_parameters}
 
-        result: dict[str, Any] = {}
+        result: dict[str, object] = {}
         for parameter_name in node_data.agent_parameters:
             parameter = agent_parameters_dictionary.get(parameter_name)
             if not parameter:
@@ -54,9 +63,10 @@ class AgentRuntimeSupport:
             agent_input = node_data.agent_parameters[parameter_name]
             match agent_input.type:
                 case "variable":
-                    variable = variable_pool.get(agent_input.value)  # type: ignore[arg-type]
+                    variable_selector = _VARIABLE_SELECTOR_ADAPTER.validate_python(agent_input.value)
+                    variable = variable_pool.get(variable_selector)
                     if variable is None:
-                        raise AgentVariableNotFoundError(str(agent_input.value))
+                        raise AgentVariableNotFoundError(str(variable_selector))
                     parameter_value = variable.value
                 case "mixed" | "constant":
                     try:
@@ -79,60 +89,38 @@ class AgentRuntimeSupport:
 
             value = parameter_value
             if parameter.type == "array[tools]":
-                value = cast(list[dict[str, Any]], value)
-                value = [tool for tool in value if tool.get("enabled", False)]
-                value = self._filter_mcp_type_tool(strategy, value)
-                for tool in value:
-                    if "schemas" in tool:
-                        tool.pop("schemas")
-                    parameters = tool.get("parameters", {})
-                    if all(isinstance(v, dict) for _, v in parameters.items()):
-                        params = {}
-                        for key, param in parameters.items():
-                            if param.get("auto", ParamsAutoGenerated.OPEN) in (
-                                ParamsAutoGenerated.CLOSE,
-                                0,
-                            ):
-                                value_param = param.get("value", {})
-                                if value_param and value_param.get("type", "") == "variable":
-                                    variable_selector = value_param.get("value")
-                                    if not variable_selector:
-                                        raise ValueError("Variable selector is missing for a variable-type parameter.")
-
-                                    variable = variable_pool.get(variable_selector)
-                                    if variable is None:
-                                        raise AgentVariableNotFoundError(str(variable_selector))
-
-                                    params[key] = variable.value
-                                else:
-                                    params[key] = value_param.get("value", "") if value_param is not None else None
-                            else:
-                                params[key] = None
-                        parameters = params
-                    tool["settings"] = {k: v.get("value", None) for k, v in tool.get("settings", {}).items()}
-                    tool["parameters"] = parameters
+                tool_payloads = _JSON_OBJECT_LIST_ADAPTER.validate_python(value)
+                value = self._normalize_tool_payloads(
+                    strategy=strategy,
+                    tools=tool_payloads,
+                    variable_pool=variable_pool,
+                )
 
             if not for_log:
                 if parameter.type == "array[tools]":
-                    value = cast(list[dict[str, Any]], value)
+                    value = _JSON_OBJECT_LIST_ADAPTER.validate_python(value)
                     tool_value = []
                     for tool in value:
-                        provider_type = ToolProviderType(tool.get("type", ToolProviderType.BUILT_IN))
-                        setting_params = tool.get("settings", {})
-                        parameters = tool.get("parameters", {})
+                        provider_type = self._coerce_tool_provider_type(tool.get("type"))
+                        setting_params = self._coerce_json_object(tool.get("settings")) or {}
+                        parameters = self._coerce_json_object(tool.get("parameters")) or {}
                         manual_input_params = [key for key, value in parameters.items() if value is not None]
 
                         parameters = {**parameters, **setting_params}
+                        provider_id = self._coerce_optional_string(tool.get("provider_name")) or ""
+                        tool_name = self._coerce_optional_string(tool.get("tool_name")) or ""
+                        plugin_unique_identifier = self._coerce_optional_string(tool.get("plugin_unique_identifier"))
+                        credential_id = self._coerce_optional_string(tool.get("credential_id"))
                         entity = AgentToolEntity(
-                            provider_id=tool.get("provider_name", ""),
+                            provider_id=provider_id,
                             provider_type=provider_type,
-                            tool_name=tool.get("tool_name", ""),
+                            tool_name=tool_name,
                             tool_parameters=parameters,
-                            plugin_unique_identifier=tool.get("plugin_unique_identifier", None),
-                            credential_id=tool.get("credential_id", None),
+                            plugin_unique_identifier=plugin_unique_identifier,
+                            credential_id=credential_id,
                         )
 
-                        extra = tool.get("extra", {})
+                        extra = self._coerce_json_object(tool.get("extra")) or {}
 
                         runtime_variable_pool: VariablePool | None = None
                         if node_data.version != "1" or node_data.tool_node_version is not None:
@@ -145,8 +133,9 @@ class AgentRuntimeSupport:
                             runtime_variable_pool,
                         )
                         if tool_runtime.entity.description:
+                            description_override = self._coerce_optional_string(extra.get("description"))
                             tool_runtime.entity.description.llm = (
-                                extra.get("description", "") or tool_runtime.entity.description.llm
+                                description_override or tool_runtime.entity.description.llm
                             )
                         for tool_runtime_params in tool_runtime.entity.parameters:
                             tool_runtime_params.form = (
@@ -167,13 +156,13 @@ class AgentRuntimeSupport:
                             {
                                 **tool_runtime.entity.model_dump(mode="json"),
                                 "runtime_parameters": runtime_parameters,
-                                "credential_id": tool.get("credential_id", None),
+                                "credential_id": credential_id,
                                 "provider_type": provider_type.value,
                             }
                         )
                     value = tool_value
                 if parameter.type == AgentStrategyParameter.AgentStrategyParameterType.MODEL_SELECTOR:
-                    value = cast(dict[str, Any], value)
+                    value = _JSON_OBJECT_ADAPTER.validate_python(value)
                     model_instance, model_schema = self.fetch_model(tenant_id=tenant_id, value=value)
                     history_prompt_messages = []
                     if node_data.memory:
@@ -199,17 +188,27 @@ class AgentRuntimeSupport:
 
         return result
 
-    def build_credentials(self, *, parameters: dict[str, Any]) -> InvokeCredentials:
+    def build_credentials(self, *, parameters: Mapping[str, object]) -> InvokeCredentials:
         credentials = InvokeCredentials()
         credentials.tool_credentials = {}
-        for tool in parameters.get("tools", []):
+        tools = parameters.get("tools")
+        if not isinstance(tools, list):
+            return credentials
+
+        for raw_tool in tools:
+            tool = self._coerce_json_object(raw_tool)
+            if tool is None:
+                continue
             if not tool.get("credential_id"):
                 continue
             try:
                 identity = ToolIdentity.model_validate(tool.get("identity", {}))
             except ValidationError:
                 continue
-            credentials.tool_credentials[identity.provider] = tool.get("credential_id", None)
+            credential_id = self._coerce_optional_string(tool.get("credential_id"))
+            if credential_id is None:
+                continue
+            credentials.tool_credentials[identity.provider] = credential_id
         return credentials
 
     def fetch_memory(
@@ -232,14 +231,14 @@ class AgentRuntimeSupport:
 
         return TokenBufferMemory(conversation=conversation, model_instance=model_instance)
 
-    def fetch_model(self, *, tenant_id: str, value: dict[str, Any]) -> tuple[ModelInstance, AIModelEntity | None]:
+    def fetch_model(self, *, tenant_id: str, value: Mapping[str, object]) -> tuple[ModelInstance, AIModelEntity | None]:
         provider_manager = ProviderManager()
         provider_model_bundle = provider_manager.get_provider_model_bundle(
             tenant_id=tenant_id,
-            provider=value.get("provider", ""),
+            provider=str(value.get("provider", "")),
             model_type=ModelType.LLM,
         )
-        model_name = value.get("model", "")
+        model_name = str(value.get("model", ""))
         model_credentials = provider_model_bundle.configuration.get_current_credentials(
             model_type=ModelType.LLM,
             model=model_name,
@@ -249,7 +248,7 @@ class AgentRuntimeSupport:
         model_instance = ModelManager().get_model_instance(
             tenant_id=tenant_id,
             provider=provider_name,
-            model_type=ModelType(value.get("model_type", "")),
+            model_type=ModelType(str(value.get("model_type", ""))),
             model=model_name,
         )
         model_schema = model_type_instance.get_model_schema(model_name, model_credentials)
@@ -268,9 +267,88 @@ class AgentRuntimeSupport:
     @staticmethod
     def _filter_mcp_type_tool(
         strategy: ResolvedAgentStrategy,
-        tools: list[dict[str, Any]],
-    ) -> list[dict[str, Any]]:
+        tools: JsonObjectList,
+    ) -> JsonObjectList:
         meta_version = strategy.meta_version
         if meta_version and Version(meta_version) > Version("0.0.1"):
             return tools
         return [tool for tool in tools if tool.get("type") != ToolProviderType.MCP]
+
+    def _normalize_tool_payloads(
+        self,
+        *,
+        strategy: ResolvedAgentStrategy,
+        tools: JsonObjectList,
+        variable_pool: VariablePool,
+    ) -> JsonObjectList:
+        enabled_tools = [dict(tool) for tool in tools if bool(tool.get("enabled", False))]
+        normalized_tools = self._filter_mcp_type_tool(strategy, enabled_tools)
+        for tool in normalized_tools:
+            tool.pop("schemas", None)
+            tool["parameters"] = self._resolve_tool_parameters(tool=tool, variable_pool=variable_pool)
+            tool["settings"] = self._resolve_tool_settings(tool)
+        return normalized_tools
+
+    def _resolve_tool_parameters(self, *, tool: Mapping[str, object], variable_pool: VariablePool) -> JsonObject:
+        parameter_configs = self._coerce_named_json_objects(tool.get("parameters"))
+        if parameter_configs is None:
+            raw_parameters = self._coerce_json_object(tool.get("parameters"))
+            return raw_parameters or {}
+
+        resolved_parameters: JsonObject = {}
+        for key, parameter_config in parameter_configs.items():
+            if parameter_config.get("auto", ParamsAutoGenerated.OPEN) in (ParamsAutoGenerated.CLOSE, 0):
+                value_param = self._coerce_json_object(parameter_config.get("value"))
+                if value_param and value_param.get("type") == "variable":
+                    variable_selector = _VARIABLE_SELECTOR_ADAPTER.validate_python(value_param.get("value"))
+                    variable = variable_pool.get(variable_selector)
+                    if variable is None:
+                        raise AgentVariableNotFoundError(str(variable_selector))
+                    resolved_parameters[key] = variable.value
+                else:
+                    resolved_parameters[key] = value_param.get("value", "") if value_param is not None else None
+            else:
+                resolved_parameters[key] = None
+
+        return resolved_parameters
+
+    @staticmethod
+    def _resolve_tool_settings(tool: Mapping[str, object]) -> JsonObject:
+        settings = AgentRuntimeSupport._coerce_named_json_objects(tool.get("settings"))
+        if settings is None:
+            return {}
+        return {key: setting.get("value") for key, setting in settings.items()}
+
+    @staticmethod
+    def _coerce_json_object(value: object) -> JsonObject | None:
+        try:
+            return _JSON_OBJECT_ADAPTER.validate_python(value)
+        except ValidationError:
+            return None
+
+    @staticmethod
+    def _coerce_optional_string(value: object) -> str | None:
+        return value if isinstance(value, str) else None
+
+    @staticmethod
+    def _coerce_tool_provider_type(value: object) -> ToolProviderType:
+        if isinstance(value, ToolProviderType):
+            return value
+        if isinstance(value, str):
+            return ToolProviderType(value)
+        return ToolProviderType.BUILT_IN
+
+    @classmethod
+    def _coerce_named_json_objects(cls, value: object) -> dict[str, JsonObject] | None:
+        if not isinstance(value, dict):
+            return None
+
+        coerced: dict[str, JsonObject] = {}
+        for key, item in value.items():
+            if not isinstance(key, str):
+                return None
+            json_object = cls._coerce_json_object(item)
+            if json_object is None:
+                return None
+            coerced[key] = json_object
+        return coerced

api/core/workflow/workflow_entry.py

@@ -1,7 +1,7 @@
 import logging
 import time
 from collections.abc import Generator, Mapping, Sequence
-from typing import Any, cast
+from typing import Any, TypeAlias, cast
 
 from configs import dify_config
 from core.app.apps.exc import GenerateTaskStoppedError
@@ -32,6 +32,13 @@ from models.workflow import Workflow
 
 logger = logging.getLogger(__name__)
 
+SpecialValueScalar: TypeAlias = str | int | float | bool | None
+SpecialValue: TypeAlias = SpecialValueScalar | File | Mapping[str, "SpecialValue"] | list["SpecialValue"]
+SerializedSpecialValue: TypeAlias = (
+    SpecialValueScalar | dict[str, "SerializedSpecialValue"] | list["SerializedSpecialValue"]
+)
+SingleNodeGraphConfig: TypeAlias = dict[str, list[dict[str, object]]]
+
 
 class _WorkflowChildEngineBuilder:
     @staticmethod
@@ -276,10 +283,10 @@ class WorkflowEntry:
     @staticmethod
     def _create_single_node_graph(
         node_id: str,
-        node_data: dict[str, Any],
+        node_data: Mapping[str, object],
         node_width: int = 114,
         node_height: int = 514,
-    ) -> dict[str, Any]:
+    ) -> SingleNodeGraphConfig:
         """
         Create a minimal graph structure for testing a single node in isolation.
 
@@ -289,14 +296,14 @@ class WorkflowEntry:
         :param node_height: height for UI layout (default: 100)
         :return: graph dictionary with start node and target node
         """
-        node_config = {
+        node_config: dict[str, object] = {
             "id": node_id,
             "width": node_width,
             "height": node_height,
             "type": "custom",
-            "data": node_data,
+            "data": dict(node_data),
         }
-        start_node_config = {
+        start_node_config: dict[str, object] = {
             "id": "start",
             "width": node_width,
             "height": node_height,
@@ -321,7 +328,12 @@ class WorkflowEntry:
 
     @classmethod
     def run_free_node(
-        cls, node_data: dict[str, Any], node_id: str, tenant_id: str, user_id: str, user_inputs: dict[str, Any]
+        cls,
+        node_data: Mapping[str, object],
+        node_id: str,
+        tenant_id: str,
+        user_id: str,
+        user_inputs: Mapping[str, object],
     ) -> tuple[Node, Generator[GraphNodeEventBase, None, None]]:
         """
         Run free node
@@ -339,6 +351,8 @@ class WorkflowEntry:
         graph_dict = cls._create_single_node_graph(node_id, node_data)
 
         node_type = node_data.get("type", "")
+        if not isinstance(node_type, str):
+            raise ValueError("Node type must be a string")
         if node_type not in {BuiltinNodeTypes.PARAMETER_EXTRACTOR, BuiltinNodeTypes.QUESTION_CLASSIFIER}:
             raise ValueError(f"Node type {node_type} not supported")
 
@@ -369,7 +383,7 @@ class WorkflowEntry:
         graph_runtime_state = GraphRuntimeState(variable_pool=variable_pool, start_at=time.perf_counter())
 
         # init workflow run state
-        node_config = NodeConfigDictAdapter.validate_python({"id": node_id, "data": node_data})
+        node_config = NodeConfigDictAdapter.validate_python({"id": node_id, "data": dict(node_data)})
         node_factory = DifyNodeFactory(
             graph_init_params=graph_init_params,
             graph_runtime_state=graph_runtime_state,
@@ -405,30 +419,34 @@ class WorkflowEntry:
             raise WorkflowNodeRunFailedError(node=node, err_msg=str(e))
 
     @staticmethod
-    def handle_special_values(value: Mapping[str, Any] | None) -> Mapping[str, Any] | None:
+    def handle_special_values(value: Mapping[str, SpecialValue] | None) -> dict[str, SerializedSpecialValue] | None:
         # NOTE(QuantumGhost): Avoid using this function in new code.
         # Keep values structured as long as possible and only convert to dict
         # immediately before serialization (e.g., JSON serialization) to maintain
         # data integrity and type information.
         result = WorkflowEntry._handle_special_values(value)
-        return result if isinstance(result, Mapping) or result is None else dict(result)
+        if result is None:
+            return None
+        if isinstance(result, dict):
+            return result
+        raise TypeError("handle_special_values expects a mapping input")
 
     @staticmethod
-    def _handle_special_values(value: Any):
+    def _handle_special_values(value: SpecialValue) -> SerializedSpecialValue:
         if value is None:
             return value
-        if isinstance(value, dict):
-            res = {}
+        if isinstance(value, Mapping):
+            res: dict[str, SerializedSpecialValue] = {}
             for k, v in value.items():
                 res[k] = WorkflowEntry._handle_special_values(v)
             return res
         if isinstance(value, list):
-            res_list = []
+            res_list: list[SerializedSpecialValue] = []
             for item in value:
                 res_list.append(WorkflowEntry._handle_special_values(item))
             return res_list
         if isinstance(value, File):
-            return value.to_dict()
+            return dict(value.to_dict())
         return value
 
     @classmethod

api/dify_graph/file/file_manager.py

@@ -112,6 +112,8 @@ def _get_encoded_string(f: File, /) -> str:
             data = _download_file_content(f.storage_key)
         case FileTransferMethod.DATASOURCE_FILE:
             data = _download_file_content(f.storage_key)
+        case _:
+            raise ValueError(f"Unsupported transfer method: {f.transfer_method}")
 
     return base64.b64encode(data).decode("utf-8")
 

api/dify_graph/graph_engine/layers/execution_limits.py

@@ -133,6 +133,8 @@ class ExecutionLimitsLayer(GraphEngineLayer):
         elif limit_type == LimitType.TIME_LIMIT:
             elapsed_time = time.time() - self.start_time if self.start_time else 0
             reason = f"Maximum execution time exceeded: {elapsed_time:.2f}s > {self.max_time}s"
+        else:
+            return
 
         self.logger.warning("Execution limit exceeded: %s", reason)
 

api/dify_graph/nodes/base/node.py

@@ -336,12 +336,7 @@ class Node(Generic[NodeDataT]):
 
     def _restore_execution_id_from_runtime_state(self) -> str | None:
         graph_execution = self.graph_runtime_state.graph_execution
-        try:
-            node_executions = graph_execution.node_executions
-        except AttributeError:
-            return None
-        if not isinstance(node_executions, dict):
-            return None
+        node_executions = graph_execution.node_executions
         node_execution = node_executions.get(self._node_id)
         if node_execution is None:
             return None
@@ -395,8 +390,7 @@ class Node(Generic[NodeDataT]):
                 if isinstance(event, NodeEventBase):  # pyright: ignore[reportUnnecessaryIsInstance]
                     yield self._dispatch(event)
                 elif isinstance(event, GraphNodeEventBase) and not event.in_iteration_id and not event.in_loop_id:  # pyright: ignore[reportUnnecessaryIsInstance]
-                    event.id = self.execution_id
-                    yield event
+                    yield event.model_copy(update={"id": self.execution_id})
                 else:
                     yield event
         except Exception as e:

api/dify_graph/nodes/document_extractor/node.py

@@ -443,7 +443,10 @@ def _extract_text_from_docx(file_content: bytes) -> str:
         # Keep track of paragraph and table positions
         content_items: list[tuple[int, str, Table | Paragraph]] = []
 
-        it = iter(doc.element.body)
+        doc_body = getattr(doc.element, "body", None)
+        if doc_body is None:
+            raise TextExtractionError("DOCX body not found")
+        it = iter(doc_body)
         part = next(it, None)
         i = 0
         while part is not None:

api/dify_graph/nodes/http_request/node.py

@@ -101,6 +101,9 @@ class HttpRequestNode(Node[HttpRequestNodeData]):
                 timeout=self._get_request_timeout(self.node_data),
                 variable_pool=self.graph_runtime_state.variable_pool,
                 http_request_config=self._http_request_config,
+                # Must be 0 to disable executor-level retries, as the graph engine handles them.
+                # This is critical to prevent nested retries.
+                max_retries=0,
                 ssl_verify=self.node_data.ssl_verify,
                 http_client=self._http_client,
                 file_manager=self._file_manager,

api/dify_graph/nodes/llm/entities.py

@@ -1,7 +1,8 @@
 from collections.abc import Mapping, Sequence
-from typing import Any, Literal
+from typing import Literal, NotRequired
 
 from pydantic import BaseModel, Field, field_validator
+from typing_extensions import TypedDict
 
 from core.prompt.entities.advanced_prompt_entities import ChatModelMessage, CompletionModelPromptTemplate, MemoryConfig
 from dify_graph.entities.base_node_data import BaseNodeData
@@ -10,11 +11,17 @@ from dify_graph.model_runtime.entities import ImagePromptMessageContent, LLMMode
 from dify_graph.nodes.base.entities import VariableSelector
 
 
+class StructuredOutputConfig(TypedDict):
+    schema: Mapping[str, object]
+    name: NotRequired[str]
+    description: NotRequired[str]
+
+
 class ModelConfig(BaseModel):
     provider: str
     name: str
     mode: LLMMode
-    completion_params: dict[str, Any] = Field(default_factory=dict)
+    completion_params: dict[str, object] = Field(default_factory=dict)
 
 
 class ContextConfig(BaseModel):
@@ -33,7 +40,7 @@ class VisionConfig(BaseModel):
 
     @field_validator("configs", mode="before")
     @classmethod
-    def convert_none_configs(cls, v: Any):
+    def convert_none_configs(cls, v: object):
         if v is None:
             return VisionConfigOptions()
         return v
@@ -44,7 +51,7 @@ class PromptConfig(BaseModel):
 
     @field_validator("jinja2_variables", mode="before")
     @classmethod
-    def convert_none_jinja2_variables(cls, v: Any):
+    def convert_none_jinja2_variables(cls, v: object):
         if v is None:
             return []
         return v
@@ -67,7 +74,7 @@ class LLMNodeData(BaseNodeData):
     memory: MemoryConfig | None = None
     context: ContextConfig
     vision: VisionConfig = Field(default_factory=VisionConfig)
-    structured_output: Mapping[str, Any] | None = None
+    structured_output: StructuredOutputConfig | None = None
     # We used 'structured_output_enabled' in the past, but it's not a good name.
     structured_output_switch_on: bool = Field(False, alias="structured_output_enabled")
     reasoning_format: Literal["separated", "tagged"] = Field(
@@ -90,11 +97,30 @@ class LLMNodeData(BaseNodeData):
 
     @field_validator("prompt_config", mode="before")
     @classmethod
-    def convert_none_prompt_config(cls, v: Any):
+    def convert_none_prompt_config(cls, v: object):
         if v is None:
             return PromptConfig()
         return v
 
+    @field_validator("structured_output", mode="before")
+    @classmethod
+    def convert_legacy_structured_output(cls, v: object) -> StructuredOutputConfig | None | object:
+        if not isinstance(v, Mapping):
+            return v
+
+        schema = v.get("schema")
+        if schema is None:
+            return None
+
+        normalized: StructuredOutputConfig = {"schema": schema}
+        name = v.get("name")
+        description = v.get("description")
+        if isinstance(name, str):
+            normalized["name"] = name
+        if isinstance(description, str):
+            normalized["description"] = description
+        return normalized
+
     @property
     def structured_output_enabled(self) -> bool:
         return self.structured_output_switch_on and self.structured_output is not None

api/dify_graph/nodes/llm/llm_utils.py

@@ -1,6 +1,9 @@
 from __future__ import annotations
 
-from collections.abc import Sequence
+import json
+import logging
+import re
+from collections.abc import Mapping, Sequence
 from typing import Any, cast
 
 from core.model_manager import ModelInstance
@@ -36,6 +39,11 @@ from .exc import (
 )
 from .protocols import TemplateRenderer
 
+logger = logging.getLogger(__name__)
+
+VARIABLE_PATTERN = re.compile(r"\{\{#[^#]+#\}\}")
+MAX_RESOLVED_VALUE_LENGTH = 1024
+
 
 def fetch_model_schema(*, model_instance: ModelInstance) -> AIModelEntity:
     model_schema = cast(LargeLanguageModel, model_instance.model_type_instance).get_model_schema(
@@ -475,3 +483,61 @@ def _append_file_prompts(
         prompt_messages[-1] = UserPromptMessage(content=file_prompts + existing_contents)
     else:
         prompt_messages.append(UserPromptMessage(content=file_prompts))
+
+
+def _coerce_resolved_value(raw: str) -> int | float | bool | str:
+    """Try to restore the original type from a resolved template string.
+
+    Variable references are always resolved to text, but completion params may
+    expect numeric or boolean values (e.g. a variable that holds "0.7" mapped to
+    the ``temperature`` parameter).  This helper attempts a JSON parse so that
+    ``"0.7"`` → ``0.7``, ``"true"`` → ``True``, etc.  Plain strings that are not
+    valid JSON literals are returned as-is.
+    """
+    stripped = raw.strip()
+    if not stripped:
+        return raw
+
+    try:
+        parsed: object = json.loads(stripped)
+    except (json.JSONDecodeError, ValueError):
+        return raw
+
+    if isinstance(parsed, (int, float, bool)):
+        return parsed
+    return raw
+
+
+def resolve_completion_params_variables(
+    completion_params: Mapping[str, Any],
+    variable_pool: VariablePool,
+) -> dict[str, Any]:
+    """Resolve variable references (``{{#node_id.var#}}``) in string-typed completion params.
+
+    Security notes:
+    - Resolved values are length-capped to ``MAX_RESOLVED_VALUE_LENGTH`` to
+      prevent denial-of-service through excessively large variable payloads.
+    - This follows the same ``VariablePool.convert_template`` pattern used across
+      Dify (Answer Node, HTTP Request Node, Agent Node, etc.).  The downstream
+      model plugin receives these values as structured JSON key-value pairs — they
+      are never concatenated into raw HTTP headers or SQL queries.
+    - Numeric/boolean coercion is applied so that variables holding ``"0.7"`` are
+      restored to their native type rather than sent as a bare string.
+    """
+    resolved: dict[str, Any] = {}
+    for key, value in completion_params.items():
+        if isinstance(value, str) and VARIABLE_PATTERN.search(value):
+            segment_group = variable_pool.convert_template(value)
+            text = segment_group.text
+            if len(text) > MAX_RESOLVED_VALUE_LENGTH:
+                logger.warning(
+                    "Resolved value for param '%s' truncated from %d to %d chars",
+                    key,
+                    len(text),
+                    MAX_RESOLVED_VALUE_LENGTH,
+                )
+                text = text[:MAX_RESOLVED_VALUE_LENGTH]
+            resolved[key] = _coerce_resolved_value(text)
+        else:
+            resolved[key] = value
+    return resolved

api/dify_graph/nodes/llm/node.py

@@ -9,6 +9,7 @@ import time
 from collections.abc import Generator, Mapping, Sequence
 from typing import TYPE_CHECKING, Any, Literal
 
+from pydantic import TypeAdapter
 from sqlalchemy import select
 
 from core.llm_generator.output_parser.errors import OutputParserError
@@ -74,6 +75,7 @@ from .entities import (
     LLMNodeChatModelMessage,
     LLMNodeCompletionModelPromptTemplate,
     LLMNodeData,
+    StructuredOutputConfig,
 )
 from .exc import (
     InvalidContextStructureError,
@@ -88,6 +90,7 @@ if TYPE_CHECKING:
     from dify_graph.runtime import GraphRuntimeState
 
 logger = logging.getLogger(__name__)
+_JSON_OBJECT_ADAPTER = TypeAdapter(dict[str, object])
 
 
 class LLMNode(Node[LLMNodeData]):
@@ -202,6 +205,10 @@ class LLMNode(Node[LLMNodeData]):
 
             # fetch model config
             model_instance = self._model_instance
+            # Resolve variable references in string-typed completion params
+            model_instance.parameters = llm_utils.resolve_completion_params_variables(
+                model_instance.parameters, variable_pool
+            )
             model_name = model_instance.model_name
             model_provider = model_instance.provider
             model_stop = model_instance.stop
@@ -354,7 +361,7 @@ class LLMNode(Node[LLMNodeData]):
         stop: Sequence[str] | None = None,
         user_id: str,
         structured_output_enabled: bool,
-        structured_output: Mapping[str, Any] | None = None,
+        structured_output: StructuredOutputConfig | None = None,
         file_saver: LLMFileSaver,
         file_outputs: list[File],
         node_id: str,
@@ -367,8 +374,10 @@ class LLMNode(Node[LLMNodeData]):
         model_schema = llm_utils.fetch_model_schema(model_instance=model_instance)
 
         if structured_output_enabled:
+            if structured_output is None:
+                raise LLMNodeError("Please provide a valid structured output schema")
             output_schema = LLMNode.fetch_structured_output_schema(
-                structured_output=structured_output or {},
+                structured_output=structured_output,
             )
             request_start_time = time.perf_counter()
 
@@ -920,6 +929,12 @@ class LLMNode(Node[LLMNodeData]):
             # Extract clean text and reasoning from <think> tags
             clean_text, reasoning_content = LLMNode._split_reasoning(full_text, reasoning_format)
 
+        structured_output = (
+            dict(invoke_result.structured_output)
+            if isinstance(invoke_result, LLMResultWithStructuredOutput) and invoke_result.structured_output is not None
+            else None
+        )
+
         event = ModelInvokeCompletedEvent(
             # Use clean_text for separated mode, full_text for tagged mode
             text=clean_text if reasoning_format == "separated" else full_text,
@@ -928,7 +943,7 @@ class LLMNode(Node[LLMNodeData]):
             # Reasoning content for workflow variables and downstream nodes
             reasoning_content=reasoning_content,
             # Pass structured output if enabled
-            structured_output=getattr(invoke_result, "structured_output", None),
+            structured_output=structured_output,
         )
         if request_latency is not None:
             event.usage.latency = round(request_latency, 3)
@@ -962,27 +977,18 @@ class LLMNode(Node[LLMNodeData]):
     @staticmethod
     def fetch_structured_output_schema(
         *,
-        structured_output: Mapping[str, Any],
-    ) -> dict[str, Any]:
+        structured_output: StructuredOutputConfig,
+    ) -> dict[str, object]:
         """
         Fetch the structured output schema from the node data.
 
         Returns:
-            dict[str, Any]: The structured output schema
+            dict[str, object]: The structured output schema
         """
-        if not structured_output:
+        schema = structured_output.get("schema")
+        if not schema:
             raise LLMNodeError("Please provide a valid structured output schema")
-        structured_output_schema = json.dumps(structured_output.get("schema", {}), ensure_ascii=False)
-        if not structured_output_schema:
-            raise LLMNodeError("Please provide a valid structured output schema")
-
-        try:
-            schema = json.loads(structured_output_schema)
-            if not isinstance(schema, dict):
-                raise LLMNodeError("structured_output_schema must be a JSON object")
-            return schema
-        except json.JSONDecodeError:
-            raise LLMNodeError("structured_output_schema is not valid JSON format")
+        return _JSON_OBJECT_ADAPTER.validate_python(schema)
 
     @staticmethod
     def _save_multimodal_output_and_convert_result_to_markdown(