AI/dify
1
0
Fork 0
mirror of https://git.mirrors.martin98.com/https://github.com/langgenius/dify.git synced 2025-08-18 22:55:58 +08:00
Code Issues Packages Projects Releases Wiki Activity

fix: support other file types for Tool

Browse Source
This commit is contained in:
Yeuoly 2024-12-04 19:26:01 +08:00
parent 0af9c4fd9d
commit 01efdee1dd
No known key found for this signature in database
GPG Key ID: A66E7E320FB19F61
4 changed files with 12 additions and 3 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

1
api/core/tools/entities/tool_entities.py
View File

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

2
api/core/tools/utils/message_transformer.py
View File

@ -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 {},
)

8
api/core/workflow/nodes/tool/tool_node.py
View File

@ -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
api/factories/file_factory.py
View File

@ -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