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>

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

@@ -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

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

@@ -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