Underground-Digital/Workflow-Engine

1import base642 3from configs import dify_config4from core.file import file_repository5from core.helper import ssrf_proxy6from core.model_runtime.entities import AudioPromptMessageContent, ImagePromptMessageContent7from extensions.ext_database import db8from extensions.ext_storage import storage9 10from . import helpers11from .enums import FileAttribute12from .models import File, FileTransferMethod, FileType13from .tool_file_parser import ToolFileParser14 15 16def get_attr(*, file: File, attr: FileAttribute):17    match attr:18        case FileAttribute.TYPE:19            return file.type.value20        case FileAttribute.SIZE:21            return file.size22        case FileAttribute.NAME:23            return file.filename24        case FileAttribute.MIME_TYPE:25            return file.mime_type26        case FileAttribute.TRANSFER_METHOD:27            return file.transfer_method.value28        case FileAttribute.URL:29            return file.remote_url30        case FileAttribute.EXTENSION:31            return file.extension32        case _:33            raise ValueError(f"Invalid file attribute: {attr}")34 35 36def to_prompt_message_content(f: File, /):37    """38    Convert a File object to an ImagePromptMessageContent object.39 40    This function takes a File object and converts it to an ImagePromptMessageContent41    object, which can be used as a prompt for image-based AI models.42 43    Args:44        file (File): The File object to convert. Must be of type FileType.IMAGE.45 46    Returns:47        ImagePromptMessageContent: An object containing the image data and detail level.48 49    Raises:50        ValueError: If the file is not an image or if the file data is missing.51 52    Note:53        The detail level of the image prompt is determined by the file's extra_config.54        If not specified, it defaults to ImagePromptMessageContent.DETAIL.LOW.55    """56    match f.type:57        case FileType.IMAGE:58            if dify_config.MULTIMODAL_SEND_IMAGE_FORMAT == "url":59                data = _to_url(f)60            else:61                data = _to_base64_data_string(f)62 63            if f._extra_config and f._extra_config.image_config and f._extra_config.image_config.detail:64                detail = f._extra_config.image_config.detail65            else:66                detail = ImagePromptMessageContent.DETAIL.LOW67 68            return ImagePromptMessageContent(data=data, detail=detail)69        case FileType.AUDIO:70            encoded_string = _file_to_encoded_string(f)71            if f.extension is None:72                raise ValueError("Missing file extension")73            return AudioPromptMessageContent(data=encoded_string, format=f.extension.lstrip("."))74        case _:75            raise ValueError(f"file type {f.type} is not supported")76 77 78def download(f: File, /):79    if f.transfer_method == FileTransferMethod.TOOL_FILE:80        tool_file = file_repository.get_tool_file(session=db.session(), file=f)81        return _download_file_content(tool_file.file_key)82    elif f.transfer_method == FileTransferMethod.LOCAL_FILE:83        upload_file = file_repository.get_upload_file(session=db.session(), file=f)84        return _download_file_content(upload_file.key)85    # remote file86    response = ssrf_proxy.get(f.remote_url, follow_redirects=True)87    response.raise_for_status()88    return response.content89 90 91def _download_file_content(path: str, /):92    """93    Download and return the contents of a file as bytes.94 95    This function loads the file from storage and ensures it's in bytes format.96 97    Args:98        path (str): The path to the file in storage.99 100    Returns:101        bytes: The contents of the file as a bytes object.102 103    Raises:104        ValueError: If the loaded file is not a bytes object.105    """106    data = storage.load(path, stream=False)107    if not isinstance(data, bytes):108        raise ValueError(f"file {path} is not a bytes object")109    return data110 111 112def _get_encoded_string(f: File, /):113    match f.transfer_method:114        case FileTransferMethod.REMOTE_URL:115            response = ssrf_proxy.get(f.remote_url)116            response.raise_for_status()117            content = response.content118            encoded_string = base64.b64encode(content).decode("utf-8")119            return encoded_string120        case FileTransferMethod.LOCAL_FILE:121            upload_file = file_repository.get_upload_file(session=db.session(), file=f)122            data = _download_file_content(upload_file.key)123            encoded_string = base64.b64encode(data).decode("utf-8")124            return encoded_string125        case FileTransferMethod.TOOL_FILE:126            tool_file = file_repository.get_tool_file(session=db.session(), file=f)127            data = _download_file_content(tool_file.file_key)128            encoded_string = base64.b64encode(data).decode("utf-8")129            return encoded_string130        case _:131            raise ValueError(f"Unsupported transfer method: {f.transfer_method}")132 133 134def _to_base64_data_string(f: File, /):135    encoded_string = _get_encoded_string(f)136    return f"data:{f.mime_type};base64,{encoded_string}"137 138 139def _file_to_encoded_string(f: File, /):140    match f.type:141        case FileType.IMAGE:142            return _to_base64_data_string(f)143        case FileType.AUDIO:144            return _get_encoded_string(f)145        case _:146            raise ValueError(f"file type {f.type} is not supported")147 148 149def _to_url(f: File, /):150    if f.transfer_method == FileTransferMethod.REMOTE_URL:151        if f.remote_url is None:152            raise ValueError("Missing file remote_url")153        return f.remote_url154    elif f.transfer_method == FileTransferMethod.LOCAL_FILE:155        if f.related_id is None:156            raise ValueError("Missing file related_id")157        return helpers.get_signed_file_url(upload_file_id=f.related_id)158    elif f.transfer_method == FileTransferMethod.TOOL_FILE:159        # add sign url160        if f.related_id is None or f.extension is None:161            raise ValueError("Missing file related_id or extension")162        return ToolFileParser.get_tool_file_manager().sign_file(tool_file_id=f.related_id, extension=f.extension)163    else:164        raise ValueError(f"Unsupported transfer method: {f.transfer_method}")165