from typing import Literal from uuid import UUID from flask import request from flask_restx import Resource from werkzeug.exceptions import Forbidden import services from configs import dify_config from constants import DOCUMENT_EXTENSIONS from controllers.common.errors import ( BlockedFileExtensionError, FilenameNotExistsError, FileTooLargeError, NoFileUploadedError, TooManyFilesError, UnsupportedFileTypeError, ) from controllers.common.fields import AllowedExtensionsResponse, TextContentResponse from controllers.common.schema import JsonResponseWithStatus, register_response_schema_models from controllers.console import console_ns from controllers.console.flask_admission import console_account_admission from controllers.console.wraps import ( account_initialization_required, cloud_edition_billing_resource_check, setup_required, with_current_user, ) from extensions.ext_application_services import application_services from fields.file_fields import FileResponse, UploadConfig from libs.helper import dump_response from libs.login import login_required from machinery.context import RequestContext from models import Account, UploadFile from services.feature_service import FeatureService register_response_schema_models( console_ns, UploadConfig, FileResponse, AllowedExtensionsResponse, TextContentResponse, ) PREVIEW_WORDS_LIMIT = 3000 FILE_UPLOAD_PARAMS = { "file": { "description": "File to upload", "in": "formData", "type": "file", "required": True, }, "source": { "description": "Optional upload source", "in": "formData", "type": "string", "enum": ["datasets"], "required": False, }, } def upload_file_from_request(*, current_user: Account, resource_tenant_id: str | None = None) -> UploadFile: """Validate the multipart request and persist the file under the requested resource tenant.""" source_str = request.args.get("source") or request.form.get("source") source: Literal["datasets"] | None = "datasets" if source_str == "datasets" else None if "file" not in request.files: raise NoFileUploadedError() if len(request.files) > 1: raise TooManyFilesError() file = request.files["file"] if not file.filename: raise FilenameNotExistsError() if source == "datasets" and not current_user.is_dataset_editor: raise Forbidden() if source not in ("datasets", None): source = None default_file_size_limit = ( FeatureService.get_knowledge_file_size_limit(resource_tenant_id or current_user.current_tenant_id) if source == "datasets" else None ) try: return application_services().files.upload_file( filename=file.filename, content=file.stream.read(), mimetype=file.mimetype, user=current_user, tenant_id=resource_tenant_id, source=source, default_file_size_limit=default_file_size_limit, ) except services.errors.file.FileTooLargeError as file_too_large_error: raise FileTooLargeError(file_too_large_error.description) from file_too_large_error except services.errors.file.UnsupportedFileTypeError as unsupported_file_type_error: raise UnsupportedFileTypeError() from unsupported_file_type_error except services.errors.file.BlockedFileExtensionError as blocked_extension_error: raise BlockedFileExtensionError(blocked_extension_error.description) from blocked_extension_error @console_ns.route("/files/upload") class FileApi(Resource): @console_ns.response(200, "Success", console_ns.models[UploadConfig.__name__]) @console_account_admission() def get(self, request_context: RequestContext) -> JsonResponseWithStatus: config = UploadConfig( file_size_limit=dify_config.UPLOAD_FILE_SIZE_LIMIT, knowledge_file_size_limit=FeatureService.get_knowledge_file_size_limit(request_context.active_workspace_id), batch_count_limit=dify_config.UPLOAD_FILE_BATCH_LIMIT, file_upload_limit=dify_config.BATCH_UPLOAD_LIMIT, image_file_size_limit=dify_config.UPLOAD_IMAGE_FILE_SIZE_LIMIT, video_file_size_limit=dify_config.UPLOAD_VIDEO_FILE_SIZE_LIMIT, audio_file_size_limit=dify_config.UPLOAD_AUDIO_FILE_SIZE_LIMIT, skill_file_size_limit=dify_config.UPLOAD_SKILL_FILE_SIZE_LIMIT, workflow_file_upload_limit=dify_config.WORKFLOW_FILE_UPLOAD_LIMIT, image_file_batch_limit=dify_config.IMAGE_FILE_BATCH_LIMIT, single_chunk_attachment_limit=dify_config.SINGLE_CHUNK_ATTACHMENT_LIMIT, attachment_image_file_size_limit=dify_config.ATTACHMENT_IMAGE_FILE_SIZE_LIMIT, ) return dump_response(UploadConfig, config), 200 @setup_required @login_required @account_initialization_required @cloud_edition_billing_resource_check("documents") @console_ns.doc(consumes=["multipart/form-data"], params=FILE_UPLOAD_PARAMS) @console_ns.response(201, "File uploaded successfully", console_ns.models[FileResponse.__name__]) @with_current_user def post(self, current_user: Account) -> JsonResponseWithStatus: upload_file = upload_file_from_request(current_user=current_user) return dump_response(FileResponse, upload_file), 201 @console_ns.route("/files//preview") class FilePreviewApi(Resource): @console_ns.response(200, "Success", console_ns.models[TextContentResponse.__name__]) @console_account_admission() def get(self, request_context: RequestContext, file_id: UUID) -> dict[str, object]: current_tenant_id = request_context.active_workspace_id file_id_str = str(file_id) text = application_services().files.get_file_preview(file_id=file_id_str, tenant_id=current_tenant_id) return dump_response(TextContentResponse, {"content": text}) @console_ns.route("/files/support-type") class FileSupportTypeApi(Resource): @console_ns.response(200, "Success", console_ns.models[AllowedExtensionsResponse.__name__]) @console_account_admission() def get(self, _request_context: RequestContext) -> dict[str, object]: return dump_response( AllowedExtensionsResponse, {"allowed_extensions": list(DOCUMENT_EXTENSIONS)}, )