diff --git a/api/core/datasource/utils/message_transformer.py b/api/core/datasource/utils/message_transformer.py index 9bc57235d8..12d2a71d69 100644 --- a/api/core/datasource/utils/message_transformer.py +++ b/api/core/datasource/utils/message_transformer.py @@ -1,6 +1,6 @@ import logging from collections.abc import Generator -from mimetypes import guess_extension +from mimetypes import guess_extension, guess_type from typing import Optional from core.datasource.datasource_file_manager import DatasourceFileManager @@ -57,10 +57,13 @@ class DatasourceFileMessageTransformer: elif message.type == DatasourceMessage.MessageType.BLOB: # get mime type and save blob to storage meta = message.meta or {} - - mimetype = meta.get("mime_type", "application/octet-stream") # get filename from meta filename = meta.get("file_name", None) + + mimetype = meta.get("mime_type") + if not mimetype: + mimetype = guess_type(filename)[0] or "application/octet-stream" + # if message is str, encode it to bytes if not isinstance(message.message, DatasourceMessage.BlobMessage): diff --git a/api/core/file/file_manager.py b/api/core/file/file_manager.py index ada19ef8ce..ee58840c84 100644 --- a/api/core/file/file_manager.py +++ b/api/core/file/file_manager.py @@ -72,7 +72,9 @@ def to_prompt_message_content( def download(f: File, /): - if f.transfer_method in (FileTransferMethod.TOOL_FILE, FileTransferMethod.LOCAL_FILE): + if f.transfer_method in (FileTransferMethod.TOOL_FILE, + FileTransferMethod.LOCAL_FILE, + FileTransferMethod.DATASOURCE_FILE): return _download_file_content(f._storage_key) elif f.transfer_method == FileTransferMethod.REMOTE_URL: response = ssrf_proxy.get(f.remote_url, follow_redirects=True)