tjdify/api/core/tools/tool_engine.py

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.plugin_tool_callback_handler import DifyPluginCallbackHandler
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.entities.tool_entities import ToolInvokeMessage, ToolInvokeMessageBinary, ToolInvokeMeta, ToolParameter
from core.tools.errors import (
    ToolEngineInvokeError,
    ToolInvokeError,
    ToolNotFoundError,
    ToolNotSupportedError,
    ToolParameterValidationError,
    ToolProviderCredentialValidationError,
    ToolProviderNotFoundError,
)
from core.tools.tool.tool import Tool
from core.tools.tool.workflow_tool import WorkflowTool
from core.tools.utils.message_transformer import ToolFileMessageTransformer
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.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.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.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,
        ) -> 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.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(user_id=user_id, tool_parameters=tool_parameters)

            # hit the callback handler
            workflow_tool_callback.on_tool_end(
                tool_name=tool.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 plugin_invoke(tool: Tool, tool_parameters: dict, user_id: str,
                      callback: DifyPluginCallbackHandler
        ) -> Generator[ToolInvokeMessage, None, None]:
        """
        Plugin invokes the tool with the given arguments.
        """
        try:
            # hit the callback handler
            callback.on_tool_start(
                tool_name=tool.identity.name, 
                tool_inputs=tool_parameters
            )

            response = tool.invoke(user_id, tool_parameters)

            # hit the callback handler
            callback.on_tool_end(
                tool_name=tool.identity.name,
                tool_inputs=tool_parameters,
                tool_outputs=response,
            )

            return response
        except Exception as e:
            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.identity.name,
            'tool_provider': tool.identity.provider,
            'tool_provider_type': tool.tool_provider_type().value,
            'tool_parameters': deepcopy(tool.runtime.runtime_parameters),
            'tool_icon': tool.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 == ToolInvokeMessage.MessageType.IMAGE_LINK or \
                 response.type == 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:
                result += f"tool response: {json.dumps(cast(ToolInvokeMessage.JsonMessage, response.message).json_object, ensure_ascii=False)}."
            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 == ToolInvokeMessage.MessageType.IMAGE_LINK or \
                response.type == 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