22from os import path
33from typing import Any
44
5+ from sqlalchemy import select
6+ from sqlalchemy .orm import Session
7+
58from core .callback_handler .workflow_tool_callback_handler import DifyWorkflowCallbackHandler
69from core .file .models import File , FileTransferMethod , FileType
710from core .tools .entities .tool_entities import ToolInvokeMessage , ToolParameter
1417from core .workflow .nodes .tool .entities import ToolNodeData
1518from core .workflow .utils .variable_template_parser import VariableTemplateParser
1619from enums import NodeType
20+ from extensions .ext_database import db
21+ from models import ToolFile
1722from models .workflow import WorkflowNodeExecutionStatus
1823
1924
@@ -167,45 +172,59 @@ def _extract_tool_response_binary(self, tool_response: list[ToolInvokeMessage])
167172 result = []
168173 for response in tool_response :
169174 if response .type in {ToolInvokeMessage .MessageType .IMAGE_LINK , ToolInvokeMessage .MessageType .IMAGE }:
170- url = response .message
171- ext = path .splitext (url )[1 ]
172- mimetype = response .meta .get ("mime_type" , "image/jpeg" )
173- tool_file_id = response .save_as or url .split ("/" )[- 1 ]
175+ url = str (response .message ) if response .message else None
176+ ext = path .splitext (url )[1 ] if url else ".bin"
177+ tool_file_id = response .save_as or str (url ).split ("/" )[- 1 ].split ("." )[0 ]
174178 transfer_method = response .meta .get ("transfer_method" , FileTransferMethod .TOOL_FILE )
175179
176- # get tool file id
177- tool_file_id = str (url ).split ("/" )[- 1 ].split ("." )[0 ]
180+ with Session (db .engine ) as session :
181+ stmt = select (ToolFile ).where (ToolFile .id == tool_file_id )
182+ tool_file = session .scalar (stmt )
183+ if tool_file is None :
184+ raise ValueError (f"tool file { tool_file_id } not exists" )
185+
178186 result .append (
179187 File (
180188 tenant_id = self .tenant_id ,
181189 type = FileType .IMAGE ,
182190 transfer_method = transfer_method ,
183191 remote_url = url ,
184- related_id = tool_file_id ,
185- filename = tool_file_id ,
192+ related_id = tool_file . id ,
193+ filename = tool_file . name ,
186194 extension = ext ,
187- mime_type = mimetype ,
195+ mime_type = tool_file .mimetype ,
196+ size = tool_file .size ,
188197 )
189198 )
190199 elif response .type == ToolInvokeMessage .MessageType .BLOB :
191200 # get tool file id
192201 tool_file_id = str (response .message ).split ("/" )[- 1 ].split ("." )[0 ]
202+ with Session (db .engine ) as session :
203+ stmt = select (ToolFile ).where (ToolFile .id == tool_file_id )
204+ tool_file = session .scalar (stmt )
205+ if tool_file is None :
206+ raise ValueError (f"tool file { tool_file_id } not exists" )
193207 result .append (
194208 File (
195209 tenant_id = self .tenant_id ,
196210 type = FileType .IMAGE ,
197211 transfer_method = FileTransferMethod .TOOL_FILE ,
198- related_id = tool_file_id ,
199- filename = response . save_as ,
212+ related_id = tool_file . id ,
213+ filename = tool_file . name ,
200214 extension = path .splitext (response .save_as )[1 ],
201- mime_type = response .meta .get ("mime_type" , "application/octet-stream" ),
215+ mime_type = tool_file .mimetype ,
216+ size = tool_file .size ,
202217 )
203218 )
204219 elif response .type == ToolInvokeMessage .MessageType .LINK :
205220 url = str (response .message )
206221 transfer_method = FileTransferMethod .TOOL_FILE
207- mimetype = response .meta .get ("mime_type" , "application/octet-stream" )
208222 tool_file_id = url .split ("/" )[- 1 ].split ("." )[0 ]
223+ with Session (db .engine ) as session :
224+ stmt = select (ToolFile ).where (ToolFile .id == tool_file_id )
225+ tool_file = session .scalar (stmt )
226+ if tool_file is None :
227+ raise ValueError (f"tool file { tool_file_id } not exists" )
209228 if "." in url :
210229 extension = "." + url .split ("/" )[- 1 ].split ("." )[1 ]
211230 else :
@@ -215,10 +234,11 @@ def _extract_tool_response_binary(self, tool_response: list[ToolInvokeMessage])
215234 type = FileType (response .save_as ),
216235 transfer_method = transfer_method ,
217236 remote_url = url ,
218- filename = tool_file_id ,
219- related_id = tool_file_id ,
237+ filename = tool_file . name ,
238+ related_id = tool_file . id ,
220239 extension = extension ,
221- mime_type = mimetype ,
240+ mime_type = tool_file .mimetype ,
241+ size = tool_file .size ,
222242 )
223243 result .append (file )
224244
0 commit comments