Merge branch 'fix/chore-fix' into dev/plugin-deploy

langgenius · Dec 4, 2024 · df68b6b · df68b6b
2 parents 3d3a429 + 01efdee
commit df68b6b
Show file tree

Hide file tree

Showing 4 changed files with 12 additions and 3 deletions.
diff --git a/api/core/tools/entities/tool_entities.py b/api/core/tools/entities/tool_entities.py
@@ -157,6 +157,7 @@ class MessageType(Enum):
         BLOB = "blob"
         JSON = "json"
         IMAGE_LINK = "image_link"
+        BINARY_LINK = "binary_link"
         VARIABLE = "variable"
         FILE = "file"
 

diff --git a/api/core/tools/utils/message_transformer.py b/api/core/tools/utils/message_transformer.py
@@ -85,7 +85,7 @@ def transform_tool_invoke_messages(
                     )
                 else:
                     yield ToolInvokeMessage(
-                        type=ToolInvokeMessage.MessageType.LINK,
+                        type=ToolInvokeMessage.MessageType.BINARY_LINK,
                         message=ToolInvokeMessage.TextMessage(text=url),
                         meta=message.meta.copy() if message.meta is not None else {},
                     )

diff --git a/api/core/workflow/nodes/tool/tool_node.py b/api/core/workflow/nodes/tool/tool_node.py
@@ -190,7 +190,11 @@ def _transform_message(
         variables: dict[str, Any] = {}
 
         for message in message_stream:
-            if message.type in {ToolInvokeMessage.MessageType.IMAGE_LINK, ToolInvokeMessage.MessageType.IMAGE}:
+            if message.type in {
+                ToolInvokeMessage.MessageType.IMAGE_LINK,
+                ToolInvokeMessage.MessageType.BINARY_LINK,
+                ToolInvokeMessage.MessageType.IMAGE,
+            }:
                 assert isinstance(message.message, ToolInvokeMessage.TextMessage)
 
                 url = message.message.text
@@ -209,7 +213,7 @@ def _transform_message(
 
                 mapping = {
                     "tool_file_id": tool_file_id,
-                    "type": FileType.IMAGE,
+                    "type": file_factory.get_file_type_by_mime_type(tool_file.mimetype),
                     "transfer_method": transfer_method,
                     "url": url,
                 }

diff --git a/api/factories/file_factory.py b/api/factories/file_factory.py
@@ -289,3 +289,7 @@ def _get_file_type_by_mimetype(mime_type: str) -> FileType | None:
     else:
         file_type = FileType.CUSTOM
     return file_type
+
+
+def get_file_type_by_mime_type(mime_type: str) -> FileType:
+    return _get_file_type_by_mimetype(mime_type) or FileType.CUSTOM