refactor(api): migrate dataset rag pipeline endpoints to BaseModel (#37958)

Co-authored-by: autofix-ci[bot] <114827586+autofix-ci[bot]@users.noreply.github.com> Co-authored-by: Asuka Minato <i@asukaminato.eu.org>
2026-09-24 23:22:26 +08:00 · 2026-07-06 07:00:04 +00:00
parent c3b1508712
commit 77ae583b44
18 changed files with 787 additions and 245 deletions
@@ -6,7 +6,7 @@ from pydantic import BaseModel, Field
 from werkzeug.exceptions import Forbidden, NotFound

 from configs import dify_config
-from controllers.common.fields import RedirectResponse, SimpleResultResponse
+from controllers.common.fields import SimpleResultResponse
 from controllers.common.schema import query_params_from_model, register_response_schema_models, register_schema_models
 from controllers.console import console_ns
 from controllers.console.wraps import (
@@ -19,11 +19,13 @@ from controllers.console.wraps import (
    with_current_tenant_id,
    with_current_user,
 )
+from core.entities.provider_entities import ProviderConfig
 from core.plugin.entities.plugin_daemon import PluginOAuthAuthorizationUrlResponse
 from core.plugin.impl.oauth import OAuthHandler
+from core.tools.entities.common_entities import I18nObject
 from fields.base import ResponseModel
 from graphon.model_runtime.errors.validate import CredentialsValidateFailedError
-from graphon.model_runtime.utils.encoders import jsonable_encoder
+from libs.helper import dump_response
 from libs.login import login_required
 from models import Account
 from models.provider_ids import DatasourceProviderID
@@ -33,7 +35,9 @@ from services.plugin.oauth_service import OAuthProxyService

 class DatasourceCredentialPayload(BaseModel):
    name: str | None = Field(default=None, max_length=100)
-    credentials: dict[str, Any]
+    credentials: dict[str, Any] = Field(
+        description="Plugin-defined credential parameters. The schema is declared by the datasource provider."
+    )


 class DatasourceCredentialDeletePayload(BaseModel):
@@ -43,11 +47,17 @@ class DatasourceCredentialDeletePayload(BaseModel):
 class DatasourceCredentialUpdatePayload(BaseModel):
    credential_id: str
    name: str | None = Field(default=None, max_length=100)
-    credentials: dict[str, Any] | None = Field(default=None)
+    credentials: dict[str, Any] | None = Field(
+        default=None,
+        description="Plugin-defined credential parameters. The schema is declared by the datasource provider.",
+    )


 class DatasourceCustomClientPayload(BaseModel):
-    client_params: dict[str, Any] | None = Field(default=None)
+    client_params: dict[str, Any] | None = Field(
+        default=None,
+        description="Plugin-defined OAuth client parameters. The schema is declared by the datasource provider.",
+    )
    enable_oauth_custom_client: bool | None = None


@@ -71,8 +81,48 @@ class DatasourceOAuthCallbackQuery(BaseModel):
    context_id: str | None = Field(default=None, description="OAuth proxy context ID")


-class DatasourceCredentialsResponse(ResponseModel):
-    result: Any
+class DatasourceCredentialResponse(ResponseModel):
+    credential: dict[str, Any] = Field(
+        description="Obfuscated plugin-defined credential parameters from the datasource provider."
+    )
+    type: str
+    name: str
+    avatar_url: str | None
+    id: str
+    is_default: bool
+
+
+class DatasourceCredentialListResponse(ResponseModel):
+    result: list[DatasourceCredentialResponse]
+
+
+class DatasourceOAuthSchemaResponse(ResponseModel):
+    client_schema: list[ProviderConfig]
+    credentials_schema: list[ProviderConfig]
+    oauth_custom_client_params: dict[str, Any] | None = Field(
+        description="Masked plugin-defined OAuth client parameters, when configured for the tenant."
+    )
+    is_oauth_custom_client_enabled: bool
+    is_system_oauth_params_exists: bool
+    redirect_uri: str
+
+
+class DatasourceProviderAuthResponse(ResponseModel):
+    author: str
+    provider: str
+    plugin_id: str
+    plugin_unique_identifier: str
+    icon: str
+    name: str
+    label: I18nObject
+    description: I18nObject
+    credential_schema: list[ProviderConfig]
+    oauth_schema: DatasourceOAuthSchemaResponse | None
+    credentials_list: list[DatasourceCredentialResponse]
+
+
+class DatasourceProviderAuthListResponse(ResponseModel):
+    result: list[DatasourceProviderAuthResponse]


 register_schema_models(
@@ -88,9 +138,9 @@ register_schema_models(
 )
 register_response_schema_models(
    console_ns,
-    DatasourceCredentialsResponse,
+    DatasourceCredentialListResponse,
+    DatasourceProviderAuthListResponse,
    PluginOAuthAuthorizationUrlResponse,
-    RedirectResponse,
    SimpleResultResponse,
 )

@@ -100,7 +150,7 @@ class DatasourcePluginOAuthAuthorizationUrl(Resource):
    @console_ns.doc(params=query_params_from_model(DatasourceOAuthAuthorizationQuery))
    @console_ns.response(
        200,
-        "Authorization URL retrieved successfully",
+        "Datasource OAuth authorization URL generated successfully",
        console_ns.models[PluginOAuthAuthorizationUrlResponse.__name__],
    )
    @setup_required
@@ -140,7 +190,8 @@ class DatasourcePluginOAuthAuthorizationUrl(Resource):
            redirect_uri=redirect_uri,
            system_credentials=oauth_config,
        )
-        response = make_response(jsonable_encoder(authorization_url_response))
+        # response-contract:ignore cookie-bearing Flask response
+        response = make_response(dump_response(PluginOAuthAuthorizationUrlResponse, authorization_url_response))
        response.set_cookie(
            "context_id",
            context_id,
@@ -154,11 +205,8 @@ class DatasourcePluginOAuthAuthorizationUrl(Resource):
@console_ns.route("/oauth/plugin/<path:provider_id>/datasource/callback")
 class DatasourceOAuthCallback(Resource):
    @console_ns.doc(params=query_params_from_model(DatasourceOAuthCallbackQuery))
-    @console_ns.response(
-        302,
-        "Redirect to console OAuth callback page",
-        console_ns.models[RedirectResponse.__name__],
-    )
+    # response-contract:ignore redirect response
+    @console_ns.response(302, "Redirect to OAuth callback page")
    @setup_required
    def get(self, provider_id: str):
        context_id = request.cookies.get("context_id") or request.args.get("context_id")
@@ -217,7 +265,9 @@ class DatasourceOAuthCallback(Resource):
@console_ns.route("/auth/plugin/datasource/<path:provider_id>")
 class DatasourceAuth(Resource):
    @console_ns.expect(console_ns.models[DatasourceCredentialPayload.__name__])
-    @console_ns.response(200, "Success", console_ns.models[SimpleResultResponse.__name__])
+    @console_ns.response(
+        200, "Datasource credential created successfully", console_ns.models[SimpleResultResponse.__name__]
+    )
    @setup_required
    @login_required
    @account_initialization_required
@@ -238,12 +288,16 @@ class DatasourceAuth(Resource):
            )
        except CredentialsValidateFailedError as ex:
            raise ValueError(str(ex))
-        return {"result": "success"}, 200
+        return SimpleResultResponse(result="success").model_dump(mode="json"), 200

+    @console_ns.response(
+        200,
+        "Datasource credentials retrieved successfully",
+        console_ns.models[DatasourceCredentialListResponse.__name__],
+    )
    @setup_required
    @login_required
    @account_initialization_required
-    @console_ns.response(200, "Success", console_ns.models[DatasourceCredentialsResponse.__name__])
    @with_current_user
    @with_current_tenant_id
    def get(self, current_tenant_id: str, user: Account, provider_id: str):
@@ -256,7 +310,7 @@ class DatasourceAuth(Resource):
            plugin_id=datasource_provider_id.plugin_id,
            user=user,
        )
-        return {"result": datasources}, 200
+        return dump_response(DatasourceCredentialListResponse, {"result": datasources}), 200


@console_ns.route("/auth/plugin/datasource/<path:provider_id>/delete")
@@ -282,13 +336,15 @@ class DatasourceAuthDeleteApi(Resource):
            provider=provider_name,
            plugin_id=plugin_id,
        )
