feat: [backend] vision support (#1510)

Co-authored-by: Garfield Dai <dai.hai@foxmail.com>

api/services/file_service.py

@@ -1,46 +1,62 @@
 import datetime
 import hashlib
-import time
 import uuid
+from typing import Generator, Tuple, Union
 
-from cachetools import TTLCache
-from flask import request, current_app
+from flask import current_app
 from flask_login import current_user
 from werkzeug.datastructures import FileStorage
 from werkzeug.exceptions import NotFound
 
 from core.data_loader.file_extractor import FileExtractor
+from core.file.upload_file_parser import UploadFileParser
 from extensions.ext_storage import storage
 from extensions.ext_database import db
-from models.model import UploadFile
+from models.account import Account
+from models.model import UploadFile, EndUser
 from services.errors.file import FileTooLargeError, UnsupportedFileTypeError
 
-ALLOWED_EXTENSIONS = ['txt', 'markdown', 'md', 'pdf', 'html', 'htm', 'xlsx', 'docx', 'csv']
+ALLOWED_EXTENSIONS = ['txt', 'markdown', 'md', 'pdf', 'html', 'htm', 'xlsx', 'docx', 'csv',
+                      'jpg', 'jpeg', 'png', 'webp', 'gif']
+IMAGE_EXTENSIONS = ['jpg', 'jpeg', 'png', 'webp', 'gif']
 PREVIEW_WORDS_LIMIT = 3000
-cache = TTLCache(maxsize=None, ttl=30)
 
 
 class FileService:
 
     @staticmethod
-    def upload_file(file: FileStorage) -> UploadFile:
+    def upload_file(file: FileStorage, user: Union[Account, EndUser], only_image: bool = False) -> UploadFile:
+        extension = file.filename.split('.')[-1]
+        if extension.lower() not in ALLOWED_EXTENSIONS:
+            raise UnsupportedFileTypeError()
+        elif only_image and extension.lower() not in IMAGE_EXTENSIONS:
+            raise UnsupportedFileTypeError()
+
         # read file content
         file_content = file.read()
+
         # get file size
         file_size = len(file_content)
 
-        file_size_limit = current_app.config.get("UPLOAD_FILE_SIZE_LIMIT") * 1024 * 1024
+        if extension.lower() in IMAGE_EXTENSIONS:
+            file_size_limit = current_app.config.get("UPLOAD_IMAGE_FILE_SIZE_LIMIT") * 1024 * 1024
+        else:
+            file_size_limit = current_app.config.get("UPLOAD_FILE_SIZE_LIMIT") * 1024 * 1024
+
         if file_size > file_size_limit:
             message = f'File size exceeded. {file_size} > {file_size_limit}'
             raise FileTooLargeError(message)
 
-        extension = file.filename.split('.')[-1]
-        if extension.lower() not in ALLOWED_EXTENSIONS:
-            raise UnsupportedFileTypeError()
-
         # user uuid as file name
         file_uuid = str(uuid.uuid4())
-        file_key = 'upload_files/' + current_user.current_tenant_id + '/' + file_uuid + '.' + extension
+
+        if isinstance(user, Account):
+            current_tenant_id = user.current_tenant_id
+        else:
+            # end_user
+            current_tenant_id = user.tenant_id
+
+        file_key = 'upload_files/' + current_tenant_id + '/' + file_uuid + '.' + extension
 
         # save file to storage
         storage.save(file_key, file_content)
@@ -48,14 +64,15 @@ class FileService:
         # save file to db
         config = current_app.config
         upload_file = UploadFile(
-            tenant_id=current_user.current_tenant_id,
+            tenant_id=current_tenant_id,
             storage_type=config['STORAGE_TYPE'],
             key=file_key,
             name=file.filename,
             size=file_size,
             extension=extension,
             mime_type=file.mimetype,
-            created_by=current_user.id,
+            created_by_role=('account' if isinstance(user, Account) else 'end_user'),
+            created_by=user.id,
             created_at=datetime.datetime.utcnow(),
             used=False,
             hash=hashlib.sha3_256(file_content).hexdigest()
@@ -99,12 +116,6 @@ class FileService:
 
     @staticmethod
     def get_file_preview(file_id: str) -> str:
-        # get file storage key
-        key = file_id + request.path
-        cached_response = cache.get(key)
-        if cached_response and time.time() - cached_response['timestamp'] < cache.ttl:
-            return cached_response['response']
-
         upload_file = db.session.query(UploadFile) \
             .filter(UploadFile.id == file_id) \
             .first()
@@ -121,3 +132,25 @@ class FileService:
         text = text[0:PREVIEW_WORDS_LIMIT] if text else ''
 
         return text
+
+    @staticmethod
+    def get_image_preview(file_id: str, timestamp: str, nonce: str, sign: str) -> Tuple[Generator, str]:
+        result = UploadFileParser.verify_image_file_signature(file_id, timestamp, nonce, sign)
+        if not result:
+            raise NotFound("File not found or signature is invalid")
+
+        upload_file = db.session.query(UploadFile) \
+            .filter(UploadFile.id == file_id) \
+            .first()
+
+        if not upload_file:
+            raise NotFound("File not found or signature is invalid")
+
+        # extract text from file
+        extension = upload_file.extension
+        if extension.lower() not in IMAGE_EXTENSIONS:
+            raise UnsupportedFileTypeError()
+
+        generator = storage.load(upload_file.key, stream=True)
+
+        return generator, upload_file.mime_type