refactor(api): deduplicate dataset controller schemas into controller_schemas.py (#34756)

api/controllers/common/controller_schemas.py

@@ -1,4 +1,5 @@
 from typing import Any, Literal
+from uuid import UUID
 
 from pydantic import BaseModel, Field, model_validator
 
@@ -69,6 +70,30 @@ class WorkflowUpdatePayload(BaseModel):
     marked_comment: str | None = Field(default=None, max_length=100)
 
 
+# --- Dataset schemas ---
+
+
+DOCUMENT_BATCH_DOWNLOAD_ZIP_MAX_DOCS = 100
+
+
+class ChildChunkCreatePayload(BaseModel):
+    content: str
+
+
+class ChildChunkUpdatePayload(BaseModel):
+    content: str
+
+
+class DocumentBatchDownloadZipPayload(BaseModel):
+    """Request payload for bulk downloading documents as a zip archive."""
+
+    document_ids: list[UUID] = Field(..., min_length=1, max_length=DOCUMENT_BATCH_DOWNLOAD_ZIP_MAX_DOCS)
+
+
+class MetadataUpdatePayload(BaseModel):
+    name: str
+
+
 # --- Audio schemas ---
 
 

api/controllers/console/datasets/datasets_document.py

@@ -4,7 +4,6 @@ from argparse import ArgumentTypeError
 from collections.abc import Sequence
 from contextlib import ExitStack
 from typing import Any, Literal, cast
-from uuid import UUID
 
 import sqlalchemy as sa
 from flask import request, send_file
@@ -16,6 +15,7 @@ from sqlalchemy import asc, desc, func, select
 from werkzeug.exceptions import Forbidden, NotFound
 
 import services
+from controllers.common.controller_schemas import DocumentBatchDownloadZipPayload
 from controllers.common.schema import get_or_create_model, register_schema_models
 from controllers.console import console_ns
 from core.errors.error import (
@@ -71,9 +71,6 @@ from ..wraps import (
 
 logger = logging.getLogger(__name__)
 
-# NOTE: Keep constants near the top of the module for discoverability.
-DOCUMENT_BATCH_DOWNLOAD_ZIP_MAX_DOCS = 100
-
 
 # Register models for flask_restx to avoid dict type issues in Swagger
 dataset_model = get_or_create_model("Dataset", dataset_fields)
@@ -110,12 +107,6 @@ class GenerateSummaryPayload(BaseModel):
     document_list: list[str]
 
 
-class DocumentBatchDownloadZipPayload(BaseModel):
-    """Request payload for bulk downloading documents as a zip archive."""
-
-    document_ids: list[UUID] = Field(..., min_length=1, max_length=DOCUMENT_BATCH_DOWNLOAD_ZIP_MAX_DOCS)
-
-
 class DocumentDatasetListParam(BaseModel):
     page: int = Field(1, title="Page", description="Page number.")
     limit: int = Field(20, title="Limit", description="Page size.")

api/controllers/console/datasets/datasets_segments.py

@@ -10,6 +10,7 @@ from werkzeug.exceptions import Forbidden, NotFound
 
 import services
 from configs import dify_config
+from controllers.common.controller_schemas import ChildChunkCreatePayload, ChildChunkUpdatePayload
 from controllers.common.schema import register_schema_models
 from controllers.console import console_ns
 from controllers.console.app.error import ProviderNotInitializeError
@@ -82,14 +83,6 @@ class BatchImportPayload(BaseModel):
     upload_file_id: str
 
 
-class ChildChunkCreatePayload(BaseModel):
-    content: str
-
-
-class ChildChunkUpdatePayload(BaseModel):
-    content: str
-
-
 class ChildChunkBatchUpdatePayload(BaseModel):
     chunks: list[ChildChunkUpdateArgs]
 

api/controllers/console/datasets/metadata.py

@@ -1,9 +1,9 @@
 from typing import Literal
 
 from flask_restx import Resource, marshal_with
-from pydantic import BaseModel
 from werkzeug.exceptions import NotFound
 
+from controllers.common.controller_schemas import MetadataUpdatePayload
 from controllers.common.schema import register_schema_models
 from controllers.console import console_ns
 from controllers.console.wraps import account_initialization_required, enterprise_license_required, setup_required
@@ -18,11 +18,6 @@ from services.entities.knowledge_entities.knowledge_entities import (
 )
 from services.metadata_service import MetadataService
 
-
-class MetadataUpdatePayload(BaseModel):
-    name: str
-
-
 register_schema_models(
     console_ns, MetadataArgs, MetadataOperationData, MetadataUpdatePayload, DocumentMetadataOperation, MetadataDetail
 )

api/controllers/service_api/dataset/document.py

@@ -10,6 +10,7 @@ from sqlalchemy import desc, func, select
 from werkzeug.exceptions import Forbidden, NotFound
 
 import services
+from controllers.common.controller_schemas import DocumentBatchDownloadZipPayload
 from controllers.common.errors import (
     FilenameNotExistsError,
     FileTooLargeError,
@@ -100,15 +101,6 @@ class DocumentListQuery(BaseModel):
     status: str | None = Field(default=None, description="Document status filter")
 
 
-DOCUMENT_BATCH_DOWNLOAD_ZIP_MAX_DOCS = 100
-
-
-class DocumentBatchDownloadZipPayload(BaseModel):
-    """Request payload for bulk downloading uploaded documents as a ZIP archive."""
-
-    document_ids: list[UUID] = Field(..., min_length=1, max_length=DOCUMENT_BATCH_DOWNLOAD_ZIP_MAX_DOCS)
-
-
 register_enum_models(service_api_ns, RetrievalMethod)
 
 register_schema_models(

api/controllers/service_api/dataset/metadata.py

@@ -2,9 +2,9 @@ from typing import Literal
 
 from flask_login import current_user
 from flask_restx import marshal
-from pydantic import BaseModel
 from werkzeug.exceptions import NotFound
 
+from controllers.common.controller_schemas import MetadataUpdatePayload
 from controllers.common.schema import register_schema_model, register_schema_models
 from controllers.service_api import service_api_ns
 from controllers.service_api.wraps import DatasetApiResource, cloud_edition_billing_rate_limit_check
@@ -18,11 +18,6 @@ from services.entities.knowledge_entities.knowledge_entities import (
 )
 from services.metadata_service import MetadataService
 
-
-class MetadataUpdatePayload(BaseModel):
-    name: str
-
-
 register_schema_model(service_api_ns, MetadataUpdatePayload)
 register_schema_models(
     service_api_ns,

api/controllers/service_api/dataset/segment.py

@@ -8,6 +8,7 @@ from sqlalchemy import select
 from werkzeug.exceptions import NotFound
 
 from configs import dify_config
+from controllers.common.controller_schemas import ChildChunkCreatePayload, ChildChunkUpdatePayload
 from controllers.common.schema import register_schema_models
 from controllers.service_api import service_api_ns
 from controllers.service_api.app.error import ProviderNotInitializeError
@@ -69,20 +70,12 @@ class SegmentUpdatePayload(BaseModel):
     segment: SegmentUpdateArgs
 
 
-class ChildChunkCreatePayload(BaseModel):
-    content: str
-
-
 class ChildChunkListQuery(BaseModel):
     limit: int = Field(default=20, ge=1)
     keyword: str | None = None
     page: int = Field(default=1, ge=1)
 
 
-class ChildChunkUpdatePayload(BaseModel):
-    content: str
-
-
 register_schema_models(
     service_api_ns,
     SegmentCreatePayload,