-        return {"result": "success"}, 200
+        return SimpleResultResponse(result="success").model_dump(mode="json"), 200


@console_ns.route("/auth/plugin/datasource/<path:provider_id>/update")
 class DatasourceAuthUpdateApi(Resource):
    @console_ns.expect(console_ns.models[DatasourceCredentialUpdatePayload.__name__])
-    @console_ns.response(201, "Success", console_ns.models[SimpleResultResponse.__name__])
+    @console_ns.response(
+        201, "Datasource credential updated successfully", console_ns.models[SimpleResultResponse.__name__]
+    )
    @setup_required
    @login_required
    @account_initialization_required
@@ -308,12 +364,16 @@ class DatasourceAuthUpdateApi(Resource):
            credentials=payload.credentials or {},
            name=payload.name,
        )
-        return {"result": "success"}, 201
+        return SimpleResultResponse(result="success").model_dump(mode="json"), 201


@console_ns.route("/auth/plugin/datasource/list")
 class DatasourceAuthListApi(Resource):
-    @console_ns.response(200, "Success", console_ns.models[DatasourceCredentialsResponse.__name__])
+    @console_ns.response(
+        200,
+        "Datasource credentials retrieved successfully",
+        console_ns.models[DatasourceProviderAuthListResponse.__name__],
+    )
    @setup_required
    @login_required
    @account_initialization_required
