import json from collections.abc import Generator, Iterable from copy import deepcopy from datetime import datetime, timezone from mimetypes import guess_type from typing import Any, Optional, Union, cast from yarl import URL from core.app.entities.app_invoke_entities import InvokeFrom from core.callback_handler.agent_tool_callback_handler import DifyAgentCallbackHandler from core.callback_handler.workflow_tool_callback_handler import DifyWorkflowCallbackHandler from core.file.file_obj import FileTransferMethod from core.ops.ops_trace_manager import TraceQueueManager from core.tools.__base.tool import Tool from core.tools.entities.tool_entities import ToolInvokeMessage, ToolInvokeMessageBinary, ToolInvokeMeta, ToolParameter from core.tools.errors import ( ToolEngineInvokeError, ToolInvokeError, ToolNotFoundError, ToolNotSupportedError, ToolParameterValidationError, ToolProviderCredentialValidationError, ToolProviderNotFoundError, ) from core.tools.utils.message_transformer import ToolFileMessageTransformer from core.tools.workflow_as_tool.tool import WorkflowTool from extensions.ext_database import db from models.model import Message, MessageFile class ToolEngine: """ Tool runtime engine take care of the tool executions. """ @staticmethod def agent_invoke( tool: Tool, tool_parameters: Union[str, dict], user_id: str, tenant_id: str, message: Message, invoke_from: InvokeFrom, agent_tool_callback: DifyAgentCallbackHandler, trace_manager: Optional[TraceQueueManager] = None, ) -> tuple[str, list[tuple[MessageFile, str]], ToolInvokeMeta]: """ Agent invokes the tool with the given arguments. """ # check if arguments is a string if isinstance(tool_parameters, str): # check if this tool has only one parameter parameters = [ parameter for parameter in tool.get_runtime_parameters() or [] if parameter.form == ToolParameter.ToolParameterForm.LLM ] if parameters and len(parameters) == 1: tool_parameters = {parameters[0].name: tool_parameters} else: raise ValueError(f"tool_parameters should be a dict, but got a string: {tool_parameters}") # invoke the tool try: # hit the callback handler agent_tool_callback.on_tool_start(tool_name=tool.entity.identity.name, tool_inputs=tool_parameters) messages = ToolEngine._invoke(tool, tool_parameters, user_id) invocation_meta_dict: dict[str, ToolInvokeMeta] = {} def message_callback( invocation_meta_dict: dict, messages: Generator[ToolInvokeMessage | ToolInvokeMeta, None, None] ): for message in messages: if isinstance(message, ToolInvokeMeta): invocation_meta_dict["meta"] = message else: yield message messages = ToolFileMessageTransformer.transform_tool_invoke_messages( messages=message_callback(invocation_meta_dict, messages), user_id=user_id, tenant_id=tenant_id, conversation_id=message.conversation_id, ) # extract binary data from tool invoke message binary_files = ToolEngine._extract_tool_response_binary(messages) # create message file message_files = ToolEngine._create_message_files( tool_messages=binary_files, agent_message=message, invoke_from=invoke_from, user_id=user_id ) plain_text = ToolEngine._convert_tool_response_to_str(messages) meta = invocation_meta_dict["meta"] # hit the callback handler agent_tool_callback.on_tool_end( tool_name=tool.entity.identity.name, tool_inputs=tool_parameters, tool_outputs=plain_text, message_id=message.id, trace_manager=trace_manager, ) # transform tool invoke message to get LLM friendly message return plain_text, message_files, meta except ToolProviderCredentialValidationError as e: error_response = "Please check your tool provider credentials" agent_tool_callback.on_tool_error(e) except (ToolNotFoundError, ToolNotSupportedError, ToolProviderNotFoundError) as e: error_response = f"there is not a tool named {tool.entity.identity.name}" agent_tool_callback.on_tool_error(e) except ToolParameterValidationError as e: error_response = f"tool parameters validation error: {e}, please check your tool parameters" agent_tool_callback.on_tool_error(e) except ToolInvokeError as e: error_response = f"tool invoke error: {e}" agent_tool_callback.on_tool_error(e) except ToolEngineInvokeError as e: meta = e.args[0] error_response = f"tool invoke error: {meta.error}" agent_tool_callback.on_tool_error(e) return error_response, [], meta except Exception as e: error_response = f"unknown error: {e}" agent_tool_callback.on_tool_error(e) return error_response, [], ToolInvokeMeta.error_instance(error_response) @staticmethod def workflow_invoke( tool: Tool, tool_parameters: dict[str, Any], user_id: str, workflow_tool_callback: DifyWorkflowCallbackHandler, workflow_call_depth: int, thread_pool_id: Optional[str] = None, ) -> Generator[ToolInvokeMessage, None, None]: """ Workflow invokes the tool with the given arguments. """ try: # hit the callback handler workflow_tool_callback.on_tool_start(tool_name=tool.entity.identity.name, tool_inputs=tool_parameters) if isinstance(tool, WorkflowTool): tool.workflow_call_depth = workflow_call_depth + 1 tool.thread_pool_id = thread_pool_id if tool.runtime and tool.runtime.runtime_parameters: tool_parameters = {**tool.runtime.runtime_parameters, **tool_parameters} response = tool.invoke(user_id=user_id, tool_parameters=tool_parameters) # hit the callback handler response = workflow_tool_callback.on_tool_execution( tool_name=tool.entity.identity.name, tool_inputs=tool_parameters, tool_outputs=response, ) return response except Exception as e: workflow_tool_callback.on_tool_error(e) raise e @staticmethod def _invoke( tool: Tool, tool_parameters: dict, user_id: str ) -> Generator[ToolInvokeMessage | ToolInvokeMeta, None, None]: """ Invoke the tool with the given arguments. """ if not tool.runtime: raise ValueError("missing runtime in tool") started_at = datetime.now(timezone.utc) meta = ToolInvokeMeta( time_cost=0.0, error=None, tool_config={ "tool_name": tool.entity.identity.name, "tool_provider": tool.entity.identity.provider, "tool_provider_type": tool.tool_provider_type().value, "tool_parameters": deepcopy(tool.runtime.runtime_parameters), "tool_icon": tool.entity.identity.icon, }, ) try: yield from tool.invoke(user_id, tool_parameters) except Exception as e: meta.error = str(e) raise ToolEngineInvokeError(meta) finally: ended_at = datetime.now(timezone.utc) meta.time_cost = (ended_at - started_at).total_seconds() yield meta @staticmethod def _convert_tool_response_to_str(tool_response: Generator[ToolInvokeMessage, None, None]) -> str: """ Handle tool response """ result = "" for response in tool_response: if response.type == ToolInvokeMessage.MessageType.TEXT: result += cast(ToolInvokeMessage.TextMessage, response.message).text elif response.type == ToolInvokeMessage.MessageType.LINK: result += ( f"result link: {cast(ToolInvokeMessage.TextMessage, response.message).text}." + " please tell user to check it." ) elif response.type in {ToolInvokeMessage.MessageType.IMAGE_LINK, ToolInvokeMessage.MessageType.IMAGE}: result += ( "image has been created and sent to user already, " + "you do not need to create it, just tell the user to check it now." ) elif response.type == ToolInvokeMessage.MessageType.JSON: text = json.dumps(cast(ToolInvokeMessage.JsonMessage, response.message).json_object, ensure_ascii=False) result += f"tool response: {text}." else: result += f"tool response: {response.message}." return result @staticmethod def _extract_tool_response_binary( tool_response: Generator[ToolInvokeMessage, None, None], ) -> Generator[ToolInvokeMessageBinary, None, None]: """ Extract tool response binary """ for response in tool_response: if response.type in {ToolInvokeMessage.MessageType.IMAGE_LINK, ToolInvokeMessage.MessageType.IMAGE}: mimetype = None if not response.meta: raise ValueError("missing meta data") if response.meta.get("mime_type"): mimetype = response.meta.get("mime_type") else: try: url = URL(cast(ToolInvokeMessage.TextMessage, response.message).text) extension = url.suffix guess_type_result, _ = guess_type(f"a{extension}") if guess_type_result: mimetype = guess_type_result except Exception: pass if not mimetype: mimetype = "image/jpeg" yield ToolInvokeMessageBinary( mimetype=response.meta.get("mime_type", "image/jpeg"), url=cast(ToolInvokeMessage.TextMessage, response.message).text, save_as=response.save_as, ) elif response.type == ToolInvokeMessage.MessageType.BLOB: if not response.meta: raise ValueError("missing meta data") yield ToolInvokeMessageBinary( mimetype=response.meta.get("mime_type", "octet/stream"), url=cast(ToolInvokeMessage.TextMessage, response.message).text, save_as=response.save_as, ) elif response.type == ToolInvokeMessage.MessageType.LINK: # check if there is a mime type in meta if response.meta and "mime_type" in response.meta: yield ToolInvokeMessageBinary( mimetype=response.meta.get("mime_type", "octet/stream") if response.meta else "octet/stream", url=cast(ToolInvokeMessage.TextMessage, response.message).text, save_as=response.save_as, ) @staticmethod def _create_message_files( tool_messages: Iterable[ToolInvokeMessageBinary], agent_message: Message, invoke_from: InvokeFrom, user_id: str ) -> list[tuple[MessageFile, str]]: """ Create message file :param messages: messages :return: message files, should save as variable """ result = [] for message in tool_messages: file_type = "bin" if "image" in message.mimetype: file_type = "image" elif "video" in message.mimetype: file_type = "video" elif "audio" in message.mimetype: file_type = "audio" elif "text" in message.mimetype: file_type = "text" elif "pdf" in message.mimetype: file_type = "pdf" elif "zip" in message.mimetype: file_type = "archive" # ... message_file = MessageFile( message_id=agent_message.id, type=file_type, transfer_method=FileTransferMethod.TOOL_FILE.value, belongs_to="assistant", url=message.url, upload_file_id=None, created_by_role=("account" if invoke_from in {InvokeFrom.EXPLORE, InvokeFrom.DEBUGGER} else "end_user"), created_by=user_id, ) db.session.add(message_file) db.session.commit() db.session.refresh(message_file) result.append((message_file.id, message.save_as)) db.session.close() return result