Underground-Digital/Workflow-Engine
0
1import json2from collections.abc import Mapping3from copy import deepcopy4from datetime import datetime, timezone5from mimetypes import guess_type6from typing import Any, Optional, Union7 8from yarl import URL9 10from core.app.entities.app_invoke_entities import InvokeFrom11from core.callback_handler.agent_tool_callback_handler import DifyAgentCallbackHandler12from core.callback_handler.workflow_tool_callback_handler import DifyWorkflowCallbackHandler13from core.file import FileType14from core.file.models import FileTransferMethod15from core.ops.ops_trace_manager import TraceQueueManager16from core.tools.entities.tool_entities import ToolInvokeMessage, ToolInvokeMessageBinary, ToolInvokeMeta, ToolParameter17from core.tools.errors import (18 ToolEngineInvokeError,19 ToolInvokeError,20 ToolNotFoundError,21 ToolNotSupportedError,22 ToolParameterValidationError,23 ToolProviderCredentialValidationError,24 ToolProviderNotFoundError,25)26from core.tools.tool.tool import Tool27from core.tools.tool.workflow_tool import WorkflowTool28from core.tools.utils.message_transformer import ToolFileMessageTransformer29from extensions.ext_database import db30from models.enums import CreatedByRole31from models.model import Message, MessageFile32 33 34class ToolEngine:35 """36 Tool runtime engine take care of the tool executions.37 """38 39 @staticmethod40 def agent_invoke(41 tool: Tool,42 tool_parameters: Union[str, dict],43 user_id: str,44 tenant_id: str,45 message: Message,46 invoke_from: InvokeFrom,47 agent_tool_callback: DifyAgentCallbackHandler,48 trace_manager: Optional[TraceQueueManager] = None,49 ) -> tuple[str, list[tuple[MessageFile, bool]], ToolInvokeMeta]:50 """51 Agent invokes the tool with the given arguments.52 """53 # check if arguments is a string54 if isinstance(tool_parameters, str):55 # check if this tool has only one parameter56 parameters = [57 parameter58 for parameter in tool.get_runtime_parameters() or []59 if parameter.form == ToolParameter.ToolParameterForm.LLM60 ]61 if parameters and len(parameters) == 1:62 tool_parameters = {parameters[0].name: tool_parameters}63 else:64 raise ValueError(f"tool_parameters should be a dict, but got a string: {tool_parameters}")65 66 # invoke the tool67 try:68 # hit the callback handler69 agent_tool_callback.on_tool_start(tool_name=tool.identity.name, tool_inputs=tool_parameters)70 71 meta, response = ToolEngine._invoke(tool, tool_parameters, user_id)72 response = ToolFileMessageTransformer.transform_tool_invoke_messages(73 messages=response, user_id=user_id, tenant_id=tenant_id, conversation_id=message.conversation_id74 )75 76 # extract binary data from tool invoke message77 binary_files = ToolEngine._extract_tool_response_binary(response)78 # create message file79 message_files = ToolEngine._create_message_files(80 tool_messages=binary_files, agent_message=message, invoke_from=invoke_from, user_id=user_id81 )82 83 plain_text = ToolEngine._convert_tool_response_to_str(response)84 85 # hit the callback handler86 agent_tool_callback.on_tool_end(87 tool_name=tool.identity.name,88 tool_inputs=tool_parameters,89 tool_outputs=plain_text,90 message_id=message.id,91 trace_manager=trace_manager,92 )93 94 # transform tool invoke message to get LLM friendly message95 return plain_text, message_files, meta96 except ToolProviderCredentialValidationError as e:97 error_response = "Please check your tool provider credentials"98 agent_tool_callback.on_tool_error(e)99 except (ToolNotFoundError, ToolNotSupportedError, ToolProviderNotFoundError) as e:100 error_response = f"there is not a tool named {tool.identity.name}"101 agent_tool_callback.on_tool_error(e)102 except ToolParameterValidationError as e:103 error_response = f"tool parameters validation error: {e}, please check your tool parameters"104 agent_tool_callback.on_tool_error(e)105 except ToolInvokeError as e:106 error_response = f"tool invoke error: {e}"107 agent_tool_callback.on_tool_error(e)108 except ToolEngineInvokeError as e:109 meta = e.args[0]110 error_response = f"tool invoke error: {meta.error}"111 agent_tool_callback.on_tool_error(e)112 return error_response, [], meta113 except Exception as e:114 error_response = f"unknown error: {e}"115 agent_tool_callback.on_tool_error(e)116 117 return error_response, [], ToolInvokeMeta.error_instance(error_response)118 119 @staticmethod120 def workflow_invoke(121 tool: Tool,122 tool_parameters: Mapping[str, Any],123 user_id: str,124 workflow_tool_callback: DifyWorkflowCallbackHandler,125 workflow_call_depth: int,126 thread_pool_id: Optional[str] = None,127 ) -> list[ToolInvokeMessage]:128 """129 Workflow invokes the tool with the given arguments.130 """131 try:132 # hit the callback handler133 assert tool.identity is not None134 workflow_tool_callback.on_tool_start(tool_name=tool.identity.name, tool_inputs=tool_parameters)135 136 if isinstance(tool, WorkflowTool):137 tool.workflow_call_depth = workflow_call_depth + 1138 tool.thread_pool_id = thread_pool_id139 140 if tool.runtime and tool.runtime.runtime_parameters:141 tool_parameters = {**tool.runtime.runtime_parameters, **tool_parameters}142 response = tool.invoke(user_id=user_id, tool_parameters=tool_parameters)143 144 # hit the callback handler145 workflow_tool_callback.on_tool_end(146 tool_name=tool.identity.name,147 tool_inputs=tool_parameters,148 tool_outputs=response,149 )150 151 return response152 except Exception as e:153 workflow_tool_callback.on_tool_error(e)154 raise e155 156 @staticmethod157 def _invoke(tool: Tool, tool_parameters: dict, user_id: str) -> tuple[ToolInvokeMeta, list[ToolInvokeMessage]]:158 """159 Invoke the tool with the given arguments.160 """161 started_at = datetime.now(timezone.utc)162 meta = ToolInvokeMeta(163 time_cost=0.0,164 error=None,165 tool_config={166 "tool_name": tool.identity.name,167 "tool_provider": tool.identity.provider,168 "tool_provider_type": tool.tool_provider_type().value,169 "tool_parameters": deepcopy(tool.runtime.runtime_parameters),170 "tool_icon": tool.identity.icon,171 },172 )173 try:174 response = tool.invoke(user_id, tool_parameters)175 except Exception as e:176 meta.error = str(e)177 raise ToolEngineInvokeError(meta)178 finally:179 ended_at = datetime.now(timezone.utc)180 meta.time_cost = (ended_at - started_at).total_seconds()181 182 return meta, response183 184 @staticmethod185 def _convert_tool_response_to_str(tool_response: list[ToolInvokeMessage]) -> str:186 """187 Handle tool response188 """189 result = ""190 for response in tool_response:191 if response.type == ToolInvokeMessage.MessageType.TEXT:192 result += response.message193 elif response.type == ToolInvokeMessage.MessageType.LINK:194 result += f"result link: {response.message}. please tell user to check it."195 elif response.type in {ToolInvokeMessage.MessageType.IMAGE_LINK, ToolInvokeMessage.MessageType.IMAGE}:196 result += (197 "image has been created and sent to user already, you do not need to create it,"198 " just tell the user to check it now."199 )200 elif response.type == ToolInvokeMessage.MessageType.JSON:201 result += f"tool response: {json.dumps(response.message, ensure_ascii=False)}."202 else:203 result += f"tool response: {response.message}."204 205 return result206 207 @staticmethod208 def _extract_tool_response_binary(tool_response: list[ToolInvokeMessage]) -> list[ToolInvokeMessageBinary]:209 """210 Extract tool response binary211 """212 result = []213 214 for response in tool_response:215 if response.type in {ToolInvokeMessage.MessageType.IMAGE_LINK, ToolInvokeMessage.MessageType.IMAGE}:216 mimetype = None217 if response.meta.get("mime_type"):218 mimetype = response.meta.get("mime_type")219 else:220 try:221 url = URL(response.message)222 extension = url.suffix223 guess_type_result, _ = guess_type(f"a{extension}")224 if guess_type_result:225 mimetype = guess_type_result226 except Exception:227 pass228 229 if not mimetype:230 mimetype = "image/jpeg"231 232 result.append(233 ToolInvokeMessageBinary(234 mimetype=response.meta.get("mime_type", "image/jpeg"),235 url=response.message,236 save_as=response.save_as,237 )238 )239 elif response.type == ToolInvokeMessage.MessageType.BLOB:240 result.append(241 ToolInvokeMessageBinary(242 mimetype=response.meta.get("mime_type", "octet/stream"),243 url=response.message,244 save_as=response.save_as,245 )246 )247 elif response.type == ToolInvokeMessage.MessageType.LINK:248 # check if there is a mime type in meta249 if response.meta and "mime_type" in response.meta:250 result.append(251 ToolInvokeMessageBinary(252 mimetype=response.meta.get("mime_type", "octet/stream")253 if response.meta254 else "octet/stream",255 url=response.message,256 save_as=response.save_as,257 )258 )259 260 return result261 262 @staticmethod263 def _create_message_files(264 tool_messages: list[ToolInvokeMessageBinary],265 agent_message: Message,266 invoke_from: InvokeFrom,267 user_id: str,268 ) -> list[tuple[Any, str]]:269 """270 Create message file271 272 :param messages: messages273 :return: message files, should save as variable274 """275 result = []276 277 for message in tool_messages:278 if "image" in message.mimetype:279 file_type = FileType.IMAGE280 elif "video" in message.mimetype:281 file_type = FileType.VIDEO282 elif "audio" in message.mimetype:283 file_type = FileType.AUDIO284 elif "text" in message.mimetype or "pdf" in message.mimetype:285 file_type = FileType.DOCUMENT286 else:287 file_type = FileType.CUSTOM288 289 # extract tool file id from url290 tool_file_id = message.url.split("/")[-1].split(".")[0]291 message_file = MessageFile(292 message_id=agent_message.id,293 type=file_type,294 transfer_method=FileTransferMethod.TOOL_FILE,295 belongs_to="assistant",296 url=message.url,297 upload_file_id=tool_file_id,298 created_by_role=(299 CreatedByRole.ACCOUNT300 if invoke_from in {InvokeFrom.EXPLORE, InvokeFrom.DEBUGGER}301 else CreatedByRole.END_USER302 ),303 created_by=user_id,304 )305 306 db.session.add(message_file)307 db.session.commit()308 db.session.refresh(message_file)309 310 result.append((message_file.id, message.save_as))311 312 db.session.close()313 314 return result315