fix(api): streamline message file handling and add support for custom…

… file types
langgenius · Oct 10, 2024 · 1bf8586 · 1bf8586
1 parent e4e03c8
commit 1bf8586
Show file tree

Hide file tree

Showing 2 changed files with 23 additions and 15 deletions.
diff --git a/api/core/tools/tool_engine.py b/api/core/tools/tool_engine.py
@@ -10,6 +10,7 @@
 from core.app.entities.app_invoke_entities import InvokeFrom
 from core.callback_handler.agent_tool_callback_handler import DifyAgentCallbackHandler
 from core.callback_handler.workflow_tool_callback_handler import DifyWorkflowCallbackHandler
+from core.file import FileType
 from core.file.models import FileTransferMethod
 from core.ops.ops_trace_manager import TraceQueueManager
 from core.tools.entities.tool_entities import ToolInvokeMessage, ToolInvokeMessageBinary, ToolInvokeMeta, ToolParameter
@@ -25,6 +26,7 @@
 from core.tools.tool.tool import Tool
 from core.tools.tool.workflow_tool import WorkflowTool
 from core.tools.utils.message_transformer import ToolFileMessageTransformer
+from enums import CreatedByRole
 from extensions.ext_database import db
 from models.model import Message, MessageFile
 
@@ -259,7 +261,10 @@ def _extract_tool_response_binary(tool_response: list[ToolInvokeMessage]) -> lis
 
     @staticmethod
     def _create_message_files(
-        tool_messages: list[ToolInvokeMessageBinary], agent_message: Message, invoke_from: InvokeFrom, user_id: str
+        tool_messages: list[ToolInvokeMessageBinary],
+        agent_message: Message,
+        invoke_from: InvokeFrom,
+        user_id: str,
     ) -> list[tuple[Any, str]]:
         """
         Create message file
@@ -270,29 +275,31 @@ def _create_message_files(
         result = []
 
         for message in tool_messages:
-            file_type = "bin"
             if "image" in message.mimetype:
-                file_type = "image"
+                file_type = FileType.IMAGE
             elif "video" in message.mimetype:
-                file_type = "video"
+                file_type = FileType.VIDEO
             elif "audio" in message.mimetype:
-                file_type = "audio"
-            elif "text" in message.mimetype:
-                file_type = "text"
-            elif "pdf" in message.mimetype:
-                file_type = "pdf"
-            elif "zip" in message.mimetype:
-                file_type = "archive"
-            # ...
+                file_type = FileType.AUDIO
+            elif "text" in message.mimetype or "pdf" in message.mimetype:
+                file_type = FileType.DOCUMENT
+            else:
+                file_type = FileType.CUSTOM
 
+            # extract tool file id from url
+            tool_file_id = message.url.split("/")[-1].split(".")[0]
             message_file = MessageFile(
                 message_id=agent_message.id,
                 type=file_type,
-                transfer_method=FileTransferMethod.TOOL_FILE.value,
+                transfer_method=FileTransferMethod.TOOL_FILE,
                 belongs_to="assistant",
                 url=message.url,
-                upload_file_id=None,
-                created_by_role=("account" if invoke_from in {InvokeFrom.EXPLORE, InvokeFrom.DEBUGGER} else "end_user"),
+                upload_file_id=tool_file_id,
+                created_by_role=(
+                    CreatedByRole.ACCOUNT
+                    if invoke_from in {InvokeFrom.EXPLORE, InvokeFrom.DEBUGGER}
+                    else CreatedByRole.END_USER
+                ),
                 created_by=user_id,
             )
 

diff --git a/api/core/tools/tool_file_manager.py b/api/core/tools/tool_file_manager.py
@@ -82,6 +82,7 @@ def create_file_by_raw(
 
         db.session.add(tool_file)
         db.session.commit()
+        db.session.refresh(tool_file)
 
         return tool_file