@@ -321,12 +381,16 @@ class DatasourceAuthListApi(Resource):
    def get(self, current_tenant_id: str):
        datasource_provider_service = DatasourceProviderService()
        datasources = datasource_provider_service.get_all_datasource_credentials(tenant_id=current_tenant_id)
-        return {"result": jsonable_encoder(datasources)}, 200
+        return dump_response(DatasourceProviderAuthListResponse, {"result": datasources}), 200


@console_ns.route("/auth/plugin/datasource/default-list")
 class DatasourceHardCodeAuthListApi(Resource):
-    @console_ns.response(200, "Success", console_ns.models[DatasourceCredentialsResponse.__name__])
+    @console_ns.response(
+        200,
+        "Default datasource credentials retrieved successfully",
+        console_ns.models[DatasourceProviderAuthListResponse.__name__],
+    )
    @setup_required
    @login_required
    @account_initialization_required
@@ -334,13 +398,15 @@ class DatasourceHardCodeAuthListApi(Resource):
    def get(self, current_tenant_id: str):
        datasource_provider_service = DatasourceProviderService()
        datasources = datasource_provider_service.get_hard_code_datasource_credentials(tenant_id=current_tenant_id)
-        return {"result": jsonable_encoder(datasources)}, 200
+        return dump_response(DatasourceProviderAuthListResponse, {"result": datasources}), 200


@console_ns.route("/auth/plugin/datasource/<path:provider_id>/custom-client")
 class DatasourceAuthOauthCustomClient(Resource):
    @console_ns.expect(console_ns.models[DatasourceCustomClientPayload.__name__])
-    @console_ns.response(200, "Success", console_ns.models[SimpleResultResponse.__name__])
+    @console_ns.response(
+        200, "Datasource OAuth custom client saved successfully", console_ns.models[SimpleResultResponse.__name__]
+    )
    @setup_required
    @login_required
    @account_initialization_required
@@ -357,7 +423,7 @@ class DatasourceAuthOauthCustomClient(Resource):
            client_params=payload.client_params or {},
            enabled=payload.enable_oauth_custom_client or False,
        )
-        return {"result": "success"}, 200
+        return SimpleResultResponse(result="success").model_dump(mode="json"), 200

    @setup_required
    @login_required
@@ -371,7 +437,7 @@ class DatasourceAuthOauthCustomClient(Resource):
            tenant_id=current_tenant_id,
            datasource_provider_id=datasource_provider_id,
        )
-        return {"result": "success"}, 200
+        return SimpleResultResponse(result="success").model_dump(mode="json"), 200


@console_ns.route("/auth/plugin/datasource/<path:provider_id>/default")
@@ -393,7 +459,7 @@ class DatasourceAuthDefaultApi(Resource):
            datasource_provider_id=datasource_provider_id,
            credential_id=payload.id,
        )
-        return {"result": "success"}, 200
+        return SimpleResultResponse(result="success").model_dump(mode="json"), 200


@console_ns.route("/auth/plugin/datasource/<path:provider_id>/update-name")
@@ -416,4 +482,4 @@ class DatasourceUpdateProviderNameApi(Resource):
            name=payload.name,
            credential_id=payload.credential_id,
        )
-        return {"result": "success"}, 200
+        return SimpleResultResponse(result="success").model_dump(mode="json"), 200
@@ -3,9 +3,9 @@ from typing import Any
 from flask_restx import (  # type: ignore
    Resource,  # type: ignore
 )
-from pydantic import BaseModel, RootModel
+from pydantic import BaseModel

-from controllers.common.schema import register_response_schema_models, register_schema_models
+from controllers.common.schema import register_schema_models
 from controllers.console import console_ns
 from controllers.console.datasets.wraps import get_rag_pipeline
 from controllers.console.wraps import account_initialization_required, setup_required, with_current_user
