Underground-Digital/Workflow-Engine

1import datetime2import hashlib3import uuid4from typing import Any, Literal, Union5 6from flask_login import current_user7from werkzeug.exceptions import NotFound8 9from configs import dify_config10from constants import (11    AUDIO_EXTENSIONS,12    DOCUMENT_EXTENSIONS,13    IMAGE_EXTENSIONS,14    VIDEO_EXTENSIONS,15)16from core.file import helpers as file_helpers17from core.rag.extractor.extract_processor import ExtractProcessor18from extensions.ext_database import db19from extensions.ext_storage import storage20from models.account import Account21from models.enums import CreatedByRole22from models.model import EndUser, UploadFile23 24from .errors.file import FileTooLargeError, UnsupportedFileTypeError25 26PREVIEW_WORDS_LIMIT = 300027 28 29class FileService:30    @staticmethod31    def upload_file(32        *,33        filename: str,34        content: bytes,35        mimetype: str,36        user: Union[Account, EndUser, Any],37        source: Literal["datasets"] | None = None,38        source_url: str = "",39    ) -> UploadFile:40        # get file extension41        extension = filename.split(".")[-1].lower()42        if len(filename) > 200:43            filename = filename.split(".")[0][:200] + "." + extension44 45        if source == "datasets" and extension not in DOCUMENT_EXTENSIONS:46            raise UnsupportedFileTypeError()47 48        # get file size49        file_size = len(content)50 51        # check if the file size is exceeded52        if not FileService.is_file_size_within_limit(extension=extension, file_size=file_size):53            raise FileTooLargeError54 55        # generate file key56        file_uuid = str(uuid.uuid4())57 58        if isinstance(user, Account):59            current_tenant_id = user.current_tenant_id60        else:61            # end_user62            current_tenant_id = user.tenant_id63 64        file_key = "upload_files/" + current_tenant_id + "/" + file_uuid + "." + extension65 66        # save file to storage67        storage.save(file_key, content)68 69        # save file to db70        upload_file = UploadFile(71            tenant_id=current_tenant_id,72            storage_type=dify_config.STORAGE_TYPE,73            key=file_key,74            name=filename,75            size=file_size,76            extension=extension,77            mime_type=mimetype,78            created_by_role=(CreatedByRole.ACCOUNT if isinstance(user, Account) else CreatedByRole.END_USER),79            created_by=user.id,80            created_at=datetime.datetime.now(datetime.timezone.utc).replace(tzinfo=None),81            used=False,82            hash=hashlib.sha3_256(content).hexdigest(),83            source_url=source_url,84        )85 86        db.session.add(upload_file)87        db.session.commit()88 89        return upload_file90 91    @staticmethod92    def is_file_size_within_limit(*, extension: str, file_size: int) -> bool:93        if extension in IMAGE_EXTENSIONS:94            file_size_limit = dify_config.UPLOAD_IMAGE_FILE_SIZE_LIMIT * 1024 * 102495        elif extension in VIDEO_EXTENSIONS:96            file_size_limit = dify_config.UPLOAD_VIDEO_FILE_SIZE_LIMIT * 1024 * 102497        elif extension in AUDIO_EXTENSIONS:98            file_size_limit = dify_config.UPLOAD_AUDIO_FILE_SIZE_LIMIT * 1024 * 102499        else:100            file_size_limit = dify_config.UPLOAD_FILE_SIZE_LIMIT * 1024 * 1024101 102        return file_size <= file_size_limit103 104    @staticmethod105    def upload_text(text: str, text_name: str) -> UploadFile:106        if len(text_name) > 200:107            text_name = text_name[:200]108        # user uuid as file name109        file_uuid = str(uuid.uuid4())110        file_key = "upload_files/" + current_user.current_tenant_id + "/" + file_uuid + ".txt"111 112        # save file to storage113        storage.save(file_key, text.encode("utf-8"))114 115        # save file to db116        upload_file = UploadFile(117            tenant_id=current_user.current_tenant_id,118            storage_type=dify_config.STORAGE_TYPE,119            key=file_key,120            name=text_name,121            size=len(text),122            extension="txt",123            mime_type="text/plain",124            created_by=current_user.id,125            created_by_role=CreatedByRole.ACCOUNT,126            created_at=datetime.datetime.now(datetime.timezone.utc).replace(tzinfo=None),127            used=True,128            used_by=current_user.id,129            used_at=datetime.datetime.now(datetime.timezone.utc).replace(tzinfo=None),130        )131 132        db.session.add(upload_file)133        db.session.commit()134 135        return upload_file136 137    @staticmethod138    def get_file_preview(file_id: str):139        upload_file = db.session.query(UploadFile).filter(UploadFile.id == file_id).first()140 141        if not upload_file:142            raise NotFound("File not found")143 144        # extract text from file145        extension = upload_file.extension146        if extension.lower() not in DOCUMENT_EXTENSIONS:147            raise UnsupportedFileTypeError()148 149        text = ExtractProcessor.load_from_upload_file(upload_file, return_text=True)150        text = text[0:PREVIEW_WORDS_LIMIT] if text else ""151 152        return text153 154    @staticmethod155    def get_image_preview(file_id: str, timestamp: str, nonce: str, sign: str):156        result = file_helpers.verify_image_signature(157            upload_file_id=file_id, timestamp=timestamp, nonce=nonce, sign=sign158        )159        if not result:160            raise NotFound("File not found or signature is invalid")161 162        upload_file = db.session.query(UploadFile).filter(UploadFile.id == file_id).first()163 164        if not upload_file:165            raise NotFound("File not found or signature is invalid")166 167        # extract text from file168        extension = upload_file.extension169        if extension.lower() not in IMAGE_EXTENSIONS:170            raise UnsupportedFileTypeError()171 172        generator = storage.load(upload_file.key, stream=True)173 174        return generator, upload_file.mime_type175 176    @staticmethod177    def get_file_generator_by_file_id(file_id: str, timestamp: str, nonce: str, sign: str):178        result = file_helpers.verify_file_signature(upload_file_id=file_id, timestamp=timestamp, nonce=nonce, sign=sign)179        if not result:180            raise NotFound("File not found or signature is invalid")181 182        upload_file = db.session.query(UploadFile).filter(UploadFile.id == file_id).first()183 184        if not upload_file:185            raise NotFound("File not found or signature is invalid")186 187        generator = storage.load(upload_file.key, stream=True)188 189        return generator, upload_file190 191    @staticmethod192    def get_public_image_preview(file_id: str):193        upload_file = db.session.query(UploadFile).filter(UploadFile.id == file_id).first()194 195        if not upload_file:196            raise NotFound("File not found or signature is invalid")197 198        # extract text from file199        extension = upload_file.extension200        if extension.lower() not in IMAGE_EXTENSIONS:201            raise UnsupportedFileTypeError()202 203        generator = storage.load(upload_file.key)204 205        return generator, upload_file.mime_type206