Feat/add-remote-file-upload-api (#9906)

api/services/file_service.py

@@ -1,10 +1,9 @@
 import datetime
 import hashlib
 import uuid
-from typing import Literal, Union
+from typing import Any, Literal, Union
 
 from flask_login import current_user
-from werkzeug.datastructures import FileStorage
 from werkzeug.exceptions import NotFound
 
 from configs import dify_config
@@ -21,7 +20,8 @@ from extensions.ext_storage import storage
 from models.account import Account
 from models.enums import CreatedByRole
 from models.model import EndUser, UploadFile
-from services.errors.file import FileNotExistsError, FileTooLargeError, UnsupportedFileTypeError
+
+from .errors.file import FileTooLargeError, UnsupportedFileTypeError
 
 PREVIEW_WORDS_LIMIT = 3000
 
@@ -29,12 +29,15 @@ PREVIEW_WORDS_LIMIT = 3000
 class FileService:
     @staticmethod
     def upload_file(
-        file: FileStorage, user: Union[Account, EndUser], source: Literal["datasets"] | None = None
+        *,
+        filename: str,
+        content: bytes,
+        mimetype: str,
+        user: Union[Account, EndUser, Any],
+        source: Literal["datasets"] | None = None,
+        source_url: str = "",
     ) -> UploadFile:
-        # get file name
-        filename = file.filename
-        if not filename:
-            raise FileNotExistsError
+        # get file extension
         extension = filename.split(".")[-1].lower()
         if len(filename) > 200:
             filename = filename.split(".")[0][:200] + "." + extension
@@ -42,25 +45,12 @@ class FileService:
         if source == "datasets" and extension not in DOCUMENT_EXTENSIONS:
             raise UnsupportedFileTypeError()
 
-        # select file size limit
-        if extension in IMAGE_EXTENSIONS:
-            file_size_limit = dify_config.UPLOAD_IMAGE_FILE_SIZE_LIMIT * 1024 * 1024
-        elif extension in VIDEO_EXTENSIONS:
-            file_size_limit = dify_config.UPLOAD_VIDEO_FILE_SIZE_LIMIT * 1024 * 1024
-        elif extension in AUDIO_EXTENSIONS:
-            file_size_limit = dify_config.UPLOAD_AUDIO_FILE_SIZE_LIMIT * 1024 * 1024
-        else:
-            file_size_limit = dify_config.UPLOAD_FILE_SIZE_LIMIT * 1024 * 1024
-
-        # read file content
-        file_content = file.read()
         # get file size
-        file_size = len(file_content)
+        file_size = len(content)
 
         # check if the file size is exceeded
-        if file_size > file_size_limit:
-            message = f"File size exceeded. {file_size} > {file_size_limit}"
-            raise FileTooLargeError(message)
+        if not FileService.is_file_size_within_limit(extension=extension, file_size=file_size):
+            raise FileTooLargeError
 
         # generate file key
         file_uuid = str(uuid.uuid4())
@@ -74,7 +64,7 @@ class FileService:
         file_key = "upload_files/" + current_tenant_id + "/" + file_uuid + "." + extension
 
         # save file to storage
-        storage.save(file_key, file_content)
+        storage.save(file_key, content)
 
         # save file to db
         upload_file = UploadFile(
@@ -84,12 +74,13 @@ class FileService:
             name=filename,
             size=file_size,
             extension=extension,
-            mime_type=file.mimetype,
+            mime_type=mimetype,
             created_by_role=(CreatedByRole.ACCOUNT if isinstance(user, Account) else CreatedByRole.END_USER),
             created_by=user.id,
             created_at=datetime.datetime.now(datetime.timezone.utc).replace(tzinfo=None),
             used=False,
-            hash=hashlib.sha3_256(file_content).hexdigest(),
+            hash=hashlib.sha3_256(content).hexdigest(),
+            source_url=source_url,
         )
 
         db.session.add(upload_file)
@@ -97,6 +88,19 @@ class FileService:
 
         return upload_file
 
+    @staticmethod
+    def is_file_size_within_limit(*, extension: str, file_size: int) -> bool:
+        if extension in IMAGE_EXTENSIONS:
+            file_size_limit = dify_config.UPLOAD_IMAGE_FILE_SIZE_LIMIT * 1024 * 1024
+        elif extension in VIDEO_EXTENSIONS:
+            file_size_limit = dify_config.UPLOAD_VIDEO_FILE_SIZE_LIMIT * 1024 * 1024
+        elif extension in AUDIO_EXTENSIONS:
+            file_size_limit = dify_config.UPLOAD_AUDIO_FILE_SIZE_LIMIT * 1024 * 1024
+        else:
+            file_size_limit = dify_config.UPLOAD_FILE_SIZE_LIMIT * 1024 * 1024
+
+        return file_size <= file_size_limit
+
     @staticmethod
     def upload_text(text: str, text_name: str) -> UploadFile:
         if len(text_name) > 200: