Переглянути джерело

fix: support other file types for Tool

Yeuoly 7 місяців тому
батько
коміт
01efdee1dd

+ 1 - 0
api/core/tools/entities/tool_entities.py

@@ -157,6 +157,7 @@ class ToolInvokeMessage(BaseModel):
         BLOB = "blob"
         JSON = "json"
         IMAGE_LINK = "image_link"
+        BINARY_LINK = "binary_link"
         VARIABLE = "variable"
         FILE = "file"
 

+ 1 - 1
api/core/tools/utils/message_transformer.py

@@ -85,7 +85,7 @@ class ToolFileMessageTransformer:
                     )
                 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 {},
                     )

+ 6 - 2
api/core/workflow/nodes/tool/tool_node.py

@@ -190,7 +190,11 @@ class ToolNode(BaseNode[ToolNodeData]):
         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 @@ class ToolNode(BaseNode[ToolNodeData]):
 
                 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,
                 }

+ 4 - 0
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