estel/dify
1
0
Fork 0
Code Issues Pull Requests Actions 1 Packages Projects Releases Wiki Activity

chore(api/services): apply ruff reformatting (#7599)

Browse Source 
Co-authored-by: -LAN- <laipz8200@outlook.com>
This commit is contained in:
Bowen Liang
2024-08-26 13:43:57 +08:00
committed by GitHub
parent 979422cdc6
commit 17fd773a30
49 changed files with 2630 additions and 2655 deletions
Download Patch File Download Diff File Expand all files Collapse all files

68
api/services/file_service.py
View File

@@ -17,27 +17,45 @@ from models.account import Account
from models.model import EndUser, UploadFile
from services.errors.file import FileTooLargeError, UnsupportedFileTypeError
IMAGE_EXTENSIONS = ['jpg', 'jpeg', 'png', 'webp', 'gif', 'svg']
IMAGE_EXTENSIONS = ["jpg", "jpeg", "png", "webp", "gif", "svg"]
IMAGE_EXTENSIONS.extend([ext.upper() for ext in IMAGE_EXTENSIONS])
ALLOWED_EXTENSIONS = ['txt', 'markdown', 'md', 'pdf', 'html', 'htm', 'xlsx', 'xls', 'docx', 'csv']
UNSTRUCTURED_ALLOWED_EXTENSIONS = ['txt', 'markdown', 'md', 'pdf', 'html', 'htm', 'xlsx', 'xls',
'docx', 'csv', 'eml', 'msg', 'pptx', 'ppt', 'xml', 'epub']
ALLOWED_EXTENSIONS = ["txt", "markdown", "md", "pdf", "html", "htm", "xlsx", "xls", "docx", "csv"]
UNSTRUCTURED_ALLOWED_EXTENSIONS = [
"txt",
"markdown",
"md",
"pdf",
"html",
"htm",
"xlsx",
"xls",
"docx",
"csv",
"eml",
"msg",
"pptx",
"ppt",
"xml",
"epub",
]
PREVIEW_WORDS_LIMIT = 3000
class FileService:
@staticmethod
def upload_file(file: FileStorage, user: Union[Account, EndUser], only_image: bool = False) -> UploadFile:
filename = file.filename
extension = file.filename.split('.')[-1]
extension = file.filename.split(".")[-1]
if len(filename) > 200:
filename = filename.split('.')[0][:200] + '.' + extension
filename = filename.split(".")[0][:200] + "." + extension
etl_type = dify_config.ETL_TYPE
allowed_extensions = UNSTRUCTURED_ALLOWED_EXTENSIONS + IMAGE_EXTENSIONS if etl_type == 'Unstructured' \
allowed_extensions = (
UNSTRUCTURED_ALLOWED_EXTENSIONS + IMAGE_EXTENSIONS
if etl_type == "Unstructured"
else ALLOWED_EXTENSIONS + IMAGE_EXTENSIONS
)
if extension.lower() not in allowed_extensions:
raise UnsupportedFileTypeError()
elif only_image and extension.lower() not in IMAGE_EXTENSIONS:
@@ -55,7 +73,7 @@ class FileService:
file_size_limit = dify_config.UPLOAD_FILE_SIZE_LIMIT * 1024 * 1024
if file_size > file_size_limit:
message = f'File size exceeded. {file_size} > {file_size_limit}'
message = f"File size exceeded. {file_size} > {file_size_limit}"
raise FileTooLargeError(message)
# user uuid as file name
@@ -67,7 +85,7 @@ class FileService:
# end_user
current_tenant_id = user.tenant_id
file_key = 'upload_files/' + current_tenant_id + '/' + file_uuid + '.' + extension
file_key = "upload_files/" + current_tenant_id + "/" + file_uuid + "." + extension
# save file to storage
storage.save(file_key, file_content)
@@ -81,11 +99,11 @@ class FileService:
size=file_size,
extension=extension,
mime_type=file.mimetype,
created_by_role=('account' if isinstance(user, Account) else 'end_user'),
created_by_role=("account" if isinstance(user, Account) else "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(file_content).hexdigest(),
)
db.session.add(upload_file)
@@ -99,10 +117,10 @@ class FileService:
text_name = text_name[:200]
# user uuid as file name
file_uuid = str(uuid.uuid4())
file_key = 'upload_files/' + current_user.current_tenant_id + '/' + file_uuid + '.txt'
file_key = "upload_files/" + current_user.current_tenant_id + "/" + file_uuid + ".txt"
# save file to storage
storage.save(file_key, text.encode('utf-8'))
storage.save(file_key, text.encode("utf-8"))
# save file to db
upload_file = UploadFile(
@@ -111,13 +129,13 @@ class FileService:
key=file_key,
name=text_name,
size=len(text),
extension='txt',
mime_type='text/plain',
extension="txt",
mime_type="text/plain",
created_by=current_user.id,
created_at=datetime.datetime.now(datetime.timezone.utc).replace(tzinfo=None),
used=True,
used_by=current_user.id,
used_at=datetime.datetime.now(datetime.timezone.utc).replace(tzinfo=None)
used_at=datetime.datetime.now(datetime.timezone.utc).replace(tzinfo=None),
)
db.session.add(upload_file)
@@ -127,9 +145,7 @@ class FileService:
@staticmethod
def get_file_preview(file_id: str) -> str:
upload_file = db.session.query(UploadFile) \
.filter(UploadFile.id == file_id) \
.first()
upload_file = db.session.query(UploadFile).filter(UploadFile.id == file_id).first()
if not upload_file:
raise NotFound("File not found")
@@ -137,12 +153,12 @@ class FileService:
# extract text from file
extension = upload_file.extension
etl_type = dify_config.ETL_TYPE
allowed_extensions = UNSTRUCTURED_ALLOWED_EXTENSIONS if etl_type == 'Unstructured' else ALLOWED_EXTENSIONS
allowed_extensions = UNSTRUCTURED_ALLOWED_EXTENSIONS if etl_type == "Unstructured" else ALLOWED_EXTENSIONS
if extension.lower() not in allowed_extensions:
raise UnsupportedFileTypeError()
text = ExtractProcessor.load_from_upload_file(upload_file, return_text=True)
text = text[0:PREVIEW_WORDS_LIMIT] if text else ''
text = text[0:PREVIEW_WORDS_LIMIT] if text else ""
return text
@@ -152,9 +168,7 @@ class FileService:
if not result:
raise NotFound("File not found or signature is invalid")
upload_file = db.session.query(UploadFile) \
.filter(UploadFile.id == file_id) \
.first()
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")
@@ -170,9 +184,7 @@ class FileService:
@staticmethod
def get_public_image_preview(file_id: str) -> tuple[Generator, str]:
upload_file = db.session.query(UploadFile) \
.filter(UploadFile.id == file_id) \
.first()
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")