Skip to content

Commit 4b2abf8

Browse files
authored
fix: create_blob_message of tool will always create image type file (#10701)
1 parent 365cb4b commit 4b2abf8

File tree

2 files changed

+16
-10
lines changed

2 files changed

+16
-10
lines changed

api/core/workflow/nodes/tool/tool_node.py

-9
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,4 @@
11
from collections.abc import Mapping, Sequence
2-
from os import path
32
from typing import Any
43

54
from sqlalchemy import select
@@ -180,7 +179,6 @@ def _extract_tool_response_binary(self, tool_response: list[ToolInvokeMessage])
180179
for response in tool_response:
181180
if response.type in {ToolInvokeMessage.MessageType.IMAGE_LINK, ToolInvokeMessage.MessageType.IMAGE}:
182181
url = str(response.message) if response.message else None
183-
ext = path.splitext(url)[1] if url else ".bin"
184182
tool_file_id = str(url).split("/")[-1].split(".")[0]
185183
transfer_method = response.meta.get("transfer_method", FileTransferMethod.TOOL_FILE)
186184

@@ -202,7 +200,6 @@ def _extract_tool_response_binary(self, tool_response: list[ToolInvokeMessage])
202200
)
203201
result.append(file)
204202
elif response.type == ToolInvokeMessage.MessageType.BLOB:
205-
# get tool file id
206203
tool_file_id = str(response.message).split("/")[-1].split(".")[0]
207204
with Session(db.engine) as session:
208205
stmt = select(ToolFile).where(ToolFile.id == tool_file_id)
@@ -211,7 +208,6 @@ def _extract_tool_response_binary(self, tool_response: list[ToolInvokeMessage])
211208
raise ValueError(f"tool file {tool_file_id} not exists")
212209
mapping = {
213210
"tool_file_id": tool_file_id,
214-
"type": FileType.IMAGE,
215211
"transfer_method": FileTransferMethod.TOOL_FILE,
216212
}
217213
file = file_factory.build_from_mapping(
@@ -228,13 +224,8 @@ def _extract_tool_response_binary(self, tool_response: list[ToolInvokeMessage])
228224
tool_file = session.scalar(stmt)
229225
if tool_file is None:
230226
raise ToolFileError(f"Tool file {tool_file_id} does not exist")
231-
if "." in url:
232-
extension = "." + url.split("/")[-1].split(".")[1]
233-
else:
234-
extension = ".bin"
235227
mapping = {
236228
"tool_file_id": tool_file_id,
237-
"type": FileType.IMAGE,
238229
"transfer_method": transfer_method,
239230
"url": url,
240231
}

api/factories/file_factory.py

+16-1
Original file line numberDiff line numberDiff line change
@@ -180,6 +180,20 @@ def _get_remote_file_info(url: str):
180180
return mime_type, filename, file_size
181181

182182

183+
def _get_file_type_by_mimetype(mime_type: str) -> FileType:
184+
if "image" in mime_type:
185+
file_type = FileType.IMAGE
186+
elif "video" in mime_type:
187+
file_type = FileType.VIDEO
188+
elif "audio" in mime_type:
189+
file_type = FileType.AUDIO
190+
elif "text" in mime_type or "pdf" in mime_type:
191+
file_type = FileType.DOCUMENT
192+
else:
193+
file_type = FileType.CUSTOM
194+
return file_type
195+
196+
183197
def _build_from_tool_file(
184198
*,
185199
mapping: Mapping[str, Any],
@@ -199,12 +213,13 @@ def _build_from_tool_file(
199213
raise ValueError(f"ToolFile {mapping.get('tool_file_id')} not found")
200214

201215
extension = "." + tool_file.file_key.split(".")[-1] if "." in tool_file.file_key else ".bin"
216+
file_type = mapping.get("type", _get_file_type_by_mimetype(tool_file.mimetype))
202217

203218
return File(
204219
id=mapping.get("id"),
205220
tenant_id=tenant_id,
206221
filename=tool_file.name,
207-
type=FileType.value_of(mapping.get("type")),
222+
type=file_type,
208223
transfer_method=transfer_method,
209224
remote_url=tool_file.original_url,
210225
related_id=tool_file.id,

0 commit comments

Comments
 (0)