diff --git a/api/core/tools/tool_file_manager.py b/api/core/tools/tool_file_manager.py index 360dfb9b4c..1573e0c219 100644 --- a/api/core/tools/tool_file_manager.py +++ b/api/core/tools/tool_file_manager.py @@ -63,11 +63,18 @@ class ToolFileManager: conversation_id: Optional[str], file_binary: bytes, mimetype: str, + filename: Optional[str] = None, ) -> ToolFile: extension = guess_extension(mimetype) or ".bin" unique_name = uuid4().hex - filename = f"{unique_name}{extension}" - filepath = f"tools/{tenant_id}/{filename}" + unique_filename = f"{unique_name}{extension}" + # default just as before + present_filename = unique_filename + if filename is not None: + has_extension = len(filename.split(".")) > 1 + # Add extension flexibly + present_filename = filename if has_extension else f"{filename}{extension}" + filepath = f"tools/{tenant_id}/{unique_filename}" storage.save(filepath, file_binary) tool_file = ToolFile( @@ -76,7 +83,7 @@ class ToolFileManager: conversation_id=conversation_id, file_key=filepath, mimetype=mimetype, - name=filename, + name=present_filename, size=len(file_binary), ) diff --git a/api/core/tools/utils/message_transformer.py b/api/core/tools/utils/message_transformer.py index ec0b3a5863..6fd0c201e3 100644 --- a/api/core/tools/utils/message_transformer.py +++ b/api/core/tools/utils/message_transformer.py @@ -59,6 +59,8 @@ class ToolFileMessageTransformer: meta = message.meta or {} mimetype = meta.get("mime_type", "application/octet-stream") + # get filename from meta + filename = meta.get("file_name", None) # if message is str, encode it to bytes if not isinstance(message.message, ToolInvokeMessage.BlobMessage): @@ -72,6 +74,7 @@ class ToolFileMessageTransformer: conversation_id=conversation_id, file_binary=message.message.blob, mimetype=mimetype, + filename=filename, ) url = cls.get_tool_file_url(tool_file_id=file.id, extension=guess_extension(file.mimetype))