import contextlib import json import logging from collections.abc import Generator, Iterable from copy import deepcopy from dataclasses import dataclass from datetime import UTC, datetime from mimetypes import guess_type from typing import Any, Union, cast from sqlalchemy.orm import Session, sessionmaker 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.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.entities.ui_entities import ( DIFY_UI_JSON_ENVELOPE_KEY, ToolUIMessage, extract_ui_message_from_json, validate_tool_ui_message_batch, ) from core.tools.errors import ( ToolEngineInvokeError, ToolInvokeError, ToolNotFoundError, ToolNotSupportedError, ToolParameterValidationError, ToolProviderCredentialValidationError, ToolProviderNotFoundError, ) from core.tools.utils.message_transformer import ToolFileMessageTransformer, safe_json_value from core.tools.workflow_as_tool.tool import WorkflowTool from extensions.ext_database import db from graphon.file import FileTransferMethod, FileType from models.enums import CreatorUserRole, MessageFileBelongsTo from models.model import Message, MessageFile logger = logging.getLogger(__name__) @dataclass(frozen=True) class ToolAgentInvokeResult: """Agent-facing tool result with UI kept outside the model observation.""" observation: str message_files: list[str] ui_messages: list[ToolUIMessage] meta: ToolInvokeMeta class ToolEngine: """ Tool runtime engine take care of the tool executions. """ @staticmethod def agent_invoke( session: Session, tool: Tool, tool_parameters: Union[str, dict[str, Any]], user_id: str, tenant_id: str, message: Message, invoke_from: InvokeFrom, agent_tool_callback: DifyAgentCallbackHandler, trace_manager: TraceQueueManager | None = None, conversation_id: str | None = None, app_id: str | None = None, message_id: str | None = None, ) -> ToolAgentInvokeResult: """ Invoke a tool for an agent. UI messages are validated and returned on a dedicated channel. They are intentionally excluded from both the LLM observation and tracing callback output. """ # 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() if parameter.form == ToolParameter.ToolParameterForm.LLM ] if parameters and len(parameters) == 1: tool_parameters = {parameters[0].name: tool_parameters} else: with contextlib.suppress(Exception): tool_parameters = json.loads(tool_parameters) if not isinstance(tool_parameters, dict): raise ValueError(f"tool_parameters should be a dict, but got a string: {tool_parameters}") try: # hit the callback handler agent_tool_callback.on_tool_start(tool_name=tool.entity.identity.name, tool_inputs=tool_parameters) messages = ToolEngine._invoke(session, tool, tool_parameters, user_id, conversation_id, app_id, message_id) invocation_meta_dict: dict[str, ToolInvokeMeta] = {} def message_callback( invocation_meta_dict: dict[str, ToolInvokeMeta], 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, ) message_list, ui_messages = ToolEngine.collect_agent_messages(messages) # extract binary data from tool invoke message binary_files = ToolEngine._extract_tool_response_binary_and_text(message_list) # 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.tool_response_to_str(message_list) 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 ToolAgentInvokeResult( observation=plain_text, message_files=message_files, ui_messages=ui_messages, meta=meta, ) except ToolProviderCredentialValidationError as e: logger.error(e, exc_info=True) 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}" logger.error(e, exc_info=True) 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) logger.error(e, exc_info=True) except ToolInvokeError as e: error_response = f"tool invoke error: {e}" agent_tool_callback.on_tool_error(e) logger.error(e, exc_info=True) except ToolEngineInvokeError as e: meta = e.meta error_response = f"tool invoke error: {meta.error}" agent_tool_callback.on_tool_error(e) logger.error(e, exc_info=True) return ToolAgentInvokeResult( observation=error_response, message_files=[], ui_messages=[], meta=meta, ) except Exception as e: error_response = f"unknown error: {e}" agent_tool_callback.on_tool_error(e) logger.error(e, exc_info=True) return ToolAgentInvokeResult( observation=error_response, message_files=[], ui_messages=[], meta=ToolInvokeMeta.error_instance(error_response), ) @staticmethod def generic_invoke( session: Session, tool: Tool, tool_parameters: dict[str, Any], user_id: str, workflow_tool_callback: DifyWorkflowCallbackHandler, workflow_call_depth: int, conversation_id: str | None = None, app_id: str | None = None, message_id: str | None = 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 if tool.runtime and tool.runtime.runtime_parameters: tool_parameters = {**tool.runtime.runtime_parameters, **tool_parameters} response = tool.invoke( session=session, user_id=user_id, tool_parameters=tool_parameters, conversation_id=conversation_id, app_id=app_id, message_id=message_id, ) # 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 ToolEngine.normalize_ui_messages(response) except Exception as e: workflow_tool_callback.on_tool_error(e) raise e @staticmethod def _invoke( session: Session, tool: Tool, tool_parameters: dict[str, Any], user_id: str, conversation_id: str | None = None, app_id: str | None = None, message_id: str | None = None, ) -> Generator[ToolInvokeMessage | ToolInvokeMeta, None, None]: """ Invoke the tool with the given arguments. """ started_at = datetime.now(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(session, user_id, tool_parameters, conversation_id, app_id, message_id) except Exception as e: meta.error = str(e) raise ToolEngineInvokeError(meta) finally: ended_at = datetime.now(UTC) meta.time_cost = (ended_at - started_at).total_seconds() yield meta @staticmethod def tool_response_to_str(tool_response: list[ToolInvokeMessage]) -> str: """Convert tool invoke messages into the plain-text observation shown to the model/user.""" parts: list[str] = [] json_parts: list[str] = [] for response in tool_response: if response.type == ToolInvokeMessage.MessageType.TEXT: parts.append(cast(ToolInvokeMessage.TextMessage, response.message).text) elif response.type == ToolInvokeMessage.MessageType.LINK: parts.append( 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}: parts.append( "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: json_message = cast(ToolInvokeMessage.JsonMessage, response.message) if json_message.suppress_output: continue json_parts.append( json.dumps( safe_json_value(cast(ToolInvokeMessage.JsonMessage, response.message).json_object), ensure_ascii=False, ) ) elif response.type in { ToolInvokeMessage.MessageType.VARIABLE, ToolInvokeMessage.MessageType.UI, }: continue else: parts.append(str(response.message)) # Add JSON parts, avoiding duplicates from text parts. if json_parts: existing_parts = set(parts) parts.extend(p for p in json_parts if p not in existing_parts) return "".join(parts) @staticmethod def normalize_ui_messages(messages: Iterable[ToolInvokeMessage]) -> Generator[ToolInvokeMessage, None, None]: """Convert reserved variable/JSON compatibility envelopes into native UI.""" for message in messages: if message.type == ToolInvokeMessage.MessageType.VARIABLE: variable_message = cast(ToolInvokeMessage.VariableMessage, message.message) if variable_message.variable_name != DIFY_UI_JSON_ENVELOPE_KEY: yield message continue yield ToolInvokeMessage( type=ToolInvokeMessage.MessageType.UI, message=ToolUIMessage.model_validate(variable_message.variable_value), meta=message.meta, ) continue if message.type != ToolInvokeMessage.MessageType.JSON: yield message continue json_message = cast(ToolInvokeMessage.JsonMessage, message.message) ui_message = extract_ui_message_from_json(json_message.json_object) if ui_message is None: yield message continue yield ToolInvokeMessage(type=ToolInvokeMessage.MessageType.UI, message=ui_message, meta=message.meta) @staticmethod def collect_agent_messages( messages: Iterable[ToolInvokeMessage], ) -> tuple[list[ToolInvokeMessage], list[ToolUIMessage]]: """Collect a bounded UI batch while preserving all non-UI observations. A malformed or over-budget UI message invalidates the complete UI batch. Remaining UI is skipped without parsing, but the input iterable is still consumed so later text and file messages are retained. """ normalized_messages: list[ToolInvokeMessage] = [] ui_messages: list[ToolUIMessage] = [] ui_batch_rejected = False for message in messages: is_ui_transport = ToolEngine._is_ui_transport_message(message) if ui_batch_rejected and is_ui_transport: continue try: normalized = ToolEngine.normalize_ui_messages([message]) except (TypeError, ValueError): if not is_ui_transport: raise # Reserved UI envelopes must never fall back to JSON/variable # observations, even when malformed. ui_batch_rejected = True ui_messages.clear() normalized_messages = [ normalized_message for normalized_message in normalized_messages if normalized_message.type != ToolInvokeMessage.MessageType.UI ] logger.warning("Ignored invalid reserved tool UI message", exc_info=True) continue try: for normalized_message in normalized: if normalized_message.type != ToolInvokeMessage.MessageType.UI: normalized_messages.append(normalized_message) continue if ui_batch_rejected: continue ui_message = cast(ToolUIMessage, normalized_message.message) candidate = [*ui_messages, ui_message] try: validate_tool_ui_message_batch(candidate) except ValueError: ui_batch_rejected = True ui_messages.clear() normalized_messages = [ collected_message for collected_message in normalized_messages if collected_message.type != ToolInvokeMessage.MessageType.UI ] logger.warning( "Ignored tool UI batch that exceeds agent invocation limits", exc_info=True, ) continue ui_messages = candidate normalized_messages.append(normalized_message) except (TypeError, ValueError): if not is_ui_transport: raise ui_batch_rejected = True ui_messages.clear() normalized_messages = [ normalized_message for normalized_message in normalized_messages if normalized_message.type != ToolInvokeMessage.MessageType.UI ] logger.warning("Ignored invalid reserved tool UI message", exc_info=True) return normalized_messages, ui_messages @staticmethod def _is_ui_transport_message(message: ToolInvokeMessage) -> bool: if message.type == ToolInvokeMessage.MessageType.UI: return True if message.type == ToolInvokeMessage.MessageType.VARIABLE: variable_message = cast(ToolInvokeMessage.VariableMessage, message.message) return variable_message.variable_name == DIFY_UI_JSON_ENVELOPE_KEY if message.type == ToolInvokeMessage.MessageType.JSON: json_message = cast(ToolInvokeMessage.JsonMessage, message.message) value = json_message.json_object return isinstance(value, dict) and set(value) == {DIFY_UI_JSON_ENVELOPE_KEY} return False @staticmethod def _extract_tool_response_binary_and_text( tool_response: list[ToolInvokeMessage], ) -> Generator[ToolInvokeMessageBinary, None, None]: """ Extract tool response binary """ for response in tool_response: if response.type in { ToolInvokeMessage.MessageType.IMAGE_LINK, ToolInvokeMessage.MessageType.IMAGE, ToolInvokeMessage.MessageType.BINARY_LINK, }: mimetype = None if not response.meta: raise ValueError("missing meta data") if response.meta.get("mime_type"): mimetype = response.meta.get("mime_type") else: with contextlib.suppress(Exception): 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 if not mimetype: mimetype = ( "image/jpeg" if response.type != ToolInvokeMessage.MessageType.BINARY_LINK else "application/octet-stream" ) yield ToolInvokeMessageBinary( mimetype=response.meta.get("mime_type", mimetype), url=cast(ToolInvokeMessage.TextMessage, response.message).text, ) elif response.type == ToolInvokeMessage.MessageType.BLOB: if not response.meta: raise ValueError("missing meta data") yield ToolInvokeMessageBinary( mimetype=response.meta.get("mime_type", "application/octet-stream"), url=cast(ToolInvokeMessage.TextMessage, response.message).text, ) 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", "application/octet-stream") if response.meta else "application/octet-stream", url=cast(ToolInvokeMessage.TextMessage, response.message).text, ) @staticmethod def _create_message_files( tool_messages: Iterable[ToolInvokeMessageBinary], agent_message: Message, invoke_from: InvokeFrom, user_id: str, ) -> list[str]: """ Create message files produced by a tool call. Tool file persistence is a side effect of agent execution. Use an independent transaction so this helper never commits or closes the caller's request-scoped session. :return: message file ids """ result = [] with sessionmaker(bind=db.engine, expire_on_commit=False).begin() as session: for message in tool_messages: # extract tool file id from url tool_file_id = message.url.split("/")[-1].split(".")[0] message_file = MessageFile( message_id=agent_message.id, type=ToolEngine._resolve_tool_file_type(message), transfer_method=FileTransferMethod.TOOL_FILE, belongs_to=MessageFileBelongsTo.ASSISTANT, url=message.url, upload_file_id=tool_file_id, created_by_role=( CreatorUserRole.ACCOUNT if invoke_from in {InvokeFrom.EXPLORE, InvokeFrom.DEBUGGER} else CreatorUserRole.END_USER ), created_by=user_id, ) session.add(message_file) result.append(message_file.id) return result @staticmethod def _resolve_tool_file_type(message: ToolInvokeMessageBinary) -> FileType: if "image" in message.mimetype: return FileType.IMAGE elif "video" in message.mimetype: return FileType.VIDEO elif "audio" in message.mimetype: return FileType.AUDIO elif "text" in message.mimetype or "pdf" in message.mimetype: return FileType.DOCUMENT else: return FileType.CUSTOM