@@ -21,18 +21,13 @@ class Parser(BaseModel):
    credential_id: str | None = None


-class DataSourceContentPreviewResponse(RootModel[Any]):
-    root: Any
-
-
 register_schema_models(console_ns, Parser)
-register_response_schema_models(console_ns, DataSourceContentPreviewResponse)


@console_ns.route("/rag/pipelines/<uuid:pipeline_id>/workflows/published/datasource/nodes/<string:node_id>/preview")
 class DataSourceContentPreviewApi(Resource):
    @console_ns.expect(console_ns.models[Parser.__name__])
-    @console_ns.response(200, "Success", console_ns.models[DataSourceContentPreviewResponse.__name__])
+    @console_ns.response(200, "Success")
    @setup_required
    @login_required
    @account_initialization_required
@@ -1,9 +1,10 @@
 from collections.abc import Generator
+from datetime import datetime
 from typing import Any
 from uuid import UUID

 from flask import request
-from pydantic import BaseModel, Field, RootModel
+from pydantic import BaseModel, Field, RootModel, field_validator
 from sqlalchemy import select
 from sqlalchemy.orm import Session
 from werkzeug.exceptions import Forbidden, NotFound
@@ -13,24 +14,20 @@ from controllers.common.errors import FilenameNotExistsError, NoFileUploadedErro
 from controllers.common.fields import GeneratedAppResponse
 from controllers.common.schema import (
    query_params_from_model,
+    query_params_from_request,
    register_response_schema_models,
    register_schema_model,
-    register_schema_models,
 )
 from controllers.console.app.wraps import with_session
 from controllers.service_api import service_api_ns
 from controllers.service_api.dataset.error import PipelineRunError
-from controllers.service_api.dataset.rag_pipeline.serializers import serialize_upload_file
-from controllers.service_api.schema import (
-    event_stream_response,
-    json_or_event_stream_response,
-    multipart_file_params,
-)
+from controllers.service_api.schema import event_stream_response, json_or_event_stream_response, multipart_file_params
 from controllers.service_api.wraps import DatasetApiResource
 from core.app.apps.pipeline.pipeline_generator import PipelineGenerator
 from core.app.entities.app_invoke_entities import InvokeFrom
 from fields.base import ResponseModel
 from libs import helper
+from libs.helper import dump_response
 from libs.login import current_user
 from models import Account
 from models.dataset import Dataset, Pipeline
@@ -84,7 +81,7 @@ class DatasourcePluginResponse(ResponseModel):
    datasource_type: str | None = None
    title: str | None = None
    user_input_variables: list[dict[str, Any]] = Field(default_factory=list)
-    credentials: list[DatasourceCredentialInfoResponse]
+    credentials: list[DatasourceCredentialInfoResponse] = Field(default_factory=list)


 class DatasourcePluginListResponse(RootModel[list[DatasourcePluginResponse]]):
@@ -100,14 +97,22 @@ class PipelineUploadFileResponse(ResponseModel):
    created_by: str
    created_at: str | None = None

+    @field_validator("created_at", mode="before")
+    @classmethod
+    def _normalize_created_at(cls, value: datetime | str | None) -> str | None:
+        if isinstance(value, datetime):
+            return value.isoformat()
+        return value
+

 register_schema_model(service_api_ns, DatasourceNodeRunPayload)
+register_schema_model(service_api_ns, DatasourcePluginsQuery)
 register_schema_model(service_api_ns, PipelineRunApiEntity)
-register_schema_models(service_api_ns, DatasourcePluginsQuery)
 register_response_schema_models(
    service_api_ns,
+    DatasourceCredentialInfoResponse,
+    DatasourcePluginResponse,
    DatasourcePluginListResponse,
-    GeneratedAppResponse,
    PipelineUploadFileResponse,
 )

