From 5a446f8200ef19b055ea928bf164c505a04edbb3 Mon Sep 17 00:00:00 2001 From: Jake Armstrong <65635253+jakearmstrong59@users.noreply.github.com> Date: Mon, 13 Apr 2026 07:33:20 +0200 Subject: [PATCH] refactor(api): deduplicate dataset controller schemas into controller_schemas.py (#34756) --- api/controllers/common/controller_schemas.py | 25 +++++++++++++++++++ .../console/datasets/datasets_document.py | 11 +------- .../console/datasets/datasets_segments.py | 9 +------ api/controllers/console/datasets/metadata.py | 7 +----- .../service_api/dataset/document.py | 10 +------- .../service_api/dataset/metadata.py | 7 +----- .../service_api/dataset/segment.py | 9 +------ 7 files changed, 31 insertions(+), 47 deletions(-) diff --git a/api/controllers/common/controller_schemas.py b/api/controllers/common/controller_schemas.py index ec5c72374d..c12d576473 100644 --- a/api/controllers/common/controller_schemas.py +++ b/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 --- diff --git a/api/controllers/console/datasets/datasets_document.py b/api/controllers/console/datasets/datasets_document.py index b7584f1f00..de8fe1c0e2 100644 --- a/api/controllers/console/datasets/datasets_document.py +++ b/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.") diff --git a/api/controllers/console/datasets/datasets_segments.py b/api/controllers/console/datasets/datasets_segments.py index c5f4e3a6e2..354c299bef 100644 --- a/api/controllers/console/datasets/datasets_segments.py +++ b/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] diff --git a/api/controllers/console/datasets/metadata.py b/api/controllers/console/datasets/metadata.py index 2e69ddc5ab..d966e1629e 100644 --- a/api/controllers/console/datasets/metadata.py +++ b/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 ) diff --git a/api/controllers/service_api/dataset/document.py b/api/controllers/service_api/dataset/document.py index db34aa408e..6db047567f 100644 --- a/api/controllers/service_api/dataset/document.py +++ b/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( diff --git a/api/controllers/service_api/dataset/metadata.py b/api/controllers/service_api/dataset/metadata.py index 52166f7fcc..21db7d0cb8 100644 --- a/api/controllers/service_api/dataset/metadata.py +++ b/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, diff --git a/api/controllers/service_api/dataset/segment.py b/api/controllers/service_api/dataset/segment.py index 5b16da81e0..9ad999b93e 100644 --- a/api/controllers/service_api/dataset/segment.py +++ b/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,