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

api/controllers/web/__init__.py

@@ -2,8 +2,17 @@ from flask import Blueprint
 
 from libs.external_api import ExternalApi
 
+from .files import FileApi
+from .remote_files import RemoteFileInfoApi, RemoteFileUploadApi
+
 bp = Blueprint("web", __name__, url_prefix="/api")
 api = ExternalApi(bp)
 
+# Files
+api.add_resource(FileApi, "/files/upload")
 
-from . import app, audio, completion, conversation, feature, file, message, passport, saved_message, site, workflow
+# Remote files
+api.add_resource(RemoteFileInfoApi, "/remote-files/<path:url>")
+api.add_resource(RemoteFileUploadApi, "/remote-files/upload")
+
+from . import app, audio, completion, conversation, feature, message, passport, saved_message, site, workflow

api/controllers/web/file.py

@@ -1,56 +0,0 @@
-import urllib.parse
-
-from flask import request
-from flask_restful import marshal_with, reqparse
-
-import services
-from controllers.web import api
-from controllers.web.error import FileTooLargeError, NoFileUploadedError, TooManyFilesError, UnsupportedFileTypeError
-from controllers.web.wraps import WebApiResource
-from core.helper import ssrf_proxy
-from fields.file_fields import file_fields, remote_file_info_fields
-from services.file_service import FileService
-
-
-class FileApi(WebApiResource):
-    @marshal_with(file_fields)
-    def post(self, app_model, end_user):
-        # get file from request
-        file = request.files["file"]
-
-        parser = reqparse.RequestParser()
-        parser.add_argument("source", type=str, required=False, location="args")
-        source = parser.parse_args().get("source")
-
-        # check file
-        if "file" not in request.files:
-            raise NoFileUploadedError()
-
-        if len(request.files) > 1:
-            raise TooManyFilesError()
-        try:
-            upload_file = FileService.upload_file(file=file, user=end_user, source=source)
-        except services.errors.file.FileTooLargeError as file_too_large_error:
-            raise FileTooLargeError(file_too_large_error.description)
-        except services.errors.file.UnsupportedFileTypeError:
-            raise UnsupportedFileTypeError()
-
-        return upload_file, 201
-
-
-class RemoteFileInfoApi(WebApiResource):
-    @marshal_with(remote_file_info_fields)
-    def get(self, url):
-        decoded_url = urllib.parse.unquote(url)
-        try:
-            response = ssrf_proxy.head(decoded_url)
-            return {
-                "file_type": response.headers.get("Content-Type", "application/octet-stream"),
-                "file_length": int(response.headers.get("Content-Length", -1)),
-            }
-        except Exception as e:
-            return {"error": str(e)}, 400
-
-
-api.add_resource(FileApi, "/files/upload")
-api.add_resource(RemoteFileInfoApi, "/remote-files/<path:url>")

api/controllers/web/files.py

@@ -0,0 +1,43 @@
+from flask import request
+from flask_restful import marshal_with
+
+import services
+from controllers.common.errors import FilenameNotExistsError
+from controllers.web.error import FileTooLargeError, NoFileUploadedError, TooManyFilesError, UnsupportedFileTypeError
+from controllers.web.wraps import WebApiResource
+from fields.file_fields import file_fields
+from services.file_service import FileService
+
+
+class FileApi(WebApiResource):
+    @marshal_with(file_fields)
+    def post(self, app_model, end_user):
+        file = request.files["file"]
+        source = request.form.get("source")
+
+        if "file" not in request.files:
+            raise NoFileUploadedError()
+
+        if len(request.files) > 1:
+            raise TooManyFilesError()
+
+        if not file.filename:
+            raise FilenameNotExistsError
+
+        if source not in ("datasets", None):
+            source = None
+
+        try:
+            upload_file = FileService.upload_file(
+                filename=file.filename,
+                content=file.read(),
+                mimetype=file.mimetype,
+                user=end_user,
+                source=source,
+            )
+        except services.errors.file.FileTooLargeError as file_too_large_error:
+            raise FileTooLargeError(file_too_large_error.description)
+        except services.errors.file.UnsupportedFileTypeError:
+            raise UnsupportedFileTypeError()
+
+        return upload_file, 201

api/controllers/web/remote_files.py

@@ -0,0 +1,69 @@
+import urllib.parse
+
+from flask_login import current_user
+from flask_restful import marshal_with, reqparse
+
+from controllers.common import helpers
+from controllers.web.wraps import WebApiResource
+from core.file import helpers as file_helpers
+from core.helper import ssrf_proxy
+from fields.file_fields import file_fields_with_signed_url, remote_file_info_fields
+from services.file_service import FileService
+
+
+class RemoteFileInfoApi(WebApiResource):
+    @marshal_with(remote_file_info_fields)
+    def get(self, url):
+        decoded_url = urllib.parse.unquote(url)
+        try:
+            response = ssrf_proxy.head(decoded_url)
+            return {
+                "file_type": response.headers.get("Content-Type", "application/octet-stream"),
+                "file_length": int(response.headers.get("Content-Length", -1)),
+            }
+        except Exception as e:
+            return {"error": str(e)}, 400
+
+
+class RemoteFileUploadApi(WebApiResource):
+    @marshal_with(file_fields_with_signed_url)
+    def post(self):
+        parser = reqparse.RequestParser()
+        parser.add_argument("url", type=str, required=True, help="URL is required")
+        args = parser.parse_args()
+
+        url = args["url"]
+
+        response = ssrf_proxy.head(url)
+        response.raise_for_status()
+
+        file_info = helpers.guess_file_info_from_response(response)
+
+        if not FileService.is_file_size_within_limit(extension=file_info.extension, file_size=file_info.size):
+            return {"error": "File size exceeded"}, 400
+
+        response = ssrf_proxy.get(url)
+        response.raise_for_status()
+        content = response.content
+
+        try:
+            upload_file = FileService.upload_file(
+                filename=file_info.filename,
+                content=content,
+                mimetype=file_info.mimetype,
+                user=current_user,
+                source_url=url,
+            )
+        except Exception as e:
+            return {"error": str(e)}, 400
+
+        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