refactor: port api/fields/file_fields.py

api/controllers/console/files.py

@@ -1,7 +1,7 @@
 from typing import Literal
 
 from flask import request
-from flask_restx import Resource, marshal_with
+from flask_restx import Resource
 from werkzeug.exceptions import Forbidden
 
 import services
@@ -15,18 +15,21 @@
  TooManyFilesError,
  UnsupportedFileTypeError,
 )
+from controllers.common.schema import register_schema_models
 from controllers.console.wraps import (
  account_initialization_required,
  cloud_edition_billing_resource_check,
  setup_required,
 )
 from extensions.ext_database import db
-from fields.file_fields import file_fields, upload_config_fields
+from fields.file_fields import FileResponse, UploadConfig
 from libs.login import current_account_with_tenant, login_required
 from services.file_service import FileService
 
 from . import console_ns
 
+register_schema_models(console_ns, UploadConfig, FileResponse)
+
 PREVIEW_WORDS_LIMIT = 3000
 
 
@@ -35,26 +38,27 @@ class FileApi(Resource):
  @setup_required
  @login_required
  @account_initialization_required
- @marshal_with(upload_config_fields)
+ @console_ns.response(200, "Success", console_ns.models[UploadConfig.__name__])
  def get(self):
- return {
- "file_size_limit": dify_config.UPLOAD_FILE_SIZE_LIMIT,
- "batch_count_limit": dify_config.UPLOAD_FILE_BATCH_LIMIT,
- "file_upload_limit": dify_config.BATCH_UPLOAD_LIMIT,
- "image_file_size_limit": dify_config.UPLOAD_IMAGE_FILE_SIZE_LIMIT,
- "video_file_size_limit": dify_config.UPLOAD_VIDEO_FILE_SIZE_LIMIT,
- "audio_file_size_limit": dify_config.UPLOAD_AUDIO_FILE_SIZE_LIMIT,
- "workflow_file_upload_limit": dify_config.WORKFLOW_FILE_UPLOAD_LIMIT,
- "image_file_batch_limit": dify_config.IMAGE_FILE_BATCH_LIMIT,
- "single_chunk_attachment_limit": dify_config.SINGLE_CHUNK_ATTACHMENT_LIMIT,
- "attachment_image_file_size_limit": dify_config.ATTACHMENT_IMAGE_FILE_SIZE_LIMIT,
- }, 200
+ config = UploadConfig(
+ file_size_limit=dify_config.UPLOAD_FILE_SIZE_LIMIT,
+ batch_count_limit=dify_config.UPLOAD_FILE_BATCH_LIMIT,
+ file_upload_limit=dify_config.BATCH_UPLOAD_LIMIT,
+ image_file_size_limit=dify_config.UPLOAD_IMAGE_FILE_SIZE_LIMIT,
+ video_file_size_limit=dify_config.UPLOAD_VIDEO_FILE_SIZE_LIMIT,
+ audio_file_size_limit=dify_config.UPLOAD_AUDIO_FILE_SIZE_LIMIT,
+ workflow_file_upload_limit=dify_config.WORKFLOW_FILE_UPLOAD_LIMIT,
+ image_file_batch_limit=dify_config.IMAGE_FILE_BATCH_LIMIT,
+ single_chunk_attachment_limit=dify_config.SINGLE_CHUNK_ATTACHMENT_LIMIT,
+ attachment_image_file_size_limit=dify_config.ATTACHMENT_IMAGE_FILE_SIZE_LIMIT,
+ )
+ return config.model_dump(mode="json"), 200
 
  @setup_required
  @login_required
  @account_initialization_required
- @marshal_with(file_fields)
  @cloud_edition_billing_resource_check("documents")
+ @console_ns.response(201, "File uploaded successfully", console_ns.models[FileResponse.__name__])
  def post(self):
  current_user, _ = current_account_with_tenant()
  source_str = request.form.get("source")
@@ -90,7 +94,8 @@ def post(self):
  except services.errors.file.BlockedFileExtensionError as blocked_extension_error:
  raise BlockedFileExtensionError(blocked_extension_error.description)
 
- return upload_file, 201
+ response = FileResponse.model_validate(upload_file, from_attributes=True)
+ return response.model_dump(mode="json"), 201
 
 
 @console_ns.route("/files/<uuid:file_id>/preview")

api/controllers/console/remote_files.py

@@ -1,7 +1,7 @@
 import urllib.parse
 
 import httpx
-from flask_restx import Resource, marshal_with
+from flask_restx import Resource
 from pydantic import BaseModel, Field
 
 import services
@@ -11,30 +11,34 @@
  RemoteFileUploadError,
  UnsupportedFileTypeError,
 )
+from controllers.common.schema import register_schema_models
 from core.file import helpers as file_helpers
 from core.helper import ssrf_proxy
 from extensions.ext_database import db
-from fields.file_fields import file_fields_with_signed_url, remote_file_info_fields
+from fields.file_fields import FileWithSignedUrl, RemoteFileInfo
 from libs.login import current_account_with_tenant
 from services.file_service import FileService
 
 from . import console_ns
 
+register_schema_models(console_ns, RemoteFileInfo, FileWithSignedUrl)
+
 
 @console_ns.route("/remote-files/<path:url>")
 class RemoteFileInfoApi(Resource):
- @marshal_with(remote_file_info_fields)
+ @console_ns.response(200, "Remote file info", console_ns.models[RemoteFileInfo.__name__])
  def get(self, url):
  decoded_url = urllib.parse.unquote(url)
  resp = ssrf_proxy.head(decoded_url)
  if resp.status_code != httpx.codes.OK:
  # failed back to get method
  resp = ssrf_proxy.get(decoded_url, timeout=3)
  resp.raise_for_status()
- return {
- "file_type": resp.headers.get("Content-Type", "application/octet-stream"),
- "file_length": int(resp.headers.get("Content-Length", 0)),
- }
+ info = RemoteFileInfo(
+ file_type=resp.headers.get("Content-Type", "application/octet-stream"),
+ file_length=int(resp.headers.get("Content-Length", 0)),
+ )
+ return info.model_dump(mode="json")
 
 
 class RemoteFileUploadPayload(BaseModel):
@@ -50,7 +54,7 @@ class RemoteFileUploadPayload(BaseModel):
 @console_ns.route("/remote-files/upload")
 class RemoteFileUploadApi(Resource):
  @console_ns.expect(console_ns.models[RemoteFileUploadPayload.__name__])
- @marshal_with(file_fields_with_signed_url)
+ @console_ns.response(201, "Remote file uploaded", console_ns.models[FileWithSignedUrl.__name__])
  def post(self):
  args = RemoteFileUploadPayload.model_validate(console_ns.payload)
  url = args.url
@@ -85,13 +89,14 @@ def post(self):
  except services.errors.file.UnsupportedFileTypeError:
  raise UnsupportedFileTypeError()
 
- return {
- "id": upload_file.id,
- "name": upload_file.name,
- "size": upload_file.size,
- "extension": upload_file.extension,
- "url": file_helpers.get_signed_file_url(upload_file_id=upload_file.id),
- "mime_type": upload_file.mime_type,
- "created_by": upload_file.created_by,
- "created_at": upload_file.created_at,
- }, 201
+ payload = FileWithSignedUrl(
+ id=upload_file.id,
+ name=upload_file.name,
+ size=upload_file.size,
+ extension=upload_file.extension,
+ url=file_helpers.get_signed_file_url(upload_file_id=upload_file.id),
+ mime_type=upload_file.mime_type,
+ created_by=upload_file.created_by,
+ created_at=int(upload_file.created_at.timestamp()),
+ )
+ return payload.model_dump(mode="json"), 201

api/controllers/files/upload.py

@@ -4,18 +4,18 @@
 from flask_restx import Resource
 from flask_restx.api import HTTPStatus
 from pydantic import BaseModel, Field
-from werkzeug.datastructures import FileStorage
 from werkzeug.exceptions import Forbidden
 
 import services
 from core.file.helpers import verify_plugin_file_signature
 from core.tools.tool_file_manager import ToolFileManager
-from fields.file_fields import build_file_model
+from fields.file_fields import FileResponse
 
 from ..common.errors import (
  FileTooLargeError,
  UnsupportedFileTypeError,
 )
+from ..common.schema import register_schema_models
 from ..console.wraps import setup_required
 from ..files import files_ns
 from ..inner_api.plugin.wraps import get_user
@@ -35,6 +35,8 @@ class PluginUploadQuery(BaseModel):
  PluginUploadQuery.__name__, PluginUploadQuery.model_json_schema(ref_template=DEFAULT_REF_TEMPLATE_SWAGGER_2_0)
 )
 
+register_schema_models(files_ns, FileResponse)
+
 
 @files_ns.route("/upload/for-plugin")
 class PluginUploadFileApi(Resource):
@@ -51,7 +53,7 @@ class PluginUploadFileApi(Resource):
  415: "Unsupported file type",
  }
  )
- @files_ns.marshal_with(build_file_model(files_ns), code=HTTPStatus.CREATED)
+ @files_ns.response(HTTPStatus.CREATED, "File uploaded", files_ns.models[FileResponse.__name__])
  def post(self):
  """Upload a file for plugin usage.
 
@@ -69,7 +71,7 @@ def post(self):
  """
  args = PluginUploadQuery.model_validate(request.args.to_dict(flat=True)) # type: ignore
 
- file: FileStorage | None = request.files.get("file")
+ file = request.files.get("file")
  if file is None:
  raise Forbidden("File is required.")
 