@@ -119,8 +124,8 @@ class DatasourcePluginsApi(DatasetApiResource):
    @service_api_ns.doc(
        summary="List Datasource Plugins",
        description=(
-            "List the datasource nodes configured in the knowledge pipeline. Each node includes the "
-            "plugin it uses plus the metadata needed to run it."
+            "List the datasource nodes configured in the knowledge pipeline. Each node includes the plugin it uses "
+            "plus the metadata needed to run it."
        ),
        tags=["Knowledge Pipeline"],
        responses={
@@ -152,14 +157,13 @@ class DatasourcePluginsApi(DatasetApiResource):
        if not dataset:
            raise NotFound("Dataset not found.")

-        # Get query parameter to determine published or draft
-        is_published: bool = request.args.get("is_published", default=True, type=bool)
+        query = query_params_from_request(DatasourcePluginsQuery)

        rag_pipeline_service: RagPipelineService = RagPipelineService()
        datasource_plugins: list[dict[Any, Any]] = rag_pipeline_service.get_datasource_plugins(
-            tenant_id=tenant_id, dataset_id=dataset_id_str, is_published=is_published
+            tenant_id=tenant_id, dataset_id=dataset_id_str, is_published=query.is_published
        )
-        return datasource_plugins, 200
+        return dump_response(DatasourcePluginListResponse, datasource_plugins), 200


@service_api_ns.route("/datasets/<uuid:dataset_id>/pipeline/datasource/nodes/<string:node_id>/run")
@@ -169,8 +173,8 @@ class DatasourceNodeRunApi(DatasetApiResource):
    @service_api_ns.doc(
        summary="Run Datasource Node",
        description=(
-            "Execute a single datasource node within the knowledge pipeline. Returns a streaming "
-            "response with the node execution results."
+            "Execute a single datasource node within the knowledge pipeline. Returns a streaming response with the "
+            "node execution results."
        ),
        tags=["Knowledge Pipeline"],
        responses={
@@ -189,11 +193,6 @@ class DatasourceNodeRunApi(DatasetApiResource):
        }
    )
    @service_api_ns.expect(service_api_ns.models[DatasourceNodeRunPayload.__name__])
-    @service_api_ns.response(
-        200,
-        "Datasource node run successfully",
-        service_api_ns.models[GeneratedAppResponse.__name__],
-    )
    def post(self, tenant_id: str, dataset_id: UUID, node_id: str):
        """Resource for getting datasource plugins."""
        dataset_id_str = str(dataset_id)
@@ -210,10 +209,11 @@ class DatasourceNodeRunApi(DatasetApiResource):
        datasource_node_run_api_entity = DatasourceNodeRunApiEntity.model_validate(
            {
                **payload.model_dump(exclude_none=True),
-                "pipeline_id": str(pipeline.id),
+                "pipeline_id": pipeline.id,
                "node_id": node_id,
            }
        )
+        # response-contract:ignore compact_generate_response
        return helper.compact_generate_response(
            PipelineGenerator.convert_to_event_stream(
                rag_pipeline_service.run_datasource_workflow_node(
@@ -236,8 +236,8 @@ class PipelineRunApi(DatasetApiResource):
    @service_api_ns.doc(
        summary="Run Pipeline",
        description=(
-            "Execute the full knowledge pipeline for a knowledge base. Supports both streaming and "
-            "blocking response modes."
+            "Execute the full knowledge pipeline for a knowledge base. Supports both streaming and blocking response "
+            "modes."
        ),
        tags=["Knowledge Pipeline"],
        responses={
@@ -293,6 +293,7 @@ class PipelineRunApi(DatasetApiResource):
                streaming=payload.response_mode == "streaming",
            )

+            # response-contract:ignore compact_generate_response
            return helper.compact_generate_response(response)
        except Exception as ex:
            raise PipelineRunError(description=str(ex))
@@ -368,4 +369,4 @@ class KnowledgebasePipelineFileUploadApi(DatasetApiResource):
        except services.errors.file.UnsupportedFileTypeError:
            raise UnsupportedFileTypeError()

-        return serialize_upload_file(upload_file), 201
+        return dump_response(PipelineUploadFileResponse, upload_file), 201
@@ -1,32 +0,0 @@
-"""
-Serialization helpers for Service API knowledge pipeline endpoints.
-"""
-
-from __future__ import annotations
-
-from typing import TYPE_CHECKING, TypedDict
-
-if TYPE_CHECKING:
-    from models.model import UploadFile
-
-
-class UploadFileDict(TypedDict):
-    id: str
-    name: str
-    size: int
-    extension: str
-    mime_type: str | None
-    created_by: str
-    created_at: str | None
-
-
-def serialize_upload_file(upload_file: UploadFile) -> UploadFileDict:
-    return {
-        "id": upload_file.id,
-        "name": upload_file.name,
-        "size": upload_file.size,
-        "extension": upload_file.extension,
-        "mime_type": upload_file.mime_type,
-        "created_by": upload_file.created_by,
-        "created_at": upload_file.created_at.isoformat() if upload_file.created_at else None,
-    }