dify
This commit is contained in:
245
dify/api/services/file_service.py
Normal file
245
dify/api/services/file_service.py
Normal file
@@ -0,0 +1,245 @@
|
||||
import hashlib
|
||||
import os
|
||||
import uuid
|
||||
from typing import Literal, Union
|
||||
|
||||
from sqlalchemy import Engine, select
|
||||
from sqlalchemy.orm import Session, sessionmaker
|
||||
from werkzeug.exceptions import NotFound
|
||||
|
||||
from configs import dify_config
|
||||
from constants import (
|
||||
AUDIO_EXTENSIONS,
|
||||
DOCUMENT_EXTENSIONS,
|
||||
IMAGE_EXTENSIONS,
|
||||
VIDEO_EXTENSIONS,
|
||||
)
|
||||
from core.file import helpers as file_helpers
|
||||
from core.rag.extractor.extract_processor import ExtractProcessor
|
||||
from extensions.ext_storage import storage
|
||||
from libs.datetime_utils import naive_utc_now
|
||||
from libs.helper import extract_tenant_id
|
||||
from models import Account
|
||||
from models.enums import CreatorUserRole
|
||||
from models.model import EndUser, UploadFile
|
||||
|
||||
from .errors.file import BlockedFileExtensionError, FileTooLargeError, UnsupportedFileTypeError
|
||||
|
||||
PREVIEW_WORDS_LIMIT = 3000
|
||||
|
||||
|
||||
class FileService:
|
||||
_session_maker: sessionmaker[Session]
|
||||
|
||||
def __init__(self, session_factory: sessionmaker | Engine | None = None):
|
||||
if isinstance(session_factory, Engine):
|
||||
self._session_maker = sessionmaker(bind=session_factory)
|
||||
elif isinstance(session_factory, sessionmaker):
|
||||
self._session_maker = session_factory
|
||||
else:
|
||||
raise AssertionError("must be a sessionmaker or an Engine.")
|
||||
|
||||
def upload_file(
|
||||
self,
|
||||
*,
|
||||
filename: str,
|
||||
content: bytes,
|
||||
mimetype: str,
|
||||
user: Union[Account, EndUser],
|
||||
source: Literal["datasets"] | None = None,
|
||||
source_url: str = "",
|
||||
) -> UploadFile:
|
||||
# get file extension
|
||||
extension = os.path.splitext(filename)[1].lstrip(".").lower()
|
||||
|
||||
# check if filename contains invalid characters
|
||||
if any(c in filename for c in ["/", "\\", ":", "*", "?", '"', "<", ">", "|"]):
|
||||
raise ValueError("Filename contains invalid characters")
|
||||
|
||||
if len(filename) > 200:
|
||||
filename = filename.split(".")[0][:200] + "." + extension
|
||||
|
||||
# check if extension is in blacklist
|
||||
if extension and extension in dify_config.UPLOAD_FILE_EXTENSION_BLACKLIST:
|
||||
raise BlockedFileExtensionError(f"File extension '.{extension}' is not allowed for security reasons")
|
||||
|
||||
if source == "datasets" and extension not in DOCUMENT_EXTENSIONS:
|
||||
raise UnsupportedFileTypeError()
|
||||
|
||||
# get file size
|
||||
file_size = len(content)
|
||||
|
||||
# check if the file size is exceeded
|
||||
if not FileService.is_file_size_within_limit(extension=extension, file_size=file_size):
|
||||
raise FileTooLargeError
|
||||
|
||||
# generate file key
|
||||
file_uuid = str(uuid.uuid4())
|
||||
|
||||
current_tenant_id = extract_tenant_id(user)
|
||||
|
||||
file_key = "upload_files/" + (current_tenant_id or "") + "/" + file_uuid + "." + extension
|
||||
|
||||
# save file to storage
|
||||
storage.save(file_key, content)
|
||||
|
||||
# save file to db
|
||||
upload_file = UploadFile(
|
||||
tenant_id=current_tenant_id or "",
|
||||
storage_type=dify_config.STORAGE_TYPE,
|
||||
key=file_key,
|
||||
name=filename,
|
||||
size=file_size,
|
||||
extension=extension,
|
||||
mime_type=mimetype,
|
||||
created_by_role=(CreatorUserRole.ACCOUNT if isinstance(user, Account) else CreatorUserRole.END_USER),
|
||||
created_by=user.id,
|
||||
created_at=naive_utc_now(),
|
||||
used=False,
|
||||
hash=hashlib.sha3_256(content).hexdigest(),
|
||||
source_url=source_url,
|
||||
)
|
||||
# The `UploadFile` ID is generated within its constructor, so flushing to retrieve the ID is unnecessary.
|
||||
# We can directly generate the `source_url` here before committing.
|
||||
if not upload_file.source_url:
|
||||
upload_file.source_url = file_helpers.get_signed_file_url(upload_file_id=upload_file.id)
|
||||
|
||||
with self._session_maker(expire_on_commit=False) as session:
|
||||
session.add(upload_file)
|
||||
session.commit()
|
||||
|
||||
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
|
||||
|
||||
def upload_text(self, text: str, text_name: str, user_id: str, tenant_id: str) -> UploadFile:
|
||||
if len(text_name) > 200:
|
||||
text_name = text_name[:200]
|
||||
# user uuid as file name
|
||||
file_uuid = str(uuid.uuid4())
|
||||
file_key = "upload_files/" + tenant_id + "/" + file_uuid + ".txt"
|
||||
|
||||
# save file to storage
|
||||
storage.save(file_key, text.encode("utf-8"))
|
||||
|
||||
# save file to db
|
||||
upload_file = UploadFile(
|
||||
tenant_id=tenant_id,
|
||||
storage_type=dify_config.STORAGE_TYPE,
|
||||
key=file_key,
|
||||
name=text_name,
|
||||
size=len(text),
|
||||
extension="txt",
|
||||
mime_type="text/plain",
|
||||
created_by=user_id,
|
||||
created_by_role=CreatorUserRole.ACCOUNT,
|
||||
created_at=naive_utc_now(),
|
||||
used=True,
|
||||
used_by=user_id,
|
||||
used_at=naive_utc_now(),
|
||||
)
|
||||
|
||||
with self._session_maker(expire_on_commit=False) as session:
|
||||
session.add(upload_file)
|
||||
session.commit()
|
||||
|
||||
return upload_file
|
||||
|
||||
def get_file_preview(self, file_id: str):
|
||||
with self._session_maker(expire_on_commit=False) as session:
|
||||
upload_file = session.query(UploadFile).where(UploadFile.id == file_id).first()
|
||||
|
||||
if not upload_file:
|
||||
raise NotFound("File not found")
|
||||
|
||||
# extract text from file
|
||||
extension = upload_file.extension
|
||||
if extension.lower() not in DOCUMENT_EXTENSIONS:
|
||||
raise UnsupportedFileTypeError()
|
||||
|
||||
text = ExtractProcessor.load_from_upload_file(upload_file, return_text=True)
|
||||
text = text[0:PREVIEW_WORDS_LIMIT] if text else ""
|
||||
|
||||
return text
|
||||
|
||||
def get_image_preview(self, file_id: str, timestamp: str, nonce: str, sign: str):
|
||||
result = file_helpers.verify_image_signature(
|
||||
upload_file_id=file_id, timestamp=timestamp, nonce=nonce, sign=sign
|
||||
)
|
||||
if not result:
|
||||
raise NotFound("File not found or signature is invalid")
|
||||
with self._session_maker(expire_on_commit=False) as session:
|
||||
upload_file = session.query(UploadFile).where(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
|
||||
|
||||
def get_file_generator_by_file_id(self, file_id: str, timestamp: str, nonce: str, sign: str):
|
||||
result = file_helpers.verify_file_signature(upload_file_id=file_id, timestamp=timestamp, nonce=nonce, sign=sign)
|
||||
if not result:
|
||||
raise NotFound("File not found or signature is invalid")
|
||||
|
||||
with self._session_maker(expire_on_commit=False) as session:
|
||||
upload_file = session.query(UploadFile).where(UploadFile.id == file_id).first()
|
||||
|
||||
if not upload_file:
|
||||
raise NotFound("File not found or signature is invalid")
|
||||
|
||||
generator = storage.load(upload_file.key, stream=True)
|
||||
|
||||
return generator, upload_file
|
||||
|
||||
def get_public_image_preview(self, file_id: str):
|
||||
with self._session_maker(expire_on_commit=False) as session:
|
||||
upload_file = session.query(UploadFile).where(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)
|
||||
|
||||
return generator, upload_file.mime_type
|
||||
|
||||
def get_file_content(self, file_id: str) -> str:
|
||||
with self._session_maker(expire_on_commit=False) as session:
|
||||
upload_file: UploadFile | None = session.query(UploadFile).where(UploadFile.id == file_id).first()
|
||||
|
||||
if not upload_file:
|
||||
raise NotFound("File not found")
|
||||
content = storage.load(upload_file.key)
|
||||
|
||||
return content.decode("utf-8")
|
||||
|
||||
def delete_file(self, file_id: str):
|
||||
with self._session_maker() as session, session.begin():
|
||||
upload_file = session.scalar(select(UploadFile).where(UploadFile.id == file_id))
|
||||
|
||||
if not upload_file:
|
||||
return
|
||||
storage.delete(upload_file.key)
|
||||
session.delete(upload_file)
|
||||
Reference in New Issue
Block a user