@@ -80,8 +82,8 @@ def post(self):
  user_id = args.user_id
  user = get_user(tenant_id, user_id)
 
- filename: str | None = file.filename
- mimetype: str | None = file.mimetype
+ filename = file.filename
+ mimetype = file.mimetype
 
  if not filename or not mimetype:
  raise Forbidden("Invalid request.")
@@ -111,22 +113,22 @@ def post(self):
  preview_url = ToolFileManager.sign_file(tool_file_id=tool_file.id, extension=extension)
 
  # Create a dictionary with all the necessary attributes
- result = {
- "id": tool_file.id,
- "user_id": tool_file.user_id,
- "tenant_id": tool_file.tenant_id,
- "conversation_id": tool_file.conversation_id,
- "file_key": tool_file.file_key,
- "mimetype": tool_file.mimetype,
- "original_url": tool_file.original_url,
- "name": tool_file.name,
- "size": tool_file.size,
- "mime_type": mimetype,
- "extension": extension,
- "preview_url": preview_url,
- }
-
- return result, 201
+ result = FileResponse(
+ id=tool_file.id,
+ name=tool_file.name,
+ size=tool_file.size,
+ extension=extension,
+ mime_type=mimetype,
+ preview_url=preview_url,
+ source_url=tool_file.original_url,
+ original_url=tool_file.original_url,
+ user_id=tool_file.user_id,
+ tenant_id=tool_file.tenant_id,
+ conversation_id=tool_file.conversation_id,
+ file_key=tool_file.file_key,
+ )
+
+ return result.model_dump(mode="json"), 201
  except services.errors.file.FileTooLargeError as file_too_large_error:
  raise FileTooLargeError(file_too_large_error.description)
  except services.errors.file.UnsupportedFileTypeError:

api/controllers/service_api/app/file.py

@@ -10,13 +10,16 @@
  TooManyFilesError,
  UnsupportedFileTypeError,
 )
+from controllers.common.schema import register_schema_models
 from controllers.service_api import service_api_ns
 from controllers.service_api.wraps import FetchUserArg, WhereisUserArg, validate_app_token
 from extensions.ext_database import db
-from fields.file_fields import build_file_model
+from fields.file_fields import FileResponse
 from models import App, EndUser
 from services.file_service import FileService
 
+register_schema_models(service_api_ns, FileResponse)
+
 
 @service_api_ns.route("/files/upload")
 class FileApi(Resource):
@@ -31,8 +34,8 @@ class FileApi(Resource):
  415: "Unsupported file type",
  }
  )
- @validate_app_token(fetch_user_arg=FetchUserArg(fetch_from=WhereisUserArg.FORM))
- @service_api_ns.marshal_with(build_file_model(service_api_ns), code=HTTPStatus.CREATED)
+ @validate_app_token(fetch_user_arg=FetchUserArg(fetch_from=WhereisUserArg.FORM)) # type: ignore
+ @service_api_ns.response(HTTPStatus.CREATED, "File uploaded", service_api_ns.models[FileResponse.__name__])
  def post(self, app_model: App, end_user: EndUser):
  """Upload a file for use in conversations.
 
@@ -64,4 +67,5 @@ def post(self, app_model: App, end_user: EndUser):
  except services.errors.file.UnsupportedFileTypeError:
  raise UnsupportedFileTypeError()
 
- return upload_file, 201
+ response = FileResponse.model_validate(upload_file, from_attributes=True)
+ return response.model_dump(mode="json"), 201

api/controllers/web/files.py

@@ -1,5 +1,4 @@
 from flask import request
-from flask_restx import marshal_with
 
 import services
 from controllers.common.errors import (
@@ -9,12 +8,15 @@
  TooManyFilesError,
  UnsupportedFileTypeError,
 )
+from controllers.common.schema import register_schema_models
 from controllers.web import web_ns
 from controllers.web.wraps import WebApiResource
 from extensions.ext_database import db
-from fields.file_fields import build_file_model
+from fields.file_fields import FileResponse
 from services.file_service import FileService
 
+register_schema_models(web_ns, FileResponse)
+
 
 @web_ns.route("/files/upload")
 class FileApi(WebApiResource):
@@ -28,7 +30,7 @@ class FileApi(WebApiResource):
  415: "Unsupported file type",
  }
  )
- @marshal_with(build_file_model(web_ns))
+ @web_ns.response(201, "File uploaded successfully", web_ns.models[FileResponse.__name__])
  def post(self, app_model, end_user):
  """Upload a file for use in web applications.
 
@@ -81,4 +83,5 @@ def post(self, app_model, end_user):
  except services.errors.file.UnsupportedFileTypeError:
  raise UnsupportedFileTypeError()
 
- return upload_file, 201
+ response = FileResponse.model_validate(upload_file, from_attributes=True)
+ return response.model_dump(mode="json"), 201