Parcourir la source

support config filename in meta for create_blob_message (#15605)

Co-authored-by: StoneFancyX <kindbin@qq.com>
Co-authored-by: crazywoola <427733928@qq.com>
StoneFancyX il y a 1 mois
Parent
commit
2524f16525

+ 10 - 3
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),
         )
 

+ 3 - 0
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))