Initial commit

api/controllers/console/datasets/file.py

@@ -0,0 +1,147 @@
+import datetime
+import hashlib
+import tempfile
+import time
+import uuid
+from pathlib import Path
+
+from cachetools import TTLCache
+from flask import request, current_app
+from flask_login import login_required, current_user
+from flask_restful import Resource, marshal_with, fields
+from werkzeug.exceptions import NotFound
+
+from controllers.console import api
+from controllers.console.datasets.error import NoFileUploadedError, TooManyFilesError, FileTooLargeError, \
+    UnsupportedFileTypeError
+from controllers.console.setup import setup_required
+from controllers.console.wraps import account_initialization_required
+from core.index.readers.html_parser import HTMLParser
+from core.index.readers.pdf_parser import PDFParser
+from extensions.ext_storage import storage
+from libs.helper import TimestampField
+from extensions.ext_database import db
+from models.model import UploadFile
+
+cache = TTLCache(maxsize=None, ttl=30)
+
+FILE_SIZE_LIMIT = 15 * 1024 * 1024  # 15MB
+ALLOWED_EXTENSIONS = ['txt', 'markdown', 'md', 'pdf', 'html', 'htm']
+PREVIEW_WORDS_LIMIT = 3000
+
+
+class FileApi(Resource):
+    file_fields = {
+        'id': fields.String,
+        'name': fields.String,
+        'size': fields.Integer,
+        'extension': fields.String,
+        'mime_type': fields.String,
+        'created_by': fields.String,
+        'created_at': TimestampField,
+    }
+
+    @setup_required
+    @login_required
+    @account_initialization_required
+    @marshal_with(file_fields)
+    def post(self):
+
+        # get file from request
+        file = request.files['file']
+
+        # check file
+        if 'file' not in request.files:
+            raise NoFileUploadedError()
+
+        if len(request.files) > 1:
+            raise TooManyFilesError()
+
+        file_content = file.read()
+        file_size = len(file_content)
+
+        if file_size > FILE_SIZE_LIMIT:
+            message = "({file_size} > {FILE_SIZE_LIMIT})"
+            raise FileTooLargeError(message)
+
+        extension = file.filename.split('.')[-1]
+        if extension 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
+
+        # save file to storage
+        storage.save(file_key, file_content)
+
+        # save file to db
+        config = current_app.config
+        upload_file = UploadFile(
+            tenant_id=current_user.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_at=datetime.datetime.utcnow(),
+            used=False,
+            hash=hashlib.sha3_256(file_content).hexdigest()
+        )
+
+        db.session.add(upload_file)
+        db.session.commit()
+
+        return upload_file, 201
+
+
+class FilePreviewApi(Resource):
+    @setup_required
+    @login_required
+    @account_initialization_required
+    def get(self, file_id):
+        file_id = str(file_id)
+
+        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()
+
+        if not upload_file:
+            raise NotFound("File not found")
+
+        # extract text from file
+        extension = upload_file.extension
+        if extension not in ALLOWED_EXTENSIONS:
+            raise UnsupportedFileTypeError()
+
+        with tempfile.TemporaryDirectory() as temp_dir:
+            suffix = Path(upload_file.key).suffix
+            filepath = f"{temp_dir}/{next(tempfile._get_candidate_names())}{suffix}"
+            storage.download(upload_file.key, filepath)
+
+            if extension == 'pdf':
+                parser = PDFParser({'upload_file': upload_file})
+                text = parser.parse_file(Path(filepath))
+            elif extension in ['html', 'htm']:
+                # Use BeautifulSoup to extract text
+                parser = HTMLParser()
+                text = parser.parse_file(Path(filepath))
+            else:
+                # ['txt', 'markdown', 'md']
+                with open(filepath, "rb") as fp:
+                    data = fp.read()
+                    text = data.decode(encoding='utf-8').strip() if data else ''
+
+        text = text[0:PREVIEW_WORDS_LIMIT] if text else ''
+        return {'content': text}
+
+
+api.add_resource(FileApi, '/files/upload')
+api.add_resource(FilePreviewApi, '/files/<uuid:file_id>/preview')