Underground-Digital/Workflow-Engine

1from flask import request2from flask_login import current_user3from flask_restful import Resource, marshal_with4 5import services6from configs import dify_config7from constants import DOCUMENT_EXTENSIONS8from controllers.common.errors import FilenameNotExistsError9from controllers.console.wraps import (10    account_initialization_required,11    cloud_edition_billing_resource_check,12    setup_required,13)14from fields.file_fields import file_fields, upload_config_fields15from libs.login import login_required16from services.file_service import FileService17 18from .errors import (19    FileTooLargeError,20    NoFileUploadedError,21    TooManyFilesError,22    UnsupportedFileTypeError,23)24 25PREVIEW_WORDS_LIMIT = 300026 27 28class FileApi(Resource):29    @setup_required30    @login_required31    @account_initialization_required32    @marshal_with(upload_config_fields)33    def get(self):34        return {35            "file_size_limit": dify_config.UPLOAD_FILE_SIZE_LIMIT,36            "batch_count_limit": dify_config.UPLOAD_FILE_BATCH_LIMIT,37            "image_file_size_limit": dify_config.UPLOAD_IMAGE_FILE_SIZE_LIMIT,38            "video_file_size_limit": dify_config.UPLOAD_VIDEO_FILE_SIZE_LIMIT,39            "audio_file_size_limit": dify_config.UPLOAD_AUDIO_FILE_SIZE_LIMIT,40            "workflow_file_upload_limit": dify_config.WORKFLOW_FILE_UPLOAD_LIMIT,41        }, 20042 43    @setup_required44    @login_required45    @account_initialization_required46    @marshal_with(file_fields)47    @cloud_edition_billing_resource_check("documents")48    def post(self):49        file = request.files["file"]50        source = request.form.get("source")51 52        if "file" not in request.files:53            raise NoFileUploadedError()54 55        if len(request.files) > 1:56            raise TooManyFilesError()57 58        if not file.filename:59            raise FilenameNotExistsError60 61        if source not in ("datasets", None):62            source = None63 64        try:65            upload_file = FileService.upload_file(66                filename=file.filename,67                content=file.read(),68                mimetype=file.mimetype,69                user=current_user,70                source=source,71            )72        except services.errors.file.FileTooLargeError as file_too_large_error:73            raise FileTooLargeError(file_too_large_error.description)74        except services.errors.file.UnsupportedFileTypeError:75            raise UnsupportedFileTypeError()76 77        return upload_file, 20178 79 80class FilePreviewApi(Resource):81    @setup_required82    @login_required83    @account_initialization_required84    def get(self, file_id):85        file_id = str(file_id)86        text = FileService.get_file_preview(file_id)87        return {"content": text}88 89 90class FileSupportTypeApi(Resource):91    @setup_required92    @login_required93    @account_initialization_required94    def get(self):95        return {"allowed_extensions": DOCUMENT_EXTENSIONS}96