refactor(services/tasks): Swtich to dify_config witch Pydantic (#6203)

This commit is contained in:
Waffle
2024-07-12 12:25:38 +08:00
committed by GitHub
parent d7a6f25c63
commit 7b225a5ab0
15 changed files with 51 additions and 53 deletions

View File

@@ -4,11 +4,11 @@ import uuid
from collections.abc import Generator
from typing import Union
from flask import current_app
from flask_login import current_user
from werkzeug.datastructures import FileStorage
from werkzeug.exceptions import NotFound
from configs import dify_config
from core.file.upload_file_parser import UploadFileParser
from core.rag.extractor.extract_processor import ExtractProcessor
from extensions.ext_database import db
@@ -35,7 +35,7 @@ class FileService:
extension = file.filename.split('.')[-1]
if len(filename) > 200:
filename = filename.split('.')[0][:200] + '.' + extension
etl_type = current_app.config['ETL_TYPE']
etl_type = dify_config.ETL_TYPE
allowed_extensions = UNSTRUCTURED_ALLOWED_EXTENSIONS + IMAGE_EXTENSIONS if etl_type == 'Unstructured' \
else ALLOWED_EXTENSIONS + IMAGE_EXTENSIONS
if extension.lower() not in allowed_extensions:
@@ -50,9 +50,9 @@ class FileService:
file_size = len(file_content)
if extension.lower() in IMAGE_EXTENSIONS:
file_size_limit = current_app.config.get("UPLOAD_IMAGE_FILE_SIZE_LIMIT") * 1024 * 1024
file_size_limit = dify_config.UPLOAD_IMAGE_FILE_SIZE_LIMIT * 1024 * 1024
else:
file_size_limit = current_app.config.get("UPLOAD_FILE_SIZE_LIMIT") * 1024 * 1024
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}'
@@ -73,10 +73,9 @@ class FileService:
storage.save(file_key, file_content)
# save file to db
config = current_app.config
upload_file = UploadFile(
tenant_id=current_tenant_id,
storage_type=config['STORAGE_TYPE'],
storage_type=dify_config.STORAGE_TYPE,
key=file_key,
name=filename,
size=file_size,
@@ -106,10 +105,9 @@ class FileService:
storage.save(file_key, text.encode('utf-8'))
# save file to db
config = current_app.config
upload_file = UploadFile(
tenant_id=current_user.current_tenant_id,
storage_type=config['STORAGE_TYPE'],
storage_type=dify_config.STORAGE_TYPE,
key=file_key,
name=text_name + '.txt',
size=len(text),
@@ -138,7 +136,7 @@ class FileService:
# extract text from file
extension = upload_file.extension
etl_type = current_app.config['ETL_TYPE']
etl_type = dify_config.ETL_TYPE
allowed_extensions = UNSTRUCTURED_ALLOWED_EXTENSIONS if etl_type == 'Unstructured' else ALLOWED_EXTENSIONS
if extension.lower() not in allowed_extensions:
raise UnsupportedFileTypeError()